<commit_before>#!/usr/bin/env python
"""Testing functions for FirstOrderHMMs. Viterbi decoding, sequence generation,
the forward-backward algorithm, posterior decoding, and re-training are all tested. 


"""
import unittest
import itertools
import numpy
import scipy.stats
from nose.tools import assert_greater_equal
from numpy.testing import (assert_array_equal,
                           assert_array_almost_equal,
                           assert_almost_equal)

from minihmm.hmm import FirstOrderHMM
from minihmm.estimators import (UnivariateGaussianEmissionEstimator,
                                DiscreteStatePriorEstimator,
                                DiscreteEmissionEstimator,
                                DiscreteTransitionEstimator,
                                PseudocountTransitionEstimator,
                                )
from minihmm.training import train_baum_welch, neg_exp_noise_gen
from minihmm.factors import (ArrayFactor,
                             MatrixFactor,
                             LogFunctionFactor,
                             ScipyDistributionFactor)



_FORWARD_SEED  = 5139284
_TRAINING_SEED = 134067



class _BaseExample():

    # define these variables in this method in subclasses    
    @classmethod
    def do_subclass_setup(cls):
        cls.name = ""
        cls.min_frac_equal = 0.8
        cls.state_prior_estimator = None
        cls.transition_estimator  = None
        cls.emission_estimator    = None
        cls.models = {
            "generating" : {
                "trans_probs"      : None,
                "state_priors"     : None,
                "emission_probs"   : None,
            },
            "naive"      : {
                "trans_probs"      : None,
                "state_priors"     : None,
                "emission_probs"   : None,
            }
        }

    @classmethod
    def setUpClass(cls):


        cls.test_obs_seqs   = []    # test observation sequencs
        cls.test_state_seqs = []    # state paths through test observation seqs. not currently used
        cls.decode_tests    = []    # tests for Viterbi and posterior decoding

        cls.expected_forward_logprobs = [] # log probabilities for cls.test_obs_seqs
        cls.found_forward_logprobs    = []
        cls.forward_scalefactors      = []


        cls.seq_lengths = [5,10]

        print("Setting up class %s" % cls.__name__)
        cls.do_subclass_setup()
        cls.generating_hmm = FirstOrderHMM(**cls.models["generating"])
        cls.naive_hmm      = FirstOrderHMM(**cls.models["naive"])

        print("    Generating sequences ...")
        numpy.random.seed(_TRAINING_SEED)


        for x in cls.seq_lengths:
            states, obs, _ = cls.generating_hmm.generate(x)
            cls.test_state_seqs.append(states)
            cls.test_obs_seqs.append(obs)

        # run forward algorithm; we'll use their results for tests below
        for obs in cls.test_obs_seqs:
            logprob, scaled_forward, scaled_backward, scale_factors, ksi = cls.generating_hmm.forward_backward(obs)
            cls.found_forward_logprobs.append(logprob)
            cls.forward_scalefactors.append(scale_factors)


        # evaluate probabilities of generated observations by brute force
        # to compare against forward algorithm calculations later
        print("    Calculating probabilities ...")
        for obs_seq in cls.test_obs_seqs:
            state_paths = itertools.product(list(range(cls.generating_hmm.num_states)),
                                            repeat=len(obs_seq))
            total_prob  = 0
            for my_path in state_paths:
                my_logprob = []
                last_state = my_path[0]
                my_logprob.append(cls.generating_hmm.state_priors.logprob(last_state))
                my_logprob.append(cls.generating_hmm.emission_probs[last_state].logprob(obs_seq[0]))
                for my_state, my_obs in zip(my_path, obs_seq)[1:]:
                    my_logprob.append(cls.generating_hmm.trans_probs.logprob(last_state,my_state))
                    my_logprob.append(cls.generating_hmm.emission_probs[my_state].logprob(my_obs))
                    last_state = my_state
                
                total_prob += numpy.exp(numpy.nansum(my_logprob))
            
            cls.expected_forward_logprobs.append(numpy.log(total_prob))

        # Generate test cases for Viterbi and posterior decoding            
        cls.decode_tests = [cls.generating_hmm.generate(1000) for _ in range(10)]
            
        print("Set up class %s" % cls.__name__)

    def test_generate(self):
        assert False

    @unittest.skip
    def test_sample(self):
        # TODO : not sure what proper test is; distribution of samples
        # should approximate distribution of HMM, but we don't know what
        # distribution of HMM actually is.
        
        #
        # Maybe calculate ML solution (Viterbi) and make sure everything sampled
        # is lower probability?
        #
        # Some test of shape of distribution? Require one mode? 
        #
        # Or, sample from a path through a known sequence of e.g. coin tosses,
        # where we can calculate the background distribution, and then make
        # sure the sampled distribution approximates the background?
        assert False
        
    def test_viterbi(self):
        # make sure viterbi calls are above accuracy threshold listed above 
        for expected_states, obs, _ in self.decode_tests:
            found_states = self.generating_hmm.viterbi(obs)["viterbi_states"]
            frac_equal = 1.0 * (expected_states == found_states).sum() / len(expected_states)
            msg = "Failed viterbi test for test case '%s'. Expected at least %s%% accuracy. Got %s%%." % (self.name,self.min_frac_equal,frac_equal)
            assert_greater_equal(frac_equal, self.min_frac_equal, msg)
         
    def test_posterior_decode(self):
        # make sure posterior decode calls are above accuracy threshold listed above
        for expected_states, obs, _ in self.decode_tests:
            found_states = self.generating_hmm.posterior_decode(obs)[0]
            frac_equal = 1.0 * (expected_states == found_states).sum() / len(expected_states)
            msg = "Failed posterior decode test for test case '%s'. Expected at least %s%% accuracy. Got %s%%." % (self.name,self.min_frac_equal,frac_equal)
            assert_greater_equal(frac_equal, self.min_frac_equal, msg)

    def test_forward_logprob(self):
        # make sure vectorized forward probability calculations match those calced by brute force
        numpy.random.seed(_FORWARD_SEED)
        for n, (expected, found) in enumerate(zip(self.expected_forward_logprobs, self.found_forward_logprobs)):
            msg = "Failed forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                  self.name,
                                                                                                  expected,
                                                                                                  found,
                                                                                                  abs(expected-found)
                                                                                                 )
            yield assert_almost_equal, expected, found, 7, msg

    def test_fast_forward(self):
        # make sure fast forward probability calculations match those calced by brute force
        numpy.random.seed(_FORWARD_SEED)
        for n, (obs, expected) in enumerate(zip(self.test_obs_seqs, self.expected_forward_logprobs)):
            found = self.generating_hmm.fast_forward(obs)
            msg = "Failed fast_forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                 self.name,
                                                                                                 expected,
                                                                                                 found,
                                                                                                 abs(expected-found)
                                                                                                )
            yield assert_almost_equal, expected, found, 7, msg

    def test_forward_backward(self):
        # test forward algorithm portion of forward_backward
        
        # TODO: test backward component
        numpy.random.seed(_FORWARD_SEED)
        for n, (obs, expected_logprob) in enumerate(zip(self.test_obs_seqs, self.expected_forward_logprobs)):
            (found_logprob,
             scaled_forward,
             scaled_backward,
             scale_factors,
             ksi)     = self.generating_hmm.forward_backward(obs)
            msg = "Failed fast_forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                 self.name,
                                                                                                 expected_logprob,
                                                                                                 found_logprob,
                                                                                                 abs(expected_logprob-found_logprob)
                                                                                                )
            yield assert_almost_equal, expected_logprob, found_logprob, 7, msg

#     def test_train(self):
#         mdict = train_baum_welch(self.naive_hmm,
#                                  self.observations,
#                                  state_prior_estimator = self.state_prior_estimator,
#                                  transition_estimator = self.transition_estimator,
#                                  emission_estimator   = self.emission_estimator,
#                                  noise_weights = neg_exp_noise_gen(),
# #                                 miniter = 200,
# #                                 pseudocount_weights  = iter([0]),
#                                  )
#         new_model = mdict["best_model"]
#         print(mdict)
# 
#         yield assert_array_almost_equal, self.generating_hmm.trans_probs.data, new_model.trans_probs.data
#         for expected, found in zip(self.generating_hmm.emission_probs, new_model.emission_probs):
#             if expected is not None:
#                 yield assert_array_almost_equal, expected.data, found.data



class TestACoin(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Coin example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        cls.emission_estimator    = DiscreteEmissionEstimator(2)
 
#        cls.transition_estimator = PseudocountTransitionEstimator(numpy.array([
#            [ 0 ,  0,  1,  1],
#            [ 0, 1.0,  0,  0],
#            [ 0,   0,  1,  1],
#            [ 0,   0,  1,  1]
#            ]))

        cls.models = {
            "generating" : {
                "state_priors"     : ArrayFactor([0.005,0.995]),
                "trans_probs"      : MatrixFactor(numpy.array([
                                                    [0.8, 0.2],
                                                    [0.3, 0.7]])),
                "emission_probs"   : [ArrayFactor([0.6,0.4]),
                                      ArrayFactor([0.15,0.85])],
            },
            "naive"      : {
                "state_priors"     : ArrayFactor([0.48,0.52]),
                "trans_probs"      : MatrixFactor(numpy.array([
                                                    [0.5, 0.5],
                                                    [0.5, 0.5]])),
                "emission_probs"   : [ArrayFactor([0.5,0.5]),
                                      ArrayFactor([0.5,0.5])],
             }
        }



class TestBTwoGaussian(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Two gaussian example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        cls.emission_estimator    = UnivariateGaussianEmissionEstimator()

        transitions = numpy.matrix([[0.9,0.1],
                                    [0.25,0.75]])

        cls.models = {
            "generating" : {
                "trans_probs"    : MatrixFactor(transitions),
                "state_priors"   : ArrayFactor([0.8,0.2]),
                "emission_probs" : [ScipyDistributionFactor(scipy.stats.norm,loc=0,scale=0.5),
                                    ScipyDistributionFactor(scipy.stats.norm,loc=5,scale=10)],
            },
            "naive"      : {
                "emission_probs" : [ScipyDistributionFactor(scipy.stats.norm,loc=0,scale=0.6),
                                    ScipyDistributionFactor(scipy.stats.norm,loc=2,scale=1),
                                   ],
                "state_priors"   : ArrayFactor([0.48,0.52]),
                "trans_probs"    : MatrixFactor([[0.5, 0.5], 
                                                 [0.5, 0.5]]),
            }
        }

    

class TestCFourPoisson(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Four-state Poisson example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        
        # FIXME
        cls.emission_estimator    = UnivariateGaussianEmissionEstimator()

        transitions = numpy.matrix([[0.8,0.05,0.05,0.1],
                                    [0.2,0.6,0.1,0.1],
                                    [0.01,0.97,0.01,0.01],
                                    [0.45,0.01,0.04,0.5],
                                   ])

        cls.models = {
            "generating" : {
                "trans_probs"    : MatrixFactor(transitions),
                "state_priors"   : ArrayFactor([0.7,0.05,0.15,0.10]),
                "emission_probs" : [
                      ScipyDistributionFactor(scipy.stats.poisson,1),
                      ScipyDistributionFactor(scipy.stats.poisson,5),
                      ScipyDistributionFactor(scipy.stats.poisson,10),
                      ScipyDistributionFactor(scipy.stats.poisson,25),
                 ]

            },
            "naive"      : {
                "trans_probs"    : MatrixFactor(numpy.full((4,4),1.0/16)),
                "state_priors"   : ArrayFactor([0.25,0.25,0.25,0.25]),
                "emission_probs" : [
                      ScipyDistributionFactor(scipy.stats.poisson,3),
                      ScipyDistributionFactor(scipy.stats.poisson,2),
                      ScipyDistributionFactor(scipy.stats.poisson,1),
                      ScipyDistributionFactor(scipy.stats.poisson,4),
                 ]
            }
        }
<commit_msg>changed test calcs from brute force to dynamic programming to assist speed<commit_after>#!/usr/bin/env python
"""Testing functions for FirstOrderHMMs. Viterbi decoding, sequence generation,
the forward-backward algorithm, posterior decoding, and re-training are all tested. 


"""
import unittest
import itertools
import numpy
import scipy.stats
from nose.tools import assert_greater_equal
from numpy.testing import (assert_array_equal,
                           assert_array_almost_equal,
                           assert_almost_equal)

from minihmm.hmm import FirstOrderHMM
from minihmm.estimators import (UnivariateGaussianEmissionEstimator,
                                DiscreteStatePriorEstimator,
                                DiscreteEmissionEstimator,
                                DiscreteTransitionEstimator,
                                PseudocountTransitionEstimator,
                                )
from minihmm.training import train_baum_welch, neg_exp_noise_gen
from minihmm.factors import (ArrayFactor,
                             MatrixFactor,
                             LogFunctionFactor,
                             ScipyDistributionFactor)



_FORWARD_SEED  = 5139284
_TRAINING_SEED = 134067



class _BaseExample():

    # define these variables in this method in subclasses    
    @classmethod
    def do_subclass_setup(cls):
        cls.name = ""
        cls.min_frac_equal = 0.8
        cls.state_prior_estimator = None
        cls.transition_estimator  = None
        cls.emission_estimator    = None
        cls.models = {
            "generating" : {
                "trans_probs"      : None,
                "state_priors"     : None,
                "emission_probs"   : None,
            },
            "naive"      : {
                "trans_probs"      : None,
                "state_priors"     : None,
                "emission_probs"   : None,
            }
        }

    @classmethod
    def setUpClass(cls):


        cls.test_obs_seqs   = []    # test observation sequences
        cls.test_state_seqs = []    # state paths through test observation seqs. not currently used
        cls.decode_tests    = []    # tests for Viterbi and posterior decoding

        cls.expected_forward_logprobs = [] # log probabilities for cls.test_obs_seqs
        cls.found_forward_logprobs    = []
        cls.forward_scalefactors      = []

        cls.expected_joint_logprobs = [] # expected joint probabilities for states and observations


        cls.seq_lengths = [5,10]

        print("Setting up class %s" % cls.__name__)
        cls.do_subclass_setup()
        hmm = FirstOrderHMM(**cls.models["generating"])
        cls.generating_hmm = hmm
        cls.naive_hmm      = FirstOrderHMM(**cls.models["naive"])

        print("    Generating sequences ...")
        numpy.random.seed(_TRAINING_SEED)


        for x in cls.seq_lengths:
            states, obs, _ = hmm.generate(x)
            cls.test_state_seqs.append(states)
            cls.test_obs_seqs.append(obs)

        # precalculate useful quantities on observation sequences
        for obs in cls.test_obs_seqs:
            # run forward algorithm; we'll use their results for tests below
            logprob, scaled_forward, scaled_backward, scale_factors, ksi = hmm.forward_backward(obs)
            cls.found_forward_logprobs.append(logprob)
            cls.forward_scalefactors.append(scale_factors)

            # manually calculate probabilities of generated observations via dynamic programming
            # the forward probability of an observation sequence is equal to the sum of its joint 
            # probabilities with all possible state paths. We calculate each joint probability here,
            # then sum these to find the total expected probability
            my_states = list(range(hmm.num_states))
            paths = { (X,) : hmm.state_priors.logprob(X) + hmm.emission_probs[X].logprob(obs[0]) for X in my_states }

            for my_obs in obs[1:]:
                new_paths = {}
                for partial, prob in paths.items():
                    for new_state in my_states:
                        new_paths[ tuple(list(partial) + [new_state])] = prob + hmm.trans_probs.logprob(partial[-1], new_state) + hmm.emission_probs[new_state].logprob(my_obs)
                        
                paths = new_paths

            dp_prob = numpy.log(numpy.exp(paths.values()).sum())
            cls.expected_forward_logprobs.append(dp_prob)
            cls.expected_joint_logprobs.append(paths)

#         # evaluate probabilities of generated observations by brute force
#         # to compare against forward algorithm calculations later
#         print("    Calculating probabilities ...")
#         for obs_seq in cls.test_obs_seqs:
#             state_paths = itertools.product(list(range(cls.generating_hmm.num_states)),
#                                             repeat=len(obs))
#             total_prob  = 0
#             for my_path in state_paths:
#                 my_logprob = []
#                 last_state = my_path[0]
#                 my_logprob.append(cls.generating_hmm.state_priors.logprob(last_state))
#                 my_logprob.append(cls.generating_hmm.emission_probs[last_state].logprob(obs[0]))
#                 for my_state, my_obs in zip(my_path, obs)[1:]:
#                     my_logprob.append(cls.generating_hmm.trans_probs.logprob(last_state,my_state))
#                     my_logprob.append(cls.generating_hmm.emission_probs[my_state].logprob(my_obs))
#                     last_state = my_state
#                 
#                 total_prob += numpy.exp(numpy.nansum(my_logprob))
#             
#            cls.expected_forward_logprobs.append(numpy.log(total_prob))

        # Generate test cases for Viterbi and posterior decoding            
        cls.decode_tests = [hmm.generate(1000) for _ in range(10)]
            
        print("Set up class %s" % cls.__name__)

    @unittest.skip
    def test_generate(self):
        assert False

    @unittest.skip
    def test_sample(self):
        # TODO : not sure what proper test is; distribution of samples
        # should approximate distribution of HMM, but we don't know what
        # distribution of HMM actually is.
        
        #
        # Maybe calculate ML solution (Viterbi) and make sure everything sampled
        # is lower probability?
        #
        # Some test of shape of distribution? Require one mode? 
        #
        # Or, sample from a path through a known sequence of e.g. coin tosses,
        # where we can calculate the background distribution, and then make
        # sure the sampled distribution approximates the background?
        assert False
        
    def test_viterbi(self):
        # make sure viterbi calls are above accuracy threshold listed above 
        for expected_states, obs, _ in self.decode_tests:
            found_states = self.generating_hmm.viterbi(obs)["viterbi_states"]
            frac_equal = 1.0 * (expected_states == found_states).sum() / len(expected_states)
            msg = "Failed viterbi test for test case '%s'. Expected at least %s%% accuracy. Got %s%%." % (self.name,self.min_frac_equal,frac_equal)
            assert_greater_equal(frac_equal, self.min_frac_equal, msg)
         
    def test_posterior_decode(self):
        # make sure posterior decode calls are above accuracy threshold listed above
        for expected_states, obs, _ in self.decode_tests:
            found_states = self.generating_hmm.posterior_decode(obs)[0]
            frac_equal = 1.0 * (expected_states == found_states).sum() / len(expected_states)
            msg = "Failed posterior decode test for test case '%s'. Expected at least %s%% accuracy. Got %s%%." % (self.name,self.min_frac_equal,frac_equal)
            assert_greater_equal(frac_equal, self.min_frac_equal, msg)

    def test_forward_logprob(self):
        # make sure vectorized forward probability calculations match those calced by brute force
        numpy.random.seed(_FORWARD_SEED)
        for n, (expected, found) in enumerate(zip(self.expected_forward_logprobs, self.found_forward_logprobs)):
            msg = "Failed forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                  self.name,
                                                                                                  expected,
                                                                                                  found,
                                                                                                  abs(expected-found)
                                                                                                 )
            yield assert_almost_equal, expected, found, 7, msg

    def test_fast_forward(self):
        # make sure fast forward probability calculations match those calced by brute force
        numpy.random.seed(_FORWARD_SEED)
        for n, (obs, expected) in enumerate(zip(self.test_obs_seqs, self.expected_forward_logprobs)):
            found = self.generating_hmm.fast_forward(obs)
            msg = "Failed fast_forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                 self.name,
                                                                                                 expected,
                                                                                                 found,
                                                                                                 abs(expected-found)
                                                                                                )
            yield assert_almost_equal, expected, found, 7, msg

    def test_forward_backward(self):
        # test forward algorithm portion of forward_backward
        
        # TODO: test backward component
        numpy.random.seed(_FORWARD_SEED)
        for n, (obs, expected_logprob) in enumerate(zip(self.test_obs_seqs, self.expected_forward_logprobs)):
            (found_logprob,
             scaled_forward,
             scaled_backward,
             scale_factors,
             ksi)     = self.generating_hmm.forward_backward(obs)
            msg = "Failed fast_forward test case '%s' on HMM '%s'. Expected: '%s'. Found '%s'. Diff: '%s'." % (n,
                                                                                                 self.name,
                                                                                                 expected_logprob,
                                                                                                 found_logprob,
                                                                                                 abs(expected_logprob-found_logprob)
                                                                                                )
            yield assert_almost_equal, expected_logprob, found_logprob, 7, msg

#     def test_train(self):
#         mdict = train_baum_welch(self.naive_hmm,
#                                  self.observations,
#                                  state_prior_estimator = self.state_prior_estimator,
#                                  transition_estimator = self.transition_estimator,
#                                  emission_estimator   = self.emission_estimator,
#                                  noise_weights = neg_exp_noise_gen(),
# #                                 miniter = 200,
# #                                 pseudocount_weights  = iter([0]),
#                                  )
#         new_model = mdict["best_model"]
#         print(mdict)
# 
#         yield assert_array_almost_equal, self.generating_hmm.trans_probs.data, new_model.trans_probs.data
#         for expected, found in zip(self.generating_hmm.emission_probs, new_model.emission_probs):
#             if expected is not None:
#                 yield assert_array_almost_equal, expected.data, found.data



class TestACoin(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Coin example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        cls.emission_estimator    = DiscreteEmissionEstimator(2)
 
#        cls.transition_estimator = PseudocountTransitionEstimator(numpy.array([
#            [ 0 ,  0,  1,  1],
#            [ 0, 1.0,  0,  0],
#            [ 0,   0,  1,  1],
#            [ 0,   0,  1,  1]
#            ]))

        cls.models = {
            "generating" : {
                "state_priors"     : ArrayFactor([0.005,0.995]),
                "trans_probs"      : MatrixFactor(numpy.array([
                                                    [0.8, 0.2],
                                                    [0.3, 0.7]])),
                "emission_probs"   : [ArrayFactor([0.6,0.4]),
                                      ArrayFactor([0.15,0.85])],
            },
            "naive"      : {
                "state_priors"     : ArrayFactor([0.48,0.52]),
                "trans_probs"      : MatrixFactor(numpy.array([
                                                    [0.5, 0.5],
                                                    [0.5, 0.5]])),
                "emission_probs"   : [ArrayFactor([0.5,0.5]),
                                      ArrayFactor([0.5,0.5])],
             }
        }



class TestBTwoGaussian(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Two gaussian example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        cls.emission_estimator    = UnivariateGaussianEmissionEstimator()

        transitions = numpy.matrix([[0.9,0.1],
                                    [0.25,0.75]])

        cls.models = {
            "generating" : {
                "trans_probs"    : MatrixFactor(transitions),
                "state_priors"   : ArrayFactor([0.8,0.2]),
                "emission_probs" : [ScipyDistributionFactor(scipy.stats.norm,loc=0,scale=0.5),
                                    ScipyDistributionFactor(scipy.stats.norm,loc=5,scale=10)],
            },
            "naive"      : {
                "emission_probs" : [ScipyDistributionFactor(scipy.stats.norm,loc=0,scale=0.6),
                                    ScipyDistributionFactor(scipy.stats.norm,loc=2,scale=1),
                                   ],
                "state_priors"   : ArrayFactor([0.48,0.52]),
                "trans_probs"    : MatrixFactor([[0.5, 0.5], 
                                                 [0.5, 0.5]]),
            }
        }

    

class TestCFourPoisson(_BaseExample):

    @classmethod
    def do_subclass_setup(cls):
        cls.name = "Four-state Poisson example"
        cls.min_frac_equal = 0.7

        cls.state_prior_estimator = DiscreteStatePriorEstimator()
        cls.transition_estimator  = DiscreteTransitionEstimator()
        
        # FIXME
        cls.emission_estimator    = UnivariateGaussianEmissionEstimator()

        transitions = numpy.matrix([[0.8,0.05,0.05,0.1],
                                    [0.2,0.6,0.1,0.1],
                                    [0.01,0.97,0.01,0.01],
                                    [0.45,0.01,0.04,0.5],
                                   ])

        cls.models = {
            "generating" : {
                "trans_probs"    : MatrixFactor(transitions),
                "state_priors"   : ArrayFactor([0.7,0.05,0.15,0.10]),
                "emission_probs" : [
                      ScipyDistributionFactor(scipy.stats.poisson,1),
                      ScipyDistributionFactor(scipy.stats.poisson,5),
                      ScipyDistributionFactor(scipy.stats.poisson,10),
                      ScipyDistributionFactor(scipy.stats.poisson,25),
                 ]

            },
            "naive"      : {
                "trans_probs"    : MatrixFactor(numpy.full((4,4),1.0/16)),
                "state_priors"   : ArrayFactor([0.25,0.25,0.25,0.25]),
                "emission_probs" : [
                      ScipyDistributionFactor(scipy.stats.poisson,3),
                      ScipyDistributionFactor(scipy.stats.poisson,2),
                      ScipyDistributionFactor(scipy.stats.poisson,1),
                      ScipyDistributionFactor(scipy.stats.poisson,4),
                 ]
            }
        }

{'fraction_non_alphanumeric': 0.08512698759801318, 'fraction_numerical': 0.014526256599293992, 'mean_word_length': 1.40565116104306, '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': 178, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7024494', 'n_tokens_mistral': 8815, 'n_tokens_neox': 8178, 'n_words': 1618}