bayesmith.marginal.chain
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
A nuisance that drifts across epochs, and the recursion that integrates it out.
An epoch-local nuisance is re-drawn every epoch and integrated away inside its
own. A LINKED one is not: it is a Markov chain, and treating it as epoch-local
marginalises a single physical fluctuation ``N`` times against independent
priors -- injecting information that is not there. This module is the exact
alternative (migration gap G3).
**The recursion.** Carry a joint square-root information factor over
``(theta, zeta_e)``. Fold in an epoch by stacking its rows and
re-triangularising. Advance by widening to ``(theta, zeta_e, zeta_{e+1})``,
appending the transition's rows, and marginalising ``zeta_e``: permute it
first, re-triangularise, drop row and column. **That drop IS the Schur
complement, in square root**, which is what keeps a thousand-epoch
accumulation inside float64 where the explicit ``(F, b)`` form goes indefinite.
``theta`` is never marginalised, so what comes back is ``log p(d_1:N | theta)``
exactly.
**Two sub-scopes, because "linear-Gaussian" is not enough.** An OU chain with
an INFERRED correlation time is still linear-Gaussian, so a caveat phrased
that way is satisfied while its claim fails: ``Q(theta)``, ``phi(theta)`` and
the Schur complement all become functions of theta, and a filter run once at
compression time pins them silently. The distinction lives in the TYPE. A
:class:`LinearGaussianTransition` holds numbers, and the theta posterior is
exact under filtering; a :class:`HyperTransition` holds a builder and is
resolved INSIDE the theta likelihood, so the whole recursion is a
differentiable ``lax.scan`` over the stored per-epoch blocks. One code path
serves both, because the recursion is traceable either way -- which is also
why the fixed case is validated by the same tests rather than by a second
implementation of the same arithmetic.
**The constant bookkeeping is not optional, and it is where this module can be
wrong while looking right.** Six constants reach the answer, and the
recursion's shape, gradient and curvature are all correct without any of them
-- so every test that checks a mean, a width or a derivative passes on a
version that has dropped one. Only a comparison against a dense joint density
notices. ``tests/marginal/test_chain.py`` builds that dense reference and
deletes each constant in turn to measure what it was worth.
Ported from ``rheplicant.inference.chain``. The containers a campaign stores
(a chain memory, its epoch bookkeeping) stay upstream, per the migration
ledger's D12 and the G6 enumeration; what is here is the recursion.HyperTransition LinearGaussianTransition chain_log_likelihood chain_marginal ornstein_uhlenbeck smooth
HyperTransition
Defined in bayesmith.marginal.chain, line 205.
class HyperTransition(eqx.Module)
build: Callable[[dict[str, jax.Array]], LinearGaussianTransition] = eqx.field(static=True)
width: int = eqx.field(static=True)
hyper: tuple[str, ...] = eqx.field(static=True, default=())A transition whose blocks are FUNCTIONS of theta.
The distinction the module docstring argues for, as a type. ``build`` is
resolved inside the theta likelihood, on every evaluation, so a correlation
time that is inferred rather than pinned is differentiated on every
leapfrog step. That is why the recursion is a ``lax.scan`` over stored
blocks rather than a filter run once at compression time: the latter would
pin ``Q(theta)`` and ``phi(theta)`` silently, and the posterior it produced
would be exact for a model nobody declared.
Attributes:
build: ``{name: value} -> LinearGaussianTransition``. Positivity of
the spreads it returns is the BUILDER's responsibility and cannot
be checked here -- see :class:`LinearGaussianTransition`.
width: how many components the chain carries. Declared rather than
derived, because deriving it would mean calling ``build`` at
construction time with values that may not exist yet.
hyper: the latents ``build`` reads.HyperTransition.at
at(self, values: dict[str, jax.Array]) -> LinearGaussianTransitionThe transition at these values.
LinearGaussianTransition
Defined in bayesmith.marginal.chain, line 67.
LinearGaussianTransition(phi: Any, process_std: Any, initial_std: Any, initial_mean: Any=None, hyper: Sequence[str]=())``zeta_{e+1} = phi zeta_e + w``, ``w ~ N(0, diag(process_std)^2)``.
Attributes:
phi: ``(n, n)``. A full matrix, because a multi-component drift can
rotate; the process and initial spreads are diagonal because a
correlated INNOVATION is a modelling claim nobody has made, and a
silently-accepted full covariance would need its own Cholesky
refusal.
process_std: ``(n,)``, strictly positive.
initial_std: ``(n,)``, strictly positive -- ``sd(zeta_1)``.
initial_mean: ``(n,)``. Zero unless declared.
hyper: empty for a fixed transition. Present so a caller can ask one
question of either type.
**Positivity is checked here and nowhere else, on purpose.** The rows this
class contributes are what constrain every ``zeta_e``, so a strictly
positive spread makes each marginalisation's block full-rank BY
CONSTRUCTION -- and that is what lets the filter call
:func:`~bayesmith.marginal.sqrtinfo.marginalise_arrays` inside a
``lax.scan`` instead of the checked
:func:`~bayesmith.marginal.sqrtinfo.marginalise`, which concretises and
therefore cannot be traced or differentiated. One eager check at
declaration, not one traced check per epoch of a thousand.
A TRACED spread is not checked, and cannot be: a :class:`HyperTransition`
builds these from theta, and under NUTS theta goes wherever it likes.
Parameterise the builder so positivity is structural -- return
``jnp.exp(log_sigma)``, never a raw sampled scale -- which is why this
class takes standard deviations rather than a covariance.LinearGaussianTransition.width
width(self) -> intHow many components the chain carries.
LinearGaussianTransition.at
at(self, values: dict[str, jax.Array]) -> LinearGaussianTransitionItself. A fixed transition does not depend on theta -- that is the claim it is making, and the reason its posterior is exact.
chain_log_likelihood
Defined in bayesmith.marginal.chain, line 449.
chain_log_likelihood(blocks: tuple[jax.Array, jax.Array, jax.Array], transition: Any, values: dict[str, jax.Array], names: tuple[str, ...], shapes: tuple[tuple[int, ...], ...]) -> jax.Array``log p(d_1:N | theta)``, the chain integrated out exactly. No prior.
chain_marginal
Defined in bayesmith.marginal.chain, line 356.
chain_marginal(blocks: tuple[jax.Array, jax.Array, jax.Array], transition: Any, values: dict[str, jax.Array], names: tuple[str, ...], shapes: tuple[tuple[int, ...], ...]) -> SqrtInfo``zeta_1:N`` integrated out exactly, leaving a quadratic form in theta.
Args:
blocks: ``(factor (N, w, w), target (N, w), offset (N,))`` -- one
square per-epoch joint form over ``(*names, zeta)``, with
``zeta``'s columns **last**. Square rather than ragged because
``lax.scan`` needs one shape per iteration.
transition: a :class:`LinearGaussianTransition` or a
:class:`HyperTransition`. One code path serves both; ``.at(values)``
is where they differ.
values: theta, for a hyper transition to resolve against. A fixed one
ignores it.
names, shapes: theta's layout, as the returned :class:`SqrtInfo` will
carry it.
Returns:
A :class:`SqrtInfo` over ``names`` alone. Its ``offset`` carries every
constant of the integral, which is what makes it an EVIDENCE rather
than a shape.
**Why the scan stops one short.** Scanning over all ``N`` blocks and then
marginalising once more integrates a ``zeta_{N+1}`` no data constrained,
and comes back exact anyway -- the extra transition density integrates to
one over its own argument, so its normalisation and the extra
marginalisation's constant cancel term for term. The reason to stop one
short is COST, one QR per call, not correctness. What it does change is the
COUNT of each constant, which is what the constant tests notice and the
exactness tests cannot.ornstein_uhlenbeck
Defined in bayesmith.marginal.chain, line 178.
ornstein_uhlenbeck(tau: Any, sigma: Any, width: int=1, hyper: Sequence[str]=()) -> LinearGaussianTransitionA stationary OU chain: correlation time ``tau`` in epochs, spread ``sigma``.
``phi = exp(-1/tau)`` and ``process_std = sigma sqrt(1 - phi^2)``, so
``var(zeta_{e+1}) = phi^2 var + Q`` returns ``sigma^2`` when it starts
there -- **stationarity is arithmetic here, not an assumption**, and a test
pins it.
A FUNCTION, not a class: the type the filter consumes -- and the type a
:class:`HyperTransition` builder must return -- is
:class:`LinearGaussianTransition`. An OU is a way of constructing one, and
this package spells constructors in lower case.smooth
Defined in bayesmith.marginal.chain, line 555.
smooth(blocks: tuple[jax.Array, jax.Array, jax.Array], transition: Any, values: dict[str, jax.Array], names: tuple[str, ...], shapes: tuple[tuple[int, ...], ...]) -> tuple[jax.Array, jax.Array]``p(zeta_e | d_1:N, theta)`` for every epoch -- mean and variance.
``theta`` is **conditioned on**, not marginalised: the question a smoother
answers is "given this receiver model, what did the drift do?", and
marginalising theta would answer a different one with the same shapes.
**Not the classical backward pass, and the same quantity.** The joint form
over ``zeta_1:N`` given theta is block-tridiagonal and small enough to
assemble: ``N * width`` is the number of epochs times the chain's width, not
the data size. Assembling and solving it gives the exact smoothed mean and
marginal variances in one step, with no forward/backward pair to keep
consistent -- and the two halves of an RTS smoother disagreeing is a defect
that reads as a physical result.
**Assembled as a square root, and :func:`_zeta_joint` says why at length.**
In short: the precision would carry ``(1 / process_std) ** 2``, and a stiff
chain then asks float64 to hold a number whose product with ``eps`` is not
small. This function used to form it that way and returned confidently
wrong answers on chains this package documents as supported.
Returns:
``(mean (N, width), variance (N, width))`` -- the smoothed marginal of
each epoch's chain state.