unify

Author: aloctavodia

docs/source/api/misc.rst

@@ -7,5 +7,6 @@ Other utils
    :toctree: generated/
 
    compute_log_likelihood
+   compute_log_prior
    find_constrained_prior
    DictToArrayBijection

pymc/backends/arviz.py

@@ -436,7 +436,7 @@ def to_inference_data(self):
             id_dict["constant_data"] = self.constant_data_to_xarray()
         idata = InferenceData(save_warmup=self.save_warmup, **id_dict)
         if self.log_likelihood:
-            from pymc.stats.log_likelihood import compute_log_likelihood
+            from pymc.stats.log_density import compute_log_likelihood
 
             idata = compute_log_likelihood(
                 idata,

pymc/stats/__init__.py

@@ -27,7 +27,6 @@
     if not attr.startswith("__"):
         setattr(sys.modules[__name__], attr, obj)
 
-from pymc.stats.log_likelihood import compute_log_likelihood
-from pymc.stats.log_prior import compute_log_prior
+from pymc.stats.log_density import compute_log_likelihood, compute_log_prior
 
 __all__ = ("compute_log_likelihood", "compute_log_prior", *az.stats.__all__)

pymc/stats/log_density.py

@@ -24,6 +24,48 @@
 from pymc.pytensorf import PointFunc
 from pymc.util import dataset_to_point_list
 
+__all__ = ("compute_log_likelihood", "compute_log_prior")
+
+
+def compute_log_likelihood(
+    idata: InferenceData,
+    *,
+    var_names: Optional[Sequence[str]] = None,
+    extend_inferencedata: bool = True,
+    model: Optional[Model] = None,
+    sample_dims: Sequence[str] = ("chain", "draw"),
+    progressbar=True,
+):
+    """Compute elemwise log_likelihood of 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.
+    extend_inferencedata : bool, default True
+        Whether to extend the original InferenceData or return a new one
+    model : Model, optional
+    sample_dims : sequence of str, default ("chain", "draw")
+    progressbar : bool, default True
+
+    Returns
+    -------
+    idata : InferenceData
+        InferenceData with log_likelihood group
+    """
+    return compute_log_density(
+        idata=idata,
+        var_names=var_names,
+        extend_inferencedata=extend_inferencedata,
+        model=model,
+        kind="likelihood",
+        sample_dims=sample_dims,
+        progressbar=progressbar,
+    )
+
 
 def compute_log_prior(
     idata: InferenceData,

pymc/stats/log_likelihood.py

@@ -20,7 +20,7 @@
 from pymc.distributions import Dirichlet, Normal
 from pymc.distributions.transforms import log
 from pymc.model import Model
-from pymc.stats.log_likelihood import compute_log_likelihood
+from pymc.stats.log_density import compute_log_likelihood, compute_log_prior
 from tests.distributions.test_multivariate import dirichlet_logpdf
 
 
@@ -132,3 +132,26 @@ def test_dims_without_coords(self):
             llike.log_likelihood["y"].values,
             st.norm.logpdf([[[0, 0, 0], [1, 1, 1]]]),
         )
+
+    @pytest.mark.parametrize("transform", (False, True))
+    def test_basic_log_prior(self, transform):
+        transform = log if transform else None
+        with Model() as m:
+            x = Normal("x", transform=transform)
+            x_value_var = m.rvs_to_values[x]
+            Normal("y", x, observed=[0, 1, 2])
+
+            idata = InferenceData(posterior=dict_to_dataset({"x": np.arange(100).reshape(4, 25)}))
+            res = compute_log_prior(idata)
+
+        # Check we didn't erase the original mappings
+        assert m.rvs_to_values[x] is x_value_var
+        assert m.rvs_to_transforms[x] is transform
+
+        assert res is idata
+        assert res.log_prior.dims == {"chain": 4, "draw": 25}
+
+        np.testing.assert_allclose(
+            res.log_prior["x"].values,
+            st.norm(0, 1).logpdf(idata.posterior["x"].values),
+        )