Change to providing logprobs for marginalised variables in

the transformed space

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -255,6 +255,24 @@ 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)
@@ -263,12 +281,14 @@ def unmarginalize(self, rvs_to_unmarginalize):
     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
+        """Computes unnormalized 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
@@ -298,18 +318,20 @@ def recover_marginals(
 
         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 = {}
 
-        # Disable all transforms
-        model = self.clone()
-        model.rvs_to_transforms = {k: None for k in model.rvs_to_transforms}
-        for rv in model.free_RVs:
-            model.rvs_to_values[rv].name = rv.name
-        var_names = [model.vars_to_clone[rv] for rv in var_names]
-
         for rv in var_names:
-            m = model.clone()
+            m = self.clone()
             rv = m.vars_to_clone[rv]
             m.unmarginalize([rv])
             joint_logp = m.logp()
@@ -331,6 +353,7 @@ def recover_marginals(
             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},

pymc_experimental/tests/model/test_marginal_model.py

@@ -283,7 +283,9 @@ def test_recover_marginals_basic():
     def true_logp(y, sigma):
         y = y.repeat(len(p)).reshape(len(y), -1)
         sigma = sigma.repeat(len(p)).reshape(len(sigma), -1)
-        return np.log(p) + norm.logpdf(y, loc=mu, scale=sigma) + halfnorm.logpdf(sigma)
+        return (
+            np.log(p) + norm.logpdf(y, loc=mu, scale=sigma) + halfnorm.logpdf(sigma) + np.log(sigma)
+        )
 
     np.testing.assert_almost_equal(
         true_logp(idata.y.values.flatten(), idata.sigma.values.flatten()),