New sample_posterior_predictive API

Author: ricardoV94

pymc/sampling/forward.py

@@ -113,6 +113,8 @@ def compile_forward_sampling_function(
     givens_dict: dict[Variable, Any] | None = None,
     constant_data: dict[str, np.ndarray] | None = None,
     constant_coords: set[str] | None = None,
+    freeze_vars: set[Variable] | None = None,
+    volatile_outputs: 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.
@@ -131,6 +133,10 @@ def compile_forward_sampling_function(
     - 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
@@ -183,6 +189,17 @@ 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.
+    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, being in the outputs list, 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.
+    volatile_outputs : Optional[Set[pytensor.graph.basic.Variable]]
+        A subset of ``outputs`` that should be considered volatile. If provided, only these
+        outputs (rather than all outputs) are marked as volatile. Outputs not in this set
+        will still be computed but won't trigger volatility propagation. If ``None`` (default),
+        all outputs are considered volatile (preserving backward compatibility).
 
     Returns
     -------
@@ -202,6 +219,11 @@ def compile_forward_sampling_function(
         constant_data = {}
     if constant_coords is None:
         constant_coords = set()
+    if freeze_vars is None:
+        freeze_vars = set()
+    # If volatile_outputs is not specified, all outputs are volatile (backward compatible)
+    # If specified, only those outputs are marked volatile
+    _volatile_outputs: set[Variable] | None = volatile_outputs
 
     # We define a helper function to check if shared values match to an array
     def shared_value_matches(var):
@@ -221,8 +243,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
+            (node in fg.outputs if _volatile_outputs is None else node in _volatile_outputs)
             or node in givens_dict
             or (  # SharedVariables, except RandomState/Generators
                 isinstance(node, SharedVariable)
@@ -504,6 +528,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(
@@ -519,6 +545,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,
@@ -533,6 +561,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`.
 
@@ -553,9 +583,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
@@ -581,6 +612,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
     -------
@@ -873,17 +914,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:
@@ -894,6 +957,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)
 
@@ -902,6 +986,10 @@ def sample_posterior_predictive(
     compile_kwargs.setdefault("allow_input_downcast", True)
     compile_kwargs.setdefault("accept_inplace", True)
 
+    # Only resample_vars should be volatile outputs. Other outputs (deterministics from
+    # var_names) are just computed from the graph without triggering volatility cascades.
+    _volatile_outputs = set(resample_vars) & set(vars_to_sample)
+
     _sampler_fn, volatile_basic_rvs = compile_forward_sampling_function(
         outputs=vars_to_sample,
         vars_in_trace=vars_in_trace,
@@ -910,11 +998,38 @@ def sample_posterior_predictive(
         random_seed=random_seed,
         constant_data=constant_data,
         constant_coords=constant_coords,
+        freeze_vars=frozen,
+        volatile_outputs=_volatile_outputs,
         **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)
+        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(
@@ -942,9 +1057,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)
 
@@ -953,6 +1073,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()