Add unit tests to check initial state probability estimation

tests/test_estimation.py

@@ -1,34 +1,31 @@
-import warnings
-
 import numpy as np
 import pymc3 as pm
-import pytest
-import scipy as sp
 import theano.tensor as tt
-from theano.graph.op import get_test_value
 
 from pymc3_hmm.distributions import DiscreteMarkovChain, PoissonZeroProcess
-from pymc3_hmm.step_methods import FFBSStep, TransMatConjugateStep, ffbs_astep
-from pymc3_hmm.utils import compute_steady_state, compute_trans_freqs
+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
+def test_gamma_0_estimation():  # annotate this test with pytest
 
     np.random.seed(2032)
 
-    poiszero_sim, _ = simulate_poiszero_hmm(30, 150, pi_0 = [.5, .5]) # test [.01, .99], [.99, .01]
+    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:
-        p_0_rv = np.array([.5, .5])
-        p_1_rv = np.array([.5, .5])
+        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_tt = tt.stack([p_0_rv, p_1_rv])  # transition matrix
         P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
 
         # how well is this estimated?
-        pi_0_rv =  pm.Dirichlet("pi_0", np.r_[1,1]) # "flat" prior for initial states
+        pi_0_rv = pm.Dirichlet("pi_0", np.r_[1, 1])  # "flat" prior for initial states
 
         S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_rv, shape=y_test.shape[0])
 
@@ -41,4 +38,4 @@ def test_gamma_0_estimation(): # annotate this test with pytest
     # grab pi_0 values from trace, compare to pi_0 value specifiedd in
     # simulate_poiszero_hmm
 
-    # 
+    #

tests/test_utils.py

@@ -3,9 +3,11 @@
 import scipy as sp
 import theano
 import theano.tensor as tt
+from aesara.tensor.var import TensorVariable
 
 from pymc3_hmm.utils import (
-    compute_steady_state, compute_trans_freqs,
+    compute_steady_state,
+    compute_trans_freqs,
     logdotexp,
     multilogit_inv,
     tt_logdotexp,
@@ -121,18 +123,10 @@ def test_multilogit_inv(test_input, test_output):
     res = res.eval()
     assert np.array_equal(res.round(2), test_output)
 
+
 def test_compute_steady_state():
-
-    P = tt.as_tensor_variable(
-        np.arange(9).reshape(3, 3).astype(np.float32))
-
-    ss_probs = compute_steady_state(P)
-    ss_probs.
-    # computing a steady state to a markov chain
-    # T = transition prob matrix 
-    # v_i = vector of initial probability states
-    # --> v_i * T "converges" to a steady state
 
-    # transition matrix option: 50-50 transitions on the diagonal
-    pass
+    P = tt.as_tensor_variable(np.arange(4).reshape(2, 2).astype(np.float32))
 
+    ss_probs = compute_steady_state(P)
+    assert isinstance(ss_probs, TensorVariable)

tests/utils.py

@@ -6,7 +6,7 @@
 
 
 def simulate_poiszero_hmm(
-    N, mu=10.0, pi_0_a=np.r_[1, 1], p_0_a=np.r_[5, 1], p_1_a=np.r_[1, 1], pi_0 = None
+    N, mu=10.0, pi_0_a=np.r_[1, 1], p_0_a=np.r_[5, 1], p_1_a=np.r_[1, 1], pi_0=None
 ):
 
     with pm.Model() as test_model:
@@ -16,7 +16,6 @@ def simulate_poiszero_hmm(
         P_tt = tt.stack([p_0_rv, p_1_rv])
         P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
 
-
         if pi_0:
             pi_0_tt = tt.as_tensor(pi_0)