bayesmith.dispatch.execute
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Running what :func:`~bayesmith.dispatch.plan.compile` decided. :mod:`bayesmith.dispatch.plan` derives the partition, measures its conditioning and prints it; nothing in that module draws a sample. This one is the other half: given an :class:`~bayesmith.dispatch.plan.InferencePlan`, it runs section 6.4's five shapes and section 6.5's two, reading ``method``, ``tol`` and ``sigma_needs_rebuild`` off the plan rather than re-deriving them, so that what runs is what ``str(plan)`` printed. The two entry points are :func:`run_sample` and :func:`run_estimate`. :meth:`~bayesmith.dispatch.plan.InferencePlan.sample` and :meth:`~bayesmith.dispatch.plan.InferencePlan.estimate` are the public spelling of them and are where their defaults and their documentation live; the methods delegate here and add nothing. **The dependency runs one way only.** ``plan`` imports this module at module scope; this module reaches ``InferencePlan`` for its annotations alone, under ``TYPE_CHECKING``, so there is no import cycle to break at runtime.
SNIS_ESS_FLOOR Posterior Estimate chain_ess CHAIN_ESS_FLOOR SiteDiagnostic r_hat_ceiling chain_diagnostics run_sample run_estimate
SNIS_ESS_FLOOR
Defined in bayesmith.dispatch.execute, line 49.
SNIS_ESS_FLOOR: float = 0.1Posterior
Defined in bayesmith.dispatch.execute, line 107.
class Posterior(NamedTuple)
samples: dict[str, jax.Array]
log_weights: jax.Array | None
ess: float
khat: float | None
unreliable: bool
method: str
reason: str
diagnostics: Mapping[str, SiteDiagnostic] | None = None
cost: CostReconciliation | None = NoneWhat :meth:`InferencePlan.sample` returns, on every path.
Attributes:
samples: ``{latent: draws}``, the leading axis being the draw. Only
latents -- numpyro's ``get_samples`` also returns Deterministic
sites (measured: ``mixed_radiometer``'s ``mu``, shape (400, 10)),
and leaving them in would silently change ``ess``.
log_weights: unnormalised SNIS log weights, or ``None`` on every path
that produced an unweighted sample.
ess: effective sample size, reduced by **MIN** over every site and
every coordinate -- see :func:`chain_ess`. The Kish ESS on the
SNIS path, and ``num_samples`` exactly on the iid one.
khat: PSIS k-hat of ``log_weights``, or ``None`` where there are none.
unreliable: k-hat past ``min(1 - 1/log10(N), 0.7)``, **or** a Kish
ESS/N under ``ess_floor`` -- either is a reason not to trust the
weights, and the field is the one place a caller looks for that.
``False`` wherever ``khat`` is ``None`` and nothing collapsed,
which is **abstention, not endorsement** -- read ``ess``, which is
always there.
method: what actually RAN: ``"gcr"``, ``"gcr+snis"``, ``"gcr+mh"``,
``"nuts"`` or ``"collapse"``. Not necessarily ``plan.exact.method``:
a collapsed SNIS reports ``"nuts"`` when ``nuts_on_collapse=True``
asked for it, because that is what produced the draws, and
``"gcr+snis"`` otherwise -- see :data:`SNIS_ESS_FLOOR` for why the
default is to annotate rather than substitute -- and the collapse
arm reports ``"collapse"``.
reason: why that, in the plan's own words, including the measured
Kish ESS/N wherever the floor fired.
diagnostics: per-site split r-hat and ESS -- see
:func:`chain_diagnostics` -- on the paths that ran a CHAIN, and
``None`` on the paths that did not. ``None`` is abstention rather
than endorsement, and here it is a stronger statement than usual:
the ``gcr`` and ``gcr+snis`` paths draw independently, so r-hat
has no referent at all on them. Reporting a number would invent
one. Read :attr:`ess`, which is always there.
cost: the P7 reconciliation ledger -- what the cost scoreboard
PREDICTED for the row that ran, what this run actually cost in
seconds per effective sample, and which of the prediction's own
terms dominated it. ``None`` on every plan compiled with the
default ``strategy="declared"``, which is every published path: a
plan with no scoreboard made no prediction, and a ledger row
against no prediction would be a measurement pretending to be a
reconciliation. See
:class:`~bayesmith.dispatch.costs.CostReconciliation`.Estimate
Defined in bayesmith.dispatch.execute, line 165.
class Estimate(NamedTuple)
values: dict[str, jax.Array]
precision: dict[str, Any]
converged: bool
residual: jax.Array
iterations: jax.ArrayWhat :meth:`InferencePlan.estimate` returns: a point, and its covariance.
Attributes:
values: the GLS/Wiener solution, ``{latent: value}``.
precision: the covariance it was solved at, as ``{observed: N^-1}``
-- for a prediction-dependent model, the fixed point. This is what
``gcr_sample`` and ``log_weight`` read.
converged: always ``True``. ``False`` is not returned; it is raised as
:class:`~bayesmith.errors.ConvergenceError`, which is the whole
point of this being the promotion site. Kept as a field because a
caller reading a stored ``Estimate`` should not have to know that.
residual: relative CG residual of the final solve. Not an accuracy --
multiply by ``plan.exact.kappa`` for the error bound.
iterations: reweighting steps taken, ``1`` for a constant sigma.Estimate.noise_std
noise_std(self) -> dict[str, jax.Array] | None:attr:`precision` as per-sample sigma VALUES, or ``None``. The same shape, and for the same reason, as :attr:`~bayesmith.exact.gls.GLSResult.noise_std`: one covariance is stored, and the spelling that does not always exist is the derived one. ``None`` for a correlated model, which has no per-sample sigma; for a diagonal one these are the arrays themselves, so ``noise_std_at(graph, values)`` reproduces them.
chain_ess
Defined in bayesmith.dispatch.execute, line 215.
chain_ess(samples: Mapping[str, Any], *, num_chains: int=1) -> floatMIN of numpyro's ESS over every site and every coordinate.
**The reduction is MIN and that has to be written down.** Section 6.3's
benchmark C has ``ESS(logw)=3.0`` and ``ESS(alm, min)=40.2`` in one run;
:attr:`Posterior.ess` exists to make dividing by N a deliberate act, so it
must report the worst of them. A mean lets a well-mixing site hide a stuck
one, which is the exact shape of the bug the field is against.
**A non-finite per-coordinate ESS becomes 1.0, not dropped.** A chain that
never moved makes ``effective_sample_size`` divide 0 by 0 and return
``nan``, and ``nan`` is the one value ``min`` steps over silently --
``min(nan, 150.2)`` is either, depending on argument order. Measured:
``mixed_radiometer`` at key 2 leaves ``w`` bitwise constant over 400
sweeps and reads exactly that. One draw's worth of information is the
smallest a non-empty sample can carry, so 1.0 is both the honest value and
the one that wins the ``min``.
Args:
samples: ``{site: draws}`` with the draw axis LEADING and chains
already concatenated along it, which is what
``MCMC.get_samples()`` returns.
num_chains: how many chains that axis holds, so it can be unstacked
before the estimator sees it. Left at 1 the array is used as is.CHAIN_ESS_FLOOR
Defined in bayesmith.dispatch.execute, line 268.
CHAIN_ESS_FLOOR: float = 100.0SiteDiagnostic
Defined in bayesmith.dispatch.execute, line 322.
class SiteDiagnostic(NamedTuple)
r_hat: float
ess: float
ceiling: float
converged: bool
worst: tuple[int, ...]
reason: strOne site's convergence, judged at its worst coordinate.
Attributes:
r_hat: split r-hat at the coordinate that decided the verdict, or
``inf`` where numpyro returned a non-finite value. See
:func:`chain_diagnostics` for why ``inf`` and not ``nan``.
ess: effective sample size at that same coordinate. The pair is
reported together because neither means anything alone -- the
ceiling r-hat is judged against is a function of this number.
ceiling: ``r_hat_ceiling(ess, total coordinates in the report)``.
converged: whether ``r_hat`` is within ``ceiling``.
worst: index of the deciding coordinate, ``()`` for a scalar site.
The attribution is the point: the pre-existing scalar reported the
worst number in the whole posterior and never said where it came
from, and "which element of a 64-element site" is a different
morning's work each time.
reason: empty where ``converged``, and otherwise which of the two
checks failed. They are different problems with different
remedies -- too little information is answered by sampling
longer, chains that disagree by finding out why -- and a single
boolean cannot tell a caller which one they have.r_hat_ceiling
Defined in bayesmith.dispatch.execute, line 354.
r_hat_ceiling(ess: float, *, coordinates: int=1) -> floatThe largest split r-hat a CONVERGED chain of this ESS plausibly shows.
**A constant threshold on r-hat is not a well-posed test**, which is the
finding that shaped this function, and it is stronger than the migration
spec's warning that the per-parameter form needs its own argument.
Split r-hat's null distribution is governed by how much independent
information the chain holds, so the same constant is a coin flip at one
end of the ESS range and a no-op at the other.
Measured with 4000 independent coordinates per cell over 27 cells --
chains in {1, 2, 4}, draws in {200, 800, 3200}, AR(1) rho in
{0, 0.5, 0.9}, so ESS from 14.8 to 12693:
* ``r_hat_99 - 1`` scales as ``ESS ** -1.05``: a straight ``1 / ESS``.
* ``(r_hat_99 - 1) * ESS`` stays within 4.87 to 11.46 while ESS moves by a
factor of 860. That product is the stable quantity, so it is the one
pinned, and the ceiling is derived from it.
* A fixed **1.05** fires on **45.9 %** of converged chains at ESS 15, and
**0.0 %** at ESS 12693. 1.01 fires on 69.9 % and 0.0 %. Neither constant
is a test at both ends.
To reproduce: draw ``(chains, draws, 4000)`` standard normals, optionally
filtered to AR(1) with ``x[t] = rho x[t-1] + sqrt(1-rho^2) e[t]``, and take
``split_gelman_rubin`` and ``effective_sample_size`` over the coordinate
axis. Coordinates are independent, so the family-wise quantile over ``P``
of them is the per-coordinate quantile at ``0.99 ** (1/P)`` -- exact, and
far cheaper than simulating each ``P``.
Args:
ess: the effective sample size AT the coordinate being judged, not a
summary over the site. Mixing coordinates here would compare one
coordinate's r-hat against another's information.
coordinates: how many coordinates the report covers in total. Reporting
the worst of ``P`` inflates the null, so the ceiling rises with
``P`` -- otherwise the alarm rate grows with the size of the model
rather than with its health. Measured to saturate near 23, and
capped there: an unbounded term would make a large model
unfalsifiable.chain_diagnostics
Defined in bayesmith.dispatch.execute, line 397.
chain_diagnostics(samples: Mapping[str, Any], *, num_chains: int=1) -> dict[str, SiteDiagnostic]Per-parameter split r-hat and ESS, judged coordinate by coordinate.
What B7 asks for. :func:`chain_ess` already reduced ESS to one number by
MIN over everything, which is the right reduction for a scalar gate and
the wrong one for a diagnosis: it reports the worst number in the whole
posterior and cannot say which site, let alone which coordinate, produced
it. This reports both, per site, and names the coordinate.
**A non-finite r-hat becomes ``inf``, and that is load-bearing.** numpyro
returns ``nan`` for a coordinate that never moved -- measured, on a site
held bitwise constant. ``nan > ceiling`` is ``False``, so the most
unconverged parameter it is possible to have would pass a naive comparison
silently. ``inf`` is both the honest value for a chain carrying no
information and the one that survives the comparison. Note the direction
differs from :func:`chain_ess`, which maps a non-finite ESS to ``1.0``:
both choices push the same way, toward the answer that does not flatter
the chain.
Args:
samples: ``{site: draws}`` with the draw axis LEADING and chains
already concatenated along it, as ``MCMC.get_samples()`` returns.
num_chains: how many chains that axis holds. Unstacked before the
estimator sees it -- handing concatenated chains to split r-hat as
one long chain hides exactly the between-chain disagreement the
statistic exists to find.
Raises:
ValueError: if any site carries fewer than 4 draws per chain, which
numpyro's ``split_gelman_rubin`` refuses with a bare
``AssertionError``. A library assertion is not a diagnosis.run_sample
Defined in bayesmith.dispatch.execute, line 517.
run_sample(plan: InferencePlan, key: jax.Array, *, num_samples: int, num_warmup: int, num_chains: int, chain_method: str, progress_bar: bool, nuts_options: Mapping[str, Any] | None, tol: float | None, maxiter: int | None, require_convergence: float | None, ess_floor: float, nuts_on_collapse: bool, collapse: bool, _control: Any=None) -> Posterior:meth:`~bayesmith.dispatch.plan.InferencePlan.sample`, as a function. Every argument is required here and defaulted there: the method is the public spelling and owns the defaults, so there is one place a caller can read them off and one place they can drift from.
run_estimate
Defined in bayesmith.dispatch.execute, line 662.
run_estimate(plan: InferencePlan, *, tol: float | None, maxiter: int | None, reweight_tol: float | None, min_reweights: int, max_reweights: int, require_convergence: float | None) -> Estimate:meth:`~bayesmith.dispatch.plan.InferencePlan.estimate`, as a function. Defaulted at the method for the same reason as :func:`run_sample`.