Diagnostics

Pure functions that answer "is the filter working correctly?" — they detect problems, they don't fix them. Live under filterax.utils.*. Group them into a standard pipeline by cost:

Phase 1 — per-cycle health checks (cheap)

Run on every analysis window; cost is at most O(N_e N_x + N_y).

Function	What it tells you	Healthy value
ensemble_spread	Per-variable σ	Stable across cycles
rms_spread	Scalar RMS spread	Non-collapsing
innovation	d = y − H x̄_f	Mean ≈ 0 over time
normalized_innovation	d_i / √S_ii	~ 𝒩(0, 1)
chi2_consistency, chi2_normalized	Mahalanobis innovation norm	χ² / N_y ≈ 1
effective_ensemble_size	Weight degeneracy (particle filters)	≈ N_e
weight_entropy	Shannon entropy of weights	≈ log N_e

Phase 2 — calibration assessment (rolling window, needs truth)

Compute over 50–100 cycles in a twin / OSSE experiment.

Function	What it tells you	Healthy value
rmse_vs_truth	Accuracy of ensemble mean	Stable / decreasing
spread_skill_ratio	Calibration of σ vs RMSE	≈ 1
rank_histogram	Distribution of truth's rank in the ensemble	Uniform
rank_histogram_chi2	Flatness test for the histogram	χ²(N_e) quantiles

Phase 3 — a-posteriori covariance diagnosis (long accumulation)

Run after a complete experiment (100+ cycles) to diagnose mis-specified P or R.

Function	What it tells you
desroziers_R_estimate	A-posteriori R̂ from E[d_a d_fᵀ]
desroziers_innovation_cov	Empirical HPHᵀ + R from E[d_f d_fᵀ]
desroziers_analysis_residual_cov	Empirical R − HAHᵀ from E[d_a d_aᵀ]
dfs_from_gain	Observation information content from tr(KH)
dfs_from_ensemble	Same, computed from ensemble perturbations

Phase 4 — forecast skill (verification against obs)

Function	What it tells you
crps_ensemble, crps_ensemble_batch	Continuous Ranked Probability Score (lower is better; proper scoring rule)

Per-cycle reference

Diagnostic and utility primitives.

Re-exports the ensemble-DA diagnostics from :mod:filterax._diagnostics so users can do filterax.utils.rmse_vs_truth(...). See the module for selection guidance (per-cycle health checks, calibration assessment, a-posteriori covariance diagnosis, and the CRPS proper scoring rule).

ensemble_spread(ensemble: Float[Array, 'N_e N_x']) -> Float[Array, ' N_x']

Per-variable ensemble standard deviation.

\[ \sigma_i = \sqrt{(N_e - 1)^{-1} \sum_{j} \big(x_i^{(j)} - \bar{x}_i\big)^2} \]

Track over assimilation cycles. A well-calibrated filter has σ ≈ RMSE; collapsing spread is the first sign of filter divergence.

Parameters:

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

Returns:

Type	Description
Float[Array, ' N_x']	Per-variable standard deviation of shape (Nₓ,).

Examples:

>>> import jax.numpy as jnp
>>> from filterax.utils import ensemble_spread
>>> ens = jnp.array([[0.0, 0.0], [2.0, 4.0]])
>>> ensemble_spread(ens)
Array([1.4142135, 2.828427 ], dtype=float32)

Reference

Whitaker, J. S. & Loughe, A. F. (1998). The Relationship between Ensemble Spread and Ensemble Mean Skill. Mon. Wea. Rev., 126(12), 3292–3302.

Source code in src/filterax/_diagnostics.py

def ensemble_spread(
    ensemble: Float[Array, "N_e N_x"],
) -> Float[Array, " N_x"]:
    r"""Per-variable ensemble standard deviation.

    $$
    \sigma_i = \sqrt{(N_e - 1)^{-1}
        \sum_{j} \big(x_i^{(j)} - \bar{x}_i\big)^2}
    $$

    Track over assimilation cycles. A well-calibrated filter has
    ``σ ≈ RMSE``; collapsing spread is the first sign of filter
    divergence.

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

    Returns:
        Per-variable standard deviation of shape ``(Nₓ,)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.utils import ensemble_spread
        >>> ens = jnp.array([[0.0, 0.0], [2.0, 4.0]])
        >>> ensemble_spread(ens)
        Array([1.4142135, 2.828427 ], dtype=float32)

    Reference:
        Whitaker, J. S. & Loughe, A. F. (1998). *The Relationship
        between Ensemble Spread and Ensemble Mean Skill.* Mon. Wea.
        Rev., 126(12), 3292–3302.
    """
    return jnp.std(ensemble, axis=0, ddof=1)

rms_spread(ensemble: Float[Array, 'N_e N_x']) -> Float[Array, '']

Scalar RMS ensemble spread.

\[ \sqrt{N_x^{-1} \sum_i \sigma_i^2} \]

Source code in src/filterax/_diagnostics.py

def rms_spread(
    ensemble: Float[Array, "N_e N_x"],
) -> Float[Array, ""]:
    r"""Scalar RMS ensemble spread.

    $$
    \sqrt{N_x^{-1} \sum_i \sigma_i^2}
    $$
    """
    sigma = ensemble_spread(ensemble)
    return jnp.sqrt(jnp.mean(sigma * sigma))

innovation(y_obs: Float[Array, ' N_y'], Hx_forecast: Float[Array, ' N_y']) -> Float[Array, ' N_y']

Innovation vector d = y − H x̄_f.

Source code in src/filterax/_diagnostics.py

def innovation(
    y_obs: Float[Array, " N_y"],
    Hx_forecast: Float[Array, " N_y"],
) -> Float[Array, " N_y"]:
    r"""Innovation vector ``d = y − H x̄_f``."""
    return y_obs - Hx_forecast

normalized_innovation(y_obs: Float[Array, ' N_y'], Hx_forecast: Float[Array, ' N_y'], innovation_var: Float[Array, ' N_y']) -> Float[Array, ' N_y']

Normalised innovation.

\[ d_i \big/ \sqrt{(H P H^{\top} + R)_{ii}} \]

For a correctly-specified filter the normalised innovation is 𝒩(0, 1) componentwise. Departures from unit variance signal misspecified P or R.

Source code in src/filterax/_diagnostics.py

def normalized_innovation(
    y_obs: Float[Array, " N_y"],
    Hx_forecast: Float[Array, " N_y"],
    innovation_var: Float[Array, " N_y"],
) -> Float[Array, " N_y"]:
    r"""Normalised innovation.

    $$
    d_i \big/ \sqrt{(H P H^{\top} + R)_{ii}}
    $$

    For a correctly-specified filter the normalised innovation is
    ``𝒩(0, 1)`` componentwise. Departures from unit variance signal
    misspecified ``P`` or ``R``.
    """
    d = y_obs - Hx_forecast
    return d / jnp.sqrt(jnp.maximum(innovation_var, 1e-30))

chi2_consistency(d: Float[Array, ' N_y'], S_inv_d: Float[Array, ' N_y']) -> Float[Array, '']

Chi-squared innovation consistency statistic.

\[ \chi^2 = d^{\top} S^{-1} d \]

Should be ≈ Nᵧ under correct specification. Pre-compute S⁻¹ d once and pass it here so the diagnostic composes with whichever solver (gaussx Woodbury, dense Cholesky) you used for the analysis. Variance under the null is 2 Nᵧ.

Reference

Mehra, R. K. (1970). On the Identification of Variances and Adaptive Kalman Filtering. IEEE Trans. Automatic Control, 15(2), 175–184.

Source code in src/filterax/_diagnostics.py

def chi2_consistency(
    d: Float[Array, " N_y"],
    S_inv_d: Float[Array, " N_y"],
) -> Float[Array, ""]:
    r"""Chi-squared innovation consistency statistic.

    $$
    \chi^2 = d^{\top} S^{-1} d
    $$

    Should be ``≈ Nᵧ`` under correct specification.
    Pre-compute ``S⁻¹ d`` once and pass it here so the diagnostic
    composes with whichever solver (gaussx Woodbury, dense Cholesky)
    you used for the analysis. Variance under the null is ``2 Nᵧ``.

    Reference:
        Mehra, R. K. (1970). *On the Identification of Variances and
        Adaptive Kalman Filtering.* IEEE Trans. Automatic Control,
        15(2), 175–184.
    """
    return d @ S_inv_d

chi2_normalized(d: Float[Array, ' N_y'], S_inv_d: Float[Array, ' N_y'], n_obs: int) -> Float[Array, '']

Compact single-number consistency check χ² / Nᵧ (target ≈ 1).

Source code in src/filterax/_diagnostics.py

def chi2_normalized(
    d: Float[Array, " N_y"],
    S_inv_d: Float[Array, " N_y"],
    n_obs: int,
) -> Float[Array, ""]:
    """Compact single-number consistency check ``χ² / Nᵧ`` (target ≈ 1)."""
    return (d @ S_inv_d) / n_obs

effective_ensemble_size(weights: Float[Array, ' N_e']) -> Float[Array, '']

Effective sample size.

\[ N_{\mathrm{eff}} = 1 \Big/ \sum_{j} w_j^2 \]

For equally-weighted ensembles (the standard EnKF) this is exactly Nₑ. For particle filters / weighted EnKF variants, Nₑff ≪ Nₑ signals weight degeneracy and the need to resample.

Reference

Liu, J. S. & Chen, R. (1998). Sequential Monte Carlo Methods for Dynamic Systems. JASA, 93(443), 1032–1044.

Source code in src/filterax/_diagnostics.py

def effective_ensemble_size(
    weights: Float[Array, " N_e"],
) -> Float[Array, ""]:
    r"""Effective sample size.

    $$
    N_{\mathrm{eff}} = 1 \Big/ \sum_{j} w_j^2
    $$

    For equally-weighted ensembles (the standard EnKF) this is
    exactly ``Nₑ``. For particle filters / weighted EnKF variants,
    ``Nₑff ≪ Nₑ`` signals weight degeneracy and the need to resample.

    Reference:
        Liu, J. S. & Chen, R. (1998). *Sequential Monte Carlo Methods
        for Dynamic Systems.* JASA, 93(443), 1032–1044.
    """
    return 1.0 / jnp.sum(weights * weights)

weight_entropy(weights: Float[Array, ' N_e']) -> Float[Array, '']

Shannon entropy of importance weights.

\[ -\sum_{j} w_j \log w_j \]

Maximum log Nₑ for uniform weights, 0 for full degeneracy.

Source code in src/filterax/_diagnostics.py

def weight_entropy(
    weights: Float[Array, " N_e"],
) -> Float[Array, ""]:
    r"""Shannon entropy of importance weights.

    $$
    -\sum_{j} w_j \log w_j
    $$

    Maximum ``log Nₑ`` for uniform weights, ``0`` for full
    degeneracy.
    """
    safe = jnp.maximum(weights, 1e-30)
    return -jnp.sum(weights * jnp.log(safe))

Calibration reference

Diagnostic and utility primitives.

Re-exports the ensemble-DA diagnostics from :mod:filterax._diagnostics so users can do filterax.utils.rmse_vs_truth(...). See the module for selection guidance (per-cycle health checks, calibration assessment, a-posteriori covariance diagnosis, and the CRPS proper scoring rule).

rmse_vs_truth(ensemble_mean_: Float[Array, ' N_x'], x_true: Float[Array, ' N_x']) -> Float[Array, '']

RMSE of an ensemble mean against the true state.

\[ \mathrm{RMSE} = \sqrt{N_x^{-1} \sum_i \big(\bar{x}_i - x_i^*\big)^2} \]

Only available in twin / OSSE experiments where x_true is known. Track over time to detect filter divergence (RMSE growing without bound).

Examples:

>>> import jax.numpy as jnp
>>> from filterax.utils import rmse_vs_truth
>>> rmse_vs_truth(jnp.array([1.0, 2.0]), jnp.array([0.0, 2.0]))
Array(0.70710677, dtype=float32)

Source code in src/filterax/_diagnostics.py

def rmse_vs_truth(
    ensemble_mean_: Float[Array, " N_x"],
    x_true: Float[Array, " N_x"],
) -> Float[Array, ""]:
    r"""RMSE of an ensemble mean against the true state.

    $$
    \mathrm{RMSE} = \sqrt{N_x^{-1} \sum_i \big(\bar{x}_i - x_i^*\big)^2}
    $$

    Only available in twin / OSSE experiments where ``x_true`` is
    known. Track over time to detect filter divergence (RMSE growing
    without bound).

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.utils import rmse_vs_truth
        >>> rmse_vs_truth(jnp.array([1.0, 2.0]), jnp.array([0.0, 2.0]))
        Array(0.70710677, dtype=float32)
    """
    diff = ensemble_mean_ - x_true
    return jnp.sqrt(jnp.mean(diff * diff))

spread_skill_ratio(ensemble: Float[Array, 'N_e N_x'], x_true: Float[Array, ' N_x']) -> Float[Array, '']

Ratio of RMS spread to RMSE.

\[ \mathrm{SSR} = \sqrt{\sum_i \sigma_i^2 / N_x} \Big/ \sqrt{\sum_i (\bar{x}_i - x_i^*)^2 / N_x} \]

• SSR ≈ 1 — well calibrated.
• SSR < 1 — underdispersive (increase inflation).
• SSR > 1 — overdispersive (decrease inflation).

Reference

Fortin, V., Abaza, M., Anctil, F., & Turcotte, R. (2014). Why Should Ensemble Spread Match the RMSE of the Ensemble Mean? J. Hydrometeorol., 15(4), 1708–1713.

Source code in src/filterax/_diagnostics.py

def spread_skill_ratio(
    ensemble: Float[Array, "N_e N_x"],
    x_true: Float[Array, " N_x"],
) -> Float[Array, ""]:
    r"""Ratio of RMS spread to RMSE.

    $$
    \mathrm{SSR} = \sqrt{\sum_i \sigma_i^2 / N_x} \Big/
        \sqrt{\sum_i (\bar{x}_i - x_i^*)^2 / N_x}
    $$

    * ``SSR ≈ 1`` — well calibrated.
    * ``SSR < 1`` — underdispersive (increase inflation).
    * ``SSR > 1`` — overdispersive (decrease inflation).

    Reference:
        Fortin, V., Abaza, M., Anctil, F., & Turcotte, R. (2014).
        *Why Should Ensemble Spread Match the RMSE of the Ensemble
        Mean?* J. Hydrometeorol., 15(4), 1708–1713.
    """
    rms = rms_spread(ensemble)
    err = rmse_vs_truth(jnp.mean(ensemble, axis=0), x_true)
    return rms / jnp.maximum(err, 1e-30)

rank_histogram(ensemble_over_time: Float[Array, 'T N_e N_x'], truth_over_time: Float[Array, 'T N_x']) -> Int[Array, ' N_eP1']

Rank histogram counts across all (t, i) pairs.

For each variable i at each time t, count how many of the Nₑ ensemble members are smaller than the truth; the resulting rank is in {0, 1, …, Nₑ} (Nₑ + 1 bins). Aggregate over time and variables. A reliable ensemble produces a roughly uniform histogram; U-shaped → underdispersive, dome-shaped → overdispersive, skewed → systematic bias.

Parameters:

Name	Type	Description	Default
ensemble_over_time	Float[Array, 'T N_e N_x']	(T, Nₑ, Nₓ) ensemble snapshots.	required
truth_over_time	Float[Array, 'T N_x']	(T, Nₓ) true states at the same times.	required

Returns:

Type	Description
Int[Array, ' N_eP1']	(Nₑ + 1,) integer counts.

Examples:

>>> import jax.numpy as jnp
>>> from filterax.utils import rank_histogram
>>> ens = jnp.array([[[0.0], [1.0]], [[0.0], [1.0]]])  # (T, N_e, N_x)
>>> truth = jnp.array([[0.5], [2.0]])
>>> rank_histogram(ens, truth)
Array([0, 1, 1], dtype=int32)

Reference

Hamill, T. M. (2001). Interpretation of Rank Histograms for Verifying Ensemble Forecasts. Mon. Wea. Rev., 129(3), 550–560.

Source code in src/filterax/_diagnostics.py

def rank_histogram(
    ensemble_over_time: Float[Array, "T N_e N_x"],
    truth_over_time: Float[Array, "T N_x"],
) -> Int[Array, " N_eP1"]:
    r"""Rank histogram counts across all ``(t, i)`` pairs.

    For each variable ``i`` at each time ``t``, count how many of the
    ``Nₑ`` ensemble members are smaller than the truth; the resulting
    rank is in ``{0, 1, …, Nₑ}`` (``Nₑ + 1`` bins). Aggregate over
    time and variables. A reliable ensemble produces a roughly uniform
    histogram; U-shaped → underdispersive, dome-shaped → overdispersive,
    skewed → systematic bias.

    Args:
        ensemble_over_time: ``(T, Nₑ, Nₓ)`` ensemble snapshots.
        truth_over_time: ``(T, Nₓ)`` true states at the same times.

    Returns:
        ``(Nₑ + 1,)`` integer counts.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.utils import rank_histogram
        >>> ens = jnp.array([[[0.0], [1.0]], [[0.0], [1.0]]])  # (T, N_e, N_x)
        >>> truth = jnp.array([[0.5], [2.0]])
        >>> rank_histogram(ens, truth)
        Array([0, 1, 1], dtype=int32)

    Reference:
        Hamill, T. M. (2001). *Interpretation of Rank Histograms for
        Verifying Ensemble Forecasts.* Mon. Wea. Rev., 129(3), 550–560.
    """
    _T, N_e, _N_x = ensemble_over_time.shape
    # Count members < truth, broadcast over (T, Nₑ, Nₓ).
    smaller = jnp.sum(
        ensemble_over_time < truth_over_time[:, None, :], axis=1
    )  # (T, Nₓ) — per-pair rank in {0, …, Nₑ}
    ranks = smaller.reshape(-1)
    return jnp.bincount(ranks, length=N_e + 1)

rank_histogram_chi2(counts: Int[Array, ' bins']) -> Float[Array, '']

χ² flatness statistic for a rank histogram.

\[ \chi^2 = \sum_k (O_k - E_k)^2 / E_k, \qquad E_k = N_{\mathrm{total}} / N_{\mathrm{bins}}. \]

Larger values indicate stronger departure from uniformity; under the null hypothesis of uniformity the test statistic is χ²(N_bins − 1).

Source code in src/filterax/_diagnostics.py

def rank_histogram_chi2(counts: Int[Array, " bins"]) -> Float[Array, ""]:
    r"""χ² flatness statistic for a rank histogram.

    $$
    \chi^2 = \sum_k (O_k - E_k)^2 / E_k, \qquad
    E_k = N_{\mathrm{total}} / N_{\mathrm{bins}}.
    $$

    Larger values indicate stronger departure from uniformity; under
    the null hypothesis of uniformity the test statistic is
    ``χ²(N_bins − 1)``.
    """
    n_bins = counts.shape[0]
    total = jnp.sum(counts)
    expected = total / n_bins
    diff = counts - expected
    return jnp.sum(diff * diff / jnp.maximum(expected, 1e-30))

Covariance-diagnosis reference

Diagnostic and utility primitives.

Re-exports the ensemble-DA diagnostics from :mod:filterax._diagnostics so users can do filterax.utils.rmse_vs_truth(...). See the module for selection guidance (per-cycle health checks, calibration assessment, a-posteriori covariance diagnosis, and the CRPS proper scoring rule).

desroziers_R_estimate(d_forecast: Float[Array, 'T N_y'], d_analysis: Float[Array, 'T N_y']) -> Float[Array, 'N_y N_y']

A-posteriori estimate of R (Desroziers et al. 2005).

\[ \hat{R} = E\big[ d_a d_f^{\top} \big] \]

estimated over time. When R̂ disagrees with the prescribed R the observation-error covariance is mis-specified. Accumulate over many cycles (100+) for stable estimates.

Source code in src/filterax/_diagnostics.py

def desroziers_R_estimate(
    d_forecast: Float[Array, "T N_y"],
    d_analysis: Float[Array, "T N_y"],
) -> Float[Array, "N_y N_y"]:
    r"""A-posteriori estimate of ``R`` (Desroziers et al. 2005).

    $$
    \hat{R} = E\big[ d_a d_f^{\top} \big]
    $$

    estimated over time. When ``R̂`` disagrees with the
    prescribed ``R`` the observation-error covariance is mis-specified.
    Accumulate over many cycles (100+) for stable estimates.
    """
    T = d_forecast.shape[0]
    return einx.dot("t i, t j -> i j", d_analysis, d_forecast) / T

desroziers_innovation_cov(d_forecast: Float[Array, 'T N_y']) -> Float[Array, 'N_y N_y']

Empirical innovation covariance.

\[ E\big[ d_f d_f^{\top} \big] \approx H P H^{\top} + R \]

Source code in src/filterax/_diagnostics.py

def desroziers_innovation_cov(
    d_forecast: Float[Array, "T N_y"],
) -> Float[Array, "N_y N_y"]:
    r"""Empirical innovation covariance.

    $$
    E\big[ d_f d_f^{\top} \big] \approx H P H^{\top} + R
    $$
    """
    T = d_forecast.shape[0]
    return einx.dot("t i, t j -> i j", d_forecast, d_forecast) / T

desroziers_analysis_residual_cov(d_analysis: Float[Array, 'T N_y']) -> Float[Array, 'N_y N_y']

Empirical analysis-residual covariance.

\[ E\big[ d_a d_a^{\top} \big] \approx R - H A H^{\top} \]

Source code in src/filterax/_diagnostics.py

def desroziers_analysis_residual_cov(
    d_analysis: Float[Array, "T N_y"],
) -> Float[Array, "N_y N_y"]:
    r"""Empirical analysis-residual covariance.

    $$
    E\big[ d_a d_a^{\top} \big] \approx R - H A H^{\top}
    $$
    """
    T = d_analysis.shape[0]
    return einx.dot("t i, t j -> i j", d_analysis, d_analysis) / T

dfs_from_gain(K: Float[Array, 'N_x N_y'], H: Float[Array, 'N_y N_x']) -> Float[Array, '']

Degrees of freedom for signal from an explicit gain.

\[ \mathrm{DFS} = \mathrm{tr}(K H) \in [0, N_y] \]

Closer to Nᵧ means observations dominate the analysis; closer to 0 means the prior dominates.

Reference

Cardinali, C., Pezzulli, S., & Andersson, E. (2004). Influence-Matrix Diagnostic of a Data Assimilation System. QJRMS, 130(603), 2767–2786.

Source code in src/filterax/_diagnostics.py

def dfs_from_gain(
    K: Float[Array, "N_x N_y"],
    H: Float[Array, "N_y N_x"],
) -> Float[Array, ""]:
    r"""Degrees of freedom for signal from an explicit gain.

    $$
    \mathrm{DFS} = \mathrm{tr}(K H) \in [0, N_y]
    $$

    Closer to ``Nᵧ`` means observations
    dominate the analysis; closer to ``0`` means the prior dominates.

    Reference:
        Cardinali, C., Pezzulli, S., & Andersson, E. (2004).
        *Influence-Matrix Diagnostic of a Data Assimilation System.*
        QJRMS, 130(603), 2767–2786.
    """
    return jnp.trace(K @ H)

dfs_from_ensemble(ensemble_forecast: Float[Array, 'N_e N_x'], ensemble_analysis: Float[Array, 'N_e N_x']) -> Float[Array, '']

Ensemble-based DFS estimate from forecast / analysis perturbations.

\[ \mathrm{DFS} \approx \mathrm{tr}\big( C^{af} / \sigma_f^2 \big) \]

where Cᵃᶠ is the per-variable forecast-vs-analysis cross-covariance and σ_f² is the forecast variance — i.e. how much the analysis perturbations are driven by the forecast perturbations. DFS = Nₓ for an uninformative observation set; DFS < Nₓ measures how much constraint the observations contributed.

Avoids forming K and H explicitly — useful in ensemble-only workflows.

Source code in src/filterax/_diagnostics.py

def dfs_from_ensemble(
    ensemble_forecast: Float[Array, "N_e N_x"],
    ensemble_analysis: Float[Array, "N_e N_x"],
) -> Float[Array, ""]:
    r"""Ensemble-based DFS estimate from forecast / analysis perturbations.

    $$
    \mathrm{DFS} \approx \mathrm{tr}\big( C^{af} / \sigma_f^2 \big)
    $$

    where ``Cᵃᶠ`` is the per-variable
    forecast-vs-analysis cross-covariance and ``σ_f²`` is the
    forecast variance — i.e. how much the analysis perturbations are
    *driven* by the forecast perturbations. ``DFS = Nₓ`` for an
    uninformative observation set; ``DFS < Nₓ`` measures how much
    constraint the observations contributed.

    Avoids forming ``K`` and ``H`` explicitly — useful in
    ensemble-only workflows.
    """
    N_e = ensemble_forecast.shape[0]
    mean_f = jnp.mean(ensemble_forecast, axis=0)
    mean_a = jnp.mean(ensemble_analysis, axis=0)
    anom_f = ensemble_forecast - mean_f[None, :]
    anom_a = ensemble_analysis - mean_a[None, :]
    # Per-variable cross-covariance and forecast variance.
    cross = jnp.sum(anom_a * anom_f, axis=0) / (N_e - 1)  # (Nₓ,)
    var_f = jnp.sum(anom_f * anom_f, axis=0) / (N_e - 1)  # (Nₓ,)
    ratio = cross / jnp.maximum(var_f, 1e-30)
    # Per-variable "signal share"; sum to a scalar DFS proxy.
    return jnp.sum(ratio)

Forecast-skill reference

Diagnostic and utility primitives.

Re-exports the ensemble-DA diagnostics from :mod:filterax._diagnostics so users can do filterax.utils.rmse_vs_truth(...). See the module for selection guidance (per-cycle health checks, calibration assessment, a-posteriori covariance diagnosis, and the CRPS proper scoring rule).

crps_ensemble(ensemble: Float[Array, ' N_e'], y_obs: Float[Array, '']) -> Float[Array, '']

Continuous Ranked Probability Score for a 1D ensemble forecast.

\[ \mathrm{CRPS} = E|X - y| - \tfrac{1}{2} E|X - X'| \]

Computed via the sorted-ensemble identity (Hersbach 2000):

\[ \tfrac{1}{2} E|X - X'| = N_e^{-2} \sum_{j} x_{(j)} (2j - 1 - N_e) \]

where x_{(j)} are the order statistics. Cost O(Nₑ log Nₑ) per observation. Lower is better; CRPS is a strictly proper scoring rule and has the same units as the observed variable.

Parameters:

Name	Type	Description	Default
ensemble	Float[Array, ' N_e']	Single-variable ensemble of shape (Nₑ,).	required
y_obs	Float[Array, '']	Scalar observed value.	required

Returns:

Type	Description
Float[Array, '']	Scalar CRPS.

Examples:

>>> import jax.numpy as jnp
>>> from filterax.utils import crps_ensemble
>>> crps_ensemble(jnp.array([0.0, 1.0]), jnp.array(0.5))
Array(0.25, dtype=float32)

Reference

Gneiting, T. & Raftery, A. E. (2007). Strictly Proper Scoring Rules, Prediction, and Estimation. JASA, 102(477), 359–378.

Source code in src/filterax/_diagnostics.py

def crps_ensemble(
    ensemble: Float[Array, " N_e"],
    y_obs: Float[Array, ""],
) -> Float[Array, ""]:
    r"""Continuous Ranked Probability Score for a 1D ensemble forecast.

    $$
    \mathrm{CRPS} = E|X - y| - \tfrac{1}{2} E|X - X'|
    $$

    Computed via the sorted-ensemble identity (Hersbach 2000):

    $$
    \tfrac{1}{2} E|X - X'| = N_e^{-2} \sum_{j} x_{(j)} (2j - 1 - N_e)
    $$

    where ``x_{(j)}`` are the order statistics. Cost
    ``O(Nₑ log Nₑ)`` per observation. Lower is better; CRPS is a
    strictly proper scoring rule and has the same units as the
    observed variable.

    Args:
        ensemble: Single-variable ensemble of shape ``(Nₑ,)``.
        y_obs: Scalar observed value.

    Returns:
        Scalar CRPS.

    Examples:
        >>> import jax.numpy as jnp
        >>> from filterax.utils import crps_ensemble
        >>> crps_ensemble(jnp.array([0.0, 1.0]), jnp.array(0.5))
        Array(0.25, dtype=float32)

    Reference:
        Gneiting, T. & Raftery, A. E. (2007). *Strictly Proper
        Scoring Rules, Prediction, and Estimation.* JASA, 102(477),
        359–378.
    """
    N_e = ensemble.shape[0]
    sorted_e = jnp.sort(ensemble)
    indices = jnp.arange(1, N_e + 1, dtype=ensemble.dtype)
    first_term = jnp.mean(jnp.abs(ensemble - y_obs))
    second_term = jnp.sum(sorted_e * (2.0 * indices - 1.0 - N_e)) / (N_e * N_e)
    return first_term - second_term

crps_ensemble_batch(ensemble: Float[Array, 'N_e N_y'], y_obs: Float[Array, ' N_y']) -> Float[Array, '']

Mean CRPS over an observation vector.

Applies :func:crps_ensemble to each observation component independently and averages the results.

Source code in src/filterax/_diagnostics.py

def crps_ensemble_batch(
    ensemble: Float[Array, "N_e N_y"],
    y_obs: Float[Array, " N_y"],
) -> Float[Array, ""]:
    r"""Mean CRPS over an observation vector.

    Applies :func:`crps_ensemble` to each observation component
    independently and averages the results.
    """
    import jax

    per_obs = jax.vmap(crps_ensemble, in_axes=(1, 0))(ensemble, y_obs)
    return jnp.mean(per_obs)