<commit_before># ## Copyright (c) 2013, GPy authors (see AUTHORS.txt).
# Licensed under the BSD 3-clause license (see LICENSE.txt)

import numpy as np
import itertools
import pylab

from ..core import Model, SparseGP
from ..util.linalg import PCA
from ..kern import Kern
from bayesian_gplvm import BayesianGPLVM
from ..core.parameterization.variational import NormalPosterior, NormalPrior
from ..inference.latent_function_inference.var_dtc import VarDTCMissingData
from ..likelihoods.gaussian import Gaussian

class MRD2(Model):
    """
    Apply MRD to all given datasets Y in Ylist.

    Y_i in [n x p_i]

    The samples n in the datasets need
    to match up, whereas the dimensionality p_d can differ.

    :param [array-like] Ylist: List of datasets to apply MRD on
    :param input_dim: latent dimensionality
    :type input_dim: int
    :param array-like X: mean of starting latent space q in [n x q]
    :param array-like X_variance: variance of starting latent space q in [n x q]
    :param initx: initialisation method for the latent space :

        * 'concat' - PCA on concatenation of all datasets
        * 'single' - Concatenation of PCA on datasets, respectively
        * 'random' - Random draw from a Normal(0,1)

    :type initx: ['concat'|'single'|'random']
    :param initz: initialisation method for inducing inputs
    :type initz: 'permute'|'random'
    :param num_inducing: number of inducing inputs to use
    :param Z: initial inducing inputs
    :param kernel: list of kernels or kernel to copy for each output
    :type kernel: [GPy.kern.kern] | GPy.kern.kern | None (default)
    :param :class:`~GPy.inference.latent_function_inference inference_method: the inference method to use
    :param :class:`~GPy.likelihoods.likelihood.Likelihood` likelihood: the likelihood to use
    :param str name: the name of this model
    """

    def __init__(self, Ylist, input_dim, X=None, X_variance=None,
                 initx = 'PCA', initz = 'permute',
                 num_inducing=10, Z=None, kernel=None,
                 inference_method=None, likelihood=None, name='mrd'):
        super(MRD2, self).__init__(name)

        # sort out the kernels
        if kernel is None:
            from ..kern import RBF
            self.kern = [RBF(input_dim, ARD=1, name='Y_{}'.format(i)) for i in range(len(Ylist))]
        elif isinstance(kernel, Kern):
            self.kern = [kernel.copy(name='Y_{}'.format(i)) for i in range(len(Ylist))]
        else:
            assert len(kernel) == len(Ylist), "need one kernel per output"
            assert all([isinstance(k, Kern) for k in kernel]), "invalid kernel object detected!"

        self.input_dim = input_dim
        self.num_inducing = num_inducing

        self._in_init_ = True
        X = self._init_X(initx, Ylist)
        self.Z = self._init_Z(initz, X)
        self.num_inducing = self.Z.shape[0] # ensure M==N if M>N

        if X_variance is None:
            X_variance = np.random.uniform(0,.2,X.shape)

        self.variational_prior = NormalPrior()
        self.X = NormalPosterior(X, X_variance)

        if likelihood is None:
            likelihood = Gaussian()

        if inference_method is None:
            if any(np.any(np.isnan(y)) for y in Ylist):
                self.inference_method = VarDTCMissingData(limit=len(Ylist))

        self.Ylist = Ylist

    def parameters_changed(self):
        for y in self.Ylist:
            pass

    def _init_X(self, init='PCA', likelihood_list=None):
        if likelihood_list is None:
            likelihood_list = self.likelihood_list
        Ylist = []
        for likelihood_or_Y in likelihood_list:
            if type(likelihood_or_Y) is np.ndarray:
                Ylist.append(likelihood_or_Y)
            else:
                Ylist.append(likelihood_or_Y.Y)
        del likelihood_list
        if init in "PCA_concat":
            X = PCA(np.hstack(Ylist), self.input_dim)[0]
        elif init in "PCA_single":
            X = np.zeros((Ylist[0].shape[0], self.input_dim))
            for qs, Y in itertools.izip(np.array_split(np.arange(self.input_dim), len(Ylist)), Ylist):
                X[:, qs] = PCA(Y, len(qs))[0]
        else: # init == 'random':
            X = np.random.randn(Ylist[0].shape[0], self.input_dim)
        self.X = X
        return X

    def _init_Z(self, init="permute", X=None):
        if X is None:
            X = self.X
        if init in "permute":
            Z = np.random.permutation(X.copy())[:self.num_inducing]
        elif init in "random":
            Z = np.random.randn(self.num_inducing, self.input_dim) * X.var()
        self.Z = Z
        return Z

class MRD(Model):
    """
    Do MRD on given Datasets in Ylist.
    All Ys in likelihood_list are in [N x Dn], where Dn can be different per Yn,
    N must be shared across datasets though.

    :param likelihood_list: list of observed datasets (:py:class:`~GPy.likelihoods.gaussian.Gaussian` if not supplied directly)
    :type likelihood_list: [:py:class:`~GPy.likelihoods.likelihood.likelihood` | :py:class:`ndarray`]
    :param names: names for different gplvm models
    :type names: [str]
    :param input_dim: latent dimensionality
    :type input_dim: int
    :param initx: initialisation method for the latent space :

        * 'concat' - PCA on concatenation of all datasets
        * 'single' - Concatenation of PCA on datasets, respectively
        * 'random' - Random draw from a normal

    :type initx: ['concat'|'single'|'random']
    :param initz: initialisation method for inducing inputs
    :type initz: 'permute'|'random'
    :param X: Initial latent space
    :param X_variance: Initial latent space variance
    :param Z: initial inducing inputs
    :param num_inducing: number of inducing inputs to use
    :param kernels: list of kernels or kernel shared for all BGPLVMS
    :type kernels: [GPy.kern.kern] | GPy.kern.kern | None (default)

    """
    def __init__(self, likelihood_or_Y_list, input_dim, num_inducing=10, names=None,
                 kernels=None, initx='PCA',
                 initz='permute', _debug=False, **kw):
        if names is None:
            self.names = ["{}".format(i) for i in range(len(likelihood_or_Y_list))]

        # sort out the kernels
        if kernels is None:
            kernels = [None] * len(likelihood_or_Y_list)
        elif isinstance(kernels, Kern):
            kernels = [kernels.copy() for i in range(len(likelihood_or_Y_list))]
        else:
            assert len(kernels) == len(likelihood_or_Y_list), "need one kernel per output"
            assert all([isinstance(k, Kern) for k in kernels]), "invalid kernel object detected!"
        assert not ('kernel' in kw), "pass kernels through `kernels` argument"

        self.input_dim = input_dim
        self._debug = _debug
        self.num_inducing = num_inducing

        self._in_init_ = True
        X = self._init_X(initx, likelihood_or_Y_list)
        Z = self._init_Z(initz, X)
        self.num_inducing = Z.shape[0] # ensure M==N if M>N

        self.bgplvms = [BayesianGPLVM(l, input_dim=input_dim, kernel=k, X=X, Z=Z, num_inducing=self.num_inducing, **kw) for l, k in zip(likelihood_or_Y_list, kernels)]
        del self._in_init_

        self.gref = self.bgplvms[0]
        nparams = np.array([0] + [SparseGP._get_params(g).size - g.Z.size for g in self.bgplvms])
        self.nparams = nparams.cumsum()

        self.num_data = self.gref.num_data

        self.NQ = self.num_data * self.input_dim
        self.MQ = self.num_inducing * self.input_dim

        Model.__init__(self)
        self.ensure_default_constraints()

    def _getstate(self):
        return Model._getstate(self) + [self.names,
                self.bgplvms,
                self.gref,
                self.nparams,
                self.input_dim,
                self.num_inducing,
                self.num_data,
                self.NQ,
                self.MQ]

    def _setstate(self, state):
        self.MQ = state.pop()
        self.NQ = state.pop()
        self.num_data = state.pop()
        self.num_inducing = state.pop()
        self.input_dim = state.pop()
        self.nparams = state.pop()
        self.gref = state.pop()
        self.bgplvms = state.pop()
        self.names = state.pop()
        Model._setstate(self, state)

    @property
    def X(self):
        return self.gref.X
    @X.setter
    def X(self, X):
        try:
            self.propagate_param(X=X)
        except AttributeError:
            if not self._in_init_:
                raise AttributeError("bgplvm list not initialized")
    @property
    def Z(self):
        return self.gref.Z
    @Z.setter
    def Z(self, Z):
        try:
            self.propagate_param(Z=Z)
        except AttributeError:
            if not self._in_init_:
                raise AttributeError("bgplvm list not initialized")
    @property
    def X_variance(self):
        return self.gref.X_variance
    @X_variance.setter
    def X_variance(self, X_var):
        try:
            self.propagate_param(X_variance=X_var)
        except AttributeError:
            if not self._in_init_:
                raise AttributeError("bgplvm list not initialized")
    @property
    def likelihood_list(self):
        return [g.likelihood.Y for g in self.bgplvms]
    @likelihood_list.setter
    def likelihood_list(self, likelihood_list):
        for g, Y in itertools.izip(self.bgplvms, likelihood_list):
            g.likelihood.Y = Y

    @property
    def auto_scale_factor(self):
        """
        set auto_scale_factor for all gplvms
        :param b: auto_scale_factor
        :type b:
        """
        return self.gref.auto_scale_factor
    @auto_scale_factor.setter
    def auto_scale_factor(self, b):
        self.propagate_param(auto_scale_factor=b)

    def propagate_param(self, **kwargs):
        for key, val in kwargs.iteritems():
            for g in self.bgplvms:
                g.__setattr__(key, val)

    def randomize(self, initx='concat', initz='permute', *args, **kw):
        super(MRD, self).randomize(*args, **kw)
        self._init_X(initx, self.likelihood_list)
        self._init_Z(initz, self.X)

    #def _get_latent_param_names(self):
    def _get_param_names(self):
        n1 = self.gref._get_param_names()
        n1var = n1[:self.NQ * 2 + self.MQ]
    #    return n1var
    #
    #def _get_kernel_names(self):
        map_names = lambda ns, name: map(lambda x: "{1}_{0}".format(*x),
                                         itertools.izip(ns,
                                                        itertools.repeat(name)))
        return list(itertools.chain(n1var, *(map_names(\
                SparseGP._get_param_names(g)[self.MQ:], n) \
                for g, n in zip(self.bgplvms, self.names))))
    #    kernel_names = (map_names(SparseGP._get_param_names(g)[self.MQ:], n) for g, n in zip(self.bgplvms, self.names))
    #    return kernel_names

    #def _get_param_names(self):
        # X_names = sum([['X_%i_%i' % (n, q) for q in range(self.input_dim)] for n in range(self.num_data)], [])
        # S_names = sum([['X_variance_%i_%i' % (n, q) for q in range(self.input_dim)] for n in range(self.num_data)], [])
    #    n1var = self._get_latent_param_names()
    #    kernel_names = self._get_kernel_names()
    #    return list(itertools.chain(n1var, *kernel_names))

    #def _get_print_names(self):
    #    return list(itertools.chain(*self._get_kernel_names()))

    def _get_params(self):
        """
        return parameter list containing private and shared parameters as follows:

        =================================================================
        | mu | S | Z || theta1 | theta2 | .. | thetaN |
        =================================================================
        """
        X = self.gref.X.ravel()
        X_var = self.gref.X_variance.ravel()
        Z = self.gref.Z.ravel()
        thetas = [SparseGP._get_params(g)[g.Z.size:] for g in self.bgplvms]
        params = np.hstack([X, X_var, Z, np.hstack(thetas)])
        return params

#     def _set_var_params(self, g, X, X_var, Z):
#         g.X = X.reshape(self.num_data, self.input_dim)
#         g.X_variance = X_var.reshape(self.num_data, self.input_dim)
#         g.Z = Z.reshape(self.num_inducing, self.input_dim)
#
#     def _set_kern_params(self, g, p):
#         g.kern._set_params(p[:g.kern.num_params])
#         g.likelihood._set_params(p[g.kern.num_params:])

    def _set_params(self, x):
        start = 0; end = self.NQ
        X = x[start:end]
        start = end; end += start
        X_var = x[start:end]
        start = end; end += self.MQ
        Z = x[start:end]
        thetas = x[end:]

        # set params for all:
        for g, s, e in itertools.izip(self.bgplvms, self.nparams, self.nparams[1:]):
            g._set_params(np.hstack([X, X_var, Z, thetas[s:e]]))
#             self._set_var_params(g, X, X_var, Z)
#             self._set_kern_params(g, thetas[s:e].copy())
#             g._compute_kernel_matrices()
#             if self.auto_scale_factor:
#                 g.scale_factor = np.sqrt(g.psi2.sum(0).mean() * g.likelihood.precision)
# #                 self.scale_factor = np.sqrt(self.psi2.sum(0).mean() * self.likelihood.precision)
#             g._computations()

    def update_likelihood_approximation(self): # TODO: object oriented vs script base
        for bgplvm in self.bgplvms:
            bgplvm.update_likelihood_approximation()

    def log_likelihood(self):
        ll = -self.gref.KL_divergence()
        for g in self.bgplvms:
            ll += SparseGP.log_likelihood(g)
        return ll

    def _log_likelihood_gradients(self):
        dLdmu, dLdS = reduce(lambda a, b: [a[0] + b[0], a[1] + b[1]], (g.dL_dmuS() for g in self.bgplvms))
        dKLmu, dKLdS = self.gref.dKL_dmuS()
        dLdmu -= dKLmu
        dLdS -= dKLdS
        dLdmuS = np.hstack((dLdmu.flatten(), dLdS.flatten())).flatten()
        dldzt1 = reduce(lambda a, b: a + b, (SparseGP._log_likelihood_gradients(g)[:self.MQ] for g in self.bgplvms))

        return np.hstack((dLdmuS,
                             dldzt1,
                np.hstack([np.hstack([g.dL_dtheta(),
                                            g.likelihood._gradients(\
                                                partial=g.partial_for_likelihood)]) \
                              for g in self.bgplvms])))

    def _handle_plotting(self, fignum, axes, plotf, sharex=False, sharey=False):
        if axes is None:
            fig = pylab.figure(num=fignum)
        sharex_ax = None
        sharey_ax = None
        for i, g in enumerate(self.bgplvms):
            try:
                if sharex:
                    sharex_ax = ax # @UndefinedVariable
                    sharex = False # dont set twice
                if sharey:
                    sharey_ax = ax # @UndefinedVariable
                    sharey = False # dont set twice
            except:
                pass
            if axes is None:
                ax = fig.add_subplot(1, len(self.bgplvms), i + 1, sharex=sharex_ax, sharey=sharey_ax)
            elif isinstance(axes, (tuple, list)):
                ax = axes[i]
            else:
                raise ValueError("Need one axes per latent dimension input_dim")
            plotf(i, g, ax)
            if sharey_ax is not None:
                pylab.setp(ax.get_yticklabels(), visible=False)
        pylab.draw()
        if axes is None:
            fig.tight_layout()
            return fig
        else:
            return pylab.gcf()

    def plot_X(self, fignum=None, ax=None):
        fig = self._handle_plotting(fignum, ax, lambda i, g, ax: ax.imshow(g.X))
        return fig

    def plot_predict(self, fignum=None, ax=None, sharex=False, sharey=False, **kwargs):
        fig = self._handle_plotting(fignum,
                                    ax,
                                    lambda i, g, ax: ax.imshow(g. predict(g.X)[0], **kwargs),
                                    sharex=sharex, sharey=sharey)
        return fig

    def plot_scales(self, fignum=None, ax=None, titles=None, sharex=False, sharey=True, *args, **kwargs):
        """

        TODO: Explain other parameters

        :param titles: titles for axes of datasets

        """
        if titles is None:
            titles = [r'${}$'.format(name) for name in self.names]
        ymax = reduce(max, [np.ceil(max(g.input_sensitivity())) for g in self.bgplvms])
        def plotf(i, g, ax):
            ax.set_ylim([0,ymax])
            g.kern.plot_ARD(ax=ax, title=titles[i], *args, **kwargs)
        fig = self._handle_plotting(fignum, ax, plotf, sharex=sharex, sharey=sharey)
        return fig

    def plot_latent(self, fignum=None, ax=None, *args, **kwargs):
        fig = self.gref.plot_latent(fignum=fignum, ax=ax, *args, **kwargs) # self._handle_plotting(fignum, ax, lambda i, g, ax: g.plot_latent(ax=ax, *args, **kwargs))
        return fig

    def _debug_plot(self):
        self.plot_X_1d()
        fig = pylab.figure("MRD DEBUG PLOT", figsize=(4 * len(self.bgplvms), 9))
        fig.clf()
        axes = [fig.add_subplot(3, len(self.bgplvms), i + 1) for i in range(len(self.bgplvms))]
        self.plot_X(ax=axes)
        axes = [fig.add_subplot(3, len(self.bgplvms), i + len(self.bgplvms) + 1) for i in range(len(self.bgplvms))]
        self.plot_latent(ax=axes)
        axes = [fig.add_subplot(3, len(self.bgplvms), i + 2 * len(self.bgplvms) + 1) for i in range(len(self.bgplvms))]
        self.plot_scales(ax=axes)
        pylab.draw()
        fig.tight_layout()

<commit_msg>mrd for new parameterize<commit_after># ## Copyright (c) 2013, GPy authors (see AUTHORS.txt).
# Licensed under the BSD 3-clause license (see LICENSE.txt)

import numpy as np
import itertools
import pylab

from ..core import Model
from ..util.linalg import PCA
from ..kern import Kern
from ..core.parameterization.variational import NormalPosterior, NormalPrior
from ..core.parameterization import Param, Parameterized
from ..inference.latent_function_inference.var_dtc import VarDTCMissingData, VarDTC
from ..likelihoods import Gaussian

class MRD(Model):
    """
    Apply MRD to all given datasets Y in Ylist.

    Y_i in [n x p_i]

    The samples n in the datasets need
    to match up, whereas the dimensionality p_d can differ.

    :param [array-like] Ylist: List of datasets to apply MRD on
    :param input_dim: latent dimensionality
    :type input_dim: int
    :param array-like X: mean of starting latent space q in [n x q]
    :param array-like X_variance: variance of starting latent space q in [n x q]
    :param initx: initialisation method for the latent space :

        * 'concat' - PCA on concatenation of all datasets
        * 'single' - Concatenation of PCA on datasets, respectively
        * 'random' - Random draw from a Normal(0,1)

    :type initx: ['concat'|'single'|'random']
    :param initz: initialisation method for inducing inputs
    :type initz: 'permute'|'random'
    :param num_inducing: number of inducing inputs to use
    :param Z: initial inducing inputs
    :param kernel: list of kernels or kernel to copy for each output
    :type kernel: [GPy.kern.kern] | GPy.kern.kern | None (default)
    :param :class:`~GPy.inference.latent_function_inference inference_method: the inference method to use
    :param :class:`~GPy.likelihoods.likelihood.Likelihood` likelihood: the likelihood to use
    :param str name: the name of this model
    :param [str] Ynames: the names for the datasets given, must be of equal length as Ylist or None
    """

    def __init__(self, Ylist, input_dim, X=None, X_variance=None,
                 initx = 'PCA', initz = 'permute',
                 num_inducing=10, Z=None, kernel=None,
                 inference_method=None, likelihood=None, name='mrd', Ynames=None):
        super(MRD, self).__init__(name)

        # sort out the kernels
        if kernel is None:
            from ..kern import RBF
            self.kern = [RBF(input_dim, ARD=1, name='rbf'.format(i)) for i in range(len(Ylist))]
        elif isinstance(kernel, Kern):
            self.kern = [kernel.copy(name='{}'.format(kernel.name, i)) for i in range(len(Ylist))]
        else:
            assert len(kernel) == len(Ylist), "need one kernel per output"
            assert all([isinstance(k, Kern) for k in kernel]), "invalid kernel object detected!"
            self.kern = kernel
        self.input_dim = input_dim
        self.num_inducing = num_inducing

        self.Ylist = Ylist
        self._in_init_ = True
        X = self._init_X(initx, Ylist)
        self.Z = Param('inducing inputs', self._init_Z(initz, X))
        self.num_inducing = self.Z.shape[0] # ensure M==N if M>N

        if X_variance is None:
            X_variance = np.random.uniform(0, .2, X.shape)

        self.variational_prior = NormalPrior()
        self.X = NormalPosterior(X, X_variance)

        if likelihood is None:
            self.likelihood = [Gaussian(name='Gaussian_noise'.format(i)) for i in range(len(Ylist))]
        else: self.likelihood = likelihood

        if inference_method is None:
            self.inference_method= []
            for y in Ylist:
                if np.any(np.isnan(y)):
                    self.inference_method.append(VarDTCMissingData(limit=1))
                else:
                    self.inference_method.append(VarDTC(limit=1))
        else:
            self.inference_method = inference_method
            self.inference_method.set_limit(len(Ylist))

        self.add_parameters(self.X, self.Z)

        if Ynames is None:
            Ynames = ['Y{}'.format(i) for i in range(len(Ylist))]

        for i, n, k, l in itertools.izip(itertools.count(), Ynames, self.kern, self.likelihood):
            p = Parameterized(name=n)
            p.add_parameter(k)
            p.add_parameter(l)
            setattr(self, 'Y{}'.format(i), p)
            self.add_parameter(p)
        self._in_init_ = False

    def parameters_changed(self):
        self._log_marginal_likelihood = 0
        self.posteriors = []
        self.Z.gradient = 0.
        self.X.mean.gradient = 0.
        self.X.variance.gradient = 0.

        for y, k, l, i in itertools.izip(self.Ylist, self.kern, self.likelihood, self.inference_method):
            posterior, lml, grad_dict = i.inference(k, self.X, self.Z, l, y)

            self.posteriors.append(posterior)
            self._log_marginal_likelihood += lml

            # likelihood gradients
            l.update_gradients(grad_dict.pop('partial_for_likelihood'))

            #gradients wrt kernel
            dL_dKmm = grad_dict.pop('dL_dKmm')
            k.update_gradients_full(dL_dKmm, self.Z, None)
            target = k.gradient.copy()
            k.update_gradients_expectations(variational_posterior=self.X, Z=self.Z, **grad_dict)
            k.gradient += target

            #gradients wrt Z
            self.Z.gradient += k.gradients_X(dL_dKmm, self.Z)
            self.Z.gradient += k.gradients_Z_expectations(
                               grad_dict['dL_dpsi1'], grad_dict['dL_dpsi2'], Z=self.Z, variational_posterior=self.X)

            dL_dmean, dL_dS = k.gradients_qX_expectations(variational_posterior=self.X, Z=self.Z, **grad_dict)
            self.X.mean.gradient += dL_dmean
            self.X.variance.gradient += dL_dS

        # update for the KL divergence
        self.variational_prior.update_gradients_KL(self.X)
        self._log_marginal_likelihood -= self.variational_prior.KL_divergence(self.X)

    def log_likelihood(self):
        return self._log_marginal_likelihood

    def _init_X(self, init='PCA', Ylist=None):
        if Ylist is None:
            Ylist = self.Ylist
        if init in "PCA_concat":
            X = PCA(np.hstack(Ylist), self.input_dim)[0]
        elif init in "PCA_single":
            X = np.zeros((Ylist[0].shape[0], self.input_dim))
            for qs, Y in itertools.izip(np.array_split(np.arange(self.input_dim), len(Ylist)), Ylist):
                X[:, qs] = PCA(Y, len(qs))[0]
        else: # init == 'random':
            X = np.random.randn(Ylist[0].shape[0], self.input_dim)
        return X

    def _init_Z(self, init="permute", X=None):
        if X is None:
            X = self.X
        if init in "permute":
            Z = np.random.permutation(X.copy())[:self.num_inducing]
        elif init in "random":
            Z = np.random.randn(self.num_inducing, self.input_dim) * X.var()
        return Z

    def _handle_plotting(self, fignum, axes, plotf, sharex=False, sharey=False):
        if axes is None:
            fig = pylab.figure(num=fignum)
        sharex_ax = None
        sharey_ax = None
        for i, g in enumerate(self.bgplvms):
            try:
                if sharex:
                    sharex_ax = ax # @UndefinedVariable
                    sharex = False # dont set twice
                if sharey:
                    sharey_ax = ax # @UndefinedVariable
                    sharey = False # dont set twice
            except:
                pass
            if axes is None:
                ax = fig.add_subplot(1, len(self.bgplvms), i + 1, sharex=sharex_ax, sharey=sharey_ax)
            elif isinstance(axes, (tuple, list)):
                ax = axes[i]
            else:
                raise ValueError("Need one axes per latent dimension input_dim")
            plotf(i, g, ax)
            if sharey_ax is not None:
                pylab.setp(ax.get_yticklabels(), visible=False)
        pylab.draw()
        if axes is None:
            fig.tight_layout()
            return fig
        else:
            return pylab.gcf()

    def plot_X(self, fignum=None, ax=None):
        fig = self._handle_plotting(fignum, ax, lambda i, g, ax: ax.imshow(g.X))
        return fig

    def plot_predict(self, fignum=None, ax=None, sharex=False, sharey=False, **kwargs):
        fig = self._handle_plotting(fignum,
                                    ax,
                                    lambda i, g, ax: ax.imshow(g. predict(g.X)[0], **kwargs),
                                    sharex=sharex, sharey=sharey)
        return fig

    def plot_scales(self, fignum=None, ax=None, titles=None, sharex=False, sharey=True, *args, **kwargs):
        """

        TODO: Explain other parameters

        :param titles: titles for axes of datasets

        """
        if titles is None:
            titles = [r'${}$'.format(name) for name in self.names]
        ymax = reduce(max, [np.ceil(max(g.input_sensitivity())) for g in self.bgplvms])
        def plotf(i, g, ax):
            ax.set_ylim([0,ymax])
            g.kern.plot_ARD(ax=ax, title=titles[i], *args, **kwargs)
        fig = self._handle_plotting(fignum, ax, plotf, sharex=sharex, sharey=sharey)
        return fig

    def plot_latent(self, fignum=None, ax=None, *args, **kwargs):
        fig = self.gref.plot_latent(fignum=fignum, ax=ax, *args, **kwargs) # self._handle_plotting(fignum, ax, lambda i, g, ax: g.plot_latent(ax=ax, *args, **kwargs))
        return fig

    def _debug_plot(self):
        self.plot_X_1d()
        fig = pylab.figure("MRD DEBUG PLOT", figsize=(4 * len(self.bgplvms), 9))
        fig.clf()
        axes = [fig.add_subplot(3, len(self.bgplvms), i + 1) for i in range(len(self.bgplvms))]
        self.plot_X(ax=axes)
        axes = [fig.add_subplot(3, len(self.bgplvms), i + len(self.bgplvms) + 1) for i in range(len(self.bgplvms))]
        self.plot_latent(ax=axes)
        axes = [fig.add_subplot(3, len(self.bgplvms), i + 2 * len(self.bgplvms) + 1) for i in range(len(self.bgplvms))]
        self.plot_scales(ax=axes)
        pylab.draw()
        fig.tight_layout()

<|endoftext|>


<commit_before>from eventkit import Event

from .objects import Object, SoftDollarTier
from .util import UNSET_DOUBLE, UNSET_INTEGER

__all__ = (
    'Trade OrderStatus Order '
    'LimitOrder MarketOrder StopOrder StopLimitOrder '
    'OrderCondition ExecutionCondition MarginCondition TimeCondition '
    'PriceCondition PercentChangeCondition VolumeCondition '
    ).split()

class Trade(Object):
    """
    Trade keeps track of an order, its status and all its fills.

    Events:
        * ``statusEvent`` (trade: :class:`.Trade`)
        * ``modifyEvent`` (trade: :class:`.Trade`)
        * ``fillEvent`` (trade: :class:`.Trade`, fill: :class:`.Fill`)
        * ``commissionReportEvent`` (trade: :class:`.Trade`,
          fill: :class:`.Fill`, commissionReport: :class:`.CommissionReport`)
        * ``filledEvent`` (trade: :class:`.Trade`)
        * ``cancelEvent`` (trade: :class:`.Trade`)
        * ``cancelledEvent`` (trade: :class:`.Trade`)
    """
    events = (
        'statusEvent', 'modifyEvent', 'fillEvent',
        'commissionReportEvent', 'filledEvent',
        'cancelEvent', 'cancelledEvent')

    defaults = dict(
        contract=None,
        order=None,
        orderStatus=None,
        fills=None,
        log=None
    )
    __slots__ = tuple(defaults.keys()) + events + ('__dict__',)

    def __init__(self, *args, **kwargs):
        Object.__init__(self, *args, **kwargs)
        self.statusEvent = Event('statusEvent')
        self.modifyEvent = Event('modifyEvent')
        self.fillEvent = Event('fillEvent')
        self.commissionReportEvent = Event('commissionReportEvent')
        self.filledEvent = Event('filledEvent')
        self.cancelEvent = Event('cancelEvent')
        self.cancelledEvent = Event('cancelledEvent')

    def isActive(self):
        """
        True if eligible for execution, false otherwise.
        """
        return self.orderStatus.status in OrderStatus.ActiveStates

    def isDone(self):
        """
        True if completely filled or cancelled, false otherwise.
        """
        return self.orderStatus.status in OrderStatus.DoneStates

    def filled(self):
        """
        Number of shares filled.
        """
        return sum(f.execution.shares for f in self.fills)

    def remaining(self):
        """
        Number of shares remaining to be filled.
        """
        return self.order.totalQuantity - self.filled()

class OrderStatus(Object):
    defaults = dict(
        orderId=0,
        status='',
        filled=0,
        remaining=0,
        avgFillPrice=0.0,
        permId=0,
        parentId=0,
        lastFillPrice=0.0,
        clientId=0,
        whyHeld='',
        mktCapPrice=0.0,
        lastLiquidity=0
    )
    __slots__ = defaults.keys()

    PendingSubmit = 'PendingSubmit'
    PendingCancel = 'PendingCancel'
    PreSubmitted = 'PreSubmitted'
    Submitted = 'Submitted'
    ApiPending = 'ApiPending'
    ApiCancelled = 'ApiCancelled'
    Cancelled = 'Cancelled'
    Filled = 'Filled'
    Inactive = 'Inactive'

    DoneStates = {'Filled', 'Cancelled', 'ApiCancelled'}
    ActiveStates = {'PendingSubmit', 'ApiPending', 'PreSubmitted', 'Submitted'}

class Order(Object):
    """
    Order for trading contracts.

    https://interactivebrokers.github.io/tws-api/available_orders.html
    """
    defaults = dict(
        orderId=0,
        clientId=0,
        permId=0,
        action='',
        totalQuantity=0,
        orderType='',
        lmtPrice=UNSET_DOUBLE,
        auxPrice=UNSET_DOUBLE,
        tif='',
        activeStartTime='',
        activeStopTime='',
        ocaGroup='',
        ocaType=0,
        orderRef='',
        transmit=True,
        parentId=0,
        blockOrder=False,
        sweepToFill=False,
        displaySize=0,
        triggerMethod=0,
        outsideRth=False,
        hidden=False,
        goodAfterTime='',
        goodTillDate='',
        rule80A='',
        allOrNone=False,
        minQty=UNSET_INTEGER,
        percentOffset=UNSET_DOUBLE,
        overridePercentageConstraints=False,
        trailStopPrice=UNSET_DOUBLE,
        trailingPercent=UNSET_DOUBLE,
        faGroup='',
        faProfile='',
        faMethod='',
        faPercentage='',
        designatedLocation='',
        openClose="O",
        origin=0,
        shortSaleSlot=0,
        exemptCode=-1,
        discretionaryAmt=0,
        eTradeOnly=True,
        firmQuoteOnly=True,
        nbboPriceCap=UNSET_DOUBLE,
        optOutSmartRouting=False,
        auctionStrategy=0,
        startingPrice=UNSET_DOUBLE,
        stockRefPrice=UNSET_DOUBLE,
        delta=UNSET_DOUBLE,
        stockRangeLower=UNSET_DOUBLE,
        stockRangeUpper=UNSET_DOUBLE,
        randomizePrice=False,
        randomizeSize=False,
        volatility=UNSET_DOUBLE,
        volatilityType=UNSET_INTEGER,
        deltaNeutralOrderType='',
        deltaNeutralAuxPrice=UNSET_DOUBLE,
        deltaNeutralConId=0,
        deltaNeutralSettlingFirm='',
        deltaNeutralClearingAccount='',
        deltaNeutralClearingIntent='',
        deltaNeutralOpenClose='',
        deltaNeutralShortSale=False,
        deltaNeutralShortSaleSlot=0,
        deltaNeutralDesignatedLocation='',
        continuousUpdate=False,
        referencePriceType=UNSET_INTEGER,
        basisPoints=UNSET_DOUBLE,
        basisPointsType=UNSET_INTEGER,
        scaleInitLevelSize=UNSET_INTEGER,
        scaleSubsLevelSize=UNSET_INTEGER,
        scalePriceIncrement=UNSET_DOUBLE,
        scalePriceAdjustValue=UNSET_DOUBLE,
        scalePriceAdjustInterval=UNSET_INTEGER,
        scaleProfitOffset=UNSET_DOUBLE,
        scaleAutoReset=False,
        scaleInitPosition=UNSET_INTEGER,
        scaleInitFillQty=UNSET_INTEGER,
        scaleRandomPercent=False,
        scaleTable='',
        hedgeType='',
        hedgeParam='',
        account='',
        settlingFirm='',
        clearingAccount='',
        clearingIntent='',
        algoStrategy='',
        algoParams=None,
        smartComboRoutingParams=None,
        algoId='',
        whatIf=False,
        notHeld=False,
        solicited=False,
        modelCode='',
        orderComboLegs=None,
        orderMiscOptions=None,
        referenceContractId=0,
        peggedChangeAmount=0.0,
        isPeggedChangeAmountDecrease=False,
        referenceChangeAmount=0.0,
        referenceExchangeId='',
        adjustedOrderType='',
        triggerPrice=UNSET_DOUBLE,
        adjustedStopPrice=UNSET_DOUBLE,
        adjustedStopLimitPrice=UNSET_DOUBLE,
        adjustedTrailingAmount=UNSET_DOUBLE,
        adjustableTrailingUnit=0,
        lmtPriceOffset=UNSET_DOUBLE,
        conditions=None,
        conditionsCancelOrder=False,
        conditionsIgnoreRth=False,
        extOperator='',
        softDollarTier=None,
        cashQty=UNSET_DOUBLE,
        mifid2DecisionMaker='',
        mifid2DecisionAlgo='',
        mifid2ExecutionTrader='',
        mifid2ExecutionAlgo='',
        dontUseAutoPriceForHedge=False,
        isOmsContainer=False,
        discretionaryUpToLimitPrice=False,
        autoCancelDate='',
        filledQuantity=UNSET_DOUBLE,
        refFuturesConId=0,
        autoCancelParent=False,
        shareholder='',
        imbalanceOnly=False,
        routeMarketableToBbo=False,
        parentPermId=0,
        usePriceMgmtAlgo=False)
    __slots__ = defaults.keys()

    def __init__(self, *args, **kwargs):
        Object.__init__(self, *args, **kwargs)
        if not self.conditions:
            self.conditions = []
        if not self.softDollarTier:
            self.softDollarTier = SoftDollarTier()

    def __repr__(self):
        attrs = self.nonDefaults()
        if self.__class__ is not Order:
            attrs.pop('orderType', None)
        clsName = self.__class__.__name__
        kwargs = ', '.join(f'{k}={v!r}' for k, v in attrs.items())
        return f'{clsName}({kwargs})'

    __str__ = __repr__

    def __eq__(self, other):
        return self is other

    def __hash__(self):
        return id(self)

class LimitOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, lmtPrice, **kwargs):
        Order.__init__(
            self, orderType='LMT', action=action,
            totalQuantity=totalQuantity, lmtPrice=lmtPrice, **kwargs)

class MarketOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, **kwargs):
        Order.__init__(
            self, orderType='MKT', action=action,
            totalQuantity=totalQuantity, **kwargs)

class StopOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, stopPrice, **kwargs):
        Order.__init__(
            self, orderType='STP', action=action,
            totalQuantity=totalQuantity, auxPrice=stopPrice, **kwargs)

class StopLimitOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, lmtPrice, stopPrice, **kwargs):
        Order.__init__(
            self, orderType='STP LMT', action=action,
            totalQuantity=totalQuantity, lmtPrice=lmtPrice,
            auxPrice=stopPrice, **kwargs)

class OrderCondition(Object):
    __slots__ = ()

    @staticmethod
    def createClass(condType):
        d = {
            1: PriceCondition,
            3: TimeCondition,
            4: MarginCondition,
            5: ExecutionCondition,
            6: VolumeCondition,
            7: PercentChangeCondition}
        return d[condType]

    def __repr__(self):
        clsName = self.__class__.__qualname__
        kwargs = ', '.join(f'{k}={v!r}' for k, v in self.dict().items())
        return f'{clsName}({kwargs})'

    def And(self):
        self.conjunction = 'a'
        return self

    def Or(self):
        self.conjunction = 'o'
        return self

class PriceCondition(OrderCondition):
    defaults = dict(
        condType=1,
        conjunction='a',
        isMore=True,
        price=0.0,
        conId=0,
        exch='',
        triggerMethod=0)
    __slots__ = defaults.keys()

class TimeCondition(OrderCondition):
    defaults = dict(
        condType=3,
        conjunction='a',
        isMore=True,
        time='')
    __slots__ = defaults.keys()

class MarginCondition(OrderCondition):
    defaults = dict(
        condType=4,
        conjunction='a',
        isMore=True,
        percent=0)
    __slots__ = defaults.keys()

class ExecutionCondition(OrderCondition):
    defaults = dict(
        condType=5,
        conjunction='a',
        secType='',
        exch='',
        symbol='')
    __slots__ = defaults.keys()

class VolumeCondition(OrderCondition):
    defaults = dict(
        condType=6,
        conjunction='a',
        isMore=True,
        volume=0,
        conId=0,
        exch='')
    __slots__ = defaults.keys()

class PercentChangeCondition(OrderCondition):
    defaults = dict(
        condType=7,
        conjunction='a',
        isMore=True,
        changePercent=0.0,
        conId=0,
        exch='')
    __slots__ = defaults.keys()
<commit_msg>Fix trade.filled() for combo contracts.<commit_after>from eventkit import Event

from .objects import Object, SoftDollarTier
from .util import UNSET_DOUBLE, UNSET_INTEGER

__all__ = (
    'Trade OrderStatus Order '
    'LimitOrder MarketOrder StopOrder StopLimitOrder '
    'OrderCondition ExecutionCondition MarginCondition TimeCondition '
    'PriceCondition PercentChangeCondition VolumeCondition '
    ).split()

class Trade(Object):
    """
    Trade keeps track of an order, its status and all its fills.

    Events:
        * ``statusEvent`` (trade: :class:`.Trade`)
        * ``modifyEvent`` (trade: :class:`.Trade`)
        * ``fillEvent`` (trade: :class:`.Trade`, fill: :class:`.Fill`)
        * ``commissionReportEvent`` (trade: :class:`.Trade`,
          fill: :class:`.Fill`, commissionReport: :class:`.CommissionReport`)
        * ``filledEvent`` (trade: :class:`.Trade`)
        * ``cancelEvent`` (trade: :class:`.Trade`)
        * ``cancelledEvent`` (trade: :class:`.Trade`)
    """
    events = (
        'statusEvent', 'modifyEvent', 'fillEvent',
        'commissionReportEvent', 'filledEvent',
        'cancelEvent', 'cancelledEvent')

    defaults = dict(
        contract=None,
        order=None,
        orderStatus=None,
        fills=None,
        log=None
    )
    __slots__ = tuple(defaults.keys()) + events + ('__dict__',)

    def __init__(self, *args, **kwargs):
        Object.__init__(self, *args, **kwargs)
        self.statusEvent = Event('statusEvent')
        self.modifyEvent = Event('modifyEvent')
        self.fillEvent = Event('fillEvent')
        self.commissionReportEvent = Event('commissionReportEvent')
        self.filledEvent = Event('filledEvent')
        self.cancelEvent = Event('cancelEvent')
        self.cancelledEvent = Event('cancelledEvent')

    def isActive(self):
        """
        True if eligible for execution, false otherwise.
        """
        return self.orderStatus.status in OrderStatus.ActiveStates

    def isDone(self):
        """
        True if completely filled or cancelled, false otherwise.
        """
        return self.orderStatus.status in OrderStatus.DoneStates

    def filled(self):
        """
        Number of shares filled.
        """
        fills = self.fills
        if self.contract.secType == 'BAG':
            # don't count fills for the leg contracts
            fills = [f for f in fills if f.contract.secType == 'BAG']
        return sum(f.execution.shares for f in fills)

    def remaining(self):
        """
        Number of shares remaining to be filled.
        """
        return self.order.totalQuantity - self.filled()

class OrderStatus(Object):
    defaults = dict(
        orderId=0,
        status='',
        filled=0,
        remaining=0,
        avgFillPrice=0.0,
        permId=0,
        parentId=0,
        lastFillPrice=0.0,
        clientId=0,
        whyHeld='',
        mktCapPrice=0.0,
        lastLiquidity=0
    )
    __slots__ = defaults.keys()

    PendingSubmit = 'PendingSubmit'
    PendingCancel = 'PendingCancel'
    PreSubmitted = 'PreSubmitted'
    Submitted = 'Submitted'
    ApiPending = 'ApiPending'
    ApiCancelled = 'ApiCancelled'
    Cancelled = 'Cancelled'
    Filled = 'Filled'
    Inactive = 'Inactive'

    DoneStates = {'Filled', 'Cancelled', 'ApiCancelled'}
    ActiveStates = {'PendingSubmit', 'ApiPending', 'PreSubmitted', 'Submitted'}

class Order(Object):
    """
    Order for trading contracts.

    https://interactivebrokers.github.io/tws-api/available_orders.html
    """
    defaults = dict(
        orderId=0,
        clientId=0,
        permId=0,
        action='',
        totalQuantity=0,
        orderType='',
        lmtPrice=UNSET_DOUBLE,
        auxPrice=UNSET_DOUBLE,
        tif='',
        activeStartTime='',
        activeStopTime='',
        ocaGroup='',
        ocaType=0,
        orderRef='',
        transmit=True,
        parentId=0,
        blockOrder=False,
        sweepToFill=False,
        displaySize=0,
        triggerMethod=0,
        outsideRth=False,
        hidden=False,
        goodAfterTime='',
        goodTillDate='',
        rule80A='',
        allOrNone=False,
        minQty=UNSET_INTEGER,
        percentOffset=UNSET_DOUBLE,
        overridePercentageConstraints=False,
        trailStopPrice=UNSET_DOUBLE,
        trailingPercent=UNSET_DOUBLE,
        faGroup='',
        faProfile='',
        faMethod='',
        faPercentage='',
        designatedLocation='',
        openClose="O",
        origin=0,
        shortSaleSlot=0,
        exemptCode=-1,
        discretionaryAmt=0,
        eTradeOnly=True,
        firmQuoteOnly=True,
        nbboPriceCap=UNSET_DOUBLE,
        optOutSmartRouting=False,
        auctionStrategy=0,
        startingPrice=UNSET_DOUBLE,
        stockRefPrice=UNSET_DOUBLE,
        delta=UNSET_DOUBLE,
        stockRangeLower=UNSET_DOUBLE,
        stockRangeUpper=UNSET_DOUBLE,
        randomizePrice=False,
        randomizeSize=False,
        volatility=UNSET_DOUBLE,
        volatilityType=UNSET_INTEGER,
        deltaNeutralOrderType='',
        deltaNeutralAuxPrice=UNSET_DOUBLE,
        deltaNeutralConId=0,
        deltaNeutralSettlingFirm='',
        deltaNeutralClearingAccount='',
        deltaNeutralClearingIntent='',
        deltaNeutralOpenClose='',
        deltaNeutralShortSale=False,
        deltaNeutralShortSaleSlot=0,
        deltaNeutralDesignatedLocation='',
        continuousUpdate=False,
        referencePriceType=UNSET_INTEGER,
        basisPoints=UNSET_DOUBLE,
        basisPointsType=UNSET_INTEGER,
        scaleInitLevelSize=UNSET_INTEGER,
        scaleSubsLevelSize=UNSET_INTEGER,
        scalePriceIncrement=UNSET_DOUBLE,
        scalePriceAdjustValue=UNSET_DOUBLE,
        scalePriceAdjustInterval=UNSET_INTEGER,
        scaleProfitOffset=UNSET_DOUBLE,
        scaleAutoReset=False,
        scaleInitPosition=UNSET_INTEGER,
        scaleInitFillQty=UNSET_INTEGER,
        scaleRandomPercent=False,
        scaleTable='',
        hedgeType='',
        hedgeParam='',
        account='',
        settlingFirm='',
        clearingAccount='',
        clearingIntent='',
        algoStrategy='',
        algoParams=None,
        smartComboRoutingParams=None,
        algoId='',
        whatIf=False,
        notHeld=False,
        solicited=False,
        modelCode='',
        orderComboLegs=None,
        orderMiscOptions=None,
        referenceContractId=0,
        peggedChangeAmount=0.0,
        isPeggedChangeAmountDecrease=False,
        referenceChangeAmount=0.0,
        referenceExchangeId='',
        adjustedOrderType='',
        triggerPrice=UNSET_DOUBLE,
        adjustedStopPrice=UNSET_DOUBLE,
        adjustedStopLimitPrice=UNSET_DOUBLE,
        adjustedTrailingAmount=UNSET_DOUBLE,
        adjustableTrailingUnit=0,
        lmtPriceOffset=UNSET_DOUBLE,
        conditions=None,
        conditionsCancelOrder=False,
        conditionsIgnoreRth=False,
        extOperator='',
        softDollarTier=None,
        cashQty=UNSET_DOUBLE,
        mifid2DecisionMaker='',
        mifid2DecisionAlgo='',
        mifid2ExecutionTrader='',
        mifid2ExecutionAlgo='',
        dontUseAutoPriceForHedge=False,
        isOmsContainer=False,
        discretionaryUpToLimitPrice=False,
        autoCancelDate='',
        filledQuantity=UNSET_DOUBLE,
        refFuturesConId=0,
        autoCancelParent=False,
        shareholder='',
        imbalanceOnly=False,
        routeMarketableToBbo=False,
        parentPermId=0,
        usePriceMgmtAlgo=False)
    __slots__ = defaults.keys()

    def __init__(self, *args, **kwargs):
        Object.__init__(self, *args, **kwargs)
        if not self.conditions:
            self.conditions = []
        if not self.softDollarTier:
            self.softDollarTier = SoftDollarTier()

    def __repr__(self):
        attrs = self.nonDefaults()
        if self.__class__ is not Order:
            attrs.pop('orderType', None)
        clsName = self.__class__.__name__
        kwargs = ', '.join(f'{k}={v!r}' for k, v in attrs.items())
        return f'{clsName}({kwargs})'

    __str__ = __repr__

    def __eq__(self, other):
        return self is other

    def __hash__(self):
        return id(self)

class LimitOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, lmtPrice, **kwargs):
        Order.__init__(
            self, orderType='LMT', action=action,
            totalQuantity=totalQuantity, lmtPrice=lmtPrice, **kwargs)

class MarketOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, **kwargs):
        Order.__init__(
            self, orderType='MKT', action=action,
            totalQuantity=totalQuantity, **kwargs)

class StopOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, stopPrice, **kwargs):
        Order.__init__(
            self, orderType='STP', action=action,
            totalQuantity=totalQuantity, auxPrice=stopPrice, **kwargs)

class StopLimitOrder(Order):
    __slots__ = ()

    def __init__(self, action, totalQuantity, lmtPrice, stopPrice, **kwargs):
        Order.__init__(
            self, orderType='STP LMT', action=action,
            totalQuantity=totalQuantity, lmtPrice=lmtPrice,
            auxPrice=stopPrice, **kwargs)

class OrderCondition(Object):
    __slots__ = ()

    @staticmethod
    def createClass(condType):
        d = {
            1: PriceCondition,
            3: TimeCondition,
            4: MarginCondition,
            5: ExecutionCondition,
            6: VolumeCondition,
            7: PercentChangeCondition}
        return d[condType]

    def __repr__(self):
        clsName = self.__class__.__qualname__
        kwargs = ', '.join(f'{k}={v!r}' for k, v in self.dict().items())
        return f'{clsName}({kwargs})'

    def And(self):
        self.conjunction = 'a'
        return self

    def Or(self):
        self.conjunction = 'o'
        return self

class PriceCondition(OrderCondition):
    defaults = dict(
        condType=1,
        conjunction='a',
        isMore=True,
        price=0.0,
        conId=0,
        exch='',
        triggerMethod=0)
    __slots__ = defaults.keys()

class TimeCondition(OrderCondition):
    defaults = dict(
        condType=3,
        conjunction='a',
        isMore=True,
        time='')
    __slots__ = defaults.keys()

class MarginCondition(OrderCondition):
    defaults = dict(
        condType=4,
        conjunction='a',
        isMore=True,
        percent=0)
    __slots__ = defaults.keys()

class ExecutionCondition(OrderCondition):
    defaults = dict(
        condType=5,
        conjunction='a',
        secType='',
        exch='',
        symbol='')
    __slots__ = defaults.keys()

class VolumeCondition(OrderCondition):
    defaults = dict(
        condType=6,
        conjunction='a',
        isMore=True,
        volume=0,
        conId=0,
        exch='')
    __slots__ = defaults.keys()

class PercentChangeCondition(OrderCondition):
    defaults = dict(
        condType=7,
        conjunction='a',
        isMore=True,
        changePercent=0.0,
        conId=0,
        exch='')
    __slots__ = defaults.keys()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import getopt, sys, os, subprocess, time

def usage():
    print """Lithium, an automated testcase reduction tool by Jesse Ruderman

Usage:

./lithium.py [options] condition [condition options] file-to-reduce

Example:

./lithium.py crashes 120 ~/tracemonkey/js/src/debug/js -j a.js
     Lithium will reduce a.js subject to the condition that the following
     crashes in 120 seconds:
     ~/tracemonkey/js/src/debug/js -j a.js

Options:
* --char (-c).
      Don't treat lines as atomic units; treat the file as a sequence
      of characters rather than a sequence of lines.
* --strategy=[minimize, remove-pair, remove-substring, check-only].
      default: minimize.
* --testcase=filename.
      default: last thing on the command line, which can double as passing in.

Additional options for the default strategy (--strategy=minimize)
* --repeat=[always, last, never]. default: last
     Whether to repeat a chunk size if chunks are removed.
* --max=n. default: about half of the file.
* --min=n. default: 1.
     What chunk sizes to test.  Must be powers of two.
* --chunksize=n
     Shortcut for "repeat=never, min=n, max=n"
* --chunkstart=n
     For the first round only, start n chars/lines into the file. Best for max to divide n.  [Mostly intended for internal use]
* --repeatfirstround
     Repeat the first round.  [Mostly intended for internal use]
* --maxruntime=n
     If reduction takes more than n seconds, stop (and print instructions for continuing).

See doc/using.html for more information.

"""

# Globals

strategy = "minimize"
minimizeRepeat = "last"
minimizeMin = 1
minimizeMax = pow(2, 30)
minimizeChunkStart = 0
minimizeRepeatFirstRound = False

atom = "line"

conditionScript = None
conditionArgs = None
testcaseFilename = None
testcaseExtension = ""

testCount = 0
testTotal = 0

tempDir = None
tempFileCount = 1

before = ""
after = ""
parts = []
allPositionalArgs = []
stopAfterTime = None

# Main and friends

def main():
    global conditionScript, conditionArgs, testcaseFilename, testcaseExtension, strategy, allPositionalArgs
    global parts

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hc", ["help", "char", "strategy=", "repeat=", "min=", "max=", "chunksize=", "chunkstart=", "testcase=", "tempdir=", "repeatfirstround", "maxruntime="])
    except getopt.GetoptError, exc:
        usageError(exc.msg)

    allPositionalArgs = args

    if len(args) == 0:
        # No arguments; not even a condition was specified
        usage()
        sys.exit(0)

    if len(args) > 1:
        testcaseFilename = args[-1] # can be overridden by --testcase in processOptions

    processOptions(opts)

    if testcaseFilename == None:
        usageError("No testcase specified (use --testcase or last condition arg)")

    conditionScript = importRelativeOrAbsolute(args[0])
    conditionArgs = args[1:]

    if hasattr(conditionScript, "init"):
        conditionScript.init(conditionArgs)

    e = testcaseFilename.rsplit(".", 1)
    if len(e) > 1:
        testcaseExtension = "." + e[1]

    readTestcase()

    if tempDir == None:
        createTempDir()
        print "Intermediate files will be stored in " + tempDir + os.sep + "."

    if strategy == "check-only":
        print 'Interesting.' if interesting(parts) else 'Not interesting.'
        sys.exit(0)

    strategyFunction = {
        'minimize': minimize,
        'remove-pair': tryRemovingPair,
        'remove-adjacent-pairs': tryRemovingAdjacentPairs,
        'remove-substring': tryRemovingSubstring
    }.get(strategy, None)

    if not strategyFunction:
        usageError("Unknown strategy!")

    print "The original testcase has " + quantity(len(parts), atom) + "."
    print "Checking that the original testcase is 'interesting'..."
    if not interesting(parts):
        usageError("The original testcase is not 'interesting'!")

    if len(parts) == 0:
        usageError("The file has " + quantity(0, atom) + " so there's nothing for Lithium to try to remove!")

    writeTestcaseTemp("original", False)
    strategyFunction()

def processOptions(opts):
    global atom, strategy, testcaseFilename, tempDir
    global minimizeRepeat, minimizeMin, minimizeMax, minimizeChunkStart, minimizeRepatFirstRound, stopAfterTime

    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit(0)
        elif o == "--testcase":
            testcaseFilename = a
        elif o == "--tempdir":
            tempDir = a
        elif o in ("-c", "--char"):
            atom = "char"
        elif o == "--strategy":
            strategy = a
        elif o == "--min":
            minimizeMin = int(a)
            if not isPowerOfTwo(minimizeMin):
                usageError("min must be a power of two.")
        elif o == "--max":
            minimizeMax = int(a)
            if not isPowerOfTwo(minimizeMax):
                usageError("max must be a power of two.")
        elif o == "--repeat":
            minimizeRepeat = a
            if not (minimizeRepeat in ("always", "last", "never")):
                usageError("repeat must be 'always', 'last', or 'never'.")
        elif o == "--chunksize":
            minimizeMin = int(a)
            minimizeMax = minimizeMin
            minimizeRepeat = "never"
            if not isPowerOfTwo(minimizeMin):
                usageError("Chunk size must be a power of two.")
        elif o == "--chunkstart":
            minimizeChunkStart = int(a)
        elif o == "--repeatfirstround":
            minimizeRepatFirstRound = True
        elif o == "--maxruntime":
            stopAfterTime = time.time() + int(a)

def usageError(s):
    print s
    print "Use --help if you need it :)"
    sys.exit(2)

# Functions for manipulating the testcase (aka the 'interesting' file)

def readTestcase():
    hasDDSection = False

    try:
        file = open(testcaseFilename, "r")
    except IOError:
        usageError("Can't read the original testcase file, " + testcaseFilename + "!")

    # Determine whether the file has a DDBEGIN..DDEND section.
    for line in file:
        if line.find("DDEND") != -1:
            usageError("The testcase (" + testcaseFilename + ") has a line containing 'DDEND' without a line containing 'DDBEGIN' before it.")
        if line.find("DDBEGIN") != -1:
            hasDDSection = True
            break

    file.seek(0)

    if hasDDSection:
        # Reduce only the part of the file between 'DDBEGIN' and 'DDEND',
        # leaving the rest unchanged.
        #print "Testcase has a DD section"
        readTestcaseWithDDSection(file)
    else:
        # Reduce the entire file.
        #print "Testcase does not have a DD section"
        for line in file:
            readTestcaseLine(line)

    file.close()

def readTestcaseWithDDSection(file):
    global before, after
    global parts

    for line in file:
        before += line
        if line.find("DDBEGIN") != -1:
            break

    for line in file:
        if line.find("DDEND") != -1:
            after += line
            break
        readTestcaseLine(line)
    else:
        usageError("The testcase (" + testcaseFilename + ") has a line containing 'DDBEGIN' but no line containing 'DDEND'.")

    for line in file:
        after += line

    if atom == "char" and len(parts) > 0:
        # Move the line break at the end of the last line out of the reducible
        # part so the "DDEND" line doesn't get combined with another line.
        parts.pop()
        after = "\n" + after

def readTestcaseLine(line):
    global atom
    global parts

    if atom == "line":
       parts.append(line)
    elif atom == "char":
        for char in line:
            parts.append(char)

def writeTestcase(filename):
    file = open(filename, "w")
    file.write(before)
    for i in range(len(parts)):
        file.write(parts[i])
    file.write(after)
    file.close()

def writeTestcaseTemp(partialFilename, useNumber):
    global tempFileCount
    if useNumber:
        partialFilename = str(tempFileCount) + "-" + partialFilename
        tempFileCount += 1
    writeTestcase(tempDir + os.sep + partialFilename + testcaseExtension)

def createTempDir():
    global tempDir
    i = 1
    while 1:
        tempDir = "tmp" + str(i)
        # To avoid race conditions, we use try/except instead of exists/create
        # Hopefully we don't get any errors other than "File exists" :)
        try:
            os.mkdir(tempDir)
            break
        except OSError, e:
            i += 1

# If the file is still interesting after the change, changes the global "parts" and returns True.
def interesting(partsSuggestion):
    global tempFileCount, testcaseFilename, conditionArgs
    global testCount, testTotal
    global parts
    oldParts = parts # would rather be less side-effecty about this, and be passing partsSuggestion around
    parts = partsSuggestion

    writeTestcase(testcaseFilename)

    testCount += 1
    testTotal += len(parts)

    tempPrefix = tempDir + os.sep + str(tempFileCount)
    inter = conditionScript.interesting(conditionArgs, tempPrefix)

    # Save an extra copy of the file inside the temp directory.
    # This is useful if you're reducing an assertion and encounter a crash:
    # it gives you a way to try to reproduce the crash.
    if tempDir != None:
        tempFileTag = "interesting" if inter else "boring"
        writeTestcaseTemp(tempFileTag, True)

    if not inter:
        parts = oldParts
    return inter

# Main reduction algorithm

def minimize():
    global parts, testCount, testTotal
    global minimizeMax, minimizeMin, minimizeChunkStart, minimizeRepeatFirstRound
    origNumParts = len(parts)
    chunkSize = min(minimizeMax, largestPowerOfTwoSmallerThan(origNumParts))
    finalChunkSize = max(minimizeMin, 1)
    chunkStart = minimizeChunkStart
    anyChunksRemoved = minimizeRepeatFirstRound

    while 1:
        if stopAfterTime != None and time.time() > stopAfterTime:
            # Not all switches will be copied!  Be sure to add --tempdir, --maxruntime if desired.
            print "To continue, run lithium.py with: " + " ".join(
                 [
                 "--testcase=" + testcaseFilename,
                 "--max=" + str(chunkSize),
                 "--chunkstart=" + str(chunkStart)] +
                (["--repeatfirstround"] if anyChunksRemoved else []) +
                (["--char"] if atom == "char" else []) +
                allPositionalArgs
                )
            break

        if chunkStart >= len(parts):
            writeTestcaseTemp("did-round-" + str(chunkSize), True);
            last = (chunkSize == finalChunkSize)
            print ""
            if anyChunksRemoved and (minimizeRepeat == "always" or (minimizeRepeat == "last" and last)):
                chunkStart = 0
                print "Starting another round of chunk size " + str(chunkSize)
            elif last:
                print "Lithium is done!"
                break
            else:
                chunkStart = 0
                chunkSize /= 2
                print "Halving chunk size to " + str(chunkSize)
            anyChunksRemoved = False

        chunkEnd = min(len(parts), chunkStart + chunkSize)
        description = "Removing a chunk of size " + str(chunkSize) + " starting at " + str(chunkStart) + " of " + str(len(parts))
        if interesting(parts[:chunkStart] + parts[chunkEnd:]):
            print description + " was a successful reduction :)"
            anyChunksRemoved = True
            # leave chunkStart the same
        else:
            print description + " made the file 'uninteresting'."
            chunkStart += chunkSize

    writeTestcase(testcaseFilename)

    if chunkSize == 1 and not anyChunksRemoved and minimizeRepeat != "never":
        print "  Removing any single " + atom + " from the final file makes it uninteresting!"

    print "  Initial size: " + quantity(origNumParts, atom)
    print "  Final size: " + quantity(len(parts), atom)
    print "  Tests performed: " + str(testCount)
    print "  Test total: " + quantity(testTotal, atom)

def tryRemovingChunks(chunkSize):

    print "Done with a round of chunk size " + str(chunkSize) + "!"
    return anyChunksRemoved

# Other reduction algorithms
# (Use these if you're really frustrated with something you know is 1-minimal.)

def tryRemovingAdjacentPairs():
    # XXX capture the idea that after removing (4,5) it might be sensible to remove (3,6)
    # but also that after removing (2,3) and (4,5) it might be sensible to remove (1,6)
    # XXX also want to remove three at a time, and two at a time that are one line apart
    for i in range(0, numParts - 2):
        if enabled[i]:
            enabled[i] = False
            enabled[i + 1] = False
            if interesting():
                print "Removed an adjacent pair based at " + str(i)
            else:
                enabled[i] = True
                enabled[i + 1] = True
    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Done with one pass of removing adjacent pairs"

def tryRemovingPair():
    for i in range(0, numParts):
        enabled[i] = False
        for j in range(i + 1, numParts):
            enabled[j] = False
            print "Trying removing the pair " + str(i) + ", " + str(j)
            if interesting():
                print "Success!  Removed a pair!  Exiting."
                sys.exit(0)
            enabled[j] = True
        enabled[i] = True

    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Failure!  No pair can be removed."

def tryRemovingSubstring():
    for i in range(0, numParts):
        for j in range(i, numParts):
            enabled[j] = False
            print "Trying removing the substring " + str(i) + ".." + str(j)
            if interesting():
                print "Success!  Removed a substring!  Exiting."
                sys.exit(0)
        for j in range(i, numParts):
            enabled[j] = True

    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Failure!  No substring can be removed."

# Helpers

def divideRoundingUp(n, d):
    return (n // d) + (1 if n % d != 0 else 0)

def isPowerOfTwo(n):
    i = 1
    while 1:
        if i == n:
            return True
        if i > n:
            return False
        i *= 2

def largestPowerOfTwoSmallerThan(n):
    i = 1
    while 1:
        if i * 2 >= n:
            return i
        i *= 2

def quantity(n, s):
    """Convert a quantity to a string, with correct pluralization."""
    r = str(n) + " " + s
    if n != 1:
        r += "s"
    return r

def importRelativeOrAbsolute(f):
    # maybe there's a way to do this more sanely with the |imp| module...
    if f.endswith(".py"):
        f = f[:-3]
    if f.rfind(os.path.sep):
        # Add the path part of the filename to the import path
        (p, _, f) = f.rpartition(os.path.sep)
        sys.path.append(p)
    else:
        # Add working directory to the import path
        sys.path.append(".")
    module = __import__(f)
    del sys.path[0]
    return module

# Run main

if __name__ == "__main__":
    main()
<commit_msg>Add note<commit_after>#!/usr/bin/env python
import getopt, sys, os, subprocess, time

def usage():
    print """Lithium, an automated testcase reduction tool by Jesse Ruderman

Usage:

./lithium.py [options] condition [condition options] file-to-reduce

Example:

./lithium.py crashes 120 ~/tracemonkey/js/src/debug/js -j a.js
     Lithium will reduce a.js subject to the condition that the following
     crashes in 120 seconds:
     ~/tracemonkey/js/src/debug/js -j a.js

Options:
* --char (-c).
      Don't treat lines as atomic units; treat the file as a sequence
      of characters rather than a sequence of lines.
* --strategy=[minimize, remove-pair, remove-substring, check-only].
      default: minimize.
* --testcase=filename.
      default: last thing on the command line, which can double as passing in.

Additional options for the default strategy (--strategy=minimize)
* --repeat=[always, last, never]. default: last
     Whether to repeat a chunk size if chunks are removed.
* --max=n. default: about half of the file.
* --min=n. default: 1.
     What chunk sizes to test.  Must be powers of two.
* --chunksize=n
     Shortcut for "repeat=never, min=n, max=n"
* --chunkstart=n
     For the first round only, start n chars/lines into the file. Best for max to divide n.  [Mostly intended for internal use]
* --repeatfirstround
     Repeat the first round.  [Mostly intended for internal use]
* --maxruntime=n
     If reduction takes more than n seconds, stop (and print instructions for continuing).

See doc/using.html for more information.

"""

# Globals

strategy = "minimize"
minimizeRepeat = "last"
minimizeMin = 1
minimizeMax = pow(2, 30)
minimizeChunkStart = 0
minimizeRepeatFirstRound = False

atom = "line"

conditionScript = None
conditionArgs = None
testcaseFilename = None
testcaseExtension = ""

testCount = 0
testTotal = 0

tempDir = None
tempFileCount = 1

before = ""
after = ""
parts = []
allPositionalArgs = []
stopAfterTime = None

# Main and friends

def main():
    global conditionScript, conditionArgs, testcaseFilename, testcaseExtension, strategy, allPositionalArgs
    global parts

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hc", ["help", "char", "strategy=", "repeat=", "min=", "max=", "chunksize=", "chunkstart=", "testcase=", "tempdir=", "repeatfirstround", "maxruntime="])
    except getopt.GetoptError, exc:
        usageError(exc.msg)

    allPositionalArgs = args

    if len(args) == 0:
        # No arguments; not even a condition was specified
        usage()
        sys.exit(0)

    if len(args) > 1:
        testcaseFilename = args[-1] # can be overridden by --testcase in processOptions

    processOptions(opts)

    if testcaseFilename == None:
        usageError("No testcase specified (use --testcase or last condition arg)")

    conditionScript = importRelativeOrAbsolute(args[0])
    conditionArgs = args[1:]

    if hasattr(conditionScript, "init"):
        conditionScript.init(conditionArgs)

    e = testcaseFilename.rsplit(".", 1)
    if len(e) > 1:
        testcaseExtension = "." + e[1]

    readTestcase()

    if tempDir == None:
        createTempDir()
        print "Intermediate files will be stored in " + tempDir + os.sep + "."

    if strategy == "check-only":
        print 'Interesting.' if interesting(parts) else 'Not interesting.'
        sys.exit(0)

    strategyFunction = {
        'minimize': minimize,
        'remove-pair': tryRemovingPair,
        'remove-adjacent-pairs': tryRemovingAdjacentPairs,
        'remove-substring': tryRemovingSubstring
    }.get(strategy, None)

    if not strategyFunction:
        usageError("Unknown strategy!")

    print "The original testcase has " + quantity(len(parts), atom) + "."
    print "Checking that the original testcase is 'interesting'..."
    if not interesting(parts):
        usageError("The original testcase is not 'interesting'!")

    if len(parts) == 0:
        usageError("The file has " + quantity(0, atom) + " so there's nothing for Lithium to try to remove!")

    writeTestcaseTemp("original", False)
    strategyFunction()

def processOptions(opts):
    global atom, strategy, testcaseFilename, tempDir
    global minimizeRepeat, minimizeMin, minimizeMax, minimizeChunkStart, minimizeRepatFirstRound, stopAfterTime

    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit(0)
        elif o == "--testcase":
            testcaseFilename = a
        elif o == "--tempdir":
            tempDir = a
        elif o in ("-c", "--char"):
            atom = "char"
        elif o == "--strategy":
            strategy = a
        elif o == "--min":
            minimizeMin = int(a)
            if not isPowerOfTwo(minimizeMin):
                usageError("min must be a power of two.")
        elif o == "--max":
            minimizeMax = int(a)
            if not isPowerOfTwo(minimizeMax):
                usageError("max must be a power of two.")
        elif o == "--repeat":
            minimizeRepeat = a
            if not (minimizeRepeat in ("always", "last", "never")):
                usageError("repeat must be 'always', 'last', or 'never'.")
        elif o == "--chunksize":
            minimizeMin = int(a)
            minimizeMax = minimizeMin
            minimizeRepeat = "never"
            if not isPowerOfTwo(minimizeMin):
                usageError("Chunk size must be a power of two.")
        elif o == "--chunkstart":
            minimizeChunkStart = int(a)
        elif o == "--repeatfirstround":
            minimizeRepatFirstRound = True
        elif o == "--maxruntime":
            stopAfterTime = time.time() + int(a)

def usageError(s):
    print s
    print "Use --help if you need it :)"
    sys.exit(2)

# Functions for manipulating the testcase (aka the 'interesting' file)

def readTestcase():
    hasDDSection = False

    try:
        file = open(testcaseFilename, "r")
    except IOError:
        usageError("Can't read the original testcase file, " + testcaseFilename + "!")

    # Determine whether the file has a DDBEGIN..DDEND section.
    for line in file:
        if line.find("DDEND") != -1:
            usageError("The testcase (" + testcaseFilename + ") has a line containing 'DDEND' without a line containing 'DDBEGIN' before it.")
        if line.find("DDBEGIN") != -1:
            hasDDSection = True
            break

    file.seek(0)

    if hasDDSection:
        # Reduce only the part of the file between 'DDBEGIN' and 'DDEND',
        # leaving the rest unchanged.
        #print "Testcase has a DD section"
        readTestcaseWithDDSection(file)
    else:
        # Reduce the entire file.
        #print "Testcase does not have a DD section"
        for line in file:
            readTestcaseLine(line)

    file.close()

def readTestcaseWithDDSection(file):
    global before, after
    global parts

    for line in file:
        before += line
        if line.find("DDBEGIN") != -1:
            break

    for line in file:
        if line.find("DDEND") != -1:
            after += line
            break
        readTestcaseLine(line)
    else:
        usageError("The testcase (" + testcaseFilename + ") has a line containing 'DDBEGIN' but no line containing 'DDEND'.")

    for line in file:
        after += line

    if atom == "char" and len(parts) > 0:
        # Move the line break at the end of the last line out of the reducible
        # part so the "DDEND" line doesn't get combined with another line.
        parts.pop()
        after = "\n" + after

def readTestcaseLine(line):
    global atom
    global parts

    if atom == "line":
       parts.append(line)
    elif atom == "char":
        for char in line:
            parts.append(char)

def writeTestcase(filename):
    file = open(filename, "w")
    file.write(before)
    for i in range(len(parts)):
        file.write(parts[i])
    file.write(after)
    file.close()

def writeTestcaseTemp(partialFilename, useNumber):
    global tempFileCount
    if useNumber:
        partialFilename = str(tempFileCount) + "-" + partialFilename
        tempFileCount += 1
    writeTestcase(tempDir + os.sep + partialFilename + testcaseExtension)

def createTempDir():
    global tempDir
    i = 1
    while 1:
        tempDir = "tmp" + str(i)
        # To avoid race conditions, we use try/except instead of exists/create
        # Hopefully we don't get any errors other than "File exists" :)
        try:
            os.mkdir(tempDir)
            break
        except OSError, e:
            i += 1

# If the file is still interesting after the change, changes the global "parts" and returns True.
def interesting(partsSuggestion):
    global tempFileCount, testcaseFilename, conditionArgs
    global testCount, testTotal
    global parts
    oldParts = parts # would rather be less side-effecty about this, and be passing partsSuggestion around
    parts = partsSuggestion

    writeTestcase(testcaseFilename)

    testCount += 1
    testTotal += len(parts)

    tempPrefix = tempDir + os.sep + str(tempFileCount)
    inter = conditionScript.interesting(conditionArgs, tempPrefix)

    # Save an extra copy of the file inside the temp directory.
    # This is useful if you're reducing an assertion and encounter a crash:
    # it gives you a way to try to reproduce the crash.
    if tempDir != None:
        tempFileTag = "interesting" if inter else "boring"
        writeTestcaseTemp(tempFileTag, True)

    if not inter:
        parts = oldParts
    return inter

# Main reduction algorithm

def minimize():
    global parts, testCount, testTotal
    global minimizeMax, minimizeMin, minimizeChunkStart, minimizeRepeatFirstRound
    origNumParts = len(parts)
    chunkSize = min(minimizeMax, largestPowerOfTwoSmallerThan(origNumParts))
    finalChunkSize = max(minimizeMin, 1)
    chunkStart = minimizeChunkStart
    anyChunksRemoved = minimizeRepeatFirstRound

    while 1:
        if stopAfterTime != None and time.time() > stopAfterTime:
            # Not all switches will be copied!  Be sure to add --tempdir, --maxruntime if desired.
            print "To continue, run lithium.py with: " + " ".join(
                 [
                 "--testcase=" + testcaseFilename,
                 "--max=" + str(chunkSize),
                 "--chunkstart=" + str(chunkStart)] +
                (["--repeatfirstround"] if anyChunksRemoved else []) +
                (["--char"] if atom == "char" else []) +
                allPositionalArgs
                )
            break

        if chunkStart >= len(parts):
            writeTestcaseTemp("did-round-" + str(chunkSize), True);
            last = (chunkSize == finalChunkSize)
            print ""
            if anyChunksRemoved and (minimizeRepeat == "always" or (minimizeRepeat == "last" and last)):
                chunkStart = 0
                print "Starting another round of chunk size " + str(chunkSize)
            elif last:
                print "Lithium is done!"
                break
            else:
                chunkStart = 0
                chunkSize /= 2
                print "Halving chunk size to " + str(chunkSize)
            anyChunksRemoved = False

        chunkEnd = min(len(parts), chunkStart + chunkSize)
        description = "Removing a chunk of size " + str(chunkSize) + " starting at " + str(chunkStart) + " of " + str(len(parts))
        if interesting(parts[:chunkStart] + parts[chunkEnd:]):
            print description + " was a successful reduction :)"
            anyChunksRemoved = True
            # leave chunkStart the same
        else:
            print description + " made the file 'uninteresting'."
            chunkStart += chunkSize

    writeTestcase(testcaseFilename)

    if chunkSize == 1 and not anyChunksRemoved and minimizeRepeat != "never":
        print "  Removing any single " + atom + " from the final file makes it uninteresting!"

    print "  Initial size: " + quantity(origNumParts, atom)
    print "  Final size: " + quantity(len(parts), atom)
    print "  Tests performed: " + str(testCount)
    print "  Test total: " + quantity(testTotal, atom)

def tryRemovingChunks(chunkSize):

    print "Done with a round of chunk size " + str(chunkSize) + "!"
    return anyChunksRemoved

# Other reduction algorithms
# (Use these if you're really frustrated with something you know is 1-minimal.)

def tryRemovingAdjacentPairs():
    # XXX capture the idea that after removing (4,5) it might be sensible to remove (3,6)
    # but also that after removing (2,3) and (4,5) it might be sensible to remove (1,6)
    # XXX also want to remove three at a time, and two at a time that are one line apart
    for i in range(0, numParts - 2):
        if enabled[i]:
            enabled[i] = False
            enabled[i + 1] = False
            if interesting():
                print "Removed an adjacent pair based at " + str(i)
            else:
                enabled[i] = True
                enabled[i + 1] = True
    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Done with one pass of removing adjacent pairs"

def tryRemovingPair():
    for i in range(0, numParts):
        enabled[i] = False
        for j in range(i + 1, numParts):
            enabled[j] = False
            print "Trying removing the pair " + str(i) + ", " + str(j)
            if interesting():
                print "Success!  Removed a pair!  Exiting."
                sys.exit(0)
            enabled[j] = True
        enabled[i] = True

    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Failure!  No pair can be removed."

def tryRemovingSubstring():
    for i in range(0, numParts):
        for j in range(i, numParts):
            enabled[j] = False
            print "Trying removing the substring " + str(i) + ".." + str(j)
            if interesting():
                print "Success!  Removed a substring!  Exiting."
                sys.exit(0)
        for j in range(i, numParts):
            enabled[j] = True

    # Restore the original testcase
    writeTestcase(testcaseFilename)
    print "Failure!  No substring can be removed."

# Helpers

def divideRoundingUp(n, d):
    return (n // d) + (1 if n % d != 0 else 0)

def isPowerOfTwo(n):
    i = 1
    while 1:
        if i == n:
            return True
        if i > n:
            return False
        i *= 2

def largestPowerOfTwoSmallerThan(n):
    i = 1
    while 1:
        if i * 2 >= n:
            return i
        i *= 2

def quantity(n, s):
    """Convert a quantity to a string, with correct pluralization."""
    r = str(n) + " " + s
    if n != 1:
        r += "s"
    return r

def importRelativeOrAbsolute(f):
    # maybe there's a way to do this more sanely with the |imp| module...
    if f.endswith(".py"):
        f = f[:-3]
    if f.rfind(os.path.sep):
        # Add the path part of the filename to the import path
        (p, _, f) = f.rpartition(os.path.sep)  # os.path.split might work here
        sys.path.append(p)
    else:
        # Add working directory to the import path
        sys.path.append(".")
    module = __import__(f)
    del sys.path[0]
    return module

# Run main

if __name__ == "__main__":
    main()
<|endoftext|>


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

"""
Process logs from syslog and throw them at Lisa via Redis.
"""

import argparse
import logging
import os
import re
import sched
import signal
import sys
import time
from subprocess import call

from pathlib import Path
from statsd import StatsClient

from smithers import conf
from smithers import redis_keys as rkeys
from smithers.redis_client import client as redis
from smithers.utils import register_signals

log = logging.getLogger('lisa')
firefox_re = re.compile(r'firefox-(?:latest|{})'.format(conf.FIREFOX_VERSION),
                        re.IGNORECASE)

# has the system requested shutdown
KILLED = False

SERVER_ENV = os.getenv('DJANGO_SERVER_ENV', 'dev')

# NOTE: Not thread safe. There must be only one bart at a time.
scheduler = sched.scheduler(time.time, time.sleep)
current_event = None

parser = argparse.ArgumentParser(description='Bart flings IPs at Lisa.')
parser.add_argument('--log', default=conf.LOG_LEVEL, metavar='LOG_LEVEL',
                    help='Log level (default: %s)' % conf.LOG_LEVEL)
parser.add_argument('--env', default=SERVER_ENV, choices=['dev', 'prod'],
                    help='Running mode: dev or prod (default: %s)' % SERVER_ENV)
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()

logging.basicConfig(level=getattr(logging, args.log.upper()),
                    format='%(asctime)s: %(message)s')

if hasattr(conf, 'STATSD_HOST'):
    statsd = StatsClient(host=conf.STATSD_HOST,
                         port=conf.STATSD_PORT,
                         prefix=conf.STATSD_PREFIX)
else:
    statsd = False

def handle_signals(signum, frame):
    # NOTE: Makes this thing non-thread-safe
    # Should not be too difficult to fix if we
    # need/want threads.
    try:
        scheduler.cancel(current_event)
    except ValueError:
        pass
    global KILLED
    KILLED = True
    log.info('Attempting to shut down')

def get_syslog_pid(pid_file=None):
    """Return the process ID for rsyslogd."""
    pid_file = Path(pid_file or conf.SYSLOG_PID_FILE)
    with pid_file.open() as fh:
        try:
            return fh.read().strip()
        except IOError:
            log.exception('Can not get syslogd PID')
            return None

def poke_syslogd(sig=signal.SIGHUP):
    pid = get_syslog_pid()
    os.kill(pid, sig)

def get_newest_shared_log():
    """Return the name of the newest non-gzipped log file in NFS."""
    return conf.ARCHIVE_LOG_PATH / redis.get(rkeys.LATEST_LOG_FILE)

def get_fresh_log_file():
    """Do the env-appropriate thing to get a new log file."""
    if args.env == 'prod':
        return rotate_syslog_file()

    if args.env == 'dev':
        return get_newest_shared_log()

    raise ValueError('Unknown server environment: %s' % args.env)

def rotate_syslog_file():
    """Move the main syslog file to a new one, signal syslogd, and return new file."""
    log.debug('Rotating Syslog')
    new_log_file = conf.TMP_DIR / 'glow.{}.log'.format(time.time())
    conf.MAIN_LOG_FILE.rename(new_log_file)
    poke_syslogd()
    return new_log_file

def filter_logs(log_file):
    """Extract valid downlods from logs and return IPs."""
    log.debug('Filtering logs in {}'.format(log_file))
    with log_file.open() as fh:
        for line in fh:
            _, timestamp, ip, status_code, request = line.strip().split()
            if status_code != conf.LOG_HTTP_STATUS:
                continue

            if not firefox_re.search(request):
                continue

            yield ip

def throw_at_lisa(log_file):
    """Put IPs on a queue in redis for Lisa to process."""
    log.debug('Throwing {} at Lisa'.format(log_file))
    count = 0
    pipe = redis.pipeline()
    for ip in filter_logs(log_file):
        pipe.lpush(rkeys.IPLOGS, '0,' + ip)
        count += 1

    pipe.execute()
    statsd.incr('bart.ips_processed', count)

def finalize_log_file(log_file):
    """Do the appropriate thing to the log file for the server env."""
    log.debug('Finalizing {}'.format(log_file))
    if args.env == 'prod':
        # move it to the shared drive
        log_file.rename(conf.ARCHIVE_LOG_PATH / log_file.name)
        redis.set(rkeys.LATEST_LOG_FILE, log_file.name)
    else:
        # gzip the file on dev
        call(['gzip', str(log_file)])

def main():
    global current_event
    if KILLED:
        log.info('Shutdown successful')
        sys.exit(0)

    current_event = scheduler.enter(60, 1, main, ())
    try:
        log_file = get_fresh_log_file()
        throw_at_lisa(log_file)
        finalize_log_file(log_file)
    except Exception:
        log.exception('An error occurred')

if __name__ == '__main__':
    register_signals(handle_signals)
    current_event = scheduler.enter(0, 1, main, ())
    scheduler.run()
<commit_msg>Add better error handling and logging to bart.<commit_after>#!/usr/bin/env python

"""
Process logs from syslog and throw them at Lisa via Redis.
"""

import argparse
import logging
import os
import re
import sched
import signal
import sys
import time
from subprocess import call

from pathlib import Path
from statsd import StatsClient

from smithers import conf
from smithers import redis_keys as rkeys
from smithers.redis_client import client as redis
from smithers.utils import register_signals

log = logging.getLogger('lisa')
firefox_re = re.compile(r'firefox-(?:latest|{})'.format(conf.FIREFOX_VERSION),
                        re.IGNORECASE)

# has the system requested shutdown
KILLED = False

SERVER_ENV = os.getenv('DJANGO_SERVER_ENV', 'dev')

# NOTE: Not thread safe. There must be only one bart at a time.
scheduler = sched.scheduler(time.time, time.sleep)
current_event = None

parser = argparse.ArgumentParser(description='Bart flings IPs at Lisa.')
parser.add_argument('--log', default=conf.LOG_LEVEL, metavar='LOG_LEVEL',
                    help='Log level (default: %s)' % conf.LOG_LEVEL)
parser.add_argument('--logfile', default=None, help='Log file to process')
parser.add_argument('--env', default=SERVER_ENV, choices=['dev', 'prod'],
                    help='Running mode (default: %s)' % SERVER_ENV)
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()

logging.basicConfig(level=getattr(logging, args.log.upper()),
                    format='%(asctime)s: %(message)s')

if hasattr(conf, 'STATSD_HOST'):
    statsd = StatsClient(host=conf.STATSD_HOST,
                         port=conf.STATSD_PORT,
                         prefix=conf.STATSD_PREFIX)
else:
    statsd = False

def handle_signals(signum, frame):
    try:
        scheduler.cancel(current_event)
    except ValueError:
        pass
    global KILLED
    KILLED = True
    log.info('Attempting to shut down')

def get_syslog_pid(pid_file=None):
    """Return the process ID for rsyslogd."""
    pid_file = Path(pid_file or conf.SYSLOG_PID_FILE)
    with pid_file.open() as fh:
        try:
            return fh.read().strip()
        except IOError:
            raise IOError('Can not get syslog PID from {}'.format(pid_file))

def poke_syslogd(sig=signal.SIGHUP):
    pid = get_syslog_pid()
    os.kill(pid, sig)

def get_newest_shared_log():
    """Return the name of the newest non-gzipped log file in NFS."""
    if args.logfile:
        return Path(args.logfile)

    return conf.ARCHIVE_LOG_PATH / redis.get(rkeys.LATEST_LOG_FILE)

def get_fresh_log_file():
    """Do the env-appropriate thing to get a new log file."""
    if args.env == 'prod':
        return rotate_syslog_file()

    if args.env == 'dev':
        return get_newest_shared_log()

    # Fatal error
    log.error('Unknown server environment: %s' % args.env)
    sys.exit(1)

def rotate_syslog_file():
    """Move the main syslog file to a new one, signal syslogd, and return new file."""
    log.debug('Rotating Syslog')
    new_log_file = conf.TMP_DIR / 'glow.{}.log'.format(time.time())
    conf.MAIN_LOG_FILE.rename(new_log_file)
    poke_syslogd()
    return new_log_file

def filter_logs(log_file):
    """Extract valid downlods from logs and return IPs."""
    log.debug('Filtering logs in {}'.format(log_file))
    with log_file.open() as fh:
        for line in fh:
            _, timestamp, ip, status_code, request = line.strip().split()
            if status_code != conf.LOG_HTTP_STATUS:
                continue

            if not firefox_re.search(request):
                continue

            yield ip

def throw_at_lisa(log_file):
    """Put IPs on a queue in redis for Lisa to process."""
    log.debug('Throwing {} at Lisa'.format(log_file))
    if not log_file.exists():
        raise IOError('Log file not found: {}'.format(log_file))

    count = 0
    pipe = redis.pipeline()
    for ip in filter_logs(log_file):
        pipe.lpush(rkeys.IPLOGS, '0,' + ip)
        count += 1

    pipe.execute()
    if statsd:
        statsd.incr('bart.ips_processed', count)

def finalize_log_file(log_file):
    """Do the appropriate thing to the log file for the server env."""
    log.debug('Finalizing {}'.format(log_file))
    if args.env == 'prod':
        # move it to the shared drive
        log_file.rename(conf.ARCHIVE_LOG_PATH / log_file.name)
        redis.set(rkeys.LATEST_LOG_FILE, log_file.name)
    else:
        # gzip the file on dev
        call(['gzip', str(log_file)])

def main():
    global current_event
    if KILLED:
        log.info('Shutdown successful')
        sys.exit(0)

    current_event = scheduler.enter(60, 1, main, ())
    try:
        log_file = get_fresh_log_file()
        throw_at_lisa(log_file)
        finalize_log_file(log_file)
    except Exception as e:
        log.error('ERROR: ' + e.message)

    log.debug('Waiting...')

if __name__ == '__main__':
    register_signals(handle_signals)
    current_event = scheduler.enter(0, 1, main, ())
    scheduler.run()
<|endoftext|>