Test initial state probability estimation

tests/test_estimation.py

@@ -0,0 +1,59 @@
+import numpy as np
+import pymc3 as pm
+import theano.tensor as tt
+
+from pymc3_hmm.distributions import DiscreteMarkovChain, SwitchingProcess
+
+# from pymc3_hmm.step_methods import FFBSStep
+from tests.utils import simulate_poiszero_hmm
+
+
+# the way we use this initial probability state vector
+def test_gamma_0_estimation():  # annotate this test with pytest
+
+    np.random.seed(2032)
+
+    poiszero_sim, _ = simulate_poiszero_hmm(
+        30, 150, pi_0=[0.5, 0.5]
+    )  # test [.01, .99], [.99, .01]
+    y_test = poiszero_sim["Y_t"]
+
+    with pm.Model() as test_model:
+
+        # create a 2x2 matrix with 50% transition probs for both states
+        p_0_rv = np.array([0.5, 0.5])
+        p_1_rv = np.array([0.5, 0.5])
+        P_tt = tt.stack([p_0_rv, p_1_rv])  # transition matrix
+        P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
+
+        # pi_0 = gamma_0 = initial state
+        pi_0_rv = pm.Dirichlet("pi_0", np.r_[1, 1])  # "flat" prior for initial states
+
+        # take the transition matrix and the flat prior for initial states
+        # creates a sequence of states and it "chooses" the state and time t based
+        # on how the trans matrix P_rv is defined and
+        S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_rv, shape=y_test.shape[0])
+
+        # takes series of states
+        Y_t = SwitchingProcess("Y_t", [pm.Constant.dist(0), pm.Constant.dist(1)], S_rv)
+
+    # simulate fake data
+    sim_dat = np.random.binomial(1, 0.5, size=y_test.shape[0])
+
+    with test_model:
+
+        prior_predictive = pm.sample_prior_predictive(samples=100)
+        assert 0.49 < prior_predictive["Y_t"].mean() < 0.51
+
+        # use model to get a posterior distribution
+        trace = pm.sample()
+
+        # condition model on fake data?
+        posterior_predictive = pm.sample_posterior_predictive(trace, sim_dat)
+
+        assert 0.49 < posterior_predictive["Y_t"].mean() < 0.51
+
+    # ffbs = FFBSStep([S_rv])
+    # trace = pm.sample()
+    # S_rv = pm.Constant("S_t", P_rv, pi_0_rv, shape=y_test.shape[0])
+    # pm.Constant("Y_t", 150, S_rv, observed=y_test)