Inflation

Counteract spread collapse from finite-ensemble, model-error, or localization-side-effect sources. Two flavours of primitive:

• Drop-in inflators — MultiplicativeInflator, RTPS, RTPP, AdditiveInflator all implement AbstractInflator and plug straight into the Layer-2 assimilation loops (filterax.ETKF, LETKF, …).
• Helper primitives — inflate_adaptive returns a posterior on the multiplicative factor λ (the caller composes it with a fresh MultiplicativeInflator(factor=μ_post) on the next cycle); ledoit_wolf_shrinkage returns a regularised covariance matrix, not an inflated ensemble. Neither is an AbstractInflator, so they are driven from custom loops rather than dropped into L2 models.

The classic trio of functional primitives — inflate_multiplicative, inflate_rtps, inflate_rtpp — delegates to the corresponding gaussx primitives (gaussx.inflate_multiplicative, gaussx.inflate_rtps, gaussx.inflate_rtpp); filterax adds the EnKF-flavoured conventions and the inflator classes on top.

Picking an inflator

Method	Mean-preserving?	Tunable	Use when
MultiplicativeInflator / inflate_multiplicative	Yes	factor	Simple scalar inflation; production default
RTPS / inflate_rtps	Yes (per-variable spread)	alpha ∈ [0, 1]	Per-variable adaptive recovery of forecast spread (Whitaker & Hamill 2012)
RTPP / inflate_rtpp	Yes	alpha ∈ [0, 1]	Full-anomaly blend that preserves inter-variable correlation (Zhang et al. 2004)
AdditiveInflator / inflate_additive	Yes (after recentring)	Q_add	Explicit model-error injection (Hamill & Whitaker 2005)
inflate_adaptive	(Anderson 2009 prior update)	prior (μ_λ, σ²_λ)	Self-tuning multiplicative inflation; carry posterior across cycles
ledoit_wolf_shrinkage	n/a (returns covariance)	none (analytic optimum)	Regularise the rank-deficient sample covariance when Nₑ ≪ Nₓ

Inflator classes

Drop-in AbstractInflator implementations, accepted by every Layer-2 assimilation loop via its inflator= argument.

filterax — differentiable ensemble Kalman filters and processes for JAX.

MultiplicativeInflator

Bases: AbstractInflator

Multiplicative inflation X′ ← λ X′.

Ignores forecast_particles and any extra kwargs.

Attributes:

Name	Type	Description
factor	float	Inflation factor λ > 0; values in [1.01, 1.10] are typical.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import MultiplicativeInflator
>>> inflator = MultiplicativeInflator(factor=2.0)
>>> inflator(jnp.array([[0.0], [2.0]]))  # mean 1 preserved
Array([[-1.],
       [ 3.]], dtype=float32)

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

class MultiplicativeInflator(AbstractInflator, strict=True):
    r"""Multiplicative inflation ``X′ ← λ X′``.

    Ignores ``forecast_particles`` and any extra kwargs.

    Attributes:
        factor: Inflation factor ``λ > 0``; values in
            ``[1.01, 1.10]`` are typical.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import MultiplicativeInflator
        >>> inflator = MultiplicativeInflator(factor=2.0)
        >>> inflator(jnp.array([[0.0], [2.0]]))  # mean 1 preserved
        Array([[-1.],
               [ 3.]], dtype=float32)
    """

    factor: float = eqx.field(static=True)

    def __call__(
        self,
        particles: Float[Array, "N_e N_x"],
        forecast_particles: Float[Array, "N_e N_x"] | None = None,
        **_: Any,
    ) -> Float[Array, "N_e N_x"]:
        del forecast_particles
        return inflate_multiplicative(particles, self.factor)

RTPS

Bases: AbstractInflator

Relaxation to Prior Spread (Whitaker & Hamill 2012).

Requires forecast_particles — raises ValueError when missing rather than silently no-opping.

Attributes:

Name	Type	Description
alpha	float	Relaxation coefficient in [0, 1]. 0 keeps the analysis spread; 1 restores the forecast spread.

Examples:

>>> import jax.numpy as jnp
>>> from filterax import RTPS
>>> inflator = RTPS(alpha=1.0)
>>> forecast = jnp.array([[-2.0], [2.0]])
>>> analysis = jnp.array([[-1.0], [1.0]])
>>> inflator(analysis, forecast)  # full prior spread restored
Array([[-2.],
       [ 2.]], dtype=float32)

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

class RTPS(AbstractInflator, strict=True):
    r"""Relaxation to Prior Spread (Whitaker & Hamill 2012).

    Requires ``forecast_particles`` — raises ``ValueError`` when missing
    rather than silently no-opping.

    Attributes:
        alpha: Relaxation coefficient in ``[0, 1]``. ``0`` keeps the
            analysis spread; ``1`` restores the forecast spread.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import RTPS
        >>> inflator = RTPS(alpha=1.0)
        >>> forecast = jnp.array([[-2.0], [2.0]])
        >>> analysis = jnp.array([[-1.0], [1.0]])
        >>> inflator(analysis, forecast)  # full prior spread restored
        Array([[-2.],
               [ 2.]], dtype=float32)
    """

    alpha: float = eqx.field(static=True)

    def __call__(
        self,
        particles: Float[Array, "N_e N_x"],
        forecast_particles: Float[Array, "N_e N_x"] | None = None,
        **_: Any,
    ) -> Float[Array, "N_e N_x"]:
        if forecast_particles is None:
            raise ValueError("RTPS requires the forecast ensemble.")
        return inflate_rtps(particles, forecast_particles, self.alpha)

RTPP

Bases: AbstractInflator

Relaxation to Prior Perturbations (Zhang et al. 2004).

Requires forecast_particles.

Attributes:

Name	Type	Description
alpha	float	Relaxation coefficient in [0, 1].

Examples:

>>> import jax.numpy as jnp
>>> from filterax import RTPP
>>> inflator = RTPP(alpha=0.5)
>>> forecast = jnp.array([[-2.0], [2.0]])
>>> analysis = jnp.array([[-1.0], [1.0]])
>>> inflator(analysis, forecast)  # blend anomalies
Array([[-1.5],
       [ 1.5]], dtype=float32)

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

class RTPP(AbstractInflator, strict=True):
    r"""Relaxation to Prior Perturbations (Zhang et al. 2004).

    Requires ``forecast_particles``.

    Attributes:
        alpha: Relaxation coefficient in ``[0, 1]``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import RTPP
        >>> inflator = RTPP(alpha=0.5)
        >>> forecast = jnp.array([[-2.0], [2.0]])
        >>> analysis = jnp.array([[-1.0], [1.0]])
        >>> inflator(analysis, forecast)  # blend anomalies
        Array([[-1.5],
               [ 1.5]], dtype=float32)
    """

    alpha: float = eqx.field(static=True)

    def __call__(
        self,
        particles: Float[Array, "N_e N_x"],
        forecast_particles: Float[Array, "N_e N_x"] | None = None,
        **_: Any,
    ) -> Float[Array, "N_e N_x"]:
        if forecast_particles is None:
            raise ValueError("RTPP requires the forecast ensemble.")
        return inflate_rtpp(particles, forecast_particles, self.alpha)

AdditiveInflator

Bases: AbstractInflator

Additive Gaussian inflation X′ ← X′ + ε, ε ~ 𝒩(0, Q_add).

Wraps :func:filterax.inflate_additive. Per-call independence:

• Inside the L2 run loop (filterax.ETKF / LETKF / …): the loop passes step= through **kwargs and this inflator computes key = jr.fold_in(self.base_key, step) so successive assimilation windows draw independent Gaussian perturbations.
• In your own loop: pass step= yourself (or pass an explicit key= kwarg, which takes precedence over base_key / step). Calling without either reuses self.base_key and yields the same draw every time — fine for one-off use, but a footgun if you reuse the inflator across cycles.

Attributes:

Name	Type	Description
noise_cov	AbstractLinearOperator	Model-error covariance Q_add.
base_key	PRNGKeyArray	PRNG key seed; folded against the per-call step (or used as-is when no step / key is supplied).

Examples:

>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import lineax as lx
>>> from filterax import AdditiveInflator
>>> inflator = AdditiveInflator(
...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(2)),
...     base_key=jr.key(0),
... )
>>> out = inflator(jnp.zeros((4, 2)), step=0)
>>> out.shape
(4, 2)
>>> bool(jnp.allclose(out.mean(axis=0), 0.0, atol=1e-6))  # mean kept
True

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

class AdditiveInflator(AbstractInflator, strict=True):
    r"""Additive Gaussian inflation ``X′ ← X′ + ε, ε ~ 𝒩(0, Q_add)``.

    Wraps :func:`filterax.inflate_additive`. Per-call independence:

    * **Inside the L2 run loop** (``filterax.ETKF`` / ``LETKF`` / …):
      the loop passes ``step=`` through ``**kwargs`` and this inflator
      computes ``key = jr.fold_in(self.base_key, step)`` so successive
      assimilation windows draw independent Gaussian perturbations.
    * **In your own loop**: pass ``step=`` yourself (or pass an explicit
      ``key=`` kwarg, which takes precedence over ``base_key`` /
      ``step``). Calling without either reuses ``self.base_key`` and
      yields the *same* draw every time — fine for one-off use, but a
      footgun if you reuse the inflator across cycles.

    Attributes:
        noise_cov: Model-error covariance ``Q_add``.
        base_key: PRNG key seed; folded against the per-call ``step``
            (or used as-is when no step / key is supplied).

    Examples:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> import lineax as lx
        >>> from filterax import AdditiveInflator
        >>> inflator = AdditiveInflator(
        ...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(2)),
        ...     base_key=jr.key(0),
        ... )
        >>> out = inflator(jnp.zeros((4, 2)), step=0)
        >>> out.shape
        (4, 2)
        >>> bool(jnp.allclose(out.mean(axis=0), 0.0, atol=1e-6))  # mean kept
        True
    """

    noise_cov: lx.AbstractLinearOperator
    base_key: PRNGKeyArray

    def __call__(
        self,
        particles: Float[Array, "N_e N_x"],
        forecast_particles: Float[Array, "N_e N_x"] | None = None,
        *,
        key: PRNGKeyArray | None = None,
        step: int | jnp.ndarray | None = None,
        **_: Any,
    ) -> Float[Array, "N_e N_x"]:
        del forecast_particles
        if key is not None:
            call_key = key
        elif step is not None:
            call_key = jr.fold_in(self.base_key, jnp.asarray(step, dtype=jnp.int32))
        else:
            call_key = self.base_key
        return inflate_additive(call_key, particles, self.noise_cov)

Functional primitives

The pure functions underneath the classes, plus the two helpers (inflate_adaptive, ledoit_wolf_shrinkage) that have no class form.

filterax — differentiable ensemble Kalman filters and processes for JAX.

inflate_multiplicative(particles: Float[Array, 'N_e N_x'], factor: float) -> Float[Array, 'N_e N_x']

Multiplicative inflation X′ ← λ X′.

For each member x⁽ʲ⁾_inflated = x̄ + λ (x⁽ʲ⁾ − x̄). The inflated covariance is λ² P; the analysis mean is preserved. Typical values λ ∈ [1.01, 1.10] for operational NWP. Delegates to :func:gaussx.inflate_multiplicative.

Parameters:

Name	Type	Description	Default
particles	Float[Array, 'N_e N_x']	Ensemble of shape (Nₑ, Nₓ).	required
factor	float	Inflation factor λ. 1.0 is the identity; λ > 1 inflates spread; λ < 1 deflates.	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	Inflated ensemble of shape (Nₑ, Nₓ).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import inflate_multiplicative
>>> ens = jnp.array([[0.0, 0.0], [2.0, 2.0]])
>>> inflate_multiplicative(ens, 1.5)  # mean (1, 1) preserved
Array([[-0.5, -0.5],
       [ 2.5,  2.5]], dtype=float32)

Reference

Anderson, J. L. & Anderson, S. L. (1999). A Monte Carlo implementation of the nonlinear filtering problem. Mon. Wea. Rev., 127, 2741–2758.

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

def inflate_multiplicative(
    particles: Float[Array, "N_e N_x"],
    factor: float,
) -> Float[Array, "N_e N_x"]:
    r"""Multiplicative inflation ``X′ ← λ X′``.

    For each member ``x⁽ʲ⁾_inflated = x̄ + λ (x⁽ʲ⁾ − x̄)``. The inflated
    covariance is ``λ² P``; the analysis mean is preserved. Typical
    values ``λ ∈ [1.01, 1.10]`` for operational NWP. Delegates to
    :func:`gaussx.inflate_multiplicative`.

    Args:
        particles: Ensemble of shape ``(Nₑ, Nₓ)``.
        factor: Inflation factor ``λ``. ``1.0`` is the identity;
            ``λ > 1`` inflates spread; ``λ < 1`` deflates.

    Returns:
        Inflated ensemble of shape ``(Nₑ, Nₓ)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import inflate_multiplicative
        >>> ens = jnp.array([[0.0, 0.0], [2.0, 2.0]])
        >>> inflate_multiplicative(ens, 1.5)  # mean (1, 1) preserved
        Array([[-0.5, -0.5],
               [ 2.5,  2.5]], dtype=float32)

    Reference:
        Anderson, J. L. & Anderson, S. L. (1999). *A Monte Carlo
        implementation of the nonlinear filtering problem.* Mon. Wea.
        Rev., 127, 2741–2758.
    """
    return gaussx.inflate_multiplicative(particles, factor)

inflate_rtps(analysis_particles: Float[Array, 'N_e N_x'], forecast_particles: Float[Array, 'N_e N_x'], alpha: float) -> Float[Array, 'N_e N_x']

Relaxation to Prior Spread (Whitaker & Hamill 2012).

Per-variable target spread is a convex blend of analysis and forecast standard deviations:

\[ \sigma^{\text{target}}_{i} = (1 - \alpha)\, \sigma^{a}_{i} + \alpha\, \sigma^{f}_{i}. \]

Each analysis anomaly is then rescaled so the resulting standard deviation matches the target:

\[ x^{(j)}_{\text{relaxed}} = \bar{x}^{a} + \frac{\sigma^{\text{target}}_{i}}{\sigma^{a}_{i}} \left( x^{(j)}_{a} - \bar{x}^{a} \right). \]

Spatially adaptive (per-variable factor), mean-preserving, and parameterised by a single coefficient α. α = 0 is the identity; α = 1 restores the full forecast spread. Delegates to :func:gaussx.inflate_rtps (which calls the coefficient beta and writes the equivalent form scale = (1 − β) + β σᶠ/σᵃ).

Parameters:

Name	Type	Description	Default
analysis_particles	Float[Array, 'N_e N_x']	Posterior ensemble (Nₑ, Nₓ).	required
forecast_particles	Float[Array, 'N_e N_x']	Prior (forecast) ensemble (Nₑ, Nₓ).	required
alpha	float	Relaxation coefficient in [0, 1].	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	Relaxed analysis ensemble of shape (Nₑ, Nₓ).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import inflate_rtps
>>> forecast = jnp.array([[-2.0], [2.0]])
>>> analysis = jnp.array([[-1.0], [1.0]])
>>> inflate_rtps(analysis, forecast, alpha=1.0)  # full prior spread
Array([[-2.],
       [ 2.]], dtype=float32)

Reference

Whitaker, J. S. & Hamill, T. M. (2012). Evaluating methods to account for system errors in ensemble data assimilation. Mon. Wea. Rev., 140, 3078–3089.

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

def inflate_rtps(
    analysis_particles: Float[Array, "N_e N_x"],
    forecast_particles: Float[Array, "N_e N_x"],
    alpha: float,
) -> Float[Array, "N_e N_x"]:
    r"""Relaxation to Prior Spread (Whitaker & Hamill 2012).

    Per-variable target spread is a convex blend of analysis and
    forecast standard deviations:

    $$
    \sigma^{\text{target}}_{i} =
        (1 - \alpha)\, \sigma^{a}_{i} + \alpha\, \sigma^{f}_{i}.
    $$

    Each analysis anomaly is then rescaled so the resulting standard
    deviation matches the target:

    $$
    x^{(j)}_{\text{relaxed}} = \bar{x}^{a}
        + \frac{\sigma^{\text{target}}_{i}}{\sigma^{a}_{i}}
          \left( x^{(j)}_{a} - \bar{x}^{a} \right).
    $$

    Spatially adaptive (per-variable factor), mean-preserving, and
    parameterised by a single coefficient ``α``. ``α = 0`` is the
    identity; ``α = 1`` restores the full forecast spread. Delegates to
    :func:`gaussx.inflate_rtps` (which calls the coefficient ``beta``
    and writes the equivalent form
    ``scale = (1 − β) + β σᶠ/σᵃ``).

    Args:
        analysis_particles: Posterior ensemble ``(Nₑ, Nₓ)``.
        forecast_particles: Prior (forecast) ensemble ``(Nₑ, Nₓ)``.
        alpha: Relaxation coefficient in ``[0, 1]``.

    Returns:
        Relaxed analysis ensemble of shape ``(Nₑ, Nₓ)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import inflate_rtps
        >>> forecast = jnp.array([[-2.0], [2.0]])
        >>> analysis = jnp.array([[-1.0], [1.0]])
        >>> inflate_rtps(analysis, forecast, alpha=1.0)  # full prior spread
        Array([[-2.],
               [ 2.]], dtype=float32)

    Reference:
        Whitaker, J. S. & Hamill, T. M. (2012). *Evaluating methods to
        account for system errors in ensemble data assimilation.*
        Mon. Wea. Rev., 140, 3078–3089.
    """
    return gaussx.inflate_rtps(analysis_particles, forecast_particles, alpha)

inflate_rtpp(analysis_particles: Float[Array, 'N_e N_x'], forecast_particles: Float[Array, 'N_e N_x'], alpha: float) -> Float[Array, 'N_e N_x']

Relaxation to Prior Perturbations (Zhang et al. 2004).

Convex combination of analysis and forecast anomaly matrices:

\[ X^{\prime}_{\text{relaxed}} = (1 - \alpha)\, X^{\prime}_{a} + \alpha\, X^{\prime}_{f}. \]

Unlike RTPS (which acts on per-variable spread), RTPP operates on the full anomaly matrix and so blends inter-variable correlation structure between forecast and analysis. Mean-preserving. Delegates to :func:gaussx.inflate_rtpp.

Parameters:

Name	Type	Description	Default
analysis_particles	Float[Array, 'N_e N_x']	Posterior ensemble (Nₑ, Nₓ).	required
forecast_particles	Float[Array, 'N_e N_x']	Prior (forecast) ensemble (Nₑ, Nₓ).	required
alpha	float	Relaxation coefficient in [0, 1].	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	Relaxed analysis ensemble of shape (Nₑ, Nₓ).

Examples:

>>> import jax.numpy as jnp
>>> from filterax import inflate_rtpp
>>> forecast = jnp.array([[-2.0], [2.0]])
>>> analysis = jnp.array([[-1.0], [1.0]])
>>> inflate_rtpp(analysis, forecast, alpha=0.5)  # blend anomalies
Array([[-1.5],
       [ 1.5]], dtype=float32)

Reference

Zhang, F., Snyder, C., & Sun, J. (2004). Impacts of initial estimate and observation availability on convective-scale data assimilation. Mon. Wea. Rev., 132, 1238–1253.

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

def inflate_rtpp(
    analysis_particles: Float[Array, "N_e N_x"],
    forecast_particles: Float[Array, "N_e N_x"],
    alpha: float,
) -> Float[Array, "N_e N_x"]:
    r"""Relaxation to Prior Perturbations (Zhang et al. 2004).

    Convex combination of analysis and forecast anomaly matrices:

    $$
    X^{\prime}_{\text{relaxed}} =
        (1 - \alpha)\, X^{\prime}_{a} + \alpha\, X^{\prime}_{f}.
    $$

    Unlike RTPS (which acts on per-variable spread), RTPP operates on
    the full anomaly matrix and so blends inter-variable correlation
    structure between forecast and analysis. Mean-preserving.
    Delegates to :func:`gaussx.inflate_rtpp`.

    Args:
        analysis_particles: Posterior ensemble ``(Nₑ, Nₓ)``.
        forecast_particles: Prior (forecast) ensemble ``(Nₑ, Nₓ)``.
        alpha: Relaxation coefficient in ``[0, 1]``.

    Returns:
        Relaxed analysis ensemble of shape ``(Nₑ, Nₓ)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax import inflate_rtpp
        >>> forecast = jnp.array([[-2.0], [2.0]])
        >>> analysis = jnp.array([[-1.0], [1.0]])
        >>> inflate_rtpp(analysis, forecast, alpha=0.5)  # blend anomalies
        Array([[-1.5],
               [ 1.5]], dtype=float32)

    Reference:
        Zhang, F., Snyder, C., & Sun, J. (2004). *Impacts of initial
        estimate and observation availability on convective-scale data
        assimilation.* Mon. Wea. Rev., 132, 1238–1253.
    """
    return gaussx.inflate_rtpp(analysis_particles, forecast_particles, alpha)

inflate_additive(key: PRNGKeyArray, particles: Float[Array, 'N_e N_x'], noise_cov: lx.AbstractLinearOperator) -> Float[Array, 'N_e N_x']

Additive inflation X′ ← X′ + ε, ε ~ 𝒩(0, Q_add).

Injects stochastic model-error perturbations into every ensemble member's anomaly. Unlike multiplicative inflation, the noise direction is set by Q_add rather than by the ensemble's own spread — useful when the dominant model-error mode is not represented by the ensemble (e.g. structurally underdispersive in a particular subspace). Mean-preserving by construction (we subtract the empirical mean of the noise).

Common sources for Q_add: climatological variability, lagged forecast differences, or stochastic-physics tendency variances.

Sampling delegates to :func:filterax.perturbed_observations's diagonal-aware path, so a diagonal Q_add never materialises.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key consumed by this call.	required
particles	Float[Array, 'N_e N_x']	Ensemble of shape (Nₑ, Nₓ).	required
noise_cov	AbstractLinearOperator	Additive-noise covariance Q_add as a linear operator.	required

Returns:

Type	Description
Float[Array, 'N_e N_x']	Inflated ensemble of shape (Nₑ, Nₓ).

Reference

Hamill, T. M. & Whitaker, J. S. (2005). Accounting for the error due to unresolved scales in ensemble data assimilation: A comparison of different approaches. Mon. Wea. Rev., 133, 3132-3147.

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

def inflate_additive(
    key: PRNGKeyArray,
    particles: Float[Array, "N_e N_x"],
    noise_cov: lx.AbstractLinearOperator,
) -> Float[Array, "N_e N_x"]:
    r"""Additive inflation ``X′ ← X′ + ε,  ε ~ 𝒩(0, Q_add)``.

    Injects stochastic model-error perturbations into every ensemble
    member's anomaly. Unlike multiplicative inflation, the noise
    direction is set by ``Q_add`` rather than by the ensemble's own
    spread — useful when the dominant model-error mode is *not*
    represented by the ensemble (e.g. structurally underdispersive in
    a particular subspace). Mean-preserving by construction (we
    subtract the empirical mean of the noise).

    Common sources for ``Q_add``: climatological variability, lagged
    forecast differences, or stochastic-physics tendency variances.

    Sampling delegates to :func:`filterax.perturbed_observations`'s
    diagonal-aware path, so a diagonal ``Q_add`` never materialises.

    Args:
        key: PRNG key consumed by this call.
        particles: Ensemble of shape ``(Nₑ, Nₓ)``.
        noise_cov: Additive-noise covariance ``Q_add`` as a linear
            operator.

    Returns:
        Inflated ensemble of shape ``(Nₑ, Nₓ)``.

    Reference:
        Hamill, T. M. & Whitaker, J. S. (2005). *Accounting for the
        error due to unresolved scales in ensemble data assimilation:
        A comparison of different approaches.* Mon. Wea. Rev., 133,
        3132-3147.
    """
    from filterax._primitives._perturbations import perturbed_observations

    N_e = particles.shape[0]
    # Sample N_e draws from 𝒩(0, Q_add).
    draws = perturbed_observations(
        key, jnp.zeros(particles.shape[1], dtype=particles.dtype), noise_cov, N_e
    )
    # Centre so the empirical mean is exactly zero — keeps the
    # ensemble mean unchanged after addition.
    draws = draws - jnp.mean(draws, axis=0, keepdims=True)
    return particles + draws

inflate_adaptive(inflation_mean: Float[Array, ''] | float, inflation_var: Float[Array, ''] | float, innovation: Float[Array, ' N_y'], innovation_cov: Float[Array, 'N_y N_y'], *, min_factor: float = 1.0, max_factor: float = 1.2) -> tuple[Float[Array, ''], Float[Array, '']]

Bayesian update of a multiplicative inflation factor (Anderson 2009).

Treats the inflation factor λ as a random variable with a Gaussian prior λ ~ 𝒩(μ_λ, σ²_λ) and updates it from the cycle's normalised innovation. The data-implied factor is the clamped Mahalanobis ratio

\[ \hat{\lambda}_{\text{obs}} = \operatorname{clip}\!\left( \frac{d^{\top} S^{-1} d}{N_{y}},\ \lambda_{\min},\ \lambda_{\max} \right) \]

(the simplified Anderson-2009 form — under correct specification E[χ²/Nᵧ] = 1; under-dispersive ensembles overshoot and pull λ̂_obs above 1). The posterior on λ is then the Gaussian product of the prior with an "observation likelihood" centred at λ̂_obs whose variance is the χ²-distribution variance σ²_obs = 2 / Nᵧ:

\[ \mu_{\text{post}} = \frac{\sigma^{2}_{\lambda}\, \hat{\lambda}_{\text{obs}} + \sigma^{2}_{\text{obs}}\, \mu_{\text{prior}}} {\sigma^{2}_{\lambda} + \sigma^{2}_{\text{obs}}}, \qquad \sigma^{2}_{\text{post}} = \frac{\sigma^{2}_{\lambda}\, \sigma^{2}_{\text{obs}}} {\sigma^{2}_{\lambda} + \sigma^{2}_{\text{obs}}}. \]

Returns the posterior (μ, σ²) as JAX scalars so callers can carry the belief through jax.jit / jax.grad / lax.scan state across assimilation windows. Actual inflation of an ensemble is one extra call to :func:inflate_multiplicative with μ_post as the factor.

Parameters:

Name	Type	Description	Default
inflation_mean	Float[Array, ''] | float	Prior mean μ_λ (scalar or 0-d JAX array).	required
inflation_var	Float[Array, ''] | float	Prior variance σ²_λ (scalar or 0-d JAX array).	required
innovation	Float[Array, ' N_y']	Innovation vector d = y − H x̄.	required
innovation_cov	Float[Array, 'N_y N_y']	Dense innovation covariance S. (Nᵧ, Nᵧ) materialisation is acceptable because the recipe needs S⁻¹ applied to d — for typical observation counts the dense solve is cheaper than assembling a structured operator. Pass the dense form when you have it; otherwise build it via S_op.as_matrix() at the call site.	required
min_factor	float	Clamp lower bound for the data-implied estimate.	1.0
max_factor	float	Clamp upper bound.	1.2

Returns:

Type	Description
Float[Array, '']	(μ_post, σ²_post) JAX-scalar posterior on λ, suitable
Float[Array, '']	for use as the prior on the next cycle.

Reference

Anderson, J. L. (2009). Spatially and temporally varying adaptive covariance inflation for ensemble filters. Tellus A, 61, 72-83.

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

def inflate_adaptive(
    inflation_mean: Float[Array, ""] | float,
    inflation_var: Float[Array, ""] | float,
    innovation: Float[Array, " N_y"],
    innovation_cov: Float[Array, "N_y N_y"],
    *,
    min_factor: float = 1.0,
    max_factor: float = 1.2,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Bayesian update of a multiplicative inflation factor (Anderson 2009).

    Treats the inflation factor ``λ`` as a random variable with a
    Gaussian prior ``λ ~ 𝒩(μ_λ, σ²_λ)`` and updates it from the
    cycle's *normalised innovation*. The data-implied factor is the
    clamped Mahalanobis ratio

    $$
    \hat{\lambda}_{\text{obs}} = \operatorname{clip}\!\left(
        \frac{d^{\top} S^{-1} d}{N_{y}},\
        \lambda_{\min},\ \lambda_{\max}
    \right)
    $$

    (the simplified Anderson-2009 form — under correct specification
    ``E[χ²/Nᵧ] = 1``; under-dispersive ensembles overshoot and pull
    ``λ̂_obs`` above 1). The posterior on ``λ`` is then the Gaussian
    product of the prior with an "observation likelihood" centred at
    ``λ̂_obs`` whose variance is the χ²-distribution variance
    ``σ²_obs = 2 / Nᵧ``:

    $$
    \mu_{\text{post}} =
        \frac{\sigma^{2}_{\lambda}\, \hat{\lambda}_{\text{obs}}
              + \sigma^{2}_{\text{obs}}\, \mu_{\text{prior}}}
             {\sigma^{2}_{\lambda} + \sigma^{2}_{\text{obs}}},
    \qquad
    \sigma^{2}_{\text{post}} =
        \frac{\sigma^{2}_{\lambda}\, \sigma^{2}_{\text{obs}}}
             {\sigma^{2}_{\lambda} + \sigma^{2}_{\text{obs}}}.
    $$

    Returns the *posterior* ``(μ, σ²)`` as **JAX scalars** so callers
    can carry the belief through ``jax.jit`` / ``jax.grad`` / ``lax.scan``
    state across assimilation windows. Actual inflation of an ensemble
    is one extra call to :func:`inflate_multiplicative` with ``μ_post``
    as the factor.

    Args:
        inflation_mean: Prior mean ``μ_λ`` (scalar or 0-d JAX array).
        inflation_var: Prior variance ``σ²_λ`` (scalar or 0-d JAX array).
        innovation: Innovation vector ``d = y − H x̄``.
        innovation_cov: Dense innovation covariance ``S``.
            ``(Nᵧ, Nᵧ)`` materialisation is acceptable because the
            recipe needs ``S⁻¹`` applied to ``d`` — for typical
            observation counts the dense solve is cheaper than
            assembling a structured operator. Pass the dense form when
            you have it; otherwise build it via ``S_op.as_matrix()`` at
            the call site.
        min_factor: Clamp lower bound for the data-implied estimate.
        max_factor: Clamp upper bound.

    Returns:
        ``(μ_post, σ²_post)`` JAX-scalar posterior on ``λ``, suitable
        for use as the prior on the next cycle.

    Reference:
        Anderson, J. L. (2009). *Spatially and temporally varying
        adaptive covariance inflation for ensemble filters.* Tellus A,
        61, 72-83.
    """
    N_y = innovation.shape[0]
    chi2 = innovation @ jnp.linalg.solve(innovation_cov, innovation)
    chi2_norm = chi2 / N_y
    # Data-implied λ; clamped to keep the estimator from running away
    # on outlier innovations.
    lambda_obs = jnp.clip(chi2_norm, min_factor, max_factor)
    sigma_obs_sq = jnp.asarray(2.0 / N_y, dtype=chi2_norm.dtype)
    prior_mean = jnp.asarray(inflation_mean, dtype=chi2_norm.dtype)
    prior_var = jnp.asarray(inflation_var, dtype=chi2_norm.dtype)
    sigma_sum = prior_var + sigma_obs_sq
    mu_post = (prior_var * lambda_obs + sigma_obs_sq * prior_mean) / sigma_sum
    var_post = prior_var * sigma_obs_sq / sigma_sum
    return mu_post, var_post

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

Ledoit-Wolf optimal covariance shrinkage (Ledoit & Wolf 2004).

Replaces the sample covariance with a convex combination toward a scalar-multiple-of-identity target:

\[ P_{\text{shrunk}} = (1 - \lambda^{*})\, P_{\text{sample}} + \lambda^{*} \mu I, \]

where μ = tr(P_sample) / Nₓ is the average sample eigenvalue and λ* ∈ [0, 1] is the optimal shrinkage intensity that minimises E ‖P_shrunk − P_true‖²_F. The closed-form Ledoit-Wolf estimator:

\[ \lambda^{*} = \min\!\left( 1,\ \frac{b^{2}}{d^{2}} \right), \qquad d^{2} = \lVert P_{\text{sample}} - \mu I \rVert^{2}_{F}, \qquad b^{2} = \frac{1}{N_{e}^{2}} \sum_{j} \lVert x^{(j)} x^{(j)\top} - P_{\text{sample}} \rVert^{2}_{F}, \]

where d² is the sample deviation from the target and b² is the oracle approximation, capped at d² so λ* ≤ 1.

Useful when Nₑ ≪ Nₓ and the rank-deficient sample covariance would otherwise contaminate downstream operations. Positive semidefinite by construction (PSD, not strictly PD — when the ensemble is collapsed in some direction both the sample covariance and the target μ I are singular there). Returns a dense (Nₓ, Nₓ) matrix — only call it when Nₓ is small enough to materialise.

Parameters:

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

Returns:

Type	Description
Float[Array, 'N_x N_x']	(P_shrunk, λ*) — shrunk covariance and the optimal
Float[Array, '']	shrinkage intensity as a scalar array, so the function composes
tuple[Float[Array, 'N_x N_x'], Float[Array, '']]	with jit, vmap and grad.

Raises:

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

Reference

Ledoit, O. & Wolf, M. (2004). A well-conditioned estimator for large-dimensional covariance matrices. J. Multivariate Anal., 88, 365-411.

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

def ledoit_wolf_shrinkage(
    particles: Float[Array, "N_e N_x"],
) -> tuple[Float[Array, "N_x N_x"], Float[Array, ""]]:
    r"""Ledoit-Wolf optimal covariance shrinkage (Ledoit & Wolf 2004).

    Replaces the sample covariance with a convex combination toward a
    scalar-multiple-of-identity target:

    $$
    P_{\text{shrunk}} =
        (1 - \lambda^{*})\, P_{\text{sample}} + \lambda^{*} \mu I,
    $$

    where ``μ = tr(P_sample) / Nₓ`` is the average sample eigenvalue
    and ``λ* ∈ [0, 1]`` is the optimal shrinkage intensity that
    minimises ``E ‖P_shrunk − P_true‖²_F``. The closed-form
    Ledoit-Wolf estimator:

    $$
    \lambda^{*} = \min\!\left( 1,\ \frac{b^{2}}{d^{2}} \right),
    \qquad
    d^{2} = \lVert P_{\text{sample}} - \mu I \rVert^{2}_{F},
    \qquad
    b^{2} = \frac{1}{N_{e}^{2}} \sum_{j}
        \lVert x^{(j)} x^{(j)\top} - P_{\text{sample}} \rVert^{2}_{F},
    $$

    where ``d²`` is the sample deviation from the target and ``b²`` is
    the oracle approximation, capped at ``d²`` so ``λ* ≤ 1``.

    Useful when ``Nₑ ≪ Nₓ`` and the rank-deficient sample covariance
    would otherwise contaminate downstream operations. **Positive
    semidefinite** by construction (PSD, not strictly PD — when the
    ensemble is collapsed in some direction both the sample covariance
    and the target ``μ I`` are singular there). Returns a dense
    ``(Nₓ, Nₓ)`` matrix — only call it when ``Nₓ`` is small enough to
    materialise.

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

    Returns:
        ``(P_shrunk, λ*)`` — shrunk covariance and the optimal
        shrinkage intensity as a scalar array, so the function composes
        with ``jit``, ``vmap`` and ``grad``.

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

    Reference:
        Ledoit, O. & Wolf, M. (2004). *A well-conditioned estimator
        for large-dimensional covariance matrices.* J. Multivariate
        Anal., 88, 365-411.
    """
    N_e, N_x = particles.shape
    check_ensemble_size(N_e)
    mean = jnp.mean(particles, axis=0)
    centered = particles - mean[None, :]
    sample = centered.T @ centered / (N_e - 1)  # (Nₓ, Nₓ)
    mu = jnp.trace(sample) / N_x
    target = mu * jnp.eye(N_x, dtype=particles.dtype)

    # d² and b² as in Ledoit-Wolf 2004 §3.2.
    diff = sample - target
    d2 = jnp.sum(diff * diff)
    # Per-sample rank-one deviation from sample covariance.
    per_sample = centered[:, :, None] * centered[:, None, :]  # (Nₑ, Nₓ, Nₓ)
    b2_terms = jnp.sum((per_sample - sample[None, :, :]) ** 2, axis=(1, 2))
    b2 = jnp.mean(b2_terms) / N_e
    shrinkage = jnp.minimum(
        jnp.asarray(1.0, dtype=particles.dtype), b2 / jnp.maximum(d2, 1e-30)
    )
    shrunk = (1.0 - shrinkage) * sample + shrinkage * target
    return shrunk, shrinkage

Composition pattern

The L2 models (filterax.ETKF, filterax.LETKF, …) accept any AbstractInflator instance. To carry stateful inflation (e.g. adaptive λ) across cycles, drive the loop yourself with the L1 components:

import filterax as flx

# Anderson 2009: carry (μ_λ, σ²_λ) across cycles.
mu, var = 1.0, 0.01
inflator = flx.MultiplicativeInflator(factor=mu)

for obs, time in observations:
    forecast = vmap_dynamics(particles, t_prev, time)
    result = flx.filters.ETKF().analysis(forecast, obs, obs_op, R)

    # Update the inflation belief from this cycle's innovation.
    S = flx.innovation_covariance(vmap(obs_op)(forecast), R).as_matrix()
    innov = obs - vmap(obs_op)(forecast).mean(axis=0)
    mu, var = flx.inflate_adaptive(mu, var, innov, S)
    inflator = flx.MultiplicativeInflator(factor=mu)

    particles = inflator(result.particles, forecast)