<commit_before># -*- coding: utf-8 -*-

# http://en.wikipedia.org/wiki/ANSI_escape_code
# http://en.wikipedia.org/wiki/Web_colors
# https://gist.github.com/MicahElliott/719710

import os
import sys

# Python 3
if sys.version_info[0] == 3:
    string_type = str
    unicode = str
else:
    string_type = basestring

def is_color_supported():
    "Find out if your terminal environment supports color."
    # shinx.util.console
    if not hasattr(sys.stdout, 'isatty'):
        return False

    if not sys.stdout.isatty():
        return False

    if sys.platform == 'win32':
        try:
            import colorama
            colorama.init()
            return True
        except ImportError:
            return False

    if 'COLORTERM' in os.environ:
        return True

    term = os.environ.get('TERM', 'dumb').lower()
    return term in ('xterm', 'linux') or 'color' in term

def is_256color_supported():
    "Find out if your terminal environment supports 256 color."
    if not is_color_supported():
        return False
    term = os.environ.get('TERM', 'dumb').lower()
    return '256' in term

def rgb2ansi(r, g, b):
    """
    Convert an RGB color to 256 ansi graphics.
    """

    # Thanks to
    # https://github.com/tehmaze/ansi/blob/master/ansi/colour/rgb.py

    grayscale = False
    poss = True
    step = 2.5

    while poss:
        if min(r, g, b) < step:
            grayscale = max(r, g, b) < step
            poss = False

        step += 42.5

    if grayscale:
        return 232 + int(float(sum((r, g, b)) / 33.0))

    m = ((r, 36), (g, 6), (b, 1))
    return 16 + sum(int(6 * float(val) / 256) * mod for val, mod in m)

def hex2ansi(code):
    """
    Convert hex code to ansi.
    """

    if code.startswith('#'):
        code = code[1:]

    if len(code) == 3:
        # efc -> eeffcc
        return rgb2ansi(*map(lambda o: int(o * 2, 16), code))

    if len(code) != 6:
        raise ValueError('invalid color code')

    rgb = (code[:2], code[2:4], code[4:])
    return rgb2ansi(*map(lambda o: int(o, 16), rgb))

_reset = '\x1b[0;39;49m'
_styles = (
    'bold', 'faint', 'italic', 'underline', 'blink',
    'overline', 'inverse', 'conceal', 'strike',
)
_colors = (
    'black', 'red', 'green', 'yellow', 'blue',
    'magenta', 'cyan', 'white'
)

def _color2ansi(color):
    if color in _colors:
        return _colors.index(color)

    if isinstance(color, string_type):
        return hex2ansi(color)
    elif isinstance(color, (tuple, list)):
        return rgb2ansi(*color)

    raise ValueError('invalid color: %s' % color)

class Color(object):
    """
    Color object for painters.

    You should always use the high-level API of colors and styles,
    such as :class:`red` and :class:`bold`.

    But if you are so interested in this module, you are welcome to
    use some advanced features::

        s = Color('text')
        print(s.bold.red.italic)

    All ANSI colors and styles are available on Color.
    """

    def __init__(self, *items):
        self.items = items

        self.styles = []
        self.fgcolor = None
        self.bgcolor = None

    def __getattr__(self, key):
        if key.endswith('_bg'):
            name = key[:-3]
            if name not in _colors:
                raise AttributeError("Color has no attribute '%s'" % key)
            self.bgcolor = _colors.index(name)
            return self

        if key in _colors:
            self.fgcolor = _colors.index(key)
            return self

        if key in _styles:
            code = _styles.index(key)
            self.styles.append(code + 1)
            return self
        try:
            return super(Color, self).__getattr__(key)
        except AttributeError:
            raise AttributeError("Color has no attribute '%s'" % key)

    def __str__(self):
        text = ''.join(unicode(item) for item in self.items)

        if unicode != str:
            text = text.encode('utf-8')

        if not is_color_supported():
            return text

        is256 = is_256color_supported()

        if self.fgcolor and not isinstance(self.fgcolor, int):
            self.fgcolor = _color2ansi(self.fgcolor)

        if self.bgcolor and not isinstance(self.bgcolor, int):
            self.bgcolor = _color2ansi(self.bgcolor)

        if is256:
            if self.fgcolor is not None:
                text = '\x1b[38;5;%im%s%s' % (self.fgcolor, text, _reset)
            if self.bgcolor is not None:
                text = '\x1b[48;5;%im%s%s' % (self.bgcolor, text, _reset)

        else:
            if self.fgcolor is not None and self.fgcolor < 8:
                text = '\x1b[%im%s%s' % (30 + self.fgcolor, text, _reset)
            if self.bgcolor is not None and self.bgcolor < 8:
                text = '\x1b[%im%s%s' % (40 + self.bgcolor, text, _reset)

        if self.styles:
            code = ';'.join(str(i) for i in self.styles)
            text = '\x1b[%sm%s%s' % (code, text, _reset)

        return text

    def __repr__(self):
        return repr(str(self))

    def __len__(self):
        return sum([len(item) for item in self.items])

    def __add__(self, s):
        if not isinstance(s, (string_type, Color)):
            msg = "Concatenatation failed: %r + %r (Not a ColorString or str)"
            raise TypeError(msg % (type(s), type(self)))
        return Color(self, s)

    def __radd__(self, s):
        if not isinstance(s, (string_type, Color)):
            msg = "Concatenatation failed: %r + %r (Not a ColorString or str)"
            raise TypeError(msg % (type(s), type(self)))
        return Color(s, self)

def colorize(text, color):
    """
    Colorize text with hex code.

    :param text: the text you want to paint
    :param color: a hex color or rgb color

    ::

        colorize('hello', 'ff0000')
        colorize('hello', '#ff0000')
        colorize('hello', (255, 0, 0))

    """

    if color in _styles:
        c = Color(text)
        c.styles = [_styles.index(color) + 1]
        return c

    c = Color(text)
    c.fgcolor = _color2ansi(color)
    return c

def _create_color_func(text, fgcolor=None, bgcolor=None, *styles):
    c = Color(text)
    c.fgcolor = fgcolor
    c.bgcolor = bgcolor
    c.styles = styles
    return c

def bold(text):
    """
    Bold style.
    """
    return _create_color_func(text, None, None, 1)

def faint(text):
    """
    Faint style.
    """
    return _create_color_func(text, None, None, 2)

def italic(text):
    """
    Italic style.
    """
    return _create_color_func(text, None, None, 3)

def underline(text):
    """
    Underline style.
    """
    return _create_color_func(text, None, None, 4)

def blink(text):
    """
    Blink style.
    """
    return _create_color_func(text, None, None, 5)

def overline(text):
    """
    Overline style.
    """
    return _create_color_func(text, None, None, 6)

def inverse(text):
    """
    Inverse style.
    """
    return _create_color_func(text, None, None, 7)

def conceal(text):
    """
    Conceal style.
    """
    return _create_color_func(text, None, None, 8)

def strike(text):
    """
    Strike style.
    """
    return _create_color_func(text, None, None, 9)

def black(text):
    """
    Black color.
    """
    return _create_color_func(text, fgcolor=0)

def red(text):
    """
    Red color.
    """
    return _create_color_func(text, fgcolor=1)

def green(text):
    """
    Green color.
    """
    return _create_color_func(text, fgcolor=2)

def yellow(text):
    """
    Yellow color.
    """
    return _create_color_func(text, fgcolor=3)

def blue(text):
    """
    Blue color.
    """
    return _create_color_func(text, fgcolor=4)

def magenta(text):
    """
    Magenta color.
    """
    return _create_color_func(text, fgcolor=5)

def cyan(text):
    """
    Cyan color.
    """
    return _create_color_func(text, fgcolor=6)

def white(text):
    """
    White color.
    """
    return _create_color_func(text, fgcolor=7)

def gray(text):
    """
    Gray color.
    """
    if is_256color_supported():
        return _create_color_func(text, fgcolor=8)
    return _create_color_func(text, 0, None, 8)

def grey(text):
    """
    Alias of gray.
    """
    return gray(text)

def black_bg(text):
    """
    Black background.
    """
    return _create_color_func(text, bgcolor=0)

def red_bg(text):
    """
    Red background.
    """
    return _create_color_func(text, bgcolor=1)

def green_bg(text):
    """
    Green background.
    """
    return _create_color_func(text, bgcolor=2)

def yellow_bg(text):
    """
    Yellow background.
    """
    return _create_color_func(text, bgcolor=3)

def blue_bg(text):
    """
    Blue background.
    """
    return _create_color_func(text, bgcolor=4)

def magenta_bg(text):
    """
    Magenta background.
    """
    return _create_color_func(text, bgcolor=5)

def cyan_bg(text):
    """
    Cyan background.
    """
    return _create_color_func(text, bgcolor=6)

def white_bg(text):
    """
    White background.
    """
    return _create_color_func(text, bgcolor=7)

def gray_bg(text):
    """
    Gray background.
    """
    if is_256color_supported():
        return _create_color_func(text, bgcolor=8)
    return _create_color_func(text, None, 0, 1)

def grey_bg(text):
    """
    Alias of gray_bg.
    """
    return gray_bg(text)
<commit_msg>colorize can paint background<commit_after># -*- coding: utf-8 -*-

# http://en.wikipedia.org/wiki/ANSI_escape_code
# http://en.wikipedia.org/wiki/Web_colors
# https://gist.github.com/MicahElliott/719710

import os
import sys

# Python 3
if sys.version_info[0] == 3:
    string_type = str
    unicode = str
else:
    string_type = basestring

def is_color_supported():
    "Find out if your terminal environment supports color."
    # shinx.util.console
    if not hasattr(sys.stdout, 'isatty'):
        return False

    if not sys.stdout.isatty():
        return False

    if sys.platform == 'win32':
        try:
            import colorama
            colorama.init()
            return True
        except ImportError:
            return False

    if 'COLORTERM' in os.environ:
        return True

    term = os.environ.get('TERM', 'dumb').lower()
    return term in ('xterm', 'linux') or 'color' in term

def is_256color_supported():
    "Find out if your terminal environment supports 256 color."
    if not is_color_supported():
        return False
    term = os.environ.get('TERM', 'dumb').lower()
    return '256' in term

def rgb2ansi(r, g, b):
    """
    Convert an RGB color to 256 ansi graphics.
    """

    # Thanks to
    # https://github.com/tehmaze/ansi/blob/master/ansi/colour/rgb.py

    grayscale = False
    poss = True
    step = 2.5

    while poss:
        if min(r, g, b) < step:
            grayscale = max(r, g, b) < step
            poss = False

        step += 42.5

    if grayscale:
        return 232 + int(float(sum((r, g, b)) / 33.0))

    m = ((r, 36), (g, 6), (b, 1))
    return 16 + sum(int(6 * float(val) / 256) * mod for val, mod in m)

def hex2ansi(code):
    """
    Convert hex code to ansi.
    """

    if code.startswith('#'):
        code = code[1:]

    if len(code) == 3:
        # efc -> eeffcc
        return rgb2ansi(*map(lambda o: int(o * 2, 16), code))

    if len(code) != 6:
        raise ValueError('invalid color code')

    rgb = (code[:2], code[2:4], code[4:])
    return rgb2ansi(*map(lambda o: int(o, 16), rgb))

_reset = '\x1b[0;39;49m'
_styles = (
    'bold', 'faint', 'italic', 'underline', 'blink',
    'overline', 'inverse', 'conceal', 'strike',
)
_colors = (
    'black', 'red', 'green', 'yellow', 'blue',
    'magenta', 'cyan', 'white'
)

def _color2ansi(color):
    if color in _colors:
        return _colors.index(color)

    if isinstance(color, string_type):
        return hex2ansi(color)
    elif isinstance(color, (tuple, list)):
        return rgb2ansi(*color)

    raise ValueError('invalid color: %s' % color)

class Color(object):
    """
    Color object for painters.

    You should always use the high-level API of colors and styles,
    such as :class:`red` and :class:`bold`.

    But if you are so interested in this module, you are welcome to
    use some advanced features::

        s = Color('text')
        print(s.bold.red.italic)

    All ANSI colors and styles are available on Color.
    """

    def __init__(self, *items):
        self.items = items

        self.styles = []
        self.fgcolor = None
        self.bgcolor = None

    def __getattr__(self, key):
        if key.endswith('_bg'):
            name = key[:-3]
            if name not in _colors:
                raise AttributeError("Color has no attribute '%s'" % key)
            self.bgcolor = _colors.index(name)
            return self

        if key in _colors:
            self.fgcolor = _colors.index(key)
            return self

        if key in _styles:
            code = _styles.index(key)
            self.styles.append(code + 1)
            return self
        try:
            return super(Color, self).__getattr__(key)
        except AttributeError:
            raise AttributeError("Color has no attribute '%s'" % key)

    def __str__(self):
        text = ''.join(unicode(item) for item in self.items)

        if unicode != str:
            text = text.encode('utf-8')

        if not is_color_supported():
            return text

        is256 = is_256color_supported()

        if self.fgcolor and not isinstance(self.fgcolor, int):
            self.fgcolor = _color2ansi(self.fgcolor)

        if self.bgcolor and not isinstance(self.bgcolor, int):
            self.bgcolor = _color2ansi(self.bgcolor)

        if is256:
            if self.fgcolor is not None:
                text = '\x1b[38;5;%im%s%s' % (self.fgcolor, text, _reset)
            if self.bgcolor is not None:
                text = '\x1b[48;5;%im%s%s' % (self.bgcolor, text, _reset)

        else:
            if self.fgcolor is not None and self.fgcolor < 8:
                text = '\x1b[%im%s%s' % (30 + self.fgcolor, text, _reset)
            if self.bgcolor is not None and self.bgcolor < 8:
                text = '\x1b[%im%s%s' % (40 + self.bgcolor, text, _reset)

        if self.styles:
            code = ';'.join(str(i) for i in self.styles)
            text = '\x1b[%sm%s%s' % (code, text, _reset)

        return text

    def __repr__(self):
        return repr(str(self))

    def __len__(self):
        return sum([len(item) for item in self.items])

    def __add__(self, s):
        if not isinstance(s, (string_type, Color)):
            msg = "Concatenatation failed: %r + %r (Not a ColorString or str)"
            raise TypeError(msg % (type(s), type(self)))
        return Color(self, s)

    def __radd__(self, s):
        if not isinstance(s, (string_type, Color)):
            msg = "Concatenatation failed: %r + %r (Not a ColorString or str)"
            raise TypeError(msg % (type(s), type(self)))
        return Color(s, self)

def colorize(text, color, background=False):
    """
    Colorize text with hex code.

    :param text: the text you want to paint
    :param color: a hex color or rgb color
    :param background: decide to colorize background

    ::

        colorize('hello', 'ff0000')
        colorize('hello', '#ff0000')
        colorize('hello', (255, 0, 0))

    """

    if color in _styles:
        c = Color(text)
        c.styles = [_styles.index(color) + 1]
        return c

    c = Color(text)
    if background:
        c.bgcolor = _color2ansi(color)
    else:
        c.fgcolor = _color2ansi(color)
    return c

def _create_color_func(text, fgcolor=None, bgcolor=None, *styles):
    c = Color(text)
    c.fgcolor = fgcolor
    c.bgcolor = bgcolor
    c.styles = styles
    return c

def bold(text):
    """
    Bold style.
    """
    return _create_color_func(text, None, None, 1)

def faint(text):
    """
    Faint style.
    """
    return _create_color_func(text, None, None, 2)

def italic(text):
    """
    Italic style.
    """
    return _create_color_func(text, None, None, 3)

def underline(text):
    """
    Underline style.
    """
    return _create_color_func(text, None, None, 4)

def blink(text):
    """
    Blink style.
    """
    return _create_color_func(text, None, None, 5)

def overline(text):
    """
    Overline style.
    """
    return _create_color_func(text, None, None, 6)

def inverse(text):
    """
    Inverse style.
    """
    return _create_color_func(text, None, None, 7)

def conceal(text):
    """
    Conceal style.
    """
    return _create_color_func(text, None, None, 8)

def strike(text):
    """
    Strike style.
    """
    return _create_color_func(text, None, None, 9)

def black(text):
    """
    Black color.
    """
    return _create_color_func(text, fgcolor=0)

def red(text):
    """
    Red color.
    """
    return _create_color_func(text, fgcolor=1)

def green(text):
    """
    Green color.
    """
    return _create_color_func(text, fgcolor=2)

def yellow(text):
    """
    Yellow color.
    """
    return _create_color_func(text, fgcolor=3)

def blue(text):
    """
    Blue color.
    """
    return _create_color_func(text, fgcolor=4)

def magenta(text):
    """
    Magenta color.
    """
    return _create_color_func(text, fgcolor=5)

def cyan(text):
    """
    Cyan color.
    """
    return _create_color_func(text, fgcolor=6)

def white(text):
    """
    White color.
    """
    return _create_color_func(text, fgcolor=7)

def gray(text):
    """
    Gray color.
    """
    if is_256color_supported():
        return _create_color_func(text, fgcolor=8)
    return _create_color_func(text, 0, None, 8)

def grey(text):
    """
    Alias of gray.
    """
    return gray(text)

def black_bg(text):
    """
    Black background.
    """
    return _create_color_func(text, bgcolor=0)

def red_bg(text):
    """
    Red background.
    """
    return _create_color_func(text, bgcolor=1)

def green_bg(text):
    """
    Green background.
    """
    return _create_color_func(text, bgcolor=2)

def yellow_bg(text):
    """
    Yellow background.
    """
    return _create_color_func(text, bgcolor=3)

def blue_bg(text):
    """
    Blue background.
    """
    return _create_color_func(text, bgcolor=4)

def magenta_bg(text):
    """
    Magenta background.
    """
    return _create_color_func(text, bgcolor=5)

def cyan_bg(text):
    """
    Cyan background.
    """
    return _create_color_func(text, bgcolor=6)

def white_bg(text):
    """
    White background.
    """
    return _create_color_func(text, bgcolor=7)

def gray_bg(text):
    """
    Gray background.
    """
    if is_256color_supported():
        return _create_color_func(text, bgcolor=8)
    return _create_color_func(text, None, 0, 1)

def grey_bg(text):
    """
    Alias of gray_bg.
    """
    return gray_bg(text)
<|endoftext|>


<commit_before>import math

import numpy as np
import mdtraj as mdj

from wepy.walker import merge
from wepy.resampling.clone_merge import CloneMergeDecision
from wepy.resampling.resampler import Resampler, ResamplingRecord

class WExplore2Resampler(Resampler):
    def __init__(self,refrence_trajectory):
        self.ref = refrence_trajectory
        self.n_walkers = None
        self.walkerwt = []
        self.walkers = []
        self.amp = []
        self.distancearray = None
        self.resampling_records = []

    def __rmsd(self, traj, ref, idx):
        return np.sqrt(3*np.sum(np.square(traj.xyz[:, idx, :] - ref.xyz[:, idx, :]),
                                axis=(1, 2))/idx.shape[0])

    def __Maketraj(self, positions):
        Newxyz = np.zeros((1, self.ref.n_atoms, 3))
        for i in range(len(positions)):
            Newxyz[0,i,:] = ([positions[i]._value[0], positions[i]._value[1],
                                                        positions[i]._value[2]])
        return mdj.Trajectory(Newxyz, self.ref.topology)

    def CalculateRmsd(self, positions1, positions2):
        self.ref =self.ref.remove_solvent()
        lig_idx = self.ref.topology.select('resname "2RV"')
        b_selection = mdj.compute_neighbors(self.ref, 0.8, lig_idx)
        b_selection = np.delete(b_selection, lig_idx)

        ref_traj =self.__Maketraj(positions1[0:self.ref.n_atoms])
        traj = self.__Maketraj(positions2[0:self.ref.n_atoms])
        traj = traj.superpose(ref_traj, atom_indices = b_selection)
        return self.__rmsd(traj, ref_traj, lig_idx)

    def __calcspread(self, lpmin, dpower):
        spread = 0
        wsum = np.zeros(self.n_walkers)
        wtfac = np.zeros(self.n_walkers)
        for i in range(self.n_walkers):
            if self.walkerwt[i] > 0 and self.amp[i] > 0:
                wtfac[i] = math.log(self.walkerwt[i]/self.amp[i]) - lpmin
            else:
                wtfac[i] = 0
            if wtfac[i] < 0:
                wtfac[i] = 0

        for i in range(self.n_walkers-1):
            if self.amp[i] > 0:
                for j in range(i+1, self.n_walkers):
                    if self.amp[j] > 0:
                        d = self.distancearray[i][j] ** dpower * wtfac[i]*wtfac[j];
                        spread += d * self.amp[i] * self.amp[j];
                        wsum[i] += d * self.amp[j];
                        wsum[j] += d * self.amp[i];

        return(spread, wsum)

    def __merge(self, i, j):
        self.walkerwt[j] += self.walkerwt[i]
        self.walkerwt[i] = 0
        self.amp[i] = 0
        self.amp[j] = 1
        self.resampling_records[i] = ResamplingRecord(
                                    decision=CloneMergeDecision.SQUASH,
                                    value = j)
        self.resampling_records[j] = ResamplingRecord(
                                    decision=CloneMergeDecision.KEEP_MERGE,
                                    value = j)

    def __copystructure(self, a, b):
        # this function will copy the simulation details from walker a to walker b
        self.walkers[b].state = self.walkers[a].state

    def decide(self,):
        pmin = 1e-12
        pmax = 0.1
        lpmin = math.log(pmin/100)
        dpower = 4
        productive = True
        mergedist = 0.25
        spread, wsum = self.__calcspread(lpmin, dpower)
        print("Starting variance:", spread)
        while productive:
            productive = False
            # find min and max wsums, alter @self.amp
            minwsum = 1.0
            maxwsum = 0.0
            minwind = None
            maxwind = None
            for i in range(0, self.n_walkers-1):
                # determine eligibility and find highest wsum (to clone) and lowest wsum (to merge)
                if self.amp[i] >= 1 or self.walkerwt[i] > pmin:  # find the most clonable walker
                    if wsum[i] > maxwsum:
                        maxwsum = wsum[i]
                        maxwind = i
                if self.amp[i] == 1 or  self.walkerwt[i] < pmax:  # find the most mergeable walker
                    if wsum[i] < minwsum: # convert or to and  (important)
                        minwsum = wsum[i]
                        minwind = i

            # does minwind have an eligible merging partner?
            closedist = mergedist
            closewalk = None
            if minwind is not None and maxwind is not None and minwind != maxwind:
                for j in range(0, self.n_walkers):
                    if j != minwind and j != maxwind:
                        if self.distancearray[minwind][j]<closedist and self.amp[j] == 1 and self.walkerwt[j]<pmax :
                            closedist = self.distancearray[minwind][j]
                            closewalk = j
        # did we find a closewalk?
            if minwind is not None and maxwind is not None and closewalk is not None:
                if self.amp[minwind] != 1:
                    die("Error! minwind", minwind, "has self.amp =", self.amp[minwind])
                if self.amp[closewalk] != 1:
                    die("Error! closewalk", closewalk, "has self.amp=", self.amp[closewalk])
                # change self.amp

                tempsum = self.walkerwt[minwind] + self.walkerwt[closewalk]
                self.amp[minwind] = self.walkerwt[minwind]/tempsum
                self.amp[closewalk] = self.walkerwt[closewalk]/tempsum
                self.amp[maxwind] += 1

                # re-determine spread function, and wsum values
                (newspread, wsum) = __calcspread(lpmin, dpower)

                if newspread > spread:
                    print("Variance move to", newspread, "accepted")
                    productive = True
                    spread = newspread

                    # make a decision on which walker to keep (minwind, or closewalk)
                    r = np.random.random(self.walkerwt[closewalk] + self.walkerwt[minwind])
                    if r < self.walkerwt[closewalk]:
                        # keep closewalk, get rid of minwind
                        self.__merge(minwind, closewalk)
                    else:
                        # keep minwind, get rid of closewalk
                        self.__merge(closewalk, minwind)

                    (newspread, wsum) = __calcspread(lpmin, dpower)
                    print("variance after selection:", newspread)
                # Not productive
                else:
                    self.amp[minwind] = 1
                    self.amp[closewalk] = 1
                    self.amp[maxwind] -= 1

        # end of while productive loop:  done with cloning and merging steps
        # set walkerdist arrays
        walkerdist = wsum

        # Perform the cloning and merging steps
        # start with an empty free walker array, add all walkers with amp = 0
        freewalkers = []
        for r1walk in range(0, self.n_walkers):
            if self.amp[r1walk] == 0:
                freewalkers.append(r1walk)

        # for each self.amp[i] > 1, clone!
        for r1walk in range(0, self.n_walkers):
            if self.amp[r1walk] > 1:
                nclone = self.amp[r1walk]-1
                inds = []
                for i in range(0, nclone):
                    try:
                        tind = freewalkers.pop()
                    except:
                        raise("Error! Not enough free walkers!")
                    if r1walk == tind:
                        raise("Error!  free walker is equal to the clone!")
                    else:
                        inds.append(tind)

                newwt = self.walkerwt[r1walk]/(nclone+1)
                self.walkerwt[r1walk] = newwt
                for tind in inds:
                    self.walkerwt[tind] = newwt
                    walkerdist[tind] = walkerdist[r1walk]
                    self.__copystructure(r1walk, tind)
                    # done cloning and meging
                # add r1walka index just for keeping tracks of cloning
                inds.append(r1walk)
                self.resampling_records[r1walk] = ResamplingRecord(
                                    decision=CloneMergeDecision.CLONE,
                                    value = inds)

            if len(freewalkers) > 0:
                raise("Error! walkers left over after merging and cloning")

    def resample(self, walkers):

        print ("Starting resampling")

        self.walkers = walkers
        self.n_walkers = len(self.walkers)
        self.walkerwt = [walker.weight for walker in self.walkers]
        self.amp = [1 for i in range(self.n_walkers)]
        self.distancearray = np.zeros((self.n_walkers, self.n_walkers))
        self.resampling_records = [ResamplingRecord(decision=CloneMergeDecision.NOTHING,                                                                                value=i)
                                   for i in range(self.n_walkers)]

        for i in range(self.n_walkers):
            for j in range (i+1, self.n_walkers):
                self.distancearray[i][j] = self.CalculateRmsd(self.walkers[i].positions, self.walkers[j].positions)

        self.decide()
        for i in range(self.n_walkers):
            self.walkers[i].weight = self.walkerwt[i]

        print ("Ending resampling")

        return self.walkers, self.resampling_records
<commit_msg>make calculation function parallel<commit_after>import math
import multiprocessing as mulproc

import numpy as np
import mdtraj as mdj

from wepy.walker import merge
from wepy.resampling.clone_merge import CloneMergeDecision
from wepy.resampling.resampler import Resampler, ResamplingRecord

class WExplore2Resampler(Resampler):
    def __init__(self,refrence_trajectory):
        self.ref = refrence_trajectory
        self.n_walkers = None
        self.walkerwt = []
        self.walkers = []
        self.amp = []
        self.distancearray = None
        self.resampling_records = []

    def __rmsd(self, traj, ref, idx):
        return np.sqrt(3*np.sum(np.square(traj.xyz[:, idx, :] - ref.xyz[:, idx, :]),
                                axis=(1, 2))/idx.shape[0])

    def __Maketraj(self, positions):
        Newxyz = np.zeros((1, self.ref.n_atoms, 3))
        for i in range(len(positions)):
            Newxyz[0,i,:] = ([positions[i]._value[0], positions[i]._value[1],
                                                        positions[i]._value[2]])
        return mdj.Trajectory(Newxyz, self.ref.topology)

    def CalculateRmsd(self, i, j):
        self.ref =self.ref.remove_solvent()
        lig_idx = self.ref.topology.select('resname "2RV"')
        b_selection = mdj.compute_neighbors(self.ref, 0.8, lig_idx)
        b_selection = np.delete(b_selection, lig_idx)

        ref_traj =self.__Maketraj(self.walkers[i].positions[0:self.ref.n_atoms])
        traj = self.__Maketraj(self.walkers[j].positions[0:self.ref.n_atoms])
        traj = traj.superpose(ref_traj, atom_indices = b_selection)
        return  self.__rmsd(traj, ref_traj, lig_idx)

    def multi_run_warpper(self, args):
        return self.CalculateRmsd(*args)

    def __calcspread(self, lpmin, dpower):
        spread = 0
        wsum = np.zeros(self.n_walkers)
        wtfac = np.zeros(self.n_walkers)
        for i in range(self.n_walkers):
            if self.walkerwt[i] > 0 and self.amp[i] > 0:
                wtfac[i] = math.log(self.walkerwt[i]/self.amp[i]) - lpmin
            else:
                wtfac[i] = 0
            if wtfac[i] < 0:
                wtfac[i] = 0

        for i in range(self.n_walkers-1):
            if self.amp[i] > 0:
                for j in range(i+1, self.n_walkers):
                    if self.amp[j] > 0:
                        d = self.distancearray[i][j] ** dpower * wtfac[i]*wtfac[j];
                        spread += d * self.amp[i] * self.amp[j];
                        wsum[i] += d * self.amp[j];
                        wsum[j] += d * self.amp[i];

        return(spread, wsum)

    def __merge(self, i, j):
        self.walkerwt[j] += self.walkerwt[i]
        self.walkerwt[i] = 0
        self.amp[i] = 0
        self.amp[j] = 1
        self.resampling_records[i] = ResamplingRecord(
                                    decision=CloneMergeDecision.SQUASH,
                                    value = j)
        self.resampling_records[j] = ResamplingRecord(
                                    decision=CloneMergeDecision.KEEP_MERGE,
                                    value = j)

    def __copystructure(self, a, b):
        # this function will copy the simulation details from walker a to walker b
        self.walkers[b].state = self.walkers[a].state

    def decide(self,):
        pmin = 1e-12
        pmax = 0.1
        lpmin = math.log(pmin/100)
        dpower = 4
        productive = True
        mergedist = 0.25
        spread, wsum = self.__calcspread(lpmin, dpower)
        print("Starting variance:", spread)
        while productive:
            productive = False
            # find min and max wsums, alter @self.amp
            minwsum = 1.0
            maxwsum = 0.0
            minwind = None
            maxwind = None
            for i in range(0, self.n_walkers-1):
                # determine eligibility and find highest wsum (to clone) and lowest wsum (to merge)
                if self.amp[i] >= 1 or self.walkerwt[i] > pmin:  # find the most clonable walker
                    if wsum[i] > maxwsum:
                        maxwsum = wsum[i]
                        maxwind = i
                if self.amp[i] == 1 or  self.walkerwt[i] < pmax:  # find the most mergeable walker
                    if wsum[i] < minwsum: # convert or to and  (important)
                        minwsum = wsum[i]
                        minwind = i

            # does minwind have an eligible merging partner?
            closedist = mergedist
            closewalk = None
            if minwind is not None and maxwind is not None and minwind != maxwind:
                for j in range(0, self.n_walkers):
                    if j != minwind and j != maxwind:
                        if self.distancearray[minwind][j]<closedist and self.amp[j] == 1 and self.walkerwt[j]<pmax :
                            closedist = self.distancearray[minwind][j]
                            closewalk = j
        # did we find a closewalk?
            if minwind is not None and maxwind is not None and closewalk is not None:
                if self.amp[minwind] != 1:
                    die("Error! minwind", minwind, "has self.amp =", self.amp[minwind])
                if self.amp[closewalk] != 1:
                    die("Error! closewalk", closewalk, "has self.amp=", self.amp[closewalk])
                # change self.amp

                tempsum = self.walkerwt[minwind] + self.walkerwt[closewalk]
                self.amp[minwind] = self.walkerwt[minwind]/tempsum
                self.amp[closewalk] = self.walkerwt[closewalk]/tempsum
                self.amp[maxwind] += 1

                # re-determine spread function, and wsum values
                (newspread, wsum) = __calcspread(lpmin, dpower)

                if newspread > spread:
                    print("Variance move to", newspread, "accepted")
                    productive = True
                    spread = newspread

                    # make a decision on which walker to keep (minwind, or closewalk)
                    r = np.random.random(self.walkerwt[closewalk] + self.walkerwt[minwind])
                    if r < self.walkerwt[closewalk]:
                        # keep closewalk, get rid of minwind
                        self.__merge(minwind, closewalk)
                    else:
                        # keep minwind, get rid of closewalk
                        self.__merge(closewalk, minwind)

                    (newspread, wsum) = __calcspread(lpmin, dpower)
                    print("variance after selection:", newspread)
                # Not productive
                else:
                    self.amp[minwind] = 1
                    self.amp[closewalk] = 1
                    self.amp[maxwind] -= 1

        # end of while productive loop:  done with cloning and merging steps
        # set walkerdist arrays
        walkerdist = wsum

        # Perform the cloning and merging steps
        # start with an empty free walker array, add all walkers with amp = 0
        freewalkers = []
        for r1walk in range(0, self.n_walkers):
            if self.amp[r1walk] == 0:
                freewalkers.append(r1walk)

        # for each self.amp[i] > 1, clone!
        for r1walk in range(0, self.n_walkers):
            if self.amp[r1walk] > 1:
                nclone = self.amp[r1walk]-1
                inds = []
                for i in range(0, nclone):
                    try:
                        tind = freewalkers.pop()
                    except:
                        raise("Error! Not enough free walkers!")
                    if r1walk == tind:
                        raise("Error!  free walker is equal to the clone!")
                    else:
                        inds.append(tind)

                newwt = self.walkerwt[r1walk]/(nclone+1)
                self.walkerwt[r1walk] = newwt
                for tind in inds:
                    self.walkerwt[tind] = newwt
                    walkerdist[tind] = walkerdist[r1walk]
                    self.__copystructure(r1walk, tind)
                    # done cloning and meging
                # add r1walka index just for keeping tracks of cloning
                inds.append(r1walk)
                self.resampling_records[r1walk] = ResamplingRecord(
                                    decision=CloneMergeDecision.CLONE,
                                    value = inds)

            if len(freewalkers) > 0:
                raise("Error! walkers left over after merging and cloning")

    def resample(self, walkers):

        calc_list = []
        print ("Starting resampling")

        self.walkers = walkers
        self.n_walkers = len(self.walkers)
        self.walkerwt = [walker.weight for walker in self.walkers]
        self.amp = [1 for i in range(self.n_walkers)]
        self.distancearray = np.zeros((self.n_walkers, self.n_walkers))
        self.resampling_records = [ResamplingRecord(decision=CloneMergeDecision.NOTHING,                                                                                value=i)
                                   for i in range(self.n_walkers)]

        for i in range(self.n_walkers):
            for j in range (i+1, self.n_walkers):
                calc_list.append((i,j))

        from multiprocessing import Pool
        pool = Pool (28)
        results = pool.map (self.multi_run_warpper, calc_list)
        pool.close()
        indexitem = zip(calc_list , results)
        for index in indexitem:
            self.distancearray[index[0][0]][index[0][1]] = index[1]

        self.decide()
        for i in range(self.n_walkers):
            self.walkers[i].weight = self.walkerwt[i]

        print ("Ending resampling")

        return self.walkers, self.resampling_records
<|endoftext|>


<commit_before>"""
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
"""
from django.test import TestCase
from django.urls import reverse
from http import HTTPStatus
from django.contrib.auth.models import Group

class SurveyRenderTestCase(TestCase):
    @classmethod
    def setUpTestData(cls):
        Group.objects.create(name='admin')

    def setUp(self):
        pass

class WellTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_fixture']

    def test_well(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = ' id="survey:'
        self.assertContains(response, survey_id, count=1)  # one and only one survey included
        survey_id = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id)  # the right survey included

class DisabledSurveyTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_disabled_survey_fixture']

    def test_disabled_survey(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=0)  # one and only one survey included

class RegistryTestCase(SurveyRenderTestCase):
    fixtures = ['survey_registry_fixture']

    def test_registry(self):
        url = reverse('registry-legacy')
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=1)  # one and only one survey included
        survey_id = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id)  # the right survey included

class SearchTestCase(SurveyRenderTestCase):
    fixtures = ['survey_search_fixture']

    def test_search(self):
        url = reverse('search')
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=1)  # one and only one survey included
        survey_id = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id)  # the right survey included

class NoSurveysTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture']

    def test_no_surveys(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=0)

class MultipleSurveysTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_fixture1', 'survey_well_fixture2']

    def test_multiple_surveys(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=2)  # 2 surveys included
        survey_id1 = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id1)  # the first survey is included
        survey_id2 = survey_id + '6ec24102-129d-44ea-ad3d-67c2517e3fb9'
        self.assertContains(response, survey_id2)  # the second survey is included
<commit_msg>removed test reference<commit_after>"""
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
"""
from django.test import TestCase
from django.urls import reverse
from http import HTTPStatus
from django.contrib.auth.models import Group

class SurveyRenderTestCase(TestCase):
    @classmethod
    def setUpTestData(cls):
        Group.objects.create(name='admin')

    def setUp(self):
        pass

class WellTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_fixture']

    def test_well(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = ' id="survey:'
        self.assertContains(response, survey_id, count=1)  # one and only one survey included
        survey_id = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id)  # the right survey included

class DisabledSurveyTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_disabled_survey_fixture']

    def test_disabled_survey(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=0)  # one and only one survey included

class SearchTestCase(SurveyRenderTestCase):
    fixtures = ['survey_search_fixture']

    def test_search(self):
        url = reverse('search')
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=1)  # one and only one survey included
        survey_id = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id)  # the right survey included

class NoSurveysTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture']

    def test_no_surveys(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=0)

class MultipleSurveysTestCase(SurveyRenderTestCase):
    fixtures = ['well_detail_fixture', 'survey_well_fixture1', 'survey_well_fixture2']

    def test_multiple_surveys(self):
        url = reverse('well_detail', kwargs={'pk': '123'})
        response = self.client.get(url)
        self.assertEqual(response.status_code, HTTPStatus.OK)
        survey_id = 'id="survey:'
        self.assertContains(response, survey_id, count=2)  # 2 surveys included
        survey_id1 = survey_id + '495a9927-5a13-490e-bf1d-08bf2048b098'
        self.assertContains(response, survey_id1)  # the first survey is included
        survey_id2 = survey_id + '6ec24102-129d-44ea-ad3d-67c2517e3fb9'
        self.assertContains(response, survey_id2)  # the second survey is included
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    werkzeug.contrib.wrappers
    ~~~~~~~~~~~~~~~~~~~~~~~~~

    Extra wrappers or mixins contributed by the community.  These wrappers can
    be mixed in into request objects to add extra functionality.

    Example::

        from werkzeug.wrappers import Request as RequestBase
        from werkzeug.contrib.wrappers import JSONRequestMixin

        class Request(RequestBase, JSONRequestMixin):
            pass

    Afterwards this request object provides the extra functionality of the
    :class:`JSONRequestMixin`.

    :copyright: (c) 2014 by the Werkzeug Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""
import codecs
try:
    from simplejson import loads
except ImportError:
    from json import loads

from werkzeug.exceptions import BadRequest
from werkzeug.utils import cached_property
from werkzeug.http import dump_options_header, parse_options_header
from werkzeug._compat import wsgi_decoding_dance

def is_known_charset(charset):
    """Checks if the given charset is known to Python."""
    try:
        codecs.lookup(charset)
    except LookupError:
        return False
    return True

class JSONRequestMixin(object):

    """Add json method to a request object.  This will parse the input data
    through simplejson if possible.

    :exc:`~werkzeug.exceptions.BadRequest` will be raised if the content-type
    is not json or if the data itself cannot be parsed as json.
    """

    @cached_property
    def json(self):
        """Get the result of simplejson.loads if possible."""
        if 'json' not in self.environ.get('CONTENT_TYPE', ''):
            raise BadRequest('Not a JSON request')
        try:
            return loads(self.data.decode(self.charset, self.encoding_errors))
        except Exception:
            raise BadRequest('Unable to read JSON request')

class ProtobufRequestMixin(object):

    """Add protobuf parsing method to a request object.  This will parse the
    input data through `protobuf`_ if possible.

    :exc:`~werkzeug.exceptions.BadRequest` will be raised if the content-type
    is not protobuf or if the data itself cannot be parsed property.

    .. _protobuf: http://code.google.com/p/protobuf/
    """

    #: by default the :class:`ProtobufRequestMixin` will raise a
    #: :exc:`~werkzeug.exceptions.BadRequest` if the object is not
    #: initialized.  You can bypass that check by setting this
    #: attribute to `False`.
    protobuf_check_initialization = True

    def parse_protobuf(self, proto_type):
        """Parse the data into an instance of proto_type."""
        if 'protobuf' not in self.environ.get('CONTENT_TYPE', ''):
            raise BadRequest('Not a Protobuf request')

        obj = proto_type()
        try:
            obj.ParseFromString(self.data)
        except Exception:
            raise BadRequest("Unable to parse Protobuf request")

        # Fail if not all required fields are set
        if self.protobuf_check_initialization and not obj.IsInitialized():
            raise BadRequest("Partial Protobuf request")

        return obj

class RoutingArgsRequestMixin(object):

    """This request mixin adds support for the wsgiorg routing args
    `specification`_.

    .. _specification: http://www.wsgi.org/wsgi/Specifications/routing_args
    """

    def _get_routing_args(self):
        return self.environ.get('wsgiorg.routing_args', (()))[0]

    def _set_routing_args(self, value):
        if self.shallow:
            raise RuntimeError('A shallow request tried to modify the WSGI '
                               'environment.  If you really want to do that, '
                               'set `shallow` to False.')
        self.environ['wsgiorg.routing_args'] = (value, self.routing_vars)

    routing_args = property(_get_routing_args, _set_routing_args, doc='''
        The positional URL arguments as `tuple`.''')
    del _get_routing_args, _set_routing_args

    def _get_routing_vars(self):
        rv = self.environ.get('wsgiorg.routing_args')
        if rv is not None:
            return rv[1]
        rv = {}
        if not self.shallow:
            self.routing_vars = rv
        return rv

    def _set_routing_vars(self, value):
        if self.shallow:
            raise RuntimeError('A shallow request tried to modify the WSGI '
                               'environment.  If you really want to do that, '
                               'set `shallow` to False.')
        self.environ['wsgiorg.routing_args'] = (self.routing_args, value)

    routing_vars = property(_get_routing_vars, _set_routing_vars, doc='''
        The keyword URL arguments as `dict`.''')
    del _get_routing_vars, _set_routing_vars

class ReverseSlashBehaviorRequestMixin(object):

    """This mixin reverses the trailing slash behavior of :attr:`script_root`
    and :attr:`path`.  This makes it possible to use :func:`~urlparse.urljoin`
    directly on the paths.

    Because it changes the behavior or :class:`Request` this class has to be
    mixed in *before* the actual request class::

        class MyRequest(ReverseSlashBehaviorRequestMixin, Request):
            pass

    This example shows the differences (for an application mounted on
    `/application` and the request going to `/application/foo/bar`):

        +---------------+-------------------+---------------------+
        |               | normal behavior   | reverse behavior    |
        +===============+===================+=====================+
        | `script_root` | ``/application``  | ``/application/``   |
        +---------------+-------------------+---------------------+
        | `path`        | ``/foo/bar``      | ``foo/bar``         |
        +---------------+-------------------+---------------------+
    """

    @cached_property
    def path(self):
        """Requested path as unicode.  This works a bit like the regular path
        info in the WSGI environment but will not include a leading slash.
        """
        path = wsgi_decoding_dance(self.environ.get('PATH_INFO') or '',
                                   self.charset, self.encoding_errors)
        return path.lstrip('/')

    @cached_property
    def script_root(self):
        """The root path of the script includling a trailing slash."""
        path = wsgi_decoding_dance(self.environ.get('SCRIPT_NAME') or '',
                                   self.charset, self.encoding_errors)
        return path.rstrip('/') + '/'

class DynamicCharsetRequestMixin(object):

    """"If this mixin is mixed into a request class it will provide
    a dynamic `charset` attribute.  This means that if the charset is
    transmitted in the content type headers it's used from there.

    Because it changes the behavior or :class:`Request` this class has
    to be mixed in *before* the actual request class::

        class MyRequest(DynamicCharsetRequestMixin, Request):
            pass

    By default the request object assumes that the URL charset is the
    same as the data charset.  If the charset varies on each request
    based on the transmitted data it's not a good idea to let the URLs
    change based on that.  Most browsers assume either utf-8 or latin1
    for the URLs if they have troubles figuring out.  It's strongly
    recommended to set the URL charset to utf-8::

        class MyRequest(DynamicCharsetRequestMixin, Request):
            url_charset = 'utf-8'

    .. versionadded:: 0.6
    """

    #: the default charset that is assumed if the content type header
    #: is missing or does not contain a charset parameter.  The default
    #: is latin1 which is what HTTP specifies as default charset.
    #: You may however want to set this to utf-8 to better support
    #: browsers that do not transmit a charset for incoming data.
    default_charset = 'latin1'

    def unknown_charset(self, charset):
        """Called if a charset was provided but is not supported by
        the Python codecs module.  By default latin1 is assumed then
        to not lose any information, you may override this method to
        change the behavior.

        :param charset: the charset that was not found.
        :return: the replacement charset.
        """
        return 'latin1'

    @cached_property
    def charset(self):
        """The charset from the content type."""
        header = self.environ.get('CONTENT_TYPE')
        if header:
            ct, options = parse_options_header(header)
            charset = options.get('charset')
            if charset:
                if is_known_charset(charset):
                    return charset
                return self.unknown_charset(charset)
        return self.default_charset

class DynamicCharsetResponseMixin(object):

    """If this mixin is mixed into a response class it will provide
    a dynamic `charset` attribute.  This means that if the charset is
    looked up and stored in the `Content-Type` header and updates
    itself automatically.  This also means a small performance hit but
    can be useful if you're working with different charsets on
    responses.

    Because the charset attribute is no a property at class-level, the
    default value is stored in `default_charset`.

    Because it changes the behavior or :class:`Response` this class has
    to be mixed in *before* the actual response class::

        class MyResponse(DynamicCharsetResponseMixin, Response):
            pass

    .. versionadded:: 0.6
    """

    #: the default charset.
    default_charset = 'utf-8'

    def _get_charset(self):
        header = self.headers.get('content-type')
        if header:
            charset = parse_options_header(header)[1].get('charset')
            if charset:
                return charset
        return self.default_charset

    def _set_charset(self, charset):
        header = self.headers.get('content-type')
        ct, options = parse_options_header(header)
        if not ct:
            raise TypeError('Cannot set charset if Content-Type '
                            'header is missing.')
        options['charset'] = charset
        self.headers['Content-Type'] = dump_options_header(ct, options)

    charset = property(_get_charset, _set_charset, doc="""
        The charset for the response.  It's stored inside the
        Content-Type header as a parameter.""")
    del _get_charset, _set_charset
<commit_msg>Update specification URL for RoutingArgsRequestMixin (#990)<commit_after># -*- coding: utf-8 -*-
"""
    werkzeug.contrib.wrappers
    ~~~~~~~~~~~~~~~~~~~~~~~~~

    Extra wrappers or mixins contributed by the community.  These wrappers can
    be mixed in into request objects to add extra functionality.

    Example::

        from werkzeug.wrappers import Request as RequestBase
        from werkzeug.contrib.wrappers import JSONRequestMixin

        class Request(RequestBase, JSONRequestMixin):
            pass

    Afterwards this request object provides the extra functionality of the
    :class:`JSONRequestMixin`.

    :copyright: (c) 2014 by the Werkzeug Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""
import codecs
try:
    from simplejson import loads
except ImportError:
    from json import loads

from werkzeug.exceptions import BadRequest
from werkzeug.utils import cached_property
from werkzeug.http import dump_options_header, parse_options_header
from werkzeug._compat import wsgi_decoding_dance

def is_known_charset(charset):
    """Checks if the given charset is known to Python."""
    try:
        codecs.lookup(charset)
    except LookupError:
        return False
    return True

class JSONRequestMixin(object):

    """Add json method to a request object.  This will parse the input data
    through simplejson if possible.

    :exc:`~werkzeug.exceptions.BadRequest` will be raised if the content-type
    is not json or if the data itself cannot be parsed as json.
    """

    @cached_property
    def json(self):
        """Get the result of simplejson.loads if possible."""
        if 'json' not in self.environ.get('CONTENT_TYPE', ''):
            raise BadRequest('Not a JSON request')
        try:
            return loads(self.data.decode(self.charset, self.encoding_errors))
        except Exception:
            raise BadRequest('Unable to read JSON request')

class ProtobufRequestMixin(object):

    """Add protobuf parsing method to a request object.  This will parse the
    input data through `protobuf`_ if possible.

    :exc:`~werkzeug.exceptions.BadRequest` will be raised if the content-type
    is not protobuf or if the data itself cannot be parsed property.

    .. _protobuf: http://code.google.com/p/protobuf/
    """

    #: by default the :class:`ProtobufRequestMixin` will raise a
    #: :exc:`~werkzeug.exceptions.BadRequest` if the object is not
    #: initialized.  You can bypass that check by setting this
    #: attribute to `False`.
    protobuf_check_initialization = True

    def parse_protobuf(self, proto_type):
        """Parse the data into an instance of proto_type."""
        if 'protobuf' not in self.environ.get('CONTENT_TYPE', ''):
            raise BadRequest('Not a Protobuf request')

        obj = proto_type()
        try:
            obj.ParseFromString(self.data)
        except Exception:
            raise BadRequest("Unable to parse Protobuf request")

        # Fail if not all required fields are set
        if self.protobuf_check_initialization and not obj.IsInitialized():
            raise BadRequest("Partial Protobuf request")

        return obj

class RoutingArgsRequestMixin(object):

    """This request mixin adds support for the wsgiorg routing args
    `specification`_.

    .. _specification: https://wsgi.readthedocs.io/en/latest/specifications/routing_args.html
    """

    def _get_routing_args(self):
        return self.environ.get('wsgiorg.routing_args', (()))[0]

    def _set_routing_args(self, value):
        if self.shallow:
            raise RuntimeError('A shallow request tried to modify the WSGI '
                               'environment.  If you really want to do that, '
                               'set `shallow` to False.')
        self.environ['wsgiorg.routing_args'] = (value, self.routing_vars)

    routing_args = property(_get_routing_args, _set_routing_args, doc='''
        The positional URL arguments as `tuple`.''')
    del _get_routing_args, _set_routing_args

    def _get_routing_vars(self):
        rv = self.environ.get('wsgiorg.routing_args')
        if rv is not None:
            return rv[1]
        rv = {}
        if not self.shallow:
            self.routing_vars = rv
        return rv

    def _set_routing_vars(self, value):
        if self.shallow:
            raise RuntimeError('A shallow request tried to modify the WSGI '
                               'environment.  If you really want to do that, '
                               'set `shallow` to False.')
        self.environ['wsgiorg.routing_args'] = (self.routing_args, value)

    routing_vars = property(_get_routing_vars, _set_routing_vars, doc='''
        The keyword URL arguments as `dict`.''')
    del _get_routing_vars, _set_routing_vars

class ReverseSlashBehaviorRequestMixin(object):

    """This mixin reverses the trailing slash behavior of :attr:`script_root`
    and :attr:`path`.  This makes it possible to use :func:`~urlparse.urljoin`
    directly on the paths.

    Because it changes the behavior or :class:`Request` this class has to be
    mixed in *before* the actual request class::

        class MyRequest(ReverseSlashBehaviorRequestMixin, Request):
            pass

    This example shows the differences (for an application mounted on
    `/application` and the request going to `/application/foo/bar`):

        +---------------+-------------------+---------------------+
        |               | normal behavior   | reverse behavior    |
        +===============+===================+=====================+
        | `script_root` | ``/application``  | ``/application/``   |
        +---------------+-------------------+---------------------+
        | `path`        | ``/foo/bar``      | ``foo/bar``         |
        +---------------+-------------------+---------------------+
    """

    @cached_property
    def path(self):
        """Requested path as unicode.  This works a bit like the regular path
        info in the WSGI environment but will not include a leading slash.
        """
        path = wsgi_decoding_dance(self.environ.get('PATH_INFO') or '',
                                   self.charset, self.encoding_errors)
        return path.lstrip('/')

    @cached_property
    def script_root(self):
        """The root path of the script includling a trailing slash."""
        path = wsgi_decoding_dance(self.environ.get('SCRIPT_NAME') or '',
                                   self.charset, self.encoding_errors)
        return path.rstrip('/') + '/'

class DynamicCharsetRequestMixin(object):

    """"If this mixin is mixed into a request class it will provide
    a dynamic `charset` attribute.  This means that if the charset is
    transmitted in the content type headers it's used from there.

    Because it changes the behavior or :class:`Request` this class has
    to be mixed in *before* the actual request class::

        class MyRequest(DynamicCharsetRequestMixin, Request):
            pass

    By default the request object assumes that the URL charset is the
    same as the data charset.  If the charset varies on each request
    based on the transmitted data it's not a good idea to let the URLs
    change based on that.  Most browsers assume either utf-8 or latin1
    for the URLs if they have troubles figuring out.  It's strongly
    recommended to set the URL charset to utf-8::

        class MyRequest(DynamicCharsetRequestMixin, Request):
            url_charset = 'utf-8'

    .. versionadded:: 0.6
    """

    #: the default charset that is assumed if the content type header
    #: is missing or does not contain a charset parameter.  The default
    #: is latin1 which is what HTTP specifies as default charset.
    #: You may however want to set this to utf-8 to better support
    #: browsers that do not transmit a charset for incoming data.
    default_charset = 'latin1'

    def unknown_charset(self, charset):
        """Called if a charset was provided but is not supported by
        the Python codecs module.  By default latin1 is assumed then
        to not lose any information, you may override this method to
        change the behavior.

        :param charset: the charset that was not found.
        :return: the replacement charset.
        """
        return 'latin1'

    @cached_property
    def charset(self):
        """The charset from the content type."""
        header = self.environ.get('CONTENT_TYPE')
        if header:
            ct, options = parse_options_header(header)
            charset = options.get('charset')
            if charset:
                if is_known_charset(charset):
                    return charset
                return self.unknown_charset(charset)
        return self.default_charset

class DynamicCharsetResponseMixin(object):

    """If this mixin is mixed into a response class it will provide
    a dynamic `charset` attribute.  This means that if the charset is
    looked up and stored in the `Content-Type` header and updates
    itself automatically.  This also means a small performance hit but
    can be useful if you're working with different charsets on
    responses.

    Because the charset attribute is no a property at class-level, the
    default value is stored in `default_charset`.

    Because it changes the behavior or :class:`Response` this class has
    to be mixed in *before* the actual response class::

        class MyResponse(DynamicCharsetResponseMixin, Response):
            pass

    .. versionadded:: 0.6
    """

    #: the default charset.
    default_charset = 'utf-8'

    def _get_charset(self):
        header = self.headers.get('content-type')
        if header:
            charset = parse_options_header(header)[1].get('charset')
            if charset:
                return charset
        return self.default_charset

    def _set_charset(self, charset):
        header = self.headers.get('content-type')
        ct, options = parse_options_header(header)
        if not ct:
            raise TypeError('Cannot set charset if Content-Type '
                            'header is missing.')
        options['charset'] = charset
        self.headers['Content-Type'] = dump_options_header(ct, options)

    charset = property(_get_charset, _set_charset, doc="""
        The charset for the response.  It's stored inside the
        Content-Type header as a parameter.""")
    del _get_charset, _set_charset
<|endoftext|>


<commit_before>from django.core import serializers
from django.contrib import messages
from django.contrib.auth import logout as auth_logout
from django.contrib.auth.models import User
from django.contrib.auth.decorators import login_required
from whatisit.apps.api.views import getToken
from django.shortcuts import render_to_response, redirect, render
from django.template.context import RequestContext

import logging
import os
import pickle

from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse
from django.http import HttpResponse, HttpResponseBadRequest, HttpResponseRedirect

from whatisit.apps.users.forms import TeamForm
from whatisit.apps.users.models import Team
from whatisit.apps.users.utils import (
    get_team,
    get_user_team,
    has_team_edit_permission,
    remove_user_teams
)

from whatisit.apps.wordfish.utils import (
    summarize_team_annotations,
    summarize_teams_annotations
)

##################################################################################
# AUTHENTICATION VIEWS ###########################################################
##################################################################################

def login(request):
    #context = {'request': request, 'user': request.user}
    #context = RequestContext(request,context)
    #return render_to_response('social/login.html', context_instance=context)
    return render(request, 'social/login.html')

@login_required(login_url='/')
def home(request):
    return render_to_response('social/home.html')

@login_required
def token(request):
    '''getToken retrieves the user's token, and returns a page with it
    for the user to use to authenticate with the API
    '''
    token = getToken(request,json_response=False)
    context = RequestContext(request, {'request': request,
                                       'user': request.user,
                                       'token': token['token'] })
    return render_to_response("users/token.html", context_instance=context)

def logout(request):
    '''log out of the social authentication backend'''
    auth_logout(request)
    return redirect('/')

##################################################################################
# TEAM VIEWS #####################################################################
##################################################################################

@login_required
def edit_team(request, tid=None):
    '''edit_team is the view to edit an existing team, or create a new team.
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    if tid:
        team = get_team(request,tid)
        edit_permission = has_team_edit_permission(request,team)
        title = "Edit Team"
    else:
        team = Team()
        edit_permission = True
        title = "New Team"

    # Only the owner is allowed to edit a team
    if edit_permission:

        if request.method == "POST":
            form = TeamForm(request.POST,request.FILES,instance=team)
            if form.is_valid():
                team = form.save(commit=False)
                team.save()
                team.members.add(request.user)
                team.save()
                return HttpResponseRedirect(team.get_absolute_url())
        else:
            form = TeamForm(instance=team)

        context = {"form": form,
                   "edit_permission": edit_permission,
                   "title":title,
                   "team":team}

        return render(request, "teams/edit_team.html", context)

    # If user makes it down here, does not have permission
    messages.info(request, "You don't have permission to edit this team")
    return redirect("teams")

@login_required
def view_teams(	request):
    '''view all teams (log in not required)
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    teams = Team.objects.all()

    context = {"teams": teams}
    user_team = get_user_team(request)
    context['user_team'] = user_team # returns None if not in team

    return render(request, "teams/all_teams.html", context)

@login_required
def view_team(request, tid):
    '''view the details about a team
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    team = get_team(request,tid)

    # Need to create annotation counts with "total" for all members
    annotation_counts = summarize_team_annotations(team.members.all())

    # Only team members are allowed to edit their team
    edit_permission = has_team_edit_permission(request,team)

    context = {"team": team,
               "edit_permission":edit_permission,
               "annotation_counts":annotation_counts}

    return render(request, "teams/team_details.html", context)

@login_required
def join_team(request, tid):
    '''add a user to a new team, and remove from previous team
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    team = get_team(request,tid)
    user = request.user

    # Get the user's current team, and remove
    users_team = Team.objects.filter(members__contains=user)
    removed_team = remove_user_teams(remove_teams=users_team,
                                     user=user)

    # If the user was removed, tell him/her
    if removed_team != None:
        messages.info(request,"You have been removed from team %s" %(removed_team))

    # Add the user to the team
    if user not in team.members.all():
        team.members.add(user)
        team.save()
        messages.info(request,"%s has been successfully added to!" %(user.username,team.name))
    else:
        messages.info(request,"You are already a member of %s!" %(team.name))

    edit_permission = has_team_edit_permission(request,team)
    context = {"team": team,
               "edit_permission":edit_permission}
    return render(request, "teams/team_details.html", context)

# Python social auth extensions

def redirect_if_no_refresh_token(backend, response, social, *args, **kwargs):
    '''http://python-social-auth.readthedocs.io/en/latest/use_cases.html
       #re-prompt-google-oauth2-users-to-refresh-the-refresh-token
    '''
    if backend.name == 'google-oauth2' and social and response.get('refresh_token') is None and social.extra_data.get('refresh_token') is None:
        return redirect('/login/google-oauth2?approval_prompt=force')

# A User should not be allowed to associate a Github (or other) account with a different
# gmail, etc.
def social_user(backend, uid, user=None, *args, **kwargs):
    '''OVERRIDED: It will give the user an error message if the
    account is already associated with a username.'''
    provider = backend.name
    social = backend.strategy.storage.user.get_social_auth(provider, uid)
    if social:
        if user and social.user != user:
            msg = 'This {0} account is already in use.'.format(provider)
            return login(request=backend.strategy.request,
                         message=msg)
            #raise AuthAlreadyAssociated(backend, msg)
        elif not user:
            user = social.user

    return {'social': social,
            'user': user,
            'is_new': user is None,
            'new_association': social is None}
<commit_msg>fixing bug with "contains" to get a user team<commit_after>from django.core import serializers
from django.contrib import messages
from django.contrib.auth import logout as auth_logout
from django.contrib.auth.models import User
from django.contrib.auth.decorators import login_required
from whatisit.apps.api.views import getToken
from django.shortcuts import render_to_response, redirect, render
from django.template.context import RequestContext

import logging
import os
import pickle

from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse
from django.http import HttpResponse, HttpResponseBadRequest, HttpResponseRedirect

from whatisit.apps.users.forms import TeamForm
from whatisit.apps.users.models import Team
from whatisit.apps.users.utils import (
    get_team,
    get_user_team,
    has_team_edit_permission,
    remove_user_teams
)

from whatisit.apps.wordfish.utils import (
    summarize_team_annotations,
    summarize_teams_annotations
)

##################################################################################
# AUTHENTICATION VIEWS ###########################################################
##################################################################################

def login(request):
    #context = {'request': request, 'user': request.user}
    #context = RequestContext(request,context)
    #return render_to_response('social/login.html', context_instance=context)
    return render(request, 'social/login.html')

@login_required(login_url='/')
def home(request):
    return render_to_response('social/home.html')

@login_required
def token(request):
    '''getToken retrieves the user's token, and returns a page with it
    for the user to use to authenticate with the API
    '''
    token = getToken(request,json_response=False)
    context = RequestContext(request, {'request': request,
                                       'user': request.user,
                                       'token': token['token'] })
    return render_to_response("users/token.html", context_instance=context)

def logout(request):
    '''log out of the social authentication backend'''
    auth_logout(request)
    return redirect('/')

##################################################################################
# TEAM VIEWS #####################################################################
##################################################################################

@login_required
def edit_team(request, tid=None):
    '''edit_team is the view to edit an existing team, or create a new team.
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    if tid:
        team = get_team(request,tid)
        edit_permission = has_team_edit_permission(request,team)
        title = "Edit Team"
    else:
        team = Team()
        edit_permission = True
        title = "New Team"

    # Only the owner is allowed to edit a team
    if edit_permission:

        if request.method == "POST":
            form = TeamForm(request.POST,request.FILES,instance=team)
            if form.is_valid():
                team = form.save(commit=False)
                team.save()
                team.members.add(request.user)
                team.save()
                return HttpResponseRedirect(team.get_absolute_url())
        else:
            form = TeamForm(instance=team)

        context = {"form": form,
                   "edit_permission": edit_permission,
                   "title":title,
                   "team":team}

        return render(request, "teams/edit_team.html", context)

    # If user makes it down here, does not have permission
    messages.info(request, "You don't have permission to edit this team")
    return redirect("teams")

@login_required
def view_teams(	request):
    '''view all teams (log in not required)
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    teams = Team.objects.all()

    context = {"teams": teams}
    user_team = get_user_team(request)
    context['user_team'] = user_team # returns None if not in team

    return render(request, "teams/all_teams.html", context)

@login_required
def view_team(request, tid):
    '''view the details about a team
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    team = get_team(request,tid)

    # Need to create annotation counts with "total" for all members
    annotation_counts = summarize_team_annotations(team.members.all())

    # Only team members are allowed to edit their team
    edit_permission = has_team_edit_permission(request,team)

    context = {"team": team,
               "edit_permission":edit_permission,
               "annotation_counts":annotation_counts}

    return render(request, "teams/team_details.html", context)

@login_required
def join_team(request, tid):
    '''add a user to a new team, and remove from previous team
    :parma tid: the team id to edit or create. If none, indicates a new team
    '''
    team = get_team(request,tid)
    user = request.user

    # Get the user's current team, and remove
    users_team = get_user_team(request)
    removed_team = remove_user_teams(remove_teams=users_team,
                                     user=user)

    # If the user was removed, tell him/her
    if removed_team != None:
        messages.info(request,"You have been removed from team %s" %(removed_team))

    # Add the user to the team
    if user not in team.members.all():
        team.members.add(user)
        team.save()
        messages.info(request,"%s has been successfully added to!" %(user.username,team.name))
    else:
        messages.info(request,"You are already a member of %s!" %(team.name))

    edit_permission = has_team_edit_permission(request,team)
    context = {"team": team,
               "edit_permission":edit_permission}
    return render(request, "teams/team_details.html", context)

# Python social auth extensions

def redirect_if_no_refresh_token(backend, response, social, *args, **kwargs):
    '''http://python-social-auth.readthedocs.io/en/latest/use_cases.html
       #re-prompt-google-oauth2-users-to-refresh-the-refresh-token
    '''
    if backend.name == 'google-oauth2' and social and response.get('refresh_token') is None and social.extra_data.get('refresh_token') is None:
        return redirect('/login/google-oauth2?approval_prompt=force')

# A User should not be allowed to associate a Github (or other) account with a different
# gmail, etc.
def social_user(backend, uid, user=None, *args, **kwargs):
    '''OVERRIDED: It will give the user an error message if the
    account is already associated with a username.'''
    provider = backend.name
    social = backend.strategy.storage.user.get_social_auth(provider, uid)
    if social:
        if user and social.user != user:
            msg = 'This {0} account is already in use.'.format(provider)
            return login(request=backend.strategy.request,
                         message=msg)
            #raise AuthAlreadyAssociated(backend, msg)
        elif not user:
            user = social.user

    return {'social': social,
            'user': user,
            'is_new': user is None,
            'new_association': social is None}
<|endoftext|>