allow for interger mu in simulate_poiszero_hmm

AmpersandTV · Jun 25, 2020 · f761419 · f761419
1 parent cd02908
commit f761419
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/tests/utils.py b/tests/utils.py
@@ -1,21 +1,18 @@
-import numpy as np
+from pymc3_hmm.distributions import HMMStateSeq, SwitchingProcess
 
+import numpy as np
 import theano.tensor as tt
-
 import pymc3 as pm
-
-from pymc3_hmm.distributions import HMMStateSeq, SwitchingProcess
-
 import pandas as pd
+import numbers
 
 
 def simulate_poiszero_hmm(N, mus=np.r_[10.0, 30.0],
-                          pi_0_a=np.r_[1, 1, 1],
-                          Gamma=np.r_['0,2,1',
-                                      [5, 1, 1],
-                                      [1, 3, 1],
-                                      [1, 1, 5]]
+                          pi_0_a=np.r_[1, 1],
+                          Gamma=np.r_['0,2', [5, 1], [1, 3]]
                           ):
+    if isinstance(mus, numbers.Number):
+        mus = np.r_[mus]
     assert pi_0_a.size == mus.size + 1 == Gamma.shape[0] == Gamma.shape[1]
 
     with pm.Model() as test_model: