Adding recover_marginals utility function

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -2,27 +2,37 @@
 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
+from pytensor.tensor.shape import Shape
+from scipy.special import log_softmax
 
 __all__ = ["MarginalModel"]
 
-from pytensor.tensor.shape import Shape
-
 
 class MarginalModel(Model):
     """Subclass of PyMC Model that implements functionality for automatic
@@ -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,176 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         # Raise errors and warnings immediately
         self.clone()._marginalize(user_warnings=True)
 
+    def _to_transformed(self):
+        "Create a function from the untransformed space to the transformed space"
+        transformed_rvs = []
+        transformed_names = []
+
+        for rv in self.free_RVs:
+            transform = self.rvs_to_transforms.get(rv)
+            if transform is None:
+                transformed_rvs.append(rv)
+                transformed_names.append(rv.name)
+            else:
+                transformed_rv = transform.forward(rv, *rv.owner.inputs)
+                transformed_rvs.append(transformed_rv)
+                transformed_names.append(self.rvs_to_values[rv].name)
+
+        fn = self.compile_fn(inputs=self.free_RVs, outs=transformed_rvs)
+        return fn, transformed_names
+
+    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, return_samples=True, extend_inferencedata=True
+    ):
+        """Computes normalized posterior probabilities 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
+
+        All log-probabilities are within the transformed space
+
+        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
+        return_samples : bool, default True
+            If True, also return 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)
+
+        # Handle Transforms
+        transform_fn, transform_names = self._to_transformed()
+
+        def transform_input(inputs):
+            return dict(zip(transform_names, transform_fn(inputs)))
+
+        posterior_pts = [transform_input(vs) for vs in posterior_pts]
+
+        rv_dict = {}
+        rv_dims_dict = {}
+
+        for rv in var_names:
+            supported_dists = (Bernoulli, Categorical, DiscreteUniform)
+            if not isinstance(rv.owner.op, supported_dists):
+                raise NotImplementedError(
+                    f"RV with distribution {rv.owner.op} cannot be marginalized. "
+                    f"Supported distribution include {supported_dists}"
+                )
+
+            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
+            joint_logps = vectorize_graph(
+                joint_logp,
+                replace={marginalized_value: rv_domain_tensor},
+            )
+
+            rv_loglike_fn = None
+            if return_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 return_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] = log_softmax(
+                    np.reshape(
+                        logps,
+                        tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:],
+                    ),
+                    axis=len(stacked_dims),
+                )
+                rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
+            else:
+                logps = np.array(logvs)
+                rv_dict["lp_" + rv.name] = log_softmax(
+                    np.reshape(
+                        logps,
+                        tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:],
+                    ),
+                    axis=len(stacked_dims),
+                )
+                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

@@ -6,11 +6,13 @@
 import pymc as pm
 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 scipy.special import logsumexp
+from scipy.special import log_softmax, logsumexp
+from scipy.stats import halfnorm, norm
 
 from pymc_experimental.model.marginal_model import (
     FiniteDiscreteMarginalRV,
@@ -251,6 +253,101 @@ def test_marginalized_change_point_model_sampling(disaster_model):
     )
 
 
+def test_recover_marginals_basic():
+    with MarginalModel() as m:
+        sigma = pm.HalfNormal("sigma")
+        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], sigma=sigma)
+
+    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, return_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, sigma):
+        y = y.repeat(len(p)).reshape(len(y), -1)
+        sigma = sigma.repeat(len(p)).reshape(len(sigma), -1)
+        return log_softmax(
+            np.log(p)
+            + norm.logpdf(y, loc=mu, scale=sigma)
+            + halfnorm.logpdf(sigma)
+            + np.log(sigma),
+            axis=1,
+        )
+
+    np.testing.assert_almost_equal(
+        true_logp(idata.y.values.flatten(), idata.sigma.values.flatten()),
+        idata.lp_k[0].values,
+    )
+
+
+def test_nested_recover_marginals():
+    """Test that marginalization works when there are nested marginalized RVs"""
+
+    with MarginalModel() as m:
+        idx = pm.Bernoulli("idx", p=0.75)
+        sub_idx = pm.Bernoulli("sub_idx", p=pt.switch(pt.eq(idx, 0), 0.15, 0.95))
+        sub_dep = pm.Normal("y", mu=idx + sub_idx, sigma=1.0)
+
+    m.marginalize([idx, sub_idx])
+
+    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, return_samples=True)
+    assert "idx" in idata
+    assert "lp_idx" in idata
+    assert idata.idx.shape == idata.y.shape
+    assert idata.lp_idx.shape == idata.idx.shape + (2,)
+    assert "sub_idx" in idata
+    assert "lp_sub_idx" in idata
+    assert idata.sub_idx.shape == idata.y.shape
+    assert idata.lp_sub_idx.shape == idata.sub_idx.shape + (2,)
+
+    def true_idx_logp(y):
+        idx_0 = np.log(0.85 * 0.25 * norm.pdf(y, loc=0) + 0.15 * 0.25 * norm.pdf(y, loc=1))
+        idx_1 = np.log(0.05 * 0.75 * norm.pdf(y, loc=1) + 0.95 * 0.75 * norm.pdf(y, loc=2))
+        return log_softmax(np.stack([idx_0, idx_1]).T, axis=1)
+
+    np.testing.assert_almost_equal(
+        true_idx_logp(idata.y.values.flatten()),
+        idata.lp_idx[0].values,
+    )
+
+    def true_sub_idx_logp(y):
+        sub_idx_0 = np.log(0.85 * 0.25 * norm.pdf(y, loc=0) + 0.05 * 0.75 * norm.pdf(y, loc=1))
+        sub_idx_1 = np.log(0.15 * 0.25 * norm.pdf(y, loc=1) + 0.95 * 0.75 * norm.pdf(y, loc=2))
+        return log_softmax(np.stack([sub_idx_0, sub_idx_1]).T, axis=1)
+
+    np.testing.assert_almost_equal(
+        true_sub_idx_logp(idata.y.values.flatten()),
+        idata.lp_sub_idx[0].values,
+    )
+
+
 @pytest.mark.filterwarnings("error")
 def test_not_supported_marginalized():
     """Marginalized graphs with non-Elemwise Operations are not supported as they