[MaxGhenis]

Do not merge on its own. Once approved on its gates, this lands with the d833 uk-data batch on Max's go. It is part of workstream W3 of the Autumn Budget 2026 CI plan (the final-fit gate measures the shipped weights, so they must be the final ones).

Fixes the "weights ship from epoch 510 of 512" item in #456.

Why

Epochs and optimiser steps are numbered from 0 here, as in the calibration log's epoch column: epoch e takes step e, so a 512-epoch run takes steps 0–511.

calibrate_local_areas saved the weight file, set the returned dataset's household_weight and logged performance only when epoch % 10 == 0. The weights it shipped were therefore the state after the last step whose epoch is a multiple of 10, and the steps after that were dropped:

• a release build (512 epochs) shipped the state after step 510 and dropped step 511;
• a TESTING build (32 epochs) shipped the state after step 30 and dropped step 31;
• the logged "final" fit described those weights, not the optimiser's end state.

What changes

• Both training loops, verbose and non-verbose, also save, log and set the weights at the last epoch. Nothing else changes: the same steps run, and the same files are written.
• A comment in log_performance says which log columns describe which weights (see "Reading the last log row").

Effect on the data

Releases now ship the state after step 511 instead of 510, and TESTING builds the state after step 31 instead of 30. So each area's log-weights take one more Adam update (learning rate 0.1). The effect on aggregates has not been measured locally; it needs a full build. Expect small diffs in calibration outputs in the batch. The batch's build-and-test run and its comparison will show them.

Reading the last log row

On every logged row, training_loss, validation_loss and the two pct_close columns come from the dropout-perturbed weights before that epoch's optimiser step. The saved weights, saved_weights_loss and the per-target columns come after it. So on the last row, the one people read for the shipped fit, saved_weights_loss is the consistent loss. Main already logs this way; the PR adds the comment and doesn't change what is logged.

Invariants (tested)

test_saved_returned_and_logged_weights_are_the_final_optimiser_state is a Hypothesis property. It draws epoch counts from 1 to 40 and either loop: 20 examples, plus the TESTING build's 32 verbose epochs, which are pinned. The toy uses a two-country mask with zeros, so the saved weights differ from exp(parameters). For every draw:

1. The saved weight file equals the final optimiser state, exp(trained parameters) × country mask, exactly.
2. The returned household weights are its column sums.
3. The log CSV records every tenth epoch and the last, so its last row has epoch == epochs - 1.
4. That row's saved_weights_loss equals the training loss of the saved weights, recomputed independently in float64 (relative tolerance 1e-5).

Checked outside the suite, with scripts kept in the reviews folder (w3-uk-data-gates-2026-10-11/594-checks/), not committed:

• Exhaustive run. All 80 cases (1–40 epochs × both loops) pass. The largest gap between the logged and recomputed loss is 1.0e-7, relative.
• Mutation runs. The property fails on each of six mutants of calibrate.py:
  • main's every-tenth-epoch schedule (it fails first on the pinned 32-epoch TESTING case);
  • * r dropped from the saved weights in both loops, in the verbose loop only, or in the non-verbose loop only;
  • the training loss logged as saved_weights_loss;
  • the last epoch not logged.
• The previous test. The parametrized test it replaces, with its all-ones mask, passed all 16 of its cases with * r dropped in both loops.

Interactions

#507 (seeded dropout) edits the same function. It changes the imports, the signature and dropout_weights. This PR changes the save conditions and adds a comment in log_performance. git merge-tree of this head with #507's head (60a87ff7) is clean, and the merged tree passes pytest policyengine_uk_data/tests/test_calibrat*.py: 28 passed.

Tests

From runs at e13f22b8, Python 3.13, serial:

• pytest policyengine_uk_data/tests/test_calibrat*.py: 18 passed. The count was 33 at 677f8039. It fell because the 16 parametrized cases are now one property test.
• pytest policyengine_uk_data/tests/test_changelog_fragment_check.py: 88 passed.
• ruff format --check .: 264 files already formatted.

axiom: n/a: calibration build; no policy change

🤖 Generated with Claude Code