<commit_before>"""Rank genes according to differential expression.
"""

import numpy as np
import pandas as pd
from math import sqrt, floor
from scipy.sparse import issparse

from .. import utils
from .. import settings
from .. import logging as logg
from ..preprocessing import simple

def rank_genes_groups(
        adata,
        groupby,
        use_raw=True,
        groups='all',
        reference='rest',
        n_genes=100,
        only_positive=True,
        key_added=None,
        copy=False,
        method='t-test_overestim_var',
        **kwds):
    """Rank genes for characterizing groups.

    Parameters
    ----------
    adata : :class:`~anndata.AnnData`
        Annotated data matrix.
    groupby : `str`
        The key of the observations grouping to consider.
    use_raw : `bool`, optional (default: `True`)
        Use `raw` attribute of `adata` if present.
    groups : `str`, `list`, optional (default: `'all'`)
        Subset of groups, e.g. `['g1', 'g2', 'g3']`, to which comparison shall
        be restricted. If not passed, a ranking will be generated for all
        groups.
    reference : `str`, optional (default: `'rest'`)
        If `'rest'`, compare each group to the union of the rest of the group.  If
        a group identifier, compare with respect to this group.
    n_genes : `int`, optional (default: 100)
        The number of genes that appear in the returned tables.
    method : {'logreg', 't-test', 'wilcoxon', 't-test_overestim_var'}, optional (default: 't-test_overestim_var')
        If 't-test', uses t-test, if 'wilcoxon', uses Wilcoxon-Rank-Sum. If
        't-test_overestim_var', overestimates variance of each group. If
        'logreg' uses logistic regression, see [Ntranos18]_, `here
        <https://github.com/theislab/scanpy/issues/95>`__ and `here
        <http://www.nxn.se/valent/2018/3/5/actionable-scrna-seq-clusters>`__, for
        why this is meaningful.
    only_positive : bool, optional (default: `True`)
        Only consider positive differences.
    **kwds : keyword parameters
        Are passed to test methods. Currently this affects only parameters that
        are passed to `sklearn.linear_model.LogisticRegression
        <http://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html>`__.
        For instance, you can pass `penalty='l1'` to try to come up with a
        minimal set of genes that are good predictors (sparse solution meaning
        few non-zero fitted coefficients).

    Returns
    -------
    Updates `adata` with the following fields.
    names : structured `np.ndarray` (`.uns['rank_genes_groups']`)
        Structured array to be indexed by group id storing the gene
        names. Ordered according to scores.
    scores : structured `np.ndarray` (`.uns['rank_genes_groups']`)
        Structured array to be indexed by group id storing the score for each
        gene for each group. Ordered according to scores.
    """
    logg.info('ranking genes', r=True)
    adata = adata.copy() if copy else adata
    utils.sanitize_anndata(adata)
    # for clarity, rename variable
    groups_order = groups
    if isinstance(groups_order, list) and isinstance(groups_order[0], int):
        groups_order = [str(n) for n in groups_order]
    if reference != 'rest' and reference not in set(groups_order):
        groups_order += [reference]
    if (reference != 'rest'
        and reference not in set(adata.obs[groupby].cat.categories)):
        raise ValueError('reference = {} needs to be one of groupby = {}.'
                         .format(reference,
                                 adata.obs[groupby].cat.categories.tolist()))
    groups_order, groups_masks = utils.select_groups(
        adata, groups_order, groupby)

    if key_added is None:
        key_added = 'rank_genes_groups'
    adata.uns[key_added] = {}
    adata.uns[key_added]['params'] = {
        'groupby': groupby,
        'reference': reference,
        'method': method,
        'use_raw': use_raw,
    }

    # adata_comp mocks an AnnData object if use_raw is True
    # otherwise it's just the AnnData object
    adata_comp = adata
    if adata.raw is not None and use_raw:
        adata_comp = adata.raw
    X = adata_comp.X

    # for clarity, rename variable
    n_genes_user = n_genes
    # make sure indices are not OoB in case there are less genes than n_genes
    if n_genes_user > X.shape[1]:
        n_genes_user = X.shape[1]
    # in the following, n_genes is simply another name for the total number of genes
    n_genes = X.shape[1]

    rankings_gene_scores = []
    rankings_gene_names = []
    n_groups = groups_masks.shape[0]
    ns = np.zeros(n_groups, dtype=int)
    for imask, mask in enumerate(groups_masks):
        ns[imask] = np.where(mask)[0].size
    logg.msg('consider \'{}\' groups:'.format(groupby), groups_order, v=4)
    logg.msg('with sizes:', ns, v=4)
    if reference != 'rest':
        ireference = np.where(groups_order == reference)[0][0]
    reference_indices = np.arange(adata_comp.n_vars, dtype=int)

    avail_methods = {'t-test', 't-test_overestim_var', 'wilcoxon', 'logreg'}
    if method not in avail_methods:
        raise ValueError('Method must be one of {}.'.format(avail_methods))

    # TODO: all of this is probably going to be removed
    if method is 't-test_correction_factors':
        if correction_factors is None:
            raise ValueError('For this test type, you need to enter correction factors manually.')
        if len(correction_factors) != 2:
            raise ValueError('We need exactly 2 correction factors, accessible via correction_factors[i], i=0,1')
        if correction_factors[0] < 0 or correction_factors[1] < 0:
            raise ValueError('Correction factors need to be positive numbers!')

    if method in {'t-test', 't-test_overestim_var',
                  't-test_double_overestim_var', 't-test_correction_factors'}:
        # loop over all masks and compute means, variances and sample numbers
        means = np.zeros((n_groups, n_genes))
        vars = np.zeros((n_groups, n_genes))
        for imask, mask in enumerate(groups_masks):
            means[imask], vars[imask] = simple._get_mean_var(X[mask])
        # test each either against the union of all other groups or against a
        # specific group
        for igroup in range(n_groups):
            if reference == 'rest':
                mask_rest = ~groups_masks[igroup]
            else:
                if igroup == ireference: continue
                else: mask_rest = groups_masks[ireference]
            mean_rest, var_rest = simple._get_mean_var(X[mask_rest])
            if method == 't-test':
                ns_rest = np.where(mask_rest)[0].size
            elif method in {'t-test_overestim_var', 't-test_double_overestim_var'}:
                ns_rest = ns[igroup]  # hack for overestimating the variance
            elif method == 't-test_correction_factors':
                # The tendency is as follows: For the comparison group (rest), overesimate variance --> smaller ns_rest
                ns_rest = np.where(mask_rest)[0].size/correction_factors[1]

            if method in {'t-test', 't-test_overestim_var'}:
                ns_group = ns[igroup]
            elif method == 't-test_correction_factors':
                # We underestimate group variance by increasing denominator, i.e. ns_group
                ns_group = ns[igroup] * correction_factors[0]
            else:
                # We do the opposite of t-test_overestim_var
                ns_group = np.where(mask_rest)[0].size

            denominator = np.sqrt(vars[igroup]/ns_group + var_rest/ns_rest)
            denominator[np.flatnonzero(denominator == 0)] = np.nan
            scores = (means[igroup] - mean_rest) / denominator
            scores[np.isnan(scores)] = 0
            scores = scores if only_positive else np.abs(scores)
            partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
            partial_indices = np.argsort(scores[partition])[::-1]
            global_indices = reference_indices[partition][partial_indices]
            rankings_gene_scores.append(scores[global_indices])
            rankings_gene_names.append(adata_comp.var_names[global_indices])
    elif method == 'logreg':
        from sklearn.linear_model import LogisticRegression
        if reference != 'rest':
            raise ValueError('\'logreg\' is only implemented for `reference==\'rest\'`.')
        clf = LogisticRegression(**kwds)
        clf.fit(X, adata.obs[groupby])
        scores_all = clf.coef_
        for igroup, group in enumerate(groups_order):
            scores = scores_all[igroup]
            partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
            partial_indices = np.argsort(scores[partition])[::-1]
            global_indices = reference_indices[partition][partial_indices]
            rankings_gene_scores.append(scores[global_indices])
            rankings_gene_names.append(adata_comp.var_names[global_indices])
    elif method == 'wilcoxon':
        # Wilcoxon-rank-sum test is usually more powerful in detecting marker genes
        # Limit maximal RAM that is required by the calculation. Currently set fixed to roughly 100 MByte
        CONST_MAX_SIZE = 10000000
        ns_rest = np.zeros(n_groups, dtype=int)
        # initialize space for z-scores
        scores = np.zeros(n_genes)
        # First loop: Loop over all genes
        if reference != 'rest':
            for imask, mask in enumerate(groups_masks):
                if imask == ireference: continue
                else: mask_rest = groups_masks[ireference]
                ns_rest[imask] = np.where(mask_rest)[0].size
                if ns_rest[imask] <= 25 or ns[imask] <= 25:
                    logg.hint('Few observations in a group for '
                              'normal approximation (<=25). Lower test accuracy.')
                n_active = ns[imask]
                m_active = ns_rest[imask]
                # Now calculate gene expression ranking in chunkes:
                chunk = []
                # Calculate chunk frames
                n_genes_max_chunk = floor(CONST_MAX_SIZE / (n_active + m_active))
                if n_genes_max_chunk < n_genes - 1:
                    chunk_index = n_genes_max_chunk
                    while chunk_index < n_genes - 1:
                        chunk.append(chunk_index)
                        chunk_index = chunk_index + n_genes_max_chunk
                    chunk.append(n_genes - 1)
                else:
                    chunk.append(n_genes - 1)
                left = 0
                # Calculate rank sums for each chunk for the current mask
                for chunk_index, right in enumerate(chunk):
                    # Check if issparse is true: AnnData objects are currently sparse.csr or ndarray.
                    if issparse(X):
                        df1 = pd.DataFrame(data=X[mask, left:right].todense())
                        df2 = pd.DataFrame(data=X[mask_rest, left:right].todense(),
                                           index=np.arange(start=n_active, stop=n_active + m_active))
                    else:
                        df1 = pd.DataFrame(data=X[mask, left:right])
                        df2 = pd.DataFrame(data=X[mask_rest, left:right],
                                           index=np.arange(start=n_active, stop=n_active + m_active))
                    df1 = df1.append(df2)
                    ranks = df1.rank()
                    # sum up adjusted_ranks to calculate W_m,n
                    scores[left:right] = np.sum(ranks.loc[0:n_active, :])
                    left = right + 1
                scores = (scores - (n_active * (n_active + m_active + 1) / 2)) / sqrt(
                    (n_active * m_active * (n_active + m_active + 1) / 12))
                scores = scores if only_positive else np.abs(scores)
                scores[np.isnan(scores)] = 0
                partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
                partial_indices = np.argsort(scores[partition])[::-1]
                global_indices = reference_indices[partition][partial_indices]
                rankings_gene_scores.append(scores[global_indices])
                rankings_gene_names.append(adata_comp.var_names[global_indices])
        # If no reference group exists, ranking needs only to be done once (full mask)
        else:
            scores = np.zeros((n_groups, n_genes))
            chunk = []
            n_cells = X.shape[0]
            n_genes_max_chunk = floor(CONST_MAX_SIZE / n_cells)
            if n_genes_max_chunk < n_genes - 1:
                chunk_index = n_genes_max_chunk
                while chunk_index < n_genes - 1:
                    chunk.append(chunk_index)
                    chunk_index = chunk_index + n_genes_max_chunk
                chunk.append(n_genes - 1)
            else:
                chunk.append(n_genes - 1)
            left = 0
            for chunk_index, right in enumerate(chunk):
                # Check if issparse is true
                if issparse(X):
                    df1 = pd.DataFrame(data=X[:, left:right].todense())
                else:
                    df1 = pd.DataFrame(data=X[:, left:right])
                ranks = df1.rank()
                # sum up adjusted_ranks to calculate W_m,n
                for imask, mask in enumerate(groups_masks):
                    scores[imask, left:right] = np.sum(ranks.loc[mask, :])
                left = right + 1

            for imask, mask in enumerate(groups_masks):
                scores[imask, :] = (scores[imask, :] - (ns[imask] * (n_cells + 1) / 2)) / sqrt(
                    (ns[imask] * (n_cells - ns[imask]) * (n_cells + 1) / 12))
                scores = scores if only_positive else np.abs(scores)
                scores[np.isnan(scores)] = 0
                partition = np.argpartition(scores[imask, :], -n_genes_user)[-n_genes_user:]
                partial_indices = np.argsort(scores[imask, partition])[::-1]
                global_indices = reference_indices[partition][partial_indices]
                rankings_gene_scores.append(scores[imask, global_indices])
                rankings_gene_names.append(adata_comp.var_names[global_indices])

    groups_order_save = [str(g) for g in groups_order]
    if reference != 'rest':
        groups_order_save = [g for g in groups_order if g != reference]
    adata.uns[key_added]['scores'] = np.rec.fromarrays(
        [n for n in rankings_gene_scores],
        dtype=[(rn, 'float32') for rn in groups_order_save])
    adata.uns[key_added]['names'] = np.rec.fromarrays(
        [n for n in rankings_gene_names],
        dtype=[(rn, 'U50') for rn in groups_order_save])
    logg.info('    finished', time=True, end=' ' if settings.verbosity > 2 else '\n')
    logg.hint(
        'added to `.uns[\'{}\']`\n'
        '    \'names\', sorted np.recarray to be indexed by group ids\n'
        '    \'scores\', sorted np.recarray to be indexed by group ids'
        .format(key_added))
    return adata if copy else None
<commit_msg>one vs. one logreg update<commit_after>"""Rank genes according to differential expression.
"""

import numpy as np
import pandas as pd
from math import sqrt, floor
from scipy.sparse import issparse

from .. import utils
from .. import settings
from .. import logging as logg
from ..preprocessing import simple

def rank_genes_groups(
        adata,
        groupby,
        use_raw=True,
        groups='all',
        reference='rest',
        n_genes=100,
        only_positive=True,
        key_added=None,
        copy=False,
        method='t-test_overestim_var',
        **kwds):
    """Rank genes for characterizing groups.

    Parameters
    ----------
    adata : :class:`~anndata.AnnData`
        Annotated data matrix.
    groupby : `str`
        The key of the observations grouping to consider.
    use_raw : `bool`, optional (default: `True`)
        Use `raw` attribute of `adata` if present.
    groups : `str`, `list`, optional (default: `'all'`)
        Subset of groups, e.g. `['g1', 'g2', 'g3']`, to which comparison shall
        be restricted. If not passed, a ranking will be generated for all
        groups.
    reference : `str`, optional (default: `'rest'`)
        If `'rest'`, compare each group to the union of the rest of the group.  If
        a group identifier, compare with respect to this group.
    n_genes : `int`, optional (default: 100)
        The number of genes that appear in the returned tables.
    method : {'logreg', 't-test', 'wilcoxon', 't-test_overestim_var'}, optional (default: 't-test_overestim_var')
        If 't-test', uses t-test, if 'wilcoxon', uses Wilcoxon-Rank-Sum. If
        't-test_overestim_var', overestimates variance of each group. If
        'logreg' uses logistic regression, see [Ntranos18]_, `here
        <https://github.com/theislab/scanpy/issues/95>`__ and `here
        <http://www.nxn.se/valent/2018/3/5/actionable-scrna-seq-clusters>`__, for
        why this is meaningful.
    only_positive : bool, optional (default: `True`)
        Only consider positive differences.
    **kwds : keyword parameters
        Are passed to test methods. Currently this affects only parameters that
        are passed to `sklearn.linear_model.LogisticRegression
        <http://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html>`__.
        For instance, you can pass `penalty='l1'` to try to come up with a
        minimal set of genes that are good predictors (sparse solution meaning
        few non-zero fitted coefficients).

    Returns
    -------
    Updates `adata` with the following fields.
    names : structured `np.ndarray` (`.uns['rank_genes_groups']`)
        Structured array to be indexed by group id storing the gene
        names. Ordered according to scores.
    scores : structured `np.ndarray` (`.uns['rank_genes_groups']`)
        Structured array to be indexed by group id storing the score for each
        gene for each group. Ordered according to scores.
    """
    logg.info('ranking genes', r=True)
    adata = adata.copy() if copy else adata
    utils.sanitize_anndata(adata)
    # for clarity, rename variable
    groups_order = groups
    if isinstance(groups_order, list) and isinstance(groups_order[0], int):
        groups_order = [str(n) for n in groups_order]
    if reference != 'rest' and reference not in set(groups_order):
        groups_order += [reference]
    if (reference != 'rest'
        and reference not in set(adata.obs[groupby].cat.categories)):
        raise ValueError('reference = {} needs to be one of groupby = {}.'
                         .format(reference,
                                 adata.obs[groupby].cat.categories.tolist()))
    groups_order, groups_masks = utils.select_groups(
        adata, groups_order, groupby)

    if key_added is None:
        key_added = 'rank_genes_groups'
    adata.uns[key_added] = {}
    adata.uns[key_added]['params'] = {
        'groupby': groupby,
        'reference': reference,
        'method': method,
        'use_raw': use_raw,
    }

    # adata_comp mocks an AnnData object if use_raw is True
    # otherwise it's just the AnnData object
    adata_comp = adata
    if adata.raw is not None and use_raw:
        adata_comp = adata.raw
    X = adata_comp.X

    # for clarity, rename variable
    n_genes_user = n_genes
    # make sure indices are not OoB in case there are less genes than n_genes
    if n_genes_user > X.shape[1]:
        n_genes_user = X.shape[1]
    # in the following, n_genes is simply another name for the total number of genes
    n_genes = X.shape[1]

    rankings_gene_scores = []
    rankings_gene_names = []
    n_groups = groups_masks.shape[0]
    ns = np.zeros(n_groups, dtype=int)
    for imask, mask in enumerate(groups_masks):
        ns[imask] = np.where(mask)[0].size
    logg.msg('consider \'{}\' groups:'.format(groupby), groups_order, v=4)
    logg.msg('with sizes:', ns, v=4)
    if reference != 'rest':
        ireference = np.where(groups_order == reference)[0][0]
    reference_indices = np.arange(adata_comp.n_vars, dtype=int)

    avail_methods = {'t-test', 't-test_overestim_var', 'wilcoxon', 'logreg'}
    if method not in avail_methods:
        raise ValueError('Method must be one of {}.'.format(avail_methods))

    # TODO: all of this is probably going to be removed
    if method is 't-test_correction_factors':
        if correction_factors is None:
            raise ValueError('For this test type, you need to enter correction factors manually.')
        if len(correction_factors) != 2:
            raise ValueError('We need exactly 2 correction factors, accessible via correction_factors[i], i=0,1')
        if correction_factors[0] < 0 or correction_factors[1] < 0:
            raise ValueError('Correction factors need to be positive numbers!')

    if method in {'t-test', 't-test_overestim_var',
                  't-test_double_overestim_var', 't-test_correction_factors'}:
        # loop over all masks and compute means, variances and sample numbers
        means = np.zeros((n_groups, n_genes))
        vars = np.zeros((n_groups, n_genes))
        for imask, mask in enumerate(groups_masks):
            means[imask], vars[imask] = simple._get_mean_var(X[mask])
        # test each either against the union of all other groups or against a
        # specific group
        for igroup in range(n_groups):
            if reference == 'rest':
                mask_rest = ~groups_masks[igroup]
            else:
                if igroup == ireference: continue
                else: mask_rest = groups_masks[ireference]
            mean_rest, var_rest = simple._get_mean_var(X[mask_rest])
            if method == 't-test':
                ns_rest = np.where(mask_rest)[0].size
            elif method in {'t-test_overestim_var', 't-test_double_overestim_var'}:
                ns_rest = ns[igroup]  # hack for overestimating the variance
            elif method == 't-test_correction_factors':
                # The tendency is as follows: For the comparison group (rest), overesimate variance --> smaller ns_rest
                ns_rest = np.where(mask_rest)[0].size/correction_factors[1]

            if method in {'t-test', 't-test_overestim_var'}:
                ns_group = ns[igroup]
            elif method == 't-test_correction_factors':
                # We underestimate group variance by increasing denominator, i.e. ns_group
                ns_group = ns[igroup] * correction_factors[0]
            else:
                # We do the opposite of t-test_overestim_var
                ns_group = np.where(mask_rest)[0].size

            denominator = np.sqrt(vars[igroup]/ns_group + var_rest/ns_rest)
            denominator[np.flatnonzero(denominator == 0)] = np.nan
            scores = (means[igroup] - mean_rest) / denominator
            scores[np.isnan(scores)] = 0
            scores = scores if only_positive else np.abs(scores)
            partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
            partial_indices = np.argsort(scores[partition])[::-1]
            global_indices = reference_indices[partition][partial_indices]
            rankings_gene_scores.append(scores[global_indices])
            rankings_gene_names.append(adata_comp.var_names[global_indices])
    elif method == 'logreg':
    #if reference is not set, then the groups listed will be compared to the rest
    #if reference is set, then the groups listed will be compared only to the other groups listed
        from sklearn.linear_model import LogisticRegression
        if reference not in groups:
            raise Exception("reference must be in groups")
        if len(groups) == 1:
            raise Exception("Cannot perform logistic regression on a single cluster.")
        if len(groups) == 2: #binary logistic regression
            adata_copy = adata[(adata.obs[groupby] == groups_order[0]) | (adata.obs[groupby] == groups_order[1])]
            X = adata_copy.raw.X
        else:
            logg.hint("Logreg is currently implemented for either all or binary classification. Running one vs. all:")
            adata_copy = adata

        clf = LogisticRegression(**kwds)
        clf.fit(X, adata_copy.obs[groupby])
        scores_all = clf.coef_
        for igroup, group in enumerate(groups_order):
            if reference != 'rest' and len(groups) <= 2: #binary logistic regression
                scores = scores_all[0]
            else:
                scores = scores_all[igroup]
            partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
            partial_indices = np.argsort(scores[partition])[::-1]
            global_indices = reference_indices[partition][partial_indices]
            rankings_gene_scores.append(scores[global_indices])
            rankings_gene_names.append(adata_copy.raw.var_names[global_indices])
            if reference != 'rest' and len(groups) <= 2:
                break

    elif method == 'wilcoxon':
        # Wilcoxon-rank-sum test is usually more powerful in detecting marker genes
        # Limit maximal RAM that is required by the calculation. Currently set fixed to roughly 100 MByte
        CONST_MAX_SIZE = 10000000
        ns_rest = np.zeros(n_groups, dtype=int)
        # initialize space for z-scores
        scores = np.zeros(n_genes)
        # First loop: Loop over all genes
        if reference != 'rest':
            for imask, mask in enumerate(groups_masks):
                if imask == ireference: continue
                else: mask_rest = groups_masks[ireference]
                ns_rest[imask] = np.where(mask_rest)[0].size
                if ns_rest[imask] <= 25 or ns[imask] <= 25:
                    logg.hint('Few observations in a group for '
                              'normal approximation (<=25). Lower test accuracy.')
                n_active = ns[imask]
                m_active = ns_rest[imask]
                # Now calculate gene expression ranking in chunkes:
                chunk = []
                # Calculate chunk frames
                n_genes_max_chunk = floor(CONST_MAX_SIZE / (n_active + m_active))
                if n_genes_max_chunk < n_genes - 1:
                    chunk_index = n_genes_max_chunk
                    while chunk_index < n_genes - 1:
                        chunk.append(chunk_index)
                        chunk_index = chunk_index + n_genes_max_chunk
                    chunk.append(n_genes - 1)
                else:
                    chunk.append(n_genes - 1)
                left = 0
                # Calculate rank sums for each chunk for the current mask
                for chunk_index, right in enumerate(chunk):
                    # Check if issparse is true: AnnData objects are currently sparse.csr or ndarray.
                    if issparse(X):
                        df1 = pd.DataFrame(data=X[mask, left:right].todense())
                        df2 = pd.DataFrame(data=X[mask_rest, left:right].todense(),
                                           index=np.arange(start=n_active, stop=n_active + m_active))
                    else:
                        df1 = pd.DataFrame(data=X[mask, left:right])
                        df2 = pd.DataFrame(data=X[mask_rest, left:right],
                                           index=np.arange(start=n_active, stop=n_active + m_active))
                    df1 = df1.append(df2)
                    ranks = df1.rank()
                    # sum up adjusted_ranks to calculate W_m,n
                    scores[left:right] = np.sum(ranks.loc[0:n_active, :])
                    left = right + 1
                scores = (scores - (n_active * (n_active + m_active + 1) / 2)) / sqrt(
                    (n_active * m_active * (n_active + m_active + 1) / 12))
                scores = scores if only_positive else np.abs(scores)
                scores[np.isnan(scores)] = 0
                partition = np.argpartition(scores, -n_genes_user)[-n_genes_user:]
                partial_indices = np.argsort(scores[partition])[::-1]
                global_indices = reference_indices[partition][partial_indices]
                rankings_gene_scores.append(scores[global_indices])
                rankings_gene_names.append(adata_comp.var_names[global_indices])
        # If no reference group exists, ranking needs only to be done once (full mask)
        else:
            scores = np.zeros((n_groups, n_genes))
            chunk = []
            n_cells = X.shape[0]
            n_genes_max_chunk = floor(CONST_MAX_SIZE / n_cells)
            if n_genes_max_chunk < n_genes - 1:
                chunk_index = n_genes_max_chunk
                while chunk_index < n_genes - 1:
                    chunk.append(chunk_index)
                    chunk_index = chunk_index + n_genes_max_chunk
                chunk.append(n_genes - 1)
            else:
                chunk.append(n_genes - 1)
            left = 0
            for chunk_index, right in enumerate(chunk):
                # Check if issparse is true
                if issparse(X):
                    df1 = pd.DataFrame(data=X[:, left:right].todense())
                else:
                    df1 = pd.DataFrame(data=X[:, left:right])
                ranks = df1.rank()
                # sum up adjusted_ranks to calculate W_m,n
                for imask, mask in enumerate(groups_masks):
                    scores[imask, left:right] = np.sum(ranks.loc[mask, :])
                left = right + 1

            for imask, mask in enumerate(groups_masks):
                scores[imask, :] = (scores[imask, :] - (ns[imask] * (n_cells + 1) / 2)) / sqrt(
                    (ns[imask] * (n_cells - ns[imask]) * (n_cells + 1) / 12))
                scores = scores if only_positive else np.abs(scores)
                scores[np.isnan(scores)] = 0
                partition = np.argpartition(scores[imask, :], -n_genes_user)[-n_genes_user:]
                partial_indices = np.argsort(scores[imask, partition])[::-1]
                global_indices = reference_indices[partition][partial_indices]
                rankings_gene_scores.append(scores[imask, global_indices])
                rankings_gene_names.append(adata_comp.var_names[global_indices])

    groups_order_save = [str(g) for g in groups_order]
    if reference != 'rest':
        groups_order_save = [g for g in groups_order if g != reference]
    adata.uns[key_added]['scores'] = np.rec.fromarrays(
        [n for n in rankings_gene_scores],
        dtype=[(rn, 'float32') for rn in groups_order_save])
    adata.uns[key_added]['names'] = np.rec.fromarrays(
        [n for n in rankings_gene_names],
        dtype=[(rn, 'U50') for rn in groups_order_save])
    logg.info('    finished', time=True, end=' ' if settings.verbosity > 2 else '\n')
    logg.hint(
        'added to `.uns[\'{}\']`\n'
        '    \'names\', sorted np.recarray to be indexed by group ids\n'
        '    \'scores\', sorted np.recarray to be indexed by group ids'
        .format(key_added))
    return adata if copy else None
<|endoftext|>


<commit_before>import os
import subprocess
import tempfile
import nbformat
import pytest
import markdown
import re
import urllib.request

from enchant.checker import SpellChecker
from enchant.tokenize import Filter, EmailFilter, URLFilter
from enchant import DictWithPWL

from lxml.html import document_fromstring, etree
try:
   # Python 3
   from urllib.request import urlopen, URLError
except ImportError:
   from urllib2 import urlopen, URLError

"""
run all tests:
pytest -v --tb=short

run python tests:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_python_notebook

run r tests:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_r_notebook

run specific Python test:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_python_notebook[00_Setup.ipynb]

run specific R test:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_r_notebook[00_Setup.ipynb]

-s : disable all capturing of output.
"""

class Test_notebooks(object):
    """
    Testing of SimpleITK Jupyter notebooks:
    1. Static analysis:
       Check that notebooks do not contain output (sanity check as these should
       not have been pushed to the repository).
       Check that all the URLs in the markdown cells are not broken.
    2. Dynamic analysis:
       Run the notebook and check for errors. In some notebooks we
       intentionally cause errors to illustrate certain features of the toolkit.
       All code cells that intentionally generate an error are expected to be
       marked using the cell's metadata. In the notebook go to
       "View->Cell Toolbar->Edit Metadata and add the following json entry:

       "simpleitk_error_expected": simpleitk_error_message

       with the appropriate "simpleitk_error_message" text.
       Cells where an error is allowed, but not necessarily expected should be
       marked with the following json:

       "simpleitk_error_allowed": simpleitk_error_message

       The simpleitk_error_message is a substring of the generated error
       message, such as 'Exception thrown in SimpleITK Show:'

    To test notebooks that use too much memory (exceed the 4Gb allocated for the testing
    machine):
    1. Create an enviornment variable named SIMPLE_ITK_MEMORY_CONSTRAINED_ENVIRONMENT
    2. Import the setup_for_testing.py at the top of the notebook. This module will
       decorate the sitk.ReadImage so that after reading the initial image it is
       resampled by a factor of 4 in each dimension.

    Adding a test:
    Simply add the new notebook file name to the list of files decorating the test_python_notebook
    or test_r_notebook functions. DON'T FORGET THE COMMA.
    """

    _allowed_error_markup = 'simpleitk_error_allowed'
    _expected_error_markup = 'simpleitk_error_expected'

    @pytest.mark.parametrize('notebook_file_name',
                             ['00_Setup.ipynb',
                              '01_Image_Basics.ipynb',
                              '02_Pythonic_Image.ipynb',
                              '03_Image_Details.ipynb',
                              '04_Image_Display.ipynb',
                              '05_Results_Visualization.ipynb',
                              '10_matplotlib\'s_imshow.ipynb',
                              '20_Expand_With_Interpolators.ipynb',
                              '21_Transforms_and_Resampling.ipynb',
                              '22_Transforms.ipynb',
                              '300_Segmentation_Overview.ipynb',
                              '30_Segmentation_Region_Growing.ipynb',
                              '31_Levelset_Segmentation.ipynb',
                              '32_Watersheds_Segmentation.ipynb',
                              '33_Segmentation_Thresholding_Edge_Detection.ipynb',
                              '34_Segmentation_Evaluation.ipynb',
                              pytest.param('35_Segmentation_Shape_Analysis.ipynb', marks=pytest.mark.skipif(os.environ.get('CIRCLECI')=='true', \
                                                                                reason="runtime too long for CircleCI")),
                              '36_Microscopy_Colocalization_Distance_Analysis.ipynb',
                              #'11_Progress.ipynb', # This notebook times out when run with nbconvert, due to javascript issues, so not tested.
                              '51_VH_Segmentation1.ipynb',
                              '55_VH_Resample.ipynb',
                              '56_VH_Registration1.ipynb',
                              '60_Registration_Introduction.ipynb',
                              '61_Registration_Introduction_Continued.ipynb',
                              '62_Registration_Tuning.ipynb',
                              '63_Registration_Initialization.ipynb',
                              '64_Registration_Memory_Time_Tradeoff.ipynb',
                              '65_Registration_FFD.ipynb',
                              '66_Registration_Demons.ipynb',
                              '67_Registration_Semiautomatic_Homework.ipynb',
                              '68_Registration_Errors.ipynb',
                              '69_x-ray-panorama.ipynb',
                              '70_Data_Augmentation.ipynb',
                              '71_Trust_But_Verify.ipynb'])
    def test_python_notebook(self, notebook_file_name):
       self.evaluate_notebook(self.absolute_path_python(notebook_file_name), 'python')

    @pytest.mark.parametrize('notebook_file_name',
                             ['00_Setup.ipynb',
                              'Image_Basics.ipynb',
                              'R_style_image.ipynb',
                              '33_Segmentation_Thresholding_Edge_Detection.ipynb',
                              '34_Segmentation_Evaluation.ipynb',
                              '35_Cell_Segmentation.ipynb',
                              '22_Transforms.ipynb',
                              '300_Segmentation_Overview.ipynb',
                              '60_Registration_Introduction.ipynb',
                              '61_Registration_Introduction_Continued.ipynb',
                              '65_Registration_FFD.ipynb',
                              '66_Registration_Demons.ipynb',
                              '70_Data_Augmentation.ipynb'])
    def test_r_notebook(self, notebook_file_name):
       self.evaluate_notebook(self.absolute_path_r(notebook_file_name), 'ir')

    def evaluate_notebook(self, path, kernel_name):
        """
        Perform static and dynamic analysis of the notebook.
        Execute a notebook via nbconvert and print the results of the test (errors etc.)
        Args:
            path (string): Name of notebook to run.
            kernel_name (string): Which jupyter kernel to use to run the test.
                                  Relevant values are:'python2', 'python3', 'ir'.
        """

        dir_name, file_name = os.path.split(path)
        if dir_name:
            os.chdir(dir_name)

        print('-------- begin (kernel {0}) {1} --------'.format(kernel_name,file_name))
        no_static_errors = self.static_analysis(path)
        no_dynamic_errors = self.dynamic_analysis(path, kernel_name)
        print('-------- end (kernel {0}) {1} --------'.format(kernel_name,file_name))
        assert(no_static_errors and no_dynamic_errors)

    def static_analysis(self, path):
        """
        Perform static analysis of the notebook.
        Read the notebook and check that there is no ouput and that the links
        in the markdown cells are not broken.
        Args:
            path (string): Name of notebook.
        Return:
            boolean: True if static analysis succeeded, otherwise False.
        """

        nb = nbformat.read(path, nbformat.current_nbformat)

        #######################
        # Check that the notebook does not contain output from code cells
        # (should not be in the repository, but well...).
        #######################
        no_unexpected_output = True

        # Check that the cell dictionary has an 'outputs' key and that it is
        # empty, relies on Python using short circuit evaluation so that we
        # don't get KeyError when retrieving the 'outputs' entry.
        cells_with_output = [c.source for c in nb.cells if 'outputs' in c and c.outputs]
        if cells_with_output:
            no_unexpected_output = False
            print('Cells with unexpected output:\n_____________________________')
            for cell in cells_with_output:
                print(cell+'\n---')
        else:
           print('no unexpected output')

        #######################
        # Check that all the links in the markdown cells are valid/accessible.
        #######################
        no_broken_links = True

        cells_and_broken_links = []
        for c in nb.cells:
           if c.cell_type == 'markdown':
              html_tree = document_fromstring(markdown.markdown(c.source))
              broken_links = []
              #iterlinks() returns tuples of the form (element, attribute, link, pos)
              for document_link in html_tree.iterlinks():
                 try:
                    if all( prefix not in document_link[2] for prefix in ['http','https']):  # Local file.
                       url = 'file:' + urllib.request.pathname2url(document_link[2])
                    else:  # Remote file.
                       url = document_link[2]
                    urlopen(url)
                 except URLError:
                    broken_links.append(url)
              if broken_links:
                 cells_and_broken_links.append((broken_links,c.source))
        if cells_and_broken_links:
           no_broken_links = False
           print('Cells with broken links:\n________________________')
           for links, cell in cells_and_broken_links:
              print(cell+'\n')
              print('\tBroken links:')
              print('\t'+'\n\t'.join(links)+'\n---')
        else:
           print('no broken links')

        #######################
        # Spell check all markdown cells and comments in code cells using the pyenchant spell checker.
        #######################
        no_spelling_mistakes = True
        simpleitk_notebooks_dictionary = DictWithPWL('en_US', os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                                     'additional_dictionary.txt'))
        spell_checker = SpellChecker(simpleitk_notebooks_dictionary, filters = [EmailFilter, URLFilter])
        cells_and_spelling_mistakes = []
        for c in nb.cells:
           spelling_mistakes = []
           if c.cell_type == 'markdown':
              # Get the text as a string from the html without the markup which is replaced by space.
              spell_checker.set_text(' '.join(etree.XPath('//text()')(document_fromstring(markdown.markdown(c.source)))))
           elif c.cell_type == 'code':
              # Get all the comments and concatenate them into a single string separated by newlines.
              comment_lines = re.findall('#+.*',c.source)
              spell_checker.set_text('\n'.join(comment_lines))
           for error in spell_checker:
              error_message = 'error: '+ '\'' + error.word +'\', ' + 'suggestions: ' + str(spell_checker.suggest())
              spelling_mistakes.append(error_message)
           if spelling_mistakes:
              cells_and_spelling_mistakes.append((spelling_mistakes, c.source))
        if cells_and_spelling_mistakes:
           no_spelling_mistakes = False
           print('Cells with spelling mistakes:\n________________________')
           for misspelled_words, cell in cells_and_spelling_mistakes:
              print(cell+'\n')
              print('\tMisspelled words and suggestions:')
              print('\t'+'\n\t'.join(misspelled_words)+'\n---')
        else:
           print('no spelling mistakes')

        return(no_unexpected_output and no_broken_links and no_spelling_mistakes)

    def dynamic_analysis(self, path, kernel_name):
        """
        Perform dynamic analysis of the notebook.
        Execute a notebook via nbconvert and print the results of the test
        (errors etc.)
        Args:
            path (string): Name of notebook to run.
            kernel_name (string): Which jupyter kernel to use to run the test.
                                  Relevant values are:'python', 'ir'.
        Return:
            boolean: True if dynamic analysis succeeded, otherwise False.
        """

        # Execute the notebook and allow errors (run all cells), output is
        # written to a temporary file which is automatically deleted.
        # The delete=False is here to address an issue that Windows that seems to have with temporary files (see https://bugs.python.org/issue14243).
        with tempfile.NamedTemporaryFile(suffix='.ipynb',delete=False) as fout:
            args = ['jupyter', 'nbconvert',
                    '--to', 'notebook',
                    '--execute',
                    '--ExecutePreprocessor.kernel_name='+kernel_name,
                    '--ExecutePreprocessor.allow_errors=True',
                    '--ExecutePreprocessor.timeout=6000', # seconds till timeout
                    '--output', fout.name, path]
            subprocess.check_call(args)
            nb = nbformat.read(fout.name, nbformat.current_nbformat)

        # Get all of the unexpected errors (logic: cell has output with an error
        # and no error is expected or the allowed/expected error is not the one which
        # was generated.)
        unexpected_errors = [(output.evalue, c.source) for c in nb.cells \
                              if 'outputs' in c for output in c.outputs \
                              if (output.output_type=='error') and \
                               (((Test_notebooks._allowed_error_markup not in c.metadata) and (Test_notebooks._expected_error_markup not in c.metadata))or \
                               ((Test_notebooks._allowed_error_markup in c.metadata) and (c.metadata[Test_notebooks._allowed_error_markup] not in output.evalue)) or \
                               ((Test_notebooks._expected_error_markup in c.metadata) and (c.metadata[Test_notebooks._expected_error_markup] not in output.evalue)))]

        no_unexpected_errors = True
        if unexpected_errors:
            no_unexpected_errors = False
            print('Cells with unexpected errors:\n_____________________________')
            for e, src in unexpected_errors:
                print(src)
                print('unexpected error: '+e)
        else:
            print('no unexpected errors')

        # Get all of the missing expected errors (logic: cell has output
        # but expected error was not generated.)
        missing_expected_errors = []
        for c in nb.cells:
            if Test_notebooks._expected_error_markup in c.metadata:
                missing_error = True
                if 'outputs' in c:
                    for output in c.outputs:
                        if (output.output_type=='error') and (c.metadata[Test_notebooks._expected_error_markup] in output.evalue):
                                missing_error = False
                if missing_error:
                    missing_expected_errors.append((c.metadata[Test_notebooks._expected_error_markup],c.source))

        no_missing_expected_errors = True
        if missing_expected_errors:
            no_missing_expected_errors = False
            print('\nCells with missing expected errors:\n___________________________________')
            for e, src in missing_expected_errors:
                print(src)
                print('missing expected error: '+e)
        else:
            print('no missing expected errors')

        return(no_unexpected_errors and no_missing_expected_errors)

    def absolute_path_python(self, notebook_file_name):
        return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../Python', notebook_file_name))

    def absolute_path_r(self, notebook_file_name):
        return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../R', notebook_file_name))
<commit_msg>Change issues with url not accessible from error to warning.<commit_after>import os
import subprocess
import tempfile
import nbformat
import pytest
import warnings
import markdown
import re
import urllib.request

from enchant.checker import SpellChecker
from enchant.tokenize import Filter, EmailFilter, URLFilter
from enchant import DictWithPWL

from lxml.html import document_fromstring, etree
try:
   # Python 3
   from urllib.request import urlopen, URLError
except ImportError:
   from urllib2 import urlopen, URLError

"""
run all tests:
pytest -v --tb=short

run python tests:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_python_notebook

run r tests:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_r_notebook

run specific Python test:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_python_notebook[00_Setup.ipynb]

run specific R test:
pytest -v --tb=short tests/test_notebooks.py::Test_notebooks::test_r_notebook[00_Setup.ipynb]

-s : disable all capturing of output.
"""

class Test_notebooks(object):
    """
    Testing of SimpleITK Jupyter notebooks:
    1. Static analysis:
       Check that notebooks do not contain output (sanity check as these should
       not have been pushed to the repository).
       Check that all the URLs in the markdown cells are not broken.
    2. Dynamic analysis:
       Run the notebook and check for errors. In some notebooks we
       intentionally cause errors to illustrate certain features of the toolkit.
       All code cells that intentionally generate an error are expected to be
       marked using the cell's metadata. In the notebook go to
       "View->Cell Toolbar->Edit Metadata and add the following json entry:

       "simpleitk_error_expected": simpleitk_error_message

       with the appropriate "simpleitk_error_message" text.
       Cells where an error is allowed, but not necessarily expected should be
       marked with the following json:

       "simpleitk_error_allowed": simpleitk_error_message

       The simpleitk_error_message is a substring of the generated error
       message, such as 'Exception thrown in SimpleITK Show:'

    To test notebooks that use too much memory (exceed the 4Gb allocated for the testing
    machine):
    1. Create an enviornment variable named SIMPLE_ITK_MEMORY_CONSTRAINED_ENVIRONMENT
    2. Import the setup_for_testing.py at the top of the notebook. This module will
       decorate the sitk.ReadImage so that after reading the initial image it is
       resampled by a factor of 4 in each dimension.

    Adding a test:
    Simply add the new notebook file name to the list of files decorating the test_python_notebook
    or test_r_notebook functions. DON'T FORGET THE COMMA.
    """

    _allowed_error_markup = 'simpleitk_error_allowed'
    _expected_error_markup = 'simpleitk_error_expected'

    @pytest.mark.parametrize('notebook_file_name',
                             ['00_Setup.ipynb',
                              '01_Image_Basics.ipynb',
                              '02_Pythonic_Image.ipynb',
                              '03_Image_Details.ipynb',
                              '04_Image_Display.ipynb',
                              '05_Results_Visualization.ipynb',
                              '10_matplotlib\'s_imshow.ipynb',
                              '20_Expand_With_Interpolators.ipynb',
                              '21_Transforms_and_Resampling.ipynb',
                              '22_Transforms.ipynb',
                              '300_Segmentation_Overview.ipynb',
                              '30_Segmentation_Region_Growing.ipynb',
                              '31_Levelset_Segmentation.ipynb',
                              '32_Watersheds_Segmentation.ipynb',
                              '33_Segmentation_Thresholding_Edge_Detection.ipynb',
                              '34_Segmentation_Evaluation.ipynb',
                              pytest.param('35_Segmentation_Shape_Analysis.ipynb', marks=pytest.mark.skipif(os.environ.get('CIRCLECI')=='true', \
                                                                                reason="runtime too long for CircleCI")),
                              '36_Microscopy_Colocalization_Distance_Analysis.ipynb',
                              #'11_Progress.ipynb', # This notebook times out when run with nbconvert, due to javascript issues, so not tested.
                              '51_VH_Segmentation1.ipynb',
                              '55_VH_Resample.ipynb',
                              '56_VH_Registration1.ipynb',
                              '60_Registration_Introduction.ipynb',
                              '61_Registration_Introduction_Continued.ipynb',
                              '62_Registration_Tuning.ipynb',
                              '63_Registration_Initialization.ipynb',
                              '64_Registration_Memory_Time_Tradeoff.ipynb',
                              '65_Registration_FFD.ipynb',
                              '66_Registration_Demons.ipynb',
                              '67_Registration_Semiautomatic_Homework.ipynb',
                              '68_Registration_Errors.ipynb',
                              '69_x-ray-panorama.ipynb',
                              '70_Data_Augmentation.ipynb',
                              '71_Trust_But_Verify.ipynb'])
    def test_python_notebook(self, notebook_file_name):
       self.evaluate_notebook(self.absolute_path_python(notebook_file_name), 'python')

    @pytest.mark.parametrize('notebook_file_name',
                             ['00_Setup.ipynb',
                              'Image_Basics.ipynb',
                              'R_style_image.ipynb',
                              '33_Segmentation_Thresholding_Edge_Detection.ipynb',
                              '34_Segmentation_Evaluation.ipynb',
                              '35_Cell_Segmentation.ipynb',
                              '22_Transforms.ipynb',
                              '300_Segmentation_Overview.ipynb',
                              '60_Registration_Introduction.ipynb',
                              '61_Registration_Introduction_Continued.ipynb',
                              '65_Registration_FFD.ipynb',
                              '66_Registration_Demons.ipynb',
                              '70_Data_Augmentation.ipynb'])
    def test_r_notebook(self, notebook_file_name):
       self.evaluate_notebook(self.absolute_path_r(notebook_file_name), 'ir')

    def evaluate_notebook(self, path, kernel_name):
        """
        Perform static and dynamic analysis of the notebook.
        Execute a notebook via nbconvert and print the results of the test (errors etc.)
        Args:
            path (string): Name of notebook to run.
            kernel_name (string): Which jupyter kernel to use to run the test.
                                  Relevant values are:'python2', 'python3', 'ir'.
        """

        dir_name, file_name = os.path.split(path)
        if dir_name:
            os.chdir(dir_name)

        print('-------- begin (kernel {0}) {1} --------'.format(kernel_name,file_name))
        no_static_errors = self.static_analysis(path)
        no_dynamic_errors = self.dynamic_analysis(path, kernel_name)
        print('-------- end (kernel {0}) {1} --------'.format(kernel_name,file_name))
        assert(no_static_errors and no_dynamic_errors)

    def static_analysis(self, path):
        """
        Perform static analysis of the notebook.
        Read the notebook and check that there is no ouput and that the links
        in the markdown cells are not broken.
        Args:
            path (string): Name of notebook.
        Return:
            boolean: True if static analysis succeeded, otherwise False.
        """

        nb = nbformat.read(path, nbformat.current_nbformat)

        #######################
        # Check that the notebook does not contain output from code cells
        # (should not be in the repository, but well...).
        #######################
        no_unexpected_output = True

        # Check that the cell dictionary has an 'outputs' key and that it is
        # empty, relies on Python using short circuit evaluation so that we
        # don't get KeyError when retrieving the 'outputs' entry.
        cells_with_output = [c.source for c in nb.cells if 'outputs' in c and c.outputs]
        if cells_with_output:
            no_unexpected_output = False
            print('Cells with unexpected output:\n_____________________________')
            for cell in cells_with_output:
                print(cell+'\n---')
        else:
           print('no unexpected output')

        #######################
        # Check that all the links in the markdown cells are valid/accessible.
        #######################
        no_broken_links = True

        cells_and_broken_links = []
        for c in nb.cells:
           if c.cell_type == 'markdown':
              html_tree = document_fromstring(markdown.markdown(c.source))
              broken_links = []
              #iterlinks() returns tuples of the form (element, attribute, link, pos)
              for document_link in html_tree.iterlinks():
                 try:
                    if all( prefix not in document_link[2] for prefix in ['http','https']):  # Local file.
                       url = 'file:' + urllib.request.pathname2url(document_link[2])
                    else:  # Remote file.
                       url = document_link[2]
                    urlopen(url)
                 except URLError:
                    broken_links.append(url)
              if broken_links:
                 cells_and_broken_links.append((broken_links,c.source))
        if cells_and_broken_links:
           no_broken_links = False
           broken_links_list = []
           print('Cells with broken links:\n________________________')
           for links, cell in cells_and_broken_links:
              print(cell+'\n')
              print('\tBroken links:')
              print('\t'+'\n\t'.join(links)+'\n---')
              broken_links_list.extend(links)
        else:
           print('no broken links')

        #######################
        # Spell check all markdown cells and comments in code cells using the pyenchant spell checker.
        #######################
        no_spelling_mistakes = True
        simpleitk_notebooks_dictionary = DictWithPWL('en_US', os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                                     'additional_dictionary.txt'))
        spell_checker = SpellChecker(simpleitk_notebooks_dictionary, filters = [EmailFilter, URLFilter])
        cells_and_spelling_mistakes = []
        for c in nb.cells:
           spelling_mistakes = []
           if c.cell_type == 'markdown':
              # Get the text as a string from the html without the markup which is replaced by space.
              spell_checker.set_text(' '.join(etree.XPath('//text()')(document_fromstring(markdown.markdown(c.source)))))
           elif c.cell_type == 'code':
              # Get all the comments and concatenate them into a single string separated by newlines.
              comment_lines = re.findall('#+.*',c.source)
              spell_checker.set_text('\n'.join(comment_lines))
           for error in spell_checker:
              error_message = 'error: '+ '\'' + error.word +'\', ' + 'suggestions: ' + str(spell_checker.suggest())
              spelling_mistakes.append(error_message)
           if spelling_mistakes:
              cells_and_spelling_mistakes.append((spelling_mistakes, c.source))
        if cells_and_spelling_mistakes:
           no_spelling_mistakes = False
           print('Cells with spelling mistakes:\n________________________')
           for misspelled_words, cell in cells_and_spelling_mistakes:
              print(cell+'\n')
              print('\tMisspelled words and suggestions:')
              print('\t'+'\n\t'.join(misspelled_words)+'\n---')
        else:
           print('no spelling mistakes')
        if not no_broken_links:
            warnings.warn('Broken links:\n'+'\n\t'.join(broken_links_list))
        return(no_unexpected_output and no_spelling_mistakes)

    def dynamic_analysis(self, path, kernel_name):
        """
        Perform dynamic analysis of the notebook.
        Execute a notebook via nbconvert and print the results of the test
        (errors etc.)
        Args:
            path (string): Name of notebook to run.
            kernel_name (string): Which jupyter kernel to use to run the test.
                                  Relevant values are:'python', 'ir'.
        Return:
            boolean: True if dynamic analysis succeeded, otherwise False.
        """

        # Execute the notebook and allow errors (run all cells), output is
        # written to a temporary file which is automatically deleted.
        # The delete=False is here to address an issue that Windows that seems to have with temporary files (see https://bugs.python.org/issue14243).
        with tempfile.NamedTemporaryFile(suffix='.ipynb',delete=False) as fout:
            args = ['jupyter', 'nbconvert',
                    '--to', 'notebook',
                    '--execute',
                    '--ExecutePreprocessor.kernel_name='+kernel_name,
                    '--ExecutePreprocessor.allow_errors=True',
                    '--ExecutePreprocessor.timeout=6000', # seconds till timeout
                    '--output', fout.name, path]
            subprocess.check_call(args)
            nb = nbformat.read(fout.name, nbformat.current_nbformat)

        # Get all of the unexpected errors (logic: cell has output with an error
        # and no error is expected or the allowed/expected error is not the one which
        # was generated.)
        unexpected_errors = [(output.evalue, c.source) for c in nb.cells \
                              if 'outputs' in c for output in c.outputs \
                              if (output.output_type=='error') and \
                               (((Test_notebooks._allowed_error_markup not in c.metadata) and (Test_notebooks._expected_error_markup not in c.metadata))or \
                               ((Test_notebooks._allowed_error_markup in c.metadata) and (c.metadata[Test_notebooks._allowed_error_markup] not in output.evalue)) or \
                               ((Test_notebooks._expected_error_markup in c.metadata) and (c.metadata[Test_notebooks._expected_error_markup] not in output.evalue)))]

        no_unexpected_errors = True
        if unexpected_errors:
            no_unexpected_errors = False
            print('Cells with unexpected errors:\n_____________________________')
            for e, src in unexpected_errors:
                print(src)
                print('unexpected error: '+e)
        else:
            print('no unexpected errors')

        # Get all of the missing expected errors (logic: cell has output
        # but expected error was not generated.)
        missing_expected_errors = []
        for c in nb.cells:
            if Test_notebooks._expected_error_markup in c.metadata:
                missing_error = True
                if 'outputs' in c:
                    for output in c.outputs:
                        if (output.output_type=='error') and (c.metadata[Test_notebooks._expected_error_markup] in output.evalue):
                                missing_error = False
                if missing_error:
                    missing_expected_errors.append((c.metadata[Test_notebooks._expected_error_markup],c.source))

        no_missing_expected_errors = True
        if missing_expected_errors:
            no_missing_expected_errors = False
            print('\nCells with missing expected errors:\n___________________________________')
            for e, src in missing_expected_errors:
                print(src)
                print('missing expected error: '+e)
        else:
            print('no missing expected errors')

        return(no_unexpected_errors and no_missing_expected_errors)

    def absolute_path_python(self, notebook_file_name):
        return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../Python', notebook_file_name))

    def absolute_path_r(self, notebook_file_name):
        return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../R', notebook_file_name))
<|endoftext|>


<commit_before># coding:utf-8
import mock
import pytest

from stellar_base.operation import *
from stellar_base.asset import Asset

class TestXdrAmount:
    def test_to_xdr_amount(self):
        assert (Operation.to_xdr_amount("20") == 20 * 10 ** 7)
        assert (Operation.to_xdr_amount("0.1234567") == 1234567)

    def test_to_xdr_amount_inexact(self):
        with pytest.raises(Exception):
            Operation.to_xdr_amount("0.12345678")

    def test_to_xdr_amount_not_number(self):
        with pytest.raises(Exception):
            # in python2.7 it will raise Exception('Invalid literal for Decimal')
            # but in python3, it will raise decimal.InvalidOperation
            Operation.to_xdr_amount("test")

    def test_to_xdr_amount_not_string(self):
        with pytest.raises(
                TypeError, match="value of type 'float' is not a string"):
            Operation.to_xdr_amount(0.1234)

    def test_from_xdr_amount(self):
        assert (Operation.from_xdr_amount(10 ** 7) == "1")
        assert (Operation.from_xdr_amount(20 * 10 ** 7) == "20")
        assert (Operation.from_xdr_amount(1234567) == "0.1234567")
        assert (Operation.from_xdr_amount(112345678) == "11.2345678")

def _load_operations():
    source = 'GDJVFDG5OCW5PYWHB64MGTHGFF57DRRJEDUEFDEL2SLNIOONHYJWHA3Z'
    seed = 'SAHPFH5CXKRMFDXEIHO6QATHJCX6PREBLCSFKYXTTCDDV6FJ3FXX4POT'
    dest = 'GCW24FUIFPC2767SOU4JI3JEAXIHYJFIJLH7GBZ2AVCBVP32SJAI53F5'
    amount = "1"
    return [
        ("create_account_min", CreateAccount({
            'source': source, 'destination': dest,
            'starting_balance': amount
        })),
        ("payment_min", Payment({
            'source': source, 'destination': dest,
            'asset': Asset.native(), 'amount': amount,
        })),
        ("allow_trust_short_asset", AllowTrust({
            'source': source, 'trustor': dest, 'asset_code': 'beer',
            'authorize': True,
        })),
        ("allow_trust_long_asset", AllowTrust({
            'source': source, 'trustor': dest,
            'asset_code': 'pocketknives', 'authorize': True,
        })),
    ]

@pytest.mark.parametrize("name, operation", _load_operations())
def test_operation(name, operation):
    operation_restored = Operation.from_xdr(operation.xdr())
    assert operation == operation_restored
    original = dict(operation.__dict__)
    original.pop("body")
    restored = dict(operation_restored.__dict__)
    restored.pop("body")
    assert original == restored

def test_from_xdr_object_raise():
    operation = mock.MagicMock(type=2561)
    pytest.raises(
        NotImplementedError, Operation.from_xdr_object, operation
    )

class TestOp:
    source = 'GDJVFDG5OCW5PYWHB64MGTHGFF57DRRJEDUEFDEL2SLNIOONHYJWHA3Z'
    seed = 'SAHPFH5CXKRMFDXEIHO6QATHJCX6PREBLCSFKYXTTCDDV6FJ3FXX4POT'
    dest = 'GCW24FUIFPC2767SOU4JI3JEAXIHYJFIJLH7GBZ2AVCBVP32SJAI53F5'
    amount = "1"

    def test_payment_short_asset(self):
        op = Payment({
            'source': self.source,
            'destination': self.dest,
            'asset': Asset('USD4', self.source),
            'amount': self.amount,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.destination == self.dest
        assert op_x.asset == Asset('USD4', self.source)
        assert op_x.amount == self.amount

    def test_payment_long_asset(self):
        op = Payment({
            'source': self.source,
            'destination': self.dest,
            'asset': Asset('SNACKS789ABC', self.source),
            'amount': self.amount,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.destination == self.dest
        assert op_x.asset == Asset('SNACKS789ABC', self.source)
        assert op_x.amount == self.amount

    def test_pathPayment_min(self):
        op = PathPayment({
            'source': self.source,
            'destination': self.dest,
            'send_asset': Asset.native(),
            'dest_asset': Asset.native(),
            'send_max': self.amount,
            'dest_amount': self.amount,
            'path': [],
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.destination == self.dest
        assert op_x.send_asset == Asset.native()
        assert op_x.dest_asset == Asset.native()
        assert op_x.send_max == self.amount
        assert op_x.dest_amount == self.amount
        assert op_x.path == []

    def test_manageOffer_min(self):
        op = ManageOffer({
            'selling': Asset('beer', self.source),
            'buying': Asset('beer', self.dest),
            'amount': "100",
            'price': 3.14159,
            'offer_id': 1,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == None
        assert op_x.selling == Asset('beer', self.source)
        assert op_x.buying == Asset('beer', self.dest)
        assert op_x.amount == "100"
        assert op_x.price == 3.14159
        assert op_x.offer_id == 1

    def test_createPassiveOffer_min(self):
        op = CreatePassiveOffer({
            'selling': Asset('beer', self.source),
            'buying': Asset('beer', self.dest),
            'amount': "100",
            'price': 3.14159,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == None
        assert op_x.selling == Asset('beer', self.source)
        assert op_x.buying == Asset('beer', self.dest)
        assert op_x.amount == "100"
        assert op_x.price == 3.14159

    def test_SetOptions_empty(self):
        op = SetOptions({})
        assert op == Operation.from_xdr(op.xdr())

    def test_changeTrust_min(self):
        op = ChangeTrust({
            'source': self.source,
            'asset': Asset('beer', self.dest),
            'limit': '100'
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.line == Asset('beer', self.dest)
        assert op_x.limit == '100'

    def test_accountMerge_min(self):
        op = AccountMerge({
            'source': self.source,
            'destination': self.dest,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.destination == self.dest

    def test_inflation(self):
        op = Inflation({
            'source': self.source,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source

    def test_manage_data(self):
        op = ManageData({
            'source': self.source,
            'data_name': '1KFHE7w8BhaENAswwryaoccDb6qcT6DbYY',
            'data_value': self.source,
        })
        op_x = Operation.from_xdr(op.xdr())
        assert op == op_x
        assert op_x.source == self.source
        assert op_x.data_name == '1KFHE7w8BhaENAswwryaoccDb6qcT6DbYY'
        assert op_x.data_value == self.source

    def test_manage_data_too_long(self):
        with pytest.raises(
                XdrLengthError,
                match='Data or value should be <= 64 bytes \(ascii encoded\).'
        ):
            ManageData({
                'data_name': '1KFHE7w8BhaENAswwryaoccDb6qcT6DbYY',
                'data_value': '1234567890' * 7
            })
<commit_msg>refactor all tests in test_operation<commit_after># coding:utf-8
import mock
import pytest

from stellar_base.operation import *
from stellar_base.asset import Asset

class TestXdrAmount:
    def test_to_xdr_amount(self):
        assert (Operation.to_xdr_amount("20") == 20 * 10 ** 7)
        assert (Operation.to_xdr_amount("0.1234567") == 1234567)

    def test_to_xdr_amount_inexact(self):
        with pytest.raises(Exception):
            Operation.to_xdr_amount("0.12345678")

    def test_to_xdr_amount_not_number(self):
        with pytest.raises(Exception):
            # in python2.7 it will raise Exception('Invalid literal for Decimal')
            # but in python3, it will raise decimal.InvalidOperation
            Operation.to_xdr_amount("test")

    def test_to_xdr_amount_not_string(self):
        with pytest.raises(
                TypeError, match="value of type 'float' is not a string"):
            Operation.to_xdr_amount(0.1234)

    def test_from_xdr_amount(self):
        assert (Operation.from_xdr_amount(10 ** 7) == "1")
        assert (Operation.from_xdr_amount(20 * 10 ** 7) == "20")
        assert (Operation.from_xdr_amount(1234567) == "0.1234567")
        assert (Operation.from_xdr_amount(112345678) == "11.2345678")

def _load_operations():
    source = 'GDJVFDG5OCW5PYWHB64MGTHGFF57DRRJEDUEFDEL2SLNIOONHYJWHA3Z'
    dest = 'GCW24FUIFPC2767SOU4JI3JEAXIHYJFIJLH7GBZ2AVCBVP32SJAI53F5'
    amount = "1"
    return [
        ("create_account_min", CreateAccount({
            'source': source, 'destination': dest,
            'starting_balance': amount
        })),
        ("payment_min", Payment({
            'source': source, 'destination': dest,
            'asset': Asset.native(), 'amount': amount,
        })),
        ("payment_short_asset", Payment({
            'source': source, 'destination': dest,
            'asset': Asset('USD4', source), 'amount': amount,
        })),
        ("payment_long_asset", Payment({
            'source': source, 'destination': dest,
            'asset': Asset('SNACKS789ABC', source), 'amount': amount,
        })),
        ("path_payment_min", PathPayment({
            'source': source, 'destination': dest,
            'send_asset': Asset.native(), 'dest_asset': Asset.native(),
            'send_max': amount, 'dest_amount': amount, 'path': [],
        })),
        ("allow_trust_short_asset", AllowTrust({
            'source': source, 'trustor': dest, 'asset_code': 'beer',
            'authorize': True,
        })),
        ("allow_trust_long_asset", AllowTrust({
            'source': source, 'trustor': dest,
            'asset_code': 'pocketknives', 'authorize': True,
        })),
        ("manage_offer_min", ManageOffer({
            'selling': Asset('beer', source), 'buying': Asset('beer', dest),
            'amount': "100", 'price': 3.14159, 'offer_id': 1,
        })),
        ("create_passive_offer_min", CreatePassiveOffer({
            'selling': Asset('beer', source), 'buying': Asset('beer', dest),
            'amount': "100", 'price': 3.14159,
        })),
        ("set_options_empty", SetOptions({})),
        ("change_trust_min", ChangeTrust({
            'source': source, 'asset': Asset('beer', dest), 'limit': '100'
        })),
        ("account_merge_min", AccountMerge({
            'source': source, 'destination': dest,
        })),
        ("inflation", Inflation({'source': source})),
        ("manage_data", ManageData({
            'source': source,
            'data_name': '1KFHE7w8BhaENAswwryaoccDb6qcT6DbYY',
            'data_value': source,
        })),
    ]

@pytest.mark.parametrize("name, operation", _load_operations())
def test_operation(name, operation):
    operation_restored = Operation.from_xdr(operation.xdr())
    assert operation == operation_restored
    original = dict(operation.__dict__)
    original.pop("body")
    restored = dict(operation_restored.__dict__)
    restored.pop("body")
    assert original == restored

def test_from_xdr_object_raise():
    operation = mock.MagicMock(type=2561)
    pytest.raises(NotImplementedError, Operation.from_xdr_object, operation)

def test_manage_data_too_long_raises():
    msg = 'Data or value should be <= 64 bytes \(ascii encoded\).'
    with pytest.raises(XdrLengthError, match=msg):
        ManageData({
            'data_name': '1KFHE7w8BhaENAswwryaoccDb6qcT6DbYY',
            'data_value': '1234567890' * 7
        })
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
test_parameter
----------------------------------

Tests Uniform and Normal parameter priors and sampling functions
"""

import unittest

import numpy as np
import scipy.stats

from enterprise.signals.parameter import UniformPrior, UniformSampler, Uniform
from enterprise.signals.parameter import NormalPrior, NormalSampler, Normal
from enterprise.signals.parameter import TruncNormalPrior, TruncNormalSampler, TruncNormal
from enterprise.signals.parameter import LinearExpPrior, LinearExpSampler

class TestParameter(unittest.TestCase):
    def test_uniform(self):
        """Test Uniform parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        p_min, p_max = 0.2, 1.1
        x = 0.5

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = UniformSampler(p_min, p_max)
        assert p_min < x1 < p_max, msg2
        assert type(x1) == float, msg2

        # vector argument
        x = np.array([0.5, 0.1])
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        x1 = UniformSampler(p_min, p_max, size=3)
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (3,), msg2

        # vector bounds
        p_min, p_max = np.array([0.2, 0.3]), np.array([1.1, 1.2])
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        x1 = UniformSampler(p_min, p_max)
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (2,), msg2

        x1 = UniformSampler(p_min, p_max, size=(3, 2))
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (3, 2), msg2

    def test_linearexp(self):
        """Test LinearExp parameter prior and sampler."""

        # scalar
        p_min, p_max = 1, 3
        x = 2
        msg1 = "Scalar prior does not match"
        assert np.allclose(LinearExpPrior(x, p_min, p_max), np.log(10) * 10**2 / (10**3 - 10**1)), msg1

        x = LinearExpSampler(p_min, p_max)
        msg1b = "Scalar sampler out of range"
        assert p_min <= x <= p_max, msg1b

        # vector argument
        x = np.array([0, 1.5, 2.5])
        msg2 = "Vector-argument prior does not match"
        assert np.allclose(
            LinearExpPrior(x, p_min, p_max), np.array([0, 10**1.5, 10**2.5]) * np.log(10) / (10**3 - 10**1)
        ), msg2

        x = LinearExpSampler(p_min, p_max, size=10)
        msg2b = "Vector-argument sampler out of range"
        assert np.all((p_min < x) & (x < p_max)), msg2b

        # vector bounds
        p_min, p_max = np.array([0, 1]), np.array([2, 3])
        x = np.array([1, 2])
        msg3 = "Vector-argument+bounds prior does not match"
        assert np.allclose(
            LinearExpPrior(x, p_min, p_max),
            np.array([10**1 / (10**2 - 10**0), 10**2 / (10**3 - 10**1)]) * np.log(10),
        ), msg3

    def test_normal(self):
        """Test Normal parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        mu, sigma = 0.1, 2
        x = 0.5

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(
            NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = NormalSampler(mu, sigma)
        # this should almost never fail
        assert -5 < (x1 - mu) / sigma < 5, msg2

        # vector argument
        x = np.array([-0.2, 0.1, 0.5])

        assert np.allclose(
            NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        x1, x2 = (
            NormalSampler(mu, sigma, size=10),
            scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma**2, size=10),
        )
        assert x1.shape == x2.shape, msg2

        # vector bounds; note the different semantics from `NormalPrior`,
        # which returns a vector consistently with `UniformPrior`
        mu, sigma = np.array([0.1, 0.15, 0.2]), np.array([2, 1, 2])
        assert np.allclose(
            np.prod(NormalPrior(x, mu, sigma)), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        x1, x2 = NormalSampler(mu, sigma), scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma**2)
        assert x1.shape == x2.shape, msg2

        # matrix covariance
        sigma = np.array([[1, 0.1, 0], [0.1, 1, 0], [0, 0, 1]])

        assert np.allclose(NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma)), msg1

        x1, x2 = NormalSampler(mu, sigma), scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma)
        assert x1.shape == x2.shape, msg2

    def test_truncnormal(self):
        """Test TruncNormal parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        mu, sigma, pmin, pmax = 0.1, 2, -2, 3
        x = 0.5
        a, b = (pmin - mu) / sigma, (pmax - mu) / sigma

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = TruncNormalSampler(mu, sigma, pmin, pmax)
        # this should almost never fail
        assert -5 < (x1 - mu) / sigma < 5, msg2

        # vector argument
        x = np.array([-0.2, 0.1, 0.5])

        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        x1, x2 = (
            TruncNormalSampler(mu, sigma, pmin, pmax, size=10),
            scipy.stats.truncnorm.rvs(a, b, mu, sigma, size=10),
        )
        assert x1.shape == x2.shape, msg2

        # vector bounds; note the different semantics from `NormalPrior`,
        # which returns a vector consistently with `UniformPrior`
        mu, sigma = np.array([0.1, 0.15, 0.2]), np.array([2, 1, 2])
        pmin, pmax = np.array([-2, -2, -3]), np.array([1, 2, 1])
        a, b = (pmin - mu) / sigma, (pmax - mu) / sigma
        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        x1, x2 = TruncNormalSampler(mu, sigma, pmin, pmax), scipy.stats.truncnorm.rvs(a, b, mu, sigma)
        assert x1.shape == x2.shape, msg2

    def test_normalandtruncnormal(self):
        mu, sigma = 0, 1

        msg3 = "Normal and [-inf, inf] TruncNormal do not match"

        paramA = Normal(mu, sigma)("A")
        paramB = TruncNormal(mu, sigma, -np.inf, np.inf)("B")
        xs = np.linspace(-3, 3, 20)
        assert np.allclose(paramA.get_pdf(xs), paramB.get_pdf(xs)), msg3

    def test_metaparam(self):
        mu = Uniform(-1, 1)("mean")
        sigma = 2
        pmin, pmax = -3, 3

        msg4 = "problem with meta-parameter in TruncNormal"
        zeros = np.zeros(2)

        paramA = TruncNormal(mu, sigma, pmin, pmax)("A")
        xs = np.array([-3.5, 3.5])
        assert np.alltrue(paramA.get_pdf(xs, mu=mu.sample()) == zeros), msg4
<commit_msg>linting<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
test_parameter
----------------------------------

Tests Uniform and Normal parameter priors and sampling functions
"""

import unittest

import numpy as np
import scipy.stats

from enterprise.signals.parameter import UniformPrior, UniformSampler, Uniform
from enterprise.signals.parameter import NormalPrior, NormalSampler, Normal
from enterprise.signals.parameter import TruncNormalPrior, TruncNormalSampler, TruncNormal
from enterprise.signals.parameter import LinearExpPrior, LinearExpSampler

class TestParameter(unittest.TestCase):
    def test_uniform(self):
        """Test Uniform parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        p_min, p_max = 0.2, 1.1
        x = 0.5

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = UniformSampler(p_min, p_max)
        assert p_min < x1 < p_max, msg2
        assert type(x1) == float, msg2

        # vector argument
        x = np.array([0.5, 0.1])
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        x1 = UniformSampler(p_min, p_max, size=3)
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (3,), msg2

        # vector bounds
        p_min, p_max = np.array([0.2, 0.3]), np.array([1.1, 1.2])
        assert np.allclose(UniformPrior(x, p_min, p_max), scipy.stats.uniform.pdf(x, p_min, p_max - p_min)), msg1

        x1 = UniformSampler(p_min, p_max)
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (2,), msg2

        x1 = UniformSampler(p_min, p_max, size=(3, 2))
        assert np.all((p_min < x1) & (x1 < p_max)), msg2
        assert x1.shape == (3, 2), msg2

    def test_linearexp(self):
        """Test LinearExp parameter prior and sampler."""

        # scalar
        p_min, p_max = 1, 3
        x = 2
        msg1 = "Scalar prior does not match"
        assert np.allclose(LinearExpPrior(x, p_min, p_max), np.log(10) * 10**2 / (10**3 - 10**1)), msg1

        x = LinearExpSampler(p_min, p_max)
        msg1b = "Scalar sampler out of range"
        assert p_min <= x <= p_max, msg1b

        # vector argument
        x = np.array([0, 1.5, 2.5])
        msg2 = "Vector-argument prior does not match"
        assert np.allclose(
            LinearExpPrior(x, p_min, p_max), np.array([0, 10**1.5, 10**2.5]) * np.log(10) / (10**3 - 10**1)
        ), msg2

        x = LinearExpSampler(p_min, p_max, size=10)
        msg2b = "Vector-argument sampler out of range"
        assert np.all((p_min < x) & (x < p_max)), msg2b

        # vector bounds
        p_min, p_max = np.array([0, 1]), np.array([2, 3])
        x = np.array([1, 2])
        msg3 = "Vector-argument+bounds prior does not match"
        assert np.allclose(
            LinearExpPrior(x, p_min, p_max),
            np.array([10**1 / (10**2 - 10**0), 10**2 / (10**3 - 10**1)]) * np.log(10),
        ), msg3

    def test_normal(self):
        """Test Normal parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        mu, sigma = 0.1, 2
        x = 0.5

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(
            NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = NormalSampler(mu, sigma)
        # this should almost never fail
        assert -5 < (x1 - mu) / sigma < 5, msg2

        # vector argument
        x = np.array([-0.2, 0.1, 0.5])

        assert np.allclose(
            NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        x1, x2 = (
            NormalSampler(mu, sigma, size=10),
            scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma**2, size=10),
        )
        assert x1.shape == x2.shape, msg2

        # vector bounds; note the different semantics from `NormalPrior`,
        # which returns a vector consistently with `UniformPrior`
        mu, sigma = np.array([0.1, 0.15, 0.2]), np.array([2, 1, 2])
        assert np.allclose(
            np.prod(NormalPrior(x, mu, sigma)), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma**2)
        ), msg1

        x1, x2 = NormalSampler(mu, sigma), scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma**2)
        assert x1.shape == x2.shape, msg2

        # matrix covariance
        sigma = np.array([[1, 0.1, 0], [0.1, 1, 0], [0, 0, 1]])

        assert np.allclose(NormalPrior(x, mu, sigma), scipy.stats.multivariate_normal.pdf(x, mean=mu, cov=sigma)), msg1

        x1, x2 = NormalSampler(mu, sigma), scipy.stats.multivariate_normal.rvs(mean=mu, cov=sigma)
        assert x1.shape == x2.shape, msg2

    def test_truncnormal(self):
        """Test TruncNormal parameter prior and sampler for various combinations of scalar and vector arguments."""

        # scalar
        mu, sigma, pmin, pmax = 0.1, 2, -2, 3
        x = 0.5
        a, b = (pmin - mu) / sigma, (pmax - mu) / sigma

        msg1 = "Enterprise and scipy priors do not match"
        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        msg2 = "Enterprise samples have wrong value, type, or size"
        x1 = TruncNormalSampler(mu, sigma, pmin, pmax)
        # this should almost never fail
        assert -5 < (x1 - mu) / sigma < 5, msg2

        # vector argument
        x = np.array([-0.2, 0.1, 0.5])

        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        x1, x2 = (
            TruncNormalSampler(mu, sigma, pmin, pmax, size=10),
            scipy.stats.truncnorm.rvs(a, b, mu, sigma, size=10),
        )
        assert x1.shape == x2.shape, msg2

        # vector bounds; note the different semantics from `NormalPrior`,
        # which returns a vector consistently with `UniformPrior`
        mu, sigma = np.array([0.1, 0.15, 0.2]), np.array([2, 1, 2])
        pmin, pmax = np.array([-2, -2, -3]), np.array([1, 2, 1])
        a, b = (pmin - mu) / sigma, (pmax - mu) / sigma
        assert np.allclose(
            TruncNormalPrior(x, mu, sigma, pmin, pmax), scipy.stats.truncnorm.pdf(x, a, b, mu, sigma)
        ), msg1

        x1, x2 = TruncNormalSampler(mu, sigma, pmin, pmax), scipy.stats.truncnorm.rvs(a, b, mu, sigma)
        assert x1.shape == x2.shape, msg2

    def test_normalandtruncnormal(self):
        mu, sigma = 0, 1

        msg3 = "Normal and [-inf, inf] TruncNormal do not match"

        paramA = Normal(mu, sigma)("A")
        paramB = TruncNormal(mu, sigma, -np.inf, np.inf)("B")
        xs = np.linspace(-3, 3, 20)
        assert np.allclose(paramA.get_pdf(xs), paramB.get_pdf(xs)), msg3

    def test_metaparam(self):
        mu = Uniform(-1, 1)("mean")
        sigma = 2
        pmin, pmax = -3, 3

        msg4 = "problem with meta-parameter in TruncNormal"
        zeros = np.zeros(2)

        paramA = TruncNormal(mu, sigma, pmin, pmax)("A")
        xs = np.array([-3.5, 3.5])
        assert np.alltrue(paramA.get_pdf(xs, mu=mu.sample()) == zeros), msg4
<|endoftext|>