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)

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.

pure_amse(prediction, target)

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

forward(prediction, target)

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

WeightedBCEWithLogitsLoss

Bases: BCEWithLogitsLoss

BCEWithLogits loss with elementwise weighting.

forward(preds, targets, sample_weights=None)

Compute weighted BCEWithLogits loss.

Parameters:

Name Type Description Default
preds Tensor

Predicted values.

 required
targets Tensor

Ground-truth values.

 required
sample_weights Tensor | None

Elementwise weighting tensor. If None, no weighting is applied.

 None

Returns:

Name Type Description
Tensor Tensor

Scalar weighted loss value.

WeightedL1Loss

Bases: L1Loss

L1 loss with elementwise weighting.

forward(preds, targets, sample_weights=None)

Compute weighted L1 loss.

Parameters:

Name Type Description Default
preds Tensor

Predicted values.

 required
targets Tensor

Ground-truth values.

 required
sample_weights Tensor | None

Elementwise weighting tensor. If None, no weighting is applied.

 None

Returns:

Name Type Description
Tensor Tensor

Scalar weighted loss value.

WeightedMSELoss

Bases: MSELoss

Mean-squared error loss with per-element weighting.

forward(preds, targets, sample_weights=None)

Compute weighted mean squared error loss.

Parameters:

Name Type Description Default
preds Tensor

Predicted values.

 required
targets Tensor

Ground-truth values.

 required
sample_weights Tensor | None

Elementwise weighting tensor. If None, no weighting is applied.

 None

Returns:

Name Type Description
Tensor Tensor

Scalar weighted loss value.