Skip to content

Average long_time_metrics time curves over validation batches - #93

Open
Adonyth wants to merge 4 commits into
PolymathicAI:masterfrom
Adonyth:fix/long-time-metrics-batch-average
Open

Average long_time_metrics time curves over validation batches#93
Adonyth wants to merge 4 commits into
PolymathicAI:masterfrom
Adonyth:fix/long-time-metrics-batch-average

Conversation

@Adonyth

@Adonyth Adonyth commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Fixes #78

Inside Trainer.validation_loop, the time curves collected for long_time_metrics are merged with time_logs |= new_time_logs, which overwrites every key on each batch — the curves handed to plot_all_time_metrics reflect only the last validation batch. A two-batch toy example makes it concrete: if a metric's curve is 100.0 on batch 1 and 1.0 on batch 2, the logged value is 1.0 instead of the 50.5 average. This is inconsistent with loss_dict three lines below, which correctly accumulates loss_dict.get(k, 0.0) + v / denom.

The fix mirrors the loss_dict accumulation exactly, so the time curves become means over the same denom (the number of validation batches actually processed):

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
        )

Test note: Trainer isn't constructible without a full model/datamodule stack, so tests/test_trainer.py builds one via Trainer.__new__, sets only the attributes validation_loop reads, stubs rollout_model to return pre-made (y_pred, y_ref) batches, and captures what reaches plot_all_time_metrics. It asserts the captured full_VRMSE_rollout curve equals the mean of the two per-batch curves (and not the last batch's curve). The test fails on current master and passes with this change; ruff check / ruff format are clean under the pre-commit pinned v0.6.4. Happy to restructure the test if you'd prefer a different fixture approach.

Found while running an independent quantified-impact study of the benchmark's metric design; happy to adjust to maintainer preferences.

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 PolymathicAI#78
Comment thread tests/test_trainer.py

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good.

Adonyth and others added 3 commits August 9, 2026 19:52
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: long_time_metrics reflect only last batch of data and not the average over all batches

2 participants