diff --git a/changelog.d/final-calibration-weights.fixed.md b/changelog.d/final-calibration-weights.fixed.md new file mode 100644 index 000000000..21210147e --- /dev/null +++ b/changelog.d/final-calibration-weights.fixed.md @@ -0,0 +1 @@ +Save and return the calibrated weights from the final optimiser step. Weights were saved only every tenth epoch, so a 512-epoch build shipped the state after epoch 510 of 0–511 and dropped the last update (#456). diff --git a/policyengine_uk_data/tests/test_calibrate_save.py b/policyengine_uk_data/tests/test_calibrate_save.py index b778596ed..0ebce898a 100644 --- a/policyengine_uk_data/tests/test_calibrate_save.py +++ b/policyengine_uk_data/tests/test_calibrate_save.py @@ -1,18 +1,24 @@ -"""Regression test for calibrate_local_areas weight saving. - -Ensures that the non-verbose branch of `calibrate_local_areas` saves the -weight file to disk — previously the save block was indented outside the -for-loop, so on a typical 512-epoch run (511 % 10 = 1) no save ever ran -and the output .h5 file was never written. +"""Tests for calibrate_local_areas: weight priors, saving and logging. + +Two save bugs are pinned here. The save block once sat outside the training +loop, so a 512-epoch run (511 % 10 = 1) never wrote the weight file. It then +saved only when epoch % 10 == 0, so a 512-epoch build shipped the state after +epoch 510 of 0-511 and dropped the final update (#456). A property test +checks, in both training loops, that the saved file, the returned weights and +the last logged row all describe the final optimiser state. """ from __future__ import annotations import importlib.util +import tempfile +from pathlib import Path import numpy as np import pandas as pd import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st if ( importlib.util.find_spec("torch") is None @@ -24,7 +30,12 @@ ) -def _make_toy_inputs(n_households: int = 4, area_count: int = 2, n_targets: int = 2): +def _make_toy_inputs( + n_households: int = 4, + area_count: int = 2, + n_targets: int = 2, + country_mask: np.ndarray | None = None, +): """Build minimal matrix/target/mask arrays for `calibrate_local_areas`. The calibration routine treats `matrix` as (n_households, n_targets), @@ -39,8 +50,9 @@ def _make_toy_inputs(n_households: int = 4, area_count: int = 2, n_targets: int matrix = pd.DataFrame(np.eye(n_households, n_targets, dtype=float)) # Per-area local targets — identical rows are fine for a smoke test. local_targets = pd.DataFrame(np.tile([100.0, 200.0][:n_targets], (area_count, 1))) - # Country mask: every area includes every household (simple case). - country_mask = np.ones((area_count, n_households), dtype=float) + # Country mask: by default every area includes every household. + if country_mask is None: + country_mask = np.ones((area_count, n_households), dtype=float) national_matrix = pd.DataFrame(np.ones((n_households, 1), dtype=float)) national_targets = pd.Series([1000.0]) @@ -54,6 +66,46 @@ def national_matrix_fn(_dataset): return matrix_fn, national_matrix_fn +# Two countries of two households each: area 0 holds households 0 and 1, area +# 1 holds 2 and 3. Its zeros make exp(parameters) differ from the saved +# weights, so a save that dropped the mask would fail. +TWO_COUNTRY_MASK = np.array([[1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 1.0]]) + + +def _national_performance(weights, _m_c, _y_c, m_n, y_n, _excluded_targets): + """A one-row `get_performance` for the toy's single national target.""" + estimates = weights.sum(axis=0) @ m_n + error = float(estimates.iloc[0] - y_n.iloc[0]) + return pd.DataFrame( + { + "name": ["UK"], + "metric": ["national_total"], + "estimate": [float(estimates.iloc[0])], + "target": [float(y_n.iloc[0])], + "error": [error], + "abs_error": [abs(error)], + "rel_abs_error": [abs(error) / float(y_n.iloc[0])], + "validation": [False], + } + ) + + +def _toy_loss(weights, matrix_fn, national_matrix_fn): + """The calibrator's training loss without dropout, recomputed in float64.""" + + def sre(x, y): + return np.minimum(((1 + x) / (1 + y) - 1) ** 2, ((1 + y) / (1 + x) - 1) ** 2) + + weights = weights.astype(np.float64) + matrix, local_targets, _ = matrix_fn(None) + national_matrix, national_targets = national_matrix_fn(None) + local = sre(weights @ matrix.to_numpy(), local_targets.to_numpy()) + national = sre( + weights.sum(axis=0) @ national_matrix.to_numpy(), national_targets.to_numpy() + ) + return local.mean() + national.mean() + + class _StubDataset: """Minimal stand-in for a UKSingleYearDataset. @@ -212,22 +264,6 @@ def test_calibrate_local_areas_logs_loss_targets_and_source_diagnostics( household_is_spi_synthetic=[False, True, False, True], ) - def get_performance(weights, _m_c, _y_c, m_n, y_n, _excluded_targets): - estimates = weights.sum(axis=0) @ m_n - error = float(estimates.iloc[0] - y_n.iloc[0]) - return pd.DataFrame( - { - "name": ["UK"], - "metric": ["national_total"], - "estimate": [float(estimates.iloc[0])], - "target": [float(y_n.iloc[0])], - "error": [error], - "abs_error": [abs(error)], - "rel_abs_error": [abs(error) / float(y_n.iloc[0])], - "validation": [False], - } - ) - weight_file = "toy_diagnostic_weights.h5" log_csv = tmp_path / "diagnostics.csv" calibrate_local_areas( @@ -239,7 +275,7 @@ def get_performance(weights, _m_c, _y_c, m_n, y_n, _excluded_targets): dataset_key="2025", epochs=1, log_csv=log_csv, - get_performance=get_performance, + get_performance=_national_performance, verbose=False, ) @@ -258,3 +294,77 @@ def get_performance(weights, _m_c, _y_c, m_n, y_n, _excluded_targets): assert row["household_is_spi_synthetic_rows"] == 2 assert row["household_is_spi_synthetic_prior_share"] == pytest.approx(0.5) assert row["household_is_spi_synthetic_household_weight"] > 0 + + +def _final_state_run(inputs, epochs, verbose): + """Calibrate the toy problem in a scratch folder. + + Returns the saved weights, the returned household weights, the final + optimiser state (exp of the trained parameters times the country mask) + and the performance log. + """ + import h5py + import torch + + from policyengine_uk_data.utils import calibrate as calibrate_module + from policyengine_uk_data.utils.calibrate import calibrate_local_areas + + parameters = [] + + class SpyAdam(torch.optim.Adam): + def __init__(self, params, *args, **kwargs): + params = list(params) + parameters.extend(params) + super().__init__(params, *args, **kwargs) + + matrix_fn, national_matrix_fn = inputs + # Hypothesis reruns the test body, so it can't use pytest's + # function-scoped tmp_path and monkeypatch fixtures. + with ( + tempfile.TemporaryDirectory() as folder, + pytest.MonkeyPatch.context() as monkeypatch, + ): + folder = Path(folder) + monkeypatch.setattr(calibrate_module, "STORAGE_FOLDER", folder) + monkeypatch.setattr(torch.optim, "Adam", SpyAdam) + returned = calibrate_local_areas( + dataset=_StubDataset(np.array([1.0, 2.0, 3.0, 4.0])), + matrix_fn=matrix_fn, + national_matrix_fn=national_matrix_fn, + area_count=2, + weight_file="final.h5", + dataset_key="2025", + epochs=epochs, + log_csv=folder / "log.csv", + get_performance=_national_performance, + verbose=verbose, + ) + with h5py.File(folder / "final.h5", "r") as f: + saved = f["2025"][:] + log = pd.read_csv(folder / "log.csv") + (weights,) = parameters + final = torch.exp(weights).detach().numpy() * TWO_COUNTRY_MASK + return saved, returned.household["household_weight"].to_numpy(), final, log + + +@settings(max_examples=20, deadline=None) +@example(epochs=32, verbose=True) # a TESTING build +@example(epochs=32, verbose=False) # the same, through the non-verbose loop +@given(epochs=st.integers(1, 40), verbose=st.booleans()) +def test_saved_returned_and_logged_weights_are_the_final_optimiser_state( + epochs, verbose +): + """#456: weights were saved only when epoch % 10 == 0, so a 512-epoch + build shipped the state after epoch 510 of 0-511 and dropped the last + update.""" + inputs = _make_toy_inputs(country_mask=TWO_COUNTRY_MASK) + saved, returned, final, log = _final_state_run(inputs, epochs, verbose) + np.testing.assert_array_equal(saved, final) + np.testing.assert_allclose(returned, final.sum(axis=0), rtol=1e-6) + # Every tenth epoch and the last are logged, so the last row is the + # final epoch, and its saved_weights_loss is the saved weights' loss. + assert log["epoch"].iloc[-1] == epochs - 1 + assert log["epoch"].tolist() == sorted({*range(0, epochs, 10), epochs - 1}) + assert log["saved_weights_loss"].iloc[-1] == pytest.approx( + _toy_loss(saved, *inputs), rel=1e-5 + ) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index f97c4e1ef..cf6e0b158 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -423,6 +423,11 @@ def log_performance( ) performance_step["epoch"] = epoch performance_step["loss"] = performance_step.rel_abs_error**2 + # training_loss, validation_loss and the two pct_close columns describe + # the dropout-perturbed weights before this epoch's optimiser step. + # saved_weights_loss and the per-target columns describe final_weights, + # the weights after the step, which are the ones saved. So on the last + # row, which describes the shipped weights, read saved_weights_loss. performance_step["training_loss"] = training_loss performance_step["saved_weights_loss"] = saved_weights_loss performance_step["validation_loss"] = validation_loss @@ -476,7 +481,9 @@ def log_performance( calculating_loss=False, ) - if epoch % 10 == 0: + # Every tenth epoch, and the last: the weights saved and + # returned are the final optimiser state (#456). + if epoch % 10 == 0 or epoch == epochs - 1: final_weights = (torch.exp(weights) * r).detach().numpy() log_performance( @@ -516,7 +523,7 @@ def log_performance( f"Loss: {loss_value.item()}, Epoch: {epoch}, {area_name}<10%: {local_close:.1%}, National<10%: {national_close:.1%}" ) - if epoch % 10 == 0: + if epoch % 10 == 0 or epoch == epochs - 1: final_weights = (torch.exp(weights) * r).detach().numpy() log_performance(