Adding unmarginalize

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -2,15 +2,19 @@
 from typing import Sequence, Tuple, Union
 
 import numpy as np
+import pymc
 import pytensor.tensor as pt
+from arviz import dict_to_dataset
 from pymc import SymbolicRandomVariable
+from pymc.backends.arviz import coords_and_dims_for_inferencedata
 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.transforms import IntervalTransform
 from pymc.model import Model
-from pymc.pytensorf import constant_fold, inputvars
+from pymc.pytensorf import compile_pymc, constant_fold, inputvars
+from pymc.util import dataset_to_point_list, treedict
 from pytensor import Mode
 from pytensor.compile import SharedVariable
 from pytensor.compile.builders import OpFromGraph
@@ -206,7 +210,7 @@ def clone(self):
         cloned_vars = clone_replace(vars)
         vars_to_clone = {var: cloned_var for var, cloned_var in zip(vars, cloned_vars)}
 
-        m.named_vars = {name: vars_to_clone[var] for name, var in self.named_vars.items()}
+        m.named_vars = treedict({name: vars_to_clone[var] for name, var in self.named_vars.items()})
         m.named_vars_to_dims = self.named_vars_to_dims
         m.values_to_rvs = {i: vars_to_clone[rv] for i, rv in self.values_to_rvs.items()}
         m.rvs_to_values = {vars_to_clone[rv]: i for rv, i in self.rvs_to_values.items()}
@@ -244,6 +248,127 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         # Raise errors and warnings immediately
         self.clone()._marginalize(user_warnings=True)
 
+    def unmarginalize(
+        self, idata, var_names=None, include_samples=False, extend_inferencedata=True
+    ):
+        """Computes log-likelihoods of marginalized variables conditioned on parameters
+        of the model given InferenceData with posterior group
+
+        Parameters
+        ----------
+        idata : InferenceData
+            InferenceData with posterior group
+        var_names : sequence of str, optional
+            List of Observed variable names for which to compute log_likelihood. Defaults to all observed variables
+        include_samples : bool, default False
+            Include samples of the marginalized variables
+        extend_inferencedata : bool, default True
+            Whether to extend the original InferenceData or return a new one
+
+        Returns
+        -------
+        idata : InferenceData
+            InferenceData with var_names added to posterior
+
+        """
+        if var_names is None:
+            var_names = self.marginalized_rvs
+
+        joint_logp = self.logp()
+        posterior = idata.posterior
+
+        # Remove Deterministics
+        posterior_values = posterior[
+            [rv.name for rv in mm.free_RVs if rv not in self.marginalized_rvs]
+        ]
+
+        sample_dims = ("chain", "draw")
+        posterior_pts, stacked_dims = dataset_to_point_list(posterior_values, sample_dims)
+        rv_dict = {}
+        rv_dims_dict = {}
+
+        for rv in var_names:
+            m = self.clone()
+            m.register_rv(rv, name=rv.name)
+            m.marginalized_rvs = [r for r in m.marginalized_rvs if r is not rv]
+
+            rv_shape = constant_fold(tuple(rv.shape))
+            rv_domain = get_domain_of_finite_discrete_rv(rv)
+            rv_domain_tensor = pt.swapaxes(
+                pt.full(
+                    (*rv_shape, len(rv_domain)),
+                    rv_domain,
+                    dtype=rv.dtype,
+                ),
+                axis1=0,
+                axis2=-1,
+            )
+
+            marginalized_value = m.rvs_to_values[rv]
+
+            other_values = [v for v in m.value_vars if v is not marginalized_value]
+
+            # TODO: Handle constants
+            # TODO: Handle transformed variables
+            joint_logp_op = OpFromGraph(
+                [marginalized_value] + other_values, [joint_logp], inline=True
+            )
+            joint_logps = [
+                joint_logp_op(rv_domain_tensor[i], *other_values) for i in range(len(rv_domain))
+            ]
+
+            rv_loglike_fn = None
+            if include_samples:
+                sample_rv_outs = pm.Categorical.dist(logit_p=joint_logps)
+                rv_loglike_fn = compile_pymc(
+                    inputs=other_values,
+                    outputs=[pt.stack(joint_logps, 0), sample_rv_outs],
+                    on_unused_input="ignore",
+                )
+            else:
+                rv_loglike_fn = compile_pymc(
+                    inputs=other_values,
+                    outputs=pt.stack(joint_logps, 0),
+                    on_unused_input="ignore",
+                )
+
+            logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
+
+            if include_samples:
+                logps, samples = zip(*logvs)
+                logps = np.array(logps)
+                rv_dict[rv.name] = np.reshape(
+                    samples, tuple(len(coord) for coord in stacked_dims.values())
+                )
+                rv_dims_dict[rv.name] = sample_dims
+                rv_dict["lp_" + rv.name] = np.reshape(
+                    logps, tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
+                )
+                rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
+            else:
+                logps = np.array(logvs)
+                rv_dict["lp_" + rv.name] = np.reshape(
+                    logps, tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
+                )
+                rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
+
+        coords, dims = coords_and_dims_for_inferencedata(model)
+        rv_dataset = dict_to_dataset(
+            rv_dict,
+            library=pymc,
+            dims=dims,
+            coords=coords,
+            default_dims=list(sample_dims),
+            skip_event_dims=True,
+        )
+
+        if extend_inferencedata:
+            rv_dict = {k: (rv_dims_dict[k], v) for (k, v) in rv_dict.items()}
+            idata = idata.posterior.assign(**rv_dict)
+            return idata
+        else:
+            return rv_dataset
+
 
 class MarginalRV(SymbolicRandomVariable):
     """Base class for Marginalized RVs"""

pymc_experimental/tests/model/test_marginal_model.py

@@ -52,7 +52,12 @@ def test_marginalized_bernoulli_logp():
 
     idx = pm.Bernoulli.dist(0.7, name="idx")
     y = pm.Normal.dist(mu=mu[idx], sigma=1.0, name="y")
-    marginal_rv_node = FiniteDiscreteMarginalRV([mu], [idx, y], ndim_supp=None, n_updates=0,)(
+    marginal_rv_node = FiniteDiscreteMarginalRV(
+        [mu],
+        [idx, y],
+        ndim_supp=None,
+        n_updates=0,
+    )(
         mu
     )[0].owner
 
@@ -251,6 +256,26 @@ def test_marginalized_change_point_model_sampling(disaster_model):
     )
 
 
+@pytest.mark.slow
+@pytest.mark.filterwarnings("error")
+def test_unmarginalized_basic(disaster_model):
+    m, years = disaster_model
+
+    with pytest.warns(UserWarning, match="There are multiple dependent variables"):
+        m.marginalize([m["switchpoint"]])
+
+    rng = np.random.default_rng(211)
+
+    with m:
+        idata = pm.sample(chains=2, random_seed=rng).posterior.stack(sample=("draw", "chain"))
+
+    idata = m.unmarginalize(idata, include_samples=True)
+    assert "switchpoint" in idata
+    assert "lp_switchpoint" in idata
+    assert idata.switchpoint.shape == idata.early_mean.shape
+    assert idata.lp_switchpoint.shape == idata.switchpoint.shape + (len(years),)
+
+
 @pytest.mark.filterwarnings("error")
 def test_not_supported_marginalized():
     """Marginalized graphs with non-Elemwise Operations are not supported as they