<commit_before>from __future__ import division
import numpy as np
from functools import partial
from builtins import zip

from pybasicbayes.distributions import DiagonalRegression, Gaussian, Regression

import pyhsmm
from pyhsmm.util.general import list_split

from pyslds.states import HMMSLDSStatesPython, HMMSLDSStatesEigen, HSMMSLDSStatesPython, HSMMSLDSStatesEigen
from pyslds.states import HMMCountSLDSStatesPython, HMMCountSLDSStatesEigen, HSMMCountSLDSStatesPython, \
    HSMMCountSLDSStatesEigen
from pyslds.util import gaussian_map_estimation, regression_map_estimation, gaussian_logprior, regression_logprior


class _SLDSMixin(object):
    def __init__(self,dynamics_distns,emission_distns,init_dynamics_distns,**kwargs):
        self.init_dynamics_distns = init_dynamics_distns
        self.dynamics_distns = dynamics_distns

        # Allow for a single, shared emission distribution
        if not isinstance(emission_distns, list):
            self._single_emission = True
            self._emission_distn = emission_distns
            self.emission_distns = [emission_distns] * len(self.dynamics_distns)
        else:
            assert isinstance(emission_distns, list) and \
                   len(emission_distns) == len(dynamics_distns)
            self._single_emission = False
            self.emission_distns = emission_distns

        super(_SLDSMixin,self).__init__(
            obs_distns=self.dynamics_distns,**kwargs)

    def generate(self, T=100, keep=True, **kwargs):
        s = self._states_class(model=self, T=T,
                               generate=True,
                               initialize_from_prior=True,
                               **kwargs)

        data = self._generate_obs(s)
        if keep:
            self.states_list.append(s)
        return data + (s.stateseq,)


    def _generate_obs(self,s):
        if s.data is None:
            s.data = s.generate_obs()
        else:
            # TODO: Handle missing data
            raise NotImplementedError

        return s.data, s.gaussian_states

    def smooth(self, data, inputs=None, mask=None):
        self.add_data(data, inputs=inputs, mask=mask)
        s = self.states_list.pop()
        return s.smooth()

    @property
    def diagonal_noise(self):
        return all([isinstance(ed, DiagonalRegression) for ed in self.emission_distns])

    @property
    def has_missing_data(self):
        return any([s.mask is not None for s in self.states_list])


class _SLDSGibbsMixin(_SLDSMixin):
    def resample_parameters(self):
        self.resample_lds_parameters()
        self.resample_hmm_parameters()

    def resample_lds_parameters(self):
        self.resample_init_dynamics_distns()
        self.resample_dynamics_distns()
        self.resample_emission_distns()

    def resample_hmm_parameters(self):
        super(_SLDSGibbsMixin,self).resample_parameters()

    def resample_init_dynamics_distns(self):
        for state, d in enumerate(self.init_dynamics_distns):
            d.resample(
                [s.gaussian_states[0] for s in self.states_list
                    if s.stateseq[0] == state])
        self._clear_caches()

    def resample_dynamics_distns(self):
        zs = [s.stateseq[:-1] for s in self.states_list]
        xs = [np.hstack((s.gaussian_states[:-1], s.inputs[:-1]))
              for s in self.states_list]
        ys = [s.gaussian_states[1:] for s in self.states_list]

        for state, d in enumerate(self.dynamics_distns):
            d.resample(
                [(x[z == state], y[z == state])
                 for x, y, z in zip(xs, ys, zs)])
        self._clear_caches()

    def resample_emission_distns(self):
        if self._single_emission:
            data = [(np.hstack((s.gaussian_states, s.inputs)), s.data)
                    for s in self.states_list]
            mask = [s.mask for s in self.states_list] if self.has_missing_data else None

            if self.has_missing_data:
                self._emission_distn.resample(data=data, mask=mask)
            else:
                self._emission_distn.resample(data=data)
        else:
            for state, d in enumerate(self.emission_distns):
                data = [(np.hstack((s.gaussian_states[s.stateseq == state],
                                    s.inputs[s.stateseq == state])),
                         s.data[s.stateseq == state])
                        for s in self.states_list]

                mask = [s.mask[s.stateseq == state] for s in self.states_list] \
                    if self.has_missing_data else None

                if self.has_missing_data:
                    d.resample(data=data, mask=mask)
                else:
                    d.resample(data=data)


        self._clear_caches()

    def resample_obs_distns(self):
        pass  # handled in resample_parameters

    ### joblib parallel

    def _joblib_resample_states(self,states_list,num_procs):
        # TODO: Update to handle inputs
        from joblib import Parallel, delayed
        import parallel

        if len(states_list) > 0:
            joblib_args = list(map(self._get_joblib_pair, states_list))

            parallel.model = self
            parallel.args = list_split(joblib_args, num_procs)

            idxs = range(len(parallel.args))
            raw_stateseqs = Parallel(n_jobs=num_procs,backend='multiprocessing')\
                    (list(map(delayed(parallel._get_sampled_stateseq), idxs)))

            flatten = lambda lst: [x for y in lst for x in y]
            raw_stateseqs = flatten(raw_stateseqs)

            # since list_split might reorder things, do the same to states_list
            states_list = flatten(list_split(states_list, num_procs))<commit_msg>updating parallel code.<commit_after>from __future__ import division
import numpy as np
from functools import partial
from builtins import zip

from pybasicbayes.distributions import DiagonalRegression, Gaussian, Regression

import pyhsmm
from pyhsmm.util.general import list_split

from pyslds.states import HMMSLDSStatesPython, HMMSLDSStatesEigen, HSMMSLDSStatesPython, HSMMSLDSStatesEigen
from pyslds.states import HMMCountSLDSStatesPython, HMMCountSLDSStatesEigen, HSMMCountSLDSStatesPython, \
    HSMMCountSLDSStatesEigen
from pyslds.util import gaussian_map_estimation, regression_map_estimation, gaussian_logprior, regression_logprior


class _SLDSMixin(object):
    def __init__(self,dynamics_distns,emission_distns,init_dynamics_distns,**kwargs):
        self.init_dynamics_distns = init_dynamics_distns
        self.dynamics_distns = dynamics_distns

        # Allow for a single, shared emission distribution
        if not isinstance(emission_distns, list):
            self._single_emission = True
            self._emission_distn = emission_distns
            self.emission_distns = [emission_distns] * len(self.dynamics_distns)
        else:
            assert isinstance(emission_distns, list) and \
                   len(emission_distns) == len(dynamics_distns)
            self._single_emission = False
            self.emission_distns = emission_distns

        super(_SLDSMixin,self).__init__(
            obs_distns=self.dynamics_distns,**kwargs)

    def generate(self, T=100, keep=True, **kwargs):
        s = self._states_class(model=self, T=T,
                               generate=True,
                               initialize_from_prior=True,
                               **kwargs)

        data = self._generate_obs(s)
        if keep:
            self.states_list.append(s)
        return data + (s.stateseq,)


    def _generate_obs(self,s):
        if s.data is None:
            s.data = s.generate_obs()
        else:
            # TODO: Handle missing data
            raise NotImplementedError

        return s.data, s.gaussian_states

    def smooth(self, data, inputs=None, mask=None):
        self.add_data(data, inputs=inputs, mask=mask)
        s = self.states_list.pop()
        return s.smooth()

    @property
    def diagonal_noise(self):
        return all([isinstance(ed, DiagonalRegression) for ed in self.emission_distns])

    @property
    def has_missing_data(self):
        return any([s.mask is not None for s in self.states_list])


class _SLDSGibbsMixin(_SLDSMixin):
    def resample_parameters(self):
        self.resample_lds_parameters()
        self.resample_hmm_parameters()

    def resample_lds_parameters(self):
        self.resample_init_dynamics_distns()
        self.resample_dynamics_distns()
        self.resample_emission_distns()

    def resample_hmm_parameters(self):
        super(_SLDSGibbsMixin,self).resample_parameters()

    def resample_init_dynamics_distns(self):
        for state, d in enumerate(self.init_dynamics_distns):
            d.resample(
                [s.gaussian_states[0] for s in self.states_list
                    if s.stateseq[0] == state])
        self._clear_caches()

    def resample_dynamics_distns(self):
        zs = [s.stateseq[:-1] for s in self.states_list]
        xs = [np.hstack((s.gaussian_states[:-1], s.inputs[:-1]))
              for s in self.states_list]
        ys = [s.gaussian_states[1:] for s in self.states_list]

        for state, d in enumerate(self.dynamics_distns):
            d.resample(
                [(x[z == state], y[z == state])
                 for x, y, z in zip(xs, ys, zs)])
        self._clear_caches()

    def resample_emission_distns(self):
        if self._single_emission:
            data = [(np.hstack((s.gaussian_states, s.inputs)), s.data)
                    for s in self.states_list]
            mask = [s.mask for s in self.states_list] if self.has_missing_data else None

            if self.has_missing_data:
                self._emission_distn.resample(data=data, mask=mask)
            else:
                self._emission_distn.resample(data=data)
        else:
            for state, d in enumerate(self.emission_distns):
                data = [(np.hstack((s.gaussian_states[s.stateseq == state],
                                    s.inputs[s.stateseq == state])),
                         s.data[s.stateseq == state])
                        for s in self.states_list]

                mask = [s.mask[s.stateseq == state] for s in self.states_list] \
                    if self.has_missing_data else None

                if self.has_missing_data:
                    d.resample(data=data, mask=mask)
                else:
                    d.resample(data=data)


        self._clear_caches()

    def resample_obs_distns(self):
        pass  # handled in resample_parameters

    ### joblib parallel

    def _joblib_resample_states(self,states_list,num_procs):
        from joblib import Parallel, delayed
        import pyslds.parallel as parallel

        if len(states_list) > 0:
            joblib_args = list(map(self._get_joblib_pair, states_list))

            parallel.model = self
            parallel.args = list_split(joblib_args, num_procs)

            idxs = range(len(parallel.args))
            raw_stateseqs = Parallel(n_jobs=num_procs,backend='multiprocessing')\
                    (list(map(delayed(parallel._get_sampled_stateseq), idxs)))

            flatten = lambda lst: [x for y in lst for x in y]
            raw_stateseqs = flatten(raw_stateseqs)

            # since list_split might reorder things, do the same to states_list
            states_list = flatten(list_split(states_list, num_procs))
{'fraction_non_alphanumeric': 0.08563173102180194, 'fraction_numerical': 0.0017511601435951317, 'mean_word_length': 2.3193839000290613, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3575442', 'n_tokens_mistral': 3586, 'n_tokens_neox': 3351, 'n_words': 716}