hs gibbs

pymc3_hmm/distributions.py

@@ -473,8 +473,6 @@ def random(self, point=None, size=None):
 
     def logp(self, values):
         """
-        Place holder for horseshoe logp as we are not be needing this func
+        Place holder for horseshoe logp as we are not needing this func
         """
         return 1
-
-

pymc3_hmm/step_methods.py

@@ -11,9 +11,12 @@
     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.sparse.basic import StructuredDot
     from aesara.tensor.elemwise import DimShuffle, Elemwise
     from aesara.tensor.subtensor import AdvancedIncSubtensor1
     from aesara.tensor.var import TensorConstant
+
 except ImportError:  # pragma: no cover
     import theano.scalar as aes
     import theano.tensor as at
@@ -26,6 +29,8 @@
     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.sparse.basic import StructuredDot
 
 import pymc3 as pm
 import scipy
@@ -34,9 +39,10 @@
 from pymc3.util import get_untransformed_name
 from scipy.stats import invgamma
 
-from pymc3_hmm.distributions import DiscreteMarkovChain, SwitchingProcess
+from pymc3_hmm.distributions import DiscreteMarkovChain, HorseShoe, SwitchingProcess
 from pymc3_hmm.utils import compute_trans_freqs
 
+
 big: float = 1e20
 small: float = 1.0 / big
 
@@ -481,20 +487,19 @@ def large_p_mvnormal_sampler(D_diag, Phi, a):
     N = a.shape[0]
     u = np.random.normal(0, np.sqrt(D_diag))
     delta = np.random.normal(size=N)
-    if not (type(Phi) is scipy.sparse.csr.csr_matrix):
-        Phi_D = Phi * D_diag
-        v = Phi @ u + delta
-        Z = Phi_D @ Phi.T + np.eye(N)
-        w = scipy.linalg.solve(Z, a - v, assume_a="sym")
-        beta = u + Phi_D.T @ w
-
-    else:
+    if scipy.sparse.issparse(Phi):
         Phi_D = Phi.multiply(D_diag)
         v = Phi * u + delta
-        Z = Phi_D * Phi.T + np.eye(N)
+        Z = (Phi_D * Phi.T + scipy.sparse.eye(N)).toarray()
         w = scipy.linalg.solve(Z, a - v, assume_a="sym")
         beta = u + Phi_D.T * w
-
+    else:
+        Phi_D = Phi * D_diag
+        v = Phi.dot(u) + delta
+        Z = Phi_D.dot(Phi.T)
+        Z.flat[:: N + 1] += 1
+        w = scipy.linalg.solve(Z, a - v, assume_a="sym")
+        beta = u + Phi_D.T @ w
     return beta
 
 
@@ -523,33 +528,69 @@ def hs_step(
 class HSStep(BlockedStep):
     name = "hsgibbs"
 
-    def __init__(self, vars, y, X, values=None, model=None):
+    def __init__(self, vars, values=None, model=None):
+        model = pm.modelcontext(model)
 
         if len(vars) > 1:
             raise ValueError("This sampler only takes one variable.")
 
-        (var,) = pm.inputvars(vars)
+        (beta,) = pm.inputvars(vars)
 
-        if not isinstance(var.distribution, pm.distributions.Normal):
-            raise TypeError("This sampler only samples `Normal`s.")
+        if not isinstance(beta.distribution, HorseShoe):
+            raise TypeError("This sampler only samples `HorseShoe`s.")
 
-        model = pm.modelcontext(model)
+        non_var_rvs = [
+            value for attr, value in model.named_vars.items() if value != beta
+        ]
 
-        self.vars = [var]
+        # 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
+            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
+            ):
+                continue
+            if i in model.observed_RVs:
+
+                def y_fn(x):
+                    return i.observations
+
+            else:
+                y_fn = model.fn(i)
 
-        M = model.test_point[var.name].shape[-1]
+            X_fn = model.fn(mu.owner.inputs[0])
+
+        self.vars = [beta]
+
+        M = model.test_point[beta.name].shape[-1]
 
         self.vi = np.full(M, 1)
         self.lambda2 = np.full(M, 1)
         self.beta = np.full(M, 1)
         self.tau2 = 1
         self.xi = 1
-        self.y = y
-        self.X = X
+        self.y_fn = y_fn
+        self.X_fn = X_fn
 
     def step(self, point):
+        y = self.y_fn(point)
+        X = self.X_fn(point)
+
         self.beta, self.lambda2, self.tau2, self.vi, self.xi = hs_step(
-            self.lambda2, self.tau2, self.vi, self.xi, self.X, self.y
+            self.lambda2, self.tau2, self.vi, self.xi, X, y
         )
         point[self.vars[0].name] = self.beta
         return point

tests/test_distributions.py

@@ -13,6 +13,7 @@
 from pymc3_hmm.distributions import (
     Constant,
     DiscreteMarkovChain,
+    HorseShoe,
     PoissonZeroProcess,
     SwitchingProcess,
     distribution_subset_args,
@@ -562,3 +563,17 @@ def test_subset_args():
     res = distribution_subset_args(test_dist, shape=[3], idx=test_idx)
     assert np.array_equal(res["mu"].eval(), np.r_[0.1, 2.3])
     assert np.array_equal(res["alpha"].eval(), np.r_[2.0, 2.0])
+
+
+def test_HorseShoe():
+    with pm.Model():
+        test_beta = HorseShoe.dist(tau=1, shape=2)
+        test_sample = test_beta.random(size=1000).eval()
+
+    assert test_sample.shape == (1000, 2)
+    assert np.isclose(np.median(test_sample), 0, atol=0.01)
+    assert np.percentile(test_sample.flatten(), 85) < 2
+    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

tests/test_step_methods.py

@@ -5,15 +5,17 @@
 try:
     import aesara.tensor as at
     from aesara.graph.op import get_test_value
+    from aesara.sparse import structured_dot as sp_dot
 except ImportError:
     import theano.tensor as at
     from theano.graph.op import get_test_value
+    from theano.sparse import structured_dot as sp_dot
 
 import pymc3 as pm
 import pytest
 import scipy as sp
 
-from pymc3_hmm.distributions import DiscreteMarkovChain, PoissonZeroProcess
+from pymc3_hmm.distributions import DiscreteMarkovChain, HorseShoe, PoissonZeroProcess
 from pymc3_hmm.step_methods import (
     FFBSStep,
     HSStep,
@@ -398,9 +400,10 @@ def test_Hsstep():
     y = np.random.normal(X.dot(beta_true), 1)
 
     M = X.shape[1]
-    with pm.Model() as _:
-        beta = pm.Normal("beta", 0, 1, shape=M)
-        hsstep = HSStep([beta], y, X)
+    with pm.Model():
+        beta = HorseShoe("beta", tau=1, shape=M)
+        pm.Normal("y", mu=X * beta, sigma=1, observed=y)
+        hsstep = HSStep([beta])
         trace = pm.sample(
             draws=20, tune=0, step=hsstep, chains=1, return_inferencedata=True
         )
@@ -410,20 +413,19 @@ def test_Hsstep():
     assert beta_samples.shape == (20, M)
 
     # test case for sparse matrix
-    X = np.random.choice([0, 1], size=25).reshape((M, M))
-    X_sp = sp.sparse.csr_matrix(X)
+    X = sp.sparse.csr_matrix(X)
 
     M = X.shape[1]
-    with pm.Model() as _:
-        beta = pm.Normal("beta", 0, 1, shape=M)
-        hsstep = HSStep([beta], y, X_sp)
+    with pm.Model():
+        beta = HorseShoe("beta", tau=1, shape=M)
+        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
         )
-
+        beta_2 = HorseShoe("beta_2", tau=1, shape=M)
         with pytest.raises(ValueError):
-            beta_2 = pm.Normal("beta_2", 0, 1)
-            hsstep = HSStep([beta, beta_2], y, X_sp)
+            HSStep([beta, beta_2])
 
     beta_samples = trace.posterior["beta"][0].values