Average long_time_metrics time curves over validation batches

tests/test_trainer.py

@@ -0,0 +1,124 @@
+import torch
+
+from the_well.benchmark.trainer.training import (
+    accumulate_batch_mean,
+    fold_batch_losses,
+    long_time_metrics,
+)
+
+
+def test_accumulate_batch_mean_averages_over_batches():
+    """Folding N batches with denom=N yields their mean, not the last one."""
+    acc = {}
+    values = [1.0, 10.0, 100.0]
+    for v in values:
+        accumulate_batch_mean(acc, {"a": v}, len(values))
+    assert acc["a"] == sum(values) / len(values)
+
+
+def test_accumulate_batch_mean_is_not_an_update():
+    """The regression this guards: merging with dict-update keeps only the last batch.
+
+    ``time_logs`` used to be merged with ``|=``, so a per-timestep curve logged over
+    several validation batches reflected whichever batch happened to come last, while
+    the scalar losses beside it were averaged. This pins the difference.
+    """
+    curves = [torch.tensor([1.0, 2.0]), torch.tensor([9.0, 10.0])]
+    folded, updated = {}, {}
+    for c in curves:
+        accumulate_batch_mean(folded, {"curve": c}, len(curves))
+        updated |= {"curve": c}
+    expected_mean = (curves[0] + curves[1]) / 2
+    assert torch.allclose(folded["curve"], expected_mean)
+    assert torch.allclose(updated["curve"], curves[-1])
+    assert not torch.allclose(updated["curve"], expected_mean)
+
+
+def test_accumulate_batch_mean_handles_tensors_and_new_keys():
+    """Tensor values accumulate elementwise, and an unseen key starts from zero."""
+    acc = {"seen": torch.tensor([1.0, 1.0])}
+    accumulate_batch_mean(acc, {"seen": torch.tensor([3.0, 3.0])}, 1.0)
+    accumulate_batch_mean(acc, {"fresh": torch.tensor([2.0, 4.0])}, 2.0)
+    assert torch.allclose(acc["seen"], torch.tensor([4.0, 4.0]))
+    assert torch.allclose(acc["fresh"], torch.tensor([1.0, 2.0]))
+
+
+def test_accumulate_batch_mean_returns_the_accumulator():
+    """It updates in place and returns the same object, so calls can be chained."""
+    acc = {}
+    out = accumulate_batch_mean(acc, {"a": 2.0}, 2.0)
+    assert out is acc
+    assert out["a"] == 1.0
+
+
+def test_accumulate_batch_mean_leaves_disjoint_keys_untouched():
+    """Batches that log different keys do not clobber one another."""
+    acc = {}
+    accumulate_batch_mean(acc, {"a": 4.0}, 2.0)
+    accumulate_batch_mean(acc, {"b": 6.0}, 2.0)
+    assert acc == {"a": 2.0, "b": 3.0}
+
+
+def test_fold_batch_losses_averages_time_curves_over_batches():
+    """The regression guard for the bug this PR fixes.
+
+    Unlike the helper tests above, this exercises the ROUTING in validation_loop: which
+    accumulator each metric's curve lands in, and whether it is averaged or overwritten.
+    Reverting the time-log path to ``time_logs |= new_time_logs`` fails this test; it
+    would pass every other test in this file.
+
+    ``split_fn`` is injected, so no Trainer is constructed and the test does not depend
+    on which attributes validation_loop happens to touch.
+    """
+    metric = long_time_metrics[0]
+    curves = [torch.tensor([1.0, 2.0]), torch.tensor([9.0, 10.0])]
+
+    def fake_split(sub_loss, k, dset_name, field_names):
+        # (scalar losses, time curves) -- shaped like split_up_losses' return
+        return {f"{dset_name}/{k}_T=all": sub_loss.mean()}, {
+            f"{dset_name}/full_{k}_rollout": sub_loss
+        }
+
+    loss_dict, time_logs = {}, {}
+    for c in curves:
+        fold_batch_losses(
+            {metric: c.unsqueeze(0)},
+            fake_split,
+            "ds",
+            ["f"],
+            len(curves),
+            loss_dict,
+            time_logs,
+        )
+
+    key = f"ds/full_{metric}_rollout"
+    expected = (curves[0] + curves[1]) / 2
+    assert (
+        key in time_logs
+    ), f"a long-horizon metric must reach time_logs; got {list(time_logs)}"
+    assert torch.allclose(time_logs[key], expected)
+    assert not torch.allclose(
+        time_logs[key], curves[-1]
+    ), "curve reflects only the last batch"
+
+
+def test_fold_batch_losses_keeps_short_metrics_out_of_time_logs():
+    """A metric that is not long-horizon contributes scalars only, no curve."""
+
+    def fake_split(sub_loss, k, dset_name, field_names):
+        return {f"{dset_name}/{k}_T=all": sub_loss.mean()}, {
+            f"{dset_name}/{k}_rollout": sub_loss
+        }
+
+    loss_dict, time_logs = {}, {}
+    fold_batch_losses(
+        {"not_a_long_time_metric": torch.tensor([[1.0, 2.0]])},
+        fake_split,
+        "ds",
+        ["f"],
+        1.0,
+        loss_dict,
+        time_logs,
+    )
+    assert time_logs == {}
+    assert loss_dict["ds/not_a_long_time_metric_T=all"] == 1.5

the_well/benchmark/trainer/training.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import time
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, Optional
 
 import numpy as np
 import torch
@@ -36,6 +36,51 @@ def param_norm(parameters):
         return total_norm**0.5
 
 
+def accumulate_batch_mean(
+    accumulator: Dict[str, Any], values: Dict[str, Any], denom: float
+) -> Dict[str, Any]:
+    """Fold one batch's values into a running mean over ``denom`` batches.
+
+    ``accumulator`` is updated in place and returned. A key seen for the first time
+    starts from zero, so callers do not need to pre-populate it.
+
+    Both the scalar validation losses and the long-horizon time curves are reduced this
+    way. Keeping them on one code path is deliberate: the curves were previously merged
+    with ``|=``, which kept only the final batch, while the scalars were averaged --- an
+    asymmetry that is easy to reintroduce if the two are written out separately.
+    """
+    for name, value in values.items():
+        accumulator[name] = accumulator.get(name, 0.0) + value / denom
+    return accumulator
+
+
+def fold_batch_losses(
+    losses: Dict[str, Any],
+    split_fn: Callable,
+    dset_name: str,
+    field_names: Any,
+    denom: float,
+    loss_dict: Dict[str, Any],
+    time_logs: Dict[str, Any],
+) -> None:
+    """Fold one validation batch's losses into the running accumulators.
+
+    ``split_fn`` is ``Trainer.split_up_losses``, injected rather than reached through
+    ``self`` so this routing can be exercised without building a Trainer.
+
+    Both accumulators are running means over ``denom`` batches. The long-horizon time
+    curves in particular must be averaged and not merged: they were previously combined
+    with ``|=``, which silently kept only the final batch.
+    """
+    for k, v in losses.items():
+        sub_loss = v.mean(0)
+        new_losses, new_time_logs = split_fn(sub_loss, k, dset_name, field_names)
+        # TODO get better way to include spectral error.
+        if k in long_time_metrics or "spectral_error" in k:
+            accumulate_batch_mean(time_logs, new_time_logs, denom)
+        accumulate_batch_mean(loss_dict, new_losses, denom)
+
+
 class Trainer:
     def __init__(
         self,
@@ -356,18 +401,15 @@ def validation_loop(
                     if not isinstance(loss, dict):
                         loss = {loss_fn.__class__.__name__: loss}
                     # Split the losses and update the logging dictionary
-                    for k, v in loss.items():
-                        sub_loss = v.mean(0)
-                        new_losses, new_time_logs = self.split_up_losses(
-                            sub_loss, k, dset_name, field_names
-                        )
-                        # TODO get better way to include spectral error.
-                        if k in long_time_metrics or "spectral_error" in k:
-                            time_logs |= new_time_logs
-                        for loss_name, loss_value in new_losses.items():
-                            loss_dict[loss_name] = (
-                                loss_dict.get(loss_name, 0.0) + loss_value / denom
-                            )
+                    fold_batch_losses(
+                        loss,
+                        self.split_up_losses,
+                        dset_name,
+                        field_names,
+                        denom,
+                        loss_dict,
+                        time_logs,
+                    )
                 count += 1
                 if not full and count >= self.short_validation_length:
                     break