bayesmith.dispatch.predictive
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Predictive forward generation: push posterior draws onto observations. The two primitives of the R2 predictive seam, both built on the SAME loc/scale :func:`bayesmith.exact.gaussian.observation_parts` reads (§0.1): replay and replicate differ only by log_prob vs sample, never by a second, parallel forward model. This module is part of the dispatch bridge -- it reads a Graph and JAX, and returns a :class:`~bayesmith.artifacts.base.NamedArray` for the pointwise log-likelihood -- and it never reaches the higher artifact protocol (results, tasks, identity, refusal): those are `dispatch.task`'s business. R3 §0.7 adds the two remaining forward primitives -- `prior_draws`, which walks the whole graph from its priors, and `forward_draws`, which pushes ONE parameter setting through it -- so that `SimulationTask`'s three parameter sources share this module rather than growing a simulator each. That sharing is structural, not asserted: the fixed arm IS the replicate arm at a parameter that happens not to vary, and `tests/dispatch/test_predictive_seam.py` compares them at `rtol=0`. **Generation law (§0.3).** A replicated draw is `Normal(loc, scale).sample` at the node's own observation distribution, one per source posterior draw (draw axis one-to-one, no resampling). This is the node's own distribution -- for a diagonal Gaussian it is exactly `apply_probabilistic(...).sample` -- so density and generator are one object, not a hand-written simulator. **Coverage domain (§0.4).** `observation_parts` is a diagonal walk, so a correlated or non-Gaussian observed node raises :class:`~bayesmith.errors.NotGaussian` here rather than being silently approximated. The caller turns that into a typed Refusal.
replicated_draws
Defined in bayesmith.dispatch.predictive, line 104.
replicated_draws(graph: Graph, latent_draws: Mapping[str, Any], key: jax.Array) -> dict[str, jax.Array]`{obs: draws}` -- one replicated draw per source draw, per observed node.
Args:
graph: the model.
latent_draws: every latent's posterior draws, each with the draw axis
leading. They must cover every latent `evaluate` needs.
key: PRNG key, split once per source draw so the replicated draws do
not share a stream.
Returns:
One array per observed node, shape `(draws, *node_shape)`, sampled from
the node's own observation distribution at each source draw's loc/scale.
Raises:
NotGaussian: for a correlated or non-Gaussian observed node (§0.4) --
a classification outcome the caller turns into a Refusal, not a
silent approximation.pointwise_log_likelihood
Defined in bayesmith.dispatch.predictive, line 143.
pointwise_log_likelihood(graph: Graph, latent_draws: Mapping[str, Any]) -> NamedArray | NoneThe observed-data replay: pointwise log-likelihood, indexed by draw first.
For every source posterior draw, evaluates `Normal(loc, scale).log_prob(data)`
at the observed node's own data -- the conditioning half, NOT a posterior
predictive (§0.1) -- zeroing masked positions. The leading axis is `draw`;
the remaining axes are observation units named by `_dims` (a single
observed node keeps its plate / `{name}_dim{i}` axes, several observed
nodes are concatenated into one `observation` axis in declaration order).
Returns:
A `NamedArray` named `log_likelihood`, or `None` when the graph has no
observed node -- there is no observation unit to score, so the caller
abstains rather than fabricating a pointwise density.
Raises:
NotGaussian: for a correlated or non-Gaussian observed node (§0.4).prior_draws
Defined in bayesmith.dispatch.predictive, line 231.
prior_draws(graph: Graph, key: jax.Array, n: int) -> dict[str, jax.Array]`{node: draws}` for EVERY node, drawn from the model's priors forward.
Observed nodes are drawn too, from their own distributions rather than
carried as data -- so this is the prior PREDICTIVE, and the same primitive
answers both of the questions R3 asks of it: the `(theta, y)` pairs SBC
replicates over (§0.6) and the prior predictive check's replicated
datasets (§0.2). A version that stopped at the latents would need a
second pass to reach the data, which is the parallel simulator invariant 1
forbids.
**Key discipline: one key per node, in ``graph.nodes`` order.** Const and
Deterministic nodes consume a key they never use. That is not tidiness --
it is what probe_28 §8 measured, and reproducing it bit for bit is what
lets the plan's ``P(max|d_prior| >= max|d_obs|) = 0.210`` be CHECKED
against this function rather than merely quoted beside it
(``tests/dispatch/test_predictive_seam.py`` does exactly that). The cost
is that ``jax.random.split(k, len(nodes))`` reshuffles every node's stream
when a node is added or removed; that is acceptable because such an edit
already moves the GRAPH_STRUCTURE fingerprint, so by the protocol's own
reckoning it is a different run and not a reproducibility claim that broke.
Args:
graph: the model.
key: PRNG key, split once per draw so the draws do not share a stream.
n: how many draws. Never defaulted: a draw count is a budget, and
§0.7 puts it in the task's ``budget.draws``.
Returns:
One array per node, shape `(n, *node_shape)`, keyed by node name.forward_draws
Defined in bayesmith.dispatch.predictive, line 285.
forward_draws(graph: Graph, values: Mapping[str, Any], key: jax.Array, n: int) -> dict[str, jax.Array]`{obs: draws}` -- `n` observations generated at ONE parameter setting.
§0.7's fixed arm. Implemented by repeating the setting along a draw axis
and calling :func:`replicated_draws`, which is the point rather than a
shortcut: the fixed source and the posterior source then share ONE
generation law by construction, and a future divergence between them is a
change to one function instead of a drift between two that were never
compared. The price is `n` identical `evaluate` calls under the vmap,
paid knowingly -- a second, cheaper path is exactly the parallel forward
model §0.1 spent R2 removing.
Args:
graph: the model.
values: the parameter setting, one entry per latent `evaluate` needs.
key: PRNG key, split once per draw.
n: how many observations to generate.
Returns:
One array per observed node, shape `(n, *node_shape)`.
Raises:
NotGaussian: for a correlated or non-Gaussian observed node (§0.4),
through the same diagonal walk its sibling refuses on.