Implement observe and do model transformations

Author: ricardoV94

docs/api_reference.rst

@@ -35,6 +35,17 @@ Distributions
    histogram_approximation
 
 
+Model Transformations
+=====================
+
+.. currentmodule:: pymc_experimental.model_transform
+.. autosummary::
+   :toctree: generated/
+
+   conditioning.do
+   conditioning.observe
+
+
 Utils
 =====
 

pymc_experimental/model_transform/__init__.py

@@ -0,0 +1,165 @@
+from typing import Any, Dict, List, Sequence, Union
+
+from pymc import Model
+from pymc.pytensorf import _replace_vars_in_graphs
+from pytensor.tensor import TensorVariable
+
+from pymc_experimental.utils.model_fgraph import (
+    ModelFreeRV,
+    extract_dims,
+    fgraph_from_model,
+    model_from_fgraph,
+    model_named,
+    model_observed_rv,
+    toposort_replace,
+)
+
+
+def observe(model: Model, vars_to_observations: Dict[Union["str", TensorVariable], Any]) -> Model:
+    """Convert free RVs to observed RVs.
+
+    Parameters
+    ----------
+    model: PyMC Model
+    vars_to_observations: Dict of variable or name to TensorLike
+        Dictionary that maps model variables (or names) to observed values.
+        Observed values must have a shape and data type that is compatible
+        with the original model variable.
+
+    Returns
+    -------
+    new_model: PyMC model
+        A distinct PyMC model with the relevant variables observed.
+        All remaining variables are cloned and can be retrieved via `new_model["var_name"]`.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import pymc as pm
+        from pymc_experimental.model_transform.conditioning import observe
+
+        with pm.Model() as m:
+            x = pm.Normal("x")
+            y = pm.Normal("y", x)
+            z = pm.Normal("z", y)
+
+        m_new = observe(m, {y: 0.5})
+
+    """
+    vars_to_observations = {
+        model[var] if isinstance(var, str) else var: obs
+        for var, obs in vars_to_observations.items()
+    }
+
+    # Note: Since PyMC can infer logprob expressions we could also allow observing Deterministics
+    if any(var not in model.free_RVs for var in vars_to_observations):
+        raise ValueError(f"At least one var is not a free variable in the model")
+
+    fgraph, memo = fgraph_from_model(model)
+
+    replacements = {}
+    for var, obs in vars_to_observations.items():
+        model_free_rv = memo[var]
+
+        # Just a sanity check
+        assert isinstance(model_free_rv.owner.op, ModelFreeRV)
+        assert model_free_rv in fgraph.variables
+
+        rv, vv, *dims = model_free_rv.owner.inputs
+        model_obs_rv = model_observed_rv(rv, rv.type.filter_variable(obs), *dims)
+        replacements[model_free_rv] = model_obs_rv
+
+    toposort_replace(fgraph, tuple(replacements.items()))
+
+    return model_from_fgraph(fgraph)
+
+
+def replace_vars_in_graphs(graphs: Sequence[TensorVariable], replacements) -> List[TensorVariable]:
+    def replacement_fn(var, inner_replacements):
+        if var in replacements:
+            inner_replacements[var] = replacements[var]
+
+        # Handle root inputs as those will never be passed to the replacement_fn
+        for inp in var.owner.inputs:
+            if inp.owner is None and inp in replacements:
+                inner_replacements[inp] = replacements[inp]
+
+        return [var]
+
+    replaced_graphs, _ = _replace_vars_in_graphs(graphs=graphs, replacement_fn=replacement_fn)
+    return replaced_graphs
+
+
+def do(model: Model, vars_to_interventions: Dict[Union["str", TensorVariable], Any]) -> Model:
+    """Replace model variables by intervention variables.
+
+    Parameters
+    ----------
+    model: PyMC Model
+    vars_to_interventions: Dict of variable or name to TensorLike
+        Dictionary that maps model variables (or names) to intervention expressions.
+        Intervention expressions must have a shape and data type that is compatible
+        with the original model variable.
+
+    Returns
+    -------
+    new_model: PyMC model
+        A distinct PyMC model with the relevant variables replaced by the intervention expressions.
+        All remaining variables are cloned and can be retrieved via `new_model["var_name"]`.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import pymc as pm
+        from pymc_experimental.model_transform.conditioning import do
+
+        with pm.Model() as m:
+            x = pm.Normal("x", 0, 1)
+            y = pm.Normal("y", x, 1)
+            z = pm.Normal("z", y + x, 1)
+
+        # Dummy posterior, same as calling `pm.sample`
+        idata_m = az.from_dict({rv.name: [pm.draw(rv, draws=500)] for rv in [x, y, z]})
+
+        # Replace `y` by a constant `100.0`
+        m_do = do(m, {y: 100.0})
+        with m_do:
+            idata_do = pm.sample_posterior_predictive(idata_m, var_names="z")
+
+    """
+    do_mapping = {}
+    for var, obs in vars_to_interventions.items():
+        if isinstance(var, str):
+            var = model[var]
+        do_mapping[var] = var.type.filter_variable(obs)
+
+    if any(var not in (model.basic_RVs + model.deterministics) for var in do_mapping):
+        raise ValueError(f"At least one var is not a variable or deterministic in the model")
+
+    fgraph, memo = fgraph_from_model(model)
+
+    # We need the interventions defined in terms of the IR fgraph representation,
+    # In case they reference other variables in the model
+    ir_interventions = replace_vars_in_graphs(list(do_mapping.values()), replacements=memo)
+
+    replacements = {}
+    for var, intervention in zip(do_mapping, ir_interventions):
+        model_var = memo[var]
+
+        # Just a sanity check
+        assert model_var in fgraph.variables
+
+        intervention.name = model_var.name
+        dims = extract_dims(model_var)
+        new_var = model_named(intervention, *dims)
+
+        replacements[model_var] = new_var
+
+    # Replace variables by interventions
+    toposort_replace(fgraph, tuple(replacements.items()))
+
+    return model_from_fgraph(fgraph)

pymc_experimental/model_transform/conditioning.py

@@ -0,0 +1,122 @@
+import arviz as az
+import numpy as np
+import pymc as pm
+from pymc.variational.minibatch_rv import create_minibatch_rv
+
+from pymc_experimental.model_transform.conditioning import do, observe
+
+
+def test_observe():
+    with pm.Model() as m_old:
+        x = pm.Normal("x")
+        y = pm.Normal("y", x)
+        z = pm.Normal("z", y)
+
+    m_new = observe(m_old, {y: 0.5})
+
+    assert len(m_new.free_RVs) == 2
+    assert len(m_new.observed_RVs) == 1
+    assert m_new["x"] in m_new.free_RVs
+    assert m_new["y"] in m_new.observed_RVs
+    assert m_new["z"] in m_new.free_RVs
+
+    np.testing.assert_allclose(
+        m_old.compile_logp()({"x": 0.9, "y": 0.5, "z": 1.4}),
+        m_new.compile_logp()({"x": 0.9, "z": 1.4}),
+    )
+
+    # Test two substitutions
+    m_new = observe(m_old, {y: 0.5, z: 1.4})
+
+    assert len(m_new.free_RVs) == 1
+    assert len(m_new.observed_RVs) == 2
+    assert m_new["x"] in m_new.free_RVs
+    assert m_new["y"] in m_new.observed_RVs
+    assert m_new["z"] in m_new.observed_RVs
+
+    np.testing.assert_allclose(
+        m_old.compile_logp()({"x": 0.9, "y": 0.5, "z": 1.4}),
+        m_new.compile_logp()({"x": 0.9}),
+    )
+
+
+def test_observe_minibatch():
+    data = np.zeros((100,))
+    batch_size = (10,)
+    with pm.Model() as m_old:
+        x = pm.Normal("x")
+        y = pm.Normal("y", x)
+        # Minibatch RVs are usually created with `total_size` kwarg
+        z_raw = pm.Normal.dist(y, shape=batch_size)
+        mb_z = create_minibatch_rv(z_raw, total_size=data.shape)
+        m_old.register_rv(mb_z, name="mb_z")
+
+    mb_data = pm.Minibatch(data, batch_size=batch_size)
+    m_new = observe(m_old, {mb_z: mb_data})
+
+    assert len(m_new.free_RVs) == 2
+    assert len(m_new.observed_RVs) == 1
+    assert m_new["x"] in m_new.free_RVs
+    assert m_new["y"] in m_new.free_RVs
+    assert m_new["mb_z"] in m_new.observed_RVs
+
+    np.testing.assert_allclose(
+        m_old.compile_logp()({"x": 0.9, "y": 0.5, "mb_z": np.zeros(10)}),
+        m_new.compile_logp()({"x": 0.9, "y": 0.5}),
+    )
+
+
+def test_do():
+    with pm.Model() as m_old:
+        x = pm.Normal("x", 0, 1e-3)
+        y = pm.Normal("y", x, 1e-3)
+        z = pm.Normal("z", y + x, 1e-3)
+
+    assert -5 < pm.draw(z) < 5
+
+    m_new = do(m_old, {y: x + 100})
+
+    assert len(m_new.free_RVs) == 2
+    assert m_new["x"] in m_new.free_RVs
+    assert m_new["y"] in m_new.named_vars.values()
+    assert m_new["z"] in m_new.free_RVs
+
+    assert 95 < pm.draw(m_new["z"]) < 105
+
+    # Test two substitutions
+    with m_old:
+        switch = pm.MutableData("switch", 1)
+    m_new = do(m_old, {y: 100 * switch, x: 100 * switch})
+
+    assert len(m_new.free_RVs) == 1
+    assert m_new["x"] in m_new.named_vars.values()
+    assert m_new["y"] in m_new.named_vars.values()
+    assert m_new["z"] in m_new.free_RVs
+
+    assert 195 < pm.draw(m_new["z"]) < 205
+    with m_new:
+        pm.set_data({"switch": 0})
+    assert -5 < pm.draw(m_new["z"]) < 5
+
+
+def test_do_posterior_predictive():
+    with pm.Model() as m:
+        x = pm.Normal("x", 0, 1)
+        y = pm.Normal("y", x, 1)
+        z = pm.Normal("z", y + x, 1e-3)
+
+    # Dummy posterior
+    idata_m = az.from_dict(
+        {
+            "x": np.full((2, 500), 25),
+            "y": np.full((2, 500), np.nan),
+            "z": np.full((2, 500), np.nan),
+        }
+    )
+
+    # Replace `y` by a constant `100.0`
+    m_do = do(m, {y: 100.0})
+    with m_do:
+        idata_do = pm.sample_posterior_predictive(idata_m, var_names="z")
+
+    np.testing.assert_allclose(idata_do.posterior_predictive["z"].mean(), 125)

pymc_experimental/tests/model_transform/__init__.py

@@ -326,3 +326,13 @@ def clone_model(model: Model) -> Tuple[Model]:
 
     """
     return model_from_fgraph(fgraph_from_model(model)[0])
+
+
+def extract_dims(var) -> Tuple:
+    dims = ()
+    if isinstance(var, ModelVar):
+        if isinstance(var, ModelValuedVar):
+            dims = var.inputs[2:]
+        else:
+            dims = var.inputs[1:]
+    return dims