New sample_posterior_predictive API

Author: ricardoV94

pymc/sampling/forward.py

@@ -32,7 +32,6 @@
 from pytensor import tensor as pt
 from pytensor.graph import vectorize_graph
 from pytensor.graph.basic import (
-    Apply,
     Constant,
     Variable,
 )
@@ -110,9 +109,10 @@ def compile_forward_sampling_function(
     outputs: list[Variable],
     vars_in_trace: list[Variable],
     basic_rvs: list[Variable] | None = None,
-    givens_dict: dict[Variable, Any] | None = None,
     constant_data: dict[str, np.ndarray] | None = None,
     constant_coords: set[str] | None = None,
+    volatile_vars: set[Variable] | None = None,
+    freeze_vars: set[Variable] | None = None,
     **kwargs,
 ) -> tuple[Callable[..., np.ndarray | list[np.ndarray]], set[Variable]]:
     """Compile a function to draw samples, conditioned on the values of some variables.
@@ -125,12 +125,15 @@ def compile_forward_sampling_function(
     Volatile variables are variables whose values could change between runs of the
     compiled function or after inference has been run. These variables are:
 
-    - Variables in the outputs list
+    - Variables in ``volatile_vars``
     - ``SharedVariable`` instances that are not ``RandomGeneratorSharedVariable``, and whose values changed with respect to what they were at inference time
     - Variables that are in the `basic_rvs` list but not in the ``vars_in_trace`` list
-    - Variables that are keys in the ``givens_dict``
     - Variables that have volatile inputs
 
+    Variables in ``freeze_vars`` are never considered volatile, regardless of the above
+    rules. They act as volatility barriers, stopping the propagation of volatility to
+    their dependents. Frozen variables are always treated as trace inputs.
+
     Concretely, this function can be used to compile a function to sample from the
     posterior predictive distribution of a model that has variables that are conditioned
     on ``Data`` instances. The variables that depend on the mutable data that have changed
@@ -139,19 +142,11 @@ def compile_forward_sampling_function(
     ignored and new values will be computed (in the case of deterministics and potentials) or
     sampled (in the case of random variables).
 
-    This function also enables a way to impute values for any variable in the computational
-    graph that produces the desired outputs: the ``givens_dict``. This dictionary can be used
-    to set the ``givens`` argument of the pytensor function compilation. This will essentially
-    replace a node in the computational graph with any other expression that has the same
-    type as the desired node. Passing variables in the givens_dict is considered an intervention
-    that might lead to different variable values from those that could have been seen during
-    inference, as such, **any variable that is passed in the ``givens_dict`` will be considered
-    volatile**.
-
     Parameters
     ----------
     outputs : List[pytensor.graph.basic.Variable]
-        The list of variables that will be returned by the compiled function
+        The list of variables that will be returned by the compiled function.
+        Outputs are not inherently volatile.
     vars_in_trace : List[pytensor.graph.basic.Variable]
         The list of variables that are assumed to have values stored in the trace
     basic_rvs : Optional[List[pytensor.graph.basic.Variable]]
@@ -160,10 +155,6 @@ def compile_forward_sampling_function(
         be considered as random variable instances. This includes variables that have
         a ``RandomVariable`` owner op, but also unpure random variables like Mixtures, or
         Censored distributions.
-    givens_dict : Optional[Dict[pytensor.graph.basic.Variable, Any]]
-        A dictionary that maps tensor variables to the values that should be used to replace them
-        in the compiled function. The types of the key and value should match or an error will be
-        raised during compilation.
     constant_data : Optional[Dict[str, numpy.ndarray]]
         A dictionary that maps the names of ``Data`` instances to their
         corresponding values at inference time. If a model was created with ``Data``, these
@@ -183,6 +174,14 @@ def compile_forward_sampling_function(
         which case, it is considered volatile. If a ``SharedVariable`` is not found
         in either ``constant_data`` or ``constant_coords``, then it is assumed to be volatile.
         Setting ``constant_coords`` to ``None`` is equivalent to passing an empty set.
+    volatile_vars : Optional[Set[pytensor.graph.basic.Variable]]
+        Variables that are unconditionally volatile. Volatility propagates from these
+        to their dependents in the graph.
+    freeze_vars : Optional[Set[pytensor.graph.basic.Variable]]
+        A set of variables that should never be considered volatile, even if they would
+        otherwise be due to having volatile inputs or depending on changed data. Frozen
+        variables act as volatility barriers: they stop the propagation of volatility to
+        their dependents and are always treated as inputs that pull values from the trace.
 
     Returns
     -------
@@ -192,16 +191,17 @@ def compile_forward_sampling_function(
         Set of all basic_rvs that were considered volatile and will be resampled when
         the function is evaluated
     """
-    if givens_dict is None:
-        givens_dict = {}
-
     if basic_rvs is None:
         basic_rvs = []
 
     if constant_data is None:
         constant_data = {}
     if constant_coords is None:
         constant_coords = set()
+    if volatile_vars is None:
+        volatile_vars = set()
+    if freeze_vars is None:
+        freeze_vars = set()
 
     # We define a helper function to check if shared values match to an array
     def shared_value_matches(var):
@@ -221,9 +221,10 @@ def shared_value_matches(var):
     )  # type: ignore[call-overload]
     volatile_nodes: set[Any] = set()
     for node in nodes:
+        if node in freeze_vars:
+            continue  # Frozen variables are never volatile, and block propagation
         if (
-            node in fg.outputs
-            or node in givens_dict
+            node in volatile_vars
             or (  # SharedVariables, except RandomState/Generators
                 isinstance(node, SharedVariable)
                 and not isinstance(node, RandomGeneratorSharedVariable)
@@ -260,20 +261,9 @@ def expand(node):
     # the entire graph
     list(walk(fg.outputs, expand))
 
-    # Populate the givens list
-    givens = [
-        (
-            node,
-            value
-            if isinstance(value, Variable | Apply)
-            else pt.constant(value, dtype=getattr(node, "dtype", None), name=node.name),
-        )
-        for node, value in givens_dict.items()
-    ]
-
     return (
-        compile(inputs, fg.outputs, givens=givens, on_unused_input="ignore", **kwargs),
-        set(basic_rvs) & (volatile_nodes - set(givens_dict)),  # Basic RVs that will be resampled
+        compile(inputs, fg.outputs, on_unused_input="ignore", **kwargs),
+        set(basic_rvs) & volatile_nodes,  # Basic RVs that will be resampled
     )
 
 
@@ -450,7 +440,6 @@ def sample_prior_predictive(
         vars_to_sample,
         vars_in_trace=[],
         basic_rvs=model.basic_RVs,
-        givens_dict=None,
         random_seed=random_seed,
         **compile_kwargs,
     )
@@ -490,6 +479,8 @@ def sample_posterior_predictive(
     predictions: bool = False,
     idata_kwargs: dict | None = None,
     compile_kwargs: dict | None = None,
+    sample_vars: list[str] | None = None,
+    freeze_vars: list[str] | None = None,
 ) -> InferenceData: ...
 @overload
 def sample_posterior_predictive(
@@ -505,6 +496,8 @@ def sample_posterior_predictive(
     predictions: bool = False,
     idata_kwargs: dict | None = None,
     compile_kwargs: dict | None = None,
+    sample_vars: list[str] | None = None,
+    freeze_vars: list[str] | None = None,
 ) -> dict[str, np.ndarray]: ...
 def sample_posterior_predictive(
     trace,
@@ -519,6 +512,8 @@ def sample_posterior_predictive(
     predictions: bool = False,
     idata_kwargs: dict | None = None,
     compile_kwargs: dict | None = None,
+    sample_vars: list[str] | None = None,
+    freeze_vars: list[str] | None = None,
 ) -> InferenceData | dict[str, np.ndarray]:
     """Generate forward samples for `var_names`, conditioned on the posterior samples of variables found in the `trace`.
 
@@ -539,9 +534,10 @@ def sample_posterior_predictive(
         Model to be used to generate the posterior predictive samples. It will
         generally be the model used to generate the `trace`, but it doesn't need to be.
     var_names : Iterable[str], optional
-        Names of variables for which to compute the posterior predictive samples.
-        By default, only observed variables are sampled.
-        See the example below for what happens when this argument is customized.
+        Names of variables to include in the returned dataset. This only controls which
+        variables appear in the output, not which variables are resampled. Use ``sample_vars``
+        to control resampling. By default, observed variables and their dependent
+        deterministics are included.
     sample_dims : list of str, optional
         Dimensions over which to loop and generate posterior predictive samples.
         When ``sample_dims`` is ``None`` (default) both "chain" and "draw" are considered sample
@@ -567,6 +563,16 @@ def sample_posterior_predictive(
         :func:`pymc.predictions_to_inference_data` otherwise.
     compile_kwargs: dict, optional
         Keyword arguments for :func:`pymc.pytensorf.compile`.
+    sample_vars : list of str, optional
+        Names of unobserved variables that should be explicitly resampled. Observed variables
+        are always resampled. Use this to request resampling of specific unobserved variables
+        (e.g., for out-of-model predictions or forecasting). Variables not in ``sample_vars``
+        will have their values taken from the trace if available.
+    freeze_vars : list of str, optional
+        Names of variables that should always be reused from the trace, even if they would
+        otherwise be resampled due to depending on changed data or other resampled variables.
+        Frozen variables act as barriers that stop the propagation of "volatility" in the
+        computational graph. Must be present in the trace. Cannot overlap with ``sample_vars``.
 
     Returns
     -------
@@ -859,17 +865,39 @@ def sample_posterior_predictive(
 
     constant_coords = get_constant_coords(trace_coords, model)
 
+    # Resolve output variables (what to return in the dataset)
     if var_names is not None:
-        vars_ = [model[x] for x in var_names]
+        output_vars = [model[x] for x in var_names]
     else:
         observed_vars = model.observed_RVs
         if observed_data is not None:
             observed_vars += [
                 model[x] for x in observed_data if x in model and x not in observed_vars
             ]
-        vars_ = observed_vars + observed_dependent_deterministics(model, observed_vars)
-
-    vars_to_sample = list(get_default_varnames(vars_, include_transformed=False))
+        output_vars = observed_vars + observed_dependent_deterministics(model, observed_vars)
+
+    # Resolve variables to resample
+    # Observed variables are always resampled, plus any explicit sample_vars
+    resample_vars: list[Variable] = list(model.observed_RVs)
+    if sample_vars is not None:
+        resample_vars += [model[x] for x in sample_vars]
+
+    # Compiled function outputs = resample_vars + dependent deterministics +
+    # deterministics from var_names (they need recomputation, not "resampling")
+    basic_rv_set = set(model.basic_RVs)
+    compiled_outputs = list(resample_vars)
+    compiled_outputs += observed_dependent_deterministics(model, resample_vars)
+    # Add deterministics from var_names that need recomputation
+    for var in output_vars:
+        if var not in basic_rv_set and var not in compiled_outputs:
+            compiled_outputs.append(var)
+
+    vars_to_sample = list(
+        get_default_varnames(
+            list({v.name: v for v in compiled_outputs}.values()),
+            include_transformed=False,
+        )
+    )
 
     if not vars_to_sample:
         if return_inferencedata and not extend_inferencedata:
@@ -880,6 +908,27 @@ def sample_posterior_predictive(
 
     vars_in_trace = get_vars_in_point_list(_trace, model)
 
+    # Resolve freeze vars
+    frozen: set[Variable] = set()
+    if freeze_vars is not None:
+        frozen = {model[x] for x in freeze_vars}
+        # Validate: freeze_vars must be in trace
+        vars_in_trace_names = {v.name for v in vars_in_trace}
+        missing = {x for x in freeze_vars if x not in vars_in_trace_names}
+        if missing:
+            raise ValueError(
+                f"freeze_vars {sorted(missing)} are not present in the trace. "
+                f"Cannot freeze variables without stored values."
+            )
+        # Validate: freeze_vars and sample_vars must be disjoint
+        if sample_vars is not None:
+            overlap = set(freeze_vars) & set(sample_vars)
+            if overlap:
+                raise ValueError(
+                    f"Variables {sorted(overlap)} are in both sample_vars and freeze_vars. "
+                    f"A variable cannot be both resampled and frozen."
+                )
+
     if random_seed is not None:
         (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
@@ -892,15 +941,41 @@ def sample_posterior_predictive(
         outputs=vars_to_sample,
         vars_in_trace=vars_in_trace,
         basic_rvs=model.basic_RVs,
-        givens_dict=None,
         random_seed=random_seed,
         constant_data=constant_data,
         constant_coords=constant_coords,
+        volatile_vars=set(resample_vars),
+        freeze_vars=frozen,
         **compile_kwargs,
     )
     sampler_fn = point_wrapper(_sampler_fn)
+
+    # Warn about implicitly volatile trace variables
+    vars_in_trace_set = set(vars_in_trace)
+    sample_var_names = set(sample_vars) if sample_vars is not None else set()
+    implicit_volatile = {
+        rv
+        for rv in volatile_basic_rvs
+        if rv in vars_in_trace_set and rv.name not in sample_var_names
+    }
+    if implicit_volatile:
+        implicit_names = sorted(rv.name for rv in implicit_volatile)  # type: ignore[type-var]
+        warnings.warn(
+            f"Variables {implicit_names} are in the trace but will be resampled because "
+            f"they depend on data/coords that changed. To silence this warning, add them "
+            f"to `sample_vars` explicitly, or add them to `freeze_vars` to reuse their "
+            f"trace values.",
+            UserWarning,
+            stacklevel=2,
+        )
+
     # All model variables have a name, but mypy does not know this
     _log.info(f"Sampling: {sorted(volatile_basic_rvs, key=lambda var: var.name)}")  # type: ignore[arg-type, return-value]
+
+    # Determine output-only variables that should be copied from trace, not sampled
+    sampled_names = {v.name for v in vars_to_sample}
+    copy_from_trace_names = [var.name for var in output_vars if var.name not in sampled_names]
+
     ppc_trace_t = _DefaultTrace(samples)
 
     progress = create_simple_progress(
@@ -928,9 +1003,14 @@ def sample_posterior_predictive(
 
                 values = sampler_fn(**param)
 
-                for k, v in zip(vars_, values):
+                for k, v in zip(vars_to_sample, values):
                     ppc_trace_t.insert(k.name, v, idx)
 
+                # Copy output-only variables from trace
+                for name in copy_from_trace_names:
+                    if name in param:
+                        ppc_trace_t.insert(name, param[name], idx)
+
                 progress.advance(task)
             progress.update(task, refresh=True, completed=samples)
 
@@ -939,6 +1019,12 @@ def sample_posterior_predictive(
 
     ppc_trace = ppc_trace_t.trace_dict
 
+    # Filter to only include requested output variables and sample_vars
+    output_var_names = {v.name for v in output_vars}
+    if sample_vars is not None:
+        output_var_names |= set(sample_vars)
+    ppc_trace = {k: v for k, v in ppc_trace.items() if k in output_var_names}
+
     for k, ary in ppc_trace.items():
         if stacked_dims is not None:
             ppc_trace[k] = ary.reshape(

tests/model/transform/test_conditioning.py

@@ -187,7 +187,7 @@ def test_do_posterior_predictive():
     # 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")
+        idata_do = pm.sample_posterior_predictive(idata_m, var_names=["z"], sample_vars=["z"])
 
     assert 120 < idata_do.posterior_predictive["z"].mean() < 130
 
@@ -315,7 +315,9 @@ def test_do_sample_posterior_predictive(make_interventions_shared):
     idata = az.from_dict({"a": [[1.0]], "b": [[2.0]], "c": [[1.0]]})
 
     with do(model, {a: 1000}, make_interventions_shared=make_interventions_shared):
-        pp = sample_posterior_predictive(idata, var_names=["c"], predictions=True).predictions
+        pp = sample_posterior_predictive(
+            idata, var_names=["c"], sample_vars=["c"], predictions=True
+        ).predictions
     assert (pp["c"] > 500).all()