.hmm

Author: ricardoV94

pymc_experimental/marginal_model.py

@@ -23,6 +23,8 @@
 
 from pytensor.tensor.shape import Shape
 
+from pymc_experimental.distributions import DiscreteMarkovChain
+
 
 class MarginalModel(Model):
     """Subclass of PyMC Model that implements functionality for automatic
@@ -226,7 +228,7 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         if not isinstance(rvs_to_marginalize, Sequence):
             rvs_to_marginalize = (rvs_to_marginalize,)
 
-        supported_dists = (Bernoulli, Categorical, DiscreteUniform)
+        supported_dists = (Bernoulli, Categorical, DiscreteUniform, DiscreteMarkovChain)
         for rv_to_marginalize in rvs_to_marginalize:
             if rv_to_marginalize not in self.free_RVs:
                 raise ValueError(
@@ -342,6 +344,7 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     ndim_supp = {rv.owner.op.ndim_supp for rv in dependent_rvs}
     if len(ndim_supp) != 1:
         raise NotImplementedError()
+    ndim_supp = tuple(ndim_supp)[0]
     # if max(ndim_supp) > 0:
     #     raise NotImplementedError(
     #         "Marginalization with dependent Multivariate RVs not implemented"
@@ -400,6 +403,9 @@ def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> Tuple[int, ...]:
     elif isinstance(op, DiscreteUniform):
         lower, upper = constant_fold(rv.owner.inputs[3:])
         return tuple(range(lower, upper + 1))
+    elif isinstance(op, DiscreteMarkovChain):
+        p = rv.owner.inputs[0]
+        return tuple(range(pt.get_vector_length(p[-1])))
 
     raise NotImplementedError(f"Cannot compute domain for op {op}")
 
@@ -457,6 +463,7 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
     )
 
     # Arbitrary cutoff to switch to Scan implementation to keep graph size under control
+    # TODO: Try vectorize here
     if len(marginalized_rv_domain) <= 10:
         joint_logps = [
             joint_logp_op(marginalized_rv_domain_tensor[i], *values, *inputs)

pymc_experimental/tests/test_marginal_model.py

@@ -12,6 +12,7 @@
 from pymc.util import UNSET
 from scipy.special import logsumexp
 
+from pymc_experimental.distributions import DiscreteMarkovChain
 from pymc_experimental.marginal_model import (
     FiniteDiscreteMarginalRV,
     MarginalModel,
@@ -470,3 +471,90 @@ def dist(idx, size):
     ):
         pt = {"norm": test_value}
         np.testing.assert_allclose(logp_fn(pt), ref_logp_fn(pt))
+
+
+def hmm_logp(values, P, steps, init_dist, state_rng):
+
+    [e_value] = values
+
+    # P = [[0, 1], [1, 0]]
+    domain = tuple(range(pt.get_vector_length(P[-1])))
+
+    # This should be done on log-scale
+    # Probability of states at t0
+    logprob_states = pm.math.stack([logp(init_dist, d) for d in domain])
+
+    logprob_emiss_ts = []
+    for e_value_t in e_value:
+        # Use vectorize
+        logprob_emiss_t = pt.sum(
+            [
+                logpprob_state + logp(clone_replace(emission_rv, replace={state_rv: state_value}), e_value_t)
+                for (logpprob_state, state_value) in zip(logprob_states, domain)
+            ]
+        )
+
+        # Probability next state
+        # prob_states = prob_states @ P
+        logprob_states = P[:, None]
+
+        logprob_emiss_ts.append(logprob_emiss_t)
+
+    return logprob_emiss_ts.sum()
+
+
+
+def test_hmm():
+
+    with MarginalModel() as m:
+        p = pt.as_tensor(np.array([1, 0]))
+
+        chain_0 = pm.Bernoulli("chain_0", p=0)
+        chain_1 = pm.Bernoulli("chain_1", p=p[chain_0])
+        chain_2 = pm.Bernoulli("chain_2", p=p[chain_1])
+        chain_3 = pm.Bernoulli("chain_3", p=p[chain_2])
+
+        pm.Normal("emission_0", chain_0 * 2 - 1, sigma=1e-1)
+        pm.Normal("emission_1", chain_1 * 2 - 1, sigma=1e-1)
+        pm.Normal("emission_2", chain_2 * 2 - 1, sigma=1e-1)
+        pm.Normal("emission_3", chain_3 * 2 - 1, sigma=1e-1)
+
+
+
+    with pytest.warns(UserWarning, match="multiple dependent variables"):
+        m.marginalize([chain_0, chain_1, chain_2, chain_3])
+    import pytensor
+    print()
+    pytensor.dprint(m.clone()._marginalize().free_RVs)
+
+    logp_fn = m.compile_logp()
+    test_value = [-1, 1, -1, 1]
+
+    expected_logp = pm.logp(pm.Normal.dist(0, 1e-1), np.zeros_like(test_value)).sum().eval()
+
+    np.testing.assert_allclose(
+        logp_fn({f"emission_{i}": test_value_i for i, test_value_i in enumerate(test_value)}),
+        expected_logp,
+    )
+    return
+
+    # with MarginalModel() as m:
+    #     P = [[0, 1], [1, 0]]
+    #     zero = pm.DiracDelta.dist(np.array(0, dtype="int64"))
+    #     chain = DiscreteMarkovChain("chain", P=P, init_dist=zero, steps=3)
+    #     emmission = pm.Normal("emission", mu=chain * 2 - 1, sigma=1e-1)
+    # np.testing.assert_equal(pm.draw(chain), [0, 1, 0, 1])
+    # m.marginalize(chain)
+
+    # test_value = [-1, 1, -1, 1]
+    # expected_logp = pm.logp(pm.Normal.dist(0, 1e-1), [0, 0, 0, 0]).eval()
+
+    # np.testing.assert_allclose(
+    #     logp_fn({"emission": test_value}),
+    #     expected_logp,
+    # )
+    #
+    # np.testing.assert_allclose(
+    #     logp_fn({f"emission{i}": test_value_i for i, test_value_i in enumerate(test_value)}),
+    #     expected_logp,
+    # )