Skip to content

Losses

AMSELoss

Bases: Module

Adjusted-MSE (anti-double-penalty) spectral loss on a flat grid.

Two modes:

  • "hybrid" (default): the repo's standard Huber loss plus spectral_weight times the AMSE excess sum_k 2 * (max(P_k(x), P_k(y)) - sqrt(P_k(x) * P_k(y))) * (1 - Coh_k), which is >= 0, zero iff per-bin spectra match or coherence is perfect, and reduces to the unmodified Huber control at spectral_weight=0.
  • "pure": the AMSE formula itself (per-bin amplitude + decorrelation terms plus a direct DC term), averaged over fields.

Both modes sum the annular bins UNWEIGHTED by default, as in the Subich reference. wavenumber_weight="fastnet" instead upweights each bin by gamma_k = max(N_k * k**sqrt(3), 1) (FastNet, arXiv:2509.17601), which moves the penalty from the large scales that dominate a steep sea-ice spectrum onto the fine scales where the ice edge lives; the total is held fixed so spectral_weight stays comparable across the two settings (see _weight_bins). Default OFF: unset, the loss is bit-for-bit unchanged.

Numerical guards follow the reference implementation: eps inside the geometric-mean square root (the coherence denominator is singular for a zero field), coherence clamped at 1 from above (Cauchy-Schwarz can be violated by the eps), the mean/DC component handled as a direct squared difference (catastrophic cancellation), and all spectral arithmetic in float32.

mode = mode instance-attribute

spectral_weight = spectral_weight instance-attribute

delta = delta instance-attribute

merge_bins_below = merge_bins_below instance-attribute

eps = eps instance-attribute

wavenumber_weight = weight_mode instance-attribute

static_ref_path = static_ref_path instance-attribute

binned_spectra(prediction, target)

Per-field, per-bin spectral statistics of the mean-removed fields.

Returns:

Type Description
Tensor

Tuple (P, T, C, dc) where P[i, b] / T[i, b] are the
Tensor

prediction/target powers of field i in |k|-bin b,
Tensor

C[i, b] is the real cross-spectrum, and dc[i] is the
Tensor

squared difference of the field means. Powers are normalised so
tuple[Tensor, Tensor, Tensor, Tensor]

that sum_b (P + T - 2C) + dc equals the per-field MSE
tuple[Tensor, Tensor, Tensor, Tensor]

(Parseval's identity).

spectral_excess(prediction, target, n_groups=1)

Per-field AMSE excess AMSE - MSE (the pure anti-blur surcharge).

Non-negative; zero iff, in every |k|-bin, the two power spectra match or the coherence is perfect. Shrinking predicted amplitude below the target's cannot reduce this term.

n_groups splits the fields (in their flattened order) into equal groups that are normalised independently by _weight_bins.

pure_amse(prediction, target, n_groups=1)

Per-field AMSE (amplitude + adjusted decorrelation terms + DC term).

n_groups is as for spectral_excess.

forward(prediction, target)

Return the scalar loss for prediction/target of shape [..., H, W].

per_lead_time_loss(prediction, target)

Return the loss at each lead time for NTCHW prediction/target, shape [T].

Equal to stack([self(prediction[:, t], target[:, t]) for t in ...]) but computed with a single FFT pass over every lead time at once.

LeadTimeWeightedLoss

Bases: Module

Weight a base loss function so that later lead times contribute more.

exponent property

forward(prediction, target)

Return the lead-time weighted scalar loss for NTCHW prediction/target.

weights(n_steps, device)

Return the per-lead-time weights for n_steps forecast steps.

RMSELoss

Bases: Module

Root Mean Squared Error loss.

mse = nn.MSELoss() instance-attribute

eps = eps instance-attribute

forward(prediction, target)