Smoothers

Backward-pass refiners that turn a sequential-filter history into the smoothing distribution p(x_t | y_{1:T}) — conditioning every state on the whole observation record rather than only the past. All smoothers consume the stacked forecast_history / analysis_history produced by an AssimilationResult and return a SmoothingResult.

Rough selection guide: EnKS (Evensen & van Leeuwen 2000) is the standard single backward pass and the default choice; EnsembleRTS is the ensemble analogue of the classical Rauch-Tung-Striebel recursion (it coincides with EnKS for zero model error); EnsembleSqrtSmoother is the deterministic square-root variant, paired with a forward square-root filter (ETKF, EnSRF, ESTKF); FixedLagSmoother smooths only a trailing window of length lag — the bounded-memory / streaming option, equivalent to EnKS when lag ≥ T − 1; and IES is the standalone iterative ensemble smoother (Chen & Oliver 2013) that solves an inverse problem over a single observation window, for strongly nonlinear cases where one backward pass is not enough.

These are importable both from the top level (filterax.EnKS) and from the themed namespace (filterax.smoothers.EnKS).

Reference

filterax — differentiable ensemble Kalman filters and processes for JAX.

EnKS

Bases: Module

Ensemble Kalman Smoother — Evensen & van Leeuwen (2000).

Standard single backward pass over a stored filter history. The smoother gain at step t is the sample cross-covariance times the pseudoinverse of the sample forecast covariance,

\[ G_t = (N_e - 1)^{-1} A_t^{\top} F_{t+1} \cdot \big((N_e - 1)^{-1} F_{t+1}^{\top} F_{t+1}\big)^{+}, \]

factored through the Nₑ × Nₑ Gram matrix F Fᵀ so the dense Nₓ × Nₓ forecast covariance is never materialised.

Compatible with any :class:AbstractSequentialFilter (StochasticEnKF, ETKF, EnSRF, LETKF, ETKF_Livings, EnSRF_Serial, ESTKF). Pass the forecast_history and analysis_history produced by an :class:AssimilationResult.

Complexity: O(T (Nₑ² Nₓ + Nₑ³)) for the whole backward pass.

Examples:

>>> import jax
>>> from filterax import EnKS
>>> k1, k2 = jax.random.split(jax.random.key(0))
>>> forecast = jax.random.normal(k1, (3, 4, 2))  # (T, N_e, N_x)
>>> analysis = jax.random.normal(k2, (3, 4, 2))
>>> result = EnKS().smooth(forecast, analysis)
>>> result.smoothed_history.shape
(3, 4, 2)

Source code in src/filterax/_smoothers.py

class EnKS(eqx.Module, strict=True):
    r"""Ensemble Kalman Smoother — Evensen & van Leeuwen (2000).

    Standard single backward pass over a stored filter history. The
    smoother gain at step ``t`` is the sample cross-covariance times the
    pseudoinverse of the sample forecast covariance,

    $$
    G_t = (N_e - 1)^{-1} A_t^{\top} F_{t+1}
        \cdot \big((N_e - 1)^{-1} F_{t+1}^{\top} F_{t+1}\big)^{+},
    $$

    factored through the ``Nₑ × Nₑ`` Gram matrix ``F Fᵀ`` so the dense
    ``Nₓ × Nₓ`` forecast covariance is never materialised.

    Compatible with any :class:`AbstractSequentialFilter` (StochasticEnKF,
    ETKF, EnSRF, LETKF, ETKF_Livings, EnSRF_Serial, ESTKF). Pass the
    ``forecast_history`` and ``analysis_history`` produced by an
    :class:`AssimilationResult`.

    Complexity: ``O(T (Nₑ² Nₓ + Nₑ³))`` for the whole backward pass.

    Examples:
        >>> import jax
        >>> from filterax import EnKS
        >>> k1, k2 = jax.random.split(jax.random.key(0))
        >>> forecast = jax.random.normal(k1, (3, 4, 2))  # (T, N_e, N_x)
        >>> analysis = jax.random.normal(k2, (3, 4, 2))
        >>> result = EnKS().smooth(forecast, analysis)
        >>> result.smoothed_history.shape
        (3, 4, 2)
    """

    def smooth(
        self,
        forecast_history: Float[Array, "T N_e N_x"],
        analysis_history: Float[Array, "T N_e N_x"],
    ) -> SmoothingResult:
        """Run the full backward pass.

        Args:
            forecast_history: Stacked forecast ensembles
                ``X^f_0, …, X^f_{T-1}`` from the filter pass.
            analysis_history: Stacked analysis ensembles
                ``X^a_0, …, X^a_{T-1}`` from the filter pass.

        Returns:
            :class:`SmoothingResult` with ``smoothed_history[t] = X^s_t``
            and ``particles = X^s_{T-1}`` — the terminal ensemble that
            chains into a follow-up forecast, matching the
            ``AssimilationResult.particles`` convention.
        """
        _validate_history(analysis_history, forecast_history)
        smoothed = _enks_backward(analysis_history, forecast_history)
        return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

smooth(forecast_history: Float[Array, 'T N_e N_x'], analysis_history: Float[Array, 'T N_e N_x']) -> SmoothingResult

Run the full backward pass.

Parameters:

Name	Type	Description	Default
forecast_history	Float[Array, 'T N_e N_x']	Stacked forecast ensembles X^f_0, …, X^f_{T-1} from the filter pass.	required
analysis_history	Float[Array, 'T N_e N_x']	Stacked analysis ensembles X^a_0, …, X^a_{T-1} from the filter pass.	required

Returns:

Type	Description
SmoothingResult	class:SmoothingResult with smoothed_history[t] = X^s_t
SmoothingResult	and particles = X^s_{T-1} — the terminal ensemble that
SmoothingResult	chains into a follow-up forecast, matching the
SmoothingResult	AssimilationResult.particles convention.

Source code in src/filterax/_smoothers.py

def smooth(
    self,
    forecast_history: Float[Array, "T N_e N_x"],
    analysis_history: Float[Array, "T N_e N_x"],
) -> SmoothingResult:
    """Run the full backward pass.

    Args:
        forecast_history: Stacked forecast ensembles
            ``X^f_0, …, X^f_{T-1}`` from the filter pass.
        analysis_history: Stacked analysis ensembles
            ``X^a_0, …, X^a_{T-1}`` from the filter pass.

    Returns:
        :class:`SmoothingResult` with ``smoothed_history[t] = X^s_t``
        and ``particles = X^s_{T-1}`` — the terminal ensemble that
        chains into a follow-up forecast, matching the
        ``AssimilationResult.particles`` convention.
    """
    _validate_history(analysis_history, forecast_history)
    smoothed = _enks_backward(analysis_history, forecast_history)
    return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

EnsembleRTS

Bases: Module

Ensemble Rauch-Tung-Striebel smoother.

Conceptual ensemble analog of the classical RTS recursion. With the cross-covariance form used here — G_t = C^{af}_{t,t+1} (C^{ff}_{t+1})⁺ — the algorithm coincides with :class:EnKS for zero model error (Evensen 2003, §5; see also the smoothers design doc). The class exists so callers can name the RTS interpretation explicitly; the implementation is shared with :class:EnKS.

Compatible with any :class:AbstractSequentialFilter. Complexity: O(T (Nₑ² Nₓ + Nₑ³)).

Source code in src/filterax/_smoothers.py

class EnsembleRTS(eqx.Module, strict=True):
    r"""Ensemble Rauch-Tung-Striebel smoother.

    Conceptual ensemble analog of the classical RTS recursion. With the
    cross-covariance form used here — ``G_t = C^{af}_{t,t+1}
    (C^{ff}_{t+1})⁺`` — the algorithm coincides with :class:`EnKS` for
    zero model error (Evensen 2003, §5; see also the smoothers design
    doc). The class exists so callers can name the RTS interpretation
    explicitly; the implementation is shared with :class:`EnKS`.

    Compatible with any :class:`AbstractSequentialFilter`. Complexity:
    ``O(T (Nₑ² Nₓ + Nₑ³))``.
    """

    def smooth(
        self,
        forecast_history: Float[Array, "T N_e N_x"],
        analysis_history: Float[Array, "T N_e N_x"],
    ) -> SmoothingResult:
        """Run the backward RTS pass — see :meth:`EnKS.smooth`."""
        _validate_history(analysis_history, forecast_history)
        smoothed = _enks_backward(analysis_history, forecast_history)
        return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

smooth(forecast_history: Float[Array, 'T N_e N_x'], analysis_history: Float[Array, 'T N_e N_x']) -> SmoothingResult

Run the backward RTS pass — see :meth:EnKS.smooth.

Source code in src/filterax/_smoothers.py

def smooth(
    self,
    forecast_history: Float[Array, "T N_e N_x"],
    analysis_history: Float[Array, "T N_e N_x"],
) -> SmoothingResult:
    """Run the backward RTS pass — see :meth:`EnKS.smooth`."""
    _validate_history(analysis_history, forecast_history)
    smoothed = _enks_backward(analysis_history, forecast_history)
    return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

EnsembleSqrtSmoother

Bases: Module

Ensemble Square Root Smoother — Tippett et al. (2003); Whitaker & Compo (2002).

Deterministic backward pass paired with a forward square-root filter (ETKF, EnSRF, ESTKF). Differs from :class:EnKS in how the smoothed perturbations are produced: instead of the raw per-member correction δ_j = G_t (X^s_{t+1}[j] - X^f_{t+1}[j]), the smoothed mean and perturbations are updated separately, with the perturbation half applied through a symmetric square root of the ensemble-space cov-update transform

\[ \Lambda = I + K_e \big(D^{\top} D - F^{\top} F\big) K_e^{\top}, \qquad K_e = (F F^{\top})^{+} F. \]

This avoids accumulating cross-member coupling across many backward steps — the same reason :class:ETKF uses a symmetric square-root transform on the forward pass.

Mathematically agrees with :class:EnKS in the infinite-ensemble limit; with finite ensembles the mean is identical (both produce m^a_t + G_t (m^s_{t+1} - m^f_{t+1})) and the perturbations differ only in their ensemble-space rotation.

Complexity: O(T (Nₑ² Nₓ + Nₑ³)) — same as :class:EnKS.

Source code in src/filterax/_smoothers.py

class EnsembleSqrtSmoother(eqx.Module, strict=True):
    r"""Ensemble Square Root Smoother — Tippett et al. (2003); Whitaker & Compo (2002).

    Deterministic backward pass paired with a forward square-root filter
    (ETKF, EnSRF, ESTKF). Differs from :class:`EnKS` in *how* the
    smoothed perturbations are produced: instead of the raw per-member
    correction ``δ_j = G_t (X^s_{t+1}[j] - X^f_{t+1}[j])``, the
    smoothed mean and perturbations are updated separately, with the
    perturbation half applied through a *symmetric square root* of the
    ensemble-space cov-update transform

    $$
    \Lambda = I + K_e \big(D^{\top} D - F^{\top} F\big) K_e^{\top}, \qquad
    K_e = (F F^{\top})^{+} F.
    $$

    This avoids accumulating cross-member coupling across many backward
    steps — the same reason :class:`ETKF` uses a symmetric square-root
    transform on the forward pass.

    Mathematically agrees with :class:`EnKS` in the infinite-ensemble
    limit; with finite ensembles the *mean* is identical (both produce
    ``m^a_t + G_t (m^s_{t+1} - m^f_{t+1})``) and the *perturbations*
    differ only in their ensemble-space rotation.

    Complexity: ``O(T (Nₑ² Nₓ + Nₑ³))`` — same as :class:`EnKS`.
    """

    def smooth(
        self,
        forecast_history: Float[Array, "T N_e N_x"],
        analysis_history: Float[Array, "T N_e N_x"],
    ) -> SmoothingResult:
        """Run the square-root backward pass — see :meth:`EnKS.smooth`."""
        _validate_history(analysis_history, forecast_history)
        smoothed = _sqrt_backward(analysis_history, forecast_history)
        return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

smooth(forecast_history: Float[Array, 'T N_e N_x'], analysis_history: Float[Array, 'T N_e N_x']) -> SmoothingResult

Run the square-root backward pass — see :meth:EnKS.smooth.

Source code in src/filterax/_smoothers.py

def smooth(
    self,
    forecast_history: Float[Array, "T N_e N_x"],
    analysis_history: Float[Array, "T N_e N_x"],
) -> SmoothingResult:
    """Run the square-root backward pass — see :meth:`EnKS.smooth`."""
    _validate_history(analysis_history, forecast_history)
    smoothed = _sqrt_backward(analysis_history, forecast_history)
    return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

FixedLagSmoother

Bases: Module

Fixed-lag ensemble smoother (Anderson & Anderson 1999; Nerger et al. 2012).

At each position s, runs the backward smoother only over the window [s, min(T-1, s + lag)]. Equivalent to :class:EnKS when lag ≥ T-1; equivalent to "no smoothing" when lag == 0.

The batch interface is the one exposed here — for an online / streaming setting (where states older than lag are finalised and discarded from memory) the user maintains the rolling buffer explicitly and calls :meth:smooth over the current window.

Attributes:

Name	Type	Description
lag	int	Number of backward steps. lag == 0 returns the filter analysis unchanged; large lag recovers the full EnKS.

Complexity: O(T · lag · (Nₑ² Nₓ + Nₑ³)) for the batch sweep.

Source code in src/filterax/_smoothers.py

class FixedLagSmoother(eqx.Module, strict=True):
    r"""Fixed-lag ensemble smoother (Anderson & Anderson 1999; Nerger et al. 2012).

    At each position ``s``, runs the backward smoother only over the
    window ``[s, min(T-1, s + lag)]``. Equivalent to :class:`EnKS` when
    ``lag ≥ T-1``; equivalent to "no smoothing" when ``lag == 0``.

    The batch interface is the one exposed here — for an online /
    streaming setting (where states older than ``lag`` are finalised and
    discarded from memory) the user maintains the rolling buffer
    explicitly and calls :meth:`smooth` over the current window.

    Attributes:
        lag: Number of backward steps. ``lag == 0`` returns the filter
            analysis unchanged; large ``lag`` recovers the full EnKS.

    Complexity: ``O(T · lag · (Nₑ² Nₓ + Nₑ³))`` for the batch sweep.
    """

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

    def __check_init__(self) -> None:
        if not isinstance(self.lag, int) or self.lag < 0:
            raise ValueError(f"lag must be a non-negative int; got {self.lag!r}.")

    def smooth(
        self,
        forecast_history: Float[Array, "T N_e N_x"],
        analysis_history: Float[Array, "T N_e N_x"],
    ) -> SmoothingResult:
        """Run the windowed backward pass.

        Args:
            forecast_history: Stacked forecast ensembles.
            analysis_history: Stacked analysis ensembles.

        Returns:
            :class:`SmoothingResult` whose ``smoothed_history[s]`` was
            produced by an EnKS-style backward pass restricted to the
            window ``[s, min(T-1, s + lag)]``.
        """
        _validate_history(analysis_history, forecast_history)
        T = analysis_history.shape[0]
        # Effective lag: clamp at ``T - 1`` (max usable lookahead) so a
        # large user-supplied ``lag`` does no extra work and the
        # ``T == 1, lag > 0`` case can never index ``forecast_history[1]``.
        lag_eff = min(self.lag, max(T - 1, 0))

        if lag_eff == 0:
            return SmoothingResult(
                particles=analysis_history[-1],
                smoothed_history=analysis_history,
            )

        def smooth_one_anchor(s: Int[Array, ""]) -> Float[Array, "N_e N_x"]:
            """Smoothed estimate at position ``s`` using lookahead ≤ ``lag_eff``."""
            end = jnp.minimum(s + lag_eff, T - 1)
            init = analysis_history[end]

            def step(
                carry: Float[Array, "N_e N_x"], k: Int[Array, ""]
            ) -> tuple[Float[Array, "N_e N_x"], None]:
                # k counts backward steps from the window end (0 = first step).
                # The nominal index processed at this step is `end - 1 - k`.
                t = end - 1 - k
                valid = t >= s
                # Gather safely even for out-of-window steps; we mask the carry.
                t_safe = jnp.maximum(t, 0)
                analysis_t = analysis_history[t_safe]
                forecast_t1 = forecast_history[t_safe + 1]
                candidate = _smoother_step(carry, analysis_t, forecast_t1)
                next_carry = jnp.where(valid, candidate, carry)
                return next_carry, None

            final, _ = jax.lax.scan(step, init, jnp.arange(lag_eff))
            return final

        smoothed = jax.vmap(smooth_one_anchor)(jnp.arange(T))
        return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

smooth(forecast_history: Float[Array, 'T N_e N_x'], analysis_history: Float[Array, 'T N_e N_x']) -> SmoothingResult

Run the windowed backward pass.

Parameters:

Name	Type	Description	Default
forecast_history	Float[Array, 'T N_e N_x']	Stacked forecast ensembles.	required
analysis_history	Float[Array, 'T N_e N_x']	Stacked analysis ensembles.	required

Returns:

Type	Description
SmoothingResult	class:SmoothingResult whose smoothed_history[s] was
SmoothingResult	produced by an EnKS-style backward pass restricted to the
SmoothingResult	window [s, min(T-1, s + lag)].

Source code in src/filterax/_smoothers.py

def smooth(
    self,
    forecast_history: Float[Array, "T N_e N_x"],
    analysis_history: Float[Array, "T N_e N_x"],
) -> SmoothingResult:
    """Run the windowed backward pass.

    Args:
        forecast_history: Stacked forecast ensembles.
        analysis_history: Stacked analysis ensembles.

    Returns:
        :class:`SmoothingResult` whose ``smoothed_history[s]`` was
        produced by an EnKS-style backward pass restricted to the
        window ``[s, min(T-1, s + lag)]``.
    """
    _validate_history(analysis_history, forecast_history)
    T = analysis_history.shape[0]
    # Effective lag: clamp at ``T - 1`` (max usable lookahead) so a
    # large user-supplied ``lag`` does no extra work and the
    # ``T == 1, lag > 0`` case can never index ``forecast_history[1]``.
    lag_eff = min(self.lag, max(T - 1, 0))

    if lag_eff == 0:
        return SmoothingResult(
            particles=analysis_history[-1],
            smoothed_history=analysis_history,
        )

    def smooth_one_anchor(s: Int[Array, ""]) -> Float[Array, "N_e N_x"]:
        """Smoothed estimate at position ``s`` using lookahead ≤ ``lag_eff``."""
        end = jnp.minimum(s + lag_eff, T - 1)
        init = analysis_history[end]

        def step(
            carry: Float[Array, "N_e N_x"], k: Int[Array, ""]
        ) -> tuple[Float[Array, "N_e N_x"], None]:
            # k counts backward steps from the window end (0 = first step).
            # The nominal index processed at this step is `end - 1 - k`.
            t = end - 1 - k
            valid = t >= s
            # Gather safely even for out-of-window steps; we mask the carry.
            t_safe = jnp.maximum(t, 0)
            analysis_t = analysis_history[t_safe]
            forecast_t1 = forecast_history[t_safe + 1]
            candidate = _smoother_step(carry, analysis_t, forecast_t1)
            next_carry = jnp.where(valid, candidate, carry)
            return next_carry, None

        final, _ = jax.lax.scan(step, init, jnp.arange(lag_eff))
        return final

    smoothed = jax.vmap(smooth_one_anchor)(jnp.arange(T))
    return SmoothingResult(particles=smoothed[-1], smoothed_history=smoothed)

IES

Bases: Module

Iterative Ensemble Smoother — Chen & Oliver (2013); Evensen et al. (2019).

Standalone iterative ensemble method that solves an inverse problem in a single window of observations. With perturbed observations d^j = y + ε^j, each ε^j ~ 𝒩(0, Γ) drawn once, each iteration is:

\[ \theta_{i+1}^j = (1 - \alpha)\, \theta_i^j + \alpha \big[\theta_0^j + K_i \big(d^j - G(\theta_i^j) + G_i (\theta_i^j - \theta_0^j)\big)\big], \qquad K_i = C^0_{\theta\theta} G_i^\top \big(G_i C^0_{\theta\theta} G_i^\top + \Gamma\big)^{-1} \]

where C⁰_θθ is the initial ensemble covariance and G_i the ensemble-average sensitivity at the current iterate. The anchor to θ_0^j (the initial ensemble member) is what distinguishes IES from :class:~filterax.EKI's drift-style iteration: rather than walking θ toward a MAP with a step-size schedule, IES re-targets θ_0 plus a Gauss-Newton correction every iteration. For a linear G the first iterate with α = 1 is already the ensemble posterior and later iterates stay there. α < 1 damps the iteration, Levenberg-Marquardt style, for strongly nonlinear G.

Use case: history matching, reservoir simulation, and any inverse problem where the forward model maps parameters to a full observation time series and a single backward pass on a sequential filter would be insufficient. For sequential filtering followed by a backward refinement, use :class:EnKS instead.

Attributes:

Name	Type	Description
n_iterations	int	Static iteration count.
step_size	float	Damping factor α ∈ (0, 1]. Default 1.0 (undamped Gauss-Newton).
seed	int	Integer for the default PRNG key of the perturbations.
base_key	PRNGKeyArray | None	Optional explicit PRNG key (takes precedence over seed).

Source code in src/filterax/_smoothers.py

class IES(eqx.Module, strict=True):
    r"""Iterative Ensemble Smoother — Chen & Oliver (2013); Evensen et al. (2019).

    Standalone iterative ensemble method that solves an inverse problem
    in a single window of observations. With perturbed observations
    ``d^j = y + ε^j``, each ``ε^j ~ 𝒩(0, Γ)`` drawn once, each iteration is:

    $$
    \theta_{i+1}^j = (1 - \alpha)\, \theta_i^j + \alpha \big[\theta_0^j
        + K_i \big(d^j - G(\theta_i^j) + G_i (\theta_i^j - \theta_0^j)\big)\big],
    \qquad
    K_i = C^0_{\theta\theta} G_i^\top
        \big(G_i C^0_{\theta\theta} G_i^\top + \Gamma\big)^{-1}
    $$

    where ``C⁰_θθ`` is the initial ensemble covariance and ``G_i`` the
    ensemble-average sensitivity at the current iterate. The anchor to
    ``θ_0^j`` (the *initial* ensemble member) is what distinguishes IES
    from :class:`~filterax.EKI`'s drift-style iteration: rather than walking
    ``θ`` toward a MAP with a step-size schedule, IES re-targets ``θ_0``
    plus a Gauss-Newton correction every iteration. For a linear ``G`` the
    first iterate with ``α = 1`` is already the ensemble posterior and later
    iterates stay there. ``α < 1`` damps the iteration, Levenberg-Marquardt
    style, for strongly nonlinear ``G``.

    Use case: history matching, reservoir simulation, and any inverse
    problem where the forward model maps parameters to a *full*
    observation time series and a single backward pass on a sequential
    filter would be insufficient. For sequential filtering followed by a
    backward refinement, use :class:`EnKS` instead.

    Attributes:
        n_iterations: Static iteration count.
        step_size: Damping factor ``α ∈ (0, 1]``. Default 1.0
            (undamped Gauss-Newton).
        seed: Integer for the default PRNG key of the perturbations.
        base_key: Optional explicit PRNG key (takes precedence over
            ``seed``).
    """

    n_iterations: int = eqx.field(static=True)
    step_size: float = eqx.field(static=True, default=1.0)
    seed: int = eqx.field(static=True, default=0)
    base_key: PRNGKeyArray | None = None

    def __check_init__(self) -> None:
        if not isinstance(self.n_iterations, int) or self.n_iterations < 1:
            raise ValueError(
                f"n_iterations must be a positive int; got {self.n_iterations!r}."
            )
        if not (0.0 < float(self.step_size) <= 1.0):
            raise ValueError(f"step_size must lie in (0, 1]; got {self.step_size!r}.")

    def solve(
        self,
        particles: Float[Array, "J N_p"],
        obs: Float[Array, " N_d"],
        noise_cov: lx.AbstractLinearOperator,
        forward_model: Callable[[Float[Array, " N_p"]], Float[Array, " N_d"]],
    ) -> AnalysisResult:
        """Run the iterative smoother to its fixed iteration count.

        Args:
            particles: Initial ensemble ``(J, Nₚ)``. Doubles as the
                anchor ``θ_0`` against which every subsequent iterate is
                compared.
            obs: Observation vector ``y ∈ ℝ^{N_d}``.
            noise_cov: Observation noise covariance ``Γ_y``.
            forward_model: ``G(θ): ℝ^{Nₚ} → ℝ^{N_d}`` applied to a single
                particle; vectorised internally with :func:`jax.vmap`.

        Returns:
            :class:`AnalysisResult` with the final iterate's particles.
        """
        check_ensemble_size(particles.shape[0])
        base_key = self.base_key if self.base_key is not None else jr.PRNGKey(self.seed)
        step_size = jnp.asarray(self.step_size, dtype=particles.dtype)

        # d^j = y + ε^j, fixed per member: redrawing per iteration changes the
        # objective each member is minimising, so the iteration cannot settle.
        perturbed_obs = perturbed_observations(
            base_key, obs, noise_cov, particles.shape[0]
        )

        def step(
            carry: Float[Array, "J N_p"], _: None
        ) -> tuple[Float[Array, "J N_p"], None]:
            new_carry = _ies_step(
                carry,
                particles,
                perturbed_obs,
                noise_cov,
                forward_model,
                step_size,
            )
            return new_carry, None

        final, _ = jax.lax.scan(step, particles, None, length=self.n_iterations)
        return AnalysisResult(particles=final)

solve(particles: Float[Array, 'J N_p'], obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator, forward_model: Callable[[Float[Array, ' N_p']], Float[Array, ' N_d']]) -> AnalysisResult

Run the iterative smoother to its fixed iteration count.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'J N_p']	Initial ensemble (J, Nₚ). Doubles as the anchor θ_0 against which every subsequent iterate is compared.	required
obs	Float[Array, ' N_d']	Observation vector y ∈ ℝ^{N_d}.	required
noise_cov	AbstractLinearOperator	Observation noise covariance Γ_y.	required
forward_model	Callable[[Float[Array, ' N_p']], Float[Array, ' N_d']]	G(θ): ℝ^{Nₚ} → ℝ^{N_d} applied to a single particle; vectorised internally with :func:jax.vmap.	required

Returns:

Type	Description
AnalysisResult	class:AnalysisResult with the final iterate's particles.

Source code in src/filterax/_smoothers.py

def solve(
    self,
    particles: Float[Array, "J N_p"],
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
    forward_model: Callable[[Float[Array, " N_p"]], Float[Array, " N_d"]],
) -> AnalysisResult:
    """Run the iterative smoother to its fixed iteration count.

    Args:
        particles: Initial ensemble ``(J, Nₚ)``. Doubles as the
            anchor ``θ_0`` against which every subsequent iterate is
            compared.
        obs: Observation vector ``y ∈ ℝ^{N_d}``.
        noise_cov: Observation noise covariance ``Γ_y``.
        forward_model: ``G(θ): ℝ^{Nₚ} → ℝ^{N_d}`` applied to a single
            particle; vectorised internally with :func:`jax.vmap`.

    Returns:
        :class:`AnalysisResult` with the final iterate's particles.
    """
    check_ensemble_size(particles.shape[0])
    base_key = self.base_key if self.base_key is not None else jr.PRNGKey(self.seed)
    step_size = jnp.asarray(self.step_size, dtype=particles.dtype)

    # d^j = y + ε^j, fixed per member: redrawing per iteration changes the
    # objective each member is minimising, so the iteration cannot settle.
    perturbed_obs = perturbed_observations(
        base_key, obs, noise_cov, particles.shape[0]
    )

    def step(
        carry: Float[Array, "J N_p"], _: None
    ) -> tuple[Float[Array, "J N_p"], None]:
        new_carry = _ies_step(
            carry,
            particles,
            perturbed_obs,
            noise_cov,
            forward_model,
            step_size,
        )
        return new_carry, None

    final, _ = jax.lax.scan(step, particles, None, length=self.n_iterations)
    return AnalysisResult(particles=final)

Exact linear-Gaussian smoother

When the dynamics and observation operator are linear and the noise is Gaussian, no ensemble is needed: KalmanSmoother runs the exact Kalman filter forward and the Rauch-Tung-Striebel smoother backward (through gaussx's kalman_filter and rts_smoother) and returns the filtered and smoothed Gaussian marginals in a KalmanSmoothingResult. It is the exact counterpart of EnKS and EnsembleRTS, and the ground truth they converge to as the ensemble grows.

Its prior sits at the first observation time and each transition spans the actual gap to the next one, so irregular observation times need no regridding; KalmanSmoother.from_sde builds the transitions from a gaussx SDE kernel such as gaussx.MaternSDE. Masked observation entries are marginalised out exactly, and a time with every entry masked only predicts.

filterax — differentiable ensemble Kalman filters and processes for JAX.

KalmanSmoother

Bases: Module

Exact Kalman filter and Rauch–Tung–Striebel smoother over irregular times.

The linear-Gaussian counterpart of :class:filterax.EnKS and :class:filterax.EnsembleRTS: a state observed at times t_1 < … < t_T evolves as a Gauss–Markov process between them,

\[ x_k = F_k x_{k-1} + \eta_k, \quad \eta_k \sim \mathcal{N}(0, \Sigma_{\eta,k}), \qquad y_k = H_k x_k + \varepsilon_k, \quad \varepsilon_k \sim \mathcal{N}(0, \Sigma_{\varepsilon,k}), \]

with x_1 ~ 𝒩(μ_x, Σ_x) at the first observation time, and :meth:smooth returns the exact marginals p(x_k | y_1:k) and p(x_k | y_1:T). F_k and Σ_η,k act over the actual gap t_k − t_{k−1}, so irregular times need no regridding; :meth:from_sde builds them from a gaussx SDE kernel.

Unlike :class:SquareRootKF, whose prior sits one transition before the first observation, the prior here is the state at t_1, so there are T − 1 transitions for T observation times.

Algorithm

1. Prepend F_1 = I and Σ_η,1 = 0, so that gaussx's predict-then-update filter observes x_1 under the prior 𝒩(μ_x, Σ_x) itself.
2. Make masked observation entries inert (zero row of H_k, unit variance in Σ_ε,k, zero datum), which equals deleting them; a time with every entry masked only predicts.
3. Filter forward with :func:gaussx.kalman_filter.
4. Smooth backward with :func:gaussx.rts_smoother, μ_x,k|T = μ_x,k|k + G_k (μ_x,k+1|T − μ_x,k+1|k) with the smoother gain G_k = Σ_x,k|k F_k+1ᵀ Σ_x,k+1|k⁻¹.

Attributes:

Name	Type	Description
transition	Float[Array, '*G N N']	F_k from time k − 1 to k, shape (T − 1, N, N), or (N, N) for every gap.
process_noise	Float[Array, '*G N N']	Σ_η,k, the same shapes as transition.
obs_model	Float[Array, '*T M N']	H_k, shape (T, M, N), or (M, N) for every time.
obs_noise	Float[Array, '*T M M']	Σ_ε,k, shape (T, M, M) or (M, M).
init_mean	Float[Array, ' N']	Prior mean μ_x of x_1, shape (N,).
init_cov	Float[Array, 'N N']	Prior covariance Σ_x of x_1, shape (N, N).

Complexity: O(T (N³ + M³)), sequential in T.

Examples:

A random walk observed at three times, the middle one missing; the smoother fills the gap from both neighbours:

>>> import jax.numpy as jnp
>>> from filterax import KalmanSmoother
>>> smoother = KalmanSmoother(
...     transition=jnp.eye(1),
...     process_noise=jnp.eye(1),
...     obs_model=jnp.eye(1),
...     obs_noise=0.01 * jnp.eye(1),
...     init_mean=jnp.zeros(1),
...     init_cov=jnp.eye(1),
... )
>>> obs = jnp.array([[1.0], [jnp.nan], [3.0]])
>>> mask = jnp.array([[True], [False], [True]])
>>> result = smoother.smooth(obs, mask)
>>> result.smoothed_means.shape
(3, 1)
>>> bool(jnp.isclose(result.smoothed_means[1, 0], 2.0, atol=0.05))
True

References

Rauch, Tung & Striebel (1965), AIAA Journal 3(8), 1445–1450. Särkkä (2013), Bayesian Filtering and Smoothing, ch. 8.

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

class KalmanSmoother(eqx.Module, strict=True):
    r"""Exact Kalman filter and Rauch–Tung–Striebel smoother over irregular times.

    The linear-Gaussian counterpart of :class:`filterax.EnKS` and
    :class:`filterax.EnsembleRTS`: a state observed at times
    ``t_1 < … < t_T`` evolves as a Gauss–Markov process between them,

    $$
    x_k = F_k x_{k-1} + \eta_k, \quad \eta_k \sim \mathcal{N}(0, \Sigma_{\eta,k}),
    \qquad
    y_k = H_k x_k + \varepsilon_k, \quad
    \varepsilon_k \sim \mathcal{N}(0, \Sigma_{\varepsilon,k}),
    $$

    with ``x_1 ~ 𝒩(μ_x, Σ_x)`` at the first observation time, and
    :meth:`smooth` returns the exact marginals ``p(x_k | y_1:k)`` and
    ``p(x_k | y_1:T)``. ``F_k`` and ``Σ_η,k`` act over the actual gap
    ``t_k − t_{k−1}``, so irregular times need no regridding;
    :meth:`from_sde` builds them from a gaussx SDE kernel.

    Unlike :class:`SquareRootKF`, whose prior sits one transition before
    the first observation, the prior here is the state at ``t_1``, so
    there are ``T − 1`` transitions for ``T`` observation times.

    Algorithm:
        1. Prepend ``F_1 = I`` and ``Σ_η,1 = 0``, so that gaussx's
           predict-then-update filter observes ``x_1`` under the prior
           ``𝒩(μ_x, Σ_x)`` itself.
        2. Make masked observation entries inert (zero row of ``H_k``,
           unit variance in ``Σ_ε,k``, zero datum), which equals
           deleting them; a time with every entry masked only predicts.
        3. Filter forward with :func:`gaussx.kalman_filter`.
        4. Smooth backward with :func:`gaussx.rts_smoother`,
           ``μ_x,k|T = μ_x,k|k + G_k (μ_x,k+1|T − μ_x,k+1|k)`` with the
           smoother gain ``G_k = Σ_x,k|k F_k+1ᵀ Σ_x,k+1|k⁻¹``.

    Attributes:
        transition: ``F_k`` from time ``k − 1`` to ``k``, shape
            ``(T − 1, N, N)``, or ``(N, N)`` for every gap.
        process_noise: ``Σ_η,k``, the same shapes as ``transition``.
        obs_model: ``H_k``, shape ``(T, M, N)``, or ``(M, N)`` for every
            time.
        obs_noise: ``Σ_ε,k``, shape ``(T, M, M)`` or ``(M, M)``.
        init_mean: Prior mean ``μ_x`` of ``x_1``, shape ``(N,)``.
        init_cov: Prior covariance ``Σ_x`` of ``x_1``, shape ``(N, N)``.

    Complexity: ``O(T (N³ + M³))``, sequential in ``T``.

    Examples:
        A random walk observed at three times, the middle one missing;
        the smoother fills the gap from both neighbours:

        >>> import jax.numpy as jnp
        >>> from filterax import KalmanSmoother
        >>> smoother = KalmanSmoother(
        ...     transition=jnp.eye(1),
        ...     process_noise=jnp.eye(1),
        ...     obs_model=jnp.eye(1),
        ...     obs_noise=0.01 * jnp.eye(1),
        ...     init_mean=jnp.zeros(1),
        ...     init_cov=jnp.eye(1),
        ... )
        >>> obs = jnp.array([[1.0], [jnp.nan], [3.0]])
        >>> mask = jnp.array([[True], [False], [True]])
        >>> result = smoother.smooth(obs, mask)
        >>> result.smoothed_means.shape
        (3, 1)
        >>> bool(jnp.isclose(result.smoothed_means[1, 0], 2.0, atol=0.05))
        True

    References:
        Rauch, Tung & Striebel (1965), AIAA Journal 3(8), 1445–1450.
        Särkkä (2013), Bayesian Filtering and Smoothing, ch. 8.
    """

    transition: Float[Array, "*G N N"]
    process_noise: Float[Array, "*G N N"]
    obs_model: Float[Array, "*T M N"]
    obs_noise: Float[Array, "*T M M"]
    init_mean: Float[Array, " N"]
    init_cov: Float[Array, "N N"]

    def __check_init__(self) -> None:
        for name in ("transition", "process_noise", "obs_model", "obs_noise"):
            ndim = getattr(self, name).ndim
            if ndim not in (2, 3):
                raise ValueError(
                    f"{name} must be one matrix or a (T, …) stack of them; "
                    f"got {ndim} dimensions."
                )

    @classmethod
    def from_sde(
        cls,
        sde: gaussx.SDEKernel,
        times: Float[Array, " T"],
        obs_model: Float[Array, "*T M N"],
        obs_noise: Float[Array, "*T M M"],
        init_mean: Float[Array, " N"] | None = None,
    ) -> KalmanSmoother:
        r"""Discretise a gaussx SDE kernel over the gaps of ``times``.

        ``F_k = exp(A Δt_k)`` and ``Σ_η,k = P_∞ − F_k P_∞ F_kᵀ`` on
        ``Δt_k = t_k − t_{k−1}``, where ``A`` is the SDE's drift and
        ``P_∞`` its stationary covariance, which is also the prior
        covariance of ``x_1``.

        Args:
            sde: A gaussx SDE kernel, e.g. ``gaussx.MaternSDE``.
            times: Observation times ``t_1 < … < t_T``, shape ``(T,)``.
            obs_model: ``H_k`` on the SDE state, shape ``(T, M, N)`` or
                ``(M, N)``.
            obs_noise: ``Σ_ε,k``, shape ``(T, M, M)`` or ``(M, M)``.
            init_mean: Prior mean ``μ_x`` of ``x_1``. Default ``None``:
                zero.

        Returns:
            The smoother for those times.
        """
        # Δt_k = t_k − t_{k−1}
        gaps = jnp.diff(times)
        # F_k = exp(A Δt_k), Σ_η,k = P_∞ − F_k P_∞ F_kᵀ
        transition, process_noise = sde.discretise_sequence(gaps)
        # Σ_x = P_∞
        init_cov = sde.sde_params().P_inf
        if init_cov is None:
            raise ValueError(
                f"{type(sde).__name__}.sde_params() reports P_inf=None (no "
                "stationary covariance), so from_sde cannot set the prior."
            )
        if init_mean is None:
            init_mean = jnp.zeros(init_cov.shape[0], init_cov.dtype)
        return cls(
            transition=transition,
            process_noise=process_noise,
            obs_model=obs_model,
            obs_noise=obs_noise,
            init_mean=init_mean,
            init_cov=init_cov,
        )

    def smooth(
        self,
        observations: Float[Array, "T M"],
        mask: Bool[Array, "T M"] | None = None,
    ) -> KalmanSmoothingResult:
        """Run the filter forward and the RTS smoother backward.

        Args:
            observations: ``y_1, …, y_T``, shape ``(T, M)``. Masked
                entries are never read, so they may be ``NaN``.
            mask: ``True`` where an entry was observed, shape ``(T, M)``.
                Default ``None``: every entry observed.

        Returns:
            :class:`KalmanSmoothingResult` with the filtered and
            smoothed marginals at every time and the log-likelihood.

        Raises:
            ValueError: If a stacked field's time axis does not match
                ``T`` (``T − 1`` for ``transition`` and
                ``process_noise``).
        """
        num_steps = observations.shape[0]
        if mask is None:
            mask = jnp.ones(observations.shape, dtype=bool)

        # --- sequences over T steps ---
        transition = _over(self.transition, num_steps - 1, "transition")
        process_noise = _over(self.process_noise, num_steps - 1, "process_noise")
        identity = jnp.eye(self.init_mean.shape[0], dtype=transition.dtype)
        # F_1 = I, Σ_η,1 = 0: x_1 is observed under the prior itself
        first_transition = einx.id("row col -> 1 row col", identity)
        transition = jnp.concatenate([first_transition, transition])
        first_noise = jnp.zeros_like(first_transition)
        process_noise = jnp.concatenate([first_noise, process_noise])
        obs_model = _over(self.obs_model, num_steps, "obs_model")
        obs_noise = _over(self.obs_noise, num_steps, "obs_noise")

        # --- missing data ---
        obs_model, obs_noise, observations = _inert_masked_entries(
            obs_model, obs_noise, observations, mask
        )
        # a time with no observed entry only predicts
        step_observed = einx.any("T M -> T", mask)

        # --- filter and smoother ---
        state = gaussx.kalman_filter(
            transition,
            obs_model,
            process_noise,
            obs_noise,
            observations,
            self.init_mean,
            self.init_cov,
            mask=step_observed,
            solver=gaussx.DenseSolver(),
        )
        smoothed_means, smoothed_covs = gaussx.rts_smoother(
            state, transition, process_noise, solver=gaussx.DenseSolver()
        )

        # --- log-likelihood ---
        # each inert entry of an updated step adds −½ log 2π to gaussx's sum
        missing_per_step = einx.sum("T M -> T", ~mask)
        missing_updated = jnp.where(step_observed, missing_per_step, 0)
        num_missing = einx.sum("T ->", missing_updated)
        log_likelihood = state.log_likelihood + 0.5 * math.log(2 * math.pi) * (
            num_missing
        )
        return KalmanSmoothingResult(
            filtered_means=state.filtered_means,
            filtered_covs=state.filtered_covs,
            smoothed_means=smoothed_means,
            smoothed_covs=smoothed_covs,
            log_likelihood=log_likelihood,
        )

from_sde(sde: gaussx.SDEKernel, times: Float[Array, ' T'], obs_model: Float[Array, '*T M N'], obs_noise: Float[Array, '*T M M'], init_mean: Float[Array, ' N'] | None = None) -> KalmanSmoother classmethod

Discretise a gaussx SDE kernel over the gaps of times.

F_k = exp(A Δt_k) and Σ_η,k = P_∞ − F_k P_∞ F_kᵀ on Δt_k = t_k − t_{k−1}, where A is the SDE's drift and P_∞ its stationary covariance, which is also the prior covariance of x_1.

Parameters:

Name	Type	Description	Default
sde	SDEKernel	A gaussx SDE kernel, e.g. gaussx.MaternSDE.	required
times	Float[Array, ' T']	Observation times t_1 < … < t_T, shape (T,).	required
obs_model	Float[Array, '*T M N']	H_k on the SDE state, shape (T, M, N) or (M, N).	required
obs_noise	Float[Array, '*T M M']	Σ_ε,k, shape (T, M, M) or (M, M).	required
init_mean	Float[Array, ' N'] | None	Prior mean μ_x of x_1. Default None: zero.	None

Returns:

Type	Description
KalmanSmoother	The smoother for those times.

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

@classmethod
def from_sde(
    cls,
    sde: gaussx.SDEKernel,
    times: Float[Array, " T"],
    obs_model: Float[Array, "*T M N"],
    obs_noise: Float[Array, "*T M M"],
    init_mean: Float[Array, " N"] | None = None,
) -> KalmanSmoother:
    r"""Discretise a gaussx SDE kernel over the gaps of ``times``.

    ``F_k = exp(A Δt_k)`` and ``Σ_η,k = P_∞ − F_k P_∞ F_kᵀ`` on
    ``Δt_k = t_k − t_{k−1}``, where ``A`` is the SDE's drift and
    ``P_∞`` its stationary covariance, which is also the prior
    covariance of ``x_1``.

    Args:
        sde: A gaussx SDE kernel, e.g. ``gaussx.MaternSDE``.
        times: Observation times ``t_1 < … < t_T``, shape ``(T,)``.
        obs_model: ``H_k`` on the SDE state, shape ``(T, M, N)`` or
            ``(M, N)``.
        obs_noise: ``Σ_ε,k``, shape ``(T, M, M)`` or ``(M, M)``.
        init_mean: Prior mean ``μ_x`` of ``x_1``. Default ``None``:
            zero.

    Returns:
        The smoother for those times.
    """
    # Δt_k = t_k − t_{k−1}
    gaps = jnp.diff(times)
    # F_k = exp(A Δt_k), Σ_η,k = P_∞ − F_k P_∞ F_kᵀ
    transition, process_noise = sde.discretise_sequence(gaps)
    # Σ_x = P_∞
    init_cov = sde.sde_params().P_inf
    if init_cov is None:
        raise ValueError(
            f"{type(sde).__name__}.sde_params() reports P_inf=None (no "
            "stationary covariance), so from_sde cannot set the prior."
        )
    if init_mean is None:
        init_mean = jnp.zeros(init_cov.shape[0], init_cov.dtype)
    return cls(
        transition=transition,
        process_noise=process_noise,
        obs_model=obs_model,
        obs_noise=obs_noise,
        init_mean=init_mean,
        init_cov=init_cov,
    )

smooth(observations: Float[Array, 'T M'], mask: Bool[Array, 'T M'] | None = None) -> KalmanSmoothingResult

Run the filter forward and the RTS smoother backward.

Parameters:

Name	Type	Description	Default
observations	Float[Array, 'T M']	y_1, …, y_T, shape (T, M). Masked entries are never read, so they may be NaN.	required
mask	Bool[Array, 'T M'] | None	True where an entry was observed, shape (T, M). Default None: every entry observed.	None

Returns:

Type	Description
KalmanSmoothingResult	class:KalmanSmoothingResult with the filtered and
KalmanSmoothingResult	smoothed marginals at every time and the log-likelihood.

Raises:

Type	Description
ValueError	If a stacked field's time axis does not match T (T − 1 for transition and process_noise).

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

def smooth(
    self,
    observations: Float[Array, "T M"],
    mask: Bool[Array, "T M"] | None = None,
) -> KalmanSmoothingResult:
    """Run the filter forward and the RTS smoother backward.

    Args:
        observations: ``y_1, …, y_T``, shape ``(T, M)``. Masked
            entries are never read, so they may be ``NaN``.
        mask: ``True`` where an entry was observed, shape ``(T, M)``.
            Default ``None``: every entry observed.

    Returns:
        :class:`KalmanSmoothingResult` with the filtered and
        smoothed marginals at every time and the log-likelihood.

    Raises:
        ValueError: If a stacked field's time axis does not match
            ``T`` (``T − 1`` for ``transition`` and
            ``process_noise``).
    """
    num_steps = observations.shape[0]
    if mask is None:
        mask = jnp.ones(observations.shape, dtype=bool)

    # --- sequences over T steps ---
    transition = _over(self.transition, num_steps - 1, "transition")
    process_noise = _over(self.process_noise, num_steps - 1, "process_noise")
    identity = jnp.eye(self.init_mean.shape[0], dtype=transition.dtype)
    # F_1 = I, Σ_η,1 = 0: x_1 is observed under the prior itself
    first_transition = einx.id("row col -> 1 row col", identity)
    transition = jnp.concatenate([first_transition, transition])
    first_noise = jnp.zeros_like(first_transition)
    process_noise = jnp.concatenate([first_noise, process_noise])
    obs_model = _over(self.obs_model, num_steps, "obs_model")
    obs_noise = _over(self.obs_noise, num_steps, "obs_noise")

    # --- missing data ---
    obs_model, obs_noise, observations = _inert_masked_entries(
        obs_model, obs_noise, observations, mask
    )
    # a time with no observed entry only predicts
    step_observed = einx.any("T M -> T", mask)

    # --- filter and smoother ---
    state = gaussx.kalman_filter(
        transition,
        obs_model,
        process_noise,
        obs_noise,
        observations,
        self.init_mean,
        self.init_cov,
        mask=step_observed,
        solver=gaussx.DenseSolver(),
    )
    smoothed_means, smoothed_covs = gaussx.rts_smoother(
        state, transition, process_noise, solver=gaussx.DenseSolver()
    )

    # --- log-likelihood ---
    # each inert entry of an updated step adds −½ log 2π to gaussx's sum
    missing_per_step = einx.sum("T M -> T", ~mask)
    missing_updated = jnp.where(step_observed, missing_per_step, 0)
    num_missing = einx.sum("T ->", missing_updated)
    log_likelihood = state.log_likelihood + 0.5 * math.log(2 * math.pi) * (
        num_missing
    )
    return KalmanSmoothingResult(
        filtered_means=state.filtered_means,
        filtered_covs=state.filtered_covs,
        smoothed_means=smoothed_means,
        smoothed_covs=smoothed_covs,
        log_likelihood=log_likelihood,
    )

KalmanSmoothingResult

Bases: Module

Output of :meth:KalmanSmoother.smooth.

The Gaussian marginals of the state at every observation time, both filtered (conditioned on the past) and smoothed (conditioned on the whole record), along a leading time axis. Cross-time covariances are not kept.

Attributes:

Name	Type	Description
filtered_means	Float[Array, 'T N']	μ_x,k|k for k = 1, …, T, shape (T, N).
filtered_covs	Float[Array, 'T N N']	Σ_x,k|k, shape (T, N, N).
smoothed_means	Float[Array, 'T N']	μ_x,k|T, shape (T, N).
smoothed_covs	Float[Array, 'T N N']	Σ_x,k|T, shape (T, N, N).
log_likelihood	Float[Array, '']	Marginal log-likelihood log p(y_1:T) of the observed entries; masked entries do not contribute.

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

class KalmanSmoothingResult(eqx.Module, strict=True):
    """Output of :meth:`KalmanSmoother.smooth`.

    The Gaussian marginals of the state at every observation time, both
    filtered (conditioned on the past) and smoothed (conditioned on the
    whole record), along a leading time axis. Cross-time covariances are
    not kept.

    Attributes:
        filtered_means: ``μ_x,k|k`` for ``k = 1, …, T``, shape ``(T, N)``.
        filtered_covs: ``Σ_x,k|k``, shape ``(T, N, N)``.
        smoothed_means: ``μ_x,k|T``, shape ``(T, N)``.
        smoothed_covs: ``Σ_x,k|T``, shape ``(T, N, N)``.
        log_likelihood: Marginal log-likelihood ``log p(y_1:T)`` of the
            observed entries; masked entries do not contribute.
    """

    filtered_means: Float[Array, "T N"]
    filtered_covs: Float[Array, "T N N"]
    smoothed_means: Float[Array, "T N"]
    smoothed_covs: Float[Array, "T N N"]
    log_likelihood: Float[Array, ""]