Progress on marginalization of DiscreteMarkovChain

Author: jessegrabowski

pymc_experimental/marginal_model.py

@@ -7,14 +7,15 @@
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.abstract import _logprob
-from pymc.logprob.basic import conditional_logp
+from pymc.logprob.basic import conditional_logp, logp
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
 from pymc.pytensorf import constant_fold, inputvars
-from pytensor import Mode
+from pytensor import Mode, scan
 from pytensor.compile import SharedVariable
 from pytensor.compile.builders import OpFromGraph
 from pytensor.graph import Constant, FunctionGraph, ancestors, clone_replace
+from pytensor.graph.replace import vectorize_graph
 from pytensor.scan import map as scan_map
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.elemwise import Elemwise
@@ -255,6 +256,10 @@ class FiniteDiscreteMarginalRV(MarginalRV):
     """Base class for Finite Discrete Marginalized RVs"""
 
 
+class DiscreteMarginalMarkovChainRV(MarginalRV):
+    """Base class for Discrete Marginal Markov Chain RVs"""
+
+
 def static_shape_ancestors(vars):
     """Identify ancestors Shape Ops of static shapes (therefore constant in a valid graph)."""
     return [
@@ -383,11 +388,17 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     replace_inputs.update({input_rv: input_rv.type() for input_rv in input_rvs})
     cloned_outputs = clone_replace(outputs, replace=replace_inputs)
 
-    marginalization_op = FiniteDiscreteMarginalRV(
+    if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
+        marginalize_constructor = DiscreteMarginalMarkovChainRV
+    else:
+        marginalize_constructor = FiniteDiscreteMarginalRV
+
+    marginalization_op = marginalize_constructor(
         inputs=list(replace_inputs.values()),
         outputs=cloned_outputs,
         ndim_supp=ndim_supp,
     )
+
     marginalized_rvs = marginalization_op(*replace_inputs.keys())
     fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))
     return rvs_to_marginalize, marginalized_rvs
@@ -435,7 +446,7 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
         values_axis_bcast = [i for i, (m, v) in enumerate(zip(mbcast, vbcast)) if m != v]
         joint_logp += logps_dict[inner_value].sum(values_axis_bcast, keepdims=True)
 
-    # Wrap the joint_logp graph in an OpFromGrah, so that we can evaluate it at different
+    # Wrap the joint_logp graph in an OpFromGraph, so that we can evaluate it at different
     # values of the marginalized RV
     # Some inputs are not root inputs (such as transformed projections of value variables)
     # Or cannot be used as inputs to an OpFromGraph (shared variables and constants)
@@ -487,3 +498,55 @@ def logp_fn(marginalized_rv_const, *non_sequences):
 
     # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
     return joint_logps, *(pt.constant(0),) * (len(values) - 1)
+
+
+@_logprob.register(DiscreteMarginalMarkovChainRV)
+def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
+    def step_alpha(log_alpha, log_P):
+        return pt.logsumexp(log_alpha[:, None] + log_P, 0)
+
+    def eval_logp(x):
+        return logp(init_dist_, x)
+
+    marginalized_rvs_node = op.make_node(*inputs)
+    inner_rvs = clone_replace(
+        op.inner_outputs,
+        replace={u: v for u, v in zip(op.inner_inputs, marginalized_rvs_node.inputs)},
+    )
+
+    chain_rv, *dependent_rvs = inner_rvs
+    P_, n_steps_, init_dist_, rng = chain_rv.owner.inputs
+
+    domain = pt.arange(P_.shape[0], dtype="int32")
+
+    vec_eval_logp = pt.vectorize(eval_logp, "()->()")
+    logp_init = vec_eval_logp(domain)
+
+    # This will break the dependency between chain and the init_dist_ random variable
+    # TODO: Make this comment more robust after I understand better.
+    chain_dummy = chain_rv.clone()
+    dependent_rvs = clone_replace(dependent_rvs, {chain_rv: chain_dummy})
+    input_dict = dict(zip(dependent_rvs, values))
+    logp_value_dict = conditional_logp(input_dict)
+
+    # TODO: Is values[0] robust to every situation?
+    sub_dict = {
+        chain_dummy: pt.moveaxis(pt.broadcast_to(domain, (*values[0].shape, domain.size)), -1, 0)
+    }
+
+    # TODO: @Ricardo: If you don't concatenate here, you get -inf in the logp (why?)
+    # TODO: I'm stacking the results (adds a batch dim to the left) and summing away the batch dim == joint probability?
+    vec_logp_emission = pt.stack(vectorize_graph(tuple(logp_value_dict.values()), sub_dict)).sum(
+        axis=0
+    )
+
+    log_alpha_seq, _ = scan(
+        step_alpha, non_sequences=[pt.log(P_)], outputs_info=[logp_init], n_steps=n_steps_
+    )
+
+    log_alpha_seq = pt.moveaxis(pt.concatenate([logp_init[None], log_alpha_seq], axis=0), -1, 0)
+    joint_log_obs_given_states = pt.logsumexp(pt.add(vec_logp_emission) + log_alpha_seq, axis=0)
+
+    # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
+    dummy_logps = (pt.constant(0.0),) * (len(values) - 1)
+    return joint_log_obs_given_states, dummy_logps

pymc_experimental/tests/test_marginal_model.py

@@ -473,88 +473,38 @@ def dist(idx, size):
         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)
+def test_marginalized_hmm_with_one_emission():
+    with MarginalModel() as m:
+        P = [[0, 1], [1, 0]]
+        init_dist = pm.Categorical.dist(p=[1, 0])
+        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, steps=3)
+        emission = pm.Normal("emission", mu=chain * 2 - 1, sigma=1e-1)
 
-    return logprob_emiss_ts.sum()
+    m.marginalize([chain])
 
+    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": test_value}), expected_logp)
 
-def test_hmm():
 
+def test_marginalized_hmm_with_many_emissions():
     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)
-
-
+        P = [[0, 1], [1, 0]]
+        init_dist = pm.Categorical.dist(p=[1, 0])
+        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, steps=3)
+        emission_1 = pm.Normal("emission_1", mu=chain * 2 - 1, sigma=1e-1)
+        emission_2 = pm.Normal("emission_2", mu=(1 - chain) * 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)
+        m.marginalize([chain])
 
     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()
+    test_point = {"emission_1": test_value, "emission_2": test_value * -1}
 
-    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,
-    # )
+    assert False
+    # np.testing.assert_allclose(logp_fn(test_point), expected_logp)