pastax.score

Proper scoring rules for probabilistic (ensemble) trajectory forecasts.

Implements four scoring rules (see Pic et al., 2025) for ensemble forecasts of shape (S, T, C) evaluated against an observed trajectory (T, C):

The trailing axis C is the feature dimension — typically 2 for [lon, lat] trajectories, or 3 for e.g. chordal-projected unit-sphere coordinates. Kernels must support the C they are given: euclidean_distance() accepts any C; pastax.metric.separation_distance() and the haversine expect C = 2.

• squared_error() — deterministic-mean squared distance.

• dawid_sebastiani() — Gaussian-likelihood-based, no kernel.

• energy_score() — kernel-based proper scoring rule (unbiased estimator).

• variogram_score() — temporal variogram score for trajectory dependence.

All scores follow the negative orientation convention: lower is better.

The pointwise scores (squared_error(), dawid_sebastiani(), energy_score()) accept a reduce argument:

• reduce=None returns the per-time score of shape (T,).

• reduce="last" returns the scalar score at the final time.

• reduce="sum" returns (weights * score).sum(), defaulting to a uniform sum when weights is None. By Proposition 2 of Pic et al., 2025, a non-negative-weighted sum of proper scoring rules is itself proper.

• reduce="joint" returns the scalar joint score of the whole trajectory, treated as one flattened (T*C,) vector; supported by squared_error() and energy_score() only.

The default distance kernel for energy_score(), squared_error(), and variogram_score() is the Euclidean distance. A user may pass any callable satisfying the broadcasting kernel contract — notably pastax.metric.separation_distance() for great-circle distances on the sphere. Under reduce="joint" the score always uses euclidean_distance() on flattened trajectories; custom kernel are ignored.

variogram_score() is trajectory-level rather than pointwise in time. It compares forecast and observed increments across selected temporal lags and returns one scalar score.

dawid_sebastiani¶

pastax.score.dawid_sebastiani(forecast, obs, *, reduce=None, weights=None)

Dawid-Sebastiani score: Gaussian log-likelihood under the ensemble (Eq. 9 of Pic et al., 2025).

The per-time score is

$$\mathrm{DS}_t = \log\det \Sigma_t + (\mu_t - y_t)^{\top}\, \Sigma_t^{-1}\, (\mu_t - y_t)$$

(1)

where $\Sigma_t$ is the unbiased (ddof=1) sample covariance of the ensemble at time $t$. Requires $S \geq C + 1$ for $\Sigma_t$ to be a.s. full-rank on $\mathbb{R}^C$; for $S \leq C$ the score is undefined (singular covariance).

Parameters

• forecast (Float[jaxlib._jax.Array, 'S T C']) – Ensemble forecast, shape (S, T, C), with S >= C + 1.

• obs (Float[jaxlib._jax.Array, 'T C']) – Observed trajectory, shape (T, C).

• reduce (Literal['last', 'sum', 'joint'] | None) – See squared_error(). "joint" is NOT supported and raises ValueError.

• weights (Float[jaxlib._jax.Array, 'T'] | None) – See squared_error().

Returns

Per-time score of shape (T,) or a scalar, per reduce.

Return type

Float[jaxlib._jax.Array, ‘T’] | Float[jaxlib._jax.Array, ‘’]

energy_score¶

pastax.score.energy_score(forecast, obs, *, kernel=euclidean_distance, alpha=1.0, reduce=None, weights=None)

Energy score (Eq. 12 of Pic et al., 2025) — unbiased estimator.

$$\mathrm{ES}_t = \frac{1}{S} \sum_s d\!\left(X_t^{(s)}, y_t\right)^{\alpha} - \frac{1}{2 S (S-1)} \sum_{s \neq s'} d\!\left(X_t^{(s)}, X_t^{(s')}\right)^{\alpha}$$

(2)

The pairwise term is computed as a full (S, S) mean (including the zero diagonal) multiplied by S/(S-1), which recovers the unbiased off-diagonal estimator exactly.

With reduce="joint", the score is computed on trajectories flattened to (S, T*C) / (T*C,): the mean euclidean_distance to the observation minus half the S/(S-1)-scaled mean pairwise euclidean_distance, each raised to alpha. forecast and obs must have matching (T, C) trailing shapes. A custom kernel is ignored (a UserWarning is emitted) and weights is ignored.

Let $d_\alpha(x,y)=\operatorname{kernel}(x,y)^\alpha$. The energy score is proper when $d_\alpha$ is of negative type, and strictly proper when it is of strong negative type. For the Euclidean L2 kernel this gives strict propriety for $\alpha\in(0,2)$.

Parameters

• forecast (Float[jaxlib._jax.Array, 'S T C']) – Ensemble forecast, shape (S, T, C), with S >= 2.

• obs (Float[jaxlib._jax.Array, 'T C']) – Observed trajectory, shape (T, C).

• kernel (Callable[[Float[jaxlib._jax.Array, '... C'], Float[jaxlib._jax.Array, '... C']], Float[jaxlib._jax.Array, '...']]) – Broadcasting distance kernel. Defaults to euclidean_distance().

• alpha (float) – Distance exponent (typically in (0, 2)). Default 1.0.

• reduce (Literal['last', 'sum', 'joint'] | None) – See squared_error().

• weights (Float[jaxlib._jax.Array, 'T'] | None) – See squared_error().

Returns

Per-time fair energy score of shape (T,) or a scalar, per reduce.

Return type

Float[jaxlib._jax.Array, ‘T’] | Float[jaxlib._jax.Array, ‘’]

squared_error¶

pastax.score.squared_error(forecast, obs, *, kernel=euclidean_distance, reduce=None, weights=None)

Squared distance between ensemble mean and observation.

$$\mathrm{SE}_t = \operatorname{kernel}\!\left( \operatorname{mean}_s \mathrm{forecast}[s, t],\ \mathrm{obs}[t]\right)^2$$

(3)

With the default L2 kernel this is the squared error of the ensemble mean (Eq. 11 of Pic et al., 2025).

Parameters

• forecast (Float[jaxlib._jax.Array, 'S T C']) – Ensemble forecast, shape (S, T, C).

• obs (Float[jaxlib._jax.Array, 'T C']) – Observed trajectory, shape (T, C).

• kernel (Callable[[Float[jaxlib._jax.Array, '... C'], Float[jaxlib._jax.Array, '... C']], Float[jaxlib._jax.Array, '...']]) – Broadcasting distance kernel. Defaults to euclidean_distance().

• reduce (Literal['last', 'sum', 'joint'] | None) – Time reduction. None returns the per-time vector; "last" returns the scalar at the final time; "sum" returns the (optionally weighted) sum over time; "joint" returns the scalar squared L2 distance between the ensemble-mean trajectory and the observation, both flattened to (T*C,) (custom kernel and weights are ignored).

• weights (Float[jaxlib._jax.Array, 'T'] | None) – Per-time weights for reduce="sum"; ignored otherwise.

Returns

Per-time score of shape (T,) or a scalar, per reduce.

Return type

Float[jaxlib._jax.Array, ‘T’] | Float[jaxlib._jax.Array, ‘’]

variogram_score¶

pastax.score.variogram_score(forecast, obs, *, kernel=euclidean_distance, p=0.5, lags=None, lag_weights=None)

Fair temporal variogram score (Eq. 13 of Pic et al., 2025) of order p for ensemble trajectories.

The population score is

$$\mathrm{VS}_{p} = \sum_{\ell\in\mathcal L} \frac{a_\ell}{T-\ell} \sum_{i=0}^{T-\ell-1} \left[ \mathbb{E}_F\!\left\{ d(X_i,X_{i+\ell})^p \right\} - d(y_i,y_{i+\ell})^p \right]^2 .$$

(4)

Here $d$ is kernel, $\mathcal L$ is the selected set of positive temporal lags, and $a_\ell$ is the corresponding lag weight.

For a finite ensemble, this function uses the fair (unbiased) estimator of each squared variogram discrepancy.

The implementation evaluates the U-statistic in $O(S)$ rather than explicitly constructing the $S\times S$ off-diagonal matrix:

$$\frac{ \left(\sum_s R_s\right)^2 - \sum_s R_s^2 }{S(S-1)}.$$

(5)

For each lag, the score is averaged over every valid starting time before the lag weight is applied, so short lags do not receive more weight merely because they contain more valid time pairs.

This is the variogram construction applied to the temporal dimensions of a trajectory rather than to coordinate components at a single time. It therefore targets temporal dependence / increment structure and complements the pointwise-in-time energy_score().

Parameters

• forecast (Float[jaxlib._jax.Array, 'S T C']) – Ensemble forecast, shape (S, T, C), with S >= 2.

• obs (Float[jaxlib._jax.Array, 'T C']) – Observed trajectory, shape (T, C).

• kernel (Callable[[Float[jaxlib._jax.Array, '... C'], Float[jaxlib._jax.Array, '... C']], Float[jaxlib._jax.Array, '...']]) – Broadcasting distance kernel between trajectory states. Defaults to euclidean_distance().

• p (float) – Variogram order, p > 0. Defaults to 0.5.

• lags (tuple[int, ...] | None) – Positive integer temporal lags. If None, uses every lag 1, ..., T - 1. Lags are Python integers and should be static when the function is wrapped in jax.jit().

• lag_weights (Float[jaxlib._jax.Array, 'L'] | None) – Non-negative weight for each entry of lags. If None, uses equal weights summing to one. Supplied weights are normalized to sum to one. Invalid lags raise ValueError (lags must stay concrete Python integers under jax.jit() — close them over or use static_argnames); invalid lag_weights raise ValueError in eager mode.

Returns

Scalar fair temporal variogram score.

Return type

Float[jaxlib._jax.Array, ‘’]