bayesmithDocumentation0.10.0 source

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

Defined in bayesmith.dispatch.execute, line 49.

SNIS_ESS_FLOOR: float = 0.1

Posterior

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 = None
What :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.Array
What :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) -> float
MIN 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.0

SiteDiagnostic

Defined in bayesmith.dispatch.execute, line 322.

class SiteDiagnostic(NamedTuple)
r_hat: float
ess: float
ceiling: float
converged: bool
worst: tuple[int, ...]
reason: str
One 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) -> float
The 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`.

Type a concept or API name.