Add utility to compute deterministics

Author: ricardoV94

.github/workflows/tests.yml

@@ -69,6 +69,7 @@ jobs:
           - |
             tests/distributions/test_censored.py
             tests/distributions/test_simulator.py
+            tests/sampling/test_deterministic.py
             tests/sampling/test_forward.py
             tests/sampling/test_population.py
             tests/stats/test_convergence.py

docs/source/api/samplers.rst

@@ -13,6 +13,9 @@ This submodule contains functions for MCMC and forward sampling.
    sample_posterior_predictive
    draw
 
+.. currentmodule:: pymc.sampling.deterministic
+   compute_deterministics
+
 
 .. currentmodule:: pymc.sampling.mcmc
 

pymc/sampling/__init__.py

@@ -12,5 +12,6 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from pymc.sampling.deterministic import compute_deterministics
 from pymc.sampling.forward import *
 from pymc.sampling.mcmc import *

pymc/sampling/deterministic.py

@@ -0,0 +1,114 @@
+#   Copyright 2024 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+from collections.abc import Sequence
+
+import xarray
+
+from xarray import Dataset
+
+from pymc.backends.arviz import apply_function_over_dataset, coords_and_dims_for_inferencedata
+from pymc.model.core import Model, modelcontext
+
+
+def compute_deterministics(
+    dataset: Dataset,
+    *,
+    var_names: Sequence[str] | None = None,
+    model: Model | None = None,
+    sample_dims: Sequence[str] = ("chain", "draw"),
+    merge_dataset: bool = False,
+    progressbar: bool = True,
+    compile_kwargs: dict | None = None,
+) -> Dataset:
+    """Compute model deterministics given a dataset with values for model variables.
+
+    Parameters
+    ----------
+    dataset : Dataset
+        Dataset with values for model variables. Commonly InferenceData["posterior"].
+    var_names : sequence of str, optional
+        List of names of deterministic variable to compute.
+        If None, compute all deterministics in the model.
+    model : Model, optional
+        Model to use. If None, use context model.
+    sample_dims : sequence of str, default ("chain", "draw")
+        Sample (batch) dimensions of the dataset over which to compute the deterministics.
+    merge_dataset : bool, default False
+        Whether to extend the original dataset or return a new one.
+    progressbar : bool, default True
+        Whether to display a progress bar in the command line.
+    progressbar_theme : Theme, optional
+        Custom theme for the progress bar.
+    compile_kwargs: dict, optional
+        Additional arguments passed to `model.compile_fn`.
+
+    Returns
+    -------
+    Dataset
+        Dataset with values for the deterministics.
+
+
+    Examples
+    --------
+    .. code:: python
+
+        import pymc as pm
+
+        with pm.Model(coords={"group": (0, 2, 4)}) as m:
+            mu_raw = pm.Normal("mu_raw", 0, 1, dims="group")
+            mu = pm.Deterministic("mu", mu_raw.cumsum(), dims="group")
+
+            trace = pm.sample(var_names=["mu_raw"], chains=2, tune=5 draws=5)
+
+        assert "mu" not in trace.posterior
+
+        with m:
+            trace.posterior = pm.compute_deterministics(trace.posterior, merge_dataset=True)
+
+        assert "mu" in trace.posterior
+
+
+    """
+    model = modelcontext(model)
+
+    if var_names is None:
+        deterministics = model.deterministics
+    else:
+        deterministics = [model[var_name] for var_name in var_names]
+        if not set(deterministics).issubset(set(model.deterministics)):
+            raise ValueError("Not all var_names corresponded to model deterministics")
+
+    fn = model.compile_fn(
+        inputs=model.free_RVs,
+        outs=deterministics,
+        on_unused_input="ignore",
+        **(compile_kwargs or {}),
+    )
+
+    coords, dims = coords_and_dims_for_inferencedata(model)
+
+    new_dataset = apply_function_over_dataset(
+        fn,
+        dataset[[rv.name for rv in model.free_RVs]],
+        output_var_names=[det.name for det in model.deterministics],
+        dims=dims,
+        coords=coords,
+        sample_dims=sample_dims,
+        progressbar=progressbar,
+    )
+
+    if merge_dataset:
+        new_dataset = xarray.merge([dataset, new_dataset], compat="override")
+
+    return new_dataset

tests/sampling/test_deterministic.py

@@ -0,0 +1,77 @@
+#   Copyright 2024 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import numpy as np
+import pytest
+
+from numpy.testing import assert_allclose
+
+from pymc.distributions import Normal
+from pymc.model.core import Deterministic, Model
+from pymc.sampling.deterministic import compute_deterministics
+from pymc.sampling.forward import sample_prior_predictive
+
+# Turn all warnings into errors for this module
+pytestmark = pytest.mark.filterwarnings("error")
+
+
+def test_compute_deterministics():
+    with Model(coords={"group": (0, 2, 4)}) as m:
+        mu_raw = Normal("mu_raw", 0, 1, dims="group")
+        mu = Deterministic("mu", mu_raw.cumsum(), dims="group")
+
+        sigma_raw = Normal("sigma_raw", 0, 1)
+        sigma = Deterministic("sigma", sigma_raw.exp())
+
+    dataset = sample_prior_predictive(
+        samples=5, model=m, var_names=["mu_raw", "sigma_raw"], random_seed=22
+    ).prior
+
+    # Test default
+    with m:
+        all_dets = compute_deterministics(dataset)
+    assert set(all_dets.data_vars.variables) == {"mu", "sigma"}
+    assert all_dets["mu"].dims == ("chain", "draw", "group")
+    assert all_dets["sigma"].dims == ("chain", "draw")
+    assert_allclose(all_dets["mu"], dataset["mu_raw"].cumsum("group"))
+    assert_allclose(all_dets["sigma"], np.exp(dataset["sigma_raw"]))
+
+    # Test custom arguments
+    extended_with_mu = compute_deterministics(
+        dataset,
+        var_names=["mu"],
+        merge_dataset=True,
+        model=m,
+        compile_kwargs={"mode": "FAST_COMPILE"},
+        progressbar=False,
+    )
+    assert set(extended_with_mu.data_vars.variables) == {"mu_raw", "sigma_raw", "mu"}
+    assert extended_with_mu["mu"].dims == ("chain", "draw", "group")
+    assert_allclose(extended_with_mu["mu"], dataset["mu_raw"].cumsum("group"))
+
+
+def test_docstring_example():
+    import pymc as pm
+
+    with pm.Model(coords={"group": (0, 2, 4)}) as m:
+        mu_raw = pm.Normal("mu_raw", 0, 1, dims="group")
+        mu = pm.Deterministic("mu", mu_raw.cumsum(), dims="group")
+
+        trace = pm.sample(var_names=["mu_raw"], chains=2, tune=5, draws=5)
+
+    assert "mu" not in trace.posterior
+
+    with m:
+        trace.posterior = pm.compute_deterministics(trace.posterior, merge_dataset=True)
+
+    assert "mu" in trace.posterior