<commit_before># coding: utf-8

""" Core functionality for doing stream inference """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
from collections import OrderedDict
import time
import random

# Third-party
from astropy import log as logger
import astropy.table as at
from emcee import EnsembleSampler, PTSampler
import numpy as np
import astropy.units as u
import h5py
import numexpr

# Project
from streamteam.inference import EmceeModel
from .. import heliocentric_names
from ..util import streamspath
from .rewinder_likelihood import rewinder_likelihood

__all__ = ["Rewinder", "RewinderSampler"]

class Rewinder(EmceeModel):

    def __init__(self, potential, progenitor, stars, t1, t2, dt=-1.):
        """ Model for tidal streams that uses backwards integration to Rewind
            the positions of stars.

            Parameters
            ----------
            potential : Potential object
            progenitor :
            stars :
        """

        # Potential
        for par_name,par in potential.parameters:
            self.add_parameter(par, "potential")
        self.potential_class = potential.c_cls

        # Progenitor
        for par_name,par in progenitor.parameters:
            self.add_parameter(par, "progenitor")

        # Stars
        for par_name,par in stars.parameters:
            self.add_parameter(par, "stars")
        self.nstars = len(stars)

        self.args = (t1,t2,dt)

    @classmethod
    def from_config(cls, config):
        """ Construct a StreamModel from a configuration dictionary.
            Typically comes from a YAML file via `streams.io.read_config`.

            Parameters
            ----------
            config : dict
        """

        if not isinstance(config, dict):
            from ..io import read_config
            config = read_config(config)

        seed = config.get('seed', np.random.randint(100000))
        logger.debug("Using seed: {}".format(seed))
        np.random.seed(seed)
        random.seed(seed)

        nlead = config.get('nlead', 0)
        ntrail = config.get('ntrail', 0)
        nstars = nlead + ntrail

        # more complicated selections
        star_idx = config.get('star_idx', None)
        expr = config.get('expr', None)

        if star_idx is not None and nstars is not 0 and len(star_idx) != 0:
            raise ValueError("Number of stars does not match length of"
                             "star index specification.")

        # load progenitor and star data
        logger.debug("Reading data from:\n\t{}".format(config['data_file']))
        progenitor = at.Table.read(config['data_file'], path='progenitor')
        stars = at.Table.read(config['data_file'], path='stars')

        if star_idx is None:
            if nstars is 0:
                raise ValueError("If not specifying indexes, must specicy number"
                                 "of stars (nstars)")

            if expr is not None:
                expr_idx = numexpr.evaluate(expr, stars)
            else:
                expr_idx = np.ones(len(stars)).astype(bool)

            # leading tail stars
            lead_stars, = np.where((stars["tail"] == -1.) & expr_idx)
            np.random.shuffle(lead_stars)
            lead_stars = lead_stars[:nlead]

            # trailing tail stars
            trail_stars, = np.where((stars["tail"] == 1.) & expr_idx)
            np.random.shuffle(trail_stars)
            trail_stars = trail_stars[:ntrail]

            star_idx = np.append(lead_stars, trail_stars)

        stars = stars[star_idx]
        logger.info("Running with {} stars.".format(len(stars)))
        logger.debug("Using stars: {}".format(list(star_idx)))

        # integration stuff
        integration = at.Table.read(config['data_file'], path='integration')
        t1 = config.get("t1", integration.meta['t1'])
        t2 = config.get("t2", integration.meta['t2'])
        dt = config.get("dt", -1.)
        logger.debug("Integration from {} to {}, ∆t={} Myr".format(t1,t2,dt))

        # get the potential object specified from the potential subpackage
        from .. import potential as sp
        Potential = getattr(sp, config["potential"]["class_name"])
        potential = Potential()
        logger.info("Using potential '{}'...".format(config["potential"]["class_name"]))
        for par in potential.parameters.values():
            if par.name not in config["potential"]["parameters"]:
                par.freeze("truth")

        return

        # TODO: turn stars and progenitor into something that has parameters, but
        #       also knows about observational errors, etc....

        # Initialize the model
        model = cls(potential, progenitor, stars, t1, t2, dt)

        # Particle parameters
        if config['particles']['parameters']:
            logger.debug("Particle properties added as parameters:")

            for ii,name in enumerate(config["particles"]["parameters"]):
                if name in heliocentric.coord_names:
                    priors = [LogPrior() for ii in range(nparticles)]
                    X = ModelParameter(name,
                                       value=particles[name].decompose(usys).value,
                                       prior=priors,
                                       truth=true_particles[name].decompose(usys).value)
                    model.add_parameter('particles', X)
                    logger.debug("\t\t{}".format(name))
                else:
                    p = getattr(particles,name)
                    logger.debug("Prior on {}: {}".format(name, p._prior))
                    model.add_parameter('particles', p)

            missing_dims = config["particles"].get("missing_dims", list())
            for missing_dim in missing_dims:
                ix = heliocentric.coord_names.index(missing_dim)
                model.particles._error_X[:,ix] = np.inf

        # Satellite parameters
        if config['satellite']['parameters']:
            logger.debug("Satellite properties added as parameters:")
            for ii,name in enumerate(config["satellite"]["parameters"]):

                if name in heliocentric.coord_names:
                    X = ModelParameter(name,
                                       value=satellite[name].decompose(usys).value,
                                       prior=LogPrior(),
                                       truth=true_satellite[name].decompose(usys).value)
                    model.add_parameter('satellite', X)
                    logger.debug("\t\t{}".format(name))

                elif name == 'alpha':
                    p = getattr(satellite,name)
                    logger.debug("Prior on {}: ".format(name))
                    model.add_parameter('satellite', p)

                else:
                    p = getattr(satellite,name)
                    logger.debug("Prior on {}: {}".format(name, p._prior))
                    model.add_parameter('satellite', p)

            missing_dims = config["satellite"].get("missing_dims", list())
            for missing_dim in missing_dims:
                ix = heliocentric.coord_names.index(missing_dim)
                model.satellite._error_X[:,ix] = np.inf

        return model

    def ln_prior(self, parameters, parameter_values, t1, t2, dt):
        """ Evaluate the log-prior at the given parameter values, but
            not for the star positions, which we need in the likelihood
            function.

            Parameters
            ----------
            parameters : dict
                Dictionary of ModelParameter objects.
            parameter_values : dict
                The dictionary of model parameter values.
            t1,t2,dt : numeric
                Integration limits.
        """
        ln_prior = 0.
        for group_name in ['potential', 'progenitor']:
            for param in parameters[group_name].keys():
                v = value_dict[group_name][param.name]
                ln_prior += param.prior(v)

        return ln_prior

    def ln_likelihood(self, parameters, parameter_values, t1, t2, dt):
        """ Evaluate the log-likelihood at the given parameter values.

            Parameters
            ----------
            parameters : dict
                Dictionary of ModelParameter objects.
            parameter_values : dict
                The dictionary of model parameter values.
            t1,t2,dt : numeric
                Integration limits.
        """

        # TODO: need to figure out how to get P(D|X) values into here...
        #   - Maybe write custom ln_prior function that does all priors
        #       except the phase-space coords, then in the likelihood
        #       loop over the priors and get probs from there?

        ######################################################################
        # potential parameters:
        #
        pparams = dict()
        for par in parameters['potential']:
            # parameter is free to vary, take value from the dictionary of
            #   variable parameter values
            if par.frozen is False:
                val = parameter_values['potential'][par.name]

            # parameter is frozen to some value - take value from what it is
            #   frozen to
            else:
                val = par.frozen

            pparams[par.name] = val
        potential = self.potential_class(**pparams)

        ######################################################################
        # progenitor parameters:
        #

        # mass
        if parameters['progenitor']['mass'].frozen is False:
            mass = parameter_values['progenitor']['mass']
        else:
            mass = parameters['progenitor']['mass'].frozen

        # mass-loss
        if parameters['progenitor']['mdot'].frozen is False:
            mdot = parameter_values['progenitor']['mdot']
        else:
            mdot = parameters['progenitor']['mdot'].frozen

        # position of effective tidal radius
        if parameters['progenitor']['alpha'].frozen is False:
            alpha = parameter_values['progenitor']['alpha']
        else:
            alpha = parameters['progenitor']['alpha'].frozen

        # satellite coordinates
        prog_helct = np.empty((1,6))
        for ii,par_name in enumerate(heliocentric_names):
            if parameters['progenitor'][par_name].frozen is False:
                prog_helct[:,i] = parameter_values['progenitor'][par_name]
            else:
                prog_helct[:,i] = parameters['progenitor'][par_name].frozen
        prog_galct = _hel_to_gc(prog_helct) # TODO: need better way to do this so can specify R_sun and V_circ if I want in the future

        ######################################################################
        # star parameters:
        #

        # tail assignment
        if parameters['stars']['beta'].frozen is False:
            beta = parameter_values['stars']['beta']
        else:
            beta = parameters['stars']['beta'].frozen

        # star coordinates
        star_helct = np.empty((self.nstars,6))
        for ii,par_name in enumerate(heliocentric_names):
            if parameters['stars'][par_name].frozen is False:
                star_helct[:,i] = parameter_values['stars'][par_name]
            else:
                star_helct[:,i] = parameters['stars'][par_name].frozen
        star_galct = _hel_to_gc(star_helct) # TODO: need better way to do this so can specify R_sun and V_circ if I want in the future

        # TODO:
        ln_like = rewinder_likelihood(t1, t2, dt, pparams,
                                      prog_galct, star_galct,
                                      logmass, logmdot, alpha, betas)

        return np.sum(ln_like + data_like) + np.squeeze(sat_like)

# TODO: check here down

class RewinderSampler(EnsembleSampler):

    def __init__(self, model, nwalkers=None, pool=None, a=2.):
        """ """

        if nwalkers is None:
            nwalkers = model.nparameters*2 + 2
        self.nwalkers = nwalkers

        super(StreamModelSampler, self).__init__(self.nwalkers, model.nparameters, model,
                                                 pool=pool, a=a)

    def write(self, filename, ii=None):
        if ii is None:
            ii = self.chain.shape[1]

        # write the sampler data to an HDF5 file
        logger.info("Writing sampler data to '{}'...".format(filename))
        with h5py.File(filename, "w") as f:
            f["last_step"] = ii
            f["chain"] = self.chain
            f["lnprobability"] = self.lnprobability
            f["acceptance_fraction"] = self.acceptance_fraction
            try:
                f["acor"] = self.acor
            except:
                logger.warn("Failed to compute autocorrelation time.")
                f["acor"] = []

    def run_inference(self, pos, nsteps, path, output_every=None,
                      output_file_fmt="inference_{:06d}.hdf5", first_step=0):
        """ Custom run MCMC that caches the sampler every specified number
            of steps.
        """
        if output_every is None:
            output_every = nsteps

        logger.info("Running {} walkers for {} steps..."\
                .format(self.nwalkers, nsteps))

        time0 = time.time()
        ii = first_step
        for outer_loop in range(nsteps // output_every):
            self.reset()
            for results in self.sample(pos, iterations=output_every):
                ii += 1

            self.write(os.path.join(path,output_file_fmt.format(ii)), ii=ii)
            pos = results[0]

        # the remainder
        remainder = nsteps % output_every
        if remainder > 0:
            self.reset()
            for results in self.sample(pos, iterations=remainder):
                ii += 1

            self.write(os.path.join(path,output_file_fmt.format(ii)), ii=ii)

        t = time.time() - time0
        logger.debug("Spent {} seconds on main sampling...".format(t))

<commit_msg>accept kinematic parameters<commit_after># coding: utf-8

""" Core functionality for doing stream inference """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
from collections import OrderedDict
import time
import random

# Third-party
from astropy import log as logger
import astropy.table as at
from emcee import EnsembleSampler, PTSampler
import numpy as np
import astropy.units as u
import h5py
import numexpr

# Project
from streamteam.inference import EmceeModel, ModelParameter, LogUniformPrior, LogPrior
from .. import heliocentric_names
from ..util import streamspath
from .rewinder_likelihood import rewinder_likelihood
from .kinematicobject import KinematicObject

__all__ = ["Rewinder", "RewinderSampler"]

class Rewinder(EmceeModel):

    def __init__(self, potential, progenitor, stars, t1, t2, dt=-1.):
        """ Model for tidal streams that uses backwards integration to Rewind
            the positions of stars.

            Parameters
            ----------
            potential : Potential object
            progenitor :
            stars :
        """

        super(Rewinder, self).__init__(lambda x: None)

        # Potential
        for par in potential.parameters.values():
            self.add_parameter(par, "potential")
        self.potential_class = potential.c_class

        # Progenitor
        for par in progenitor.parameters.values():
            self.add_parameter(par, "progenitor")

        # Stars
        for par in stars.parameters.values():
            self.add_parameter(par, "stars")
        self.nstars = len(stars)

        self.args = (t1,t2,dt)

    @classmethod
    def from_config(cls, config):
        """ Construct a StreamModel from a configuration dictionary.
            Typically comes from a YAML file via `streams.io.read_config`.

            Parameters
            ----------
            config : dict
        """

        if not isinstance(config, dict):
            from ..io import read_config
            config = read_config(config)

        seed = config.get('seed', np.random.randint(100000))
        logger.debug("Using seed: {}".format(seed))
        np.random.seed(seed)
        random.seed(seed)

        nlead = config.get('nlead', 0)
        ntrail = config.get('ntrail', 0)
        nstars = nlead + ntrail

        # more complicated selections
        star_idx = config.get('star_idx', None)
        expr = config.get('expr', None)

        if star_idx is not None and nstars is not 0 and len(star_idx) != 0:
            raise ValueError("Number of stars does not match length of"
                             "star index specification.")

        # load progenitor and star data
        logger.debug("Reading data from:\n\t{}".format(config['data_file']))

        progenitor = at.Table.read(config['data_file'], path='progenitor')
        true_progenitor = at.Table.read(config['data_file'], path='true_progenitor')
        progenitor_err = at.Table.read(config['data_file'], path='error_progenitor')

        stars = at.Table.read(config['data_file'], path='stars')
        true_stars = at.Table.read(config['data_file'], path='true_stars')
        stars_err = at.Table.read(config['data_file'], path='error_stars')

        # progenitor parameter container
        # TODO: handle missing dimensions
        prog_ko = KinematicObject(hel=dict([(k,progenitor[k].data) for k in heliocentric_names]),
                          errors=dict([(k,progenitor_err[k].data) for k in heliocentric_names]),
                          truths=dict([(k,true_progenitor[k].data) for k in heliocentric_names]))

        prog_ko.parameters['m0'] = ModelParameter('m0', truth=float(progenitor['m0']),
                                                  prior=LogPrior()) # TODO: logspace?
        true_mdot = np.log(3.2*10**(np.floor(np.log10(float(progenitor['m0'])))-4))# THIS IS A HACK FOR 2.5e8!
        prog_ko.parameters['mdot'] = ModelParameter('mdot', truth=true_mdot,
                                                    prior=LogPrior()) # TODO: logspace?
        prog_ko.parameters['alpha'] = ModelParameter('alpha', shape=(1,),
                                                     prior=LogUniformPrior(1., 2.))

        # deal with extra star selection crap
        if star_idx is None:
            if nstars is 0:
                raise ValueError("If not specifying indexes, must specicy number"
                                 "of stars (nstars)")

            if expr is not None:
                expr_idx = numexpr.evaluate(expr, stars)
            else:
                expr_idx = np.ones(len(stars)).astype(bool)

            # leading tail stars
            lead_stars, = np.where((stars["tail"] == -1.) & expr_idx)
            np.random.shuffle(lead_stars)
            lead_stars = lead_stars[:nlead]

            # trailing tail stars
            trail_stars, = np.where((stars["tail"] == 1.) & expr_idx)
            np.random.shuffle(trail_stars)
            trail_stars = trail_stars[:ntrail]

            star_idx = np.append(lead_stars, trail_stars)

        stars = stars[star_idx]
        true_stars = true_stars[star_idx]
        stars_err = stars_err[star_idx]
        # TODO: handle missing dimensions
        logger.info("Running with {} stars.".format(len(stars)))
        logger.debug("Using stars: {}".format(list(star_idx)))

        star_ko = KinematicObject(hel=dict([(k,stars[k].data) for k in heliocentric_names]),
                                errors=dict([(k,stars_err[k].data) for k in heliocentric_names]),
                                truths=dict([(k,true_stars[k].data) for k in heliocentric_names]))
        star_ko.parameters['tail'] = ModelParameter('tail', truth=stars['tail'])

        # integration stuff
        integration = at.Table.read(config['data_file'], path='integration')
        t1 = config.get("t1", integration.meta['t1'])
        t2 = config.get("t2", integration.meta['t2'])
        dt = config.get("dt", -1.)
        logger.debug("Integration from {} to {}, ∆t={} Myr".format(t1,t2,dt))

        # get the potential object specified from the potential subpackage
        from .. import potential as sp
        Potential = getattr(sp, config["potential"]["class_name"])
        potential = Potential()
        logger.info("Using potential '{}'...".format(config["potential"]["class_name"]))

        for par in prog_ko.parameters.values():
            if par.name not in config["progenitor"]["parameters"]:
                par.freeze("truth")

        for par in star_ko.parameters.values():
            if par.name not in config["stars"]["parameters"]:
                par.freeze("truth")

        for par in potential.parameters.values():
            if par.name not in config["potential"]["parameters"]:
                par.freeze("truth")

        # Initialize the model
        model = cls(potential, prog_ko, star_ko, t1, t2, dt)

        return model

    def ln_prior(self, parameters, parameter_values, t1, t2, dt):
        """ Evaluate the log-prior at the given parameter values, but
            not for the star positions, which we need in the likelihood
            function.

            Parameters
            ----------
            parameters : dict
                Dictionary of ModelParameter objects.
            parameter_values : dict
                The dictionary of model parameter values.
            t1,t2,dt : numeric
                Integration limits.
        """
        ln_prior = 0.
        for group_name in ['potential', 'progenitor']:
            for param in parameters[group_name].keys():
                v = value_dict[group_name][param.name]
                ln_prior += param.prior(v)

        return ln_prior

    def ln_likelihood(self, parameters, parameter_values, t1, t2, dt):
        """ Evaluate the log-likelihood at the given parameter values.

            Parameters
            ----------
            parameters : dict
                Dictionary of ModelParameter objects.
            parameter_values : dict
                The dictionary of model parameter values.
            t1,t2,dt : numeric
                Integration limits.
        """

        # TODO: need to figure out how to get P(D|X) values into here...
        #   - Maybe write custom ln_prior function that does all priors
        #       except the phase-space coords, then in the likelihood
        #       loop over the priors and get probs from there?

        ######################################################################
        # potential parameters:
        #
        pparams = dict()
        for par in parameters['potential']:
            # parameter is free to vary, take value from the dictionary of
            #   variable parameter values
            if par.frozen is False:
                val = parameter_values['potential'][par.name]

            # parameter is frozen to some value - take value from what it is
            #   frozen to
            else:
                val = par.frozen

            pparams[par.name] = val
        potential = self.potential_class(**pparams)

        ######################################################################
        # progenitor parameters:
        #

        # mass
        if parameters['progenitor']['m0'].frozen is False:
            mass = parameter_values['progenitor']['m0']
        else:
            mass = parameters['progenitor']['m0'].frozen

        # mass-loss
        if parameters['progenitor']['mdot'].frozen is False:
            mdot = parameter_values['progenitor']['mdot']
        else:
            mdot = parameters['progenitor']['mdot'].frozen

        # position of effective tidal radius
        if parameters['progenitor']['alpha'].frozen is False:
            alpha = parameter_values['progenitor']['alpha']
        else:
            alpha = parameters['progenitor']['alpha'].frozen

        # satellite coordinates
        prog_helct = np.empty((1,6))
        for ii,par_name in enumerate(heliocentric_names):
            if parameters['progenitor'][par_name].frozen is False:
                prog_helct[:,i] = parameter_values['progenitor'][par_name]
            else:
                prog_helct[:,i] = parameters['progenitor'][par_name].frozen
        prog_galct = _hel_to_gc(prog_helct) # TODO: need better way to do this so can specify R_sun and V_circ if I want in the future

        ######################################################################
        # star parameters:
        #

        # tail assignment
        if parameters['stars']['beta'].frozen is False:
            beta = parameter_values['stars']['beta']
        else:
            beta = parameters['stars']['beta'].frozen

        # star coordinates
        star_helct = np.empty((self.nstars,6))
        for ii,par_name in enumerate(heliocentric_names):
            if parameters['stars'][par_name].frozen is False:
                star_helct[:,i] = parameter_values['stars'][par_name]
            else:
                star_helct[:,i] = parameters['stars'][par_name].frozen
        star_galct = _hel_to_gc(star_helct) # TODO: need better way to do this so can specify R_sun and V_circ if I want in the future

        # TODO:
        ln_like = rewinder_likelihood(t1, t2, dt, pparams,
                                      prog_galct, star_galct,
                                      logmass, logmdot, alpha, betas)

        return np.sum(ln_like + data_like) + np.squeeze(sat_like)

# TODO: check here down

class RewinderSampler(EnsembleSampler):

    def __init__(self, model, nwalkers=None, pool=None, a=2.):
        """ """

        if nwalkers is None:
            nwalkers = model.nparameters*2 + 2
        self.nwalkers = nwalkers

        super(StreamModelSampler, self).__init__(self.nwalkers, model.nparameters, model,
                                                 pool=pool, a=a)

    def write(self, filename, ii=None):
        if ii is None:
            ii = self.chain.shape[1]

        # write the sampler data to an HDF5 file
        logger.info("Writing sampler data to '{}'...".format(filename))
        with h5py.File(filename, "w") as f:
            f["last_step"] = ii
            f["chain"] = self.chain
            f["lnprobability"] = self.lnprobability
            f["acceptance_fraction"] = self.acceptance_fraction
            try:
                f["acor"] = self.acor
            except:
                logger.warn("Failed to compute autocorrelation time.")
                f["acor"] = []

    def run_inference(self, pos, nsteps, path, output_every=None,
                      output_file_fmt="inference_{:06d}.hdf5", first_step=0):
        """ Custom run MCMC that caches the sampler every specified number
            of steps.
        """
        if output_every is None:
            output_every = nsteps

        logger.info("Running {} walkers for {} steps..."\
                .format(self.nwalkers, nsteps))

        time0 = time.time()
        ii = first_step
        for outer_loop in range(nsteps // output_every):
            self.reset()
            for results in self.sample(pos, iterations=output_every):
                ii += 1

            self.write(os.path.join(path,output_file_fmt.format(ii)), ii=ii)
            pos = results[0]

        # the remainder
        remainder = nsteps % output_every
        if remainder > 0:
            self.reset()
            for results in self.sample(pos, iterations=remainder):
                ii += 1

            self.write(os.path.join(path,output_file_fmt.format(ii)), ii=ii)

        t = time.time() - time0
        logger.debug("Spent {} seconds on main sampling...".format(t))

<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import logging

# Third-party
import numpy as np
import astropy.units as u

__all__ = ["LogPrior", "LogUniformPrior", "LogNormal1DPrior"]

logger = logging.getLogger(__name__)

class LogPrior(object):

    def eval(self, value):
        return 0.

    def __call__(self, value):
        return self.eval(value)

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        raise ValueError("Cannot sample from a LogPrior object.")

    def __str__(self):
        return "<LogPrior>"

    def __repr__(self):
        return self.__str__()

class LogUniformPrior(LogPrior):

    def __init__(self, a, b):
        """ Log of a uniform distribution. Returns 0 if value is
            outside of the range defined by a < value < b.

            Parameters
            ----------
            a : numeric, quantity_like, array_like
                Lower bound.
            b : numeric, quantity_like, array_like
                Lower bound.
        """
        self.a = np.array(a)
        self.b = np.array(b)

        if self.a.shape != self.b.shape:
            raise ValueError("a and b must match in shape!")

    def __str__(self):
        return "<Uniform a={}, b={}>".format(self.a, self.b)

    def eval(self, value):
        p = np.atleast_1d(np.log(1 / (self.b - self.a)))
        ix = np.atleast_1d((value < self.a) | (value > self.b))
        p[ix] = -np.inf
        return np.squeeze(p)

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        if n is not None:
            return np.random.uniform(self.a, self.b, size=(n,) + self.a.shape)
        else:
            return np.random.uniform(self.a, self.b)

class LogNormal1DPrior(LogPrior):

    def __init__(self, mean, stddev):
        self.mean = np.array(mean)
        self.stddev = np.array(stddev)
        self._norm = -0.5*np.log(2*np.pi) - np.log(self.stddev)

    def __str__(self):
        return "<Normal1D μ={}, σ={}>".format(self.mean, self.stddev)

    def eval(self, value):
        X = self.mean - value
        return self._norm - 0.5*(X / self.stddev)**2

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        if n is not None:
            return np.random.normal(self.mean, self.stddev,
                                    size=(n,) + self.mean.shape)
        else:
            return np.random.normal(self.mean, self.stddev)

# class LogNormalPrior(LogPrior):

#     def __call__(self, value):
#         value = np.atleast_1d(value)
#         if value.ndim <= 1:
#             value = np.atleast_2d(value).T.copy()
#         else:
#             value = np.atleast_2d(value)

#         X = self.mean - value
#         arg = np.array([np.dot(XX, np.dot(IC,XX)) for XX,IC in zip(X,self.icov)])
#         return self._norm - 0.5*arg

#     def __init__(self, mean, stddev=None, cov=None):
#         """ Log of a normal distribution. Need to specify either
#             stddev, an array of standard deviations, or cov, an
#             array of covariance matrices.

#             Shape guide:
#                 (n,)  -> n 1-D Gaussians
#                 (n,m) -> n m-D Gaussians

#             Parameters
#             ----------
#             mean : array_like
#             stddev : array_like
#             cov : array_like
#         """
#         mean = np.atleast_1d(mean)
#         if mean.ndim <= 1:
#             self.mean = np.atleast_2d(mean).T.copy()
#         else:
#             self.mean = np.atleast_2d(mean)

#         self.shape = self.mean.shape
#         self.ndim = self.shape[-1]

#         if stddev is not None:
#             if mean.ndim <= 1:
#                 stddev = np.atleast_2d(stddev).T.copy()
#             else:
#                 stddev = np.atleast_2d(stddev)

#             if stddev.shape != self.shape:
#                 raise ValueError("stddev has wrong shape {}, should be {}"
#                                  .format(stddev.shape,self.shape))

#             cov = np.array([np.diag(s**2) for s in stddev])

#         elif cov is not None:
#             if cov.ndim < 2:
#                 raise ValueError("Invalid covariance matrix - must have 2 dims.")
#             elif cov.ndim == 2:
#                 cov = cov.reshape((1,)+cov.shape)

#         else:
#             raise ValueError("You must specify stddev or cov.")

#         self.stddev = stddev
#         self.cov = cov
#         self.icov = np.linalg.inv(cov)

#         dets = np.linalg.slogdet(self.cov)[1]
#         self._norm = -0.5*self.ndim*np.log(2*np.pi) - 0.5*dets

#     def sample(self, size=1):
#         s = np.array([np.random.multivariate_normal(m,C,size=size)
#                         for m,C in zip(self.mean, self.cov)])

#         return np.rollaxis(s, 1)<commit_msg>newline<commit_after># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import logging

# Third-party
import numpy as np
import astropy.units as u

__all__ = ["LogPrior", "LogUniformPrior", "LogNormal1DPrior"]

logger = logging.getLogger(__name__)

class LogPrior(object):

    def eval(self, value):
        return 0.

    def __call__(self, value):
        return self.eval(value)

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        raise ValueError("Cannot sample from a LogPrior object.")

    def __str__(self):
        return "<LogPrior>"

    def __repr__(self):
        return self.__str__()

class LogUniformPrior(LogPrior):

    def __init__(self, a, b):
        """ Log of a uniform distribution. Returns 0 if value is
            outside of the range defined by a < value < b.

            Parameters
            ----------
            a : numeric, quantity_like, array_like
                Lower bound.
            b : numeric, quantity_like, array_like
                Lower bound.
        """
        self.a = np.array(a)
        self.b = np.array(b)

        if self.a.shape != self.b.shape:
            raise ValueError("a and b must match in shape!")

    def __str__(self):
        return "<Uniform a={}, b={}>".format(self.a, self.b)

    def eval(self, value):
        p = np.atleast_1d(np.log(1 / (self.b - self.a)))
        ix = np.atleast_1d((value < self.a) | (value > self.b))
        p[ix] = -np.inf
        return np.squeeze(p)

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        if n is not None:
            return np.random.uniform(self.a, self.b, size=(n,) + self.a.shape)
        else:
            return np.random.uniform(self.a, self.b)

class LogNormal1DPrior(LogPrior):

    def __init__(self, mean, stddev):
        self.mean = np.array(mean)
        self.stddev = np.array(stddev)
        self._norm = -0.5*np.log(2*np.pi) - np.log(self.stddev)

    def __str__(self):
        return "<Normal1D μ={}, σ={}>".format(self.mean, self.stddev)

    def eval(self, value):
        X = self.mean - value
        return self._norm - 0.5*(X / self.stddev)**2

    def sample(self, n=None):
        """ Sample from this prior. The returned array axis=0 is the
            sample axis.

            Parameters
            ----------
            n : int (optional)
                Number of samples to draw
        """
        if n is not None:
            return np.random.normal(self.mean, self.stddev,
                                    size=(n,) + self.mean.shape)
        else:
            return np.random.normal(self.mean, self.stddev)

# class LogNormalPrior(LogPrior):

#     def __call__(self, value):
#         value = np.atleast_1d(value)
#         if value.ndim <= 1:
#             value = np.atleast_2d(value).T.copy()
#         else:
#             value = np.atleast_2d(value)

#         X = self.mean - value
#         arg = np.array([np.dot(XX, np.dot(IC,XX)) for XX,IC in zip(X,self.icov)])
#         return self._norm - 0.5*arg

#     def __init__(self, mean, stddev=None, cov=None):
#         """ Log of a normal distribution. Need to specify either
#             stddev, an array of standard deviations, or cov, an
#             array of covariance matrices.

#             Shape guide:
#                 (n,)  -> n 1-D Gaussians
#                 (n,m) -> n m-D Gaussians

#             Parameters
#             ----------
#             mean : array_like
#             stddev : array_like
#             cov : array_like
#         """
#         mean = np.atleast_1d(mean)
#         if mean.ndim <= 1:
#             self.mean = np.atleast_2d(mean).T.copy()
#         else:
#             self.mean = np.atleast_2d(mean)

#         self.shape = self.mean.shape
#         self.ndim = self.shape[-1]

#         if stddev is not None:
#             if mean.ndim <= 1:
#                 stddev = np.atleast_2d(stddev).T.copy()
#             else:
#                 stddev = np.atleast_2d(stddev)

#             if stddev.shape != self.shape:
#                 raise ValueError("stddev has wrong shape {}, should be {}"
#                                  .format(stddev.shape,self.shape))

#             cov = np.array([np.diag(s**2) for s in stddev])

#         elif cov is not None:
#             if cov.ndim < 2:
#                 raise ValueError("Invalid covariance matrix - must have 2 dims.")
#             elif cov.ndim == 2:
#                 cov = cov.reshape((1,)+cov.shape)

#         else:
#             raise ValueError("You must specify stddev or cov.")

#         self.stddev = stddev
#         self.cov = cov
#         self.icov = np.linalg.inv(cov)

#         dets = np.linalg.slogdet(self.cov)[1]
#         self._norm = -0.5*self.ndim*np.log(2*np.pi) - 0.5*dets

#     def sample(self, size=1):
#         s = np.array([np.random.multivariate_normal(m,C,size=size)
#                         for m,C in zip(self.mean, self.cov)])

#         return np.rollaxis(s, 1)
<|endoftext|>


<commit_before>"""Analyzes backwards-optimization experiment."""

import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot
from generalexam.machine_learning import evaluation_utils
from gewittergefahr.deep_learning import backwards_optimization as backwards_opt
from gewittergefahr.plotting import plotting_utils

L2_WEIGHTS = numpy.logspace(-4, 1, num=11)
MINMAX_WEIGHTS = numpy.logspace(-4, 1, num=11)

TOP_EXPERIMENT_DIR_NAME = (
    '/glade/p/univ/uokl0026/myrorss_experiment/conus_only/'
    'dropout=0.750_l2=0.003162_num-dense-layers=2_data-aug=1/testing/'
    'extreme_examples/unique_storm_cells/bwo_experiment_best_hits'
)

COLOUR_MAP_OBJECT = pyplot.get_cmap('viridis')
FIGURE_RESOLUTION_DPI = 300

FONT_SIZE = 30
pyplot.rc('font', size=FONT_SIZE)
pyplot.rc('axes', titlesize=FONT_SIZE)
pyplot.rc('axes', labelsize=FONT_SIZE)
pyplot.rc('xtick', labelsize=FONT_SIZE)
pyplot.rc('ytick', labelsize=FONT_SIZE)
pyplot.rc('legend', fontsize=FONT_SIZE)
pyplot.rc('figure', titlesize=FONT_SIZE)

def _run():
    """Analyzes backwards-optimization experiment.

    This is effectively the main method.
    """

    num_l2_weights = len(L2_WEIGHTS)
    num_minmax_weights = len(MINMAX_WEIGHTS)

    mean_final_activation_matrix = numpy.full(
        (num_l2_weights, num_minmax_weights), numpy.nan
    )

    for i in range(num_l2_weights):
        for j in range(num_minmax_weights):
            this_file_name = (
                '{0:s}/bwo_pmm_l2-weight={1:.10f}_minmax-weight={2:.10f}.p'
            ).format(
                TOP_EXPERIMENT_DIR_NAME, L2_WEIGHTS[i], MINMAX_WEIGHTS[j]
            )

            print('Reading data from: "{0:s}"...'.format(this_file_name))
            this_bwo_dict = backwards_opt.read_file(this_file_name)[0]
            mean_final_activation_matrix[i, j] = this_bwo_dict[
                backwards_opt.MEAN_FINAL_ACTIVATION_KEY
            ]

    x_tick_labels = ['{0:.1f}'.format(r) for r in numpy.log10(MINMAX_WEIGHTS)]
    y_tick_labels = ['{0:.1f}'.format(w) for w in numpy.log10(L2_WEIGHTS)]

    axes_object = evaluation_utils.plot_scores_2d(
        score_matrix=mean_final_activation_matrix,
        x_tick_label_strings=x_tick_labels, y_tick_label_strings=y_tick_labels,
        colour_map_object=COLOUR_MAP_OBJECT,
        min_colour_value=0., max_colour_value=1.
    )

    axes_object.set_xlabel(r'Min-max weight (log$_{10}$)')
    axes_object.set_ylabel(r'L$_2$ weight (log$_{10}$)')

    plotting_utils.plot_linear_colour_bar(
        axes_object_or_matrix=axes_object,
        data_matrix=mean_final_activation_matrix,
        colour_map_object=COLOUR_MAP_OBJECT, min_value=0., max_value=1.,
        orientation_string='vertical', extend_min=False, extend_max=False,
        font_size=FONT_SIZE
    )

    output_file_name = '{0:s}/mean_final_activations.jpg'.format(
        TOP_EXPERIMENT_DIR_NAME
    )

    print('Saving figure to: "{0:s}"...'.format(output_file_name))
    pyplot.savefig(
        output_file_name, dpi=FIGURE_RESOLUTION_DPI,
        pad_inches=0, bbox_inches='tight'
    )
    pyplot.close()

if __name__ == '__main__':
    _run()
<commit_msg>Fixed stupid bug in MYRORSS version of analyze_bwo_experiment.py.<commit_after>"""Analyzes backwards-optimization experiment."""

import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot
from gewittergefahr.gg_utils import model_evaluation
from gewittergefahr.deep_learning import backwards_optimization as backwards_opt
from gewittergefahr.plotting import plotting_utils

L2_WEIGHTS = numpy.logspace(-4, 1, num=11)
MINMAX_WEIGHTS = numpy.logspace(-4, 1, num=11)

TOP_EXPERIMENT_DIR_NAME = (
    '/glade/p/univ/uokl0026/myrorss_experiment/conus_only/'
    'dropout=0.750_l2=0.003162_num-dense-layers=2_data-aug=1/testing/'
    'extreme_examples/unique_storm_cells/bwo_experiment_best_hits'
)

COLOUR_MAP_OBJECT = pyplot.get_cmap('viridis')
FIGURE_RESOLUTION_DPI = 300

FONT_SIZE = 30
pyplot.rc('font', size=FONT_SIZE)
pyplot.rc('axes', titlesize=FONT_SIZE)
pyplot.rc('axes', labelsize=FONT_SIZE)
pyplot.rc('xtick', labelsize=FONT_SIZE)
pyplot.rc('ytick', labelsize=FONT_SIZE)
pyplot.rc('legend', fontsize=FONT_SIZE)
pyplot.rc('figure', titlesize=FONT_SIZE)

def _run():
    """Analyzes backwards-optimization experiment.

    This is effectively the main method.
    """

    num_l2_weights = len(L2_WEIGHTS)
    num_minmax_weights = len(MINMAX_WEIGHTS)

    mean_final_activation_matrix = numpy.full(
        (num_l2_weights, num_minmax_weights), numpy.nan
    )

    for i in range(num_l2_weights):
        for j in range(num_minmax_weights):
            this_file_name = (
                '{0:s}/bwo_pmm_l2-weight={1:.10f}_minmax-weight={2:.10f}.p'
            ).format(
                TOP_EXPERIMENT_DIR_NAME, L2_WEIGHTS[i], MINMAX_WEIGHTS[j]
            )

            print('Reading data from: "{0:s}"...'.format(this_file_name))
            this_bwo_dict = backwards_opt.read_file(this_file_name)[0]
            mean_final_activation_matrix[i, j] = this_bwo_dict[
                backwards_opt.MEAN_FINAL_ACTIVATION_KEY
            ]

    x_tick_labels = ['{0:.1f}'.format(r) for r in numpy.log10(MINMAX_WEIGHTS)]
    y_tick_labels = ['{0:.1f}'.format(w) for w in numpy.log10(L2_WEIGHTS)]

    axes_object = model_evaluation.plot_hyperparam_grid(
        score_matrix=mean_final_activation_matrix,
        colour_map_object=COLOUR_MAP_OBJECT,
        min_colour_value=0., max_colour_value=1.
    )

    axes_object.set_xticklabels(x_tick_labels, rotation=90.)
    axes_object.set_yticklabels(y_tick_labels)

    axes_object.set_xlabel(r'Min-max weight (log$_{10}$)')
    axes_object.set_ylabel(r'L$_2$ weight (log$_{10}$)')

    plotting_utils.plot_linear_colour_bar(
        axes_object_or_matrix=axes_object,
        data_matrix=mean_final_activation_matrix,
        colour_map_object=COLOUR_MAP_OBJECT, min_value=0., max_value=1.,
        orientation_string='vertical', extend_min=False, extend_max=False,
        font_size=FONT_SIZE
    )

    output_file_name = '{0:s}/mean_final_activations.jpg'.format(
        TOP_EXPERIMENT_DIR_NAME
    )

    print('Saving figure to: "{0:s}"...'.format(output_file_name))
    pyplot.savefig(
        output_file_name, dpi=FIGURE_RESOLUTION_DPI,
        pad_inches=0, bbox_inches='tight'
    )
    pyplot.close()

if __name__ == '__main__':
    _run()
<|endoftext|>


<commit_before>import sublime
import sublime_plugin

from queue import Queue
from threading import Thread
from zipfile import ZipFile
import os

from FSharp.lib import const
from FSharp.lib.fsac import get_server
from FSharp.lib import fs

tasks = Queue()

SIG_STOP = '__STOP__'

def plugin_loaded():
    """
    Initializes plugin.
    """

    # Install binaries if needed.
    if not installation.check_binaries():
        installation.install_binaries()
        print('FSharp: Binaries installed. Everything ok.')
    else:
        print('FSharp: Binaries found. Everything ok.')

    # Start the pipe server.
    AsyncPipe()

def plugin_unloaded():
    tasks.put((SIG_STOP, ()))

class AsyncPipe(object):
    """
    Wraps the fsac server to make it asynchronous.
    """
    def __init__(self):
        self.server = get_server()
        self.tasks = tasks

        writer = Thread(target=self.write, daemon=True)
        reader = Thread(target=self.read, daemon=True)

        writer.start()
        reader.start()

    def write(self):
        while True:
            action, args = self.tasks.get()
            method = getattr(self.server, action, None)

            if not method:
                process_output({'Kind': 'ERROR', 'Data': 'Not a valid call.'})
                continue

            if action == SIG_STOP:
                # Give the other thread a chance to exit.
                self.tasks.put((action, args))
                break

            # Write to server's stdin.
            method(*args)

    def read(self):
        while True:
            data = self.server.read_line()
            process_output(data)

            try:
                # Don't block here so we can read all the remaining output.
                action, args = self.tasks.get(timeout=0.01)
            except:
                continue

            if action == SIG_STOP:
                #  Give the other thread a chance to exit.
                self.tasks.put((action, args))
                break

            self.tasks.put((action, args))

class actions:
    """
    Groups methods that process data received from the autocomplete server.
    """
    @staticmethod
    def generic_action(data=None):
        sublime.status_message("RECEIVED: " + str(data))
        print("RECEIVED: " + str(data))

    @staticmethod
    def show_info(data):
        print(data)

    @staticmethod
    def find_declaration(data):
        data = data['Data']
        fname = data['File']
        row = data['Line'] + 1
        col = data['Column'] + 1
        encoded = "{0}:{1}:{2}".format(fname, row, col)
        sublime.active_window().open_file(encoded, sublime.ENCODED_POSITION)

    @staticmethod
    def declarations(data):
        decls = data['Data']
        print(decls)

    @staticmethod
    def show_completions(data):
        v = sublime.active_window().active_view()
        v.show_popup_menu(data['Data'], None)

    @staticmethod
    def show_tooltip(data):
        v = sublime.active_window().active_view()
        heading = list(data['Data'].keys())[0]
        body = data['Data'][heading]
        v.show_popup_menu([heading, body], None)

def process_output(data):
    action = None
    if data['Kind'] == 'completion':
        # Completions should normally be processed via events.
        # action = actions.show_completions
        pass
    elif data['Kind'] == 'helptext':
        action = actions.show_tooltip
    elif data['Kind'] == 'INFO':
        action = actions.show_info
    elif data['Kind'] == 'finddecl':
        action = actions.find_declaration
    elif data['Kind'] == 'declarations':
        action = actions.declarations
    elif data['Kind'] == 'project':
        for fname in data['Data']:
            tasks.put(('parse', (fname, True)))
    else:
        action = actions.generic_action

    if action:
        # Run action on the main UI thread to make ST happy.
        sublime.set_timeout(lambda: action(data), 0)

class installation:
    @staticmethod
    def check_binaries():
        print('FSharp: Checking installed files')
        return os.path.exists(const.path_to_fs_ac_binary())

    @staticmethod
    def install_binaries():
        print('FSharp: Installing files to Packages/FSharp_Binaries...')
        sublime.status_message('FSharp: Installing files to Packages/FSharp_Binaries...')
        try:
            os.mkdir(const.path_to_fs_binaries())
        except IOError:
            pass

        zipped_bytes = sublime.load_binary_resource('Packages/FSharp/bundled/fsautocomplete.zip')
        target = os.path.join(const.path_to_fs_binaries(), 'fsautocomplete.zip')
        with open(target, 'wb') as f:
            f.write(zipped_bytes)

        with open(target, 'rb') as f:
            ZipFile(f).extractall(path=const.path_to_fs_binaries())
        os.unlink(target)

class FsSetProjectFile(sublime_plugin.WindowCommand):
    def run(self):
        v = self.window.active_view()
        if not (v and v.file_name().endswith('.fsproj')):
            msg = 'FSharp: Not a project file.'
            print(msg)
            sublime.status_message(msg)
            return

        sublime.status_message('FSharp: Loading project...')
        tasks.put(('project', (v.file_name(),)))

class FsParseFile(sublime_plugin.WindowCommand):
    def run(self):
        v = self.window.active_view()
        fname = v.file_name()
        if not (v and fname.endswith(('.fs', '.fsi', '.fsx'))):
            msg = 'FSharp: Not a fsharp file.'
            print(msg)
            sublime.status_message(msg)
            return

        tasks.put(('parse', (fname, True)))

class FsFindDeclaration(sublime_plugin.WindowCommand):
    def run(self):
        v = self.window.active_view()
        if not (v and v.file_name().endswith(('.fs', '.fsx'))):
            msg = 'FSharp: Not an F# file.'
            print(msg)
            sublime.status_message(msg)
            return

        row, col = v.rowcol(v.sel()[0].b)
        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('find_declaration', (v.file_name(), row, col)))

class FsDeclarations(sublime_plugin.WindowCommand):
    def run(self):
        v = self.window.active_view()
        if not (v and v.file_name().endswith(('.fs', '.fsx'))):
            msg = 'FSharp: Not an F# file.'
            print(msg)
            sublime.status_message(msg)
            return

        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('declarations', (v.file_name(),)))

class FsFindCompletions(sublime_plugin.WindowCommand):
    def run(self):
        v = self.window.active_view()
        if not (v and v.file_name().endswith(('.fs', '.fsx'))):
            msg = 'FSharp: Not an F# file.'
            print(msg)
            sublime.status_message(msg)
            return

        row, col = v.rowcol(v.sel()[0].b)
        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('completions', (v.file_name(), row, col)))
<commit_msg>Use helper functions<commit_after>import sublime
import sublime_plugin

from queue import Queue
from threading import Thread
from zipfile import ZipFile
import os

from FSharp.lib import const
from FSharp.lib.fsac import get_server
from FSharp.lib import fs

tasks = Queue()

SIG_STOP = '__STOP__'

def plugin_loaded():
    """
    Initializes plugin.
    """

    # Install binaries if needed.
    if not installation.check_binaries():
        installation.install_binaries()
        print('FSharp: Binaries installed. Everything ok.')
    else:
        print('FSharp: Binaries found. Everything ok.')

    # Start the pipe server.
    AsyncPipe()

def plugin_unloaded():
    tasks.put((SIG_STOP, ()))

class AsyncPipe(object):
    """
    Wraps the fsac server to make it asynchronous.
    """
    def __init__(self):
        self.server = get_server()
        self.tasks = tasks

        writer = Thread(target=self.write, daemon=True)
        reader = Thread(target=self.read, daemon=True)

        writer.start()
        reader.start()

    def write(self):
        while True:
            action, args = self.tasks.get()
            method = getattr(self.server, action, None)

            if not method:
                process_output({'Kind': 'ERROR', 'Data': 'Not a valid call.'})
                continue

            if action == SIG_STOP:
                # Give the other thread a chance to exit.
                self.tasks.put((action, args))
                break

            # Write to server's stdin.
            method(*args)

    def read(self):
        while True:
            data = self.server.read_line()
            process_output(data)

            try:
                # Don't block here so we can read all the remaining output.
                action, args = self.tasks.get(timeout=0.01)
            except:
                continue

            if action == SIG_STOP:
                #  Give the other thread a chance to exit.
                self.tasks.put((action, args))
                break

            self.tasks.put((action, args))

class actions:
    """
    Groups methods that process data received from the autocomplete server.
    """
    @staticmethod
    def generic_action(data=None):
        sublime.status_message("RECEIVED: " + str(data))
        print("RECEIVED: " + str(data))

    @staticmethod
    def show_info(data):
        print(data)

    @staticmethod
    def find_declaration(data):
        data = data['Data']
        fname = data['File']
        row = data['Line'] + 1
        col = data['Column'] + 1
        encoded = "{0}:{1}:{2}".format(fname, row, col)
        sublime.active_window().open_file(encoded, sublime.ENCODED_POSITION)

    @staticmethod
    def declarations(data):
        decls = data['Data']
        print(decls)

    @staticmethod
    def show_completions(data):
        v = sublime.active_window().active_view()
        v.show_popup_menu(data['Data'], None)

    @staticmethod
    def show_tooltip(data):
        v = sublime.active_window().active_view()
        heading = list(data['Data'].keys())[0]
        body = data['Data'][heading]
        v.show_popup_menu([heading, body], None)

def process_output(data):
    action = None
    if data['Kind'] == 'completion':
        # Completions should normally be processed via events.
        # action = actions.show_completions
        pass
    elif data['Kind'] == 'helptext':
        action = actions.show_tooltip
    elif data['Kind'] == 'INFO':
        action = actions.show_info
    elif data['Kind'] == 'finddecl':
        action = actions.find_declaration
    elif data['Kind'] == 'declarations':
        action = actions.declarations
    elif data['Kind'] == 'project':
        for fname in data['Data']:
            tasks.put(('parse', (fname, True)))
    else:
        action = actions.generic_action

    if action:
        # Run action on the main UI thread to make ST happy.
        sublime.set_timeout(lambda: action(data), 0)

class installation:
    @staticmethod
    def check_binaries():
        print('FSharp: Checking installed files')
        return os.path.exists(const.path_to_fs_ac_binary())

    @staticmethod
    def install_binaries():
        print('FSharp: Installing files to Packages/FSharp_Binaries...')
        sublime.status_message('FSharp: Installing files to Packages/FSharp_Binaries...')
        try:
            os.mkdir(const.path_to_fs_binaries())
        except IOError:
            pass

        zipped_bytes = sublime.load_binary_resource('Packages/FSharp/bundled/fsautocomplete.zip')
        target = os.path.join(const.path_to_fs_binaries(), 'fsautocomplete.zip')
        with open(target, 'wb') as f:
            f.write(zipped_bytes)

        with open(target, 'rb') as f:
            ZipFile(f).extractall(path=const.path_to_fs_binaries())
        os.unlink(target)

class FsSetProjectFile(sublime_plugin.WindowCommand):
    def is_enabled(self):
        v = self.window.active_view()
        if v and fs.is_fsharp_project(v.file_name()):
            return True

        msg = 'FSharp: Not a project file.'
        print(msg)
        sublime.status_message(msg)
        return False

    def run(self):
        v = self.window.active_view()
        sublime.status_message('FSharp: Loading project...')
        tasks.put(('project', (v.file_name(),)))

class FsParseFile(sublime_plugin.WindowCommand):
    def is_enabled(self):
        v = self.window.active_view()
        if v and fs.is_fsharp_code(v.file_name()):
            return True

        msg = 'FSharp: Not an F# code file.'
        print(msg)
        sublime.status_message(msg)
        return False

    def run(self):
        v = self.window.active_view()
        tasks.put(('parse', (v.file_name(), True)))

class FsFindDeclaration(sublime_plugin.WindowCommand):
    def is_enabled(self):
        v = self.window.active_view()
        if v and fs.is_fsharp_code(v.file_name()):
            return True

        msg = 'FSharp: Not an F# code file.'
        print(msg)
        sublime.status_message(msg)
        return False

    def run(self):
        v = self.window.active_view()
        row, col = v.rowcol(v.sel()[0].b)
        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('find_declaration', (v.file_name(), row, col)))

class FsDeclarations(sublime_plugin.WindowCommand):
    def is_enabled(self):
        v = self.window.active_view()
        if v and fs.is_fsharp_code(v.file_name()):
            return True

        msg = 'FSharp: Not an F# code file.'
        print(msg)
        sublime.status_message(msg)
        return False

    def run(self):
        v = self.window.active_view()
        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('declarations', (v.file_name(),)))

class FsFindCompletions(sublime_plugin.WindowCommand):
    def is_enabled(self):
        v = self.window.active_view()
        if v and fs.is_fsharp_code(v.file_name()):
            return True

        msg = 'FSharp: Not an F# code file.'
        print(msg)
        sublime.status_message(msg)
        return False

    def run(self):
        v = self.window.active_view()
        row, col = v.rowcol(v.sel()[0].b)
        tasks.put(('parse', (v.file_name(), True)))
        tasks.put(('completions', (v.file_name(), row, col)))
<|endoftext|>


<commit_before>#!/usr/bin/python

import sys

bad_PDUs_hex = [
	# generic errors
	'01 02 00 00 00 00 00 08', # invalid protocol version = 1, otherwise valid Reset Query
	'00 ff 00 00 00 00 00 08', # invalid PDU type, otherwise valid Reset Query
	'00 00 f0 0f ff ff ff ff', # length too long TODO: this should try to send 0xFFFFFFFF bytes, maybe
	# TODO: PDUs with truncated headers?

	# Serial Notify
	'00 00 f0 0f 00 00 00 08', # length too short
	'00 00 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 00 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Serial Query
	'00 01 f0 0f 00 00 00 08', # length too short
	'00 01 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 01 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Reset Query
	'00 02 00 00 00 00 00 09 f0', # length too long

	# Cache Response
	'00 03 f0 0f 00 00 00 09 f0', # length too long

	# IPv4 Prefix
	'00 04 00 00 00 00 00 08', # length too short
	'00 04 00 00 00 00 00 0c f0 0f f0 0f', # length too short
	'00 04 00 00 ' + '00 00 00 13 ' + '00 00 00 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0', # length too short
	'00 04 00 00 ' + '00 00 00 15 ' + '00 00 00 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f ' + 'f0', # length too long
	'00 04 00 00 ' + '00 00 00 14 ' + '00 21 21 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f', # prefix length too large
	'00 04 00 00 ' + '00 00 00 14 ' + '00 00 21 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f', # max length too large
	# TODO: prefix length > max length

	# IPv6 Prefix
	'00 06 00 00 00 00 00 08', # length too short
	'00 06 00 00 00 00 00 0c f0 0f f0 0f', # length too short
	'00 06 00 00 ' + '00 00 00 1f ' + '00 00 00 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0', # length too short
	'00 06 00 00 ' + '00 00 00 21 ' + '00 00 00 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f ' + 'f0', # length too long
	'00 06 00 00 ' + '00 00 00 20 ' + '00 81 81 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f', # prefix length too large
	'00 06 00 00 ' + '00 00 00 20 ' + '00 00 81 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f', # max length too large
	# TODO: prefix length > max length

	# End of Data
	'00 07 f0 0f 00 00 00 08', # length too short
	'00 07 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 07 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Cache Reset
	'00 08 00 00 00 00 00 09 f0', # length too long

	# Error Report
	'00 0a 00 01 00 00 00 08', # Internal Error, length too short
	'00 0a 00 01 00 00 00 0c 00 00 00 00', # Internal Error, length too short
	'00 0a 00 01 00 00 00 0f 00 00 00 00 00 00 00', # Internal Error, length too short
	'00 0a 00 01 00 00 00 10 ' + '00 00 00 01 ' + 'f0 ' + '00 00 00', # Internal Error, truncated Length of Error Text field
	'00 0a 00 01 ' + '00 00 00 11 ' + '00 00 00 00 ' + '00 00 00 00 ' + 'f0', # Internal Error with internal lengths less than total length
	'00 0a 00 01 ' + '00 00 00 10 ' + '00 00 00 00 ' + '00 00 00 01', # Internal Error with internal lengths greater than total length
]

bad_PDUs = [''.join([chr(int(c, 0x10)) for c in s.split(' ')]) for s in bad_PDUs_hex]

if __name__ == '__main__':
	def print_usage():
		print("Usage: %s <index> | length" % sys.argv[0])
		print("    %s <index>: output the bad PDU at the 1-based index specified" % sys.argv[0])
		print("    %s length: output the total number of bad PDUs" % sys.argv[0])
		sys.exit(1)

	if len(sys.argv) != 2:
		print_usage()

	if sys.argv[1] == 'length':
		print len(bad_PDUs)
		sys.exit()

	try:
		index = int(sys.argv[1])
	except ValueError:
		print_usage()

	if index < 1:
		print_usage()
	if index > len(bad_PDUs):
		sys.exit("Error: index %d is greater than the total number of PDUs (%d)" % (index, len(bad_PDUs)))

	print bad_PDUs[index - 1]
<commit_msg>remove TODO's for prefix length > max length<commit_after>#!/usr/bin/python

import sys

bad_PDUs_hex = [
	# generic errors
	'01 02 00 00 00 00 00 08', # invalid protocol version = 1, otherwise valid Reset Query
	'00 ff 00 00 00 00 00 08', # invalid PDU type, otherwise valid Reset Query
	'00 00 f0 0f ff ff ff ff', # length too long TODO: this should try to send 0xFFFFFFFF bytes, maybe
	# TODO: PDUs with truncated headers?

	# Serial Notify
	'00 00 f0 0f 00 00 00 08', # length too short
	'00 00 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 00 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Serial Query
	'00 01 f0 0f 00 00 00 08', # length too short
	'00 01 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 01 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Reset Query
	'00 02 00 00 00 00 00 09 f0', # length too long

	# Cache Response
	'00 03 f0 0f 00 00 00 09 f0', # length too long

	# IPv4 Prefix
	'00 04 00 00 00 00 00 08', # length too short
	'00 04 00 00 00 00 00 0c f0 0f f0 0f', # length too short
	'00 04 00 00 ' + '00 00 00 13 ' + '00 00 00 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0', # length too short
	'00 04 00 00 ' + '00 00 00 15 ' + '00 00 00 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f ' + 'f0', # length too long
	'00 04 00 00 ' + '00 00 00 14 ' + '00 21 21 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f', # prefix length too large
	'00 04 00 00 ' + '00 00 00 14 ' + '00 00 21 00 ' + 'f0 0f f0 0f ' + 'f0 0f f0 0f', # max length too large

	# IPv6 Prefix
	'00 06 00 00 00 00 00 08', # length too short
	'00 06 00 00 00 00 00 0c f0 0f f0 0f', # length too short
	'00 06 00 00 ' + '00 00 00 1f ' + '00 00 00 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0', # length too short
	'00 06 00 00 ' + '00 00 00 21 ' + '00 00 00 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f ' + 'f0', # length too long
	'00 06 00 00 ' + '00 00 00 20 ' + '00 81 81 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f', # prefix length too large
	'00 06 00 00 ' + '00 00 00 20 ' + '00 00 81 00 ' + 'f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f f0 0f ' + 'f0 0f f0 0f', # max length too large

	# End of Data
	'00 07 f0 0f 00 00 00 08', # length too short
	'00 07 f0 0f 00 00 00 0b f0 0f f0', # length too short
	'00 07 f0 0f 00 00 00 0d f0 0f f0 0f f0', # length too long

	# Cache Reset
	'00 08 00 00 00 00 00 09 f0', # length too long

	# Error Report
	'00 0a 00 01 00 00 00 08', # Internal Error, length too short
	'00 0a 00 01 00 00 00 0c 00 00 00 00', # Internal Error, length too short
	'00 0a 00 01 00 00 00 0f 00 00 00 00 00 00 00', # Internal Error, length too short
	'00 0a 00 01 00 00 00 10 ' + '00 00 00 01 ' + 'f0 ' + '00 00 00', # Internal Error, truncated Length of Error Text field
	'00 0a 00 01 ' + '00 00 00 11 ' + '00 00 00 00 ' + '00 00 00 00 ' + 'f0', # Internal Error with internal lengths less than total length
	'00 0a 00 01 ' + '00 00 00 10 ' + '00 00 00 00 ' + '00 00 00 01', # Internal Error with internal lengths greater than total length
]

bad_PDUs = [''.join([chr(int(c, 0x10)) for c in s.split(' ')]) for s in bad_PDUs_hex]

if __name__ == '__main__':
	def print_usage():
		print("Usage: %s <index> | length" % sys.argv[0])
		print("    %s <index>: output the bad PDU at the 1-based index specified" % sys.argv[0])
		print("    %s length: output the total number of bad PDUs" % sys.argv[0])
		sys.exit(1)

	if len(sys.argv) != 2:
		print_usage()

	if sys.argv[1] == 'length':
		print len(bad_PDUs)
		sys.exit()

	try:
		index = int(sys.argv[1])
	except ValueError:
		print_usage()

	if index < 1:
		print_usage()
	if index > len(bad_PDUs):
		sys.exit("Error: index %d is greater than the total number of PDUs (%d)" % (index, len(bad_PDUs)))

	print bad_PDUs[index - 1]
<|endoftext|>


<commit_before>"""Flag feature tests."""

import pytest
from pytest_bdd import (
    given,
    scenario,
    then,
)

from ensign import BinaryFlag, DefaultStorage

@pytest.fixture
def flag(db):
    """
    Create a flag for the current feature under test, clean up when done.
    """
    return BinaryFlag.create("flag0")

@given('a feature flag named "flag0"')
def a_feature_flag_named_flag0():
    """a feature flag named "flag0"."""

@pytest.mark.spec
@scenario('flag.feature', 'Enable the feature')
def test_enable_the_feature():
    """Enable the feature."""

@given('the flag is enabled')
def the_flag_is_enabled(flag):
    """the flag is enabled."""
    flag.set()

@then('the code should run')
def the_code_should_run(flag):
    """the code should run."""
    @flag
    def testfun():
        return "executed"
    assert testfun() == "executed"

@pytest.mark.spec
@scenario('flag.feature', 'Disable the feature')
def test_disable_the_feature():
    """Disable the feature."""

@given('the flag is disabled')
def the_flag_is_disabled(flag):
    """the flag is disabled."""
    flag.unset()

@then('the code should not run')
def the_code_should_not_run(flag):
    """the code should not run."""
    @flag
    def testfun():
        return "executed"
    assert testfun() is None
<commit_msg>Remove unused import.<commit_after>"Flag feature tests."""

import pytest
from pytest_bdd import (
    given,
    scenario,
    then,
)

from ensign import BinaryFlag

@pytest.fixture
def flag(db):
    """
    Create a flag for the current feature under test, clean up when done.
    """
    return BinaryFlag.create("flag0")

@given('a feature flag named "flag0"')
def a_feature_flag_named_flag0():
    """a feature flag named "flag0"."""

@pytest.mark.spec
@scenario('flag.feature', 'Enable the feature')
def test_enable_the_feature():
    """Enable the feature."""

@given('the flag is enabled')
def the_flag_is_enabled(flag):
    """the flag is enabled."""
    flag.set()

@then('the code should run')
def the_code_should_run(flag):
    """the code should run."""
    @flag
    def testfun():
        return "executed"
    assert testfun() == "executed"

@pytest.mark.spec
@scenario('flag.feature', 'Disable the feature')
def test_disable_the_feature():
    """Disable the feature."""

@given('the flag is disabled')
def the_flag_is_disabled(flag):
    """the flag is disabled."""
    flag.unset()

@then('the code should not run')
def the_code_should_not_run(flag):
    """the code should not run."""
    @flag
    def testfun():
        return "executed"
    assert testfun() is None
<|endoftext|>


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

import region_proposal as rp
import file_io
import cv2
import os
import annotation as ann
import show

import matplotlib.pyplot as plt

N_IMAGES = 2
DIR = '../../datasets/svhn/train'

# 1. file 을 load
files = file_io.list_files(directory='../../datasets/svhn/train', n_files_to_sample=N_IMAGES, random_order=False)
annotation_file = "../../datasets/svhn/train/digitStruct.json"
detector = rp.MserDetector()

for image_file in files:
    image = cv2.imread(image_file)
    candidates = detector.detect(image, True)

    gts, _ = ann.get_annotation(image_file, annotation_file)
    show.plot_bounding_boxes(image, gts)

    # 3. Ground Truth 와의 overlap 이 5% 미만인 모든 sample 을 negative set 으로 저장
    # gts, candidates 의 overlap 을 구한다.

<commit_msg>negative sample 저장하는 Test code 완성<commit_after>#-*- coding: utf-8 -*-

import region_proposal as rp
import file_io
import cv2
import numpy as np
import os
import annotation as ann
import show
import eval

import matplotlib.pyplot as plt

N_IMAGES = 2
DIR = '../../datasets/svhn/train'

# 1. file 을 load
files = file_io.list_files(directory='../../datasets/svhn/train', n_files_to_sample=N_IMAGES, random_order=False)
annotation_file = "../../datasets/svhn/train/digitStruct.json"
detector = rp.MserDetector()

negative_samples = []
for image_file in files:
    image = cv2.imread(image_file)
    candidates = detector.detect(image, False)

    gts, _ = ann.get_annotation(image_file, annotation_file)
    #show.plot_bounding_boxes(image, gts)

    # gts, candidates 의 overlap 을 구한다.
    overlaps = eval.calc_iou(candidates, gts)
    #show.plot_bounding_boxes(image, candidates[overlaps>0.05])

    # Ground Truth 와의 overlap 이 5% 미만인 모든 sample 을 negative set 으로 저장
    negative_boxes = candidates[overlaps<0.05]

    # negative box 를 crop, resize to 32x32 해서 sample 에 추가
    import utils
    for bb in negative_boxes:
        sample = utils.crop_bb(image, bb, padding=0, dst_size=(32,32))
        negative_samples.append(sample)

negative_samples = np.array(negative_samples)
print negative_samples.shape

<|endoftext|>


<commit_before>#! /usr/bin/env python

import argparse
import os
import string
import time
import traceback
import subprocess
import signal

default_repo_url = 'https://github.com/Cyan4973/zstd.git'
working_dir_name = 'speedTest'
working_path = os.getcwd() + '/' + working_dir_name     # /path/to/zstd/tests/speedTest
clone_path = working_path + '/' + 'zstd'                # /path/to/zstd/tests/speedTest/zstd
email_header = '[ZSTD_speedTest]'
pid = str(os.getpid())

def log(text):
    print(time.strftime("%Y/%m/%d %H:%M:%S") + ' - ' + text)

def execute(command, print_output=False, print_error=True, param_shell=True):
    log("> " + command)
    popen = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=param_shell, cwd=execute.cwd)
    stdout = popen.communicate()[0]
    stdout_lines = stdout.splitlines()
    if print_output:
        print('\n'.join(stdout_lines))
    if popen.returncode is not None and popen.returncode != 0:
        if not print_output and print_error:
            print('\n'.join(stdout_lines))
        raise RuntimeError('\n'.join(stdout_lines))
    return stdout_lines
execute.cwd = None

def does_command_exist(command):
    try:
        execute(command, False, False);
    except Exception as e:
        return False
    return True

def send_email(emails, topic, text, have_mutt, have_mail):
    logFileName = working_path + '/' + 'tmpEmailContent'
    with open(logFileName, "w") as myfile:
        myfile.writelines(text)
        myfile.close()
        if have_mutt:
            execute('mutt -s "' + topic + '" ' + emails + ' < ' + logFileName)
        elif have_mail:
            execute('mail -s "' + topic + '" ' + emails + ' < ' + logFileName)
        else:
            log("e-mail cannot be sent (mail or mutt not found)")

def send_email_with_attachments(branch, commit, last_commit, emails, text, results_files, logFileName, lower_limit, have_mutt, have_mail):
    with open(logFileName, "w") as myfile:
        myfile.writelines(text)
        myfile.close()
        email_topic = '%s:%s Warning for %s:%s last_commit=%s speed<%s' % (email_header, pid, branch, commit, last_commit, lower_limit)
        if have_mutt:
            execute('mutt -s "' + email_topic + '" ' + emails + ' -a ' + results_files + ' < ' + logFileName)
        elif have_mail:
            execute('mail -s "' + email_topic + '" ' + emails + ' < ' + logFileName)
        else:
            log("e-mail cannot be sent (mail or mutt not found)")

def git_get_branches():
    execute('git fetch -p')
    output = execute('git branch -rl')
    for line in output:
        if "HEAD" in line:
            output.remove(line)  # remove "origin/HEAD -> origin/dev"
    return map(lambda l: l.strip(), output)

def git_get_changes(branch, commit, last_commit):
    fmt = '--format="%h: (%an) %s, %ar"'
    if last_commit is None:
        commits = execute('git log -n 10 %s %s' % (fmt, commit))
    else:
        commits = execute('git --no-pager log %s %s..%s' % (fmt, last_commit, commit))
    return str('Changes in %s since %s:\n' % (branch, last_commit)) + '\n'.join(commits)

def compile(branch, commit, last_commit, dry_run):
    local_branch = string.split(branch, '/')[1]
    version = local_branch.rpartition('-')[2]
    version = version + '_' + commit
    execute('git checkout -- . && git checkout ' + branch)
    print(git_get_changes(branch, commit, last_commit))
    if not dry_run:
        execute('make clean zstdprogram MOREFLAGS="-DZSTD_GIT_COMMIT=%s"' % version)

def get_last_results(resultsFileName):
    if not os.path.isfile(resultsFileName):
        return None, None, None
    commit = None
    cspeed = []
    dspeed = []
    with open(resultsFileName,'r') as f:
        for line in f:
            words = line.split()
            if len(words) == 2:   # branch + commit
                commit = words[1];
                cspeed = []
                dspeed = []
            if (len(words) == 8):  # results
                cspeed.append(float(words[3]))
                dspeed.append(float(words[5]))
    return commit, cspeed, dspeed

def benchmark_and_compare(branch, commit, resultsFileName, lastCLevel, testFilePath, fileName, last_cspeed, last_dspeed, lower_limit, maxLoadAvg, message):
    sleepTime = 30
    while os.getloadavg()[0] > maxLoadAvg:
        log("WARNING: bench loadavg=%.2f is higher than %s, sleeping for %s seconds" % (os.getloadavg()[0], maxLoadAvg, sleepTime))
        time.sleep(sleepTime)
    start_load = str(os.getloadavg())
    result = execute('programs/zstd -qi5b1e%s %s' % (lastCLevel, testFilePath), print_output=True)
    end_load = str(os.getloadavg())
    linesExpected = lastCLevel + 2;
    if len(result) != linesExpected:
        raise RuntimeError("ERROR: number of result lines=%d is different that expected %d\n%s" % (len(result), linesExpected, '\n'.join(result)))
    with open(resultsFileName, "a") as myfile:
        myfile.write(branch + " " + commit + "\n")
        myfile.write('\n'.join(result) + '\n')
        myfile.close()
        if (last_cspeed == None):
            log("WARNING: No data for comparison for branch=%s file=%s " % (branch, fileName))
            return ""
        commit, cspeed, dspeed = get_last_results(resultsFileName)
        text = ""
        for i in range(0, min(len(cspeed), len(last_cspeed))):
            print("%s:%s -%d cspeed=%6.2f clast=%6.2f cdiff=%1.4f dspeed=%6.2f dlast=%6.2f ddiff=%1.4f %s" % (branch, commit, i+1, cspeed[i], last_cspeed[i], cspeed[i]/last_cspeed[i], dspeed[i], last_dspeed[i], dspeed[i]/last_dspeed[i], fileName))
            if (cspeed[i]/last_cspeed[i] < lower_limit):
                text += "WARNING: -%d cspeed=%.2f clast=%.2f cdiff=%.4f %s\n" % (i+1, cspeed[i], last_cspeed[i], cspeed[i]/last_cspeed[i], fileName)
            if (dspeed[i]/last_dspeed[i] < lower_limit):
                text += "WARNING: -%d dspeed=%.2f dlast=%.2f ddiff=%.4f %s\n" % (i+1, dspeed[i], last_dspeed[i], dspeed[i]/last_dspeed[i], fileName)
        if text:
            text = message + ("\nmaxLoadAvg=%s  load average at start=%s end=%s\n" % (maxLoadAvg, start_load, end_load)) + text
        return text

def update_config_file(branch, commit):
    last_commit = None
    commitFileName = working_path + "/commit_" + branch.replace("/", "_") + ".txt"
    if os.path.isfile(commitFileName):
        last_commit = file(commitFileName, 'r').read()
    file(commitFileName, 'w').write(commit)
    return last_commit

def test_commit(branch, commit, last_commit, args, testFilePaths, have_mutt, have_mail):
    compile(branch, commit, last_commit, args.dry_run)
    logFileName = working_path + "/log_" + branch.replace("/", "_") + ".txt"
    text_to_send = []
    results_files = ""
    for filePath in testFilePaths:
        fileName = filePath.rpartition('/')[2]
        resultsFileName = working_path + "/results_" + branch.replace("/", "_") + "_" + fileName.replace(".", "_") + ".txt"
        last_commit, cspeed, dspeed = get_last_results(resultsFileName)

        if not args.dry_run:
            text = benchmark_and_compare(branch, commit, resultsFileName, args.lastCLevel, filePath, fileName, cspeed, dspeed, args.lowerLimit, args.maxLoadAvg, args.message)
            if text:
                log("WARNING: redoing tests for branch %s: commit %s" % (branch, commit))
                text = benchmark_and_compare(branch, commit, resultsFileName, args.lastCLevel, filePath, fileName, cspeed, dspeed, args.lowerLimit, args.maxLoadAvg, args.message)
                if text:
                    text_to_send.append(text)
                    results_files += resultsFileName + " "
    if text_to_send:
        send_email_with_attachments(branch, commit, last_commit, args.emails, text_to_send, results_files, logFileName, args.lowerLimit, have_mutt, have_mail)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('testFileNames', help='file names list for speed benchmark')
    parser.add_argument('emails', help='list of e-mail addresses to send warnings')
    parser.add_argument('--message', help='attach an additional message to e-mail', default="")
    parser.add_argument('--repoURL', help='changes default repository URL', default=default_repo_url)
    parser.add_argument('--lowerLimit', type=float, help='send email if speed is lower than given limit', default=0.98)
    parser.add_argument('--maxLoadAvg', type=float, help='maximum load average to start testing', default=0.75)
    parser.add_argument('--lastCLevel', type=int, help='last compression level for testing', default=5)
    parser.add_argument('--sleepTime', type=int, help='frequency of repository checking in seconds', default=300)
    parser.add_argument('--dry-run', dest='dry_run', action='store_true', help='not build', default=False)
    args = parser.parse_args()

    # check if test files are accessible
    testFileNames = args.testFileNames.split()
    testFilePaths = []
    for fileName in testFileNames:
        fileName = os.path.expanduser(fileName)
        if os.path.isfile(fileName):
            testFilePaths.append(os.path.abspath(fileName))
        else:
            log("ERROR: File not found: " + fileName)
            exit(1)

    # check availability of e-mail senders
    have_mutt = does_command_exist("mutt -h");
    have_mail = does_command_exist("mail -V");
    if not have_mutt and not have_mail:
        log("ERROR: e-mail senders 'mail' or 'mutt' not found")
        exit(1)

    print("PARAMETERS:\nrepoURL=%s" % args.repoURL)
    print("working_path=%s" % working_path)
    print("clone_path=%s" % clone_path)
    print("testFilePath(%s)=%s" % (len(testFilePaths), testFilePaths))
    print("message=%s" % args.message)
    print("emails=%s" % args.emails)
    print("maxLoadAvg=%s" % args.maxLoadAvg)
    print("lowerLimit=%s" % args.lowerLimit)
    print("lastCLevel=%s" % args.lastCLevel)
    print("sleepTime=%s" % args.sleepTime)
    print("dry_run=%s" % args.dry_run)
    print("have_mutt=%s have_mail=%s" % (have_mutt, have_mail))

    # clone ZSTD repo if needed
    if not os.path.isdir(working_path):
        os.mkdir(working_path)
    if not os.path.isdir(clone_path):
        execute.cwd = working_path
        execute('git clone ' + args.repoURL)
    if not os.path.isdir(clone_path):
        log("ERROR: ZSTD clone not found: " + clone_path)
        exit(1)
    execute.cwd = clone_path

    # check if speedTest.pid already exists
    pidfile = "./speedTest.pid"
    if os.path.isfile(pidfile):
        log("ERROR: %s already exists, exiting" % pidfile)
        exit(1)

    send_email(args.emails, email_header + ':%s test-zstd-speed.py has been started' % pid, args.message, have_mutt, have_mail)
    file(pidfile, 'w').write(pid)

    while True:
        try:
            loadavg = os.getloadavg()[0]
            if (loadavg <= args.maxLoadAvg):
                branches = git_get_branches()
                for branch in branches:
                    commit = execute('git show -s --format=%h ' + branch)[0]
                    last_commit = update_config_file(branch, commit)
                    if commit == last_commit:
                        log("skipping branch %s: head %s already processed" % (branch, commit))
                    else:
                        log("build branch %s: head %s is different from prev %s" % (branch, commit, last_commit))
                        test_commit(branch, commit, last_commit, args, testFilePaths, have_mutt, have_mail)
            else:
                log("WARNING: main loadavg=%.2f is higher than %s" % (loadavg, args.maxLoadAvg))
            log("sleep for %s seconds" % args.sleepTime)
            time.sleep(args.sleepTime)
        except Exception as e:
            stack = traceback.format_exc()
            email_topic = '%s:%s ERROR in %s:%s' % (email_header, pid, branch, commit)
            send_email(args.emails, email_topic, stack, have_mutt, have_mail)
            print(stack)
        except KeyboardInterrupt:
            os.unlink(pidfile)
            send_email(args.emails, email_header + ':%s test-zstd-speed.py has been stopped' % pid, args.message, have_mutt, have_mail)
            exit(0)
<commit_msg>test-zstd-speed.py: reordering of functions<commit_after>#! /usr/bin/env python

import argparse
import os
import string
import time
import traceback
import subprocess
import signal

default_repo_url = 'https://github.com/Cyan4973/zstd.git'
working_dir_name = 'speedTest'
working_path = os.getcwd() + '/' + working_dir_name     # /path/to/zstd/tests/speedTest
clone_path = working_path + '/' + 'zstd'                # /path/to/zstd/tests/speedTest/zstd
email_header = '[ZSTD_speedTest]'
pid = str(os.getpid())

def log(text):
    print(time.strftime("%Y/%m/%d %H:%M:%S") + ' - ' + text)

def execute(command, print_output=False, print_error=True, param_shell=True):
    log("> " + command)
    popen = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=param_shell, cwd=execute.cwd)
    stdout = popen.communicate()[0]
    stdout_lines = stdout.splitlines()
    if print_output:
        print('\n'.join(stdout_lines))
    if popen.returncode is not None and popen.returncode != 0:
        if not print_output and print_error:
            print('\n'.join(stdout_lines))
        raise RuntimeError('\n'.join(stdout_lines))
    return stdout_lines
execute.cwd = None

def does_command_exist(command):
    try:
        execute(command, False, False);
    except Exception as e:
        return False
    return True

def send_email(emails, topic, text, have_mutt, have_mail):
    logFileName = working_path + '/' + 'tmpEmailContent'
    with open(logFileName, "w") as myfile:
        myfile.writelines(text)
        myfile.close()
        if have_mutt:
            execute('mutt -s "' + topic + '" ' + emails + ' < ' + logFileName)
        elif have_mail:
            execute('mail -s "' + topic + '" ' + emails + ' < ' + logFileName)
        else:
            log("e-mail cannot be sent (mail or mutt not found)")

def send_email_with_attachments(branch, commit, last_commit, emails, text, results_files, logFileName, lower_limit, have_mutt, have_mail):
    with open(logFileName, "w") as myfile:
        myfile.writelines(text)
        myfile.close()
        email_topic = '%s:%s Warning for %s:%s last_commit=%s speed<%s' % (email_header, pid, branch, commit, last_commit, lower_limit)
        if have_mutt:
            execute('mutt -s "' + email_topic + '" ' + emails + ' -a ' + results_files + ' < ' + logFileName)
        elif have_mail:
            execute('mail -s "' + email_topic + '" ' + emails + ' < ' + logFileName)
        else:
            log("e-mail cannot be sent (mail or mutt not found)")

def git_get_branches():
    execute('git fetch -p')
    output = execute('git branch -rl')
    for line in output:
        if "HEAD" in line:
            output.remove(line)  # remove "origin/HEAD -> origin/dev"
    return map(lambda l: l.strip(), output)

def git_get_changes(branch, commit, last_commit):
    fmt = '--format="%h: (%an) %s, %ar"'
    if last_commit is None:
        commits = execute('git log -n 10 %s %s' % (fmt, commit))
    else:
        commits = execute('git --no-pager log %s %s..%s' % (fmt, last_commit, commit))
    return str('Changes in %s since %s:\n' % (branch, last_commit)) + '\n'.join(commits)

def get_last_results(resultsFileName):
    if not os.path.isfile(resultsFileName):
        return None, None, None
    commit = None
    cspeed = []
    dspeed = []
    with open(resultsFileName,'r') as f:
        for line in f:
            words = line.split()
            if len(words) == 2:   # branch + commit
                commit = words[1];
                cspeed = []
                dspeed = []
            if (len(words) == 8):  # results
                cspeed.append(float(words[3]))
                dspeed.append(float(words[5]))
    return commit, cspeed, dspeed

def benchmark_and_compare(branch, commit, resultsFileName, lastCLevel, testFilePath, fileName, last_cspeed, last_dspeed, lower_limit, maxLoadAvg, message):
    sleepTime = 30
    while os.getloadavg()[0] > maxLoadAvg:
        log("WARNING: bench loadavg=%.2f is higher than %s, sleeping for %s seconds" % (os.getloadavg()[0], maxLoadAvg, sleepTime))
        time.sleep(sleepTime)
    start_load = str(os.getloadavg())
    result = execute('programs/zstd -qi5b1e%s %s' % (lastCLevel, testFilePath), print_output=True)
    end_load = str(os.getloadavg())
    linesExpected = lastCLevel + 2;
    if len(result) != linesExpected:
        raise RuntimeError("ERROR: number of result lines=%d is different that expected %d\n%s" % (len(result), linesExpected, '\n'.join(result)))
    with open(resultsFileName, "a") as myfile:
        myfile.write(branch + " " + commit + "\n")
        myfile.write('\n'.join(result) + '\n')
        myfile.close()
        if (last_cspeed == None):
            log("WARNING: No data for comparison for branch=%s file=%s " % (branch, fileName))
            return ""
        commit, cspeed, dspeed = get_last_results(resultsFileName)
        text = ""
        for i in range(0, min(len(cspeed), len(last_cspeed))):
            print("%s:%s -%d cspeed=%6.2f clast=%6.2f cdiff=%1.4f dspeed=%6.2f dlast=%6.2f ddiff=%1.4f %s" % (branch, commit, i+1, cspeed[i], last_cspeed[i], cspeed[i]/last_cspeed[i], dspeed[i], last_dspeed[i], dspeed[i]/last_dspeed[i], fileName))
            if (cspeed[i]/last_cspeed[i] < lower_limit):
                text += "WARNING: -%d cspeed=%.2f clast=%.2f cdiff=%.4f %s\n" % (i+1, cspeed[i], last_cspeed[i], cspeed[i]/last_cspeed[i], fileName)
            if (dspeed[i]/last_dspeed[i] < lower_limit):
                text += "WARNING: -%d dspeed=%.2f dlast=%.2f ddiff=%.4f %s\n" % (i+1, dspeed[i], last_dspeed[i], dspeed[i]/last_dspeed[i], fileName)
        if text:
            text = message + ("\nmaxLoadAvg=%s  load average at start=%s end=%s\n" % (maxLoadAvg, start_load, end_load)) + text
        return text

def update_config_file(branch, commit):
    last_commit = None
    commitFileName = working_path + "/commit_" + branch.replace("/", "_") + ".txt"
    if os.path.isfile(commitFileName):
        last_commit = file(commitFileName, 'r').read()
    file(commitFileName, 'w').write(commit)
    return last_commit

def test_commit(branch, commit, last_commit, args, testFilePaths, have_mutt, have_mail):
    local_branch = string.split(branch, '/')[1]
    version = local_branch.rpartition('-')[2] + '_' + commit
    if not args.dry_run:
        execute('make clean zstdprogram MOREFLAGS="-DZSTD_GIT_COMMIT=%s"' % version)
    logFileName = working_path + "/log_" + branch.replace("/", "_") + ".txt"
    text_to_send = []
    results_files = ""
    for filePath in testFilePaths:
        fileName = filePath.rpartition('/')[2]
        resultsFileName = working_path + "/results_" + branch.replace("/", "_") + "_" + fileName.replace(".", "_") + ".txt"
        last_commit, cspeed, dspeed = get_last_results(resultsFileName)
        if not args.dry_run:
            text = benchmark_and_compare(branch, commit, resultsFileName, args.lastCLevel, filePath, fileName, cspeed, dspeed, args.lowerLimit, args.maxLoadAvg, args.message)
            if text:
                log("WARNING: redoing tests for branch %s: commit %s" % (branch, commit))
                text = benchmark_and_compare(branch, commit, resultsFileName, args.lastCLevel, filePath, fileName, cspeed, dspeed, args.lowerLimit, args.maxLoadAvg, args.message)
                if text:
                    text_to_send.append(text)
                    results_files += resultsFileName + " "
    if text_to_send:
        send_email_with_attachments(branch, commit, last_commit, args.emails, text_to_send, results_files, logFileName, args.lowerLimit, have_mutt, have_mail)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('testFileNames', help='file names list for speed benchmark')
    parser.add_argument('emails', help='list of e-mail addresses to send warnings')
    parser.add_argument('--message', help='attach an additional message to e-mail', default="")
    parser.add_argument('--repoURL', help='changes default repository URL', default=default_repo_url)
    parser.add_argument('--lowerLimit', type=float, help='send email if speed is lower than given limit', default=0.98)
    parser.add_argument('--maxLoadAvg', type=float, help='maximum load average to start testing', default=0.75)
    parser.add_argument('--lastCLevel', type=int, help='last compression level for testing', default=5)
    parser.add_argument('--sleepTime', type=int, help='frequency of repository checking in seconds', default=300)
    parser.add_argument('--dry-run', dest='dry_run', action='store_true', help='not build', default=False)
    args = parser.parse_args()

    # check if test files are accessible
    testFileNames = args.testFileNames.split()
    testFilePaths = []
    for fileName in testFileNames:
        fileName = os.path.expanduser(fileName)
        if os.path.isfile(fileName):
            testFilePaths.append(os.path.abspath(fileName))
        else:
            log("ERROR: File not found: " + fileName)
            exit(1)

    # check availability of e-mail senders
    have_mutt = does_command_exist("mutt -h");
    have_mail = does_command_exist("mail -V");
    if not have_mutt and not have_mail:
        log("ERROR: e-mail senders 'mail' or 'mutt' not found")
        exit(1)

    print("PARAMETERS:\nrepoURL=%s" % args.repoURL)
    print("working_path=%s" % working_path)
    print("clone_path=%s" % clone_path)
    print("testFilePath(%s)=%s" % (len(testFilePaths), testFilePaths))
    print("message=%s" % args.message)
    print("emails=%s" % args.emails)
    print("maxLoadAvg=%s" % args.maxLoadAvg)
    print("lowerLimit=%s" % args.lowerLimit)
    print("lastCLevel=%s" % args.lastCLevel)
    print("sleepTime=%s" % args.sleepTime)
    print("dry_run=%s" % args.dry_run)
    print("have_mutt=%s have_mail=%s" % (have_mutt, have_mail))

    # clone ZSTD repo if needed
    if not os.path.isdir(working_path):
        os.mkdir(working_path)
    if not os.path.isdir(clone_path):
        execute.cwd = working_path
        execute('git clone ' + args.repoURL)
    if not os.path.isdir(clone_path):
        log("ERROR: ZSTD clone not found: " + clone_path)
        exit(1)
    execute.cwd = clone_path

    # check if speedTest.pid already exists
    pidfile = "./speedTest.pid"
    if os.path.isfile(pidfile):
        log("ERROR: %s already exists, exiting" % pidfile)
        exit(1)

    send_email(args.emails, email_header + ':%s test-zstd-speed.py has been started' % pid, args.message, have_mutt, have_mail)
    file(pidfile, 'w').write(pid)

    while True:
        try:
            loadavg = os.getloadavg()[0]
            if (loadavg <= args.maxLoadAvg):
                branches = git_get_branches()
                for branch in branches:
                    commit = execute('git show -s --format=%h ' + branch)[0]
                    last_commit = update_config_file(branch, commit)
                    if commit == last_commit:
                        log("skipping branch %s: head %s already processed" % (branch, commit))
                    else:
                        log("build branch %s: head %s is different from prev %s" % (branch, commit, last_commit))
                        execute('git checkout -- . && git checkout ' + branch)
                        print(git_get_changes(branch, commit, last_commit))
                        test_commit(branch, commit, last_commit, args, testFilePaths, have_mutt, have_mail)
            else:
                log("WARNING: main loadavg=%.2f is higher than %s" % (loadavg, args.maxLoadAvg))
            log("sleep for %s seconds" % args.sleepTime)
            time.sleep(args.sleepTime)
        except Exception as e:
            stack = traceback.format_exc()
            email_topic = '%s:%s ERROR in %s:%s' % (email_header, pid, branch, commit)
            send_email(args.emails, email_topic, stack, have_mutt, have_mail)
            print(stack)
        except KeyboardInterrupt:
            os.unlink(pidfile)
            send_email(args.emails, email_header + ':%s test-zstd-speed.py has been stopped' % pid, args.message, have_mutt, have_mail)
            exit(0)
<|endoftext|>