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 linear regression

Closed-form Gaussian posterior updates — full covariance or diagonal-only — plus the marginal likelihood and expected log-likelihood that score them.

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, *, solver: AbstractSolverStrategy | None = None) -> 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	Learning rate / damping factor.	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

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

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,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> 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.

    Returns:
        Tuple ``(nat1_new, nat2_new)`` — updated natural parameters.
    """
    # 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, nat2_new

blr_diag_update(nat1: Float[Array, ' d'], nat2_diag: Float[Array, ' d'], grad: Float[Array, ' d'], hessian_diag: Float[Array, ' d'], lr: float) -> 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	Learning rate / damping factor.	required

Returns:

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

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,
) -> 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.

    Returns:
        Tuple ``(nat1_new, nat2_new)`` — updated natural parameters.
    """
    # 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, nat2_new

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)
    from gaussx._linalg._linalg import trace_product

    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 and gaussx.damped_natural_update. This differs from gaussx.mean_cov_to_natural / gaussx.natural_to_mean_cov, which use the exponential-family convention η₂ = −Λ/2.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	Current mean, shape (N,) or (D,).	required
jacobian	Float[Array, ' N']	First derivative of log-likelihood, shape (N,).	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` and
        `gaussx.damped_natural_update`. This differs from
        `gaussx.mean_cov_to_natural` / `gaussx.natural_to_mean_cov`, which
        use the exponential-family convention ``η₂ = −Λ/2``.

    Args:
        mean: Current mean, shape ``(N,)`` or ``(D,)``.
        jacobian: First derivative of log-likelihood, shape ``(N,)``.
        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 = 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}

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	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 = 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}

    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

    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 be the same type"
        raise TypeError(msg)

    return nat1_new, nat2_new

gauss_newton_precision(jacobian: Float[Array, 'D_obs D_latent']) -> lx.AbstractLinearOperator

Gauss-Newton precision matrix J^T J.

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).

When D_{obs} < D_{latent}, returns a LowRankUpdate to enable efficient Woodbury-based solves downstream.

Parameters:

Name	Type	Description	Default
jacobian	Float[Array, 'D_obs D_latent']	Jacobian of the residual, shape (D_obs, D_latent).	required

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"],
) -> lx.AbstractLinearOperator:
    r"""Gauss-Newton precision matrix ``J^T J``.

    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).

    When ``D_{obs} < D_{latent}``, returns a `LowRankUpdate`
    to enable efficient Woodbury-based solves downstream.

    Args:
        jacobian: Jacobian of the residual, shape ``(D_obs, D_latent)``.

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

    if D_obs < D_latent:
        base = lx.DiagonalLinearOperator(jnp.zeros(D_latent))
        return LowRankUpdate(
            base=base,
            U=jacobian.T,
            d=jnp.ones(D_obs),
            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 = 1.0) -> Float[Array, 'd d']

Riemannian gradient correction for PSD precision updates.

Ensures the corrected Hessian remains negative semi-definite, stabilizing 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.

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	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 = 1.0,
) -> Float[Array, "d d"]:
    r"""Riemannian gradient correction for PSD precision updates.

    Ensures the corrected Hessian remains negative semi-definite,
    stabilizing 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.

    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) -> 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 and gaussx.damped_natural_update. This differs from gaussx.mean_cov_to_natural / gaussx.natural_to_mean_cov, which use the exponential-family convention η₂ = −Λ/2.

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

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,
) -> 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` and `gaussx.damped_natural_update`.
        This differs from `gaussx.mean_cov_to_natural` /
        `gaussx.natural_to_mean_cov`, which use the exponential-family
        convention ``η₂ = −Λ/2``.

    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).

    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
        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
    from gaussx._linalg._linalg import solve_rows

    W = solve_rows(K_zz, K_xz, solver=solver)  # (N, M)
    tr_approx = jnp.sum(K_xz * W)

    return tr_full - tr_approx

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.

Source code in src/gaussx/_inference/_ensemble.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.
    """
    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).

Source code in src/gaussx/_inference/_ensemble.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)``.
    """
    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) -> 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 as a LowRankUpdate so solve_rows can use structural dispatch via the Woodbury identity.

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.	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

Returns:

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

Source code in src/gaussx/_inference/_ensemble.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,
) -> 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 as a ``LowRankUpdate`` so
    ``solve_rows`` can use structural dispatch via the Woodbury identity.

    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.
        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.

    Returns:
        Dense 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,
    )
    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 a Cholesky factor of obs_noise.
• 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/_ensemble.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 a Cholesky
      factor of ``obs_noise``.
    - ``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 = 1.0) -> 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 1, which makes the transform exactly mean-preserving (sum_j X'^a_j = 0).

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	Multiplicative covariance inflation lambda >= 1, applied to the prior term (J - 1) / lambda.	1.0

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.

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

def etkf_transform(
    obs_particles: Float[Array, "J M"],
    y: Float[Array, " M"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    inflation: float = 1.0,
) -> 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 ``1``, which makes the transform
    exactly mean-preserving (``sum_j X'^a_j = 0``).

    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``.

    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``.
    """
    n_ens = obs_particles.shape[0]
    obs_mean = jnp.mean(obs_particles, axis=0)
    obs_pert = obs_particles - obs_mean[None, :]  # (J, M), zero-mean rows

    r_matrix = obs_noise.as_matrix()
    # R^{-1} applied to the (M, .) right-hand sides.
    rinv_pert = jnp.linalg.solve(r_matrix, obs_pert.T)  # (M, J)
    rinv_d = jnp.linalg.solve(r_matrix, y - obs_mean)  # (M,)

    precision = (n_ens - 1) / inflation * jnp.eye(n_ens) + obs_pert @ rinv_pert
    precision = symmetrize(precision)
    analysis_cov = jnp.linalg.inv(precision)  # tilde A, (J, J)

    w_mean = analysis_cov @ (obs_pert @ rinv_d)  # (J,)
    transform = _symmetric_sqrt((n_ens - 1) * analysis_cov)
    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 for the gain computation. Note: etkf_transform materializes obs_noise via as_matrix() when deterministic=True.	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.	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/_ensemble.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` for the gain
            computation. Note: `etkf_transform` materializes ``obs_noise`` via
            ``as_matrix()`` when ``deterministic=True``.
        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``.
        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)

    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 deterministic:
        obs_mean = jnp.mean(obs_particles, axis=0)  # (M,)
        anomalies = particles - jnp.mean(particles, axis=0, keepdims=True)  # (J, N)
        _, transform = etkf_transform(obs_particles, observation, tempered_noise)
        mean_increment = gain @ (observation - obs_mean)  # (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

    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/_ensemble.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/_ensemble.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, 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	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/_ensemble.py

def localization_matrix(
    coords_a: Float[Array, "Na D"],
    coords_b: Float[Array, "Nb D"],
    c: float,
    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/_ensemble.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, '*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	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/_ensemble.py

def gaspari_cohn(r: Float[Array, "*shape"], c: float) -> 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, '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	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/_ensemble.py

def inflate_multiplicative(
    ensemble: Float[Array, "J N"],
    factor: float,
) -> 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, '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	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/_ensemble.py

def inflate_rtpp(
    posterior: Float[Array, "J N"],
    prior: Float[Array, "J N"],
    alpha: float,
) -> 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, 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	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/_ensemble.py

def inflate_rtps(
    posterior: Float[Array, "J N"],
    prior: Float[Array, "J N"],
    beta: float,
    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/_ensemble.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/_ensemble.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)))