From e76f135019d09a22727033b7e2531e18f2f6daf2 Mon Sep 17 00:00:00 2001 From: Adonyth Date: Sun, 12 Jul 2026 10:30:46 -0400 Subject: [PATCH 1/4] Average long_time_metrics time curves over validation batches Inside Trainer.validation_loop, 'time_logs |= new_time_logs' overwrote the long_time_metrics time curves on every batch, so the curves passed to plot_all_time_metrics reflected only the last validation batch. Accumulate them as running means over the batch count instead, exactly mirroring how loss_dict is accumulated a few lines below. Add a lightweight validation_loop test with a stubbed rollout. Fixes #78 --- tests/test_trainer.py | 78 ++++++++++++++++++++++++++ the_well/benchmark/trainer/training.py | 5 +- 2 files changed, 82 insertions(+), 1 deletion(-) create mode 100644 tests/test_trainer.py diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 00000000..a61f33e5 --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,78 @@ +import torch + +from the_well.benchmark.metrics import VRMSE +from the_well.benchmark.trainer import training +from the_well.data.datasets import WellMetadata + + +def test_validation_loop_averages_time_logs_over_batches(monkeypatch, tmp_path): + """The time curves logged for long_time_metrics must be averaged over all + validation batches, mirroring how loss_dict is accumulated, instead of + reflecting only the last batch.""" + meta = WellMetadata( + dataset_name="test", + n_spatial_dims=1, + spatial_resolution=(8,), + scalar_names=[], + constant_scalar_names=[], + field_names={0: ["test"]}, + constant_field_names={}, + boundary_condition_types=["periodic"], + n_files=1, + n_trajectories_per_file=[10], + n_steps_per_trajectory=[100], + ) + metric = VRMSE() + + # Build a Trainer without running __init__ (which requires a full + # model/datamodule stack) and set only the attributes validation_loop uses. + trainer = training.Trainer.__new__(training.Trainer) + trainer.model = torch.nn.Identity() + trainer.device = torch.device("cpu") + trainer.enable_amp = False + trainer.amp_type = torch.float16 + trainer.formatter = None + trainer.dset_metadata = meta + trainer.validation_suite = [metric] + trainer.loss_fn = metric + trainer.num_time_intervals = 2 + trainer.short_validation_length = 10 + trainer.viz_folder = str(tmp_path) + trainer.make_rollout_videos = False + trainer.is_distributed = False + # The rollout itself is not under test: each "batch" already holds + # a (y_pred, y_ref) pair with shape (batch, time, space, field). + trainer.rollout_model = lambda model, batch, formatter, train=False: batch + + torch.manual_seed(0) + batches = [] + for scale in (1.0, 10.0): + y_ref = torch.randn(2, 4, 8, 1) + y_pred = y_ref + scale * torch.randn(2, 4, 8, 1) + batches.append((y_pred, y_ref)) + + captured = {} + + def capture_time_logs(time_logs, *args, **kwargs): + captured.update(time_logs) + + monkeypatch.setattr(training, "validation_plots", []) + monkeypatch.setattr(training, "plot_all_time_metrics", capture_time_logs) + + trainer.validation_loop(batches, full=True) + + # Expected: the average over batches of the per-batch rollout curves. + per_batch_curves = [] + for y_pred, y_ref in batches: + sub_loss = metric(y_pred, y_ref, meta).mean(0) + _, batch_time_logs = trainer.split_up_losses( + sub_loss, "VRMSE", meta.dataset_name, ["test"] + ) + per_batch_curves.append(batch_time_logs) + + key = f"{meta.dataset_name}/full_VRMSE_rollout" + assert key in captured + expected = (per_batch_curves[0][key] + per_batch_curves[1][key]) / 2 + last_batch_only = per_batch_curves[1][key] + assert not torch.allclose(expected, last_batch_only) + assert torch.allclose(captured[key], expected) diff --git a/the_well/benchmark/trainer/training.py b/the_well/benchmark/trainer/training.py index 7d9e0ad7..b5e545af 100755 --- a/the_well/benchmark/trainer/training.py +++ b/the_well/benchmark/trainer/training.py @@ -363,7 +363,10 @@ def validation_loop( ) # 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 log_name, log_value in new_time_logs.items(): + time_logs[log_name] = ( + time_logs.get(log_name, 0.0) + log_value / denom + ) for loss_name, loss_value in new_losses.items(): loss_dict[loss_name] = ( loss_dict.get(loss_name, 0.0) + loss_value / denom From d70e8ebf54c9ab4e9cfddb061d52287212616e25 Mon Sep 17 00:00:00 2001 From: Adonyth Date: Sun, 9 Aug 2026 19:52:28 -0400 Subject: [PATCH 2/4] Extract the batch-mean accumulator and test it directly Review feedback: the previous test built a partially initialized Trainer via __new__ and reproduced by hand the attributes validation_loop happens to touch, which is brittle against unrelated changes in that method. Extract the accumulation into a module-level helper, accumulate_batch_mean, and test that directly. No Trainer, no metadata, no monkeypatching. The helper is now used for BOTH accumulators in validation_loop. That is the point rather than a tidy-up: the scalar losses were already averaged over batches while the time curves were merged with |=, so the curves reflected whichever batch came last. Routing both through one function removes the asymmetry that produced the bug instead of only correcting its symptom. Verified the helper reproduces the previous inline loop exactly on the same inputs. Five tests cover the mean over N batches, the fold-versus-update difference that is the actual regression, tensor values, unseen keys, in-place return, and disjoint keys. Co-Authored-By: Claude Opus 5 --- tests/test_trainer.py | 129 ++++++++++--------------- the_well/benchmark/trainer/training.py | 30 ++++-- 2 files changed, 74 insertions(+), 85 deletions(-) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index a61f33e5..3fd22aad 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,78 +1,55 @@ import torch -from the_well.benchmark.metrics import VRMSE -from the_well.benchmark.trainer import training -from the_well.data.datasets import WellMetadata - - -def test_validation_loop_averages_time_logs_over_batches(monkeypatch, tmp_path): - """The time curves logged for long_time_metrics must be averaged over all - validation batches, mirroring how loss_dict is accumulated, instead of - reflecting only the last batch.""" - meta = WellMetadata( - dataset_name="test", - n_spatial_dims=1, - spatial_resolution=(8,), - scalar_names=[], - constant_scalar_names=[], - field_names={0: ["test"]}, - constant_field_names={}, - boundary_condition_types=["periodic"], - n_files=1, - n_trajectories_per_file=[10], - n_steps_per_trajectory=[100], - ) - metric = VRMSE() - - # Build a Trainer without running __init__ (which requires a full - # model/datamodule stack) and set only the attributes validation_loop uses. - trainer = training.Trainer.__new__(training.Trainer) - trainer.model = torch.nn.Identity() - trainer.device = torch.device("cpu") - trainer.enable_amp = False - trainer.amp_type = torch.float16 - trainer.formatter = None - trainer.dset_metadata = meta - trainer.validation_suite = [metric] - trainer.loss_fn = metric - trainer.num_time_intervals = 2 - trainer.short_validation_length = 10 - trainer.viz_folder = str(tmp_path) - trainer.make_rollout_videos = False - trainer.is_distributed = False - # The rollout itself is not under test: each "batch" already holds - # a (y_pred, y_ref) pair with shape (batch, time, space, field). - trainer.rollout_model = lambda model, batch, formatter, train=False: batch - - torch.manual_seed(0) - batches = [] - for scale in (1.0, 10.0): - y_ref = torch.randn(2, 4, 8, 1) - y_pred = y_ref + scale * torch.randn(2, 4, 8, 1) - batches.append((y_pred, y_ref)) - - captured = {} - - def capture_time_logs(time_logs, *args, **kwargs): - captured.update(time_logs) - - monkeypatch.setattr(training, "validation_plots", []) - monkeypatch.setattr(training, "plot_all_time_metrics", capture_time_logs) - - trainer.validation_loop(batches, full=True) - - # Expected: the average over batches of the per-batch rollout curves. - per_batch_curves = [] - for y_pred, y_ref in batches: - sub_loss = metric(y_pred, y_ref, meta).mean(0) - _, batch_time_logs = trainer.split_up_losses( - sub_loss, "VRMSE", meta.dataset_name, ["test"] - ) - per_batch_curves.append(batch_time_logs) - - key = f"{meta.dataset_name}/full_VRMSE_rollout" - assert key in captured - expected = (per_batch_curves[0][key] + per_batch_curves[1][key]) / 2 - last_batch_only = per_batch_curves[1][key] - assert not torch.allclose(expected, last_batch_only) - assert torch.allclose(captured[key], expected) +from the_well.benchmark.trainer.training import accumulate_batch_mean + + +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} diff --git a/the_well/benchmark/trainer/training.py b/the_well/benchmark/trainer/training.py index b5e545af..e35babec 100755 --- a/the_well/benchmark/trainer/training.py +++ b/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,24 @@ 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 + + class Trainer: def __init__( self, @@ -363,14 +381,8 @@ def validation_loop( ) # TODO get better way to include spectral error. if k in long_time_metrics or "spectral_error" in k: - for log_name, log_value in new_time_logs.items(): - time_logs[log_name] = ( - time_logs.get(log_name, 0.0) + log_value / denom - ) - for loss_name, loss_value in new_losses.items(): - loss_dict[loss_name] = ( - loss_dict.get(loss_name, 0.0) + loss_value / denom - ) + accumulate_batch_mean(time_logs, new_time_logs, denom) + accumulate_batch_mean(loss_dict, new_losses, denom) count += 1 if not full and count >= self.short_validation_length: break From 2068c925e9179613333e6ebf0baff97c64dc19be Mon Sep 17 00:00:00 2001 From: Adonyth Date: Sun, 9 Aug 2026 20:04:31 -0400 Subject: [PATCH 3/4] Restore a regression guard for the routing, without a Trainer My previous commit answered the review but dropped something the brittle test had been providing: coverage of validation_loop's routing. All five helper tests would pass with `time_logs |= new_time_logs` restored, so nothing failed if the bug came back. Lift the per-batch fold into fold_batch_losses(), taking split_fn as a parameter instead of reaching through self. That keeps the reviewer's point -- no __new__, no reproducing validation_loop's attribute surface -- while making the routing testable: which accumulator each metric lands in, and whether the curve is averaged or overwritten. Mutation-checked: reverting the time-log path to the original `|=` fails test_fold_batch_losses_averages_time_curves_over_batches and nothing else. Co-Authored-By: Claude Opus 5 --- tests/test_trainer.py | 54 +++++++++++++++++++++++++- the_well/benchmark/trainer/training.py | 45 ++++++++++++++++----- 2 files changed, 89 insertions(+), 10 deletions(-) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 3fd22aad..90b1a745 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,6 +1,10 @@ import torch -from the_well.benchmark.trainer.training import accumulate_batch_mean +from the_well.benchmark.trainer.training import ( + accumulate_batch_mean, + fold_batch_losses, + long_time_metrics, +) def test_accumulate_batch_mean_averages_over_batches(): @@ -53,3 +57,51 @@ def test_accumulate_batch_mean_leaves_disjoint_keys_untouched(): 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 diff --git a/the_well/benchmark/trainer/training.py b/the_well/benchmark/trainer/training.py index e35babec..f935006d 100755 --- a/the_well/benchmark/trainer/training.py +++ b/the_well/benchmark/trainer/training.py @@ -54,6 +54,33 @@ def accumulate_batch_mean( 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, @@ -374,15 +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: - accumulate_batch_mean(time_logs, new_time_logs, denom) - accumulate_batch_mean(loss_dict, new_losses, 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 From e6c05cea4af38dd3be7c0e4c0f726c67bbcccd0e Mon Sep 17 00:00:00 2001 From: Chen Jiaxuan Date: Sun, 9 Aug 2026 20:10:15 -0400 Subject: [PATCH 4/4] style: apply ruff-format to test_trainer.py The pre-commit ruff-format hook reflows the new fold_batch_losses call sites and assert messages. No behaviour change; the 7 tests pass before and after, and reintroducing the time_logs bug still fails exactly one. --- tests/test_trainer.py | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 90b1a745..01cbe299 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -82,26 +82,43 @@ def fake_split(sub_loss, k, dset_name, field_names): 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, + {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 ( + 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" + 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} + 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, + 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