diff --git a/pymc3_hmm/distributions.py b/pymc3_hmm/distributions.py index 7a12f53..93a32a1 100644 --- a/pymc3_hmm/distributions.py +++ b/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 - - diff --git a/pymc3_hmm/step_methods.py b/pymc3_hmm/step_methods.py index 0baecb9..2a7ba5f 100644 --- a/pymc3_hmm/step_methods.py +++ b/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 diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 677dc03..be546a0 100644 --- a/tests/test_distributions.py +++ b/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 diff --git a/tests/test_step_methods.py b/tests/test_step_methods.py index f8db321..f773210 100644 --- a/tests/test_step_methods.py +++ b/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