Add deterministic tests

pymc3_hmm/distributions.py

@@ -473,6 +473,13 @@ def random(self, point=None, size=None):
 
     def logp(self, values):
         """
-        Place holder for horseshoe logp as we are not needing this func
+        XXX: This is only a placeholder for the log-probability.
+        It is required to get past PyMC3 design restrictions.
+        Do not use this distribution without the `HSStep` sampler!
+
         """
-        return 1
+        warnings.warn(
+            "logp of HorseShoe class is not implemented."
+            "Do not use this distribution without the `HSStep` sampler!"
+        )
+        return 0

pymc3_hmm/step_methods.py

@@ -11,7 +11,7 @@
     from aesara.graph.op import get_test_value as test_value
     from aesara.graph.opt import OpRemove, pre_greedy_local_optimizer
     from aesara.graph.optdb import Query
-    from aesara.scalar.basic import Mul
+    from aesara.scalar.basic import Dot
     from aesara.sparse.basic import StructuredDot
     from aesara.tensor.elemwise import DimShuffle, Elemwise
     from aesara.tensor.subtensor import AdvancedIncSubtensor1
@@ -29,7 +29,7 @@
     from theano.tensor.elemwise import DimShuffle, Elemwise
     from theano.tensor.subtensor import AdvancedIncSubtensor1
     from theano.tensor.var import TensorConstant
-    from theano.scalar.basic import Mul
+    from theano.tensor.basic import Dot
     from theano.sparse.basic import StructuredDot
 
 import pymc3 as pm
@@ -541,36 +541,48 @@ def __init__(self, vars, values=None, model=None):
         non_var_rvs = [
             value for attr, value in model.named_vars.items() if value != beta
         ]
-
+        y_fn, X_fn = None, None
         # loop through all the rvs excpet for the stepping rv
         for i in non_var_rvs:
             # looking thru all the attributes of the rv and see if any of the
             # parameters are consist of multiplication relationship of the horseshoe var
             if hasattr(i, "distribution") and isinstance(i.distribution, pm.Normal):
                 mu = i.distribution.mu
             elif isinstance(i, pm.model.DeterministicWrapper):
-                mu = i
+                mu = i.owner.inputs[0]
             else:
                 continue
-            dense_dot = isinstance(mu.owner.op, Elemwise) and isinstance(
-                mu.owner.op.scalar_op, Mul
-            )
-            sparse_dot = isinstance(mu.owner.op, StructuredDot)
-            if not (
-                mu.owner
-                and (dense_dot or sparse_dot)
-                and beta in mu.owner.inputs[1].owner.inputs
-            ):
+            dense_dot = mu.owner and isinstance(mu.owner.op, Dot)
+            sparse_dot = mu.owner and isinstance(mu.owner.op, StructuredDot)
+
+            dense_inputs = dense_dot and beta in mu.owner.inputs
+            sparse_inputs = sparse_dot and beta in mu.owner.inputs[1].owner.inputs
+
+            if not (dense_inputs or sparse_inputs):
                 continue
-            if i in model.observed_RVs:
 
-                def y_fn(x):
-                    return i.observations
+            y_fn = None
+            for j in model.observed_RVs:
 
-            else:
+                if i == j or (
+                    isinstance(j.distribution, pm.Normal) and i == j.distribution.mu
+                ):
+
+                    def y_fn(x):
+                        return j.observations
+
+            if not y_fn:
                 y_fn = model.fn(i)
 
-            X_fn = model.fn(mu.owner.inputs[0])
+            if dense_inputs:
+                X_fn = model.fn(mu.owner.inputs[1].T)
+            else:
+                X_fn = model.fn(mu.owner.inputs[0])
+
+        if not (X_fn and y_fn):
+            raise RuntimeError(
+                f"Cannot find Design matrix or dependent variable assoicate with {beta}"
+            )
 
         self.vars = [beta]
 
@@ -585,6 +597,7 @@ def y_fn(x):
         self.X_fn = X_fn
 
     def step(self, point):
+        # breakpoint()
         y = self.y_fn(point)
         X = self.X_fn(point)
 

tests/test_distributions.py

@@ -576,4 +576,4 @@ def test_HorseShoe():
     assert np.percentile(test_sample.flatten(), 15) > -2
     assert np.percentile(test_sample.flatten(), 99) > 10
     assert np.percentile(test_sample.flatten(), 1) < -10
-    assert test_beta.logp(1) == 1
+    assert test_beta.logp(1) == 0

tests/test_step_methods.py

@@ -352,24 +352,30 @@ def test_large_p_mvnormal_sampler():
 
     # test case for dense matrix
     np.random.seed(2032)
-    X = np.random.choice([0, 1], size=25).reshape((5, 5))
-    y = np.array(range(5))
+    X = np.random.normal(size=250).reshape((50, 5))
+    beta_true = np.ones(5)
+    y = np.random.normal(X.dot(beta_true), 1)
 
     samples = large_p_mvnormal_sampler(np.ones(5), X, y)
     assert samples.shape == (5,)
 
+    sample_sigma = np.linalg.inv(X.T.dot(X) + np.eye(5))
+    sample_mu = sample_sigma.dot(X.T).dot(y)
+    np.testing.assert_allclose(sample_mu.mean(), 1, rtol=0.1, atol=0)
+
     # test case for sparse matrix
     samples_sp = large_p_mvnormal_sampler(np.ones(5), sp.sparse.csr_matrix(X), y)
     assert samples_sp.shape == (5,)
+    np.testing.assert_allclose(samples_sp.mean(), 1, rtol=0.1, atol=0)
 
 
 def test_hs_step():
     # test case for dense matrix
     np.random.seed(2032)
     M = 5
-    X = np.random.choice([0, 1], size=25).reshape((M, M))
+    X = np.random.normal(size=250).reshape((50, M))
     beta_true = np.random.normal(size=M)
-    y = np.random.normal(X.dot(beta_true), 1)
+    v = np.random.normal(X.dot(beta_true), 1)
 
     vi = np.full(M, 1)
     lambda2 = np.full(M, 1)
@@ -395,22 +401,56 @@ def test_hs_step():
 def test_Hsstep():
     np.random.seed(2032)
     M = 5
-    X = np.random.choice([0, 1], size=25).reshape((M, M))
+    N = 50
+    X = np.random.normal(size=N * M).reshape((N, M))
     beta_true = np.random.normal(size=M)
     y = np.random.normal(X.dot(beta_true), 1)
 
     M = X.shape[1]
     with pm.Model():
         beta = HorseShoe("beta", tau=1, shape=M)
-        pm.Normal("y", mu=X * beta, sigma=1, observed=y)
+        pm.Normal("y", mu=beta.dot(X.T), sigma=1, observed=y)
+        hsstep = HSStep([beta])
+        trace = pm.sample(
+            draws=50, tune=0, step=hsstep, chains=1, return_inferencedata=True
+        )
+
+    beta_samples = trace.posterior["beta"][0].values
+
+    assert beta_samples.shape == (50, M)
+    np.testing.assert_allclose(beta_samples.mean(0), beta_true, atol=0.3)
+
+    with pm.Model():
+        beta = HorseShoe("beta", tau=1, shape=M)
+        mu = pm.Deterministic("mu", beta.dot(X.T))
+        pm.Normal("y", mu=mu, sigma=1, observed=y)
         hsstep = HSStep([beta])
         trace = pm.sample(
-            draws=20, tune=0, step=hsstep, chains=1, return_inferencedata=True
+            draws=50, tune=0, step=hsstep, chains=1, return_inferencedata=True
         )
 
     beta_samples = trace.posterior["beta"][0].values
+    assert beta_samples.shape == (50, M)
+    np.testing.assert_allclose(beta_samples.mean(0), beta_true, atol=0.3)
 
-    assert beta_samples.shape == (20, M)
+    with pm.Model():
+        beta = HorseShoe("beta", tau=1, shape=M)
+        pm.TruncatedNormal("alpha", mu=beta.dot(X.T), sigma=1, observed=y)
+
+        with pytest.raises(RuntimeError):
+            hsstep = HSStep([beta])
+
+    with pm.Model():
+        beta = HorseShoe("beta", tau=1, shape=M)
+        mu = pm.Deterministic("mu", beta.dot(X.T))
+        pm.TruncatedNormal("y",mu = mu, sigma = 1 , observed=y)
+        hsstep = HSStep([beta])
+        trace = pm.sample(
+            draws=50, tune=0, step=hsstep, chains=1, return_inferencedata=True
+        )
+
+    beta_samples = trace.posterior["beta"][0].values
+    assert beta_samples.shape == (50, M)
 
     # test case for sparse matrix
     X = sp.sparse.csr_matrix(X)
@@ -421,12 +461,13 @@ def test_Hsstep():
         pm.Normal("y", mu=sp_dot(X, at.shape_padright(beta)), sigma=1, observed=y)
         hsstep = HSStep([beta])
         trace = pm.sample(
-            draws=20, tune=0, step=hsstep, chains=1, return_inferencedata=True
+            draws=50, tune=0, step=hsstep, chains=1, return_inferencedata=True
         )
         beta_2 = HorseShoe("beta_2", tau=1, shape=M)
         with pytest.raises(ValueError):
             HSStep([beta, beta_2])
 
     beta_samples = trace.posterior["beta"][0].values
 
-    assert beta_samples.shape == (20, M)
+    assert beta_samples.shape == (50, M)
+    np.testing.assert_allclose(beta_samples.mean(0), beta_true, atol=0.3)