<commit_before>from __future__ import division
import numpy as np

from pybasicbayes.abstractions import Model, ModelGibbsSampling, \
    ModelEM, ModelMeanField, ModelMeanFieldSVI

from states import LDSStates

# TODO make separate versions for stationary, nonstationary,
# nonstationary-and-distinct-for-each-sequence

# NOTE: dynamics_distn should probably be an instance of AutoRegression,
# emission_distn should probably be an instance of Regression, and
# init_dynamics_distn should probably be an instance of Gaussian


######################
#  algorithm mixins  #
######################

class _LDSBase(Model):
    def __init__(self,dynamics_distn,emission_distn):
        self.dynamics_distn = dynamics_distn
        self.emission_distn = emission_distn
        self.states_list = []

    def add_data(self,data,**kwargs):
        assert isinstance(data,np.ndarray)
        self.states_list.append(LDSStates(model=self,data=data,**kwargs))
        return self

    def log_likelihood(self,data=None):
        if data is not None:
            assert isinstance(data,(list,np.ndarray))
            if isinstance(data,np.ndarray):
                self.add_data(data=data,generate=False)
                return self.states_list.pop().log_likelihood()
            else:
                return sum(self.log_likelihood(d) for d in data)
        else:
            return sum(s.log_likelihood() for s in self.states_list)

    def generate(self,T,keep=True):
        s = LDSStates(model=self,T=T,initialize_from_prior=True)
        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 = self.emission_distn.rvs(x=s.stateseq,return_xy=False)
        else:
            # filling in missing data
            raise NotImplementedError
        return s.data

    def predict(self, data, Tpred):
        # return means and covariances
        raise NotImplementedError

    def sample_predictions(self, data, Tpred, states_noise=True, obs_noise=True):
        self.add_data(data, generate=False)
        s = self.states_list.pop()
        return s.sample_predictions(Tpred, states_noise, obs_noise)

    # convenience properties

    @property
    def n(self):
        'latent dimension'
        return self.emission_distn.D_in

    @property
    def p(self):
        'emission dimension'
        return self.emission_distn.D_out

    @property
    def mu_init(self):
        return np.zeros(self.n) if not hasattr(self,'_mu_init') \
            else self._mu_init

    @mu_init.setter
    def mu_init(self,mu_init):
        self._mu_init = mu_init

    @property
    def sigma_init(self):
        if hasattr(self,'_sigma_init'):
            return self._sigma_init

        try:
            import pydare
            return pydare.dlyap(self.A, self.sigma_states)
        except ImportError:<commit_msg>replace dependency on pydare+slycot with scipy.linalg.solve_discrete_lyapunov<commit_after>from __future__ import division
import numpy as np

from pybasicbayes.abstractions import Model, ModelGibbsSampling, \
    ModelEM, ModelMeanField, ModelMeanFieldSVI

from states import LDSStates

# TODO make separate versions for stationary, nonstationary,
# nonstationary-and-distinct-for-each-sequence

# NOTE: dynamics_distn should probably be an instance of AutoRegression,
# emission_distn should probably be an instance of Regression, and
# init_dynamics_distn should probably be an instance of Gaussian


######################
#  algorithm mixins  #
######################

class _LDSBase(Model):
    def __init__(self,dynamics_distn,emission_distn):
        self.dynamics_distn = dynamics_distn
        self.emission_distn = emission_distn
        self.states_list = []

    def add_data(self,data,**kwargs):
        assert isinstance(data,np.ndarray)
        self.states_list.append(LDSStates(model=self,data=data,**kwargs))
        return self

    def log_likelihood(self,data=None):
        if data is not None:
            assert isinstance(data,(list,np.ndarray))
            if isinstance(data,np.ndarray):
                self.add_data(data=data,generate=False)
                return self.states_list.pop().log_likelihood()
            else:
                return sum(self.log_likelihood(d) for d in data)
        else:
            return sum(s.log_likelihood() for s in self.states_list)

    def generate(self,T,keep=True):
        s = LDSStates(model=self,T=T,initialize_from_prior=True)
        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 = self.emission_distn.rvs(x=s.stateseq,return_xy=False)
        else:
            # filling in missing data
            raise NotImplementedError
        return s.data

    def predict(self, data, Tpred):
        # return means and covariances
        raise NotImplementedError

    def sample_predictions(self, data, Tpred, states_noise=True, obs_noise=True):
        self.add_data(data, generate=False)
        s = self.states_list.pop()
        return s.sample_predictions(Tpred, states_noise, obs_noise)

    # convenience properties

    @property
    def n(self):
        'latent dimension'
        return self.emission_distn.D_in

    @property
    def p(self):
        'emission dimension'
        return self.emission_distn.D_out

    @property
    def mu_init(self):
        return np.zeros(self.n) if not hasattr(self,'_mu_init') \
            else self._mu_init

    @mu_init.setter
    def mu_init(self,mu_init):
        self._mu_init = mu_init

    @property
    def sigma_init(self):
        if hasattr(self,'_sigma_init'):
            return self._sigma_init

        try:
            from scipy.linalg import solve_discrete_lyapunov as dtlyap
            return dtlyap(self.A, self.sigma_states)
        except ImportError:
{'fraction_non_alphanumeric': 0.09604142308334725, 'fraction_numerical': 0.0, 'mean_word_length': 2.747183979974969, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1725159', 'n_tokens_mistral': 1862, 'n_tokens_neox': 1776, 'n_words': 406}