<commit_before># Standard library
import os

# Third-party
import numpy as np
import pymc3 as pm
import exoplanet as xo
import exoplanet.units as xu
import theano.tensor as tt

# Project
from .logging import logger
from .data_helpers import validate_prepare_data
from .samples_helpers import is_P_unimodal
from .likelihood_helpers import get_trend_design_matrix
from .prior_helpers import validate_n_offsets, validate_poly_trend
from .prior import JokerPrior, _validate_model
from .src.fast_likelihood import CJokerHelper
from .utils import tempfile_decorator
from .samples import JokerSamples

__all__ = ['TheJoker']

class TheJoker:
    """
    A custom Monte-Carlo sampler for two-body systems.

    Parameters
    ----------
    prior : `~thejoker.JokerPrior`
        The specification of the prior probability distribution over all
        parameters used in The Joker.
    pool : `schwimmbad.BasePool` (optional)
        A processing pool (default is a `schwimmbad.SerialPool` instance).
    random_state : `numpy.random.RandomState` (optional)
        A `numpy.random.RandomState` instance for controlling random number
        generation.
    tempfile_path : str (optional)
        A location on disk where The Joker may store some temporary files. Any
        files written here by The Joker should be cleaned up: If any files in
        this path persist, something must have gone wrong within The Joker.
        Default: ``~/.thejoker``
    """

    def __init__(self, prior, pool=None, random_state=None, tempfile_path=None):

        # set the processing pool
        if pool is None:
            import schwimmbad
            pool = schwimmbad.SerialPool()
        elif not hasattr(pool, 'map') or not hasattr(pool, 'close'):
            raise TypeError("Input pool object must have .map() and .close() "
                            "methods. We recommend using `schwimmbad` pools.")
        self.pool = pool

        # Set the parent random state - child processes get different states
        # based on the parent
        if random_state is None:
            random_state = np.random.RandomState()
        elif not isinstance(random_state, np.random.RandomState):
            raise TypeError("Random state object must be a numpy RandomState "
                            "instance, not '{0}'".format(type(random_state)))
        self.random_state = random_state

        # check if a JokerParams instance was passed in to specify the state
        if not isinstance(prior, JokerPrior):
            raise TypeError("The input prior must be a JokerPrior instance.")
        self.prior = prior

        if tempfile_path is None:
            self._tempfile_path = os.path.expanduser('~/.thejoker/')
        else:
            self._tempfile_path = os.path.abspath(
                os.path.expanduser(tempfile_path))

    @property
    def tempfile_path(self):
        os.makedirs(self._tempfile_path, exist_ok=True)
        return self._tempfile_path

    def _make_joker_helper(self, data):
        all_data, ids, trend_M = validate_prepare_data(data,
                                                       self.prior.poly_trend,
                                                       self.prior.n_offsets)
        joker_helper = CJokerHelper(all_data, self.prior, trend_M)
        return joker_helper

    @tempfile_decorator
    def marginal_ln_likelihood(self, data, prior_samples, n_batches=None):
        f"""
        Compute the marginal log-likelihood at each of the input prior samples.

        Parameters
        ----------
        data : `thejoker.RVData`, iterable, dict
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.

        Returns
        -------
        ln_likelihood : `numpy.ndarray`
            The marginal log-likelihood computed at the location of each prior
            sample.
        """
        from .multiproc_helpers import marginal_ln_likelihood_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        return marginal_ln_likelihood_helper(joker_helper, prior_samples,
                                             self.pool, n_batches=n_batches)

    @tempfile_decorator
    def rejection_sample(self, data, prior_samples,
                         n_prior_samples=None,
                         max_posterior_samples=None,
                         n_linear_samples=1,
                         return_logprobs=False,
                         n_batches=None,
                         randomize_prior_order=False):
        f"""
        Run The Joker's rejection sampling on prior samples to get posterior
        samples for the input data.

        You must either specify the number of prior samples to generate and
        use for rejection sampling, ``n_prior_samples``, or the path to a file
        containing prior samples, ``prior_cache_file``.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `~thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_prior_samples : int (optional)
            The number of prior samples to run on. This is only used if passing
            in a string filename: If the file contains a large number of prior
            samples, you may want to set this to only run on a subset.
        max_posterior_samples : int (optional)
            The maximum number of posterior samples to generate. If using a
            large library of prior samples, and running on uninformative data,
            you may want to set this to a small but reasonable number (like,
            256).
        n_linear_samples : int (optional)
            The number of linear parameter samples to generate for each
            nonlinear parameter sample returned from the rejection sampling
            step.
        return_logprobs : bool (optional)
            Also return the log-prior and (marginal) log-likelihood values
            evaluated at each sample.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.
        randomize_prior_order : bool (optional)
            Randomly shuffle the prior samples before reading and running the
            rejection sampler. This is only useful if you are using a large
            library of prior samples, and choosing to run on a subset of those
            samples.

        Returns
        -------
        samples : `~thejoker.JokerSamples`
            The posterior samples produced from The Joker.

        """
        from .multiproc_helpers import rejection_sample_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        samples = rejection_sample_helper(
            joker_helper,
            prior_samples,
            pool=self.pool,
            random_state=self.random_state,
            n_prior_samples=n_prior_samples,
            max_posterior_samples=max_posterior_samples,
            n_linear_samples=n_linear_samples,
            return_logprobs=return_logprobs,
            n_batches=n_batches,
            randomize_prior_order=randomize_prior_order)

        return samples

    @tempfile_decorator
    def iterative_rejection_sample(self, data, prior_samples,
                                   n_requested_samples,
                                   max_prior_samples=None,
                                   n_linear_samples=1,
                                   return_logprobs=False,
                                   n_batches=None,
                                   randomize_prior_order=False,
                                   init_batch_size=None,
                                   growth_factor=128):

        """This is an experimental sampling method that adaptively generates
        posterior samples given a large library of prior samples. The advantage
        of this function over the standard ``rejection_sample`` method is that
        it will try to adaptively figure out how many prior samples it needs to
        evaluate the likelihood at in order to return the desired number of
        posterior samples.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `~thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_requested_samples : int (optional)
            The number of posterior samples desired.
        max_prior_samples : int (optional)
            The maximum number of prior samples to process.
        n_linear_samples : int (optional)
            The number of linear parameter samples to generate for each
            nonlinear parameter sample returned from the rejection sampling
            step.
        return_logprobs : bool (optional)
            Also return the log-prior and (marginal) log-likelihood values
            evaluated at each sample.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.
        randomize_prior_order : bool (optional)
            Randomly shuffle the prior samples before reading and running the
            rejection sampler. This is only useful if you are using a large
            library of prior samples, and choosing to run on a subset of those
            samples.
        init_batch_size : int (optional)
            The initial batch size of likelihoods to compute, before growing
            the batches using the multiplicative growth factor, below.
        growth_factor : int (optional)
            A factor used to adaptively grow the number of prior samples to
            evaluate on. Larger numbers make the trial batches grow faster.

        Returns
        -------
        samples : `~thejoker.JokerSamples`
            The posterior samples produced from The Joker.
        """
        from .multiproc_helpers import iterative_rejection_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        samples = iterative_rejection_helper(
            joker_helper,
            prior_samples,
            init_batch_size=init_batch_size,
            growth_factor=growth_factor,
            pool=self.pool,
            random_state=self.random_state,
            n_requested_samples=n_requested_samples,
            max_prior_samples=max_prior_samples,
            n_linear_samples=n_linear_samples,
            return_logprobs=return_logprobs,
            n_batches=n_batches,
            randomize_prior_order=randomize_prior_order)

        return samples

    def setup_mcmc(self, data, joker_samples, model=None):
        """
        Setup the model to run MCMC using pymc3.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        joker_samples : `~thejoker.JokerSamples`
            If a single sample is passed in, this is packed into a pymc3
            initialization dictionary and returned after setting up. If
            multiple samples are passed in, the median (along period) sample is
            taken and returned after setting up for MCMC.
        model : `pymc3.Model`
            This is either required, or this function must be called within a
            pymc3 model context.

        Returns
        -------
        mcmc_init : dict

        """

        model = _validate_model(model)

        # Reduce data, strip units:
        data, ids, _ = validate_prepare_data(data,
                                             self.prior.poly_trend,
                                             self.prior.n_offsets)
        x = data._t_bmjd - data._t0_bmjd
        y = data.rv.value
        err = data.rv_err.to_value(data.rv.unit)

        # First, prepare the joker_samples:
        if not isinstance(joker_samples, JokerSamples):
            raise TypeError("You must pass in a JokerSamples instance to the "
                            "joker_samples argument.")

        if len(joker_samples) > 1:
            # check if unimodal in P, if not, warn
            if not is_P_unimodal(joker_samples, data):
                logger.warn("TODO: samples ain't unimodal")

            joker_samples = joker_samples.median()

        mcmc_init = dict()
        for name in self.prior.par_names:
            unit = getattr(self.prior.pars[name], xu.UNIT_ATTR_NAME)
            mcmc_init[name] = joker_samples[name].to_value(unit)[0]

        if 't_peri' in model.named_vars:
            logger.info('pymc3 model has already been setup for running MCMC - '
                        'using the previously setup model parameters.')
            return mcmc_init

        p = self.prior.pars
        with model:
            t_peri = pm.Deterministic('t_peri',
                                      p['P'] * p['M0'] / (2*np.pi))

            # Set up the orbit model
            orbit = xo.orbits.KeplerianOrbit(period=p['P'],
                                             ecc=p['e'],
                                             omega=p['omega'],
                                             t_periastron=t_peri)

        # design matrix
        M = get_trend_design_matrix(data, ids, self.prior.poly_trend)

        # deal with v0_offsets, trend here:
        _, offset_names = validate_n_offsets(self.prior.n_offsets)
        _, vtrend_names = validate_poly_trend(self.prior.poly_trend)

        with model:
            v_pars = ([p['v0']]
                      + [p[name] for name in offset_names]
                      + [p[name] for name in vtrend_names[1:]])  # skip v0
            v_trend_vec = tt.stack(v_pars, axis=0)
            trend = tt.dot(M, v_trend_vec)

            rv_model = orbit.get_radial_velocity(x, K=p['K']) + trend

            err = tt.sqrt(err**2 + p['s']**2)
            pm.Normal("obs", mu=rv_model, sd=err, observed=y)

        return mcmc_init

    def trace_to_samples(self, trace, remove_constants=True):
        """
        Create a ``JokerSamples`` instance from a pymc3 trace object.

        Parameters
        ----------
        trace : `~pymc3.backends.base.MultiTrace`
        """
        from pymc3.model import Factor, TransformedRV

        df = pm.trace_to_dataframe(trace)

        samples = JokerSamples(poly_trend=self.prior.poly_trend,
                               n_offsets=self.prior.n_offsets)
        for name in self.prior.par_names:
            par = self.prior.pars[name]
            if (remove_constants
                and isinstance(par, tt.TensorVariable)
                and not (isinstance(par, Factor)
                         or isinstance(par, TransformedRV))):
                continue

            unit = getattr(par, xu.UNIT_ATTR_NAME)
            samples[name] = df[name].values * unit

        return samples
<commit_msg>fix trace to samples<commit_after># Standard library
import os

# Third-party
import numpy as np
import pymc3 as pm
import exoplanet as xo
import exoplanet.units as xu
import theano.tensor as tt

# Project
from .logging import logger
from .data_helpers import validate_prepare_data
from .samples_helpers import is_P_unimodal
from .likelihood_helpers import get_trend_design_matrix
from .prior_helpers import validate_n_offsets, validate_poly_trend
from .prior import JokerPrior, _validate_model
from .src.fast_likelihood import CJokerHelper
from .utils import tempfile_decorator
from .samples import JokerSamples

__all__ = ['TheJoker']

class TheJoker:
    """
    A custom Monte-Carlo sampler for two-body systems.

    Parameters
    ----------
    prior : `~thejoker.JokerPrior`
        The specification of the prior probability distribution over all
        parameters used in The Joker.
    pool : `schwimmbad.BasePool` (optional)
        A processing pool (default is a `schwimmbad.SerialPool` instance).
    random_state : `numpy.random.RandomState` (optional)
        A `numpy.random.RandomState` instance for controlling random number
        generation.
    tempfile_path : str (optional)
        A location on disk where The Joker may store some temporary files. Any
        files written here by The Joker should be cleaned up: If any files in
        this path persist, something must have gone wrong within The Joker.
        Default: ``~/.thejoker``
    """

    def __init__(self, prior, pool=None, random_state=None, tempfile_path=None):

        # set the processing pool
        if pool is None:
            import schwimmbad
            pool = schwimmbad.SerialPool()
        elif not hasattr(pool, 'map') or not hasattr(pool, 'close'):
            raise TypeError("Input pool object must have .map() and .close() "
                            "methods. We recommend using `schwimmbad` pools.")
        self.pool = pool

        # Set the parent random state - child processes get different states
        # based on the parent
        if random_state is None:
            random_state = np.random.RandomState()
        elif not isinstance(random_state, np.random.RandomState):
            raise TypeError("Random state object must be a numpy RandomState "
                            "instance, not '{0}'".format(type(random_state)))
        self.random_state = random_state

        # check if a JokerParams instance was passed in to specify the state
        if not isinstance(prior, JokerPrior):
            raise TypeError("The input prior must be a JokerPrior instance.")
        self.prior = prior

        if tempfile_path is None:
            self._tempfile_path = os.path.expanduser('~/.thejoker/')
        else:
            self._tempfile_path = os.path.abspath(
                os.path.expanduser(tempfile_path))

    @property
    def tempfile_path(self):
        os.makedirs(self._tempfile_path, exist_ok=True)
        return self._tempfile_path

    def _make_joker_helper(self, data):
        all_data, ids, trend_M = validate_prepare_data(data,
                                                       self.prior.poly_trend,
                                                       self.prior.n_offsets)
        joker_helper = CJokerHelper(all_data, self.prior, trend_M)
        return joker_helper

    @tempfile_decorator
    def marginal_ln_likelihood(self, data, prior_samples, n_batches=None):
        f"""
        Compute the marginal log-likelihood at each of the input prior samples.

        Parameters
        ----------
        data : `thejoker.RVData`, iterable, dict
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.

        Returns
        -------
        ln_likelihood : `numpy.ndarray`
            The marginal log-likelihood computed at the location of each prior
            sample.
        """
        from .multiproc_helpers import marginal_ln_likelihood_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        return marginal_ln_likelihood_helper(joker_helper, prior_samples,
                                             self.pool, n_batches=n_batches)

    @tempfile_decorator
    def rejection_sample(self, data, prior_samples,
                         n_prior_samples=None,
                         max_posterior_samples=None,
                         n_linear_samples=1,
                         return_logprobs=False,
                         n_batches=None,
                         randomize_prior_order=False):
        f"""
        Run The Joker's rejection sampling on prior samples to get posterior
        samples for the input data.

        You must either specify the number of prior samples to generate and
        use for rejection sampling, ``n_prior_samples``, or the path to a file
        containing prior samples, ``prior_cache_file``.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `~thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_prior_samples : int (optional)
            The number of prior samples to run on. This is only used if passing
            in a string filename: If the file contains a large number of prior
            samples, you may want to set this to only run on a subset.
        max_posterior_samples : int (optional)
            The maximum number of posterior samples to generate. If using a
            large library of prior samples, and running on uninformative data,
            you may want to set this to a small but reasonable number (like,
            256).
        n_linear_samples : int (optional)
            The number of linear parameter samples to generate for each
            nonlinear parameter sample returned from the rejection sampling
            step.
        return_logprobs : bool (optional)
            Also return the log-prior and (marginal) log-likelihood values
            evaluated at each sample.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.
        randomize_prior_order : bool (optional)
            Randomly shuffle the prior samples before reading and running the
            rejection sampler. This is only useful if you are using a large
            library of prior samples, and choosing to run on a subset of those
            samples.

        Returns
        -------
        samples : `~thejoker.JokerSamples`
            The posterior samples produced from The Joker.

        """
        from .multiproc_helpers import rejection_sample_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        samples = rejection_sample_helper(
            joker_helper,
            prior_samples,
            pool=self.pool,
            random_state=self.random_state,
            n_prior_samples=n_prior_samples,
            max_posterior_samples=max_posterior_samples,
            n_linear_samples=n_linear_samples,
            return_logprobs=return_logprobs,
            n_batches=n_batches,
            randomize_prior_order=randomize_prior_order)

        return samples

    @tempfile_decorator
    def iterative_rejection_sample(self, data, prior_samples,
                                   n_requested_samples,
                                   max_prior_samples=None,
                                   n_linear_samples=1,
                                   return_logprobs=False,
                                   n_batches=None,
                                   randomize_prior_order=False,
                                   init_batch_size=None,
                                   growth_factor=128):

        """This is an experimental sampling method that adaptively generates
        posterior samples given a large library of prior samples. The advantage
        of this function over the standard ``rejection_sample`` method is that
        it will try to adaptively figure out how many prior samples it needs to
        evaluate the likelihood at in order to return the desired number of
        posterior samples.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        prior_samples : str, `~thejoker.JokerSamples`
            Either a path to a file containing prior samples generated from The
            Joker, or a `~thejoker.JokerSamples` instance containing the prior
            samples.
        n_requested_samples : int (optional)
            The number of posterior samples desired.
        max_prior_samples : int (optional)
            The maximum number of prior samples to process.
        n_linear_samples : int (optional)
            The number of linear parameter samples to generate for each
            nonlinear parameter sample returned from the rejection sampling
            step.
        return_logprobs : bool (optional)
            Also return the log-prior and (marginal) log-likelihood values
            evaluated at each sample.
        n_batches : int (optional)
            The number of batches to split the prior samples into before
            distributing for computation. If using the (default) serial
            computation pool, this doesn't have any impact. If using
            multiprocessing or MPI, this determines how many batches to split
            the samples into before scattering over all workers.
        randomize_prior_order : bool (optional)
            Randomly shuffle the prior samples before reading and running the
            rejection sampler. This is only useful if you are using a large
            library of prior samples, and choosing to run on a subset of those
            samples.
        init_batch_size : int (optional)
            The initial batch size of likelihoods to compute, before growing
            the batches using the multiplicative growth factor, below.
        growth_factor : int (optional)
            A factor used to adaptively grow the number of prior samples to
            evaluate on. Larger numbers make the trial batches grow faster.

        Returns
        -------
        samples : `~thejoker.JokerSamples`
            The posterior samples produced from The Joker.
        """
        from .multiproc_helpers import iterative_rejection_helper
        joker_helper = self._make_joker_helper(data)  # also validates data
        samples = iterative_rejection_helper(
            joker_helper,
            prior_samples,
            init_batch_size=init_batch_size,
            growth_factor=growth_factor,
            pool=self.pool,
            random_state=self.random_state,
            n_requested_samples=n_requested_samples,
            max_prior_samples=max_prior_samples,
            n_linear_samples=n_linear_samples,
            return_logprobs=return_logprobs,
            n_batches=n_batches,
            randomize_prior_order=randomize_prior_order)

        return samples

    def setup_mcmc(self, data, joker_samples, model=None):
        """
        Setup the model to run MCMC using pymc3.

        Parameters
        ----------
        data : `~thejoker.RVData`
            The radial velocity data, or an iterable containing ``RVData``
            objects for each data source.
        joker_samples : `~thejoker.JokerSamples`
            If a single sample is passed in, this is packed into a pymc3
            initialization dictionary and returned after setting up. If
            multiple samples are passed in, the median (along period) sample is
            taken and returned after setting up for MCMC.
        model : `pymc3.Model`
            This is either required, or this function must be called within a
            pymc3 model context.

        Returns
        -------
        mcmc_init : dict

        """

        model = _validate_model(model)

        # Reduce data, strip units:
        data, ids, _ = validate_prepare_data(data,
                                             self.prior.poly_trend,
                                             self.prior.n_offsets)
        x = data._t_bmjd - data._t0_bmjd
        y = data.rv.value
        err = data.rv_err.to_value(data.rv.unit)

        # First, prepare the joker_samples:
        if not isinstance(joker_samples, JokerSamples):
            raise TypeError("You must pass in a JokerSamples instance to the "
                            "joker_samples argument.")

        if len(joker_samples) > 1:
            # check if unimodal in P, if not, warn
            if not is_P_unimodal(joker_samples, data):
                logger.warn("TODO: samples ain't unimodal")

            joker_samples = joker_samples.median()

        mcmc_init = dict()
        for name in self.prior.par_names:
            unit = getattr(self.prior.pars[name], xu.UNIT_ATTR_NAME)
            mcmc_init[name] = joker_samples[name].to_value(unit)[0]

        if 't_peri' in model.named_vars:
            logger.info('pymc3 model has already been setup for running MCMC - '
                        'using the previously setup model parameters.')
            return mcmc_init

        p = self.prior.pars
        with model:
            t_peri = pm.Deterministic('t_peri',
                                      p['P'] * p['M0'] / (2*np.pi))

            # Set up the orbit model
            orbit = xo.orbits.KeplerianOrbit(period=p['P'],
                                             ecc=p['e'],
                                             omega=p['omega'],
                                             t_periastron=t_peri)

        # design matrix
        M = get_trend_design_matrix(data, ids, self.prior.poly_trend)

        # deal with v0_offsets, trend here:
        _, offset_names = validate_n_offsets(self.prior.n_offsets)
        _, vtrend_names = validate_poly_trend(self.prior.poly_trend)

        with model:
            v_pars = ([p['v0']]
                      + [p[name] for name in offset_names]
                      + [p[name] for name in vtrend_names[1:]])  # skip v0
            v_trend_vec = tt.stack(v_pars, axis=0)
            trend = tt.dot(M, v_trend_vec)

            rv_model = orbit.get_radial_velocity(x, K=p['K']) + trend

            err = tt.sqrt(err**2 + p['s']**2)
            pm.Normal("obs", mu=rv_model, sd=err, observed=y)

        return mcmc_init

    def trace_to_samples(self, trace, names=None):
        """
        Create a ``JokerSamples`` instance from a pymc3 trace object.

        Parameters
        ----------
        trace : `~pymc3.backends.base.MultiTrace`
        """
        df = pm.trace_to_dataframe(trace)

        samples = JokerSamples(poly_trend=self.prior.poly_trend,
                               n_offsets=self.prior.n_offsets)

        if names is None:
            names = self.prior.par_names

        for name in names:
            par = self.prior.pars[name]
            unit = getattr(par, xu.UNIT_ATTR_NAME)
            samples[name] = df[name].values * unit

        return samples
<|endoftext|>


<commit_before>
import schedule as sched_module

class ScheduleScorer(object):

    def __init__(self):
        self.schedule = None
        self.score_values = dict()
        self.score_info = {
            'no-marathons': {
                'weight': 1,
                'function': self._no_marathons
            },
            'nine-to-five': {
                'weight': 50,
                'function': self._nine_to_five
            },
            'start-time': {
                'weight': 1,
                'function': self._start_time
            }
        }

    def read(self, name='all'):
        if name == 'all':
            return self.score_values
        else:
            return self.score_values.get(name)

    def score(self, schedule):
        self.schedule = schedule
        for name in self.score_info.keys():
            self.score_values.update({
                name: self._weight(name) * self._score(name)
            })
        self.score_values.update({
            'overall': sum(self.score_values.values())
        })

    def _weight(self, name):
        info = self.score_info.get(name)
        if info is not None:
            return info.get('weight', 1)
        return None

    def _score(self, name):
        info = self.score_info.get(name)
        if info is not None:
            return info.get('function', lambda: 0)()
        else:
            return 0

    def _no_marathons(self):
        """
        A schedule is better if it has a more even distribution
        of breaks throughout the day, and also if it squishes sections
        together so as to not waste time. Constantly alternating between
        class and break is bad. Having more than 3 hours in a row is bad.
        """
        max_consecutive_blocks = 3*2 # 3 hours x 2 blocks/hour
        total_marathon_blocks = 0
        for day_bitmap in self.schedule.conflict_bitmap[:]:
            consecutive_blocks = 0
            while day_bitmap:
                day_bitmap &= (day_bitmap << 1)
                consecutive_blocks += 1
            marathon_blocks = consecutive_blocks - max_consecutive_blocks
            if marathon_blocks > 0:
                total_marathon_blocks += marathon_blocks
        score = -1*total_marathon_blocks
        return 5 * score

    def _nine_to_five(self):
        #               0 1 2 3 4 5 6 7 8 9 A B C 1 2 3 4 5 6 7 8 9 A B
        bad_zone = int('111111111111111111000000000000000011111111111111', 2)
        num_outside = 0
        for day_bitmap in self.schedule.conflict_bitmap:
            in_bad_zone = day_bitmap & bad_zone
            num_outside += bin(in_bad_zone).count('1')
        score = -1*num_outside
        return 1 * 50 * score

    def _start_time(self):
        score = 0
        early_block = 8*2
        ideal_start_bitmap = (1 << (sched_module.Schedule.NUM_BLOCKS - early_block -1))
        for day_bitmap in self.schedule.conflict_bitmap[:]:
            if day_bitmap == 0:
                continue
            while day_bitmap < ideal_start_bitmap:
                day_bitmap <<= 1
                score -= 1
        return 1 * score
<commit_msg>chore(scorer-semantics): describe what positive, zero, and negative weightings mean for each score-function<commit_after>
import schedule as sched_module

class ScheduleScorer(object):

    def __init__(self):
        self.schedule = None
        self.score_values = dict()
        self.score_info = {
            'no-marathons': {
                'weight': -100,
                'function': self._no_marathons
            },
            'nine-to-five': {
                'weight': 0,
                'function': self._nine_to_five
            },
            'start-early': {
                'weight': 0,
                'function': self._start_early
            }
        }

    def read(self, name='all'):
        if name == 'all':
            return self.score_values
        else:
            return self.score_values.get(name)

    def score(self, schedule):
        """Score a schedule using each scoring function in self.score_info
        """
        self.schedule = schedule
        for name in self.score_info.keys():
            self.score_values.update({
                name: self._weight(name) * self._score(name)
            })
        self.score_values.update({
            'overall': sum(self.score_values.values())
        })

    def _weight(self, name):
        """Return the weight of a particular scoring function
        """
        info = self.score_info.get(name)
        if info is not None:
            return info.get('weight', 1)
        return None

    def _score(self, name):
        """Run a particular scoring function, and return its result
        """
        info = self.score_info.get(name)
        if info is not None:
            return info.get('function', lambda: 0)()
        else:
            return 0

    def _no_marathons(self):
        """
        + weight: spread out. Less than 3hrs of consecutive classes
        0 weight: -no effect-
        - weight: clumped up. More than 3hrs of consecutive classes
        """
        max_consecutive_blocks = 3*2 # 3 hours x 2 blocks/hour
        total_marathon_blocks = 0
        for day_bitmap in self.schedule.conflict_bitmap[:]:
            consecutive_blocks = 0
            while day_bitmap:
                day_bitmap &= (day_bitmap << 1)
                consecutive_blocks += 1
            marathon_blocks = consecutive_blocks - max_consecutive_blocks
            if marathon_blocks > 0:
                total_marathon_blocks += marathon_blocks
        score = -1*total_marathon_blocks
        return 5 * score

    def _nine_to_five(self):
        """
        + weight: classes inside 9a-5p
        0 weight: -no effect-
        - weight: classes outside 9a-5p
        """
        #               0 1 2 3 4 5 6 7 8 9 A B C 1 2 3 4 5 6 7 8 9 A B
        bad_zone = int('111111111111111111000000000000000011111111111111', 2)
        num_outside = 0
        for day_bitmap in self.schedule.conflict_bitmap:
            in_bad_zone = day_bitmap & bad_zone
            num_outside += bin(in_bad_zone).count('1')
        score = -1*num_outside
        return 1 * 50 * score

    def _start_early(self):
        """
        + weight: start early
        0 weight: -no effect-
        - weight: start late
        """
        score = 0
        early_block = 8*2
        ideal_start_bitmap = (1 << (sched_module.Schedule.NUM_BLOCKS - early_block -1))
        for day_bitmap in self.schedule.conflict_bitmap[:]:
            if day_bitmap == 0:
                continue
            while day_bitmap < ideal_start_bitmap:
                day_bitmap <<= 1
                score -= 1
        return 1 * score
<|endoftext|>


<commit_before>import OpenPNM
import numpy as np
import OpenPNM.Physics.models as pm

class GenericAlgorithmTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[5, 5, 5])
        self.phase = OpenPNM.Phases.GenericPhase(network=self.net)
        self.alg = OpenPNM.Algorithms.GenericLinearTransport(network=self.net,
                                                             phase=self.phase)

    def test_set_BC_modes(self):
        BC1_pores = np.arange(25, 35)
        BC1_pores = np.arange(25, 35)
        self.alg.set_boundary_conditions(component=self.phase,
                                         bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC1_pores)
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == BC1_pores)
        BC2_pores = np.arange(43, 50)
        self.alg.set_boundary_conditions(component=self.phase,
                                         bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC2_pores,
                                         mode='merge')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == np.concatenate((BC1_pores, BC2_pores)))
        BC3_pores = np.arange(4, 9)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC3_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == BC3_pores)
        BC4_pores = [11, 90]
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         pores=BC4_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Neumann')
        assert np.all(ptest == BC4_pores)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=BC1_pores,
                                         bcvalue=0.3)
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=self.alg.Ps,
                                         mode='remove')
        Dp = np.sum(self.alg['pore.' + self.phase.name + '_Dirichlet'])
        assert Dp == 0
        alg.set_boundary_conditions(bctype='Neumann',
                                    mode='remove')
        assert ('pore.' + selfphase.name + '_Neumann' not in alg.labels())

<commit_msg>Minor change for pep8<commit_after>import OpenPNM
import numpy as np

class GenericAlgorithmTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[5, 5, 5])
        self.phase = OpenPNM.Phases.GenericPhase(network=self.net)
        self.alg = OpenPNM.Algorithms.GenericLinearTransport(network=self.net,
                                                             phase=self.phase)

    def test_set_BC_modes(self):
        BC1_pores = np.arange(25, 35)
        BC1_pores = np.arange(25, 35)
        self.alg.set_boundary_conditions(component=self.phase,
                                         bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC1_pores)
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == BC1_pores)
        BC2_pores = np.arange(43, 50)
        self.alg.set_boundary_conditions(component=self.phase,
                                         bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC2_pores,
                                         mode='merge')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == np.concatenate((BC1_pores, BC2_pores)))
        BC3_pores = np.arange(4, 9)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC3_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        assert np.all(ptest == BC3_pores)
        BC4_pores = [11, 90]
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         pores=BC4_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.' + self.phase.name + '_Neumann')
        assert np.all(ptest == BC4_pores)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=BC1_pores,
                                         bcvalue=0.3)
        ptest = self.alg.pores('pore.' + self.phase.name + '_Dirichlet')
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=self.alg.Ps,
                                         mode='remove')
        Dp = np.sum(self.alg['pore.' + self.phase.name + '_Dirichlet'])
        assert Dp == 0
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         mode='remove')
        label = 'pore.' + self.phase.name + '_Neumann'
        assert (label not in self.alg.labels())
<|endoftext|>


<commit_before>import urllib
import sys
import os
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO
from django.conf import settings
from django.contrib.auth import authenticate, login
from django.core.handlers.base import BaseHandler
from django.core.handlers.wsgi import WSGIRequest
from django.core.signals import got_request_exception
from django.dispatch import dispatcher
from django.http import SimpleCookie, HttpRequest
from django.template import TemplateDoesNotExist
from django.test import signals
from django.utils.functional import curry
from django.utils.encoding import smart_str
from django.utils.http import urlencode
from django.utils.itercompat import is_iterable

BOUNDARY = 'BoUnDaRyStRiNg'
MULTIPART_CONTENT = 'multipart/form-data; boundary=%s' % BOUNDARY

class FakePayload(object):
    """
    A wrapper around StringIO that restricts what can be read since data from
    the network can't be seeked and cannot be read outside of its content
    length. This makes sure that views can't do anything under the test client
    that wouldn't work in Real Life.
    """
    def __init__(self, content):
        self.__content = StringIO(content)
        self.__len = len(content)

    def read(self, num_bytes=None):
        if num_bytes is None:
            num_bytes = self.__len or 1
        assert self.__len >= num_bytes, "Cannot read more than the available bytes from the HTTP incoming data."
        content = self.__content.read(num_bytes)
        self.__len -= num_bytes
        return content

class ClientHandler(BaseHandler):
    """
    A HTTP Handler that can be used for testing purposes.
    Uses the WSGI interface to compose requests, but returns
    the raw HttpResponse object
    """
    def __call__(self, environ):
        from django.conf import settings
        from django.core import signals

        # Set up middleware if needed. We couldn't do this earlier, because
        # settings weren't available.
        if self._request_middleware is None:
            self.load_middleware()

        dispatcher.send(signal=signals.request_started)
        try:
            request = WSGIRequest(environ)
            response = self.get_response(request)

            # Apply response middleware
            for middleware_method in self._response_middleware:
                response = middleware_method(request, response)
            response = self.apply_response_fixes(request, response)
        finally:
            dispatcher.send(signal=signals.request_finished)

        return response

def store_rendered_templates(store, signal, sender, template, context):
    "A utility function for storing templates and contexts that are rendered"
    store.setdefault('template',[]).append(template)
    store.setdefault('context',[]).append(context)

def encode_multipart(boundary, data):
    """
    A simple method for encoding multipart POST data from a dictionary of
    form values.

    The key will be used as the form data name; the value will be transmitted
    as content. If the value is a file, the contents of the file will be sent
    as an application/octet-stream; otherwise, str(value) will be sent.
    """
    lines = []
    to_str = lambda s: smart_str(s, settings.DEFAULT_CHARSET)
    for (key, value) in data.items():
        if isinstance(value, file):
            lines.extend([
                '--' + boundary,
                'Content-Disposition: form-data; name="%s"; filename="%s"' % (to_str(key), to_str(os.path.basename(value.name))),
                'Content-Type: application/octet-stream',
                '',
                value.read()
            ])
        else:
            if not isinstance(value, basestring) and is_iterable(value):
                for item in value:
                    lines.extend([
                        '--' + boundary,
                        'Content-Disposition: form-data; name="%s"' % to_str(key),
                        '',
                        to_str(item)
                    ])
            else:
                lines.extend([
                    '--' + boundary,
                    'Content-Disposition: form-data; name="%s"' % to_str(key),
                    '',
                    to_str(value)
                ])

    lines.extend([
        '--' + boundary + '--',
        '',
    ])
    return '\r\n'.join(lines)

class Client:
    """
    A class that can act as a client for testing purposes.

    It allows the user to compose GET and POST requests, and
    obtain the response that the server gave to those requests.
    The server Response objects are annotated with the details
    of the contexts and templates that were rendered during the
    process of serving the request.

    Client objects are stateful - they will retain cookie (and
    thus session) details for the lifetime of the Client instance.

    This is not intended as a replacement for Twill/Selenium or
    the like - it is here to allow testing against the
    contexts and templates produced by a view, rather than the
    HTML rendered to the end-user.
    """
    def __init__(self, **defaults):
        self.handler = ClientHandler()
        self.defaults = defaults
        self.cookies = SimpleCookie()
        self.exc_info = None

    def store_exc_info(self, *args, **kwargs):
        """
        Utility method that can be used to store exceptions when they are
        generated by a view.
        """
        self.exc_info = sys.exc_info()

    def _session(self):
        "Obtain the current session variables"
        if 'django.contrib.sessions' in settings.INSTALLED_APPS:
            engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])
            cookie = self.cookies.get(settings.SESSION_COOKIE_NAME, None)
            if cookie:
                return engine.SessionStore(cookie.value)
        return {}
    session = property(_session)

    def request(self, **request):
        """
        The master request method. Composes the environment dictionary
        and passes to the handler, returning the result of the handler.
        Assumes defaults for the query environment, which can be overridden
        using the arguments to the request.
        """

        environ = {
            'HTTP_COOKIE':      self.cookies,
            'PATH_INFO':         '/',
            'QUERY_STRING':      '',
            'REQUEST_METHOD':    'GET',
            'SCRIPT_NAME':       None,
            'SERVER_NAME':       'testserver',
            'SERVER_PORT':       80,
            'SERVER_PROTOCOL':   'HTTP/1.1',
        }
        environ.update(self.defaults)
        environ.update(request)

        # Curry a data dictionary into an instance of
        # the template renderer callback function
        data = {}
        on_template_render = curry(store_rendered_templates, data)
        dispatcher.connect(on_template_render, signal=signals.template_rendered)

        # Capture exceptions created by the handler
        dispatcher.connect(self.store_exc_info, signal=got_request_exception)

        try:
            response = self.handler(environ)
        except TemplateDoesNotExist, e:
            # If the view raises an exception, Django will attempt to show
            # the 500.html template. If that template is not available,
            # we should ignore the error in favor of re-raising the
            # underlying exception that caused the 500 error. Any other
            # template found to be missing during view error handling
            # should be reported as-is.
            if e.args != ('500.html',):
                raise

        # Look for a signalled exception, clear the current context
        # exception data, then re-raise the signalled exception.
        # Also make sure that the signalled exception is cleared from
        # the local cache!
        if self.exc_info:
            exc_info = self.exc_info
            self.exc_info = None
            raise exc_info[1], None, exc_info[2]

        # Save the client and request that stimulated the response
        response.client = self
        response.request = request

        # Add any rendered template detail to the response
        # If there was only one template rendered (the most likely case),
        # flatten the list to a single element
        for detail in ('template', 'context'):
            if data.get(detail):
                if len(data[detail]) == 1:
                    setattr(response, detail, data[detail][0]);
                else:
                    setattr(response, detail, data[detail])
            else:
                setattr(response, detail, None)

        # Update persistent cookie data
        if response.cookies:
            self.cookies.update(response.cookies)

        return response

    def get(self, path, data={}, **extra):
        "Request a response from the server using GET."
        r = {
            'CONTENT_LENGTH':  None,
            'CONTENT_TYPE':    'text/html; charset=utf-8',
            'PATH_INFO':       urllib.unquote(path),
            'QUERY_STRING':    urlencode(data, doseq=True),
            'REQUEST_METHOD': 'GET',
        }
        r.update(extra)

        return self.request(**r)

    def post(self, path, data={}, content_type=MULTIPART_CONTENT, **extra):
        "Request a response from the server using POST."

        if content_type is MULTIPART_CONTENT:
            post_data = encode_multipart(BOUNDARY, data)
        else:
            post_data = data

        r = {
            'CONTENT_LENGTH': len(post_data),
            'CONTENT_TYPE':   content_type,
            'PATH_INFO':      urllib.unquote(path),
            'REQUEST_METHOD': 'POST',
            'wsgi.input':     FakePayload(post_data),
        }
        r.update(extra)

        return self.request(**r)

    def login(self, **credentials):
        """Set the Client to appear as if it has sucessfully logged into a site.

        Returns True if login is possible; False if the provided credentials
        are incorrect, or the user is inactive, or if the sessions framework is
        not available.
        """
        user = authenticate(**credentials)
        if user and user.is_active and 'django.contrib.sessions' in settings.INSTALLED_APPS:
            engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])

            # Create a fake request to store login details
            request = HttpRequest()
            request.session = engine.SessionStore()
            login(request, user)

            # Set the cookie to represent the session
            self.cookies[settings.SESSION_COOKIE_NAME] = request.session.session_key
            self.cookies[settings.SESSION_COOKIE_NAME]['max-age'] = None
            self.cookies[settings.SESSION_COOKIE_NAME]['path'] = '/'
            self.cookies[settings.SESSION_COOKIE_NAME]['domain'] = settings.SESSION_COOKIE_DOMAIN
            self.cookies[settings.SESSION_COOKIE_NAME]['secure'] = settings.SESSION_COOKIE_SECURE or None
            self.cookies[settings.SESSION_COOKIE_NAME]['expires'] = None

            # Save the session values
            request.session.save()

            return True
        else:
            return False

    def logout(self):
        """Removes the authenticated user's cookies.

        Causes the authenticated user to be logged out.
        """
        session = __import__(settings.SESSION_ENGINE, {}, {}, ['']).SessionStore()
        session.delete(session_key=self.cookies[settings.SESSION_COOKIE_NAME].value)
        self.cookies = SimpleCookie()
<commit_msg>Minor style fixes.<commit_after>import urllib
import sys
import os
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.contrib.auth import authenticate, login
from django.core.handlers.base import BaseHandler
from django.core.handlers.wsgi import WSGIRequest
from django.core.signals import got_request_exception
from django.dispatch import dispatcher
from django.http import SimpleCookie, HttpRequest
from django.template import TemplateDoesNotExist
from django.test import signals
from django.utils.functional import curry
from django.utils.encoding import smart_str
from django.utils.http import urlencode
from django.utils.itercompat import is_iterable

BOUNDARY = 'BoUnDaRyStRiNg'
MULTIPART_CONTENT = 'multipart/form-data; boundary=%s' % BOUNDARY

class FakePayload(object):
    """
    A wrapper around StringIO that restricts what can be read since data from
    the network can't be seeked and cannot be read outside of its content
    length. This makes sure that views can't do anything under the test client
    that wouldn't work in Real Life.
    """
    def __init__(self, content):
        self.__content = StringIO(content)
        self.__len = len(content)

    def read(self, num_bytes=None):
        if num_bytes is None:
            num_bytes = self.__len or 1
        assert self.__len >= num_bytes, "Cannot read more than the available bytes from the HTTP incoming data."
        content = self.__content.read(num_bytes)
        self.__len -= num_bytes
        return content

class ClientHandler(BaseHandler):
    """
    A HTTP Handler that can be used for testing purposes.
    Uses the WSGI interface to compose requests, but returns
    the raw HttpResponse object
    """
    def __call__(self, environ):
        from django.conf import settings
        from django.core import signals

        # Set up middleware if needed. We couldn't do this earlier, because
        # settings weren't available.
        if self._request_middleware is None:
            self.load_middleware()

        dispatcher.send(signal=signals.request_started)
        try:
            request = WSGIRequest(environ)
            response = self.get_response(request)

            # Apply response middleware.
            for middleware_method in self._response_middleware:
                response = middleware_method(request, response)
            response = self.apply_response_fixes(request, response)
        finally:
            dispatcher.send(signal=signals.request_finished)

        return response

def store_rendered_templates(store, signal, sender, template, context):
    """
    Stores templates and contexts that are rendered.
    """
    store.setdefault('template',[]).append(template)
    store.setdefault('context',[]).append(context)

def encode_multipart(boundary, data):
    """
    Encodes multipart POST data from a dictionary of form values.

    The key will be used as the form data name; the value will be transmitted
    as content. If the value is a file, the contents of the file will be sent
    as an application/octet-stream; otherwise, str(value) will be sent.
    """
    lines = []
    to_str = lambda s: smart_str(s, settings.DEFAULT_CHARSET)
    for (key, value) in data.items():
        if isinstance(value, file):
            lines.extend([
                '--' + boundary,
                'Content-Disposition: form-data; name="%s"; filename="%s"' \
                    % (to_str(key), to_str(os.path.basename(value.name))),
                'Content-Type: application/octet-stream',
                '',
                value.read()
            ])
        else:
            if not isinstance(value, basestring) and is_iterable(value):
                for item in value:
                    lines.extend([
                        '--' + boundary,
                        'Content-Disposition: form-data; name="%s"' % to_str(key),
                        '',
                        to_str(item)
                    ])
            else:
                lines.extend([
                    '--' + boundary,
                    'Content-Disposition: form-data; name="%s"' % to_str(key),
                    '',
                    to_str(value)
                ])

    lines.extend([
        '--' + boundary + '--',
        '',
    ])
    return '\r\n'.join(lines)

class Client:
    """
    A class that can act as a client for testing purposes.

    It allows the user to compose GET and POST requests, and
    obtain the response that the server gave to those requests.
    The server Response objects are annotated with the details
    of the contexts and templates that were rendered during the
    process of serving the request.

    Client objects are stateful - they will retain cookie (and
    thus session) details for the lifetime of the Client instance.

    This is not intended as a replacement for Twill/Selenium or
    the like - it is here to allow testing against the
    contexts and templates produced by a view, rather than the
    HTML rendered to the end-user.
    """
    def __init__(self, **defaults):
        self.handler = ClientHandler()
        self.defaults = defaults
        self.cookies = SimpleCookie()
        self.exc_info = None

    def store_exc_info(self, *args, **kwargs):
        """
        Stores exceptions when they are generated by a view.
        """
        self.exc_info = sys.exc_info()

    def _session(self):
        """
        Obtains the current session variables.
        """
        if 'django.contrib.sessions' in settings.INSTALLED_APPS:
            engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])
            cookie = self.cookies.get(settings.SESSION_COOKIE_NAME, None)
            if cookie:
                return engine.SessionStore(cookie.value)
        return {}
    session = property(_session)

    def request(self, **request):
        """
        The master request method. Composes the environment dictionary
        and passes to the handler, returning the result of the handler.
        Assumes defaults for the query environment, which can be overridden
        using the arguments to the request.
        """
        environ = {
            'HTTP_COOKIE':      self.cookies,
            'PATH_INFO':         '/',
            'QUERY_STRING':      '',
            'REQUEST_METHOD':    'GET',
            'SCRIPT_NAME':       None,
            'SERVER_NAME':       'testserver',
            'SERVER_PORT':       80,
            'SERVER_PROTOCOL':   'HTTP/1.1',
        }
        environ.update(self.defaults)
        environ.update(request)

        # Curry a data dictionary into an instance of the template renderer
        # callback function.
        data = {}
        on_template_render = curry(store_rendered_templates, data)
        dispatcher.connect(on_template_render, signal=signals.template_rendered)

        # Capture exceptions created by the handler.
        dispatcher.connect(self.store_exc_info, signal=got_request_exception)

        try:
            response = self.handler(environ)
        except TemplateDoesNotExist, e:
            # If the view raises an exception, Django will attempt to show
            # the 500.html template. If that template is not available,
            # we should ignore the error in favor of re-raising the
            # underlying exception that caused the 500 error. Any other
            # template found to be missing during view error handling
            # should be reported as-is.
            if e.args != ('500.html',):
                raise

        # Look for a signalled exception, clear the current context
        # exception data, then re-raise the signalled exception.
        # Also make sure that the signalled exception is cleared from
        # the local cache!
        if self.exc_info:
            exc_info = self.exc_info
            self.exc_info = None
            raise exc_info[1], None, exc_info[2]

        # Save the client and request that stimulated the response.
        response.client = self
        response.request = request

        # Add any rendered template detail to the response.
        # If there was only one template rendered (the most likely case),
        # flatten the list to a single element.
        for detail in ('template', 'context'):
            if data.get(detail):
                if len(data[detail]) == 1:
                    setattr(response, detail, data[detail][0]);
                else:
                    setattr(response, detail, data[detail])
            else:
                setattr(response, detail, None)

        # Update persistent cookie data.
        if response.cookies:
            self.cookies.update(response.cookies)

        return response

    def get(self, path, data={}, **extra):
        """
        Requests a response from the server using GET.
        """
        r = {
            'CONTENT_LENGTH':  None,
            'CONTENT_TYPE':    'text/html; charset=utf-8',
            'PATH_INFO':       urllib.unquote(path),
            'QUERY_STRING':    urlencode(data, doseq=True),
            'REQUEST_METHOD': 'GET',
        }
        r.update(extra)

        return self.request(**r)

    def post(self, path, data={}, content_type=MULTIPART_CONTENT, **extra):
        """
        Requests a response from the server using POST.
        """
        if content_type is MULTIPART_CONTENT:
            post_data = encode_multipart(BOUNDARY, data)
        else:
            post_data = data

        r = {
            'CONTENT_LENGTH': len(post_data),
            'CONTENT_TYPE':   content_type,
            'PATH_INFO':      urllib.unquote(path),
            'REQUEST_METHOD': 'POST',
            'wsgi.input':     FakePayload(post_data),
        }
        r.update(extra)

        return self.request(**r)

    def login(self, **credentials):
        """
        Sets the Client to appear as if it has successfully logged into a site.

        Returns True if login is possible; False if the provided credentials
        are incorrect, or the user is inactive, or if the sessions framework is
        not available.
        """
        user = authenticate(**credentials)
        if user and user.is_active \
                and 'django.contrib.sessions' in settings.INSTALLED_APPS:
            engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])

            # Create a fake request to store login details.
            request = HttpRequest()
            request.session = engine.SessionStore()
            login(request, user)

            # Set the cookie to represent the session.
            self.cookies[settings.SESSION_COOKIE_NAME] = request.session.session_key
            self.cookies[settings.SESSION_COOKIE_NAME]['max-age'] = None
            self.cookies[settings.SESSION_COOKIE_NAME]['path'] = '/'
            self.cookies[settings.SESSION_COOKIE_NAME]['domain'] = settings.SESSION_COOKIE_DOMAIN
            self.cookies[settings.SESSION_COOKIE_NAME]['secure'] = settings.SESSION_COOKIE_SECURE or None
            self.cookies[settings.SESSION_COOKIE_NAME]['expires'] = None

            # Save the session values.
            request.session.save()

            return True
        else:
            return False

    def logout(self):
        """
        Removes the authenticated user's cookies.

        Causes the authenticated user to be logged out.
        """
        session = __import__(settings.SESSION_ENGINE, {}, {}, ['']).SessionStore()
        session.delete(session_key=self.cookies[settings.SESSION_COOKIE_NAME].value)
        self.cookies = SimpleCookie()
<|endoftext|>


<commit_before>from ..misc.enum import Enum
import threading
from ..data.value import Value, Unknown
import collections

class Variable:
    VarTypes = Enum("None", "Discrete", "Continuous", "String")
    MakeStatus = Enum("OK", "MissingValues", "NoRecognizedValues",
        "Incompatible", "NotFound")
    DefaultUnknownStr =  {"?", ".", "", "NA", "~", None}

    def __init__(self, var_type, name="", ordered=False, default_col = -1):
        self.var_type = var_type
        self.name = name
        self.ordered = ordered
        self.default_col = default_col
        self.random_generator = None
        self.source_variable = None
        self.get_value_from = None
        self.get_value_lock = threading.Lock()
        self.unknown_str = set(Variable.DefaultUnknownStr)

    def computeValue(self, inst):
        if self.get_value_from is None:
            return Value(self)
        with self.get_value_lock:
            return Value(self, self.get_value_from(inst))

    def is_primitive(self):
        return False

    def repr_val(self, val):
        return str(val)

    str_val = repr_val

    def __str__(self):
        return "{}('{}')".format(self.__class__.__name__, self.name)

    __repr__ = __str__

class DiscreteVariable(Variable):
    all_discrete_vars = collections.defaultdict(set)
    presorted_values = [["no", "yes"], ["No", "Yes"]]

    def __init__(self, name="", values=(), ordered=False, default_col=-1, base_value=-1):
        super().__init__(Variable.VarTypes.Discrete,
            name, ordered, default_col)
        self.values = list(values)
        self.base_value = base_value
        DiscreteVariable.all_discrete_vars[name].add(self)

    def __str__(self):
        args = "values=[" + ", ".join(self.values[:5]) + \
               "..."*(len(self.values)>5) + "]"
        if self.ordered:
            args += ", ordered=True"
        if self.base_value >= 0:
            args += ", base_value={}".format(self.base_value)
        return "{}('{}', {})".format(self.__class__.__name__, self.name, args)

    @staticmethod
    def is_primitive(self):
        return True

    def to_val(self, s):
        if isinstance(s, int):
            return s
        if s in self.unknown_str:
            return Unknown
        if not isinstance(s, str):
            raise TypeError('Cannot convert {} to value of "{}"'
                            .format(type(s).__name__, self.name))
        return self.values.index(s)

    def val_from_str_add(self, s):
        try:
            return Unknown if s in self.unknown_str else self.values.index(s)
        except ValueError:
            self.values.append(s)
            return len(self.values) - 1

    def repr_val(self, val):
        return '{}'.format(self.values[int(val)])

    str_val = repr_val

    def check_values_order(self, values):
        return self.values[:len(values)] == values[:len(self.values)]

    @staticmethod
    def make(name, values=(), ordered=False, default_meta_id=-1, base_value=-1):
        var = DiscreteVariable.find_compatible(
            name, values, ordered, default_meta_id, base_value)
        if var:
            return var
        return DiscreteVariable(name, values, ordered, default_meta_id, base_value)

    @staticmethod
    def find_compatible(name, values=(), ordered=False, default_col=-1, base_value=-1):
        existing = DiscreteVariable.all_discrete_vars.get(name)
        if existing:
            for var in existing:
                if var.n_existing_values(values) >= 0 and var.ordered == ordered:
                    if base_value != -1:
                        if var.base_value == -1:
                            var.base_value = base_value
                        elif var.base_value != base_value:
                            continue
                    if default_col != -1:
                        if var.default_col == -1:
                            var.default_col = default_col
                        elif var.default_col != default_col:
                            continue
                    return var

    @staticmethod
    def order_values(values):
        values = set(values)
        for presorted in DiscreteVariable.presorted_values:
            if values == set(presorted):
                return presorted
        return sorted(values)

    def add_values(self, unordered, ordered=()):
        used = self.n_existing_values(ordered)
        if used < 0:
            raise ValueError("Mismatching order of values")
        self.values += ordered[used:]
        remaining = set(unordered) - set(self.values)
        self.values += self.order_values(remaining)

    def n_existing_values(self, ordered):
        i = 0
        for val in self.values:
            if i < len(ordered) and ordered[i] == val:
                i += 1
        if set(ordered[i:]) & set(self.values):
            return -1
        return i

class ContinuousVariable(Variable):
    all_continuous_vars = {}

    def __init__(self, name="", default_col=-1):
        super().__init__(Variable.VarTypes.Continuous, name, default_col)
        self.number_of_decimals = 3
        self.scientific_format = False
        self.adjustDecimals = 2
        ContinuousVariable.all_continuous_vars[name] = self

    @staticmethod
    def make(name, default_meta_id=-1):
        return ContinuousVariable.all_continuous_vars.get(name) or \
               ContinuousVariable(name, default_meta_id)

    def is_primitive(self):
        return False

    def to_val(self, s):
        return Unknown if s in self.unknown_str else float(s)

    val_from_str_add = to_val

    def repr_val(self, val):
        return "%.*f" % (self.number_of_decimals, val)

    str_val = repr_val

class StringVariable(Variable):
    all_string_vars = {}

    def __init__(self, name="", default_col=-1):
        super().__init__(Variable.VarTypes.String, name, default_col)
        StringVariable.all_string_vars[name] = self

    def to_val(self, s):
        if s is None:
            return ""
        if isinstance(s, str):
            return s
        return str(s)

    val_from_str_add = to_val

    def str_val(self, val):
        if isinstance(val, Value):
            if val.value is None:
                return "None"
            val = val.value
        return str(val)

    def repr_val(self, val):
        return '"{}"'.format(self.str_val(val))

    @staticmethod
    def make(name, default_meta_id=-1):
        return StringVariable.all_string_vars.get(name) or \
               StringVariable(name, default_meta_id)
<commit_msg>Fixed printing of nans for variable.Discrete and variable.Continuous<commit_after>from ..misc.enum import Enum
import threading
from ..data.value import Value, Unknown
import collections
from math import isnan

class Variable:
    VarTypes = Enum("None", "Discrete", "Continuous", "String")
    MakeStatus = Enum("OK", "MissingValues", "NoRecognizedValues",
        "Incompatible", "NotFound")
    DefaultUnknownStr =  {"?", ".", "", "NA", "~", None}

    def __init__(self, var_type, name="", ordered=False, default_col = -1):
        self.var_type = var_type
        self.name = name
        self.ordered = ordered
        self.default_col = default_col
        self.random_generator = None
        self.source_variable = None
        self.get_value_from = None
        self.get_value_lock = threading.Lock()
        self.unknown_str = set(Variable.DefaultUnknownStr)

    def computeValue(self, inst):
        if self.get_value_from is None:
            return Value(self)
        with self.get_value_lock:
            return Value(self, self.get_value_from(inst))

    def is_primitive(self):
        return False

    def repr_val(self, val):
        return str(val)

    str_val = repr_val

    def __str__(self):
        return "{}('{}')".format(self.__class__.__name__, self.name)

    __repr__ = __str__

class DiscreteVariable(Variable):
    all_discrete_vars = collections.defaultdict(set)
    presorted_values = [["no", "yes"], ["No", "Yes"]]

    def __init__(self, name="", values=(), ordered=False, default_col=-1, base_value=-1):
        super().__init__(Variable.VarTypes.Discrete,
            name, ordered, default_col)
        self.values = list(values)
        self.base_value = base_value
        DiscreteVariable.all_discrete_vars[name].add(self)

    def __str__(self):
        args = "values=[" + ", ".join(self.values[:5]) + \
               "..."*(len(self.values)>5) + "]"
        if self.ordered:
            args += ", ordered=True"
        if self.base_value >= 0:
            args += ", base_value={}".format(self.base_value)
        return "{}('{}', {})".format(self.__class__.__name__, self.name, args)

    @staticmethod
    def is_primitive(self):
        return True

    def to_val(self, s):
        if isinstance(s, int):
            return s
        if s in self.unknown_str:
            return Unknown
        if not isinstance(s, str):
            raise TypeError('Cannot convert {} to value of "{}"'
                            .format(type(s).__name__, self.name))
        return self.values.index(s)

    def val_from_str_add(self, s):
        try:
            return Unknown if s in self.unknown_str else self.values.index(s)
        except ValueError:
            self.values.append(s)
            return len(self.values) - 1

    def repr_val(self, val):
        if isnan(val):
            return "?"
        return '{}'.format(self.values[int(val)])

    str_val = repr_val

    def check_values_order(self, values):
        return self.values[:len(values)] == values[:len(self.values)]

    @staticmethod
    def make(name, values=(), ordered=False, default_meta_id=-1, base_value=-1):
        var = DiscreteVariable.find_compatible(
            name, values, ordered, default_meta_id, base_value)
        if var:
            return var
        return DiscreteVariable(name, values, ordered, default_meta_id, base_value)

    @staticmethod
    def find_compatible(name, values=(), ordered=False, default_col=-1, base_value=-1):
        existing = DiscreteVariable.all_discrete_vars.get(name)
        if existing:
            for var in existing:
                if var.n_existing_values(values) >= 0 and var.ordered == ordered:
                    if base_value != -1:
                        if var.base_value == -1:
                            var.base_value = base_value
                        elif var.base_value != base_value:
                            continue
                    if default_col != -1:
                        if var.default_col == -1:
                            var.default_col = default_col
                        elif var.default_col != default_col:
                            continue
                    return var

    @staticmethod
    def order_values(values):
        values = set(values)
        for presorted in DiscreteVariable.presorted_values:
            if values == set(presorted):
                return presorted
        return sorted(values)

    def add_values(self, unordered, ordered=()):
        used = self.n_existing_values(ordered)
        if used < 0:
            raise ValueError("Mismatching order of values")
        self.values += ordered[used:]
        remaining = set(unordered) - set(self.values)
        self.values += self.order_values(remaining)

    def n_existing_values(self, ordered):
        i = 0
        for val in self.values:
            if i < len(ordered) and ordered[i] == val:
                i += 1
        if set(ordered[i:]) & set(self.values):
            return -1
        return i

class ContinuousVariable(Variable):
    all_continuous_vars = {}

    def __init__(self, name="", default_col=-1):
        super().__init__(Variable.VarTypes.Continuous, name, default_col)
        self.number_of_decimals = 3
        self.scientific_format = False
        self.adjustDecimals = 2
        ContinuousVariable.all_continuous_vars[name] = self

    @staticmethod
    def make(name, default_meta_id=-1):
        return ContinuousVariable.all_continuous_vars.get(name) or \
               ContinuousVariable(name, default_meta_id)

    def is_primitive(self):
        return False

    def to_val(self, s):
        return Unknown if s in self.unknown_str else float(s)

    val_from_str_add = to_val

    def repr_val(self, val):
        if isnan(val):
            return "?"
        return "%.*f" % (self.number_of_decimals, val)

    str_val = repr_val

class StringVariable(Variable):
    all_string_vars = {}

    def __init__(self, name="", default_col=-1):
        super().__init__(Variable.VarTypes.String, name, default_col)
        StringVariable.all_string_vars[name] = self

    def to_val(self, s):
        if s is None:
            return ""
        if isinstance(s, str):
            return s
        return str(s)

    val_from_str_add = to_val

    def str_val(self, val):
        if isinstance(val, Value):
            if val.value is None:
                return "None"
            val = val.value
        return str(val)

    def repr_val(self, val):
        return '"{}"'.format(self.str_val(val))

    @staticmethod
    def make(name, default_meta_id=-1):
        return StringVariable.all_string_vars.get(name) or \
               StringVariable(name, default_meta_id)
<|endoftext|>


<commit_before>#!/usr/bin/env python
##
##  Copyright 2015 SRI International
##  License: https://ncanda.sri.com/software-license.txt
##
##  $Revision: 2110 $
##  $LastChangedBy: nicholsn $
##  $LastChangedDate: 2015-08-07 09:10:29 -0700 (Fri, 07 Aug 2015) $
##
"""
NP Release Script
======================
Generates a CSV file for all subjects included in NP analysis
"""
import os

import pandas as pd

directory = "/fs/ncanda-share/releases/NCANDA_DATA_00019/summaries"

nps_file = ["ataxia.csv", "cddr.csv", "clinical.csv", "cnp.csv", "dd100.csv",
            "dd1000.csv", "grooved_pegboard.csv", "ishihara.csv",
            "landoltc.csv", "rey-o.csv", "wais4.csv", "wrat4.csv",
            "mr_session_report.csv", "parentreport.csv"]

fields = ['site', 'sex', 'visit_age','exceed_or_mrianomaly',
          'exceeds_bl_drinking', 'cddr_past_month_binge','cddr_past_year_binge',
          'binge_groups_1', 'binge_groups_month', 'hispanic', 'race',
          'race_label', 'mri_analysisanomalies', 'hi_ed',
          'african_american_black', 'asian','caucasian_white',
          'native_american_american_indian', 'pacific_islander',
          'cddrheight_inches', 'cddrweight_pounds', 'ses_parent_yoe',
          'fh_alc_density','fh_drug_density', 'bmi_value', 'bmi_zscore',
          'bmi_percentile', 'pds_score','np_atax_sht_sum', 'np_atax_steps_sum',
          'np_atax_standr_sum','np_atax_standl_sum', 'cnp_mpract_mp2rtcr',
          'cnp_cpf_ifac_tot','cnp_cpf_ifac_rtc', 'cnp_cpw_iwrd_tot',
          'cnp_cpw_iwrd_rtc','cnp_spcptnl_scpl_tprt', 'cnp_spcptnl_scpt_tp',
          'cnp_spcptnl_scpt_tprt','cnp_sfnb2_sfnb_rtc', 'cnp_sfnb2_sfnb_rtc0',
          'cnp_sfnb2_sfnb_rtc1','cnp_sfnb2_sfnb_rtc2', 'cnp_sfnb2_sfnb_mcr',
          'cnp_sfnb2_sfnb_mrtc','cnp_pmat24a_pmat24_a_cr',
          'cnp_pmat24a_pmat24_a_rtcr', 'cnp_cpfd_cpfdtp','cnp_cpfd_cpfdtprt',
          'cnp_cpwd_dwrd_tot', 'cnp_cpwd_dwrd_rtc','cnp_shortvolt_svt',
          'cnp_shortvolt_svtcrt', 'cnp_shortvolt_svt_eff','cnp_er40d_er40_cr',
          'cnp_er40d_er40_crt', 'cnp_er40d_er40angrt','cnp_er40d_er40fearrt',
          'cnp_er40d_er40haprt', 'cnp_er40d_er40noert','cnp_er40d_er40sadrt',
          'cnp_pcet_pcetcr', 'cnp_pcet_pcetrtcr','cnp_pcet_pcet_eff',
          'cnp_pcet_pcet_acc', 'cnp_pcet_pcet_acc2',
          'cnp_medf36_medf36_hap_rtcr', 'cnp_medf36_medf36_sad_rtcr',
          'cnp_medf36_medf36_ang_rtcr', 'cnp_medf36_medf36_fear_rtcr',
          'cnp_medf36_medf36_a', 'cnp_medf36_medf36_t', 'cnp_pvoc_pvoccr',
          'cnp_pvoc_pvocrtcr', 'cnp_pvrt_pvrtcr', 'cnp_pvrt_pvrtrtcr',
          'cnp_pvrt_pvrt_eff', 'cnp_svdelay_svt_ld', 'cnp_svdelay_svtldrtc',
          'cnp_svdelay_svtldrter', 'cnp_svdelay_svt_effd', 'cnp_sfnb2_sfnb_mrt',
          'cnp_sfnb2_sfnb_meff', 'np_gpeg_dh_time', 'np_gpeg_ndh_time',
          'np_wais4_rawscore', 'np_wrat4_wr_raw', 'np_wrat4_mc_raw',
          'cnp_spcptnl_scpn_tp', 'cnp_spcptnl_scpn_tprt', 'cnp_spcptnl_scpl_tp',
          'cnp_cpfd_dfac_tot', 'cnp_cpfd_dfac_rtc']

race_map = dict(native_american_american_indian=1,
                asian=2,
                pacific_islander=3,
                african_american_black=4,
                caucasian_white=5);

def replace_binge_groups_1(x):
	"""
	 A binary variable from CDDR_PastYr_Binge (i.e., 0 or blank,
	 but not missing = 0, and 1 or more = 1)
	"""
	if x > 0:
		result = 1
	elif x == 0:
		result = 0
	else:
		result = pd.np.NaN
	return result

def replace_binge_groups_month(x):
	"""
	 #3groups based on past month binge drinking (0 - is o or blank, but not missing)
	"""
	if x == 0:
		result = 'None'
	elif x == 1:
		result = '1'
	elif x > 1:
		result = '2+'
	else:
		result = ''
	return result

def main(args=None):
    final_df = pd.read_csv(os.path.join(directory, "demographics.csv"),
                                        index_col=['subject','arm','visit'])

    # Create a series for each race where there is a 1 or 0 indicating if the
    # participant belongs to the race or not and append to the final_df
    for k, v in race_map.iteritems():
        race_filter = final_df.race == v
        final_df[k] = race_filter.apply(lambda x: 1 if x == True else 0)

        for i in nps_file:
            df = pd.read_csv(os.path.join(directory, i),
                        index_col=['subject','arm','visit'])
                        final_df = pd.concat([final_df, df], axis=1)

        final_df = final_df.rename(columns={'cddr31':'cddr_past_month_binge',
                                    'cddr30':'cddr_past_year_binge'})

        final_df['binge_groups_1'] = final_df['cddr_past_year_binge'].apply(replace_binge_groups_1)
        final_df['binge_groups_month'] = final_df['cddr_past_month_binge'].apply(replace_binge_groups_month)

        final_df['exceed_or_mrianomaly'] = pd.np.NaN
        final_df['hi_ed'] = pd.np.NaN

        final_df = final_df[fields]

        final_df.to_csv('np_release.csv')

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    argv = parser.parse_args()
    sys.exit(main(args=argv))
<commit_msg>UPDATE: Data Quality Control Scrpit<commit_after>#!/usr/bin/env python
##
##  Copyright 2015 SRI International
##  License: https://ncanda.sri.com/software-license.txt
##
##  $Revision: 2110 $
##  $LastChangedBy: nicholsn $
##  $LastChangedDate: 2015-08-07 09:10:29 -0700 (Fri, 07 Aug 2015) $
##
"""
NP Release Script
======================
Generates a CSV file for all subjects included in NP analysis
"""
import os

import pandas as pd

directory = "/fs/ncanda-share/releases/NCANDA_DATA_00010/summaries"

nps_file = ["ataxia.csv", "cddr.csv", "clinical.csv", "cnp.csv", "dd100.csv",
            "dd1000.csv", "grooved_pegboard.csv", "ishihara.csv",
            "landoltc.csv", "rey-o.csv", "wais4.csv", "wrat4.csv",
            "mr_session_report.csv", "parentreport.csv"]

fields = ['site', 'sex', 'visit_age','exceed_or_mrianomaly',
          'exceeds_bl_drinking', 'cddr_past_month_binge','cddr_past_year_binge',
          'binge_groups_1', 'binge_groups_month', 'hispanic', 'race',
          'race_label', 'mri_analysisanomalies', 'hi_ed',
          'african_american_black', 'asian','caucasian_white',
          'native_american_american_indian', 'pacific_islander',
          'cddrheight_inches', 'cddrweight_pounds', 'ses_parent_yoe',
          'fh_alc_density','fh_drug_density', 'bmi_value', 'bmi_zscore',
          'bmi_percentile', 'pds_score','np_atax_sht_sum', 'np_atax_steps_sum',
          'np_atax_standr_sum','np_atax_standl_sum', 'cnp_mpract_mp2rtcr',
          'cnp_cpf_ifac_tot','cnp_cpf_ifac_rtc', 'cnp_cpw_iwrd_tot',
          'cnp_cpw_iwrd_rtc','cnp_spcptnl_scpl_tprt', 'cnp_spcptnl_scpt_tp',
          'cnp_spcptnl_scpt_tprt','cnp_sfnb2_sfnb_rtc', 'cnp_sfnb2_sfnb_rtc0',
          'cnp_sfnb2_sfnb_rtc1','cnp_sfnb2_sfnb_rtc2', 'cnp_sfnb2_sfnb_mcr',
          'cnp_sfnb2_sfnb_mrtc','cnp_pmat24a_pmat24_a_cr',
          'cnp_pmat24a_pmat24_a_rtcr', 'cnp_cpfd_cpfdtp','cnp_cpfd_cpfdtprt',
          'cnp_cpwd_dwrd_tot', 'cnp_cpwd_dwrd_rtc','cnp_shortvolt_svt',
          'cnp_shortvolt_svtcrt', 'cnp_shortvolt_svt_eff','cnp_er40d_er40_cr',
          'cnp_er40d_er40_crt', 'cnp_er40d_er40angrt','cnp_er40d_er40fearrt',
          'cnp_er40d_er40haprt', 'cnp_er40d_er40noert','cnp_er40d_er40sadrt',
          'cnp_pcet_pcetcr', 'cnp_pcet_pcetrtcr','cnp_pcet_pcet_eff',
          'cnp_pcet_pcet_acc', 'cnp_pcet_pcet_acc2',
          'cnp_medf36_medf36_hap_rtcr', 'cnp_medf36_medf36_sad_rtcr',
          'cnp_medf36_medf36_ang_rtcr', 'cnp_medf36_medf36_fear_rtcr',
          'cnp_medf36_medf36_a', 'cnp_medf36_medf36_t', 'cnp_pvoc_pvoccr',
          'cnp_pvoc_pvocrtcr', 'cnp_pvrt_pvrtcr', 'cnp_pvrt_pvrtrtcr',
          'cnp_pvrt_pvrt_eff', 'cnp_svdelay_svt_ld', 'cnp_svdelay_svtldrtc',
          'cnp_svdelay_svtldrter', 'cnp_svdelay_svt_effd', 'cnp_sfnb2_sfnb_mrt',
          'cnp_sfnb2_sfnb_meff', 'np_gpeg_dh_time', 'np_gpeg_ndh_time',
          'np_wais4_rawscore', 'np_wrat4_wr_raw', 'np_wrat4_mc_raw',
          'cnp_spcptnl_scpn_tp', 'cnp_spcptnl_scpn_tprt', 'cnp_spcptnl_scpl_tp',
          'cnp_cpfd_dfac_tot', 'cnp_cpfd_dfac_rtc']

race_map = dict(native_american_american_indian=1,
                asian=2,
                pacific_islander=3,
                african_american_black=4,
                caucasian_white=5);

def replace_binge_groups_1(x):
	"""
	 A binary variable from CDDR_PastYr_Binge (i.e., 0 or blank,
	 but not missing = 0, and 1 or more = 1)
	"""
	if x > 0:
		result = 1
	elif x == 0:
		result = 0
	else:
		result = pd.np.NaN
	return result

def replace_binge_groups_month(x):
	"""
	 #3groups based on past month binge drinking (0 - is o or blank, but not missing)
	"""
	if x == 0:
		result = 'None'
	elif x == 1:
		result = '1'
	elif x > 1:
		result = '2+'
	else:
		result = ''
	return result

def main(args=None):
    final_df = pd.read_csv(os.path.join(directory, "demographics.csv"),
                                        index_col=['subject','arm','visit'])

    # Create a series for each race where there is a 1 or 0 indicating if the
    # participant belongs to the race or not and append to the final_df
    for k, v in race_map.iteritems():
        race_filter = final_df.race == v
        final_df[k] = race_filter.apply(lambda x: 1 if x == True else 0)

        for i in nps_file:
            df = pd.read_csv(os.path.join(directory, i),
                        index_col=['subject','arm','visit'])
                        final_df = pd.concat([final_df, df], axis=1)

        final_df = final_df.rename(columns={'cddr31':'cddr_past_month_binge',
                                    'cddr30':'cddr_past_year_binge'})

        final_df['binge_groups_1'] = final_df['cddr_past_year_binge'].apply(replace_binge_groups_1)
        final_df['binge_groups_month'] = final_df['cddr_past_month_binge'].apply(replace_binge_groups_month)

        final_df['exceed_or_mrianomaly'] = pd.np.NaN
        final_df['hi_ed'] = pd.np.NaN

        final_df = final_df[fields]

        final_df.to_csv('np_release.csv')

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    argv = parser.parse_args()
    sys.exit(main(args=argv))
<|endoftext|>


<commit_before>"""
Created on 13 Nov 2016

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

south-coast-science-dev
43308b72-ad41-4555-b075-b4245c1971db
"""

import urllib.parse

from scs_core.osio.client.rest_client import RESTClient
from scs_core.osio.data.topic import Topic

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

class TopicManager(object):
    """
    classdocs
    """
    __FINDER_BATCH_SIZE = 100

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

    def __init__(self, http_client, api_key):
        """
        Constructor
        """
        self.__rest_client = RESTClient(http_client, api_key)

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

    def find(self, topic_path):
        path = '/v1/topics' + urllib.parse.quote(topic_path, '')

        # request...
        self.__rest_client.connect()

        try:
            response_jdict = self.__rest_client.get(path)
        except RuntimeError:
            response_jdict = None

        self.__rest_client.close()

        topic = Topic.construct_from_jdict(response_jdict)

        return topic

    def find_for_org(self, org_id):
        topics = []

        self.__rest_client.connect()

        try:
            for batch in self.__get(org_id):
                topics.extend(batch)

        finally:
            self.__rest_client.close()

        return topics

    def create(self, topic):
        path = '/v2/topics'

        # request...
        self.__rest_client.connect()

        try:
            response = self.__rest_client.post(path, topic.as_json())

        finally:
            self.__rest_client.close()

        success = response == topic.path

        return success

    def delete(self, topic_path):
        path = '/v1/topics/' + urllib.parse.quote(topic_path, '')

        # request...
        self.__rest_client.connect()

        try:
            response = self.__rest_client.delete(path)

        finally:
            self.__rest_client.close()

        success = response == ''

        return success

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

    def __get(self, org_id):
        path = '/v2/orgs/' + org_id + '/topics'
        params = {'offset': 0, 'count': self.__FINDER_BATCH_SIZE}

        while True:
            # request...
            response_jdict = self.__rest_client.get(path, params)

            topics_jdict = response_jdict.get('topics')
            topics = [Topic.construct_from_jdict(topic_jdict) for topic_jdict in topics_jdict] if topics_jdict else []

            yield topics

            if len(topics) == 0:
                break

            params['offset'] = len(topics)

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

    def __str__(self, *args, **kwargs):
        return "TopicManager:{rest_client:%s}" % self.__rest_client
<commit_msg>Added exception handling to OSIO managers.<commit_after>"""
Created on 13 Nov 2016

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

south-coast-science-dev
43308b72-ad41-4555-b075-b4245c1971db
"""

import urllib.parse

from scs_core.osio.client.rest_client import RESTClient
from scs_core.osio.data.topic import Topic

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

class TopicManager(object):
    """
    classdocs
    """
    __FINDER_BATCH_SIZE = 100

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

    def __init__(self, http_client, api_key):
        """
        Constructor
        """
        self.__rest_client = RESTClient(http_client, api_key)

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

    def find(self, topic_path):
        path = '/v1/topics' + urllib.parse.quote(topic_path, '')

        # request...
        self.__rest_client.connect()

        try:
            response_jdict = self.__rest_client.get(path)
        except RuntimeError:
            response_jdict = None

        self.__rest_client.close()

        topic = Topic.construct_from_jdict(response_jdict)

        return topic

    def find_for_org(self, org_id):
        topics = []

        self.__rest_client.connect()

        try:
            for batch in self.__get(org_id):
                topics.extend(batch)

        finally:
            self.__rest_client.close()

        return topics

    def create(self, topic):
        path = '/v2/topics'

        # request...
        self.__rest_client.connect()

        try:
            response = self.__rest_client.post(path, topic.as_json())

        finally:
            self.__rest_client.close()

        success = response == topic.path

        return success

    def delete(self, topic_path):
        path = '/v1/topics/' + urllib.parse.quote(topic_path, '')

        # request...
        self.__rest_client.connect()

        try:
            response = self.__rest_client.delete(path)

        finally:
            self.__rest_client.close()

        success = response == ''

        return success

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

    def __get(self, org_id):
        path = '/v2/orgs/' + org_id + '/topics'
        params = {'offset': 0, 'count': self.__FINDER_BATCH_SIZE}

        while True:
            # request...
            response_jdict = self.__rest_client.get(path, params)

            # topics...
            topics_jdict = response_jdict.get('topics')
            topics = [Topic.construct_from_jdict(topic_jdict) for topic_jdict in topics_jdict] if topics_jdict else []

            yield topics

            if len(topics) == 0:
                break

            # next...
            params['offset'] = len(topics)

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

    def __str__(self, *args, **kwargs):
        return "TopicManager:{rest_client:%s}" % self.__rest_client
<|endoftext|>