From a9e59711cd9b589145de4111dc2bad2b0d5d4fa5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:01:10 +0100 Subject: [PATCH 1/3] Run UK simulations on datasets without area codes and guard the weight-matrix fallback A UK dataset without `la_code_oa` (a national-only file) now has no local-authority metadata instead of failing every run and the UK worker check, and its constituency and local-authority breakdowns are omitted. Stage 12 accepts a baseline/reform pair that both lack the metadata. Constituency and local-authority regions built on the enhanced-FRS weight matrices now stop before the simulations are built when the matrix has no weights for the run year or was built for a different number of households, instead of failing inside the run with an out-of-date-matrix error. Refs #725. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 2 + .../simulation_output_geographic.py | 16 ++- .../simulation_runtime.py | 59 ++++++++++ .../stage12_runtime/aggregation.py | 9 +- .../stage12_runtime/simulation.py | 2 + .../uk_local_authority_metadata.py | 24 ++-- .../tests/test_simulation_output_builder.py | 10 +- .../test_simulation_output_geographic.py | 35 +++++- .../tests/test_stage12_runtime.py | 4 + .../tests/test_uk_local_authority_metadata.py | 24 ++++ .../tests/test_uk_weight_matrix_guard.py | 105 ++++++++++++++++++ 11 files changed, 272 insertions(+), 18 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 08bf97ee9..b006f820e 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -6,3 +6,5 @@ - Ensured Stage 12-only deployments run authenticated candidate tests and verify every required deployment phase before reporting success - Allowed Cloud Run stable endpoint checks to wait for documented traffic propagation before restoring the previous revision - Added a dedicated read-only Hugging Face credential for private UK runtime datasets, validated it before synchronization, and tested its worker configuration + - Allowed UK simulations and workers on datasets without area codes, omitting their constituency and local-authority breakdowns instead of failing + - Stopped UK constituency and local-authority runs before they start when the enhanced-FRS weight matrices do not match the selected dataset diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index fb3fb6cda..572f10412 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -210,11 +210,19 @@ def compute_and_format() -> CongressionalDistrictImpactOutput | None: return _try_compute_output("congressional district impacts", compute_and_format) +def _output_household(simulation) -> pd.DataFrame: + return pd.DataFrame(simulation.output_dataset.data.household) + + def build_uk_constituency_impact( country: str, baseline, reform ) -> GeographicImpactOutput | None: if country != "uk": return None + # A UK dataset without area codes (a national-only file) has no + # constituency breakdown. + if "constituency_code_oa" not in _output_household(baseline).columns: + return None lookup_csv_path = _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC) impact = _output_module_function( @@ -242,6 +250,11 @@ def build_uk_local_authority_impact( if country != "uk": return None + baseline_household = _output_household(baseline) + # A UK dataset without area codes (a national-only file) has no + # local-authority breakdown. + if "la_code_oa" not in baseline_household.columns: + return None if uk_local_authority_metadata is None: raise ValueError("UK local-authority boundary metadata is required") @@ -252,8 +265,7 @@ def build_uk_local_authority_impact( load_uk_local_authority_resources, ) - baseline_household = pd.DataFrame(baseline.output_dataset.data.household) - reform_household = pd.DataFrame(reform.output_dataset.data.household) + reform_household = _output_household(reform) numeric_records = compute_longwise_uk_geography_impacts( baseline_household=baseline_household, reform_household=reform_household, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index cf430b0f9..a7ad204f8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -273,6 +273,62 @@ def _build_uk_weight_replacement_region(region_code: str): ) +def _require_uk_weight_matrix_matches_dataset(scoping_strategy, dataset) -> None: + """Stop a weight-matrix region whose matrix was built for another dataset. + + UK constituency and local-authority regions that policyengine.py does not + list reweight households with enhanced-FRS matrices, one column per + household of that file. On any other dataset ``WeightReplacementStrategy`` + fails inside the run and reports the matrix as out of date; this check + stops before the simulations are built and names the cause. + """ + + from policyengine.core.scoping_strategy import WeightReplacementStrategy + + if not isinstance(scoping_strategy, WeightReplacementStrategy): + return + + import h5py + import pandas as pd + from policyengine.data.uk_geography_assets import ( + UKGeographyAssetSpec, + resolve_uk_geography_asset_paths, + ) + + paths = resolve_uk_geography_asset_paths( + UKGeographyAssetSpec( + geography_type="weight replacement", + weight_matrix_filename=scoping_strategy.weight_matrix_key, + lookup_csv_filename=scoping_strategy.lookup_csv_key, + bucket=scoping_strategy.weight_matrix_bucket, + weight_matrix_bucket=scoping_strategy.weight_matrix_bucket, + lookup_csv_bucket=scoping_strategy.lookup_csv_bucket, + ), + download_missing_assets=scoping_strategy.download_missing_assets, + ) + region = scoping_strategy.region_code + matrix_name = scoping_strategy.weight_matrix_key + year = str(dataset.year) + with h5py.File(paths.weight_matrix_path, "r") as matrix: + weights = matrix.get(year) + if not isinstance(weights, h5py.Dataset): + covered = ", ".join(sorted(matrix.keys())) + raise ValueError( + f"UK region {region!r} reweights households with {matrix_name}, " + f"which has no weights for {year} (it covers {covered})." + ) + matrix_households = weights.shape[-1] + households = len(pd.DataFrame(dataset.data.entity_data["household"])) + if households != matrix_households: + raise ValueError( + f"UK region {region!r} reweights households with {matrix_name}, " + f"which was built for {matrix_households} households; the selected " + f"UK dataset has {households}. Constituency and local-authority " + "runs on this dataset need a local-area dataset that carries " + "constituency and local-authority codes." + ) + + def _region_parent_dataset_reference( country_module, country: str, @@ -606,6 +662,9 @@ def _run_simulation_impl_core( ) uk_local_authority_metadata = detect_uk_local_authority_metadata(country, dataset) + _require_uk_weight_matrix_matches_dataset( + region_resolution.scoping_strategy, dataset + ) with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): baseline_policy = _normalise_policy(simulation_params.get("baseline")) reform_policy = _normalise_policy(simulation_params.get("reform")) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index b60db9243..cee7fc177 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -22,11 +22,16 @@ def validate_uk_local_authority_metadata( baseline: UKLocalAuthorityMetadata | None, reform: UKLocalAuthorityMetadata | None, ) -> UKLocalAuthorityMetadata | None: - """Require matching UK authority metadata and reject it elsewhere.""" + """Require matching UK authority metadata and reject it elsewhere. + + Both UK artifacts lack it when the dataset carries no local-authority codes. + """ if country == "uk": + if baseline is None and reform is None: + return None if baseline is None or reform is None: - raise ValueError("UK simulation artifact metadata is missing") + raise ValueError("UK simulation artifact metadata is missing on one side") if baseline != reform: raise ValueError("UK simulation artifact boundary versions do not match") return baseline diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 46c499f1a..a7dbcfd44 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -130,6 +130,7 @@ def calculate_simulation_frames( _country_module, _load_dataset, _normalise_policy, + _require_uk_weight_matrix_matches_dataset, _resolve_dataset_selection, _resolve_region, setup_gcp_credentials, @@ -184,6 +185,7 @@ def calculate_simulation_frames( country, dataset, ) + _require_uk_weight_matrix_matches_dataset(region.scoping_strategy, dataset) policy_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.POLICY_NORMALIZATION)) if runtime is not None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py index 5ddf1a665..4d766f0a1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py @@ -281,7 +281,12 @@ def detect_uk_local_authority_metadata( country: str, dataset: object, ) -> UKLocalAuthorityMetadata | None: - """Inspect a complete dataset before any requested regional scoping.""" + """Inspect a complete dataset before any requested regional scoping. + + A UK dataset without ``la_code_oa`` (a national file that carries no area + codes) has no local-authority metadata, so this returns ``None`` and the + constituency and local-authority outputs are omitted for it. + """ if country != "uk": return None @@ -294,7 +299,7 @@ def detect_uk_local_authority_metadata( raise ValueError("UK dataset contains no household table") household_frame = pd.DataFrame(household) if "la_code_oa" not in household_frame: - raise ValueError("UK dataset household table contains no la_code_oa column") + return None return detect_uk_local_authority_boundary_version( household_frame["la_code_oa"].tolist() ) @@ -302,8 +307,11 @@ def detect_uk_local_authority_metadata( def detect_uk_local_authority_metadata_from_hdf( dataset_path: str, -) -> UKLocalAuthorityMetadata: - """Validate an installed UK HDF dataset against the packaged metadata.""" +) -> UKLocalAuthorityMetadata | None: + """Validate an installed UK HDF dataset against the packaged metadata. + + Returns ``None`` when the household table carries no ``la_code_oa``. + """ observed_codes: set[object] = set() with pd.HDFStore(dataset_path, mode="r") as store: @@ -317,14 +325,10 @@ def detect_uk_local_authority_metadata_from_hdf( ) for chunk in chunks: if "la_code_oa" not in chunk: - raise ValueError( - "UK dataset household table contains no la_code_oa column" - ) + return None observed_codes.update(chunk["la_code_oa"].drop_duplicates().tolist()) except (KeyError, TypeError, ValueError) as error: if "la_code_oa" in str(error): - raise ValueError( - "UK dataset household table contains no la_code_oa column" - ) from error + return None raise return detect_uk_local_authority_boundary_version(observed_codes) diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 869489048..4f0f9e51d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1752,8 +1752,14 @@ def fail_change_output_variable(*args, **kwargs): def test_uk_constituency_impact_uses_policyengine_output_function(monkeypatch): - baseline = object() - reform = object() + def coded_simulation() -> SimpleNamespace: + household = pd.DataFrame({"constituency_code_oa": ["E14001063"]}) + return SimpleNamespace( + output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) + ) + + baseline = coded_simulation() + reform = coded_simulation() expected = [_constituency_impact_record().model_dump(mode="json")] def fake_output_module_function(module_name, name): diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py index 9fbb1007b..e7178f64d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -200,6 +200,37 @@ def test_local_authority_output_requires_detected_boundary_metadata( with pytest.raises(ValueError, match="boundary metadata is required"): simulation_output_geographic.build_uk_local_authority_impact( "uk", - object(), - object(), + _uk_simulation("E06000063", 100.0), + _uk_simulation("E06000063", 110.0), ) + + +def test_uk_geographic_outputs_are_omitted_without_area_codes(monkeypatch) -> None: + monkeypatch.setattr( + simulation_output_geographic, + "_required_uk_geography_lookup_csv_path", + lambda _: pytest.fail("a dataset without area codes needs no lookup"), + ) + household = pd.DataFrame( + { + "region": ["LONDON"], + "household_net_income": [100.0], + "household_weight": [2.0], + } + ) + national = SimpleNamespace( + output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) + ) + + assert ( + simulation_output_geographic.build_uk_constituency_impact( + "uk", national, national + ) + is None + ) + assert ( + simulation_output_geographic.build_uk_local_authority_impact( + "uk", national, national + ) + is None + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index d291a555e..29300f7c2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -564,6 +564,10 @@ def test_country_metadata_alignment_rejects_missing_or_mismatched_uk_values() -> assert validate_uk_local_authority_metadata("us", None, None) is None +def test_country_metadata_alignment_accepts_uk_artifacts_without_area_codes() -> None: + assert validate_uk_local_authority_metadata("uk", None, None) is None + + def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) diff --git a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py index 9a92c2c56..29eea047e 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py @@ -133,6 +133,16 @@ def test_dataset_detector_skips_non_uk_datasets() -> None: assert detect_uk_local_authority_metadata("us", object()) is None +def test_dataset_detector_returns_none_without_local_authority_codes() -> None: + dataset = SimpleNamespace( + data=SimpleNamespace( + entity_data={"household": pd.DataFrame({"region": ["LONDON", "WALES"]})} + ) + ) + + assert detect_uk_local_authority_metadata("uk", dataset) is None + + def test_installed_hdf_detector_reads_local_authority_codes(tmp_path) -> None: dataset_path = tmp_path / "uk-dataset.h5" pd.DataFrame( @@ -152,6 +162,20 @@ def test_installed_hdf_detector_reads_local_authority_codes(tmp_path) -> None: assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD22 +def test_installed_hdf_detector_returns_none_without_local_authority_codes( + tmp_path, +) -> None: + dataset_path = tmp_path / "uk-national-dataset.h5" + pd.DataFrame({"household_id": [1, 2], "region": ["LONDON", "WALES"]}).to_hdf( + dataset_path, + key="household", + format="table", + data_columns=True, + ) + + assert detect_uk_local_authority_metadata_from_hdf(str(dataset_path)) is None + + def test_detector_rejects_mixed_authority_configurations() -> None: with pytest.raises(ValueError, match="mixes LAD22 and LAD23"): detect_uk_local_authority_boundary_version(["E07000026", "E06000063"]) diff --git a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py new file mode 100644 index 000000000..33969e5e9 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py @@ -0,0 +1,105 @@ +"""Unit tests for the UK weight-matrix fallback guard.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import h5py +import numpy as np +import pandas as pd +import pytest +from policyengine.core.scoping_strategy import ( + RowFilterStrategy, + WeightReplacementStrategy, +) + +from policyengine_simulation_executor import simulation_runtime as sr + +_STRATEGY = WeightReplacementStrategy( + weight_matrix_bucket="policyengine-uk-data-private", + weight_matrix_key="parliamentary_constituency_weights.h5", + lookup_csv_bucket="policyengine-uk-data-private", + lookup_csv_key="constituencies_2024.csv", + region_code="E14001063", +) + + +def _dataset(households: int, year: int = 2025) -> SimpleNamespace: + household = pd.DataFrame({"household_id": range(1, households + 1)}) + return SimpleNamespace( + year=year, + data=SimpleNamespace(entity_data={"household": household}), + ) + + +def _matrix(tmp_path, households: int) -> str: + path = tmp_path / "parliamentary_constituency_weights.h5" + with h5py.File(path, "w") as matrix: + matrix.create_dataset("2025", data=np.ones((2, households))) + return str(path) + + +def _resolver(path: str, calls: list): + def resolve(spec, **kwargs): + calls.append((spec, kwargs)) + return SimpleNamespace(weight_matrix_path=path, lookup_csv_path="unused") + + return resolve + + +@pytest.mark.parametrize( + "strategy", + [None, RowFilterStrategy(variable_name="country", variable_value="ENGLAND")], +) +def test_guard_ignores_regions_without_weight_matrices(monkeypatch, strategy) -> None: + monkeypatch.setattr( + "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", + lambda *args, **kwargs: pytest.fail("no weight matrix should be resolved"), + ) + + sr._require_uk_weight_matrix_matches_dataset(strategy, _dataset(3)) + + +def test_guard_accepts_the_dataset_the_matrix_was_built_for( + monkeypatch, tmp_path +) -> None: + calls: list = [] + monkeypatch.setattr( + "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", + _resolver(_matrix(tmp_path, households=3), calls), + ) + + sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(3)) + + [(spec, kwargs)] = calls + assert spec.weight_matrix_filename == _STRATEGY.weight_matrix_key + assert spec.lookup_csv_filename == _STRATEGY.lookup_csv_key + assert spec.resolved_weight_matrix_bucket == _STRATEGY.weight_matrix_bucket + assert kwargs == {"download_missing_assets": True} + + +def test_guard_rejects_a_dataset_of_another_size(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", + _resolver(_matrix(tmp_path, households=3), []), + ) + + with pytest.raises( + ValueError, + match=( + r"'E14001063' reweights households with " + r"parliamentary_constituency_weights\.h5, which was built for 3 " + r"households; the selected UK dataset has 4" + ), + ): + sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(4)) + + +def test_guard_rejects_a_year_the_matrix_does_not_cover(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", + _resolver(_matrix(tmp_path, households=3), []), + ) + + with pytest.raises(ValueError, match=r"has no weights for 2026 \(it covers 2025\)"): + sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(3, 2026)) From 9821812e2d5b4c9a8d3cc34f571716904c1fc8e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:15:06 +0100 Subject: [PATCH 2/3] Admit UK datasets without area codes only when local-area regions are routed Review on #727: treating a missing `la_code_oa` as "this dataset is national-only" let a local-area dataset that lost its area columns pass the worker check and every run, with its breakdowns silently dropped. A UK dataset may now lack area codes only for a national or nation-level run whose policyengine.py bundle routes both constituency and local-authority regions to another dataset (`region_datasets` lists them). Constituency and local-authority runs always require the codes, the installed-dataset check keeps main's strict detection unless that routing exists, and the constituency and local-authority breakdowns raise a clear error instead of returning nothing when the codes are missing. With today's bundle (policyengine 6.2.1) nothing is routed, so the executor admits exactly what main admits. Refs #725, #728. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 3 +- .../simulation_output_geographic.py | 30 +++++-- .../simulation_runtime.py | 12 ++- .../stage12_runtime/aggregation.py | 4 +- .../stage12_runtime/simulation.py | 1 + .../stage12_worker_validation.py | 9 +++ .../uk_local_authority_metadata.py | 74 ++++++++++++++--- .../test_simulation_output_geographic.py | 21 +++-- .../tests/test_stage12_worker_validation.py | 55 +++++++++++++ .../tests/test_uk_local_authority_metadata.py | 81 +++++++++++++++++-- 10 files changed, 251 insertions(+), 39 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index b006f820e..b22504ba7 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -6,5 +6,6 @@ - Ensured Stage 12-only deployments run authenticated candidate tests and verify every required deployment phase before reporting success - Allowed Cloud Run stable endpoint checks to wait for documented traffic propagation before restoring the previous revision - Added a dedicated read-only Hugging Face credential for private UK runtime datasets, validated it before synchronization, and tested its worker configuration - - Allowed UK simulations and workers on datasets without area codes, omitting their constituency and local-authority breakdowns instead of failing + - Accepted a UK dataset without area codes only when the policyengine.py bundle routes constituency and local-authority regions to another dataset, and kept requiring area codes on every constituency and local-authority run + - Raised a clear error, never an omission, when a UK constituency or local-authority breakdown is built over a dataset without area codes - Stopped UK constituency and local-authority runs before they start when the enhanced-FRS weight matrices do not match the selected dataset diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 572f10412..21dff1d00 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -214,15 +214,32 @@ def _output_household(simulation) -> pd.DataFrame: return pd.DataFrame(simulation.output_dataset.data.household) +def _require_uk_area_column( + household: pd.DataFrame, column: str, breakdown: str +) -> None: + """Refuse a UK breakdown over a dataset that lacks its area codes. + + A breakdown is never dropped silently: a dataset without area codes has + none of its own, and one built from a routed local-area dataset is a + separate run. + """ + + if column not in household.columns: + raise ValueError( + f"UK {breakdown} breakdowns need {column} on the simulated dataset, " + "which carries no area codes; for such a dataset they need a run on " + "the local-area dataset its bundle routes these regions to." + ) + + def build_uk_constituency_impact( country: str, baseline, reform ) -> GeographicImpactOutput | None: if country != "uk": return None - # A UK dataset without area codes (a national-only file) has no - # constituency breakdown. - if "constituency_code_oa" not in _output_household(baseline).columns: - return None + _require_uk_area_column( + _output_household(baseline), "constituency_code_oa", "constituency" + ) lookup_csv_path = _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC) impact = _output_module_function( @@ -251,10 +268,7 @@ def build_uk_local_authority_impact( return None baseline_household = _output_household(baseline) - # A UK dataset without area codes (a national-only file) has no - # local-authority breakdown. - if "la_code_oa" not in baseline_household.columns: - return None + _require_uk_area_column(baseline_household, "la_code_oa", "local-authority") if uk_local_authority_metadata is None: raise ValueError("UK local-authority boundary metadata is required") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index a7ad204f8..31b5b9e82 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -310,13 +310,15 @@ def _require_uk_weight_matrix_matches_dataset(scoping_strategy, dataset) -> None matrix_name = scoping_strategy.weight_matrix_key year = str(dataset.year) with h5py.File(paths.weight_matrix_path, "r") as matrix: - weights = matrix.get(year) - if not isinstance(weights, h5py.Dataset): - covered = ", ".join(sorted(matrix.keys())) + if year not in matrix: + covered = ", ".join(sorted(matrix)) raise ValueError( f"UK region {region!r} reweights households with {matrix_name}, " f"which has no weights for {year} (it covers {covered})." ) + weights = matrix[year] + if not isinstance(weights, h5py.Dataset): + raise TypeError(f"{matrix_name} entry {year} is not a weight matrix") matrix_households = weights.shape[-1] households = len(pd.DataFrame(dataset.data.entity_data["household"])) if households != matrix_households: @@ -661,7 +663,9 @@ def _run_simulation_impl_core( detect_uk_local_authority_metadata, ) - uk_local_authority_metadata = detect_uk_local_authority_metadata(country, dataset) + uk_local_authority_metadata = detect_uk_local_authority_metadata( + country, dataset, region_code=region_resolution.code + ) _require_uk_weight_matrix_matches_dataset( region_resolution.scoping_strategy, dataset ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index cee7fc177..a2a08462c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -24,7 +24,9 @@ def validate_uk_local_authority_metadata( ) -> UKLocalAuthorityMetadata | None: """Require matching UK authority metadata and reject it elsewhere. - Both UK artifacts lack it when the dataset carries no local-authority codes. + Both UK artifacts lack it for a national dataset without area codes, which + detection admits only when the bundle routes constituency and + local-authority regions to another dataset. """ if country == "uk": diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index a7dbcfd44..2989a6eee 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -184,6 +184,7 @@ def calculate_simulation_frames( uk_local_authority_metadata = detect_uk_local_authority_metadata( country, dataset, + region_code=region.code, ) _require_uk_weight_matrix_matches_dataset(region.scoping_strategy, dataset) policy_span = ( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index abae8e113..008bfae40 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -80,9 +80,18 @@ def _check_uk_local_authority_resources() -> None: def _check_uk_local_authority_dataset(dataset_path: str) -> None: from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata_from_hdf, + uk_area_regions_routed, + uk_hdf_household_has_local_authority_codes, ) _check_uk_local_authority_resources() + # A national default may carry no area codes only when the bundle routes + # constituency and local-authority regions to another dataset; anything + # else is checked exactly as before. + if uk_area_regions_routed() and not uk_hdf_household_has_local_authority_codes( + dataset_path + ): + return detect_uk_local_authority_metadata_from_hdf(dataset_path) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py index 4d766f0a1..5f5114237 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py @@ -277,15 +277,54 @@ def detect_uk_local_authority_boundary_version( return UKLocalAuthorityMetadata(boundary_version=matches[0].boundary_version) +UK_AREA_REGION_TYPES = ("constituency", "local_authority") +_MISSING_AREA_CODES = ( + "UK dataset household table contains no la_code_oa column; a UK dataset " + "may omit area codes only when the policyengine.py bundle routes " + "constituency and local-authority regions to another dataset" +) + + +def uk_area_regions_routed() -> bool: + """Whether the bundle routes UK constituency and local-authority regions. + + policyengine.py lists a region type in ``region_datasets`` when its + ``regional_dataset_defaults`` names a dataset for it. Only when both area + region types have one may the UK default, a national file, carry no area + codes: their runs and breakdowns then belong to that local-area dataset. + """ + + from policyengine.provenance.manifest import get_release_manifest + + region_datasets = get_release_manifest("uk").region_datasets + return all(region_type in region_datasets for region_type in UK_AREA_REGION_TYPES) + + +def uk_area_codes_required(region_code: str | None) -> bool: + """Whether a UK run's dataset must carry area codes. + + A constituency or local-authority run always needs them. A national or + nation-level run may use a dataset without them only when the bundle + routes both area region types to another dataset. + """ + + region_type = (region_code or "").split("/", maxsplit=1)[0] + if region_type in UK_AREA_REGION_TYPES: + return True + return not uk_area_regions_routed() + + def detect_uk_local_authority_metadata( country: str, dataset: object, + *, + region_code: str | None = None, ) -> UKLocalAuthorityMetadata | None: """Inspect a complete dataset before any requested regional scoping. - A UK dataset without ``la_code_oa`` (a national file that carries no area - codes) has no local-authority metadata, so this returns ``None`` and the - constituency and local-authority outputs are omitted for it. + A UK dataset without ``la_code_oa`` passes, as ``None``, only for a run + that does not need area codes (see ``uk_area_codes_required``); a missing + column is never read as "this dataset is national-only" on its own. """ if country != "uk": @@ -299,19 +338,30 @@ def detect_uk_local_authority_metadata( raise ValueError("UK dataset contains no household table") household_frame = pd.DataFrame(household) if "la_code_oa" not in household_frame: + if uk_area_codes_required(region_code): + raise ValueError(_MISSING_AREA_CODES) return None return detect_uk_local_authority_boundary_version( household_frame["la_code_oa"].tolist() ) +def uk_hdf_household_has_local_authority_codes(dataset_path: str) -> bool: + """Whether an installed UK HDF dataset's household table has ``la_code_oa``.""" + + with pd.HDFStore(dataset_path, mode="r") as store: + if "household" not in store: + raise ValueError("UK dataset contains no household table") + household = store.select("household", stop=0) + if not isinstance(household, pd.DataFrame): + raise TypeError("UK dataset household table is not a data frame") + return "la_code_oa" in household.columns + + def detect_uk_local_authority_metadata_from_hdf( dataset_path: str, -) -> UKLocalAuthorityMetadata | None: - """Validate an installed UK HDF dataset against the packaged metadata. - - Returns ``None`` when the household table carries no ``la_code_oa``. - """ +) -> UKLocalAuthorityMetadata: + """Validate an installed UK HDF dataset against the packaged metadata.""" observed_codes: set[object] = set() with pd.HDFStore(dataset_path, mode="r") as store: @@ -325,10 +375,14 @@ def detect_uk_local_authority_metadata_from_hdf( ) for chunk in chunks: if "la_code_oa" not in chunk: - return None + raise ValueError( + "UK dataset household table contains no la_code_oa column" + ) observed_codes.update(chunk["la_code_oa"].drop_duplicates().tolist()) except (KeyError, TypeError, ValueError) as error: if "la_code_oa" in str(error): - return None + raise ValueError( + "UK dataset household table contains no la_code_oa column" + ) from error raise return detect_uk_local_authority_boundary_version(observed_codes) diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py index e7178f64d..d8a09cc9a 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -205,7 +205,9 @@ def test_local_authority_output_requires_detected_boundary_metadata( ) -def test_uk_geographic_outputs_are_omitted_without_area_codes(monkeypatch) -> None: +def test_uk_geographic_outputs_are_never_dropped_for_missing_area_codes( + monkeypatch, +) -> None: monkeypatch.setattr( simulation_output_geographic, "_required_uk_geography_lookup_csv_path", @@ -222,15 +224,18 @@ def test_uk_geographic_outputs_are_omitted_without_area_codes(monkeypatch) -> No output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) ) - assert ( + with pytest.raises( + ValueError, match="constituency breakdowns need constituency_code_oa" + ): simulation_output_geographic.build_uk_constituency_impact( "uk", national, national ) - is None - ) - assert ( + with pytest.raises(ValueError, match="local-authority breakdowns need la_code_oa"): simulation_output_geographic.build_uk_local_authority_impact( - "uk", national, national + "uk", + national, + national, + uk_local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), ) - is None - ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index d952df443..a1abf01d9 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -172,3 +172,58 @@ def test_dataset_check_verifies_installed_content(tmp_path) -> None: with pytest.raises(RuntimeError, match="digest differs"): _check_dataset_access(str(dataset), "0" * 64) + + +def _uk_household_hdf(tmp_path, columns: dict[str, list]) -> str: + import pandas as pd + + dataset = tmp_path / "uk-dataset.h5" + pd.DataFrame({"household_id": [1, 2], **columns}).to_hdf( + dataset, key="household", format="table", data_columns=True + ) + return str(dataset) + + +@pytest.mark.parametrize("routed", [False, True]) +def test_uk_dataset_check_validates_codes_whenever_present( + monkeypatch, tmp_path, routed: bool +) -> None: + from policyengine_simulation_executor.stage12_worker_validation import ( + _check_uk_local_authority_dataset, + ) + + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + _check_uk_local_authority_dataset( + _uk_household_hdf(tmp_path, {"la_code_oa": ["E06000001", "E06000063"]}) + ) + with pytest.raises(ValueError, match="unsupported local-authority code"): + _check_uk_local_authority_dataset( + _uk_household_hdf(tmp_path, {"la_code_oa": ["E06000001", "E06000999"]}) + ) + + +def test_uk_dataset_check_admits_missing_codes_only_when_area_regions_are_routed( + monkeypatch, tmp_path +) -> None: + from policyengine_simulation_executor.stage12_worker_validation import ( + _check_uk_local_authority_dataset, + ) + + national = _uk_household_hdf(tmp_path, {"region": ["LONDON", "WALES"]}) + routed = {"value": False} + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed["value"], + ) + + with pytest.raises(ValueError, match="contains no la_code_oa column"): + _check_uk_local_authority_dataset(national) + + routed["value"] = True + _check_uk_local_authority_dataset(national) diff --git a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py index 29eea047e..7bef14702 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_local_authority_metadata.py @@ -16,6 +16,8 @@ detect_uk_local_authority_boundary_version, detect_uk_local_authority_metadata_from_hdf, load_uk_local_authority_resources, + uk_area_regions_routed, + uk_hdf_household_has_local_authority_codes, ) LAD22_ONLY_CODES = { @@ -133,14 +135,78 @@ def test_dataset_detector_skips_non_uk_datasets() -> None: assert detect_uk_local_authority_metadata("us", object()) is None -def test_dataset_detector_returns_none_without_local_authority_codes() -> None: - dataset = SimpleNamespace( +def _national_dataset() -> SimpleNamespace: + return SimpleNamespace( data=SimpleNamespace( entity_data={"household": pd.DataFrame({"region": ["LONDON", "WALES"]})} ) ) - assert detect_uk_local_authority_metadata("uk", dataset) is None + +def _routes_area_regions(monkeypatch, routed: bool) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + +@pytest.mark.parametrize("region_code", [None, "uk", "country/england"]) +def test_dataset_detector_requires_codes_unless_area_regions_are_routed( + monkeypatch, region_code: str | None +) -> None: + _routes_area_regions(monkeypatch, False) + + with pytest.raises(ValueError, match="contains no la_code_oa column"): + detect_uk_local_authority_metadata( + "uk", _national_dataset(), region_code=region_code + ) + + +def test_dataset_detector_accepts_a_national_run_when_area_regions_are_routed( + monkeypatch, +) -> None: + _routes_area_regions(monkeypatch, True) + + assert ( + detect_uk_local_authority_metadata("uk", _national_dataset(), region_code="uk") + is None + ) + + +@pytest.mark.parametrize( + "region_code", ["constituency/E14001063", "local_authority/E06000063"] +) +def test_dataset_detector_requires_codes_for_area_runs_even_when_routed( + monkeypatch, region_code: str +) -> None: + _routes_area_regions(monkeypatch, True) + + with pytest.raises(ValueError, match="contains no la_code_oa column"): + detect_uk_local_authority_metadata( + "uk", _national_dataset(), region_code=region_code + ) + + +@pytest.mark.parametrize( + ("region_types", "routed"), + [ + (("national",), False), + (("national", "constituency"), False), + (("national", "constituency", "local_authority"), True), + ], +) +def test_area_regions_count_as_routed_only_when_both_have_a_dataset( + monkeypatch, region_types: tuple[str, ...], routed: bool +) -> None: + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", + lambda country: SimpleNamespace( + region_datasets={region_type: object() for region_type in region_types} + ), + ) + + assert uk_area_regions_routed() is routed def test_installed_hdf_detector_reads_local_authority_codes(tmp_path) -> None: @@ -159,12 +225,11 @@ def test_installed_hdf_detector_reads_local_authority_codes(tmp_path) -> None: metadata = detect_uk_local_authority_metadata_from_hdf(str(dataset_path)) + assert uk_hdf_household_has_local_authority_codes(str(dataset_path)) assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD22 -def test_installed_hdf_detector_returns_none_without_local_authority_codes( - tmp_path, -) -> None: +def test_installed_hdf_detector_rejects_a_dataset_without_codes(tmp_path) -> None: dataset_path = tmp_path / "uk-national-dataset.h5" pd.DataFrame({"household_id": [1, 2], "region": ["LONDON", "WALES"]}).to_hdf( dataset_path, @@ -173,7 +238,9 @@ def test_installed_hdf_detector_returns_none_without_local_authority_codes( data_columns=True, ) - assert detect_uk_local_authority_metadata_from_hdf(str(dataset_path)) is None + assert not uk_hdf_household_has_local_authority_codes(str(dataset_path)) + with pytest.raises(ValueError, match="contains no la_code_oa column"): + detect_uk_local_authority_metadata_from_hdf(str(dataset_path)) def test_detector_rejects_mixed_authority_configurations() -> None: From 14421976e0c6f93d86988dee8f9c84ad6d6168f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:16:20 +0100 Subject: [PATCH 3/3] Bind UK weight-matrix regions to the certified release and require evidence for artifacts without metadata Review on #727: - The weight-matrix guard compared only the year key and the household count, but WeightReplacementStrategy gives households their weights by column position, so a matrix from another release would pass and attach weights to the wrong records (uk-data 1.57.4's matrices have the same 650 x 52,846 shape as the certified 1.56.16 ones, and its households are in a different order). The fallback regions no longer download the unversioned bucket copy. `_require_certified_uk_weight_matrix` materializes, digest-checked, the matrix the policyengine.py bundle certifies in the same release (repository and revision) as the selected dataset, where the strategy looks first, confirms the strategy resolves exactly that file, then checks the run year and household dimension. - `validate_uk_local_authority_metadata` accepted any UK artifact pair without metadata. It now takes the report's region and accepts such a pair only for a national or nation-level request whose bundle routes both constituency and local-authority regions to another dataset: the same evidence per-run detection requires. Refs #725. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 2 +- .../simulation_runtime.py | 97 ++++++++--- .../stage12_qualification.py | 1 + .../stage12_runtime/aggregation.py | 15 +- .../stage12_runtime/coordination.py | 1 + .../stage12_runtime/simulation.py | 6 +- .../tests/test_stage12_runtime.py | 34 +++- .../tests/test_uk_weight_matrix_guard.py | 157 +++++++++++++----- 8 files changed, 241 insertions(+), 72 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index b22504ba7..f11ae4b5b 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -8,4 +8,4 @@ - Added a dedicated read-only Hugging Face credential for private UK runtime datasets, validated it before synchronization, and tested its worker configuration - Accepted a UK dataset without area codes only when the policyengine.py bundle routes constituency and local-authority regions to another dataset, and kept requiring area codes on every constituency and local-authority run - Raised a clear error, never an omission, when a UK constituency or local-authority breakdown is built over a dataset without area codes - - Stopped UK constituency and local-authority runs before they start when the enhanced-FRS weight matrices do not match the selected dataset + - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 31b5b9e82..a6cf23a3c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -269,18 +269,29 @@ def _build_uk_weight_replacement_region(region_code: str): lookup_csv_bucket=asset_spec.bucket, lookup_csv_key=asset_spec.lookup_csv_filename, region_code=value, + # The bucket copy is replaced by every data release; the matrix + # comes from the certified bundle instead, placed locally by + # ``_require_certified_uk_weight_matrix`` before the run. + download_missing_assets=False, ), ) -def _require_uk_weight_matrix_matches_dataset(scoping_strategy, dataset) -> None: - """Stop a weight-matrix region whose matrix was built for another dataset. - - UK constituency and local-authority regions that policyengine.py does not - list reweight households with enhanced-FRS matrices, one column per - household of that file. On any other dataset ``WeightReplacementStrategy`` - fails inside the run and reports the matrix as out of date; this check - stops before the simulations are built and names the cause. +def _require_certified_uk_weight_matrix( + scoping_strategy, dataset, dataset_selection +) -> None: + """Bind a weight-matrix region to the release of the dataset it reweights. + + ``WeightReplacementStrategy`` gives each household the weight in its + column of the matrix, by position. Only the matrix certified with the + selected dataset (same repository and revision in the policyengine.py + bundle) lines up with its households: another release can have the same + shape and a different household order. This places that certified, + digest-checked matrix where the strategy looks first, confirms the + strategy will read exactly that file, and checks the run year and the + household dimension. The unversioned bucket copy, which every data + release replaces, is never downloaded (see + ``_build_uk_weight_replacement_region``). """ from policyengine.core.scoping_strategy import WeightReplacementStrategy @@ -288,26 +299,70 @@ def _require_uk_weight_matrix_matches_dataset(scoping_strategy, dataset) -> None if not isinstance(scoping_strategy, WeightReplacementStrategy): return + from pathlib import Path + import h5py import pandas as pd from policyengine.data.uk_geography_assets import ( UKGeographyAssetSpec, + default_download_dir, resolve_uk_geography_asset_paths, ) + from policyengine.provenance.dataset_materialization import materialize_dataset + from policyengine.provenance.manifest import get_release_manifest + + from policyengine_simulation_executor import simulation_output_geographic - paths = resolve_uk_geography_asset_paths( - UKGeographyAssetSpec( - geography_type="weight replacement", - weight_matrix_filename=scoping_strategy.weight_matrix_key, - lookup_csv_filename=scoping_strategy.lookup_csv_key, - bucket=scoping_strategy.weight_matrix_bucket, - weight_matrix_bucket=scoping_strategy.weight_matrix_bucket, - lookup_csv_bucket=scoping_strategy.lookup_csv_bucket, - ), - download_missing_assets=scoping_strategy.download_missing_assets, - ) region = scoping_strategy.region_code matrix_name = scoping_strategy.weight_matrix_key + manifest = get_release_manifest("uk") + package = manifest.data_package + + def release(reference) -> tuple[str, str]: + return ( + reference.repo_id or package.repo_id, + reference.revision or package.release_manifest_revision or package.version, + ) + + selected = manifest.datasets.get(dataset_selection.name) + if selected is None: + raise ValueError( + f"UK dataset {dataset_selection.name!r} is not in the certified bundle" + ) + certified_name = next( + ( + name + for name, reference in manifest.datasets.items() + if reference.path == matrix_name and release(reference) == release(selected) + ), + None, + ) + if certified_name is None: + raise ValueError( + f"UK region {region!r} reweights households with {matrix_name}, but the " + f"certified bundle has no {matrix_name} from the release of " + f"{dataset_selection.name!r}; a matrix from another release does not " + "line up with its households." + ) + certified = materialize_dataset( + "uk", certified_name, data_dir=default_download_dir() + ) + spec = UKGeographyAssetSpec( + geography_type="weight replacement", + weight_matrix_filename=matrix_name, + lookup_csv_filename=scoping_strategy.lookup_csv_key, + bucket=scoping_strategy.weight_matrix_bucket, + weight_matrix_bucket=scoping_strategy.weight_matrix_bucket, + lookup_csv_bucket=scoping_strategy.lookup_csv_bucket, + ) + simulation_output_geographic._required_uk_geography_lookup_csv_path(spec) + paths = resolve_uk_geography_asset_paths(spec, download_missing_assets=False) + if Path(paths.weight_matrix_path).resolve() != Path(certified.path).resolve(): + raise ValueError( + f"UK region {region!r} would read {paths.weight_matrix_path}, not the " + f"{matrix_name} certified with {dataset_selection.name!r} " + f"({certified.path})." + ) year = str(dataset.year) with h5py.File(paths.weight_matrix_path, "r") as matrix: if year not in matrix: @@ -666,8 +721,8 @@ def _run_simulation_impl_core( uk_local_authority_metadata = detect_uk_local_authority_metadata( country, dataset, region_code=region_resolution.code ) - _require_uk_weight_matrix_matches_dataset( - region_resolution.scoping_strategy, dataset + _require_certified_uk_weight_matrix( + region_resolution.scoping_strategy, dataset, dataset_selection ) with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): baseline_policy = _normalise_policy(simulation_params.get("baseline")) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index fc2ec3c38..473d1ea82 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -308,6 +308,7 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: report.baseline.geography.country, baseline_calculation.uk_local_authority_metadata, reform_calculation.uk_local_authority_metadata, + region_code=report.baseline.geography.region, ), ) v2_result = v2_report.get("result") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index a2a08462c..ee92f8fbf 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -21,16 +21,25 @@ def validate_uk_local_authority_metadata( country: CountryId, baseline: UKLocalAuthorityMetadata | None, reform: UKLocalAuthorityMetadata | None, + *, + region_code: str | None = None, ) -> UKLocalAuthorityMetadata | None: """Require matching UK authority metadata and reject it elsewhere. - Both UK artifacts lack it for a national dataset without area codes, which - detection admits only when the bundle routes constituency and - local-authority regions to another dataset. + Retained artifacts are read without rerunning dataset detection, so a UK + pair without metadata passes only with the same evidence detection needs: + the report's region is national or nation-level and the bundle routes + both constituency and local-authority regions to another dataset. """ if country == "uk": if baseline is None and reform is None: + from policyengine_simulation_executor.uk_local_authority_metadata import ( + uk_area_codes_required, + ) + + if region_code is None or uk_area_codes_required(region_code): + raise ValueError("UK simulation artifact metadata is missing") return None if baseline is None or reform is None: raise ValueError("UK simulation artifact metadata is missing on one side") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index e6d285916..63fca98e2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -420,6 +420,7 @@ def coordinate_report( report.baseline.geography.country, deserialize_uk_local_authority_metadata(baseline_payload), deserialize_uk_local_authority_metadata(reform_payload), + region_code=report.baseline.geography.region, ) validate_output_frames(baseline_frames, output_plan) validate_output_frames(reform_frames, output_plan) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 2989a6eee..cc483437f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -130,7 +130,7 @@ def calculate_simulation_frames( _country_module, _load_dataset, _normalise_policy, - _require_uk_weight_matrix_matches_dataset, + _require_certified_uk_weight_matrix, _resolve_dataset_selection, _resolve_region, setup_gcp_credentials, @@ -186,7 +186,9 @@ def calculate_simulation_frames( dataset, region_code=region.code, ) - _require_uk_weight_matrix_matches_dataset(region.scoping_strategy, dataset) + _require_certified_uk_weight_matrix( + region.scoping_strategy, dataset, dataset_selection + ) policy_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.POLICY_NORMALIZATION)) if runtime is not None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 29300f7c2..d3f41a15f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -564,8 +564,38 @@ def test_country_metadata_alignment_rejects_missing_or_mismatched_uk_values() -> assert validate_uk_local_authority_metadata("us", None, None) is None -def test_country_metadata_alignment_accepts_uk_artifacts_without_area_codes() -> None: - assert validate_uk_local_authority_metadata("uk", None, None) is None +@pytest.mark.parametrize( + ("region_code", "routed", "accepted"), + [ + ("uk", True, True), + ("country/england", True, True), + ("uk", False, False), + ("constituency/E14001063", True, False), + ("local_authority/E06000063", True, False), + (None, True, False), + ], +) +def test_uk_artifacts_without_metadata_need_a_routed_national_request( + monkeypatch, region_code: str | None, routed: bool, accepted: bool +) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + if accepted: + assert ( + validate_uk_local_authority_metadata( + "uk", None, None, region_code=region_code + ) + is None + ) + else: + with pytest.raises(ValueError, match="metadata is missing"): + validate_uk_local_authority_metadata( + "uk", None, None, region_code=region_code + ) def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py index 33969e5e9..c633f2930 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py @@ -1,7 +1,8 @@ -"""Unit tests for the UK weight-matrix fallback guard.""" +"""Unit tests for binding UK weight-matrix regions to the certified release.""" from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import h5py @@ -13,15 +14,73 @@ WeightReplacementStrategy, ) +from policyengine_simulation_executor import simulation_output_geographic from policyengine_simulation_executor import simulation_runtime as sr +_REPO = "policyengine/policyengine-uk-data-private" +_MATRIX = "parliamentary_constituency_weights.h5" _STRATEGY = WeightReplacementStrategy( weight_matrix_bucket="policyengine-uk-data-private", - weight_matrix_key="parliamentary_constituency_weights.h5", + weight_matrix_key=_MATRIX, lookup_csv_bucket="policyengine-uk-data-private", lookup_csv_key="constituencies_2024.csv", region_code="E14001063", + download_missing_assets=False, ) +_SELECTION = SimpleNamespace(name="enhanced_frs_2024_25") + + +def _reference(path: str, revision: str = "1.56.16") -> SimpleNamespace: + return SimpleNamespace(path=path, repo_id=_REPO, revision=revision) + + +def _bundle(monkeypatch, *, matrix_revision: str = "1.56.16") -> None: + manifest = SimpleNamespace( + data_package=SimpleNamespace( + repo_id=_REPO, version="1.56.16", release_manifest_revision=None + ), + datasets={ + "enhanced_frs_2024_25": _reference("enhanced_frs_2024_25.h5"), + "parliamentary_constituency_weights": _reference(_MATRIX, matrix_revision), + }, + ) + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", + lambda country: manifest, + ) + + +def _write_matrix(path: Path, households: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as matrix: + matrix.create_dataset("2025", data=np.ones((2, households))) + + +def _assets(monkeypatch, tmp_path, *, households: int = 3, certified_dir=None): + """Search UK geography assets in tmp_path and stub the certified download.""" + + monkeypatch.setenv("POLICYENGINE_UK_GEOGRAPHY_DATA_DIR", str(tmp_path)) + calls: list = [] + + def materialize(country, dataset, *, data_dir, **kwargs): + target = Path(certified_dir or data_dir) / _MATRIX + _write_matrix(target, households) + calls.append((country, dataset, Path(data_dir))) + return SimpleNamespace(path=str(target)) + + def lookup(spec): + path = tmp_path / spec.lookup_csv_filename + path.write_text("code,name,x,y\nE14001063,Aldershot,56,-40\n") + return str(path) + + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.materialize_dataset", + materialize, + ) + monkeypatch.setattr( + simulation_output_geographic, "_required_uk_geography_lookup_csv_path", lookup + ) + return calls def _dataset(households: int, year: int = 2025) -> SimpleNamespace: @@ -32,19 +91,11 @@ def _dataset(households: int, year: int = 2025) -> SimpleNamespace: ) -def _matrix(tmp_path, households: int) -> str: - path = tmp_path / "parliamentary_constituency_weights.h5" - with h5py.File(path, "w") as matrix: - matrix.create_dataset("2025", data=np.ones((2, households))) - return str(path) - - -def _resolver(path: str, calls: list): - def resolve(spec, **kwargs): - calls.append((spec, kwargs)) - return SimpleNamespace(weight_matrix_path=path, lookup_csv_path="unused") +def test_fallback_regions_never_download_the_bucket_matrix() -> None: + region = sr._build_uk_weight_replacement_region("constituency/E14001063") - return resolve + assert region is not None + assert region.scoping_strategy.download_missing_assets is False @pytest.mark.parametrize( @@ -53,53 +104,73 @@ def resolve(spec, **kwargs): ) def test_guard_ignores_regions_without_weight_matrices(monkeypatch, strategy) -> None: monkeypatch.setattr( - "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", - lambda *args, **kwargs: pytest.fail("no weight matrix should be resolved"), + "policyengine.provenance.manifest.get_release_manifest", + lambda country: pytest.fail("no weight matrix should be resolved"), ) - sr._require_uk_weight_matrix_matches_dataset(strategy, _dataset(3)) + sr._require_certified_uk_weight_matrix(strategy, _dataset(3), _SELECTION) -def test_guard_accepts_the_dataset_the_matrix_was_built_for( +def test_guard_reads_the_matrix_certified_with_the_selected_dataset( monkeypatch, tmp_path ) -> None: - calls: list = [] - monkeypatch.setattr( - "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", - _resolver(_matrix(tmp_path, households=3), calls), - ) + _bundle(monkeypatch) + calls = _assets(monkeypatch, tmp_path) - sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(3)) + sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3), _SELECTION) - [(spec, kwargs)] = calls - assert spec.weight_matrix_filename == _STRATEGY.weight_matrix_key - assert spec.lookup_csv_filename == _STRATEGY.lookup_csv_key - assert spec.resolved_weight_matrix_bucket == _STRATEGY.weight_matrix_bucket - assert kwargs == {"download_missing_assets": True} + assert calls == [("uk", "parliamentary_constituency_weights", tmp_path)] -def test_guard_rejects_a_dataset_of_another_size(monkeypatch, tmp_path) -> None: - monkeypatch.setattr( - "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", - _resolver(_matrix(tmp_path, households=3), []), - ) +def test_guard_rejects_a_matrix_from_another_release(monkeypatch, tmp_path) -> None: + _bundle(monkeypatch, matrix_revision="1.57.4") + _assets(monkeypatch, tmp_path) with pytest.raises( ValueError, match=( - r"'E14001063' reweights households with " - r"parliamentary_constituency_weights\.h5, which was built for 3 " - r"households; the selected UK dataset has 4" + r"has no parliamentary_constituency_weights\.h5 from the release of " + r"'enhanced_frs_2024_25'" ), ): - sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(4)) + sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3), _SELECTION) + + +def test_guard_rejects_a_dataset_outside_the_bundle(monkeypatch, tmp_path) -> None: + _bundle(monkeypatch) + _assets(monkeypatch, tmp_path) + + with pytest.raises(ValueError, match="is not in the certified bundle"): + sr._require_certified_uk_weight_matrix( + _STRATEGY, _dataset(3), SimpleNamespace(name="populace_uk_2023") + ) + + +def test_guard_rejects_a_copy_that_shadows_the_certified_matrix( + monkeypatch, tmp_path +) -> None: + _bundle(monkeypatch) + _assets(monkeypatch, tmp_path, certified_dir=tmp_path / "certified") + _write_matrix(tmp_path / _MATRIX, households=3) + + with pytest.raises(ValueError, match="would read .*, not the parliamentary"): + sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3), _SELECTION) + + +def test_guard_rejects_a_dataset_of_another_size(monkeypatch, tmp_path) -> None: + _bundle(monkeypatch) + _assets(monkeypatch, tmp_path, households=3) + + with pytest.raises( + ValueError, + match=(r"which was built for 3 households; the selected UK dataset has 4"), + ): + sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(4), _SELECTION) def test_guard_rejects_a_year_the_matrix_does_not_cover(monkeypatch, tmp_path) -> None: - monkeypatch.setattr( - "policyengine.data.uk_geography_assets.resolve_uk_geography_asset_paths", - _resolver(_matrix(tmp_path, households=3), []), - ) + _bundle(monkeypatch) + _assets(monkeypatch, tmp_path) with pytest.raises(ValueError, match=r"has no weights for 2026 \(it covers 2025\)"): - sr._require_uk_weight_matrix_matches_dataset(_STRATEGY, _dataset(3, 2026)) + sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3, 2026), _SELECTION)