Bayesian Inference & Ensembles

Layer 3 recipes for conjugate updates, second-order variational steps, and ensemble data assimilation. All covariances are operators, so the updates inherit structured solves; all stochastic routines take explicit PRNG keys.

Bayesian learning rule

Damped natural-parameter updates of the Bayesian learning rule (BLR) — full covariance or diagonal-only — plus the marginal likelihood and expected log-likelihood that score Gaussian posteriors.

Structured linear algebra and Gaussian primitives for JAX.

blr_full_update(nat1: Float[Array, ' d'], nat2: Float[Array, 'd d'], grad: Float[Array, ' d'], hessian: Float[Array, 'd d'], lr: float | Float[Array, ''], *, solver: AbstractSolverStrategy | None = None, convention: NaturalConvention = 'expfam') -> tuple[Float[Array, ' d'], Float[Array, 'd d']]

Full-rank natural parameter BLR update step.

Computes the damped update for full-rank variational parameters:

nat2_{new} = (1 - lr) \cdot nat2 + lr \cdot (-\tfrac{1}{2}(-H))
\mu = solve(-2 \cdot nat2, nat1)
nat1_{new} = (1 - lr) \cdot nat1 + lr \cdot (grad - H \mu)

Parameters:

Name	Type	Description	Default
nat1	Float[Array, ' d']	Current natural location, shape (d,).	required
nat2	Float[Array, 'd d']	Current natural precision matrix (eta2), shape (d, d).	required
grad	Float[Array, ' d']	Gradient of log-likelihood, shape (d,).	required
hessian	Float[Array, 'd d']	Hessian of log-likelihood (negative for log-concave), shape (d, d).	required
lr	float | Float[Array, '']	Learning rate / damping factor.	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None
convention	NaturalConvention	Convention of nat2 and of the returned nat2_new. "expfam" (default) is eta2 = -Lambda/2; "precision" is nat2 = +Lambda, the convention of gaussx.newton_update and gaussx.cavity_distribution. nat1 is Lambda mu in both.	'expfam'

Returns:

Type	Description
tuple[Float[Array, ' d'], Float[Array, 'd d']]	Tuple (nat1_new, nat2_new) — updated natural parameters.

Note

By default nat2 is the exponential-family eta2 = -Lambda/2, matching gaussx.mean_cov_to_natural. gaussx.newton_update and gaussx.cavity_distribution use nat2 = +Lambda instead; convert with nat2_plus = -2 * eta2, or pass convention="precision" to read and return +Lambda directly. gaussx.damped_natural_update is linear, so it works in either convention as long as both of its arguments share it.

Source code in src/gaussx/_inference/_blr.py

def blr_full_update(
    nat1: Float[Array, " d"],
    nat2: Float[Array, "d d"],
    grad: Float[Array, " d"],
    hessian: Float[Array, "d d"],
    lr: float | Float[Array, ""],
    *,
    solver: AbstractSolverStrategy | None = None,
    convention: NaturalConvention = "expfam",
) -> tuple[Float[Array, " d"], Float[Array, "d d"]]:
    r"""Full-rank natural parameter BLR update step.

    Computes the damped update for full-rank variational parameters:

        nat2_{new} = (1 - lr) \cdot nat2 + lr \cdot (-\tfrac{1}{2}(-H))
        \mu = solve(-2 \cdot nat2, nat1)
        nat1_{new} = (1 - lr) \cdot nat1 + lr \cdot (grad - H \mu)

    Args:
        nat1: Current natural location, shape ``(d,)``.
        nat2: Current natural precision matrix (eta2), shape ``(d, d)``.
        grad: Gradient of log-likelihood, shape ``(d,)``.
        hessian: Hessian of log-likelihood (negative for log-concave),
            shape ``(d, d)``.
        lr: Learning rate / damping factor.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.
        convention: Convention of ``nat2`` and of the returned ``nat2_new``.
            ``"expfam"`` (default) is ``eta2 = -Lambda/2``; ``"precision"``
            is ``nat2 = +Lambda``, the convention of `gaussx.newton_update`
            and `gaussx.cavity_distribution`. ``nat1`` is ``Lambda mu`` in
            both.

    Returns:
        Tuple ``(nat1_new, nat2_new)`` — updated natural parameters.

    Note:
        By default ``nat2`` is the exponential-family ``eta2 = -Lambda/2``,
        matching `gaussx.mean_cov_to_natural`. `gaussx.newton_update` and
        `gaussx.cavity_distribution` use ``nat2 = +Lambda`` instead; convert
        with ``nat2_plus = -2 * eta2``, or pass ``convention="precision"`` to
        read and return ``+Lambda`` directly. `gaussx.damped_natural_update`
        is linear, so it works in either convention as long as both of its
        arguments share it.
    """
    nat2 = _to_expfam(nat2, convention)
    # Current mean from natural parameters: mu = solve(-2*eta2, eta1)
    Lambda = -2.0 * nat2
    Lambda_op = lx.MatrixLinearOperator(Lambda, lx.positive_semidefinite_tag)
    mu = dispatch_solve(Lambda_op, nat1, solver)

    # Target natural parameters from Newton step
    nat1_target = grad - hessian @ mu
    nat2_target = 0.5 * hessian  # eta2 = -0.5 * (-H) = 0.5 * H

    # Damped update
    nat1_new = (1.0 - lr) * nat1 + lr * nat1_target
    nat2_new = (1.0 - lr) * nat2 + lr * nat2_target

    return nat1_new, _from_expfam(nat2_new, convention)

blr_diag_update(nat1: Float[Array, ' d'], nat2_diag: Float[Array, ' d'], grad: Float[Array, ' d'], hessian_diag: Float[Array, ' d'], lr: float | Float[Array, ''], *, convention: NaturalConvention = 'expfam') -> tuple[Float[Array, ' d'], Float[Array, ' d']]

Diagonal natural parameter BLR update step.

Computes the damped update for diagonal variational parameters:

\mu = nat1 / (-2 \cdot nat2)
eta2_{target} = -\tfrac{1}{2}(-hessian\_diag) = 0.5 \cdot hessian\_diag
eta1_{target} = grad - hessian\_diag \cdot \mu
nat1_{new} = (1 - lr) \cdot nat1 + lr \cdot eta1_{target}
nat2_{new} = (1 - lr) \cdot nat2 + lr \cdot eta2_{target}

where nat2 (eta2) stores -\tfrac{1}{2} \lambda with \lambda = -hessian\_diag (diagonal precision).

Parameters:

Name	Type	Description	Default
nat1	Float[Array, ' d']	Current natural location, shape (d,).	required
nat2_diag	Float[Array, ' d']	Current diagonal natural precision (eta2), shape (d,).	required
grad	Float[Array, ' d']	Gradient of log-likelihood, shape (d,).	required
hessian_diag	Float[Array, ' d']	Diagonal of Hessian (negative for log-concave), shape (d,).	required
lr	float | Float[Array, '']	Learning rate / damping factor.	required
convention	NaturalConvention	Convention of nat2_diag and of the returned nat2_new. "expfam" (default) is eta2 = -Lambda/2; "precision" is nat2 = +Lambda, the convention of gaussx.newton_update and gaussx.cavity_distribution. nat1 is Lambda mu in both.	'expfam'

Returns:

Type	Description
tuple[Float[Array, ' d'], Float[Array, ' d']]	Tuple (nat1_new, nat2_new) — updated natural parameters.

Note

By default nat2 is the exponential-family eta2 = -Lambda/2, matching gaussx.mean_cov_to_natural. gaussx.newton_update and gaussx.cavity_distribution use nat2 = +Lambda instead; convert with nat2_plus = -2 * eta2, or pass convention="precision" to read and return +Lambda directly. gaussx.damped_natural_update is linear, so it works in either convention as long as both of its arguments share it.

Source code in src/gaussx/_inference/_blr.py

def blr_diag_update(
    nat1: Float[Array, " d"],
    nat2_diag: Float[Array, " d"],
    grad: Float[Array, " d"],
    hessian_diag: Float[Array, " d"],
    lr: float | Float[Array, ""],
    *,
    convention: NaturalConvention = "expfam",
) -> tuple[Float[Array, " d"], Float[Array, " d"]]:
    r"""Diagonal natural parameter BLR update step.

    Computes the damped update for diagonal variational parameters:

        \mu = nat1 / (-2 \cdot nat2)
        eta2_{target} = -\tfrac{1}{2}(-hessian\_diag) = 0.5 \cdot hessian\_diag
        eta1_{target} = grad - hessian\_diag \cdot \mu
        nat1_{new} = (1 - lr) \cdot nat1 + lr \cdot eta1_{target}
        nat2_{new} = (1 - lr) \cdot nat2 + lr \cdot eta2_{target}

    where ``nat2`` (eta2) stores ``-\tfrac{1}{2} \lambda`` with
    ``\lambda = -hessian\_diag`` (diagonal precision).

    Args:
        nat1: Current natural location, shape ``(d,)``.
        nat2_diag: Current diagonal natural precision (eta2), shape ``(d,)``.
        grad: Gradient of log-likelihood, shape ``(d,)``.
        hessian_diag: Diagonal of Hessian (negative for log-concave),
            shape ``(d,)``.
        lr: Learning rate / damping factor.
        convention: Convention of ``nat2_diag`` and of the returned
            ``nat2_new``. ``"expfam"`` (default) is ``eta2 = -Lambda/2``;
            ``"precision"`` is ``nat2 = +Lambda``, the convention of
            `gaussx.newton_update` and `gaussx.cavity_distribution`.
            ``nat1`` is ``Lambda mu`` in both.

    Returns:
        Tuple ``(nat1_new, nat2_new)`` — updated natural parameters.

    Note:
        By default ``nat2`` is the exponential-family ``eta2 = -Lambda/2``,
        matching `gaussx.mean_cov_to_natural`. `gaussx.newton_update` and
        `gaussx.cavity_distribution` use ``nat2 = +Lambda`` instead; convert
        with ``nat2_plus = -2 * eta2``, or pass ``convention="precision"`` to
        read and return ``+Lambda`` directly. `gaussx.damped_natural_update`
        is linear, so it works in either convention as long as both of its
        arguments share it.
    """
    nat2_diag = _to_expfam(nat2_diag, convention)
    # Current mean from natural parameters
    mu = nat1 / (-2.0 * nat2_diag)

    # Target natural parameters from Newton step
    nat1_target = grad - hessian_diag * mu
    nat2_target = -0.5 * (-hessian_diag)  # eta2 = -0.5 * (-H) = 0.5 * H

    # Damped update
    nat1_new = (1.0 - lr) * nat1 + lr * nat1_target
    nat2_new = (1.0 - lr) * nat2_diag + lr * nat2_target

    return nat1_new, _from_expfam(nat2_new, convention)

log_marginal_likelihood(loc: Float[Array, ' N'], cov_operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']

GP log marginal likelihood.

Computes:

log p(y) = -0.5 * (y-mu)^T K^{-1} (y-mu) - 0.5 * log|K| - N/2 * log(2pi)

Delegates to gaussx.gaussian_log_prob.

Parameters:

Name	Type	Description	Default
loc	Float[Array, ' N']	Prior mean, shape (N,).	required
cov_operator	AbstractLinearOperator	Covariance operator K, shape (N, N).	required
y	Float[Array, ' N']	Observations, shape (N,).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, '']	Scalar log marginal likelihood.

Source code in src/gaussx/_inference/_inference.py

def log_marginal_likelihood(
    loc: Float[Array, " N"],
    cov_operator: lx.AbstractLinearOperator,
    y: Float[Array, " N"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> Float[Array, ""]:
    """GP log marginal likelihood.

    Computes:

        log p(y) = -0.5 * (y-mu)^T K^{-1} (y-mu) - 0.5 * log|K| - N/2 * log(2pi)

    Delegates to `gaussx.gaussian_log_prob`.

    Args:
        loc: Prior mean, shape ``(N,)``.
        cov_operator: Covariance operator K, shape ``(N, N)``.
        y: Observations, shape ``(N,)``.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Scalar log marginal likelihood.
    """
    return gaussian_log_prob(loc, cov_operator, y, solver=solver)

gaussian_expected_log_lik(y: Float[Array, ' N'], q_mu: Float[Array, ' N'], q_cov: lx.AbstractLinearOperator, noise: lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']

Expected log-likelihood E_q[log N(y | f, R)].

Computes:

E_q[log N(y|f,R)] = log N(y | q_mu, R) - 0.5 * tr(R^{-1} q_cov)

Core to variational inference ELBO computation.

Parameters:

Name	Type	Description	Default
y	Float[Array, ' N']	Observations, shape (N,).	required
q_mu	Float[Array, ' N']	Variational mean, shape (N,).	required
q_cov	AbstractLinearOperator	Variational covariance operator, shape (N, N).	required
noise	AbstractLinearOperator	Noise covariance operator R, shape (N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, '']	Scalar expected log-likelihood.

Source code in src/gaussx/_inference/_inference.py

def gaussian_expected_log_lik(
    y: Float[Array, " N"],
    q_mu: Float[Array, " N"],
    q_cov: lx.AbstractLinearOperator,
    noise: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> Float[Array, ""]:
    r"""Expected log-likelihood ``E_q[log N(y | f, R)]``.

    Computes:

        E_q[log N(y|f,R)] = log N(y | q_mu, R) - 0.5 * tr(R^{-1} q_cov)

    Core to variational inference ELBO computation.

    Args:
        y: Observations, shape ``(N,)``.
        q_mu: Variational mean, shape ``(N,)``.
        q_cov: Variational covariance operator, shape ``(N, N)``.
        noise: Noise covariance operator R, shape ``(N, N)``.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Scalar expected log-likelihood.
    """
    N = y.shape[-1]
    residual = y - q_mu
    alpha = dispatch_solve(noise, residual, solver)
    quad = residual @ alpha
    ld = dispatch_logdet(noise, solver)

    # Trace correction: tr(R^{-1} q_cov)
    R_inv = inv(noise)
    tr_term = trace_product(R_inv, q_cov)

    return -0.5 * (N * _LOG_2PI + ld + quad + tr_term)

Newton & natural-gradient updates

Second-order variational steps: Newton's method on the variational objective, Gauss-Newton curvature (exact diagonal or Hutchinson-estimated), damped natural-gradient steps, and the PSD projection that keeps Riemannian updates on the manifold.

Structured linear algebra and Gaussian primitives for JAX.

newton_update(mean: Float[Array, ' N'], jacobian: Float[Array, ' N'], hessian: Float[Array, 'N N'] | Float[Array, ' N'], *, precision_floor: float = 1e-06) -> tuple[Float[Array, ' N'], Float[Array, 'N N'] | Float[Array, ' N']]

Convert a Newton step to natural pseudo-likelihood parameters.

Computes:

nat1 = jacobian - hessian @ mean
nat2 = -hessian

Used in Laplace/Newton-based approximate inference to convert function-space derivatives into site natural parameters.

Passing hessian as an (N,) array of per-site second derivatives — the shape site-based EP / Laplace inference over N scalar latents actually has — takes an elementwise O(N) path instead of forming the (N, N) matrix product:

\[ \Lambda = \max(-h, \varepsilon), \qquad \lambda_1 = g + \Lambda f, \qquad \lambda_2 = \Lambda. \]

Note

Both forms use the nat2 = +Λ (positive precision) convention, matching gaussx.cavity_distribution. This differs from gaussx.mean_cov_to_natural / gaussx.natural_to_mean_cov and the default of gaussx.blr_diag_update / gaussx.blr_full_update, which use the exponential-family convention η₂ = −Λ/2; convert with nat2 = −2 η₂, or call the BLR updates with convention="precision". gaussx.damped_natural_update is convention-agnostic.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	Current mean, shape (N,) or (D,).	required
jacobian	Float[Array, ' N']	Gradient (first derivative) of the log-likelihood with respect to the latent, shape (N,). The name is historical; it is a gradient, not a Jacobian matrix.	required
hessian	Float[Array, 'N N'] | Float[Array, ' N']	Second derivative (negative definite), either the full (N, N) matrix or an (N,) diagonal.	required
precision_floor	float	Lower bound on the returned precision, applied only on the diagonal path. Keeps sites from a non-log-concave likelihood (positive hessian entries) from producing a negative precision. The full-matrix path returns -hessian unmodified, since flooring it would require an eigendecomposition.	1e-06

Returns:

Type	Description
Float[Array, ' N']	Tuple (nat1, nat2) — site natural parameters. nat2 matches
Float[Array, 'N N'] | Float[Array, ' N']	the shape of hessian.

Source code in src/gaussx/_inference/_inference.py

def newton_update(
    mean: Float[Array, " N"],
    jacobian: Float[Array, " N"],
    hessian: Float[Array, "N N"] | Float[Array, " N"],
    *,
    precision_floor: float = 1e-6,
) -> tuple[Float[Array, " N"], Float[Array, "N N"] | Float[Array, " N"]]:
    r"""Convert a Newton step to natural pseudo-likelihood parameters.

    Computes:

        nat1 = jacobian - hessian @ mean
        nat2 = -hessian

    Used in Laplace/Newton-based approximate inference to convert
    function-space derivatives into site natural parameters.

    Passing ``hessian`` as an ``(N,)`` array of per-site second
    derivatives — the shape site-based EP / Laplace inference over ``N``
    scalar latents actually has — takes an elementwise ``O(N)`` path
    instead of forming the ``(N, N)`` matrix product:

    $$
    \Lambda = \max(-h, \varepsilon), \qquad
    \lambda_1 = g + \Lambda f, \qquad
    \lambda_2 = \Lambda.
    $$

    Note:
        Both forms use the ``nat2 = +Λ`` (positive precision) convention,
        matching `gaussx.cavity_distribution`. This differs from
        `gaussx.mean_cov_to_natural` / `gaussx.natural_to_mean_cov` and the
        default of `gaussx.blr_diag_update` / `gaussx.blr_full_update`, which
        use the exponential-family convention ``η₂ = −Λ/2``; convert with
        ``nat2 = −2 η₂``, or call the BLR updates with
        ``convention="precision"``. `gaussx.damped_natural_update` is
        convention-agnostic.

    Args:
        mean: Current mean, shape ``(N,)`` or ``(D,)``.
        jacobian: Gradient (first derivative) of the log-likelihood with
            respect to the latent, shape ``(N,)``. The name is historical; it
            is a gradient, not a Jacobian matrix.
        hessian: Second derivative (negative definite), either the full
            ``(N, N)`` matrix or an ``(N,)`` diagonal.
        precision_floor: Lower bound on the returned precision, applied
            only on the diagonal path. Keeps sites from a non-log-concave
            likelihood (positive ``hessian`` entries) from producing a
            negative precision. The full-matrix path returns
            ``-hessian`` unmodified, since flooring it would require an
            eigendecomposition.

    Returns:
        Tuple ``(nat1, nat2)`` — site natural parameters. ``nat2`` matches
        the shape of ``hessian``.
    """
    if hessian.ndim == 1:
        precision = jnp.maximum(-hessian, precision_floor)
        return jacobian + precision * mean, precision

    nat1 = jacobian - hessian @ mean
    nat2 = -hessian
    return nat1, nat2

damped_natural_update(nat1_old: Float[Array, ' d'], nat2_old: lx.AbstractLinearOperator | Float[Array, 'd d'], nat1_target: Float[Array, ' d'], nat2_target: lx.AbstractLinearOperator | Float[Array, 'd d'], lr: float | Float[Array, ''] = 1.0) -> tuple[Float[Array, ' d'], lx.AbstractLinearOperator | Float[Array, 'd d']]

Damped update in natural parameter space.

The universal primitive for iterative approximate inference (EP, VI, Newton, PL). Every method reduces to computing target natural parameters and applying this damped update:

nat1_{new} = (1 - lr) \cdot nat1_{old} + lr \cdot nat1_{target}
nat2_{new} = (1 - lr) \cdot nat2_{old} + lr \cdot nat2_{target}

Note

The update is linear, so it is convention-agnostic: it works with the nat2 = +Λ convention of gaussx.newton_update / gaussx.cavity_distribution and with the exponential-family η₂ = −Λ/2 of gaussx.blr_diag_update / gaussx.mean_cov_to_natural alike. Both arguments must use the same one; mixing them silently flips the sign and halves the precision.

Parameters:

Name	Type	Description	Default
nat1_old	Float[Array, ' d']	Current natural location parameter.	required
nat2_old	AbstractLinearOperator | Float[Array, 'd d']	Current natural precision-like parameter. Can be an array, BlockTriDiag, or any linear operator.	required
nat1_target	Float[Array, ' d']	Target natural location parameter.	required
nat2_target	AbstractLinearOperator | Float[Array, 'd d']	Target natural precision-like parameter.	required
lr	float | Float[Array, '']	Learning rate / damping factor. lr=1 gives the undamped update. Default 1.0.	1.0

Returns:

Type	Description
tuple[Float[Array, ' d'], AbstractLinearOperator | Float[Array, 'd d']]	Tuple (nat1_new, nat2_new) with same types as inputs.

Source code in src/gaussx/_inference/_natural_gradient.py

def damped_natural_update(
    nat1_old: Float[Array, " d"],
    nat2_old: lx.AbstractLinearOperator | Float[Array, "d d"],
    nat1_target: Float[Array, " d"],
    nat2_target: lx.AbstractLinearOperator | Float[Array, "d d"],
    lr: float | Float[Array, ""] = 1.0,
) -> tuple[Float[Array, " d"], lx.AbstractLinearOperator | Float[Array, "d d"]]:
    r"""Damped update in natural parameter space.

    The universal primitive for iterative approximate inference
    (EP, VI, Newton, PL). Every method reduces to computing target
    natural parameters and applying this damped update:

        nat1_{new} = (1 - lr) \cdot nat1_{old} + lr \cdot nat1_{target}
        nat2_{new} = (1 - lr) \cdot nat2_{old} + lr \cdot nat2_{target}

    Note:
        The update is linear, so it is convention-agnostic: it works with
        the ``nat2 = +Λ`` convention of `gaussx.newton_update` /
        `gaussx.cavity_distribution` and with the exponential-family
        ``η₂ = −Λ/2`` of `gaussx.blr_diag_update` /
        `gaussx.mean_cov_to_natural` alike. Both arguments must use the same
        one; mixing them silently flips the sign and halves the precision.

    Args:
        nat1_old: Current natural location parameter.
        nat2_old: Current natural precision-like parameter.
            Can be an array, ``BlockTriDiag``, or any linear operator.
        nat1_target: Target natural location parameter.
        nat2_target: Target natural precision-like parameter.
        lr: Learning rate / damping factor. ``lr=1`` gives the
            undamped update. Default ``1.0``.

    Returns:
        Tuple ``(nat1_new, nat2_new)`` with same types as inputs.
    """
    nat1_new = (1.0 - lr) * nat1_old + lr * nat1_target

    # Anything that is not an operator is an array-like (NumPy arrays and
    # nested lists included); normalise those before dispatching on type.
    if not isinstance(nat2_old, lx.AbstractLinearOperator):
        nat2_old = jnp.asarray(nat2_old)
    if not isinstance(nat2_target, lx.AbstractLinearOperator):
        nat2_target = jnp.asarray(nat2_target)

    if isinstance(nat2_old, jax.Array) and isinstance(nat2_target, jax.Array):
        nat2_new: lx.AbstractLinearOperator | Float[Array, "d d"] = (
            1.0 - lr
        ) * nat2_old + lr * nat2_target
    elif isinstance(nat2_old, BlockTriDiag) and isinstance(nat2_target, BlockTriDiag):
        nat2_new = (1.0 - lr) * nat2_old + lr * nat2_target
    elif isinstance(nat2_old, lx.AbstractLinearOperator) and isinstance(
        nat2_target, lx.AbstractLinearOperator
    ):
        nat2_new_mat = (1.0 - lr) * nat2_old.as_matrix() + lr * nat2_target.as_matrix()
        nat2_new = lx.MatrixLinearOperator(nat2_new_mat)
    else:
        msg = (
            "nat2_old and nat2_target must both be arrays or both be linear "
            f"operators, got {type(nat2_old).__name__} and "
            f"{type(nat2_target).__name__}."
        )
        raise TypeError(msg)

    return nat1_new, nat2_new

gauss_newton_precision(jacobian: Float[Array, 'D_obs D_latent'], *, base: lx.AbstractLinearOperator | None = None) -> lx.AbstractLinearOperator

Gauss-Newton precision matrix J^T J, optionally plus a prior.

For likelihoods with residual structure r(f), the Gauss-Newton Hessian approximation is -J_r^T J_r which gives precision \Lambda = J^T J (always PSD).

With base (a prior precision \Lambda_0), returns the posterior precision \Lambda_0 + J^T J as a LowRankUpdate on base, for any D_obs. gaussx.solve and gaussx.logdet then use the Woodbury identity and base's own structured solve. This is the form to use whenever the precision will be solved against.

Without base, J^T J has rank at most D_obs. When D_{obs} < D_{latent} it is returned as a LowRankUpdate on a zero base, which is singular: it is meant as a summand or for mv, and gaussx.solve / gaussx.logdet on it are undefined (NaN). Adding a prior afterwards (gauss_newton_precision(J) + prior) gives a plain lineax sum that is solved densely; pass base=prior instead.

Parameters:

Name	Type	Description	Default
jacobian	Float[Array, 'D_obs D_latent']	Jacobian of the residual, shape (D_obs, D_latent).	required
base	AbstractLinearOperator | None	Optional prior precision of shape (D_latent, D_latent). Its symmetry and positive-semidefiniteness tags carry over to the result.	None

Returns:

Type	Description
AbstractLinearOperator	PSD precision operator of shape (D_latent, D_latent).

Source code in src/gaussx/_inference/_natural_gradient.py

def gauss_newton_precision(
    jacobian: Float[Array, "D_obs D_latent"],
    *,
    base: lx.AbstractLinearOperator | None = None,
) -> lx.AbstractLinearOperator:
    r"""Gauss-Newton precision matrix ``J^T J``, optionally plus a prior.

    For likelihoods with residual structure ``r(f)``, the Gauss-Newton
    Hessian approximation is ``-J_r^T J_r`` which gives precision
    ``\Lambda = J^T J`` (always PSD).

    With ``base`` (a prior precision ``\Lambda_0``), returns the posterior
    precision ``\Lambda_0 + J^T J`` as a `LowRankUpdate` on ``base``, for
    any ``D_obs``. `gaussx.solve` and `gaussx.logdet` then use the Woodbury
    identity and ``base``'s own structured solve. This is the form to use
    whenever the precision will be solved against.

    Without ``base``, ``J^T J`` has rank at most ``D_obs``. When
    ``D_{obs} < D_{latent}`` it is returned as a `LowRankUpdate` on a zero
    base, which is **singular**: it is meant as a summand or for ``mv``,
    and `gaussx.solve` / `gaussx.logdet` on it are undefined (``NaN``).
    Adding a prior afterwards (``gauss_newton_precision(J) + prior``) gives
    a plain lineax sum that is solved densely; pass ``base=prior`` instead.

    Args:
        jacobian: Jacobian of the residual, shape ``(D_obs, D_latent)``.
        base: Optional prior precision of shape ``(D_latent, D_latent)``.
            Its symmetry and positive-semidefiniteness tags carry over to
            the result.

    Returns:
        PSD precision operator of shape ``(D_latent, D_latent)``.
    """
    D_obs, D_latent = jacobian.shape
    ones = jnp.ones(D_obs, dtype=jacobian.dtype)

    if base is not None:
        if base.in_size() != D_latent or base.out_size() != D_latent:
            raise ValueError(
                f"base must be ({D_latent}, {D_latent}) to match the Jacobian's "
                f"D_latent, got ({base.out_size()}, {base.in_size()})."
            )
        # J^T J is PSD, so the sum keeps whatever the prior can claim.
        tags = frozenset(
            tag
            for query, tag in (
                (lx.is_symmetric, lx.symmetric_tag),
                (lx.is_positive_semidefinite, lx.positive_semidefinite_tag),
            )
            if query(base)
        )
        return LowRankUpdate(base=base, U=jacobian.T, d=ones, tags=tags)

    if D_obs < D_latent:
        base = lx.DiagonalLinearOperator(jnp.zeros(D_latent, dtype=jacobian.dtype))
        return LowRankUpdate(
            base=base,
            U=jacobian.T,
            d=ones,
            tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
        )

    return lx.MatrixLinearOperator(
        jacobian.T @ jacobian,
        lx.positive_semidefinite_tag,
    )

ggn_diagonal(jacobian: Float[Array, 'N d']) -> Float[Array, ' d']

Generalized Gauss-Newton diagonal approximation.

Computes \mathrm{diag}(J^T J) = \sum_i J_{i,:}^2, the diagonal of the Gauss-Newton Hessian approximation. Always non-negative, guaranteeing PSD precision updates.

Parameters:

Name	Type	Description	Default
jacobian	Float[Array, 'N d']	Jacobian matrix, shape (N, d) where N is the number of observations and d is the parameter dimension.	required

Returns:

Type	Description
Float[Array, ' d']	Diagonal of J^T J, shape (d,).

Source code in src/gaussx/_inference/_blr.py

def ggn_diagonal(
    jacobian: Float[Array, "N d"],
) -> Float[Array, " d"]:
    r"""Generalized Gauss-Newton diagonal approximation.

    Computes ``\mathrm{diag}(J^T J) = \sum_i J_{i,:}^2``, the diagonal
    of the Gauss-Newton Hessian approximation. Always non-negative,
    guaranteeing PSD precision updates.

    Args:
        jacobian: Jacobian matrix, shape ``(N, d)`` where N is the
            number of observations and d is the parameter dimension.

    Returns:
        Diagonal of ``J^T J``, shape ``(d,)``.
    """
    return reduce(jacobian**2, "K D -> D", "sum")

hutchinson_hessian_diag(hvp_fn: Callable[[Float[Array, ' d']], Float[Array, ' d']], key: jax.Array, d: int, n_samples: int = 1, dtype: DTypeLike | None = None) -> Float[Array, ' d']

Stochastic Hessian diagonal via Hutchinson with Rademacher probes.

Estimates \mathrm{diag}(H) using the identity \mathrm{diag}(H) = E[z \odot (H z)] where z is a Rademacher random vector (entries \pm 1 with equal probability).

Parameters:

Name	Type	Description	Default
hvp_fn	Callable[[Float[Array, ' d']], Float[Array, ' d']]	Hessian-vector product function v -> H @ v.	required
key	Array	PRNG key for random probe generation.	required
d	int	Dimension of the Hessian.	required
n_samples	int	Number of random probes. More samples give better estimates. Default 1.	1
dtype	DTypeLike | None	Floating-point dtype for the Rademacher probes. Defaults to the current JAX default floating dtype.	None

Returns:

Type	Description
Float[Array, ' d']	Estimated diagonal of the Hessian, shape (d,).

Source code in src/gaussx/_inference/_blr.py

def hutchinson_hessian_diag(
    hvp_fn: Callable[[Float[Array, " d"]], Float[Array, " d"]],
    key: jax.Array,
    d: int,
    n_samples: int = 1,
    dtype: DTypeLike | None = None,
) -> Float[Array, " d"]:
    r"""Stochastic Hessian diagonal via Hutchinson with Rademacher probes.

    Estimates ``\mathrm{diag}(H)`` using the identity
    ``\mathrm{diag}(H) = E[z \odot (H z)]`` where ``z`` is a
    Rademacher random vector (entries ``\pm 1`` with equal probability).

    Args:
        hvp_fn: Hessian-vector product function ``v -> H @ v``.
        key: PRNG key for random probe generation.
        d: Dimension of the Hessian.
        n_samples: Number of random probes. More samples give better
            estimates. Default ``1``.
        dtype: Floating-point dtype for the Rademacher probes. Defaults to the
            current JAX default floating dtype.

    Returns:
        Estimated diagonal of the Hessian, shape ``(d,)``.
    """

    probe_dtype = jnp.dtype(jnp.asarray(0.0).dtype if dtype is None else dtype)

    def _single_probe(k):
        z = jnp.where(
            jax.random.bernoulli(k, shape=(d,)),
            jnp.array(1.0, dtype=probe_dtype),
            jnp.array(-1.0, dtype=probe_dtype),
        )
        return z * hvp_fn(z)

    keys = jax.random.split(key, n_samples)
    estimates = jax.vmap(_single_probe)(keys)
    return jnp.mean(estimates, axis=0)

riemannian_psd_correction(hessian: Float[Array, 'd d'], site_precision: Float[Array, 'd d'], site_covariance: Float[Array, 'd d'], lr: float | Float[Array, ''] = 1.0) -> Float[Array, 'd d']

Riemannian gradient correction for PSD precision updates.

Stabilises Newton/EP/VI when the raw Hessian is indefinite:

G = site\_precision + hessian
H_{psd} = hessian - 0.5 \cdot lr \cdot G \cdot S \cdot G

where S is the site covariance.

The corrected Hessian itself is not guaranteed to be negative semi-definite (with H = diag(-1, 2), Lambda = I and lr = 0.1 it has eigenvalues [-1, 1.55]). What the correction guarantees is that the damped precision update stays positive definite: with S = Lambda^{-1},

\[ (1 - \mathrm{lr})\,\Lambda - \mathrm{lr}\,H_{psd} = \tfrac12 \Lambda + \tfrac12 (I - \mathrm{lr}\,G S)\,\Lambda\,(I - \mathrm{lr}\,S G), \]

which follows by expanding both sides to Lambda - lr G + lr^2/2 G S G. The right-hand side is a positive definite matrix plus a positive semi-definite one.

Parameters:

Name	Type	Description	Default
hessian	Float[Array, 'd d']	Raw second derivative, shape (d, d).	required
site_precision	Float[Array, 'd d']	Current site precision, shape (d, d).	required
site_covariance	Float[Array, 'd d']	Current site covariance, shape (d, d).	required
lr	float | Float[Array, '']	Learning rate. Default 1.0.	1.0

Returns:

Type	Description
Float[Array, 'd d']	Corrected Hessian, shape (d, d).

Source code in src/gaussx/_inference/_natural_gradient.py

def riemannian_psd_correction(
    hessian: Float[Array, "d d"],
    site_precision: Float[Array, "d d"],
    site_covariance: Float[Array, "d d"],
    lr: float | Float[Array, ""] = 1.0,
) -> Float[Array, "d d"]:
    r"""Riemannian gradient correction for PSD precision updates.

    Stabilises Newton/EP/VI when the raw Hessian is indefinite:

        G = site\_precision + hessian
        H_{psd} = hessian - 0.5 \cdot lr \cdot G \cdot S \cdot G

    where ``S`` is the site covariance.

    The corrected Hessian itself is **not** guaranteed to be negative
    semi-definite (with ``H = diag(-1, 2)``, ``Lambda = I`` and ``lr = 0.1``
    it has eigenvalues ``[-1, 1.55]``). What the correction guarantees is
    that the damped precision update stays positive definite: with
    ``S = Lambda^{-1}``,

    $$
    (1 - \mathrm{lr})\,\Lambda - \mathrm{lr}\,H_{psd}
      = \tfrac12 \Lambda
        + \tfrac12 (I - \mathrm{lr}\,G S)\,\Lambda\,(I - \mathrm{lr}\,S G),
    $$

    which follows by expanding both sides to
    ``Lambda - lr G + lr^2/2 G S G``. The right-hand side is a positive
    definite matrix plus a positive semi-definite one.

    Args:
        hessian: Raw second derivative, shape ``(d, d)``.
        site_precision: Current site precision, shape ``(d, d)``.
        site_covariance: Current site covariance, shape ``(d, d)``.
        lr: Learning rate. Default ``1.0``.

    Returns:
        Corrected Hessian, shape ``(d, d)``.
    """
    G = site_precision + hessian
    correction = G @ site_covariance @ G
    return hessian - 0.5 * lr * correction

cavity_distribution(post_mean: Float[Array, ' N'], post_cov: lx.AbstractLinearOperator | Float[Array, ' N'], site_nat1: Float[Array, ' N'], site_nat2: lx.AbstractLinearOperator | Float[Array, ' N'], power: float = 1.0, *, precision_floor: float | None = None) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator | Float[Array, ' N']]

Compute EP cavity distribution by removing a site.

Computes:

cav_prec = post_prec - power * site_nat2
cav_cov  = inv(cav_prec)
cav_mean = cav_cov @ (post_prec @ post_mean - power * site_nat1)

Two forms are dispatched on the argument types. Passing post_cov and site_nat2 as operators takes the full-covariance path. Passing both as (N,) arrays — the marginal variances and per-site precisions of N scalar latents, as site-based EP over GPs represents them — takes an elementwise fast path costing O(N) rather than the O(N²) of wrapping them in a lineax.DiagonalLinearOperator:

\[ v_{\mathrm{cav}}^{-1} = v^{-1} - \alpha \lambda_2, \qquad m_{\mathrm{cav}} = v_{\mathrm{cav}} \left( \frac{m}{v} - \alpha \lambda_1 \right). \]

Note

Both forms use the nat2 = +Λ (positive precision) convention, matching gaussx.newton_update. This differs from gaussx.mean_cov_to_natural / gaussx.natural_to_mean_cov and the default of gaussx.blr_diag_update / gaussx.blr_full_update, which use the exponential-family convention η₂ = −Λ/2; convert with nat2 = −2 η₂, or call the BLR updates with convention="precision". gaussx.damped_natural_update is convention-agnostic.

Parameters:

Name	Type	Description	Default
post_mean	Float[Array, ' N']	Posterior mean, shape (N,).	required
post_cov	AbstractLinearOperator | Float[Array, ' N']	Posterior covariance operator, or (N,) marginal variances for the diagonal path.	required
site_nat1	Float[Array, ' N']	Site natural parameter (precision-weighted mean), shape (N,).	required
site_nat2	AbstractLinearOperator | Float[Array, ' N']	Site natural parameter (precision) as an operator, or (N,) per-site precisions for the diagonal path.	required
power	float	Power EP fraction (default 1.0 for standard EP).	1.0
precision_floor	float | None	Optional lower bound on the cavity precision, applied only on the diagonal path. A site that has absorbed more precision than the posterior holds leaves a negative cavity precision, which EP loops guard against by clipping; None (the default) returns the raw cavity, whose negative variances ep_tilted_moments treats as invalid and passes through unchanged (a zero site update). The operator path ignores it, since flooring a matrix would require an eigendecomposition — the same split as newton_update.	None

Returns:

Type	Description
Float[Array, ' N']	Tuple (cav_mean, cav_cov). cav_cov is an operator for the
AbstractLinearOperator | Float[Array, ' N']	operator path and an (N,) array of variances for the diagonal
tuple[Float[Array, ' N'], AbstractLinearOperator | Float[Array, ' N']]	path.

Raises:

Type	Description
TypeError	If post_cov and site_nat2 are not both arrays or both operators.

Source code in src/gaussx/_inference/_inference.py

def cavity_distribution(
    post_mean: Float[Array, " N"],
    post_cov: lx.AbstractLinearOperator | Float[Array, " N"],
    site_nat1: Float[Array, " N"],
    site_nat2: lx.AbstractLinearOperator | Float[Array, " N"],
    power: float = 1.0,
    *,
    precision_floor: float | None = None,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator | Float[Array, " N"]]:
    r"""Compute EP cavity distribution by removing a site.

    Computes:

        cav_prec = post_prec - power * site_nat2
        cav_cov  = inv(cav_prec)
        cav_mean = cav_cov @ (post_prec @ post_mean - power * site_nat1)

    Two forms are dispatched on the argument types. Passing ``post_cov``
    and ``site_nat2`` as operators takes the full-covariance path. Passing
    both as ``(N,)`` arrays — the marginal variances and per-site
    precisions of ``N`` scalar latents, as site-based EP over GPs
    represents them — takes an elementwise fast path costing ``O(N)``
    rather than the ``O(N²)`` of wrapping them in a
    `lineax.DiagonalLinearOperator`:

    $$
    v_{\mathrm{cav}}^{-1} = v^{-1} - \alpha \lambda_2, \qquad
    m_{\mathrm{cav}} = v_{\mathrm{cav}}
        \left( \frac{m}{v} - \alpha \lambda_1 \right).
    $$

    Note:
        Both forms use the ``nat2 = +Λ`` (positive precision) convention,
        matching `gaussx.newton_update`. This differs from
        `gaussx.mean_cov_to_natural` / `gaussx.natural_to_mean_cov` and the
        default of `gaussx.blr_diag_update` / `gaussx.blr_full_update`, which
        use the exponential-family convention ``η₂ = −Λ/2``; convert with
        ``nat2 = −2 η₂``, or call the BLR updates with
        ``convention="precision"``. `gaussx.damped_natural_update` is
        convention-agnostic.

    Args:
        post_mean: Posterior mean, shape ``(N,)``.
        post_cov: Posterior covariance operator, or ``(N,)`` marginal
            variances for the diagonal path.
        site_nat1: Site natural parameter (precision-weighted mean),
            shape ``(N,)``.
        site_nat2: Site natural parameter (precision) as an operator, or
            ``(N,)`` per-site precisions for the diagonal path.
        power: Power EP fraction (default 1.0 for standard EP).
        precision_floor: Optional lower bound on the cavity precision,
            applied only on the diagonal path. A site that has absorbed
            more precision than the posterior holds leaves a negative
            cavity precision, which EP loops guard against by clipping;
            ``None`` (the default) returns the raw cavity, whose negative
            variances `ep_tilted_moments` treats as invalid and passes
            through unchanged (a zero site update). The operator
            path ignores it, since flooring a matrix would require an
            eigendecomposition — the same split as `newton_update`.

    Returns:
        Tuple ``(cav_mean, cav_cov)``. ``cav_cov`` is an operator for the
        operator path and an ``(N,)`` array of variances for the diagonal
        path.

    Raises:
        TypeError: If ``post_cov`` and ``site_nat2`` are not both arrays
            or both operators.
    """
    if isinstance(post_cov, jax.Array) and isinstance(site_nat2, jax.Array):
        cav_prec = 1.0 / post_cov - power * site_nat2
        if precision_floor is not None:
            cav_prec = jnp.maximum(cav_prec, precision_floor)
        cav_var = 1.0 / cav_prec
        cav_mean = cav_var * (post_mean / post_cov - power * site_nat1)
        return cav_mean, cav_var
    if isinstance(post_cov, jax.Array) or isinstance(site_nat2, jax.Array):
        msg = "post_cov and site_nat2 must both be arrays or both be operators"
        raise TypeError(msg)

    post_prec = inv(post_cov)
    cav_prec_mat = post_prec.as_matrix() - power * site_nat2.as_matrix()
    cav_prec = lx.MatrixLinearOperator(cav_prec_mat)
    cav_cov = inv(cav_prec)

    eta1_cav = post_prec.mv(post_mean) - power * site_nat1
    cav_mean = cav_cov.mv(eta1_cav)

    return cav_mean, cav_cov

trace_correction(K_xx: lx.AbstractLinearOperator, K_xz: Float[Array, 'N M'], K_zz: lx.AbstractLinearOperator, *, solver: AbstractSolveStrategy | None = None) -> Float[Array, '']

Trace term in Titsias collapsed ELBO.

Computes:

tr(K_xx) - tr(K_xz^T K_zz^{-1} K_xz)

This is the "trace correction" that penalizes the Nystrom approximation error.

Parameters:

Name	Type	Description	Default
K_xx	AbstractLinearOperator	Full covariance, shape (N, N).	required
K_xz	Float[Array, 'N M']	Cross-covariance, shape (N, M).	required
K_zz	AbstractLinearOperator	Inducing covariance, shape (M, M).	required
solver	AbstractSolveStrategy | None	Optional solve strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, '']	Scalar trace correction.

Source code in src/gaussx/_inference/_inference.py

def trace_correction(
    K_xx: lx.AbstractLinearOperator,
    K_xz: Float[Array, "N M"],
    K_zz: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, ""]:
    """Trace term in Titsias collapsed ELBO.

    Computes:

        tr(K_xx) - tr(K_xz^T K_zz^{-1} K_xz)

    This is the "trace correction" that penalizes the Nystrom
    approximation error.

    Args:
        K_xx: Full covariance, shape ``(N, N)``.
        K_xz: Cross-covariance, shape ``(N, M)``.
        K_zz: Inducing covariance, shape ``(M, M)``.
        solver: Optional solve strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Scalar trace correction.
    """
    tr_full = trace(K_xx)

    # tr(K_xz^T K_zz^{-1} K_xz) = sum_ij W_ij * K_xz_ij
    # where W = K_zz^{-1} K_xz^T reshaped, but easier:
    # tr(A^T B) = sum(A * B), so tr(K_xz^T W) where W_col = K_zz^{-1} K_xz_col
    W = solve_rows(K_zz, K_xz, solver=solver)  # (N, M)
    tr_approx = jnp.sum(K_xz * W)

    return tr_full - tr_approx

Natural-parameter conventions

Two conventions for the second natural parameter coexist. Both share the first one, \(\eta_1 = \Lambda\mu\). Feeding the output of one family into the other silently flips the sign and halves the precision. Convert with \(\text{nat2} = -2\,\eta_2\).

Function	Convention of nat2
blr_diag_update, blr_full_update	\(\eta_2 = -\tfrac12\Lambda\) by default; \(+\Lambda\) with convention="precision"
mean_cov_to_natural, natural_to_mean_cov	\(\eta_2 = -\tfrac12\Lambda\)
newton_update, cavity_distribution	\(+\Lambda\)
damped_natural_update	either (it is linear), as long as both arguments share it

Precision-form Laplace (INLA)

laplace_mode is the inner loop of INLA: the Gaussian approximation \(\mathcal N(\hat x, H^{-1})\) of a latent Gaussian model at fixed hyperparameters \(\theta\), computed entirely in precision form. With \(\eta = Ax + o\), \(g = \partial_\eta\log p(y\mid\eta)\) and \(W = -\operatorname{diag}(\partial^2_\eta\log p)\), each Newton step solves

\[ (Q + A^\top WA)\,x_{t+1} = Q\mu + A^\top(g + WAx_t), \]

with a matrix whose structure never changes: an identity or row-selection projector keeps a BlockTriDiag prior banded, and anything else becomes a SparseOperator on a pattern whose sparse Cholesky analysis is shared by every step and every \(\theta\). Hard sum-to-zero (null-space) constraints of an IntrinsicGMRF are applied by kriging each solve. At the mode,

\[ \log\tilde\pi(y\mid\theta) = \log p(y\mid A\hat x + o) - \tfrac12(\hat x-\mu)^\top Q(\hat x-\mu) + \tfrac12\log|Q| - \tfrac12\log|H|, \]

with the constraint corrections of Rue et al. (2009) for intrinsic priors. The mode is differentiated implicitly (\(\partial_\theta\hat x = H^{-1}\partial_\theta\nabla\ell\), via jax.lax.custom_root) and the log-determinants through their structured or sparse (Takahashi) VJPs, so jax.grad and the reverse-over-reverse Hessian that theta_design uses are exact.

import jax
import jax.numpy as jnp
import numpy as np

import gaussx as gx

# Poisson counts with an RW2 seasonal effect; θ = log τ
n_days = 364  # even: rw2_structure(n) pads odd n with one node (see below)
t = jnp.arange(n_days, dtype=float)
counts = jnp.asarray(np.random.default_rng(0).poisson(np.exp(1 + np.sin(t / 58))))
rw2_null = jnp.column_stack([jnp.ones(n_days), t - t.mean()])


def log_marginal(log_tau):
    prior = gx.IntrinsicGMRF(
        jnp.zeros(n_days),
        jnp.exp(log_tau),
        gx.rw2_structure(n_days),
        null_space=rw2_null,
        constraint="hard",
    )
    result = gx.laplace_mode(prior, gx.PoissonLikelihood(counts))
    return result.log_marginal  # H stays BlockTriDiag


value, grad = jax.value_and_grad(log_marginal)(
    0.0
)  # exact: implicit diff + block Cholesky

The likelihood holds the observations, so laplace_mode(prior, likelihood) takes no separate y. For an odd number of nodes rw2_structure(n) returns n + 1 rows (a decoupled padding node): build the field on n + 1 nodes with null_space zero on the padding row, observe the first n through a row selection SparseOperator.from_coo(np.arange(n), np.arange(n), jnp.ones(n), (n, n + 1)) (which keeps \(H\) banded), and drop mode[n].

Structured linear algebra and Gaussian primitives for JAX.

LaplaceResult

Bases: Module

The Gaussian approximation \(\mathcal N(\hat x, H^{-1})\) at the mode.

Attributes:

Name	Type	Description
mode	Float[Array, ' N']	\(\hat x\), shape (N,). Implicitly differentiable in every array the prior and likelihood hold (\(\theta\)).
hessian	AbstractLinearOperator	\(H = Q + A^\top WA\) at the mode: a BlockTriDiag, SparseOperator, diagonal or dense operator like \(Q\) where the projector allows (see laplace_mode).
factor	Any	The factorisation of \(H\), with solve(b) and logdet(): a SparseCholeskyFactor (whose diag_inv gives the marginal variances by Takahashi) for a sparse \(H\), else the structured Cholesky factor's solves.
log_marginal	Float[Array, '']	\(\log\tilde\pi(y\mid\theta)\); see laplace_mode for the constants it includes.
n_iter	Int[Array, '']	Newton steps taken.
converged	Bool[Array, '']	Whether the last step moved \(x\) by less than the tolerance. Implicit gradients assume it; re-run or discard a \(\theta\) where it is False.

Source code in src/gaussx/_inference/_laplace.py

class LaplaceResult(eqx.Module):
    r"""The Gaussian approximation $\mathcal N(\hat x, H^{-1})$ at the mode.

    Attributes:
        mode: $\hat x$, shape ``(N,)``. Implicitly differentiable in every
            array the prior and likelihood hold ($\theta$).
        hessian: $H = Q + A^\top WA$ at the mode: a `BlockTriDiag`,
            `SparseOperator`, diagonal or dense operator like $Q$ where the
            projector allows (see `laplace_mode`).
        factor: The factorisation of $H$, with ``solve(b)`` and ``logdet()``:
            a `SparseCholeskyFactor` (whose ``diag_inv`` gives the marginal
            variances by Takahashi) for a sparse $H$, else the structured
            Cholesky factor's solves.
        log_marginal: $\log\tilde\pi(y\mid\theta)$; see `laplace_mode` for
            the constants it includes.
        n_iter: Newton steps taken.
        converged: Whether the last step moved $x$ by less than the
            tolerance. Implicit gradients assume it; re-run or discard a
            $\theta$ where it is ``False``.
    """

    mode: Float[Array, " N"]
    hessian: lx.AbstractLinearOperator
    factor: Any
    log_marginal: Float[Array, ""]
    n_iter: Int[Array, ""]
    converged: Bool[Array, ""]

laplace_mode(prior: GaussianMRF | IntrinsicGMRF, likelihood: AbstractLikelihood, *, projector: lx.AbstractLinearOperator | None = None, offset: Float[ArrayLike, ' M'] | Float[ArrayLike, ''] | None = None, init: Float[ArrayLike, ' N'] | None = None, max_iter: int = 50, tol: float = 1e-08, damping: float = 1.0) -> LaplaceResult

Mode, Hessian and log-marginal of the Laplace approximation, in precision form.

The latent Gaussian model is \(x \sim\) prior (precision \(Q\), mean \(\mu\)), \(\eta = Ax + o\) and \(y \mid \eta \sim\) likelihood, which holds \(y\) and must factorise over sites (its Hessian in \(\eta\) is diagonal). Newton on \(\ell(x) = \log p(y\mid Ax+o) + \log\pi(x)\):

\[ H_t = Q + A^\top W_tA,\qquad H_t\,\tilde x = Q\mu + A^\top(g_t + W_tAx_t),\qquad x_{t+1} = x_t + \text{damping}\,(\tilde x - x_t), \]

with \(g_t\), \(W_t = \max(-f''_t, 10^{-6})\) from the likelihood's site_derivatives (newton_update's diagonal path), until \(\|x_{t+1}-x_t\|_\infty \le \text{tol}\,(1 + \|x_{t+1}\|_\infty)\) (tol is floored at 64 machine epsilons of the dtype). \(H_t\) is rebuilt values-only on a structure fixed by \(Q\) and \(A\):

• \(A\) the identity (None) or one entry per row (a row selection, or a scaled one): \(A^\top WA\) is diagonal, so a BlockTriDiag \(Q\) (RW1, RW2, AR(1)) stays BlockTriDiag, and a SparseOperator, diagonal or dense \(Q\) keeps its class;
• a general SparseOperator \(A\) (a FEM projector, an intercept and covariates): SparseOperator.congruence on the union pattern, with a banded or diagonal \(Q\) converted to sparse;
• a dense \(A\) or \(Q\): dense; any other \(Q\) (a grid SpectralFunction or KroneckerSum with a non-diagonal update): the matrix-free sum, solved and log-determined by gaussx.solve / gaussx.logdet dispatch.

A sparse \(H\) is factored by SparseCholeskySolver, whose symbolic analysis is cached per pattern and shared by every step and every \(\theta\). A Gaussian likelihood converges in one step (a second confirms it).

Intrinsic priors. An IntrinsicGMRF contributes \(Q = \tau R\). With constraint="hard" the field lives on \(V^\top x = 0\) (\(V\) = null_space): each Newton solve is kriged onto it (c extra solves with the same factor), and the Gaussian approximation is the constrained density (Rue et al., 2009, eq. 3; Rue & Held, 2005, eq. 2.30). \(H\) must be positive definite on the whole space, which holds when the observations see \(\ker R\) (an RW2 observed at two or more times, a Besag field observed in every connected component). constraint="none" runs unconstrained; "soft" is not supported.

The log-marginal. At the mode \(\hat x\), \(\log\tilde\pi(y\mid\theta) = \log p(y\mid A\hat x + o) + \log\pi(\hat x) - \log\tilde\pi_G(\hat x)\), where \(\log\pi\) is the prior's log_prob and

\[ \log\tilde\pi_G(\hat x) = \tfrac12\log|H| - \tfrac N2\log 2\pi \;\Big[+ \tfrac12\log|V^\top H^{-1}V| - \tfrac12\log|V^\top V| + \tfrac c2\log 2\pi\Big], \]

the bracket for the hard constraint. For a GaussianMRF this is \(f(A\hat x + o) - \tfrac12(\hat x-\mu)^\top Q(\hat x-\mu) + \tfrac12\log|Q| - \tfrac12\log|H|\), the exact log-evidence for a Gaussian likelihood. For an IntrinsicGMRF it is exact with include_normalizer=True and otherwise off by the θ-free constant \(\tfrac12\log|R|_+ - \tfrac{N-c}2\log 2\pi\). The likelihood's own normalising terms are included.

Gradients. Every array in prior, likelihood, projector and offset is part of \(\theta\). The mode is differentiated by the implicit function theorem through jax.lax.custom_root (\(\partial_\theta\hat x = H^{-1}\partial_\theta\nabla\ell\), one solve with a factor of \(H\) at the mode), and the log-determinants through their structured or sparse (Takahashi) VJPs. Reverse-over-reverse (jax.jacrev(jax.jacrev(...)), theta_design's default Hessian) is exact too. The loop itself is not differentiated, so gradients assume converged.

Parameters:

Name	Type	Description	Default
prior	GaussianMRF | IntrinsicGMRF	The latent field, a GaussianMRF or IntrinsicGMRF.	required
likelihood	AbstractLikelihood	A site-factorising AbstractLikelihood holding \(y\) (PoissonLikelihood, BinomialLikelihood, NegativeBinomialLikelihood, BernoulliLikelihood, GaussianLikelihood, ...).	required
projector	AbstractLinearOperator | None	\(A\), shape (M, N); None is the identity.	None
offset	Float[ArrayLike, ' M'] | Float[ArrayLike, ''] | None	\(o\), a scalar or shape (M,) (e.g. log expected counts).	None
init	Float[ArrayLike, ' N'] | None	Starting point, shape (N,); zeros by default.	None
max_iter	int	Maximum Newton steps.	50
tol	float	Relative tolerance on the step's sup-norm.	1e-08
damping	float	Step size in (0, 1]; below 1 for hard, non-log-concave problems.	1.0

Returns:

Type	Description
LaplaceResult	The LaplaceResult.

Raises:

Type	Description
TypeError	For another prior class.
NotImplementedError	For an IntrinsicGMRF with constraint="soft".
ValueError	If projector does not have N columns.

Examples:

import jax
import jax.numpy as jnp
import numpy as np
import gaussx as gx

# Poisson counts with an RW2 seasonal effect; θ = log τ. Even n keeps
# rw2_structure(n) at size n (odd n adds a padding node: see below).
n = 60
t = jnp.arange(n, dtype=float)
counts = jnp.asarray(np.random.default_rng(0).poisson(np.exp(np.sin(t / 9.0))))
null = jnp.column_stack([jnp.ones(n), t - t.mean()])

def log_marginal(log_tau):
    prior = gx.IntrinsicGMRF(
        jnp.zeros(n), jnp.exp(log_tau), gx.rw2_structure(n), null
    )
    result = gx.laplace_mode(prior, gx.PoissonLikelihood(counts))
    return result.log_marginal  # H stays BlockTriDiag

value, grad = jax.value_and_grad(log_marginal)(2.0)  # exact

# Odd n: rw2_structure(n) has n + 1 rows. Observe the first n nodes
# through a row selection (H stays BlockTriDiag), with the null
# space zero on the padding node.
m = 61
s = jnp.arange(m, dtype=float)
V = jnp.zeros((m + 1, 2)).at[:m, 0].set(1.0).at[:m, 1].set(s - s.mean())
select = gx.SparseOperator.from_coo(
    np.arange(m), np.arange(m), jnp.ones(m), (m, m + 1)
)
prior = gx.IntrinsicGMRF(jnp.zeros(m + 1), 10.0, gx.rw2_structure(m), V)
y = jnp.ones(m)
fit = gx.laplace_mode(prior, gx.PoissonLikelihood(y), projector=select)
trend = fit.mode[:m]  # drop the padding node

Source code in src/gaussx/_inference/_laplace.py

def laplace_mode(
    prior: GaussianMRF | IntrinsicGMRF,
    likelihood: AbstractLikelihood,
    *,
    projector: lx.AbstractLinearOperator | None = None,
    offset: Float[ArrayLike, " M"] | Float[ArrayLike, ""] | None = None,
    init: Float[ArrayLike, " N"] | None = None,
    max_iter: int = 50,
    tol: float = 1e-8,
    damping: float = 1.0,
) -> LaplaceResult:
    r"""Mode, Hessian and log-marginal of the Laplace approximation, in precision form.

    The latent Gaussian model is $x \sim$ ``prior`` (precision $Q$, mean
    $\mu$), $\eta = Ax + o$ and $y \mid \eta \sim$ ``likelihood``, which
    holds $y$ and must factorise over sites (its Hessian in $\eta$ is
    diagonal). Newton on $\ell(x) = \log p(y\mid Ax+o) + \log\pi(x)$:

    $$
    H_t = Q + A^\top W_tA,\qquad
    H_t\,\tilde x = Q\mu + A^\top(g_t + W_tAx_t),\qquad
    x_{t+1} = x_t + \text{damping}\,(\tilde x - x_t),
    $$

    with $g_t$, $W_t = \max(-f''_t, 10^{-6})$ from the likelihood's
    `site_derivatives` (`newton_update`'s diagonal path), until
    $\|x_{t+1}-x_t\|_\infty \le \text{tol}\,(1 + \|x_{t+1}\|_\infty)$
    (``tol`` is floored at 64 machine epsilons of the dtype). $H_t$ is
    rebuilt values-only on a structure fixed by $Q$ and $A$:

    - $A$ the identity (``None``) or one entry per row (a row selection,
      or a scaled one): $A^\top WA$ is diagonal, so a `BlockTriDiag` $Q$
      (RW1, RW2, AR(1)) stays `BlockTriDiag`, and a `SparseOperator`, diagonal
      or dense $Q$ keeps its class;
    - a general `SparseOperator` $A$ (a FEM projector, an intercept and
      covariates): `SparseOperator.congruence` on the union pattern, with a
      banded or diagonal $Q$ converted to sparse;
    - a dense $A$ or $Q$: dense; any other $Q$ (a grid `SpectralFunction` or
      `KroneckerSum` with a non-diagonal update): the matrix-free sum, solved
      and log-determined by `gaussx.solve` / `gaussx.logdet` dispatch.

    A sparse $H$ is factored by `SparseCholeskySolver`, whose symbolic
    analysis is cached per pattern and shared by every step and every
    $\theta$. A Gaussian likelihood converges in one step (a second confirms
    it).

    **Intrinsic priors.** An `IntrinsicGMRF` contributes $Q = \tau R$. With
    ``constraint="hard"`` the field lives on $V^\top x = 0$ ($V$ =
    ``null_space``): each Newton solve is kriged onto it (``c`` extra
    solves with the same factor), and the Gaussian approximation is the
    constrained density (Rue et al., 2009, eq. 3; Rue & Held, 2005,
    eq. 2.30). $H$ must be positive definite on the whole space, which holds
    when the observations see $\ker R$ (an RW2 observed at two or more
    times, a Besag field observed in every connected component).
    ``constraint="none"`` runs unconstrained; ``"soft"`` is not supported.

    **The log-marginal.** At the mode $\hat x$,
    $\log\tilde\pi(y\mid\theta) = \log p(y\mid A\hat x + o) + \log\pi(\hat x)
    - \log\tilde\pi_G(\hat x)$, where $\log\pi$ is the prior's ``log_prob``
    and

    $$
    \log\tilde\pi_G(\hat x) = \tfrac12\log|H| - \tfrac N2\log 2\pi
        \;\Big[+ \tfrac12\log|V^\top H^{-1}V| - \tfrac12\log|V^\top V|
        + \tfrac c2\log 2\pi\Big],
    $$

    the bracket for the hard constraint. For a `GaussianMRF` this is
    $f(A\hat x + o) - \tfrac12(\hat x-\mu)^\top Q(\hat x-\mu)
    + \tfrac12\log|Q| - \tfrac12\log|H|$, the exact log-evidence for a
    Gaussian likelihood. For an `IntrinsicGMRF` it is exact with
    ``include_normalizer=True`` and otherwise off by the θ-free constant
    $\tfrac12\log|R|_+ - \tfrac{N-c}2\log 2\pi$. The likelihood's own
    normalising terms are included.

    **Gradients.** Every array in ``prior``, ``likelihood``, ``projector``
    and ``offset`` is part of $\theta$. The mode is differentiated by the
    implicit function theorem through `jax.lax.custom_root`
    ($\partial_\theta\hat x = H^{-1}\partial_\theta\nabla\ell$, one solve
    with a factor of $H$ at the mode), and the log-determinants through
    their structured or sparse (Takahashi) VJPs. Reverse-over-reverse
    (``jax.jacrev(jax.jacrev(...))``, `theta_design`'s default Hessian)
    is exact too. The loop itself is not differentiated, so gradients assume
    ``converged``.

    Args:
        prior: The latent field, a `GaussianMRF` or `IntrinsicGMRF`.
        likelihood: A site-factorising `AbstractLikelihood` holding $y$
            (`PoissonLikelihood`, `BinomialLikelihood`,
            `NegativeBinomialLikelihood`, `BernoulliLikelihood`,
            `GaussianLikelihood`, ...).
        projector: $A$, shape ``(M, N)``; ``None`` is the identity.
        offset: $o$, a scalar or shape ``(M,)`` (e.g. log expected counts).
        init: Starting point, shape ``(N,)``; zeros by default.
        max_iter: Maximum Newton steps.
        tol: Relative tolerance on the step's sup-norm.
        damping: Step size in ``(0, 1]``; below 1 for hard, non-log-concave
            problems.

    Returns:
        The `LaplaceResult`.

    Raises:
        TypeError: For another prior class.
        NotImplementedError: For an `IntrinsicGMRF` with
            ``constraint="soft"``.
        ValueError: If ``projector`` does not have ``N`` columns.

    Examples:
        ```python
        import jax
        import jax.numpy as jnp
        import numpy as np
        import gaussx as gx

        # Poisson counts with an RW2 seasonal effect; θ = log τ. Even n keeps
        # rw2_structure(n) at size n (odd n adds a padding node: see below).
        n = 60
        t = jnp.arange(n, dtype=float)
        counts = jnp.asarray(np.random.default_rng(0).poisson(np.exp(np.sin(t / 9.0))))
        null = jnp.column_stack([jnp.ones(n), t - t.mean()])

        def log_marginal(log_tau):
            prior = gx.IntrinsicGMRF(
                jnp.zeros(n), jnp.exp(log_tau), gx.rw2_structure(n), null
            )
            result = gx.laplace_mode(prior, gx.PoissonLikelihood(counts))
            return result.log_marginal  # H stays BlockTriDiag

        value, grad = jax.value_and_grad(log_marginal)(2.0)  # exact

        # Odd n: rw2_structure(n) has n + 1 rows. Observe the first n nodes
        # through a row selection (H stays BlockTriDiag), with the null
        # space zero on the padding node.
        m = 61
        s = jnp.arange(m, dtype=float)
        V = jnp.zeros((m + 1, 2)).at[:m, 0].set(1.0).at[:m, 1].set(s - s.mean())
        select = gx.SparseOperator.from_coo(
            np.arange(m), np.arange(m), jnp.ones(m), (m, m + 1)
        )
        prior = gx.IntrinsicGMRF(jnp.zeros(m + 1), 10.0, gx.rw2_structure(m), V)
        y = jnp.ones(m)
        fit = gx.laplace_mode(prior, gx.PoissonLikelihood(y), projector=select)
        trend = fit.mode[:m]  # drop the padding node
        ```
    """
    model = _model(prior, likelihood, projector, offset)
    n = prior.event_shape[0]
    dtype = prior.loc.dtype
    x0 = jnp.zeros(n, dtype=dtype) if init is None else jnp.asarray(init, dtype=dtype)
    x0 = x0 - model.null_component(x0)

    x_star, n_iter, converged = _newton_loop(
        _stop_gradient(model), jax.lax.stop_gradient(x0), max_iter, tol, damping
    )
    mode = _implicit_mode(model, x_star)

    _, H = model.hessian(mode)
    factor = _factor(H)
    eta = model.project(mode) + model.offset
    log_gaussian = 0.5 * factor.logdet() - 0.5 * n * _LOG_2PI
    _, S = model.krige(factor, mode)
    if S is not None:
        V = model.null_space
        assert V is not None
        c = V.shape[1]
        gram = einsum(V, V, "n c, n d -> c d")
        log_gaussian = (
            log_gaussian
            + jnp.sum(jnp.log(jnp.diagonal(jnp.linalg.cholesky(S))))
            - 0.5 * jnp.linalg.slogdet(gram)[1]
            + 0.5 * c * _LOG_2PI
        )
    log_marginal = likelihood.log_prob(eta) + prior.log_prob(mode) - log_gaussian
    return LaplaceResult(mode, H, factor, log_marginal, n_iter, converged)

Low-rank VB mean correction (INLA)

Under a skewed likelihood (Bernoulli, Poisson with small counts) the Laplace mode is a biased estimate of the posterior mean. vb_mean_correction (Van Niekerk & Rue, 2024; R-INLA's default control.vb correction) keeps the Laplace covariance \(H^{-1}\) and shifts the mean within a \(p\)-dimensional subspace, \(\bar x = \hat x + H^{-1}S\delta\), with \(\delta\) maximising

\[ \mathcal L(\delta) = \sum_i\mathbb E_{\eta_i\sim\mathcal N(m_i(\delta),\,v_i)} \big[\log p(y_i\mid\eta_i)\big] - \tfrac12(\bar x-\mu)^\top Q(\bar x-\mu), \]

\(m = A\bar x + o\), \(v_i = [AH^{-1}A^\top]_{ii}\). The expectations are 1-D Gauss-Hermite per site, the predictor variances come from the same Takahashi selected inverse as the marginal variances, and the Newton steps on \(\delta\) are \(p\)-dimensional, so it costs \(p\) solves with the Laplace factor. For a Gaussian likelihood the correction is exactly zero. Typically \(S\) spans the fixed effects:

import jax.numpy as jnp
import numpy as np

import gaussx as gx

# Rare detections along a transect: logit p_i = β₀ + u_i, u an RW2 under
# Σu = Σ(t − t̄)u = 0 and β₀ ~ N(0, 10³); latent x = (u, β₀)
n = 40
t = np.arange(n, dtype=float)
p_true = 1 / (1 + np.exp(2 - np.sin(t / 5)))
detected = jnp.asarray(np.random.default_rng(0).random(n) < p_true, float)
R = np.asarray(gx.rw2_structure(n).as_matrix())
rows, cols = np.nonzero(np.tril(R))
Q = gx.SparseOperator.from_coo(
    np.r_[rows, n],
    np.r_[cols, n],
    jnp.r_[2.0 * R[rows, cols], 1e-3],
    (n + 1, n + 1),
    symmetric=True,
)
V = jnp.zeros((n + 1, 2)).at[:n, 0].set(1.0).at[:n, 1].set(t - t.mean())
prior = gx.IntrinsicGMRF(jnp.zeros(n + 1), 1.0, Q, V)
A = gx.SparseOperator.from_coo(  # η = u + β₀
    np.r_[np.arange(n), np.arange(n)],
    np.r_[np.arange(n), np.full(n, n)],
    jnp.ones(2 * n),
    (n, n + 1),
)
detection = gx.BernoulliLikelihood(detected)

res = gx.laplace_mode(prior, detection, projector=A)
fixed_effect_idx = jnp.array([n])
mean_vb = gx.vb_mean_correction(
    res, prior, detection, projector=A, subspace=fixed_effect_idx
)

As with laplace_mode, the likelihood holds the observations, so there is no separate y argument; pass the same projector and offset as to laplace_mode.

Structured linear algebra and Gaussian primitives for JAX.

vb_mean_correction(result: LaplaceResult, prior: GaussianMRF | IntrinsicGMRF, likelihood: AbstractLikelihood, *, subspace: Int[ArrayLike, ' p'] | Float[ArrayLike, 'N p'], projector: lx.AbstractLinearOperator | None = None, offset: Float[ArrayLike, ' M'] | Float[ArrayLike, ''] | None = None, n_iter: int = 5, integrator: AbstractIntegrator = _GAUSS_HERMITE_20) -> Float[Array, ' N']

Low-rank VB correction of the Laplace mean (Van Niekerk & Rue, 2024).

The Laplace approximation \(\mathcal N(\hat x, H^{-1})\) is centred at the mode, which under a skewed likelihood (Poisson with small counts, Bernoulli) is a biased estimate of the posterior mean. Keeping the covariance \(H^{-1}\), shift the mean within the span of \(S\) (subspace, \(N\times p\)),

\[ q_\delta = \mathcal N\big(\bar x(\delta),\ H^{-1}\big),\qquad \bar x(\delta) = \hat x + H^{-1}S\delta, \]

and choose \(\delta\) to maximise the evidence lower bound, whose covariance terms do not depend on \(\delta\):

\[ \mathcal L(\delta) = \sum_i\mathbb E_{\eta_i\sim\mathcal N(m_i(\delta),\,v_i)} \big[\log p(y_i\mid\eta_i)\big] - \tfrac12\big(\bar x(\delta)-\mu\big)^\top Q\big(\bar x(\delta)-\mu\big), \]

with \(m(\delta) = A\bar x(\delta) + o\) and \(v_i = [AH^{-1}A^\top]_{ii}\). With \(U = H^{-1}S\), \(B = AU\) and the site expectations \(\bar g_i = \mathbb E[\partial_\eta\log p(y_i\mid\eta_i)]\), \(\bar h_i = \mathbb E[\partial^2_\eta\log p(y_i\mid\eta_i)]\) (1-D quadrature by integrator on each site), n_iter Newton steps on the \(p\)-dimensional \(\delta\) (from \(\delta = 0\)) use

\[ \nabla\mathcal L = B^\top\bar g - U^\top Q(\bar x-\mu),\qquad \nabla^2\mathcal L = B^\top\operatorname{diag}(\bar h)B - U^\top QU. \]

For a Gaussian likelihood the mode is the exact posterior mean and the correction is zero.

Cost. \(p\) solves with result.factor (plus \(c\) for a hard constraint), then \(p\)-dimensional Newton steps whose site expectations are order × M likelihood derivatives. The predictor variances come from the selected inverse, not from solves: for an identity or row-selection projector they are \(A_{ij}^2\Sigma_{jj}\) (a BlockTriDiag, sparse or diagonal \(H\)), and for a general sparse projector (a FEM projector, fixed-effect columns) every pair of nodes in a row of \(A\) is coupled in \(A^\top WA\), so each \(\Sigma_{ab}\) needed lies in Takahashi's pattern(L + Lᵀ). A dense \(A\) or \(H\) costs one solve per row of \(A\). The Newton step is the minimum-norm one, so directions of \(S\) that do not move \(\bar x\) (dependent columns, or ones in the constrained-out span of \(V\)) are harmless.

Intrinsic priors. With constraint="hard" both the covariance and the shift are those of the constrained approximation, kriged onto \(V^\top x = 0\): \(\Sigma = H^{-1} - H^{-1}V(V^\top H^{-1}V)^{-1}V^\top H^{-1}\) and \(\bar x(\delta) = \hat x + \Sigma S\delta\).

Gradients. The result is an ordinary JAX function of the arrays in result, prior, likelihood and projector, so it can be differentiated in reverse mode (through laplace_mode's implicit mode and the factor's solves); n_iter fixed Newton steps are differentiated as unrolled.

Parameters:

Name	Type	Description	Default
result	LaplaceResult	laplace_mode's output for the same model.	required
prior	GaussianMRF | IntrinsicGMRF	The latent field passed to laplace_mode.	required
likelihood	AbstractLikelihood	The site-factorising likelihood passed to laplace_mode (it holds \(y\)).	required
subspace	Int[ArrayLike, ' p'] | Float[ArrayLike, 'N p']	The span of the mean shift: an integer index array (S is those columns of the identity, e.g. the fixed effects' nodes) or a dense basis of shape (N, p).	required
projector	AbstractLinearOperator | None	\(A\) as passed to laplace_mode; None is the identity.	None
offset	Float[ArrayLike, ' M'] | Float[ArrayLike, ''] | None	\(o\) as passed to laplace_mode.	None
n_iter	int	Newton steps on \(\delta\).	5
integrator	AbstractIntegrator	A point-based 1-D rule for the site expectations; its points and weights for \(\mathcal N(0, 1)\) are reused for every site.	_GAUSS_HERMITE_20

Returns:

Type	Description
Float[Array, ' N']	The corrected mean \(\bar x(\delta)\), shape (N,).

Raises:

Type	Description
ValueError	If subspace has the wrong shape.

Examples:

import jax.numpy as jnp
import numpy as np
import gaussx as gx

# Rare detections along a transect: logit p_i = β₀ + u_i, u an RW2
# under Σu = Σ(t − t̄)u = 0 and β₀ ~ N(0, 10³), latent x = (u, β₀).
n = 40
t = np.arange(n, dtype=float)
p_true = 1 / (1 + np.exp(2 - np.sin(t / 5)))
detected = jnp.asarray(np.random.default_rng(0).random(n) < p_true, float)
R = np.asarray(gx.rw2_structure(n).as_matrix())
rows, cols = np.nonzero(np.tril(R))
Q = gx.SparseOperator.from_coo(
    np.r_[rows, n], np.r_[cols, n], jnp.r_[2.0 * R[rows, cols], 1e-3],
    (n + 1, n + 1), symmetric=True,
)
V = jnp.zeros((n + 1, 2)).at[:n, 0].set(1.0).at[:n, 1].set(t - t.mean())
prior = gx.IntrinsicGMRF(jnp.zeros(n + 1), 1.0, Q, V)
A = gx.SparseOperator.from_coo(  # η = u + β₀
    np.r_[np.arange(n), np.arange(n)], np.r_[np.arange(n), np.full(n, n)],
    jnp.ones(2 * n), (n, n + 1),
)
lik = gx.BernoulliLikelihood(detected)

res = gx.laplace_mode(prior, lik, projector=A)
fixed_effect_idx = jnp.array([n])
mean_vb = gx.vb_mean_correction(
    res, prior, lik, projector=A, subspace=fixed_effect_idx
)

Source code in src/gaussx/_inference/_vb_correction.py

def vb_mean_correction(
    result: LaplaceResult,
    prior: GaussianMRF | IntrinsicGMRF,
    likelihood: AbstractLikelihood,
    *,
    subspace: Int[ArrayLike, " p"] | Float[ArrayLike, "N p"],
    projector: lx.AbstractLinearOperator | None = None,
    offset: Float[ArrayLike, " M"] | Float[ArrayLike, ""] | None = None,
    n_iter: int = 5,
    integrator: AbstractIntegrator = _GAUSS_HERMITE_20,
) -> Float[Array, " N"]:
    r"""Low-rank VB correction of the Laplace mean (Van Niekerk & Rue, 2024).

    The Laplace approximation $\mathcal N(\hat x, H^{-1})$ is centred at the
    mode, which under a skewed likelihood (Poisson with small counts,
    Bernoulli) is a biased estimate of the posterior mean. Keeping the
    covariance $H^{-1}$, shift the mean within the span of $S$ (``subspace``,
    $N\times p$),

    $$
    q_\delta = \mathcal N\big(\bar x(\delta),\ H^{-1}\big),\qquad
    \bar x(\delta) = \hat x + H^{-1}S\delta,
    $$

    and choose $\delta$ to maximise the evidence lower bound, whose
    covariance terms do not depend on $\delta$:

    $$
    \mathcal L(\delta) = \sum_i\mathbb E_{\eta_i\sim\mathcal N(m_i(\delta),\,v_i)}
        \big[\log p(y_i\mid\eta_i)\big]
        - \tfrac12\big(\bar x(\delta)-\mu\big)^\top Q\big(\bar x(\delta)-\mu\big),
    $$

    with $m(\delta) = A\bar x(\delta) + o$ and $v_i = [AH^{-1}A^\top]_{ii}$.
    With $U = H^{-1}S$, $B = AU$ and the site expectations
    $\bar g_i = \mathbb E[\partial_\eta\log p(y_i\mid\eta_i)]$,
    $\bar h_i = \mathbb E[\partial^2_\eta\log p(y_i\mid\eta_i)]$ (1-D
    quadrature by ``integrator`` on each site), ``n_iter`` Newton steps
    on the $p$-dimensional $\delta$ (from $\delta = 0$) use

    $$
    \nabla\mathcal L = B^\top\bar g - U^\top Q(\bar x-\mu),\qquad
    \nabla^2\mathcal L = B^\top\operatorname{diag}(\bar h)B - U^\top QU.
    $$

    For a Gaussian likelihood the mode is the exact posterior mean and the
    correction is zero.

    **Cost.** $p$ solves with ``result.factor`` (plus $c$ for a hard
    constraint), then $p$-dimensional Newton steps whose site expectations
    are ``order × M`` likelihood derivatives. The predictor variances come
    from the selected inverse, not from solves: for an identity or
    row-selection projector they are $A_{ij}^2\Sigma_{jj}$ (a
    `BlockTriDiag`, sparse or diagonal $H$), and for a general sparse
    projector (a FEM projector, fixed-effect columns) every pair of nodes
    in a row of $A$ is coupled in $A^\top WA$, so each $\Sigma_{ab}$ needed
    lies in Takahashi's ``pattern(L + Lᵀ)``. A dense $A$ or $H$ costs one
    solve per row of $A$. The Newton step is the minimum-norm one, so
    directions of $S$ that do not move $\bar x$ (dependent columns, or
    ones in the constrained-out span of $V$) are harmless.

    **Intrinsic priors.** With ``constraint="hard"`` both the covariance
    and the shift are those of the constrained approximation, kriged onto
    $V^\top x = 0$: $\Sigma = H^{-1} - H^{-1}V(V^\top H^{-1}V)^{-1}V^\top H^{-1}$
    and $\bar x(\delta) = \hat x + \Sigma S\delta$.

    **Gradients.** The result is an ordinary JAX function of the arrays in
    ``result``, ``prior``, ``likelihood`` and ``projector``, so it can be
    differentiated in reverse mode (through ``laplace_mode``'s implicit
    mode and the factor's solves); ``n_iter`` fixed Newton steps are
    differentiated as unrolled.

    Args:
        result: `laplace_mode`'s output for the same model.
        prior: The latent field passed to `laplace_mode`.
        likelihood: The site-factorising likelihood passed to `laplace_mode`
            (it holds $y$).
        subspace: The span of the mean shift: an integer index array
            (``S`` is those columns of the identity, e.g. the fixed effects'
            nodes) or a dense basis of shape ``(N, p)``.
        projector: $A$ as passed to `laplace_mode`; ``None`` is the identity.
        offset: $o$ as passed to `laplace_mode`.
        n_iter: Newton steps on $\delta$.
        integrator: A point-based 1-D rule for the site expectations; its
            points and weights for $\mathcal N(0, 1)$ are reused for every
            site.

    Returns:
        The corrected mean $\bar x(\delta)$, shape ``(N,)``.

    Raises:
        ValueError: If ``subspace`` has the wrong shape.

    Examples:
        ```python
        import jax.numpy as jnp
        import numpy as np
        import gaussx as gx

        # Rare detections along a transect: logit p_i = β₀ + u_i, u an RW2
        # under Σu = Σ(t − t̄)u = 0 and β₀ ~ N(0, 10³), latent x = (u, β₀).
        n = 40
        t = np.arange(n, dtype=float)
        p_true = 1 / (1 + np.exp(2 - np.sin(t / 5)))
        detected = jnp.asarray(np.random.default_rng(0).random(n) < p_true, float)
        R = np.asarray(gx.rw2_structure(n).as_matrix())
        rows, cols = np.nonzero(np.tril(R))
        Q = gx.SparseOperator.from_coo(
            np.r_[rows, n], np.r_[cols, n], jnp.r_[2.0 * R[rows, cols], 1e-3],
            (n + 1, n + 1), symmetric=True,
        )
        V = jnp.zeros((n + 1, 2)).at[:n, 0].set(1.0).at[:n, 1].set(t - t.mean())
        prior = gx.IntrinsicGMRF(jnp.zeros(n + 1), 1.0, Q, V)
        A = gx.SparseOperator.from_coo(  # η = u + β₀
            np.r_[np.arange(n), np.arange(n)], np.r_[np.arange(n), np.full(n, n)],
            jnp.ones(2 * n), (n, n + 1),
        )
        lik = gx.BernoulliLikelihood(detected)

        res = gx.laplace_mode(prior, lik, projector=A)
        fixed_effect_idx = jnp.array([n])
        mean_vb = gx.vb_mean_correction(
            res, prior, lik, projector=A, subspace=fixed_effect_idx
        )
        ```
    """
    model = _model(prior, likelihood, projector, offset)
    St = _basis(subspace, result.mode.shape[0], result.mode.dtype)
    return _correct(model, result, St, n_iter, integrator)

Ensemble covariances, gain & analysis

Bessel-corrected empirical (cross-)covariances from ensemble members, the ensemble Kalman gain built from them, and the analysis step that applies it.

The gain functions are the pieces; enkf_analysis is the step -- the stochastic (perturbed-observation) update that turns a prior ensemble and an observation into a posterior ensemble. etkf_transform is its deterministic square-root counterpart.

A caveat worth stating up front: the Gaussian assumption in an ensemble Kalman filter is a property of the coordinates, not of the algorithm. Applied to a non-Gaussian prior the update is biased, and the bias does not shrink with ensemble size. Conjugating the update with a bijection that Gaussianises the prior -- warp, analyse, warp back -- removes it.

That conjugated update is exact Bayes only under conditions worth stating precisely, since they are easy to over-claim. It holds in the population limit -- with a finite ensemble the gain is empirical and the perturbations are Monte Carlo, so the result is an estimate regardless -- and only when the observation model is affine with additive Gaussian noise in the same latent coordinates that Gaussianise the prior. A Gaussian conditional likelihood is not sufficient on its own: y = z² + ε has Gaussian noise and a non-Gaussian posterior that no Kalman update reproduces. Outside those conditions conjugation is an approximation with no guaranteed ordering against the physical-space update -- usually much better, but a badly matched warp can make the latent joint less Gaussian and do worse.

Structured linear algebra and Gaussian primitives for JAX.

ensemble_covariance(particles: Float[Array, 'J N'], *, bessel: bool = False) -> LowRankUpdate

Empirical covariance from an ensemble as a low-rank operator.

Returns C = c X'^T X' with c = 1 / J when bessel=False (default, maximum likelihood) and c = 1 / (J - 1) when bessel=True (unbiased / ensemble Kalman filter convention). The result is a LowRankUpdate of rank <= J-1 rather than materializing the full (N, N) matrix. Efficient when J << N.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Ensemble of shape (J, N).	required
bessel	bool	If True, apply the 1 / (J - 1) Bessel correction used throughout the ensemble Kalman filter literature. This lower-level helper defaults to False for backwards compatibility; ensemble_kalman_gain defaults to True for the EnKF convention.	False

Returns:

Type	Description
LowRankUpdate	A LowRankUpdate operator representing the empirical
LowRankUpdate	covariance, with a zero base and J-column low-rank factor.

Note

The bessel default differs across this module: these two covariance helpers default to bessel=False (1 / J), while ensemble_kalman_gain, localized_kalman_gain, enkf_analysis, eki_step and discrepancy_step_size default to True (1 / (J - 1)) and etkf_transform is 1 / (J - 1) throughout. A gain assembled by hand from these helpers with their defaults therefore differs from ensemble_kalman_gain's; pass bessel=True to match it.

Source code in src/gaussx/_inference/_ensemble_stats.py

def ensemble_covariance(
    particles: Float[Array, "J N"],
    *,
    bessel: bool = False,
) -> LowRankUpdate:
    r"""Empirical covariance from an ensemble as a low-rank operator.

    Returns ``C = c X'^T X'`` with ``c = 1 / J`` when ``bessel=False``
    (default, maximum likelihood) and ``c = 1 / (J - 1)`` when
    ``bessel=True`` (unbiased / ensemble Kalman filter convention).
    The result is a ``LowRankUpdate`` of rank ``<= J-1`` rather than
    materializing the full ``(N, N)`` matrix.  Efficient when
    ``J << N``.

    Args:
        particles: Ensemble of shape ``(J, N)``.
        bessel: If True, apply the ``1 / (J - 1)`` Bessel correction
            used throughout the ensemble Kalman filter literature. This
            lower-level helper defaults to False for backwards compatibility;
            `ensemble_kalman_gain` defaults to True for the EnKF
            convention.

    Returns:
        A ``LowRankUpdate`` operator representing the empirical
        covariance, with a zero base and ``J``-column low-rank factor.

    Note:
        The ``bessel`` default differs across this module: these two
        covariance helpers default to ``bessel=False`` (``1 / J``), while
        `ensemble_kalman_gain`, `localized_kalman_gain`, `enkf_analysis`,
        `eki_step` and `discrepancy_step_size` default to ``True``
        (``1 / (J - 1)``) and `etkf_transform` is ``1 / (J - 1)``
        throughout. A gain assembled by hand from these helpers with their
        defaults therefore differs from `ensemble_kalman_gain`'s; pass
        ``bessel=True`` to match it.
    """
    J, N = particles.shape
    _check_ensemble_size(J, bessel)
    mean = jnp.mean(particles, axis=0)
    deviations = particles - mean[None, :]  # (J, N)

    divisor = J - 1 if bessel else J
    U = deviations.T / jnp.sqrt(divisor)  # (N, J)

    base = lx.DiagonalLinearOperator(jnp.zeros(N, dtype=particles.dtype))
    return LowRankUpdate(base, U)

ensemble_cross_covariance(particles_theta: Float[Array, 'J N'], particles_G: Float[Array, 'J M'], *, bessel: bool = False) -> Float[Array, 'N M']

Cross-covariance between two ensemble sets.

Computes C^{theta,G} = c sum_j (theta_j - bar)(G_j - bar)^T with c = 1 / J by default or c = 1 / (J - 1) when bessel=True.

Parameters:

Name	Type	Description	Default
particles_theta	Float[Array, 'J N']	First ensemble, shape (J, N).	required
particles_G	Float[Array, 'J M']	Second ensemble, shape (J, M).	required
bessel	bool	If True, apply the 1 / (J - 1) Bessel correction used by ensemble Kalman filter recipes. This lower-level helper defaults to False for backwards compatibility; ensemble_kalman_gain defaults to True for the EnKF convention.	False

Returns:

Type	Description
Float[Array, 'N M']	Cross-covariance array of shape (N, M).

Note

The bessel default differs across this module: these two covariance helpers default to bessel=False (1 / J), while ensemble_kalman_gain, localized_kalman_gain, enkf_analysis, eki_step and discrepancy_step_size default to True (1 / (J - 1)) and etkf_transform is 1 / (J - 1) throughout. A gain assembled by hand from these helpers with their defaults therefore differs from ensemble_kalman_gain's; pass bessel=True to match it.

Source code in src/gaussx/_inference/_ensemble_stats.py

def ensemble_cross_covariance(
    particles_theta: Float[Array, "J N"],
    particles_G: Float[Array, "J M"],
    *,
    bessel: bool = False,
) -> Float[Array, "N M"]:
    r"""Cross-covariance between two ensemble sets.

    Computes ``C^{theta,G} = c sum_j (theta_j - bar)(G_j - bar)^T``
    with ``c = 1 / J`` by default or ``c = 1 / (J - 1)`` when
    ``bessel=True``.

    Args:
        particles_theta: First ensemble, shape ``(J, N)``.
        particles_G: Second ensemble, shape ``(J, M)``.
        bessel: If True, apply the ``1 / (J - 1)`` Bessel correction
            used by ensemble Kalman filter recipes. This lower-level helper
            defaults to False for backwards compatibility; `ensemble_kalman_gain`
            defaults to True for the EnKF convention.

    Returns:
        Cross-covariance array of shape ``(N, M)``.

    Note:
        The ``bessel`` default differs across this module: these two
        covariance helpers default to ``bessel=False`` (``1 / J``), while
        `ensemble_kalman_gain`, `localized_kalman_gain`, `enkf_analysis`,
        `eki_step` and `discrepancy_step_size` default to ``True``
        (``1 / (J - 1)``) and `etkf_transform` is ``1 / (J - 1)``
        throughout. A gain assembled by hand from these helpers with their
        defaults therefore differs from `ensemble_kalman_gain`'s; pass
        ``bessel=True`` to match it.
    """
    J = particles_theta.shape[0]
    _check_ensemble_size(J, bessel)
    dev_theta = particles_theta - jnp.mean(particles_theta, axis=0, keepdims=True)
    dev_G = particles_G - jnp.mean(particles_G, axis=0, keepdims=True)
    divisor = J - 1 if bessel else J
    return (dev_theta.T @ dev_G) / divisor

ensemble_kalman_gain(particles: Float[Array, 'J N'], obs_particles: Float[Array, 'J M'], obs_noise: lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None, bessel: bool = True, dense_innovation: bool | None = None) -> Float[Array, 'N M']

Kalman gain from an ensemble and its image in observation space.

Computes K = C^{xH} (C^{HH} + R)^{-1}, where C^{xH} is the state-observation cross-covariance and C^{HH} is the observation-space ensemble covariance. The innovation covariance S = C^{HH} + R is assembled by one of two routes, which solve the same system and agree to round-off:

• Woodbury (dense_innovation=False): S is a LowRankUpdate of R, and solve_rows inverts a (J, J) capacitance through R's own structured solve. Costs O(J^3 + J^2 M) and never materialises R -- the route for a few dozen members against many observations. It solves against R itself, so R must be positive definite.
• Dense (dense_innovation=True): S is formed as an (M, M) matrix, calling obs_noise.as_matrix(), and solved directly. Costs O(J M^2 + M^3) -- the route for many members against few observations, where the capacitance would be the larger matrix (3.2 GB in float64 at J = 20 000). It also handles a positive semi-definite R.

dense_innovation=None (the default) picks dense when J >= M and Woodbury otherwise, the same rule as enkf_analysis and eki_step.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Prior ensemble in state space, shape (J, N).	required
obs_particles	Float[Array, 'J M']	Prior ensemble in observation space, shape (J, M).	required
obs_noise	AbstractLinearOperator	Observation error covariance operator, shape (M, M).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. None uses structural dispatch. A matrix-free strategy (e.g. CGSolver) wants dense_innovation=False so it is handed the structured operator.	None
bessel	bool	Defaults to True, unlike the lower-level covariance helpers, because this recipe follows the unbiased EnKF convention. Use False for maximum-likelihood recipes with a 1 / J divisor.	True
dense_innovation	bool | None	Whether to form the (M, M) innovation densely. None chooses by shape (dense when J >= M); True / False force the dense / Woodbury route.	None

Returns:

Type	Description
Float[Array, 'N M']	Dense Kalman gain of shape (N, M).

Source code in src/gaussx/_inference/_enkf.py

def ensemble_kalman_gain(
    particles: Float[Array, "J N"],
    obs_particles: Float[Array, "J M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
    bessel: bool = True,
    dense_innovation: bool | None = None,
) -> Float[Array, "N M"]:
    r"""Kalman gain from an ensemble and its image in observation space.

    Computes ``K = C^{xH} (C^{HH} + R)^{-1}``, where ``C^{xH}`` is the
    state-observation cross-covariance and ``C^{HH}`` is the
    observation-space ensemble covariance. The innovation covariance
    ``S = C^{HH} + R`` is assembled by one of two routes, which solve the
    same system and agree to round-off:

    - **Woodbury** (``dense_innovation=False``): ``S`` is a `LowRankUpdate`
      of ``R``, and ``solve_rows`` inverts a ``(J, J)`` capacitance through
      ``R``'s own structured solve. Costs ``O(J^3 + J^2 M)`` and never
      materialises ``R`` -- the route for a few dozen members against many
      observations. It solves against ``R`` itself, so ``R`` must be
      positive **definite**.
    - **Dense** (``dense_innovation=True``): ``S`` is formed as an ``(M, M)``
      matrix, calling ``obs_noise.as_matrix()``, and solved directly. Costs
      ``O(J M^2 + M^3)`` -- the route for many members against few
      observations, where the capacitance would be the larger matrix
      (3.2 GB in float64 at ``J = 20 000``). It also handles a positive
      *semi*-definite ``R``.

    ``dense_innovation=None`` (the default) picks dense when ``J >= M`` and
    Woodbury otherwise, the same rule as `enkf_analysis` and `eki_step`.

    Args:
        particles: Prior ensemble in state space, shape ``(J, N)``.
        obs_particles: Prior ensemble in observation space, shape ``(J, M)``.
        obs_noise: Observation error covariance operator, shape ``(M, M)``.
        solver: Optional solver strategy. ``None`` uses structural dispatch.
            A matrix-free strategy (e.g. `CGSolver`) wants
            ``dense_innovation=False`` so it is handed the structured operator.
        bessel: Defaults to True, unlike the lower-level covariance helpers,
            because this recipe follows the unbiased EnKF convention. Use
            False for maximum-likelihood recipes with a ``1 / J`` divisor.
        dense_innovation: Whether to form the ``(M, M)`` innovation densely.
            ``None`` chooses by shape (dense when ``J >= M``); ``True`` /
            ``False`` force the dense / Woodbury route.

    Returns:
        Dense Kalman gain of shape ``(N, M)``.
    """
    n_ens, n_obs = obs_particles.shape
    if particles.shape[0] != n_ens:
        raise ValueError(
            "particles and obs_particles must share the same ensemble size, "
            f"got J={particles.shape[0]} and J={n_ens}."
        )
    use_dense = n_ens >= n_obs if dense_innovation is None else dense_innovation
    cross_cov = ensemble_cross_covariance(
        particles,
        obs_particles,
        bessel=bessel,
    )
    if use_dense:
        # A `LowRankUpdate` innovation would send `solve_rows` through Woodbury
        # and form a (J, J) capacitance matrix -- 320 GB at J = 200_000 -- so
        # assemble the (M, M) innovation densely instead. Same solve, same
        # answer.
        obs_cov = ensemble_cross_covariance(obs_particles, obs_particles, bessel=bessel)
        innovation = symmetrize(obs_cov + obs_noise.as_matrix())  # (M, M)
        innovation_op = lx.MatrixLinearOperator(
            innovation, lx.positive_semidefinite_tag
        )
        return solve_rows(innovation_op, cross_cov, solver=solver)  # (N, M)
    # Fewer members than observations: the Woodbury capacitance is (J, J) and
    # cheap, so keep the ensemble term low-rank.
    innovation_cov = ensemble_covariance(obs_particles, bessel=bessel)
    innovation_cov = LowRankUpdate(obs_noise, innovation_cov.U)
    return solve_rows(innovation_cov, cross_cov, solver=solver)

enkf_analysis(particles: Float[Array, 'J N'], obs_particles: Float[Array, 'J M'], observation: Float[Array, ' M'], obs_noise: lx.AbstractLinearOperator, *, key: PRNGKeyArray | None = None, perturbed_obs: Float[Array, 'J M'] | None = None, localization: Float[Array, 'N M'] | None = None, obs_localization: Float[Array, 'M M'] | None = None, solver: AbstractSolverStrategy | None = None, dense_innovation: bool | None = None, bessel: bool = True) -> Float[Array, 'J N']

Stochastic (perturbed-observation) ensemble Kalman analysis step.

Updates a prior ensemble \(X^f\) toward an observation \(y\):

\[ X^a_j = X^f_j + K\,(y + \varepsilon_j - \mathcal{H}(X^f_j)), \qquad \varepsilon_j \sim N(0, R), \]

with \(K\) from ensemble_kalman_gain (or localized_kalman_gain when localization is given). The observation operator enters only through obs_particles -- the image \(\mathcal{H}(X^f)\) of the prior ensemble in observation space -- so nonlinear operators need no special handling.

The perturbation \(\varepsilon_j\) is what keeps the analysis spread correct. The deterministic update \(X^a_j = X^f_j + K(y - \mathcal{H}(X^f_j))\) drives the ensemble covariance to \((I - KH)P(I - KH)^\top\) instead of \((I - KH)P\), i.e. under-dispersive. There is deliberately no perturb=False flag: the deterministic alternative is a different filter (the square-root / ETKF family, see etkf_transform), not an option on this one.

Two ways to supply the observation perturbations:

• key -- draw \(\varepsilon_j \sim N(0, R)\) internally, via an exact PSD factor \(L L^\top = R\) of obs_noise that keeps its structure (a symmetric square root at dense leaves, so a singular \(R\) is fine).
• perturbed_obs -- pass a pre-built perturbed-observation ensemble \(y + \varepsilon_j\). Preferred when the same noise realisation must be reused across filters, and when the perturbations come from a nonlinear observation model rather than an additive \(R\).

Exactly one of key / perturbed_obs must be given.

Known limitation. The update is a Gaussian one, applied in whatever coordinates the caller supplies. For a non-Gaussian prior it is biased, and the bias does not shrink with ensemble size -- it is an error of coordinates, not of sampling. On the lognormal / logit-normal prior of Chipilski (2025), whose exact posterior mean is [0.548062, 0.353937], the physical-space update plateaus several percent off that value and stays there as \(J\) grows by two orders of magnitude.

The fix is to conjugate the update with a bijection \(\Gamma\) that Gaussianises the prior -- call this function on \(\Gamma^{-1}(X^f)\) and map the result back through \(\Gamma\): the ensemble Kalman filter's Gaussian assumption is a statement about coordinates, not about the algorithm. Pass the same perturbed_obs through both routes to compare them on one noise realisation.

That conjugated update is exact Bayes only under conditions worth stating precisely, because it is easy to over-claim. It needs the population limit -- with a finite ensemble the gain is empirical and the perturbations are Monte Carlo, so the result is an estimate either way -- and it needs the observation model to be affine with additive Gaussian noise in the same latent coordinates that Gaussianise the prior. A merely "Gaussian likelihood" is not enough: \(y = \zeta^2 + \varepsilon\) has Gaussian noise and a non-Gaussian posterior that no Kalman update reproduces. Outside those conditions conjugation is an approximation with no guaranteed ordering against the physical-space update -- usually much better, but a badly matched \(\Gamma\) can make the latent joint less Gaussian and do worse.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Prior ensemble in state space, shape (J, N).	required
obs_particles	Float[Array, 'J M']	Prior ensemble in observation space, shape (J, M).	required
observation	Float[Array, ' M']	The observation, shape (M,).	required
obs_noise	AbstractLinearOperator	Observation error covariance \(R\), shape (M, M).	required
key	PRNGKeyArray | None	PRNG key for internally drawn perturbations. Mutually exclusive with perturbed_obs.	None
perturbed_obs	Float[Array, 'J M'] | None	Pre-built perturbed observation ensemble, shape (J, M). Mutually exclusive with key.	None
localization	Float[Array, 'N M'] | None	Optional state-observation taper \(\rho_{xy}\), shape (N, M), e.g. from localization_matrix. When given, the gain comes from localized_kalman_gain instead of ensemble_kalman_gain.	None
obs_localization	Float[Array, 'M M'] | None	Optional observation-observation taper \(\rho_{yy}\), shape (M, M). Only consulted when localization is given; defaults to all-ones, i.e. no tapering of the innovation covariance.	None
solver	AbstractSolverStrategy | None	Optional solver strategy for the innovation solve. None uses structural dispatch. A matrix-free strategy (e.g. CGSolver) wants dense_innovation=False so it is handed the structured operator instead of a materialised one.	None
dense_innovation	bool | None	Whether to form the (M, M) innovation covariance densely. None (default) chooses by shape, as described in the note below. False keeps the structured LowRankUpdate no matter the shapes -- what a matrix-free solver wants. True forces the dense assembly, which is the way out when obs_noise is only positive semi-definite, since the structured route solves against obs_noise itself.	None
bessel	bool	Use the \(1/(J-1)\) divisor. Defaults to True, matching ensemble_kalman_gain.	True

Returns:

Type	Description
Float[Array, 'J N']	Analysis ensemble, shape (J, N).

Note

How the innovation covariance \(C^{HH} + R\) is assembled defaults to a choice made from the static shapes, because the two regimes have wildly different costs. With \(J < M\) the gain comes from ensemble_kalman_gain, which keeps the ensemble term low-rank and inverts a \((J, J)\) Woodbury capacitance -- the right choice for the geoscience regime of a few dozen members against many observations. With \(J \ge M\) that capacitance is the larger of the two (320 GB at \(J = 200{,}000\)), so the \((M, M)\) innovation is formed densely instead. Both routes solve the same system and agree to round-off.

Shapes are the wrong criterion in two cases, which is why dense_innovation exists to override it:

• A matrix-free solver. With \(J \ge M\) and \(M\) still large, the dense assembly allocates an \((M, M)\) array before the solver is ever called -- around 40 GB at \(M = 100{,}000\) in float32 -- even though an iterative strategy could work through matvecs on the structured operator. Pass dense_innovation=False.
• Singular observation noise. The Woodbury route solves against \(R\) itself, so a positive semi-definite \(R\) divides by zero and returns infinities or NaN even when \(C^{HH} + R\) is perfectly invertible -- e.g. \(R = \mathrm{diag}(1, 1, 0)\) with ensemble anomalies spanning the third observation direction. The \(J < M\) path therefore requires \(R\) to be positive definite; with a singular \(R\), pass dense_innovation=True to solve the full innovation instead. This is not checked: PSD-ness of an arbitrary operator is not something this function can establish cheaply, and certainly not under jit.

Raises:

Type	Description
ValueError	If neither or both of key / perturbed_obs are given, if the ensemble sizes disagree, or if the observation-space shapes disagree.

Example

import jax.numpy as jnp import jax.random as jr import lineax as lx from gaussx import enkf_analysis key, subkey = jr.split(jr.key(0)) prior = jr.normal(subkey, (500, 3)) # (J, N) H = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) obs_prior = prior @ H.T # (J, M) R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2)) posterior = enkf_analysis( ... prior, obs_prior, jnp.array([1.0, -1.0]), R, key=key ... ) posterior.shape (500, 3)

Source code in src/gaussx/_inference/_enkf.py

def enkf_analysis(
    particles: Float[Array, "J N"],
    obs_particles: Float[Array, "J M"],
    observation: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    key: PRNGKeyArray | None = None,
    perturbed_obs: Float[Array, "J M"] | None = None,
    localization: Float[Array, "N M"] | None = None,
    obs_localization: Float[Array, "M M"] | None = None,
    solver: AbstractSolverStrategy | None = None,
    dense_innovation: bool | None = None,
    bessel: bool = True,
) -> Float[Array, "J N"]:
    r"""Stochastic (perturbed-observation) ensemble Kalman analysis step.

    Updates a prior ensemble $X^f$ toward an observation $y$:

    $$
    X^a_j = X^f_j + K\,(y + \varepsilon_j - \mathcal{H}(X^f_j)),
    \qquad \varepsilon_j \sim N(0, R),
    $$

    with $K$ from `ensemble_kalman_gain` (or `localized_kalman_gain` when
    ``localization`` is given). The observation operator enters only through
    ``obs_particles`` -- the image $\mathcal{H}(X^f)$ of the prior ensemble in
    observation space -- so nonlinear operators need no special handling.

    The perturbation $\varepsilon_j$ is what keeps the analysis spread correct.
    The deterministic update $X^a_j = X^f_j + K(y - \mathcal{H}(X^f_j))$ drives
    the ensemble covariance to $(I - KH)P(I - KH)^\top$ instead of $(I - KH)P$,
    i.e. under-dispersive. There is deliberately no ``perturb=False`` flag: the
    deterministic alternative is a different filter (the square-root / ETKF
    family, see `etkf_transform`), not an option on this one.

    Two ways to supply the observation perturbations:

    - ``key`` -- draw $\varepsilon_j \sim N(0, R)$ internally, via an exact
      PSD factor $L L^\top = R$ of ``obs_noise`` that keeps its structure (a
      symmetric square root at dense leaves, so a singular $R$ is fine).
    - ``perturbed_obs`` -- pass a pre-built perturbed-observation ensemble
      $y + \varepsilon_j$. Preferred when the same noise realisation must be
      reused across filters, and when the perturbations come from a nonlinear
      observation model rather than an additive $R$.

    Exactly one of ``key`` / ``perturbed_obs`` must be given.

    Known limitation. The update is a Gaussian one, applied in whatever
    coordinates the caller supplies. For a non-Gaussian prior it is biased, and
    the bias does **not** shrink with ensemble size -- it is an error of
    coordinates, not of sampling. On the lognormal / logit-normal prior of
    Chipilski (2025), whose exact posterior mean is ``[0.548062, 0.353937]``,
    the physical-space update plateaus several percent off that value and stays
    there as $J$ grows by two orders of magnitude.

    The fix is to conjugate the update with a bijection $\Gamma$ that
    Gaussianises the prior -- call this function on $\Gamma^{-1}(X^f)$ and map
    the result back through $\Gamma$: the ensemble Kalman filter's Gaussian
    assumption is a statement about coordinates, not about the algorithm. Pass
    the same ``perturbed_obs`` through both routes to compare them on one noise
    realisation.

    That conjugated update is *exact Bayes* only under conditions worth stating
    precisely, because it is easy to over-claim. It needs the population limit
    -- with a finite ensemble the gain is empirical and the perturbations are
    Monte Carlo, so the result is an estimate either way -- and it needs the
    observation model to be **affine with additive Gaussian noise** in the same
    latent coordinates that Gaussianise the prior. A merely "Gaussian
    likelihood" is not enough: $y = \zeta^2 + \varepsilon$ has Gaussian noise
    and a non-Gaussian posterior that no Kalman update reproduces. Outside
    those conditions conjugation is an approximation with no guaranteed
    ordering against the physical-space update -- usually much better, but a
    badly matched $\Gamma$ can make the latent joint less Gaussian and do
    worse.

    Args:
        particles: Prior ensemble in state space, shape ``(J, N)``.
        obs_particles: Prior ensemble in observation space, shape ``(J, M)``.
        observation: The observation, shape ``(M,)``.
        obs_noise: Observation error covariance $R$, shape ``(M, M)``.
        key: PRNG key for internally drawn perturbations. Mutually exclusive
            with ``perturbed_obs``.
        perturbed_obs: Pre-built perturbed observation ensemble, shape
            ``(J, M)``. Mutually exclusive with ``key``.
        localization: Optional state-observation taper $\rho_{xy}$, shape
            ``(N, M)``, e.g. from `localization_matrix`. When given, the gain
            comes from `localized_kalman_gain` instead of
            `ensemble_kalman_gain`.
        obs_localization: Optional observation-observation taper $\rho_{yy}$,
            shape ``(M, M)``. Only consulted when ``localization`` is given;
            defaults to all-ones, i.e. no tapering of the innovation
            covariance.
        solver: Optional solver strategy for the innovation solve. ``None``
            uses structural dispatch. A matrix-free strategy (e.g. `CGSolver`)
            wants ``dense_innovation=False`` so it is handed the structured
            operator instead of a materialised one.
        dense_innovation: Whether to form the ``(M, M)`` innovation covariance
            densely. ``None`` (default) chooses by shape, as described in the
            note below. ``False`` keeps the structured `LowRankUpdate` no
            matter the shapes -- what a matrix-free solver wants. ``True``
            forces the dense assembly, which is the way out when ``obs_noise``
            is only positive *semi*-definite, since the structured route
            solves against ``obs_noise`` itself.
        bessel: Use the $1/(J-1)$ divisor. Defaults to ``True``, matching
            `ensemble_kalman_gain`.

    Returns:
        Analysis ensemble, shape ``(J, N)``.

    Note:
        How the innovation covariance $C^{HH} + R$ is assembled defaults to a
        choice made from the static shapes, because the two regimes have wildly
        different costs. With $J < M$ the gain comes from
        `ensemble_kalman_gain`, which keeps the ensemble term low-rank and
        inverts a $(J, J)$ Woodbury capacitance -- the right choice for the
        geoscience regime of a few dozen members against many observations.
        With $J \ge M$ that capacitance is the larger of the two (320 GB at
        $J = 200{,}000$), so the $(M, M)$ innovation is formed densely instead.
        Both routes solve the same system and agree to round-off.

        Shapes are the wrong criterion in two cases, which is why
        ``dense_innovation`` exists to override it:

        - **A matrix-free solver.** With $J \ge M$ and $M$ still large, the
          dense assembly allocates an $(M, M)$ array before the solver is ever
          called -- around 40 GB at $M = 100{,}000$ in float32 -- even though
          an iterative strategy could work through matvecs on the structured
          operator. Pass ``dense_innovation=False``.
        - **Singular observation noise.** The Woodbury route solves against
          $R$ itself, so a positive *semi*-definite $R$ divides by zero and
          returns infinities or ``NaN`` even when $C^{HH} + R$ is perfectly
          invertible -- e.g. $R = \mathrm{diag}(1, 1, 0)$ with ensemble
          anomalies spanning the third observation direction. The $J < M$ path
          therefore requires $R$ to be positive **definite**; with a singular
          $R$, pass ``dense_innovation=True`` to solve the full innovation
          instead. This is not checked: PSD-ness of an arbitrary operator is
          not something this function can establish cheaply, and certainly not
          under ``jit``.

    Raises:
        ValueError: If neither or both of ``key`` / ``perturbed_obs`` are
            given, if the ensemble sizes disagree, or if the observation-space
            shapes disagree.

    Example:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from gaussx import enkf_analysis
        >>> key, subkey = jr.split(jr.key(0))
        >>> prior = jr.normal(subkey, (500, 3))           # (J, N)
        >>> H = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
        >>> obs_prior = prior @ H.T                       # (J, M)
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> posterior = enkf_analysis(
        ...     prior, obs_prior, jnp.array([1.0, -1.0]), R, key=key
        ... )
        >>> posterior.shape
        (500, 3)
    """
    if (key is None) == (perturbed_obs is None):
        raise ValueError(
            "Pass exactly one of 'key' (draw perturbations from obs_noise) or "
            "'perturbed_obs' (supply them directly)."
        )

    n_ens, n_state, n_obs = _check_analysis_shapes(
        particles, obs_particles, observation, obs_noise, bessel
    )

    # Whether to form the (M, M) innovation densely. `None` picks by shape; an
    # explicit value overrides that, which is what a matrix-free solver needs.
    # Decided here rather than at the gain, because the factor below needs it.
    use_dense = n_ens >= n_obs if dense_innovation is None else dense_innovation

    # Branch on `key` rather than on `perturbed_obs`: the check above makes the
    # two equivalent, but this way each branch narrows the argument it uses.
    if key is not None:
        # eps_j = L n_j with R = L L^T. The factor dispatches on structure, so
        # a DiagonalLinearOperator stays diagonal and its matvec stays O(M) --
        # materialising R here would allocate a dense (M, M) and cost O(M^3),
        # which for the low-rank branch below (J < M, M possibly enormous) would
        # OOM before the gain is ever formed.
        #
        # Unless the gain is about to allocate that (M, M) anyway: every route
        # but Woodbury calls `obs_noise.as_matrix()`, and once the dense cost
        # is being paid regardless there is nothing left to protect by holding
        # on to a positive-definite-only factor.
        factor = _noise_factor(
            obs_noise, allow_dense=use_dense or localization is not None
        )
        noise = jr.normal(key, (n_ens, n_obs), dtype=particles.dtype)  # (J, M)
        perturbed = observation[None, :] + jax.vmap(factor.mv)(noise)  # (J, M)
    else:
        perturbed = perturbed_obs
        if perturbed is None or perturbed.shape != (n_ens, n_obs):
            raise ValueError(
                f"perturbed_obs must have shape ({n_ens}, {n_obs}) to match "
                f"obs_particles, got "
                f"{None if perturbed is None else perturbed.shape}."
            )

    _check_localization_shapes(n_state, n_obs, localization, obs_localization)

    gain = _analysis_gain(
        particles,
        obs_particles,
        obs_noise,
        localization=localization,
        obs_localization=obs_localization,
        solver=solver,
        use_dense=use_dense,
        bessel=bessel,
    )  # (N, M)

    innovation = perturbed - obs_particles  # (J, M)
    return particles + innovation @ gain.T  # (J, M) @ (M, N) -> (J, N)

etkf_transform(obs_particles: Float[Array, 'J M'], y: Float[Array, ' M'], obs_noise: lx.AbstractLinearOperator, *, inflation: float | Float[Array, ''] = 1.0, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' J'], Float[Array, 'J J']]

Ensemble Transform Kalman Filter (ETKF) analysis weights.

Deterministic (perturbed-obs-free) ensemble square-root analysis in the J-dimensional ensemble space (Bishop et al. 2001; Hunt et al. 2007). With raw observation perturbations Y = H X'^f (columns are members) and d = y - H x_bar^f,

\[ \tilde{A}^{-1} = \tfrac{J-1}{\lambda} I + Y^T R^{-1} Y, \qquad \bar{w} = \tilde{A}\, Y^T R^{-1} d, \qquad W = \big((J-1)\,\tilde{A}\big)^{1/2}, \]

where lambda is the (multiplicative) inflation and W is the symmetric square root. The analysis ensemble is reconstructed as

\[ \bar{x}^a = \bar{x}^f + X'^f \bar{w}, \qquad X'^a = X'^f\, W. \]

The symmetric (eigendecomposition) square root -- not a Cholesky factor -- is required: because the observation perturbations are zero-mean, 1 is an eigenvector of W (with eigenvalue sqrt(lambda)), which makes the transform exactly mean-preserving (sum_j X'^a_j = 0).

Cost. R^{-1} is applied to the J + 1 right-hand sides [Y^T, d] in one solve_rows call, so a structured R (diagonal, gaussx.BlockDiag, gaussx.Kronecker, ...) keeps its own solve and is never materialised, and a dense R is factored once. The ensemble-space algebra then takes a single (J, J) eigendecomposition \tilde{A}^{-1} = V \operatorname{diag}(s) V^T, from which both \tilde{A} = V \operatorname{diag}(1/s) V^T and W = V \operatorname{diag}(\sqrt{(J-1)/s}) V^T are read off (Hunt et al. 2007), with no explicit inverse. Its derivative is taken in closed form in the eigenbasis (Daleckii-Krein), which stays finite at the repeated prior eigenvalue (J - 1)/\lambda that has multiplicity J - M whenever M < J.

Parameters:

Name	Type	Description	Default
obs_particles	Float[Array, 'J M']	Forecast ensemble in observation space, shape (J, M).	required
y	Float[Array, ' M']	Observation vector, shape (M,).	required
obs_noise	AbstractLinearOperator	Observation error covariance operator R, shape (M, M).	required
inflation	float | Float[Array, '']	Multiplicative covariance inflation lambda >= 1, applied to the prior term (J - 1) / lambda.	1.0
solver	AbstractSolverStrategy | None	Optional solver strategy for the R^{-1} solves. None uses structural dispatch.	None

Returns:

Type	Description
Float[Array, ' J']	(w_mean, transform) where w_mean has shape (J,) and
Float[Array, 'J J']	transform has shape (J, J). Apply to forecast state
tuple[Float[Array, ' J'], Float[Array, 'J J']]	perturbations Xp (shape (J, N)) as
tuple[Float[Array, ' J'], Float[Array, 'J J']]	x_bar^a = x_bar^f + w_mean @ Xp and X'^a = transform @ Xp.

Raises:

Type	Description
ValueError	If J < 2 (the prior term (J - 1) / lambda is then zero and the ensemble-space precision singular), if y or obs_noise do not match M, or if a concrete inflation is not positive.

Source code in src/gaussx/_inference/_etkf.py

def etkf_transform(
    obs_particles: Float[Array, "J M"],
    y: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    inflation: float | Float[Array, ""] = 1.0,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " J"], Float[Array, "J J"]]:
    r"""Ensemble Transform Kalman Filter (ETKF) analysis weights.

    Deterministic (perturbed-obs-free) ensemble square-root analysis in the
    ``J``-dimensional ensemble space (Bishop et al. 2001; Hunt et al. 2007).
    With raw observation perturbations ``Y = H X'^f`` (columns are members) and
    ``d = y - H x_bar^f``,

    $$
    \tilde{A}^{-1} = \tfrac{J-1}{\lambda} I + Y^T R^{-1} Y, \qquad
    \bar{w} = \tilde{A}\, Y^T R^{-1} d, \qquad
    W = \big((J-1)\,\tilde{A}\big)^{1/2},
    $$

    where ``lambda`` is the (multiplicative) ``inflation`` and ``W`` is the
    **symmetric** square root. The analysis ensemble is reconstructed as

    $$
    \bar{x}^a = \bar{x}^f + X'^f \bar{w}, \qquad X'^a = X'^f\, W.
    $$

    The symmetric (eigendecomposition) square root -- not a Cholesky factor --
    is required: because the observation perturbations are zero-mean, ``1`` is
    an eigenvector of ``W`` (with eigenvalue ``sqrt(lambda)``), which makes the
    transform exactly mean-preserving (``sum_j X'^a_j = 0``).

    Cost. ``R^{-1}`` is applied to the ``J + 1`` right-hand sides
    ``[Y^T, d]`` in one `solve_rows` call, so a structured ``R`` (diagonal,
    `gaussx.BlockDiag`, `gaussx.Kronecker`, ...) keeps its own solve and is
    never materialised, and a dense ``R`` is factored once. The ensemble-space
    algebra then takes a single ``(J, J)`` eigendecomposition
    ``\tilde{A}^{-1} = V \operatorname{diag}(s) V^T``, from which both
    ``\tilde{A} = V \operatorname{diag}(1/s) V^T`` and
    ``W = V \operatorname{diag}(\sqrt{(J-1)/s}) V^T`` are read off (Hunt et
    al. 2007), with no explicit inverse. Its derivative is taken in closed
    form in the eigenbasis (Daleckii-Krein), which stays finite at the
    repeated prior eigenvalue ``(J - 1)/\lambda`` that has multiplicity
    ``J - M`` whenever ``M < J``.

    Args:
        obs_particles: Forecast ensemble in observation space, shape ``(J, M)``.
        y: Observation vector, shape ``(M,)``.
        obs_noise: Observation error covariance operator ``R``, shape ``(M, M)``.
        inflation: Multiplicative covariance inflation ``lambda >= 1``, applied
            to the prior term ``(J - 1) / lambda``.
        solver: Optional solver strategy for the ``R^{-1}`` solves. ``None``
            uses structural dispatch.

    Returns:
        ``(w_mean, transform)`` where ``w_mean`` has shape ``(J,)`` and
        ``transform`` has shape ``(J, J)``. Apply to forecast state
        perturbations ``Xp`` (shape ``(J, N)``) as
        ``x_bar^a = x_bar^f + w_mean @ Xp`` and ``X'^a = transform @ Xp``.

    Raises:
        ValueError: If ``J < 2`` (the prior term ``(J - 1) / lambda`` is then
            zero and the ensemble-space precision singular), if ``y`` or
            ``obs_noise`` do not match ``M``, or if a concrete ``inflation``
            is not positive.
    """
    _check_observation_shapes(obs_particles, y, obs_noise, observation_name="y")
    n_ens = obs_particles.shape[0]
    if n_ens < 2:
        raise ValueError(
            "etkf_transform requires J >= 2 ensemble members (the prior term "
            f"is (J - 1) / inflation); got J={n_ens}."
        )
    if isinstance(inflation, (int, float)) and inflation <= 0:
        raise ValueError(f"inflation must be positive, got {inflation}.")
    obs_mean = reduce(obs_particles, "J M -> M", "mean")
    obs_pert = einx.subtract("J M, M -> J M", obs_particles, obs_mean)  # zero-mean

    # R^{-1} applied to all J + 1 right-hand sides at once: one structured
    # solve, so a dense R is factored once and a structured R never densified.
    rhs = jnp.vstack([obs_pert, y - obs_mean])  # (J + 1, M)
    weighted = solve_rows(obs_noise, rhs, solver=solver)  # (J + 1, M)
    rinv_pert, rinv_d = weighted[:-1], weighted[-1]

    eye = jnp.eye(n_ens, dtype=rinv_pert.dtype)
    precision = (n_ens - 1) / inflation * eye + einsum(
        obs_pert, rinv_pert, "J M, K M -> J K"
    )
    analysis_cov, inv_sqrt = _inverse_and_inverse_sqrt(symmetrize(precision))

    w_mean = einsum(
        analysis_cov, einsum(obs_pert, rinv_d, "J M, M -> J"), "J K, K -> J"
    )
    transform = jnp.sqrt(jnp.asarray(n_ens - 1, dtype=inv_sqrt.dtype)) * inv_sqrt
    return w_mean, transform

Ensemble Kalman inversion

eki_step is the same Kalman update as enkf_analysis with two knobs added, and reduces to it exactly at dt=1. It is the inverse problem reading of the ensemble filter: one fixed observation, no time axis, and a schedule of tempered steps instead of a sequence of assimilation windows.

dt is the observation-side tempering step, replacing R by R / dt. Over a schedule with sum(dt) = 1 the composition is exactly one Bayesian update in the linear-Gaussian population limit -- the precisions add -- so the sum condition is what makes a schedule a tempering path rather than a heuristic. step is the state-side operator in the gradient-flow view: it multiplies each member's increment, so a BlockDiag of scaled identities gives a different rate per state block. It changes the trajectory, not the fixed point.

The two helpers cover the standard variations. tikhonov_augment puts a prior N(m0, C0) into the step by observation augmentation (TEKI) -- a helper rather than a flag, so C0 stays an operator and the step itself knows nothing about priors. discrepancy_step_size is the tuning-parameter-free data misfit controller of Iglesias & Yang (2021): a pure function of the ensemble misfits, so it belongs here rather than in whatever drives the iteration.

The iteration loop, the stopping rule, and the forward model itself are all out of scope: these are array-in / array-out steps.

Structured linear algebra and Gaussian primitives for JAX.

eki_step(particles: Float[Array, 'J N'], obs_particles: Float[Array, 'J M'], observation: Float[Array, ' M'], obs_noise: lx.AbstractLinearOperator, *, dt: float | Float[Array, ''] = 1.0, step: lx.AbstractLinearOperator | None = None, key: PRNGKeyArray | None = None, perturbed_obs: Float[Array, 'J M'] | None = None, deterministic: bool = False, localization: Float[Array, 'N M'] | None = None, obs_localization: Float[Array, 'M M'] | None = None, solver: AbstractSolverStrategy | None = None, dense_innovation: bool | None = None, bessel: bool = True) -> Float[Array, 'J N']

One ensemble Kalman inversion (EKI) update.

A single tempered Kalman update of an ensemble against a fixed observation (Iglesias, Law & Stuart 2013). The gain is \(K = C^{uG}(C^{GG} + R/\Delta t)^{-1}\), and the stochastic (perturbed-observation) update is

\[ u_j \leftarrow u_j + \Lambda\,K\, \big(y + \varepsilon_j/\sqrt{\Delta t} - \mathcal{G}(u_j)\big), \qquad \varepsilon_j \sim N(0, R). \]

This is enkf_analysis with two knobs added, and reduces to it exactly at dt=1, step=None. Everything the forward model does enters through obs_particles, so this function is pure array-in / array-out: there is no iteration, no stopping rule, and no \(\mathcal{G}\). The driver that supplies the schedule lives outside this package.

Tempering (dt). One iteration replaces \(R\) by \(R/\Delta t\), i.e. a likelihood raised to the power \(\Delta t\). Over a schedule with \(\sum_n \Delta t_n = 1\) the composition is exactly one Bayesian update in the linear-Gaussian population limit -- the precisions add, \(C_N^{-1} = C_0^{-1} + \sum_n \Delta t_n\, A^\top R^{-1} A\) -- which is what makes an EKI schedule a tempering path rather than a heuristic. The sum condition is load-bearing: at \(\sum \Delta t_n \neq 1\) the result is the posterior of a different problem, over- or under-weighting the data. dt may be a traced scalar, so an adaptive schedule (see discrepancy_step_size) stays inside jit.

State-side step (step). In the gradient-flow view \(\dot{u} = -C^{uu}\nabla\Phi(u)\), step is the operator \(\Lambda\) in the Euler step \(u \leftarrow u + \Lambda C^{uG} S^{-1}(y - \mathcal{G}(u))\). It is applied by step.mv to each member's increment, so a gaussx.BlockDiag of scaled identities gives a different rate per state block -- parameters, latents, initial conditions -- without densifying an \((N, N)\) matrix. \(\Lambda\) changes the trajectory, not the fixed point: where \(K(y - \bar{\mathcal{G}}) = 0\) the increment is zero for every \(\Lambda\), invertible or not.

Deterministic variant. deterministic=True replaces the perturbed observations with an ETKF square-root transform (etkf_transform at \(R/\Delta t\)), applied to the increment so that \(\Lambda\) still acts on a difference:

\[ \bar{u} \leftarrow \bar{u} + \Lambda K (y - \bar{\mathcal{G}}), \qquad U'^a = U'^f + \Lambda\,(W U'^f - U'^f). \]

At \(\Lambda = I\) the anomaly update collapses to \(U'^a = W U'^f\), i.e. plain etkf_transform. This is the variant to use for the exactness property above: the stochastic one is exact only in expectation, so a finite ensemble carries Monte Carlo error on top of the tempering.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Prior ensemble in state space, shape (J, N).	required
obs_particles	Float[Array, 'J M']	Its image \(\mathcal{G}(u_j)\) in observation space, shape (J, M).	required
observation	Float[Array, ' M']	The observation \(y\), shape (M,).	required
obs_noise	AbstractLinearOperator	Observation error covariance \(R\), shape (M, M). Scaled to \(R/\Delta t\) as a lazy gaussx.ScaledOperator, so a structured \(R\) keeps its own solve in both variants.	required
dt	float | Float[Array, '']	Observation-side tempering step \(\Delta t > 0\). Positivity is not checked -- it may be traced.	1.0
step	AbstractLinearOperator | None	State-side operator \(\Lambda\), shape (N, N). None is the identity, and skips the matvec rather than building one.	None
key	PRNGKeyArray | None	PRNG key for internally drawn perturbations \(\varepsilon_j \sim N(0, R)\), which are then scaled by \(1/\sqrt{\Delta t}\) to match \(R/\Delta t\). Mutually exclusive with perturbed_obs; both must be None when deterministic.	None
perturbed_obs	Float[Array, 'J M'] | None	Pre-built perturbed observation ensemble, shape (J, M), used as given -- the caller owns the \(1/\sqrt{\Delta t}\) scaling. Mutually exclusive with key.	None
deterministic	bool	Use the ETKF square-root transform instead of perturbed observations.	False
localization	Float[Array, 'N M'] | None	Optional state-observation taper \(\rho_{xy}\), shape (N, M). Stochastic variant only.	None
obs_localization	Float[Array, 'M M'] | None	Optional observation-observation taper \(\rho_{yy}\), shape (M, M). Only consulted when localization is given.	None
solver	AbstractSolverStrategy | None	Optional solver strategy for the innovation solve. None uses structural dispatch.	None
dense_innovation	bool | None	Whether to form the (M, M) innovation densely. None chooses by shape. Same contract as enkf_analysis, including that a positive semi-definite \(R\) needs True. Stochastic variant only: the deterministic one forms no gain -- etkf_transform supplies the mean update as well -- and solves against \(R\) itself.	None
bessel	bool	Use the \(1/(J-1)\) divisor. Must stay True when deterministic, since etkf_transform is \(1/(J-1)\) throughout and a mismatched gain would move the mean and the anomalies by inconsistent amounts.	True

Returns:

Type	Description
Float[Array, 'J N']	The updated ensemble, shape (J, N).

Raises:

Type	Description
ValueError	If the shapes disagree; if key / perturbed_obs are not given exactly once in the stochastic variant, or given at all in the deterministic one; or if deterministic is combined with localization or with bessel=False.

Note

deterministic=True rejects localization rather than ignoring it. Schur-product localization has no square-root analogue: tapering the gain but not the transform would localize the mean update and leave the anomalies unlocalized, an inconsistent analysis that looks like a working one. (The LETKF localizes by domain decomposition instead, which is a different construction, not this argument.)

Example

import jax.numpy as jnp import jax.random as jr import lineax as lx from gaussx import eki_step key, subkey = jr.split(jr.key(0)) u = jr.normal(subkey, (200, 3)) # (J, N) A = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) G = u @ A.T # (J, M) R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2)) eki_step(u, G, jnp.array([1.0, -1.0]), R, dt=0.5, key=key).shape (200, 3)

Source code in src/gaussx/_inference/_eki.py

def eki_step(
    particles: Float[Array, "J N"],
    obs_particles: Float[Array, "J M"],
    observation: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    dt: float | Float[Array, ""] = 1.0,
    step: lx.AbstractLinearOperator | None = None,
    key: PRNGKeyArray | None = None,
    perturbed_obs: Float[Array, "J M"] | None = None,
    deterministic: bool = False,
    localization: Float[Array, "N M"] | None = None,
    obs_localization: Float[Array, "M M"] | None = None,
    solver: AbstractSolverStrategy | None = None,
    dense_innovation: bool | None = None,
    bessel: bool = True,
) -> Float[Array, "J N"]:
    r"""One ensemble Kalman inversion (EKI) update.

    A single tempered Kalman update of an ensemble against a fixed
    observation (Iglesias, Law & Stuart 2013). The gain is
    $K = C^{uG}(C^{GG} + R/\Delta t)^{-1}$, and the stochastic
    (perturbed-observation) update is

    $$
    u_j \leftarrow u_j + \Lambda\,K\,
        \big(y + \varepsilon_j/\sqrt{\Delta t} - \mathcal{G}(u_j)\big),
    \qquad \varepsilon_j \sim N(0, R).
    $$

    This is `enkf_analysis` with two knobs added, and reduces to it exactly at
    ``dt=1``, ``step=None``. Everything the forward model does enters through
    ``obs_particles``, so this function is pure array-in / array-out: there is
    no iteration, no stopping rule, and no $\mathcal{G}$. The driver that
    supplies the schedule lives outside this package.

    **Tempering (``dt``).** One iteration replaces $R$ by $R/\Delta t$, i.e.
    a likelihood raised to the power $\Delta t$. Over a schedule with
    $\sum_n \Delta t_n = 1$ the composition is *exactly* one Bayesian update
    in the linear-Gaussian population limit -- the precisions add,
    $C_N^{-1} = C_0^{-1} + \sum_n \Delta t_n\, A^\top R^{-1} A$ -- which is
    what makes an EKI schedule a tempering path rather than a heuristic. The
    sum condition is load-bearing: at $\sum \Delta t_n \neq 1$ the result is
    the posterior of a different problem, over- or under-weighting the data.
    ``dt`` may be a traced scalar, so an adaptive schedule (see
    `discrepancy_step_size`) stays inside ``jit``.

    **State-side step (``step``).** In the gradient-flow view
    $\dot{u} = -C^{uu}\nabla\Phi(u)$, ``step`` is the operator $\Lambda$ in
    the Euler step $u \leftarrow u + \Lambda C^{uG} S^{-1}(y - \mathcal{G}(u))$.
    It is applied by ``step.mv`` to each member's *increment*, so a
    `gaussx.BlockDiag` of scaled identities gives a different rate per state
    block -- parameters, latents, initial conditions -- without densifying an
    $(N, N)$ matrix. $\Lambda$ changes the trajectory, not the fixed point:
    where $K(y - \bar{\mathcal{G}}) = 0$ the increment is zero for every
    $\Lambda$, invertible or not.

    **Deterministic variant.** ``deterministic=True`` replaces the perturbed
    observations with an ETKF square-root transform (`etkf_transform` at
    $R/\Delta t$), applied to the increment so that $\Lambda$ still acts on a
    difference:

    $$
    \bar{u} \leftarrow \bar{u} + \Lambda K (y - \bar{\mathcal{G}}),
    \qquad
    U'^a = U'^f + \Lambda\,(W U'^f - U'^f).
    $$

    At $\Lambda = I$ the anomaly update collapses to $U'^a = W U'^f$, i.e.
    plain `etkf_transform`. This is the variant to use for the exactness
    property above: the stochastic one is exact only in expectation, so a
    finite ensemble carries Monte Carlo error on top of the tempering.

    Args:
        particles: Prior ensemble in state space, shape ``(J, N)``.
        obs_particles: Its image $\mathcal{G}(u_j)$ in observation space,
            shape ``(J, M)``.
        observation: The observation $y$, shape ``(M,)``.
        obs_noise: Observation error covariance $R$, shape ``(M, M)``. Scaled
            to $R/\Delta t$ as a lazy `gaussx.ScaledOperator`, so a
            structured $R$ keeps its own solve in both variants.
        dt: Observation-side tempering step $\Delta t > 0$. Positivity is not
            checked -- it may be traced.
        step: State-side operator $\Lambda$, shape ``(N, N)``. ``None`` is the
            identity, and skips the matvec rather than building one.
        key: PRNG key for internally drawn perturbations
            $\varepsilon_j \sim N(0, R)$, which are then scaled by
            $1/\sqrt{\Delta t}$ to match $R/\Delta t$. Mutually exclusive with
            ``perturbed_obs``; both must be ``None`` when ``deterministic``.
        perturbed_obs: Pre-built perturbed observation ensemble, shape
            ``(J, M)``, used **as given** -- the caller owns the
            $1/\sqrt{\Delta t}$ scaling. Mutually exclusive with ``key``.
        deterministic: Use the ETKF square-root transform instead of perturbed
            observations.
        localization: Optional state-observation taper $\rho_{xy}$, shape
            ``(N, M)``. Stochastic variant only.
        obs_localization: Optional observation-observation taper $\rho_{yy}$,
            shape ``(M, M)``. Only consulted when ``localization`` is given.
        solver: Optional solver strategy for the innovation solve. ``None``
            uses structural dispatch.
        dense_innovation: Whether to form the ``(M, M)`` innovation densely.
            ``None`` chooses by shape. Same contract as `enkf_analysis`,
            including that a positive *semi*-definite $R$ needs ``True``.
            Stochastic variant only: the deterministic one forms no gain --
            `etkf_transform` supplies the mean update as well -- and solves
            against $R$ itself.
        bessel: Use the $1/(J-1)$ divisor. Must stay ``True`` when
            ``deterministic``, since `etkf_transform` is $1/(J-1)$ throughout
            and a mismatched gain would move the mean and the anomalies by
            inconsistent amounts.

    Returns:
        The updated ensemble, shape ``(J, N)``.

    Raises:
        ValueError: If the shapes disagree; if ``key`` / ``perturbed_obs`` are
            not given exactly once in the stochastic variant, or given at all
            in the deterministic one; or if ``deterministic`` is combined with
            ``localization`` or with ``bessel=False``.

    Note:
        ``deterministic=True`` rejects ``localization`` rather than ignoring
        it. Schur-product localization has no square-root analogue: tapering
        the gain but not the transform would localize the mean update and
        leave the anomalies unlocalized, an inconsistent analysis that looks
        like a working one. (The LETKF localizes by domain decomposition
        instead, which is a different construction, not this argument.)

    Example:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from gaussx import eki_step
        >>> key, subkey = jr.split(jr.key(0))
        >>> u = jr.normal(subkey, (200, 3))                 # (J, N)
        >>> A = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
        >>> G = u @ A.T                                     # (J, M)
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> eki_step(u, G, jnp.array([1.0, -1.0]), R, dt=0.5, key=key).shape
        (200, 3)
    """
    n_ens, n_state, n_obs = _check_analysis_shapes(
        particles, obs_particles, observation, obs_noise, bessel
    )
    if deterministic:
        if key is not None or perturbed_obs is not None:
            raise ValueError(
                "The deterministic variant draws no perturbations; pass "
                "neither 'key' nor 'perturbed_obs' with deterministic=True."
            )
        if localization is not None:
            raise ValueError(
                "deterministic=True does not support 'localization': tapering "
                "the gain without tapering the ETKF transform would localize "
                "the mean update and not the anomalies. Use the stochastic "
                "variant, or localize by domain decomposition (LETKF)."
            )
        if not bessel:
            raise ValueError(
                "deterministic=True requires bessel=True, because "
                "etkf_transform uses the 1 / (J - 1) divisor throughout and a "
                "1 / J gain would update the mean and the anomalies by "
                "inconsistent amounts."
            )
    elif (key is None) == (perturbed_obs is None):
        raise ValueError(
            "Pass exactly one of 'key' (draw perturbations from obs_noise) or "
            "'perturbed_obs' (supply them directly)."
        )
    if step is not None and (step.in_size(), step.out_size()) != (n_state, n_state):
        raise ValueError(
            f"step must be ({n_state}, {n_state}) to match particles, got "
            f"({step.out_size()}, {step.in_size()})."
        )
    _check_localization_shapes(n_state, n_obs, localization, obs_localization)

    dt = jnp.asarray(dt)
    if dt.ndim != 0:
        raise ValueError(f"dt must be a scalar, got shape {dt.shape}.")
    # R / dt as a lazy scale. `solve` unwraps a MulLinearOperator by dividing
    # the inner solve, so a Diagonal / BlockDiag / Kronecker R keeps its own
    # solve rather than falling back to a dense (M, M). At dt = 1 the scalar is
    # exactly 1.0 and every route is bit-for-bit `enkf_analysis`.
    tempered_noise = ScaledOperator(obs_noise, 1.0 / dt)

    if deterministic:
        # By the push-through identity K (y - G_bar) = U'^T w_mean, so the
        # ETKF weights give the mean update too, and no separate innovation
        # solve for the gain is needed.
        anomalies = particles - jnp.mean(particles, axis=0, keepdims=True)  # (J, N)
        w_mean, transform = etkf_transform(
            obs_particles, observation, tempered_noise, solver=solver
        )
        mean_increment = einsum(w_mean, anomalies, "J, J N -> N")  # (N,)
        # The increment, not the transformed anomalies: Lambda acts on
        # differences, so Lambda = I leaves `transform @ anomalies` exactly.
        anomaly_increment = transform @ anomalies - anomalies  # (J, N)
        if step is not None:
            mean_increment = step.mv(mean_increment)
            anomaly_increment = jax.vmap(step.mv)(anomaly_increment)
        return particles + mean_increment[None, :] + anomaly_increment

    use_dense = n_ens >= n_obs if dense_innovation is None else dense_innovation
    gain = _analysis_gain(
        particles,
        obs_particles,
        tempered_noise,
        localization=localization,
        obs_localization=obs_localization,
        solver=solver,
        use_dense=use_dense,
        bessel=bessel,
    )  # (N, M)

    if key is not None:
        # eps_j ~ N(0, R) drawn against the *unscaled* R, then scaled by
        # 1/sqrt(dt) to give N(0, R/dt). Drawing against R/dt instead would
        # hand `_noise_factor` a MulLinearOperator it has no structured branch
        # for, and densify an (M, M) for nothing.
        factor = _noise_factor(
            obs_noise, allow_dense=use_dense or localization is not None
        )
        noise = jr.normal(key, (n_ens, n_obs), dtype=particles.dtype)  # (J, M)
        perturbation = jax.vmap(factor.mv)(noise) / jnp.sqrt(dt)  # (J, M)
        perturbed = observation[None, :] + perturbation  # (J, M)
    else:
        perturbed = perturbed_obs
        if perturbed is None or perturbed.shape != (n_ens, n_obs):
            raise ValueError(
                f"perturbed_obs must have shape ({n_ens}, {n_obs}) to match "
                f"obs_particles, got "
                f"{None if perturbed is None else perturbed.shape}."
            )

    increment = (perturbed - obs_particles) @ gain.T  # (J, M) @ (M, N) -> (J, N)
    if step is not None:
        increment = jax.vmap(step.mv)(increment)
    return particles + increment

tikhonov_augment(particles: Float[Array, 'J N'], obs_particles: Float[Array, 'J M'], observation: Float[Array, ' M'], obs_noise: lx.AbstractLinearOperator, prior_mean: Float[Array, ' N'], prior_cov: lx.AbstractLinearOperator) -> tuple[Float[Array, 'J M+N'], Float[Array, ' M+N'], lx.AbstractLinearOperator]

Observation augmentation for Tikhonov-regularised EKI (TEKI).

Puts the prior \(N(m_0, C_0)\) into an EKI step by treating the state as its own observation (Chada, Stuart & Tong 2020):

\[ y_{\text{aug}} = \begin{bmatrix} y \\ m_0 \end{bmatrix}, \qquad \mathcal{G}_{\text{aug}}(u) = \begin{bmatrix} \mathcal{G}(u) \\ u \end{bmatrix}, \qquad R_{\text{aug}} = \operatorname{blockdiag}(R, C_0), \]

so the augmented least-squares functional is the regularised one, \(\tfrac12\|y - \mathcal{G}(u)\|_R^2 + \tfrac12\|u - m_0\|_{C_0}^2\). Unregularised EKI collapses onto the data-misfit minimiser and, for an ill-posed problem, keeps going; the prior term is what stops it.

A helper rather than a flag on eki_step: the triple goes straight into eki_step (and into discrepancy_step_size, whose \(M\) is then \(M + N\)), nothing inside the step knows about priors, and \(C_0\) stays an operator, so a gaussx.Kronecker prior keeps its structured solve inside the gaussx.BlockDiag.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Ensemble in state space, shape (J, N).	required
obs_particles	Float[Array, 'J M']	Its image \(\mathcal{G}(u_j)\), shape (J, M).	required
observation	Float[Array, ' M']	The observation \(y\), shape (M,).	required
obs_noise	AbstractLinearOperator	Observation error covariance \(R\), shape (M, M).	required
prior_mean	Float[Array, ' N']	Prior mean \(m_0\), shape (N,).	required
prior_cov	AbstractLinearOperator	Prior covariance \(C_0\), shape (N, N).	required

Returns:

Type	Description
Float[Array, 'J M+N']	(obs_particles_aug, observation_aug, obs_noise_aug) with shapes
Float[Array, ' M+N']	(J, M + N), (M + N,) and (M + N, M + N). Pass them to
AbstractLinearOperator	eki_step in place of obs_particles, observation and
tuple[Float[Array, 'J M+N'], Float[Array, ' M+N'], AbstractLinearOperator]	obs_noise; particles is unchanged.

Raises:

Type	Description
ValueError	If any of the shapes disagree.

Example

import jax.numpy as jnp import jax.random as jr import lineax as lx from gaussx import eki_step, tikhonov_augment key, subkey = jr.split(jr.key(0)) u = jr.normal(subkey, (200, 3)) A = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2)) C0 = lx.DiagonalLinearOperator(jnp.ones(3)) G_aug, y_aug, R_aug = tikhonov_augment( ... u, u @ A.T, jnp.array([1.0, -1.0]), R, jnp.zeros(3), C0 ... ) eki_step(u, G_aug, y_aug, R_aug, key=key).shape (200, 3)

Source code in src/gaussx/_inference/_eki.py

def tikhonov_augment(
    particles: Float[Array, "J N"],
    obs_particles: Float[Array, "J M"],
    observation: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    prior_mean: Float[Array, " N"],
    prior_cov: lx.AbstractLinearOperator,
) -> tuple[
    Float[Array, "J M+N"],
    Float[Array, " M+N"],
    lx.AbstractLinearOperator,
]:
    r"""Observation augmentation for Tikhonov-regularised EKI (TEKI).

    Puts the prior $N(m_0, C_0)$ into an EKI step by treating the state as its
    own observation (Chada, Stuart & Tong 2020):

    $$
    y_{\text{aug}} = \begin{bmatrix} y \\ m_0 \end{bmatrix},
    \qquad
    \mathcal{G}_{\text{aug}}(u) = \begin{bmatrix} \mathcal{G}(u) \\ u
        \end{bmatrix},
    \qquad
    R_{\text{aug}} = \operatorname{blockdiag}(R, C_0),
    $$

    so the augmented least-squares functional is the regularised one,
    $\tfrac12\|y - \mathcal{G}(u)\|_R^2 + \tfrac12\|u - m_0\|_{C_0}^2$.
    Unregularised EKI collapses onto the data-misfit minimiser and, for an
    ill-posed problem, keeps going; the prior term is what stops it.

    A helper rather than a flag on `eki_step`: the triple goes straight into
    `eki_step` (and into `discrepancy_step_size`, whose $M$ is then $M + N$),
    nothing inside the step knows about priors, and $C_0$ stays an operator, so
    a `gaussx.Kronecker` prior keeps its structured solve inside the
    `gaussx.BlockDiag`.

    Args:
        particles: Ensemble in state space, shape ``(J, N)``.
        obs_particles: Its image $\mathcal{G}(u_j)$, shape ``(J, M)``.
        observation: The observation $y$, shape ``(M,)``.
        obs_noise: Observation error covariance $R$, shape ``(M, M)``.
        prior_mean: Prior mean $m_0$, shape ``(N,)``.
        prior_cov: Prior covariance $C_0$, shape ``(N, N)``.

    Returns:
        ``(obs_particles_aug, observation_aug, obs_noise_aug)`` with shapes
        ``(J, M + N)``, ``(M + N,)`` and ``(M + N, M + N)``. Pass them to
        `eki_step` in place of ``obs_particles``, ``observation`` and
        ``obs_noise``; ``particles`` is unchanged.

    Raises:
        ValueError: If any of the shapes disagree.

    Example:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from gaussx import eki_step, tikhonov_augment
        >>> key, subkey = jr.split(jr.key(0))
        >>> u = jr.normal(subkey, (200, 3))
        >>> A = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> C0 = lx.DiagonalLinearOperator(jnp.ones(3))
        >>> G_aug, y_aug, R_aug = tikhonov_augment(
        ...     u, u @ A.T, jnp.array([1.0, -1.0]), R, jnp.zeros(3), C0
        ... )
        >>> eki_step(u, G_aug, y_aug, R_aug, key=key).shape
        (200, 3)
    """
    n_ens, n_state = particles.shape
    n_obs = obs_particles.shape[1]
    if obs_particles.shape[0] != n_ens:
        raise ValueError(
            "particles and obs_particles must share the same ensemble size, "
            f"got J={n_ens} and J={obs_particles.shape[0]}."
        )
    if observation.shape != (n_obs,):
        raise ValueError(
            f"observation must have shape ({n_obs},) to match obs_particles, "
            f"got {observation.shape}."
        )
    if (obs_noise.in_size(), obs_noise.out_size()) != (n_obs, n_obs):
        raise ValueError(
            f"obs_noise must be ({n_obs}, {n_obs}) to match obs_particles, got "
            f"({obs_noise.out_size()}, {obs_noise.in_size()})."
        )
    if prior_mean.shape != (n_state,):
        raise ValueError(
            f"prior_mean must have shape ({n_state},) to match particles, got "
            f"{prior_mean.shape}."
        )
    if (prior_cov.in_size(), prior_cov.out_size()) != (n_state, n_state):
        raise ValueError(
            f"prior_cov must be ({n_state}, {n_state}) to match particles, got "
            f"({prior_cov.out_size()}, {prior_cov.in_size()})."
        )

    obs_particles_aug = jnp.concatenate([obs_particles, particles], axis=1)
    observation_aug = jnp.concatenate([observation, prior_mean])
    return obs_particles_aug, observation_aug, BlockDiag(obs_noise, prior_cov)

discrepancy_step_size(obs_particles: Float[Array, 'J M'], observation: Float[Array, ' M'], obs_noise: lx.AbstractLinearOperator, *, remaining: Float[Array, ''], bessel: bool = True) -> Float[Array, '']

Adaptive EKI tempering step: the data misfit controller.

Iglesias & Yang (2021), eq. (14) -- the selection rule of their EKI-DMC (Algorithm 3), which needs no tuning parameter. With the per-particle least-squares functional

\[ \Phi_j = \tfrac12 \big\| R^{-1/2}(y - \mathcal{G}(u_j)) \big\|^2 = \tfrac12 (y - \mathcal{G}(u_j))^\top R^{-1} (y - \mathcal{G}(u_j)), \]

and \(\bar\Phi\), \(\sigma^2_\Phi\) its empirical mean and variance across the ensemble,

\[ \Delta t = \min\!\left( \max\!\left(\frac{M}{2\bar\Phi},\; \sqrt{\frac{M}{2\sigma^2_\Phi}}\right),\; \text{remaining}\right). \]

The two candidates are the paper's statistical discrepancy principle applied to the tempered sub-problem \(y = \mathcal{G}(u) + \sqrt{\alpha}\eta\) with \(\alpha = 1/\Delta t\): since \(\|R^{-1/2}(y - \mathcal{G}(u))\|^2\) is \(\chi^2_M\) under the correct model, it has mean \(M\) (their C1, accuracy, giving \(M/2\bar\Phi\)) and variance \(2M\) (their C2, uncertainty, giving \(\sqrt{M/2\sigma^2_\Phi}\)). The max is deliberate and enforces at least one of the two, not both: their Remark 2 notes that a wide prior makes \(\bar\Phi \ll \sigma_\Phi\), so C1 binds; a narrow prior centred far from the truth flips it, and C2 then licenses the larger step. The outer min is the tempering budget \(1 - t_n\), which is also the stopping rule -- the driver halts on the iteration where it binds.

Note

gh-230 specified both terms in the misfit of the ensemble mean, \(\|R^{-1/2}(y - \bar{\mathcal{G}})\|^2\). This follows the paper instead: eq. (13) defines \(\Phi_n\) as the set of per-particle functionals and eq. (14) takes their mean and variance. The distinction is not cosmetic -- the second term's \(\sigma^2_\Phi\) is identically zero for any single vector, so a mean-misfit reading would make C2 infinite and the max vacuous.

Only obs_noise solves are used; \(R^{-1/2}\) is never formed.

Parameters:

Name	Type	Description	Default
obs_particles	Float[Array, 'J M']	Ensemble in observation space \(\mathcal{G}(u_j)\), shape (J, M). Under tikhonov_augment pass the augmented ensemble, so that \(M\) counts the augmented observations.	required
observation	Float[Array, ' M']	The observation \(y\), shape (M,).	required
obs_noise	AbstractLinearOperator	Observation error covariance \(R\), shape (M, M).	required
remaining	Float[Array, '']	Tempering budget left, \(1 - \sum_{k<n}\Delta t_k\). May be traced.	required
bessel	bool	Use the \(1/(J-1)\) divisor for \(\sigma^2_\Phi\). Defaults to True, matching the rest of this module.	True

Returns:

Type	Description
Float[Array, '']	The step \(\Delta t\), a scalar. Never exceeds remaining, and is
Float[Array, '']	positive whenever remaining is: a degenerate ensemble
Float[Array, '']	(\(\sigma^2_\Phi = 0\)) sends the second candidate to \(+\infty\), so the
Float[Array, '']	max saturates and remaining is returned.

Raises:

Type	Description
ValueError	If the shapes disagree, or if bessel is set with J < 2.

Example

import jax.numpy as jnp import jax.random as jr import lineax as lx from gaussx import discrepancy_step_size G = jr.normal(jr.key(0), (50, 4)) R = lx.DiagonalLinearOperator(jnp.ones(4)) dt = discrepancy_step_size( ... G, jnp.zeros(4), R, remaining=jnp.asarray(1.0) ... ) bool(0.0 < dt <= 1.0) True

Source code in src/gaussx/_inference/_eki.py

def discrepancy_step_size(
    obs_particles: Float[Array, "J M"],
    observation: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    remaining: Float[Array, ""],
    bessel: bool = True,
) -> Float[Array, ""]:
    r"""Adaptive EKI tempering step: the data misfit controller.

    Iglesias & Yang (2021), eq. (14) -- the selection rule of their EKI-DMC
    (Algorithm 3), which needs no tuning parameter. With the per-particle
    least-squares functional

    $$
    \Phi_j = \tfrac12 \big\| R^{-1/2}(y - \mathcal{G}(u_j)) \big\|^2
           = \tfrac12 (y - \mathcal{G}(u_j))^\top R^{-1}
             (y - \mathcal{G}(u_j)),
    $$

    and $\bar\Phi$, $\sigma^2_\Phi$ its empirical mean and variance across the
    ensemble,

    $$
    \Delta t = \min\!\left(
        \max\!\left(\frac{M}{2\bar\Phi},\;
                    \sqrt{\frac{M}{2\sigma^2_\Phi}}\right),\;
        \text{remaining}\right).
    $$

    The two candidates are the paper's statistical discrepancy principle
    applied to the tempered sub-problem $y = \mathcal{G}(u) + \sqrt{\alpha}\eta$
    with $\alpha = 1/\Delta t$: since $\|R^{-1/2}(y - \mathcal{G}(u))\|^2$ is
    $\chi^2_M$ under the correct model, it has mean $M$ (their C1, accuracy,
    giving $M/2\bar\Phi$) and variance $2M$ (their C2, uncertainty, giving
    $\sqrt{M/2\sigma^2_\Phi}$). The **max** is deliberate and enforces *at
    least one* of the two, not both: their Remark 2 notes that a wide prior
    makes $\bar\Phi \ll \sigma_\Phi$, so C1 binds; a narrow prior centred far
    from the truth flips it, and C2 then licenses the larger step. The outer
    min is the tempering budget $1 - t_n$, which is also the stopping rule --
    the driver halts on the iteration where it binds.

    Note:
        gh-230 specified both terms in the misfit of the ensemble *mean*,
        $\|R^{-1/2}(y - \bar{\mathcal{G}})\|^2$. This follows the paper
        instead: eq. (13) defines $\Phi_n$ as the set of per-particle
        functionals and eq. (14) takes their mean and variance. The
        distinction is not cosmetic -- the second term's $\sigma^2_\Phi$ is
        identically zero for any single vector, so a mean-misfit reading would
        make C2 infinite and the ``max`` vacuous.

    Only ``obs_noise`` solves are used; $R^{-1/2}$ is never formed.

    Args:
        obs_particles: Ensemble in observation space $\mathcal{G}(u_j)$, shape
            ``(J, M)``. Under `tikhonov_augment` pass the augmented ensemble,
            so that $M$ counts the augmented observations.
        observation: The observation $y$, shape ``(M,)``.
        obs_noise: Observation error covariance $R$, shape ``(M, M)``.
        remaining: Tempering budget left, $1 - \sum_{k<n}\Delta t_k$. May be
            traced.
        bessel: Use the $1/(J-1)$ divisor for $\sigma^2_\Phi$. Defaults to
            ``True``, matching the rest of this module.

    Returns:
        The step $\Delta t$, a scalar. Never exceeds ``remaining``, and is
        positive whenever ``remaining`` is: a degenerate ensemble
        ($\sigma^2_\Phi = 0$) sends the second candidate to $+\infty$, so the
        ``max`` saturates and ``remaining`` is returned.

    Raises:
        ValueError: If the shapes disagree, or if ``bessel`` is set with
            ``J < 2``.

    Example:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from gaussx import discrepancy_step_size
        >>> G = jr.normal(jr.key(0), (50, 4))
        >>> R = lx.DiagonalLinearOperator(jnp.ones(4))
        >>> dt = discrepancy_step_size(
        ...     G, jnp.zeros(4), R, remaining=jnp.asarray(1.0)
        ... )
        >>> bool(0.0 < dt <= 1.0)
        True
    """
    n_ens, n_obs = obs_particles.shape
    _check_ensemble_size(n_ens, bessel)
    if observation.shape != (n_obs,):
        raise ValueError(
            f"observation must have shape ({n_obs},) to match obs_particles, "
            f"got {observation.shape}."
        )
    if (obs_noise.in_size(), obs_noise.out_size()) != (n_obs, n_obs):
        raise ValueError(
            f"obs_noise must be ({n_obs}, {n_obs}) to match obs_particles, got "
            f"({obs_noise.out_size()}, {obs_noise.in_size()})."
        )

    residuals = observation[None, :] - obs_particles  # (J, M)
    weighted = solve_rows(obs_noise, residuals)  # (J, M), R^{-1} r_j
    misfit = 0.5 * jnp.sum(residuals * weighted, axis=-1)  # (J,), Phi_j

    mean_misfit = jnp.mean(misfit)
    var_misfit = jnp.var(misfit, ddof=1 if bessel else 0)
    accuracy = n_obs / (2.0 * mean_misfit)  # C1
    uncertainty = jnp.sqrt(n_obs / (2.0 * var_misfit))  # C2
    return jnp.minimum(jnp.maximum(accuracy, uncertainty), remaining)

Localization & inflation

The standard fixes for small-ensemble rank deficiency: Schur-product localization with a taper (Gaspari-Cohn by default) and multiplicative / RTPP / RTPS inflation.

Structured linear algebra and Gaussian primitives for JAX.

localization_matrix(coords_a: Float[Array, 'Na D'], coords_b: Float[Array, 'Nb D'], c: float | Float[Array, ''], metric: Callable[[Float[Array, 'Na D'], Float[Array, 'Nb D']], Float[Array, 'Na Nb']] = euclidean_distance) -> Float[Array, 'Na Nb']

Pairwise Gaspari-Cohn taper rho(dist(a_i, b_j); c).

Use this to build the rho_xy (state-obs) and rho_yy (obs-obs) localization matrices consumed by localized_kalman_gain.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'Na D']	First set of points, shape (Na, D).	required
coords_b	Float[Array, 'Nb D']	Second set of points, shape (Nb, D).	required
c	float | Float[Array, '']	Gaspari-Cohn compact-support radius.	required
metric	Callable[[Float[Array, 'Na D'], Float[Array, 'Nb D']], Float[Array, 'Na Nb']]	Pairwise distance function returning an (Na, Nb) matrix. Defaults to euclidean_distance; pass haversine_distance for spherical coordinates.	euclidean_distance

Returns:

Type	Description
Float[Array, 'Na Nb']	Localization matrix of shape (Na, Nb) with entries in [0, 1].

Source code in src/gaussx/_inference/_localization.py

def localization_matrix(
    coords_a: Float[Array, "Na D"],
    coords_b: Float[Array, "Nb D"],
    c: float | Float[Array, ""],
    metric: Callable[
        [Float[Array, "Na D"], Float[Array, "Nb D"]], Float[Array, "Na Nb"]
    ] = euclidean_distance,
) -> Float[Array, "Na Nb"]:
    """Pairwise Gaspari-Cohn taper ``rho(dist(a_i, b_j); c)``.

    Use this to build the ``rho_xy`` (state-obs) and ``rho_yy`` (obs-obs)
    localization matrices consumed by `localized_kalman_gain`.

    Args:
        coords_a: First set of points, shape ``(Na, D)``.
        coords_b: Second set of points, shape ``(Nb, D)``.
        c: Gaspari-Cohn compact-support radius.
        metric: Pairwise distance function returning an ``(Na, Nb)`` matrix.
            Defaults to `euclidean_distance`; pass
            `haversine_distance` for spherical coordinates.

    Returns:
        Localization matrix of shape ``(Na, Nb)`` with entries in ``[0, 1]``.
    """
    return gaspari_cohn(metric(coords_a, coords_b), c)

localized_kalman_gain(particles: Float[Array, 'J N'], obs_particles: Float[Array, 'J M'], obs_noise: lx.AbstractLinearOperator, rho_xy: Float[Array, 'N M'], rho_yy: Float[Array, 'M M'], *, solver: AbstractSolverStrategy | None = None, bessel: bool = True) -> Float[Array, 'N M']

Ensemble Kalman gain with Hadamard (Schur-product) localization.

Computes

\[ K = (\rho_{xy} \circ P_{xy})\,(\rho_{yy} \circ P_{yy} + R)^{-1}, \]

where P_xy is the state-observation cross-covariance and P_yy the observation-space ensemble covariance. Tapering kills spurious long-range sample correlations; because Gaspari-Cohn is positive-definite, the Schur product theorem keeps rho_yy . P_yy PSD, so the innovation covariance stays invertible.

This is the localized counterpart of ensemble_kalman_gain. Unlike that routine, the Hadamard product destroys the low-rank structure, so the innovation covariance is materialized densely and the solve is O(N M + M^3). Recover the unlocalized gain as the c -> inf limit (rho_xy = rho_yy = 1).

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N']	Prior ensemble in state space, shape (J, N).	required
obs_particles	Float[Array, 'J M']	Prior ensemble in observation space, shape (J, M).	required
obs_noise	AbstractLinearOperator	Observation error covariance operator R, shape (M, M).	required
rho_xy	Float[Array, 'N M']	State-observation localization matrix, shape (N, M).	required
rho_yy	Float[Array, 'M M']	Observation-observation localization matrix, shape (M, M).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for the dense innovation solve.	None
bessel	bool	Use the 1 / (J - 1) divisor (EnKF convention, default).	True

Returns:

Type	Description
Float[Array, 'N M']	Dense localized Kalman gain of shape (N, M).

Source code in src/gaussx/_inference/_localization.py

def localized_kalman_gain(
    particles: Float[Array, "J N"],
    obs_particles: Float[Array, "J M"],
    obs_noise: lx.AbstractLinearOperator,
    rho_xy: Float[Array, "N M"],
    rho_yy: Float[Array, "M M"],
    *,
    solver: AbstractSolverStrategy | None = None,
    bessel: bool = True,
) -> Float[Array, "N M"]:
    r"""Ensemble Kalman gain with Hadamard (Schur-product) localization.

    Computes

    $$
    K = (\rho_{xy} \circ P_{xy})\,(\rho_{yy} \circ P_{yy} + R)^{-1},
    $$

    where ``P_xy`` is the state-observation cross-covariance and ``P_yy`` the
    observation-space ensemble covariance. Tapering kills spurious long-range
    sample correlations; because Gaspari-Cohn is positive-definite, the Schur
    product theorem keeps ``rho_yy . P_yy`` PSD, so the innovation covariance
    stays invertible.

    This is the localized counterpart of `ensemble_kalman_gain`. Unlike
    that routine, the Hadamard product destroys the low-rank structure, so the
    innovation covariance is materialized densely and the solve is
    ``O(N M + M^3)``. Recover the unlocalized gain as the ``c -> inf`` limit
    (``rho_xy = rho_yy = 1``).

    Args:
        particles: Prior ensemble in state space, shape ``(J, N)``.
        obs_particles: Prior ensemble in observation space, shape ``(J, M)``.
        obs_noise: Observation error covariance operator ``R``, shape ``(M, M)``.
        rho_xy: State-observation localization matrix, shape ``(N, M)``.
        rho_yy: Observation-observation localization matrix, shape ``(M, M)``.
        solver: Optional solver strategy for the dense innovation solve.
        bessel: Use the ``1 / (J - 1)`` divisor (EnKF convention, default).

    Returns:
        Dense localized Kalman gain of shape ``(N, M)``.
    """
    if particles.shape[0] != obs_particles.shape[0]:
        raise ValueError(
            "particles and obs_particles must share the same ensemble size, "
            f"got J={particles.shape[0]} and J={obs_particles.shape[0]}."
        )
    cross_cov = ensemble_cross_covariance(particles, obs_particles, bessel=bessel)
    obs_cov = ensemble_cross_covariance(obs_particles, obs_particles, bessel=bessel)

    localized_cross = rho_xy * cross_cov
    innovation = rho_yy * obs_cov + obs_noise.as_matrix()
    innovation = symmetrize(innovation)
    innovation_op = lx.MatrixLinearOperator(innovation, lx.positive_semidefinite_tag)
    return solve_rows(innovation_op, localized_cross, solver=solver)

gaspari_cohn(r: Float[Array, '*shape'], c: float | Float[Array, '']) -> Float[Array, '*shape']

Gaspari-Cohn (1999) fifth-order compactly-supported taper.

The standard positive-definite, approximately-Gaussian localization function. With z = 2 |r| / c it is the piecewise-rational

\[ \begin{aligned} \rho = \begin{cases} -\tfrac14 z^5 + \tfrac12 z^4 + \tfrac58 z^3 - \tfrac53 z^2 + 1 & 0 \le z \le 1 \\ \tfrac1{12} z^5 - \tfrac12 z^4 + \tfrac58 z^3 + \tfrac53 z^2 - 5 z + 4 - \tfrac{2}{3 z} & 1 < z \le 2 \\ 0 & z > 2. \end{cases} \end{aligned} \]

so rho(0) = 1 and rho = 0 for |r| >= c (c is the compact-support radius, not a Gaussian length scale). The taper is only \(C^1\) at the knots z = 1, 2.

Differentiability: the 2 / (3 z) term in the middle branch is guarded with a safe denominator so reverse-mode gradients are finite at r = 0 (which would otherwise produce NaN via the standard where pitfall).

Parameters:

Name	Type	Description	Default
r	Float[Array, '*shape']	Distances (any shape), e.g. a pairwise distance matrix.	required
c	float | Float[Array, '']	Compact-support radius; rho = 0 beyond |r| = c.	required

Returns:

Type	Description
Float[Array, '*shape']	Taper values in [0, 1], same shape as r.

Source code in src/gaussx/_inference/_localization.py

def gaspari_cohn(
    r: Float[Array, "*shape"], c: float | Float[Array, ""]
) -> Float[Array, "*shape"]:
    r"""Gaspari-Cohn (1999) fifth-order compactly-supported taper.

    The standard positive-definite, approximately-Gaussian localization
    function. With ``z = 2 |r| / c`` it is the piecewise-rational

    $$
    \begin{aligned}
    \rho = \begin{cases}
      -\tfrac14 z^5 + \tfrac12 z^4 + \tfrac58 z^3 - \tfrac53 z^2 + 1
        & 0 \le z \le 1 \\
      \tfrac1{12} z^5 - \tfrac12 z^4 + \tfrac58 z^3 + \tfrac53 z^2
        - 5 z + 4 - \tfrac{2}{3 z}
        & 1 < z \le 2 \\
      0 & z > 2.
    \end{cases}
    \end{aligned}
    $$

    so ``rho(0) = 1`` and ``rho = 0`` for ``|r| >= c`` (``c`` is the
    compact-support radius, **not** a Gaussian length scale). The taper is
    only $C^1$ at the knots ``z = 1, 2``.

    Differentiability: the ``2 / (3 z)`` term in the middle branch is guarded
    with a safe denominator so reverse-mode gradients are finite at ``r = 0``
    (which would otherwise produce ``NaN`` via the standard ``where`` pitfall).

    Args:
        r: Distances (any shape), e.g. a pairwise distance matrix.
        c: Compact-support radius; ``rho = 0`` beyond ``|r| = c``.

    Returns:
        Taper values in ``[0, 1]``, same shape as ``r``.
    """
    z = 2.0 * jnp.abs(r) / c
    # Guard the 1 / z term: at z = 0 the near branch is selected, but JAX still
    # traces the middle branch, so an unguarded 1 / z poisons the gradient.
    z_safe = jnp.where(z > 0.0, z, 1.0)

    near = -0.25 * z**5 + 0.5 * z**4 + 0.625 * z**3 - (5.0 / 3.0) * z**2 + 1.0
    mid = (
        (1.0 / 12.0) * z**5
        - 0.5 * z**4
        + 0.625 * z**3
        + (5.0 / 3.0) * z**2
        - 5.0 * z
        + 4.0
        - 2.0 / (3.0 * z_safe)
    )
    return jnp.where(z <= 1.0, near, jnp.where(z < 2.0, mid, 0.0))

inflate_multiplicative(ensemble: Float[Array, 'J N'], factor: float | Float[Array, '']) -> Float[Array, 'J N']

Multiplicative ensemble inflation about the mean.

Restores ensemble spread lost to sampling error / model collapse by scaling perturbations: x_j <- x_bar + factor (x_j - x_bar).

Parameters:

Name	Type	Description	Default
ensemble	Float[Array, 'J N']	Ensemble of shape (J, N).	required
factor	float | Float[Array, '']	Inflation factor >= 1 (e.g. 1.02-1.10).	required

Returns:

Type	Description
Float[Array, 'J N']	Inflated ensemble, shape (J, N). The mean is unchanged.

Source code in src/gaussx/_inference/_inflation.py

def inflate_multiplicative(
    ensemble: Float[Array, "J N"],
    factor: float | Float[Array, ""],
) -> Float[Array, "J N"]:
    r"""Multiplicative ensemble inflation about the mean.

    Restores ensemble spread lost to sampling error / model collapse by scaling
    perturbations: ``x_j <- x_bar + factor (x_j - x_bar)``.

    Args:
        ensemble: Ensemble of shape ``(J, N)``.
        factor: Inflation factor ``>= 1`` (e.g. ``1.02``-``1.10``).

    Returns:
        Inflated ensemble, shape ``(J, N)``. The mean is unchanged.
    """
    mean = jnp.mean(ensemble, axis=0, keepdims=True)
    return mean + factor * (ensemble - mean)

inflate_rtpp(posterior: Float[Array, 'J N'], prior: Float[Array, 'J N'], alpha: float | Float[Array, '']) -> Float[Array, 'J N']

Relaxation to prior perturbations (RTPP; Zhang et al. 2004).

Relaxes posterior perturbations toward the prior perturbations while keeping the posterior mean: x'^a <- (1 - alpha) x'^a + alpha x'^f, where the perturbations are taken about each ensemble's own mean.

Parameters:

Name	Type	Description	Default
posterior	Float[Array, 'J N']	Analysis ensemble, shape (J, N).	required
prior	Float[Array, 'J N']	Forecast ensemble, shape (J, N).	required
alpha	float | Float[Array, '']	Relaxation weight in [0, 1].	required

Returns:

Type	Description
Float[Array, 'J N']	Relaxed analysis ensemble, shape (J, N). The posterior mean is
Float[Array, 'J N']	preserved.

Source code in src/gaussx/_inference/_inflation.py

def inflate_rtpp(
    posterior: Float[Array, "J N"],
    prior: Float[Array, "J N"],
    alpha: float | Float[Array, ""],
) -> Float[Array, "J N"]:
    r"""Relaxation to prior perturbations (RTPP; Zhang et al. 2004).

    Relaxes posterior perturbations toward the prior perturbations while keeping
    the posterior mean: ``x'^a <- (1 - alpha) x'^a + alpha x'^f``, where the
    perturbations are taken about each ensemble's own mean.

    Args:
        posterior: Analysis ensemble, shape ``(J, N)``.
        prior: Forecast ensemble, shape ``(J, N)``.
        alpha: Relaxation weight in ``[0, 1]``.

    Returns:
        Relaxed analysis ensemble, shape ``(J, N)``. The posterior mean is
        preserved.
    """
    post_mean = jnp.mean(posterior, axis=0, keepdims=True)
    post_pert = posterior - post_mean
    prior_pert = prior - jnp.mean(prior, axis=0, keepdims=True)
    return post_mean + (1.0 - alpha) * post_pert + alpha * prior_pert

inflate_rtps(posterior: Float[Array, 'J N'], prior: Float[Array, 'J N'], beta: float | Float[Array, ''], eps: float = 1e-12) -> Float[Array, 'J N']

Relaxation to prior spread (RTPS; Whitaker & Hamill 2012).

Scales each posterior perturbation, per coordinate, so the analysis spread relaxes back toward the prior spread: x'^a <- x'^a [ (1 - beta) + beta sigma^f / sigma^a ], with sigma the per-coordinate ensemble standard deviation.

Parameters:

Name	Type	Description	Default
posterior	Float[Array, 'J N']	Analysis ensemble, shape (J, N).	required
prior	Float[Array, 'J N']	Forecast ensemble, shape (J, N).	required
beta	float | Float[Array, '']	Relaxation weight in [0, 1].	required
eps	float	Floor on the posterior std to avoid division by zero.	1e-12

Returns:

Type	Description
Float[Array, 'J N']	Spread-restored analysis ensemble, shape (J, N). The posterior mean
Float[Array, 'J N']	is preserved.

Source code in src/gaussx/_inference/_inflation.py

def inflate_rtps(
    posterior: Float[Array, "J N"],
    prior: Float[Array, "J N"],
    beta: float | Float[Array, ""],
    eps: float = 1e-12,
) -> Float[Array, "J N"]:
    r"""Relaxation to prior spread (RTPS; Whitaker & Hamill 2012).

    Scales each posterior perturbation, per coordinate, so the analysis spread
    relaxes back toward the prior spread:
    ``x'^a <- x'^a [ (1 - beta) + beta sigma^f / sigma^a ]``, with ``sigma`` the
    per-coordinate ensemble standard deviation.

    Args:
        posterior: Analysis ensemble, shape ``(J, N)``.
        prior: Forecast ensemble, shape ``(J, N)``.
        beta: Relaxation weight in ``[0, 1]``.
        eps: Floor on the posterior std to avoid division by zero.

    Returns:
        Spread-restored analysis ensemble, shape ``(J, N)``. The posterior mean
        is preserved.
    """
    post_mean = jnp.mean(posterior, axis=0, keepdims=True)
    post_pert = posterior - post_mean
    sigma_post = jnp.std(posterior, axis=0)
    sigma_prior = jnp.std(prior, axis=0)
    scale = (1.0 - beta) + beta * sigma_prior / (sigma_post + eps)
    return post_mean + post_pert * scale[None, :]

Distances

Structured linear algebra and Gaussian primitives for JAX.

euclidean_distance(coords_a: Float[Array, 'Na D'], coords_b: Float[Array, 'Nb D']) -> Float[Array, 'Na Nb']

Pairwise Euclidean distances ||a_i - b_j||.

A default metric for localization_matrix. Builds on stable_squared_distances and takes a gradient-safe square root so zero distances (e.g. the diagonal of a self-distance matrix) do not produce NaN gradients.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'Na D']	First set of points, shape (Na, D).	required
coords_b	Float[Array, 'Nb D']	Second set of points, shape (Nb, D).	required

Returns:

Type	Description
Float[Array, 'Na Nb']	Distance matrix of shape (Na, Nb).

Source code in src/gaussx/_inference/_localization.py

def euclidean_distance(
    coords_a: Float[Array, "Na D"],
    coords_b: Float[Array, "Nb D"],
) -> Float[Array, "Na Nb"]:
    """Pairwise Euclidean distances ``||a_i - b_j||``.

    A default ``metric`` for `localization_matrix`. Builds on
    `stable_squared_distances` and takes a gradient-safe square root so
    zero distances (e.g. the diagonal of a self-distance matrix) do not produce
    ``NaN`` gradients.

    Args:
        coords_a: First set of points, shape ``(Na, D)``.
        coords_b: Second set of points, shape ``(Nb, D)``.

    Returns:
        Distance matrix of shape ``(Na, Nb)``.
    """
    sq = stable_squared_distances(
        coords_a,
        coords_b,
        compute_dtype=coords_a.dtype,
        accumulate_dtype=coords_a.dtype,
    )
    sq_safe = jnp.where(sq > 0.0, sq, 1.0)
    return jnp.where(sq > 0.0, jnp.sqrt(sq_safe), 0.0)

haversine_distance(coords_a: Float[Array, 'Na 2'], coords_b: Float[Array, 'Nb 2'], radius: float = 6371000.0) -> Float[Array, 'Na Nb']

Pairwise great-circle (haversine) distances on a sphere.

A metric for localization_matrix on geophysical grids. Coordinates are (latitude, longitude) in radians.

Parameters:

Name	Type	Description	Default
coords_a	Float[Array, 'Na 2']	First set of points (lat, lon) in radians, shape (Na, 2).	required
coords_b	Float[Array, 'Nb 2']	Second set of points (lat, lon) in radians, shape (Nb, 2).	required
radius	float	Sphere radius in the units of the returned distance (default the Earth mean radius, 6.371e6 m).	6371000.0

Returns:

Type	Description
Float[Array, 'Na Nb']	Great-circle distance matrix of shape (Na, Nb).

Source code in src/gaussx/_inference/_localization.py

def haversine_distance(
    coords_a: Float[Array, "Na 2"],
    coords_b: Float[Array, "Nb 2"],
    radius: float = 6.371e6,
) -> Float[Array, "Na Nb"]:
    """Pairwise great-circle (haversine) distances on a sphere.

    A ``metric`` for `localization_matrix` on geophysical grids.
    Coordinates are ``(latitude, longitude)`` in **radians**.

    Args:
        coords_a: First set of points ``(lat, lon)`` in radians, shape ``(Na, 2)``.
        coords_b: Second set of points ``(lat, lon)`` in radians, shape ``(Nb, 2)``.
        radius: Sphere radius in the units of the returned distance (default the
            Earth mean radius, ``6.371e6`` m).

    Returns:
        Great-circle distance matrix of shape ``(Na, Nb)``.
    """
    lat_a = coords_a[:, 0][:, None]
    lon_a = coords_a[:, 1][:, None]
    lat_b = coords_b[:, 0][None, :]
    lon_b = coords_b[:, 1][None, :]
    dlat = lat_b - lat_a
    dlon = lon_b - lon_a
    h = (
        jnp.sin(dlat / 2.0) ** 2
        + jnp.cos(lat_a) * jnp.cos(lat_b) * jnp.sin(dlon / 2.0) ** 2
    )
    return 2.0 * radius * jnp.arcsin(jnp.sqrt(jnp.clip(h, 0.0, 1.0)))