<commit_before># coding: utf8
import base64
import codecs
import hashlib
import json
import math
import re
import string
import sys

import sublime

import sublime_plugin

try:
    from .encodingutils.escape_table import (
        html_escape_table,
        html5_escape_table,
        html_reserved_list,
        xml_escape_table
    )
except ValueError:
    from encodingutils.escape_table import (
        html5_escape_table,
        html_escape_table,
        html_reserved_list,
        xml_escape_table
    )

try:
    import urllib.parse
    quote_plus = urllib.parse.quote_plus
    unquote_plus = urllib.parse.unquote_plus
except ImportError:
    import urllib

    def quote_plus(text):
        return urllib.quote_plus(text.encode('utf8'))

    def unquote_plus(text):
        return urllib.unquote_plus(text.encode('utf8'))

try:
    unichr(32)
except NameError:
    def unichr(val):
        return chr(val)

class StringEncodePaste(sublime_plugin.WindowCommand):
    def run(self, **kwargs):
        items = [
            ('NCR Encode', 'panos_ncr'),
            ('NCR Decode', 'dencr'),
            ('Html Entitize', 'html_entitize'),
            ('Html Deentitize', 'html_deentitize'),
            ('Css Escape', 'css_escape'),
            ('Css Unescape', 'css_unescape'),
            ('Safe Html Entitize', 'safe_html_entitize'),
            ('Safe Html Deentitize', 'safe_html_deentitize'),
            ('Xml Entitize', 'xml_entitize'),
            ('Xml Deentitize', 'xml_deentitize'),
            ('Json Escape', 'json_escape'),
            ('Json Unescape', 'json_unescape'),
            ('Url Encode', 'url_encode'),
            ('Url Decode', 'url_decode'),
            ('Base64 Encode', 'base64_encode'),
            ('Base64 Decode', 'base64_decode'),
            ('Md5 Encode', 'md5_encode'),
            ('Sha256 Encode', 'sha256_encode'),
            ('Sha512 Encode', 'sha512_encode'),
            ('Escape Regex', 'escape_regex'),
            ('Escape Like', 'escape_like'),
            ('Hex Dec', 'hex_dec'),
            ('Dec Hex', 'dec_hex'),
            ('Unicode Hex', 'unicode_hex'),
            ('Hex Unicode', 'hex_unicode'),
        ]

        lines = list(map(lambda line: line[0], items))
        commands = list(map(lambda line: line[1], items))
        view = self.window.active_view()
        if not view:
            return

        def on_done(item):
            if item == -1:
                return
            view.run_command(commands[item], {'source': 'clipboard'})

        self.window.show_quick_panel(lines, on_done)

class StringEncode(sublime_plugin.TextCommand):
    def run(self, edit, **kwargs):
        regions = self.view.sel()

        if kwargs.get('source') == 'clipboard':
            del kwargs['source']
            text = sublime.get_clipboard()
            replacement = self.encode(text, **kwargs)
            for region in regions:
                if region.empty():
                    self.view.insert(edit, region.begin(), replacement)
                else:
                    self.view.replace(edit, region, replacement)
            return

        elif 'source' in kwargs:
            sublime.status_message('Unsupported source {0!r}'.format(kwargs['source']))
            return

        if any(map(lambda region: region.empty(), regions)):
            regions = [sublime.Region(0, self.view.size())]
        for region in regions:
            text = self.view.substr(region)
            replacement = self.encode(text, **kwargs)
            self.view.replace(edit, region, replacement)

class ShannonCommand(StringEncode):
    def encode(self, text):
        def range_bytes():
            return range(256)

        def range_printable():
            return (ord(c) for c in string.printable)

        def h(data, iterator=range_bytes):
            if not data:
                return 0
            entropy = 0
            for x in iterator():
                p_x = float(data.count(chr(x))) / len(data)
                if p_x > 0:
                    entropy += - p_x * math.log(p_x, 2)
            return entropy
        ret = 0
        ret = str(h(text, range_printable))
        return ret

class StrengthCommand(StringEncode):
    def encode(self, text):

        def read_str(psw):
            self.numeric = re.compile('\d')
            self.loweralpha = re.compile('[a-z]')
            self.upperalpha = re.compile('[A-Z]')
            # self.symbols = re.compile('[-_.:,;<>?"#$%&/()!@~]')
            self.symbols = re.compile('[-~`@#$%^&*()_+=/?>.<,;:"]')
            self.extended = re.compile('[^\x00-\x7F]+')
            self.num_of_symbols = 20  # adjust accordingly...
            from math import log, pow
            charset = 0
            if self.numeric.search(psw):
                charset += 10
            if self.loweralpha.search(psw):
                charset += 26
            if self.upperalpha.search(psw):
                charset += 26
            if self.symbols.search(psw):
                charset += self.num_of_symbols
            # if ((len(psw) > 0) and (flag == 0)):
                # str_entropy = str(float(log(pow(255, len(psw)), 2)))
                str_entropy = str(float(log(pow(charset, len(psw)), 2)))
            return str_entropy
            # a!
        ret = 0
        ret = str(read_str(text))
        return ret

class CheckOrdCommand(StringEncode):
    def encode(self, text):
        ret = ''
        for i, c in enumerate(text):
            ret += str(ord(c)) + '.'
        return ret

class PanosRotCommand(StringEncode):
    def encode(self, text):
        return codecs.encode(text, 'rot_13')

class PanosNcrCommand(StringEncode):
    def encode(self, text):
        ret = ''
        for i, c in enumerate(text):
            # if ord(c) > 127:
                ret += '&#' + str(ord(c)) + ';'
            # else:
                # ret += c
        return ret

class DencrCommand(StringEncode):
    def encode(self, text):
        while re.search('&#[0-9]+;', text):
            match = re.search('&#([0-9]+);', text)
            text = text.replace(match.group(0), unichr(int(match.group(1), 10)))
        text = text.replace('&', '&')
        return text

class HtmlEntitizeCommand(StringEncode):

    def encode(self, text):
        text = text.replace('&', '&')
        for k in html_escape_table:
            v = html_escape_table[k]
            text = text.replace(k, v)
        ret = ''
        for i, c in enumerate(text):
            if ord(c) > 127:
                ret += hex(ord(c)).replace('0x', '&#x') + ';'
            else:
                ret += c
        return ret

class HtmlDeentitizeCommand(StringEncode):

    def encode(self, text):
        for k in html_escape_table:
            v = html_escape_table[k]
            text = text.replace(v, k)
        for k in html5_escape_table:
            v = html5_escape_table[k]
            text = text.replace(v, k)
        while re.search('&#[xX][a-fA-F0-9]+;', text):
            match = re.search('&#[xX]([a-fA-F0-9]+);', text)
            text = text.replace(
                match.group(0), unichr(int('0x' + match.group(1), 16)))
        text = text.replace('&', '&')
        return text

class CssEscapeCommand(StringEncode):

    def encode(self, text):
        ret = ''
        for i, c in enumerate(text):
            if ord(c) > 127:
                ret += hex(ord(c)).replace('0x', '\\')
            else:
                ret += c
        return ret

class CssUnescapeCommand(StringEncode):

    def encode(self, text):
        while re.search(r'\\[a-fA-F0-9]+', text):
            match = re.search(r'\\([a-fA-F0-9]+)', text)
            text = text.replace(
                match.group(0), unichr(int('0x' + match.group(1), 16)))
        return text

class SafeHtmlEntitizeCommand(StringEncode):

    def encode(self, text):
        for k in html_escape_table:
            # skip HTML reserved characters
            if k in html_reserved_list:
                continue
            v = html_escape_table[k]
            text = text.replace(k, v)
        ret = ''
        for i, c in enumerate(text):
            if ord(c) > 127:
                ret += hex(ord(c)).replace('0x', '&#x') + ';'
            else:
                ret += c
        return ret

class SafeHtmlDeentitizeCommand(StringEncode):

    def encode(self, text):
        for k in html_escape_table:
            # skip HTML reserved characters
            if k in html_reserved_list:
                continue
            v = html_escape_table[k]
            text = text.replace(v, k)
        while re.search('&#[xX][a-fA-F0-9]+;', text):
            match = re.search('&#[xX]([a-fA-F0-9]+);', text)
            text = text.replace(
                match.group(0), unichr(int('0x' + match.group(1), 16)))
        text = text.replace('&', '&')
        return text

class XmlEntitizeCommand(StringEncode):

    def encode(self, text):
        text = text.replace('&', '&')
        for k in xml_escape_table:
            v = xml_escape_table[k]
            text = text.replace(k, v)
        ret = ''
        for i, c in enumerate(text):
            if ord(c) > 127:
                ret += hex(ord(c)).replace('0x', '&#x') + ';'
            else:
                ret += c
        return ret

class XmlDeentitizeCommand(StringEncode):

    def encode(self, text):
        for k in xml_escape_table:
            v = xml_escape_table[k]
            text = text.replace(v, k)
        text = text.replace('&', '&')
        return text

class JsonEscapeCommand(StringEncode):

    def encode(self, text):
        return json.dumps(text)

class JsonUnescapeCommand(StringEncode):

    def encode(self, text):
        return json.loads(text)

class UrlEncodeCommand(StringEncode):

    def encode(self, text, old_school=True):
        quoted = quote_plus(text)
        if old_school:
            return quoted.replace("+", "%20")
        return quoted

class UrlDecodeCommand(StringEncode):

    def encode(self, text):
        return unquote_plus(text)

class Base64EncodeCommand(StringEncode):

    def encode(self, text):
        return base64.b64encode(text.encode('raw_unicode_escape')).decode('ascii')

class Base64DecodeCommand(StringEncode):

    def encode(self, text):
        return base64.b64decode(text).decode('raw_unicode_escape')

class Md5EncodeCommand(StringEncode):

    def encode(self, text):
        hasher = hashlib.md5()
        hasher.update(bytes(text, 'utf-8'))
        return hasher.hexdigest()

class Sha256EncodeCommand(StringEncode):

    def encode(self, text):
        hasher = hashlib.sha256()
        hasher.update(bytes(text, 'utf-8'))
        return hasher.hexdigest()

class Sha512EncodeCommand(StringEncode):

    def encode(self, text):
        hasher = hashlib.sha512()
        hasher.update(bytes(text, 'utf-8'))
        return hasher.hexdigest()

class Escaper(StringEncode):

    def encode(self, text):
        return re.sub(r'(?<!\\)(%s)' % self.meta, r'\\\1', text)

class EscapeRegexCommand(Escaper):
    meta = r'[\\*.+^$()\[\]\{\}]'

class EscapeLikeCommand(Escaper):
    meta = r'[%_]'

class HexDecCommand(StringEncode):

    def encode(self, text):
        return str(int(text, 16))

class DecHexCommand(StringEncode):

    def encode(self, text):
        return hex(int(text))

class UnicodeHexCommand(StringEncode):

    def encode(self, text):
        hex_text = u''
        text_bytes = bytes(text, 'utf-16')

        if text_bytes[0:2] == b'\xff\xfe':
            endian = 'little'
            text_bytes = text_bytes[2:]
        elif text_bytes[0:2] == b'\xfe\xff':
            endian = 'big'
            text_bytes = text_bytes[2:]

        char_index = 0
        for c in text_bytes:
            if char_index == 0:
                c1 = c
                char_index += 1
            elif char_index == 1:
                c2 = c
                if endian == 'little':
                    c1, c2 = c2, c1
                tmp = (c1 << 8) + c2
                if tmp < 0x80:
                    hex_text += chr(tmp)
                    char_index = 0
                elif tmp >= 0xd800 and tmp <= 0xdbff:
                    char_index += 1
                else:
                    hex_text += '\\u' + '{0:04x}'.format(tmp)
                    char_index = 0
            elif char_index == 2:
                c3 = c
                char_index += 1
            elif char_index == 3:
                c4 = c
                if endian == 'little':
                    c3, c4 = c4, c3
                tmp1 = ((c1 << 8) + c2) - 0xd800
                tmp2 = ((c3 << 8) + c4) - 0xdc00
                tmp = (tmp1 * 0x400) + tmp2 + 0x10000
                hex_text += '\\U' + '{0:08x}'.format(tmp)
                char_index = 0
        return hex_text

class HexUnicodeCommand(StringEncode):

    def encode(self, text):
        uni_text = text

        endian = sys.byteorder

        r = re.compile(r'\\u([0-9a-fA-F]{2})([0-9a-fA-F]{2})')
        rr = r.search(uni_text)
        while rr:
            first_byte = int(rr.group(1), 16)

            if first_byte >= 0xd8 and first_byte <= 0xdf:
                # Surrogate pair
                pass
            else:
                if endian == 'little':
                    b1 = int(rr.group(2), 16)
                    b2 = int(rr.group(1), 16)
                else:
                    b1 = int(rr.group(1), 16)
                    b2 = int(rr.group(2), 16)

                ch = bytes([b1, b2]).decode('utf-16')

                uni_text = uni_text.replace(rr.group(0), ch)
            rr = r.search(uni_text, rr.start(0) + 1)

        # Surrogate pair (2 bytes + 2 bytes)
        r = re.compile(
            r'\\u([0-9a-fA-F]{2})([0-9a-fA-F]{2})\\u([0-9a-fA-F]{2})([0-9a-fA-F]{2})')
        rr = r.search(uni_text)
        while rr:
            if endian == 'little':
                b1 = int(rr.group(2), 16)
                b2 = int(rr.group(1), 16)
                b3 = int(rr.group(4), 16)
                b4 = int(rr.group(3), 16)
            else:
                b1 = int(rr.group(1), 16)
                b2 = int(rr.group(2), 16)
                b3 = int(rr.group(3), 16)
                b4 = int(rr.group(4), 16)

            ch = bytes([b1, b2, b3, b4]).decode('utf-16')

            uni_text = uni_text.replace(rr.group(0), ch)
            rr = r.search(uni_text)

        # Surrogate pair (4 bytes)
        r = re.compile(
            r'\\U([0-9a-fA-F]{2})([0-9a-fA-F]{2})([0-9a-fA-F]{2})([0-9a-fA-F]{2})')
        rr = r.search(uni_text)
        while rr:
            tmp = (int(rr.group(1), 16) << 24) \
                + (int(rr.group(2), 16) << 16) \
                + (int(rr.group(3), 16) << 8) \
                + (int(rr.group(4), 16))

            if (tmp <= 0xffff):
                ch = chr(tmp)
            else:
                tmp -= 0x10000
                c1 = 0xd800 + int(tmp / 0x400)
                c2 = 0xdc00 + int(tmp % 0x400)
                if endian == 'little':
                    b1 = c1 & 0xff
                    b2 = c1 >> 8
                    b3 = c2 & 0xff
                    b4 = c2 >> 8
                else:
                    b1 = c1 >> 8
                    b2 = c1 & 0xff
                    b3 = c2 >> 8
                    b4 = c2 & 0xff

                ch = bytes([b1, b2, b3, b4]).decode('utf-16')

            uni_text = uni_text.replace(rr.group(0), ch)
            rr = r.search(uni_text)

        return uni_text
<commit_msg>Delete encoding_utils.py<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
checks.py

Each function in here should

- Take a DataFrame as its first argument, maybe optional arguments
- Makes its assert on the result
- Return the original DataFrame
"""
import numpy as np
import pandas as pd
import pandas.util.testing as tm
import six

from engarde import generic
from engarde.generic import verify, verify_all, verify_any

def none_missing(df, columns=None):
    """
    Asserts that there are no missing values (NaNs) in the DataFrame.

    Parameters
    ----------
    df : DataFrame
    columns : list
      list of columns to restrict the check to

    Returns
    -------
    df : DataFrame
      same as the original
    """
    if columns is None:
        columns = df.columns
    try:
        assert not df[columns].isnull().any().any()
    except AssertionError as e:
        missing = df[columns].isnull()
        msg = generic.bad_locations(missing)
        e.args = msg
        raise
    return df

def is_monotonic(df, items=None, increasing=None, strict=False):
    """
    Asserts that the DataFrame is monotonic.

    Parameters
    ==========

    df : Series or DataFrame
    items : dict
        mapping columns to conditions (increasing, strict)
    increasing : None or bool
        None is either increasing or decreasing.
    strict : whether the comparison should be strict

    Returns
    =======
    df : DataFrame
    """
    if items is None:
        items = {k: (increasing, strict) for k in df}

    for col, (increasing, strict) in items.items():
        s = pd.Index(df[col])
        if increasing:
            good = getattr(s, 'is_monotonic_increasing')
        elif increasing is None:
            good = getattr(s, 'is_monotonic') | getattr(s, 'is_monotonic_decreasing')
        else:
            good = getattr(s, 'is_monotonic_decreasing')
        if strict:
            if increasing:
                good = good & (s.to_series().diff().dropna() > 0).all()
            elif increasing is None:
                good = good & ((s.to_series().diff().dropna() > 0).all() |
                               (s.to_series().diff().dropna() < 0).all())
            else:
                good = good & (s.to_series().diff().dropna() < 0).all()
        if not good:
            raise AssertionError
    return df

def is_shape(df, shape):
    """
    Asserts that the DataFrame is of a known shape.

    Parameters
    ==========

    df : DataFrame
    shape : tuple
      (n_rows, n_columns). Use None or -1 if you don't care
      about a dimension.

    Returns
    =======
    df : DataFrame
    """
    try:
        check = np.all(np.equal(df.shape, shape) | (np.equal(shape, [-1, -1]) |
                                                    np.equal(shape, [None, None])))
        assert check
    except AssertionError as e:
        msg = ("Expected shape: {}\n"
               "\t\tActual shape:   {}".format(shape, df.shape))
        e.args = (msg,)
        raise
    return df

def unique(df, columns=None):
    """
    Asserts that columns in the DataFrame only have unique values.

    Parameters
    ----------
    df : DataFrame
    columns : list
      list of columns to restrict the check to. If None, check all columns.

    Returns
    -------
    df : DataFrame
      same as the original
    """
    if columns is None:
        columns = df.columns
    for col in columns:
        if not df[col].is_unique:
            raise AssertionError("Column {!r} contains non-unique values".format(col))
    return df

def unique_index(df):
    """
    Assert that the index is unique

    Parameters
    ==========
    df : DataFrame

    Returns
    =======
    df : DataFrame
    """
    try:
        assert df.index.is_unique
    except AssertionError as e:
        e.args = df.index.get_duplicates()
        raise
    return df

def within_set(df, items=None):
    """
    Assert that df is a subset of items

    Parameters
    ==========
    df : DataFrame
    items : dict
      mapping of columns (k) to array-like of values (v) that
      ``df[k]`` is expected to be a subset of

    Returns
    =======
    df : DataFrame
    """
    for k, v in items.items():
        if not df[k].isin(v).all():
            bad = df.loc[~df[k].isin(v), k]
            raise AssertionError('Not in set', bad)
    return df

def within_range(df, items=None):
    """
    Assert that a DataFrame is within a range.

    Parameters
    ==========
    df : DataFame
    items : dict
      mapping of columns (k) to a (low, high) tuple (v)
      that ``df[k]`` is expected to be between.

    Returns
    =======
    df : DataFrame
    """
    for k, (lower, upper) in items.items():
        if (lower > df[k]).any() or (upper < df[k]).any():
            bad = (lower > df[k]) | (upper < df[k])
            raise AssertionError("Outside range", bad)
    return df

def within_n_std(df, n=3):
    """
    Assert that every value is within ``n`` standard
    deviations of its column's mean.

    Parameters
    ==========
    df : DataFame
    n : int
      number of standard deviations from the mean

    Returns
    =======
    df : DataFrame
    """
    means = df.mean()
    stds = df.std()
    inliers = (np.abs(df[means.index] - means) < n * stds)
    if not np.all(inliers):
        msg = generic.bad_locations(~inliers)
        raise AssertionError(msg)
    return df

def has_dtypes(df, items):
    """
    Assert that a DataFrame has ``dtypes``

    Parameters
    ==========
    df: DataFrame
    items: dict
      mapping of columns to dtype.

    Returns
    =======
    df : DataFrame
    """
    dtypes = df.dtypes
    for k, v in items.items():
        if not dtypes[k] == v:
            raise AssertionError("{} has the wrong dtype ({})".format(k, v))
    return df

def one_to_many(df, unitcol, manycol):
    """
    Assert that a many-to-one relationship is preserved between two
    columns. For example, a retail store will have have distinct
    departments, each with several employees. If each employee may
    only work in a single department, then the relationship of the
    department to the employees is one to many.

    Parameters
    ==========
    df : DataFrame
    unitcol : str
        The column that encapulates the groups in ``manycol``.
    manycol : str
        The column that must remain unique in the distict pairs
        between ``manycol`` and ``unitcol``

    Returns
    =======
    df : DataFrame

    """
    subset = df[[manycol, unitcol]].drop_duplicates()
    for many in subset[manycol].unique():
        if subset[subset[manycol] == many].shape[0] > 1:
            msg = "{} in {} has multiple values for {}".format(many, manycol, unitcol)
            raise AssertionError(msg)

    return df

def is_same_as(df, df_to_compare, **kwargs):
    """
    Assert that two pandas dataframes are the equal

    Parameters
    ==========
    df : pandas DataFrame
    df_to_compare : pandas DataFrame
    **kwargs : dict
        keyword arguments passed through to panda's ``assert_frame_equal``

    Returns
    =======
    df : DataFrame

    """
    try:
        tm.assert_frame_equal(df, df_to_compare, **kwargs)
    except AssertionError as exc:
        six.raise_from(AssertionError("DataFrames are not equal"), exc)
    return df

__all__ = ['is_monotonic', 'is_same_as', 'is_shape', 'none_missing',
           'unique_index', 'within_n_std', 'within_range', 'within_set',
           'has_dtypes', 'verify', 'verify_all', 'verify_any',
           'one_to_many','is_same_as',]

<commit_msg>Add dtypes[k] to has_dtypes error message<commit_after># -*- coding: utf-8 -*-
"""
checks.py

Each function in here should

- Take a DataFrame as its first argument, maybe optional arguments
- Makes its assert on the result
- Return the original DataFrame
"""
import numpy as np
import pandas as pd
import pandas.util.testing as tm
import six

from engarde import generic
from engarde.generic import verify, verify_all, verify_any

def none_missing(df, columns=None):
    """
    Asserts that there are no missing values (NaNs) in the DataFrame.

    Parameters
    ----------
    df : DataFrame
    columns : list
      list of columns to restrict the check to

    Returns
    -------
    df : DataFrame
      same as the original
    """
    if columns is None:
        columns = df.columns
    try:
        assert not df[columns].isnull().any().any()
    except AssertionError as e:
        missing = df[columns].isnull()
        msg = generic.bad_locations(missing)
        e.args = msg
        raise
    return df

def is_monotonic(df, items=None, increasing=None, strict=False):
    """
    Asserts that the DataFrame is monotonic.

    Parameters
    ==========

    df : Series or DataFrame
    items : dict
        mapping columns to conditions (increasing, strict)
    increasing : None or bool
        None is either increasing or decreasing.
    strict : whether the comparison should be strict

    Returns
    =======
    df : DataFrame
    """
    if items is None:
        items = {k: (increasing, strict) for k in df}

    for col, (increasing, strict) in items.items():
        s = pd.Index(df[col])
        if increasing:
            good = getattr(s, 'is_monotonic_increasing')
        elif increasing is None:
            good = getattr(s, 'is_monotonic') | getattr(s, 'is_monotonic_decreasing')
        else:
            good = getattr(s, 'is_monotonic_decreasing')
        if strict:
            if increasing:
                good = good & (s.to_series().diff().dropna() > 0).all()
            elif increasing is None:
                good = good & ((s.to_series().diff().dropna() > 0).all() |
                               (s.to_series().diff().dropna() < 0).all())
            else:
                good = good & (s.to_series().diff().dropna() < 0).all()
        if not good:
            raise AssertionError
    return df

def is_shape(df, shape):
    """
    Asserts that the DataFrame is of a known shape.

    Parameters
    ==========

    df : DataFrame
    shape : tuple
      (n_rows, n_columns). Use None or -1 if you don't care
      about a dimension.

    Returns
    =======
    df : DataFrame
    """
    try:
        check = np.all(np.equal(df.shape, shape) | (np.equal(shape, [-1, -1]) |
                                                    np.equal(shape, [None, None])))
        assert check
    except AssertionError as e:
        msg = ("Expected shape: {}\n"
               "\t\tActual shape:   {}".format(shape, df.shape))
        e.args = (msg,)
        raise
    return df

def unique(df, columns=None):
    """
    Asserts that columns in the DataFrame only have unique values.

    Parameters
    ----------
    df : DataFrame
    columns : list
      list of columns to restrict the check to. If None, check all columns.

    Returns
    -------
    df : DataFrame
      same as the original
    """
    if columns is None:
        columns = df.columns
    for col in columns:
        if not df[col].is_unique:
            raise AssertionError("Column {!r} contains non-unique values".format(col))
    return df

def unique_index(df):
    """
    Assert that the index is unique

    Parameters
    ==========
    df : DataFrame

    Returns
    =======
    df : DataFrame
    """
    try:
        assert df.index.is_unique
    except AssertionError as e:
        e.args = df.index.get_duplicates()
        raise
    return df

def within_set(df, items=None):
    """
    Assert that df is a subset of items

    Parameters
    ==========
    df : DataFrame
    items : dict
      mapping of columns (k) to array-like of values (v) that
      ``df[k]`` is expected to be a subset of

    Returns
    =======
    df : DataFrame
    """
    for k, v in items.items():
        if not df[k].isin(v).all():
            bad = df.loc[~df[k].isin(v), k]
            raise AssertionError('Not in set', bad)
    return df

def within_range(df, items=None):
    """
    Assert that a DataFrame is within a range.

    Parameters
    ==========
    df : DataFame
    items : dict
      mapping of columns (k) to a (low, high) tuple (v)
      that ``df[k]`` is expected to be between.

    Returns
    =======
    df : DataFrame
    """
    for k, (lower, upper) in items.items():
        if (lower > df[k]).any() or (upper < df[k]).any():
            bad = (lower > df[k]) | (upper < df[k])
            raise AssertionError("Outside range", bad)
    return df

def within_n_std(df, n=3):
    """
    Assert that every value is within ``n`` standard
    deviations of its column's mean.

    Parameters
    ==========
    df : DataFame
    n : int
      number of standard deviations from the mean

    Returns
    =======
    df : DataFrame
    """
    means = df.mean()
    stds = df.std()
    inliers = (np.abs(df[means.index] - means) < n * stds)
    if not np.all(inliers):
        msg = generic.bad_locations(~inliers)
        raise AssertionError(msg)
    return df

def has_dtypes(df, items):
    """
    Assert that a DataFrame has ``dtypes``

    Parameters
    ==========
    df: DataFrame
    items: dict
      mapping of columns to dtype.

    Returns
    =======
    df : DataFrame
    """
    dtypes = df.dtypes
    for k, v in items.items():
        if not dtypes[k] == v:
            raise AssertionError("{} has the wrong dtype. Should be ({}), is ({})".format(k, v,dtypes[k]))
    return df

def one_to_many(df, unitcol, manycol):
    """
    Assert that a many-to-one relationship is preserved between two
    columns. For example, a retail store will have have distinct
    departments, each with several employees. If each employee may
    only work in a single department, then the relationship of the
    department to the employees is one to many.

    Parameters
    ==========
    df : DataFrame
    unitcol : str
        The column that encapulates the groups in ``manycol``.
    manycol : str
        The column that must remain unique in the distict pairs
        between ``manycol`` and ``unitcol``

    Returns
    =======
    df : DataFrame

    """
    subset = df[[manycol, unitcol]].drop_duplicates()
    for many in subset[manycol].unique():
        if subset[subset[manycol] == many].shape[0] > 1:
            msg = "{} in {} has multiple values for {}".format(many, manycol, unitcol)
            raise AssertionError(msg)

    return df

def is_same_as(df, df_to_compare, **kwargs):
    """
    Assert that two pandas dataframes are the equal

    Parameters
    ==========
    df : pandas DataFrame
    df_to_compare : pandas DataFrame
    **kwargs : dict
        keyword arguments passed through to panda's ``assert_frame_equal``

    Returns
    =======
    df : DataFrame

    """
    try:
        tm.assert_frame_equal(df, df_to_compare, **kwargs)
    except AssertionError as exc:
        six.raise_from(AssertionError("DataFrames are not equal"), exc)
    return df

__all__ = ['is_monotonic', 'is_same_as', 'is_shape', 'none_missing',
           'unique_index', 'within_n_std', 'within_range', 'within_set',
           'has_dtypes', 'verify', 'verify_all', 'verify_any',
           'one_to_many','is_same_as',]
<|endoftext|>


<commit_before>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 03-07-2017 16:50
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module contains Ensembl models used as Entities and DAO/Services
"""

import config_manager

class Species:
    def __init__(self, ensembl_species_entry):
        self.__ensembl_species_entry = ensembl_species_entry
        self.__logger = config_manager.get_app_config_manager().get_logger_for(__name__)

    def _get_logger(self):
        return self.__logger

    def get_division(self):
        # TODO
        pass

    def get_ncbi_taxonomy_id(self):
        # TODO
        pass

    def get_name(self):
        # TODO
        pass

    def get_ensembl_release(self):
        # TODO
        pass

    def get_display_name(self):
        # TODO
        pass

    def get_assembly_accession(self):
        # TODO
        pass

    def get_strain_collection(self):
        # TODO
        pass

    def get_common_name(self):
        # TODO
        pass

    def get_strain(self):
        # TODO
        pass

    def get_aliases(self):
        # TODO
        pass

    def get_groups(self):
        # TODO
        pass

    def get_assembly(self):
        # TODO
        pass

class SpeciesService:
    # TODO
    pass
<commit_msg>ensembl_species_entry getter<commit_after>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 03-07-2017 16:50
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module contains Ensembl models used as Entities and DAO/Services
"""

import config_manager

class Species:
    def __init__(self, ensembl_species_entry):
        self.__ensembl_species_entry = ensembl_species_entry
        self.__logger = config_manager.get_app_config_manager().get_logger_for(__name__)

    def _get_logger(self):
        return self.__logger

    def get_ensembl_species_entry(self):
        return self.__ensembl_species_entry

    def get_division(self):
        # TODO
        pass

    def get_ncbi_taxonomy_id(self):
        # TODO
        pass

    def get_name(self):
        # TODO
        pass

    def get_ensembl_release(self):
        # TODO
        pass

    def get_display_name(self):
        # TODO
        pass

    def get_assembly_accession(self):
        # TODO
        pass

    def get_strain_collection(self):
        # TODO
        pass

    def get_common_name(self):
        # TODO
        pass

    def get_strain(self):
        # TODO
        pass

    def get_aliases(self):
        # TODO
        pass

    def get_groups(self):
        # TODO
        pass

    def get_assembly(self):
        # TODO
        pass

class SpeciesService:
    # TODO
    pass
<|endoftext|>


<commit_before>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 03-07-2017 16:50
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module contains Ensembl models used as Entities and DAO/Services
"""

import config_manager

class Species:
    def __init__(self, ensembl_species_entry):
        self.__ensembl_species_entry = ensembl_species_entry
        self.__logger = config_manager.get_app_config_manager().get_logger_for(__name__)

    def _get_logger(self):
        return self.__logger

    def _get_value_for_key_or_default(self, key, default='-not_available-'):
        if key in self.get_ensembl_species_entry():
            return self.get_ensembl_species_entry()[key]
        return default

    def get_ensembl_species_entry(self):
        return self.__ensembl_species_entry

    def get_division(self):
        return self._get_value_for_key_or_default('division')

    def get_ncbi_taxonomy_id(self):
        return self._get_value_for_key_or_default('taxon_id')

    def get_name(self):
        return self._get_value_for_key_or_default('name')

    def get_ensembl_release(self):
        return self._get_value_for_key_or_default('release')

    def get_display_name(self):
        return self._get_value_for_key_or_default('display_name')

    def get_assembly_accession(self):
        return self._get_value_for_key_or_default('accession')

    def get_strain_collection(self):
        if self._get_value_for_key_or_default('strain_collection', 'null') == 'null':
            return None
        return self._get_value_for_key_or_default('strain_collection')

    def get_common_name(self):
        return self._get_value_for_key_or_default('common_name')

    def get_strain(self):
        return self._get_value_for_key_or_default('strain')

    def get_aliases(self):
        return self._get_value_for_key_or_default('aliases', [])

    def get_groups(self):
        return self._get_value_for_key_or_default('groups', [])

    def get_assembly(self):
        # TODO
        pass

class SpeciesService:
    # TODO
    pass
<commit_msg>get assembly from ensembl species entry<commit_after>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 03-07-2017 16:50
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module contains Ensembl models used as Entities and DAO/Services
"""

import config_manager

class Species:
    def __init__(self, ensembl_species_entry):
        self.__ensembl_species_entry = ensembl_species_entry
        self.__logger = config_manager.get_app_config_manager().get_logger_for(__name__)

    def _get_logger(self):
        return self.__logger

    def _get_value_for_key_or_default(self, key, default='-not_available-'):
        if key in self.get_ensembl_species_entry():
            return self.get_ensembl_species_entry()[key]
        return default

    def get_ensembl_species_entry(self):
        return self.__ensembl_species_entry

    def get_division(self):
        return self._get_value_for_key_or_default('division')

    def get_ncbi_taxonomy_id(self):
        return self._get_value_for_key_or_default('taxon_id')

    def get_name(self):
        return self._get_value_for_key_or_default('name')

    def get_ensembl_release(self):
        return self._get_value_for_key_or_default('release')

    def get_display_name(self):
        return self._get_value_for_key_or_default('display_name')

    def get_assembly_accession(self):
        return self._get_value_for_key_or_default('accession')

    def get_strain_collection(self):
        if self._get_value_for_key_or_default('strain_collection', 'null') == 'null':
            return None
        return self._get_value_for_key_or_default('strain_collection')

    def get_common_name(self):
        return self._get_value_for_key_or_default('common_name')

    def get_strain(self):
        return self._get_value_for_key_or_default('strain')

    def get_aliases(self):
        return self._get_value_for_key_or_default('aliases', [])

    def get_groups(self):
        return self._get_value_for_key_or_default('groups', [])

    def get_assembly(self):
        return self._get_value_for_key_or_default('assembly')

class SpeciesService:
    # TODO
    pass
<|endoftext|>


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

# Copyright (c) 2017 Ansible Tower by Red Hat
# All Rights Reserved.

# Django
from django.conf import settings

# AWX
from awx.main.models import Instance

def _add_remove_celery_worker_queues(app, controlled_instances, worker_queues, worker_name):
    removed_queues = []
    added_queues = []
    ig_names = set(['tower_instance_router'])
    hostnames = set([instance.hostname for instance in controlled_instances])
    for instance in controlled_instances:
        ig_names.update(instance.rampart_groups.values_list('name', flat=True))
    worker_queue_names = set([q['name'] for q in worker_queues])

    # Remove queues that aren't in the instance group
    for queue in worker_queues:
        if queue['name'] in settings.AWX_CELERY_QUEUES_STATIC or \
                queue['alias'] in settings.AWX_CELERY_QUEUES_STATIC:
            continue

        if queue['name'] not in ig_names | hostnames or not instance.enabled:
            app.control.cancel_consumer(queue['name'], reply=True, destination=[worker_name])
            removed_queues.append(queue['name'])

    # Add queues for instance and instance groups
    for queue_name in ig_names | hostnames:
        if queue_name not in worker_queue_names:
            app.control.add_consumer(queue_name, reply=True, destination=[worker_name])
            added_queues.append(queue_name)

    return (added_queues, removed_queues)

def update_celery_worker_routes(instance, conf):
    tasks = [
        'awx.main.tasks.cluster_node_heartbeat',
        'awx.main.tasks.purge_old_stdout_files',
    ]
    routes_updated = {}
    # Instance is, effectively, a controller node
    if instance.is_controller():
        tasks.append('awx.main.tasks.awx_isolated_heartbeat')
    else:
        if 'awx.main.tasks.awx_isolated_heartbeat' in conf.CELERY_ROUTES:
            del conf.CELERY_ROUTES['awx.main.tasks.awx_isolated_heartbeat']

    for t in tasks:
        conf.CELERY_ROUTES[t] = {'queue': instance.hostname, 'routing_key': instance.hostname}
        routes_updated[t] = conf.CELERY_ROUTES[t]

    return routes_updated

def register_celery_worker_queues(app, celery_worker_name):
    instance = Instance.objects.me()
    controlled_instances = [instance]
    if instance.is_controller():
        controlled_instances.extend(Instance.objects.filter(
            rampart_groups__controller__instances__hostname=instance.hostname
        ))
    added_queues = []
    removed_queues = []

    celery_host_queues = app.control.inspect([celery_worker_name]).active_queues()

    celery_worker_queues = celery_host_queues[celery_worker_name] if celery_host_queues else []
    (added_queues, removed_queues) = _add_remove_celery_worker_queues(app, controlled_instances,
                                                                      celery_worker_queues, celery_worker_name)

    return (controlled_instances, removed_queues, added_queues)

<commit_msg>use non-unicode queue names<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2017 Ansible Tower by Red Hat
# All Rights Reserved.

# Python
import six

# Django
from django.conf import settings

# AWX
from awx.main.models import Instance

def _add_remove_celery_worker_queues(app, controlled_instances, worker_queues, worker_name):
    removed_queues = []
    added_queues = []
    ig_names = set([six.text_type('tower_instance_router')])
    hostnames = set([instance.hostname for instance in controlled_instances])
    for instance in controlled_instances:
        ig_names.update(instance.rampart_groups.values_list('name', flat=True))
    worker_queue_names = set([q['name'] for q in worker_queues])

    # Remove queues that aren't in the instance group
    for queue in worker_queues:
        if queue['name'] in settings.AWX_CELERY_QUEUES_STATIC or \
                queue['alias'] in settings.AWX_CELERY_QUEUES_STATIC:
            continue

        if queue['name'] not in ig_names | hostnames or not instance.enabled:
            app.control.cancel_consumer(queue['name'].encode("utf8"), reply=True, destination=[worker_name])
            removed_queues.append(queue['name'].encode("utf8"))

    # Add queues for instance and instance groups
    for queue_name in ig_names | hostnames:
        if queue_name not in worker_queue_names:
            app.control.add_consumer(queue_name.encode("utf8"), reply=True, destination=[worker_name])
            added_queues.append(queue_name.encode("utf8"))

    return (added_queues, removed_queues)

def update_celery_worker_routes(instance, conf):
    tasks = [
        'awx.main.tasks.cluster_node_heartbeat',
        'awx.main.tasks.purge_old_stdout_files',
    ]
    routes_updated = {}
    # Instance is, effectively, a controller node
    if instance.is_controller():
        tasks.append('awx.main.tasks.awx_isolated_heartbeat')
    else:
        if 'awx.main.tasks.awx_isolated_heartbeat' in conf.CELERY_ROUTES:
            del conf.CELERY_ROUTES['awx.main.tasks.awx_isolated_heartbeat']

    for t in tasks:
        conf.CELERY_ROUTES[t] = {'queue': instance.hostname.encode("utf8"), 'routing_key': instance.hostname.encode("utf8")}
        routes_updated[t] = conf.CELERY_ROUTES[t]

    return routes_updated

def register_celery_worker_queues(app, celery_worker_name):
    instance = Instance.objects.me()
    controlled_instances = [instance]
    if instance.is_controller():
        controlled_instances.extend(Instance.objects.filter(
            rampart_groups__controller__instances__hostname=instance.hostname
        ))
    added_queues = []
    removed_queues = []

    celery_host_queues = app.control.inspect([celery_worker_name]).active_queues()

    celery_worker_queues = celery_host_queues[celery_worker_name] if celery_host_queues else []
    (added_queues, removed_queues) = _add_remove_celery_worker_queues(app, controlled_instances,
                                                                      celery_worker_queues, celery_worker_name)

    return (controlled_instances, removed_queues, added_queues)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Tue Nov 01 15:17:52 2011

Author: Mike
Author: Josef

mainly script for checking Kernel Regression
"""

if __name__ == "__main__":
    from scikits.statsmodels.sandbox.nonparametric import smoothers as s
    import matplotlib.pyplot as plt
    from numpy import sin, array, random

    import time
    random.seed(500)
    x = random.normal(size = 250)
    y = array([sin(i*5)/i + 2*i + (3+i)*random.normal() for i in x])

    K = s.kernels.Biweight(0.25)
    K2 = s.kernels.CustomKernel(lambda x: (1 - x*x)**2, 0.25, domain = [-1.0,
                               1.0])

    KS = s.KernelSmoother(x, y, K)
    KS2 = s.KernelSmoother(x, y, K2)

    KSx = np.arange(-3, 3, 0.1)
    start = time.time()
    KSy = KS.conf(KSx)
    KVar = KS.std(KSx)
    print time.time() - start    # This should be significantly quicker...
    start = time.time()          #
    KS2y = KS2.conf(KSx)         #
    K2Var = KS2.std(KSx)         #
    print time.time() - start    # ...than this.

    KSConfIntx, KSConfInty = KS.conf(15)

    print "Norm const should be 0.9375"
    print K2.norm_const

    print "L2 Norms Should Match:"
    print K.L2Norm
    print K2.L2Norm

    print "Fit values should match:"
    #print zip(KSy, KS2y)
    print KSy[28]
    print KS2y[28]

    print "Var values should match:"
    #print zip(KVar, K2Var)
    print KVar[39]
    print K2Var[39]

    fig = plt.figure()
    ax = fig.add_subplot(221)
    ax.plot(x, y, "+")
    ax.plot(KSx, KSy, "-o")
    ax.set_ylim(-20, 30)
    ax2 = fig.add_subplot(222)
    ax2.plot(KSx, KVar, "-o")

    ax3 = fig.add_subplot(223)
    ax3.plot(x, y, "+")
    ax3.plot(KSx, KS2y, "-o")
    ax3.set_ylim(-20, 30)
    ax4 = fig.add_subplot(224)
    ax4.plot(KSx, K2Var, "-o")

    fig2 = plt.figure()
    ax5 = fig2.add_subplot(111)
    ax5.plot(x, y, "+")
    ax5.plot(KSConfIntx, KSConfInty, "-o")

    from scikits.statsmodels.nonparametric import lowess as lo
    ys = lo.lowess(y, x)
    ax5.plot(ys[:,0], ys[:,1], 'b-')
    ys2 = lo.lowess(y, x, frac=0.25)
    ax5.plot(ys2[:,0], ys2[:,1], 'b--', lw=2)

    xind = np.argsort(x)
    pmod = s.PolySmoother(5, x[xind])
    pmod.fit(y[xind])

    yp = pmod(x[xind])
    ax5.plot(x[xind], yp, 'k-')
    ax5.set_title('Kernel regression, lowess - blue, polysmooth - black')

    #plt.show()
<commit_msg>REF: smoothers example<commit_after># -*- coding: utf-8 -*-
"""
Created on Tue Nov 01 15:17:52 2011

Author: Mike
Author: Josef

mainly script for checking Kernel Regression
"""
import numpy as np

if __name__ == "__main__":
    #from scikits.statsmodels.sandbox.nonparametric import smoothers as s
    from scikits.statsmodels.sandbox.nonparametric import smoothers, kernels
    import matplotlib.pyplot as plt
    #from numpy import sin, array, random

    import time
    np.random.seed(500)
    x = np.random.normal(size = 250)
    y = np.array([np.sin(i*5)/i + 2*i + (3+i)*np.random.normal() for i in x])

    K = kernels.Biweight(0.25)
    K2 = kernels.CustomKernel(lambda x: (1 - x*x)**2, 0.25, domain = [-1.0,
                               1.0])

    KS = smoothers.KernelSmoother(x, y, K)
    KS2 = smoothers.KernelSmoother(x, y, K2)

    KSx = np.arange(-3, 3, 0.1)
    start = time.time()
    KSy = KS.conf(KSx)
    KVar = KS.std(KSx)
    print time.time() - start    # This should be significantly quicker...
    start = time.time()          #
    KS2y = KS2.conf(KSx)         #
    K2Var = KS2.std(KSx)         #
    print time.time() - start    # ...than this.

    KSConfIntx, KSConfInty = KS.conf(15)

    print "Norm const should be 0.9375"
    print K2.norm_const

    print "L2 Norms Should Match:"
    print K.L2Norm
    print K2.L2Norm

    print "Fit values should match:"
    #print zip(KSy, KS2y)
    print KSy[28]
    print KS2y[28]

    print "Var values should match:"
    #print zip(KVar, K2Var)
    print KVar[39]
    print K2Var[39]

    fig = plt.figure()
    ax = fig.add_subplot(221)
    ax.plot(x, y, "+")
    ax.plot(KSx, KSy, "-o")
    ax.set_ylim(-20, 30)
    ax2 = fig.add_subplot(222)
    ax2.plot(KSx, KVar, "-o")

    ax3 = fig.add_subplot(223)
    ax3.plot(x, y, "+")
    ax3.plot(KSx, KS2y, "-o")
    ax3.set_ylim(-20, 30)
    ax4 = fig.add_subplot(224)
    ax4.plot(KSx, K2Var, "-o")

    fig2 = plt.figure()
    ax5 = fig2.add_subplot(111)
    ax5.plot(x, y, "+")
    ax5.plot(KSConfIntx, KSConfInty, "-o")

    from scikits.statsmodels.nonparametric import lowess as lo
    ys = lo.lowess(y, x)
    ax5.plot(ys[:,0], ys[:,1], 'b-')
    ys2 = lo.lowess(y, x, frac=0.25)
    ax5.plot(ys2[:,0], ys2[:,1], 'b--', lw=2)

    #need to sort for matplolib plot ?
    xind = np.argsort(x)
    pmod = smoothers.PolySmoother(5, x[xind])
    pmod.fit(y[xind])

    yp = pmod(x[xind])
    ax5.plot(x[xind], yp, 'k-')
    ax5.set_title('Kernel regression, lowess - blue, polysmooth - black')

    #plt.show()
<|endoftext|>


<commit_before># Copyright (c) Arnaud Cassan.
# Distributed under the terms of the MIT license.

import os
os.environ["OMP_NUM_THREADS"] = "1"
import numpy as np
import emcee
import matplotlib.pylab as plt
from microlensing.utils import checkandtimeit, verbosity, printi, printd, printw
from multiprocessing import Pool
from microlensing.interferometry.ESPL import mod_VIS
from microlensing.interferometry.data import obs_VIS2

class InterferolensingModels():
    """List of interferometric microlensing models

        Parameters
        ----------
        model : string
            Model name.

        Attributes
        ----------
        model : string
            Model name.
        ndim : int
            Number of parameters.
        params : dict
            Dictionary of parameters (name and value, default values are None).
        pnames : string
            Print format for parameter names.
        lxnames : list
            Parameter names in LaTeX format.

        Usage
        -----
        >>> ilm = InterferolensingModels(model)
        >>> ilm.pnames
        >>> ilm.params['thetaE']
        >>> ilm.ndim
        >>> ilm.lxnames
    """

    def __init__(self, model):
        self.model = model

        # ESPL - single epoch, contouring
        if self.model == 'ESPL_CONT_SINGEP':
            self.ndim = 4
            self.params = {'alpha1':None, 'thetaE':None, 'eta1':None, 'u1':None}
            self.lxnames = [r'$\alpha_1$', r'$\theta_E$', r'$\eta_1$', r'$u_1$']

        # ESPL - single epoch, full integration
        elif self.model == 'ESPL_FULLINT_SINGEP':
            self.ndim = 4
            self.params = {'alpha1':None, 'thetaE':None, 'eta1':None, 'u1':None}
            self.lxnames = [r'$\alpha_1$', r'$\theta_E$', r'$\eta_1$', r'$u_1$']

        # ESPL - single epoch, point-source approximation
        elif self.model == 'ESPL_PS_SINGEP':
            self.ndim = 3
            self.params = {'alpha1':None, 'thetaE':None, 'u1':None}
            self.lxnames = [r'$\alpha_1$', r'$\theta_E$', r'$u_1$']

        # ESPL - single epoch, flat-images approsimation
        elif self.model == 'ESPL_FLAT_SINGEP':
            self.ndim = 3
            self.params = {'alpha1':None, 'thetaE':None, 'eta1':None}
            self.lxnames = [r'$\alpha_1$', r'$\theta_E$', r'$\eta_1$']

        # ESPL - multi-epochs, contouring
        elif self.model == 'ESPL_CONT_MULTIEP':
            self.ndim = 6
            self.params = {'alpha':None, 'thetaE':None, 'eta0':None, 't*':None, 't0':None, 'u0':None}
            self.lxnames = [r'$\alpha$', r'$\theta_E$', r'$\eta_0$', r'$t_*$', r'$t_0$', r'$u_0$']

        # ESPL - multi-epochs, full integration
        elif self.model == 'ESPL_FULLINT_MULTIEP':
            self.ndim = 6
            self.params = {'alpha':None, 'thetaE':None, 'eta0':None, 't*':None, 't0':None, 'u0':None}
            self.lxnames = [r'$\alpha$', r'$\theta_E$', r'$\eta_0$', r'$t_*$', r'$t_0$', r'$u_0$']

        # ESPL - multi-epochs, full integration, microlensing parameters
        elif self.model == 'ESPL_MLENSING_MULTIEP':
            self.ndim = 6
            self.params = {'alpha':None, 'thetaE':None, 'rho':None, 'tE':None, 't0':None, 'u0':None}
            self.lxnames = [r'$\alpha$', r'$\theta_E$', r'$\rho$', r'$t_E$', r'$t_0$', r'$u_0$']

        # ESPL - multi-epochs, point-source approximation
        elif self.model == 'ESPL_PS_MULTIEP':
            self.ndim = 5
            self.params = {'alpha':None, 'thetaE':None, 'tE':None, 't0':None, 'u0':None}
            self.lxnames = [r'$\alpha$', r'$\theta_E$', r'$t_E$', r'$t_0$', r'$u_0$']

        # ESPL - multi-epochs, flat-images approsimation
        elif self.model == 'ESPL_FLAT_MULTIEP':
            self.ndim = 5
            self.params = {'alpha':None, 'thetaE':None, 'eta0':None, 't*':None, 't0':None}
            self.lxnames = [r'$\alpha$', r'$\theta_E$', r'$\eta_0$', r'$t_*$', r'$t_0$']

        else:
            raise ValueError("Wrong model ID ({}) in InterferolensingModels".format(model))

        # print format for parameter names
        self.pnames = '(' + ', '.join([p for p in self.params]) + ')'

def fit_emcee(oifitslist, model, params, priorfun, samplerfile, nwalkers, chainlen, ncpu, Nc=20000, LLD=0., Na=8, tol=1e-5, ref_MJD=0., ifconverged=False, resume=False):
    """Compute MCMC chains using EMCEE-3

        Parameters
        ----------
        oifitslist : list of strings
            List of input OIFITS files
        params : list
            Constains 2-elements lists [[pi_low, pi_high], ...] of parameters bounds
            Ialpha, Irho, IthetaE, Iu0, ItE, It0 : arrays of [float, float]
            Intervals for drawing MCMC samples for each parameter
        Nc : int
            Common source contour initial sampling (global Nc_init)
            and final samplig (global Nc_fin)
        samplerfile : string
            Name of file where the chains are stored (HDF5 format)
        nwalkers : int
            Number of chains in parallel
        chainlen : int
            Individual chain length for one MCMC run
        ncpu : int
            Number of chains computed in parallel (threads)
        ref_MJD : float
            Reference MJD
        ifconverged : boolean
            Automatic stop if convergence
        resume : boolean
            Resume previous run
        hdulists (global) : list
            Contains a list [u, v, VIS2, VIS2ERR, MJD] for each (ep, B, lbd)

        Outputs
        -------
        samplerfile : string
            Name of file where the chains are stored (HDF5 format)
    """
    # set I/O shell display
    tcol, tun, tend, tit = "\033[0m\033[34m", "\033[0m\033[1m\033[34m", "\033[0m", "\033[0m\033[3m"

    # get model
    ilm = InterferolensingModels(model)
    ndim = ilm.ndim
    mnames = ilm.pnames
    printi(tcol + "Chosen model : " + tend + "{0}, {1}".format(model, mnames))

    # setup
    initial_length = chainlen
    resume_length = chainlen

    # read OIFITS files and create data fit-lists
    global hdulists
    hdulists = []
    for oifits in oifitslist:
        for b in range(6):
            for l in range(6):
                hdulists = hdulists + obs_VIS2(oifits, b + 1, l + 1, ref_MJD=ref_MJD)

    # initialize run
    if not resume:
        pos = []
        # Initialize the walkers
        for p in params:
            pos.append(np.random.uniform(p[0], p[1], nwalkers))
        pos = np.array(pos).T

        # Set up the backend
        backend = emcee.backends.HDFBackend(samplerfile)
        backend.reset(nwalkers, ndim)

        # Initialize the sampler
        with Pool(processes=ncpu) as pool:
            sampler = emcee.EnsembleSampler(nwalkers, ndim, logprob, args=[priorfun], kwargs={'model':model, 'Nc':Nc, 'LLD':LLD ,'Na':Na, 'tol':tol}, backend=backend, pool=pool)

            # Compute average changes of autocorrelation time
            index = 0
            autocorr = np.empty(initial_length)

            # Initialize the converge variable
            taup = np.inf

            # Now we'll sample for up to max_n steps
            for sample in sampler.sample(pos, iterations=initial_length, progress=True):
                if ifconverged:
                    # Convergence test every 300 steps
                    if sampler.iteration % 300: continue

                    # Force computing autocorrelation
                    tau = sampler.get_autocorr_time(tol=0)
                    autocorr[index] = np.mean(tau)
                    index = index + 1

                    # Test convergence
                    converged = np.all(tau * 100 < sampler.iteration)
                    converged &= np.all(np.abs(taup - tau) / tau < 0.01)
                    if converged: break
                    taup = tau

    # resume / continue run
    else:
        backend = emcee.backends.HDFBackend(samplerfile)
        printi(tcol + "Initial size :" + tend + " {0}".format(backend.iteration))
        with Pool(processes=ncpu) as pool:
            sampler = emcee.EnsembleSampler(nwalkers, ndim, logprob, args=[priorfun], kwargs={'model':model, 'Nc':Nc, 'LLD':LLD ,'Na':Na, 'tol':tol}, backend=backend, pool=pool)
            sampler.run_mcmc(None, resume_length, progress=True)
        printi(tcol + "Final size :" + tend + " {0}".format(backend.iteration))

def logprob(x, priorfun, model=None, Nc=None, LLD=None, Na=None, tol=None):
    """Compute gobal chi2 for a set epochs, baselines and wavelenghts,
    and fit for (alpha [deg], rho, thetaE [mas], u0, tE, t0) for time series
    or (alpha [deg], rho, thetaE [mas], u0) for a single epoch

        Parameters
        ----------
        hdulists (global) : list
            Contains a list [u, v, VIS2, VIS2ERR, MJD] for each (ep, B, lbd)
        x : list
            Model parameters (alpha, rho, thetaE, u0, tE, t0) for time series
            or (alpha, rho, thetaE, u0) for a single epoch

        Returns
        -------
        logp : float
            log(Likelihood) = - chi^2 / 2
    """
    # set I/O shell display
    tcol, tun, tend, tit = "\033[0m\033[34m", "\033[0m\033[1m\033[34m", "\033[0m", "\033[0m\033[3m"

    # conversion mas/deg
    mas = np.pi / 180. / 3600. / 1000.

    # sum chi2 over all epochs, baselines and wavelenghts
    chi2 = 0.
    for hdulist in hdulists:
        # get individual data
        u, v, VIS2, VIS2ERR, MJD = hdulist

        # conversion mas -> rad (thetaE is always x[1])
        u = u * mas * x[1]
        v = v * mas * x[1]

        # single epoch, contouring or full integration (alpha1, thetaE, eta1, u1)
        if model == 'ESPL_CONT_SINGEP' or model == 'ESPL_FULLINT_SINGEP':
            zetac = np.complex(0., x[3]) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = np.abs(mod_VIS(zetac, x[2] * x[3], u, v, model, Nc=Nc, LLD=LLD, Na=Na, tol=tol))**2

        # single epoch, point-source approximation (alpha1, thetaE, u1)
        elif model == 'ESPL_PS_SINGEP':
            zetac = np.complex(0., x[2]) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = mod_VIS(zetac, 0, u, v, model)**2

        # single epoch, flat-images approsimation (alpha1, thetaE, eta1)
        elif model == 'ESPL_FLAT_SINGEP':
            zetac = np.complex(0., 1.) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = mod_VIS(zetac, x[2], u, v, model, LLD=LLD, Na=Na, tol=tol)**2

        # multi-epochs, contouring or full integration (alpha, thetaE, eta0, t*, t0, u0)
        elif model == 'ESPL_CONT_MULTIEP' or model == 'ESPL_FULLINT_MULTIEP':
            zetac = np.complex((MJD - x[4]) * x[2] * x[5] / x[3], x[5]) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = np.abs(mod_VIS(zetac, x[2] * x[5], u, v, model, Nc=Nc, LLD=LLD, Na=Na, tol=tol))**2

        # multi-epochs, full integration, microlensing parameters (alpha, thetaE, rho, tE, t0, u0)
        elif model == 'ESPL_MLENSING_MULTIEP':
            zetac = np.complex((MJD - x[4]) / x[3], x[5]) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = np.abs(mod_VIS(zetac, x[2], u, v, model, Nc=Nc, LLD=LLD, Na=Na, tol=tol))**2

        # multi-epochs, point-source approximation (alpha, thetaE, tE, t0, u0)
        elif model == 'ESPL_PS_MULTIEP':
            zetac = np.complex((MJD - x[3]) / x[2], x[4]) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = mod_VIS(zetac, 0, u, v, model)**2

        # multi-epochs, flat-images approsimation (alpha, thetaE, eta0, t*, t0)
        elif model == 'ESPL_FLAT_MULTIEP':
            zetac = np.complex((MJD - x[4]) * x[2] / x[3], 1.) * np.exp(np.complex(0., np.pi / 180. * x[0]))
            VISE2 = mod_VIS(zetac, x[2], u, v, model, LLD=LLD, Na=Na, tol=tol)**2

        else:
            raise ValueError("Wrong model ID ({}) in logprob".format(model))

        # add chi2/pt
        chi2 +=  (VIS2 - VISE2)**2 / VIS2ERR**2

    # log probability
    logp = -0.5 * chi2 + priorfun(x)

    # print current values of parameters and chi2
    printd(tcol + "Parameters | chi2 | log(P): " + tend + "{0} | {1} | {2}".format(x, chi2, logp))

    # return log(prob)
    return logp
<commit_msg>changes<commit_after><|endoftext|>


<commit_before>
def calculate_threshold(sweep, trusted_range):
    from scipy.ndimage.measurements import histogram
    from thresholding import maximum_deviation
    import numpy

    threshold_list = []
    for i, flex_image in enumerate(sweep):

        # Get image as numpy array
        image = flex_image.as_numpy_array()

        # Cap pixels to within trusted range
        image.shape = -1
        ind = numpy.where(image < trusted_range[0])
        image[ind] = trusted_range[0]
        ind = numpy.where(image > trusted_range[1])
        image[ind] = trusted_range[1]

        # Histogram the pixels
        histo = histogram(image, trusted_range[0], trusted_range[1], trusted_range[1])
        histo = histo / numpy.sum(histo)

        # Calculate the threshold and add to list
        threshold = maximum_deviation(histo)
        threshold_list.append(threshold)

    # Return mean threshold
    return numpy.mean(threshold_list)

def select_strong_pixels(sweep, trusted_range):

    import numpy

    # Calculate the threshold
    print "Calculating a threshold."
    threshold = calculate_threshold(sweep, trusted_range)

    # Select only those pixels with counts > threshold
    print "Selecting pixels"
    image = sweep.to_array().as_numpy_array()
    mask = image >= threshold
    return image, mask

def group_pixels(mask):
    from scipy.ndimage.measurements import label, find_objects

    # Label the indices in the mask
    regions, nregions = label(mask)#, structure)

    # Get the bounding box of each object
    objects = find_objects(regions)

    # Return the list of objects
    return objects

def filter_objects(mask, objects, min_pixels):
    import numpy

    # Select objects with more pixels than min_pixels
    new_obj = []
    for obj in objects:
        ind = numpy.where(mask[obj] != 0)[0]
        if len(ind) >= min_pixels:
            new_obj.append(obj)
        else:
            mask[obj] = 0

    return new_obj

def direction_var(values, weights):
    import numpy
    from scitbx import matrix
    weights = numpy.array(weights)
    valx = numpy.array([x for x, y, z in values])
    valy = numpy.array([y for x, y, z in values])
    valz = numpy.array([z for x, y, z in values])

    # Calculate avergae x, y, z
    avrx = numpy.average(valx, weights=weights)
    avry = numpy.average(valy, weights=weights)
    avrz = numpy.average(valz, weights=weights)

    # Calculate mean direction vector
    s1m = matrix.col((avrx, avry, avrz)).normalize()

    # Calculate angles between vectors
    angles = []
    for s in values:
        angles.append(s1m.angle(s))

    # Calculate variance of angles
    angles = numpy.array(angles)
    var = numpy.dot(weights, (angles)**2)/numpy.sum(weights)
    return var

def centroid(image, mask, detector):
    from numpy import zeros, int32, argmax, where, average
    from scitbx.array_family import flex
    from scitbx import matrix
    import numpy

    var = []
    cent = []
    for obj in objects:
        xs = []
        ys = []
        zs = []
        s1s = []
        weights = []
        bbox = [obj[2].start, obj[2].stop,
                obj[1].start, obj[1].stop,
                obj[0].start, obj[0].stop]

        # Calcyulate beam vector for each point
        for i in range(bbox[4], bbox[5]):
            for s in range(bbox[2], bbox[3]):
                for f in range(bbox[0], bbox[1]):
                    if mask[i, s, f]:
                        s1 = matrix.col(detector.get_pixel_lab_coord((f, s)))
                        s1 = s1.normalize()
                        xs.append(f)
                        ys.append(s)
                        zs.append(i)
                        s1s.append(s1)
                        weights.append(image[i, s ,f])

        v = direction_var(s1s, weights)
        var.append(v)
        avrx = average(xs, weights=weights)
        avry = average(ys, weights=weights)
        avrz = average(zs, weights=weights)
        cent.append((avrx, avry, avrz))

    # Return a list of centroids and variances
    return numpy.array(cent), numpy.array(var)

def calculate_sigma_beam_divergence(var):
    '''Calculate the beam divergence as the sum of centroid variance of the
    intensity weighted diffracted beam directions.'''
    from math import sqrt

    # Calculate the sum of s^2
    sum_variance = reduce(lambda x, y: x + y, var)

    # Return the beam divergence as the sum / num reflections
    return sqrt(sum_variance / len(var))

def find_nearest_neighbour(obs_xyz, reflections):
    from annlib_ext import AnnAdaptor
    from scitbx.array_family import flex
    from math import sqrt

    # Get the predicted coords
    pred_xyz = []
    for r in reflections:
        x = r.image_coord_px[0]
        y = r.image_coord_px[1]
        z = r.frame_number
        pred_xyz.append((x, y, z))

    # Create the KD Tree
    ann = AnnAdaptor(flex.double(pred_xyz).as_1d(), 3)

    ann.query(flex.double(obs_xyz).as_1d())

#    for i in xrange(len(ann.nn)):
#        print "Neighbor of {0}, index {1} distance {2}".format(
#        obs_xyz[i], ann.nn[i], sqrt(ann.distances[i]))

    return ann.nn, ann.distances

def filter_objects_by_distance(ref_index, distance2, max_distance):

    import numpy
    from math import sqrt

    # Only accept objects closer than max distance
    indices = []
    for i in range(len(ref_index)):
        if sqrt(distance2[i]) > max_distance:
            continue
        else:
            indices.append(i)

    return numpy.array(indices)

if __name__ == '__main__':

    import os
    import libtbx.load_env
    from dials.util.nexus import NexusFile
    from glob import glob
    from dxtbx.sweep import SweepFactory
    from math import pi

    try:
        dials_regression = libtbx.env.dist_path('dials_regression')
    except KeyError, e:
        print 'FAIL: dials_regression not configured'
        raise

    # The XDS values
    xds_sigma_d = 0.060
    xds_sigma_m = 0.154

    # Read the reflections file
    filename = os.path.join(dials_regression,
        'centroid_test_data', 'reflections.h5')
    handle = NexusFile(filename, 'r')

    # Get the reflection list
    print 'Reading reflections.'
    reflections = handle.get_reflections()
    print 'Read {0} reflections.'.format(len(reflections))

    # Read images
    template = os.path.join(dials_regression,
        'centroid_test_data', 'centroid_*.cbf')
    filenames = glob(template)

    # Load the sweep
    print 'Loading sweep'
    sweep = SweepFactory.sweep(filenames)
    print 'Loaded sweep of {0} images.'.format(len(sweep))

    # Select the strong pixels to use in the divergence calculation
    print 'Select the strong pixels from the images.'
    trusted_range = (0, 20000)
    image, mask = select_strong_pixels(sweep, trusted_range)

    # Putting pixels into groups
    print 'Putting pixels into groups'
    objects = group_pixels(mask)
    print 'Found {0} objects'.format(len(objects))

    print 'Filtering objects'
    min_pixels = 6
    objects = filter_objects(mask, objects, min_pixels)
    print '{0} remaining objects'.format(len(objects))

    print 'Calculating centroid and variance.'
    cent, var = centroid(image, mask, sweep.get_detector())

    print 'Find object nearest predicted reflection'
    ref_index, distance2 = find_nearest_neighbour(cent, reflections)

    print 'Filter objects by distance from nearest reflection.'
    max_distance = 2
    indices = filter_objects_by_distance(ref_index, distance2, max_distance)
    print '{0} remaining objects'.format(len(indices))

    print 'Calculate the e.s.d of the beam divergence'
    sigma_d = calculate_sigma_beam_divergence(var[indices])
    print 'Sigma_d = {0} deg'.format(sigma_d * 180.0 / pi)
    print 'XDS Sigma_d = {0} deg'.format(xds_sigma_d)
<commit_msg>Added calculating phi/zeta.<commit_after>
def calculate_threshold(sweep, trusted_range):
    from scipy.ndimage.measurements import histogram
    from thresholding import maximum_deviation
    import numpy

    threshold_list = []
    for i, flex_image in enumerate(sweep):

        # Get image as numpy array
        image = flex_image.as_numpy_array()

        # Cap pixels to within trusted range
        image.shape = -1
        ind = numpy.where(image < trusted_range[0])
        image[ind] = trusted_range[0]
        ind = numpy.where(image > trusted_range[1])
        image[ind] = trusted_range[1]

        # Histogram the pixels
        histo = histogram(image, trusted_range[0], trusted_range[1], trusted_range[1])
        histo = histo / numpy.sum(histo)

        # Calculate the threshold and add to list
        threshold = maximum_deviation(histo)
        threshold_list.append(threshold)

    # Return mean threshold
    return numpy.mean(threshold_list)

def select_strong_pixels(sweep, trusted_range):

    import numpy

    # Calculate the threshold
    print "Calculating a threshold."
    threshold = calculate_threshold(sweep, trusted_range)

    # Select only those pixels with counts > threshold
    print "Selecting pixels"
    image = sweep.to_array().as_numpy_array()
    mask = image >= threshold
    return image, mask

def group_pixels(mask):
    from scipy.ndimage.measurements import label, find_objects

    # Label the indices in the mask
    regions, nregions = label(mask)#, structure)

    # Get the bounding box of each object
    objects = find_objects(regions)

    # Return the list of objects
    return objects

def filter_objects(mask, objects, min_pixels):
    import numpy

    # Select objects with more pixels than min_pixels
    new_obj = []
    for obj in objects:
        ind = numpy.where(mask[obj] != 0)[0]
        if len(ind) >= min_pixels:
            new_obj.append(obj)
        else:
            mask[obj] = 0

    return new_obj

def direction_var(values, weights):
    import numpy
    from scitbx import matrix
    weights = numpy.array(weights)
    valx = numpy.array([x for x, y, z in values])
    valy = numpy.array([y for x, y, z in values])
    valz = numpy.array([z for x, y, z in values])

    # Calculate avergae x, y, z
    avrx = numpy.average(valx, weights=weights)
    avry = numpy.average(valy, weights=weights)
    avrz = numpy.average(valz, weights=weights)

    # Calculate mean direction vector
    s1m = matrix.col((avrx, avry, avrz)).normalize()

    # Calculate angles between vectors
    angles = []
    for s in values:
        angles.append(s1m.angle(s))

    # Calculate variance of angles
    angles = numpy.array(angles)
    var = numpy.dot(weights, (angles)**2)/numpy.sum(weights)
    return var

def centroid(image, mask, detector):
    from numpy import zeros, int32, argmax, where, average
    from scitbx.array_family import flex
    from scitbx import matrix
    import numpy

    var = []
    cent = []
    for obj in objects:
        xs = []
        ys = []
        zs = []
        s1s = []
        weights = []
        bbox = [obj[2].start, obj[2].stop,
                obj[1].start, obj[1].stop,
                obj[0].start, obj[0].stop]

        # Calcyulate beam vector for each point
        for i in range(bbox[4], bbox[5]):
            for s in range(bbox[2], bbox[3]):
                for f in range(bbox[0], bbox[1]):
                    if mask[i, s, f]:
                        s1 = matrix.col(detector.get_pixel_lab_coord((f, s)))
                        s1 = s1.normalize()
                        xs.append(f)
                        ys.append(s)
                        zs.append(i)
                        s1s.append(s1)
                        weights.append(image[i, s ,f])

        v = direction_var(s1s, weights)
        var.append(v)
        avrx = average(xs, weights=weights)
        avry = average(ys, weights=weights)
        avrz = average(zs, weights=weights)
        cent.append((avrx, avry, avrz))

    # Return a list of centroids and variances
    return numpy.array(cent), numpy.array(var)

def calculate_sigma_beam_divergence(var):
    '''Calculate the beam divergence as the sum of centroid variance of the
    intensity weighted diffracted beam directions.'''
    from math import sqrt

    # Calculate the sum of s^2
    sum_variance = reduce(lambda x, y: x + y, var)

    # Return the beam divergence as the sum / num reflections
    return sqrt(sum_variance / len(var))

def find_nearest_neighbour(obs_xyz, reflections):
    from annlib_ext import AnnAdaptor
    from scitbx.array_family import flex
    from math import sqrt

    # Get the predicted coords
    pred_xyz = []
    for r in reflections:
        x = r.image_coord_px[0]
        y = r.image_coord_px[1]
        z = r.frame_number
        pred_xyz.append((x, y, z))

    # Create the KD Tree
    ann = AnnAdaptor(flex.double(pred_xyz).as_1d(), 3)

    ann.query(flex.double(obs_xyz).as_1d())

#    for i in xrange(len(ann.nn)):
#        print "Neighbor of {0}, index {1} distance {2}".format(
#        obs_xyz[i], ann.nn[i], sqrt(ann.distances[i]))

    return ann.nn, ann.distances

def filter_objects_by_distance(ref_index, distance2, max_distance):

    import numpy
    from math import sqrt

    # Only accept objects closer than max distance
    indices = []
    for i in range(len(ref_index)):
        if sqrt(distance2[i]) > max_distance:
            continue
        else:
            indices.append(i)

    return numpy.array(indices)

if __name__ == '__main__':

    import os
    import libtbx.load_env
    from dials.util.nexus import NexusFile
    from glob import glob
    from dxtbx.sweep import SweepFactory
    from math import pi
    from scitbx import matrix

    try:
        dials_regression = libtbx.env.dist_path('dials_regression')
    except KeyError, e:
        print 'FAIL: dials_regression not configured'
        raise

    # The XDS values
    xds_sigma_d = 0.060
    xds_sigma_m = 0.154

    # Read the reflections file
    filename = os.path.join(dials_regression,
        'centroid_test_data', 'reflections.h5')
    handle = NexusFile(filename, 'r')

    # Get the reflection list
    print 'Reading reflections.'
    reflections = handle.get_reflections()
    print 'Read {0} reflections.'.format(len(reflections))

    # Read images
    template = os.path.join(dials_regression,
        'centroid_test_data', 'centroid_*.cbf')
    filenames = glob(template)

    # Load the sweep
    print 'Loading sweep'
    sweep = SweepFactory.sweep(filenames)
    print 'Loaded sweep of {0} images.'.format(len(sweep))

    # Select the strong pixels to use in the divergence calculation
    print 'Select the strong pixels from the images.'
    trusted_range = (0, 20000)
    image, mask = select_strong_pixels(sweep, trusted_range)

    # Putting pixels into groups
    print 'Putting pixels into groups'
    objects = group_pixels(mask)
    print 'Found {0} objects'.format(len(objects))

    print 'Filtering objects'
    min_pixels = 6
    objects = filter_objects(mask, objects, min_pixels)
    print '{0} remaining objects'.format(len(objects))

    print 'Calculating centroid and variance.'
    cent, var = centroid(image, mask, sweep.get_detector())

    print 'Find object nearest predicted reflection'
    ref_index, distance2 = find_nearest_neighbour(cent, reflections)

    print 'Filter objects by distance from nearest reflection.'
    max_distance = 2
    indices = filter_objects_by_distance(ref_index, distance2, max_distance)
    print '{0} remaining objects'.format(len(indices))

    print 'Calculate the e.s.d of the beam divergence'
    sigma_d = calculate_sigma_beam_divergence(var[indices])
    print 'Sigma_d = {0} deg'.format(sigma_d * 180.0 / pi)
    print 'XDS Sigma_d = {0} deg'.format(xds_sigma_d)

    # Calculate rotation angles
    phi = []
    for x, y, z in cent:
        phi.append(sweep.get_scan().get_angle_from_array_index(z))

    # Calculate the zetas
    s0 = matrix.col(sweep.get_beam().get_s0())
    m2 = matrix.col(sweep.get_goniometer().get_rotation_axis())
    zeta = []
    for x, y, z in cent:
        s1 = sweep.get_detector().get_pixel_lab_coord((x, y))
        s1 = matrix.col(s1).normalize() * s0.length()
        e1 = s1.cross(s0).normalize()
        zeta.append(m2.dot(e1))
<|endoftext|>


<commit_before>"""
Created on 27 Sep 2020

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)
"""

import os
import time

from abc import ABC
from subprocess import Popen

# --------------------------------------------------------------------------------------------------------------------

class Network(ABC):
    """
    classdocs
    """
    __TEST_RESOURCE =       'google.com'
    __NETWORK_WAIT_TIME =   5.0                         # seconds

    # ----------------------------------------------------------------------------------------------------------------

    @staticmethod
    def resource_is_available(resource):
        with open(os.devnull, 'wb') as devnull:
            return Popen(['ping', '-c', '1', resource], stdout=devnull, stderr=devnull).wait() == 0

    # ----------------------------------------------------------------------------------------------------------------

    @classmethod
    def is_available(cls):
        return cls.resource_is_available(cls.__TEST_RESOURCE)

    @classmethod
    def wait(cls):
        cls.wait_for_resource(cls.__TEST_RESOURCE)

    @classmethod
    def wait_for_resource(cls, resource):
        while not cls.resource_is_available(resource):
            time.sleep(cls.__NETWORK_WAIT_TIME)
<commit_msg>Added networkMonitor class<commit_after>"""
Created on 27 Sep 2020

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)
"""

import os
import time

from abc import ABC
from multiprocessing import Process
from subprocess import Popen

# --------------------------------------------------------------------------------------------------------------------

class Network(ABC):
    """
    classdocs
    """
    __TEST_RESOURCE =       'google.com'
    __NETWORK_WAIT_TIME =   5.0                         # seconds

    # ----------------------------------------------------------------------------------------------------------------

    @staticmethod
    def resource_is_available(resource):
        with open(os.devnull, 'wb') as devnull:
            return Popen(['ping', '-c', '1', resource], stdout=devnull, stderr=devnull).wait() == 0

    # ----------------------------------------------------------------------------------------------------------------

    @classmethod
    def is_available(cls):
        return cls.resource_is_available(cls.__TEST_RESOURCE)

    @classmethod
    def wait(cls):
        cls.wait_for_resource(cls.__TEST_RESOURCE)

    @classmethod
    def wait_for_resource(cls, resource):
        while not cls.resource_is_available(resource):
            time.sleep(cls.__NETWORK_WAIT_TIME)

# ----------------------------------------------------------------------------------------------------------------

class NetworkMonitor(object):
    """
    classdocs
    """

    # ----------------------------------------------------------------------------------------------------------------

    def __init__(self, interval, network_unavailable_handler):
        self.__interval = interval                                                  # float seconds
        self.__network_unavailable_handler = network_unavailable_handler            # function

        self.__proc = None                                                          # Process

    # ----------------------------------------------------------------------------------------------------------------

    def start(self):
        self.__proc = Process(target=self.run)
        self.__proc.start()

    def run(self):
        Network.wait()

        while True:
            if not Network.is_available():
                self.__network_unavailable_handler()

            time.sleep(self.__interval)

    def join(self):
        self.__proc.join()

    # ----------------------------------------------------------------------------------------------------------------

    def __str__(self, *args, **kwargs):
        return "NetworkMonitor:{interval:%s, network_unavailable_handler:%s, proc:%s}" % \
               (self.__interval, self.__network_unavailable_handler, self.__proc)

<|endoftext|>


<commit_before>"""
Created on 27 Sep 2016

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)
"""

import json
import re

from collections import OrderedDict
from copy import deepcopy

from scs_core.data.json import JSONable

# --------------------------------------------------------------------------------------------------------------------

class PathDict(JSONable):
    """
    classdocs
    """

    # ----------------------------------------------------------------------------------------------------------------

    @staticmethod
    def union(*pairs):
        union = PathDict()

        for name, value in pairs:
            if value is not None:
                union.append(name, value)

        return union

    @staticmethod
    def sub_path_includes_path(sub_path, path):
        sub_path_nodes = [node[0] for node in re.findall(r'([^.:]+)([.:]*)', sub_path)]
        path_nodes = [node[0] for node in re.findall(r'([^.:]+)([.:]*)', path)]

        for i in range(len(sub_path_nodes)):
            try:
                if sub_path_nodes[i] != path_nodes[i]:
                    return False
            except IndexError:
                return False

        return True

    # ----------------------------------------------------------------------------------------------------------------

    @classmethod
    def construct_from_jstr(cls, jstr, sort_paths=()):
        try:
            jdict = json.loads(jstr)
        except ValueError:
            return None

        return cls(dictionary=jdict, sort_paths=sort_paths)

    # ----------------------------------------------------------------------------------------------------------------

    def __init__(self, dictionary=None, sort_paths=()):
        """
        Constructor
        """
        self.__dictionary = OrderedDict() if dictionary is None else dictionary
        self.__sort_paths = sort_paths

    def __len__(self):
        return len(self.__dictionary)

    def __lt__(self, other):
        for sort_path in self.__sort_paths:
            self_value = self.node(sub_path=sort_path)
            other_value = other.node(sub_path=sort_path)

            if self_value < other_value:
                return True

            if self_value > other_value:
                return False

        return False

    # -----------------------------------------------------------------------------------------------------------------
    # source...

    # Tests whether a leaf node path is present...

    def has_path(self, path):
        return path in self.paths()

    # Tests whether a leaf node or internal path is present...

    def has_sub_path(self, sub_path=None):
        if sub_path is None:
            return True

        try:
            self.node(sub_path)
            return True

        except (KeyError, TypeError):
            return False

    # Returns all the leaf node paths...

    def paths(self, sub_path=None):
        node = self if sub_path is None else PathDict(self.node(sub_path))

        return node.__paths(node.__dictionary, sub_path)

    # Returns a leaf node or internal node...

    def node(self, sub_path=None):
        if sub_path is None:
            return self.__dictionary

        try:
            return self.__node(self.__dictionary, re.split(r"[.:]", sub_path))

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # Appends or replaces value at sub_path...

    def append(self, sub_path, value):
        nodes = re.findall('([^.:]+)([.:]*)', sub_path)

        try:
            self.__append(self.__dictionary, nodes, value)

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # Sets an internal node to the given value...

    def set_node(self, sub_path, value):
        try:
            return self.__set_node(self.__dictionary, re.split(r"[.:]", sub_path), value)

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # ----------------------------------------------------------------------------------------------------------------
    # target...

    # Copies from other at sub_path...

    def copy(self, other, sub_path=None):
        if sub_path is None:
            self.__dictionary = deepcopy(other.__dictionary)
            return

        nodes = re.findall('([^.:]+)([.:]*)', sub_path)

        try:
            self.__append(self.__dictionary, nodes, other.node(sub_path))

        except KeyError:
            raise KeyError(sub_path)

    # ----------------------------------------------------------------------------------------------------------------

    def __paths(self, container, prefix=None):
        # dict...
        if isinstance(container, dict):
            separator = '.'
            keys = container.keys()

        # list...
        elif isinstance(container, list):
            separator = ':'
            keys = range(len(container))

        # scalar...
        else:
            return [prefix]

        prefix = prefix + separator if prefix else ''

        paths = []

        for key in keys:
            path = prefix + str(key)
            value = container[key]

            # dict...
            if isinstance(value, dict):
                paths.extend(self.__paths(value, path))

            # list...
            elif isinstance(value, list):
                for i in range(len(value)):
                    if isinstance(value[i], (list, dict)):
                        paths.extend(self.__paths(value[i], path + ':' + str(i)))
                    else:
                        paths.append(path + ':' + str(i))

            # scalar...
            else:
                paths.append(path)

        return paths

    def __node(self, container, nodes):
        # key...
        if isinstance(container, list):
            try:
                key = int(nodes[0])

            except ValueError:
                raise KeyError(nodes[0])            # a non-integer key on an array is a KeyError

        else:
            key = nodes[0]

        # value...
        value = container[key]                      # raises TypeError if container is None

        # leaf...
        if len(nodes) == 1:
            return value

        # deep...
        return self.__node(value, nodes[1:])

    def __set_node(self, container, nodes, value):
        # key...
        if isinstance(container, list):
            try:
                key = int(nodes[0])

            except ValueError:
                raise KeyError(nodes[0])            # a non-integer key on an array is a KeyError

        else:
            key = nodes[0]

        # value...
        container[key] = value

        # leaf...
        if len(nodes) == 1:
            return value

        # deep...
        return self.__set_node(value, nodes[1:], value)

    def __append(self, container, nodes, value):
        # key...
        key = nodes[0][0]
        separator = nodes[0][1]

        # dict...
        if isinstance(container, dict):
            if key not in container:
                container[key] = [] if separator == ':' else OrderedDict()

        # list...
        elif isinstance(container, list):
            try:
                key = int(key)
            except ValueError:
                raise KeyError(key)

            while key >= len(container):
                container.append(None)

            if container[key] is None:
                container[key] = [] if separator == ':' else OrderedDict()

        # leaf...
        if len(nodes) == 1:
            container[key] = deepcopy(value)
            return

        # deep...
        self.__append(container[key], nodes[1:], value)

    # ----------------------------------------------------------------------------------------------------------------

    @property
    def dictionary(self):
        return self.node()

    @property
    def sort_paths(self):
        return self.__sort_paths

    # ----------------------------------------------------------------------------------------------------------------

    def as_json(self):
        return self.__dictionary

    # ----------------------------------------------------------------------------------------------------------------

    def __str__(self, *args, **kwargs):
        return "PathDict:{sort_paths:%s, dictionary:%s}" % (self.sort_paths, self.node())
<commit_msg>Cleaning<commit_after>"""
Created on 27 Sep 2016

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)
"""

import json
import re

from collections import OrderedDict
from copy import deepcopy

from scs_core.data.json import JSONable

# --------------------------------------------------------------------------------------------------------------------

class PathDict(JSONable):
    """
    classdocs
    """

    # ----------------------------------------------------------------------------------------------------------------

    @staticmethod
    def union(*pairs):
        union = PathDict()

        for name, value in pairs:
            if value is not None:
                union.append(name, value)

        return union

    @staticmethod
    def sub_path_includes_path(sub_path, path):
        sub_path_nodes = [node[0] for node in re.findall(r'([^.:]+)([.:]*)', sub_path)]
        path_nodes = [node[0] for node in re.findall(r'([^.:]+)([.:]*)', path)]

        for i in range(len(sub_path_nodes)):
            try:
                if sub_path_nodes[i] != path_nodes[i]:
                    return False
            except IndexError:
                return False

        return True

    # ----------------------------------------------------------------------------------------------------------------

    @classmethod
    def construct_from_jstr(cls, jstr, sort_paths=()):
        try:
            jdict = json.loads(jstr)
        except ValueError:
            return None

        return cls(dictionary=jdict, sort_paths=sort_paths)

    # ----------------------------------------------------------------------------------------------------------------

    def __init__(self, dictionary=None, sort_paths=()):
        """
        Constructor
        """
        self.__dictionary = OrderedDict() if dictionary is None else dictionary
        self.__sort_paths = sort_paths

    def __len__(self):
        return len(self.__dictionary)

    def __lt__(self, other):
        for sort_path in self.__sort_paths:
            value = self.node(sub_path=sort_path)
            other_value = other.node(sub_path=sort_path)

            if value < other_value:
                return True

            if value > other_value:
                return False

        return False

    # -----------------------------------------------------------------------------------------------------------------
    # source...

    # Tests whether a leaf node path is present...

    def has_path(self, path):
        return path in self.paths()

    # Tests whether a leaf node or internal path is present...

    def has_sub_path(self, sub_path=None):
        if sub_path is None:
            return True

        try:
            self.node(sub_path)
            return True

        except (KeyError, TypeError):
            return False

    # Returns all the leaf node paths...

    def paths(self, sub_path=None):
        node = self if sub_path is None else PathDict(self.node(sub_path))

        return node.__paths(node.__dictionary, sub_path)

    # Returns a leaf node or internal node...

    def node(self, sub_path=None):
        if sub_path is None:
            return self.__dictionary

        try:
            return self.__node(self.__dictionary, re.split(r"[.:]", sub_path))

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # Appends or replaces value at sub_path...

    def append(self, sub_path, value):
        nodes = re.findall('([^.:]+)([.:]*)', sub_path)

        try:
            self.__append(self.__dictionary, nodes, value)

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # Sets an internal node to the given value...

    def set_node(self, sub_path, value):
        try:
            return self.__set_node(self.__dictionary, re.split(r"[.:]", sub_path), value)

        except (KeyError, TypeError):
            raise KeyError(sub_path)

    # ----------------------------------------------------------------------------------------------------------------
    # target...

    # Copies from other at sub_path...

    def copy(self, other, sub_path=None):
        if sub_path is None:
            self.__dictionary = deepcopy(other.__dictionary)
            return

        nodes = re.findall('([^.:]+)([.:]*)', sub_path)

        try:
            self.__append(self.__dictionary, nodes, other.node(sub_path))

        except KeyError:
            raise KeyError(sub_path)

    # ----------------------------------------------------------------------------------------------------------------

    def __paths(self, container, prefix=None):
        # dict...
        if isinstance(container, dict):
            separator = '.'
            keys = container.keys()

        # list...
        elif isinstance(container, list):
            separator = ':'
            keys = range(len(container))

        # scalar...
        else:
            return [prefix]

        prefix = prefix + separator if prefix else ''

        paths = []

        for key in keys:
            path = prefix + str(key)
            value = container[key]

            # dict...
            if isinstance(value, dict):
                paths.extend(self.__paths(value, path))

            # list...
            elif isinstance(value, list):
                for i in range(len(value)):
                    if isinstance(value[i], (list, dict)):
                        paths.extend(self.__paths(value[i], path + ':' + str(i)))
                    else:
                        paths.append(path + ':' + str(i))

            # scalar...
            else:
                paths.append(path)

        return paths

    def __node(self, container, nodes):
        # key...
        if isinstance(container, list):
            try:
                key = int(nodes[0])

            except ValueError:
                raise KeyError(nodes[0])            # a non-integer key on an array is a KeyError

        else:
            key = nodes[0]

        # value...
        value = container[key]                      # raises TypeError if container is None

        # leaf...
        if len(nodes) == 1:
            return value

        # deep...
        return self.__node(value, nodes[1:])

    def __set_node(self, container, nodes, value):
        # key...
        if isinstance(container, list):
            try:
                key = int(nodes[0])

            except ValueError:
                raise KeyError(nodes[0])            # a non-integer key on an array is a KeyError

        else:
            key = nodes[0]

        # value...
        container[key] = value

        # leaf...
        if len(nodes) == 1:
            return value

        # deep...
        return self.__set_node(value, nodes[1:], value)

    def __append(self, container, nodes, value):
        # key...
        key = nodes[0][0]
        separator = nodes[0][1]

        # dict...
        if isinstance(container, dict):
            if key not in container:
                container[key] = [] if separator == ':' else OrderedDict()

        # list...
        elif isinstance(container, list):
            try:
                key = int(key)
            except ValueError:
                raise KeyError(key)

            while key >= len(container):
                container.append(None)

            if container[key] is None:
                container[key] = [] if separator == ':' else OrderedDict()

        # leaf...
        if len(nodes) == 1:
            container[key] = deepcopy(value)
            return

        # deep...
        self.__append(container[key], nodes[1:], value)

    # ----------------------------------------------------------------------------------------------------------------

    @property
    def dictionary(self):
        return self.node()

    @property
    def sort_paths(self):
        return self.__sort_paths

    # ----------------------------------------------------------------------------------------------------------------

    def as_json(self):
        return self.__dictionary

    # ----------------------------------------------------------------------------------------------------------------

    def __str__(self, *args, **kwargs):
        return "PathDict:{sort_paths:%s, dictionary:%s}" % (self.sort_paths, self.node())
<|endoftext|>