bayesmith.dispatch
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Structural dispatch: deriving what to run from the graph computation. :mod:`bayesmith.dispatch.classify` answers one question -- which latents an exact linear-Gaussian method applies to, how groups form, and which method the group needs. Affinity is discovered structurally and checked numerically; support and prediction-dependent covariance then constrain the method. The classifier produces no samples of its own.
partition Classification block_at prior_environment SIGMA_RTOL FACTOR_METHODS FactorPlan factor_partition first_fit sample_factors
partition
Defined in bayesmith.dispatch.classify, line 526.
partition(graph: Graph, *, key: jax.Array | None=None) -> ClassificationDerive the exact block and its method from the graph's structure.
Classification
Defined in bayesmith.dispatch.classify, line 81.
class Classification
exact: tuple[str, ...]
nuts: tuple[str, ...]
method: str
reason: str
linearity: dict[int, dict[float, float]] | None
sigma_movement: float | None
sigma_needs_rebuild: bool
structure: dict | None = None
discovery: Any = NoneThe partition of one graph, and the evidence behind it.
Attributes:
exact: the single exact block, sorted. Empty if there is none.
nuts: every other latent, sorted.
method: ``"gcr"``, ``"gcr+snis"``, ``"gcr+mh"`` or ``"nuts"``.
reason: why -- named members on a refusal, so a user one declaration
away from an exact solve can see that from the plan.
linearity: ``check_linearity``'s per-at-point errors, or ``None``.
sigma_movement: the largest relative movement of sigma with the
block, or ``None`` if there is no block. Measured from TWO
anchors -- prior-scale probes around the prior centre, and the
block's own Wiener solution, which is where the data put the
posterior. See :func:`_sigma_movement` for why one of them is not
enough.
sigma_needs_rebuild: whether any observed node's scale has a latent
ANCESTOR outside the block, in which case ``noise_std`` must be
recomputed every sweep rather than hoisted. Distinct from
``sigma_movement``, which only sees movement with the BLOCK.block_at
Defined in bayesmith.dispatch.classify, line 120.
block_at(graph: Graph, names: tuple[str, ...], *, env: dict[str, Any] | None=None) -> dict[str, Any]``{outside latent: prior centre}`` -- where a block is built.
A block is affine *given* the latents outside it, so it has to be
anchored somewhere, and the prior's centre is the "before seeing
anything" point. Anchoring at zero instead is not equivalent and is not
merely a different starting guess: ``affine_only_at_zero``'s prediction is
affine exactly where its outside latent is zero, so a zero anchor accepts
at the caller's own at-point a model that the prior's centre refuses
there.prior_environment
Defined in bayesmith.exact._environment, line 59.
prior_environment(graph: Graph) -> dict[str, Any]Every node's value, with each latent at the centre of its own prior. Repeats :func:`~bayesmith.graph.evaluate.evaluate`'s isinstance ladder for the same reason :func:`~bayesmith.exact.block._env_before` does: the values it substitutes have to be derived DURING the scan, from what is already in hand. A latent whose prior width is another latent's value -- ``shared_ancestor``, ``mixed_radiometer`` -- cannot be centred before its ancestor has been. Public because :mod:`bayesmith.dispatch.plan` must anchor the block it BUILDS at exactly the point :func:`partition` classified it at. Two independent spellings of "the prior mean" would let the plan solve a different problem from the one that was checked.
SIGMA_RTOL
Defined in bayesmith.dispatch.classify, line 63.
SIGMA_RTOL = 1e-08FACTOR_METHODS
Defined in bayesmith.dispatch.factor, line 85.
FACTOR_METHODS: tuple[str, ...] = ('gcr', 'log-gcr', 'nuts')FactorPlan
Defined in bayesmith.dispatch.factor, line 114.
class FactorPlan(eqx.Module)
blocks: tuple[Block, ...]
log_space: LogSpace | NoneThe factor partition of one graph, ready to sweep.
Attributes:
blocks: closed-form blocks first (``"gcr"`` then ``"log-gcr"``, each
in discovery order), then at most one ``"nuts"`` block holding
everything neither probe accepted -- whose ``reason`` names, per
latent, why not.
log_space: the transformed graph the ``"log-gcr"`` blocks solve
against, or ``None`` when there are none. Carried on the plan
because every rebuild in a sweep needs it, and reconstructing it
per sweep would re-run the scenario probes hundreds of times for
an answer that cannot change.FactorPlan.exact
exact(self) -> tuple[Block, ...]The closed-form blocks, in sweep order.
FactorPlan.nuts
nuts(self) -> tuple[str, ...]The latents left to NUTS.
factor_partition
Defined in bayesmith.dispatch.factor, line 186.
factor_partition(graph: Graph, *, key: jax.Array | None=None, scales: Sequence[float]=DEFAULT_SCALES, log_scales: Sequence[float]=LOG_DEFAULT_SCALES, rtol: float | None=None, _discovery=None, _linear_only: bool=False) -> FactorPlanDerive the factor partition: exact blocks by structural discovery, NUTS rest.
The linear pass uses shared automatic affinity discovery: Gaussian prior
eligibility, latent ancestry, a primal symbolic-complement certificate,
and independent numerical probes. Final candidate groups are certified
jointly rather than inferred from pairwise probes. The log pass takes what
remains, transforms the graph once, and repeats its numerical checks;
nothing is declared for it because the probe answers the question a
declaration would merely assert.
Args:
graph: the model.
key: PRNG key for every probe and kappa sweep. Fixed by default, so a
partition is reproducible.
scales, log_scales: probe magnitudes for the two passes, in units of
each latent's own prior width. Separate arguments because the
linear default's ``1e3`` entry, pushed through an exponential,
measures the dtype rather than the model -- see
:data:`~bayesmith.exact.loglinear.LOG_DEFAULT_SCALES`.
rtol: forwarded to every affinity check.
Returns:
A :class:`FactorPlan`. Its ``str`` is the readable form, in the same
shape the package front page sketches.
Raises:
GraphError: if an evidence term covers a latent assigned to an exact
block, whose conditional solve would omit that term. Put the
covered latent in NUTS or keep its likelihood explicit.first_fit
Defined in bayesmith.dispatch.factor, line 88.
first_fit(names: Sequence[str], compatible: Callable[[str, str], bool]) -> list[list[str]]Group ``names`` so every pair within a group is compatible. First-fit in the caller's order: each name joins the first existing group whose every member it is compatible with, and otherwise opens a new one. One function for however many probes call it -- the linear pass, the log-linear pass, and any other package that partitions by a pairwise predicate -- so the grouping RULE cannot drift between them. No pair is asked twice, and no cache is needed to say so: the groups are disjoint, so a name meets any other name in at most one group, and ``all()`` short-circuits the moment a group is ruled out.
sample_factors
Defined in bayesmith.dispatch.factor, line 711.
sample_factors(graph: Graph, plan: FactorPlan, key: jax.Array, *, num_warmup: int=1000, num_samples: int=2000, maxiter: int | None=None, nuts_options: dict[str, Any] | None=None, on_sweep: Callable[[SweepReport], None] | None=None) -> dict[str, jax.Array]Draw from the posterior by sweeping the plan's blocks.
Every closed-form block is updated by an exact conditional draw --
:func:`~bayesmith.exact.solve.gcr_sample` on its own source graph --
conditioning on the LATEST values of every other latent, in plan order.
With no NUTS remainder the sweep is an exact Gibbs sampler and runs as a
plain loop here; with one, the sweep becomes ``HMCGibbs``'s ``gibbs_fn``
and NUTS advances the remainder between sweeps, exactly as the
single-block :func:`~bayesmith.exact.gibbs.assemble` arranges.
Args:
graph: the model the plan was derived from.
plan: from :func:`factor_partition`. Its per-block ``tol`` is used
for every solve -- the compile-time choice standing in for the
in-sweep guard, for the reason
:func:`~bayesmith.exact.gibbs.gibbs_factory` gives.
key: PRNG key for the whole run.
num_warmup, num_samples: sweeps discarded and kept.
maxiter: CG iteration cap per solve.
nuts_options: forwarded to the inner NUTS kernel, when there is one.
on_sweep: called with a :class:`SweepReport` after every sweep, warmup
included. **Refused on a plan with a NUTS remainder**, and the
reason is measured rather than defensive: with one, the sweep
becomes ``HMCGibbs``'s ``gibbs_fn``, which numpyro TRACES -- on a
two-block plan over five sweeps it is entered twice at the Python
level. A callback there would fire once, at trace time, and report
a sweep that never happened, with plausible values in it.
The inner NUTS kernel, when there is one, is initialised at the latents'
prior centres unless ``nuts_options`` names its own ``init_strategy`` --
numpyro's default draws the initial point from the priors in
unconstrained space, and on a model whose prediction can go negative out
there (a radiometer likelihood under a wide sky prior) that manufactures
an invalid scale before the first sweep runs.
Returns:
``{latent: draws}`` with ``num_samples`` rows each, closed-form and
NUTS latents alike.
Raises:
GraphError: if the plan has no closed-form block at all -- run
:func:`~bayesmith.bridge.numpyro_bridge.nuts` instead, and this
refusal names it -- or if a block carries a method this module
does not sweep; or if an evidence term covers one of the plan's
exact latents, including when a cached plan came from another
graph state.
StructureError: if ``on_sweep`` is given for a plan with a NUTS
remainder.