Primitives

Pure functions on arrays and operators. Stateless, differentiable, composable; compatible with jax.jit, jax.grad, and eqx.filter_vmap. Every filter and smoother in the package is assembled from these.

Heavy linear algebra is delegated to gaussx: ensemble_covariance returns a gaussx.LowRankUpdate operator (rank N_e, never a dense (N_x, N_x) matrix), and the solves inside kalman_gain and log_likelihood dispatch through gaussx's structural solver — a Woodbury / matrix-determinant-lemma path for low-rank updates — with an optional solver= keyword to override the strategy.

Localization tapers and inflation primitives have their own pages: Localization and Inflation.

Ensemble statistics

Bessel-corrected moments of an (N_e, N_x) ensemble: mean, centred anomalies, and the (cross-)covariances built from them.

filterax — differentiable ensemble Kalman filters and processes for JAX.

ensemble_mean(particles: Float[Array, 'N_e N_x']) -> Float[Array, ' N_x']

Ensemble mean x̄ = Nₑ⁻¹ Σⱼ x⁽ʲ⁾.

Computed as an einx reduction over the ensemble axis. O(Nₑ Nₓ).

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Ensemble of shape (Nₑ, Nₓ) with rows as members.	required

Returns:

Type	Description
Float[Array, ' N_x']	Mean vector of shape (Nₓ,).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import ensemble_mean
>>> ens = jnp.array([[0.0, 1.0], [2.0, 3.0]])
>>> ensemble_mean(ens)
Array([1., 2.], dtype=float32)

Source code in src/filterax/_primitives/_statistics.py

def ensemble_mean(
    particles: Float[Array, "N_e N_x"],
) -> Float[Array, " N_x"]:
    r"""Ensemble mean ``x̄ = Nₑ⁻¹ Σⱼ x⁽ʲ⁾``.

    Computed as an ``einx`` reduction over the ensemble axis. ``O(Nₑ Nₓ)``.

    Args:
        particles: Ensemble of shape ``(Nₑ, Nₓ)`` with rows as members.

    Returns:
        Mean vector of shape ``(Nₓ,)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import ensemble_mean
        >>> ens = jnp.array([[0.0, 1.0], [2.0, 3.0]])
        >>> ensemble_mean(ens)
        Array([1., 2.], dtype=float32)
    """
    return einx.mean("e x -> x", particles)

ensemble_anomalies(particles: Float[Array, 'N_e N_x']) -> Float[Array, 'N_e N_x']

Centred perturbations X′ = X − 𝟙 x̄ᵀ.

The rows of X′ sum to zero by construction. X′ is the fundamental input to the cross-covariance and Kalman-gain recipes: every ensemble Kalman filter is some choice of square root applied to these anomalies.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Ensemble of shape (Nₑ, Nₓ).	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	Centred anomaly matrix of shape (Nₑ, Nₓ).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import ensemble_anomalies
>>> ens = jnp.array([[0.0, 1.0], [2.0, 3.0]])
>>> ensemble_anomalies(ens)  # rows sum to zero
Array([[-1., -1.],
       [ 1.,  1.]], dtype=float32)

Source code in src/filterax/_primitives/_statistics.py

def ensemble_anomalies(
    particles: Float[Array, "N_e N_x"],
) -> Float[Array, "N_e N_x"]:
    r"""Centred perturbations ``X′ = X − 𝟙 x̄ᵀ``.

    The rows of ``X′`` sum to zero by construction. ``X′`` is the
    fundamental input to the cross-covariance and Kalman-gain recipes:
    every ensemble Kalman filter is some choice of square root applied to
    these anomalies.

    Args:
        particles: Ensemble of shape ``(Nₑ, Nₓ)``.

    Returns:
        Centred anomaly matrix of shape ``(Nₑ, Nₓ)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import ensemble_anomalies
        >>> ens = jnp.array([[0.0, 1.0], [2.0, 3.0]])
        >>> ensemble_anomalies(ens)  # rows sum to zero
        Array([[-1., -1.],
               [ 1.,  1.]], dtype=float32)
    """
    mean = ensemble_mean(particles)
    return einx.subtract("e x, x -> e x", particles, mean)

ensemble_covariance(particles: Float[Array, 'N_e N_x']) -> gaussx.LowRankUpdate

Sample covariance P = (Nₑ − 1)⁻¹ X′ᵀ X′ as a low-rank operator.

Delegates to :func:gaussx.ensemble_covariance with the EnKF Bessel convention. Returns a :class:gaussx.LowRankUpdate of rank ≤ Nₑ − 1; the dense (Nₓ, Nₓ) matrix is never materialised. Downstream solve / logdet calls exploit the low-rank structure via the Woodbury identity and matrix-determinant lemma — see :mod:gaussx dispatch.

Complexity: O(Nₑ Nₓ) to construct; structure-aware solves cost O(Nₑ² Nₓ + Nₑ³) instead of O(Nₓ³).

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Ensemble of shape (Nₑ, Nₓ).	required

Returns:

Type	Description
LowRankUpdate	class:gaussx.LowRankUpdate representing P.

Raises:

Type	Description
ValueError	if Nₑ < 2 (the Bessel divisor is undefined).

Examples:

>>> import gaussx
>>> import jax.numpy as jnp
>>> from filterax import ensemble_covariance
>>> ens = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
>>> P = ensemble_covariance(ens)
>>> isinstance(P, gaussx.LowRankUpdate)
True
>>> P.as_matrix().shape
(2, 2)

Source code in src/filterax/_primitives/_statistics.py

def ensemble_covariance(
    particles: Float[Array, "N_e N_x"],
) -> gaussx.LowRankUpdate:
    r"""Sample covariance ``P = (Nₑ − 1)⁻¹ X′ᵀ X′`` as a low-rank operator.

    Delegates to :func:`gaussx.ensemble_covariance` with the EnKF Bessel
    convention. Returns a :class:`gaussx.LowRankUpdate` of rank
    ``≤ Nₑ − 1``; the dense ``(Nₓ, Nₓ)`` matrix is *never* materialised.
    Downstream ``solve`` / ``logdet`` calls exploit the low-rank structure
    via the Woodbury identity and matrix-determinant lemma — see
    :mod:`gaussx` dispatch.

    Complexity: ``O(Nₑ Nₓ)`` to construct; structure-aware solves cost
    ``O(Nₑ² Nₓ + Nₑ³)`` instead of ``O(Nₓ³)``.

    Args:
        particles: Ensemble of shape ``(Nₑ, Nₓ)``.

    Returns:
        :class:`gaussx.LowRankUpdate` representing ``P``.

    Raises:
        ValueError: if ``Nₑ < 2`` (the Bessel divisor is undefined).

    Examples:
        >>> import gaussx
        >>> import jax.numpy as jnp
        >>> from filterax import ensemble_covariance
        >>> ens = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
        >>> P = ensemble_covariance(ens)
        >>> isinstance(P, gaussx.LowRankUpdate)
        True
        >>> P.as_matrix().shape
        (2, 2)
    """
    check_ensemble_size(particles.shape[0])
    return gaussx.ensemble_covariance(particles, bessel=True)

cross_covariance(particles: Float[Array, 'N_e N_x'], obs_particles: Float[Array, 'N_e N_y']) -> Float[Array, 'N_x N_y']

Cross-covariance Cˣᴴ = (Nₑ − 1)⁻¹ X′ᵀ (HX)′.

Returned as a dense (Nₓ, Nᵧ) array because Nᵧ is typically small (instrument footprint, point observations). For nonlinear observation operators this is the ensemble's implicit derivative-free linearisation of ∇H — see Evensen (2003) §5.

Complexity: O(Nₑ Nₓ Nᵧ).

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Prior ensemble in state space, shape (Nₑ, Nₓ).	required
obs_particles	Float[Array, 'N_e N_y']	Prior ensemble mapped to obs space, shape (Nₑ, Nᵧ).	required

Returns:

Type	Description
Float[Array, 'N_x N_y']	Dense cross-covariance of shape (Nₓ, Nᵧ).

Raises:

Type	Description
ValueError	if Nₑ < 2.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import cross_covariance
>>> particles = jnp.array([[0.0, 0.0], [2.0, 2.0]])
>>> obs_particles = jnp.array([[0.0], [1.0]])
>>> cross_covariance(particles, obs_particles)
Array([[1.],
       [1.]], dtype=float32)

Source code in src/filterax/_primitives/_statistics.py

def cross_covariance(
    particles: Float[Array, "N_e N_x"],
    obs_particles: Float[Array, "N_e N_y"],
) -> Float[Array, "N_x N_y"]:
    r"""Cross-covariance ``Cˣᴴ = (Nₑ − 1)⁻¹ X′ᵀ (HX)′``.

    Returned as a dense ``(Nₓ, Nᵧ)`` array because ``Nᵧ`` is typically
    small (instrument footprint, point observations). For nonlinear
    observation operators this *is* the ensemble's implicit derivative-free
    linearisation of ``∇H`` — see Evensen (2003) §5.

    Complexity: ``O(Nₑ Nₓ Nᵧ)``.

    Args:
        particles: Prior ensemble in state space, shape ``(Nₑ, Nₓ)``.
        obs_particles: Prior ensemble mapped to obs space, shape
            ``(Nₑ, Nᵧ)``.

    Returns:
        Dense cross-covariance of shape ``(Nₓ, Nᵧ)``.

    Raises:
        ValueError: if ``Nₑ < 2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import cross_covariance
        >>> particles = jnp.array([[0.0, 0.0], [2.0, 2.0]])
        >>> obs_particles = jnp.array([[0.0], [1.0]])
        >>> cross_covariance(particles, obs_particles)
        Array([[1.],
               [1.]], dtype=float32)
    """
    check_ensemble_size(particles.shape[0])
    return gaussx.ensemble_cross_covariance(particles, obs_particles, bessel=True)

Kalman gain

For the Schur-tapered variant that suppresses spurious long-range correlations, see localized_kalman_gain on the Localization page.

filterax — differentiable ensemble Kalman filters and processes for JAX.

kalman_gain(particles: Float[Array, 'N_e N_x'], obs_particles: Float[Array, 'N_e N_y'], obs_noise: lx.AbstractLinearOperator, *, solver: gaussx.AbstractSolverStrategy | None = None) -> Float[Array, 'N_x N_y']

Ensemble Kalman gain K = Cˣᴴ (Cᴴᴴ + R)⁻¹.

Thin wrapper over :func:gaussx.ensemble_kalman_gain with the EnKF Bessel convention. The innovation covariance S = Cᴴᴴ + R is assembled as a :class:gaussx.LowRankUpdate so structural dispatch routes the solve through the Woodbury identity:

\[ S^{-1} = R^{-1} - R^{-1} U \left( (N_{e} - 1) I + U^{\top} R^{-1} U \right)^{-1} U^{\top} R^{-1}, \qquad U = \frac{(HX)^{\prime\top}}{\sqrt{N_{e} - 1}}. \]

Total cost O(Nₑ² Nᵧ + Nₑ³) for low-rank Woodbury, instead of O(Nᵧ³) for the dense fallback. The Woodbury route is not unconditional: from gaussx v0.6.5, :func:gaussx.ensemble_kalman_gain forms the (Nᵧ, Nᵧ) innovation densely and solves it directly when Nₑ ≥ Nᵧ, and uses Woodbury only otherwise; the two routes agree to rounding. On the Woodbury route the dense (Nᵧ, Nᵧ) matrix is never materialised when R carries structure (diagonal, low-rank, Toeplitz, …).

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Prior ensemble in state space, shape (Nₑ, Nₓ).	required
obs_particles	Float[Array, 'N_e N_y']	H applied to each member, shape (Nₑ, Nᵧ). For linear H, this is X H ᵀ; for nonlinear H, the ensemble provides an implicit linearisation.	required
obs_noise	AbstractLinearOperator	Observation error covariance R as a linear operator.	required
solver	AbstractSolverStrategy | None	Optional :class:gaussx.AbstractSolverStrategy. None lets gaussx pick by operator type (Woodbury for low-rank + structured base, dense Cholesky otherwise).	None

Returns:

Type	Description
Float[Array, 'N_x N_y']	Dense Kalman gain of shape (Nₓ, Nᵧ). Materialising K is
Float[Array, 'N_x N_y']	intentional — it is consumed once by the analysis update and is
Float[Array, 'N_x N_y']	usually small (Nᵧ is the observation count).

Raises:

Type	Description
ValueError	if Nₑ < 2.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import kalman_gain
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> obs_particles = particles[:, :1]  # H observes the first component
>>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(1))
>>> kalman_gain(particles, obs_particles, R)
Array([[0.6666666],
       [0.       ]], dtype=float32)

Source code in src/filterax/_primitives/_gain.py

def kalman_gain(
    particles: Float[Array, "N_e N_x"],
    obs_particles: Float[Array, "N_e N_y"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    solver: gaussx.AbstractSolverStrategy | None = None,
) -> Float[Array, "N_x N_y"]:
    r"""Ensemble Kalman gain ``K = Cˣᴴ (Cᴴᴴ + R)⁻¹``.

    Thin wrapper over :func:`gaussx.ensemble_kalman_gain` with the EnKF
    Bessel convention. The innovation covariance ``S = Cᴴᴴ + R`` is
    assembled as a :class:`gaussx.LowRankUpdate` so structural dispatch
    routes the solve through the Woodbury identity:

    $$
    S^{-1} = R^{-1} - R^{-1} U
        \left( (N_{e} - 1) I + U^{\top} R^{-1} U \right)^{-1}
        U^{\top} R^{-1},
    \qquad
    U = \frac{(HX)^{\prime\top}}{\sqrt{N_{e} - 1}}.
    $$

    Total cost
    ``O(Nₑ² Nᵧ + Nₑ³)`` for low-rank Woodbury, instead of ``O(Nᵧ³)`` for
    the dense fallback. The Woodbury route is not unconditional: from
    gaussx v0.6.5, :func:`gaussx.ensemble_kalman_gain` forms the ``(Nᵧ, Nᵧ)``
    innovation densely and solves it directly when ``Nₑ ≥ Nᵧ``, and uses
    Woodbury only otherwise; the two routes agree to rounding. On the
    Woodbury route the dense ``(Nᵧ, Nᵧ)`` matrix is never materialised when
    ``R`` carries structure (diagonal, low-rank, Toeplitz, …).

    Args:
        particles: Prior ensemble in state space, shape ``(Nₑ, Nₓ)``.
        obs_particles: ``H`` applied to each member, shape ``(Nₑ, Nᵧ)``.
            For linear ``H``, this is ``X H ᵀ``; for nonlinear ``H``, the
            ensemble provides an implicit linearisation.
        obs_noise: Observation error covariance ``R`` as a linear operator.
        solver: Optional :class:`gaussx.AbstractSolverStrategy`. ``None``
            lets gaussx pick by operator type (Woodbury for low-rank +
            structured base, dense Cholesky otherwise).

    Returns:
        Dense Kalman gain of shape ``(Nₓ, Nᵧ)``. Materialising ``K`` is
        intentional — it is consumed once by the analysis update and is
        usually small (``Nᵧ`` is the observation count).

    Raises:
        ValueError: if ``Nₑ < 2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import kalman_gain
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> obs_particles = particles[:, :1]  # H observes the first component
        >>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(1))
        >>> kalman_gain(particles, obs_particles, R)
        Array([[0.6666666],
               [0.       ]], dtype=float32)
    """
    check_ensemble_size(particles.shape[0])
    return gaussx.ensemble_kalman_gain(
        particles, obs_particles, obs_noise, solver=solver, bessel=True
    )

Likelihood & innovation diagnostics

The innovation covariance S = Cᴴᴴ + R as a low-rank update, the Gaussian log-likelihood of an innovation under it, and the bundled per-cycle innovation statistics.

filterax — differentiable ensemble Kalman filters and processes for JAX.

InnovationStatistics

Bases: TypedDict

Return type of :func:innovation_statistics.

Keys

innovation: v = y − H x̄. innovation_cov: S = Cᴴᴴ + R as a :class:gaussx.LowRankUpdate. normalized_innovation: S⁻¹ᐟ² v (whitened residual). log_likelihood: Scalar Gaussian log-probability of v under S.

Source code in src/filterax/_primitives/_likelihood.py

class InnovationStatistics(TypedDict):
    """Return type of :func:`innovation_statistics`.

    Keys:
        innovation: ``v = y − H x̄``.
        innovation_cov: ``S = Cᴴᴴ + R`` as a :class:`gaussx.LowRankUpdate`.
        normalized_innovation: ``S⁻¹ᐟ² v`` (whitened residual).
        log_likelihood: Scalar Gaussian log-probability of ``v`` under ``S``.
    """

    innovation: Float[Array, " N_y"]
    innovation_cov: lx.AbstractLinearOperator
    normalized_innovation: Float[Array, " N_y"]
    log_likelihood: Float[Array, ""]

log_likelihood(innovation: Float[Array, ' N_y'], innovation_cov: lx.AbstractLinearOperator, *, solver: gaussx.AbstractSolverStrategy | None = None) -> Float[Array, '']

Gaussian log-probability of an innovation vector.

\[ \log p(y \mid \text{forecast}) = -\tfrac{1}{2} \left[ N_{y} \log(2\pi) + \log\lvert S \rvert + v^{\top} S^{-1} v \right], \]

where v is the innovation and S the innovation covariance. Both log|S| and S⁻¹ v flow through :mod:gaussx dispatch — when S is a :class:gaussx.LowRankUpdate (the standard case), log-determinants use the matrix-determinant lemma and solves use Woodbury, both at O(Nₑ² Nᵧ + Nₑ³).

This is the training signal for differentiable data assimilation: the gradient of this scalar with respect to dynamics or observation-model parameters provides an end-to-end learning loop without ever forming S explicitly. See Decision D9 in the design docs.

Parameters:

Name	Type	Description	Default
innovation	Float[Array, ' N_y']	Innovation v = y − H x̄, shape (Nᵧ,).	required
innovation_cov	AbstractLinearOperator	Innovation covariance S = H P H ᵀ + R as a linear operator (typically :class:gaussx.LowRankUpdate).	required
solver	AbstractSolverStrategy | None	Optional :class:gaussx.AbstractSolverStrategy. None uses structural dispatch.	None

Returns:

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

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import log_likelihood
>>> innovation = jnp.array([0.5, -0.5])
>>> S = lx.DiagonalLinearOperator(jnp.ones(2))
>>> log_likelihood(innovation, S)
Array(-2.087877, dtype=float32)

Source code in src/filterax/_primitives/_likelihood.py

def log_likelihood(
    innovation: Float[Array, " N_y"],
    innovation_cov: lx.AbstractLinearOperator,
    *,
    solver: gaussx.AbstractSolverStrategy | None = None,
) -> Float[Array, ""]:
    r"""Gaussian log-probability of an innovation vector.

    $$
    \log p(y \mid \text{forecast}) = -\tfrac{1}{2} \left[
        N_{y} \log(2\pi) + \log\lvert S \rvert + v^{\top} S^{-1} v
    \right],
    $$

    where ``v`` is the innovation and ``S`` the innovation covariance.
    Both ``log|S|`` and ``S⁻¹ v`` flow through :mod:`gaussx` dispatch —
    when ``S`` is a :class:`gaussx.LowRankUpdate` (the standard case),
    log-determinants use the matrix-determinant lemma and solves use
    Woodbury, both at ``O(Nₑ² Nᵧ + Nₑ³)``.

    This is the training signal for differentiable data assimilation: the
    gradient of this scalar with respect to dynamics or observation-model
    parameters provides an end-to-end learning loop without ever forming
    ``S`` explicitly. See Decision D9 in the design docs.

    Args:
        innovation: Innovation ``v = y − H x̄``, shape ``(Nᵧ,)``.
        innovation_cov: Innovation covariance ``S = H P H ᵀ + R`` as a
            linear operator (typically :class:`gaussx.LowRankUpdate`).
        solver: Optional :class:`gaussx.AbstractSolverStrategy`. ``None``
            uses structural dispatch.

    Returns:
        Scalar log-probability.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import log_likelihood
        >>> innovation = jnp.array([0.5, -0.5])
        >>> S = lx.DiagonalLinearOperator(jnp.ones(2))
        >>> log_likelihood(innovation, S)
        Array(-2.087877, dtype=float32)
    """
    zero = jax.numpy.zeros_like(innovation)
    return gaussx.gaussian_log_prob(zero, innovation_cov, innovation, solver=solver)

innovation_covariance(obs_particles: Float[Array, 'N_e N_y'], obs_noise: lx.AbstractLinearOperator) -> gaussx.LowRankUpdate

Innovation covariance S = Cᴴᴴ + R as a low-rank update.

Wraps :func:gaussx.ensemble_covariance (Bessel-corrected) to build the low-rank factor U = (HX)′ᵀ / √(Nₑ − 1) and composes it with the structured R base. The result is a :class:gaussx.LowRankUpdate so all downstream linear algebra (solve, logdet, sample, …) takes advantage of the rank-(Nₑ − 1) ensemble term without densifying.

Parameters:

Name	Type	Description	Default
obs_particles	Float[Array, 'N_e N_y']	H applied to each ensemble member, shape (Nₑ, Nᵧ).	required
obs_noise	AbstractLinearOperator	Observation error covariance R.	required

Returns:

Type	Description
LowRankUpdate	class:gaussx.LowRankUpdate representing S.

Raises:

Type	Description
ValueError	if Nₑ < 2.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import innovation_covariance
>>> obs_particles = jnp.array([[0.0], [1.0]])  # Cᴴᴴ = 0.5
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(1))
>>> S = innovation_covariance(obs_particles, R)
>>> S.as_matrix()  # 0.5 + 0.1
Array([[0.6]], dtype=float32)

Source code in src/filterax/_primitives/_likelihood.py

def innovation_covariance(
    obs_particles: Float[Array, "N_e N_y"],
    obs_noise: lx.AbstractLinearOperator,
) -> gaussx.LowRankUpdate:
    r"""Innovation covariance ``S = Cᴴᴴ + R`` as a low-rank update.

    Wraps :func:`gaussx.ensemble_covariance` (Bessel-corrected) to build
    the low-rank factor ``U = (HX)′ᵀ / √(Nₑ − 1)`` and composes it with the
    structured ``R`` base. The result is a :class:`gaussx.LowRankUpdate`
    so all downstream linear algebra (solve, logdet, sample, …) takes
    advantage of the rank-``(Nₑ − 1)`` ensemble term without densifying.

    Args:
        obs_particles: ``H`` applied to each ensemble member, shape
            ``(Nₑ, Nᵧ)``.
        obs_noise: Observation error covariance ``R``.

    Returns:
        :class:`gaussx.LowRankUpdate` representing ``S``.

    Raises:
        ValueError: if ``Nₑ < 2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import innovation_covariance
        >>> obs_particles = jnp.array([[0.0], [1.0]])  # Cᴴᴴ = 0.5
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(1))
        >>> S = innovation_covariance(obs_particles, R)
        >>> S.as_matrix()  # 0.5 + 0.1
        Array([[0.6]], dtype=float32)
    """
    check_ensemble_size(obs_particles.shape[0])
    cov = gaussx.ensemble_covariance(obs_particles, bessel=True)
    return gaussx.LowRankUpdate(obs_noise, cov.U)

innovation_statistics(particles: Float[Array, 'N_e N_x'], obs: Float[Array, ' N_y'], obs_op: Callable[[Float[Array, ' N_x']], Float[Array, ' N_y']], obs_noise: lx.AbstractLinearOperator, *, solver: gaussx.AbstractSolverStrategy | None = None) -> InnovationStatistics

Innovation diagnostics for a forecast ensemble.

Returns v, S, the whitened residual S⁻¹ᐟ² v, and log p(y | forecast). The whitened residual is useful for χ² / Desroziers-style consistency checks (under correctly-specified S each component is unit-variance Gaussian).

The whitening solve goes through S.root_inv_decomposition() — structural dispatch in gaussx picks either Woodbury for :class:gaussx.LowRankUpdate or Cholesky for dense S, so the dense matrix is not materialised in the common case.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Forecast ensemble, shape (Nₑ, Nₓ).	required
obs	Float[Array, ' N_y']	Observation vector, shape (Nᵧ,).	required
obs_op	Callable[[Float[Array, ' N_x']], Float[Array, ' N_y']]	H applied to a single state vector. vmap'd over the ensemble internally.	required
obs_noise	AbstractLinearOperator	Observation error covariance R.	required
solver	AbstractSolverStrategy | None	Optional solver strategy (used for both logdet and solve).	None

Returns:

Type	Description
InnovationStatistics	class:InnovationStatistics mapping with innovation,
InnovationStatistics	innovation_cov, normalized_innovation, log_likelihood.

Raises:

Type	Description
ValueError	if Nₑ < 2.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import innovation_statistics
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0]])
>>> obs = jnp.array([0.5])
>>> R = lx.DiagonalLinearOperator(jnp.ones(1))
>>> stats = innovation_statistics(particles, obs, lambda x: x[:1], R)
>>> stats["innovation"]  # y − H x̄ = 0.5 − 0.5
Array([0.], dtype=float32)
>>> stats["log_likelihood"]
Array(-1.1216711, dtype=float32)

Source code in src/filterax/_primitives/_likelihood.py

def innovation_statistics(
    particles: Float[Array, "N_e N_x"],
    obs: Float[Array, " N_y"],
    obs_op: Callable[[Float[Array, " N_x"]], Float[Array, " N_y"]],
    obs_noise: lx.AbstractLinearOperator,
    *,
    solver: gaussx.AbstractSolverStrategy | None = None,
) -> InnovationStatistics:
    r"""Innovation diagnostics for a forecast ensemble.

    Returns ``v``, ``S``, the whitened residual ``S⁻¹ᐟ² v``, and
    ``log p(y | forecast)``. The whitened residual is useful for χ² /
    Desroziers-style consistency checks (under correctly-specified ``S``
    each component is unit-variance Gaussian).

    The whitening solve goes through ``S.root_inv_decomposition()`` —
    structural dispatch in gaussx picks either Woodbury for
    :class:`gaussx.LowRankUpdate` or Cholesky for dense ``S``, so the
    dense matrix is *not* materialised in the common case.

    Args:
        particles: Forecast ensemble, shape ``(Nₑ, Nₓ)``.
        obs: Observation vector, shape ``(Nᵧ,)``.
        obs_op: ``H`` applied to a single state vector. ``vmap``'d over
            the ensemble internally.
        obs_noise: Observation error covariance ``R``.
        solver: Optional solver strategy (used for both logdet and solve).

    Returns:
        :class:`InnovationStatistics` mapping with ``innovation``,
        ``innovation_cov``, ``normalized_innovation``, ``log_likelihood``.

    Raises:
        ValueError: if ``Nₑ < 2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import innovation_statistics
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0]])
        >>> obs = jnp.array([0.5])
        >>> R = lx.DiagonalLinearOperator(jnp.ones(1))
        >>> stats = innovation_statistics(particles, obs, lambda x: x[:1], R)
        >>> stats["innovation"]  # y − H x̄ = 0.5 − 0.5
        Array([0.], dtype=float32)
        >>> stats["log_likelihood"]
        Array(-1.1216711, dtype=float32)
    """
    check_ensemble_size(particles.shape[0])
    obs_particles = jax.vmap(obs_op)(particles)
    innovation = obs - ensemble_mean(obs_particles)

    S = innovation_covariance(obs_particles, obs_noise)
    log_prob = log_likelihood(innovation, S, solver=solver)

    # z = L⁻¹ v with L the Cholesky of the (small) Nᵧ × Nᵧ matrix.
    # Cholesky gives a true whitening (z has identity covariance) — the
    # gaussx structural-sqrt path returns a factor with L Lᵀ = S⁻¹ but
    # ``M v`` for that M is *not* a proper whitener. The (Nᵧ, Nᵧ) cost
    # is acceptable here: ``Nᵧ`` is the per-window observation count,
    # not the state dimension.
    L = jax.numpy.linalg.cholesky(S.as_matrix())
    normalized = jax.scipy.linalg.solve_triangular(L, innovation, lower=True)

    return InnovationStatistics(
        innovation=innovation,
        innovation_cov=S,
        normalized_innovation=normalized,
        log_likelihood=log_prob,
    )

Perturbed observations

Stochastic-EnKF observation perturbations, drawn with an explicit PRNG key.

filterax — differentiable ensemble Kalman filters and processes for JAX.

perturbed_observations(key: PRNGKeyArray, obs: Float[Array, ' N_y'], obs_noise: lx.AbstractLinearOperator, n_ensemble: int) -> Float[Array, 'N_e N_y']

Generate perturbed observations for the stochastic EnKF.

\[ y^{(j)}_{\text{pert}} = y + \varepsilon^{(j)}, \qquad \varepsilon^{(j)} \sim \mathcal{N}(0, R), \qquad j = 1, \dots, N_{e}. \]

Sampling delegates to :func:gaussx.sample_mvn. Every path draws ε = L z with z ~ 𝒩(0, I) and L Lᵀ = R exactly, and only the factor depends on the structure of R:

• Diagonal R: ε = √diag(R) ⊙ z. No matrix is ever formed; cost is O(Nₑ Nᵧ).
• gaussx block-diagonal, Kronecker, Toeplitz and low-rank-update R: a factor that keeps the structure, avoiding the cubic cost.
• Any other R: a dense symmetric square root. It costs O(Nᵧ³) once per call and also handles a singular R.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key consumed by this call.	required
obs	Float[Array, ' N_y']	Observation vector (Nᵧ,).	required
obs_noise	AbstractLinearOperator	Observation error covariance R as a linear operator.	required
n_ensemble	int	Number of perturbed copies Nₑ.	required

Returns:

Type	Description
Float[Array, 'N_e N_y']	Perturbed observation matrix of shape (Nₑ, Nᵧ).

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import perturbed_observations
>>> obs = jnp.array([1.0, 2.0])
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> y_pert = perturbed_observations(jax.random.key(0), obs, R, 4)
>>> y_pert.shape
(4, 2)

Reference

Burgers, G., van Leeuwen, P. J., & Evensen, G. (1998). Analysis scheme in the ensemble Kalman filter. Mon. Wea. Rev., 126, 1719–1724.

Source code in src/filterax/_primitives/_perturbations.py

def perturbed_observations(
    key: PRNGKeyArray,
    obs: Float[Array, " N_y"],
    obs_noise: lx.AbstractLinearOperator,
    n_ensemble: int,
) -> Float[Array, "N_e N_y"]:
    r"""Generate perturbed observations for the stochastic EnKF.

    $$
    y^{(j)}_{\text{pert}} = y + \varepsilon^{(j)},
    \qquad
    \varepsilon^{(j)} \sim \mathcal{N}(0, R),
    \qquad
    j = 1, \dots, N_{e}.
    $$

    Sampling delegates to :func:`gaussx.sample_mvn`. Every path draws
    ``ε = L z`` with ``z ~ 𝒩(0, I)`` and ``L Lᵀ = R`` exactly, and only the
    factor depends on the structure of ``R``:

    * Diagonal ``R``: ``ε = √diag(R) ⊙ z``. No matrix is ever formed; cost
      is ``O(Nₑ Nᵧ)``.
    * ``gaussx`` block-diagonal, Kronecker, Toeplitz and low-rank-update
      ``R``: a factor that keeps the structure, avoiding the cubic cost.
    * Any other ``R``: a dense symmetric square root. It costs ``O(Nᵧ³)``
      once per call and also handles a singular ``R``.

    Args:
        key: PRNG key consumed by this call.
        obs: Observation vector ``(Nᵧ,)``.
        obs_noise: Observation error covariance ``R`` as a linear operator.
        n_ensemble: Number of perturbed copies ``Nₑ``.

    Returns:
        Perturbed observation matrix of shape ``(Nₑ, Nᵧ)``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import perturbed_observations
        >>> obs = jnp.array([1.0, 2.0])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> y_pert = perturbed_observations(jax.random.key(0), obs, R, 4)
        >>> y_pert.shape
        (4, 2)

    Reference:
        Burgers, G., van Leeuwen, P. J., & Evensen, G. (1998). *Analysis
        scheme in the ensemble Kalman filter.* Mon. Wea. Rev., 126,
        1719–1724.
    """
    # y⁽ʲ⁾ = y + L z⁽ʲ⁾, L Lᵀ = R
    return gaussx.sample_mvn(obs, obs_noise, key=key, num_samples=n_ensemble)