import abc
import dataclasses
from collections.abc import Callable, Mapping
from typing import Any, Literal
import torch
import torch.linalg
import torch.nn.functional as F
from fme.core.device import get_device
from fme.core.ensemble import get_crps, get_energy_score
from fme.core.gridded_ops import GriddedOperations
from fme.core.normalizer import StandardNormalizer
from fme.core.packer import Packer
from fme.core.typing_ import TensorMapping
@dataclasses.dataclass
class ChannelLossInfo:
"""Per-channel loss value and the number of batch samples that contributed."""
loss: torch.Tensor
count: int
class LossComponent(abc.ABC):
"""A pre-weighted loss tensor that knows how to reduce itself to ``(B, C)``.
All loss tensors are pre-weighted so that ``.mean()`` over trailing
(non-batch, non-channel) dimensions gives the correct per-sample,
per-channel loss. Subclasses encode the tensor layout (where the
channel dimension lives) and implement :meth:`reduce_to_channel`.
"""
def __init__(self, loss: torch.Tensor):
"""
Args:
loss: The loss tensor. Can be a scalar (``ndim == 0``),
a partially-reduced tensor like ``(B, C)``, or a
full element-wise tensor like ``(B, C, lat, lon)``.
"""
self.loss = loss
@abc.abstractmethod
def reduce_to_channel(self) -> torch.Tensor:
"""Reduce to ``(B, C)`` by meaning over non-batch, non-channel dims."""
class StandardLoss(LossComponent):
"""Standard ``(B, C, ...)`` layout with channel at dim 1."""
def reduce_to_channel(self) -> torch.Tensor:
if self.loss.ndim <= 2:
return self.loss
return self.loss.mean(dim=tuple(range(2, self.loss.ndim)))
class EnsembleComponentLoss(LossComponent):
"""Ensemble ``(B, E, C, ...)`` layout with channel at dim 2."""
def reduce_to_channel(self) -> torch.Tensor:
dims = tuple(i for i in range(self.loss.ndim) if i not in (0, 2))
return self.loss.mean(dim=dims) if dims else self.loss
class LossOutput:
"""Container for loss values returned by WeightedMappingLoss/StepLoss.
Holds one or more :class:`LossComponent` instances and provides
convenience methods for the scalar total and per-channel breakdowns.
Reduction is computed once and cached: ``total()`` derives from
the cached per-channel values so they are always consistent.
"""
def __init__(
self,
losses: list[LossComponent],
channel_names: list[str],
mask: torch.Tensor | None = None,
):
self._losses = losses
self._channel_names = channel_names
self._mask = mask
self._per_channel: torch.Tensor | None = None
self._counts: list[int] | None = None
def _reduce(self) -> tuple[torch.Tensor, list[int]]:
"""Return ``(per_channel, counts)`` tensors, computed once.
When a mask is present (shape ``(B, C)``), each channel's loss
is averaged only over the batch samples where that channel is
present, so masked-out variables never dilute the result.
"""
if self._per_channel is None:
bc = sum(c.reduce_to_channel() for c in self._losses)
assert isinstance(bc, torch.Tensor)
if bc.ndim == 0:
self._per_channel = bc.expand(len(self._channel_names))
self._counts = [1] * len(self._channel_names)
elif self._mask is not None:
masked_sum = (bc * self._mask).sum(dim=0)
self._per_channel = masked_sum / self._mask.sum(dim=0).clamp(min=1)
self._counts = [int(c.item()) for c in self._mask.sum(dim=0)]
else:
self._per_channel = bc.mean(dim=0)
self._counts = [bc.shape[0]] * len(self._channel_names)
assert self._per_channel is not None and self._counts is not None
return self._per_channel, self._counts
def total(self) -> torch.Tensor:
"""Scalar loss used as the optimization target.
This is the mean of the per-channel losses across channels (over
active channels only when a mask is present), not a sum. Adding
or removing channels therefore does not change the scale of the
returned value.
"""
pc, _ = self._reduce()
if self._mask is not None:
active = self._mask.sum(dim=0) > 0
if active.any():
return pc[active].mean()
return pc.mean()
def get_channel_losses(self) -> dict[str, ChannelLossInfo]:
"""Per-channel mean losses with active-sample counts.
Each :class:`ChannelLossInfo` carries the mean loss for that
channel (averaged over active samples only) and the number of
batch samples that contributed. Downstream aggregators should
use the counts to compute properly weighted means across
batches.
"""
pc, counts = self._reduce()
n_channels = len(self._channel_names)
if pc.ndim > 0 and pc.shape[0] != n_channels:
raise RuntimeError(
f"Per-channel loss has {pc.shape[0]} elements but "
f"{n_channels} channel names were provided."
)
return {
name: ChannelLossInfo(loss=pc[i], count=counts[i])
for i, name in enumerate(self._channel_names)
}
def scale(self, weight: float) -> "LossOutput":
"""Return a new ``LossOutput`` with every component scaled."""
return LossOutput(
[type(c)(c.loss * weight) for c in self._losses],
self._channel_names,
mask=self._mask,
)
class _MSELoss(torch.nn.Module):
"""MSE with ``reduction="none"`` that returns ``list[LossComponent]``."""
def __init__(self):
super().__init__()
self._loss = torch.nn.MSELoss(reduction="none")
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
return [StandardLoss(self._loss(x, y))]
class _L1Loss(torch.nn.Module):
"""L1 with ``reduction="none"`` that returns ``list[LossComponent]``."""
def __init__(self):
super().__init__()
self._loss = torch.nn.L1Loss(reduction="none")
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
return [StandardLoss(self._loss(x, y))]
class NaNLoss(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, input: torch.Tensor, target: torch.Tensor) -> list[LossComponent]:
return [StandardLoss(torch.tensor(torch.nan))]
class WeightedMappingLoss:
def __init__(
self,
loss: Callable[
[torch.Tensor, torch.Tensor], list[LossComponent] | torch.Tensor
],
weights: dict[str, float],
out_names: list[str],
normalizer: StandardNormalizer,
channel_dim: int = -3,
):
"""
Args:
loss: The loss function to apply. Should return a
``list[LossComponent]``. Element-wise losses (e.g.
``torch.nn.MSELoss``) that return a raw tensor are also
accepted and will be wrapped automatically based on
*channel_dim*.
weights: A dictionary of variable names with individual
weights to apply to their normalized losses
out_names: The names of the output variables.
normalizer: The normalizer to use.
channel_dim: The channel dimension of the input tensors.
"""
self._weight_tensor = _construct_weight_tensor(
weights, out_names, channel_dim=channel_dim
)
self.loss = VariableWeightingLoss(
weights=self._weight_tensor,
loss=loss,
)
if self._weight_tensor.flatten().shape[0] != len(out_names):
raise RuntimeError(
"The number of weights must match the number of output names, "
"behavior of _construct_weight_tensor has changed."
)
self.packer = Packer(out_names)
self.channel_dim = channel_dim
self.normalizer = normalizer
def __call__(
self,
predict_dict: TensorMapping,
target_dict: TensorMapping,
data_mask: TensorMapping | None = None,
) -> LossOutput:
"""
Args:
predict_dict: The predicted data.
target_dict: The target data.
data_mask: Optional per-variable boolean masks of shape
``[batch]`` indicating which samples have each variable
present. Used to exclude masked channels from the loss
average.
Returns:
A ``LossOutput`` wrapping pre-weighted loss component tensors.
"""
predict_tensors = self.packer.pack(
self.normalizer.normalize(predict_dict), axis=self.channel_dim
)
target_tensors = self.packer.pack(
self.normalizer.normalize(target_dict), axis=self.channel_dim
)
nan_mask = target_tensors.isnan()
if nan_mask.any():
predict_tensors = torch.where(nan_mask, 0.0, predict_tensors)
target_tensors = torch.where(nan_mask, 0.0, target_tensors)
result = self.loss(predict_tensors, target_tensors)
input_ndim = predict_tensors.ndim
cdim = (
input_ndim + self.channel_dim if self.channel_dim < 0 else self.channel_dim
)
def _wrap_elementwise(t: torch.Tensor) -> StandardLoss:
# Element-wise loss tensors have the same shape as the input;
# the channel position depends on the data layout (ensemble,
# tile, etc.). Reduce non-(batch, channel) dims here so the
# downstream component carries a canonical ``(B, C)`` tensor.
dims = tuple(i for i in range(t.ndim) if i not in (0, cdim))
reduced = t.mean(dim=dims) if dims else t
return StandardLoss(reduced)
if isinstance(result, list):
# Inner losses that return raw element-wise tensors (e.g. MSE,
# L1) wrap themselves in StandardLoss but don't know the input
# channel layout, so reduce around the actual channel dim here.
losses = [
_wrap_elementwise(c.loss)
if c.loss.ndim == input_ndim and type(c) is StandardLoss
else c
for c in result
]
else:
losses = [_wrap_elementwise(result)]
mask = None
if data_mask is not None:
batch_size = predict_tensors.shape[0]
device = predict_tensors.device
filled: dict[str, torch.Tensor] = {}
for name in self.packer.names:
if name in data_mask:
filled[name] = data_mask[name].to(device=device, dtype=torch.float)
else:
filled[name] = torch.ones(
batch_size, device=device, dtype=torch.float
)
mask = self.packer.pack(filled, axis=1)
return LossOutput(
losses=losses,
channel_names=list(self.packer.names),
mask=mask,
)
def get_normalizer_state(self) -> dict[str, float]:
return self.normalizer.get_state()
def _construct_weight_tensor(
weights: dict[str, float],
out_names: list[str],
n_dim: int = 4,
channel_dim: int = -3,
) -> torch.Tensor:
"""Creates a packed weight tensor with the appropriate dimensions for
broadcasting with generated or target output tensors. When used in
the n_forward_steps loop in the stepper's run_on_batch, the channel dim is
-3 and the n_dim is 4 (sample, channel, lat, lon).
Args:
weights: dict of variable names with individual weights to apply
to their normalized loss
out_names: list of output variable names
n_dim: number of dimensions of the output tensor
channel_dim: the channel dimension of the output tensor
"""
weights_tensor = torch.tensor([weights.get(key, 1.0) for key in out_names])
# positive index of the channel dimension
_channel_dim = n_dim + channel_dim if channel_dim < 0 else channel_dim
reshape_dim = (
len(weights_tensor) if i == _channel_dim else 1 for i in range(n_dim)
)
return weights_tensor.reshape(*reshape_dim).to(get_device(), dtype=torch.float)
class LpLoss(torch.nn.Module):
def __init__(self, p=2):
"""
Args:
p: Lp-norm type. For example, p=1 for L1-norm, p=2 for L2-norm.
"""
super().__init__()
if p <= 0:
raise ValueError("Lp-norm type should be positive")
self.p = p
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
B, C = x.shape[0], x.shape[1]
x_flat = x.reshape(B, C, -1)
y_flat = y.reshape(B, C, -1)
diff_norms = torch.linalg.norm(x_flat - y_flat, ord=self.p, dim=2)
y_norms = torch.linalg.norm(y_flat, ord=self.p, dim=2)
return [StandardLoss(diff_norms / y_norms)]
class AreaWeightedMSELoss(torch.nn.Module):
def __init__(self, area_weighted_mean: Callable[[torch.Tensor], torch.Tensor]):
super().__init__()
self._area_weighted_mean = area_weighted_mean
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
return [StandardLoss(self._area_weighted_mean((x - y) ** 2))]
class WeightedSum(torch.nn.Module):
"""
A module which applies multiple loss-function modules (taking two inputs)
and returns their weighted components as a flat list.
"""
def __init__(self, modules: list[torch.nn.Module], weights: list[float]):
"""
Args:
modules: A list of modules, each of which takes two tensors and
returns a ``list[LossComponent]``.
weights: A list of weights to apply to the outputs of the modules.
"""
super().__init__()
if len(modules) != len(weights):
raise ValueError("modules and weights must have the same length")
self._wrapped = modules
self._weights = weights
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
components: list[LossComponent] = []
for w, module in zip(self._weights, self._wrapped):
for c in module(x, y):
components.append(type(c)(c.loss * w))
return components
class GlobalMeanLoss(torch.nn.Module):
"""
A module which computes a loss on the global mean of each sample.
"""
def __init__(
self,
area_weighted_mean: Callable[[torch.Tensor], torch.Tensor],
loss: torch.nn.Module,
):
"""
Args:
area_weighted_mean: Computes an area-weighted mean, removing the
horizontal dimensions.
loss: A loss function which takes two tensors of shape
(n_samples, n_channels) and returns a
``list[LossComponent]``.
"""
super().__init__()
self.global_mean = GlobalMean(area_weighted_mean)
self.loss = loss
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
x = self.global_mean(x)
y = self.global_mean(y)
return self.loss(x, y)
class GlobalMean(torch.nn.Module):
def __init__(self, area_weighted_mean: Callable[[torch.Tensor], torch.Tensor]):
"""
Args:
area_weighted_mean: Computes an area-weighted mean, removing the
horizontal dimensions.
"""
super().__init__()
self._area_weighted_mean = area_weighted_mean
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: A tensor with spatial dimensions in shape (n_samples, n_timesteps,
n_channels, n_lat, n_lon).
"""
return self._area_weighted_mean(x)
class VariableWeightingLoss(torch.nn.Module):
def __init__(self, weights: torch.Tensor, loss: torch.nn.Module):
"""
Args:
weights: A tensor of shape (n_samples, n_channels, n_lat, n_lon)
containing the weights to apply to each channel.
loss: A loss function which takes two tensors.
"""
super().__init__()
self.loss = loss
self.weights = weights
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
return self.loss(self.weights * x, self.weights * y)
class SpectralWhitening:
"""Per-sample spectral-whitening reweight for :class:`EnergyScoreLoss`.
``factor`` returns a per-(sample, channel, degree) multiplier (broadcast over
order m) computed from the *detached* target coefficients, so it only
reweights the energy score across degrees and adds no gradient path. The
unnormalized weight is ``(1 / amp_l) ** exponent`` where ``amp_l`` is the
per-degree RMS amplitude over valid orders (m <= l), floored at ``eps_frac``
of the per-sample mean amplitude to bound the boost at near-zero-power
degrees. ``exponent`` gamma in (0, 1] sets the strength: gamma=1 fully
flattens the spectrum, gamma<1 whitens partially. The factor is rescaled per
(sample, channel) to preserve the amplitude-weighted total, so the overall
energy-score magnitude (and the meaning of ``energy_score_weight``) is
unchanged. A white-spectrum target yields a uniform factor (no-op).
"""
def __init__(self, eps_frac: float = 0.02, exponent: float = 1.0):
self.eps_frac = eps_frac
self.exponent = exponent
def factor(self, y_hat: torch.Tensor) -> torch.Tensor:
# The target carries a singleton ensemble dim (B, 1, [C], L, M); drop it
# up front so the per-(sample, [channel], degree) factor aligns with the
# energy score, whose ensemble dim get_energy_score has already reduced.
yt = y_hat.detach()[:, 0] # (B, [C], L, M)
amp_mode = yt.abs() # (B, [C], L, M)
real_dtype = amp_mode.dtype
n_l, n_m = yt.shape[-2], yt.shape[-1]
l_idx = torch.arange(n_l, device=yt.device).unsqueeze(-1) # (L, 1)
m_idx = torch.arange(n_m, device=yt.device).unsqueeze(0) # (1, M)
valid = (m_idx <= l_idx).to(real_dtype) # (L, M); zero where m > l
# real-SHT redundancy (m>0 counts double) x validity, clean (L, M) shape
redundancy = 2.0 * torch.ones(n_l, n_m, device=yt.device, dtype=real_dtype)
redundancy[:, 0] = 1.0
w = redundancy * valid # (L, M)
tiny = torch.finfo(real_dtype).tiny
# per-mode mean power at degree l (over valid m) -> per-mode RMS amplitude
meanpow_l = (amp_mode**2 * w).sum(dim=-1) / w.sum(dim=-1).clamp_min(tiny)
amp_l = torch.sqrt(meanpow_l) # (B, [C], L)
mean_amp = amp_l.mean(dim=-1, keepdim=True)
f = 1.0 / torch.clamp(amp_l, min=self.eps_frac * mean_amp)
if self.exponent != 1.0:
f = f**self.exponent
f_m = f.unsqueeze(-1) # (B, [C], L, 1), broadcast over m
# Magnitude preservation: the per-mode energy score scales ~ |y_hat|, so
# rescale so sum_lm w * |y_hat| is unchanged by the reweight.
num = (w * amp_mode).sum(dim=(-2, -1), keepdim=True)
den = (w * f_m * amp_mode).sum(dim=(-2, -1), keepdim=True)
alpha = num / (den + tiny)
return alpha * f_m # (B, [C], L, 1)
#: Default whitening strength when ``kind='per_sample'`` and ``exponent`` is
#: unset. A two-seed gamma in {0, 0.2, 0.5, 0.8} sweep selected 0.5 as the
#: stable knee: it captures most of the small-scale spectral gain while staying
#: rollout-stable, whereas full whitening (gamma=1) over-upweights
#: noise-dominated low-amplitude high-l degrees and destabilizes residual
#: rollouts. See the whitening-gamma-selection report linked from PR #1303.
_DEFAULT_WHITENING_EXPONENT = 0.5
#: Default per-degree amplitude floor (fraction of the per-sample mean degree
#: amplitude) when ``kind='per_sample'`` and ``eps_frac`` is unset.
_DEFAULT_WHITENING_EPS_FRAC = 0.02
@dataclasses.dataclass
class SpectralWhiteningConfig:
"""Configures per-sample spectral whitening of the energy score (see
:class:`SpectralWhitening` for what the reweight does).
Args:
kind: ``'none'`` (the default) disables whitening; ``'per_sample'``
enables the per-sample reweight.
eps_frac: floor on the per-degree amplitude, as a fraction of the
per-sample mean degree-amplitude. It bounds the boost applied to
near-zero-power degrees (where ``1 / amp_l`` would blow up).
Unset (``None``) defaults to ``0.02`` when whitening is enabled.
Requires ``kind='per_sample'``.
exponent: whitening strength gamma in (0, 1]; the per-degree weight is
``(1 / amp_l) ** gamma``. gamma=1 fully flattens the target
spectrum; smaller gamma whitens partially, taming the upweighting of
noise-dominated low-amplitude degrees. Unset (``None``) defaults to
``0.5`` when whitening is enabled -- the validated stable knee; full
whitening (gamma=1) destabilizes residual rollouts, so it is opt-in
rather than the default. Requires ``kind='per_sample'``.
"""
kind: Literal["none", "per_sample"] = "none"
eps_frac: float | None = None
exponent: float | None = None
def __post_init__(self):
if self.kind not in ("none", "per_sample"):
raise NotImplementedError(
f"spectral whitening kind={self.kind!r} not supported; "
"use 'none' or 'per_sample'."
)
if self.kind == "none":
if self.eps_frac is not None or self.exponent is not None:
raise ValueError(
"eps_frac and exponent require kind='per_sample'; "
"got kind='none'."
)
return
# kind='per_sample': resolve unset fields to their validated defaults,
# then validate. Storing the resolved values keeps build() total and the
# config introspectable.
if self.eps_frac is None:
self.eps_frac = _DEFAULT_WHITENING_EPS_FRAC
if self.exponent is None:
self.exponent = _DEFAULT_WHITENING_EXPONENT
if self.eps_frac <= 0:
raise ValueError(f"eps_frac must be positive, got {self.eps_frac}")
if not 0.0 < self.exponent <= 1.0:
raise ValueError(f"exponent must be in (0, 1], got {self.exponent}")
def build(self) -> SpectralWhitening | None:
if self.kind == "none":
return None
assert self.eps_frac is not None and self.exponent is not None
return SpectralWhitening(eps_frac=self.eps_frac, exponent=self.exponent)
class EnergyScoreLoss(torch.nn.Module):
"""
Compute the energy score over the complex-valued spectral coefficients.
The energy score is defined as
.. math::
E[||X - y||^{beta}] - 1/2 E[||X - X'||^{beta}]
where :math:`X` is the ensemble, :math:`y` is the target, and :math:`||.||`
is the complex modulus. It is a proper scoring rule for beta in (0, 2). Here
we use beta=1. See Gneiting and Raftery (2007) [1]_ Section 4.3 for more details.
We use a scaling factor of 2 * sqrt(n_l * n_m) to bring its magnitude in
line with the real-valued CRPS loss, and to prevent its value depending on domain
size for Gaussian distributed random data where n_lon = 2 * n_lat.
Returns a pre-weighted ``(B, C, L, M)`` tensor where ``.mean(dim=(-2, -1))``
reproduces the old scalar value per ``(B, C)`` pair.
.. [1] https://sites.stat.washington.edu/people/raftery/Research/PDF/Gneiting2007jasa.pdf
"""
def __init__(
self,
sht: Callable[[torch.Tensor], torch.Tensor],
whitening: SpectralWhitening | None = None,
):
super().__init__()
self.sht = sht
# None (whitening disabled) makes forward() skip the reweight entirely.
self._whitening = whitening
self.scaling: float | None = None
self.n_spectral: int | None = None
self.mode_weights: torch.Tensor | None = None
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
x_hat = self.sht(x)
y_hat = self.sht(y)
n_l, n_m = x_hat.shape[-2], x_hat.shape[-1]
if self.scaling is None:
self.scaling = 2 * (n_l * n_m) ** 0.5
self.n_spectral = n_l * n_m
if self.mode_weights is None:
self.mode_weights = 2 * torch.ones(
(*([1] * (x_hat.ndim - 1)), n_l, n_m),
device=x_hat.device,
)
self.mode_weights[..., 0] = 1
assert self.n_spectral is not None
es = get_energy_score(x_hat, y_hat) * self.mode_weights
if self._whitening is not None:
es = es * self._whitening.factor(y_hat)
# Old path: .sum(dim=(-2,-1)).mean() / scaling
# New path: StandardLoss does .mean(dim=(-2,-1)) i.e. sum/(L*M)
# Multiply by L*M/scaling so mean gives the same result as sum/scaling.
pre_weighted = es * (self.n_spectral / self.scaling)
return [StandardLoss(pre_weighted)]
class CRPSLoss(torch.nn.Module):
"""
Compute the CRPS loss.
Supports almost-fair modification to CRPS from
https://arxiv.org/html/2412.15832v1, which claims to be helpful in
avoiding numerical issues with fair CRPS.
"""
def __init__(self, alpha: float):
super().__init__()
self.alpha = alpha
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
return [StandardLoss(get_crps(x, y, alpha=self.alpha))]
class FiniteDifferenceCRPSLoss(torch.nn.Module):
"""
Computes the CRPS of the x and y finite differences of the input tensors,
which helps with representations of horizontal stochastic structures.
Returns a ``(B, C)`` tensor (spatial dims reduced internally because
lat and lon diffs have incompatible shapes).
"""
def __init__(self, alpha: float, levels: int = 1):
super().__init__()
if levels < 1:
raise ValueError(f"levels must be at least 1, got {levels}")
self.alpha = alpha
self.levels = levels
def forward(self, x: torch.Tensor, y: torch.Tensor) -> list[LossComponent]:
result = _get_finite_difference_crps_loss(x, y, self.alpha, levels=self.levels)
return [StandardLoss(result / self.levels)]
def _reduce_spatial(t: torch.Tensor) -> torch.Tensor:
"""Reduce trailing (non-batch, non-channel) dims of a ``(B, C, ...)`` tensor."""
if t.ndim <= 2:
return t
return t.mean(dim=tuple(range(2, t.ndim)))
def _get_finite_difference_crps_loss(
x: torch.Tensor, y: torch.Tensor, alpha: float, levels: int
) -> torch.Tensor:
"""Returns a ``(B, C)`` tensor summing contributions from each level."""
x_diff_lat = x[..., 1:, :] - x[..., :-1, :]
y_diff_lat = y[..., 1:, :] - y[..., :-1, :]
crps_lat = _reduce_spatial(get_crps(x_diff_lat, y_diff_lat, alpha=alpha))
x_diff_lon = torch.roll(x, shifts=-1, dims=-1) - x
y_diff_lon = torch.roll(y, shifts=-1, dims=-1) - y
crps_lon = _reduce_spatial(get_crps(x_diff_lon, y_diff_lon, alpha=alpha))
level_crps = 0.5 * (crps_lat + crps_lon)
if levels > 1:
x_flat = x.reshape(-1, 1, x.shape[-2], x.shape[-1])
y_flat = y.reshape(-1, 1, y.shape[-2], y.shape[-1])
x_pooled = F.avg_pool2d(x_flat, kernel_size=2, stride=2, ceil_mode=True)
y_pooled = F.avg_pool2d(y_flat, kernel_size=2, stride=2, ceil_mode=True)
x_coarse = x_pooled.reshape(
*x.shape[:-2], x_pooled.shape[-2], x_pooled.shape[-1]
)
y_coarse = y_pooled.reshape(
*y.shape[:-2], y_pooled.shape[-2], y_pooled.shape[-1]
)
return level_crps + _get_finite_difference_crps_loss(
x_coarse, y_coarse, alpha=alpha, levels=levels - 1
)
return level_crps
class EnsembleLoss(torch.nn.Module):
def __init__(
self,
crps_weight: float,
energy_score_weight: float,
sht: Callable[[torch.Tensor], torch.Tensor],
finite_difference_crps_weight: float = 0.0,
finite_difference_crps_levels: int = 1,
almost_fair_crps_alpha: float = 1.0,
energy_score_whitening: SpectralWhitening | None = None,
):
super().__init__()
if crps_weight < 0 or energy_score_weight < 0:
raise ValueError(
"crps_weight and energy_score_weight must be non-negative, "
f"got {crps_weight} and {energy_score_weight}"
)
if finite_difference_crps_weight < 0:
raise ValueError(
"finite_difference_crps_weight must be non-negative, "
f"got {finite_difference_crps_weight}"
)
if crps_weight + energy_score_weight == 0:
raise ValueError(
"crps_weight and energy_score_weight must sum to a positive value, "
f"got {crps_weight} and {energy_score_weight}"
)
self.crps_loss = CRPSLoss(alpha=almost_fair_crps_alpha)
if finite_difference_crps_weight > 0:
self.diff_crps_loss: FiniteDifferenceCRPSLoss | None = (
FiniteDifferenceCRPSLoss(
alpha=almost_fair_crps_alpha,
levels=finite_difference_crps_levels,
)
)
else:
self.diff_crps_loss = None
self.energy_score_loss = EnergyScoreLoss(
sht=sht,
whitening=energy_score_whitening,
)
self.crps_weight = crps_weight
self.diff_crps_weight = finite_difference_crps_weight
self.energy_score_weight = energy_score_weight
def forward(
self,
gen_norm: torch.Tensor,
target_norm: torch.Tensor,
) -> list[LossComponent]:
components: list[LossComponent] = []
if self.crps_weight > 0:
for c in self.crps_loss(gen_norm, target_norm):
components.append(type(c)(c.loss * self.crps_weight))
if self.energy_score_weight > 0:
for c in self.energy_score_loss(gen_norm, target_norm):
components.append(type(c)(c.loss * self.energy_score_weight))
if self.diff_crps_loss is not None:
for c in self.diff_crps_loss(gen_norm, target_norm):
components.append(type(c)(c.loss * self.diff_crps_weight))
return components
@dataclasses.dataclass
class LossConfig:
"""
A dataclass containing all the information needed to build a loss function,
including the type of the loss function and the data needed to build it.
Args:
type: the type of the loss function
kwargs: data for a loss function instance of the indicated type
global_mean_type: the type of the loss function to apply to the global
mean of each sample, by default no loss is applied
global_mean_kwargs: data for a loss function instance of the indicated
type to apply to the global mean of each sample
global_mean_weight: the weight to apply to the global mean loss
relative to the main loss
"""
type: Literal["LpLoss", "L1", "MSE", "AreaWeightedMSE", "NaN", "EnsembleLoss"] = (
"MSE"
)
kwargs: Mapping[str, Any] = dataclasses.field(default_factory=lambda: {})
global_mean_type: Literal["LpLoss"] | None = None
global_mean_kwargs: Mapping[str, Any] = dataclasses.field(
default_factory=lambda: {}
)
global_mean_weight: float = 1.0
def __post_init__(self):
if self.type not in (
"LpLoss",
"L1",
"MSE",
"AreaWeightedMSE",
"NaN",
"EnsembleLoss",
):
raise NotImplementedError(self.type)
if self.global_mean_type is not None and self.global_mean_type != "LpLoss":
raise NotImplementedError(self.global_mean_type)
def build(
self,
gridded_operations: GriddedOperations | None,
) -> Any:
"""
Args:
gridded_operations: The gridded operations to use in the case that
the loss function requires use of the horizontal dimensions.
"""
if self.type == "LpLoss":
main_loss = LpLoss(**self.kwargs)
elif self.type == "L1":
main_loss = _L1Loss()
elif self.type == "MSE":
main_loss = _MSELoss()
elif self.type == "AreaWeightedMSE":
if gridded_operations is None:
raise ValueError("gridded_operations is required for AreaWeightedMSE")
main_loss = AreaWeightedMSELoss(gridded_operations.area_weighted_mean)
elif self.type == "NaN":
main_loss = NaNLoss()
elif self.type == "EnsembleLoss":
if gridded_operations is None:
raise ValueError("gridded_operations is required for EnsembleLoss")
kwargs = dict(self.kwargs)
crps_weight = kwargs.pop("crps_weight", 1.0)
energy_score_weight = kwargs.pop("energy_score_weight", 0.0)
# kwargs is opaque (Mapping[str, Any]), so dacite does not descend
# into it; build the nested whitening config from its dict here, then
# pass the built operator (or None) down to EnsembleLoss.
whitening_config = kwargs.pop("energy_score_whitening", None)
if isinstance(whitening_config, Mapping):
whitening_config = SpectralWhiteningConfig(**whitening_config)
whitening = (
whitening_config.build() if whitening_config is not None else None
)
main_loss = EnsembleLoss(
sht=gridded_operations.get_real_sht(),
crps_weight=crps_weight,
energy_score_weight=energy_score_weight,
energy_score_whitening=whitening,
**kwargs,
)
if self.global_mean_type is not None:
if gridded_operations is None:
raise ValueError("gridded_operations is required for global mean loss")
global_mean_loss = GlobalMeanLoss(
area_weighted_mean=gridded_operations.area_weighted_mean,
loss=LpLoss(**self.global_mean_kwargs),
)
final_loss = WeightedSum(
modules=[main_loss, global_mean_loss],
weights=[1.0, self.global_mean_weight],
)
else:
final_loss = main_loss
return final_loss.to(device=get_device())
class StepLoss(torch.nn.Module):
def __init__(
self, loss: WeightedMappingLoss, sqrt_loss_decay_constant: float = 0.0
):
super().__init__()
self.loss = loss
self.sqrt_loss_decay_constant = sqrt_loss_decay_constant
@property
def _normalizer(self) -> StandardNormalizer:
# private because this is only used in unit tests
return self.loss.normalizer
def forward(
self,
predict_dict: TensorMapping,
target_dict: TensorMapping,
step: int,
data_mask: TensorMapping | None = None,
) -> LossOutput:
"""
Args:
predict_dict: The predicted data.
target_dict: The target data.
step: The step number, indexed from 0 for the first step.
data_mask: Optional per-variable boolean masks forwarded to
the underlying :class:`WeightedMappingLoss`.
Returns:
A ``LossOutput`` wrapping the step-weighted loss tensor.
"""
step_weight = (1.0 + self.sqrt_loss_decay_constant * step) ** (-0.5)
return self.loss(predict_dict, target_dict, data_mask=data_mask).scale(
step_weight
)
[docs]@dataclasses.dataclass
class StepLossConfig:
"""
Loss configuration class that has the same fields as LossConfig but also
has additional weights field, and optional step loss decay.
The build method will apply the weights to
the inputs of the loss function. The loss returned by build will be a
MappingLoss, which takes Dict[str, tensor] as inputs instead of packed
tensors.
Args:
type: the type of the loss function
kwargs: data for a loss function instance of the indicated type
global_mean_type: the type of the loss function to apply to the global
mean of each sample, by default no loss is applied
global_mean_kwargs: data for a loss function instance of the indicated
type to apply to the global mean of each sample
global_mean_weight: the weight to apply to the global mean loss
relative to the main loss
sqrt_loss_step_decay_constant: the constant to use for the square root
loss step decay, alpha in 1/sqrt(1.0 + alpha * step) where step is
indexed from 0 for the first step.
weights: A dictionary of variable names with individual
weights to apply to their normalized losses
"""
type: Literal["LpLoss", "MSE", "AreaWeightedMSE", "EnsembleLoss"] = "MSE"
kwargs: Mapping[str, Any] = dataclasses.field(default_factory=lambda: {})
global_mean_type: Literal["LpLoss"] | None = None
global_mean_kwargs: Mapping[str, Any] = dataclasses.field(
default_factory=lambda: {}
)
global_mean_weight: float = 1.0
sqrt_loss_step_decay_constant: float = 0.0
weights: dict[str, float] = dataclasses.field(default_factory=lambda: {})
def __post_init__(self):
self.loss_config = LossConfig(
type=self.type,
kwargs=self.kwargs,
global_mean_type=self.global_mean_type,
global_mean_kwargs=self.global_mean_kwargs,
global_mean_weight=self.global_mean_weight,
)
def build(
self,
gridded_ops: GriddedOperations | None,
out_names: list[str],
normalizer: StandardNormalizer,
channel_dim: int = -3,
) -> StepLoss:
loss = self.loss_config.build(
gridded_operations=gridded_ops,
)
return StepLoss(
WeightedMappingLoss(
loss=loss,
weights=self.weights,
out_names=out_names,
channel_dim=channel_dim,
normalizer=normalizer,
),
sqrt_loss_decay_constant=self.sqrt_loss_step_decay_constant,
)