<commit_before>import collections
import contextlib
import copy
import json
import warnings

import numpy
import six

from chainer import backend
from chainer.backends import cuda
from chainer import configuration
from chainer import serializer as serializer_module
from chainer import variable

def _copy_variable(value):
    if isinstance(value, variable.Variable):
        return copy.copy(value)
    return value

class Reporter(object):

    """Object to which observed values are reported.

    Reporter is used to collect values that users want to watch. The reporter
    object holds a mapping from value names to the actually observed values.
    We call this mapping `observations`.

    When a value is passed to the reporter, an object called `observer` can be
    optionally attached. In this case, the name of the observer is added as the
    prefix of the value name. The observer name should be registered
    beforehand.

    See the following example::

       >>> from chainer import Reporter, report, report_scope
       >>>
       >>> reporter = Reporter()
       >>> observer = object()  # it can be an arbitrary (reference) object
       >>> reporter.add_observer('my_observer', observer)
       >>> observation = {}
       >>> with reporter.scope(observation):
       ...     reporter.report({'x': 1}, observer)
       ...
       >>> observation
       {'my_observer/x': 1}

    There are also a global API to add values::

       >>> observation = {}
       >>> with report_scope(observation):
       ...     report({'x': 1}, observer)
       ...
       >>> observation
       {'my_observer/x': 1}

    The most important application of Reporter is to report observed values
    from each link or chain in the training and validation procedures.
    :class:`~chainer.training.Trainer` and some extensions prepare their own
    Reporter object with the hierarchy of the target link registered as
    observers. We can use :func:`report` function inside any links and chains
    to report the observed values (e.g., training loss, accuracy, activation
    statistics, etc.).

    Attributes:
        observation: Dictionary of observed values.

    """

    def __init__(self):
        self._observer_names = {}
        self.observation = {}

    def __enter__(self):
        """Makes this reporter object current."""
        _reporters.append(self)

    def __exit__(self, exc_type, exc_value, traceback):
        """Recovers the previous reporter object to the current."""
        _reporters.pop()

    @contextlib.contextmanager
    def scope(self, observation):
        """Creates a scope to report observed values to ``observation``.

        This is a context manager to be passed to ``with`` statements. In this
        scope, the observation dictionary is changed to the given one.

        It also makes this reporter object current.

        Args:
            observation (dict): Observation dictionary. All observations
                reported inside of the ``with`` statement are written to this
                dictionary.

        """
        old = self.observation
        self.observation = observation
        self.__enter__()
        yield
        self.__exit__(None, None, None)
        self.observation = old

    def add_observer(self, name, observer):
        """Registers an observer of values.

        Observer defines a scope of names for observed values. Values observed
        with the observer are registered with names prefixed by the observer
        name.

        Args:
            name (str): Name of the observer.
            observer: The observer object. Note that the reporter distinguishes
                the observers by their object ids (i.e., ``id(owner)``), rather
                than the object equality.

        """
        self._observer_names[id(observer)] = name

    def add_observers(self, prefix, observers):
        """Registers multiple observers at once.

        This is a convenient method to register multiple objects at once.

        Args:
            prefix (str): Prefix of each name of observers.
            observers: Iterator of name and observer pairs.

        """
        for name, observer in observers:
            self._observer_names[id(observer)] = prefix + name

    def report(self, values, observer=None):
        """Reports observed values.

        The values are written with the key, prefixed by the name of the
        observer object if given.

        .. note::
           As of v2.0.0, if a value is of type :class:`~chainer.Variable`, the
           variable is copied without preserving the computational graph and
           the new variable object purged from the graph is stored to the
           observer. This behavior can be changed by setting
           ``chainer.config.keep_graph_on_report`` to ``True``.

        Args:
            values (dict): Dictionary of observed values.
            observer: Observer object. Its object ID is used to retrieve the
                observer name, which is used as the prefix of the registration
                name of the observed value.

        """
        if not configuration.config.keep_graph_on_report:
            values = {k: _copy_variable(v) for k, v in six.iteritems(values)}

        if observer is not None:
            observer_id = id(observer)
            if observer_id not in self._observer_names:
                raise KeyError(
                    'Given observer is not registered to the reporter.')
            observer_name = self._observer_names[observer_id]
            for key, value in six.iteritems(values):
                name = '%s/%s' % (observer_name, key)
                self.observation[name] = value
        else:
            self.observation.update(values)

_reporters = []

def get_current_reporter():
    """Returns the current reporter object."""
    return _reporters[-1]

def report(values, observer=None):
    """Reports observed values with the current reporter object.

    Any reporter object can be set current by the ``with`` statement. This
    function calls the :meth:`Report.report` method of the current reporter.
    If no reporter object is current, this function does nothing.

    .. admonition:: Example

       The most typical example is a use within links and chains. Suppose that
       a link is registered to the current reporter as an observer (for
       example, the target link of the optimizer is automatically registered to
       the reporter of the :class:`~chainer.training.Trainer`). We can report
       some values from the link as follows::

          class MyRegressor(chainer.Chain):
              def __init__(self, predictor):
                  super(MyRegressor, self).__init__(predictor=predictor)

              def __call__(self, x, y):
                  # This chain just computes the mean absolute and squared
                  # errors between the prediction and y.
                  pred = self.predictor(x)
                  abs_error = F.sum(F.abs(pred - y)) / len(x)
                  loss = F.mean_squared_error(pred, y)

                  # Report the mean absolute and squared errors.
                  report({'abs_error': abs_error, 'squared_error': loss}, self)

                  return loss

       If the link is named ``'main'`` in the hierarchy (which is the default
       name of the target link in the
       :class:`~chainer.training.updaters.StandardUpdater`),
       these reported values are
       named ``'main/abs_error'`` and ``'main/squared_error'``. If these values
       are reported inside the :class:`~chainer.training.extension.Evaluator`
       extension, ``'validation/'`` is added at the head of the link name, thus
       the item names are changed to ``'validation/main/abs_error'`` and
       ``'validation/main/squared_error'`` (``'validation'`` is the default
       name of the Evaluator extension).

    Args:
        values (dict): Dictionary of observed values.
        observer: Observer object. Its object ID is used to retrieve the
            observer name, which is used as the prefix of the registration name
            of the observed value.

    """
    if _reporters:
        current = _reporters[-1]
        current.report(values, observer)

@contextlib.contextmanager
def report_scope(observation):
    """Returns a report scope with the current reporter.

    This is equivalent to ``get_current_reporter().scope(observation)``,
    except that it does not make the reporter current redundantly.

    """
    current = _reporters[-1]
    old = current.observation
    current.observation = observation
    yield
    current.observation = old

def _get_device(x):
    if numpy.isscalar(x):
        return cuda.DummyDevice
    else:
        return cuda.get_device_from_array(x)

class Summary(object):

    """Online summarization of a sequence of scalars.

    Summary computes the statistics of given scalars online.

    """

    def __init__(self):
        self._x = 0
        self._x2 = 0
        self._n = 0

    def add(self, value, weight=1):
        """Adds a scalar value.

        Args:
            value: Scalar value to accumulate. It is either a NumPy scalar or
                a zero-dimensional array (on CPU or GPU).
            weight: An optional weight for the value. It is a NumPy scalar or
                a zero-dimensional array (on CPU or GPU).
                Default is 1 (integer).

        """
        with _get_device(value):
            self._x += weight * value
            self._x2 += weight * value * value
            self._n += weight

    def compute_mean(self):
        """Computes the mean."""
        x, n = self._x, self._n
        with _get_device(x):
            return x / n

    def make_statistics(self):
        """Computes and returns the mean and standard deviation values.

        Returns:
            tuple: Mean and standard deviation values.

        """
        x, n = self._x, self._n
        xp = backend.get_array_module(x)
        with _get_device(x):
            mean = x / n
            var = self._x2 / n - mean * mean
            std = xp.sqrt(var)
            return mean, std

    def serialize(self, serializer):
        try:
            self._x = serializer('_x', self._x)
            self._x2 = serializer('_x2', self._x2)
            self._n = serializer('_n', self._n)
        except KeyError:
            warnings.warn('The previous statistics are not saved.')

class DictSummary(object):

    """Online summarization of a sequence of dictionaries.

    ``DictSummary`` computes the statistics of a given set of scalars online.
    It only computes the statistics for scalar values and variables of scalar
    values in the dictionaries.

    """

    def __init__(self):
        self._summaries = collections.defaultdict(Summary)

    def add(self, d):
        """Adds a dictionary of scalars.

        Args:
            d (dict): Dictionary of scalars to accumulate. Only elements of
               scalars, zero-dimensional arrays, and variables of
               zero-dimensional arrays are accumulated. When the value
               is a tuple, the second element is interpreted as a weight.

        """
        summaries = self._summaries
        for k, v in six.iteritems(d):
            w = 1
            if isinstance(v, tuple):
                w = v[1]
                v = v[0]
                if isinstance(w, variable.Variable):
                    w = w.array
                if not numpy.isscalar(w) and not getattr(w, 'ndim', -1) == 0:
                    raise ValueError(
                        'Given weight to {} was not scalar.'.format(k))
            if isinstance(v, variable.Variable):
                v = v.array
            if numpy.isscalar(v) or getattr(v, 'ndim', -1) == 0:
                summaries[k].add(v, weight=w)

    def compute_mean(self):
        """Creates a dictionary of mean values.

        It returns a single dictionary that holds a mean value for each entry
        added to the summary.

        Returns:
            dict: Dictionary of mean values.

        """
        return {name: summary.compute_mean()
                for name, summary in six.iteritems(self._summaries)}

    def make_statistics(self):
        """Creates a dictionary of statistics.

        It returns a single dictionary that holds mean and standard deviation
        values for every entry added to the summary. For an entry of name
        ``'key'``, these values are added to the dictionary by names ``'key'``
        and ``'key.std'``, respectively.

        Returns:
            dict: Dictionary of statistics of all entries.

        """
        stats = {}
        for name, summary in six.iteritems(self._summaries):
            mean, std = summary.make_statistics()
            stats[name] = mean
            stats[name + '.std'] = std

        return stats

    def serialize(self, serializer):
        if isinstance(serializer, serializer_module.Serializer):
            names = list(self._summaries.keys())
            serializer('_names', json.dumps(names))
            for index, name in enumerate(names):
                self._summaries[name].serialize(
                    serializer['_summaries'][str(index)])
        else:
            self._summaries.clear()
            try:
                names = json.loads(serializer('_names', ''))
            except KeyError:
                warnings.warn('The names of statistics are not saved.')
                return
            for index, name in enumerate(names):
                self._summaries[name].serialize(
                    serializer['_summaries'][str(index)])
<commit_msg>Fix typos in `chainer.report`<commit_after>import collections
import contextlib
import copy
import json
import warnings

import numpy
import six

from chainer import backend
from chainer.backends import cuda
from chainer import configuration
from chainer import serializer as serializer_module
from chainer import variable

def _copy_variable(value):
    if isinstance(value, variable.Variable):
        return copy.copy(value)
    return value

class Reporter(object):

    """Object to which observed values are reported.

    Reporter is used to collect values that users want to watch. The reporter
    object holds a mapping from value names to the actually observed values.
    We call this mapping `observations`.

    When a value is passed to the reporter, an object called `observer` can be
    optionally attached. In this case, the name of the observer is added as the
    prefix of the value name. The observer name should be registered
    beforehand.

    See the following example::

       >>> from chainer import Reporter, report, report_scope
       >>>
       >>> reporter = Reporter()
       >>> observer = object()  # it can be an arbitrary (reference) object
       >>> reporter.add_observer('my_observer', observer)
       >>> observation = {}
       >>> with reporter.scope(observation):
       ...     reporter.report({'x': 1}, observer)
       ...
       >>> observation
       {'my_observer/x': 1}

    There are also a global API to add values::

       >>> observation = {}
       >>> with report_scope(observation):
       ...     report({'x': 1}, observer)
       ...
       >>> observation
       {'my_observer/x': 1}

    The most important application of Reporter is to report observed values
    from each link or chain in the training and validation procedures.
    :class:`~chainer.training.Trainer` and some extensions prepare their own
    Reporter object with the hierarchy of the target link registered as
    observers. We can use :func:`report` function inside any links and chains
    to report the observed values (e.g., training loss, accuracy, activation
    statistics, etc.).

    Attributes:
        observation: Dictionary of observed values.

    """

    def __init__(self):
        self._observer_names = {}
        self.observation = {}

    def __enter__(self):
        """Makes this reporter object current."""
        _reporters.append(self)

    def __exit__(self, exc_type, exc_value, traceback):
        """Recovers the previous reporter object to the current."""
        _reporters.pop()

    @contextlib.contextmanager
    def scope(self, observation):
        """Creates a scope to report observed values to ``observation``.

        This is a context manager to be passed to ``with`` statements. In this
        scope, the observation dictionary is changed to the given one.

        It also makes this reporter object current.

        Args:
            observation (dict): Observation dictionary. All observations
                reported inside of the ``with`` statement are written to this
                dictionary.

        """
        old = self.observation
        self.observation = observation
        self.__enter__()
        yield
        self.__exit__(None, None, None)
        self.observation = old

    def add_observer(self, name, observer):
        """Registers an observer of values.

        Observer defines a scope of names for observed values. Values observed
        with the observer are registered with names prefixed by the observer
        name.

        Args:
            name (str): Name of the observer.
            observer: The observer object. Note that the reporter distinguishes
                the observers by their object ids (i.e., ``id(owner)``), rather
                than the object equality.

        """
        self._observer_names[id(observer)] = name

    def add_observers(self, prefix, observers):
        """Registers multiple observers at once.

        This is a convenient method to register multiple objects at once.

        Args:
            prefix (str): Prefix of each name of observers.
            observers: Iterator of name and observer pairs.

        """
        for name, observer in observers:
            self._observer_names[id(observer)] = prefix + name

    def report(self, values, observer=None):
        """Reports observed values.

        The values are written with the key, prefixed by the name of the
        observer object if given.

        .. note::
           As of v2.0.0, if a value is of type :class:`~chainer.Variable`, the
           variable is copied without preserving the computational graph and
           the new variable object purged from the graph is stored to the
           observer. This behavior can be changed by setting
           ``chainer.config.keep_graph_on_report`` to ``True``.

        Args:
            values (dict): Dictionary of observed values.
            observer: Observer object. Its object ID is used to retrieve the
                observer name, which is used as the prefix of the registration
                name of the observed value.

        """
        if not configuration.config.keep_graph_on_report:
            values = {k: _copy_variable(v) for k, v in six.iteritems(values)}

        if observer is not None:
            observer_id = id(observer)
            if observer_id not in self._observer_names:
                raise KeyError(
                    'Given observer is not registered to the reporter.')
            observer_name = self._observer_names[observer_id]
            for key, value in six.iteritems(values):
                name = '%s/%s' % (observer_name, key)
                self.observation[name] = value
        else:
            self.observation.update(values)

_reporters = []

def get_current_reporter():
    """Returns the current reporter object."""
    return _reporters[-1]

def report(values, observer=None):
    """Reports observed values with the current reporter object.

    Any reporter object can be set current by the ``with`` statement. This
    function calls the :meth:`Reporter.report` method of the current reporter.
    If no reporter object is current, this function does nothing.

    .. admonition:: Example

       The most typical example is a use within links and chains. Suppose that
       a link is registered to the current reporter as an observer (for
       example, the target link of the optimizer is automatically registered to
       the reporter of the :class:`~chainer.training.Trainer`). We can report
       some values from the link as follows::

          class MyRegressor(chainer.Chain):
              def __init__(self, predictor):
                  super(MyRegressor, self).__init__(predictor=predictor)

              def __call__(self, x, y):
                  # This chain just computes the mean absolute and squared
                  # errors between the prediction and y.
                  pred = self.predictor(x)
                  abs_error = F.sum(abs(pred - y)) / len(x)
                  loss = F.mean_squared_error(pred, y)

                  # Report the mean absolute and squared errors.
                  chainer.report({
                      'abs_error': abs_error,
                      'squared_error': loss,
                  }, self)

                  return loss

       If the link is named ``'main'`` in the hierarchy (which is the default
       name of the target link in the
       :class:`~chainer.training.updaters.StandardUpdater`),
       these reported values are
       named ``'main/abs_error'`` and ``'main/squared_error'``. If these values
       are reported inside the :class:`~chainer.training.extensions.Evaluator`
       extension, ``'validation/'`` is added at the head of the link name, thus
       the item names are changed to ``'validation/main/abs_error'`` and
       ``'validation/main/squared_error'`` (``'validation'`` is the default
       name of the Evaluator extension).

    Args:
        values (dict): Dictionary of observed values.
        observer: Observer object. Its object ID is used to retrieve the
            observer name, which is used as the prefix of the registration name
            of the observed value.

    """
    if _reporters:
        current = _reporters[-1]
        current.report(values, observer)

@contextlib.contextmanager
def report_scope(observation):
    """Returns a report scope with the current reporter.

    This is equivalent to ``get_current_reporter().scope(observation)``,
    except that it does not make the reporter current redundantly.

    """
    current = _reporters[-1]
    old = current.observation
    current.observation = observation
    yield
    current.observation = old

def _get_device(x):
    if numpy.isscalar(x):
        return cuda.DummyDevice
    else:
        return cuda.get_device_from_array(x)

class Summary(object):

    """Online summarization of a sequence of scalars.

    Summary computes the statistics of given scalars online.

    """

    def __init__(self):
        self._x = 0
        self._x2 = 0
        self._n = 0

    def add(self, value, weight=1):
        """Adds a scalar value.

        Args:
            value: Scalar value to accumulate. It is either a NumPy scalar or
                a zero-dimensional array (on CPU or GPU).
            weight: An optional weight for the value. It is a NumPy scalar or
                a zero-dimensional array (on CPU or GPU).
                Default is 1 (integer).

        """
        with _get_device(value):
            self._x += weight * value
            self._x2 += weight * value * value
            self._n += weight

    def compute_mean(self):
        """Computes the mean."""
        x, n = self._x, self._n
        with _get_device(x):
            return x / n

    def make_statistics(self):
        """Computes and returns the mean and standard deviation values.

        Returns:
            tuple: Mean and standard deviation values.

        """
        x, n = self._x, self._n
        xp = backend.get_array_module(x)
        with _get_device(x):
            mean = x / n
            var = self._x2 / n - mean * mean
            std = xp.sqrt(var)
            return mean, std

    def serialize(self, serializer):
        try:
            self._x = serializer('_x', self._x)
            self._x2 = serializer('_x2', self._x2)
            self._n = serializer('_n', self._n)
        except KeyError:
            warnings.warn('The previous statistics are not saved.')

class DictSummary(object):

    """Online summarization of a sequence of dictionaries.

    ``DictSummary`` computes the statistics of a given set of scalars online.
    It only computes the statistics for scalar values and variables of scalar
    values in the dictionaries.

    """

    def __init__(self):
        self._summaries = collections.defaultdict(Summary)

    def add(self, d):
        """Adds a dictionary of scalars.

        Args:
            d (dict): Dictionary of scalars to accumulate. Only elements of
               scalars, zero-dimensional arrays, and variables of
               zero-dimensional arrays are accumulated. When the value
               is a tuple, the second element is interpreted as a weight.

        """
        summaries = self._summaries
        for k, v in six.iteritems(d):
            w = 1
            if isinstance(v, tuple):
                w = v[1]
                v = v[0]
                if isinstance(w, variable.Variable):
                    w = w.array
                if not numpy.isscalar(w) and not getattr(w, 'ndim', -1) == 0:
                    raise ValueError(
                        'Given weight to {} was not scalar.'.format(k))
            if isinstance(v, variable.Variable):
                v = v.array
            if numpy.isscalar(v) or getattr(v, 'ndim', -1) == 0:
                summaries[k].add(v, weight=w)

    def compute_mean(self):
        """Creates a dictionary of mean values.

        It returns a single dictionary that holds a mean value for each entry
        added to the summary.

        Returns:
            dict: Dictionary of mean values.

        """
        return {name: summary.compute_mean()
                for name, summary in six.iteritems(self._summaries)}

    def make_statistics(self):
        """Creates a dictionary of statistics.

        It returns a single dictionary that holds mean and standard deviation
        values for every entry added to the summary. For an entry of name
        ``'key'``, these values are added to the dictionary by names ``'key'``
        and ``'key.std'``, respectively.

        Returns:
            dict: Dictionary of statistics of all entries.

        """
        stats = {}
        for name, summary in six.iteritems(self._summaries):
            mean, std = summary.make_statistics()
            stats[name] = mean
            stats[name + '.std'] = std

        return stats

    def serialize(self, serializer):
        if isinstance(serializer, serializer_module.Serializer):
            names = list(self._summaries.keys())
            serializer('_names', json.dumps(names))
            for index, name in enumerate(names):
                self._summaries[name].serialize(
                    serializer['_summaries'][str(index)])
        else:
            self._summaries.clear()
            try:
                names = json.loads(serializer('_names', ''))
            except KeyError:
                warnings.warn('The names of statistics are not saved.')
                return
            for index, name in enumerate(names):
                self._summaries[name].serialize(
                    serializer['_summaries'][str(index)])
<|endoftext|>


<commit_before># Author: Pearu Peterson
# Created: March 2008

from ..utils import MATRIX, MATRIX_DICT, MATRIX_DICT_T
from ..arithmetic.numbers import div
from .algebra import MatrixDict, Matrix

from ..core import init_module
init_module.import_lowlevel_operations()

def MATRIX_DICT_gauss_jordan_elimination(self, overwrite=False):
    """ Perform Gauss-Jordan elimination of a m x n matrix A.

    Outputs::

      B - m x n matrix where the lhs part is unit matrix
    """
    head, data = self.pair
    m, n = head.shape
    k = min(m,n)
    if overwrite and self.is_writable:
        udata = data
    else:
        udata = dict(data)
    if head.is_transpose:
        B = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        gauss_jordan_elimination_MATRIX_T(m, n, udata)
    elif head.is_diagonal:
        raise NotImplementedError(`head, head.is_diagonal`)
    else:
        B = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        gauss_jordan_elimination_MATRIX(m, n, udata)
    return B

def MATRIX_DICT_lu(self, overwrite=False):
    """ Perform LU factorization of a m x n matrix A.

    Parameters::

      overwrite=False - if True then discard the content of matrix A

    Outputs::

      P, L, U - LU decomposition matrices of A.

    Definitions::

      P - m x m permuation matrix
      L - m x k lower triangular or trapezoidal matrix with unit-diagonal
      U - k x n upper triangular or trapezoidal matrix
      k = min(m,n)

      A = P * L * U
    """
    head, data = self.pair
    m, n = head.shape
    k = min(m, n)
    ldata = {}
    if overwrite and self.is_writable:
        udata = data
    else:
        udata = dict(data)
    L = MatrixDict(MATRIX(m, k, MATRIX_DICT), ldata)
    if head.is_transpose:
        U = MatrixDict(MATRIX(k, n, MATRIX_DICT_T), udata)
        pivot_table = lu_MATRIX_T(m, n, k, ldata, udata)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        U = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        pivot_table = lu_MATRIX(m, n, k, ldata, udata)
    P = Matrix(pivot_table, permutation=True).T
    return P, L, U

def gauss_jordan_elimination_MATRIX(m, n, data):
    print 'gauss_jordan_elimination_MATRIX'
    data_get = data.get
    data_has = data.has_key
    rows = get_rc_map(data)
    for i in xrange(m):
        if 0:
            ncols, j = n, None
            for j1 in range(i,m):
                if data_has((j1,i)):
                    l = len(rows[j1])
                    if l <= ncols:
                        ncols, j = l, j1
            if j is None:
                continue
            a_ii = data[j,i]
            if i!=j:
                swap_rows_MATRIX(data, i, j)
                swap_rc_map(rows, i, j)
        else:
            a_ii = data_get((i,i))
            if a_ii is None:
                for j in xrange(i+1, m):
                    a_ii = data_get((j,i))
                    if a_ii is not None:
                        break
                if a_ii is None:
                    continue
                swap_rows_MATRIX(data, i, j)
                swap_rc_map(rows, i, j)
        irow = rows[i]
        for j in range(m):
            if j==i:
                continue
            u_ji = data_get((j,i))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = rows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = data[i,p] * c
                jp = j,p
                b = data_get(jp)
                if b is None:
                    data[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if u_ip_c==b:
                        del data[jp]
                        jrow_remove(p)
                    else:
                        data[jp] = b - u_ip_c
            if not jrow:
                del rows[j]
        data[i,i] = 1

        for p in irow:
            if p <= i: continue
            ip = i,p
            u_ip = data_get(ip)
            if u_ip is None:
                continue
            data[ip] = div(u_ip, a_ii)

def gauss_jordan_elimination_MATRIX_T(m, n, data):
    data_get = data.get
    rows = get_rc_map_T(data)
    for i in xrange(m):
        a_ii = data_get((i,i))
        if a_ii is None:
            for j in xrange(i+1, m):
                a_ii = data_get((i,j))
                if a_ii is not None:
                    break
            if a_ii is None:
                continue
            swap_rows_MATRIX_T(data, i, j)
            swap_rc_map(rows, i, j)
        irow = rows[i]
        for j in range(m):
            if j==i:
                continue
            u_ji = data_get((i,j))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = rows[j]
            for p in range(i,n):
                u_ip = data_get((p,i))
                if u_ip is None:
                    continue
                jp = p,j
                b = data_get(jp)
                if b is None:
                    data[jp] = -u_ip*c
                    jrow.add(p)
                else:
                    b -= u_ip*c
                    if b:
                        data[jp] = b
                    else:
                        del data[jp]
                        jrow.remove(p)
            if not jrow:
                del rows[j]
        data[i,i] = 1
        for p in range(i+1, n):
            ip = p,i
            u_ip = data_get(ip)
            if u_ip is None:
                continue
            data[ip] = div(u_ip, a_ii)

def get_rc_map(data):
    rows = {}
    for i, j in data:
        s = rows.get (i)
        if s is None:
            s = rows[i] = set ()
        s.add(j)
    return rows

def get_rc_map_T(data):
    rows = {}
    for j, i in data:
        s = rows.get (i)
        if s is None:
            s = rows[i] = set ()
        s.add(j)
    return rows

def swap_rc_map(rows, i, j):
    ri = rows.get(i)
    rj = rows.get(j)
    if rj is not None: del rows[j]
    if ri is not None: del rows[i]
    if rj is not None:
        rows[i] = rj
    if ri is not None:
        rows[j] = ri

def lu_MATRIX(m, n, k, ldata, udata):
    if m>n:
        for i in xrange(n,m):
            udata[(i,i)] = 1
    pivot_table = range(m)
    udata_get = udata.get
    udata_has = udata.has_key

    urows = get_rc_map(udata)
    for i in xrange(m-1):
        ncols, j = n, None
        for j1 in range(i,m):
            if udata_has((j1,i)):
                l = len(urows[j1])
                if l <= ncols:
                    ncols, j = l, j1
        if j is None:
            continue
        a_ii = udata[j,i]
        if i!=j:
            pivot_table[i], pivot_table[j] = pivot_table[j], pivot_table[i]
            swap_rows_MATRIX(udata, i, j)
            swap_rows_MATRIX(ldata, i, j)
            swap_rc_map(urows, i, j)
        irow = urows[i]
        for j in range (i+1,m):
            u_ji = udata_get((j,i))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = urows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = udata[i,p] * c
                jp = j,p
                b = udata_get(jp)
                if b is None:
                    udata[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if b==u_ip_c:
                        del udata[jp]
                        jrow_remove(p)
                    else:
                        udata[jp] = b - u_ip_c
            if not jrow:
                del urows[j]
            if i<k and j<m:
                ldata[j, i] = c
    for i in xrange(min(m, k)):
        ldata[i,i] = 1
    if m>n:
        crop_MATRIX(k, n, udata)
    return pivot_table

def lu_MATRIX_T(m, n, k, ldata, udata):
    if m>n:
        for i in xrange(n,m):
            udata[(i,i)] = 1
    pivot_table = range(m)
    udata_get = udata.get
    udata_has = udata.has_key
    urows = get_rc_map_T(udata)
    for i in xrange(m-1):
        ncols, j = n, None
        for j1 in range(i, m):
            if udata_has((i,j1)):
                l = len(urows[j1])
                if l <= ncols:
                    ncols, j = l, j1
        if j is None:
            continue
        a_ii = udata[i,j]
        if i!=j:
            pivot_table[i], pivot_table[j] = pivot_table[j], pivot_table[i]
            swap_rows_MATRIX_T(udata, i, j)
            swap_rows_MATRIX(ldata, i, j)
            swap_rc_map(urows, i, j)

        irow = urows[i]
        for j in xrange(i+1,m):
            u_ji = udata_get((i,j))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = urows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = udata[p,i] * c
                jp = p,j
                b = udata_get(jp)
                if b is None:
                    udata[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if b==u_ip_c:
                        del udata[jp]
                        jrow_remove(p)
                    else:
                        udata[jp] = b - u_ip_c
            if not jrow:
                del urows[j]
            if i<k and j<m:
                ldata[j, i] = c
    for i in xrange(min(m, k)):
        ldata[i,i] = 1
    if m>n:
        crop_MATRIX_T(k, n, udata)
    return pivot_table

def MATRIX_DICT_crop(self):
    """ Remove matrix elements that are out of dimensions inplace and return the matrix.
    """
    if not self.is_writable:
        raise TypeError('Cannot crop read-only matrix inplace')
    head, data = self.pair
    m, n = head.shape
    if head.is_transpose:
        crop_MATRIX_T(m, n, data)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        crop_MATRIX(m, n, data)
    return self

def crop_MATRIX(m, n, data):
    for (i,j) in data.keys():
        if 0<=i<m and 0<=j<n:
            continue
        del data[i,j]

def crop_MATRIX_T(m, n, data):
    for (j, i) in data.keys():
        if 0<=i<m and 0<=j<n:
            continue
        del data[j, i]

def swap_rows_MATRIX(data, i, j):
    if i==j:
        return
    row_i = []
    row_j = []
    for index, element in data.items():
        i0 = index[0]
        if i0==i:
            row_i.append((index[1], element))
            del data[index]
        elif i0==j:
            row_j.append((index[1], element))
            del data[index]
    for k, element in row_i:
        data[j, k] = element
    for k, element in row_j:
        data[i, k] = element

def swap_cols_MATRIX(data, i, j):
    if i==j:
        return
    column_i = []
    column_j = []
    for index, element in data.items():
        i0 = index[1]
        if i0==i:
            column_i.append((index[0], element))
            del data[index]
        elif i0==j:
            column_j.append((index[0], element))
            del data[index]
    for k, element in column_i:
        data[k, j] = element
    for k,element in column_j:
        data[k, i] = element

swap_rows_MATRIX_T = swap_cols_MATRIX
swap_cols_MATRIX_T = swap_rows_MATRIX

def MATRIX_DICT_swap_rows(self, i, j):
    if not self.is_writable:
        raise TypeError('Cannot swap rows of a read-only matrix')
    head, data = self.pair
    if head.is_transpose:
        swap_rows_MATRIX_T(data, i, j)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        swap_rows_MATRIX(data, i, j)

def MATRIX_DICT_swap_cols(self, i, j):
    if not self.is_writable:
        raise TypeError('Cannot swap columns of a read-only matrix')
    head, data = self.pair
    if head.is_transpose:
        swap_cols_MATRIX_T(data, i, j)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        swap_cols_MATRIX(data, i, j)

def MATRIX_DICT_trace(self):
    """ Return trace of a matrix.
    """
    head, data = self.pair
    m, n = head.shape
    s = 0
    if head.is_diagonal:
        dget = data.get
        for i in xrange(min(n,m)):
            s += dget((i, i), 0)
        return s
    if m != n:
        raise ValueError("matrix trace is only defined for square matrices but got %sx%s" % (m,n))
    sparse = len(data) < m
    if sparse:
        for (i, j), element in data.items():
            if i == j:
                s += element
    else:
        dget = data.get
        for i in xrange(n):
            s += dget((i, i), 0)
    return s
<commit_msg>Another 2x speed up patch for sparse lu.<commit_after># Author: Pearu Peterson
# Created: March 2008

from ..utils import MATRIX, MATRIX_DICT, MATRIX_DICT_T
from ..arithmetic.numbers import div
from .algebra import MatrixDict, Matrix

from ..core import init_module
init_module.import_lowlevel_operations()

def MATRIX_DICT_gauss_jordan_elimination(self, overwrite=False):
    """ Perform Gauss-Jordan elimination of a m x n matrix A.

    Outputs::

      B - m x n matrix where the lhs part is unit matrix
    """
    head, data = self.pair
    m, n = head.shape
    k = min(m,n)
    if overwrite and self.is_writable:
        udata = data
    else:
        udata = dict(data)
    if head.is_transpose:
        B = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        gauss_jordan_elimination_MATRIX_T(m, n, udata)
    elif head.is_diagonal:
        raise NotImplementedError(`head, head.is_diagonal`)
    else:
        B = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        gauss_jordan_elimination_MATRIX(m, n, udata)
    return B

def MATRIX_DICT_lu(self, overwrite=False):
    """ Perform LU factorization of a m x n matrix A.

    Parameters::

      overwrite=False - if True then discard the content of matrix A

    Outputs::

      P, L, U - LU decomposition matrices of A.

    Definitions::

      P - m x m permuation matrix
      L - m x k lower triangular or trapezoidal matrix with unit-diagonal
      U - k x n upper triangular or trapezoidal matrix
      k = min(m,n)

      A = P * L * U
    """
    head, data = self.pair
    m, n = head.shape
    k = min(m, n)
    ldata = {}
    if overwrite and self.is_writable:
        udata = data
    else:
        udata = dict(data)
    L = MatrixDict(MATRIX(m, k, MATRIX_DICT), ldata)
    if head.is_transpose:
        U = MatrixDict(MATRIX(k, n, MATRIX_DICT_T), udata)
        pivot_table = lu_MATRIX_T(m, n, k, ldata, udata)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        U = MatrixDict(MATRIX(k, n, MATRIX_DICT), udata)
        pivot_table = lu_MATRIX(m, n, k, ldata, udata)
    P = Matrix(pivot_table, permutation=True).T
    return P, L, U

def gauss_jordan_elimination_MATRIX(m, n, data):
    data_get = data.get
    data_has = data.has_key
    rows = get_rc_map(data)
    for i in xrange(m):
        if 0:
            # for upper diagonal sparse matrix this is a bit slower..
            ncols, j = n, None
            for j1 in range(i,m):
                if data_has((j1,i)):
                    l = len(rows[j1])
                    if l <= ncols:
                        ncols, j = l, j1
            if j is None:
                continue
            a_ii = data[j,i]
            if i!=j:
                swap_rows_MATRIX(data, i, j, rows)
                swap_rc_map(rows, i, j)
        else:
            a_ii = data_get((i,i))
            if a_ii is None:
                for j in xrange(i+1, m):
                    a_ii = data_get((j,i))
                    if a_ii is not None:
                        break
                if a_ii is None:
                    continue
                swap_rows_MATRIX(data, i, j, rows)
                swap_rc_map(rows, i, j)
        irow = rows[i]
        for j in range(m):
            if j==i:
                continue
            u_ji = data_get((j,i))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = rows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = data[i,p] * c
                jp = j,p
                b = data_get(jp)
                if b is None:
                    data[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if u_ip_c==b:
                        del data[jp]
                        jrow_remove(p)
                    else:
                        data[jp] = b - u_ip_c
            if not jrow:
                del rows[j]
        data[i,i] = 1

        for p in irow:
            if p <= i: continue
            ip = i,p
            data[ip] = div(data[ip], a_ii)

def gauss_jordan_elimination_MATRIX_T(m, n, data):
    data_get = data.get
    rows = get_rc_map_T(data)
    for i in xrange(m):
        a_ii = data_get((i,i))
        if a_ii is None:
            for j in xrange(i+1, m):
                a_ii = data_get((i,j))
                if a_ii is not None:
                    break
            if a_ii is None:
                continue
            swap_rows_MATRIX_T(data, i, j, rows)
            swap_rc_map(rows, i, j)
        irow = rows[i]
        for j in range(m):
            if j==i:
                continue
            u_ji = data_get((i,j))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = rows[j]
            for p in range(i,n):
                u_ip = data_get((p,i))
                if u_ip is None:
                    continue
                jp = p,j
                b = data_get(jp)
                if b is None:
                    data[jp] = -u_ip*c
                    jrow.add(p)
                else:
                    b -= u_ip*c
                    if b:
                        data[jp] = b
                    else:
                        del data[jp]
                        jrow.remove(p)
            if not jrow:
                del rows[j]
        data[i,i] = 1
        for p in range(i+1, n):
            ip = p,i
            data[ip] = div(data[ip], a_ii)

def get_rc_map(data):
    rows = {}
    for i, j in data:
        s = rows.get (i)
        if s is None:
            s = rows[i] = set ()
        s.add(j)
    return rows

def get_rc_map_T(data):
    rows = {}
    for j, i in data:
        s = rows.get (i)
        if s is None:
            s = rows[i] = set ()
        s.add(j)
    return rows

def swap_rc_map(rows, i, j):
    ri = rows.get(i)
    rj = rows.get(j)
    if rj is not None: del rows[j]
    if ri is not None: del rows[i]
    if rj is not None:
        rows[i] = rj
    if ri is not None:
        rows[j] = ri

def lu_MATRIX(m, n, k, ldata, udata):
    if m>n:
        for i in xrange(n,m):
            udata[(i,i)] = 1
    pivot_table = range(m)
    udata_get = udata.get
    udata_has = udata.has_key

    urows = get_rc_map(udata)
    lrows = get_rc_map(ldata)
    for i in xrange(m-1):
        ncols, j = n, None
        for j1 in range(i,m):
            if udata_has((j1,i)):
                l = len(urows[j1])
                if l <= ncols:
                    ncols, j = l, j1
        if j is None:
            continue
        a_ii = udata[j,i]
        if i!=j:
            pivot_table[i], pivot_table[j] = pivot_table[j], pivot_table[i]
            swap_rows_MATRIX(udata, i, j, urows)
            swap_rows_MATRIX(ldata, i, j, lrows)
            swap_rc_map(urows, i, j)
            swap_rc_map(lrows, i, j)
        irow = urows[i]
        for j in range (i+1,m):
            u_ji = udata_get((j,i))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = urows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = udata[i,p] * c
                jp = j,p
                b = udata_get(jp)
                if b is None:
                    udata[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if b==u_ip_c:
                        del udata[jp]
                        jrow_remove(p)
                    else:
                        udata[jp] = b - u_ip_c
            if not jrow:
                del urows[j]
            if i<k and j<m:
                ldata[j, i] = c
                s = lrows.get(j)
                if s is None:
                    s = lrows[j] = set()
                s.add(i)
    for i in xrange(min(m, k)):
        ldata[i,i] = 1
    if m>n:
        crop_MATRIX(k, n, udata)
    return pivot_table

def lu_MATRIX_T(m, n, k, ldata, udata):
    if m>n:
        for i in xrange(n,m):
            udata[(i,i)] = 1
    pivot_table = range(m)
    udata_get = udata.get
    udata_has = udata.has_key
    urows = get_rc_map_T(udata)
    lrows = get_rc_map(ldata)
    for i in xrange(m-1):
        ncols, j = n, None
        for j1 in range(i, m):
            if udata_has((i,j1)):
                l = len(urows[j1])
                if l <= ncols:
                    ncols, j = l, j1
        if j is None:
            continue
        a_ii = udata[i,j]
        if i!=j:
            pivot_table[i], pivot_table[j] = pivot_table[j], pivot_table[i]
            swap_rows_MATRIX_T(udata, i, j, urows)
            swap_rows_MATRIX(ldata, i, j, lrows)
            swap_rc_map(urows, i, j)
            swap_rc_map(lrows, i, j)
        irow = urows[i]
        for j in xrange(i+1,m):
            u_ji = udata_get((i,j))
            if u_ji is None:
                continue
            c = div(u_ji, a_ii)
            jrow = urows[j]
            jrow_add = jrow.add
            jrow_remove = jrow.remove
            for p in irow:
                if p < i: continue
                u_ip_c = udata[p,i] * c
                jp = p,j
                b = udata_get(jp)
                if b is None:
                    udata[jp] = -u_ip_c
                    jrow_add(p)
                else:
                    if b==u_ip_c:
                        del udata[jp]
                        jrow_remove(p)
                    else:
                        udata[jp] = b - u_ip_c
            if not jrow:
                del urows[j]
            if i<k and j<m:
                ldata[j, i] = c
                s = lrows.get(j)
                if s is None:
                    s = lrows[j] = set()
                s.add(i)
    for i in xrange(min(m, k)):
        ldata[i,i] = 1
    if m>n:
        crop_MATRIX_T(k, n, udata)
    return pivot_table

def MATRIX_DICT_crop(self):
    """ Remove matrix elements that are out of dimensions inplace and return the matrix.
    """
    if not self.is_writable:
        raise TypeError('Cannot crop read-only matrix inplace')
    head, data = self.pair
    m, n = head.shape
    if head.is_transpose:
        crop_MATRIX_T(m, n, data)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        crop_MATRIX(m, n, data)
    return self

def crop_MATRIX(m, n, data):
    for (i,j) in data.keys():
        if 0<=i<m and 0<=j<n:
            continue
        del data[i,j]

def crop_MATRIX_T(m, n, data):
    for (j, i) in data.keys():
        if 0<=i<m and 0<=j<n:
            continue
        del data[j, i]

def swap_rows_MATRIX(data, i, j, rows=None):
    if i==j:
        return
    if rows is not None:
        d = {}
        data_pop = data.pop
        for k in rows.get(i,[]):
            d[j,k] = data_pop((i,k))
        for k in rows.get(j,[]):
            d[i,k] = data_pop((j,k))
        data.update(d)
        return

    row_i = []
    row_j = []
    for index, element in data.items():
        i0 = index[0]
        if i0==i:
            row_i.append((index[1], element))
            del data[index]
        elif i0==j:
            row_j.append((index[1], element))
            del data[index]
    for k, element in row_i:
        data[j, k] = element
    for k, element in row_j:
        data[i, k] = element

def swap_cols_MATRIX(data, i, j, cols=None):
    if i==j:
        return
    if cols is not None:
        d = {}
        data_pop = data.pop
        for k in cols.get(i,[]):
            d[k,j] = data_pop((k,i))
        for k in cols.get(j,[]):
            d[k,i] = data_pop((k,j))
        data.update(d)
        return
    column_i = []
    column_j = []
    for index, element in data.items():
        i0 = index[1]
        if i0==i:
            column_i.append((index[0], element))
            del data[index]
        elif i0==j:
            column_j.append((index[0], element))
            del data[index]
    for k, element in column_i:
        data[k, j] = element
    for k,element in column_j:
        data[k, i] = element

swap_rows_MATRIX_T = swap_cols_MATRIX
swap_cols_MATRIX_T = swap_rows_MATRIX

def MATRIX_DICT_swap_rows(self, i, j):
    if not self.is_writable:
        raise TypeError('Cannot swap rows of a read-only matrix')
    head, data = self.pair
    if head.is_transpose:
        swap_rows_MATRIX_T(data, i, j)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        swap_rows_MATRIX(data, i, j)

def MATRIX_DICT_swap_cols(self, i, j):
    if not self.is_writable:
        raise TypeError('Cannot swap columns of a read-only matrix')
    head, data = self.pair
    if head.is_transpose:
        swap_cols_MATRIX_T(data, i, j)
    elif head.is_diagonal:
        raise NotImplementedError(`head`)
    else:
        swap_cols_MATRIX(data, i, j)

def MATRIX_DICT_trace(self):
    """ Return trace of a matrix.
    """
    head, data = self.pair
    m, n = head.shape
    s = 0
    if head.is_diagonal:
        dget = data.get
        for i in xrange(min(n,m)):
            s += dget((i, i), 0)
        return s
    if m != n:
        raise ValueError("matrix trace is only defined for square matrices but got %sx%s" % (m,n))
    sparse = len(data) < m
    if sparse:
        for (i, j), element in data.items():
            if i == j:
                s += element
    else:
        dget = data.get
        for i in xrange(n):
            s += dget((i, i), 0)
    return s
<|endoftext|>


<commit_before>import numpy as np
import matplotlib.pyplot as plt

from astropy.table import hstack, Table
import agasc
import Quaternion
from Ska.quatutil import radec2yagzag

from .transform import pixels_to_yagzag

# rc definitions
frontcolor = 'black'
backcolor = 'white'
plt.rcParams['lines.color'] = frontcolor
plt.rcParams['patch.edgecolor'] = frontcolor
plt.rcParams['text.color'] = frontcolor
plt.rcParams['axes.facecolor'] = backcolor
plt.rcParams['axes.edgecolor'] = frontcolor
plt.rcParams['axes.labelcolor'] = frontcolor
plt.rcParams['xtick.color'] = frontcolor
plt.rcParams['ytick.color'] = frontcolor
plt.rcParams['grid.color'] = frontcolor
plt.rcParams['figure.facecolor'] = backcolor
plt.rcParams['figure.edgecolor'] = backcolor
plt.rcParams['savefig.facecolor'] = backcolor
plt.rcParams['savefig.edgecolor'] = backcolor

BAD_STAR_COLOR = 'tomato'
BAD_STAR_ALPHA = .75
FAINT_STAR_COLOR = 'lightseagreen'
FAINT_STAR_ALPHA = .75

def symsize(mag):
    # map mags to figsizes, defining
    # mag 6 as 40 and mag 11 as 3
    # interp should leave it at the bounding value outside
    # the range
    return np.interp(mag, [6.0, 11.0], [40.0, 3.0])

def _plot_catalog_items(ax, catalog):
    """
    Plot catalog items (guide, acq, bot, mon, fid) in yang and zang on the supplied
    axes object in place.

    :param ax: matplotlib axes
    :param catalog: data structure containing starcheck-style columns/attributes
                    catalog records.  This can be anything that will work with
                    astropy.table.Table(catalog).  A list of dicts is the convention.
    """
    cat = Table(catalog)
    gui = cat[(cat['type'] == 'GUI') | (cat['type'] == 'BOT')]
    acq = cat[(cat['type'] == 'ACQ') | (cat['type'] == 'BOT')]
    fid = cat[cat['type'] == 'FID']
    mon = cat[cat['type'] == 'MON']
    for row in cat:
        ax.annotate("%s" % row['idx'],
                    xy=(row['yang'] - 120, row['zang'] + 60),
                    color='red',
                    fontsize=12)
    ax.scatter(gui['yang'], gui['zang'],
               facecolors='none',
               edgecolors='green',
               s=100)
    for acq_star in acq:
        acq_box = plt.Rectangle(
            (acq_star['yang'] - acq_star['halfw'],
             acq_star['zang'] - acq_star['halfw']),
            width=acq_star['halfw'] * 2,
            height=acq_star['halfw'] * 2,
            color='blue',
            fill=False)
        ax.add_patch(acq_box)
    for mon_box in mon:
        # starcheck convention was to plot monitor boxes at 2X halfw
        box = plt.Rectangle(
            (mon_box['yang'] - (mon_box['halfw'] * 2),
             mon_box['zang'] - (mon_box['halfw'] * 2)),
            width=mon_box['halfw'] * 4,
            height=mon_box['halfw'] * 4,
            color='orange',
            fill=False)
        ax.add_patch(box)
    ax.scatter(fid['yang'], fid['zang'],
               facecolors='none',
               edgecolors='red',
               linewidth=.5,
               marker='o',
               s=175)
    ax.scatter(fid['yang'], fid['zang'],
               facecolors='none',
               edgecolors='red',
               marker='+',
               linewidth=.5,
               s=175)

def _plot_field_stars(ax, stars, attitude, red_mag_lim=None, bad_stars=None):
    """
    Plot plot field stars in yang and zang on the supplied
    axes object in place.

    :param ax: matplotlib axes
    :param stars: astropy.table compatible set of records of agasc entries of stars
    :param attitude: Quaternion-compatible attitude
    :param red_mag_lim: faint limit
    :param bad_stars: boolean mask of stars to be plotted in red
    """
    stars = Table(stars)
    quat = Quaternion.Quat(attitude)

    # Add star Y angle and Z angle in arcsec to the stars table
    yagzags = (radec2yagzag(star['RA_PMCORR'], star['DEC_PMCORR'], quat)
               for star in stars)
    yagzags = Table(rows=[(y * 3600, z * 3600) for y, z in yagzags], names=['yang', 'zang'])
    stars = hstack([stars, yagzags])
    ok = np.ones(len(stars), dtype='bool')

    if red_mag_lim:
        # Mark stars with the FAINT_STAR_COLOR if they have MAG_ACA
        # that is fainter than red_mag_lim but MAG_ACA - a rough mag error
        # is brighter than that limit.  The rough mag error calculation is
        # based on the SAUSAGE acq stage 1 check, which uses nsigma of
        # 3.0, a mag low limit of 1.5, and a random error of 0.26.
        nsigma = 3.0
        mag_error_low_limit = 1.5
        randerr = 0.26
        caterr = stars['MAG_ACA_ERR'] / 100.
        error = nsigma * np.sqrt(randerr*randerr + caterr*caterr)
        error[error < mag_error_low_limit] = mag_error_low_limit
        faint = ((stars['MAG_ACA'] >= red_mag_lim)
                 & ((stars['MAG_ACA'] - error) < red_mag_lim))
        too_dim_to_plot = ((stars['MAG_ACA'] >= red_mag_lim)
                            & ((stars['MAG_ACA'] - error) >= red_mag_lim))
        # All of the indices are on stars, so reindex without too_dim_to_plot
        stars = stars[~too_dim_to_plot]
        faint = faint[~too_dim_to_plot]
        ok = ok[~too_dim_to_plot]
        ok[faint] = False
        if bad_stars is not None:
            bad_stars = bad_stars[~too_dim_to_plot]
    if bad_stars is not None:
        ok[bad_stars] = False

    size = symsize(stars['MAG_ACA'])
    ax.scatter(stars[ok]['yang'], stars[ok]['zang'],
               c='black', s=size[ok])
    if red_mag_lim:
        ax.scatter(stars[faint]['yang'], stars[faint]['zang'],
                   c=FAINT_STAR_COLOR, s=size[faint], edgecolor='none',
                   alpha=FAINT_STAR_ALPHA)
    if bad_stars is not None:
        ax.scatter(stars[bad_stars]['yang'], stars[bad_stars]['zang'],
                   c=BAD_STAR_COLOR, s=size[bad_stars], edgecolor='none',
                   alpha=BAD_STAR_ALPHA)

def star_plot(catalog=None, attitude=None, stars=None, title=None, starcat_time=None,
              red_mag_lim=None, quad_bound=True, grid=True, bad_stars=None):
    """
    Plot a catalog, a star field, or both in a matplotlib figure.
    If supplying a star field, an attitude must also be supplied.

    :param catalog: Records describing catalog.  Must be astropy table compatible.
    :param attitude: A Quaternion compatible attitude for the pointing
    :param stars: astropy table compatible set of agasc records of stars
    :param title: string to be used as suptitle for the figure
    :param red_mag_lim: faint limit for field star plotting.
    :param quad_bound: boolean, plot inner quadrant boundaries
    :param grid: boolean, plot axis grid
    :param bad_stars: boolean mask on 'stars' of those that don't meet minimum requirements
                      to be selected as acq stars
    :returns: matplotlib figure
    """
    if stars is None:
        quat = Quaternion.Quat(attitude)
        stars = agasc.get_agasc_cone(quat.ra, quat.dec,
                                     radius=1.5,
                                     date=starcat_time)
    if bad_stars is None:
        bad_stars = bad_acq_stars(stars)

    fig = plt.figure(figsize=(5.325, 5.325))
    ax = fig.add_subplot(1, 1, 1)
    plt.subplots_adjust(top=0.95)
    ax.set_aspect('equal')

    # plot the box and set the labels
    plt.xlim(2900, -2900)
    plt.ylim(-2900, 2900)
    b1hw = 2560
    box1 = plt.Rectangle((b1hw, -b1hw), -2 * b1hw, 2 * b1hw,
                         fill=False)
    ax.add_patch(box1)
    b2w = 2600
    box2 = plt.Rectangle((b2w, -b1hw), -4 + -2 * b2w, 2 * b1hw,
                         fill=False)
    ax.add_patch(box2)

    ax.scatter([-2700, -2700, -2700, -2700, -2700],
               [2400, 2100, 1800, 1500, 1200],
               c='orange', edgecolors='none',
               s=symsize(np.array([10.0, 9.0, 8.0, 7.0, 6.0])))

    [l.set_rotation(90) for l in ax.get_yticklabels()]
    ax.grid(grid)
    ax.set_ylabel("Zag (arcsec)")
    ax.set_xlabel("Yag (arcsec)")

    if quad_bound:
        pix_range = np.linspace(-510, 510, 50)
        minus_half_pix = -0.5 * np.ones_like(pix_range)
        # plot the row = -0.5 line
        yag, zag = pixels_to_yagzag(minus_half_pix, pix_range)
        ax.plot(yag, zag, color='magenta', alpha=.4)
        # plot the col = -0.5 line
        yag, zag = pixels_to_yagzag(pix_range, minus_half_pix)
        ax.plot(yag, zag, color='magenta', alpha=.4)

    # plot field if present
    if stars is not None:
        if attitude is None:
            raise ValueError("Must supply attitude to plot field stars")
        _plot_field_stars(ax, stars, attitude=attitude,
                          bad_stars=bad_stars, red_mag_lim=red_mag_lim)
    # plot starcheck catalog
    if catalog is not None:
        _plot_catalog_items(ax, catalog)
    if title is not None:
        fig.suptitle(title, fontsize='small')
    return fig

def bad_acq_stars(stars):
    """
    Return mask of 'bad' stars, by evaluating AGASC star parameters.
    :param stars: astropy table-compatible set of agasc records of stars
    :returns: boolean mask true for 'bad' stars
    """
    return ((stars['CLASS'] != 0) |
            (stars['MAG_ACA_ERR'] > 100) |
            (stars['POS_ERR'] > 3000) |
            (stars['ASPQ1'] > 0) |
            (stars['ASPQ2'] > 0) |
            (stars['ASPQ3'] > 999) |
            (stars['VAR'] > -9999))

def plot_compass(roll):
    """
    Make a compass plot.

    :param roll: Attitude roll for compass plot.
    :returns: matplotlib figure
    """
    fig = plt.figure(figsize=(3, 3))
    ax = plt.subplot(polar=True)
    ax.annotate("", xy=(0, 0), xytext=(0, 1),
                arrowprops=dict(arrowstyle="<-", color="k"))
    ax.annotate("", xy=(0, 0), xytext=(np.radians(90), 1),
                arrowprops=dict(arrowstyle="<-", color="k"))
    ax.annotate("N", xy=(0, 0), xytext=(0, 1.2))
    ax.annotate("E", xy=(0, 0), xytext=(np.radians(90), 1.2))
    ax.set_theta_offset(np.radians(90 + roll))
    ax.grid(False)
    ax.set_yticklabels([])
    plt.ylim(0, 1.4)
    plt.tight_layout()
    return fig
<commit_msg>Improve style in red_mag_lim determination<commit_after>import numpy as np
import matplotlib.pyplot as plt

from astropy.table import hstack, Table
import agasc
import Quaternion
from Ska.quatutil import radec2yagzag

from .transform import pixels_to_yagzag

# rc definitions
frontcolor = 'black'
backcolor = 'white'
plt.rcParams['lines.color'] = frontcolor
plt.rcParams['patch.edgecolor'] = frontcolor
plt.rcParams['text.color'] = frontcolor
plt.rcParams['axes.facecolor'] = backcolor
plt.rcParams['axes.edgecolor'] = frontcolor
plt.rcParams['axes.labelcolor'] = frontcolor
plt.rcParams['xtick.color'] = frontcolor
plt.rcParams['ytick.color'] = frontcolor
plt.rcParams['grid.color'] = frontcolor
plt.rcParams['figure.facecolor'] = backcolor
plt.rcParams['figure.edgecolor'] = backcolor
plt.rcParams['savefig.facecolor'] = backcolor
plt.rcParams['savefig.edgecolor'] = backcolor

BAD_STAR_COLOR = 'tomato'
BAD_STAR_ALPHA = .75
FAINT_STAR_COLOR = 'lightseagreen'
FAINT_STAR_ALPHA = .75

def symsize(mag):
    # map mags to figsizes, defining
    # mag 6 as 40 and mag 11 as 3
    # interp should leave it at the bounding value outside
    # the range
    return np.interp(mag, [6.0, 11.0], [40.0, 3.0])

def _plot_catalog_items(ax, catalog):
    """
    Plot catalog items (guide, acq, bot, mon, fid) in yang and zang on the supplied
    axes object in place.

    :param ax: matplotlib axes
    :param catalog: data structure containing starcheck-style columns/attributes
                    catalog records.  This can be anything that will work with
                    astropy.table.Table(catalog).  A list of dicts is the convention.
    """
    cat = Table(catalog)
    gui = cat[(cat['type'] == 'GUI') | (cat['type'] == 'BOT')]
    acq = cat[(cat['type'] == 'ACQ') | (cat['type'] == 'BOT')]
    fid = cat[cat['type'] == 'FID']
    mon = cat[cat['type'] == 'MON']
    for row in cat:
        ax.annotate("%s" % row['idx'],
                    xy=(row['yang'] - 120, row['zang'] + 60),
                    color='red',
                    fontsize=12)
    ax.scatter(gui['yang'], gui['zang'],
               facecolors='none',
               edgecolors='green',
               s=100)
    for acq_star in acq:
        acq_box = plt.Rectangle(
            (acq_star['yang'] - acq_star['halfw'],
             acq_star['zang'] - acq_star['halfw']),
            width=acq_star['halfw'] * 2,
            height=acq_star['halfw'] * 2,
            color='blue',
            fill=False)
        ax.add_patch(acq_box)
    for mon_box in mon:
        # starcheck convention was to plot monitor boxes at 2X halfw
        box = plt.Rectangle(
            (mon_box['yang'] - (mon_box['halfw'] * 2),
             mon_box['zang'] - (mon_box['halfw'] * 2)),
            width=mon_box['halfw'] * 4,
            height=mon_box['halfw'] * 4,
            color='orange',
            fill=False)
        ax.add_patch(box)
    ax.scatter(fid['yang'], fid['zang'],
               facecolors='none',
               edgecolors='red',
               linewidth=.5,
               marker='o',
               s=175)
    ax.scatter(fid['yang'], fid['zang'],
               facecolors='none',
               edgecolors='red',
               marker='+',
               linewidth=.5,
               s=175)

def _plot_field_stars(ax, stars, attitude, red_mag_lim=None, bad_stars=None):
    """
    Plot plot field stars in yang and zang on the supplied
    axes object in place.

    :param ax: matplotlib axes
    :param stars: astropy.table compatible set of records of agasc entries of stars
    :param attitude: Quaternion-compatible attitude
    :param red_mag_lim: faint limit
    :param bad_stars: boolean mask of stars to be plotted in red
    """
    stars = Table(stars)
    quat = Quaternion.Quat(attitude)

    # Add star Y angle and Z angle in arcsec to the stars table
    yagzags = (radec2yagzag(star['RA_PMCORR'], star['DEC_PMCORR'], quat)
               for star in stars)
    yagzags = Table(rows=[(y * 3600, z * 3600) for y, z in yagzags], names=['yang', 'zang'])
    stars = hstack([stars, yagzags])
    ok = np.ones(len(stars), dtype='bool')

    if red_mag_lim:
        # Mark stars with the FAINT_STAR_COLOR if they have MAG_ACA
        # that is fainter than red_mag_lim but MAG_ACA - a rough mag error
        # is brighter than that limit.  The rough mag error calculation is
        # based on the SAUSAGE acq stage 1 check, which uses nsigma of
        # 3.0, a mag low limit of 1.5, and a random error of 0.26.
        nsigma = 3.0
        mag_error_low_limit = 1.5
        randerr = 0.26
        caterr = stars['MAG_ACA_ERR'] / 100.
        error = nsigma * np.sqrt(randerr**2 + caterr**2)
        error = np.clip(error, a_min=mag_error_low_limit, a_max=None)
        faint = ((stars['MAG_ACA'] >= red_mag_lim)
                 & ((stars['MAG_ACA'] - error) < red_mag_lim))
        too_dim_to_plot = ((stars['MAG_ACA'] >= red_mag_lim)
                            & ((stars['MAG_ACA'] - error) >= red_mag_lim))
        # All of the indices are on stars, so reindex without too_dim_to_plot
        stars = stars[~too_dim_to_plot]
        faint = faint[~too_dim_to_plot]
        ok = ok[~too_dim_to_plot]
        ok[faint] = False
        if bad_stars is not None:
            bad_stars = bad_stars[~too_dim_to_plot]
    if bad_stars is not None:
        ok[bad_stars] = False

    size = symsize(stars['MAG_ACA'])
    ax.scatter(stars[ok]['yang'], stars[ok]['zang'],
               c='black', s=size[ok])
    if red_mag_lim:
        ax.scatter(stars[faint]['yang'], stars[faint]['zang'],
                   c=FAINT_STAR_COLOR, s=size[faint], edgecolor='none',
                   alpha=FAINT_STAR_ALPHA)
    if bad_stars is not None:
        ax.scatter(stars[bad_stars]['yang'], stars[bad_stars]['zang'],
                   c=BAD_STAR_COLOR, s=size[bad_stars], edgecolor='none',
                   alpha=BAD_STAR_ALPHA)

def star_plot(catalog=None, attitude=None, stars=None, title=None, starcat_time=None,
              red_mag_lim=None, quad_bound=True, grid=True, bad_stars=None):
    """
    Plot a catalog, a star field, or both in a matplotlib figure.
    If supplying a star field, an attitude must also be supplied.

    :param catalog: Records describing catalog.  Must be astropy table compatible.
    :param attitude: A Quaternion compatible attitude for the pointing
    :param stars: astropy table compatible set of agasc records of stars
    :param title: string to be used as suptitle for the figure
    :param red_mag_lim: faint limit for field star plotting.
    :param quad_bound: boolean, plot inner quadrant boundaries
    :param grid: boolean, plot axis grid
    :param bad_stars: boolean mask on 'stars' of those that don't meet minimum requirements
                      to be selected as acq stars
    :returns: matplotlib figure
    """
    if stars is None:
        quat = Quaternion.Quat(attitude)
        stars = agasc.get_agasc_cone(quat.ra, quat.dec,
                                     radius=1.5,
                                     date=starcat_time)
    if bad_stars is None:
        bad_stars = bad_acq_stars(stars)

    fig = plt.figure(figsize=(5.325, 5.325))
    ax = fig.add_subplot(1, 1, 1)
    plt.subplots_adjust(top=0.95)
    ax.set_aspect('equal')

    # plot the box and set the labels
    plt.xlim(2900, -2900)
    plt.ylim(-2900, 2900)
    b1hw = 2560
    box1 = plt.Rectangle((b1hw, -b1hw), -2 * b1hw, 2 * b1hw,
                         fill=False)
    ax.add_patch(box1)
    b2w = 2600
    box2 = plt.Rectangle((b2w, -b1hw), -4 + -2 * b2w, 2 * b1hw,
                         fill=False)
    ax.add_patch(box2)

    ax.scatter([-2700, -2700, -2700, -2700, -2700],
               [2400, 2100, 1800, 1500, 1200],
               c='orange', edgecolors='none',
               s=symsize(np.array([10.0, 9.0, 8.0, 7.0, 6.0])))

    [l.set_rotation(90) for l in ax.get_yticklabels()]
    ax.grid(grid)
    ax.set_ylabel("Zag (arcsec)")
    ax.set_xlabel("Yag (arcsec)")

    if quad_bound:
        pix_range = np.linspace(-510, 510, 50)
        minus_half_pix = -0.5 * np.ones_like(pix_range)
        # plot the row = -0.5 line
        yag, zag = pixels_to_yagzag(minus_half_pix, pix_range)
        ax.plot(yag, zag, color='magenta', alpha=.4)
        # plot the col = -0.5 line
        yag, zag = pixels_to_yagzag(pix_range, minus_half_pix)
        ax.plot(yag, zag, color='magenta', alpha=.4)

    # plot field if present
    if stars is not None:
        if attitude is None:
            raise ValueError("Must supply attitude to plot field stars")
        _plot_field_stars(ax, stars, attitude=attitude,
                          bad_stars=bad_stars, red_mag_lim=red_mag_lim)
    # plot starcheck catalog
    if catalog is not None:
        _plot_catalog_items(ax, catalog)
    if title is not None:
        fig.suptitle(title, fontsize='small')
    return fig

def bad_acq_stars(stars):
    """
    Return mask of 'bad' stars, by evaluating AGASC star parameters.
    :param stars: astropy table-compatible set of agasc records of stars
    :returns: boolean mask true for 'bad' stars
    """
    return ((stars['CLASS'] != 0) |
            (stars['MAG_ACA_ERR'] > 100) |
            (stars['POS_ERR'] > 3000) |
            (stars['ASPQ1'] > 0) |
            (stars['ASPQ2'] > 0) |
            (stars['ASPQ3'] > 999) |
            (stars['VAR'] > -9999))

def plot_compass(roll):
    """
    Make a compass plot.

    :param roll: Attitude roll for compass plot.
    :returns: matplotlib figure
    """
    fig = plt.figure(figsize=(3, 3))
    ax = plt.subplot(polar=True)
    ax.annotate("", xy=(0, 0), xytext=(0, 1),
                arrowprops=dict(arrowstyle="<-", color="k"))
    ax.annotate("", xy=(0, 0), xytext=(np.radians(90), 1),
                arrowprops=dict(arrowstyle="<-", color="k"))
    ax.annotate("N", xy=(0, 0), xytext=(0, 1.2))
    ax.annotate("E", xy=(0, 0), xytext=(np.radians(90), 1.2))
    ax.set_theta_offset(np.radians(90 + roll))
    ax.grid(False)
    ax.set_yticklabels([])
    plt.ylim(0, 1.4)
    plt.tight_layout()
    return fig
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division
import itertools
import logging

from sqlalchemy.sql import func

from chanjo.compat import itervalues
from chanjo.utils import list_get
from .core import Store
from .converter import ChanjoConverterMixin
from .models import Exon, ExonStatistic, Sample, Transcript
from .utils import filter_samples, group_by_field, predict_gender

logger = logging.getLogger(__name__)

class ChanjoAPI(Store, ChanjoConverterMixin):

    """Interface to chanjo database for asking interesting questions.

    Attributes:
        weighten_average (BinaryExpression): weighter mean for metrics
    """

    def init_app(self, app, key_base='CHANJO_'):
        """Configure API (Flask style) after lazy initialization.

        Args:
            app (Flask): Flask app instance
            key_base (str): namespace to look for under ``app.config``

        Returns:
            ChanjoAPI: ``self``
        """
        uri = app.config["{}URI".format(key_base)]
        self.connect(db_uri=uri)
        return self

    @property
    def weighted_average(self):
        """Build a dynamic field to calculate a weighted average.

        Works for any of the exon metrics.

        Returns:
            BinaryExpression: dynamic field for a query
        """
        # length of the exon
        weight = Exon.end - Exon.start
        total_weight = func.sum(weight)
        total_value = func.sum(ExonStatistic.value * weight)
        weighted_mean = (total_value / total_weight).label('weighted_mean')
        return weighted_mean

    def samples(self, group_id=None, sample_ids=None):
        """Get samples from the database.

        Args:
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            List[Sample]: samples objects from the database
        """
        query = self.query(Sample).order_by(Sample.sample_id)
        query = filter_samples(query, group_id=group_id, sample_ids=sample_ids)
        return query

    def means(self, query=None):
        """Calculate means on a per sample basis.

        Args:
            query (Optional[Query]): initialized and filtered query

        Returns:
            List[tuple]: sample id and grouped mean metrics
        """
        query = query or self.query()
        ready_query = (query.add_columns(Sample.sample_id,
                                         ExonStatistic.metric,
                                         self.weighted_average)
                            .join(ExonStatistic.sample, ExonStatistic.exon)
                            .group_by(Sample.sample_id, ExonStatistic.metric)
                            .order_by(Sample.sample_id))
        results = group_by_field(ready_query)
        return results

    def region_alt(self, region_id, sample_id=None, per=None):
        """Parse region id as input for `region` method.

        Args:
            region_id (str): for example 1:1242314-1243419
            sample_id (Optional[str]): filter to only a single sample
            per (Optional[str]): report on a per exon basis ("exon")

        Returns:
            dict: weighted metrics (across samples)
        """
        chromosome, pos_str = region_id.split(':')
        start, end = [int(pos) for pos in pos_str.split('-')]
        data = self.region(chromosome, start, end, sample_id=sample_id,
                           per=per)
        return data

    def region(self, chromosome, start, end, sample_id=None, per=None):
        """Report coverage across a genomics region (of exons).

        Args:
            chromosome (str): chromosome id
            start (int): start position (inclusive)
            end (int): end position (inclusive)
            sample_id (Optional[str]): filter to only a single sample
            per (Optional[str]): report on a per exon basis ("exon")

        Returns:
            dict: weighted metrics (across samples)
        """
        results = (self.query(Exon.exon_id, ExonStatistic.metric,
                              self.weighted_average)
                       .join(ExonStatistic.exon)
                       .filter(Exon.chromosome == chromosome,
                               Exon.start >= start,
                               Exon.end <= end)
                       .group_by(ExonStatistic.metric))

        if sample_id:
            results = (results.join(ExonStatistic.sample)
                              .filter(Sample.sample_id == sample_id))

        if per == 'exon':
            results = results.group_by(Exon.exon_id)
            exon_groups = group_by_field(results, name='exon_id')
            return exon_groups
        else:
            data = {metric: value for exon_id, metric, value in results}
            return data

    def gene(self, *gene_ids):
        r"""Report aggregate statistics for particular genes.

        Args:
            \*gene_ids (List[str]): gene ids

        Returns:
            List[dict]: metrics grouped by sample and gene
        """
        samples = {}
        for gene_id in gene_ids:
            logger.debug('figure out which transcripts the gene belongs to')
            exon_objs = self.gene_exons(gene_id).all()
            if len(exon_objs) == 0:
                raise AttributeError("gene id not in database: {}"
                                     .format(gene_id))
            exon_ids = [exon_obj.exon_id for exon_obj in exon_objs]
            query = (self.query().filter(Exon.exon_id.in_(exon_ids)))
            for sample_id, data in self.means(query):
                if sample_id not in samples:
                    samples[sample_id] = {'sample_id': sample_id, 'genes': {}}
                samples[sample_id]['genes'][gene_id] = data

        return itervalues(samples)

    def transcripts(self, *transcript_ids):
        r"""Report metrics on specific transcripts.

        Args:
            \*transcript_ids (List[str]): transcript ids

        Returns:
            List[tuple]: sample, transcript, metric, value
        """
        results = (self.query(Sample.sample_id,
                              Transcript.transcript_id,
                              ExonStatistic.metric,
                              self.weighted_average)
                       .join(ExonStatistic.sample, ExonStatistic.exon,
                             Exon.transcripts)
                       .filter(Transcript.transcript_id.in_(transcript_ids))
                       .group_by(Sample.sample_id, ExonStatistic.metric,
                                 Transcript.transcript_id))
        return results

    def incomplete_exons(self, query=None, level=10, threshold=100):
        """Fetch exons with incomplete coverage at a specifc level.

        Args:
            level (Optional[int]): completeness level, default: 10
            threshold (Optional[int]): % completeness cutoff, default: 100
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            List[tuple]: sample_id, dict with exons and completeness
        """
        completeness = "completeness_{}".format(level)
        query = query or self.query()
        ready_query = (self.query(Sample.sample_id,
                                  Exon.exon_id,
                                  ExonStatistic.value)
                           .join(ExonStatistic.exon, ExonStatistic.sample)
                           .filter(ExonStatistic.metric == completeness,
                                   ExonStatistic.value < threshold)
                           .order_by(Sample.sample_id))
        exon_samples = group_by_field(ready_query)
        return exon_samples

    def sex_coverage(self, sex_chromosomes=('X', 'Y')):
        """Query for average on X/Y chromsomes.

        Args:
            sex_chromosomes (Optional[tuple]): chromosome ids

        Returns:
            List[tuple]: sample, chromosome, weighted average coverage
        """
        query = (self.query(Sample.sample_id,
                            Exon.chromosome,
                            self.weighted_average)
                     .join(ExonStatistic.exon, ExonStatistic.sample)
                     .filter(Exon.chromosome.in_(sex_chromosomes),
                             ExonStatistic.metric == 'mean_coverage')
                     .group_by(Sample.sample_id, Exon.chromosome))
        return query

    def sex_check(self, group_id=None, sample_ids=None):
        """Predict gender based on coverage of sex chromosomes.

        Args:
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            tuple: sample, gender, coverage X, coverage, Y
        """
        query = self.sex_coverage()
        query = filter_samples(query, group_id=group_id, sample_ids=sample_ids)
        logger.debug('group rows based on sample')
        samples = itertools.groupby(query, lambda row: row[0])
        for sample_id, chromosomes in samples:
            sex_coverage = [coverage for _, _, coverage in chromosomes]
            logger.debug('predict gender')
            # run the predictor
            gender = predict_gender(*sex_coverage)
            yield sample_id, gender, sex_coverage[0], sex_coverage[1]

    def diagnostic_yield(self, sample_id, query=None, level=10, threshold=100):
        """Calculate transcripts that aren't completely covered."""
        query = query or self.query()
        level = "completeness_{}".format(level)

        all_query = (query.add_columns(Transcript.transcript_id)
                          .join(Exon.transcripts))
        all_count = all_query.count()

        logger.debug('find out which exons failed')
        yield_query = (query.add_columns(Exon.exon_id)
                            .join(ExonStatistic.exon)
                            .group_by(ExonStatistic.metric)
                            .filter(ExonStatistic.metric == level,
                                    ExonStatistic.value < threshold))
        exon_ids = [row[0] for row in yield_query]
        transcripts = self.exon_transcripts(exon_ids)

        tx_count = transcripts.count()
        diagnostic_yield = 100 - (tx_count/all_count * 100)
        return {
            "diagnostic_yield": diagnostic_yield,
            "count": tx_count,
            "total_count": all_count,
            "transcripts": transcripts
        }

    def completeness_levels(self):
        """Return a list of included completeness levels."""
        metrics = (self.query(ExonStatistic.metric)
                       .distinct(ExonStatistic.metric))
        # for all completeness levels, extract the level as int and full name
        levels = ((int(field[0].split('_')[-1]), field[0])
                  for field in metrics
                  if field[0].startswith('completeness'))
        # sort them based on the level int
        sorted_levels = sorted(levels, key=lambda level: list_get(level, 0))
        return sorted_levels
<commit_msg>fix bug in diag. yield<commit_after># -*- coding: utf-8 -*-
from __future__ import division
import itertools
import logging

from sqlalchemy.sql import func

from chanjo.compat import itervalues
from chanjo.utils import list_get
from .core import Store
from .converter import ChanjoConverterMixin
from .models import Exon, ExonStatistic, Sample, Transcript
from .utils import filter_samples, group_by_field, predict_gender

logger = logging.getLogger(__name__)

class ChanjoAPI(Store, ChanjoConverterMixin):

    """Interface to chanjo database for asking interesting questions.

    Attributes:
        weighten_average (BinaryExpression): weighter mean for metrics
    """

    def init_app(self, app, key_base='CHANJO_'):
        """Configure API (Flask style) after lazy initialization.

        Args:
            app (Flask): Flask app instance
            key_base (str): namespace to look for under ``app.config``

        Returns:
            ChanjoAPI: ``self``
        """
        uri = app.config["{}URI".format(key_base)]
        self.connect(db_uri=uri)
        return self

    @property
    def weighted_average(self):
        """Build a dynamic field to calculate a weighted average.

        Works for any of the exon metrics.

        Returns:
            BinaryExpression: dynamic field for a query
        """
        # length of the exon
        weight = Exon.end - Exon.start
        total_weight = func.sum(weight)
        total_value = func.sum(ExonStatistic.value * weight)
        weighted_mean = (total_value / total_weight).label('weighted_mean')
        return weighted_mean

    def samples(self, group_id=None, sample_ids=None):
        """Get samples from the database.

        Args:
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            List[Sample]: samples objects from the database
        """
        query = self.query(Sample).order_by(Sample.sample_id)
        query = filter_samples(query, group_id=group_id, sample_ids=sample_ids)
        return query

    def means(self, query=None):
        """Calculate means on a per sample basis.

        Args:
            query (Optional[Query]): initialized and filtered query

        Returns:
            List[tuple]: sample id and grouped mean metrics
        """
        query = query or self.query()
        ready_query = (query.add_columns(Sample.sample_id,
                                         ExonStatistic.metric,
                                         self.weighted_average)
                            .join(ExonStatistic.sample, ExonStatistic.exon)
                            .group_by(Sample.sample_id, ExonStatistic.metric)
                            .order_by(Sample.sample_id))
        results = group_by_field(ready_query)
        return results

    def region_alt(self, region_id, sample_id=None, per=None):
        """Parse region id as input for `region` method.

        Args:
            region_id (str): for example 1:1242314-1243419
            sample_id (Optional[str]): filter to only a single sample
            per (Optional[str]): report on a per exon basis ("exon")

        Returns:
            dict: weighted metrics (across samples)
        """
        chromosome, pos_str = region_id.split(':')
        start, end = [int(pos) for pos in pos_str.split('-')]
        data = self.region(chromosome, start, end, sample_id=sample_id,
                           per=per)
        return data

    def region(self, chromosome, start, end, sample_id=None, per=None):
        """Report coverage across a genomics region (of exons).

        Args:
            chromosome (str): chromosome id
            start (int): start position (inclusive)
            end (int): end position (inclusive)
            sample_id (Optional[str]): filter to only a single sample
            per (Optional[str]): report on a per exon basis ("exon")

        Returns:
            dict: weighted metrics (across samples)
        """
        results = (self.query(Exon.exon_id, ExonStatistic.metric,
                              self.weighted_average)
                       .join(ExonStatistic.exon)
                       .filter(Exon.chromosome == chromosome,
                               Exon.start >= start,
                               Exon.end <= end)
                       .group_by(ExonStatistic.metric))

        if sample_id:
            results = (results.join(ExonStatistic.sample)
                              .filter(Sample.sample_id == sample_id))

        if per == 'exon':
            results = results.group_by(Exon.exon_id)
            exon_groups = group_by_field(results, name='exon_id')
            return exon_groups
        else:
            data = {metric: value for exon_id, metric, value in results}
            return data

    def gene(self, *gene_ids):
        r"""Report aggregate statistics for particular genes.

        Args:
            \*gene_ids (List[str]): gene ids

        Returns:
            List[dict]: metrics grouped by sample and gene
        """
        samples = {}
        for gene_id in gene_ids:
            logger.debug('figure out which transcripts the gene belongs to')
            exon_objs = self.gene_exons(gene_id).all()
            if len(exon_objs) == 0:
                raise AttributeError("gene id not in database: {}"
                                     .format(gene_id))
            exon_ids = [exon_obj.exon_id for exon_obj in exon_objs]
            query = (self.query().filter(Exon.exon_id.in_(exon_ids)))
            for sample_id, data in self.means(query):
                if sample_id not in samples:
                    samples[sample_id] = {'sample_id': sample_id, 'genes': {}}
                samples[sample_id]['genes'][gene_id] = data

        return itervalues(samples)

    def transcripts(self, *transcript_ids):
        r"""Report metrics on specific transcripts.

        Args:
            \*transcript_ids (List[str]): transcript ids

        Returns:
            List[tuple]: sample, transcript, metric, value
        """
        results = (self.query(Sample.sample_id,
                              Transcript.transcript_id,
                              ExonStatistic.metric,
                              self.weighted_average)
                       .join(ExonStatistic.sample, ExonStatistic.exon,
                             Exon.transcripts)
                       .filter(Transcript.transcript_id.in_(transcript_ids))
                       .group_by(Sample.sample_id, ExonStatistic.metric,
                                 Transcript.transcript_id))
        return results

    def incomplete_exons(self, query=None, level=10, threshold=100):
        """Fetch exons with incomplete coverage at a specifc level.

        Args:
            level (Optional[int]): completeness level, default: 10
            threshold (Optional[int]): % completeness cutoff, default: 100
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            List[tuple]: sample_id, dict with exons and completeness
        """
        completeness = "completeness_{}".format(level)
        query = query or self.query()
        ready_query = (self.query(Sample.sample_id,
                                  Exon.exon_id,
                                  ExonStatistic.value)
                           .join(ExonStatistic.exon, ExonStatistic.sample)
                           .filter(ExonStatistic.metric == completeness,
                                   ExonStatistic.value < threshold)
                           .order_by(Sample.sample_id))
        exon_samples = group_by_field(ready_query)
        return exon_samples

    def sex_coverage(self, sex_chromosomes=('X', 'Y')):
        """Query for average on X/Y chromsomes.

        Args:
            sex_chromosomes (Optional[tuple]): chromosome ids

        Returns:
            List[tuple]: sample, chromosome, weighted average coverage
        """
        query = (self.query(Sample.sample_id,
                            Exon.chromosome,
                            self.weighted_average)
                     .join(ExonStatistic.exon, ExonStatistic.sample)
                     .filter(Exon.chromosome.in_(sex_chromosomes),
                             ExonStatistic.metric == 'mean_coverage')
                     .group_by(Sample.sample_id, Exon.chromosome))
        return query

    def sex_check(self, group_id=None, sample_ids=None):
        """Predict gender based on coverage of sex chromosomes.

        Args:
            group_id (Optional[str]): sample group identifier
            sample_ids (Optional[List[str]]): sample ids

        Returns:
            tuple: sample, gender, coverage X, coverage, Y
        """
        query = self.sex_coverage()
        query = filter_samples(query, group_id=group_id, sample_ids=sample_ids)
        logger.debug('group rows based on sample')
        samples = itertools.groupby(query, lambda row: row[0])
        for sample_id, chromosomes in samples:
            sex_coverage = [coverage for _, _, coverage in chromosomes]
            logger.debug('predict gender')
            # run the predictor
            gender = predict_gender(*sex_coverage)
            yield sample_id, gender, sex_coverage[0], sex_coverage[1]

    def diagnostic_yield(self, sample_id, query=None, level=10, threshold=100):
        """Calculate transcripts that aren't completely covered."""
        query = query or self.query()
        level = "completeness_{}".format(level)

        all_query = (query.add_columns(Transcript.transcript_id)
                          .distinct(Transcript.transcript_id)
                          .join(Exon.transcripts))
        all_count = all_query.count()

        logger.debug('find out which exons failed')
        yield_query = (query.add_columns(Exon.exon_id)
                            .join(ExonStatistic.exon)
                            .filter(ExonStatistic.metric == level,
                                    ExonStatistic.value < threshold))
        exon_ids = [row[0] for row in yield_query]
        transcripts = self.exon_transcripts(exon_ids)

        tx_count = transcripts.count()
        diagnostic_yield = 100 - (tx_count/all_count * 100)
        return {
            "diagnostic_yield": diagnostic_yield,
            "count": tx_count,
            "total_count": all_count,
            "transcripts": transcripts
        }

    def completeness_levels(self):
        """Return a list of included completeness levels."""
        metrics = (self.query(ExonStatistic.metric)
                       .distinct(ExonStatistic.metric))
        # for all completeness levels, extract the level as int and full name
        levels = ((int(field[0].split('_')[-1]), field[0])
                  for field in metrics
                  if field[0].startswith('completeness'))
        # sort them based on the level int
        sorted_levels = sorted(levels, key=lambda level: list_get(level, 0))
        return sorted_levels
<|endoftext|>