Fix UnboundLocalError in the CRPS losses for single member forecasts with spatial weights - #204

CRPSLoss, GradientCRPSLoss, VortDivCRPSLoss and KernelScoreLoss raise UnboundLocalError for a single member forecast with spatial weights. The E == 1 path skips the ensemble branch, and that branch is the only place that sets spatial_weights_split. MetricsHandler passes weights as soon as a target contains NaN, so the CRPS metric fails this way for deterministic models. I noticed this in #202 and split it out from there as requested.

import torch
from makani.utils.losses import CRPSLoss

crps = CRPSLoss(img_shape=(32, 64), crop_shape=(32, 64), crop_offset=(0, 0), channel_names=["t2m"], grid_type="equiangular")
fc = torch.randn(2, 1, 1, 32, 64)  # a single member
obs = torch.randn(2, 1, 32, 64)
crps(fc, obs, torch.ones_like(obs))
# main: UnboundLocalError: cannot access local variable 'spatial_weights_split' where it is not associated with a value

The fix flattens the weights before the branch, so both paths have them, and splits them in ensemble direction only when the loss is ensemble distributed. That is what CRPSLoss already did, and it is how the observations are handled. GradientCRPSLoss, VortDivCRPSLoss and KernelScoreLoss used to split the weights whenever the ensemble group had more than one rank, even when the observations stayed whole. With ensemble_distributed=True nothing changes.

test_e1_spatial_weights in tests/test_losses.py checks two things for all four losses. The weighted score of a single member must equal that of two identical members, whose spread is zero. Weights of one must leave the score unchanged. On main it fails with the UnboundLocalError for all four.

Testing, on CPU only: tests/test_losses.py passes (533 passed). I ran it with TORCHDYNAMO_DISABLE=1 because torch.compile cannot invoke the compiler in my environment. tests/distributed/test_distributed_losses.py passes on two gloo ranks, with the ensemble split and with the latitude split. black, ruff and the license header check are clean on the changed files.

Not changed here: with weights per channel, KernelScoreLoss and GradientCRPSLoss(absolute=False) fail with a shape mismatch for any ensemble size. They expand the channels, into the kernel basis and into the two gradient components, but the weights keep the original channel count. Weights that broadcast over the channels work. I can look at that separately if it is useful.

🤖 Generated with Claude Code