Expand observe transformation to Deterministics

Author: ricardoV94

pymc_experimental/model_transform/conditioning.py

@@ -5,6 +5,7 @@
 from pytensor.tensor import TensorVariable
 
 from pymc_experimental.utils.model_fgraph import (
+    ModelDeterministic,
     ModelFreeRV,
     extract_dims,
     fgraph_from_model,
@@ -16,7 +17,7 @@
 
 
 def observe(model: Model, vars_to_observations: Dict[Union["str", TensorVariable], Any]) -> Model:
-    """Convert free RVs to observed RVs.
+    """Convert free RVs or Deterministics to observed RVs.
 
     Parameters
     ----------
@@ -47,29 +48,48 @@ def observe(model: Model, vars_to_observations: Dict[Union["str", TensorVariable
 
         m_new = observe(m, {y: 0.5})
 
+    Deterministic variables can also be observed.
+    This relies on PyMC ability to infer the logp of the underlying expression
+
+    .. 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.dist(x, shape=(5,))
+            y_censored = pm.Deterministic("y_censored", pm.math.clip(y, -1, 1))
+
+        new_m = observe(m, {y_censored: [0.9, 0.5, 0.3, 1, 1]})
+
+
     """
     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")
+    valid_model_vars = set(model.free_RVs + model.deterministics)
+    if any(var not in valid_model_vars for var in vars_to_observations):
+        raise ValueError(f"At least one var is not a free variable or deterministic in the model")
 
     fgraph, memo = fgraph_from_model(model)
 
     replacements = {}
     for var, obs in vars_to_observations.items():
-        model_free_rv = memo[var]
+        model_var = memo[var]
 
         # Just a sanity check
-        assert isinstance(model_free_rv.owner.op, ModelFreeRV)
-        assert model_free_rv in fgraph.variables
+        assert isinstance(model_var.owner.op, (ModelFreeRV, ModelDeterministic))
+        assert model_var 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
+        var = model_var.owner.inputs[0]
+        var.name = model_var.name
+        dims = extract_dims(var)
+        model_obs_rv = model_observed_rv(var, var.type.filter_variable(obs), *dims)
+        replacements[model_var] = model_obs_rv
 
     toposort_replace(fgraph, tuple(replacements.items()))
 

pymc_experimental/tests/model_transform/test_conditioning.py

@@ -66,6 +66,26 @@ def test_observe_minibatch():
     )
 
 
+def test_observe_deterministic():
+    y_censored_obs = [0.9, 0.5, 0.3, 1, 1]
+
+    with pm.Model() as m_old:
+        x = pm.Normal("x")
+        y = pm.Normal.dist(x, shape=(5,))
+        y_censored = pm.Deterministic("y_censored", pm.math.clip(y, -1, 1))
+
+    m_new = observe(m_old, {y_censored: y_censored_obs})
+
+    with pm.Model() as m_ref:
+        x = pm.Normal("x")
+        pm.Censored("y_censored", pm.Normal.dist(x), lower=-1, upper=1, observed=y_censored_obs)
+
+    np.testing.assert_allclose(
+        m_new.compile_logp()({"x": 0.9}),
+        m_ref.compile_logp()({"x": 0.9}),
+    )
+
+
 def test_do():
     with pm.Model() as m_old:
         x = pm.Normal("x", 0, 1e-3)