<commit_before>"""Partial dependence plots for tree ensembles. """

# Authors: Peter Prettenhofer
# License: BSD 3 clause

from itertools import count

from sklearn.externals.six.moves import zip

import numbers

import numpy as np

from scipy.stats.mstats import mquantiles

from ..utils.extmath import cartesian
from ..externals.joblib import Parallel, delayed
from ..externals import six
from ..externals.six.moves import xrange
from ..utils import array2d
from ..tree._tree import DTYPE

from ._gradient_boosting import _partial_dependence_tree
from .gradient_boosting import BaseGradientBoosting

def _grid_from_X(X, percentiles=(0.05, 0.95), grid_resolution=100):
    """Generate a grid of points based on the ``percentiles of ``X``.

    The grid is generated by placing ``grid_resolution`` equally
    spaced points between the ``percentiles`` of each column
    of ``X``.

    Parameters
    ----------
    X : ndarray
        The data
    percentiles : tuple of floats
        The percentiles which are used to construct the extreme
        values of the grid axes.
    grid_resolution : int
        The number of equally spaced points that are placed
        on the grid.

    Returns
    -------
    grid : ndarray
        All data points on the grid; ``grid.shape[1] == X.shape[1]``
        and ``grid.shape[0] == grid_resolution * X.shape[1]``.
    axes : seq of ndarray
        The axes with which the grid has been created.
    """
    if len(percentiles) != 2:
        raise ValueError('percentile must be tuple of len 2')
    if not all(map(lambda x: 0.0 <= x <= 1.0, percentiles)):
        raise ValueError('percentile values must be in [0, 1]')

    axes = []
    for col in range(X.shape[1]):
        uniques = np.unique(X[:, col])
        if uniques.shape[0] < grid_resolution:
            # feature has low resolution use unique vals
            axis = uniques
        else:
            emp_percentiles = mquantiles(X, prob=percentiles, axis=0)
            # create axis based on percentiles and grid resolution
            axis = np.linspace(emp_percentiles[0, col],
                               emp_percentiles[1, col],
                               num=grid_resolution, endpoint=True)
        axes.append(axis)

    return cartesian(axes), axes

def partial_dependence(gbrt, target_variables, grid=None, X=None,
                       percentiles=(0.05, 0.95), grid_resolution=100):
    """Partial dependence of ``target_variables``.

    Partial dependence plots show the dependence between the joint values
    of the ``target_variables`` and the function represented
    by the ``gbrt``.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    target_variables : array-like, dtype=int
        The target features for which the partial dependecy should be
        computed (size should be smaller than 3 for visual renderings).
    grid : array-like, shape=(n_points, len(target_variables))
        The grid of ``target_variables`` values for which the
        partial dependecy should be evaluated (either ``grid`` or ``X``
        must be specified).
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained. It is used to generate
        a ``grid`` for the ``target_variables``. The ``grid`` comprises
        ``grid_resolution`` equally spaced points between the two
        ``percentiles``.
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used create the extreme values
        for the ``grid``. Only if ``X`` is not None.
    grid_resolution : int, default=100
        The number of equally spaced points on the ``grid``.

    Returns
    -------
    pdp : array, shape=(n_classes, n_points)
        The partial dependence function evaluated on the ``grid``.
        For regression and binary classification ``n_classes==1``.
    axes : seq of ndarray or None
        The axes with which the grid has been created or None if
        the grid has been given.

    Examples
    --------
    >>> samples = [[0, 0, 2], [1, 0, 0]]
    >>> labels = [0, 1]
    >>> from sklearn.ensemble import GradientBoostingClassifier
    >>> gb = GradientBoostingClassifier().fit(samples, labels)
    >>> kwargs = dict(X=samples, percentiles=(0, 1), grid_resolution=2)
    >>> partial_dependence(gb, [0], **kwargs)
    (array([[-10.72892297,  10.72892297]]), [array([ 0.,  1.])])
    """
    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)
    if (grid is None and X is None) or (grid is not None and X is not None):
        raise ValueError('Either grid or X must be specified')

    target_variables = np.asarray(target_variables, dtype=np.int32,
                                  order='C').ravel()

    if any([not (0 <= fx < gbrt.n_features) for fx in target_variables]):
        raise ValueError('target_variables must be in [0, %d]'
                         % (gbrt.n_features - 1))

    if X is not None:
        X = array2d(X, dtype=DTYPE, order='C')
        grid, axes = _grid_from_X(X[:, target_variables], percentiles,
                                  grid_resolution)
    else:
        assert grid is not None
        # dont return axes if grid is given
        axes = None
        # grid must be 2d
        if grid.ndim == 1:
            grid = grid[:, np.newaxis]
        if grid.ndim != 2:
            raise ValueError('grid must be 2d but is %dd' % grid.ndim)

    grid = np.asarray(grid, dtype=DTYPE, order='C')
    assert grid.shape[1] == target_variables.shape[0]

    n_trees_per_stage = gbrt.estimators_.shape[1]
    n_estimators = gbrt.estimators_.shape[0]
    pdp = np.zeros((n_trees_per_stage, grid.shape[0],), dtype=np.float64,
                   order='C')
    for stage in xrange(n_estimators):
        for k in range(n_trees_per_stage):
            tree = gbrt.estimators_[stage, k].tree_
            _partial_dependence_tree(tree, grid, target_variables,
                                     gbrt.learning_rate, pdp[k])

    return pdp, axes

def plot_partial_dependence(gbrt, X, features, feature_names=None,
                            label=None, n_cols=3, grid_resolution=100,
                            percentiles=(0.05, 0.95), n_jobs=1,
                            verbose=0, ax=None, line_kw=None,
                            contour_kw=None, **fig_kw):
    """Partial dependence plots for ``features``.

    The ``len(features)`` plots are arranged in a grid with ``n_cols``
    columns. Two-way partial dependence plots are plotted as contour
    plots.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained.
    features : seq of tuples or ints
        If seq[i] is an int or a tuple with one int value, a one-way
        PDP is created; if seq[i] is a tuple of two ints, a two-way
        PDP is created.
    feature_names : seq of str
        Name of each feature; feature_names[i] holds
        the name of the feature with index i.
    label : object
        The class label for which the PDPs should be computed.
        Only if gbrt is a multi-class model. Must be in ``gbrt.classes_``.
    n_cols : int
        The number of columns in the grid plot (default: 3).
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used create the extreme values
        for the PDP axes.
    grid_resolution : int, default=100
        The number of equally spaced points on the axes.
    n_jobs : int
        The number of CPUs to use to compute the PDs. -1 means 'all CPUs'.
        Defaults to 1.
    verbose : int
        Verbose output during PD computations. Defaults to 0.
    ax : Matplotlib axis object, default None
        An axis object onto which the plots will be drawn.
    line_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For one-way partial dependence plots.
    contour_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For two-way partial dependence plots.
    fig_kw : dict
        Dict with keywords passed to the figure() call.
        Note that all keywords not recognized above will be automatically
        included here.

    Returns
    -------
    fig : figure
        The Matplotlib Figure object.
    axs : seq of Axis objects
        A seq of Axis objects, one for each subplot.

    Examples
    --------
    >>> from sklearn.datasets import make_friedman1
    >>> from sklearn.ensemble import GradientBoostingRegressor
    >>> X, y = make_friedman1()
    >>> clf = GradientBoostingRegressor(n_estimators=10).fit(X, y)
    >>> fig, axs = plot_partial_dependence(clf, X, [0, (0, 1)]) #doctest: +SKIP
    ...
    """
    import matplotlib.pyplot as plt
    from matplotlib import transforms
    from matplotlib.ticker import MaxNLocator
    from matplotlib.ticker import ScalarFormatter

    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)

    # set label_idx for multi-class GBRT
    if hasattr(gbrt, 'classes_') and np.size(gbrt.classes_) > 2:
        if label is None:
            raise ValueError('label is not given for multi-class PDP')
        label_idx = np.searchsorted(gbrt.classes_, label)
        if gbrt.classes_[label_idx] != label:
            raise ValueError('label %s not in ``gbrt.classes_``' % str(label))
    else:
        # regression and binary classification
        label_idx = 0

    X = array2d(X, dtype=DTYPE, order='C')
    if gbrt.n_features != X.shape[1]:
        raise ValueError('X.shape[1] does not match gbrt.n_features')

    if line_kw is None:
        line_kw = {'color': 'green'}
    if contour_kw is None:
        contour_kw = {}

    # convert feature_names to list
    if feature_names is None:
        # if not feature_names use fx indices as name
        feature_names = map(str, range(gbrt.n_features))
    elif isinstance(feature_names, np.ndarray):
        feature_names = feature_names.tolist()

    def convert_feature(fx):
        if isinstance(fx, six.string_types):
            try:
                fx = feature_names.index(fx)
            except ValueError:
                raise ValueError('Feature %s not in feature_names' % fx)
        return fx

    # convert features into a seq of int tuples
    tmp_features = []
    for fxs in features:
        if isinstance(fxs, (numbers.Integral,) + six.string_types):
            fxs = (fxs,)
        try:
            fxs = np.array([convert_feature(fx) for fx in fxs], dtype=np.int32)
        except TypeError:
            raise ValueError('features must be either int, str, or tuple '
                             'of int/str')
        if not (1 <= np.size(fxs) <= 2):
            raise ValueError('target features must be either one or two')

        tmp_features.append(fxs)

    features = tmp_features

    names = []
    try:
        for fxs in features:
            names.append([feature_names[i] for i in fxs])
    except IndexError:
        raise ValueError('features[i] must be in [0, n_features) '
                         'but was %d' % i)

    # compute PD functions
    pd_result = Parallel(n_jobs=n_jobs, verbose=verbose)(
        delayed(partial_dependence)(gbrt, fxs, X=X,
                                    grid_resolution=grid_resolution)
        for fxs in features)

    # get global min and max values of PD grouped by plot type
    pdp_lim = {}
    for pdp, axes in pd_result:
        min_pd, max_pd = pdp[label_idx].min(), pdp[label_idx].max()
        n_fx = len(axes)
        old_min_pd, old_max_pd = pdp_lim.get(n_fx, (min_pd, max_pd))
        min_pd = min(min_pd, old_min_pd)
        max_pd = max(max_pd, old_max_pd)
        pdp_lim[n_fx] = (min_pd, max_pd)

    # create contour levels for two-way plots
    if 2 in pdp_lim:
        Z_level = np.linspace(*pdp_lim[2], num=8)

    if ax is None:
        fig = plt.figure(**fig_kw)
    else:
        fig = ax.get_figure()
        fig.clear()

    n_cols = min(n_cols, len(features))
    n_rows = int(np.ceil(len(features) / float(n_cols)))
    axs = []
    for i, fx, name, (pdp, axes) in zip(count(), features, names,
                                        pd_result):
        ax = fig.add_subplot(n_rows, n_cols, i + 1)

        if len(axes) == 1:
            ax.plot(axes[0], pdp[label_idx].ravel(), **line_kw)
        else:
            # make contour plot
            assert len(axes) == 2
            XX, YY = np.meshgrid(axes[0], axes[1])
            Z = pdp[label_idx].reshape(map(np.size, axes)).T
            CS = ax.contour(XX, YY, Z, levels=Z_level, linewidths=0.5,
                            colors='k')
            ax.contourf(XX, YY, Z, levels=Z_level, vmax=Z_level[-1],
                        vmin=Z_level[0], alpha=0.75, **contour_kw)
            ax.clabel(CS, fmt='%2.2f', colors='k', fontsize=10, inline=True)

        # plot data deciles + axes labels
        deciles = mquantiles(X[:, fx[0]], prob=np.arange(0.1, 1.0, 0.1))
        trans = transforms.blended_transform_factory(ax.transData,
                                                     ax.transAxes)
        ylim = ax.get_ylim()
        ax.vlines(deciles, [0], 0.05, transform=trans, color='k')
        ax.set_xlabel(name[0])
        ax.set_ylim(ylim)

        # prevent x-axis ticks from overlapping
        ax.xaxis.set_major_locator(MaxNLocator(nbins=6, prune='lower'))
        tick_formatter = ScalarFormatter()
        tick_formatter.set_powerlimits((-3, 4))
        ax.xaxis.set_major_formatter(tick_formatter)

        if len(axes) > 1:
            # two-way PDP - y-axis deciles + labels
            deciles = mquantiles(X[:, fx[1]], prob=np.arange(0.1, 1.0, 0.1))
            trans = transforms.blended_transform_factory(ax.transAxes,
                                                         ax.transData)
            xlim = ax.get_xlim()
            ax.hlines(deciles, [0], 0.05, transform=trans, color='k')
            ax.set_ylabel(name[1])
            # hline erases xlim
            ax.set_xlim(xlim)
        else:
            ax.set_ylabel('Partial dependence')

        if len(axes) == 1:
            ax.set_ylim(pdp_lim[1])
        axs.append(ax)

    fig.subplots_adjust(bottom=0.15, top=0.7, left=0.1, right=0.95, wspace=0.4,
                        hspace=0.3)
    return fig, axs
<commit_msg>FIX: doctests in partial_dependence<commit_after>"""Partial dependence plots for tree ensembles. """

# Authors: Peter Prettenhofer
# License: BSD 3 clause

from itertools import count

from sklearn.externals.six.moves import zip

import numbers

import numpy as np

from scipy.stats.mstats import mquantiles

from ..utils.extmath import cartesian
from ..externals.joblib import Parallel, delayed
from ..externals import six
from ..externals.six.moves import xrange
from ..utils import array2d
from ..tree._tree import DTYPE

from ._gradient_boosting import _partial_dependence_tree
from .gradient_boosting import BaseGradientBoosting

def _grid_from_X(X, percentiles=(0.05, 0.95), grid_resolution=100):
    """Generate a grid of points based on the ``percentiles of ``X``.

    The grid is generated by placing ``grid_resolution`` equally
    spaced points between the ``percentiles`` of each column
    of ``X``.

    Parameters
    ----------
    X : ndarray
        The data
    percentiles : tuple of floats
        The percentiles which are used to construct the extreme
        values of the grid axes.
    grid_resolution : int
        The number of equally spaced points that are placed
        on the grid.

    Returns
    -------
    grid : ndarray
        All data points on the grid; ``grid.shape[1] == X.shape[1]``
        and ``grid.shape[0] == grid_resolution * X.shape[1]``.
    axes : seq of ndarray
        The axes with which the grid has been created.
    """
    if len(percentiles) != 2:
        raise ValueError('percentile must be tuple of len 2')
    if not all(map(lambda x: 0.0 <= x <= 1.0, percentiles)):
        raise ValueError('percentile values must be in [0, 1]')

    axes = []
    for col in range(X.shape[1]):
        uniques = np.unique(X[:, col])
        if uniques.shape[0] < grid_resolution:
            # feature has low resolution use unique vals
            axis = uniques
        else:
            emp_percentiles = mquantiles(X, prob=percentiles, axis=0)
            # create axis based on percentiles and grid resolution
            axis = np.linspace(emp_percentiles[0, col],
                               emp_percentiles[1, col],
                               num=grid_resolution, endpoint=True)
        axes.append(axis)

    return cartesian(axes), axes

def partial_dependence(gbrt, target_variables, grid=None, X=None,
                       percentiles=(0.05, 0.95), grid_resolution=100):
    """Partial dependence of ``target_variables``.

    Partial dependence plots show the dependence between the joint values
    of the ``target_variables`` and the function represented
    by the ``gbrt``.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    target_variables : array-like, dtype=int
        The target features for which the partial dependecy should be
        computed (size should be smaller than 3 for visual renderings).
    grid : array-like, shape=(n_points, len(target_variables))
        The grid of ``target_variables`` values for which the
        partial dependecy should be evaluated (either ``grid`` or ``X``
        must be specified).
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained. It is used to generate
        a ``grid`` for the ``target_variables``. The ``grid`` comprises
        ``grid_resolution`` equally spaced points between the two
        ``percentiles``.
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used create the extreme values
        for the ``grid``. Only if ``X`` is not None.
    grid_resolution : int, default=100
        The number of equally spaced points on the ``grid``.

    Returns
    -------
    pdp : array, shape=(n_classes, n_points)
        The partial dependence function evaluated on the ``grid``.
        For regression and binary classification ``n_classes==1``.
    axes : seq of ndarray or None
        The axes with which the grid has been created or None if
        the grid has been given.

    Examples
    --------
    >>> samples = [[0, 0, 2], [1, 0, 0]]
    >>> labels = [0, 1]
    >>> from sklearn.ensemble import GradientBoostingClassifier
    >>> gb = GradientBoostingClassifier(random_state=0).fit(samples, labels)
    >>> kwargs = dict(X=samples, percentiles=(0, 1), grid_resolution=2)
    >>> partial_dependence(gb, [0], **kwargs)
    (array([[-5.67647953,  5.67647953]]), [array([ 0.,  1.])])
    """
    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)
    if (grid is None and X is None) or (grid is not None and X is not None):
        raise ValueError('Either grid or X must be specified')

    target_variables = np.asarray(target_variables, dtype=np.int32,
                                  order='C').ravel()

    if any([not (0 <= fx < gbrt.n_features) for fx in target_variables]):
        raise ValueError('target_variables must be in [0, %d]'
                         % (gbrt.n_features - 1))

    if X is not None:
        X = array2d(X, dtype=DTYPE, order='C')
        grid, axes = _grid_from_X(X[:, target_variables], percentiles,
                                  grid_resolution)
    else:
        assert grid is not None
        # dont return axes if grid is given
        axes = None
        # grid must be 2d
        if grid.ndim == 1:
            grid = grid[:, np.newaxis]
        if grid.ndim != 2:
            raise ValueError('grid must be 2d but is %dd' % grid.ndim)

    grid = np.asarray(grid, dtype=DTYPE, order='C')
    assert grid.shape[1] == target_variables.shape[0]

    n_trees_per_stage = gbrt.estimators_.shape[1]
    n_estimators = gbrt.estimators_.shape[0]
    pdp = np.zeros((n_trees_per_stage, grid.shape[0],), dtype=np.float64,
                   order='C')
    for stage in xrange(n_estimators):
        for k in range(n_trees_per_stage):
            tree = gbrt.estimators_[stage, k].tree_
            _partial_dependence_tree(tree, grid, target_variables,
                                     gbrt.learning_rate, pdp[k])

    return pdp, axes

def plot_partial_dependence(gbrt, X, features, feature_names=None,
                            label=None, n_cols=3, grid_resolution=100,
                            percentiles=(0.05, 0.95), n_jobs=1,
                            verbose=0, ax=None, line_kw=None,
                            contour_kw=None, **fig_kw):
    """Partial dependence plots for ``features``.

    The ``len(features)`` plots are arranged in a grid with ``n_cols``
    columns. Two-way partial dependence plots are plotted as contour
    plots.

    Parameters
    ----------
    gbrt : BaseGradientBoosting
        A fitted gradient boosting model.
    X : array-like, shape=(n_samples, n_features)
        The data on which ``gbrt`` was trained.
    features : seq of tuples or ints
        If seq[i] is an int or a tuple with one int value, a one-way
        PDP is created; if seq[i] is a tuple of two ints, a two-way
        PDP is created.
    feature_names : seq of str
        Name of each feature; feature_names[i] holds
        the name of the feature with index i.
    label : object
        The class label for which the PDPs should be computed.
        Only if gbrt is a multi-class model. Must be in ``gbrt.classes_``.
    n_cols : int
        The number of columns in the grid plot (default: 3).
    percentiles : (low, high), default=(0.05, 0.95)
        The lower and upper percentile used create the extreme values
        for the PDP axes.
    grid_resolution : int, default=100
        The number of equally spaced points on the axes.
    n_jobs : int
        The number of CPUs to use to compute the PDs. -1 means 'all CPUs'.
        Defaults to 1.
    verbose : int
        Verbose output during PD computations. Defaults to 0.
    ax : Matplotlib axis object, default None
        An axis object onto which the plots will be drawn.
    line_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For one-way partial dependence plots.
    contour_kw : dict
        Dict with keywords passed to the ``pylab.plot`` call.
        For two-way partial dependence plots.
    fig_kw : dict
        Dict with keywords passed to the figure() call.
        Note that all keywords not recognized above will be automatically
        included here.

    Returns
    -------
    fig : figure
        The Matplotlib Figure object.
    axs : seq of Axis objects
        A seq of Axis objects, one for each subplot.

    Examples
    --------
    >>> from sklearn.datasets import make_friedman1
    >>> from sklearn.ensemble import GradientBoostingRegressor
    >>> X, y = make_friedman1()
    >>> clf = GradientBoostingRegressor(n_estimators=10).fit(X, y)
    >>> fig, axs = plot_partial_dependence(clf, X, [0, (0, 1)]) #doctest: +SKIP
    ...
    """
    import matplotlib.pyplot as plt
    from matplotlib import transforms
    from matplotlib.ticker import MaxNLocator
    from matplotlib.ticker import ScalarFormatter

    if not isinstance(gbrt, BaseGradientBoosting):
        raise ValueError('gbrt has to be an instance of BaseGradientBoosting')
    if gbrt.estimators_.shape[0] == 0:
        raise ValueError('Call %s.fit before partial_dependence' %
                         gbrt.__class__.__name__)

    # set label_idx for multi-class GBRT
    if hasattr(gbrt, 'classes_') and np.size(gbrt.classes_) > 2:
        if label is None:
            raise ValueError('label is not given for multi-class PDP')
        label_idx = np.searchsorted(gbrt.classes_, label)
        if gbrt.classes_[label_idx] != label:
            raise ValueError('label %s not in ``gbrt.classes_``' % str(label))
    else:
        # regression and binary classification
        label_idx = 0

    X = array2d(X, dtype=DTYPE, order='C')
    if gbrt.n_features != X.shape[1]:
        raise ValueError('X.shape[1] does not match gbrt.n_features')

    if line_kw is None:
        line_kw = {'color': 'green'}
    if contour_kw is None:
        contour_kw = {}

    # convert feature_names to list
    if feature_names is None:
        # if not feature_names use fx indices as name
        feature_names = map(str, range(gbrt.n_features))
    elif isinstance(feature_names, np.ndarray):
        feature_names = feature_names.tolist()

    def convert_feature(fx):
        if isinstance(fx, six.string_types):
            try:
                fx = feature_names.index(fx)
            except ValueError:
                raise ValueError('Feature %s not in feature_names' % fx)
        return fx

    # convert features into a seq of int tuples
    tmp_features = []
    for fxs in features:
        if isinstance(fxs, (numbers.Integral,) + six.string_types):
            fxs = (fxs,)
        try:
            fxs = np.array([convert_feature(fx) for fx in fxs], dtype=np.int32)
        except TypeError:
            raise ValueError('features must be either int, str, or tuple '
                             'of int/str')
        if not (1 <= np.size(fxs) <= 2):
            raise ValueError('target features must be either one or two')

        tmp_features.append(fxs)

    features = tmp_features

    names = []
    try:
        for fxs in features:
            names.append([feature_names[i] for i in fxs])
    except IndexError:
        raise ValueError('features[i] must be in [0, n_features) '
                         'but was %d' % i)

    # compute PD functions
    pd_result = Parallel(n_jobs=n_jobs, verbose=verbose)(
        delayed(partial_dependence)(gbrt, fxs, X=X,
                                    grid_resolution=grid_resolution)
        for fxs in features)

    # get global min and max values of PD grouped by plot type
    pdp_lim = {}
    for pdp, axes in pd_result:
        min_pd, max_pd = pdp[label_idx].min(), pdp[label_idx].max()
        n_fx = len(axes)
        old_min_pd, old_max_pd = pdp_lim.get(n_fx, (min_pd, max_pd))
        min_pd = min(min_pd, old_min_pd)
        max_pd = max(max_pd, old_max_pd)
        pdp_lim[n_fx] = (min_pd, max_pd)

    # create contour levels for two-way plots
    if 2 in pdp_lim:
        Z_level = np.linspace(*pdp_lim[2], num=8)

    if ax is None:
        fig = plt.figure(**fig_kw)
    else:
        fig = ax.get_figure()
        fig.clear()

    n_cols = min(n_cols, len(features))
    n_rows = int(np.ceil(len(features) / float(n_cols)))
    axs = []
    for i, fx, name, (pdp, axes) in zip(count(), features, names,
                                        pd_result):
        ax = fig.add_subplot(n_rows, n_cols, i + 1)

        if len(axes) == 1:
            ax.plot(axes[0], pdp[label_idx].ravel(), **line_kw)
        else:
            # make contour plot
            assert len(axes) == 2
            XX, YY = np.meshgrid(axes[0], axes[1])
            Z = pdp[label_idx].reshape(map(np.size, axes)).T
            CS = ax.contour(XX, YY, Z, levels=Z_level, linewidths=0.5,
                            colors='k')
            ax.contourf(XX, YY, Z, levels=Z_level, vmax=Z_level[-1],
                        vmin=Z_level[0], alpha=0.75, **contour_kw)
            ax.clabel(CS, fmt='%2.2f', colors='k', fontsize=10, inline=True)

        # plot data deciles + axes labels
        deciles = mquantiles(X[:, fx[0]], prob=np.arange(0.1, 1.0, 0.1))
        trans = transforms.blended_transform_factory(ax.transData,
                                                     ax.transAxes)
        ylim = ax.get_ylim()
        ax.vlines(deciles, [0], 0.05, transform=trans, color='k')
        ax.set_xlabel(name[0])
        ax.set_ylim(ylim)

        # prevent x-axis ticks from overlapping
        ax.xaxis.set_major_locator(MaxNLocator(nbins=6, prune='lower'))
        tick_formatter = ScalarFormatter()
        tick_formatter.set_powerlimits((-3, 4))
        ax.xaxis.set_major_formatter(tick_formatter)

        if len(axes) > 1:
            # two-way PDP - y-axis deciles + labels
            deciles = mquantiles(X[:, fx[1]], prob=np.arange(0.1, 1.0, 0.1))
            trans = transforms.blended_transform_factory(ax.transAxes,
                                                         ax.transData)
            xlim = ax.get_xlim()
            ax.hlines(deciles, [0], 0.05, transform=trans, color='k')
            ax.set_ylabel(name[1])
            # hline erases xlim
            ax.set_xlim(xlim)
        else:
            ax.set_ylabel('Partial dependence')

        if len(axes) == 1:
            ax.set_ylim(pdp_lim[1])
        axs.append(ax)

    fig.subplots_adjust(bottom=0.15, top=0.7, left=0.1, right=0.95, wspace=0.4,
                        hspace=0.3)
    return fig, axs
<|endoftext|>


<commit_before>"""Tools for handling ion mutation."""

import itertools
import json
import os
import re
from copy import deepcopy
from operator import itemgetter
from typing import Callable, Generator, Optional, Tuple

import numpy as np
import pandas as pd
import pymatgen.analysis.structure_prediction as pymatgen_sp
from .utilities import parse_spec
from .structure import SmactStructure

class CationMutator:
    """Handles cation mutation of SmactStructures based on substitution probability.

    Based on the algorithm presented in:
        Hautier, G., Fischer, C., Ehrlacher, V., Jain, A., and Ceder, G. (2011)
        Data Mined Ionic Substitutions for the Discovery of New Compounds.
        Inorganic Chemistry, 50(2), 656-663.
        `doi:10.1021/ic102031h <https://pubs.acs.org/doi/10.1021/ic102031h>`_

    """

    def __init__(
      self,
      lambda_df: pd.DataFrame,
      alpha: Optional[Callable[[str, str], float]] = (lambda s1, s2: -5.0),
    ):
        """Assign attributes and get lambda table."""
        self.lambda_tab = lambda_df

        self.specs = set(
          itertools.chain.from_iterable(
            set(getattr(self.lambda_tab, x)) for x in ["columns", "index"]
          )
        )

        self.alpha = alpha

        # Make sure table is fully populated
        self._populate_lambda()

        self.Z = np.exp(self.lambda_tab.to_numpy()).sum()

    @staticmethod
    def from_json(
      lambda_json: Optional[str] = None,
      alpha: Optional[Callable[[str, str], float]] = (lambda s1, s2: -5.0),
    ):
        """Create a CationMutator instance from a DataFrame."""
        if lambda_json is not None:
            with open(lambda_json, 'r') as f:
                lambda_dat = json.load(f)
        else:
            # Get pymatgen lambda table
            py_sp_dir = os.path.dirname(pymatgen_sp.__file__)
            pymatgen_lambda = os.path.join(py_sp_dir, "data", "lambda.json")
            with open(pymatgen_lambda, 'r') as f:
                lambda_dat = json.load(f)

            # Get rid of 'D1+' values to reflect pymatgen
            # implementation
            lambda_dat = [x for x in lambda_dat if 'D1+' not in x]

        # Convert lambda table to pandas DataFrame
        lambda_dat = [tuple(x) for x in lambda_dat]
        lambda_df = pd.DataFrame(lambda_dat)

        lambda_df = lambda_df.pivot(index=0, columns=1, values=2)

        return CationMutator(lambda_df, alpha)

    def _populate_lambda(self):
        """Populate lambda table.

        Ensures no values are NaN and performs alpha calculations,
        such that an entry exists for every possible species
        combination in the lambda table.
        Also ensures lambda table symmetry.

        """
        pairs = itertools.combinations_with_replacement(self.specs, 2)

        def add_alpha(s1, s2):
            """Add an alpha value to the lambda table at both coordinates."""
            a = self.alpha(s1, s2)
            self.lambda_tab.at[s1, s2] = a
            self.lambda_tab.at[s2, s1] = a

        def mirror_lambda(s1, s2):
            """Mirror the lambda value at (s2, s1) into (s1, s2)."""
            self.lambda_tab.at[s1, s2] = self.lambda_tab.at[s2, s1]

        for s1, s2 in pairs:
            try:
                if np.isnan(self.lambda_tab.at[s1, s2]):
                    try:
                        if not np.isnan(self.lambda_tab.at[s2, s1]):
                            mirror_lambda(s1, s2)
                        else:
                            add_alpha(s1, s2)
                    except KeyError:
                        add_alpha(s1, s2)
                else:
                    mirror_lambda(s2, s1)
            except KeyError:
                try:
                    if np.isnan(self.lambda_tab.at[s2, s1]):
                        add_alpha(s1, s2)
                    else:
                        mirror_lambda(s1, s2)
                except KeyError:
                    add_alpha(s1, s2)

        # Ensure symmetry
        idx = self.lambda_tab.index
        self.lambda_tab = self.lambda_tab[idx]

    def get_lambda(self, s1: str, s2: str) -> float:
        """Get lambda values corresponding to a pair of species."""
        if {s1, s2} <= self.specs:
            return self.lambda_tab.at[s1, s2]

        return self.alpha(s1, s2)

    def get_lambdas(self, species: str) -> pd.Series:
        """Get all the lambda values associated with a species."""
        if not {species} <= self.specs:
            raise ValueError(f"{species} not in lambda table.")

        return self.lambda_tab.loc[species]

    @staticmethod
    def _mutate_structure(
      structure: SmactStructure,
      init_species: str,
      final_species: str, ) -> SmactStructure:
        """Mutate an ion within a SmactStructure."""
        struct_buff = deepcopy(structure)
        init_spec_tup = parse_spec(init_species)
        struct_spec_tups = list(map(itemgetter(0, 1), struct_buff.species))
        spec_loc = struct_spec_tups.index(init_spec_tup)

        final_spec_tup = parse_spec(final_species)

        # Replace species tuple
        struct_buff.species[spec_loc] = (*final_spec_tup, struct_buff.species[spec_loc][2])

        # Check for charge neutrality
        if sum(x[1] * x[2] for x in struct_buff.species) != 0:
            raise ValueError("New structure is not charge neutral.")

        # Sort species again
        struct_buff.species.sort(key=itemgetter(1), reverse=True)
        struct_buff.species.sort(key=itemgetter(0))

        # Replace sites
        struct_buff.sites[final_species] = struct_buff.sites.pop(init_species)
        # And sort
        species_strs = struct_buff._format_style("{ele}{charge}{sign}").split(" ")
        struct_buff.sites = {spec: struct_buff.sites[spec] for spec in species_strs}

        return struct_buff

    def sub_prob(self, s1: str, s2: str) -> float:
        """Calculate the probability of substitution of two species."""
        return np.exp(self.get_lambda(s1, s2)) / self.Z

    def sub_probs(self, s1: str) -> pd.Series:
        """Determine the substitution probabilities of a species with others.

        Determines the probability of substitution of the species with every
        species in the lambda table.

        """
        probs = self.get_lambdas(s1)
        probs = np.exp(probs)
        probs /= self.Z
        return probs

    def complete_sub_probs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the substitution probabilities."""
        return np.exp(self.lambda_tab) / self.Z

    def complete_cond_probs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the conditional substitution probabilities."""
        lambda_exp = np.exp(self.lambda_tab)
        return lambda_exp / lambda_exp.sum(axis=1)

    def complete_pair_corrs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the pair correlations."""
        corr = self.complete_sub_probs()

        # Sum each row (symmetry means this is the same as column sums)
        sums = corr.sum(axis=0)
        # Stack into matrix
        mat_sums = np.vstack([sums] * len(sums))
        # Make each element the product of (row_sum * col_sum)
        mat_sums *= sums[:, None]

        corr /= mat_sums

        return corr

    def same_spec_probs(self) -> pd.Series:
        """Calculate the same species substiution probabilities."""
        return np.exp(
          pd.Series(
            np.diag(self.lambda_tab),
            index=[self.lambda_tab.index, self.lambda_tab.columns], )
        ) / self.Z

    def same_spec_cond_probs(self) -> pd.Series:
        """Calculate the same species conditional substiution probabilities."""
        return (
          np.exp(self.lambda_tab.to_numpy().diagonal()) / np.exp(self.lambda_tab).sum(axis=0)
        )

    def pair_corr(self, s1: str, s2: str) -> float:
        """Determine the pair correlation of two ionic species."""
        corr = self.sub_prob(s1, s2)
        corr /= self.sub_probs(s1).sum()
        corr /= self.sub_probs(s2).sum()
        return corr

    def cond_sub_prob(self, s1: str, s2: str) -> float:
        """Calculate the probability of substitution of one species with another."""
        return np.exp(self.get_lambda(s1, s2)) / np.exp(self.get_lambdas(s2)).sum()

    def cond_sub_probs(self, s1: str) -> pd.Series:
        """Calculate the probabilities of substitution of a given species.

        Calculates probabilities of substitution of given species with all
        others in the lambda table.

        """
        probs = self.get_lambdas(s1)
        probs = np.exp(probs)
        probs /= np.exp(self.lambda_tab).sum()
        return probs

    def unary_substitute(
      self,
      structure: SmactStructure,
      thresh: Optional[float] = 1e-5, ) -> Generator[Tuple[SmactStructure, float], None, None]:
        """Find all structures with 1 substitution with probability above a threshold.

        Args:
            structure: A :class:`SmactStructure` instance from which to generate compounds.
            thresh (float): The probability threshold; discard all substitutions that have
                a probability to generate a naturally-occuring compound less than this.

        Yields:
            Tuples of (:class:`SmactStructure`, probability).

        """
        for specie in structure.get_spec_strs():
            cond_probs = self.cond_sub_probs(specie)
            likely_probs = cond_probs.loc[cond_probs > thresh]

            for new_spec, prob in likely_probs.items():
                if any([
                  new_spec == specie,
                  parse_spec(specie)[1] != parse_spec(new_spec)[1], ]):
                    continue
                yield (self._mutate_structure(structure, specie, new_spec), prob)
<commit_msg>Docstring update: mutation.py<commit_after>"""Tools for handling ion mutation."""

import itertools
import json
import os
import re
from copy import deepcopy
from operator import itemgetter
from typing import Callable, Generator, Optional, Tuple

import numpy as np
import pandas as pd
import pymatgen.analysis.structure_prediction as pymatgen_sp
from .utilities import parse_spec
from .structure import SmactStructure

class CationMutator:
    """Handles cation mutation of SmactStructures based on substitution probability.

    Based on the algorithm presented in:
        Hautier, G., Fischer, C., Ehrlacher, V., Jain, A., and Ceder, G. (2011)
        Data Mined Ionic Substitutions for the Discovery of New Compounds.
        Inorganic Chemistry, 50(2), 656-663.
        `doi:10.1021/ic102031h <https://pubs.acs.org/doi/10.1021/ic102031h>`_

    """

    def __init__(
      self,
      lambda_df: pd.DataFrame,
      alpha: Optional[Callable[[str, str], float]] = (lambda s1, s2: -5.0),
    ):
        """Assign attributes and get lambda table.

        Args:
            lambda_df: A pandas DataFrame, with column and index labels
                as species strings and lambda values as entries.
            alpha: A function to call to fill in missing lambda values.
                The function must take the two species' strings as
                arguments, and return a floating point lambda
                value.
                Defaults to a function that unconditionally returns
                -5.0.

        """
        self.lambda_tab = lambda_df

        self.specs = set(
          itertools.chain.from_iterable(
            set(getattr(self.lambda_tab, x)) for x in ["columns", "index"]
          )
        )

        self.alpha = alpha

        # Make sure table is fully populated
        self._populate_lambda()

        self.Z = np.exp(self.lambda_tab.to_numpy()).sum()

    @staticmethod
    def from_json(
      lambda_json: Optional[str] = None,
      alpha: Optional[Callable[[str, str], float]] = (lambda s1, s2: -5.0),
    ):
        """Create a CationMutator instance from a DataFrame.

        Args:
            lambda_json (str, optional): JSON-style representation of the
                lambda table. This is a list of entries, containing pairs
                and their associated lambda values.
                Each entry is a list of [species1, species2, lambda].
                If not supplied, defaults to the lambda table
                included with pymatgen.
            alpha: See :meth:`__init__`.

        Returns:
            A :class:`CationMutator` instance.

        """
        if lambda_json is not None:
            with open(lambda_json, 'r') as f:
                lambda_dat = json.load(f)
        else:
            # Get pymatgen lambda table
            py_sp_dir = os.path.dirname(pymatgen_sp.__file__)
            pymatgen_lambda = os.path.join(py_sp_dir, "data", "lambda.json")
            with open(pymatgen_lambda, 'r') as f:
                lambda_dat = json.load(f)

            # Get rid of 'D1+' values to reflect pymatgen
            # implementation
            lambda_dat = [x for x in lambda_dat if 'D1+' not in x]

        # Convert lambda table to pandas DataFrame
        lambda_dat = [tuple(x) for x in lambda_dat]
        lambda_df = pd.DataFrame(lambda_dat)

        lambda_df = lambda_df.pivot(index=0, columns=1, values=2)

        return CationMutator(lambda_df, alpha)

    def _populate_lambda(self):
        """Populate lambda table.

        Ensures no values are NaN and performs alpha calculations,
        such that an entry exists for every possible species
        combination in the lambda table.
        Also ensures lambda table symmetry.

        """
        pairs = itertools.combinations_with_replacement(self.specs, 2)

        def add_alpha(s1, s2):
            """Add an alpha value to the lambda table at both coordinates."""
            a = self.alpha(s1, s2)
            self.lambda_tab.at[s1, s2] = a
            self.lambda_tab.at[s2, s1] = a

        def mirror_lambda(s1, s2):
            """Mirror the lambda value at (s2, s1) into (s1, s2)."""
            self.lambda_tab.at[s1, s2] = self.lambda_tab.at[s2, s1]

        for s1, s2 in pairs:
            try:
                if np.isnan(self.lambda_tab.at[s1, s2]):
                    try:
                        if not np.isnan(self.lambda_tab.at[s2, s1]):
                            mirror_lambda(s1, s2)
                        else:
                            add_alpha(s1, s2)
                    except KeyError:
                        add_alpha(s1, s2)
                else:
                    mirror_lambda(s2, s1)
            except KeyError:
                try:
                    if np.isnan(self.lambda_tab.at[s2, s1]):
                        add_alpha(s1, s2)
                    else:
                        mirror_lambda(s1, s2)
                except KeyError:
                    add_alpha(s1, s2)

        # Ensure symmetry
        idx = self.lambda_tab.index
        self.lambda_tab = self.lambda_tab[idx]

    def get_lambda(self, s1: str, s2: str) -> float:
        """Get lambda values corresponding to a pair of species.

        Args:
            s1 (str): One of the species.
            s2 (str): The other species.

        Returns:
            lambda (float): The lambda value, if it exists
                in the table. Otherwise, the alpha value
                for the two species.

        """
        if {s1, s2} <= self.specs:
            return self.lambda_tab.at[s1, s2]

        return self.alpha(s1, s2)

    def get_lambdas(self, species: str) -> pd.Series:
        """Get all the lambda values associated with a species.

        Args:
            species (str): The species for which to get the
                lambda values.

        Returns:
            A pandas Series, with index-labelled lambda entries.

        """
        if not {species} <= self.specs:
            raise ValueError(f"{species} not in lambda table.")

        return self.lambda_tab.loc[species]

    @staticmethod
    def _mutate_structure(
      structure: SmactStructure,
      init_species: str,
      final_species: str, ) -> SmactStructure:
        """Mutate a species within a SmactStructure.

        Replaces all instances of the species within the
        structure. Every site occupied by the species
        has its label changed to the new species,
        and the list of species of the structure is
        changed to reflect the mutation.
        Stoichiometry is maintained.
        Requires the species to have the same charge.

        Note:
            Creates a deepcopy of the supplied structure,
            such that the original instance is not modified.

        Args:
            init_species (str): The species within the structure
                to mutate.
            final_species (str): The species to replace the
                initial species with.

        Returns:
            A :class:`.~SmactStructure`, with the species
                mutated.

        """
        struct_buff = deepcopy(structure)
        init_spec_tup = parse_spec(init_species)
        struct_spec_tups = list(map(itemgetter(0, 1), struct_buff.species))
        spec_loc = struct_spec_tups.index(init_spec_tup)

        final_spec_tup = parse_spec(final_species)

        # Replace species tuple
        struct_buff.species[spec_loc] = (*final_spec_tup, struct_buff.species[spec_loc][2])

        # Check for charge neutrality
        if sum(x[1] * x[2] for x in struct_buff.species) != 0:
            raise ValueError("New structure is not charge neutral.")

        # Sort species again
        struct_buff.species.sort(key=itemgetter(1), reverse=True)
        struct_buff.species.sort(key=itemgetter(0))

        # Replace sites
        struct_buff.sites[final_species] = struct_buff.sites.pop(init_species)
        # And sort
        species_strs = struct_buff._format_style("{ele}{charge}{sign}").split(" ")
        struct_buff.sites = {spec: struct_buff.sites[spec] for spec in species_strs}

        return struct_buff

    def sub_prob(self, s1: str, s2: str) -> float:
        """Calculate the probability of substitution of two species."""
        return np.exp(self.get_lambda(s1, s2)) / self.Z

    def sub_probs(self, s1: str) -> pd.Series:
        """Determine the substitution probabilities of a species with others.

        Determines the probability of substitution of the species with every
        species in the lambda table.

        """
        probs = self.get_lambdas(s1)
        probs = np.exp(probs)
        probs /= self.Z
        return probs

    def complete_sub_probs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the substitution probabilities."""
        return np.exp(self.lambda_tab) / self.Z

    def complete_cond_probs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the conditional substitution probabilities."""
        lambda_exp = np.exp(self.lambda_tab)
        return lambda_exp / lambda_exp.sum(axis=1)

    def complete_pair_corrs(self) -> pd.DataFrame:
        """Generate a DataFrame with all the pair correlations."""
        corr = self.complete_sub_probs()

        # Sum each row (symmetry means this is the same as column sums)
        sums = corr.sum(axis=0)
        # Stack into matrix
        mat_sums = np.vstack([sums] * len(sums))
        # Make each element the product of (row_sum * col_sum)
        mat_sums *= sums[:, None]

        corr /= mat_sums

        return corr

    def same_spec_probs(self) -> pd.Series:
        """Calculate the same species substiution probabilities."""
        return np.exp(
          pd.Series(
            np.diag(self.lambda_tab),
            index=[self.lambda_tab.index, self.lambda_tab.columns], )
        ) / self.Z

    def same_spec_cond_probs(self) -> pd.Series:
        """Calculate the same species conditional substiution probabilities."""
        return (
          np.exp(self.lambda_tab.to_numpy().diagonal()) / np.exp(self.lambda_tab).sum(axis=0)
        )

    def pair_corr(self, s1: str, s2: str) -> float:
        """Determine the pair correlation of two ionic species."""
        corr = self.sub_prob(s1, s2)
        corr /= self.sub_probs(s1).sum()
        corr /= self.sub_probs(s2).sum()
        return corr

    def cond_sub_prob(self, s1: str, s2: str) -> float:
        """Calculate the probability of substitution of one species with another."""
        return np.exp(self.get_lambda(s1, s2)) / np.exp(self.get_lambdas(s2)).sum()

    def cond_sub_probs(self, s1: str) -> pd.Series:
        """Calculate the probabilities of substitution of a given species.

        Calculates probabilities of substitution of given species with all
        others in the lambda table.

        """
        probs = self.get_lambdas(s1)
        probs = np.exp(probs)
        probs /= np.exp(self.lambda_tab).sum()
        return probs

    def unary_substitute(
      self,
      structure: SmactStructure,
      thresh: Optional[float] = 1e-5, ) -> Generator[Tuple[SmactStructure, float], None, None]:
        """Find all structures with 1 substitution with probability above a threshold.

        Args:
            structure: A :class:`SmactStructure` instance from which to generate compounds.
            thresh (float): The probability threshold; discard all substitutions that have
                a probability to generate a naturally-occuring compound less than this.

        Yields:
            Tuples of (:class:`SmactStructure`, probability).

        """
        for specie in structure.get_spec_strs():
            cond_probs = self.cond_sub_probs(specie)
            likely_probs = cond_probs.loc[cond_probs > thresh]

            for new_spec, prob in likely_probs.items():
                if any([
                  new_spec == specie,
                  parse_spec(specie)[1] != parse_spec(new_spec)[1], ]):
                    continue
                yield (self._mutate_structure(structure, specie, new_spec), prob)
<|endoftext|>


<commit_before>from __future__ import absolute_import
import re

from django.test import TestCase
from django.core import mail

from ..email import TemplateAttachmentEmail, send_plea_email
from ..models import CourtEmailPlea, CourtEmailCount

class EmailGenerationTests(TestCase):
    def setUp(self):
        mail.outbox = []

    def test_template_attachment_sends_email(self):
        email_context = {"URN": "1A2B3C4D5E"}
        email = TemplateAttachmentEmail("test_from@example.org",
                                        "test.html",
                                        "plea/plea_email_attachment.html",
                                        email_context,
                                        "
Test Content

{{ URN }}
")

        email.send(["test_to@example.org", ],
                   "Subject line",
                   "Body Text")

        self.assertEqual(mail.outbox[0].subject, "Subject line")
        self.assertEqual(mail.outbox[0].body, "Body Text")

    def test_plea_email_sends(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89"},
                        "your_details": {"name": "vcx", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx",
                                         "email": "test@test.com"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data)

        # only expecting 1 email for now, PLP email cancelled for the time being
        self.assertEqual(len(mail.outbox), 3)

    def test_plea_email_body_contains_plea_and_count_ids(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89"},
                        "your_details": {"name": "vcx", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx",
                                         "email": "test@test.com"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data)

        plea_obj = CourtEmailPlea.objects.latest('date_sent')
        count_obj = CourtEmailCount.objects.latest('date_sent')

        matches = re.search("<<<makeaplea-ref:\s*(\d+)/(\d+)>>>", mail.outbox[0].body)

        try:
            matches.groups()
        except AttributeError:
            self.fail('Body makeaplea-ref tag not found!')

        plea_id, count_id = matches.groups()

        self.assertEqual(int(plea_id), plea_obj.id)
        self.assertEqual(int(count_id), count_obj.id)

    def test_user_confirmation_sends_email(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89"},
                        "your_details": {"name": "vcx", "email": "lyndon@antlyn.com", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data, send_user_email=True)

        self.assertEqual(len(mail.outbox), 3)
        self.assertIn(context_data['case']['urn'].upper(), mail.outbox[-1].body)
        self.assertIn(context_data['your_details']['email'], mail.outbox[-1].to)

    def test_user_confirmation_sends_email_opt_out(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89"},
                        "your_details": {"name": "vcx", "email": "lyndon@antlyn.com", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data, send_user_email=False)

        self.assertEqual(len(mail.outbox), 3)<commit_msg>Fix the email generation unit tests<commit_after>from __future__ import absolute_import
import re

from django.test import TestCase
from django.core import mail

from ..email import TemplateAttachmentEmail, send_plea_email
from ..models import CourtEmailPlea, CourtEmailCount

class EmailGenerationTests(TestCase):
    def setUp(self):
        mail.outbox = []

    def test_template_attachment_sends_email(self):
        email_context = {"URN": "1A2B3C4D5E"}
        email = TemplateAttachmentEmail("test_from@example.org",
                                        "test.html",
                                        "plea/plea_email_attachment.html",
                                        email_context,
                                        "
Test Content

{{ URN }}
")

        email.send(["test_to@example.org", ],
                   "Subject line",
                   "Body Text")

        self.assertEqual(mail.outbox[0].subject, "Subject line")
        self.assertEqual(mail.outbox[0].body, "Body Text")

    def test_plea_email_sends(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89",
                                 "number_of_charges": 2},
                        "your_details": {"name": "vcx", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx",
                                         "email": "test@test.com"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data)

        # only expecting 1 email for now, PLP email cancelled for the time being
        self.assertEqual(len(mail.outbox), 3)

    def test_plea_email_body_contains_plea_and_count_ids(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89",
                                 "number_of_charges": 2},
                        "your_details": {"name": "vcx", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx",
                                         "email": "test@test.com"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data)

        plea_obj = CourtEmailPlea.objects.latest('date_sent')
        count_obj = CourtEmailCount.objects.latest('date_sent')

        matches = re.search("<<<makeaplea-ref:\s*(\d+)/(\d+)>>>", mail.outbox[0].body)

        try:
            matches.groups()
        except AttributeError:
            self.fail('Body makeaplea-ref tag not found!')

        plea_id, count_id = matches.groups()

        self.assertEqual(int(plea_id), plea_obj.id)
        self.assertEqual(int(count_id), count_obj.id)

    def test_user_confirmation_sends_email(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89",
                                 "number_of_charges": 2},
                        "your_details": {"name": "vcx", "email": "lyndon@antlyn.com", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data, send_user_email=True)

        self.assertEqual(len(mail.outbox), 3)
        self.assertIn(context_data['case']['urn'].upper(), mail.outbox[-1].body)
        self.assertIn(context_data['your_details']['email'], mail.outbox[-1].to)

    def test_user_confirmation_sends_email_opt_out(self):
        context_data = {"case": {"date_of_hearing": "2014-06-30",
                                 "time_of_hearing": "12:00:00",
                                 "urn": "cvxcvx89",
                                 "number_of_charges": 2},
                        "your_details": {"name": "vcx", "email": "lyndon@antlyn.com", "national_insurance_number": "xxx",
                                         "driving_licence_number": "xxx", "registration_number": "xxx"},
                        "plea": {"PleaForms": [{"mitigations": "test1", "guilty": "guilty"},
                                               {"mitigations": "test2", "guilty": "guilty"}],
                                 "understand": True}}

        send_plea_email(context_data, send_user_email=False)

        self.assertEqual(len(mail.outbox), 3)<|endoftext|>


<commit_before>import os
import pytest
import shutil
import tempfile

try:
    import attic.repository
    import attic.key
    import attic.helpers
except ImportError:
    attic = None
pytestmark = pytest.mark.skipif(attic is None,
                                reason = 'cannot find an attic install')

from ..converter import AtticRepositoryConverter, NotImplementedException
from ..helpers import get_keys_dir
from ..key import KeyfileKey
from ..repository import Repository, MAGIC
from . import BaseTestCase

class ConversionTestCase(BaseTestCase):

    class MockArgs:
        def __init__(self, path):
            self.repository = attic.helpers.Location(path)

    def open(self, path, repo_type  = Repository, create=False):
        return repo_type(os.path.join(path, 'repository'), create = create)

    def setUp(self):
        self.tmppath = tempfile.mkdtemp()
        self.attic_repo = self.open(self.tmppath,
                                    repo_type = attic.repository.Repository,
                                    create = True)
        # throw some stuff in that repo, copied from `RepositoryTestCase.test1`_
        for x in range(100):
            self.attic_repo.put(('%-32d' % x).encode('ascii'), b'SOMEDATA')
        self.attic_repo.close()

        # we use the repo dir for the created keyfile, because we do
        # not want to clutter existing keyfiles
        os.environ['ATTIC_KEYS_DIR'] = self.tmppath

        # we use the same directory for the converted files, which
        # will clutter the previously created one, which we don't care
        # about anyways. in real runs, the original key will be retained.
        os.environ['BORG_KEYS_DIR'] = self.tmppath
        os.environ['ATTIC_PASSPHRASE'] = 'test'
        self.key = attic.key.KeyfileKey.create(self.attic_repo, self.MockArgs(self.tmppath))

    def tearDown(self):
        shutil.rmtree(self.tmppath)

    def test_convert(self):
        self.repository = self.open(self.tmppath)
        # check should fail because of magic number
        print("this will show an error, it is expected")
        assert not self.repository.check() # can't check raises() because check() handles the error
        self.repository.close()
        print("opening attic repository with borg and converting")
        with pytest.raises(NotImplementedException):
            self.open(self.tmppath, repo_type = AtticRepositoryConverter).convert(dryrun=False)
        # check that the new keyfile is alright
        keyfile = os.path.join(get_keys_dir(),
                               os.path.basename(self.key.path))
        with open(keyfile, 'r') as f:
            assert f.read().startswith(KeyfileKey.FILE_ID)
        self.repository = self.open(self.tmppath)
        assert self.repository.check()
        self.repository.close()
<commit_msg>refactor: group test repo subroutine<commit_after>import os
import pytest
import shutil
import tempfile

try:
    import attic.repository
    import attic.key
    import attic.helpers
except ImportError:
    attic = None
pytestmark = pytest.mark.skipif(attic is None,
                                reason = 'cannot find an attic install')

from ..converter import AtticRepositoryConverter, NotImplementedException
from ..helpers import get_keys_dir
from ..key import KeyfileKey
from ..repository import Repository, MAGIC
from . import BaseTestCase

class ConversionTestCase(BaseTestCase):

    class MockArgs:
        def __init__(self, path):
            self.repository = attic.helpers.Location(path)

    def open(self, path, repo_type  = Repository, create=False):
        return repo_type(os.path.join(path, 'repository'), create = create)

    def setUp(self):
        self.tmppath = tempfile.mkdtemp()
        self.attic_repo = self.open(self.tmppath,
                                    repo_type = attic.repository.Repository,
                                    create = True)
        # throw some stuff in that repo, copied from `RepositoryTestCase.test1`_
        for x in range(100):
            self.attic_repo.put(('%-32d' % x).encode('ascii'), b'SOMEDATA')
        self.attic_repo.close()

        # we use the repo dir for the created keyfile, because we do
        # not want to clutter existing keyfiles
        os.environ['ATTIC_KEYS_DIR'] = self.tmppath

        # we use the same directory for the converted files, which
        # will clutter the previously created one, which we don't care
        # about anyways. in real runs, the original key will be retained.
        os.environ['BORG_KEYS_DIR'] = self.tmppath
        os.environ['ATTIC_PASSPHRASE'] = 'test'
        self.key = attic.key.KeyfileKey.create(self.attic_repo, self.MockArgs(self.tmppath))

    def tearDown(self):
        shutil.rmtree(self.tmppath)

    def check_repo(self, state = True):
        if not state:
            print("this will show an error, this is expected")
        self.repository = self.open(self.tmppath)
        assert self.repository.check() is state # can't check raises() because check() handles the error
        self.repository.close()

    def test_convert(self):
        # check should fail because of magic number
        self.check_repo(False)
        print("opening attic repository with borg and converting")
        with pytest.raises(NotImplementedException):
            self.open(self.tmppath, repo_type = AtticRepositoryConverter).convert(dryrun=False)
        # check that the new keyfile is alright
        keyfile = os.path.join(get_keys_dir(),
                               os.path.basename(self.key.path))
        with open(keyfile, 'r') as f:
            assert f.read().startswith(KeyfileKey.FILE_ID)
        self.check_repo()
<|endoftext|>


<commit_before>#!/usr/bin/python3

import os
import random
import time
# import re
#import sys
import urllib
# import json
import pprint

import twitter

# import pudb
# pu.db

import pymongo

client = pymongo.MongoClient(os.getenv('mongo_functal'))

# --- get_tweets

def get_tweets(topic):
    print('topic : ' + topic)

    try:
        search_results = twit.search.tweets(q=topic, lang='en', result_type='popular')
        # print('search_results')
        # pp.pprint(search_results)

        db = client['topics']
        tweets = db['tweets']

        # 'user': tweet['user']['name']
        texts = [{'_id': tweet['id_str'], 'text': tweet['text'], 'topic': topic}
                 for tweet in search_results['statuses'] if topic in tweet['text'] and tweets.find({'_id': tweet['id_str']}).count() == 0]

        if len(texts) > 0:
            pp.pprint(texts)

            # store
            try:
                result = tweets.insert(texts, {'ordered': False})
                print(str(len(result)) + ' tweets inserted')
            except pymongo.errors.PyMongoError as e:
                print(type(e))
                print(e)

    except urllib.error.URLError as e:
        print(e)

#--- global

pp = pprint.PrettyPrinter(indent=4)
auth = twitter.oauth.OAuth(os.environ['token'], os.environ['token_secret'], os.environ[
    'consumer_key'], os.environ['consumer_secret'])
twit = twitter.Twitter(auth=auth)

#--- main

def main():

    topics = ["red", "orange", "yellow", "green", "blue", "purple",
              "pink", "triangle", "square", "circle", "arrow", "asterisk", "wavy", "star",
              "sunset", "gold", "golden"]

    while True:
        try:
            get_tweets(topics[random.randint(0, len(topics) - 1)])
        # pylint: disable=W0703
        except Exception as e:
            print(type(e))
            print(e)
        # pylint: enable=W0703
        finally:
            time.sleep(60)

#--- run

main()
<commit_msg>no new tweets found<commit_after>#!/usr/bin/python3

import os
import random
import time
# import re
#import sys
import urllib
# import json
import pprint

import twitter

# import pudb
# pu.db

import pymongo

client = pymongo.MongoClient(os.getenv('mongo_functal'))

# --- get_tweets

def get_tweets(topic):
    print('topic : ' + topic)

    try:
        search_results = twit.search.tweets(q=topic, lang='en', result_type='popular')
        # print('search_results')
        # pp.pprint(search_results)

        db = client['topics']
        tweets = db['tweets']

        # 'user': tweet['user']['name']
        texts = [{'_id': tweet['id_str'], 'text': tweet['text'], 'topic': topic}
                 for tweet in search_results['statuses'] if topic in tweet['text'] and tweets.find({'_id': tweet['id_str']}).count() == 0]

        if len(texts) > 0:
            pp.pprint(texts)

            # store
            try:
                result = tweets.insert(texts, {'ordered': False})
                print(str(len(result)) + ' tweets inserted')
            except pymongo.errors.PyMongoError as e:
                print(type(e))
                print(e)
        else:
            print('No new tweets')

    except urllib.error.URLError as e:
        print(e)

#--- global

pp = pprint.PrettyPrinter(indent=4)
auth = twitter.oauth.OAuth(os.environ['token'], os.environ['token_secret'], os.environ[
    'consumer_key'], os.environ['consumer_secret'])
twit = twitter.Twitter(auth=auth)

#--- main

def main():

    # pylint: disable=W0703

    topics = ["red", "orange", "yellow", "green", "blue", "purple",
              "pink", "triangle", "square", "circle", "arrow", "asterisk", "wavy", "star",
              "sunset", "gold", "golden"]

    while True:
        try:
            get_tweets(topics[random.randint(0, len(topics) - 1)])
        except Exception as e:
            print(type(e))
            print(e)
        finally:
            time.sleep(60)

    # pylint: enable=W0703

#--- run

main()
<|endoftext|>


<commit_before>from bot.action.core.action import IntermediateAction

class TextMessageAction(IntermediateAction):
    def process(self, event):
        text = event.message.text
        if text is not None:
            event.text = text
            self._continue(event)

class VoiceMessageAction(IntermediateAction):
    def process(self, event):
        voice = event.message.voice
        if voice is not None:
            event.voice = voice
            self._continue(event)

class NoForwardedMessage(IntermediateAction):
    def process(self, event):
        forwarded = event.message.forward_date
        if forwarded is None:
            self._continue(event)

class EditedMessageAction(IntermediateAction):
    def process(self, event):
        edited_message = event.update.edited_message
        if edited_message is not None:
            event.message = edited_message
            event.edited_message = edited_message
            event.chat = edited_message.chat
            self._continue(event)

class MessageAction(IntermediateAction):
    def process(self, event):
        message = event.update.message
        if message is not None:
            event.message = message
            event.chat = message.chat
            self._continue(event)

class NoPendingAction(IntermediateAction):
    def process(self, event):
        if not event.is_pending:
            self._continue(event)

class PendingAction(IntermediateAction):
    def process(self, event):
        if event.is_pending:
            self._continue(event)
<commit_msg>Add InlineQueryAction to filter.py<commit_after>from bot.action.core.action import IntermediateAction

class TextMessageAction(IntermediateAction):
    def process(self, event):
        text = event.message.text
        if text is not None:
            event.text = text
            self._continue(event)

class VoiceMessageAction(IntermediateAction):
    def process(self, event):
        voice = event.message.voice
        if voice is not None:
            event.voice = voice
            self._continue(event)

class NoForwardedMessage(IntermediateAction):
    def process(self, event):
        forwarded = event.message.forward_date
        if forwarded is None:
            self._continue(event)

class EditedMessageAction(IntermediateAction):
    def process(self, event):
        edited_message = event.update.edited_message
        if edited_message is not None:
            event.message = edited_message
            event.edited_message = edited_message
            event.chat = edited_message.chat
            self._continue(event)

class MessageAction(IntermediateAction):
    def process(self, event):
        message = event.update.message
        if message is not None:
            event.message = message
            event.chat = message.chat
            self._continue(event)

class InlineQueryAction(IntermediateAction):
    def process(self, event):
        inline_query = event.update.inline_query
        if inline_query is not None:
            event.query = inline_query
            self._continue(event)

class NoPendingAction(IntermediateAction):
    def process(self, event):
        if not event.is_pending:
            self._continue(event)

class PendingAction(IntermediateAction):
    def process(self, event):
        if event.is_pending:
            self._continue(event)
<|endoftext|>


<commit_before>from bot.action.core.action import IntermediateAction

class TextMessageAction(IntermediateAction):
    def process(self, event):
        text = event.message.text
        if text is not None:
            event.text = text
            self._continue(event)

class MessageAction(IntermediateAction):
    def process(self, event):
        message = event.update.message
        if message is not None:
            event.message = message
            event.chat = message.chat
            self._continue(event)

class NoPendingAction(IntermediateAction):
    def process(self, event):
        if not event.is_pending:
            self._continue(event)
<commit_msg>Add EditedMessageAction<commit_after>from bot.action.core.action import IntermediateAction

class TextMessageAction(IntermediateAction):
    def process(self, event):
        text = event.message.text
        if text is not None:
            event.text = text
            self._continue(event)

class EditedMessageAction(IntermediateAction):
    def process(self, event):
        edited_message = event.update.edited_message
        if edited_message is not None:
            event.message = edited_message
            event.edited_message = edited_message
            event.chat = edited_message.chat
            self._continue(event)

class MessageAction(IntermediateAction):
    def process(self, event):
        message = event.update.message
        if message is not None:
            event.message = message
            event.chat = message.chat
            self._continue(event)

class NoPendingAction(IntermediateAction):
    def process(self, event):
        if not event.is_pending:
            self._continue(event)
<|endoftext|>