<commit_before>from __future__ import division
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import cm
from warnings import warn

from pyhsmm.internals import states, initial_state, transitions

class hmm(object):
    '''
    The HMM class is a convenient wrapper that provides useful constructors and
    packages all the components.
    '''

    def __init__(self,obs_distns,alpha=None,gamma=None,**kwargs):
        self.state_dim = len(obs_distns)

        self.obs_distns = obs_distns

        if alpha is None or gamma is None:
            assert 'transitions' in kwargs, 'must specify transition distribution to initialize %s without concentration parameters' % type(self)

        self.trans_distn = transitions.hdphmm_transitions(state_dim=self.state_dim,alpha=alpha,gamma=gamma,**kwargs)\
                if 'transitions' not in kwargs else kwargs['transitions']

        self.init_state_distn = initial_state.initial_state(state_dim=self.state_dim,**kwargs)\
                if 'initial_state_distn' not in kwargs else kwargs['initial_state_distn']

        self.states_list = []

    def add_data(self,data,stateseq=None):
        self.states_list.append(states.hmm_states(len(data),self.state_dim,self.obs_distns,self.trans_distn,
                self.init_state_distn,data=data,stateseq=stateseq))

    def resample(self):
        # resample obsparams
        for state, distn in enumerate(self.obs_distns):
            # TODO make obs distns take lols
            all_obs = [s.data[s.stateseq == state] for s in self.states_list]
            if len(all_obs) > 0:
                distn.resample(np.concatenate(all_obs))
            else:
                distn.resample()

        # resample transitions
        self.trans_distn.resample([s.stateseq for s in self.states_list])

        # resample pi_0
        self.init_state_distn.resample([s.stateseq[0] for s in self.states_list])

        # resample states
        for s in self.states_list:
            s.resample()

    def generate(self,T,keep=True):
        '''
        Generates a forward sample using the current values of all parameters.
        Returns an observation sequence and a state sequence of length T.

        If keep is True, the states object created is appended to the
        states_list. This is mostly useful for generating synthetic data and
        keeping it around in an HSMM object as the latent truth.

        To construct a posterior sample, one must call both the add_data and
        resample methods first. Then, calling generate() will produce a sample
        from the posterior (as long as the Gibbs sampling has converged). In
        these cases, the keep argument should be False.
        '''
        tempstates = states.hmm_states(T,self.state_dim,self.obs_distns,self.trans_distn,
                self.init_state_distn)

        return self._generate(tempstates,keep)

    def _generate(self,tempstates,keep):
        # TODO probably already generated when tempstates was added, so this
        # call to generate only needs to be a call to generate_obs
        obs,labels = tempstates.generate(), tempstates.stateseq

        if keep:
            tempstates.added_with_generate = True # I love Python
            tempstates.data = obs
            self.states_list.append(tempstates)

        return obs, labels

    def plot(self,color=None):
        assert len(self.obs_distns) != 0
        assert len(set([type(o) for o in self.obs_distns])) == 1, \
                'plot can only be used when all observation distributions are the same class'

        # set up figure and state-color mapping dict
        fig = plt.gcf()
        fig.set_size_inches((6,6))
        state_colors = {}
        cmap = cm.get_cmap()
        used_states = reduce(set.union,[set(s.stateseq) for s in self.states_list])
        num_states = len(used_states)
        num_subfig_cols = len(self.states_list)
        for idx,state in enumerate(used_states):
            state_colors[state] = idx/(num_states-1) if color is None else color
<commit_msg>changed generate not to re-generate state sequences that are generated when the states object is instantiated anyway<commit_after>from __future__ import division
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import cm
from warnings import warn

from pyhsmm.internals import states, initial_state, transitions

class hmm(object):
    '''
    The HMM class is a convenient wrapper that provides useful constructors and
    packages all the components.
    '''

    def __init__(self,obs_distns,alpha=None,gamma=None,**kwargs):
        self.state_dim = len(obs_distns)

        self.obs_distns = obs_distns

        if alpha is None or gamma is None:
            assert 'transitions' in kwargs, 'must specify transition distribution to initialize %s without concentration parameters' % type(self)

        self.trans_distn = transitions.hdphmm_transitions(state_dim=self.state_dim,alpha=alpha,gamma=gamma,**kwargs)\
                if 'transitions' not in kwargs else kwargs['transitions']

        self.init_state_distn = initial_state.initial_state(state_dim=self.state_dim,**kwargs)\
                if 'initial_state_distn' not in kwargs else kwargs['initial_state_distn']

        self.states_list = []

    def add_data(self,data,stateseq=None):
        self.states_list.append(states.hmm_states(len(data),self.state_dim,self.obs_distns,self.trans_distn,
                self.init_state_distn,data=data,stateseq=stateseq))

    def resample(self):
        # resample obsparams
        for state, distn in enumerate(self.obs_distns):
            # TODO make obs distns take lols
            all_obs = [s.data[s.stateseq == state] for s in self.states_list]
            if len(all_obs) > 0:
                distn.resample(np.concatenate(all_obs))
            else:
                distn.resample()

        # resample transitions
        self.trans_distn.resample([s.stateseq for s in self.states_list])

        # resample pi_0
        self.init_state_distn.resample([s.stateseq[0] for s in self.states_list])

        # resample states
        for s in self.states_list:
            s.resample()

    def generate(self,T,keep=True):
        '''
        Generates a forward sample using the current values of all parameters.
        Returns an observation sequence and a state sequence of length T.

        If keep is True, the states object created is appended to the
        states_list. This is mostly useful for generating synthetic data and
        keeping it around in an HSMM object as the latent truth.

        To construct a posterior sample, one must call both the add_data and
        resample methods first. Then, calling generate() will produce a sample
        from the posterior (as long as the Gibbs sampling has converged). In
        these cases, the keep argument should be False.
        '''
        tempstates = states.hmm_states(T,self.state_dim,self.obs_distns,self.trans_distn,
                self.init_state_distn)

        return self._generate(tempstates,keep)

    def _generate(self,tempstates,keep):
        obs,labels = tempstates.generate_obs(), tempstates.stateseq

        if keep:
            tempstates.added_with_generate = True # I love Python
            tempstates.data = obs
            self.states_list.append(tempstates)

        return obs, labels

    def plot(self,color=None):
        assert len(self.obs_distns) != 0
        assert len(set([type(o) for o in self.obs_distns])) == 1, \
                'plot can only be used when all observation distributions are the same class'

        # set up figure and state-color mapping dict
        fig = plt.gcf()
        fig.set_size_inches((6,6))
        state_colors = {}
        cmap = cm.get_cmap()
        used_states = reduce(set.union,[set(s.stateseq) for s in self.states_list])
        num_states = len(used_states)
        num_subfig_cols = len(self.states_list)
        for idx,state in enumerate(used_states):
            state_colors[state] = idx/(num_states-1) if color is None else color

{'fraction_non_alphanumeric': 0.077389046853752, 'fraction_numerical': 0.0019779948077636297, 'mean_word_length': 2.7025171624713957, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1133820', 'n_tokens_mistral': 2324, 'n_tokens_neox': 2248, 'n_words': 695}