Filters

Sequential ensemble filters come in two layers:

• Layer 1 — analysis steps (filterax.filters.*). A one-shot posterior update: forecast ensemble + observation + observation operator + noise covariance in, analysed ensemble out (AnalysisResult). All implement AbstractSequentialFilter and share the gaussx-backed ensemble-statistics primitives; they differ only in how the ensemble is updated.
• Layer 2 — assimilation models (top-level filterax.*, same names). The full forecast → analyse → inflate loop: each composes an AbstractDynamics forecast (via eqx.filter_vmap), the matching L1 analysis step, and an optional AbstractInflator into a single assimilate() call over a sequence of observation windows, returning an AssimilationResult.

So filterax.filters.ETKF is the one-shot analysis step, while filterax.ETKF is the loop built around it. Drop down to Layer 1 when the standard cycle does not fit (custom inflation schedules, pipekit cycling, differentiable training).

ETKF analysis core

filterax's ETKF keeps its own rank-N_y QR + eigh analysis core — a thin QR of the observation anomalies followed by a small (N_y, N_y) symmetric eigendecomposition — rather than delegating to gaussx.etkf_transform. The decomposition is of a symmetric PSD matrix, which keeps the transform differentiability-safe under jax.grad (see Differentiable training).

For the specialised deterministic variants (ESTKF, ETKF_Livings, EnSRF_Serial) see Advanced filters.

Layer-1 analysis steps

The everyday defaults: symmetric square-root ETKF, observation-space localized LETKF, serial-free EnSRF, and the perturbed-observations stochastic EnKF.

Layer-1 sequential filter components.

Re-exports the L1 filter classes from filterax._filters._sequential so the canonical filterax.filters.ETKF() access path matches the design docs.

Wave 4 additions: * :class:ETKF_Livings — ETKF + mean-preserving random rotation * :class:EnSRF_Serial — scalar serial sqrt updates * :class:ESTKF — error-subspace transform * :class:SquareRootKF — parametric Cholesky filter (lives in filterax._filters._parametric for namespacing; re-exported here for convenience).

ETKF

Bases: AbstractSequentialFilter

Ensemble Transform Kalman Filter (Bishop, Etherton & Majumdar 2001).

Deterministic square-root update working entirely in the Nₑ-dimensional ensemble subspace. Define the transform precision in the ensemble space and its symmetric inverse square root:

\[ \tilde{C} = (N_e - 1) I + Y' R^{-1} Y'^{\top}, \qquad T = \tilde{C}^{-1}, \qquad W_a = \sqrt{(N_e - 1)\, T} \]

Mean and perturbation weights:

\[ \bar{w}_a = T\, Y' R^{-1} v, \qquad X_a = \bar{x} \mathbf{1}^{\top} + X'^{\top} \big(\bar{w}_a \mathbf{1}^{\top} + W_a\big) \]

The eigendecomposition is taken on the symmetrised C̃ with a positive eigenvalue floor to keep the square root well-defined under floating-point error.

Complexity O(Nₑ² Nᵧ + Nₑ³) per analysis — the inner solve is routed through gaussx so structured R does not densify.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax.filters import ETKF
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> result = ETKF().analysis(
...     particles, jnp.array([1.0, 0.5]), lambda x: x, R
... )
>>> result.particles.shape
(3, 2)

Source code in src/filterax/_filters/_sequential.py

class ETKF(AbstractSequentialFilter, strict=True):
    r"""Ensemble Transform Kalman Filter (Bishop, Etherton & Majumdar 2001).

    Deterministic square-root update working entirely in the
    ``Nₑ``-dimensional ensemble subspace. Define the transform precision
    in the ensemble space and its symmetric inverse square root:

    $$
    \tilde{C} = (N_e - 1) I + Y' R^{-1} Y'^{\top}, \qquad
    T = \tilde{C}^{-1}, \qquad
    W_a = \sqrt{(N_e - 1)\, T}
    $$

    Mean and perturbation weights:

    $$
    \bar{w}_a = T\, Y' R^{-1} v, \qquad
    X_a = \bar{x} \mathbf{1}^{\top}
        + X'^{\top} \big(\bar{w}_a \mathbf{1}^{\top} + W_a\big)
    $$

    The eigendecomposition is taken on the symmetrised ``C̃`` with a
    positive eigenvalue floor to keep the square root well-defined under
    floating-point error.

    Complexity ``O(Nₑ² Nᵧ + Nₑ³)`` per analysis — the inner solve is
    routed through gaussx so structured ``R`` does not densify.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax.filters import ETKF
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> result = ETKF().analysis(
        ...     particles, jnp.array([1.0, 0.5]), lambda x: x, R
        ... )
        >>> result.particles.shape
        (3, 2)
    """

    def analysis(
        self,
        particles: Float[Array, "N_e N_x"],
        obs: Float[Array, " N_y"],
        obs_op: AbstractObsOperator | ObsCallable,
        obs_noise: lx.AbstractLinearOperator,
        **_: Any,
    ) -> AnalysisResult:
        N_e = particles.shape[0]
        check_ensemble_size(N_e)
        obs_particles = _apply_obs_op(obs_op, particles)

        anom = ensemble_anomalies(particles)  # X′ ∈ ℝ^{Nₑ×Nₓ}
        obs_anom = ensemble_anomalies(obs_particles)  # Y′ ∈ ℝ^{Nₑ×Nᵧ}
        x_bar = ensemble_mean(particles)
        innovation = obs - ensemble_mean(obs_particles)

        # R⁻¹ Y′ — gaussx structural dispatch (no dense Nᵧ × Nᵧ matrix
        # for DiagonalLinearOperator, Toeplitz, etc.).
        R_inv_Y = gaussx.solve_rows(obs_noise, obs_anom)  # (Nₑ, Nᵧ)

        # Rank-Nᵧ spectrum of Y′ R⁻¹ Y′ᵀ via QR + small eigh; avoids the
        # gradient-NaN trap of an eigh on the full (Nₑ, Nₑ) C̃ matrix.
        eigvals, U_y = _etkf_inner_spectrum(obs_anom, R_inv_Y)

        # T = C̃⁻¹ = (1/(Nₑ−1)) I + U_y diag(inv_diff) U_yᵀ.
        inv_base = jnp.asarray(1.0 / (N_e - 1), dtype=anom.dtype)
        inv_diff = 1.0 / ((N_e - 1) + eigvals) - inv_base
        # Wₐ = √((Nₑ−1) C̃⁻¹) = I + U_y diag(sqrt_diff) U_yᵀ.
        sqrt_base = jnp.asarray(1.0, dtype=anom.dtype)
        sqrt_diff = jnp.sqrt((N_e - 1) / ((N_e - 1) + eigvals)) - sqrt_base

        # Mean: w̄ = T (Y′ R⁻¹ v); then mean_correction = w̄ᵀ X′.
        Y_R_inv_v = einx.dot("e y, y -> e", R_inv_Y, innovation)
        w_bar = _apply_ctilde_func(U_y, inv_diff, inv_base, Y_R_inv_v)
        mean_correction = einx.dot("e, e x -> x", w_bar, anom)

        # Perturbations: Wₐ X′.
        pert_correction = _apply_ctilde_func(U_y, sqrt_diff, sqrt_base, anom)

        particles_a = (
            einx.add("x, x -> x", x_bar, mean_correction)[None, :] + pert_correction
        )

        S = innovation_covariance(obs_particles, obs_noise)
        log_p = log_likelihood(innovation, S)
        return AnalysisResult(particles=particles_a, log_likelihood=log_p)

LETKF

Bases: AbstractSequentialFilter

Local Ensemble Transform Kalman Filter (Hunt, Kostelich & Szunyogh 2007).

Runs an independent ETKF analysis at each state grid point using only observations whose distance to that grid point is ≤ radius (R-localization). For grid point i:

1. Compute distances dᵢₖ = ‖xᵢ − yₖ‖₂ and tapers ρᵢₖ = ρ(dᵢₖ / r) via taper_fn (default Gaspari-Cohn).
2. Hard cutoff at r — observations with dᵢₖ > r are dropped entirely. (Gaspari-Cohn has compact support out to 2r so a purely taper-based mask would let every nonzero-tapered observation participate. We use the explicit distance cutoff to match the documented selection rule.)
3. Inflate the local observation variance by 1/ρᵢₖ: R⁻¹_loc[k] = ρᵢₖ / Rₖₖ (zero when ρᵢₖ = 0).
4. Solve the local Nₑ × Nₑ ETKF transform and write back the analysis at grid point i.

R-localization assumes R is diagonal (Hunt et al. 2007 §2). We accept :class:lineax.DiagonalLinearOperator directly via :func:lineax.diagonal (no densification); anything else raises :class:NotImplementedError.

Local analyses are embarrassingly parallel; we jax.vmap over state grid points and accumulate the per-point (w̄ᵢ, Wᵢ) weights in a single call.

Attributes:

Name	Type	Description
radius	float	Localization half-width r. Observations beyond r are excluded entirely.
taper_fn	Callable[[Float[Array, '...'], float], Float[Array, '...']]	Distance-to-weight function with signature (distances, radius) -> weights. Defaults to :func:filterax.gaspari_cohn.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax.filters import LETKF
>>> particles = jax.random.normal(jax.random.key(0), (4, 3))
>>> state_coords = jnp.arange(3.0)[:, None]  # 1-D grid
>>> obs_coords = jnp.array([[0.0], [2.0]])
>>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(2))
>>> result = LETKF(radius=1.5).analysis(
...     particles,
...     jnp.array([0.1, -0.2]),
...     lambda x: x[::2],  # observe grid points 0 and 2
...     R,
...     state_coords=state_coords,
...     obs_coords=obs_coords,
... )
>>> result.particles.shape
(4, 3)

Source code in src/filterax/_filters/_sequential.py

class LETKF(AbstractSequentialFilter, strict=True):
    r"""Local Ensemble Transform Kalman Filter (Hunt, Kostelich & Szunyogh 2007).

    Runs an independent ETKF analysis at each state grid point using only
    observations whose distance to that grid point is ``≤ radius``
    (R-localization). For grid point ``i``:

    1. Compute distances ``dᵢₖ = ‖xᵢ − yₖ‖₂`` and tapers
       ``ρᵢₖ = ρ(dᵢₖ / r)`` via ``taper_fn`` (default Gaspari-Cohn).
    2. **Hard cutoff at ``r``** — observations with ``dᵢₖ > r`` are
       dropped entirely. (Gaspari-Cohn has compact support out to ``2r``
       so a purely taper-based mask would let *every* nonzero-tapered
       observation participate. We use the explicit distance cutoff to
       match the documented selection rule.)
    3. Inflate the local observation variance by ``1/ρᵢₖ``:
       ``R⁻¹_loc[k] = ρᵢₖ / Rₖₖ`` (zero when ``ρᵢₖ = 0``).
    4. Solve the local ``Nₑ × Nₑ`` ETKF transform and write back the
       analysis at grid point ``i``.

    R-localization assumes ``R`` is **diagonal** (Hunt et al. 2007 §2).
    We accept :class:`lineax.DiagonalLinearOperator` directly via
    :func:`lineax.diagonal` (no densification); anything else raises
    :class:`NotImplementedError`.

    Local analyses are embarrassingly parallel; we ``jax.vmap`` over
    state grid points and accumulate the per-point ``(w̄ᵢ, Wᵢ)`` weights
    in a single call.

    Attributes:
        radius: Localization half-width ``r``. Observations beyond ``r``
            are excluded entirely.
        taper_fn: Distance-to-weight function with signature
            ``(distances, radius) -> weights``. Defaults to
            :func:`filterax.gaspari_cohn`.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax.filters import LETKF
        >>> particles = jax.random.normal(jax.random.key(0), (4, 3))
        >>> state_coords = jnp.arange(3.0)[:, None]  # 1-D grid
        >>> obs_coords = jnp.array([[0.0], [2.0]])
        >>> R = lx.DiagonalLinearOperator(0.5 * jnp.ones(2))
        >>> result = LETKF(radius=1.5).analysis(
        ...     particles,
        ...     jnp.array([0.1, -0.2]),
        ...     lambda x: x[::2],  # observe grid points 0 and 2
        ...     R,
        ...     state_coords=state_coords,
        ...     obs_coords=obs_coords,
        ... )
        >>> result.particles.shape
        (4, 3)
    """

    radius: float = eqx.field(static=True)
    taper_fn: Callable[[Float[Array, "..."], float], Float[Array, "..."]] = eqx.field(
        static=True, default=gaspari_cohn
    )

    def analysis(
        self,
        particles: Float[Array, "N_e N_x"],
        obs: Float[Array, " N_y"],
        obs_op: AbstractObsOperator | ObsCallable,
        obs_noise: lx.AbstractLinearOperator,
        *,
        state_coords: Float[Array, "N_x D"] | None = None,
        obs_coords: Float[Array, "N_y D"] | None = None,
        **_: Any,
    ) -> AnalysisResult:
        N_e, N_x = particles.shape
        check_ensemble_size(N_e)
        if state_coords is None or obs_coords is None:
            raise ValueError(
                "LETKF.analysis requires state_coords and obs_coords keyword "
                "arguments — pass them through the L2 model or the call site."
            )
        if state_coords.shape[0] != N_x:
            raise ValueError(
                "state_coords must have one row per state component; got "
                f"{state_coords.shape[0]} vs N_x={N_x}."
            )
        if obs_coords.shape[0] != obs.shape[0]:
            raise ValueError(
                "obs_coords must have one row per observation; got "
                f"{obs_coords.shape[0]} vs N_y={obs.shape[0]}."
            )
        if not isinstance(obs_noise, lx.DiagonalLinearOperator):
            raise NotImplementedError(
                "LETKF R-localization assumes diagonal observation noise "
                "(Hunt et al. 2007 §2). Pass a lineax.DiagonalLinearOperator "
                "or decorrelate observations beforehand."
            )

        obs_particles = _apply_obs_op(obs_op, particles)  # Y ∈ ℝ^{Nₑ×Nᵧ}
        obs_anom = ensemble_anomalies(obs_particles)
        anom = ensemble_anomalies(particles)
        x_bar = ensemble_mean(particles)
        innovation = obs - ensemble_mean(obs_particles)

        # Diagonal of R extracted without ever materialising the dense matrix.
        R_diag = lx.diagonal(obs_noise)
        radius = self.radius

        def per_point(
            state_point: Float[Array, " D"],
        ) -> tuple[Float[Array, " N_e"], Float[Array, "N_e N_e"]]:
            """Local ETKF weights at one grid point.

            Returns ``(w̄ᵢ, Wᵢ)`` with shapes ``(Nₑ,)`` and ``(Nₑ, Nₑ)``.
            """
            # ‖xᵢ − yₖ‖₂ — Euclidean distance to every observation.
            diff = einx.subtract("y d, d -> y d", obs_coords, state_point)
            dist = jnp.sqrt(einx.sum("y d -> y", diff * diff))

            # Distance-based hard cutoff at the radius; taper inside.
            within = dist <= radius
            rho = jnp.where(within, self.taper_fn(dist, radius), 0.0)

            # R⁻¹_loc = ρ ⊙ R⁻¹ (diagonal, no matrix materialisation).
            inv_R_loc = jnp.where(within, rho / R_diag, 0.0)  # (Nᵧ,)

            # Local ETKF in ensemble space.
            # R⁻¹_loc Y′ — broadcasted elementwise over the obs axis.
            R_inv_Y = einx.multiply("e y, y -> e y", obs_anom, inv_R_loc)
            eigvals_loc, U_y_loc = _etkf_inner_spectrum(obs_anom, R_inv_Y)

            inv_base = jnp.asarray(1.0 / (N_e - 1), dtype=anom.dtype)
            inv_diff = 1.0 / ((N_e - 1) + eigvals_loc) - inv_base
            sqrt_base = jnp.asarray(1.0, dtype=anom.dtype)
            sqrt_diff = jnp.sqrt((N_e - 1) / ((N_e - 1) + eigvals_loc)) - sqrt_base

            # w̄ᵢ = T (Y′ R⁻¹_loc v); T = (1/(Nₑ−1)) I + rank-Nᵧ correction.
            Y_R_inv_v = einx.dot("e y, y -> e", R_inv_Y, innovation)
            w_bar_i = _apply_ctilde_func(U_y_loc, inv_diff, inv_base, Y_R_inv_v)
            # Wᵢ = I + U_y_loc diag(sqrt_diff) U_y_locᵀ — materialise the
            # (Nₑ, Nₑ) transform so the outer ``einx`` can fuse over grid
            # points without re-invoking per-point spectra.
            scaled_U = einx.multiply("e a, a -> e a", U_y_loc, sqrt_diff)
            correction = einx.dot("e a, f a -> e f", scaled_U, U_y_loc)
            W_i = jnp.eye(N_e, dtype=anom.dtype) + correction
            return w_bar_i, W_i

        # vmap over state grid points → per-point (w̄ᵢ, Wᵢ).
        w_bars, Ws = jax.vmap(per_point)(state_coords)  # (Nₓ, Nₑ), (Nₓ, Nₑ, Nₑ)

        # Apply per-point updates with anomaly columns shared globally:
        # ``x_a[i] = x̄[i] + Σₖ w̄[i, k] X′[k, i]`` for the mean and
        # ``x⁽ʲ⁾_a[i] = x̄_a[i] + Σₖ W[i, j, k] X′[k, i]`` for members.
        mean_correction = einx.dot("x e, e x -> x", w_bars, anom)
        pert_correction = einx.dot("x j k, k x -> j x", Ws, anom)

        mean_a = einx.add("x, x -> x", x_bar, mean_correction)
        particles_a = mean_a[None, :] + pert_correction

        S = innovation_covariance(obs_particles, obs_noise)
        log_p = log_likelihood(innovation, S)
        return AnalysisResult(particles=particles_a, log_likelihood=log_p)

EnSRF

Bases: AbstractSequentialFilter

Ensemble Square Root Filter (Whitaker & Hamill 2002, batch form).

Separate mean and perturbation updates that produce the exact analysis covariance Pₐ = (I − KH) P_f in the linear-Gaussian limit, without perturbed observations:

• Mean from the Kalman gain: x̄ₐ = x̄_f + K (y − ȳ_f).
• Perturbations via the symmetric ETKF transform applied to anomalies.

Why "EnSRF" and not "ETKF" then? In batch form (this class) the two are numerically equivalent on the analysis ensemble — both choices of ensemble square root produce the same posterior covariance and the symmetric sqrt is the only PSD- and mean-preserving choice for rows-as-members anomalies (Tippett et al. 2003). The classical Whitaker & Hamill (2002) distinction — a one-sided reduced gain on perturbations — only matters when observations are assimilated one at a time. That serial variant lives in Wave 4 as filterax.filters.EnSRF_Serial.

For now, :class:EnSRF gives users the more familiar API (mean from K, perturbations from a sqrt-of-T) while remaining mathematically equivalent to :class:ETKF. Pick whichever phrasing is more idiomatic for your domain.

Complexity matches ETKF: O(Nₑ² Nᵧ + Nₑ³).

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax.filters import EnSRF
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> result = EnSRF().analysis(
...     particles, jnp.array([1.0, 0.5]), lambda x: x, R
... )
>>> result.particles.shape
(3, 2)

Source code in src/filterax/_filters/_sequential.py

class EnSRF(AbstractSequentialFilter, strict=True):
    r"""Ensemble Square Root Filter (Whitaker & Hamill 2002, batch form).

    Separate mean and perturbation updates that produce the exact
    analysis covariance ``Pₐ = (I − KH) P_f`` in the linear-Gaussian
    limit, without perturbed observations:

    * Mean from the Kalman gain: ``x̄ₐ = x̄_f + K (y − ȳ_f)``.
    * Perturbations via the symmetric ETKF transform applied to anomalies.

    Why "EnSRF" and not "ETKF" then? In *batch* form (this class) the two
    are numerically equivalent on the analysis ensemble — both choices
    of ensemble square root produce the same posterior covariance and
    the symmetric sqrt is the only PSD- and mean-preserving choice for
    rows-as-members anomalies (Tippett et al. 2003). The classical
    Whitaker & Hamill (2002) distinction — a *one-sided* reduced gain on
    perturbations — only matters when observations are assimilated one at
    a time. That serial variant lives in Wave 4 as
    ``filterax.filters.EnSRF_Serial``.

    For now, :class:`EnSRF` gives users the more familiar API (mean from
    ``K``, perturbations from a sqrt-of-``T``) while remaining
    mathematically equivalent to :class:`ETKF`. Pick whichever phrasing
    is more idiomatic for your domain.

    Complexity matches ETKF: ``O(Nₑ² Nᵧ + Nₑ³)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax.filters import EnSRF
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> result = EnSRF().analysis(
        ...     particles, jnp.array([1.0, 0.5]), lambda x: x, R
        ... )
        >>> result.particles.shape
        (3, 2)
    """

    def analysis(
        self,
        particles: Float[Array, "N_e N_x"],
        obs: Float[Array, " N_y"],
        obs_op: AbstractObsOperator | ObsCallable,
        obs_noise: lx.AbstractLinearOperator,
        **_: Any,
    ) -> AnalysisResult:
        N_e = particles.shape[0]
        check_ensemble_size(N_e)
        obs_particles = _apply_obs_op(obs_op, particles)

        x_bar = ensemble_mean(particles)
        anom = ensemble_anomalies(particles)  # (Nₑ, Nₓ)
        obs_anom = ensemble_anomalies(obs_particles)  # (Nₑ, Nᵧ)
        innovation = obs - ensemble_mean(obs_particles)

        # Mean update via the standard Bessel-corrected gain.
        K = kalman_gain(particles, obs_particles, obs_noise)
        mean_a = einx.add(
            "x, x -> x",
            x_bar,
            einx.dot("x y, y -> x", K, innovation),
        )

        # Perturbation update via the symmetric ETKF sqrt
        # ``Wₐ = U diag(√((Nₑ−1)/λ)) Uᵀ`` applied to anomalies. Symmetric
        # because that is the only sqrt of ``T`` that maps mean-zero
        # anomalies to mean-zero anomalies, leaving the analysis mean
        # equal to ``x̄ + K v`` (see Tippett et al. 2003 §3).
        R_inv_Y = gaussx.solve_rows(obs_noise, obs_anom)
        eigvals, U_y = _etkf_inner_spectrum(obs_anom, R_inv_Y)
        sqrt_base = jnp.asarray(1.0, dtype=anom.dtype)
        sqrt_diff = jnp.sqrt((N_e - 1) / ((N_e - 1) + eigvals)) - sqrt_base
        pert_correction = _apply_ctilde_func(U_y, sqrt_diff, sqrt_base, anom)

        particles_a = mean_a[None, :] + pert_correction

        S = innovation_covariance(obs_particles, obs_noise)
        log_p = log_likelihood(innovation, S)
        return AnalysisResult(particles=particles_a, log_likelihood=log_p)

StochasticEnKF

Bases: AbstractSequentialFilter

Stochastic Ensemble Kalman Filter (Evensen 1994).

Perturbed-observation update — each member sees an independent draw from the observation noise:

\[ \epsilon^{(j)} \sim \mathcal{N}(0, R), \qquad x^{(j)}_a = x^{(j)}_f + K \big(y + \epsilon^{(j)} - H x^{(j)}_f\big) \]

Simple and robust; the perturbations introduce extra Monte Carlo variance whose magnitude scales as 1/√Nₑ.

PRNG key handling. The key stored on the filter is the default source of randomness. analysis(..., key=subkey) overrides it for a single call — this is how the Layer-2 :class:filterax.StochasticEnKF threads a fresh sub-key per assimilation window so consecutive perturbations stay independent. If you call analysis directly, pass a freshly split key each cycle or you will get identical observation perturbations every window.

Attributes:

Name	Type	Description
key	PRNGKeyArray	Default PRNG key used when no key is supplied in the analysis kwargs.

Examples:

>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import lineax as lx
>>> from filterax.filters import StochasticEnKF
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> filt = StochasticEnKF(key=jr.key(0))
>>> result = filt.analysis(
...     particles, jnp.array([1.0, 0.5]), lambda x: x, R, key=jr.key(1)
... )
>>> result.particles.shape
(3, 2)

Source code in src/filterax/_filters/_sequential.py

class StochasticEnKF(AbstractSequentialFilter, strict=True):
    r"""Stochastic Ensemble Kalman Filter (Evensen 1994).

    Perturbed-observation update — each member sees an independent draw
    from the observation noise:

    $$
    \epsilon^{(j)} \sim \mathcal{N}(0, R), \qquad
    x^{(j)}_a = x^{(j)}_f + K \big(y + \epsilon^{(j)} - H x^{(j)}_f\big)
    $$

    Simple and robust; the perturbations introduce extra Monte Carlo
    variance whose magnitude scales as ``1/√Nₑ``.

    PRNG key handling. The key stored on the filter is the *default*
    source of randomness. ``analysis(..., key=subkey)`` overrides it for
    a single call — this is how the Layer-2 :class:`filterax.StochasticEnKF`
    threads a fresh sub-key per assimilation window so consecutive
    perturbations stay independent. If you call ``analysis`` directly,
    pass a freshly split key each cycle or you will get identical
    observation perturbations every window.

    Attributes:
        key: Default PRNG key used when no ``key`` is supplied in the
            ``analysis`` kwargs.

    Examples:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from filterax.filters import StochasticEnKF
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> filt = StochasticEnKF(key=jr.key(0))
        >>> result = filt.analysis(
        ...     particles, jnp.array([1.0, 0.5]), lambda x: x, R, key=jr.key(1)
        ... )
        >>> result.particles.shape
        (3, 2)
    """

    key: PRNGKeyArray

    def __init__(self, key: PRNGKeyArray | int = 0):
        if isinstance(key, int):
            key = jr.PRNGKey(key)
        self.key = key

    def analysis(
        self,
        particles: Float[Array, "N_e N_x"],
        obs: Float[Array, " N_y"],
        obs_op: AbstractObsOperator | ObsCallable,
        obs_noise: lx.AbstractLinearOperator,
        *,
        key: PRNGKeyArray | None = None,
        **_: Any,
    ) -> AnalysisResult:
        r"""Stochastic EnKF analysis step.

        Args:
            particles: Forecast ensemble ``(Nₑ, Nₓ)``.
            obs: Observation ``y`` of shape ``(Nᵧ,)``.
            obs_op: ``H`` applied to a single state vector.
            obs_noise: Observation covariance ``R``.
            key: Per-step PRNG key. When ``None``, falls back to
                ``self.key`` — the same draw is repeated across calls
                unless callers split externally.

        Returns:
            :class:`AnalysisResult` with the perturbed-observation
            posterior ensemble and the Gaussian log-likelihood of the
            innovation under ``S = Cᴴᴴ + R``.
        """
        N_e = particles.shape[0]
        check_ensemble_size(N_e)
        step_key = self.key if key is None else key

        obs_particles = _apply_obs_op(obs_op, particles)  # Y ∈ ℝ^{Nₑ×Nᵧ}

        # K = Cˣᴴ (Cᴴᴴ + R)⁻¹ via gaussx-dispatched Woodbury solve.
        K = kalman_gain(particles, obs_particles, obs_noise)  # (Nₓ, Nᵧ)

        # y_pert[j] = y + ε⁽ʲ⁾ with ε⁽ʲ⁾ ~ 𝒩(0, R) — sample structure-
        # aware so diagonal R never densifies.
        y_pert = perturbed_observations(step_key, obs, obs_noise, N_e)
        innovations = y_pert - obs_particles  # v⁽ʲ⁾ = y + ε⁽ʲ⁾ − H x⁽ʲ⁾
        # x⁽ʲ⁾_a = x⁽ʲ⁾ + K v⁽ʲ⁾, summed over the obs axis.
        increment = einx.dot("e y, x y -> e x", innovations, K)
        particles_a = einx.add("e x, e x -> e x", particles, increment)

        S = innovation_covariance(obs_particles, obs_noise)
        log_p = log_likelihood(obs - ensemble_mean(obs_particles), S)
        return AnalysisResult(particles=particles_a, log_likelihood=log_p)

analysis(particles: Float[Array, 'N_e N_x'], obs: Float[Array, ' N_y'], obs_op: AbstractObsOperator | ObsCallable, obs_noise: lx.AbstractLinearOperator, *, key: PRNGKeyArray | None = None, **_: Any) -> AnalysisResult

Stochastic EnKF analysis step.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Forecast ensemble (Nₑ, Nₓ).	required
obs	Float[Array, ' N_y']	Observation y of shape (Nᵧ,).	required
obs_op	AbstractObsOperator | ObsCallable	H applied to a single state vector.	required
obs_noise	AbstractLinearOperator	Observation covariance R.	required
key	PRNGKeyArray | None	Per-step PRNG key. When None, falls back to self.key — the same draw is repeated across calls unless callers split externally.	None

Returns:

Type	Description
AnalysisResult	class:AnalysisResult with the perturbed-observation
AnalysisResult	posterior ensemble and the Gaussian log-likelihood of the
AnalysisResult	innovation under S = Cᴴᴴ + R.

Source code in src/filterax/_filters/_sequential.py

def analysis(
    self,
    particles: Float[Array, "N_e N_x"],
    obs: Float[Array, " N_y"],
    obs_op: AbstractObsOperator | ObsCallable,
    obs_noise: lx.AbstractLinearOperator,
    *,
    key: PRNGKeyArray | None = None,
    **_: Any,
) -> AnalysisResult:
    r"""Stochastic EnKF analysis step.

    Args:
        particles: Forecast ensemble ``(Nₑ, Nₓ)``.
        obs: Observation ``y`` of shape ``(Nᵧ,)``.
        obs_op: ``H`` applied to a single state vector.
        obs_noise: Observation covariance ``R``.
        key: Per-step PRNG key. When ``None``, falls back to
            ``self.key`` — the same draw is repeated across calls
            unless callers split externally.

    Returns:
        :class:`AnalysisResult` with the perturbed-observation
        posterior ensemble and the Gaussian log-likelihood of the
        innovation under ``S = Cᴴᴴ + R``.
    """
    N_e = particles.shape[0]
    check_ensemble_size(N_e)
    step_key = self.key if key is None else key

    obs_particles = _apply_obs_op(obs_op, particles)  # Y ∈ ℝ^{Nₑ×Nᵧ}

    # K = Cˣᴴ (Cᴴᴴ + R)⁻¹ via gaussx-dispatched Woodbury solve.
    K = kalman_gain(particles, obs_particles, obs_noise)  # (Nₓ, Nᵧ)

    # y_pert[j] = y + ε⁽ʲ⁾ with ε⁽ʲ⁾ ~ 𝒩(0, R) — sample structure-
    # aware so diagonal R never densifies.
    y_pert = perturbed_observations(step_key, obs, obs_noise, N_e)
    innovations = y_pert - obs_particles  # v⁽ʲ⁾ = y + ε⁽ʲ⁾ − H x⁽ʲ⁾
    # x⁽ʲ⁾_a = x⁽ʲ⁾ + K v⁽ʲ⁾, summed over the obs axis.
    increment = einx.dot("e y, x y -> e x", innovations, K)
    particles_a = einx.add("e x, e x -> e x", particles, increment)

    S = innovation_covariance(obs_particles, obs_noise)
    log_p = log_likelihood(obs - ensemble_mean(obs_particles), S)
    return AnalysisResult(particles=particles_a, log_likelihood=log_p)

Parametric square-root filter

A non-ensemble reference: propagates (μ, S) with P = S Sᵀ instead of particles, wrapping gaussx's parallel Kalman filter. Use it for linear-Gaussian baselines and twin-experiment ground truths.

Layer-1 sequential filter components.

Re-exports the L1 filter classes from filterax._filters._sequential so the canonical filterax.filters.ETKF() access path matches the design docs.

Wave 4 additions: * :class:ETKF_Livings — ETKF + mean-preserving random rotation * :class:EnSRF_Serial — scalar serial sqrt updates * :class:ESTKF — error-subspace transform * :class:SquareRootKF — parametric Cholesky filter (lives in filterax._filters._parametric for namespacing; re-exported here for convenience).

SquareRootKF

Bases: Module

Square-root parametric Kalman filter (Maybeck 1979 / gaussx).

Propagates a Gaussian belief through a linear-Gaussian model:

\[ \begin{aligned} x_n &= \Phi x_{n-1} + w_n, & w_n &\sim \mathcal{N}(0, Q) \\ y_n &= H x_n + v_n, & v_n &\sim \mathcal{N}(0, R) \end{aligned} \]

The forecast and analysis steps are both rewritten in square-root form on the Cholesky factor S of P so the covariance stays PSD by construction. Internally this delegates to :func:gaussx.parallel_kalman_filter (form='sqrt'), which runs Blelloch's parallel scan along the time axis — cost O(log T · (Nₓ² Nᵧ + Nₓ³)) on hardware that can exploit the parallelism.

The constructor stores the (time-invariant) Φ, H, Q, R and lets the user push observations through :meth:filter. For time-varying systems pass leading-time-axis arrays directly to :func:gaussx.parallel_kalman_filter and skip this wrapper.

Attributes:

Name	Type	Description
transition	Float[Array, 'N N'] | AbstractLinearOperator	State transition operator Φ (matrix or :class:lineax.AbstractLinearOperator).
obs_model	Float[Array, 'M N'] | AbstractLinearOperator	Observation operator H.
process_noise	Float[Array, 'N N'] | AbstractLinearOperator	Process-noise covariance Q.
obs_noise	Float[Array, 'M M'] | AbstractLinearOperator	Observation-noise covariance R.

Source code in src/filterax/_filters/_parametric.py

class SquareRootKF(eqx.Module, strict=True):
    r"""Square-root parametric Kalman filter (Maybeck 1979 / gaussx).

    Propagates a Gaussian belief through a linear-Gaussian model:

    $$
    \begin{aligned}
    x_n &= \Phi x_{n-1} + w_n, & w_n &\sim \mathcal{N}(0, Q) \\
    y_n &= H x_n + v_n, & v_n &\sim \mathcal{N}(0, R)
    \end{aligned}
    $$

    The forecast and analysis steps are both rewritten in *square-root*
    form on the Cholesky factor ``S`` of ``P`` so the covariance stays
    PSD by construction. Internally this delegates to
    :func:`gaussx.parallel_kalman_filter` (``form='sqrt'``), which runs
    Blelloch's parallel scan along the time axis — cost
    ``O(log T · (Nₓ² Nᵧ + Nₓ³))`` on hardware that can exploit the
    parallelism.

    The constructor stores the (time-invariant) ``Φ``, ``H``, ``Q``,
    ``R`` and lets the user push observations through :meth:`filter`.
    For time-varying systems pass leading-time-axis arrays directly to
    :func:`gaussx.parallel_kalman_filter` and skip this wrapper.

    Attributes:
        transition: State transition operator ``Φ`` (matrix or
            :class:`lineax.AbstractLinearOperator`).
        obs_model: Observation operator ``H``.
        process_noise: Process-noise covariance ``Q``.
        obs_noise: Observation-noise covariance ``R``.
    """

    transition: Float[Array, "N N"] | lx.AbstractLinearOperator
    obs_model: Float[Array, "M N"] | lx.AbstractLinearOperator
    process_noise: Float[Array, "N N"] | lx.AbstractLinearOperator
    obs_noise: Float[Array, "M M"] | lx.AbstractLinearOperator

    def filter(
        self,
        observations: Float[Array, "T M"],
        init_mean: Float[Array, " N"],
        init_cov: Float[Array, "N N"],
    ) -> "SquareRootFilterResult":
        r"""Run the square-root Kalman filter over ``observations``.

        Args:
            observations: Observation sequence ``(T, M)``.
            init_mean: Initial belief mean ``μ₀ ∈ ℝᴺ``.
            init_cov: Initial belief covariance ``P₀ ∈ ℝ^{N×N}``.

        Returns:
            :class:`SquareRootFilterResult` containing the per-step
            filtered means, filtered covariances, and the marginal
            log-likelihood ``log p(y₁…y_T)``.
        """
        state = gaussx.parallel_kalman_filter(
            self.transition,
            self.obs_model,
            self.process_noise,
            self.obs_noise,
            observations,
            init_mean,
            init_cov,
            form="sqrt",
        )
        return SquareRootFilterResult(
            filtered_means=state.filtered_means,
            filtered_covs=state.filtered_covs,
            predicted_means=state.predicted_means,
            predicted_covs=state.predicted_covs,
            log_likelihood=state.log_likelihood,
        )

filter(observations: Float[Array, 'T M'], init_mean: Float[Array, ' N'], init_cov: Float[Array, 'N N']) -> 'SquareRootFilterResult'

Run the square-root Kalman filter over observations.

Parameters:

Name	Type	Description	Default
observations	Float[Array, 'T M']	Observation sequence (T, M).	required
init_mean	Float[Array, ' N']	Initial belief mean μ₀ ∈ ℝᴺ.	required
init_cov	Float[Array, 'N N']	Initial belief covariance P₀ ∈ ℝ^{N×N}.	required

Returns:

Type	Description
'SquareRootFilterResult'	class:SquareRootFilterResult containing the per-step
'SquareRootFilterResult'	filtered means, filtered covariances, and the marginal
'SquareRootFilterResult'	log-likelihood log p(y₁…y_T).

Source code in src/filterax/_filters/_parametric.py

def filter(
    self,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " N"],
    init_cov: Float[Array, "N N"],
) -> "SquareRootFilterResult":
    r"""Run the square-root Kalman filter over ``observations``.

    Args:
        observations: Observation sequence ``(T, M)``.
        init_mean: Initial belief mean ``μ₀ ∈ ℝᴺ``.
        init_cov: Initial belief covariance ``P₀ ∈ ℝ^{N×N}``.

    Returns:
        :class:`SquareRootFilterResult` containing the per-step
        filtered means, filtered covariances, and the marginal
        log-likelihood ``log p(y₁…y_T)``.
    """
    state = gaussx.parallel_kalman_filter(
        self.transition,
        self.obs_model,
        self.process_noise,
        self.obs_noise,
        observations,
        init_mean,
        init_cov,
        form="sqrt",
    )
    return SquareRootFilterResult(
        filtered_means=state.filtered_means,
        filtered_covs=state.filtered_covs,
        predicted_means=state.predicted_means,
        predicted_covs=state.predicted_covs,
        log_likelihood=state.log_likelihood,
    )

SquareRootFilterResult

Bases: Module

Output of :meth:SquareRootKF.filter.

Mirrors :class:gaussx.FilterState so callers don't depend on the internal gaussx structure: per-step filtered / predicted means and covariances along a leading time axis, plus the marginal log-likelihood of the full observation sequence.

Attributes:

Name	Type	Description
filtered_means	Float[Array, 'T N']	μ_{n|n} for n = 1, …, T, shape (T, N).
filtered_covs	Float[Array, 'T N N']	P_{n|n} for n = 1, …, T, shape (T, N, N).
predicted_means	Float[Array, 'T N']	μ_{n|n−1} (forecast means before update), shape (T, N).
predicted_covs	Float[Array, 'T N N']	P_{n|n−1} (forecast covariances), shape (T, N, N).
log_likelihood	Float[Array, '']	Scalar marginal log-likelihood.

Source code in src/filterax/_filters/_parametric.py

class SquareRootFilterResult(eqx.Module, strict=True):
    """Output of :meth:`SquareRootKF.filter`.

    Mirrors :class:`gaussx.FilterState` so callers don't depend on the
    internal gaussx structure: per-step filtered / predicted means and
    covariances along a leading time axis, plus the marginal
    log-likelihood of the full observation sequence.

    Attributes:
        filtered_means: ``μ_{n|n}`` for ``n = 1, …, T``, shape ``(T, N)``.
        filtered_covs: ``P_{n|n}`` for ``n = 1, …, T``, shape ``(T, N, N)``.
        predicted_means: ``μ_{n|n−1}`` (forecast means before update),
            shape ``(T, N)``.
        predicted_covs: ``P_{n|n−1}`` (forecast covariances), shape
            ``(T, N, N)``.
        log_likelihood: Scalar marginal log-likelihood.
    """

    filtered_means: Float[Array, "T N"]
    filtered_covs: Float[Array, "T N N"]
    predicted_means: Float[Array, "T N"]
    predicted_covs: Float[Array, "T N N"]
    log_likelihood: Float[Array, ""]

Layer-2 assimilation models

The full forecast–analyse–inflate loops. Each assimilate() call stacks per-window forecasts, posteriors, and log-likelihoods along a leading time axis in the returned AssimilationResult.

filterax — differentiable ensemble Kalman filters and processes for JAX.

ETKF

Bases: Module

Ensemble Transform Kalman Filter — full forecast-analysis loop.

Composes dynamics → L1 ETKF analysis → optional inflator over each observation window. Localization is not applied at this level; use :class:LETKF for localized assimilation.

Attributes:

Name	Type	Description
dynamics	AbstractDynamics	Forward model applied with :func:eqx.filter_vmap over the ensemble.
obs_op	AbstractObsOperator	Observation operator H.
inflator	AbstractInflator | None	Optional posterior inflator (e.g. :class:filterax.RTPS, :class:filterax.MultiplicativeInflator).
config	FilterConfig | None	Optional :class:FilterConfig reserved for future static configuration (ensemble size, diagnostics toggle, …).

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from filterax import ETKF
>>> model = ETKF(dynamics=lambda x, t0, t1: x, obs_op=lambda x: x)
>>> init = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> obs = [(jnp.array([0.5, 0.5]), 1.0), (jnp.array([0.6, 0.4]), 2.0)]
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
>>> result = model.assimilate(init, obs, R)
>>> result.particles.shape, result.analysis_history.shape
((3, 2), (2, 3, 2))

Source code in src/filterax/_filters/_models.py

class ETKF(eqx.Module, strict=True):
    r"""Ensemble Transform Kalman Filter — full forecast-analysis loop.

    Composes ``dynamics → L1 ETKF analysis → optional inflator`` over
    each observation window. Localization is not applied at this level;
    use :class:`LETKF` for localized assimilation.

    Attributes:
        dynamics: Forward model applied with :func:`eqx.filter_vmap` over
            the ensemble.
        obs_op: Observation operator ``H``.
        inflator: Optional posterior inflator (e.g.
            :class:`filterax.RTPS`, :class:`filterax.MultiplicativeInflator`).
        config: Optional :class:`FilterConfig` reserved for future
            static configuration (ensemble size, diagnostics toggle, …).

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from filterax import ETKF
        >>> model = ETKF(dynamics=lambda x, t0, t1: x, obs_op=lambda x: x)
        >>> init = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> obs = [(jnp.array([0.5, 0.5]), 1.0), (jnp.array([0.6, 0.4]), 2.0)]
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(2))
        >>> result = model.assimilate(init, obs, R)
        >>> result.particles.shape, result.analysis_history.shape
        ((3, 2), (2, 3, 2))
    """

    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    inflator: AbstractInflator | None = None
    config: FilterConfig | None = None

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_x"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        t0: float | Float[Array, ""] = 0.0,
    ) -> AssimilationResult:
        """Run forecast → ETKF analysis → optional inflation over each window."""
        return _run_loop(
            lambda _step: _ETKFFilter(),
            self.dynamics,
            self.obs_op,
            self.inflator,
            init_ensemble,
            observations,
            obs_noise,
            t0,
        )

assimilate(init_ensemble: Float[Array, 'N_e N_x'], observations: Sequence[tuple[Float[Array, ' N_y'], float | Float[Array, '']]], obs_noise: lx.AbstractLinearOperator, t0: float | Float[Array, ''] = 0.0) -> AssimilationResult

Run forecast → ETKF analysis → optional inflation over each window.

Source code in src/filterax/_filters/_models.py

def assimilate(
    self,
    init_ensemble: Float[Array, "N_e N_x"],
    observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
    obs_noise: lx.AbstractLinearOperator,
    t0: float | Float[Array, ""] = 0.0,
) -> AssimilationResult:
    """Run forecast → ETKF analysis → optional inflation over each window."""
    return _run_loop(
        lambda _step: _ETKFFilter(),
        self.dynamics,
        self.obs_op,
        self.inflator,
        init_ensemble,
        observations,
        obs_noise,
        t0,
    )

LETKF

Bases: Module

Local ETKF — full forecast-analysis loop with R-localization.

state_coords and obs_coords are passed at assimilate time because they describe the assimilation problem, not the filter configuration. radius (and the optional taper_fn) are configured here.

Source code in src/filterax/_filters/_models.py

class LETKF(eqx.Module, strict=True):
    r"""Local ETKF — full forecast-analysis loop with R-localization.

    ``state_coords`` and ``obs_coords`` are passed at ``assimilate`` time
    because they describe the assimilation problem, not the filter
    configuration. ``radius`` (and the optional ``taper_fn``) are
    configured here.
    """

    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    radius: float = eqx.field(static=True)
    inflator: AbstractInflator | None = None
    config: FilterConfig | None = None

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_x"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        state_coords: Float[Array, "N_x D"],
        obs_coords: Float[Array, "N_y D"],
        t0: float | Float[Array, ""] = 0.0,
    ) -> AssimilationResult:
        return _run_loop(
            lambda _step: _LETKFFilter(radius=self.radius),
            self.dynamics,
            self.obs_op,
            self.inflator,
            init_ensemble,
            observations,
            obs_noise,
            t0,
            analysis_extra={
                "state_coords": state_coords,
                "obs_coords": obs_coords,
            },
        )

EnSRF

Bases: Module

Ensemble Square Root Filter — full forecast-analysis loop.

Source code in src/filterax/_filters/_models.py

class EnSRF(eqx.Module, strict=True):
    """Ensemble Square Root Filter — full forecast-analysis loop."""

    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    inflator: AbstractInflator | None = None
    config: FilterConfig | None = None

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_x"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        t0: float | Float[Array, ""] = 0.0,
    ) -> AssimilationResult:
        return _run_loop(
            lambda _step: _EnSRFFilter(),
            self.dynamics,
            self.obs_op,
            self.inflator,
            init_ensemble,
            observations,
            obs_noise,
            t0,
        )

StochasticEnKF

Bases: Module

Stochastic Ensemble Kalman Filter — full forecast-analysis loop.

The PRNG key for window step is derived as jr.fold_in(base_key, step) so successive windows draw independent observation perturbations even when called inside the same assimilation. Pass seed (int) or base_key (PRNG array) at construction; base_key overrides seed when both are provided.

Attributes:

Name	Type	Description
dynamics	AbstractDynamics	Forward model.
obs_op	AbstractObsOperator	Observation operator H.
inflator	AbstractInflator | None	Optional posterior inflator.
seed	int	Integer used as jr.PRNGKey(seed) when base_key is not supplied.
base_key	PRNGKeyArray | None	Explicit PRNG array. Takes precedence over seed.
config	FilterConfig | None	Optional :class:FilterConfig.

Source code in src/filterax/_filters/_models.py

class StochasticEnKF(eqx.Module, strict=True):
    r"""Stochastic Ensemble Kalman Filter — full forecast-analysis loop.

    The PRNG key for window ``step`` is derived as
    ``jr.fold_in(base_key, step)`` so successive windows draw
    independent observation perturbations even when called inside the
    same assimilation. Pass ``seed`` (int) or ``base_key`` (PRNG array)
    at construction; ``base_key`` overrides ``seed`` when both are
    provided.

    Attributes:
        dynamics: Forward model.
        obs_op: Observation operator ``H``.
        inflator: Optional posterior inflator.
        seed: Integer used as ``jr.PRNGKey(seed)`` when ``base_key`` is
            not supplied.
        base_key: Explicit PRNG array. Takes precedence over ``seed``.
        config: Optional :class:`FilterConfig`.
    """

    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    inflator: AbstractInflator | None = None
    seed: int = eqx.field(static=True, default=0)
    base_key: PRNGKeyArray | None = None
    config: FilterConfig | None = None

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_x"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        t0: float | Float[Array, ""] = 0.0,
    ) -> AssimilationResult:
        base_key = self.base_key if self.base_key is not None else jr.PRNGKey(self.seed)

        def factory(step: int) -> AbstractSequentialFilter:
            # jr.fold_in derives a fresh sub-key from (base, step) so each
            # window's perturbed observations are independent.
            return _StochasticEnKFFilter(key=jr.fold_in(base_key, step))

        return _run_loop(
            factory,
            self.dynamics,
            self.obs_op,
            self.inflator,
            init_ensemble,
            observations,
            obs_noise,
            t0,
        )

Latent-space assimilation

Run the analysis in a learned (or prescribed) latent space z = E(x): LatentETKF / LatentLETKF are L2 loops that forecast and analyse in z-space and decode back, returning a LatentAssimilationResult. The helper wrappers lift the pieces of a latent problem into the standard protocols — LatentDynamics wraps a z-space forward model, LiftedObs composes the decoder with an x-space observation operator, and EncodedDynamics runs x-space dynamics through an encode/decode round-trip.

filterax — differentiable ensemble Kalman filters and processes for JAX.

LatentETKF

Bases: Module

ETKF analysis applied to a latent-space ensemble (D17).

Thin composition over the existing :class:ETKF Layer-1 step: encode the initial ensemble (when supplied in x-space), run the standard forecast → ETKF → optional inflation cycle on the latent ensemble using LiftedObs(decoder, x_obs_op) as the observation operator, then decode the histories back to x-space alongside the latent representations.

The ensemble math is unchanged — the wrapper exists for ergonomics and for compile-time correctness (it would be hard to accidentally feed an x-space ensemble into a latent-space analysis step once the user builds a LatentETKF).

Attributes:

Name	Type	Description
latent_map	Any	Object exposing .encode and .decode. Typically a :class:pipekit_cycle.LatentMap-shaped autoencoder or :func:identity_latent_map for regression tests.
dynamics	AbstractDynamics	Latent-space dynamics — either an :class:AbstractDynamics (typically :class:LatentDynamics or :class:EncodedDynamics) or a structural LatentForwardModel exposing .step(z, dt) → z. Raw .step-only objects are auto-wrapped in :class:LatentDynamics at construction.
obs_op	AbstractObsOperator	x-space observation operator. Lifted internally via :class:LiftedObs.
inflator	AbstractInflator | None	Optional deterministic inflator applied in latent space.
config	FilterConfig | None	Reserved :class:FilterConfig.

Source code in src/filterax/_filters/_models.py

class LatentETKF(eqx.Module, strict=True):
    r"""ETKF analysis applied to a latent-space ensemble (D17).

    Thin composition over the existing :class:`ETKF` Layer-1 step:
    encode the initial ensemble (when supplied in ``x``-space), run
    the standard forecast → ETKF → optional inflation cycle on the
    latent ensemble using ``LiftedObs(decoder, x_obs_op)`` as the
    observation operator, then decode the histories back to
    ``x``-space alongside the latent representations.

    The ensemble math is unchanged — the wrapper exists for ergonomics
    and for compile-time correctness (it would be hard to accidentally
    feed an ``x``-space ensemble into a latent-space analysis step
    once the user builds a ``LatentETKF``).

    Attributes:
        latent_map: Object exposing ``.encode`` and ``.decode``.
            Typically a :class:`pipekit_cycle.LatentMap`-shaped autoencoder
            or :func:`identity_latent_map` for regression tests.
        dynamics: Latent-space dynamics — either an
            :class:`AbstractDynamics` (typically :class:`LatentDynamics`
            or :class:`EncodedDynamics`) or a structural
            ``LatentForwardModel`` exposing ``.step(z, dt) → z``. Raw
            ``.step``-only objects are auto-wrapped in
            :class:`LatentDynamics` at construction.
        obs_op: ``x``-space observation operator. Lifted internally
            via :class:`LiftedObs`.
        inflator: Optional deterministic inflator applied in latent
            space.
        config: Reserved :class:`FilterConfig`.
    """

    latent_map: Any
    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    inflator: AbstractInflator | None
    config: FilterConfig | None

    def __init__(
        self,
        latent_map: Any,
        dynamics: Any,
        obs_op: AbstractObsOperator,
        inflator: AbstractInflator | None = None,
        config: FilterConfig | None = None,
    ):
        self.latent_map = latent_map
        self.dynamics = _coerce_latent_dynamics(dynamics)
        self.obs_op = obs_op
        self.inflator = inflator
        self.config = config

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_d"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        t0: float | Float[Array, ""] = 0.0,
        *,
        in_x_space: bool = True,
    ) -> LatentAssimilationResult:
        """Run forecast → ETKF analysis → optional inflation in latent space.

        Args:
            init_ensemble: Prior ensemble. Shape ``(Nₑ, Nₓ)`` when
                ``in_x_space=True`` (default); ``(Nₑ, N_z)`` when the
                caller already encoded.
            observations: Sequence of ``(y_t, t)`` pairs in
                ``x``-space coordinates.
            obs_noise: Observation error covariance ``R`` (operates in
                ``y``-space, untouched).
            t0: Starting time for the first forecast window.
            in_x_space: Encode the initial ensemble first. ``False``
                lets the caller skip the round-trip when their initial
                ensemble already lives in ``z``-space.

        Returns:
            :class:`LatentAssimilationResult` carrying both the
            latent and decoded ``x``-space ensembles.
        """
        z0 = (
            latent_ensemble(self.latent_map.encode, init_ensemble)
            if in_x_space
            else init_ensemble
        )
        result_z = _run_loop(
            lambda _step: _ETKFFilter(),
            self.dynamics,
            LiftedObs(decoder=self.latent_map, inner=self.obs_op),
            self.inflator,
            z0,
            observations,
            obs_noise,
            t0,
        )
        return _wrap_latent_result(self.latent_map, result_z)

assimilate(init_ensemble: Float[Array, 'N_e N_d'], observations: Sequence[tuple[Float[Array, ' N_y'], float | Float[Array, '']]], obs_noise: lx.AbstractLinearOperator, t0: float | Float[Array, ''] = 0.0, *, in_x_space: bool = True) -> LatentAssimilationResult

Run forecast → ETKF analysis → optional inflation in latent space.

Parameters:

Name	Type	Description	Default
init_ensemble	Float[Array, 'N_e N_d']	Prior ensemble. Shape (Nₑ, Nₓ) when in_x_space=True (default); (Nₑ, N_z) when the caller already encoded.	required
observations	Sequence[tuple[Float[Array, ' N_y'], float | Float[Array, '']]]	Sequence of (y_t, t) pairs in x-space coordinates.	required
obs_noise	AbstractLinearOperator	Observation error covariance R (operates in y-space, untouched).	required
t0	float | Float[Array, '']	Starting time for the first forecast window.	0.0
in_x_space	bool	Encode the initial ensemble first. False lets the caller skip the round-trip when their initial ensemble already lives in z-space.	True

Returns:

Type	Description
LatentAssimilationResult	class:LatentAssimilationResult carrying both the
LatentAssimilationResult	latent and decoded x-space ensembles.

Source code in src/filterax/_filters/_models.py

def assimilate(
    self,
    init_ensemble: Float[Array, "N_e N_d"],
    observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
    obs_noise: lx.AbstractLinearOperator,
    t0: float | Float[Array, ""] = 0.0,
    *,
    in_x_space: bool = True,
) -> LatentAssimilationResult:
    """Run forecast → ETKF analysis → optional inflation in latent space.

    Args:
        init_ensemble: Prior ensemble. Shape ``(Nₑ, Nₓ)`` when
            ``in_x_space=True`` (default); ``(Nₑ, N_z)`` when the
            caller already encoded.
        observations: Sequence of ``(y_t, t)`` pairs in
            ``x``-space coordinates.
        obs_noise: Observation error covariance ``R`` (operates in
            ``y``-space, untouched).
        t0: Starting time for the first forecast window.
        in_x_space: Encode the initial ensemble first. ``False``
            lets the caller skip the round-trip when their initial
            ensemble already lives in ``z``-space.

    Returns:
        :class:`LatentAssimilationResult` carrying both the
        latent and decoded ``x``-space ensembles.
    """
    z0 = (
        latent_ensemble(self.latent_map.encode, init_ensemble)
        if in_x_space
        else init_ensemble
    )
    result_z = _run_loop(
        lambda _step: _ETKFFilter(),
        self.dynamics,
        LiftedObs(decoder=self.latent_map, inner=self.obs_op),
        self.inflator,
        z0,
        observations,
        obs_noise,
        t0,
    )
    return _wrap_latent_result(self.latent_map, result_z)

LatentLETKF

Bases: Module

Local ETKF in latent space (D17).

Sister of :class:LatentETKF over :class:LETKF. v0.1 implements option (a) of latent_da.md §9.2 — no localization in latent space — because the latent dimensions of a typical autoencoder are global modes without an intrinsic notion of locality. Pseudo-localization-via-mode-coords (option c) is deferred.

For the no-localization case, this class is provided for API parity with :class:LatentETKF: it gives users a LatentLETKF name that they can wire into pipekit alongside LatentETKF without having to fall back to LatentETKF themselves. The body is identical to :class:LatentETKF.

Like :class:LatentETKF, dynamics may be either an :class:AbstractDynamics or a raw structural LatentForwardModel (.step(z, dt) → z); the latter is auto-wrapped in :class:LatentDynamics at construction.

Source code in src/filterax/_filters/_models.py

class LatentLETKF(eqx.Module, strict=True):
    r"""Local ETKF in latent space (D17).

    Sister of :class:`LatentETKF` over :class:`LETKF`. v0.1 implements
    option (a) of `latent_da.md` §9.2 — *no* localization in latent
    space — because the latent dimensions of a typical autoencoder are
    global modes without an intrinsic notion of locality.
    Pseudo-localization-via-mode-coords (option c) is deferred.

    For the no-localization case, this class is provided for API
    parity with :class:`LatentETKF`: it gives users a ``LatentLETKF``
    name that they can wire into pipekit alongside ``LatentETKF``
    without having to fall back to ``LatentETKF`` themselves. The
    body is identical to :class:`LatentETKF`.

    Like :class:`LatentETKF`, ``dynamics`` may be either an
    :class:`AbstractDynamics` or a raw structural ``LatentForwardModel``
    (``.step(z, dt) → z``); the latter is auto-wrapped in
    :class:`LatentDynamics` at construction.
    """

    latent_map: Any
    dynamics: AbstractDynamics
    obs_op: AbstractObsOperator
    inflator: AbstractInflator | None
    config: FilterConfig | None

    def __init__(
        self,
        latent_map: Any,
        dynamics: Any,
        obs_op: AbstractObsOperator,
        inflator: AbstractInflator | None = None,
        config: FilterConfig | None = None,
    ):
        self.latent_map = latent_map
        self.dynamics = _coerce_latent_dynamics(dynamics)
        self.obs_op = obs_op
        self.inflator = inflator
        self.config = config

    def assimilate(
        self,
        init_ensemble: Float[Array, "N_e N_d"],
        observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
        obs_noise: lx.AbstractLinearOperator,
        t0: float | Float[Array, ""] = 0.0,
        *,
        in_x_space: bool = True,
    ) -> LatentAssimilationResult:
        """Run forecast → ETKF analysis (no localization) → inflation in z-space."""
        z0 = (
            latent_ensemble(self.latent_map.encode, init_ensemble)
            if in_x_space
            else init_ensemble
        )
        result_z = _run_loop(
            lambda _step: _ETKFFilter(),
            self.dynamics,
            LiftedObs(decoder=self.latent_map, inner=self.obs_op),
            self.inflator,
            z0,
            observations,
            obs_noise,
            t0,
        )
        return _wrap_latent_result(self.latent_map, result_z)

assimilate(init_ensemble: Float[Array, 'N_e N_d'], observations: Sequence[tuple[Float[Array, ' N_y'], float | Float[Array, '']]], obs_noise: lx.AbstractLinearOperator, t0: float | Float[Array, ''] = 0.0, *, in_x_space: bool = True) -> LatentAssimilationResult

Run forecast → ETKF analysis (no localization) → inflation in z-space.

Source code in src/filterax/_filters/_models.py

def assimilate(
    self,
    init_ensemble: Float[Array, "N_e N_d"],
    observations: Sequence[tuple[Float[Array, " N_y"], float | Float[Array, ""]]],
    obs_noise: lx.AbstractLinearOperator,
    t0: float | Float[Array, ""] = 0.0,
    *,
    in_x_space: bool = True,
) -> LatentAssimilationResult:
    """Run forecast → ETKF analysis (no localization) → inflation in z-space."""
    z0 = (
        latent_ensemble(self.latent_map.encode, init_ensemble)
        if in_x_space
        else init_ensemble
    )
    result_z = _run_loop(
        lambda _step: _ETKFFilter(),
        self.dynamics,
        LiftedObs(decoder=self.latent_map, inner=self.obs_op),
        self.inflator,
        z0,
        observations,
        obs_noise,
        t0,
    )
    return _wrap_latent_result(self.latent_map, result_z)

IdentityLatentMap

Bases: Module

φ = ψ = id regression-test fixture.

Used to verify that the latent filters reduce to their state-space counterparts when the codec is the identity — pinning the wrapping layer without depending on a real autoencoder. Carries a single static dim field so callers can sanity-check shapes.

Source code in src/filterax/_filters/_latent.py

class IdentityLatentMap(eqx.Module, strict=True):
    r"""``φ = ψ = id`` regression-test fixture.

    Used to verify that the latent filters reduce to their state-space
    counterparts when the codec is the identity — pinning the wrapping
    layer without depending on a real autoencoder. Carries a single
    static ``dim`` field so callers can sanity-check shapes.
    """

    dim: int = eqx.field(static=True)

    def encode(self, x: Float[Array, " N_x"]) -> Float[Array, " N_x"]:
        return x

    def decode(self, z: Float[Array, " N_x"]) -> Float[Array, " N_x"]:
        return z

LatentDynamics

Bases: AbstractDynamics

Wrap a structural LatentForwardModel as an AbstractDynamics.

The wrapped object only needs a .step(z, dt) method; nothing else about its type is constrained (no isinstance check). The wrapper exists so the latent dynamics can drop into any existing Layer-2 forecast loop alongside other AbstractDynamics instances.

Attributes:

Name	Type	Description
inner	Any	The wrapped latent forward model.

Source code in src/filterax/_filters/_latent.py

class LatentDynamics(AbstractDynamics, strict=True):
    r"""Wrap a structural ``LatentForwardModel`` as an ``AbstractDynamics``.

    The wrapped object only needs a ``.step(z, dt)`` method; nothing
    else about its type is constrained (no ``isinstance`` check). The
    wrapper exists so the latent dynamics can drop into any existing
    Layer-2 forecast loop alongside other ``AbstractDynamics``
    instances.

    Attributes:
        inner: The wrapped latent forward model.
    """

    inner: Any

    def __call__(
        self,
        state: Float[Array, " N_z"],
        t0: Float[Array, ""],
        t1: Float[Array, ""],
    ) -> Float[Array, " N_z"]:
        return self.inner.step(state, t1 - t0)

LiftedObs

Bases: AbstractObsOperator

Compose a decoder with an x-space observation operator.

Maps z → y via z ─ψ→ x ─H→ y. Satisfies :class:AbstractObsOperator, so it slots into ETKF / EnSRF / LETKF / EnKS without changes — they only ever see the (N_z, N_y) interface and don't care what's inside.

Attributes:

Name	Type	Description
decoder	Any	Anything callable as decoder.decode(z) → x.
inner	AbstractObsOperator	Standard x-space observation operator.

Source code in src/filterax/_filters/_latent.py

class LiftedObs(AbstractObsOperator, strict=True):
    r"""Compose a decoder with an ``x``-space observation operator.

    Maps ``z → y`` via ``z ─ψ→ x ─H→ y``. Satisfies
    :class:`AbstractObsOperator`, so it slots into ETKF / EnSRF /
    LETKF / EnKS without changes — they only ever see the
    ``(N_z, N_y)`` interface and don't care what's inside.

    Attributes:
        decoder: Anything callable as ``decoder.decode(z) → x``.
        inner: Standard ``x``-space observation operator.
    """

    decoder: Any
    inner: AbstractObsOperator

    def __call__(self, state: Float[Array, " N_z"]) -> Float[Array, " N_y"]:
        return self.inner(self.decoder.decode(state))

EncodedDynamics

Bases: AbstractDynamics

Lift an x-space dynamics into z-space via codec round-trip.

For users with an x-space :class:AbstractDynamics (a physics model or a diffrax ODE wrapper) but no learned M_z. One ensemble step costs decode → inner → encode. Used as the dynamics argument to :class:filterax.LatentETKF when no latent dynamics is available.

Attributes:

Name	Type	Description
latent_map	Any	.encode / .decode codec.
inner	AbstractDynamics	x-space dynamics.

Source code in src/filterax/_filters/_latent.py

class EncodedDynamics(AbstractDynamics, strict=True):
    r"""Lift an ``x``-space dynamics into ``z``-space via codec round-trip.

    For users with an ``x``-space :class:`AbstractDynamics` (a physics
    model or a ``diffrax`` ODE wrapper) but no learned ``M_z``. One
    ensemble step costs ``decode → inner → encode``. Used as the
    ``dynamics`` argument to :class:`filterax.LatentETKF` when no
    latent dynamics is available.

    Attributes:
        latent_map: ``.encode`` / ``.decode`` codec.
        inner: ``x``-space dynamics.
    """

    latent_map: Any
    inner: AbstractDynamics

    def __call__(
        self,
        state: Float[Array, " N_z"],
        t0: Float[Array, ""],
        t1: Float[Array, ""],
    ) -> Float[Array, " N_z"]:
        x = self.latent_map.decode(state)
        x_next = self.inner(x, t0, t1)
        return self.latent_map.encode(x_next)

latent_ensemble(encoder: Any, x_ensemble: Float[Array, 'N_e N_x']) -> Float[Array, 'N_e N_z']

Encode each ensemble member into latent space.

vmaps the encoder over the ensemble leading axis. The encoder is anything callable as encoder(x) → z — typically the .encode method bound off a :class:pipekit_cycle.LatentMap or a bare encoder eqx.Module.

Parameters:

Name	Type	Description	Default
encoder	Any	Callable x → z.	required
x_ensemble	Float[Array, 'N_e N_x']	(Nₑ, Nₓ) ensemble in state space.	required

Returns:

Type	Description
Float[Array, 'N_e N_z']	(Nₑ, N_z) ensemble in latent space.

Source code in src/filterax/_filters/_latent.py

def latent_ensemble(
    encoder: Any,
    x_ensemble: Float[Array, "N_e N_x"],
) -> Float[Array, "N_e N_z"]:
    """Encode each ensemble member into latent space.

    ``vmap``s the encoder over the ensemble leading axis. The encoder
    is anything callable as ``encoder(x) → z`` — typically the
    ``.encode`` method bound off a :class:`pipekit_cycle.LatentMap` or
    a bare encoder ``eqx.Module``.

    Args:
        encoder: Callable ``x → z``.
        x_ensemble: ``(Nₑ, Nₓ)`` ensemble in state space.

    Returns:
        ``(Nₑ, N_z)`` ensemble in latent space.
    """
    return eqx.filter_vmap(encoder)(x_ensemble)

decode_ensemble(decoder: Any, z_ensemble: Float[Array, 'N_e N_z']) -> Float[Array, 'N_e N_x']

Decode each ensemble member back into state space.

Mirror of :func:latent_ensemble for the decode direction.

Parameters:

Name	Type	Description	Default
decoder	Any	Callable z → x.	required
z_ensemble	Float[Array, 'N_e N_z']	(Nₑ, N_z) ensemble in latent space.	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	(Nₑ, Nₓ) ensemble in state space.

Source code in src/filterax/_filters/_latent.py

def decode_ensemble(
    decoder: Any,
    z_ensemble: Float[Array, "N_e N_z"],
) -> Float[Array, "N_e N_x"]:
    """Decode each ensemble member back into state space.

    Mirror of :func:`latent_ensemble` for the decode direction.

    Args:
        decoder: Callable ``z → x``.
        z_ensemble: ``(Nₑ, N_z)`` ensemble in latent space.

    Returns:
        ``(Nₑ, Nₓ)`` ensemble in state space.
    """
    return eqx.filter_vmap(decoder)(z_ensemble)

identity_latent_map(dim: int) -> IdentityLatentMap

Construct an :class:IdentityLatentMap of the given dimension.

The returned object satisfies the structural LatentMap contract (.encode + .decode) and is the canonical fixture for regression tests that pin :class:filterax.LatentETKF to plain :class:filterax.ETKF.

Parameters:

Name	Type	Description	Default
dim	int	Common state / latent dimension.	required

Returns:

Type	Description
IdentityLatentMap	Pass-through latent map.

Source code in src/filterax/_filters/_latent.py

def identity_latent_map(dim: int) -> IdentityLatentMap:
    """Construct an :class:`IdentityLatentMap` of the given dimension.

    The returned object satisfies the structural ``LatentMap`` contract
    (``.encode`` + ``.decode``) and is the canonical fixture for
    regression tests that pin :class:`filterax.LatentETKF` to plain
    :class:`filterax.ETKF`.

    Args:
        dim: Common state / latent dimension.

    Returns:
        Pass-through latent map.
    """
    return IdentityLatentMap(dim=dim)