Adding recover_marginals utility function

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -2,19 +2,29 @@
 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
-from pytensor.graph import Constant, FunctionGraph, ancestors, clone_replace
+from pytensor.graph import (
+    Constant,
+    FunctionGraph,
+    ancestors,
+    clone_replace,
+    vectorize_graph,
+)
 from pytensor.scan import map as scan_map
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.elemwise import Elemwise
@@ -205,8 +215,9 @@ def clone(self):
         vars = self.basic_RVs + self.potentials + self.deterministics + self.marginalized_rvs
         cloned_vars = clone_replace(vars)
         vars_to_clone = {var: cloned_var for var, cloned_var in zip(vars, cloned_vars)}
+        m.vars_to_clone = vars_to_clone
 
-        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 +255,134 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         # Raise errors and warnings immediately
         self.clone()._marginalize(user_warnings=True)
 
+    def unmarginalize(self, rvs_to_unmarginalize):
+        for rv in rvs_to_unmarginalize:
+            self.marginalized_rvs.remove(rv)
+            self.register_rv(rv, name=rv.name)
+
+    def recover_marginals(
+        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
+
+        When there are multiple marginalized variables, each marginalized variable is
+        conditioned on both the parameters and the other variables still marginalized
+
+        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
+
+        posterior = idata.posterior
+
+        # Remove Deterministics
+        posterior_values = posterior[
+            [rv.name for rv in self.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()
+            rv = m.vars_to_clone[rv]
+            m.unmarginalize([rv])
+            joint_logp = m.logp()
+
+            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_logps = vectorize_graph(
+                joint_logp,
+                replace={marginalized_value: rv_domain_tensor},
+            )
+
+            rv_loglike_fn = None
+            if include_samples:
+                sample_rv_outs = pymc.Categorical.dist(logit_p=joint_logps)
+                rv_loglike_fn = compile_pymc(
+                    inputs=other_values,
+                    outputs=[joint_logps, sample_rv_outs],
+                    on_unused_input="ignore",
+                )
+            else:
+                rv_loglike_fn = compile_pymc(
+                    inputs=other_values,
+                    outputs=joint_logps,
+                    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(self)
+        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

@@ -4,13 +4,17 @@
 import numpy as np
 import pandas as pd
 import pymc as pm
+import pytensor
 import pytensor.tensor as pt
 import pytest
+from arviz import InferenceData, dict_to_dataset
 from pymc import ImputationWarning, inputvars
 from pymc.distributions import transforms
 from pymc.logprob.abstract import _logprob
 from pymc.util import UNSET
+from pytensor.graph import vectorize_graph
 from scipy.special import logsumexp
+from scipy.stats import norm
 
 from pymc_experimental.model.marginal_model import (
     FiniteDiscreteMarginalRV,
@@ -251,6 +255,42 @@ def test_marginalized_change_point_model_sampling(disaster_model):
     )
 
 
+def test_recover_marginals_basic():
+    with MarginalModel() as m:
+        p = np.array([0.5, 0.2, 0.3])
+        k = pm.Categorical("k", p=p)
+        mu = np.array([-3.0, 0.0, 3.0])
+        mu_ = pt.as_tensor_variable(mu)
+        y = pm.Normal("y", mu=mu_[k])
+
+    m.marginalize([k])
+
+    rng = np.random.default_rng(211)
+
+    with m:
+        prior = pm.sample_prior_predictive(
+            samples=20,
+            random_seed=rng,
+            return_inferencedata=False,
+        )
+        idata = InferenceData(posterior=dict_to_dataset(prior))
+
+    idata = m.recover_marginals(idata, include_samples=True)
+    assert "k" in idata
+    assert "lp_k" in idata
+    assert idata.k.shape == idata.y.shape
+    assert idata.lp_k.shape == idata.k.shape + (len(p),)
+
+    def true_logp(y):
+        y = y.repeat(len(p)).reshape(len(y), -1)
+        return np.log(p) + norm.logpdf(y, loc=mu)
+
+    np.testing.assert_almost_equal(
+        true_logp(idata.y.values.flatten()),
+        idata.lp_k[0].values,
+    )
+
+
 @pytest.mark.filterwarnings("error")
 def test_not_supported_marginalized():
     """Marginalized graphs with non-Elemwise Operations are not supported as they