Count each sample once in the ACC, SSR and spread of masked data - #202

MetricRollout.update gives missing (NaN) target values zero weight, and compute_counts then counts every sample with the area its weights integrate to. That is right for the area averages (L1, RMSE, CRPS, rank histogram): their result for a sample is an integral over the weights, and finalize renormalizes it to the valid area. The ACC and the SSR of a sample, however, are ratios of integrals over the same weights and are normalized to the valid area already, so they get divided by it a second time. The spread, the square root of such an integral, comes out too large by the square root of that factor. For SST missing over land, say, the reported ACC is its actual value divided by the ocean fraction, roughly 1.4 times too large.

Minimal reproduction on CPU:

import torch
from makani.utils.metrics.functions import GeometricACC

torch.manual_seed(0)
x = torch.randn(4, 1, 32, 64)
y = x + 0.5 * torch.randn_like(x)
acc = GeometricACC("equiangular", img_shape=(32, 64), normalize=True, channel_reduction="none", batch_reduction="sum")

mask = torch.ones_like(x)
mask[..., :32] = 0.0  # the target is missing over half of the grid
print(acc.finalize(acc(x, y, mask), acc.compute_counts(x, mask)))
# main: tensor([1.7753]); with this change: tensor([0.8877]), the ACC over the valid half

Through MetricsHandler, with the target of one channel NaN over half of the grid and an 8 member ensemble, that channel reports ACC 1.30, SSR 1.93 and spread 1.39 on main, and 0.65, 0.97 and 0.98 with this change, in line with the other channels, which do not change.

The change adds GeometricBaseMetric.weighted_counts, which counts a weighted sample with the area of its weights as before, or once if the metric sets normalized_by_weights, as the ACC, SSR and spread now do. GeometricSpread.forward divides its weighted integral by the area of the weights, so that it is the spread over the valid area. A sample without any valid point still counts zero, so a field that is missing entirely stays out of the average, as before. Unweighted metrics and the area averages are unchanged, and so are all metrics for weights that integrate to one, such as the normalized masks the inferencer passes.

test_masked_metrics in tests/test_metrics.py masks about a third of the grid plus one channel of one sample, with the weights scaled by 1 and by 0.25, and compares the ACC, SSR and spread with the metrics over the valid area of every sample, computed in numpy. On main it fails: ACC and SSR come out 1.45 times the reference (5.8 times with the scaled weights) and the spread 1.21 times (2.41 times). With the change it passes.

Testing, on CPU only: tests/test_metrics.py passes (148 passed, 12 skipped for weatherbench2), and so do tests/test_inferencer.py, tests/test_inference.py and tests/test_trainers.py, which run the metrics through the inferencer and the trainers (my environment does not allow worker processes, so I ran those with the data loader workers and the visualizer's process pool switched off from outside the repository), and tests/distributed/test_distributed_metrics.py on two gloo ranks. black, ruff and the license header check are clean on the changed files.

Something I noticed but left out to keep this focused: CRPSLoss.forward raises UnboundLocalError for a single member forecast with spatial weights, because spatial_weights_split is only set in the ensemble branch. MetricsHandler hits it with a one member prediction as soon as a target contains NaN, so the inferencer, which enables the CRPS metric, should fail that way for deterministic models.

🤖 Generated with Claude Code