Advanced Sequential Filters

Specialised deterministic variants of the everyday Layer-1 filters (ETKF, EnSRF, StochasticEnKF, LETKF). Use those as the defaults; reach for these when you need the specific properties below. All three implement AbstractSequentialFilter, so they slot into the same Layer-2 loops and pipekit adapters.

Picking an advanced filter

Filter	Family	When to use
filterax.filters.ETKF_Livings	Deterministic, randomised	When the symmetric-sqrt ETKF develops preferred ensemble directions over long runs — the random mean-preserving rotation averages out the degeneracy.
filterax.filters.EnSRF_Serial	Deterministic, scalar-serial	When R is diagonal and you want to avoid any N_y × N_y solve. Cost O(N_e N_x N_y) with no matrix inversion.
filterax.filters.ESTKF	Deterministic, reduced-rank	When you want exact mean preservation by construction and a modest constant-factor speedup over ETKF — the eigendecomposition is (N_e − 1)³ rather than N_e³.

For the parametric (non-ensemble) SquareRootKF — linear-Gaussian baselines and twin-experiment ground truths — see the Filters page.

Note that ETKF_Livings is stochastic (the rotation draws from a PRNG key), so it is rejected by the differentiable training entry points, like StochasticEnKF.

Reference

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

ESTKF

Bases: AbstractSequentialFilter

Error Subspace Transform Kalman Filter (Nerger et al. 2012).

Projects ETKF into the (Nₑ − 1)-dimensional error subspace via a mean-preserving orthogonal map L ∈ ℝ^{Nₑ × (Nₑ − 1)} with Lᵀ L = I_{Nₑ − 1} and Lᵀ 𝟙 = 0. Anomalies are rank ≤ Nₑ − 1 to begin with (their rows sum to zero), so L captures them losslessly.

Reduced anomalies and transform precision:

\[ \tilde{X} = L^{\top} X' \in \mathbb{R}^{(N_e-1) \times N_x}, \qquad \tilde{Y} = L^{\top} Y' \in \mathbb{R}^{(N_e-1) \times N_y} \]

\[ A = (N_e - 1) I + \tilde{Y} R^{-1} \tilde{Y}^{\top} \in \mathbb{R}^{(N_e-1) \times (N_e-1)} \]

Eigendecompose A = U Λ Uᵀ:

\[ \tilde{w} = U \Lambda^{-1} U^{\top} \tilde{Y} R^{-1} d, \qquad \tilde{W} = U \sqrt{(N_e - 1) \Lambda^{-1}}\, U^{\top} \]

Lift back to the full ensemble:

\[ \bar{x}_a = \bar{x}_f + \tilde{w}^{\top} \tilde{X}, \qquad X_a = \bar{x}_a \mathbf{1}^{\top} + L \tilde{W} \tilde{X} \]

Mean-preserving by construction; PSD analysis covariance; eigendecomposition is (Nₑ − 1)³ rather than Nₑ³.

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

class ESTKF(AbstractSequentialFilter, strict=True):
    r"""Error Subspace Transform Kalman Filter (Nerger et al. 2012).

    Projects ETKF into the ``(Nₑ − 1)``-dimensional error subspace via
    a mean-preserving orthogonal map ``L ∈ ℝ^{Nₑ × (Nₑ − 1)}`` with
    ``Lᵀ L = I_{Nₑ − 1}`` and ``Lᵀ 𝟙 = 0``. Anomalies are rank
    ``≤ Nₑ − 1`` to begin with (their rows sum to zero), so ``L``
    captures them losslessly.

    Reduced anomalies and transform precision:

    $$
    \tilde{X} = L^{\top} X' \in \mathbb{R}^{(N_e-1) \times N_x}, \qquad
    \tilde{Y} = L^{\top} Y' \in \mathbb{R}^{(N_e-1) \times N_y}
    $$

    $$
    A = (N_e - 1) I + \tilde{Y} R^{-1} \tilde{Y}^{\top}
        \in \mathbb{R}^{(N_e-1) \times (N_e-1)}
    $$

    Eigendecompose ``A = U Λ Uᵀ``:

    $$
    \tilde{w} = U \Lambda^{-1} U^{\top} \tilde{Y} R^{-1} d, \qquad
    \tilde{W} = U \sqrt{(N_e - 1) \Lambda^{-1}}\, U^{\top}
    $$

    Lift back to the full ensemble:

    $$
    \bar{x}_a = \bar{x}_f + \tilde{w}^{\top} \tilde{X}, \qquad
    X_a = \bar{x}_a \mathbf{1}^{\top} + L \tilde{W} \tilde{X}
    $$

    Mean-preserving by construction; PSD analysis covariance;
    eigendecomposition is ``(Nₑ − 1)³`` rather than ``Nₑ³``.
    """

    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-preserving projection L ∈ ℝ^{Nₑ × (Nₑ−1)}: Householder
        # reflector mapping 𝟙/√Nₑ → e₀, then drop the first column.
        dtype = anom.dtype
        ones = jnp.ones(N_e, dtype=dtype) / jnp.sqrt(N_e)
        e0 = jnp.zeros(N_e, dtype=dtype).at[0].set(1.0)
        v = ones - e0
        v_norm = jnp.linalg.norm(v)
        v = jnp.where(v_norm > 0, v / jnp.maximum(v_norm, 1e-30), v)
        H_mat = jnp.eye(N_e, dtype=dtype) - 2.0 * jnp.outer(v, v)
        L = H_mat[:, 1:]  # (Nₑ, Nₑ−1)

        # Reduce both anomaly matrices to (Nₑ−1) ensemble coordinates.
        X_tilde = einx.dot("e r, e x -> r x", L, anom)  # (Nₑ−1, Nₓ)
        Y_tilde = einx.dot("e r, e y -> r y", L, obs_anom)  # (Nₑ−1, Nᵧ)

        # Rank-Nᵧ spectrum of Ỹ R⁻¹ Ỹᵀ in the reduced subspace; same
        # rank-deficiency story as ETKF, so the QR + small-eigh form
        # keeps gradients finite.
        R_inv_Y = gaussx.solve_rows(obs_noise, Y_tilde)  # (Nₑ−1, Nᵧ)
        eigvals, U_y = _etkf_inner_spectrum(Y_tilde, R_inv_Y)

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

        # Reduced weights w̃ = T̃ (Ỹ R⁻¹ d)  ∈ ℝ^{Nₑ−1}.
        rhs = einx.dot("r y, y -> r", R_inv_Y, innovation)
        w_tilde = _apply_ctilde_func(U_y, inv_diff, inv_base, rhs)

        # Symmetric square-root applied to X̃: W̃ X̃ = (I + correction) X̃.
        WX_tilde = _apply_ctilde_func(U_y, sqrt_diff, sqrt_base, X_tilde)

        # Lift back to Nₑ space and assemble the analysis.
        mean_shift = einx.dot("r, r x -> x", w_tilde, X_tilde)  # (Nₓ,)
        pert_full = einx.dot("e r, r x -> e x", L, WX_tilde)  # (Nₑ, Nₓ)

        particles_a = (x_bar + mean_shift)[None, :] + pert_full

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

ETKF_Livings

Bases: AbstractSequentialFilter

ETKF with a mean-preserving random rotation (Livings et al. 2008).

ETKF's symmetric square root W_a = √((Nₑ − 1) T) is the unique PSD- and mean-preserving choice, but it produces a deterministic transform that can develop preferred directions over many cycles — the ensemble loses rank to a small invariant subspace. Livings et al. break this symmetry by composing with a random orthogonal matrix:

\[ W_a^{\mathrm{rot}} = W_a\, \Theta, \qquad \Theta \in O(N_e), \quad \Theta \mathbf{1} = \mathbf{1} \]

The mean-preserving constraint Θ 𝟙 = 𝟙 keeps the analysis mean and the rank-deficiency-against-𝟙 property; the random factor averages out preferred-direction artefacts over time.

PRNG key handling. The constructor's self.key is the default source of randomness; analysis(..., key=subkey) overrides it for a single call. Pattern-match the :class:StochasticEnKF API: if you call analysis directly inside a loop, pass a freshly split key each cycle or you will get the same rotation every window — which defeats the whole purpose of the Livings recipe.

Attributes:

Name	Type	Description
key	PRNGKeyArray	Default PRNG key for the random rotation. Used when no explicit key= kwarg is passed to analysis.

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

class ETKF_Livings(AbstractSequentialFilter, strict=True):
    r"""ETKF with a mean-preserving random rotation (Livings et al. 2008).

    ETKF's symmetric square root ``W_a = √((Nₑ − 1) T)`` is the unique
    PSD- and mean-preserving choice, but it produces a *deterministic*
    transform that can develop preferred directions over many cycles —
    the ensemble loses rank to a small invariant subspace. Livings et
    al. break this symmetry by composing with a random orthogonal
    matrix:

    $$
    W_a^{\mathrm{rot}} = W_a\, \Theta, \qquad
    \Theta \in O(N_e), \quad \Theta \mathbf{1} = \mathbf{1}
    $$

    The mean-preserving constraint ``Θ 𝟙 = 𝟙`` keeps the analysis
    mean and the rank-deficiency-against-𝟙 property; the random factor
    averages out preferred-direction artefacts over time.

    PRNG key handling. The constructor's ``self.key`` is the *default*
    source of randomness; ``analysis(..., key=subkey)`` overrides it
    for a single call. Pattern-match the :class:`StochasticEnKF` API:
    if you call ``analysis`` directly inside a loop, pass a freshly
    split key each cycle or you will get the *same* rotation every
    window — which defeats the whole purpose of the Livings recipe.

    Attributes:
        key: Default PRNG key for the random rotation. Used when no
            explicit ``key=`` kwarg is passed to ``analysis``.
    """

    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:
        N_e = particles.shape[0]
        check_ensemble_size(N_e)
        rot_key = self.key if key is None else key
        obs_particles = _apply_obs_op(obs_op, particles)

        anom = ensemble_anomalies(particles)
        obs_anom = ensemble_anomalies(obs_particles)
        x_bar = ensemble_mean(particles)
        innovation = obs - ensemble_mean(obs_particles)

        R_inv_Y = gaussx.solve_rows(obs_noise, obs_anom)
        eigvals, U_y = _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) - inv_base
        sqrt_base = jnp.asarray(1.0, dtype=anom.dtype)
        sqrt_diff = jnp.sqrt((N_e - 1) / ((N_e - 1) + eigvals)) - sqrt_base

        # ETKF weight + symmetric square root, same as the base ETKF.
        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)

        W_anom = _apply_ctilde_func(U_y, sqrt_diff, sqrt_base, anom)  # (Nₑ, Nₓ)

        # Apply the mean-preserving rotation Θ on the *left* of W_a X′:
        # Θ leaves 𝟙 fixed, so the column-mean (== analysis mean) is
        # invariant under the rotation.
        Theta = _mean_preserving_rotation(rot_key, N_e, dtype=anom.dtype)
        pert_correction = einx.dot("i j, j x -> i x", Theta, W_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)

EnSRF_Serial

Bases: AbstractSequentialFilter

Serial Ensemble Square Root Filter (Whitaker & Hamill 2002).

Processes observations one at a time with the classical W&H scalar reduced-gain formula:

\[ \begin{aligned} K_k &= C^{x H_k} \big/ \big(C^{H_k H_k} + R_{kk}\big) \\ \bar{x}_a^{(k)} &= \bar{x}_a^{(k-1)} + K_k \big(y_k - H_k \bar{x}_a^{(k-1)}\big) \\ \alpha_k &= 1 \Big/ \Big(1 + \sqrt{R_{kk} / (C^{H_k H_k} + R_{kk})}\Big) \\ X'^{(k)}_a &= X'^{(k-1)}_a - \alpha_k K_k \big(H_k X'^{(k-1)}_a\big) \end{aligned} \]

No matrix inversion is required — each scalar update is an inner product. Cost O(Nₑ Nₓ Nᵧ) total. Requires diagonal R (uncorrelated obs); use :class:ETKF / :class:EnSRF for general R or pre-decorrelate observations.

The implementation reapplies obs_op to the running ensemble inside each scalar update. That keeps the serial trajectory consistent for nonlinear H (where the relationship between state-space and obs-space anomalies is updated by H itself, not by linear book-keeping). Cost is O(Nᵧ × Nₑ × cost(H)) — fractionally more than the linear-only shortcut but correct for every H filterax supports.

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

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

    Processes observations **one at a time** with the classical W&H
    scalar reduced-gain formula:

    $$
    \begin{aligned}
    K_k &= C^{x H_k} \big/ \big(C^{H_k H_k} + R_{kk}\big) \\
    \bar{x}_a^{(k)} &= \bar{x}_a^{(k-1)}
        + K_k \big(y_k - H_k \bar{x}_a^{(k-1)}\big) \\
    \alpha_k &= 1 \Big/ \Big(1 + \sqrt{R_{kk} / (C^{H_k H_k} + R_{kk})}\Big) \\
    X'^{(k)}_a &= X'^{(k-1)}_a - \alpha_k K_k \big(H_k X'^{(k-1)}_a\big)
    \end{aligned}
    $$

    No matrix inversion is required — each scalar update is an inner
    product. Cost ``O(Nₑ Nₓ Nᵧ)`` total. Requires **diagonal** ``R``
    (uncorrelated obs); use :class:`ETKF` / :class:`EnSRF` for general
    ``R`` or pre-decorrelate observations.

    The implementation reapplies ``obs_op`` to the running ensemble
    inside each scalar update. That keeps the serial trajectory
    consistent for **nonlinear** ``H`` (where the relationship between
    state-space and obs-space anomalies is updated by ``H`` itself, not
    by linear book-keeping). Cost is ``O(Nᵧ × Nₑ × cost(H))`` —
    fractionally more than the linear-only shortcut but correct for
    every ``H`` filterax supports.
    """

    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)
        if not isinstance(obs_noise, lx.DiagonalLinearOperator):
            raise NotImplementedError(
                "EnSRF_Serial requires diagonal observation noise; pre-"
                "decorrelate or use filterax.filters.EnSRF for general Γ."
            )
        R_diag = lx.diagonal(obs_noise)
        initial_obs_particles = _apply_obs_op(obs_op, particles)

        def step(parts, k):
            # Re-evaluate H on the *current* ensemble so nonlinear obs
            # ops remain consistent across the serial sweep.
            parts_obs = _apply_obs_op(obs_op, parts)  # (Nₑ, Nᵧ)
            y_k = obs[k]
            R_kk = R_diag[k]
            obs_col = parts_obs[:, k]  # (Nₑ,)
            obs_mean = jnp.mean(obs_col)
            obs_anom_col = obs_col - obs_mean
            obs_var = jnp.sum(obs_anom_col * obs_anom_col) / (N_e - 1)  # Cᴴₖᴴₖ
            innovation_var = obs_var + R_kk

            mean_parts = jnp.mean(parts, axis=0)
            state_anom = parts - mean_parts[None, :]  # (Nₑ, Nₓ)
            C_xH = einx.dot("e, e x -> x", obs_anom_col, state_anom) / (N_e - 1)

            K_k = C_xH / innovation_var  # (Nₓ,)

            # Mean update.
            mean_parts_new = mean_parts + K_k * (y_k - obs_mean)
            # Reduced-gain factor for the perturbation half of the
            # square-root update.
            alpha = 1.0 / (1.0 + jnp.sqrt(R_kk / innovation_var))
            pert_update_state = alpha * einx.dot("x, e -> e x", K_k, obs_anom_col)
            parts_new = mean_parts_new[None, :] + (state_anom - pert_update_state)
            return parts_new, None

        N_y = obs.shape[0]
        particles_a, _ = jax.lax.scan(step, particles, jnp.arange(N_y))

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