From 5ad72b0fcacfe4ca3c38bc9cde211b250cafffeb Mon Sep 17 00:00:00 2001 From: Maria Juaristi Date: Wed, 21 Oct 2026 18:18:49 +0100 Subject: [PATCH 01/31] Build routed UK area results from the local-area dataset When the policyengine.py bundle routes UK constituency and local-authority regions to a local-area dataset and the national default carries no area codes (policyengine.py 6.10.0), a national or nation-level run builds its two area breakdowns from a second baseline and reform run: the same policies and nation filter, on the routed dataset. That dataset must carry constituency_code_oa and la_code_oa, and its authority codes must match one packaged boundary version. The national figures still come from the run's own pair, and the builder computes the second pair when the first breakdown needs it. Once routed, constituency and local-authority regions resolve by name or by GSS code in any case to the row-filter regions in policyengine.py's registry, so policyengine-api's constituency/{name} and local_authority/{name} requests keep working, and no weight-matrix region is built for them. The 17 districts replaced in April 2023 are unsupported. Nothing changes on policyengine.py 6.7.5, which routes no area regions. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_builder.py | 53 ++- .../simulation_runtime.py | 120 +++++- .../uk_local_authority_metadata.py | 103 ++++- .../tests/test_simulation_output_builder.py | 377 +++++++++++++++++- .../tests/test_uk_local_authority_metadata.py | 153 +++++++ 5 files changed, 792 insertions(+), 14 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 5a86a0482..a04e2c3c3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -48,12 +48,31 @@ from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata +@dataclass(frozen=True) +class UKAreaSimulations: + """A baseline and reform on the dataset the bundle routes UK area regions to. + + A UK national or nation-level run on a dataset without area codes builds + its constituency and local-authority breakdowns from this pair, which has + the same policies and nation filter, and from the authority boundaries + detected on that dataset. + """ + + baseline: Any + reform: Any + local_authority_metadata: UKLocalAuthorityMetadata + + @dataclass class SimulationOutputBuilder: """Assemble one year's macro output from a computed baseline and reform. Each output segment is built inside its own observability span when a - ``runtime`` is given. + ``runtime`` is given. The UK constituency and local-authority breakdowns + come from ``uk_area_simulations`` when it is given, and otherwise from + ``baseline`` and ``reform``. ``omit_uk_area_breakdowns`` leaves them out: + only Stage 12 sets it, for a routed national report whose stored frames + carry no area codes. """ country: str @@ -66,6 +85,8 @@ class SimulationOutputBuilder: resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + uk_area_simulations: UKAreaSimulations | None = None + omit_uk_area_breakdowns: bool = False progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: Any = field(default=None, init=False) @@ -257,19 +278,41 @@ def _build_congressional_district_impact( region_code=self.resolved_region_code, ) + def _uk_area_breakdown_inputs( + self, + ) -> tuple[Any, Any, UKLocalAuthorityMetadata | None]: + """Return the baseline, reform and boundaries the UK area breakdowns use. + + The impact analysis computes the run's own pair before any breakdown; + a separate area pair is computed here, on first use (``ensure()`` + returns a computed simulation from policyengine.py's cache). + """ + area = self.uk_area_simulations + if area is None: + return self.baseline, self.reform, self.uk_local_authority_metadata + area.baseline.ensure() + area.reform.ensure() + return area.baseline, area.reform, area.local_authority_metadata + def _build_uk_constituency_impact(self) -> GeographicImpactOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY)): + if self.omit_uk_area_breakdowns: + return None + baseline, reform, _ = self._uk_area_breakdown_inputs() return simulation_output_geographic.build_uk_constituency_impact( - self.country, self.baseline, self.reform + self.country, baseline, reform ) def _build_uk_local_authority_impact(self) -> GeographicImpactOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY)): + if self.omit_uk_area_breakdowns: + return None + baseline, reform, metadata = self._uk_area_breakdown_inputs() return simulation_output_geographic.build_uk_local_authority_impact( self.country, - self.baseline, - self.reform, - uk_local_authority_metadata=self.uk_local_authority_metadata, + baseline, + reform, + uk_local_authority_metadata=metadata, ) def _model_version(self) -> str: 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 39a30ae20..83cede408 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 @@ -39,6 +39,7 @@ ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, + UKAreaSimulations, ) if TYPE_CHECKING: @@ -353,6 +354,40 @@ def _normalise_region_code(country: str, region: object) -> str: return f"{prefix}/{value}" +def _uk_area_regions_routed() -> bool: + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + uk_area_regions_routed, + ) + + return uk_area_regions_routed() + + +def _registry_uk_area_region_code(country_module: ModuleType, region_code: str) -> str: + """Spell a UK constituency or local authority as the region registry keys it. + + When the bundle routes UK area regions, policyengine.py's registry holds + every constituency and local authority under its GSS code, labelled with + its name. A request may name the area instead, as policyengine-api's + ``constituency/{name}`` does, or write the code in lower case; both map to + the registry's code. Any other value is returned unchanged, so an area + outside the registry, such as one of the 17 districts replaced in April + 2023, stays unsupported. + """ + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + UK_AREA_REGION_TYPES, + ) + + region_type, _, value = region_code.partition("/") + if region_type not in UK_AREA_REGION_TYPES or not _uk_area_regions_routed(): + return region_code + requested = value.casefold() + for region in country_module.model.region_registry.get_by_type(region_type): + code = region.code.partition("/")[2] + if requested in {code.casefold(), region.label.casefold()}: + return region.code + return region_code + + def _build_uk_weight_replacement_region(region_code: str) -> Region | None: if "/" not in region_code: return None @@ -576,8 +611,10 @@ def _resolve_region( dataset_reference=_resolve_dataset_reference(country), ) + if country == "uk": + region_code = _registry_uk_area_region_code(country_module, region_code) region = country_module.model.get_region(region_code) - if region is None and country == "uk": + if region is None and country == "uk" and not _uk_area_regions_routed(): region = _build_uk_weight_replacement_region(region_code) if region is None: msg = f"Unsupported {country.upper()} region: {region_code}" @@ -813,6 +850,72 @@ def _report_calculation_start( simulation.on_first_ensure(lambda: progress.enter(stage)) +def _uk_area_simulations( + params: dict[str, Any], + *, + runtime: ObservabilityRuntime, + country_module: ModuleType, + region: RegionResolution, + policies: tuple[dict[str, Any] | None, dict[str, Any] | None], +) -> UKAreaSimulations: + """Build the UK area breakdowns' baseline and reform on the routed dataset. + + A UK national or nation-level run may load a dataset without area codes + only when the bundle routes constituency and local-authority regions to + another dataset. Its two area breakdowns then come from this pair: the + run's baseline and reform policies and its nation filter, applied to the + routed dataset, whose area codes are checked as it loads. The national + figures still come from the run's own pair. + """ + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + detect_uk_area_dataset_metadata, + routed_uk_area_dataset_name, + ) + + area_role = {"dataset_role": "uk_area"} + with runtime.span( + ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION), attributes=area_role + ): + selection = _resolve_dataset_selection( + params, + region_resolution=RegionResolution( + code=region.code, dataset_reference=routed_uk_area_dataset_name() + ), + ) + with runtime.span( + ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD), attributes=area_role + ): + dataset = _load_dataset( + params, selection=selection, country_module=country_module + ) + local_authority_metadata = detect_uk_area_dataset_metadata( + dataset, dataset_name=selection.name + ) + simulations = [] + for kind, policy in zip( + ("uk_area_baseline", "uk_area_reform"), policies, strict=True + ): + with runtime.span( + ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), + attributes={"simulation_kind": kind}, + ): + simulations.append( + _build_simulation( + params, + dataset=dataset, + dataset_selection=selection, + policy=policy, + region=region, + ) + ) + baseline, reform = simulations + return UKAreaSimulations( + baseline=baseline, + reform=reform, + local_authority_metadata=local_authority_metadata, + ) + + def _run_simulation_impl_core( params: dict, *, @@ -908,6 +1011,20 @@ def _run_simulation_impl_core( _report_calculation_start(baseline, progress, JobStage.BASELINE) _report_calculation_start(reform, progress, JobStage.REFORM) + # Detection returns None for a UK dataset only when it carries no area + # codes and the bundle routes area regions to another dataset. + uk_area_simulations = ( + _uk_area_simulations( + simulation_params, + runtime=runtime, + country_module=country_module, + region=region_resolution, + policies=(baseline_policy, reform_policy), + ) + if country == "uk" and uk_local_authority_metadata is None + else None + ) + logger.info("Calculating economic impact") builder = SimulationOutputBuilder( country=country, @@ -920,6 +1037,7 @@ def _run_simulation_impl_core( resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, + uk_area_simulations=uk_area_simulations, progress=progress, ) output = builder.serialize() 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 0ce031c54..caa46006e 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 @@ -9,6 +9,7 @@ from dataclasses import dataclass from functools import lru_cache from importlib.resources import files +from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING @@ -281,6 +282,7 @@ def detect_uk_local_authority_boundary_version( UK_AREA_REGION_TYPES = ("constituency", "local_authority") +UK_AREA_CODE_COLUMNS = ("constituency_code_oa", "la_code_oa") _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 " @@ -304,6 +306,41 @@ def uk_area_regions_routed() -> bool: return all(region_type in region_datasets for region_type in UK_AREA_REGION_TYPES) +def routed_uk_area_dataset_name() -> str: + """Name the certified dataset the bundle routes UK area regions to. + + Like policyengine.py's ``region_dataset_name``, this is the bundle dataset + whose path the region type's template names. Both area region types must + share it: a national run builds both of its area breakdowns from one + baseline and reform pair on that dataset. + """ + from policyengine.provenance.manifest import ( + get_release_manifest, + ) + + manifest = get_release_manifest("uk") + names: set[str] = set() + for region_type in UK_AREA_REGION_TYPES: + template = manifest.region_datasets.get(region_type) + routed = [ + name + for name, reference in manifest.datasets.items() + if template is not None and reference.path == template.path_template + ] + if not routed: + msg = f"The UK bundle routes {region_type} regions to no certified dataset" + raise ValueError(msg) + names.add(routed[0]) + if len(names) > 1: + msg = ( + "The UK bundle routes constituency and local-authority regions to " + f"different datasets ({', '.join(sorted(names))}); a national run " + "builds both area breakdowns from one dataset" + ) + raise ValueError(msg) + return names.pop() + + def uk_area_codes_required(region_code: str | None) -> bool: """Whether a UK run's dataset must carry area codes. @@ -331,6 +368,40 @@ def detect_uk_local_authority_metadata( """ if country != "uk": return None + household_frame = _uk_household_frame(dataset) + 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 detect_uk_area_dataset_metadata( + dataset: object, *, dataset_name: str +) -> UKLocalAuthorityMetadata: + """Check a loaded dataset that UK area regions are routed to. + + That dataset serves constituency and local-authority results, so unlike a + national default it must carry both area code columns, and its authority + codes must match one packaged boundary version. + """ + household_frame = _uk_household_frame(dataset) + for column in UK_AREA_CODE_COLUMNS: + if column not in household_frame: + msg = ( + f"UK dataset {dataset_name!r}, which the bundle routes " + f"constituency and local-authority regions to, has no {column} " + "column" + ) + raise ValueError(msg) + return detect_uk_local_authority_boundary_version( + household_frame["la_code_oa"].tolist() + ) + + +def _uk_household_frame(dataset: object) -> pd.DataFrame: data = getattr(dataset, "data", None) entity_data = getattr(data, "entity_data", None) if not isinstance(entity_data, Mapping): @@ -340,14 +411,32 @@ def detect_uk_local_authority_metadata( if household is None: msg = "UK dataset contains no household table" raise ValueError(msg) - 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() + return pd.DataFrame(household) + + +def uk_area_hdf_dataset_path(installed_default_path: str) -> str: + """Return the UK HDF file whose area codes the worker's results use. + + That is the installed default, unless it carries no area codes and the + bundle routes constituency and local-authority regions to another + dataset. Then the routed dataset is downloaded next to the default and + checked against its bundle digest, and its path is returned. The download + needs the Hugging Face token, as the data is private. + """ + if not uk_area_regions_routed() or uk_hdf_household_has_local_authority_codes( + installed_default_path + ): + return installed_default_path + from policyengine.provenance.dataset_materialization import ( + materialize_dataset, + ) + + routed = materialize_dataset( + "uk", + routed_uk_area_dataset_name(), + data_dir=Path(installed_default_path).parent, ) + return routed.path def uk_hdf_household_has_local_authority_codes(dataset_path: str) -> bool: 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 9d95eab0b..cab8a5152 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -29,6 +29,10 @@ fake_analysis, ) from fixtures.weighted_frames import weighted_entity_frame +from policyengine.core.scoping_strategy import ( + RowFilterStrategy, + WeightReplacementStrategy, +) from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, @@ -60,6 +64,7 @@ ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, + UKAreaSimulations, ) from policyengine_simulation_executor.simulation_output_geographic import ( build_congressional_district_impact, @@ -79,7 +84,7 @@ from policyengine_simulation_executor.simulation_runtime import ( _load_dataset as _load_selected_dataset, ) -from policyengine_simulation_observability.stages import Stage +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage CONSTITUENCY_CSV_PATH = "/tmp/constituencies_2024.csv" @@ -216,6 +221,8 @@ class BuilderOptions: include_cliffs: bool | None = None resolved_region_code: str | None = "us" uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + uk_area_simulations: UKAreaSimulations | None = None + omit_uk_area_breakdowns: bool = False def _simulation_output_builder( @@ -255,6 +262,8 @@ def economic_impact_analysis( reform=reform, resolved_region_code=options.resolved_region_code, uk_local_authority_metadata=options.uk_local_authority_metadata, + uk_area_simulations=options.uk_area_simulations, + omit_uk_area_breakdowns=options.omit_uk_area_breakdowns, ) @@ -1149,6 +1158,7 @@ def serialize(self): "resolved_region_code": "us", "runtime": runtime, "uk_local_authority_metadata": None, + "uk_area_simulations": None, "progress": NULL_JOB_PROGRESS, } ] @@ -1284,6 +1294,281 @@ def serialize(self): assert metadata == UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 ) + assert builder_calls[0]["uk_area_simulations"] is None + + +def _routes_uk_area_regions(monkeypatch, *, routed: bool) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + +def _uk_dataset(**household: list[object]) -> SimpleNamespace: + return SimpleNamespace( + data=SimpleNamespace( + entity_data={"household": weighted_entity_frame(household)} + ) + ) + + +_NATIONAL_SELECTION = DatasetSelection( + "microcosm_uk_2024_25", "hf://policyengine/uk/national.h5@cut", is_default=True +) +_LOCAL_AREA_SELECTION = DatasetSelection( + "microcosm_uk_2024_25_dense", "hf://policyengine/uk/local.h5@cut", is_default=False +) + + +@dataclass +class _UKAreaRun: + """A UK run through ``_run_simulation_impl_core`` with every seam faked.""" + + builds: list[dict] + builder_calls: list[dict] + runtime: _TrackingRuntime + + +def _run_routed_uk_national(monkeypatch, *, local_area_dataset) -> _UKAreaRun: + national_dataset = _uk_dataset(region=["WALES", "WALES"]) + nation = RowFilterStrategy(variable_name="region", variable_value="WALES") + builds: list[dict] = [] + builder_calls: list[dict] = [] + _routes_uk_area_regions(monkeypatch, routed=True) + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "routed_uk_area_dataset_name", + lambda: _LOCAL_AREA_SELECTION.name, + ) + + def resolve_selection(_params, *, region_resolution): + if region_resolution.dataset_reference == _LOCAL_AREA_SELECTION.name: + return _LOCAL_AREA_SELECTION + return _NATIONAL_SELECTION + + def build_simulation(_params, *, dataset, dataset_selection, policy, region): + simulation = SimpleNamespace(on_first_ensure=lambda _callback: None) + builds.append( + { + "dataset": dataset, + "selection": dataset_selection, + "policy": policy, + "scoping_strategy": region.scoping_strategy, + "simulation": simulation, + } + ) + return simulation + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + for name, value in { + "_country_module": lambda _country: SimpleNamespace( + model=SimpleNamespace(version="2.132.6") + ), + "_resolve_region": lambda **_kwargs: RegionResolution( + code="country/wales", dataset_reference=None, scoping_strategy=nation + ), + "_resolve_dataset_selection": resolve_selection, + "_load_dataset": lambda _params, *, selection, country_module: ( + local_area_dataset + if selection is _LOCAL_AREA_SELECTION + else national_dataset + ), + "_build_simulation": build_simulation, + "SimulationOutputBuilder": FakeSimulationOutputBuilder, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_runtime.{name}", value + ) + + runtime = _TrackingRuntime() + _run_simulation_impl_core( + { + "country": "uk", + "region": "country/wales", + "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, + }, + runtime=runtime, + ) + return _UKAreaRun(builds=builds, builder_calls=builder_calls, runtime=runtime) + + +def test_routed_uk_national_run_builds_area_breakdowns_on_the_local_area_dataset( + monkeypatch, +): + local_area_dataset = _uk_dataset( + region=["WALES", "WALES"], + constituency_code_oa=["E14001063", "W07000112"], + la_code_oa=["E06000063", "W06000001"], + ) + + run = _run_routed_uk_national(monkeypatch, local_area_dataset=local_area_dataset) + + national, area = run.builds[:2], run.builds[2:] + assert [build["selection"] for build in national] == [_NATIONAL_SELECTION] * 2 + assert [build["selection"] for build in area] == [_LOCAL_AREA_SELECTION] * 2 + assert all(build["dataset"] is local_area_dataset for build in area) + # The same policies and the same nation filter as the national pair. + assert [build["policy"] for build in area] == [ + build["policy"] for build in national + ] + assert all( + build["scoping_strategy"] is national[0]["scoping_strategy"] + for build in run.builds + ) + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["baseline"] is national[0]["simulation"] + assert builder_kwargs["reform"] is national[1]["simulation"] + assert builder_kwargs["uk_local_authority_metadata"] is None + assert builder_kwargs["uk_area_simulations"] == UKAreaSimulations( + baseline=area[0]["simulation"], + reform=area[1]["simulation"], + local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), + ) + loads = [ + node["attrs"] + for node in run.runtime.roots + if node["name"] == ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD) + ] + assert loads == [{}, {"dataset_role": "uk_area"}] + + +def test_routed_uk_national_run_requires_area_codes_on_the_local_area_dataset( + monkeypatch, +): + without_authorities = _uk_dataset( + region=["WALES"], constituency_code_oa=["W07000112"] + ) + + with pytest.raises( + ValueError, + match=r"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " + r"la_code_oa column", + ): + _run_routed_uk_national(monkeypatch, local_area_dataset=without_authorities) + + +_LOCAL_AREA_URI = ( + "hf://policyengine/populace-uk-private/microcosm_uk_2024_25_dense.h5" + "@microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" +) + + +def _registry_area(region_type: str, code: str, label: str) -> SimpleNamespace: + return SimpleNamespace( + code=f"{region_type}/{code}", + label=label, + region_type=region_type, + parent_code="uk", + dataset_path=_LOCAL_AREA_URI, + scoping_strategy=object(), + ) + + +def _routed_uk_country_module() -> SimpleNamespace: + """Return a UK model whose registry lists area regions, as routing gives.""" + areas = [ + _registry_area("constituency", "E14001063", "Aldershot"), + _registry_area("constituency", "W07000112", "Ynys Môn"), + _registry_area("local_authority", "E07000110", "Maidstone"), + _registry_area("local_authority", "N09000011", "Ards and North Down"), + ] + by_code = {area.code: area for area in areas} + registry = SimpleNamespace( + get_by_type=lambda region_type: [ + area for area in areas if area.region_type == region_type + ] + ) + return SimpleNamespace( + model=SimpleNamespace(get_region=by_code.get, region_registry=registry) + ) + + +def _forbid_weight_matrix_regions(monkeypatch) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime." + "_build_uk_weight_replacement_region", + lambda _code: pytest.fail("a routed area never uses a weight matrix"), + ) + + +@pytest.mark.parametrize( + ("requested", "expected"), + [ + ("constituency/E14001063", "constituency/E14001063"), + ("constituency/e14001063", "constituency/E14001063"), + ("constituency/Aldershot", "constituency/E14001063"), + ("constituency/ynys môn", "constituency/W07000112"), + ("local_authority/Maidstone", "local_authority/E07000110"), + ("local_authority/n09000011", "local_authority/N09000011"), + ("local_authority/Ards and North Down", "local_authority/N09000011"), + ], +) +def test_resolve_region_maps_routed_uk_area_names_and_codes_to_the_registry( + monkeypatch, requested, expected +): + _routes_uk_area_regions(monkeypatch, routed=True) + _forbid_weight_matrix_regions(monkeypatch) + country_module = _routed_uk_country_module() + + resolution = _resolve_region( + country_module=country_module, country="uk", params={"region": requested} + ) + + region = country_module.model.get_region(expected) + assert resolution.code == expected + assert resolution.scoping_strategy is region.scoping_strategy + assert resolution.dataset_reference == _LOCAL_AREA_URI + + +@pytest.mark.parametrize( + "requested", + [ + # A district that April 2023 replaced: LAD22 only. + "local_authority/E07000026", + "local_authority/Carlisle", + "constituency/Maidstone", + ], +) +def test_resolve_region_rejects_uk_areas_outside_the_routed_registry( + monkeypatch, requested +): + _routes_uk_area_regions(monkeypatch, routed=True) + _forbid_weight_matrix_regions(monkeypatch) + + with pytest.raises(ValueError, match="Unsupported UK region"): + _resolve_region( + country_module=_routed_uk_country_module(), + country="uk", + params={"region": requested}, + ) + + +def test_resolve_region_keeps_weight_matrix_uk_areas_until_areas_are_routed( + monkeypatch, +): + _routes_uk_area_regions(monkeypatch, routed=False) + # No region registry: an unrouted bundle must not read one. + country_module = SimpleNamespace(model=SimpleNamespace(get_region=lambda _: None)) + + resolution = _resolve_region( + country_module=country_module, + country="uk", + params={"region": "constituency/Aldershot"}, + ) + + assert resolution.code == "constituency/Aldershot" + assert isinstance(resolution.scoping_strategy, WeightReplacementStrategy) + assert resolution.scoping_strategy.region_code == "Aldershot" def test_resolve_dataset_reference_uses_bundle_default(): @@ -1863,3 +2148,93 @@ def build(country, baseline_simulation, reform_simulation, **kwargs): )._build_uk_local_authority_impact() is None ) + + +class _AreaSimulation: + """A routed-area simulation that counts how often it is computed.""" + + def __init__(self): + self.ensured = 0 + + def ensure(self): + self.ensured += 1 + + +def test_uk_area_breakdowns_use_the_computed_routed_area_pair(monkeypatch): + area = UKAreaSimulations( + baseline=_AreaSimulation(), + reform=_AreaSimulation(), + local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), + ) + calls: list[tuple[str, object, object, object]] = [] + + def constituencies(country, baseline_simulation, reform_simulation): + assert country == "uk" + assert baseline_simulation.ensured + assert reform_simulation.ensured + calls.append(("constituency", baseline_simulation, reform_simulation, None)) + return GeographicImpactOutput.model_validate( + [_constituency_impact_record().model_dump(mode="json")] + ) + + def local_authorities(country, baseline_simulation, reform_simulation, **kwargs): + assert country == "uk" + assert baseline_simulation.ensured + assert reform_simulation.ensured + calls.append( + ( + "local_authority", + baseline_simulation, + reform_simulation, + kwargs["uk_local_authority_metadata"], + ) + ) + return GeographicImpactOutput.model_validate( + [_local_authority_impact_record().model_dump(mode="json")] + ) + + for name, function in { + "build_uk_constituency_impact": constituencies, + "build_uk_local_authority_impact": local_authorities, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_output_geographic.{name}", + function, + ) + builder = _simulation_output_builder( + "uk", + object(), + object(), + BuilderOptions(resolved_region_code="uk", uk_area_simulations=area), + ) + + assert builder._build_uk_constituency_impact() is not None + assert builder._build_uk_local_authority_impact() is not None + assert calls == [ + ("constituency", area.baseline, area.reform, None), + ( + "local_authority", + area.baseline, + area.reform, + area.local_authority_metadata, + ), + ] + + +def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): + for name in ("build_uk_constituency_impact", "build_uk_local_authority_impact"): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_output_geographic.{name}", + lambda *_args, **_kwargs: pytest.fail("no area breakdown is built"), + ) + builder = _simulation_output_builder( + "uk", + object(), + object(), + BuilderOptions(resolved_region_code="uk", omit_uk_area_breakdowns=True), + ) + + assert builder._build_uk_constituency_impact() is None + assert builder._build_uk_local_authority_impact() is None 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 5bbf0cad8..c9475751a 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 @@ -10,10 +10,13 @@ UKLocalAuthorityBoundaryVersion, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( + detect_uk_area_dataset_metadata, detect_uk_local_authority_boundary_version, detect_uk_local_authority_metadata, detect_uk_local_authority_metadata_from_hdf, load_uk_local_authority_resources, + routed_uk_area_dataset_name, + uk_area_hdf_dataset_path, uk_area_regions_routed, uk_hdf_household_has_local_authority_codes, ) @@ -207,6 +210,156 @@ def test_area_regions_count_as_routed_only_when_both_have_a_dataset( assert uk_area_regions_routed() is routed +_NATIONAL_FILE = "microcosm_uk_2024_25.h5" +_LOCAL_AREA_FILE = "microcosm_uk_2024_25_dense.h5" + + +def _bundle_routes(monkeypatch, region_datasets: dict[str, str]) -> None: + """Install a UK manifest routing each region type to a bundle file.""" + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", + lambda _country: SimpleNamespace( + region_datasets={ + region_type: SimpleNamespace(path_template=path) + for region_type, path in region_datasets.items() + }, + datasets={ + "microcosm_uk_2024_25": SimpleNamespace(path=_NATIONAL_FILE), + "microcosm_uk_2024_25_dense": SimpleNamespace(path=_LOCAL_AREA_FILE), + }, + ), + ) + + +def _routes_both_area_types_to(monkeypatch, path: str) -> None: + _bundle_routes( + monkeypatch, + { + "national": _NATIONAL_FILE, + "constituency": path, + "local_authority": path, + }, + ) + + +def test_routed_area_dataset_is_the_one_both_area_types_name(monkeypatch) -> None: + _routes_both_area_types_to(monkeypatch, _LOCAL_AREA_FILE) + + assert routed_uk_area_dataset_name() == "microcosm_uk_2024_25_dense" + + +def test_routed_area_dataset_must_serve_both_area_types(monkeypatch) -> None: + _bundle_routes( + monkeypatch, + {"constituency": _LOCAL_AREA_FILE, "local_authority": _NATIONAL_FILE}, + ) + + with pytest.raises(ValueError, match="to different datasets"): + routed_uk_area_dataset_name() + + +@pytest.mark.parametrize( + "region_datasets", + [ + {"constituency": _LOCAL_AREA_FILE}, + {"constituency": _LOCAL_AREA_FILE, "local_authority": "uncertified.h5"}, + ], +) +def test_routed_area_dataset_must_be_a_certified_bundle_dataset( + monkeypatch, region_datasets: dict[str, str] +) -> None: + _bundle_routes(monkeypatch, region_datasets) + + with pytest.raises( + ValueError, match="routes local_authority regions to no certified dataset" + ): + routed_uk_area_dataset_name() + + +def _area_dataset(**household: list[str]) -> SimpleNamespace: + return SimpleNamespace( + data=SimpleNamespace(entity_data={"household": pd.DataFrame(household)}) + ) + + +def test_routed_dataset_detection_identifies_the_authority_boundaries() -> None: + dataset = _area_dataset( + constituency_code_oa=["E14001063", "W07000112"], + la_code_oa=["E06000063", "W06000001"], + ) + + metadata = detect_uk_area_dataset_metadata(dataset, dataset_name="local") + + assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD23 + + +@pytest.mark.parametrize("missing", ["constituency_code_oa", "la_code_oa"]) +def test_routed_dataset_detection_requires_both_area_code_columns( + missing: str, +) -> None: + columns = { + "constituency_code_oa": ["E14001063"], + "la_code_oa": ["E06000063"], + } + del columns[missing] + + with pytest.raises( + ValueError, + match=f"'local', which the bundle routes .* has no {missing} column", + ): + detect_uk_area_dataset_metadata(_area_dataset(**columns), dataset_name="local") + + +def _hdf_household(path, **columns: list[object]) -> str: + pd.DataFrame(columns).to_hdf( + path, key="household", format="table", data_columns=True + ) + return str(path) + + +def test_area_hdf_dataset_is_the_installed_default_unless_routed_without_codes( + monkeypatch, tmp_path +) -> None: + national = _hdf_household( + tmp_path / _NATIONAL_FILE, household_id=[1], region=["WALES"] + ) + coded = _hdf_household( + tmp_path / "coded.h5", household_id=[1], la_code_oa=["E06000001"] + ) + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.materialize_dataset", + lambda *_args, **_kwargs: pytest.fail("no dataset should be downloaded"), + ) + + _bundle_routes(monkeypatch, {"national": _NATIONAL_FILE}) + assert uk_area_hdf_dataset_path(national) == national + + _routes_both_area_types_to(monkeypatch, _LOCAL_AREA_FILE) + assert uk_area_hdf_dataset_path(coded) == coded + + +def test_area_hdf_dataset_downloads_the_routed_dataset_beside_the_default( + monkeypatch, tmp_path +) -> None: + national = _hdf_household( + tmp_path / _NATIONAL_FILE, household_id=[1], region=["WALES"] + ) + _routes_both_area_types_to(monkeypatch, _LOCAL_AREA_FILE) + downloads: list[tuple[str, str, object]] = [] + + def materialize(country, dataset, *, data_dir, **_kwargs): + downloads.append((country, dataset, data_dir)) + return SimpleNamespace(path=str(data_dir / _LOCAL_AREA_FILE)) + + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.materialize_dataset", + materialize, + ) + + assert uk_area_hdf_dataset_path(national) == str(tmp_path / _LOCAL_AREA_FILE) + assert downloads == [("uk", "microcosm_uk_2024_25_dense", tmp_path)] + + def test_installed_hdf_detector_reads_local_authority_codes(tmp_path) -> None: dataset_path = tmp_path / "uk-dataset.h5" pd.DataFrame( From 6aad65205ea69fdbc6ac8a2a708280688a6cb6bf Mon Sep 17 00:00:00 2001 From: Maria Juaristi Date: Wed, 21 Oct 2026 18:18:49 +0100 Subject: [PATCH 02/31] Check the routed area dataset in the image smokes The executor and Stage 12 image smokes read the installed UK default's local-authority codes. When the default carries none and the bundle routes area regions to another dataset, they now download that dataset next to the default, check it against the bundle digest, and read its codes instead. Both smoke functions mount the Hugging Face credential for that download, and the Stage 12 smoke asks for 2 GiB instead of 512 MiB to read the routed file. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/smoke_app.py | 16 ++++-- .../modal/v2_static_runtime_files_smoke.py | 15 +++-- .../tests/test_smoke_app.py | 57 ++++++++++++++++--- .../tests/test_stage12_modal_app.py | 41 +++++++++++++ 4 files changed, 113 insertions(+), 16 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 1b655bdd5..3cb4f820e 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -25,8 +25,9 @@ ) from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata_from_hdf, + uk_area_hdf_dataset_path, ) -from src.modal.app import build_runtime_simulation_image +from src.modal.app import build_runtime_simulation_image, hf_secret from src.modal.static_runtime_files import add_static_runtime_files import modal @@ -96,16 +97,23 @@ def _summarize_bundle_status(status: dict) -> dict: def _validate_installed_uk_local_authority_dataset() -> str: - """Require the installed UK dataset to match one packaged boundary version.""" + """Require the UK area results' dataset to match one packaged boundary version. + + That is the installed default, unless it carries no area codes and the + bundle routes area regions to another dataset; then the routed dataset is + downloaded and checked instead (``uk_area_hdf_dataset_path``). + """ dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: msg = "installed certified UK dataset is unavailable" raise RuntimeError(msg) - metadata = detect_uk_local_authority_metadata_from_hdf(dataset_path) + metadata = detect_uk_local_authority_metadata_from_hdf( + uk_area_hdf_dataset_path(dataset_path) + ) return metadata.boundary_version.value -@app.function(image=smoke_image, timeout=600, memory=8192) +@app.function(image=smoke_image, timeout=600, memory=8192, secrets=[hf_secret]) def smoke_import_executor() -> dict: """Check the installed bundle, then import the workers' lazy request-time modules. diff --git a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py index 3c022295a..9a2c7efbd 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py @@ -10,8 +10,9 @@ from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata_from_hdf, load_uk_local_authority_resources, + uk_area_hdf_dataset_path, ) -from src.modal.v2_app import uk_worker_image +from src.modal.v2_app import hf_secret, uk_worker_image import modal @@ -19,20 +20,24 @@ def validate_installed_uk_local_authority_dataset() -> str: - """Return the local authority boundary version of the installed UK dataset. + """Return the local authority boundary version of the UK area results' dataset. - Raises if the image has no installed certified UK dataset. + That is the installed UK dataset, unless it carries no area codes and the + bundle routes area regions to another dataset; then the routed dataset is + downloaded and checked instead (``uk_area_hdf_dataset_path``). Raises if + the image has no installed certified UK dataset. """ dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: msg = "installed certified UK dataset is unavailable" raise RuntimeError(msg) return detect_uk_local_authority_metadata_from_hdf( - dataset_path + uk_area_hdf_dataset_path(dataset_path) ).boundary_version.value -@app.function(image=uk_worker_image, timeout=600, memory=512) +# The memory covers reading a routed dataset's authority codes. +@app.function(image=uk_worker_image, timeout=600, memory=2048, secrets=[hf_secret]) def smoke_stage12_static_runtime_files() -> dict[str, int | str]: """Count the UK local authority files inside the UK worker image.""" resources = load_uk_local_authority_resources() diff --git a/projects/policyengine-simulation-executor/tests/test_smoke_app.py b/projects/policyengine-simulation-executor/tests/test_smoke_app.py index 3e778e46e..aac34be58 100644 --- a/projects/policyengine-simulation-executor/tests/test_smoke_app.py +++ b/projects/policyengine-simulation-executor/tests/test_smoke_app.py @@ -91,27 +91,70 @@ def test_bundle_status_summary_rejects_mismatch_or_incomplete_install( summarize_bundle_status(status) -def test_uk_dataset_smoke_validates_installed_codes_against_packaged_resources( +def _uk_area_dataset_smoke( monkeypatch, smoke_module, -) -> None: + *, + installed: str, + area_dataset: str, + boundary_version: UKLocalAuthorityBoundaryVersion, +) -> list[str]: + """Fake the installed default and the dataset serving UK area results.""" observed: list[str] = [] - monkeypatch.setattr( smoke_module, "resolve_local_bundle_dataset_path", - lambda _country, _requested_data: "/installed/enhanced_frs_2024_25.h5", + lambda _country, _requested_data: installed, + ) + monkeypatch.setattr( + smoke_module, + "uk_area_hdf_dataset_path", + lambda path: area_dataset if path == installed else pytest.fail(path), ) monkeypatch.setattr( smoke_module, "detect_uk_local_authority_metadata_from_hdf", lambda path: ( observed.append(path) - or UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 - ) + or UKLocalAuthorityMetadata(boundary_version=boundary_version) ), ) + return observed + + +def test_uk_dataset_smoke_validates_installed_codes_against_packaged_resources( + monkeypatch, + smoke_module, +) -> None: + observed = _uk_area_dataset_smoke( + monkeypatch, + smoke_module, + installed="/installed/enhanced_frs_2024_25.h5", + area_dataset="/installed/enhanced_frs_2024_25.h5", + boundary_version=UKLocalAuthorityBoundaryVersion.LAD22, + ) assert smoke_module._validate_installed_uk_local_authority_dataset() == "lad22" assert observed == ["/installed/enhanced_frs_2024_25.h5"] + + +def test_uk_dataset_smoke_checks_the_routed_dataset_of_a_national_default( + monkeypatch, + smoke_module, +) -> None: + observed = _uk_area_dataset_smoke( + monkeypatch, + smoke_module, + installed="/installed/microcosm_uk_2024_25.h5", + area_dataset="/installed/microcosm_uk_2024_25_dense.h5", + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23, + ) + + assert smoke_module._validate_installed_uk_local_authority_dataset() == "lad23" + assert observed == ["/installed/microcosm_uk_2024_25_dense.h5"] + + +def test_executor_smoke_can_download_a_routed_private_dataset(smoke_module) -> None: + functions = dict(smoke_module.app.function_calls) + + assert functions["smoke_import_executor"]["secrets"] == [smoke_module.hf_secret] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index d7b47b0e6..a93e3ee96 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -183,6 +183,10 @@ def test_v2_static_runtime_file_smoke_uses_exact_uk_worker_image(monkeypatch) -> assert functions["smoke_stage12_static_runtime_files"]["image"] is ( module.uk_worker_image ) + # A routed area dataset is private: checking it needs the Hugging Face token. + assert functions["smoke_stage12_static_runtime_files"]["secrets"] == [ + module.hf_secret + ] monkeypatch.setattr( smoke, "validate_installed_uk_local_authority_dataset", @@ -196,6 +200,43 @@ def test_v2_static_runtime_file_smoke_uses_exact_uk_worker_image(monkeypatch) -> } +def test_v2_static_runtime_file_smoke_checks_the_routed_area_dataset( + monkeypatch, +) -> None: + from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, + ) + + _load(monkeypatch) + sys.modules.pop("src.modal.v2_static_runtime_files_smoke", None) + smoke = importlib.import_module("src.modal.v2_static_runtime_files_smoke") + installed = "/opt/policyengine/stage12-data/microcosm_uk_2024_25.h5" + routed = "/opt/policyengine/stage12-data/microcosm_uk_2024_25_dense.h5" + checked: list[str] = [] + monkeypatch.setattr( + smoke, "resolve_local_bundle_dataset_path", lambda _country, _data: installed + ) + monkeypatch.setattr( + smoke, + "uk_area_hdf_dataset_path", + lambda path: routed if path == installed else pytest.fail(path), + ) + monkeypatch.setattr( + smoke, + "detect_uk_local_authority_metadata_from_hdf", + lambda path: ( + checked.append(path) + or UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + ), + ) + + assert smoke.validate_installed_uk_local_authority_dataset() == "lad23" + assert checked == [routed] + + def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: module = _load(monkeypatch) From c2728bceb8c4047a8ec35c243c5d8da05ff3b5ab Mon Sep 17 00:00:00 2001 From: Maria Juaristi Date: Wed, 21 Oct 2026 18:18:49 +0100 Subject: [PATCH 03/31] Copy no area codes into Stage 12 for a routed UK national report A Stage 12 UK national or nation-level report whose bundle routes area regions to another dataset plans no constituency_code_oa or la_code_oa columns, since its national dataset may carry none, and its aggregate leaves the constituency and local-authority breakdowns out. Stage 12 has no second run on the routed dataset yet, so its shadow comparison reports the two breakdowns as missing for such reports. The planning tests that describe unrouted reports now pin routing off. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 7 +++ .../stage12_runtime/aggregation.py | 12 +++- .../stage12_runtime/output_planning.py | 23 ++++++-- .../tests/test_stage12_output_planning.py | 41 +++++++++++++- .../tests/test_stage12_runtime.py | 56 +++++++++++++++++++ 5 files changed, 132 insertions(+), 7 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index fed0d09e4..de7814761 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -217,6 +217,13 @@ the source dataset. UK geographic reports require the household columns calculated PolicyEngine variables, but they must be present in each materialized artifact; otherwise validation fails before geographic aggregation begins. +The exception is a UK national or nation-level report whose bundle routes +constituency and local-authority regions to another dataset. Its national +dataset may carry no area codes, so the plan copies none, and its aggregate +leaves the constituency and local-authority breakdowns out. The production path +builds them from a second baseline and reform run on the routed dataset; Stage +12 has no such run yet. + Aggregation uses the retained output tables without rerunning either model. The precomputed baseline and reform simulations retain their original policy metadata so PolicyEngine can correctly recognize conditional labor-supply 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 ca4f2d8f5..2df76ed61 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 @@ -172,7 +172,10 @@ def build_aggregate_report( """Build the aggregate report payload from retained baseline and reform frames. Runs the standard macro output builder over stand-in simulations that load - the stored frames, so no model is recomputed. + the stored frames, so no model is recomputed. A UK national or nation-level + report whose bundle routes area regions to another dataset has no area + codes in its frames (see ``output_planning``), so its constituency and + local-authority breakdowns are left out. """ from policyengine_simulation_executor.simulation_output_builder import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_output_builder.SimulationOutputBuilder SimulationOutputBuilder, @@ -181,6 +184,9 @@ def build_aggregate_report( _country_module, _runtime_policy, ) + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + uk_area_codes_required, + ) country = report.baseline.geography.country country_module = _country_module(country) @@ -232,6 +238,10 @@ def stand_in( resolved_data_version=report.baseline.bundle.dataset.artifact_revision, resolved_region_code=report.baseline.geography.region, uk_local_authority_metadata=uk_local_authority_metadata, + omit_uk_area_breakdowns=( + country == "uk" + and not uk_area_codes_required(report.baseline.geography.region) + ), ).serialize() output.update( build_spm_result( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py index cd0cbee0b..51016f5b7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py @@ -23,6 +23,10 @@ ) from policyengine_simulation_observability.stages import Stage +from policyengine_simulation_executor.uk_local_authority_metadata import ( + uk_area_codes_required, +) + if TYPE_CHECKING: import pandas as pd from policyengine_simulation_contract.stage12_bundle import CountryId @@ -105,12 +109,22 @@ def _include_cliff_impacts(report: ReportExecutionInput) -> bool: def _required_dataset_variables( *, country: CountryId, + region: str, aggregates: tuple[ReportAggregate, ...], ) -> dict[str, tuple[str, ...]]: - """Return input-dataset columns copied into the materialized output.""" - if country == "uk" and ReportAggregate.GEOGRAPHIC in aggregates: - return UK_GEOGRAPHIC_DATASET_VARIABLES - return {} + """Return input-dataset columns copied into the materialized output. + + A UK report's geographic aggregate reads the area codes of its own + simulations, except in a national or nation-level report whose bundle + routes UK area regions to another dataset. Its national dataset may then + carry no area codes, so none are planned, and the aggregate leaves the + area breakdowns out (``build_aggregate_report``). + """ + if country != "uk" or ReportAggregate.GEOGRAPHIC not in aggregates: + return {} + if not uk_area_codes_required(region): + return {} + return UK_GEOGRAPHIC_DATASET_VARIABLES def _configure_country_outputs( @@ -171,6 +185,7 @@ def resolve_report_output_plan( reform_variables = model.resolve_entity_variables(reform) dataset_variables = _required_dataset_variables( country=country, + region=report.baseline.geography.region, aggregates=requirements.aggregates, ) entities = [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py index c0e5a196d..3f7498516 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py @@ -63,6 +63,7 @@ def _simulation( *, policy: dict[str, JsonValue] | None = None, options: dict[str, JsonValue] | None = None, + region: str | None = None, ) -> SimulationExecutionInput: bundle = _bundle(country) return SimulationExecutionInput( @@ -82,7 +83,7 @@ def _simulation( ) ), year=2026, - geography=GeographySelection(country=country, region=country), + geography=GeographySelection(country=country, region=region or country), options=options or {}, bundle=bundle, ) @@ -94,6 +95,7 @@ def _report( reform: dict[str, JsonValue] | None = None, include_cliffs: bool = False, aggregates: tuple[ReportAggregate, ...] = tuple(ReportAggregate), + region: str | None = None, ) -> ReportExecutionInput: options: dict[str, JsonValue] = {"include_cliffs": True} if include_cliffs else {} return ReportExecutionInput( @@ -102,12 +104,14 @@ def _report( country, SimulationRole.BASELINE, options=options, + region=region, ), reform=_simulation( country, SimulationRole.REFORM, policy=reform, options=options, + region=region, ), requested_aggregates=aggregates, ) @@ -139,7 +143,16 @@ def test_us_plan_adds_budget_variables_to_the_country_defaults() -> None: ) -def test_uk_plan_uses_uk_defaults_without_us_budget_variables() -> None: +def _route_uk_area_regions(monkeypatch, *, routed: bool) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + +def test_uk_plan_uses_uk_defaults_without_us_budget_variables(monkeypatch) -> None: + _route_uk_area_regions(monkeypatch, routed=False) plan = resolve_report_output_plan(_report("uk")) assert "benunit" in {entity.entity for entity in plan.entities} @@ -152,6 +165,28 @@ def test_uk_plan_uses_uk_defaults_without_us_budget_variables() -> None: assert _dataset_variables(plan, "household").issubset(_variables(plan, "household")) +@pytest.mark.parametrize( + ("region", "routed", "planned"), + [ + ("uk", False, True), + ("uk", True, False), + ("country/wales", True, False), + ("constituency/E14001063", True, True), + ("local_authority/E06000063", True, True), + ], +) +def test_uk_plan_copies_area_codes_unless_a_routed_national_report_lacks_them( + monkeypatch, region: str, *, routed: bool, planned: bool +) -> None: + _route_uk_area_regions(monkeypatch, routed=routed) + + plan = resolve_report_output_plan(_report("uk", region=region)) + + assert _dataset_variables(plan, "household") == ( + {"constituency_code_oa", "la_code_oa"} if planned else set() + ) + + def test_cliff_variables_are_conditional() -> None: without_cliffs = resolve_report_output_plan(_report()) with_cliffs = resolve_report_output_plan(_report(include_cliffs=True)) @@ -258,10 +293,12 @@ def test_apply_and_validate_output_plan_reject_missing_materialized_columns() -> @pytest.mark.parametrize("missing", ["constituency_code_oa", "la_code_oa"]) def test_uk_output_validation_requires_geographic_dataset_columns( + monkeypatch, missing: str, ) -> None: from policyengine.tax_benefit_models import uk + _route_uk_area_regions(monkeypatch, routed=False) plan = resolve_report_output_plan(_report("uk")) simulation = Simulation.model_construct( policy={}, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index e91f4d759..08139860a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1487,6 +1487,60 @@ def serialize(self): assert json.loads(stored_aggregate)["result"]["spm_provenance"] is None +@pytest.mark.parametrize( + ("region", "routed", "omitted"), + [ + ("uk", True, True), + ("country/wales", True, True), + ("uk", False, False), + ("constituency/E14001063", True, False), + ], +) +def test_uk_aggregate_omits_area_breakdowns_only_for_routed_national_reports( + monkeypatch, region: str, *, routed: bool, omitted: bool +) -> None: + from policyengine_simulation_executor import simulation_output_builder + + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {} + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder + ) + payload = _uk_report().model_dump(mode="json") + for role in ("baseline", "reform"): + payload[role]["geography"] = {"country": "uk", "region": region} + report = ReportExecutionInput.model_validate(payload) + frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ("person", "benunit", "household") + } + artifacts = FakeArtifacts() + + build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), + ) + + assert observed["omit_uk_area_breakdowns"] is omitted + + def test_aggregate_report_wraps_builder_output_with_report_provenance( monkeypatch, ) -> None: @@ -1549,6 +1603,7 @@ def serialize(self): "resolved_data_version", "resolved_region_code", "uk_local_authority_metadata", + "omit_uk_area_breakdowns", ] assert observed["country"] == "us" assert observed["simulation_params"] == { @@ -1568,6 +1623,7 @@ def serialize(self): ) assert observed["resolved_region_code"] == report.baseline.geography.region assert observed["uk_local_authority_metadata"] is None + assert observed["omit_uk_area_breakdowns"] is False def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: From 24690a27f7de2fac014f97cdbea4d12f2cc975b7 Mon Sep 17 00:00:00 2001 From: Maria Juaristi Date: Wed, 21 Oct 2026 18:18:49 +0100 Subject: [PATCH 04/31] Record the routed UK area changes in the changelog Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index ffa169491..b065763fa 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -33,6 +33,9 @@ - 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 + - Built the constituency and local-authority breakdowns of a UK national or nation-level run on a dataset without area codes from a second baseline and reform run, with the same policies and nation filter, on the dataset the policyengine.py bundle routes those regions to, after checking that dataset's area codes and authority boundaries; the national figures still come from the run's own dataset + - Resolved UK constituency and local-authority regions by name or by GSS code in any case once policyengine.py routes them, as row filters from its region registry, so policyengine-api's constituency/{name} and local_authority/{name} requests keep working without weight matrices, and the 17 districts replaced in April 2023 are unsupported regions + - Made the executor and Stage 12 image smokes check the routed dataset's area codes, downloading it with the Hugging Face credential, when the installed UK default carries none, and stopped Stage 12 copying area codes for a routed UK national or nation-level report, whose aggregate leaves the two area breakdowns out - 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 - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated From 7dad9206b1749a88f5b94cffc404e7e61302c1f6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:18:08 +0400 Subject: [PATCH 05/31] Load the routed UK area pair through main's dataset and span helpers After the merge of main, two routed-national tests and the executor's Pyright check failed: - main's run passes data_folder= to _load_dataset, and the routed tests' fake did not take it (TypeError: unexpected keyword argument 'data_folder'). The fake is now a named function that takes data_folder=None. - _uk_area_simulations opened its three spans with runtime.span(...), which Pyright reads as possibly suppressing exceptions, so it reported selection and dataset as possibly unbound. The spans now open with scopes.span(runtime, ...), as docs/engineering/skills/type-checking.md requires; the span names and attributes are unchanged. The executor suite passes (1789 passed), and Pyright is back to the 77 errors its baseline lists. Co-Authored-By: Claude Opus 5.5 --- .../simulation_runtime.py | 13 ++++++++----- .../tests/test_simulation_output_builder.py | 11 ++++++----- 2 files changed, 14 insertions(+), 10 deletions(-) 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 e9463e54e..271fc194e 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 @@ -908,8 +908,10 @@ def _uk_area_simulations( ) area_role = {"dataset_role": "uk_area"} - with runtime.span( - ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION), attributes=area_role + with scopes.span( + runtime, + ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION), + attributes=area_role, ): selection = _resolve_dataset_selection( params, @@ -917,8 +919,8 @@ def _uk_area_simulations( code=region.code, dataset_reference=routed_uk_area_dataset_name() ), ) - with runtime.span( - ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD), attributes=area_role + with scopes.span( + runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD), attributes=area_role ): dataset = _load_dataset( params, selection=selection, country_module=country_module @@ -930,7 +932,8 @@ def _uk_area_simulations( for kind, policy in zip( ("uk_area_baseline", "uk_area_reform"), policies, strict=True ): - with runtime.span( + with scopes.span( + runtime, ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": kind}, ): 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 7760fa6b9..11b77c186 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1588,6 +1588,11 @@ def build_simulation(_params, *, dataset, dataset_selection, policy, region): ) return simulation + def load_dataset(_params, *, selection, country_module, data_folder=None): + if selection is _LOCAL_AREA_SELECTION: + return local_area_dataset + return national_dataset + class FakeSimulationOutputBuilder: def __init__(self, **kwargs): builder_calls.append(kwargs) @@ -1603,11 +1608,7 @@ def serialize(self): code="country/wales", dataset_reference=None, scoping_strategy=nation ), "_resolve_dataset_selection": resolve_selection, - "_load_dataset": lambda _params, *, selection, country_module: ( - local_area_dataset - if selection is _LOCAL_AREA_SELECTION - else national_dataset - ), + "_load_dataset": load_dataset, "_build_simulation": build_simulation, "SimulationOutputBuilder": FakeSimulationOutputBuilder, }.items(): From acc1d2bf1779afa73880fec8bdf9a8badfe8baa8 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 20 Oct 2026 21:03:15 -0400 Subject: [PATCH 06/31] Use disclosure-controlled UK geography outputs for constituencies and local authorities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The executor built UK constituency and local-authority results with policyengine's deprecated compute_uk_constituency_impacts and compute_longwise_uk_geography_impacts. They apply no minimum sample, so a small area with negative-income households shows an exploding relative change (Cardiff West: -£192 average, +1372% relative), and their output is not disclosure-controlled for licensed survey microdata. Both levels now go through compute_uk_geography_disclosure in one release, nested in regions and the nation, with the constituency/local-authority overlap acknowledged. Blocking findings are resolved by withholding the suggested published cells; region and national reference cells are never withheld, and a release that cannot be resolved publishes no area. The share families, which the output does not publish, are withheld up front. Contract: existing fields keep their types. A withheld area, or one whose relative change is undefined, is left out of the list. Published records gain optional disclosure_tier and 90% interval fields. An empty list is now valid. Stage 12 keeps region_code_oa and source_household_id. The local-authority breakdown no longer strips weights, so the last recorded weight-stripping exception is removed and its baseline is 0. Adapted for #881: the builder keeps #881's uk_area_simulations and omit_uk_area_breakdowns. _uk_geography_disclosure_output() computes the release over the pair _uk_area_breakdown_inputs() returns (the routed local-area pair when given, otherwise the run's own pair), the local-authority impact takes its boundary metadata from _uk_area_breakdown_inputs()[2], and both builders return None before computing anything when the breakdowns are omitted. #881's builder tests now patch compute_uk_geography_disclosure_output(country, baseline, reform, *, year), test_omitted_uk_area_breakdowns_compute_nothing fails if it is called, and #881's planning test expects the four household columns. The changelog keeps both sides' entries. #860 also removes six executor Pyright errors in simulation_output_common.py and simulation_output_geographic.py, so the executor baseline shrinks from 77 to 71 entries, and docs/generated/rules.md is rendered again. Fixes #859 Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 3c99e59dce14a09ca54e187db14e713dbc48e4d9) --- changelog_entry.yaml | 2 + docs/engineering/skills/microdata-weights.md | 7 +- docs/generated/rules.md | 4 +- .../macro_output.py | 39 +- .../tests/test_macro_output_uk_geography.py | 74 +++ .../simulation_output_builder.py | 48 +- .../simulation_output_common.py | 19 +- .../simulation_output_geographic.py | 476 ++++++++++++++++-- .../stage12_runtime/output_planning.py | 10 +- .../tests/test_simulation_output_builder.py | 164 +++--- .../test_simulation_output_geographic.py | 438 +++++++++++++--- .../tests/test_stage12_output_planning.py | 16 +- .../tests/golden/openapi.json | 122 ++++- rules/baselines/pyright-executor.txt | 6 - .../sim-microdata-weight-stripping.txt | 7 - rules/registry.toml | 4 +- 16 files changed, 1177 insertions(+), 259 deletions(-) create mode 100644 libs/policyengine-simulation-contract/tests/test_macro_output_uk_geography.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 0478a3cea..01a7f0d2b 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -2,6 +2,7 @@ changes: added: - Added four registered dependency rules, checked by the locked-installs quality guard with proof fixtures; every project commits a uv.lock that PR CI checks with uv lock --check, every install in workflows, scripts, Makefiles, Dockerfiles, Compose files and Modal images installs the locked versions, every test command runs through uv run --frozen, and the daily update-dependencies workflow relocks every project; every uv sync and uv run now passes --frozen, and the five sites that cannot install from a lock are registered exceptions + - Added optional disclosure fields to UK constituency and local-authority records (disclosure_tier, reported or approximate, and the 90% interval bounds of the average and relative income change); records stored before this change still validate - Added the rule sim.boundaries.sqlalchemy-and-alembic-only, with proof from the psycopg adapter of 08ebe1ca; code connects to databases only through SQLAlchemy, never by importing a database driver, and runtime code never changes the schema, which changes only through Alembic migrations; a guard parses every file under projects/*/src, libs/*/src and .github/scripts and rejects driver imports, alembic imports outside a migrations directory, create_all and drop_all, create and drop on tables, SQLAlchemy DDL statements and DDL text in string literals, which extends the Stage 12 DDL check's pattern to the gateway and every library - Made the budget's reform-minus-baseline totals log a warning when the baseline and reform household frames carry different weights, since the change is then weighted by the reform's weights alone - Added weighted-microdata rules to the registry, with proof from the #767 household count; the executor's code never aggregates a weight column as a value, reads raw values from weighted frames only in simulation_output_common.py (the UK local-authority breakdown's plain household frame is its one recorded exception there, in a list that may only shrink), and its output-builder tests build entity frames through one shared factory with non-uniform weights, while a doc-only rule records that weights, values, labels and counts each allow different operations @@ -27,6 +28,7 @@ - Ran the full PR workflow and the Modal image smoke only on the lowest unmerged pull request in a stack, behind a new CI complete gate job that fails on the others fixed: - Stopped trusting the compromised eMudhra (emSign) intermediate CA certificates; every lockfile now pins certifi 2026.10.19, which drops emSign ECC Root CA - G3, every Modal image upgrades Debian's ca-certificates package right after its base layer, and the Modal images and the Cloud Run entrypoint image set GRPC_DEFAULT_SSL_ROOTS_FILE_PATH so grpcio verifies servers against the operating-system CA bundle instead of the CA list compiled into it + - Built UK constituency and local-authority results with policyengine's disclosure-controlled compute_uk_geography_disclosure instead of the deprecated, unprotected functions, in one release nested in regions and the nation; an area that rests on fewer than 10 survey households, is dominated by one, would be recovered by subtraction, or has a non-positive baseline income is left out of the list, so small areas no longer show exploding relative changes such as +1372% for a £192 loss - Preserved bounded reform stop dates in UK and US simulations while retaining current-law baseline artifact reuse - Resolved PolicyEngine-managed regional datasets using bundled data versions instead of release-manifest commit hashes when loading from GCS - Added a Modal app-release bundle backfill for missing data artifact revisions diff --git a/docs/engineering/skills/microdata-weights.md b/docs/engineering/skills/microdata-weights.md index 9a84a2d87..ed430851e 100644 --- a/docs/engineering/skills/microdata-weights.md +++ b/docs/engineering/skills/microdata-weights.md @@ -82,7 +82,6 @@ Each finding needs a decision, recorded in the PR: Recorded exceptions are listed in `rules/baselines/sim-microdata-weight-stripping.txt`. That list may only shrink, so adding one is a deliberate change to the rule's baseline. Today it -holds one: `_unweighted_household_frame`, the plain household frame the UK -local-authority breakdown hands policyengine's longwise geography impact. It -goes once that breakdown hands policyengine the simulations, as the -constituency and US congressional district breakdowns already do. +is empty: the last one, `_unweighted_household_frame`, went when the UK +constituency and local-authority breakdowns moved to policyengine's +disclosure-controlled geography output, which takes the simulations. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 686c1ea01..7a58a25f8 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -66,9 +66,9 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf ### Ratcheted (existing violations still tolerated) - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` -- `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` +- `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt` -- `sim.quality.pyright-executor`: 77 remaining in `rules/baselines/pyright-executor.txt` +- `sim.quality.pyright-executor`: 71 remaining in `rules/baselines/pyright-executor.txt` - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt` ### Exceptions and review dates diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index a45cddec4..5fcd2bb60 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Generic, TypeVar +from typing import Generic, Literal, TypeVar from pydantic import BaseModel, ConfigDict, RootModel @@ -234,12 +234,38 @@ class CliffImpactOutput(MacroOutputModel): reform: CliffImpactInSimulation -class ConstituencyImpactRecord(MacroOutputModel): +UKGeographyDisclosureTier = Literal["reported", "approximate"] + + +class UKAreaDisclosureFields(MacroOutputModel): + """Disclosure tier and 90% intervals of a published UK area's income changes. + + UK area results come from licensed survey microdata and pass policyengine's + disclosure control. An area whose income change is withheld is left out of + the list, so every record here has values. ``disclosure_tier`` is + ``reported``, or ``approximate`` when the area's effective sample is thin + and the interval understates the uncertainty. The ``_lower`` and ``_upper`` + fields bound each change's 90% bootstrap interval, rounded outward. + + All of these fields are optional, so results stored before disclosure + control still validate; there they are null. + """ + + disclosure_tier: UKGeographyDisclosureTier | None = None + average_household_income_change_lower: float | None = None + average_household_income_change_upper: float | None = None + relative_household_income_change_lower: float | None = None + relative_household_income_change_upper: float | None = None + + +class ConstituencyImpactRecord(UKAreaDisclosureFields): """Reform impact on household net income in one UK parliamentary constituency. ``average_household_income_change`` is the average change per household and ``relative_household_income_change`` the relative change in the area's total - household net income. ``x`` and ``y`` are integer map-rendering coordinates. + household net income, both rounded to the order of their standard error. + ``population`` is the weighted number of households. ``x`` and ``y`` are + integer map-rendering coordinates. """ constituency_code: str @@ -251,7 +277,7 @@ class ConstituencyImpactRecord(MacroOutputModel): population: float -class LocalAuthorityImpactRecord(MacroOutputModel): +class LocalAuthorityImpactRecord(UKAreaDisclosureFields): """Reform impact on household net income in one UK local authority. Fields mean the same as in ``ConstituencyImpactRecord``. @@ -269,7 +295,10 @@ class LocalAuthorityImpactRecord(MacroOutputModel): class GeographicImpactOutput( MacroRootModel[list[ConstituencyImpactRecord | LocalAuthorityImpactRecord]] ): - """UK constituency or local-authority impact records.""" + """UK constituency or local-authority impact records. + + Areas whose income change disclosure control withholds are left out. + """ class CongressionalDistrictImpactRecord(MacroOutputModel): diff --git a/libs/policyengine-simulation-contract/tests/test_macro_output_uk_geography.py b/libs/policyengine-simulation-contract/tests/test_macro_output_uk_geography.py new file mode 100644 index 000000000..87d8ae48a --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_macro_output_uk_geography.py @@ -0,0 +1,74 @@ +"""UK area records stay readable before and after disclosure control.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.macro_output import ( + ConstituencyImpactRecord, + GeographicImpactOutput, + LocalAuthorityImpactRecord, +) +from pydantic import ValidationError + +STORED_BEFORE_DISCLOSURE = { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "x": 56, + "y": -40, + "average_household_income_change": -150.95, + "relative_household_income_change": -0.0026, + "population": 40_530.65, +} +DISCLOSURE_FIELDS = { + "disclosure_tier": "approximate", + "average_household_income_change_lower": -170.0, + "average_household_income_change_upper": -130.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.0022, +} + + +def test_records_stored_before_disclosure_control_still_validate() -> None: + record = ConstituencyImpactRecord.model_validate(STORED_BEFORE_DISCLOSURE) + + assert record.disclosure_tier is None + assert record.average_household_income_change_lower is None + assert record.relative_household_income_change_upper is None + + +def test_records_carry_the_disclosure_tier_and_intervals() -> None: + record = LocalAuthorityImpactRecord.model_validate( + { + "local_authority_code": "E06000063", + "local_authority_name": "Cumberland", + "x": 4, + "y": 18, + "average_household_income_change": -150.0, + "relative_household_income_change": -0.0026, + "population": 120_000.0, + **DISCLOSURE_FIELDS, + } + ) + + assert record.model_dump(mode="json")["disclosure_tier"] == "approximate" + assert record.average_household_income_change_lower == -170.0 + + +def test_a_withheld_tier_is_not_a_record() -> None: + # Withheld areas are left out of the list, so no record holds that tier. + with pytest.raises(ValidationError): + ConstituencyImpactRecord.model_validate( + {**STORED_BEFORE_DISCLOSURE, "disclosure_tier": "withheld"} + ) + + +def test_published_values_stay_required() -> None: + record = {**STORED_BEFORE_DISCLOSURE, **DISCLOSURE_FIELDS} + record["average_household_income_change"] = None + + with pytest.raises(ValidationError): + ConstituencyImpactRecord.model_validate(record) + + +def test_an_empty_area_list_is_valid() -> None: + assert GeographicImpactOutput.model_validate([]).root == [] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index a04e2c3c3..ce945dabe 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -89,6 +89,7 @@ class SimulationOutputBuilder: omit_uk_area_breakdowns: bool = False progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: Any = field(default=None, init=False) + _uk_geography_disclosure: Any = field(default=None, init=False) def __post_init__(self) -> None: """Normalize ``country`` to lower case.""" @@ -294,24 +295,61 @@ def _uk_area_breakdown_inputs( area.reform.ensure() return area.baseline, area.reform, area.local_authority_metadata + def _uk_geography_disclosure_output( + self, + ) -> simulation_output_geographic.UKGeographyDisclosureOutput | None: + """Return UK area changes under disclosure control, computed once. + + Constituencies and local authorities go through one release, so the + disclosure check sees both. The release reads the pair the area + breakdowns use (see ``_uk_area_breakdown_inputs``). None for other + countries. + """ + if self.country != "uk": + return None + if self._uk_geography_disclosure is None: + baseline, reform, _ = self._uk_area_breakdown_inputs() + self._uk_geography_disclosure = ( + simulation_output_geographic.compute_uk_geography_disclosure_output( + self.country, + baseline, + reform, + year=self._simulation_year(), + ) + ) + return self._uk_geography_disclosure + + def _simulation_year(self) -> int: + """Return the simulated year, which labels the disclosure release.""" + for value in ( + self.simulation_params.get("time_period"), + self.simulation_params.get("year"), + getattr(self.dataset, "year", None), + ): + if value is None: + continue + try: + return int(value) + except (TypeError, ValueError): + continue + return 0 + def _build_uk_constituency_impact(self) -> GeographicImpactOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY)): if self.omit_uk_area_breakdowns: return None - baseline, reform, _ = self._uk_area_breakdown_inputs() return simulation_output_geographic.build_uk_constituency_impact( - self.country, baseline, reform + self.country, self._uk_geography_disclosure_output() ) def _build_uk_local_authority_impact(self) -> GeographicImpactOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY)): if self.omit_uk_area_breakdowns: return None - baseline, reform, metadata = self._uk_area_breakdown_inputs() + metadata = self._uk_area_breakdown_inputs()[2] return simulation_output_geographic.build_uk_local_authority_impact( self.country, - baseline, - reform, + self._uk_geography_disclosure_output(), uk_local_authority_metadata=metadata, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py index 4ce0d2feb..ae5c73308 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py @@ -8,10 +8,10 @@ from importlib import import_module from typing import TYPE_CHECKING, Any -import pandas as pd from policyengine_simulation_contract.spm import spm_error_detail if TYPE_CHECKING: + import pandas as pd from _typeshed import ConvertibleToFloat from policyengine.core import Simulation @@ -77,23 +77,6 @@ def _count_output_entities(simulation: Simulation, entity: str, weight: str) -> return float(np.asarray(data[weight], dtype=float).sum()) -def _unweighted_household_frame(simulation: Simulation) -> pd.DataFrame: - """Return the household output as a plain DataFrame, without its weights. - - policyengine's ``compute_longwise_uk_geography_impacts`` takes plain frames - and weights each area by the baseline's ``household_weight`` column itself. - The UK local-authority breakdown calls it directly, so that it can apply - this service's boundary metadata. Handing it the weighted frame would weight - twice if it ever reduced a column directly. - - This is a recorded exception to sim.microdata.weight-stripping-confined, - listed in rules/baselines/sim-microdata-weight-stripping.txt, which may only - shrink. It goes once the breakdown hands policyengine the simulations, as - the constituency and US congressional district breakdowns already do. - """ - return pd.DataFrame(simulation.output_dataset.data.household) - - def _warn_on_mismatched_weights( baseline_data: pd.DataFrame, reform_data: pd.DataFrame, entity: str ) -> None: 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 849c6c7a1..9c978ad37 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 @@ -2,13 +2,29 @@ from __future__ import annotations +import logging +import math from collections.abc import Mapping +from dataclasses import dataclass, field from functools import lru_cache -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast +import pandas as pd from policyengine.data.uk_geography_assets import ( CONSTITUENCY_ASSET_SPEC, ) +from policyengine.disclosure import UNASSIGNED, DisclosureError +from policyengine.disclosure.areas import normalise_area_codes +from policyengine.disclosure.schema import ( + APPROXIMATE, + AVERAGE_CHANGE, + FAMILY_CHANGE, + FAMILY_GAINING, + FAMILY_LOSING, + RELATIVE_CHANGE, + REPORTED, + UNCHECKED_OVERLAP, +) from policyengine_simulation_executor.simulation_macro_output import ( CongressionalDistrictImpactOutput, @@ -16,19 +32,24 @@ GeographicImpactOutput, ) from policyengine_simulation_executor.simulation_output_common import ( + _entity_data, _number, _output_model_dump, - _output_module_function, _try_compute_output, - _unweighted_household_frame, ) if TYPE_CHECKING: - import pandas as pd from policyengine.core import Simulation from policyengine.data.uk_geography_assets import UKGeographyAssetSpec + from policyengine.disclosure import PublicTable + from policyengine.outputs.uk_geography_disclosure import UKGeographyDisclosure + from policyengine_simulation_contract.macro_output import ( + UKGeographyDisclosureTier, + ) from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata +logger = logging.getLogger(__name__) + def _required_uk_geography_lookup_csv_path(spec: UKGeographyAssetSpec) -> str: """Resolve one required lookup under the runtime GCP identity.""" @@ -60,10 +81,13 @@ def _complete_uk_geography_output( code_field: str, name_field: str, ) -> GeographicImpactOutput: - """Require every UK geography record to contain lookup-owned metadata.""" + """Require every UK geography record to contain lookup-owned metadata. + + An empty list is valid: disclosure control may withhold every area. + """ output = build_geographic_impact_output(value) - if output is None or not output.root: - msg = "UK geography output did not contain result records" + if output is None: + msg = "UK geography output did not contain a list of result records" raise ValueError(msg) for record in output.root: record_values = record.model_dump(mode="python") @@ -236,6 +260,77 @@ def compute_and_format() -> CongressionalDistrictImpactOutput | None: return _try_compute_output("congressional district impacts", compute_and_format) +# The household output column that holds each published UK level's area code. +UK_DISCLOSURE_AREA_COLUMNS: dict[str, str] = { + "constituency": "constituency_code_oa", + "local_authority": "la_code_oa", +} +UK_DISCLOSURE_BREAKDOWN_NAMES: dict[str, str] = { + "constituency": "constituency", + "local_authority": "local-authority", +} +# The survey household each row was copied or cloned from. The count floor +# counts support by this id, so copies of one survey household count once. +UK_SOURCE_HOUSEHOLD_ID_COLUMN = "source_household_id" +# Bootstrap replicates behind each 90% interval: the disclosure layer's +# minimum. They set the interval and the rounding, never whether a cell is +# published, and the layer's default of 1,000 would add minutes to every UK run. +UK_DISCLOSURE_REPLICATES = 200 +UK_DISCLOSURE_REFORM_LABEL = "reform" +# Each round withholds at least one more cell, so this bounds the check runs. +UK_DISCLOSURE_MAX_ROUNDS = 25 +# The output publishes the change family only. The share families are +# withheld before the check, so the check treats them as secret, which they +# are, and never asks for cells to protect shares nobody publishes. +_UNPUBLISHED_FAMILY_OUTCOMES = (FAMILY_GAINING, FAMILY_LOSING) +_FAMILY_OUTCOME = { + FAMILY_CHANGE: AVERAGE_CHANGE, + FAMILY_GAINING: FAMILY_GAINING, + FAMILY_LOSING: FAMILY_LOSING, +} + + +@dataclass(frozen=True) +class UKAreaIncomeChange: + """One published UK area's disclosure-controlled household income change.""" + + tier: UKGeographyDisclosureTier + average: float + average_lower: float + average_upper: float + relative: float + relative_lower: float + relative_upper: float + population: float + + def record_fields(self) -> dict[str, object]: + """Return the fields this change contributes to an area record.""" + return { + "average_household_income_change": self.average, + "relative_household_income_change": self.relative, + "population": self.population, + "disclosure_tier": self.tier, + "average_household_income_change_lower": self.average_lower, + "average_household_income_change_upper": self.average_upper, + "relative_household_income_change_lower": self.relative_lower, + "relative_household_income_change_upper": self.relative_upper, + } + + +@dataclass(frozen=True) +class UKGeographyDisclosureOutput: + """The published areas of each UK level, keyed by area code. + + An area the disclosure check withholds is absent. + """ + + areas: Mapping[str, Mapping[str, UKAreaIncomeChange]] = field(default_factory=dict) + + def level(self, level: str) -> Mapping[str, UKAreaIncomeChange]: + """Return the published areas of one level.""" + return self.areas.get(level, {}) + + def _require_uk_area_column( household: pd.DataFrame, column: str, breakdown: str ) -> None: @@ -254,31 +349,335 @@ def _require_uk_area_column( raise ValueError(msg) -def build_uk_constituency_impact( - country: str, baseline: Simulation, reform: Simulation -) -> GeographicImpactOutput | None: - """Compute UK parliamentary constituency impacts; None for other countries. +def _require_uk_source_household_ids(household: pd.DataFrame) -> None: + """Refuse UK area outputs when support cannot be counted by survey household.""" + if UK_SOURCE_HOUSEHOLD_ID_COLUMN not in household.columns: + msg = ( + f"UK constituency and local-authority breakdowns need " + f"{UK_SOURCE_HOUSEHOLD_ID_COLUMN} on the simulated dataset: disclosure " + "control counts the survey households behind each area by it." + ) + raise ValueError(msg) + + +def _normalised_area_code(value: object) -> str: + """Return an area code as the disclosure layer keys it.""" + return str(normalise_area_codes([value])[0]) + + +def _uk_area_household_counts(household: pd.DataFrame, column: str) -> dict[str, float]: + """Return the weighted number of households in each area of one level. + + ``household`` is the weighted output frame, so a grouped count is the sum + of household weights in each area. Areas without weight are left out, as + the disclosure layer leaves them out. + """ + counts = household.groupby(column)["household_net_income"].count() + by_code: dict[str, float] = {} + for raw_code, count in counts.items(): + code = _normalised_area_code(raw_code) + if code == UNASSIGNED: + continue + by_code[code] = by_code.get(code, 0.0) + float(cast("float", count)) + return {code: count for code, count in by_code.items() if count > 0} + + +def _withhold_unpublished_families( + impact: UKGeographyDisclosure, + area_codes: Mapping[str, Mapping[str, float]], + *, + year: int, +) -> None: + release = impact.release + for level, codes in area_codes.items(): + for code in codes: + for outcome in _UNPUBLISHED_FAMILY_OUTCOMES: + try: + release.withhold( + level=level, + area_code=code, + reform=UK_DISCLOSURE_REFORM_LABEL, + year=year, + outcome=outcome, + ) + except DisclosureError: + # Not an area of the release, for example one whose + # households all have zero weight. + continue + + +def _cells_resolving_findings( + impact: UKGeographyDisclosure, published_levels: tuple[str, ...] +) -> set[tuple[str, str, str, str, int]]: + """Return the published cells to withhold for the blocking findings. + + A finding names a cell to withhold, or else the group it is about. Only + cells of published levels are withheld: regions and the nation are + references whose numbers are public elsewhere, so withholding them here + would only hide the risk. + """ + cells: set[tuple[str, str, str, str, int]] = set() + for finding in impact.findings or []: + if not finding.blocking or finding.family is None: + continue + if finding.reform is None or finding.year is None: + continue + outcome = _FAMILY_OUTCOME.get(finding.family) + if outcome is None: + continue + if finding.suggested and finding.level in published_levels: + cells.add( + ( + finding.level, + finding.suggested, + outcome, + finding.reform, + finding.year, + ) + ) + elif finding.group and finding.group_level in published_levels: + cells.add( + ( + finding.group_level, + finding.group, + outcome, + finding.reform, + finding.year, + ) + ) + return cells + + +def _resolved_public_table( + impact: UKGeographyDisclosure, + area_codes: Mapping[str, Mapping[str, float]], + *, + year: int, +) -> PublicTable | None: + """Withhold cells until no finding blocks, and return the public table. + + Returns None when the findings cannot be resolved by withholding published + cells; then nothing at these levels may be published. + """ + published_levels = tuple(area_codes) + _withhold_unpublished_families(impact, area_codes, year=year) + public = impact.refresh() + withheld: set[tuple[str, str, str, str, int]] = set() + for _ in range(UK_DISCLOSURE_MAX_ROUNDS): + if public is not None: + return public + cells = _cells_resolving_findings(impact, published_levels) - withheld + if not cells: + break + for level, code, outcome, reform, cell_year in sorted(cells): + try: + impact.release.withhold( + level=level, + area_code=code, + reform=reform, + year=cell_year, + outcome=outcome, + ) + except DisclosureError: + # Not a cell of the release. The finding stays, so the loop + # ends without progress and nothing is published. + continue + withheld |= cells + public = impact.refresh() + if public is not None: + return public + blocking = sorted( + {finding.code for finding in impact.findings or [] if finding.blocking} + ) + # Findings carry area codes; the log names only their kinds. + logger.warning( + "UK geography disclosure findings remain after withholding %d cells " + "(%s), so no constituency or local-authority area is published", + len(withheld), + ", ".join(blocking), + ) + return None + + +def _published_area_changes( + public: PublicTable, + level: str, + populations: Mapping[str, float], +) -> dict[str, UKAreaIncomeChange]: + """Return the areas of one level whose average and relative change are public. - Downloads the constituency lookup if needed and raises when the dataset - lacks constituency codes or a record lacks lookup name or coordinates. + An area is published only when both changes are: the contract has a value + for each. A relative change is withheld on its own, for example, when the + area's baseline income is not positive. + """ + rows: dict[str, dict[str, Any]] = {} + for row in public.rows: + if row.level != level or row.outcome not in {AVERAGE_CHANGE, RELATIVE_CHANGE}: + continue + if row.tier not in {REPORTED, APPROXIMATE} or row.value is None: + continue + if row.lo is None or row.hi is None: + continue + rows.setdefault(row.area_code, {})[row.outcome] = row + changes: dict[str, UKAreaIncomeChange] = {} + for code, outcomes in rows.items(): + average = outcomes.get(AVERAGE_CHANGE) + relative = outcomes.get(RELATIVE_CHANGE) + if average is None or relative is None or code not in populations: + continue + tier: UKGeographyDisclosureTier = ( + APPROXIMATE if APPROXIMATE in {average.tier, relative.tier} else REPORTED + ) + changes[code] = UKAreaIncomeChange( + tier=tier, + average=float(average.value), + average_lower=float(average.lo), + average_upper=float(average.hi), + relative=float(relative.value), + relative_lower=float(relative.lo), + relative_upper=float(relative.hi), + population=populations[code], + ) + return changes + + +def compute_uk_geography_disclosure_output( + country: str, + baseline: Simulation, + reform: Simulation, + *, + year: int, +) -> UKGeographyDisclosureOutput | None: + """Compute disclosure-controlled constituency and local-authority changes. + + None for other countries. Both levels go through one policyengine release, + nested in regions and the nation, so the check sees what subtraction + between them recovers. Raises when the dataset lacks area codes or survey + household ids. """ if country != "uk": return None - _require_uk_area_column( - baseline.output_dataset.data.household, "constituency_code_oa", "constituency" + household = _entity_data(baseline, "household") + _entity_data(reform, "household") + for level, column in UK_DISCLOSURE_AREA_COLUMNS.items(): + _require_uk_area_column(household, column, UK_DISCLOSURE_BREAKDOWN_NAMES[level]) + _require_uk_source_household_ids(household) + + from policyengine.disclosure import ( # noqa: PLC0415 # bound per call, like this module's other policyengine.outputs imports + DisclosureSettings, + ) + from policyengine.outputs.uk_geography_disclosure import ( # noqa: PLC0415 # bound per call, like this module's other policyengine.outputs imports + compute_uk_geography_disclosure, ) - lookup_csv_path = _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC) - impact = _output_module_function( - "constituency_impact", "compute_uk_constituency_impacts" - )( + area_codes = { + level: _uk_area_household_counts(household, column) + for level, column in UK_DISCLOSURE_AREA_COLUMNS.items() + } + impact = compute_uk_geography_disclosure( baseline, reform, - constituency_csv_path=lookup_csv_path, - download_missing_assets=False, + year=year, + levels=tuple(UK_DISCLOSURE_AREA_COLUMNS), + reform=UK_DISCLOSURE_REFORM_LABEL, + settings=DisclosureSettings(n_replicates=UK_DISCLOSURE_REPLICATES), + # Constituencies and local authorities overlap without nesting. The + # release checks each area against the areas inside and around it; + # sim-api publishes both levels from every UK run, so it accepts that + # other differences between them stay unchecked. + acknowledge=(UNCHECKED_OVERLAP,), + ) + public = _resolved_public_table(impact, area_codes, year=year) + if public is None: + return UKGeographyDisclosureOutput() + return UKGeographyDisclosureOutput( + areas={ + level: _published_area_changes(public, level, area_codes[level]) + for level in UK_DISCLOSURE_AREA_COLUMNS + } + ) + + +def _uk_constituency_lookup( + lookup_csv_path: str, +) -> dict[str, tuple[str, int | None, int | None]]: + """Return each constituency's name and map coordinates, in lookup order.""" + lookup = pd.read_csv(lookup_csv_path) + if "code" not in lookup.columns: + msg = ( + "UK constituency lookup CSV must contain a 'code' column: " + f"{lookup_csv_path}" + ) + raise ValueError(msg) + entries: dict[str, tuple[str, int | None, int | None]] = {} + for row in lookup.to_dict("records"): + code = _normalised_area_code(row.get("code")) + if not code: + continue + name = _normalised_area_code(row.get("name")) or code + entries[code] = ( + name, + _lookup_coordinate(row.get("x")), + _lookup_coordinate(row.get("y")), + ) + return entries + + +def _lookup_coordinate(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + number = float(value) + except ValueError: + return None + return int(number) if math.isfinite(number) else None + + +def _require_uk_disclosure_output( + disclosure: UKGeographyDisclosureOutput | None, breakdown: str +) -> UKGeographyDisclosureOutput: + if disclosure is None: + msg = ( + f"UK {breakdown} output needs the disclosure-controlled UK geography output" + ) + raise ValueError(msg) + return disclosure + + +def build_uk_constituency_impact( + country: str, disclosure: UKGeographyDisclosureOutput | None +) -> GeographicImpactOutput | None: + """Return the published UK parliamentary constituencies; None for other countries. + + Areas that disclosure control withholds are left out. Downloads the + constituency lookup if needed and raises when a published area lacks a + lookup name or coordinates. + """ + if country != "uk": + return None + changes = _require_uk_disclosure_output(disclosure, "constituency").level( + "constituency" ) + lookup = _uk_constituency_lookup( + _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC) + ) + ordered = [code for code in lookup if code in changes] + sorted( + set(changes) - set(lookup) + ) + records: list[dict[str, object]] = [] + for code in ordered: + name, x, y = lookup.get(code, (code, None, None)) + records.append( + { + "constituency_code": code, + "constituency_name": name, + "x": x, + "y": y, + **changes[code].record_fields(), + } + ) return _complete_uk_geography_output( - getattr(impact, "constituency_results", None), + records, code_field="constituency_code", name_field="constituency_name", ) @@ -286,51 +685,35 @@ def build_uk_constituency_impact( def build_uk_local_authority_impact( country: str, - baseline: Simulation, - reform: Simulation, + disclosure: UKGeographyDisclosureOutput | None, *, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> GeographicImpactOutput | None: - """Compute UK local authority impacts, named and placed for the detected boundaries. + """Return the published UK local authorities, named and placed for the boundaries. None for other countries. For the UK, ``uk_local_authority_metadata`` is required, and a code outside its boundary version raises ``ValueError``. + Areas that disclosure control withholds are left out. """ if country != "uk": return None - - baseline_household = _unweighted_household_frame(baseline) - _require_uk_area_column(baseline_household, "la_code_oa", "local-authority") + changes = _require_uk_disclosure_output(disclosure, "local-authority").level( + "local_authority" + ) if uk_local_authority_metadata is None: msg = "UK local-authority boundary metadata is required" raise ValueError(msg) - from policyengine.outputs.uk_geography_impact import ( # noqa: PLC0415 # bound per call, like this module's patched policyengine.outputs import above - compute_longwise_uk_geography_impacts, - ) - from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module loads only for UK local-authority output load_uk_local_authority_resources, ) - reform_household = _unweighted_household_frame(reform) - numeric_records = compute_longwise_uk_geography_impacts( - baseline_household=baseline_household, - reform_household=reform_household, - geography_column="la_code_oa", - result_key_prefix="local_authority", - lookup_csv_path=None, - ) boundary_version = uk_local_authority_metadata.boundary_version display_metadata = load_uk_local_authority_resources().metadata_for( boundary_version ) records: list[dict[str, object]] = [] - for numeric_record in numeric_records: - code = numeric_record.get("local_authority_code") - if not isinstance(code, str) or not code: - msg = "UK local-authority output contains an invalid code" - raise ValueError(msg) + for code in sorted(changes): metadata = display_metadata.get(code) if metadata is None: msg = ( @@ -340,10 +723,11 @@ def build_uk_local_authority_impact( raise ValueError(msg) records.append( { - **numeric_record, + "local_authority_code": code, "local_authority_name": metadata.name, "x": metadata.x, "y": metadata.y, + **changes[code].record_fields(), } ) return _complete_uk_geography_output( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py index 51016f5b7..eb68727c6 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py @@ -31,8 +31,16 @@ import pandas as pd from policyengine_simulation_contract.stage12_bundle import CountryId +# Input-dataset columns the UK constituency and local-authority breakdowns +# read: the area codes, the region the disclosure check nests them in, and the +# survey household each row was copied from, by which it counts support. UK_GEOGRAPHIC_DATASET_VARIABLES: dict[str, tuple[str, ...]] = { - "household": ("constituency_code_oa", "la_code_oa"), + "household": ( + "constituency_code_oa", + "la_code_oa", + "region_code_oa", + "source_household_id", + ), } 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 11b77c186..fcdcbbb8c 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -37,6 +37,7 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) +from policyengine_simulation_executor import simulation_output_geographic from policyengine_simulation_executor.job_progress import NULL_JOB_PROGRESS from policyengine_simulation_executor.release_bundle import ( BUNDLE_RECEIPT_FILENAME, @@ -87,8 +88,6 @@ ) from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage -CONSTITUENCY_CSV_PATH = "/tmp/constituencies_2024.csv" - def _params_with_spm(params, selection): """Return request parameters after the runtime resolves SPM settings.""" @@ -2475,77 +2474,82 @@ def fail_change_output_variable(*_args, **_kwargs): _simulation_output_builder("us", baseline, reform)._build_budgetary_impact() -def test_uk_constituency_impact_uses_policyengine_output_function(monkeypatch): - def coded_simulation() -> SimpleNamespace: - household = weighted_entity_frame({"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 test_uk_area_impacts_share_one_disclosure_release(monkeypatch): + baseline = object() + reform = object() + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + computed: list[tuple[str, object, object, int]] = [] + built: list[tuple[str, object]] = [] - def fake_output_module_function(module_name, name): - assert module_name == "constituency_impact" - assert name == "compute_uk_constituency_impacts" + def compute(country, baseline_simulation, reform_simulation, *, year): + computed.append((country, baseline_simulation, reform_simulation, year)) + return disclosure - def compute(baseline_simulation, reform_simulation, **kwargs): - assert baseline_simulation is baseline - assert reform_simulation is reform - assert kwargs == { - "constituency_csv_path": CONSTITUENCY_CSV_PATH, - "download_missing_assets": False, - } - return SimpleNamespace(constituency_results=expected) + def build_constituencies(country, value): + built.append(("constituency", value)) + return GeographicImpactOutput.model_validate([]) - return compute + def build_authorities(country, value, **kwargs): + built.append(("local_authority", value)) + return GeographicImpactOutput.model_validate([]) monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic._output_module_function", - fake_output_module_function, + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + compute, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "build_uk_constituency_impact", + build_constituencies, ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_output_geographic." - "_required_uk_geography_lookup_csv_path", - lambda _spec: CONSTITUENCY_CSV_PATH, + "build_uk_local_authority_impact", + build_authorities, ) - result = _simulation_output_builder( - "uk", baseline, reform - )._build_uk_constituency_impact() - assert result is not None - assert result.model_dump(mode="json") == expected - assert ( - _simulation_output_builder( - "us", baseline, reform - )._build_uk_constituency_impact() - is None + builder = _simulation_output_builder("uk", baseline, reform) + builder.simulation_params["time_period"] = "2026" + builder._build_uk_constituency_impact() + builder._build_uk_local_authority_impact() + + # Both levels come from one release, computed once. + assert computed == [("uk", baseline, reform, 2026)] + assert built == [("constituency", disclosure), ("local_authority", disclosure)] + + +def test_uk_disclosure_is_not_computed_for_other_countries(monkeypatch): + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + lambda *_args, **_kwargs: pytest.fail("US runs need no UK disclosure"), ) + builder = _simulation_output_builder("us", object(), object()) + + assert builder._build_uk_constituency_impact() is None + assert builder._build_uk_local_authority_impact() is None def test_uk_local_authority_impact_receives_detected_boundary_metadata(monkeypatch): - baseline = object() - reform = object() expected = [_local_authority_impact_record().model_dump(mode="json")] metadata = UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ) - calls: list[tuple[object, object, UKLocalAuthorityMetadata | None]] = [] + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + calls: list[tuple[object, UKLocalAuthorityMetadata | None]] = [] - def build(country, baseline_simulation, reform_simulation, **kwargs): + def build(country, value, **kwargs): if country != "uk": return None - assert country == "uk" - calls.append( - ( - baseline_simulation, - reform_simulation, - kwargs["uk_local_authority_metadata"], - ) - ) + calls.append((value, kwargs["uk_local_authority_metadata"])) return GeographicImpactOutput.model_validate(expected) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + lambda *_args, **_kwargs: disclosure, + ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_output_geographic." "build_uk_local_authority_impact", @@ -2554,19 +2558,13 @@ def build(country, baseline_simulation, reform_simulation, **kwargs): result = _simulation_output_builder( "uk", - baseline, - reform, + object(), + object(), BuilderOptions(uk_local_authority_metadata=metadata), )._build_uk_local_authority_impact() assert result is not None assert result.model_dump(mode="json") == expected - assert calls == [(baseline, reform, metadata)] - assert ( - _simulation_output_builder( - "us", baseline, reform - )._build_uk_local_authority_impact() - is None - ) + assert calls == [(disclosure, metadata)] class _AreaSimulation: @@ -2587,34 +2585,30 @@ def test_uk_area_breakdowns_use_the_computed_routed_area_pair(monkeypatch): boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ), ) - calls: list[tuple[str, object, object, object]] = [] + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + computed: list[tuple[str, object, object, int]] = [] + built: list[tuple[str, object, object]] = [] - def constituencies(country, baseline_simulation, reform_simulation): - assert country == "uk" + def compute(country, baseline_simulation, reform_simulation, *, year): assert baseline_simulation.ensured assert reform_simulation.ensured - calls.append(("constituency", baseline_simulation, reform_simulation, None)) + computed.append((country, baseline_simulation, reform_simulation, year)) + return disclosure + + def constituencies(country, value): + built.append(("constituency", value, None)) return GeographicImpactOutput.model_validate( [_constituency_impact_record().model_dump(mode="json")] ) - def local_authorities(country, baseline_simulation, reform_simulation, **kwargs): - assert country == "uk" - assert baseline_simulation.ensured - assert reform_simulation.ensured - calls.append( - ( - "local_authority", - baseline_simulation, - reform_simulation, - kwargs["uk_local_authority_metadata"], - ) - ) + def local_authorities(country, value, **kwargs): + built.append(("local_authority", value, kwargs["uk_local_authority_metadata"])) return GeographicImpactOutput.model_validate( [_local_authority_impact_record().model_dump(mode="json")] ) for name, function in { + "compute_uk_geography_disclosure_output": compute, "build_uk_constituency_impact": constituencies, "build_uk_local_authority_impact": local_authorities, }.items(): @@ -2628,22 +2622,24 @@ def local_authorities(country, baseline_simulation, reform_simulation, **kwargs) object(), BuilderOptions(resolved_region_code="uk", uk_area_simulations=area), ) + builder.simulation_params["time_period"] = "2026" assert builder._build_uk_constituency_impact() is not None assert builder._build_uk_local_authority_impact() is not None - assert calls == [ - ("constituency", area.baseline, area.reform, None), - ( - "local_authority", - area.baseline, - area.reform, - area.local_authority_metadata, - ), + # One disclosure release over the routed pair serves both levels. + assert computed == [("uk", area.baseline, area.reform, 2026)] + assert built == [ + ("constituency", disclosure, None), + ("local_authority", disclosure, area.local_authority_metadata), ] def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): - for name in ("build_uk_constituency_impact", "build_uk_local_authority_impact"): + for name in ( + "compute_uk_geography_disclosure_output", + "build_uk_constituency_impact", + "build_uk_local_authority_impact", + ): monkeypatch.setattr( f"policyengine_simulation_executor.simulation_output_geographic.{name}", lambda *_args, **_kwargs: pytest.fail("no area breakdown is built"), 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 5e5e20cdf..e3d8320ff 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -5,6 +5,7 @@ from contextlib import nullcontext from types import SimpleNamespace +import numpy as np import pytest from fixtures.weighted_frames import weighted_entity_frame from policyengine.data.uk_geography_assets import CONSTITUENCY_ASSET_SPEC @@ -29,22 +30,6 @@ def _complete_constituency_record() -> dict[str, object]: } -def _uk_simulation( - code: str, - income: float, -) -> SimpleNamespace: - household = weighted_entity_frame( - { - "la_code_oa": [code], - "household_net_income": [income], - "household_weight": [2.0], - } - ) - return SimpleNamespace( - output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) - ) - - def test_required_uk_lookup_uses_gcp_credentials(monkeypatch) -> None: observed: list[str] = [] @@ -104,10 +89,21 @@ def test_complete_uk_geography_output_accepts_lookup_metadata() -> None: assert record["y"] == -40 -def test_complete_uk_geography_output_rejects_empty_result() -> None: - with pytest.raises(ValueError, match="did not contain result records"): +def test_complete_uk_geography_output_accepts_an_empty_list() -> None: + # Disclosure control may withhold every area of a level. + result = simulation_output_geographic._complete_uk_geography_output( + [], + code_field="constituency_code", + name_field="constituency_name", + ) + + assert result.root == [] + + +def test_complete_uk_geography_output_rejects_a_missing_list() -> None: + with pytest.raises(ValueError, match="did not contain a list"): simulation_output_geographic._complete_uk_geography_output( - [], + None, code_field="constituency_code", name_field="constituency_name", ) @@ -138,18 +134,169 @@ def test_complete_uk_geography_output_rejects_fallback_metadata( ) -@pytest.mark.parametrize( - ("boundary_version", "code", "expected_name"), - [ - (UKLocalAuthorityBoundaryVersion.LAD22, "E07000026", "Allerdale"), - (UKLocalAuthorityBoundaryVersion.LAD23, "E06000063", "Cumberland"), - ], -) -def test_local_authority_output_uses_detected_boundary_metadata( - monkeypatch, - boundary_version: UKLocalAuthorityBoundaryVersion, - code: str, - expected_name: str, +# Disclosure-controlled UK breakdowns. The households below sit in one region +# and two LAD23 local authorities. Each published seat rests on 30 survey +# households, well above the floor of 10; THIN has 3, so it is withheld, and +# the check then withholds one more seat so that subtraction from the region +# or from its local authority cannot recover it. NEGATIVE has a negative total +# baseline income, the Cardiff West case: its relative change is undefined, so +# the seat is left out although its average change passes. +NORTH_WEST = "E12000002" +CUMBERLAND = "E06000063" +WESTMORLAND = "E06000064" +SEATS = { + "E14000001": (CUMBERLAND, 30), + "E14000002": (CUMBERLAND, 30), + "E14000003": (WESTMORLAND, 30), + "E14000004": (WESTMORLAND, 30), + "E14000005": (WESTMORLAND, 30), +} +THIN = "E14000006" +NEGATIVE = "E14000007" +LAD23 = UKLocalAuthorityMetadata(boundary_version=UKLocalAuthorityBoundaryVersion.LAD23) + + +def _uk_disclosure_columns(*, with_source_ids: bool = True) -> dict[str, list]: + """Return the households' columns, with baseline and reformed incomes.""" + rng = np.random.default_rng(7) + rows: list[tuple[str, str, float, float]] = [] + for seat, (authority, households) in SEATS.items(): + for _ in range(households): + income = float(rng.uniform(20_000, 60_000)) + rows.append((seat, authority, income, -float(rng.uniform(50, 150)))) + rows.extend((THIN, WESTMORLAND, 30_000.0, -100.0) for _ in range(3)) + for index in range(30): + income = -40_000.0 if index < 20 else 30_000.0 + rows.append((NEGATIVE, WESTMORLAND, income, -float(rng.uniform(50, 150)))) + columns: dict[str, list[object]] = { + "household_id": list(range(len(rows))), + "constituency_code_oa": [row[0] for row in rows], + "la_code_oa": [row[1] for row in rows], + "region_code_oa": [NORTH_WEST] * len(rows), + "household_weight": [float(rng.uniform(50, 150)) for _ in rows], + } + if with_source_ids: + columns["source_household_id"] = list(range(len(rows))) + columns["baseline_income"] = [row[2] for row in rows] + columns["reformed_income"] = [row[2] + row[3] for row in rows] + return columns + + +def _uk_disclosure_simulations( + *, with_source_ids: bool = True +) -> tuple[SimpleNamespace, SimpleNamespace]: + columns = _uk_disclosure_columns(with_source_ids=with_source_ids) + baseline_income = columns.pop("baseline_income") + reformed_income = columns.pop("reformed_income") + baseline = weighted_entity_frame( + {**columns, "household_net_income": baseline_income} + ) + reform = weighted_entity_frame({**columns, "household_net_income": reformed_income}) + + def simulation(household): + return SimpleNamespace( + output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) + ) + + return simulation(baseline), simulation(reform) + + +@pytest.fixture(scope="module") +def uk_disclosure() -> simulation_output_geographic.UKGeographyDisclosureOutput: + baseline, reform = _uk_disclosure_simulations() + output = simulation_output_geographic.compute_uk_geography_disclosure_output( + "uk", baseline, reform, year=2026 + ) + assert output is not None + return output + + +@pytest.fixture +def constituency_lookup(tmp_path, monkeypatch) -> None: + lookup = tmp_path / "constituencies_2024.csv" + codes = [*SEATS, THIN, NEGATIVE] + lookup.write_text( + "code,name,x,y\n" + + "".join( + f"{code},Seat {index},{index},{-index}\n" + for index, code in enumerate(codes) + ) + ) + monkeypatch.setattr( + simulation_output_geographic, + "_required_uk_geography_lookup_csv_path", + lambda spec: str(lookup), + ) + + +def test_uk_disclosure_withholds_thin_and_undefined_seats(uk_disclosure) -> None: + seats = uk_disclosure.level("constituency") + + assert THIN not in seats + assert NEGATIVE not in seats + # Complementary suppression: with THIN withheld, at least one more seat + # of its region and its local authority is withheld too. + assert len(seats) < len(SEATS) + assert set(seats) < set(SEATS) + + +def test_uk_disclosure_publishes_rounded_changes_with_intervals( + uk_disclosure, +) -> None: + seats = uk_disclosure.level("constituency") + + assert seats + for change in seats.values(): + assert change.tier in {"reported", "approximate"} + assert change.average_lower <= change.average <= change.average_upper + assert change.relative_lower <= change.relative <= change.relative_upper + # Every household loses between 50 and 150 pounds. + assert -150 <= change.average <= -50 + assert change.relative < 0 + assert change.population > 0 + + +def test_uk_disclosure_population_is_the_weighted_household_count( + uk_disclosure, +) -> None: + columns = _uk_disclosure_columns() + expected: dict[str, float] = {} + for code, weight in zip( + columns["constituency_code_oa"], columns["household_weight"], strict=True + ): + expected[code] = expected.get(code, 0.0) + weight + + for code, change in uk_disclosure.level("constituency").items(): + assert change.population == pytest.approx(expected[code]) + + +def test_uk_constituency_output_lists_only_published_seats( + uk_disclosure, constituency_lookup +) -> None: + result = simulation_output_geographic.build_uk_constituency_impact( + "uk", uk_disclosure + ) + + assert result is not None + records = [record.model_dump(mode="json") for record in result.root] + published = uk_disclosure.level("constituency") + assert [record["constituency_code"] for record in records] == [ + code for code in SEATS if code in published + ] + for record in records: + change = published[record["constituency_code"]] + assert record["constituency_name"].startswith("Seat ") + assert record["average_household_income_change"] == change.average + assert record["relative_household_income_change"] == change.relative + assert record["disclosure_tier"] == change.tier + assert record["average_household_income_change_lower"] == (change.average_lower) + assert record["relative_household_income_change_upper"] == ( + change.relative_upper + ) + + +def test_uk_local_authority_output_uses_detected_boundary_metadata( + uk_disclosure, monkeypatch ) -> None: monkeypatch.setattr( simulation_output_geographic, @@ -158,67 +305,115 @@ def test_local_authority_output_uses_detected_boundary_metadata( ) result = simulation_output_geographic.build_uk_local_authority_impact( - "uk", - _uk_simulation(code, 100.0), - _uk_simulation(code, 110.0), - uk_local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=boundary_version - ), + "uk", uk_disclosure, uk_local_authority_metadata=LAD23 ) assert result is not None - record = result.root[0].model_dump(mode="python") - assert record["local_authority_code"] == code - assert record["local_authority_name"] == expected_name - assert isinstance(record["x"], int) - assert isinstance(record["y"], int) - assert record["average_household_income_change"] == 10.0 + records = { + record.model_dump(mode="json")["local_authority_code"]: record.model_dump( + mode="json" + ) + for record in result.root + } + assert set(records) == set(uk_disclosure.level("local_authority")) + assert records + for code, record in records.items(): + assert record["local_authority_name"] in { + "Cumberland", + "Westmorland and Furness", + } + assert isinstance(record["x"], int) + assert isinstance(record["y"], int) + assert record["average_household_income_change"] < 0 + assert record["population"] == ( + uk_disclosure.level("local_authority")[code].population + ) -def test_local_authority_output_rejects_code_outside_boundary_version() -> None: +def test_uk_local_authority_output_rejects_code_outside_boundary_version( + uk_disclosure, +) -> None: with pytest.raises( ValueError, match="not part of the detected LAD22 boundary version" ): simulation_output_geographic.build_uk_local_authority_impact( "uk", - _uk_simulation("E06000063", 100.0), - _uk_simulation("E06000063", 110.0), + uk_disclosure, uk_local_authority_metadata=UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 ), ) -def test_local_authority_output_requires_detected_boundary_metadata( - monkeypatch, +def test_uk_local_authority_output_requires_detected_boundary_metadata( + uk_disclosure, ) -> None: - monkeypatch.setattr( - simulation_output_geographic, - "_required_uk_geography_lookup_csv_path", - lambda _: pytest.fail("local-authority output must not resolve a GCS lookup"), - ) - with pytest.raises(ValueError, match="boundary metadata is required"): simulation_output_geographic.build_uk_local_authority_impact( - "uk", - _uk_simulation("E06000063", 100.0), - _uk_simulation("E06000063", 110.0), + "uk", uk_disclosure + ) + + +def test_uk_outputs_require_the_disclosure_output() -> None: + with pytest.raises(ValueError, match="disclosure-controlled"): + simulation_output_geographic.build_uk_constituency_impact("uk", None) + with pytest.raises(ValueError, match="disclosure-controlled"): + simulation_output_geographic.build_uk_local_authority_impact( + "uk", None, uk_local_authority_metadata=LAD23 + ) + + +def test_uk_outputs_are_none_for_other_countries() -> None: + baseline, reform = _uk_disclosure_simulations() + + assert ( + simulation_output_geographic.compute_uk_geography_disclosure_output( + "us", baseline, reform, year=2026 + ) + is None + ) + assert simulation_output_geographic.build_uk_constituency_impact("us", None) is None + assert ( + simulation_output_geographic.build_uk_local_authority_impact("us", None) is None + ) + + +def test_uk_outputs_list_no_area_when_every_area_is_withheld( + constituency_lookup, +) -> None: + empty = simulation_output_geographic.UKGeographyDisclosureOutput() + + constituencies = simulation_output_geographic.build_uk_constituency_impact( + "uk", empty + ) + authorities = simulation_output_geographic.build_uk_local_authority_impact( + "uk", empty, uk_local_authority_metadata=LAD23 + ) + + assert constituencies is not None + assert constituencies.root == [] + assert authorities is not None + assert authorities.root == [] + + +def test_uk_disclosure_requires_survey_household_ids() -> None: + baseline, reform = _uk_disclosure_simulations(with_source_ids=False) + + with pytest.raises(ValueError, match="need source_household_id"): + simulation_output_geographic.compute_uk_geography_disclosure_output( + "uk", baseline, reform, year=2026 ) 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", - lambda _: pytest.fail("a dataset without area codes needs no lookup"), - ) household = weighted_entity_frame( { - "region": ["LONDON"], - "household_net_income": [100.0], - "household_weight": [2.0], + "region": ["LONDON", "LONDON"], + "source_household_id": [1, 2], + "household_net_income": [100.0, 200.0], + "household_weight": [2.0, 3.0], } ) national = SimpleNamespace( @@ -228,15 +423,108 @@ def test_uk_geographic_outputs_are_never_dropped_for_missing_area_codes( with pytest.raises( ValueError, match="constituency breakdowns need constituency_code_oa" ): - simulation_output_geographic.build_uk_constituency_impact( - "uk", national, national + simulation_output_geographic.compute_uk_geography_disclosure_output( + "uk", national, national, year=2026 ) - with pytest.raises(ValueError, match="local-authority breakdowns need la_code_oa"): - simulation_output_geographic.build_uk_local_authority_impact( - "uk", - national, - national, - uk_local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ), + + +def _finding(**fields: object) -> SimpleNamespace: + return SimpleNamespace( + **{ + "blocking": True, + "code": "recoverable_remainder", + "family": "change", + "reform": "reform", + "year": 2026, + "level": "constituency", + "suggested": None, + "group": None, + "group_level": None, + **fields, + } + ) + + +class _FakeRelease: + def __init__(self) -> None: + self.withheld: list[tuple[str, str, str]] = [] + + def withhold(self, *, level, area_code, reform, year, outcome) -> None: + self.withheld.append((level, area_code, outcome)) + + +class _FakeImpact: + """A disclosure result whose findings clear once the named cells are withheld.""" + + def __init__(self, findings, *, clears_after: set[tuple[str, str]]) -> None: + self.release = _FakeRelease() + self.findings = findings + self._clears_after = clears_after + self.public = None + + def refresh(self): + withheld = { + (level, code) + for level, code, outcome in self.release.withheld + if outcome == "average_change" + } + if self._clears_after <= withheld: + self.findings = [] + self.public = SimpleNamespace(rows=[]) + return self.public + + +AREAS = {"constituency": {"S1": 1.0, "S2": 1.0}, "local_authority": {"L1": 2.0}} + + +def test_findings_are_resolved_by_withholding_published_cells() -> None: + impact = _FakeImpact( + [ + _finding(suggested="S2"), + # A containment finding without a suggestion names its group. + _finding(level="constituency", group="L1", group_level="local_authority"), + ], + clears_after={("constituency", "S2"), ("local_authority", "L1")}, + ) + + public = simulation_output_geographic._resolved_public_table( + impact, AREAS, year=2026 + ) + + assert public is not None + change_cells = [ + cell for cell in impact.release.withheld if cell[2] == "average_change" + ] + assert sorted(change_cells) == [ + ("constituency", "S2", "average_change"), + ("local_authority", "L1", "average_change"), + ] + # The shares are never published, so every area's are withheld up front. + share_cells = { + cell for cell in impact.release.withheld if cell[2] != "average_change" + } + assert share_cells == { + (level, code, outcome) + for level, codes in AREAS.items() + for code in codes + for outcome in ("share_gaining", "share_losing") + } + + +def test_reference_cells_are_never_withheld_and_publish_nothing(caplog) -> None: + # A finding that only a region cell could resolve: regions are public + # elsewhere, so the breakdowns publish no area at all. + impact = _FakeImpact( + [_finding(code="recoverable_parent", level="region", group="E12000002")], + clears_after={("region", "E12000002")}, + ) + + with caplog.at_level("WARNING"): + public = simulation_output_geographic._resolved_public_table( + impact, AREAS, year=2026 ) + + assert public is None + assert all(level != "region" for level, _, _ in impact.release.withheld) + assert "recoverable_parent" in caplog.text + assert "E12000002" not in caplog.text diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py index 3f7498516..76a22461f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py @@ -161,6 +161,8 @@ def test_uk_plan_uses_uk_defaults_without_us_budget_variables(monkeypatch) -> No assert _dataset_variables(plan, "household") == { "constituency_code_oa", "la_code_oa", + "region_code_oa", + "source_household_id", } assert _dataset_variables(plan, "household").issubset(_variables(plan, "household")) @@ -183,7 +185,14 @@ def test_uk_plan_copies_area_codes_unless_a_routed_national_report_lacks_them( plan = resolve_report_output_plan(_report("uk", region=region)) assert _dataset_variables(plan, "household") == ( - {"constituency_code_oa", "la_code_oa"} if planned else set() + { + "constituency_code_oa", + "la_code_oa", + "region_code_oa", + "source_household_id", + } + if planned + else set() ) @@ -291,7 +300,10 @@ def test_apply_and_validate_output_plan_reject_missing_materialized_columns() -> validate_output_frames(frames, plan) -@pytest.mark.parametrize("missing", ["constituency_code_oa", "la_code_oa"]) +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) def test_uk_output_validation_requires_geographic_dataset_columns( monkeypatch, missing: str, diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 76d3f626b..c1a6b7a35 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1033,6 +1033,65 @@ }, "ConstituencyImpactRecord": { "properties": { + "disclosure_tier": { + "anyOf": [ + { + "type": "string", + "enum": [ + "reported", + "approximate" + ] + }, + { + "type": "null" + } + ], + "title": "Disclosure Tier" + }, + "average_household_income_change_lower": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Average Household Income Change Lower" + }, + "average_household_income_change_upper": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Average Household Income Change Upper" + }, + "relative_household_income_change_lower": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Relative Household Income Change Lower" + }, + "relative_household_income_change_upper": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Relative Household Income Change Upper" + }, "constituency_code": { "type": "string", "title": "Constituency Code" @@ -1088,7 +1147,7 @@ "population" ], "title": "ConstituencyImpactRecord", - "description": "Reform impact on household net income in one UK parliamentary constituency.\n\n``average_household_income_change`` is the average change per household and\n``relative_household_income_change`` the relative change in the area's total\nhousehold net income. ``x`` and ``y`` are integer map-rendering coordinates." + "description": "Reform impact on household net income in one UK parliamentary constituency.\n\n``average_household_income_change`` is the average change per household and\n``relative_household_income_change`` the relative change in the area's total\nhousehold net income, both rounded to the order of their standard error.\n``population`` is the weighted number of households. ``x`` and ``y`` are\ninteger map-rendering coordinates." }, "DecileOutput": { "properties": { @@ -1180,7 +1239,7 @@ }, "type": "array", "title": "GeographicImpactOutput", - "description": "UK constituency or local-authority impact records." + "description": "UK constituency or local-authority impact records.\n\nAreas whose income change disclosure control withholds are left out." }, "HTTPValidationError": { "properties": { @@ -1682,6 +1741,65 @@ }, "LocalAuthorityImpactRecord": { "properties": { + "disclosure_tier": { + "anyOf": [ + { + "type": "string", + "enum": [ + "reported", + "approximate" + ] + }, + { + "type": "null" + } + ], + "title": "Disclosure Tier" + }, + "average_household_income_change_lower": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Average Household Income Change Lower" + }, + "average_household_income_change_upper": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Average Household Income Change Upper" + }, + "relative_household_income_change_lower": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Relative Household Income Change Lower" + }, + "relative_household_income_change_upper": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Relative Household Income Change Upper" + }, "local_authority_code": { "type": "string", "title": "Local Authority Code" diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt index 77651b94b..35fed4bd6 100644 --- a/rules/baselines/pyright-executor.txt +++ b/rules/baselines/pyright-executor.txt @@ -56,14 +56,8 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "model_dump" for class "object" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "region_type" of type "RegionType" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAssignmentType :: Type "dict[str, Any] | Policy | None" is not assignable to declared type "dict[str, Any] | None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" diff --git a/rules/baselines/sim-microdata-weight-stripping.txt b/rules/baselines/sim-microdata-weight-stripping.txt index 044b1e724..b396d381b 100644 --- a/rules/baselines/sim-microdata-weight-stripping.txt +++ b/rules/baselines/sim-microdata-weight-stripping.txt @@ -4,10 +4,3 @@ # ::::: a read of raw values from a weighted # frame in simulation_output_common.py for something other than counting # entities. Reads anywhere else fail outright. -# -# The UK local-authority breakdown hands policyengine's longwise geography -# impact a plain DataFrame of the household frame, which weights by its -# household_weight column itself. It goes once the breakdown hands -# policyengine the simulations, as the constituency and US congressional -# district breakdowns already do. -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py::_unweighted_household_frame::pd.DataFrame(simulation.output_dataset.data.household) diff --git a/rules/registry.toml b/rules/registry.toml index e2ac39930..b7a59bfb3 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -791,7 +791,7 @@ project = "projects/policyengine-simulation-executor" test = "tests/test_microdata_weight_boundaries.py::test_rule_proof_weight_stripping_confined" [rule.baseline] file = "rules/baselines/sim-microdata-weight-stripping.txt" -count = 1 +count = 0 [[rule]] id = "sim.microdata.builder-tests-use-weighted-frames" @@ -1143,7 +1143,7 @@ job = "pyright" invocation = "python -m scripts.rules pyright" [rule.baseline] file = "rules/baselines/pyright-executor.txt" -count = 77 +count = 71 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good" From 897e4b56ea203f8108190003b486f58ec1ed585b Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 20 Oct 2026 21:08:22 -0400 Subject: [PATCH 07/31] Withhold the inner areas that sum to a withheld UK area A containment finding without a suggested cell names only the outer area. When that area is withheld and the public areas of the other level inside it add up to it (Orkney and Shetland is Orkney plus Shetland), the loop withheld the outer cell again, made no progress, and published no area at all. It now withholds the inner areas that lie wholly inside the outer one, read from the household area pairs as the release reads them. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 8fb0b788207128f7ab276a99aca0b8f981e1a849) --- .../simulation_output_geographic.py | 93 ++++++++++++++----- .../test_simulation_output_geographic.py | 46 ++++++++- 2 files changed, 112 insertions(+), 27 deletions(-) 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 9c978ad37..19fcdebc9 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 @@ -21,6 +21,7 @@ FAMILY_CHANGE, FAMILY_GAINING, FAMILY_LOSING, + RECOVERABLE_PARENT, RELATIVE_CHANGE, REPORTED, UNCHECKED_OVERLAP, @@ -382,6 +383,49 @@ def _uk_area_household_counts(household: pd.DataFrame, column: str) -> dict[str, return {code: count for code, count in by_code.items() if count > 0} +# For each pair of overlapping published levels (inner, outer): the inner +# areas that lie wholly inside each outer area, as the release reads them. +_UKAreasInside = Mapping[tuple[str, str], Mapping[str, tuple[str, ...]]] + + +def _uk_areas_inside(household: pd.DataFrame) -> _UKAreasInside: + """Return, for each level, the areas of the other level wholly inside each area. + + An inner area lies wholly inside an outer one when every weighted + household of it has that outer area, as the disclosure release decides. + A containment finding without a suggested cell names only the outer area, + and these are the cells that resolve it. + """ + (first, first_column), (second, second_column) = UK_DISCLOSURE_AREA_COLUMNS.items() + counts = household.groupby([first_column, second_column], dropna=False)[ + "household_net_income" + ].count() + pairs: set[tuple[str, str]] = set() + for (first_raw, second_raw), count in counts.items(): + if float(cast("float", count)) > 0: + pairs.add( + (_normalised_area_code(first_raw), _normalised_area_code(second_raw)) + ) + inside: dict[tuple[str, str], dict[str, tuple[str, ...]]] = {} + for inner_level, outer_level, inner_at, outer_at in ( + (first, second, 0, 1), + (second, first, 1, 0), + ): + outers: dict[str, set[str]] = {} + for pair in pairs: + if pair[inner_at] != UNASSIGNED: + outers.setdefault(pair[inner_at], set()).add(pair[outer_at]) + by_outer: dict[str, list[str]] = {} + for inner, met in outers.items(): + outer = next(iter(met)) if len(met) == 1 else UNASSIGNED + if outer != UNASSIGNED: + by_outer.setdefault(outer, []).append(inner) + inside[(inner_level, outer_level)] = { + outer: tuple(sorted(inners)) for outer, inners in by_outer.items() + } + return inside + + def _withhold_unpublished_families( impact: UKGeographyDisclosure, area_codes: Mapping[str, Mapping[str, float]], @@ -407,14 +451,17 @@ def _withhold_unpublished_families( def _cells_resolving_findings( - impact: UKGeographyDisclosure, published_levels: tuple[str, ...] + impact: UKGeographyDisclosure, + published_levels: tuple[str, ...], + inside: _UKAreasInside, ) -> set[tuple[str, str, str, str, int]]: """Return the published cells to withhold for the blocking findings. - A finding names a cell to withhold, or else the group it is about. Only - cells of published levels are withheld: regions and the nation are - references whose numbers are public elsewhere, so withholding them here - would only hide the risk. + A finding names a cell to withhold, or else the group it is about. When a + withheld area is the sum of the public areas of the other level inside + it, those areas are withheld. Only cells of published levels are + withheld: regions and the nation are references whose numbers are public + elsewhere, so withholding them here would only hide the risk. """ cells: set[tuple[str, str, str, str, int]] = set() for finding in impact.findings or []: @@ -425,26 +472,19 @@ def _cells_resolving_findings( outcome = _FAMILY_OUTCOME.get(finding.family) if outcome is None: continue + key = (outcome, finding.reform, finding.year) if finding.suggested and finding.level in published_levels: - cells.add( - ( - finding.level, - finding.suggested, - outcome, - finding.reform, - finding.year, - ) - ) - elif finding.group and finding.group_level in published_levels: - cells.add( - ( - finding.group_level, - finding.group, - outcome, - finding.reform, - finding.year, - ) + cells.add((finding.level, finding.suggested, *key)) + elif not finding.group or finding.group_level not in published_levels: + continue + elif finding.code == RECOVERABLE_PARENT and finding.level in published_levels: + members = inside.get((finding.level, finding.group_level), {}) + cells.update( + (finding.level, member, *key) + for member in members.get(finding.group, ()) ) + else: + cells.add((finding.group_level, finding.group, *key)) return cells @@ -453,6 +493,7 @@ def _resolved_public_table( area_codes: Mapping[str, Mapping[str, float]], *, year: int, + inside: _UKAreasInside, ) -> PublicTable | None: """Withhold cells until no finding blocks, and return the public table. @@ -466,7 +507,7 @@ def _resolved_public_table( for _ in range(UK_DISCLOSURE_MAX_ROUNDS): if public is not None: return public - cells = _cells_resolving_findings(impact, published_levels) - withheld + cells = _cells_resolving_findings(impact, published_levels, inside) - withheld if not cells: break for level, code, outcome, reform, cell_year in sorted(cells): @@ -587,7 +628,9 @@ def compute_uk_geography_disclosure_output( # other differences between them stay unchecked. acknowledge=(UNCHECKED_OVERLAP,), ) - public = _resolved_public_table(impact, area_codes, year=year) + public = _resolved_public_table( + impact, area_codes, year=year, inside=_uk_areas_inside(household) + ) if public is None: return UKGeographyDisclosureOutput() return UKGeographyDisclosureOutput( 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 e3d8320ff..84636b533 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -488,7 +488,7 @@ def test_findings_are_resolved_by_withholding_published_cells() -> None: ) public = simulation_output_geographic._resolved_public_table( - impact, AREAS, year=2026 + impact, AREAS, year=2026, inside={} ) assert public is not None @@ -521,10 +521,52 @@ def test_reference_cells_are_never_withheld_and_publish_nothing(caplog) -> None: with caplog.at_level("WARNING"): public = simulation_output_geographic._resolved_public_table( - impact, AREAS, year=2026 + impact, AREAS, year=2026, inside={} ) assert public is None assert all(level != "region" for level, _, _ in impact.release.withheld) assert "recoverable_parent" in caplog.text assert "E12000002" not in caplog.text + + +def test_uk_areas_inside_follow_the_household_area_pairs() -> None: + baseline, _ = _uk_disclosure_simulations() + household = baseline.output_dataset.data.household + + inside = simulation_output_geographic._uk_areas_inside(household) + + # Every seat lies in one local authority; no authority lies in one seat. + assert inside[("constituency", "local_authority")] == { + CUMBERLAND: ("E14000001", "E14000002"), + WESTMORLAND: ("E14000003", "E14000004", "E14000005", THIN, NEGATIVE), + } + assert inside[("local_authority", "constituency")] == {} + + +def test_a_withheld_area_summed_by_public_areas_inside_withholds_them() -> None: + # Orkney and Shetland is the sum of two local authorities. If the seat is + # withheld, the two authorities inside it must be withheld too. + inside = {("local_authority", "constituency"): {"S1": ("L1", "L2")}} + impact = _FakeImpact( + [ + _finding( + code="recoverable_parent", + level="local_authority", + group="S1", + group_level="constituency", + ) + ], + clears_after={("local_authority", "L1"), ("local_authority", "L2")}, + ) + + public = simulation_output_geographic._resolved_public_table( + impact, AREAS, year=2026, inside=inside + ) + + assert public is not None + assert { + (level, code) + for level, code, outcome in impact.release.withheld + if outcome == "average_change" + } == {("local_authority", "L1"), ("local_authority", "L2")} From 0e132423b7d26154ce4984d832c16d80d72a015b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:28:50 +0400 Subject: [PATCH 08/31] Type the UK area pair counts that the containment helper reads The previous commit's _uk_areas_inside unpacks each key of a two-column groupby count as a pair of area codes. pandas types those keys as Hashable, so the executor's Pyright check reported a new error ("Hashable" is not iterable). The key is now checked to be a tuple before it is unpacked; a two-column groupby always gives one, so results do not change. Pyright is back to the 71 errors its baseline lists, and the executor suite passes (1804 passed). Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_geographic.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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 19fcdebc9..4ce4a49b8 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 @@ -401,7 +401,12 @@ def _uk_areas_inside(household: pd.DataFrame) -> _UKAreasInside: "household_net_income" ].count() pairs: set[tuple[str, str]] = set() - for (first_raw, second_raw), count in counts.items(): + for key, count in counts.items(): + # Grouping by two columns keys each count by a pair of area codes. + if not isinstance(key, tuple): + msg = f"UK area pair counts are keyed by {key!r}, not a pair of codes" + raise TypeError(msg) + first_raw, second_raw = key if float(cast("float", count)) > 0: pairs.add( (_normalised_area_code(first_raw), _normalised_area_code(second_raw)) From 73dfd89522d49869a82abf31c146655f066c838a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 21 Oct 2026 18:07:19 +0000 Subject: [PATCH 09/31] chore(deps): update policyengine to 6.10.0 (cherry picked from commit 2e98847d1464263535d0d5e227729caa2d3cb680) --- .../pyproject.toml | 4 ++-- .../policyengine-simulation-executor/uv.lock | 22 +++++++++---------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/projects/policyengine-simulation-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index 007a155d4..a2147ca9f 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/projects/policyengine-simulation-executor/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "policyengine-stage12-persistence", # PolicyEngine.py's models extra is the only selector for core, country # models, and the SPM calculator. Do not duplicate those versions here. - "policyengine[models]==6.7.5", + "policyengine[models]==6.10.0", "tables>=3.10.2", "modal>=0.73.0", "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", @@ -46,7 +46,7 @@ packages = ["src/policyengine_simulation_executor"] # certified datasets and writes their receipt. [dependency-groups] modal-simulation-image = [ - "policyengine[models]==6.7.5", + "policyengine[models]==6.10.0", "fastapi>=0.115.0", "tables>=3.10.2", "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", diff --git a/projects/policyengine-simulation-executor/uv.lock b/projects/policyengine-simulation-executor/uv.lock index 494cde1e1..7453cde31 100644 --- a/projects/policyengine-simulation-executor/uv.lock +++ b/projects/policyengine-simulation-executor/uv.lock @@ -1747,7 +1747,7 @@ wheels = [ [[package]] name = "policyengine" -version = "6.7.5" +version = "6.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "h5py" }, @@ -1759,9 +1759,9 @@ dependencies = [ { name = "pydantic" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dc/db/3c18644cc814dcfc1abcea6f58df67d012712d7c6b2d9c8e01b746c18dc0/policyengine-6.7.5.tar.gz", hash = "sha256:872457de79f51dc57f33efb85274cb461d14aa42017e39aabf9d6611a673b30a", size = 1465755, upload-time = "2026-10-18T22:31:35.177Z" } +sdist = { url = "https://files.pythonhosted.org/packages/af/97/b78da2200ea502a1ee624f1dd8d17edda467d96eaaf3d004509e544b9c08/policyengine-6.10.0.tar.gz", hash = "sha256:b9e621f23162c5ce99d004fcba0d34259442244bd59951e4bd19a1dfb1cc4e35", size = 1526638, upload-time = "2026-10-21T16:41:04.984Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4f/21/8faef0956e0c9a0f4ac73d16e7aae7330088e25ed3b6f81323572c4b8440/policyengine-6.7.5-py3-none-any.whl", hash = "sha256:acb921a6cd78c5bd5692827d3830fad9225a25f7e93e5742b92b2909624f76d3", size = 508363, upload-time = "2026-10-18T22:31:33.068Z" }, + { url = "https://files.pythonhosted.org/packages/a9/77/c65f5f5bbe90d5443c608a8d35a4865a565c5d7b7de68ae3120270acf38d/policyengine-6.10.0-py3-none-any.whl", hash = "sha256:104270016f9967705d30111edaa0015c997d68a28cdd595ea949431a654b200b", size = 536520, upload-time = "2026-10-21T16:41:03.271Z" }, ] [package.optional-dependencies] @@ -1774,7 +1774,7 @@ models = [ [[package]] name = "policyengine-core" -version = "3.33.0" +version = "3.33.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dpath" }, @@ -1794,9 +1794,9 @@ dependencies = [ { name = "standard-imghdr" }, { name = "wheel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/90/d7/70936abfab724399ea3d8db04bbd5808339ee1534819bf25ab46d00c8d32/policyengine_core-3.33.0.tar.gz", hash = "sha256:5c923828dc84a7a1ff8f6f407f7a704fe067246dd277ede04c15f409804b9c19", size = 674785, upload-time = "2026-10-10T09:43:12.12Z" } +sdist = { url = "https://files.pythonhosted.org/packages/17/af/e8449e8a488468e5ce28b6b9066a2372e05bc7971ad9a441eb4e6efbb740/policyengine_core-3.33.3.tar.gz", hash = "sha256:7da78ae9cc7b1d6833dc30651a0a4775edd2d5e205a73bbafcc0b15c988c880e", size = 693203, upload-time = "2026-10-12T20:53:17.299Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/02/190e06717785ff974120a9136c6c5408a00beaa8cddd792f557a6bf974c4/policyengine_core-3.33.0-py3-none-any.whl", hash = "sha256:2e32d02a8910ab683a2b7600987904b351d058fc2bd5910a8ea6641438bcd789", size = 299495, upload-time = "2026-10-10T09:43:10.318Z" }, + { url = "https://files.pythonhosted.org/packages/c8/d6/253cc99d853e36996f264138f4f8a2c662047da3eda8e37ddf39959397fc/policyengine_core-3.33.3-py3-none-any.whl", hash = "sha256:5847d9938412095fe9848c2ce535838a7dc8fac9238d8fd9e6c9d0a0648e0888", size = 301905, upload-time = "2026-10-12T20:53:15.706Z" }, ] [[package]] @@ -1951,7 +1951,7 @@ requires-dist = [ { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.65b0,<0.66" }, { name = "opentelemetry-instrumentation-sqlalchemy", specifier = ">=0.65b0,<0.66" }, - { name = "policyengine", extras = ["models"], specifier = "==6.7.5" }, + { name = "policyengine", extras = ["models"], specifier = "==6.10.0" }, { name = "policyengine-fastapi", editable = "../../libs/policyengine-fastapi" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, { name = "policyengine-simulation-contract", extras = ["modal"], editable = "../../libs/policyengine-simulation-contract" }, @@ -1973,7 +1973,7 @@ modal-simulation-image = [ { name = "fastapi", specifier = ">=0.115.0" }, { name = "google-cloud-storage", specifier = ">=2" }, { name = "huggingface-hub", specifier = ">=1,<2" }, - { name = "policyengine", extras = ["models"], specifier = "==6.7.5" }, + { name = "policyengine", extras = ["models"], specifier = "==6.10.0" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.2,<4" }, { name = "pyarrow", specifier = ">=20,<24" }, @@ -2069,7 +2069,7 @@ provides-extras = ["test", "build"] [[package]] name = "policyengine-uk" -version = "2.130.1" +version = "2.132.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "microdf-python" }, @@ -2077,9 +2077,9 @@ dependencies = [ { name = "pydantic" }, { name = "tables" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/27/82a97f5ec344d4be34db17bd26fe35077f1fe22f15c99818514d704262e9/policyengine_uk-2.130.1.tar.gz", hash = "sha256:1ce0a924a847b4e3e8ea8ba9702a84b267a0490c84a55ef94da6ed2efd965b34", size = 2496922, upload-time = "2026-10-14T01:58:47.232Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3f/88/bc06328a739b335d4ef5f2d07fd36cc6a3a5e0f6ec1b2c774a14e77e43a3/policyengine_uk-2.132.6.tar.gz", hash = "sha256:f9ce261d9881eee32fab96341c7080cd0c86717b40785b5cdaff44b07262a771", size = 2560464, upload-time = "2026-10-15T10:08:14.263Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1a/21/88e8addba0eabd9a7d3a29fa5293e141feb405baf0432209387f9a1f467a/policyengine_uk-2.130.1-py3-none-any.whl", hash = "sha256:f8b1b8b917c9297ff66937811dbebd9a62d3b1796ce894c326e98051cac4c873", size = 3804651, upload-time = "2026-10-14T01:58:44.937Z" }, + { url = "https://files.pythonhosted.org/packages/11/f2/c003cd7a95cb499737a5b8af0d619ba4e759c8bc7e55a2bb5c7efa46abd1/policyengine_uk-2.132.6-py3-none-any.whl", hash = "sha256:cc54025f2a3bc54be4f8d024d1e6f11ed29fed75ca513b214f2d2963c5bf4354", size = 3882707, upload-time = "2026-10-15T10:08:11.572Z" }, ] [[package]] From 0ec2f4d471900b93c8a3adbfc2ac47e883316142 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:53:07 +0400 Subject: [PATCH 10/31] Require the disclosure columns on the routed area dataset before simulating #860's disclosure-controlled area outputs read region_code_oa and source_household_id as well as the two area codes. The check a routed local-area dataset gets as it loads (detect_uk_area_dataset_metadata) required only the area codes, so a dataset without the other two would simulate both pairs and fail only when the outputs were built. UK_AREA_DATASET_COLUMNS in uk_local_authority_metadata.py now lists the four household columns those outputs read. The load check requires each of them and names the dataset and the missing column, and Stage 12's UK_GEOGRAPHIC_DATASET_VARIABLES copies the same tuple, so the two cannot drift apart. Tests, written first: - test_routed_area_dataset_needs_every_column_area_outputs_read, over the four columns (replacing the two-column test). Before the change, the region_code_oa and source_household_id cases failed with "DID NOT RAISE ". - test_uk_geographic_plan_copies_the_routed_dataset_columns. Before the change it failed with AttributeError: module ... has no attribute 'UK_AREA_DATASET_COLUMNS'. The routed-national runtime test's local-area fixture gains the two columns. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/output_planning.py | 11 +++------- .../uk_local_authority_metadata.py | 21 ++++++++++++++----- .../tests/test_simulation_output_builder.py | 2 ++ .../tests/test_stage12_output_planning.py | 14 +++++++++++++ .../tests/test_uk_local_authority_metadata.py | 21 +++++++++++++------ 5 files changed, 50 insertions(+), 19 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py index eb68727c6..a1b4d32aa 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py @@ -24,6 +24,7 @@ from policyengine_simulation_observability.stages import Stage from policyengine_simulation_executor.uk_local_authority_metadata import ( + UK_AREA_DATASET_COLUMNS, uk_area_codes_required, ) @@ -32,15 +33,9 @@ from policyengine_simulation_contract.stage12_bundle import CountryId # Input-dataset columns the UK constituency and local-authority breakdowns -# read: the area codes, the region the disclosure check nests them in, and the -# survey household each row was copied from, by which it counts support. +# read, the same ones a routed local-area dataset is checked for as it loads. UK_GEOGRAPHIC_DATASET_VARIABLES: dict[str, tuple[str, ...]] = { - "household": ( - "constituency_code_oa", - "la_code_oa", - "region_code_oa", - "source_household_id", - ), + "household": UK_AREA_DATASET_COLUMNS, } 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 caa46006e..7d24f6cd9 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 @@ -282,7 +282,15 @@ def detect_uk_local_authority_boundary_version( UK_AREA_REGION_TYPES = ("constituency", "local_authority") -UK_AREA_CODE_COLUMNS = ("constituency_code_oa", "la_code_oa") +# Household columns the UK constituency and local-authority outputs read: the +# two area codes, the region the disclosure check nests them in, and the +# survey household each row was copied from, by which it counts support. +UK_AREA_DATASET_COLUMNS = ( + "constituency_code_oa", + "la_code_oa", + "region_code_oa", + "source_household_id", +) _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 " @@ -384,16 +392,19 @@ def detect_uk_area_dataset_metadata( """Check a loaded dataset that UK area regions are routed to. That dataset serves constituency and local-authority results, so unlike a - national default it must carry both area code columns, and its authority - codes must match one packaged boundary version. + national default it must carry every column those results read + (``UK_AREA_DATASET_COLUMNS``), and its authority codes must match one + packaged boundary version. Checking here, as the dataset loads, stops a + run before it simulates anything on a dataset its area outputs cannot + read. """ household_frame = _uk_household_frame(dataset) - for column in UK_AREA_CODE_COLUMNS: + for column in UK_AREA_DATASET_COLUMNS: if column not in household_frame: msg = ( f"UK dataset {dataset_name!r}, which the bundle routes " f"constituency and local-authority regions to, has no {column} " - "column" + "column, which its constituency and local-authority outputs read" ) raise ValueError(msg) return detect_uk_local_authority_boundary_version( 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 fcdcbbb8c..905e9809f 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1635,6 +1635,8 @@ def test_routed_uk_national_run_builds_area_breakdowns_on_the_local_area_dataset region=["WALES", "WALES"], constituency_code_oa=["E14001063", "W07000112"], la_code_oa=["E06000063", "W06000001"], + region_code_oa=["WALES", "WALES"], + source_household_id=[1, 2], ) run = _run_routed_uk_national(monkeypatch, local_area_dataset=local_area_dataset) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py index 76a22461f..b1dd1b800 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_output_planning.py @@ -196,6 +196,20 @@ def test_uk_plan_copies_area_codes_unless_a_routed_national_report_lacks_them( ) +def test_uk_geographic_plan_copies_the_routed_dataset_columns(monkeypatch) -> None: + from policyengine_simulation_executor import uk_local_authority_metadata + + _route_uk_area_regions(monkeypatch, routed=True) + + plan = resolve_report_output_plan(_report("uk", region="constituency/E14001063")) + + # Stage 12 copies exactly the columns a routed area dataset must carry. + household = next(item for item in plan.entities if item.entity == "household") + assert tuple(household.dataset_variables) == ( + uk_local_authority_metadata.UK_AREA_DATASET_COLUMNS + ) + + def test_cliff_variables_are_conditional() -> None: without_cliffs = resolve_report_output_plan(_report()) with_cliffs = resolve_report_output_plan(_report(include_cliffs=True)) 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 c9475751a..eabdb0ff2 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 @@ -282,10 +282,22 @@ def _area_dataset(**household: list[str]) -> SimpleNamespace: ) +# One household of a routed local-area dataset: the area codes, the region +# the disclosure check nests them in, and the survey household it came from. +_ROUTED_AREA_HOUSEHOLD = { + "constituency_code_oa": ["E14001063"], + "la_code_oa": ["E06000063"], + "region_code_oa": ["SOUTH_EAST"], + "source_household_id": ["1"], +} + + def test_routed_dataset_detection_identifies_the_authority_boundaries() -> None: dataset = _area_dataset( constituency_code_oa=["E14001063", "W07000112"], la_code_oa=["E06000063", "W06000001"], + region_code_oa=["SOUTH_EAST", "WALES"], + source_household_id=["1", "2"], ) metadata = detect_uk_area_dataset_metadata(dataset, dataset_name="local") @@ -293,14 +305,11 @@ def test_routed_dataset_detection_identifies_the_authority_boundaries() -> None: assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD23 -@pytest.mark.parametrize("missing", ["constituency_code_oa", "la_code_oa"]) -def test_routed_dataset_detection_requires_both_area_code_columns( +@pytest.mark.parametrize("missing", list(_ROUTED_AREA_HOUSEHOLD)) +def test_routed_area_dataset_needs_every_column_area_outputs_read( missing: str, ) -> None: - columns = { - "constituency_code_oa": ["E14001063"], - "la_code_oa": ["E06000063"], - } + columns = dict(_ROUTED_AREA_HOUSEHOLD) del columns[missing] with pytest.raises( From 3f031fd28efe6da4f71a945eca18ede90a8186e0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:02:49 +0400 Subject: [PATCH 11/31] Add breakdown=local_areas to the simulation request and its result model A UK report's constituency and local-authority results will come from a separate run, requested on demand, instead of from every national run. This adds the contract for that run: - SimulationRequest gains breakdown: Literal["local_areas"] | None. It is accepted only for a UK request whose region is absent, uk or country/, and never with include_cliffs=true, segmented=true or a region_group. include_cliffs=false stays accepted, because the Public API always sends it. The rules live in local_area_breakdown_error() in the new local_area_breakdown module, so the worker can apply the same ones. BudgetWindowBatchRequest does not take the field. - LocalAreaBreakdownOutput (extra="forbid") is the result of such a run: kind "local_areas", model_version, policyengine_version, dataset, data_version, and the constituency_impact and local_authority_impact lists, which may be empty. - JobStatusResponse.result is SingleYearMacroOutput | LocalAreaBreakdownOutput | None. Both models forbid extra fields, so a payload validates as only one of them. - local_area_breakdown also defines the routing manifest capability (key local_area_breakdown, version 1) and supported_local_area_breakdown_capability(), copied from job_progress. The gateway OpenAPI golden is regenerated; its diff adds the field, the result model and the result union. Tests, written first, and how they failed before the change: - contract: the request tests failed with extra_forbidden where they expect value_error or literal_error, or with AttributeError: no attribute 'breakdown'; test_local_area_breakdown.py failed to import (No module named 'policyengine_simulation_contract.local_area_breakdown'). - gateway: polling a complete breakdown result returned 500 (the result failed SingleYearMacroOutput validation); it now returns 200. - entry: forwarding a breakdown got 422; the constituency-region refusal was 422 extra_forbidden instead of value_error; the OpenAPI test failed with KeyError: 'breakdown'. - executor: the published result schema had no LocalAreaBreakdownOutput. The entry tests that national submissions carry no breakdown and that the Stage 12 adapter skips a breakdown (unsupported_options) but ignores an absent one passed before the change: the adapter already skips any field it does not know. They pin that behaviour. The budget-window rejection test passes before and after for the same reason. Suites: contract 312, gateway 262, entry 668 and executor 1808 passed; make check (ruff, Pyright in every package, after regenerating the client) passes, and the generated-client tests pass. Co-Authored-By: Claude Opus 5.5 --- .../gateway_models.py | 39 ++++- .../local_area_breakdown.py | 63 +++++++ .../macro_output.py | 24 +++ .../tests/test_gateway_models.py | 79 +++++++++ .../tests/test_local_area_breakdown.py | 163 ++++++++++++++++++ .../tests/test_app.py | 44 +++++ .../tests/test_openapi.py | 16 ++ .../tests/test_stage12_adapter.py | 33 ++++ .../tests/test_simulation_api_contracts.py | 6 +- .../tests/golden/openapi.json | 64 ++++++- .../tests/test_endpoints.py | 39 +++++ 11 files changed, 564 insertions(+), 6 deletions(-) create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py create mode 100644 libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index a56d899a9..4da0a6f9a 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -19,7 +19,13 @@ JobStage, ) from policyengine_simulation_contract.json_types import JsonObject -from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.local_area_breakdown import ( + local_area_breakdown_error, +) +from policyengine_simulation_contract.macro_output import ( + LocalAreaBreakdownOutput, + SingleYearMacroOutput, +) from policyengine_simulation_contract.spm import ( SPMCapability, SPMComparisonProvenance, @@ -153,9 +159,17 @@ def enforce_max_payload_size(cls, value: object) -> object: class SimulationRequest(GatewayRequestBase): - """Request model for simulation submission.""" + """Request model for simulation submission. + + ``breakdown: "local_areas"`` asks a UK-wide or nation run for its + constituency and local-authority impacts alone, computed on the + local-area dataset; the job's result is then a + ``LocalAreaBreakdownOutput``. Without it, the job runs the usual + comparison. + """ region: str | None = None + breakdown: Literal["local_areas"] | None = None @model_validator(mode="after") def region_xor_region_group(self) -> "SimulationRequest": @@ -165,6 +179,22 @@ def region_xor_region_group(self) -> "SimulationRequest": raise ValueError(msg) return self + @model_validator(mode="after") + def local_area_breakdown_scope(self) -> "SimulationRequest": + """Reject a breakdown outside a UK-wide or nation run of its own.""" + if self.breakdown is None: + return self + error = local_area_breakdown_error( + country=self.country, + region=self.region, + region_group=self.region_group, + include_cliffs=self.include_cliffs, + segmented=self.segmented, + ) + if error is not None: + raise ValueError(error) + return self + class PolicyEngineBundle(BaseModel): """Resolved runtime provenance returned by the gateway.""" @@ -212,7 +242,10 @@ class JobStatusResponse(BaseModel): """Response model for job status polling.""" status: Literal["running", "complete", "failed"] - result: SingleYearMacroOutput | None = None + # A breakdown request's result is a LocalAreaBreakdownOutput; every other + # request's is a SingleYearMacroOutput. Both forbid extra fields, so a + # payload validates as one of them only. + result: SingleYearMacroOutput | LocalAreaBreakdownOutput | None = None errors: list[SPMErrorDetail] | None = None error: str | None = None resolved_app_name: str | None = None diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py new file mode 100644 index 000000000..c4a99e6b5 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py @@ -0,0 +1,63 @@ +"""The worker capability that serves UK local-area breakdown requests. + +A ``SimulationRequest`` with ``breakdown: "local_areas"`` runs only the +constituency and local-authority impacts on the UK local-area dataset, and +returns a ``LocalAreaBreakdownOutput``. Worker apps built before this existed +ignore the field and would return a national result, so the routing bundle +manifest of each worker that serves it declares +``LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION`` under +``LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY``, and the gateway refuses the field for +any other worker. +""" + +from __future__ import annotations + +LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY = "local_area_breakdown" +LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION = 1 + + +def supported_local_area_breakdown_capability(value: object) -> int | None: + """Return the declared capability version when this gateway reads it.""" + if isinstance(value, bool) or not isinstance(value, int): + return None + if value != LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION: + return None + return value + + +def local_area_breakdown_error( + *, + country: str, + region: str | None, + region_group: object, + include_cliffs: object, + segmented: object, +) -> str | None: + """Return why a request cannot have a local-area breakdown, or None. + + A breakdown runs one UK-wide or nation region on the local-area dataset, + so it needs a UK request whose region is absent, ``uk`` or + ``country/``. It computes no cliff impacts and no region + segments, and takes no region group. ``include_cliffs`` and ``segmented`` + may still be sent as false: the Public API always sends + ``include_cliffs: false``. The gateway contract and the worker both apply + these rules. + """ + if country.lower() != "uk": + return "breakdown is supported only for UK requests" + if region is not None and region.lower() != "uk": + region_type, _, nation = region.partition("/") + if region_type.lower() != "country" or not nation.strip(): + return ( + "breakdown needs a UK-wide or nation region (uk or " + f"country/), not {region!r}" + ) + for flag, value in ( + ("include_cliffs", include_cliffs), + ("segmented", segmented), + ): + if value is True: + return f"breakdown cannot be combined with {flag}=true" + if region_group is not None: + return "breakdown cannot be combined with region_group" + return None diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index 5fcd2bb60..bc2f7e020 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -349,3 +349,27 @@ class SingleYearMacroOutput(MacroOutputModel): local_authority_impact: GeographicImpactOutput | None congressional_district_impact: CongressionalDistrictImpactOutput | None cliff_impact: CliffImpactOutput | None = None + + +class LocalAreaBreakdownOutput(MacroOutputModel): + """Completed result of a UK local-area breakdown run. + + A request with ``breakdown: "local_areas"`` returns this instead of a + ``SingleYearMacroOutput``. The run covers the requested UK-wide or nation + region on the local-area dataset alone, and computes only the + constituency and local-authority impacts; every other figure comes from + the national run. ``kind`` is always ``local_areas``. ``dataset`` names the + local-area dataset, ``data_version`` its release revision, and + ``model_version`` and ``policyengine_version`` the country model and the + policyengine.py release that computed it. Areas whose income change + disclosure control withholds are left out of each list, so a list may be + empty. + """ + + kind: Literal["local_areas"] + model_version: str + policyengine_version: str + dataset: str + data_version: str + constituency_impact: GeographicImpactOutput + local_authority_impact: GeographicImpactOutput diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index 8fb126fe4..a7f58ade3 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -341,6 +341,85 @@ def test_job_submit_response_creates_with_all_fields(self): ) +class TestSimulationRequestBreakdown: + """``breakdown`` asks a UK-wide or nation run for its local areas only.""" + + def test_breakdown_defaults_to_none(self): + request = SimulationRequest(country="uk", scope="macro") + + assert request.breakdown is None + assert "breakdown" not in request.model_dump(exclude_none=True) + + @pytest.mark.parametrize("region", [None, "uk", "UK", "country/wales"]) + def test_local_areas_are_accepted_for_a_uk_wide_or_nation_region(self, region): + # The Public API always sends include_cliffs: false. + request = SimulationRequest( + country="uk", + scope="macro", + region=region, + include_cliffs=False, + breakdown="local_areas", + ) + + assert request.breakdown == "local_areas" + assert request.model_dump(exclude_none=True)["breakdown"] == "local_areas" + + @pytest.mark.parametrize( + ("fields", "message"), + [ + ({"country": "us"}, "only for UK requests"), + ({"region": "constituency/E14001063"}, "'constituency/E14001063'"), + ({"region": "local_authority/E06000063"}, "'local_authority/E06000063'"), + ({"region": "wales"}, "'wales'"), + ({"region": "country/"}, "'country/'"), + ({"include_cliffs": True}, "include_cliffs"), + ({"segmented": True}, "segmented"), + ({"region_group": ["country/wales"]}, "region_group"), + ], + ) + def test_local_areas_are_rejected_outside_a_uk_wide_or_nation_run( + self, fields, message + ): + payload = { + "country": "uk", + "scope": "macro", + "breakdown": "local_areas", + **fields, + } + + with pytest.raises(ValidationError) as caught: + SimulationRequest.model_validate(payload) + + (error,) = caught.value.errors() + assert error["type"] == "value_error" + assert "breakdown" in error["msg"] + assert message in error["msg"] + + def test_an_unknown_breakdown_is_rejected(self): + with pytest.raises(ValidationError) as caught: + SimulationRequest(country="uk", scope="macro", breakdown="districts") + + (error,) = caught.value.errors() + assert error["type"] == "literal_error" + assert error["loc"] == ("breakdown",) + + def test_budget_window_requests_take_no_breakdown(self): + with pytest.raises(ValidationError) as caught: + BudgetWindowBatchRequest.model_validate( + { + "country": "uk", + "region": "uk", + "start_year": "2026", + "window_size": 1, + "breakdown": "local_areas", + } + ) + + (error,) = caught.value.errors() + assert error["type"] == "extra_forbidden" + assert error["loc"] == ("breakdown",) + + class TestJobStatusResponse: """Tests for JobStatusResponse model.""" diff --git a/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py new file mode 100644 index 000000000..7c3421594 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py @@ -0,0 +1,163 @@ +"""The local-area breakdown result and the worker capability that serves it.""" + +from __future__ import annotations + +from copy import deepcopy + +import pytest +from policyengine_simulation_contract.gateway_models import JobStatusResponse +from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + supported_local_area_breakdown_capability, +) +from policyengine_simulation_contract.macro_output import ( + LocalAreaBreakdownOutput, + SingleYearMacroOutput, +) +from pydantic import ValidationError + +CONSTITUENCY = { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "x": 56, + "y": -40, + "average_household_income_change": -150.0, + "relative_household_income_change": -0.0026, + "population": 40_530.0, + "disclosure_tier": "reported", + "average_household_income_change_lower": -170.0, + "average_household_income_change_upper": -130.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.0022, +} +LOCAL_AUTHORITY = { + "local_authority_code": "E06000063", + "local_authority_name": "Cumberland", + "x": 4, + "y": 18, + "average_household_income_change": -120.0, + "relative_household_income_change": -0.002, + "population": 120_000.0, + "disclosure_tier": "approximate", + "average_household_income_change_lower": -160.0, + "average_household_income_change_upper": -80.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.001, +} +BREAKDOWN = { + "kind": "local_areas", + "model_version": "2.132.6", + "policyengine_version": "6.10.0", + "dataset": "microcosm_uk_2024_25_dense", + "data_version": "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe", + "constituency_impact": [CONSTITUENCY], + "local_authority_impact": [LOCAL_AUTHORITY], +} +_ZERO = {"baseline": 0.0, "reform": 0.0} +_AGE_POVERTY = {"child": _ZERO, "adult": _ZERO, "senior": _ZERO, "all": _ZERO} +_GENDER_POVERTY = {"male": _ZERO, "female": _ZERO} +NATIONAL = { + "spm_provenance": None, + "model_version": "2.132.6", + "data_version": "microcosm-uk-2024-25-national-20261019T083323Z-6cba5d71", + "budget": { + "tax_revenue_impact": 0.0, + "state_tax_revenue_impact": 0.0, + "benefit_spending_impact": 0.0, + "budgetary_impact": 0.0, + "households": 1.0, + "baseline_net_income": 1.0, + }, + "detailed_budget": {}, + "decile": {"average": {}, "relative": {}}, + "inequality": { + "gini": _ZERO, + "top_10_pct_share": _ZERO, + "top_1_pct_share": _ZERO, + }, + "poverty": {"poverty": _AGE_POVERTY, "deep_poverty": _AGE_POVERTY}, + "poverty_by_gender": {"poverty": _GENDER_POVERTY, "deep_poverty": _GENDER_POVERTY}, + "poverty_by_race": None, + "intra_decile": {"deciles": {}, "all": {}}, + "wealth_decile": None, + "intra_wealth_decile": None, + "labor_supply_response": None, + "constituency_impact": None, + "local_authority_impact": None, + "congressional_district_impact": None, + "cliff_impact": None, +} + + +def test_the_capability_is_declared_under_its_own_key_at_version_one(): + assert LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY == "local_area_breakdown" + assert LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION == 1 + + +def test_the_capability_accepts_only_version_one(): + assert supported_local_area_breakdown_capability(1) == 1 + + +@pytest.mark.parametrize("declared", [None, 0, 2, "1", 1.0, True]) +def test_the_capability_rejects_every_other_declaration(declared): + assert supported_local_area_breakdown_capability(declared) is None + + +def test_a_breakdown_result_carries_both_area_levels_and_their_provenance(): + output = LocalAreaBreakdownOutput.model_validate(BREAKDOWN) + + assert output.kind == "local_areas" + assert output.dataset == "microcosm_uk_2024_25_dense" + assert output.model_dump(mode="json") == BREAKDOWN + + +def test_a_job_status_reads_a_breakdown_result_as_a_breakdown(): + response = JobStatusResponse(status="complete", result=BREAKDOWN) + + assert isinstance(response.result, LocalAreaBreakdownOutput) + assert response.model_dump(mode="json")["result"] == BREAKDOWN + + +def test_a_job_status_still_reads_a_national_result_as_national(): + response = JobStatusResponse(status="complete", result=NATIONAL) + + assert isinstance(response.result, SingleYearMacroOutput) + assert response.model_dump(mode="json")["result"] == NATIONAL + + +def test_a_breakdown_payload_is_never_a_national_result(): + with pytest.raises(ValidationError): + SingleYearMacroOutput.model_validate(BREAKDOWN) + + +def test_a_national_payload_is_never_a_breakdown(): + with pytest.raises(ValidationError): + LocalAreaBreakdownOutput.model_validate(NATIONAL) + + +def test_a_breakdown_may_publish_no_area_of_either_level(): + # Disclosure control can withhold every area of a level. + payload = {**BREAKDOWN, "constituency_impact": [], "local_authority_impact": []} + + output = LocalAreaBreakdownOutput.model_validate(payload) + + assert output.constituency_impact.root == [] + assert output.local_authority_impact.root == [] + + +@pytest.mark.parametrize("field", ["constituency_impact", "local_authority_impact"]) +def test_a_breakdown_names_both_area_levels(field): + payload = deepcopy(BREAKDOWN) + payload[field] = None + + with pytest.raises(ValidationError): + LocalAreaBreakdownOutput.model_validate(payload) + + +@pytest.mark.parametrize( + "change", [{"kind": "national"}, {"budget": NATIONAL["budget"]}] +) +def test_a_breakdown_has_its_own_kind_and_no_other_fields(change): + with pytest.raises(ValidationError): + LocalAreaBreakdownOutput.model_validate({**BREAKDOWN, **change}) diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index 267d756c6..200ea4f41 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -911,6 +911,50 @@ def test_request_validation_matches_shared_contract(client): assert result.status_code == 422 +def _uk_submission(**fields: object) -> dict: + return { + "country": "uk", + "scope": "macro", + "reform": {}, + "region": "uk", + "include_cliffs": False, + "time_period": "2026", + **fields, + } + + +def test_a_local_area_breakdown_is_forwarded_to_the_gateway(client, backend): + result = client.post( + "/simulate/economy/comparison", + json=_uk_submission(breakdown="local_areas"), + ) + + assert result.status_code == 200 + assert backend.requests[-1].json_body["breakdown"] == "local_areas" + + +def test_a_breakdown_for_one_constituency_is_refused_before_the_gateway( + client, backend +): + result = client.post( + "/simulate/economy/comparison", + json=_uk_submission(region="constituency/E14001063", breakdown="local_areas"), + ) + + assert result.status_code == 422 + (error,) = result.json()["detail"] + assert error["type"] == "value_error" + assert "breakdown" in error["msg"] + assert backend.requests == [] + + +def test_a_national_submission_carries_no_breakdown(client, backend): + result = client.post("/simulate/economy/comparison", json=_uk_submission()) + + assert result.status_code == 200 + assert "breakdown" not in backend.requests[-1].json_body + + @pytest.mark.parametrize( ("error", "expected_status", "expected_detail"), [ diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index e293b7e0c..060b6bb8c 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -106,6 +106,22 @@ def test_canonical_spm_extensions_are_public(): assert "400" in spec["paths"][path]["get"]["responses"] +def test_local_area_breakdowns_are_public(): + schemas = create_app().openapi()["components"]["schemas"] + + breakdown = schemas["SimulationRequest"]["properties"]["breakdown"] + assert breakdown["anyOf"] == [ + {"type": "string", "const": "local_areas"}, + {"type": "null"}, + ] + assert schemas["LocalAreaBreakdownOutput"]["properties"]["kind"]["const"] == ( + "local_areas" + ) + assert {"$ref": "#/components/schemas/LocalAreaBreakdownOutput"} in schemas[ + "JobStatusResponse" + ]["properties"]["result"]["anyOf"] + + def test_route_table_is_frozen(): assert methods(cast("OpenAPIDocument", create_app().openapi())) == EXPECTED_ROUTES diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 5751a9b24..46ed669f4 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -155,6 +155,39 @@ def test_adapter_returns_bounded_skip_reason(update, reason) -> None: assert result.skip_reason is reason +def _uk_payload(**fields: object) -> dict: + return { + **eligible_payload(), + "country": "uk", + "region": "uk", + "version": "2.90.2", + **fields, + } + + +def test_a_local_area_breakdown_gets_no_comparison_run() -> None: + # No Stage 12 report computes a local-area breakdown alone. + result = adapt_annual_comparison( + _uk_payload(breakdown="local_areas"), + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.report is None + assert result.skip_reason is ComparisonSkipReason.UNSUPPORTED_OPTIONS + + +def test_an_absent_breakdown_keeps_a_uk_request_eligible() -> None: + result = adapt_annual_comparison( + _uk_payload(breakdown=None), + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.skip_reason is None + assert result.report is not None + + def _expected_simulation(role: str, execution_id: str, policy: dict) -> dict: dataset_uri = "hf://policyengine/data/populace_us_2024.h5@populace_us_2024-revision" return { diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py index 76bcf9ee7..03d639d93 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py @@ -52,14 +52,18 @@ def test_openapi_publishes_the_single_year_macro_result_schema(): ) result_schema = schemas["JobStatusResponse"]["properties"]["result"] assert result_schema == { + "title": "Result", "anyOf": [ { "$ref": "#/components/schemas/SingleYearMacroOutput", }, + { + "$ref": "#/components/schemas/LocalAreaBreakdownOutput", + }, { "type": "null", }, - ] + ], } diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index c1a6b7a35..8df638f92 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1368,10 +1368,14 @@ { "$ref": "#/components/schemas/SingleYearMacroOutput" }, + { + "$ref": "#/components/schemas/LocalAreaBreakdownOutput" + }, { "type": "null" } - ] + ], + "title": "Result" }, "errors": { "anyOf": [ @@ -1739,6 +1743,50 @@ "title": "LaborSupplyResponseOutput", "description": "Behavioral labor supply response to the reform.\n\n``income_lsr`` and ``substitution_lsr`` are the total earnings changes from\nthe income and substitution effects, and ``total_change`` is their sum." }, + "LocalAreaBreakdownOutput": { + "properties": { + "kind": { + "type": "string", + "const": "local_areas", + "title": "Kind" + }, + "model_version": { + "type": "string", + "title": "Model Version" + }, + "policyengine_version": { + "type": "string", + "title": "Policyengine Version" + }, + "dataset": { + "type": "string", + "title": "Dataset" + }, + "data_version": { + "type": "string", + "title": "Data Version" + }, + "constituency_impact": { + "$ref": "#/components/schemas/GeographicImpactOutput" + }, + "local_authority_impact": { + "$ref": "#/components/schemas/GeographicImpactOutput" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "kind", + "model_version", + "policyengine_version", + "dataset", + "data_version", + "constituency_impact", + "local_authority_impact" + ], + "title": "LocalAreaBreakdownOutput", + "description": "Completed result of a UK local-area breakdown run.\n\nA request with ``breakdown: \"local_areas\"`` returns this instead of a\n``SingleYearMacroOutput``. The run covers the requested UK-wide or nation\nregion on the local-area dataset alone, and computes only the\nconstituency and local-authority impacts; every other figure comes from\nthe national run. ``kind`` is always ``local_areas``. ``dataset`` names the\nlocal-area dataset, ``data_version`` its release revision, and\n``model_version`` and ``policyengine_version`` the country model and the\npolicyengine.py release that computed it. Areas whose income change\ndisclosure control withholds are left out of each list, so a list may be\nempty." + }, "LocalAuthorityImpactRecord": { "properties": { "disclosure_tier": { @@ -2471,6 +2519,18 @@ } ], "title": "Region" + }, + "breakdown": { + "anyOf": [ + { + "type": "string", + "const": "local_areas" + }, + { + "type": "null" + } + ], + "title": "Breakdown" } }, "additionalProperties": false, @@ -2479,7 +2539,7 @@ "country" ], "title": "SimulationRequest", - "description": "Request model for simulation submission." + "description": "Request model for simulation submission.\n\n``breakdown: \"local_areas\"`` asks a UK-wide or nation run for its\nconstituency and local-authority impacts alone, computed on the\nlocal-area dataset; the job's result is then a\n``LocalAreaBreakdownOutput``. Without it, the job runs the usual\ncomparison." }, "SingleYearMacroOutput": { "properties": { diff --git a/projects/policyengine-simulation-gateway/tests/test_endpoints.py b/projects/policyengine-simulation-gateway/tests/test_endpoints.py index 1037eedfa..72024828f 100644 --- a/projects/policyengine-simulation-gateway/tests/test_endpoints.py +++ b/projects/policyengine-simulation-gateway/tests/test_endpoints.py @@ -18,6 +18,30 @@ from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" +LOCAL_AREA_BREAKDOWN_RESULT = { + "kind": "local_areas", + "model_version": "2.132.6", + "policyengine_version": "6.10.0", + "dataset": "microcosm_uk_2024_25_dense", + "data_version": "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe", + "constituency_impact": [ + { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "x": 56, + "y": -40, + "average_household_income_change": -150.0, + "relative_household_income_change": -0.0026, + "population": 40530.0, + "disclosure_tier": "reported", + "average_household_income_change_lower": -170.0, + "average_household_income_change_upper": -130.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.0022, + } + ], + "local_authority_impact": [], +} def expected_bundle( @@ -972,6 +996,21 @@ def test__given_legacy_job_without_identifier__then_poll_does_not_create_one( assert response.status_code == 200 assert OBSERVABILITY_ID_HEADER not in response.headers + def test__given_complete_local_area_breakdown__then_polling_returns_it( + self, mock_modal, client: TestClient + ): + submit_response = client.post( + "/simulate/economy/comparison", + json={"country": "uk", "scope": "macro", "region": "uk", "reform": {}}, + ) + mock_modal["func"].last_call.result = LOCAL_AREA_BREAKDOWN_RESULT + + response = client.get(f"/jobs/{submit_response.json()['job_id']}") + + assert response.status_code == 200 + assert response.json()["status"] == "complete" + assert response.json()["result"] == LOCAL_AREA_BREAKDOWN_RESULT + def test__given_running_job__then_polling_returns_202( self, mock_modal, client: TestClient ): From f14cc7babb706791f7efc2a712f2d408286554e9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:12:09 +0400 Subject: [PATCH 12/31] Run a local-area breakdown on the local-area dataset alone A request with breakdown=local_areas now takes its own path through the executor (_run_local_area_breakdown_core), which _run_simulation_impl_core branches to before loading anything: - It checks the request again with the contract's local_area_breakdown_error() (UK only; region absent, uk or country/; no include_cliffs=true, segmented=true or region_group), refuses any other breakdown value, and refuses a bundle that does not route UK area regions to a local-area dataset. Each of these raises ValueError before the country module or any dataset loads. - It reuses _uk_area_simulations to load only the routed local-area dataset (microcosm_uk_2024_25_dense on .py 6.10.0) and build one baseline and reform with the request's policies and nation filter, and _report_calculation_start for the baseline and reform stages. The national file is never loaded, and the dataset_load span carries dataset_role=uk_area. - The builder's new serialize_local_area_breakdown() computes the baseline, then the reform, enters the outputs stage, and builds the two existing area breakdowns from one disclosure release under simulation_output_build. It returns a LocalAreaBreakdownOutput with the model and policyengine.py versions, the dataset name and its revision. No national impact analysis runs, and no new stage name is needed. - _selected_dataset_revision is split so _manifest_dataset_revision(country, name) reads the local-area dataset's revision from the release manifest. - UKAreaSimulations also carries the loaded dataset and its name, and the start-of-run logging moved into _log_simulation_start to keep _run_simulation_impl_core within ruff's statement limit. Tests, written first (the routed-run helper is generalised into _run_routed_uk(monkeypatch, *, params, local_area_dataset)), and how they failed before the change: - the breakdown run test: AttributeError, no _manifest_dataset_revision; the UK-wide run loaded the national dataset and then the dense one; - the invalid-request cases: DID NOT RAISE ValueError (the US case raised SPMInputError from the national path instead); - the unrouted bundle: the national path's missing la_code_oa error, not the breakdown refusal; - the three builder tests: TypeError, unexpected keyword 'resolved_dataset_name'. The test that a breakdown needs all four routed-dataset columns, and the pin in test_baseline_artifacts.py that a non-default UK dataset gets no deterministic baseline id, passed before the change: the load check and the id rule already hold. They pin both for the breakdown path. Suites: executor 1826, contract 312, gateway 262 and entry 668 passed; Pyright stays at the executor baseline's 71 errors; ruff and the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- .../local_area_breakdown.py | 2 + .../simulation_macro_output.py | 2 + .../simulation_output_builder.py | 60 ++- .../simulation_runtime.py | 139 ++++- .../tests/test_baseline_artifacts.py | 12 + .../tests/test_simulation_output_builder.py | 479 +++++++++++++++--- 6 files changed, 624 insertions(+), 70 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py index c4a99e6b5..cee1eac49 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py @@ -12,6 +12,8 @@ from __future__ import annotations +# The one value of ``SimulationRequest.breakdown``. +LOCAL_AREAS_BREAKDOWN = "local_areas" LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY = "local_area_breakdown" LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION = 1 diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_macro_output.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_macro_output.py index cfcd2e367..f1f402586 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_macro_output.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_macro_output.py @@ -22,6 +22,7 @@ LaborSupplyHoursOutput, LaborSupplyRelativeResponse, LaborSupplyResponseOutput, + LocalAreaBreakdownOutput, LocalAuthorityImpactRecord, MacroOutputModel, MacroRootModel, @@ -55,6 +56,7 @@ "LaborSupplyHoursOutput", "LaborSupplyRelativeResponse", "LaborSupplyResponseOutput", + "LocalAreaBreakdownOutput", "LocalAuthorityImpactRecord", "MacroOutputModel", "MacroRootModel", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index ce945dabe..74390035d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -36,6 +36,7 @@ InequalityOutput, IntraDecileOutput, LaborSupplyResponseOutput, + LocalAreaBreakdownOutput, PovertyByGenderOutput, PovertyByRaceOutput, PovertyModuleOutputs, @@ -55,12 +56,15 @@ class UKAreaSimulations: A UK national or nation-level run on a dataset without area codes builds its constituency and local-authority breakdowns from this pair, which has the same policies and nation filter, and from the authority boundaries - detected on that dataset. + detected on that dataset. ``dataset`` is that loaded dataset and + ``dataset_name`` its bundle name; a local-area breakdown run reports both. """ baseline: Any reform: Any local_authority_metadata: UKLocalAuthorityMetadata + dataset: Any + dataset_name: str @dataclass @@ -73,6 +77,11 @@ class SimulationOutputBuilder: ``baseline`` and ``reform``. ``omit_uk_area_breakdowns`` leaves them out: only Stage 12 sets it, for a routed national report whose stored frames carry no area codes. + + ``serialize()`` builds a ``SingleYearMacroOutput``. A local-area breakdown + run calls ``serialize_local_area_breakdown()`` instead: its ``baseline`` + and ``reform`` are the local-area pair, ``resolved_dataset_name`` names + that dataset and ``resolved_data_version`` its revision. """ country: str @@ -82,6 +91,7 @@ class SimulationOutputBuilder: baseline: Any reform: Any resolved_data_version: str | None = None + resolved_dataset_name: str | None = None resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None @@ -162,6 +172,50 @@ def serialize(self) -> dict[str, Any]: ): return output.model_dump(mode="json") + def build_local_area_breakdown(self) -> LocalAreaBreakdownOutput: + """Build the UK constituency and local-authority impacts alone. + + No national impact analysis runs. Both simulations are computed first, + baseline then reform, so a status poll moves through those stages + before the outputs stage, and both levels then come from one + disclosure release over the pair. + """ + self.baseline.ensure() + self.reform.ensure() + self.progress.enter(JobStage.OUTPUTS) + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD)): + constituency_impact = self._build_uk_constituency_impact() + local_authority_impact = self._build_uk_local_authority_impact() + if ( + constituency_impact is None + or local_authority_impact is None + or not self.resolved_dataset_name + ): + msg = ( + "A local-area breakdown is a UK output built on a named " + "local-area dataset" + ) + raise ValueError(msg) + return LocalAreaBreakdownOutput( + kind="local_areas", + model_version=self._model_version(), + policyengine_version=self._policyengine_version(), + dataset=self.resolved_dataset_name, + data_version=self._data_version(), + constituency_impact=constituency_impact, + local_authority_impact=local_authority_impact, + ) + + def serialize_local_area_breakdown(self) -> dict[str, Any]: + """Build a local-area breakdown and dump it to JSON-compatible data.""" + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): + output = self.build_local_area_breakdown() + with ( + self._span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)), + self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_MODEL_DUMP)), + ): + return output.model_dump(mode="json") + def _build_detailed_budget(self) -> DetailedBudgetOutput: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DETAILED_BUDGET)): return simulation_output_budget.build_detailed_budget(self.analysis) @@ -357,6 +411,10 @@ def _model_version(self) -> str: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): return str(getattr(self.country_module.model, "version", "")) + def _policyengine_version(self) -> str: + """Return the policyengine.py release whose bundle this run used.""" + return get_country_release_bundle(self.country).policyengine_version + def _data_version(self) -> str: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): if self.resolved_data_version: 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 271fc194e..b9f58a10f 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 @@ -22,6 +22,10 @@ from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREAS_BREAKDOWN, + local_area_breakdown_error, +) from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( @@ -769,13 +773,19 @@ def _selected_dataset_revision(country: str, selection: DatasetSelection) -> str """ if country != "us" or selection.is_default: return None + return _manifest_dataset_revision(country, selection.name) + + +def _manifest_dataset_revision(country: str, name: str) -> str | None: + """Return the revision the release manifest records for a named dataset. + + None when the manifest has no such dataset or records no revision for it. + """ from policyengine.provenance.manifest import ( get_release_manifest, ) - reference = getattr(get_release_manifest(country), "datasets", {}).get( - selection.name - ) + reference = getattr(get_release_manifest(country), "datasets", {}).get(name) revision = getattr(reference, "revision", None) return revision if isinstance(revision, str) and revision else None @@ -951,7 +961,104 @@ def _uk_area_simulations( baseline=baseline, reform=reform, local_authority_metadata=local_authority_metadata, + dataset=dataset, + dataset_name=selection.name, + ) + + +def _require_local_area_breakdown(params: dict[str, Any], *, country: str) -> None: + """Refuse a local-area breakdown this worker cannot run, before loading data. + + The gateway contract applies the same request rules; they are checked + again here because the worker can be called with any payload. The bundle + must also route UK area regions to a local-area dataset, which every + breakdown runs on. + """ + breakdown = params.get("breakdown") + if breakdown != LOCAL_AREAS_BREAKDOWN: + msg = ( + f"Unsupported breakdown {breakdown!r}; the only breakdown is " + f"{LOCAL_AREAS_BREAKDOWN!r}" + ) + raise ValueError(msg) + region = params.get("region") + error = local_area_breakdown_error( + country=country, + region=None if region is None else str(region), + region_group=params.get("region_group"), + include_cliffs=params.get("include_cliffs"), + segmented=params.get("segmented"), + ) + if error is not None: + raise ValueError(error) + if not _uk_area_regions_routed(): + msg = ( + "This policyengine.py bundle does not route UK constituency and " + "local-authority regions to a local-area dataset, so it cannot run " + "a local-area breakdown" + ) + raise ValueError(msg) + + +def _run_local_area_breakdown_core( + simulation_params: dict[str, Any], + *, + runtime: ObservabilityRuntime, + progress: JobProgressReporter, + country: str, +) -> dict: + """Run a UK local-area breakdown on the local-area dataset alone. + + The request's UK-wide or nation region filters one baseline and reform + on the dataset the bundle routes area regions to; the national file is + never loaded. The result holds the constituency and local-authority + impacts only (``LocalAreaBreakdownOutput``). + """ + _require_local_area_breakdown(simulation_params, country=country) + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): + country_module = _country_module(country) + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): + region_resolution = _resolve_region( + country_module=country_module, + country=country, + params=simulation_params, + ) + runtime.set_context(region=region_resolution.code) + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): + policies = ( + _normalise_policy(simulation_params.get("baseline")), + _normalise_policy(simulation_params.get("reform")), + ) + + logger.info("Initialising the local-area baseline and reform simulations") + area = _uk_area_simulations( + simulation_params, + runtime=runtime, + country_module=country_module, + region=region_resolution, + policies=policies, ) + _report_calculation_start(area.baseline, progress, JobStage.BASELINE) + _report_calculation_start(area.reform, progress, JobStage.REFORM) + + logger.info("Calculating the local-area breakdown") + builder = SimulationOutputBuilder( + country=country, + simulation_params=simulation_params, + country_module=country_module, + dataset=area.dataset, + baseline=area.baseline, + reform=area.reform, + resolved_data_version=_manifest_dataset_revision(country, area.dataset_name), + resolved_dataset_name=area.dataset_name, + resolved_region_code=region_resolution.code, + runtime=runtime, + uk_local_authority_metadata=area.local_authority_metadata, + progress=progress, + ) + output = builder.serialize_local_area_breakdown() + logger.info("Local-area breakdown complete") + return output def _run_simulation_impl_core( @@ -972,13 +1079,7 @@ def _run_simulation_impl_core( if selection is not None: simulation_params["spm"] = selection - logger.info( - "Starting simulation for country=%s submission_claim_id=%s", - simulation_params.get("country", "unknown"), - getattr(telemetry, "submission_claim_id", None), - ) - if metadata: - logger.info("Received simulation metadata keys: %s", sorted(metadata)) + _log_simulation_start(simulation_params, telemetry, metadata) country = simulation_params.get("country", "us").lower() _set_runtime_attributes( @@ -988,6 +1089,10 @@ def _run_simulation_impl_core( metadata=metadata, country=country, ) + if simulation_params.get("breakdown") is not None: + return _run_local_area_breakdown_core( + simulation_params, runtime=runtime, progress=progress, country=country + ) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) @@ -1118,6 +1223,20 @@ def _run_simulation_impl_core( return output +def _log_simulation_start( + simulation_params: dict[str, Any], + telemetry: TelemetryEnvelope | None, + metadata: dict[str, Any], +) -> None: + logger.info( + "Starting simulation for country=%s submission_claim_id=%s", + simulation_params.get("country", "unknown"), + getattr(telemetry, "submission_claim_id", None), + ) + if metadata: + logger.info("Received simulation metadata keys: %s", sorted(metadata)) + + def _set_runtime_attributes( *, runtime: ObservabilityRuntime, diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 22d3247f9..a6c8bdda7 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -124,6 +124,18 @@ def test_non_macro_scope_disqualifies(self, collected, scope): def test_nondefault_dataset_disqualifies(self, collected): assert self._id(collected, dataset_is_default=False) is None + def test_a_uk_local_area_baseline_gets_no_deterministic_id(self, collected): + # A UK-wide local-area breakdown simulates the whole local-area + # dataset with no scoping, as a national run does. Its baseline must + # never load the national file's precomputed artifact. + assert ( + self._id( + collected, country="uk", region_code="uk", dataset_is_default=False + ) + is None + ) + assert collected == [] + def test_missing_region_code_disqualifies(self, collected): assert self._id(collected, region_code=None) is None 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 905e9809f..285268faa 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -33,6 +33,7 @@ RowFilterStrategy, WeightReplacementStrategy, ) +from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, @@ -1548,6 +1549,41 @@ def _uk_dataset(**household: list[object]) -> SimpleNamespace: ) +_LOCAL_AREA_REVISION = "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" +_WALES = RowFilterStrategy(variable_name="region", variable_value="WALES") +_TEST_POLICIES = { + "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, +} +_WALES_NATIONAL_PARAMS = {"country": "uk", "region": "country/wales", **_TEST_POLICIES} +_WALES_BREAKDOWN_PARAMS = { + "country": "uk", + "scope": "macro", + "region": "country/wales", + "time_period": "2026", + # The Public API always sends include_cliffs: false. + "include_cliffs": False, + "breakdown": "local_areas", + **_TEST_POLICIES, +} + + +def _routed_area_dataset(*, without: str | None = None) -> SimpleNamespace: + """Return a local-area dataset with every column its area outputs read. + + ``without`` names one column to leave out. + """ + household: dict[str, list[object]] = { + "region": ["WALES", "WALES"], + "constituency_code_oa": ["E14001063", "W07000112"], + "la_code_oa": ["E06000063", "W06000001"], + "region_code_oa": ["WALES", "WALES"], + "source_household_id": [1, 2], + } + household.pop(without or "", None) + return _uk_dataset(**household) + + @dataclass class _UKAreaRun: """A UK run through ``_run_simulation_impl_core`` with every seam faked.""" @@ -1555,28 +1591,54 @@ class _UKAreaRun: builds: list[dict] builder_calls: list[dict] runtime: _TrackingRuntime + # Each selection _load_dataset was asked for, and each country module load. + loads: list[DatasetSelection] + country_modules: list[str] + # The builder methods the run called to produce its output, and the output. + serialized: list[str] + output: object = None + + def dataset_load_attributes(self) -> list[dict]: + return [ + node["attrs"] + for node in self.runtime.roots + if node["name"] == ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD) + ] -def _run_routed_uk_national(monkeypatch, *, local_area_dataset) -> _UKAreaRun: - national_dataset = _uk_dataset(region=["WALES", "WALES"]) - nation = RowFilterStrategy(variable_name="region", variable_value="WALES") - builds: list[dict] = [] - builder_calls: list[dict] = [] - _routes_uk_area_regions(monkeypatch, routed=True) - monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "routed_uk_area_dataset_name", - lambda: _LOCAL_AREA_SELECTION.name, - ) +class _RoutedUKSeams: + """Fakes for each seam of a routed UK run, recording into one ``_UKAreaRun``.""" + + def __init__(self, run, *, local_area_dataset): + self.run = run + self.local_area_dataset = local_area_dataset + self.national_dataset = _uk_dataset(region=["WALES", "WALES"]) + + def country_module(self, country): + self.run.country_modules.append(country) + return SimpleNamespace(model=SimpleNamespace(version="2.132.6")) - def resolve_selection(_params, *, region_resolution): + def resolve_region(self, *, params, **_kwargs): + if params.get("region") == "country/wales": + return RegionResolution( + code="country/wales", dataset_reference=None, scoping_strategy=_WALES + ) + return RegionResolution(code="uk", dataset_reference=None) + + def resolve_selection(self, _params, *, region_resolution): if region_resolution.dataset_reference == _LOCAL_AREA_SELECTION.name: return _LOCAL_AREA_SELECTION return _NATIONAL_SELECTION - def build_simulation(_params, *, dataset, dataset_selection, policy, region): + def load_dataset(self, _params, *, selection, country_module, data_folder=None): + self.run.loads.append(selection) + if selection is _LOCAL_AREA_SELECTION: + return self.local_area_dataset + return self.national_dataset + + def build_simulation(self, _params, *, dataset, dataset_selection, policy, region): simulation = SimpleNamespace(on_first_ensure=lambda _callback: None) - builds.append( + self.run.builds.append( { "dataset": dataset, "selection": dataset_selection, @@ -1587,59 +1649,73 @@ def build_simulation(_params, *, dataset, dataset_selection, policy, region): ) return simulation - def load_dataset(_params, *, selection, country_module, data_folder=None): - if selection is _LOCAL_AREA_SELECTION: - return local_area_dataset - return national_dataset + def output_builder(self, **kwargs): + self.run.builder_calls.append(kwargs) + return _RecordingOutputBuilder(self.run) - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - builder_calls.append(kwargs) - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) +class _RecordingOutputBuilder: + def __init__(self, run): + self.run = run + + def serialize(self): + self.run.serialized.append("serialize") + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + def serialize_local_area_breakdown(self): + self.run.serialized.append("serialize_local_area_breakdown") + return {"kind": "local_areas"} + + +def _install_routed_uk_fakes( + monkeypatch, *, local_area_dataset, routed: bool = True +) -> _UKAreaRun: + """Fake every seam of a UK run whose national file carries no area codes.""" + run = _UKAreaRun( + builds=[], + builder_calls=[], + runtime=_TrackingRuntime(), + loads=[], + country_modules=[], + serialized=[], + ) + seams = _RoutedUKSeams(run, local_area_dataset=local_area_dataset) + _routes_uk_area_regions(monkeypatch, routed=routed) + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "routed_uk_area_dataset_name", + lambda: _LOCAL_AREA_SELECTION.name, + ) for name, value in { - "_country_module": lambda _country: SimpleNamespace( - model=SimpleNamespace(version="2.132.6") - ), - "_resolve_region": lambda **_kwargs: RegionResolution( - code="country/wales", dataset_reference=None, scoping_strategy=nation - ), - "_resolve_dataset_selection": resolve_selection, - "_load_dataset": load_dataset, - "_build_simulation": build_simulation, - "SimulationOutputBuilder": FakeSimulationOutputBuilder, + "_country_module": seams.country_module, + "_resolve_region": seams.resolve_region, + "_resolve_dataset_selection": seams.resolve_selection, + "_load_dataset": seams.load_dataset, + "_build_simulation": seams.build_simulation, + "SimulationOutputBuilder": seams.output_builder, }.items(): monkeypatch.setattr( f"policyengine_simulation_executor.simulation_runtime.{name}", value ) + return run - runtime = _TrackingRuntime() - _run_simulation_impl_core( - { - "country": "uk", - "region": "country/wales", - "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, - }, - runtime=runtime, - ) - return _UKAreaRun(builds=builds, builder_calls=builder_calls, runtime=runtime) + +def _run_routed_uk(monkeypatch, *, params, local_area_dataset) -> _UKAreaRun: + run = _install_routed_uk_fakes(monkeypatch, local_area_dataset=local_area_dataset) + run.output = _run_simulation_impl_core(dict(params), runtime=run.runtime) + return run def test_routed_uk_national_run_builds_area_breakdowns_on_the_local_area_dataset( monkeypatch, ): - local_area_dataset = _uk_dataset( - region=["WALES", "WALES"], - constituency_code_oa=["E14001063", "W07000112"], - la_code_oa=["E06000063", "W06000001"], - region_code_oa=["WALES", "WALES"], - source_household_id=[1, 2], - ) + local_area_dataset = _routed_area_dataset() - run = _run_routed_uk_national(monkeypatch, local_area_dataset=local_area_dataset) + run = _run_routed_uk( + monkeypatch, + params=_WALES_NATIONAL_PARAMS, + local_area_dataset=local_area_dataset, + ) national, area = run.builds[:2], run.builds[2:] assert [build["selection"] for build in national] == [_NATIONAL_SELECTION] * 2 @@ -1663,13 +1739,10 @@ def test_routed_uk_national_run_builds_area_breakdowns_on_the_local_area_dataset local_authority_metadata=UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ), + dataset=local_area_dataset, + dataset_name=_LOCAL_AREA_SELECTION.name, ) - loads = [ - node["attrs"] - for node in run.runtime.roots - if node["name"] == ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD) - ] - assert loads == [{}, {"dataset_role": "uk_area"}] + assert run.dataset_load_attributes() == [{}, {"dataset_role": "uk_area"}] def test_routed_uk_national_run_requires_area_codes_on_the_local_area_dataset( @@ -1684,7 +1757,151 @@ def test_routed_uk_national_run_requires_area_codes_on_the_local_area_dataset( match=r"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " r"la_code_oa column", ): - _run_routed_uk_national(monkeypatch, local_area_dataset=without_authorities) + _run_routed_uk( + monkeypatch, + params=_WALES_NATIONAL_PARAMS, + local_area_dataset=without_authorities, + ) + + +def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch): + local_area_dataset = _routed_area_dataset() + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime." + "_manifest_dataset_revision", + lambda country, name: { + ("uk", _LOCAL_AREA_SELECTION.name): _LOCAL_AREA_REVISION + }[(country, name)], + ) + + run = _run_routed_uk( + monkeypatch, + params=_WALES_BREAKDOWN_PARAMS, + local_area_dataset=local_area_dataset, + ) + + # The national file is never loaded: one load, one baseline and reform. + assert run.loads == [_LOCAL_AREA_SELECTION] + assert [build["selection"] for build in run.builds] == [_LOCAL_AREA_SELECTION] * 2 + assert all(build["dataset"] is local_area_dataset for build in run.builds) + assert all(build["scoping_strategy"] is _WALES for build in run.builds) + assert [build["policy"] for build in run.builds] == [ + _normalise_policy(_TEST_POLICIES["baseline"]), + _normalise_policy(_TEST_POLICIES["reform"]), + ] + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["baseline"] is run.builds[0]["simulation"] + assert builder_kwargs["reform"] is run.builds[1]["simulation"] + assert builder_kwargs["dataset"] is local_area_dataset + assert builder_kwargs["uk_local_authority_metadata"] == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + assert builder_kwargs["resolved_dataset_name"] == _LOCAL_AREA_SELECTION.name + assert builder_kwargs["resolved_data_version"] == _LOCAL_AREA_REVISION + assert builder_kwargs["resolved_region_code"] == "country/wales" + assert builder_kwargs.get("uk_area_simulations") is None + assert run.serialized == ["serialize_local_area_breakdown"] + assert run.output == {"kind": "local_areas"} + assert run.dataset_load_attributes() == [{"dataset_role": "uk_area"}] + + +def test_a_dataset_revision_is_read_from_the_release_manifest(monkeypatch): + from policyengine_simulation_executor import simulation_runtime + + manifests = { + "uk": SimpleNamespace( + datasets={ + _LOCAL_AREA_SELECTION.name: SimpleNamespace( + revision=_LOCAL_AREA_REVISION + ), + "unrevised": SimpleNamespace(revision=""), + } + ) + } + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", manifests.__getitem__ + ) + + revision = simulation_runtime._manifest_dataset_revision + assert revision("uk", _LOCAL_AREA_SELECTION.name) == _LOCAL_AREA_REVISION + assert revision("uk", "unrevised") is None + assert revision("uk", "absent") is None + + +def test_a_uk_wide_local_area_breakdown_simulates_the_whole_dataset(monkeypatch): + run = _run_routed_uk( + monkeypatch, + params={**_WALES_BREAKDOWN_PARAMS, "region": "uk"}, + local_area_dataset=_routed_area_dataset(), + ) + + assert run.loads == [_LOCAL_AREA_SELECTION] + assert [build["scoping_strategy"] for build in run.builds] == [None, None] + assert run.builder_calls[0]["resolved_region_code"] == "uk" + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"country": "us"}, "only for UK requests"), + ({"region": "constituency/E14001063"}, "'constituency/E14001063'"), + ({"region": "local_authority/E06000063"}, "'local_authority/E06000063'"), + ({"include_cliffs": True}, "include_cliffs=true"), + ({"region": None, "region_group": ["country/wales"]}, "region_group"), + ({"breakdown": "districts"}, "'districts'"), + ], +) +def test_an_invalid_local_area_breakdown_is_refused_before_any_load( + monkeypatch, change, message +): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset() + ) + + with pytest.raises(ValueError, match=message): + _run_simulation_impl_core( + {**_WALES_BREAKDOWN_PARAMS, **change}, runtime=run.runtime + ) + + assert run.country_modules == [] + assert run.loads == [] + assert run.builds == [] + + +def test_a_local_area_breakdown_needs_a_bundle_that_routes_area_regions(monkeypatch): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset(), routed=False + ) + + with pytest.raises( + ValueError, + match="does not route UK constituency and local-authority regions", + ): + _run_simulation_impl_core(dict(_WALES_BREAKDOWN_PARAMS), runtime=run.runtime) + + assert run.loads == [] + assert run.builds == [] + + +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_a_local_area_breakdown_needs_every_column_its_outputs_read( + monkeypatch, missing +): + local_area_dataset = _routed_area_dataset(without=missing) + + with pytest.raises( + ValueError, + match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " + rf"{missing} column", + ): + _run_routed_uk( + monkeypatch, + params=_WALES_BREAKDOWN_PARAMS, + local_area_dataset=local_area_dataset, + ) _LOCAL_AREA_URI = ( @@ -2586,6 +2803,8 @@ def test_uk_area_breakdowns_use_the_computed_routed_area_pair(monkeypatch): local_authority_metadata=UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ), + dataset=object(), + dataset_name="microcosm_uk_2024_25_dense", ) disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() computed: list[tuple[str, object, object, int]] = [] @@ -2655,3 +2874,145 @@ def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): assert builder._build_uk_constituency_impact() is None assert builder._build_uk_local_authority_impact() is None + + +class _BreakdownSimulation: + """A local-area simulation that records when it is computed.""" + + def __init__(self, label, events): + self.label = label + self.events = events + + def ensure(self): + self.events.append(self.label) + + +class _RecordingProgress: + def __init__(self, events): + self.events = events + + def enter(self, stage): + self.events.append(stage) + + +def _breakdown_builder(monkeypatch, *, events, runtime=None): + """Return a UK breakdown builder whose area outputs are faked. + + The country module's impact analysis fails the test if it is called. + Each fake records into ``events``. + """ + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + + def compute(country, baseline_simulation, reform_simulation, *, year): + events.append( + ("release", country, baseline_simulation, reform_simulation, year) + ) + return disclosure + + def constituencies(country, value): + events.append(("constituency", value)) + return GeographicImpactOutput.model_validate( + [_constituency_impact_record().model_dump(mode="json")] + ) + + def local_authorities(country, value, **kwargs): + events.append(("local_authority", value, kwargs["uk_local_authority_metadata"])) + return GeographicImpactOutput.model_validate([]) + + for name, function in { + "compute_uk_geography_disclosure_output": compute, + "build_uk_constituency_impact": constituencies, + "build_uk_local_authority_impact": local_authorities, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_output_geographic.{name}", + function, + ) + + def no_impact_analysis(*_args, **_kwargs): + pytest.fail("a local-area breakdown computes no national impact analysis") + + builder = SimulationOutputBuilder( + country="uk", + simulation_params={ + "country": "uk", + "time_period": "2026", + "breakdown": "local_areas", + }, + country_module=SimpleNamespace( + model=SimpleNamespace(version="2.132.6"), + economic_impact_analysis=no_impact_analysis, + ), + dataset=SimpleNamespace(metadata={}), + baseline=_BreakdownSimulation("baseline", events), + reform=_BreakdownSimulation("reform", events), + resolved_data_version=_LOCAL_AREA_REVISION, + resolved_dataset_name="microcosm_uk_2024_25_dense", + resolved_region_code="uk", + runtime=runtime, + uk_local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), + progress=_RecordingProgress(events), + ) + return builder, disclosure + + +def test_a_local_area_breakdown_publishes_both_levels_from_one_release(monkeypatch): + events: list[object] = [] + builder, disclosure = _breakdown_builder(monkeypatch, events=events) + + output = builder.serialize_local_area_breakdown() + + releases = [event for event in events if event[0] == "release"] + assert releases == [("release", "uk", builder.baseline, builder.reform, 2026)] + assert [event for event in events if event[0] != "release"][-2:] == [ + ("constituency", disclosure), + ("local_authority", disclosure, builder.uk_local_authority_metadata), + ] + assert output == { + "kind": "local_areas", + "model_version": "2.132.6", + "policyengine_version": get_country_release_bundle("uk").policyengine_version, + "dataset": "microcosm_uk_2024_25_dense", + "data_version": _LOCAL_AREA_REVISION, + "constituency_impact": [_constituency_impact_record().model_dump(mode="json")], + "local_authority_impact": [], + } + + +def test_a_local_area_breakdown_enters_outputs_after_both_simulations(monkeypatch): + events: list[object] = [] + builder, _ = _breakdown_builder(monkeypatch, events=events) + + builder.serialize_local_area_breakdown() + + simulations_and_stages = [event for event in events if not isinstance(event, tuple)] + assert simulations_and_stages == ["baseline", "reform", JobStage.OUTPUTS] + assert events.index(JobStage.OUTPUTS) < next( + index for index, event in enumerate(events) if isinstance(event, tuple) + ) + + +def test_a_local_area_breakdown_opens_only_registered_spans(monkeypatch): + events: list[object] = [] + runtime = _TrackingRuntime() + builder, _ = _breakdown_builder(monkeypatch, events=events, runtime=runtime) + + builder.serialize_local_area_breakdown() + + registered = {stage.value for stage in ANNUAL_IMPACT_STAGES.stages} + assert set(runtime.names) <= registered + assert ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS) not in ( + runtime.names + ) + (calculation,) = runtime.roots[:1] + assert calculation["name"] == ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION) + (output_build,) = calculation["children"] + assert output_build["name"] == ANNUAL_IMPACT_STAGES.name( + Stage.SIMULATION_OUTPUT_BUILD + ) + assert { + ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY), + ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY), + } <= {child["name"] for child in output_build["children"]} From 841fc732dcb7632278a5019c9f03de272172e2a2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:17:17 +0400 Subject: [PATCH 13/31] Stop building the area pair in UK-wide and nation runs on routed bundles Until now a UK-wide or nation-level run on a bundle that routes area regions to the local-area dataset loaded that dataset as well and built a second baseline and reform there, only for the constituency and local-authority results: about 70-90 s and 12 GiB more per request, in every budget-window year and in the UK release sentinel. Those results now come from the separate breakdown=local_areas run, so: - uk_area_breakdowns_omitted(country, region) in uk_local_authority_metadata.py is True for a UK-wide or nation run on a routing bundle, and False for area regions, unrouted bundles and other countries. The runtime passes it to the builder as omit_uk_area_breakdowns, and Stage 12's aggregation uses the same helper, so the runtime and the Stage 12 shadow results agree again. - The runtime no longer calls _uk_area_simulations for national runs, and the builder loses uk_area_simulations and _uk_area_breakdown_inputs: its disclosure release reads its own baseline and reform. Such runs return null for constituency_impact and local_authority_impact. - UKAreaSimulations, now used only by the breakdown run, moves to simulation_runtime.py. Tests, written first, and how they failed before the change: - test_area_breakdowns_are_omitted_only_from_routed_uk_wide_and_nation_runs and the Stage 12 omission test's new check: AttributeError, no uk_area_breakdowns_omitted. - test_a_routed_uk_wide_or_nation_run_builds_no_area_pair (uk and country/wales): the run loaded the national and the local-area datasets instead of one. #881's test that national runs build the area pair is deleted, with its companion that checked the local-area dataset's columns during a national run (the breakdown run's tests cover that check), and so is the builder test of the routed pair. The expected builder arguments now carry omit_uk_area_breakdowns, and the test of LAD22 detection on a national file pins routing off. Suites: executor 1832 passed; Pyright stays at 71; ruff and the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_builder.py | 57 ++----- .../simulation_runtime.py | 57 ++++--- .../stage12_runtime/aggregation.py | 7 +- .../uk_local_authority_metadata.py | 14 ++ .../tests/test_simulation_output_builder.py | 153 ++++-------------- .../tests/test_stage12_runtime.py | 6 + .../tests/test_uk_local_authority_metadata.py | 25 +++ 7 files changed, 120 insertions(+), 199 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 74390035d..7dd225c90 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -49,34 +49,16 @@ from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata -@dataclass(frozen=True) -class UKAreaSimulations: - """A baseline and reform on the dataset the bundle routes UK area regions to. - - A UK national or nation-level run on a dataset without area codes builds - its constituency and local-authority breakdowns from this pair, which has - the same policies and nation filter, and from the authority boundaries - detected on that dataset. ``dataset`` is that loaded dataset and - ``dataset_name`` its bundle name; a local-area breakdown run reports both. - """ - - baseline: Any - reform: Any - local_authority_metadata: UKLocalAuthorityMetadata - dataset: Any - dataset_name: str - - @dataclass class SimulationOutputBuilder: """Assemble one year's macro output from a computed baseline and reform. Each output segment is built inside its own observability span when a - ``runtime`` is given. The UK constituency and local-authority breakdowns - come from ``uk_area_simulations`` when it is given, and otherwise from - ``baseline`` and ``reform``. ``omit_uk_area_breakdowns`` leaves them out: - only Stage 12 sets it, for a routed national report whose stored frames - carry no area codes. + ``runtime`` is given. ``omit_uk_area_breakdowns`` leaves the UK + constituency and local-authority breakdowns out, for a UK-wide or + nation-level run on a bundle that routes area regions to a local-area + dataset (``uk_area_breakdowns_omitted``); those results come from a + separate local-area breakdown run. ``serialize()`` builds a ``SingleYearMacroOutput``. A local-area breakdown run calls ``serialize_local_area_breakdown()`` instead: its ``baseline`` @@ -95,7 +77,6 @@ class SimulationOutputBuilder: resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None - uk_area_simulations: UKAreaSimulations | None = None omit_uk_area_breakdowns: bool = False progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: Any = field(default=None, init=False) @@ -333,41 +314,22 @@ def _build_congressional_district_impact( region_code=self.resolved_region_code, ) - def _uk_area_breakdown_inputs( - self, - ) -> tuple[Any, Any, UKLocalAuthorityMetadata | None]: - """Return the baseline, reform and boundaries the UK area breakdowns use. - - The impact analysis computes the run's own pair before any breakdown; - a separate area pair is computed here, on first use (``ensure()`` - returns a computed simulation from policyengine.py's cache). - """ - area = self.uk_area_simulations - if area is None: - return self.baseline, self.reform, self.uk_local_authority_metadata - area.baseline.ensure() - area.reform.ensure() - return area.baseline, area.reform, area.local_authority_metadata - def _uk_geography_disclosure_output( self, ) -> simulation_output_geographic.UKGeographyDisclosureOutput | None: """Return UK area changes under disclosure control, computed once. Constituencies and local authorities go through one release, so the - disclosure check sees both. The release reads the pair the area - breakdowns use (see ``_uk_area_breakdown_inputs``). None for other - countries. + disclosure check sees both. None for other countries. """ if self.country != "uk": return None if self._uk_geography_disclosure is None: - baseline, reform, _ = self._uk_area_breakdown_inputs() self._uk_geography_disclosure = ( simulation_output_geographic.compute_uk_geography_disclosure_output( self.country, - baseline, - reform, + self.baseline, + self.reform, year=self._simulation_year(), ) ) @@ -400,11 +362,10 @@ def _build_uk_local_authority_impact(self) -> GeographicImpactOutput | None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY)): if self.omit_uk_area_breakdowns: return None - metadata = self._uk_area_breakdown_inputs()[2] return simulation_output_geographic.build_uk_local_authority_impact( self.country, self._uk_geography_disclosure_output(), - uk_local_authority_metadata=metadata, + uk_local_authority_metadata=self.uk_local_authority_metadata, ) def _model_version(self) -> str: 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 b9f58a10f..6021e5026 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 @@ -45,7 +45,6 @@ ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, - UKAreaSimulations, ) if TYPE_CHECKING: @@ -56,6 +55,9 @@ from policyengine.core.scoping_strategy import ScopingStrategy from policyengine.provenance.manifest import ArtifactPathReference from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityMetadata, + ) from policyengine_simulation_observability.telemetry import TelemetryEnvelope from policyengine_simulation_executor.stage_reporting_simulation import ( @@ -100,6 +102,24 @@ class DatasetSelection: is_default: bool +@dataclass(frozen=True) +class UKAreaSimulations: + """A baseline and reform on the dataset the bundle routes UK area regions to. + + A local-area breakdown run builds its constituency and local-authority + impacts from this pair, with the request's policies and nation filter, + and from the authority boundaries detected on that dataset. ``dataset`` + is the loaded dataset and ``dataset_name`` its bundle name, which the + breakdown reports. + """ + + baseline: Any + reform: Any + local_authority_metadata: UKLocalAuthorityMetadata + dataset: Any + dataset_name: str + + def _normalize_credentials_blob(creds_json: str) -> str: """Return the raw JSON blob, decoding the outer escape if present. @@ -903,14 +923,12 @@ def _uk_area_simulations( region: RegionResolution, policies: tuple[dict[str, Any] | None, dict[str, Any] | None], ) -> UKAreaSimulations: - """Build the UK area breakdowns' baseline and reform on the routed dataset. - - A UK national or nation-level run may load a dataset without area codes - only when the bundle routes constituency and local-authority regions to - another dataset. Its two area breakdowns then come from this pair: the - run's baseline and reform policies and its nation filter, applied to the - routed dataset, whose area codes are checked as it loads. The national - figures still come from the run's own pair. + """Build a local-area breakdown's baseline and reform on the routed dataset. + + The bundle routes constituency and local-authority regions to a + local-area dataset. The pair applies the request's baseline and reform + policies and its UK-wide or nation filter to that dataset, whose area and + disclosure columns are checked as it loads. """ from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call detect_uk_area_dataset_metadata, @@ -1131,6 +1149,7 @@ def _run_simulation_impl_core( ) from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call detect_uk_local_authority_metadata, + uk_area_breakdowns_omitted, ) uk_local_authority_metadata = detect_uk_local_authority_metadata( @@ -1172,20 +1191,6 @@ def _run_simulation_impl_core( _report_calculation_start(baseline, progress, JobStage.BASELINE) _report_calculation_start(reform, progress, JobStage.REFORM) - # Detection returns None for a UK dataset only when it carries no area - # codes and the bundle routes area regions to another dataset. - uk_area_simulations = ( - _uk_area_simulations( - simulation_params, - runtime=runtime, - country_module=country_module, - region=region_resolution, - policies=(baseline_policy, reform_policy), - ) - if country == "uk" and uk_local_authority_metadata is None - else None - ) - logger.info("Calculating economic impact") builder = SimulationOutputBuilder( country=country, @@ -1198,7 +1203,11 @@ def _run_simulation_impl_core( resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, - uk_area_simulations=uk_area_simulations, + # A routed UK-wide or nation run computes no area results; a + # breakdown=local_areas run on the local-area dataset does. + omit_uk_area_breakdowns=uk_area_breakdowns_omitted( + country, region_resolution.code + ), progress=progress, ) output = builder.serialize() 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 2df76ed61..657414357 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 @@ -185,7 +185,7 @@ def build_aggregate_report( _runtime_policy, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call - uk_area_codes_required, + uk_area_breakdowns_omitted, ) country = report.baseline.geography.country @@ -238,9 +238,8 @@ def stand_in( resolved_data_version=report.baseline.bundle.dataset.artifact_revision, resolved_region_code=report.baseline.geography.region, uk_local_authority_metadata=uk_local_authority_metadata, - omit_uk_area_breakdowns=( - country == "uk" - and not uk_area_codes_required(report.baseline.geography.region) + omit_uk_area_breakdowns=uk_area_breakdowns_omitted( + country, report.baseline.geography.region ), ).serialize() output.update( 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 7d24f6cd9..d30b4e095 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 @@ -362,6 +362,20 @@ def uk_area_codes_required(region_code: str | None) -> bool: return not uk_area_regions_routed() +def uk_area_breakdowns_omitted(country: str, region_code: str | None) -> bool: + """Whether a run leaves out the UK constituency and local-authority breakdowns. + + A UK-wide or nation-level run on a bundle that routes area regions to a + local-area dataset computes none: its national file carries no area + codes, and those results come from a separate local-area breakdown run on + the routed dataset (``breakdown: "local_areas"``). Constituency and + local-authority runs, UK runs on a bundle that does not route area + regions, and other countries' runs are unaffected. The runtime and Stage + 12 both decide by this. + """ + return country == "uk" and not uk_area_codes_required(region_code) + + def detect_uk_local_authority_metadata( country: str, dataset: object, 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 285268faa..5e84f1dd6 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -66,7 +66,6 @@ ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, - UKAreaSimulations, ) from policyengine_simulation_executor.simulation_output_geographic import ( build_congressional_district_impact, @@ -222,7 +221,6 @@ class BuilderOptions: include_cliffs: bool | None = None resolved_region_code: str | None = "us" uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None - uk_area_simulations: UKAreaSimulations | None = None omit_uk_area_breakdowns: bool = False @@ -263,7 +261,6 @@ def economic_impact_analysis( reform=reform, resolved_region_code=options.resolved_region_code, uk_local_authority_metadata=options.uk_local_authority_metadata, - uk_area_simulations=options.uk_area_simulations, omit_uk_area_breakdowns=options.omit_uk_area_breakdowns, ) @@ -1159,7 +1156,7 @@ def serialize(self): "resolved_region_code": "us", "runtime": runtime, "uk_local_authority_metadata": None, - "uk_area_simulations": None, + "omit_uk_area_breakdowns": False, "progress": NULL_JOB_PROGRESS, } ] @@ -1471,6 +1468,9 @@ def test_a_national_run_keeps_the_release_data_version(monkeypatch): def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): + # A bundle that does not route area regions: the national file carries + # the area codes and the run builds the area breakdowns from them. + _routes_uk_area_regions(monkeypatch, routed=False) dataset = SimpleNamespace( data=SimpleNamespace( entity_data={ @@ -1522,7 +1522,7 @@ def serialize(self): assert metadata == UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 ) - assert builder_calls[0]["uk_area_simulations"] is None + assert builder_calls[0]["omit_uk_area_breakdowns"] is False def _routes_uk_area_regions(monkeypatch, *, routed: bool) -> None: @@ -1706,62 +1706,38 @@ def _run_routed_uk(monkeypatch, *, params, local_area_dataset) -> _UKAreaRun: return run -def test_routed_uk_national_run_builds_area_breakdowns_on_the_local_area_dataset( - monkeypatch, -): - local_area_dataset = _routed_area_dataset() +def _span_attributes(nodes) -> list[dict]: + """Return the attributes of every span in a tracked span tree.""" + attributes = [] + for node in nodes: + attributes.append(node["attrs"]) + attributes.extend(_span_attributes(node["children"])) + return attributes + + +@pytest.mark.parametrize("region", ["uk", "country/wales"]) +def test_a_routed_uk_wide_or_nation_run_builds_no_area_pair(monkeypatch, region): + from policyengine_simulation_executor import uk_local_authority_metadata run = _run_routed_uk( monkeypatch, - params=_WALES_NATIONAL_PARAMS, - local_area_dataset=local_area_dataset, + params={**_WALES_NATIONAL_PARAMS, "region": region}, + local_area_dataset=_routed_area_dataset(), ) - national, area = run.builds[:2], run.builds[2:] - assert [build["selection"] for build in national] == [_NATIONAL_SELECTION] * 2 - assert [build["selection"] for build in area] == [_LOCAL_AREA_SELECTION] * 2 - assert all(build["dataset"] is local_area_dataset for build in area) - # The same policies and the same nation filter as the national pair. - assert [build["policy"] for build in area] == [ - build["policy"] for build in national - ] - assert all( - build["scoping_strategy"] is national[0]["scoping_strategy"] - for build in run.builds - ) + # One dataset, the national file, and one baseline and reform on it. + assert run.loads == [_NATIONAL_SELECTION] + assert [build["selection"] for build in run.builds] == [_NATIONAL_SELECTION] * 2 (builder_kwargs,) = run.builder_calls - assert builder_kwargs["baseline"] is national[0]["simulation"] - assert builder_kwargs["reform"] is national[1]["simulation"] - assert builder_kwargs["uk_local_authority_metadata"] is None - assert builder_kwargs["uk_area_simulations"] == UKAreaSimulations( - baseline=area[0]["simulation"], - reform=area[1]["simulation"], - local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ), - dataset=local_area_dataset, - dataset_name=_LOCAL_AREA_SELECTION.name, - ) - assert run.dataset_load_attributes() == [{}, {"dataset_role": "uk_area"}] - - -def test_routed_uk_national_run_requires_area_codes_on_the_local_area_dataset( - monkeypatch, -): - without_authorities = _uk_dataset( - region=["WALES"], constituency_code_oa=["W07000112"] + assert builder_kwargs["baseline"] is run.builds[0]["simulation"] + assert builder_kwargs["reform"] is run.builds[1]["simulation"] + assert builder_kwargs["omit_uk_area_breakdowns"] is True + assert builder_kwargs["omit_uk_area_breakdowns"] is ( + uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) ) - - with pytest.raises( - ValueError, - match=r"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " - r"la_code_oa column", - ): - _run_routed_uk( - monkeypatch, - params=_WALES_NATIONAL_PARAMS, - local_area_dataset=without_authorities, - ) + assert "uk_area_simulations" not in builder_kwargs + assert run.serialized == ["serialize"] + assert {"dataset_role": "uk_area"} not in _span_attributes(run.runtime.roots) def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch): @@ -2786,75 +2762,6 @@ def build(country, value, **kwargs): assert calls == [(disclosure, metadata)] -class _AreaSimulation: - """A routed-area simulation that counts how often it is computed.""" - - def __init__(self): - self.ensured = 0 - - def ensure(self): - self.ensured += 1 - - -def test_uk_area_breakdowns_use_the_computed_routed_area_pair(monkeypatch): - area = UKAreaSimulations( - baseline=_AreaSimulation(), - reform=_AreaSimulation(), - local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ), - dataset=object(), - dataset_name="microcosm_uk_2024_25_dense", - ) - disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() - computed: list[tuple[str, object, object, int]] = [] - built: list[tuple[str, object, object]] = [] - - def compute(country, baseline_simulation, reform_simulation, *, year): - assert baseline_simulation.ensured - assert reform_simulation.ensured - computed.append((country, baseline_simulation, reform_simulation, year)) - return disclosure - - def constituencies(country, value): - built.append(("constituency", value, None)) - return GeographicImpactOutput.model_validate( - [_constituency_impact_record().model_dump(mode="json")] - ) - - def local_authorities(country, value, **kwargs): - built.append(("local_authority", value, kwargs["uk_local_authority_metadata"])) - return GeographicImpactOutput.model_validate( - [_local_authority_impact_record().model_dump(mode="json")] - ) - - for name, function in { - "compute_uk_geography_disclosure_output": compute, - "build_uk_constituency_impact": constituencies, - "build_uk_local_authority_impact": local_authorities, - }.items(): - monkeypatch.setattr( - f"policyengine_simulation_executor.simulation_output_geographic.{name}", - function, - ) - builder = _simulation_output_builder( - "uk", - object(), - object(), - BuilderOptions(resolved_region_code="uk", uk_area_simulations=area), - ) - builder.simulation_params["time_period"] = "2026" - - assert builder._build_uk_constituency_impact() is not None - assert builder._build_uk_local_authority_impact() is not None - # One disclosure release over the routed pair serves both levels. - assert computed == [("uk", area.baseline, area.reform, 2026)] - assert built == [ - ("constituency", disclosure, None), - ("local_authority", disclosure, area.local_authority_metadata), - ] - - def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): for name in ( "compute_uk_geography_disclosure_output", diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 7db073140..6c9388d2d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1884,6 +1884,12 @@ def serialize(self): ) assert observed["omit_uk_area_breakdowns"] is omitted + # The runtime omits them from exactly the same runs. + from policyengine_simulation_executor import uk_local_authority_metadata + + assert uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) is ( + omitted + ) def test_aggregate_report_wraps_builder_output_with_report_provenance( 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 eabdb0ff2..4a2418bb2 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 @@ -214,6 +214,31 @@ def test_area_regions_count_as_routed_only_when_both_have_a_dataset( _LOCAL_AREA_FILE = "microcosm_uk_2024_25_dense.h5" +@pytest.mark.parametrize( + ("country", "region", "routed", "omitted"), + [ + ("uk", "uk", True, True), + ("uk", "country/wales", True, True), + ("uk", "uk", False, False), + ("uk", "country/wales", False, False), + ("uk", "constituency/E14001063", True, False), + ("uk", "local_authority/E06000063", True, False), + ("us", "us", True, False), + ], +) +def test_area_breakdowns_are_omitted_only_from_routed_uk_wide_and_nation_runs( + monkeypatch, country: str, region: str, *, routed: bool, omitted: bool +) -> None: + from policyengine_simulation_executor import uk_local_authority_metadata + + _routes_area_regions(monkeypatch, routed=routed) + + assert ( + uk_local_authority_metadata.uk_area_breakdowns_omitted(country, region) + is omitted + ) + + def _bundle_routes(monkeypatch, region_datasets: dict[str, str]) -> None: """Install a UK manifest routing each region type to a bundle file.""" monkeypatch.setattr( From 362a904cb4274ad734cbbd903fafa0ced331157e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:20:42 +0400 Subject: [PATCH 14/31] Declare the local-area breakdown capability in the routing manifest build_bundle_manifest_metadata now adds local_area_breakdown: 1 to the routing bundle manifest of an app whose installed policyengine.py bundle routes UK area regions to a local-area dataset (uk_area_regions_routed()), and leaves the key out otherwise. The routing publisher builds the manifest only for the app it deploys, so frozen older apps, such as py6-7-5, keep manifests without the key. The next commit makes the gateway refuse breakdown requests for those apps. The value comes from the bundle manifest packaged inside the installed policyengine package; the change reads no new environment variable. Test, written first: test_bundle_manifest_declares_local_area_breakdowns_only_for_routed_area_regions. Before the change its routed case failed with KeyError: 'local_area_breakdown'; the unrouted case passed, as the key was never written. Suites: executor 1834 passed; Pyright stays at 71; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../modal/utils/update_version_registry.py | 12 +++++++ .../tests/test_update_version_registry.py | 33 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index b11acff9b..1499ba3c3 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -64,6 +64,10 @@ class BundleManifestMetadata(TypedDict): # Apps deployed before stage reporting have no key; the gateway then # reports no stage for their jobs. job_progress: NotRequired[int] + # Present only for apps whose bundle routes UK area regions to a + # local-area dataset; the gateway refuses breakdown=local_areas for any + # other app, which would ignore the field and return a national result. + local_area_breakdown: NotRequired[int] class LatestVersions(TypedDict): @@ -140,9 +144,15 @@ def build_bundle_manifest_metadata( from policyengine_simulation_contract.job_progress import ( # noqa: PLC0415 # the contract package loads only when bundle metadata is built JOB_PROGRESS_CAPABILITY_VERSION, ) + from policyengine_simulation_contract.local_area_breakdown import ( # noqa: PLC0415 # the contract package loads only when bundle metadata is built + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + ) from policyengine_simulation_executor.spm import ( runtime_spm_capability, ) + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the executor package loads only when bundle metadata is built + uk_area_regions_routed, + ) capability = runtime_spm_capability() metadata: BundleManifestMetadata = { @@ -154,6 +164,8 @@ def build_bundle_manifest_metadata( } if capability is not None: metadata["spm"] = capability.model_dump(mode="json") + if uk_area_regions_routed(): + metadata["local_area_breakdown"] = LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION return metadata diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index e4b56a658..cbe8fd3a5 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -190,6 +190,39 @@ def test_bundle_manifest_declares_job_progress_capability(monkeypatch): assert metadata["job_progress"] == JOB_PROGRESS_CAPABILITY_VERSION +@pytest.mark.usefixtures("fake_bundle_metadata") +@pytest.mark.parametrize("routed", [True, False]) +def test_bundle_manifest_declares_local_area_breakdowns_only_for_routed_area_regions( + monkeypatch, routed +): + from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + ) + from policyengine_simulation_executor import spm as executor_spm + + monkeypatch.setattr(executor_spm, "runtime_spm_capability", lambda: None) + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + metadata = registry.build_bundle_manifest_metadata( + app_name="policyengine-simulation-py6-10-0", + policyengine_version="6.10.0", + ) + + # A breakdown runs on the local-area dataset the bundle routes area + # regions to; the gateway sends breakdowns only to workers declaring it. + if routed: + assert metadata[LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY] == ( + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION + ) + else: + assert LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY not in metadata + + @pytest.mark.usefixtures("fake_bundle_metadata") def test_validate_routing_state_rejects_missing_latest_route(): state = registry.build_next_routing_state( From ef6b7eca96fe8888d42259a4c91b0b74dc4c717c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:22:29 +0400 Subject: [PATCH 15/31] Refuse breakdown requests for workers that do not declare serving them Frozen worker apps, such as py6-7-5, accept and ignore unknown payload fields, so a breakdown=local_areas request routed to one would silently come back as a national result. submit_simulation now calls _require_route_serves_breakdown(request, route) inside its existing route-resolution try block: a breakdown request whose route's bundle manifest does not declare local_area_breakdown version 1 gets a 400 before anything is spawned. The message names the policyengine.py version and app the request resolved to. Routes from the legacy version dicts have no manifest and are refused the same way. Requests without breakdown are unchanged, and the submission's observability context now records breakdown. Tests, written first (new tests/test_local_area_breakdown.py): a worker without the capability, an older pinned worker (3.9.0) and legacy-only routes all returned 200 and spawned the job before the change; they now get 400 with no spawn. The tests that a declaring worker receives the breakdown and that a national request needs no capability passed before the change, since the previous commits already forward the field. Suites: gateway 267 and entry 668 passed; gateway Pyright has no errors; ruff and the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- .../endpoints.py | 35 +++++++ .../tests/test_local_area_breakdown.py | 97 +++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py index 45f89b517..1ae74680a 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -47,6 +47,10 @@ ) from policyengine_simulation_contract.job_progress_state import get_job_progress from policyengine_simulation_contract.json_types import JsonObject +from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, + supported_local_area_breakdown_capability, +) from policyengine_simulation_contract.spm import ( SPMCapability, SPMInputError, @@ -744,6 +748,35 @@ def _route_job_progress_capability(route: RouteResolution) -> int | None: return supported_job_progress_capability(manifest.get(JOB_PROGRESS_CAPABILITY_KEY)) +def _require_route_serves_breakdown( + request: SimulationRequest, route: RouteResolution +) -> None: + """Refuse a breakdown for a route whose worker does not declare serving it. + + Workers built before local-area breakdowns ignore the field and would + return a national result, so only a bundle manifest that declares + ``LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY`` may receive one. Routes from the + legacy version dicts have no manifest and are refused too. + """ + if request.breakdown is None: + return + manifest = route.bundle_manifest + declared = ( + manifest.get(LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY) + if isinstance(manifest, dict) + else None + ) + if supported_local_area_breakdown_capability(declared) is not None: + return + version = route.policyengine_version or "an unknown version" + msg = ( + f"breakdown={request.breakdown} is not supported by the worker for " + f"policyengine.py {version} ({route.app_name}); its routing manifest " + f"declares no {LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY} capability" + ) + raise ValueError(msg) + + def _bind_observability_id( http_request: Request, runtime: ObservabilityRuntime, @@ -903,6 +936,7 @@ async def submit_simulation( country=request.country, scope=request.scope, observability_id=observability_id, + breakdown=request.breakdown, ) try: with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): @@ -911,6 +945,7 @@ async def submit_simulation( request.version, request.policyengine_version, ) + _require_route_serves_breakdown(request, route) except ValueError as e: runtime.record_exception(e, handled=True, status_code=400) raise HTTPException(status_code=400, detail=str(e)) from e diff --git a/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py b/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py new file mode 100644 index 000000000..96aed95c4 --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py @@ -0,0 +1,97 @@ +"""The gateway sends breakdown=local_areas only to workers that declare it. + +Worker apps built before local-area breakdowns existed ignore the field and +would return a national result, so the gateway refuses a breakdown for any +route whose bundle manifest does not declare the capability. +""" + +import pytest +from fastapi.testclient import TestClient +from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, +) + +ROUTING_STATE_DICT_NAME = "simulation-api-routing-state" +BREAKDOWN_REQUEST = { + "country": "uk", + "scope": "macro", + "region": "uk", + "time_period": "2026", + "reform": {}, + "include_cliffs": False, + "breakdown": "local_areas", +} + + +@pytest.fixture +def breakdown_worker(mock_modal): + """Declare the capability on the routing state's latest bundle only.""" + bundles = mock_modal["dicts"][ROUTING_STATE_DICT_NAME]["active"]["bundles"] + bundles["4.10.0"] = { + **bundles["4.10.0"], + LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY: LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + } + return mock_modal + + +def test_a_declaring_worker_gets_the_breakdown(breakdown_worker, client: TestClient): + response = client.post("/simulate/economy/comparison", json=BREAKDOWN_REQUEST) + + assert response.status_code == 200 + assert breakdown_worker["func"].last_from_name_call == ( + "policyengine-simulation-py4-10-0", + "run_simulation", + ) + assert breakdown_worker["func"].last_payload["breakdown"] == "local_areas" + + +def test_a_worker_without_the_capability_is_refused_before_spawning( + mock_modal, client: TestClient +): + response = client.post("/simulate/economy/comparison", json=BREAKDOWN_REQUEST) + + assert response.status_code == 400 + detail = response.json()["detail"] + assert "local_areas" in detail + assert "policyengine.py 4.10.0" in detail + assert mock_modal["func"].last_payload is None + + +def test_an_older_pinned_worker_is_refused_before_spawning( + breakdown_worker, client: TestClient +): + response = client.post( + "/simulate/economy/comparison", + json={**BREAKDOWN_REQUEST, "policyengine_version": "3.9.0"}, + ) + + assert response.status_code == 400 + assert "policyengine.py 3.9.0" in response.json()["detail"] + assert breakdown_worker["func"].last_payload is None + + +def test_legacy_only_routes_are_refused_before_spawning(mock_modal, client: TestClient): + # Without a routing state no bundle manifest can declare the capability. + del mock_modal["dicts"][ROUTING_STATE_DICT_NAME] + mock_modal["dicts"]["simulation-api-uk-versions"] = { + "latest": "2.66.0", + "2.66.0": "policyengine-simulation-py4-10-0", + } + + response = client.post("/simulate/economy/comparison", json=BREAKDOWN_REQUEST) + + assert response.status_code == 400 + assert "policyengine.py 4.10.0" in response.json()["detail"] + assert mock_modal["func"].last_payload is None + + +def test_a_national_request_needs_no_capability(mock_modal, client: TestClient): + national = { + key: value for key, value in BREAKDOWN_REQUEST.items() if key != "breakdown" + } + + response = client.post("/simulate/economy/comparison", json=national) + + assert response.status_code == 200 + assert "breakdown" not in mock_modal["func"].last_payload From ad9c35dac5bc280ea65af2b88a1e5e9c7c183aa9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:23:56 +0400 Subject: [PATCH 16/31] Check UK local-area breakdowns and area-free national results in beta In projects/policyengine-apis-integ/tests/simulation/test_calculate.py: - test_calculate_uk_local_area_breakdown (beta_only) submits a UK-wide breakdown=local_areas request pinned to the deployed UK model and checks the result: kind local_areas, a dataset whose name ends in _dense, at least one constituency and one local authority, and a disclosure_tier of reported or approximate on every published area. This also checks that disclosure control does not withhold every area of a level on the real local-area dataset. - test_calculate_uk_model now also checks that the national result has no constituency_impact or local_authority_impact. - assert_uk_local_area_breakdown and assert_no_uk_area_results accept the generated client's objects as well as plain mappings. Tests, written first: the local tests of the two helpers (a generated LocalAreaBreakdownOutput parsed from a job status, four incomplete breakdowns, and national results with and without area data) failed with NameError before the helpers existed, and pass now. The two beta tests need a deployment and run in the beta deploy's integration suite. The local apis-integ tests pass, except test_ping, which needs the docker-compose services and could not connect here; Pyright stays at the apis-integ baseline's 4 errors. Co-Authored-By: Claude Opus 5.5 --- .../tests/simulation/test_calculate.py | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index fbabb28ed..60de0c923 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -21,6 +21,7 @@ JobStage, JobStatusResponse, JobSubmitResponse, + LocalAreaBreakdownOutput, SimulationRequest, ) @@ -253,6 +254,45 @@ def assert_congressional_district_results( ) +_UK_AREA_RESULT_KEYS = ("constituency_impact", "local_authority_impact") +_PUBLISHED_DISCLOSURE_TIERS = {"reported", "approximate"} + + +def assert_no_uk_area_results(economy_result: object) -> None: + """Check a UK national result carries no constituency or local-authority data.""" + economy_result = _as_mapping(economy_result, label="economy result") + for key in _UK_AREA_RESULT_KEYS: + assert economy_result.get(key) is None, ( + f"Expected no {key} in a UK national result, got {economy_result[key]!r}" + ) + + +def assert_uk_local_area_breakdown(breakdown_result: object) -> None: + """Check a local-area breakdown came from the local-area dataset. + + Both area levels must publish at least one area, and every published area + must carry a disclosure tier. + """ + breakdown = _as_mapping(breakdown_result, label="local-area breakdown") + assert breakdown.get("kind") == "local_areas", ( + f"Expected a local-area breakdown, got kind {breakdown.get('kind')!r}" + ) + dataset = breakdown.get("dataset") + assert isinstance(dataset, str), f"Expected a dataset name, got {dataset!r}" + assert dataset.endswith("_dense"), ( + f"Expected the local-area dataset, got {dataset!r}" + ) + for key in _UK_AREA_RESULT_KEYS: + areas = breakdown.get(key) + assert isinstance(areas, list), f"Expected {key} to be a list, got {areas!r}" + assert areas, f"Expected {key} to publish at least one area" + for area in areas: + record = _as_mapping(area, label=f"{key} record") + assert record.get("disclosure_tier") in _PUBLISHED_DISCLOSURE_TIERS, ( + f"Expected a published disclosure_tier in {key}, got {record}" + ) + + class _GeneratedClientObject: def __init__(self, payload: Mapping) -> None: self._payload = dict(payload) @@ -291,6 +331,89 @@ def test_result_assertions_accept_generated_client_objects() -> None: assert_congressional_district_results(result, expected_district_prefix="UT-") +_BREAKDOWN_CONSTITUENCY = { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "x": 56, + "y": -40, + "average_household_income_change": -150.0, + "relative_household_income_change": -0.0026, + "population": 40530.0, + "disclosure_tier": "reported", + "average_household_income_change_lower": -170.0, + "average_household_income_change_upper": -130.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.0022, +} +_BREAKDOWN_LOCAL_AUTHORITY = { + "local_authority_code": "E06000063", + "local_authority_name": "Cumberland", + "x": 4, + "y": 18, + "average_household_income_change": -120.0, + "relative_household_income_change": -0.002, + "population": 120000.0, + "disclosure_tier": "approximate", + "average_household_income_change_lower": -160.0, + "average_household_income_change_upper": -80.0, + "relative_household_income_change_lower": -0.003, + "relative_household_income_change_upper": -0.001, +} +_BREAKDOWN_RESULT = { + "kind": "local_areas", + "model_version": "2.132.6", + "policyengine_version": "6.10.0", + "dataset": "microcosm_uk_2024_25_dense", + "data_version": "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe", + "constituency_impact": [_BREAKDOWN_CONSTITUENCY], + "local_authority_impact": [_BREAKDOWN_LOCAL_AUTHORITY], +} + + +def test_breakdown_assertions_accept_generated_client_objects() -> None: + status = JobStatusResponse.from_dict( + {"status": "complete", "result": _BREAKDOWN_RESULT} + ) + + assert isinstance(status.result, LocalAreaBreakdownOutput) + assert_uk_local_area_breakdown(status.result) + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"dataset": "microcosm_uk_2024_25"}, "local-area dataset"), + ({"constituency_impact": []}, "constituency_impact"), + ({"local_authority_impact": []}, "local_authority_impact"), + ( + { + "constituency_impact": [ + {**_BREAKDOWN_CONSTITUENCY, "disclosure_tier": None} + ] + }, + "disclosure_tier", + ), + ], +) +def test_breakdown_assertions_reject_an_incomplete_breakdown(change, message) -> None: + with pytest.raises(AssertionError, match=message): + assert_uk_local_area_breakdown({**_BREAKDOWN_RESULT, **change}) + + +def test_national_uk_assertions_accept_a_result_without_area_results() -> None: + assert_no_uk_area_results( + _GeneratedClientObject({"budget": {}, "poverty": {}, "inequality": {}}) + ) + assert_no_uk_area_results( + {"constituency_impact": None, "local_authority_impact": None} + ) + + +def test_national_uk_assertions_reject_area_results() -> None: + with pytest.raises(AssertionError, match="constituency_impact"): + assert_no_uk_area_results({"constituency_impact": [_BREAKDOWN_CONSTITUENCY]}) + + @dataclass class _StubJobStatusHttpResponse: status_code: HTTPStatus @@ -705,3 +828,48 @@ def test_calculate_uk_model( economy_result = result.result assert_economy_result_sections(economy_result) + # A UK-wide run on a bundle that routes area regions computes no area + # results; a breakdown=local_areas run does. + assert_no_uk_area_results(economy_result) + + +@pytest.mark.beta_only +def test_calculate_uk_local_area_breakdown( + client: Client | AuthenticatedClient, + uk_model_version: str, + max_wait_seconds: float, + poll_interval: float, +): + """A UK local-area breakdown returns both area levels from the local-area dataset. + + Given a UK-wide request with breakdown=local_areas + When the simulation is submitted and polled to completion + Then the result is a local-area breakdown with published constituencies + and local authorities, each with a disclosure tier. + """ + # Given + request = SimulationRequest.from_dict( + { + "country": "uk", + "version": uk_model_version, + "scope": "macro", + "region": "uk", + "time_period": "2026", + "include_cliffs": False, + "breakdown": "local_areas", + "reform": { + "gov.hmrc.income_tax.rates.uk[0].rate": {"2023-01-01.2100-12-31": 0.21} + }, + } + ) + + # When + submit_response = submit_simulation_request(client, request) + assert submit_response.version == uk_model_version + result = poll_for_completion( + client, submit_response.job_id, max_wait_seconds, poll_interval + ) + + # Then + assert result.status == "complete" + assert_uk_local_area_breakdown(result.result) From 3cdc814921c9bf1c4112c3e18bf3799c67795965 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:25:16 +0400 Subject: [PATCH 17/31] Document UK local-area breakdowns and record them in the changelog - The gateway README gains a "UK local-area breakdowns" section: the breakdown field and its rules, the LocalAreaBreakdownOutput fields, the local_area_breakdown manifest capability and the gateway's 400, the null area fields of UK-wide and nation runs, the reported stages, and the Stage 12 skip. - The Stage 12 migration doc lists the four household columns the area outputs read, and says the production national run now leaves the area breakdowns out too, by the same uk_area_breakdowns_omitted rule, with breakdown requests skipped by Stage 12. - docs/engineering/skills/microdata-weights.md says which pair, and so which dataset's weights, the UK area breakdowns aggregate. - changelog_entry.yaml replaces #881's entry about the second baseline and reform run with one about national runs no longer building it, and adds the field and result, the manifest key and gateway 400, the earlier column check, and the Stage 12 skip. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 6 ++- docs/engineering/skills/microdata-weights.md | 9 +++++ .../stage-12-modal-worker-foundation.md | 20 ++++++---- .../policyengine-simulation-gateway/README.md | 39 +++++++++++++++++++ 4 files changed, 66 insertions(+), 8 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 01a7f0d2b..44ea710ec 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,6 +1,10 @@ - bump: patch changes: added: + - Added breakdown "local_areas" to the simulation request for UK requests whose region is absent, uk or country/, refused with 422 for other regions, other countries, include_cliffs=true, segmented=true or a region_group (include_cliffs=false stays accepted); the job runs one baseline and reform on the local-area dataset alone, with the region's filter, and returns a LocalAreaBreakdownOutput (kind local_areas, model_version, policyengine_version, dataset, data_version, constituency_impact and local_authority_impact), which the job status result now accepts alongside SingleYearMacroOutput + - Declared local_area_breakdown version 1 in the routing bundle manifest of apps whose policyengine.py bundle routes UK area regions, and made the gateway refuse breakdown with 400, before spawning, for any route whose manifest does not declare it, including legacy routes, since older apps ignore the field and would return a national result + - Made a routed local-area dataset fail as it loads, before any simulation, when it lacks any of the four household columns the area outputs read (constituency_code_oa, la_code_oa, region_code_oa and source_household_id); Stage 12 copies the same columns + - Kept breakdown requests out of the Stage 12 shadow comparison, which skips them as unsupported_options, and added beta integration checks of a UK local-area breakdown and of a UK national result without area data - Added four registered dependency rules, checked by the locked-installs quality guard with proof fixtures; every project commits a uv.lock that PR CI checks with uv lock --check, every install in workflows, scripts, Makefiles, Dockerfiles, Compose files and Modal images installs the locked versions, every test command runs through uv run --frozen, and the daily update-dependencies workflow relocks every project; every uv sync and uv run now passes --frozen, and the five sites that cannot install from a lock are registered exceptions - Added optional disclosure fields to UK constituency and local-authority records (disclosure_tier, reported or approximate, and the 90% interval bounds of the average and relative income change); records stored before this change still validate - Added the rule sim.boundaries.sqlalchemy-and-alembic-only, with proof from the psycopg adapter of 08ebe1ca; code connects to databases only through SQLAlchemy, never by importing a database driver, and runtime code never changes the schema, which changes only through Alembic migrations; a guard parses every file under projects/*/src, libs/*/src and .github/scripts and rejects driver imports, alembic imports outside a migrations directory, create_all and drop_all, create and drop on tables, SQLAlchemy DDL statements and DDL text in string literals, which extends the Stage 12 DDL check's pattern to the gateway and every library @@ -37,7 +41,7 @@ - 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 - - Built the constituency and local-authority breakdowns of a UK national or nation-level run on a dataset without area codes from a second baseline and reform run, with the same policies and nation filter, on the dataset the policyengine.py bundle routes those regions to, after checking that dataset's area codes and authority boundaries; the national figures still come from the run's own dataset + - Stopped UK-wide and nation-level runs on a policyengine.py bundle that routes constituency and local-authority regions to a local-area dataset from computing those breakdowns, which they had built from a second baseline and reform run on that dataset (about 70-90 s and 12 GiB per request, in every budget-window year and the UK release sentinel); they now return null for constituency_impact and local_authority_impact, Stage 12 leaves the same breakdowns out of the same reports, and constituency and local-authority runs are unchanged - Resolved UK constituency and local-authority regions by name or by GSS code in any case once policyengine.py routes them, as row filters from its region registry, so policyengine-api's constituency/{name} and local_authority/{name} requests keep working without weight matrices, and the 17 districts replaced in April 2023 are unsupported regions - Made the executor and Stage 12 image smokes check the routed dataset's area codes, downloading it with the Hugging Face credential, when the installed UK default carries none, and stopped Stage 12 copying area codes for a routed UK national or nation-level report, whose aggregate leaves the two area breakdowns out - 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/docs/engineering/skills/microdata-weights.md b/docs/engineering/skills/microdata-weights.md index ed430851e..f1bd2fe03 100644 --- a/docs/engineering/skills/microdata-weights.md +++ b/docs/engineering/skills/microdata-weights.md @@ -85,3 +85,12 @@ shrink, so adding one is a deliberate change to the rule's baseline. Today it is empty: the last one, `_unweighted_household_frame`, went when the UK constituency and local-authority breakdowns moved to policyengine's disclosure-controlled geography output, which takes the simulations. + +The UK constituency and local-authority breakdowns are built only from the +pair they are given: a constituency or local-authority run's own pair, or, +for a UK-wide or nation-level region, the pair of a separate local-area +breakdown run (`serialize_local_area_breakdown`) on the local-area dataset. +That pair's frames carry the local-area dataset's own weights, so the +breakdowns aggregate them like any other run's. A UK-wide or nation-level run +on a bundle that routes area regions builds no breakdown and never mixes the +national file's weights with the local-area file's. diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 0eee33b10..4ecb5fc0e 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -212,17 +212,23 @@ aggregate calculations. Extra columns are permitted; missing planned columns fail the temporary Stage 12 report without affecting the production result. The plan distinguishes calculated variables from columns copied directly from -the source dataset. UK geographic reports require the household columns -`constituency_code_oa` and `la_code_oa`. Workers do not request those columns as -calculated PolicyEngine variables, but they must be present in each materialized -artifact; otherwise validation fails before geographic aggregation begins. +the source dataset. UK geographic reports require the household columns that +the constituency and local-authority outputs read (`UK_AREA_DATASET_COLUMNS`): +`constituency_code_oa`, `la_code_oa`, `region_code_oa` and +`source_household_id`. Workers do not request those columns as calculated +PolicyEngine variables, but they must be present in each materialized artifact; +otherwise validation fails before geographic aggregation begins. The exception is a UK national or nation-level report whose bundle routes constituency and local-authority regions to another dataset. Its national dataset may carry no area codes, so the plan copies none, and its aggregate -leaves the constituency and local-authority breakdowns out. The production path -builds them from a second baseline and reform run on the routed dataset; Stage -12 has no such run yet. +leaves the constituency and local-authority breakdowns out. The production +national run leaves them out too: both decide with +`uk_area_breakdowns_omitted`, so their results still compare. Production +computes those breakdowns in a separate `breakdown: "local_areas"` request on +the routed dataset. The entry adapts no such request into a Stage 12 report +(skip reason `unsupported_options`), and Stage 12 has no local-area report +yet, so the shadow comparison does not cover the breakdown numbers. Aggregation uses the retained output tables without rerunning either model. The precomputed baseline and reform simulations retain their original policy diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index adea361f5..82270c8a4 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -89,6 +89,45 @@ The worker also adds `job_stage_seconds_` attributes to its `run_simulation` operation, so telemetry can split the baseline from the reform calculation. +## UK local-area breakdowns + +A UK request to `POST /simulate/economy/comparison` may set +`breakdown: "local_areas"`. The job then computes only the constituency and +local-authority impacts, on the local-area dataset that the policyengine.py +bundle routes those regions to (`microcosm_uk_2024_25_dense` in .py 6.10.0), +filtered to the request's region. Its result is a `LocalAreaBreakdownOutput` +instead of a `SingleYearMacroOutput`: + +| Field | Meaning | +|---|---| +| `kind` | Always `local_areas`. | +| `model_version`, `policyengine_version` | The country model and the policyengine.py release that computed it. | +| `dataset`, `data_version` | The local-area dataset's name and release revision. | +| `constituency_impact`, `local_authority_impact` | The published areas of each level. Disclosure control leaves withheld areas out, so a list may be empty; each record carries a `disclosure_tier` (`reported` or `approximate`) and 90% interval bounds. | + +Rules: + +- The request must be for the UK, with `region` absent, `uk` or + `country/`. A constituency, local-authority or bare-name region, + `include_cliffs: true`, `segmented: true` or a `region_group` is refused + with 422, and so is any other `breakdown` value. `include_cliffs: false` + is accepted, because the Public API always sends it. Budget-window + requests do not take the field. +- The gateway sends a breakdown only to a worker app whose routing bundle + manifest declares `local_area_breakdown: 1`; the routing publisher writes + it for apps whose bundle routes UK area regions. Any other route, including + one from the legacy version dicts, gets 400 before a job is spawned, and + the message names the policyengine.py version. Older apps ignore fields + they do not know, so they would otherwise return a national result. +- On those bundles, a UK-wide or nation-level request without `breakdown` + computes no area results: `constituency_impact` and + `local_authority_impact` are null, and the status body omits them. + Constituency and local-authority requests are unchanged. +- The job reports the usual monolithic stages: `loading_data` while the + local-area dataset loads, then `baseline`, `reform` and `outputs`. +- The entry starts no Stage 12 comparison for a breakdown request (skip + reason `unsupported_options`). + ## Image dependencies The Modal image installs with `uv_sync(frozen=True)` from this project's From 895d20dc08b9eb07fbe109a890c9f3a85575211e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:28:39 +0400 Subject: [PATCH 18/31] Remove the UK weight-matrix region fallback On a bundle that did not route UK area regions, _resolve_region built a weight-matrix region for a constituency or local authority (_build_uk_weight_replacement_region), and the runtime and the Stage 12 worker then checked the matrix against the certified release (_require_certified_uk_weight_matrix). From this branch onward, every executor app is built with policyengine.py 6.10.0 or later, which routes those regions to the local-area dataset, and apps already deployed for older releases are never rebuilt (only the gateway and entry are redeployed from main), so nothing uses the fallback any more. The two functions, the runtime's and Stage 12's calls, tests/test_uk_weight_matrix_guard.py and the _forbid_weight_matrix_regions test helper are deleted. Test, written first: test_resolve_region_rejects_uk_area_regions_on_an_unrouted_bundle replaces the test that kept weight-matrix regions until routing. Before the change both of its cases failed with DID NOT RAISE : the fallback built a WeightReplacementStrategy region. They now get "Unsupported UK region". Three executor Pyright errors were in the deleted code, so the baseline shrinks from 71 to 68 entries and registry.toml's count follows (a change under rules/, which needs code-owner approval); docs/generated rules are rendered again. The changelog records the removal. Suites: executor 1826 passed; Pyright matches the 68-entry baseline; ruff and the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + docs/generated/rules.md | 2 +- .../simulation_runtime.py | 171 ----------------- .../stage12_runtime/simulation.py | 4 - .../tests/test_simulation_output_builder.py | 39 ++-- .../tests/test_stage12_runtime.py | 3 - .../tests/test_uk_weight_matrix_guard.py | 174 ------------------ rules/baselines/pyright-executor.txt | 3 - rules/registry.toml | 2 +- 9 files changed, 16 insertions(+), 383 deletions(-) delete mode 100644 projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 44ea710ec..23d4d55c0 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -15,6 +15,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Removed the UK constituency and local-authority weight-matrix region fallback (_build_uk_weight_replacement_region, _require_certified_uk_weight_matrix and its Stage 12 call); executor apps built from now on bundle policyengine.py 6.10.0, which routes those regions to the local-area dataset, while apps already deployed for older releases are never rebuilt and keep their own code; on a bundle that does not route them, an area region is now an unsupported region - Enabled every ruff rule (select ALL), ignoring only COM812 and ISC001, which conflict with ruff format, and the D203 and D213 halves of the mutually exclusive docstring pairs; complex functions, long parameter lists and magic numbers were split into smaller helpers and grouped parameters, each behind new characterization tests, across the contract and persistence validation, entry settings, app factory and Stage 12 adapter, gateway route resolution, Stage 12 coordination, workers, parity and qualification, the precompute and national reduce, the routing publisher and its command line, the release sentinel, the package update helper and the shared rules tool; Modal function, FastAPI route, Pydantic validator and command-line signatures and every OpenAPI document are unchanged; the contract library's batch-state builder and SPM selection resolver take their route values as one grouped argument, the simulation observability initializer drops a service_version parameter no caller passed, and no PLR0913 suppression is needed - Enabled ruff's argument, idiom and pytest-style rules (ARG, ASYNC, B, C4, FBT, PD, PERF, PGH, PLW, PT, PTH, RUF, SIM); unused arguments are removed, underscored or explained, FastAPI auth dependencies use Annotated, re-raised errors keep their cause, and boolean options of in-repo helpers are keyword-only; OpenAPI documents are unchanged, and the shared rules tool's unittest tests are exempt from pytest-style rules; tests take the unused-argument rules (ARG001, ARG002, ARG005) as file-level exceptions, since their fakes and stubs mirror real signatures; best-effort telemetry and cleanup calls that swallow any error use contextlib.suppress, with the reason as a comment - Enabled ruff's structure and import rules (INP, PLC, SLF, N, A, FAST, E); long lines are wrapped, deferred imports that Modal images, test seams, import cycles or heavy dependencies need carry a reason, script directories are declared to ruff as namespace packages, four internal exceptions gained an Error suffix, and published field names, OpenAPI documents and logged error types are unchanged; tests are also exempt from SLF001, INP001 and PLC0415; heavy or optional packages that load only inside functions are listed in ruff's banned-module-level-imports, which keeps them out of module-level imports (TID253) and lets their deferred imports go without a suppression, and the Modal utility programs take INP001 file-wide diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 7a58a25f8..7fa416bbc 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt` -- `sim.quality.pyright-executor`: 71 remaining in `rules/baselines/pyright-executor.txt` +- `sim.quality.pyright-executor`: 68 remaining in `rules/baselines/pyright-executor.txt` - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt` ### Exceptions and review dates 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 6021e5026..569cceec8 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 @@ -17,7 +17,6 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module -from pathlib import Path from typing import TYPE_CHECKING, Any from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri @@ -52,8 +51,6 @@ from types import ModuleType from policyengine.core import Dataset, Policy, Region - from policyengine.core.scoping_strategy import ScopingStrategy - from policyengine.provenance.manifest import ArtifactPathReference from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityMetadata, @@ -417,169 +414,6 @@ def _registry_uk_area_region_code(country_module: ModuleType, region_code: str) return region_code -def _build_uk_weight_replacement_region(region_code: str) -> Region | None: - if "/" not in region_code: - return None - - prefix, value = region_code.split("/", maxsplit=1) - if prefix not in {"constituency", "local_authority"}: - return None - - from policyengine.core.region import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - Region, - ) - from policyengine.core.scoping_strategy import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - WeightReplacementStrategy, - ) - from policyengine.data.uk_geography_assets import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - CONSTITUENCY_ASSET_SPEC, - LOCAL_AUTHORITY_ASSET_SPEC, - ) - - asset_spec = ( - CONSTITUENCY_ASSET_SPEC - if prefix == "constituency" - else LOCAL_AUTHORITY_ASSET_SPEC - ) - return Region( - code=region_code, - label=value, - region_type=prefix, - parent_code="uk", - scoping_strategy=WeightReplacementStrategy( - weight_matrix_bucket=asset_spec.bucket, - weight_matrix_key=asset_spec.weight_matrix_filename, - 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_certified_uk_weight_matrix( - scoping_strategy: ScopingStrategy | None, - dataset: Dataset, - dataset_selection: DatasetSelection, -) -> 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 ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - WeightReplacementStrategy, - ) - - if not isinstance(scoping_strategy, WeightReplacementStrategy): - return - - import h5py - import pandas as pd # noqa: PLC0415 # UK weight-matrix check: kept with its lazily loaded data dependencies - from policyengine.data.uk_geography_assets import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - 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 ( # noqa: PLC0415 # simulation_output_geographic imports this module back lazily; kept lazy on both sides - simulation_output_geographic, - ) - - region = scoping_strategy.region_code - matrix_name = scoping_strategy.weight_matrix_key - manifest = get_release_manifest("uk") - package = manifest.data_package - - def release(reference: ArtifactPathReference) -> 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: - msg = f"UK dataset {dataset_selection.name!r} is not in the certified bundle" - raise ValueError(msg) - 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: - msg = ( - 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." - ) - raise ValueError(msg) - 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) # noqa: SLF001 # package-internal helper, called through its module so tests can patch it - paths = resolve_uk_geography_asset_paths(spec, download_missing_assets=False) - if Path(paths.weight_matrix_path).resolve() != Path(certified.path).resolve(): - msg = ( - 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})." - ) - raise ValueError(msg) - year = str(dataset.year) - with h5py.File(paths.weight_matrix_path, "r") as matrix: - if year not in matrix: - covered = ", ".join(sorted(matrix)) - msg = ( - f"UK region {region!r} reweights households with {matrix_name}, " - f"which has no weights for {year} (it covers {covered})." - ) - raise ValueError(msg) - weights = matrix[year] - if not isinstance(weights, h5py.Dataset): - msg = f"{matrix_name} entry {year} is not a weight matrix" - raise TypeError(msg) - matrix_households = weights.shape[-1] - households = len(pd.DataFrame(dataset.data.entity_data["household"])) - if households != matrix_households: - msg = ( - 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." - ) - raise ValueError(msg) - - def _region_parent_dataset_reference( country_module: ModuleType, country: str, @@ -643,8 +477,6 @@ def _resolve_region( if country == "uk": region_code = _registry_uk_area_region_code(country_module, region_code) region = country_module.model.get_region(region_code) - if region is None and country == "uk" and not _uk_area_regions_routed(): - region = _build_uk_weight_replacement_region(region_code) if region is None: msg = f"Unsupported {country.upper()} region: {region_code}" raise ValueError(msg) @@ -1155,9 +987,6 @@ def _run_simulation_impl_core( uk_local_authority_metadata = detect_uk_local_authority_metadata( country, dataset, region_code=region_resolution.code ) - _require_certified_uk_weight_matrix( - region_resolution.scoping_strategy, dataset, dataset_selection - ) with scopes.span(runtime, 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/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 3b2bb5b56..d94f2d503 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 @@ -279,7 +279,6 @@ def calculate_simulation_frames( _country_module, _load_dataset, _normalise_policy, - _require_certified_uk_weight_matrix, _resolve_dataset_selection, _resolve_region, resolve_data_folder, @@ -325,9 +324,6 @@ def calculate_simulation_frames( dataset, region_code=region.code, ) - _require_certified_uk_weight_matrix( - region.scoping_strategy, dataset, dataset_selection - ) with stage_span(Stage.POLICY_NORMALIZATION): policy = _normalise_policy(simulation.policy) with stage_span(Stage.SIMULATION_BUILD): 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 5e84f1dd6..624e9222d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -29,10 +29,7 @@ fake_analysis, ) from fixtures.weighted_frames import weighted_entity_frame -from policyengine.core.scoping_strategy import ( - RowFilterStrategy, - WeightReplacementStrategy, -) +from policyengine.core.scoping_strategy import RowFilterStrategy from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, @@ -1916,14 +1913,6 @@ def _routed_uk_country_module() -> SimpleNamespace: ) -def _forbid_weight_matrix_regions(monkeypatch) -> None: - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime." - "_build_uk_weight_replacement_region", - lambda _code: pytest.fail("a routed area never uses a weight matrix"), - ) - - @pytest.mark.parametrize( ("requested", "expected"), [ @@ -1940,7 +1929,6 @@ def test_resolve_region_maps_routed_uk_area_names_and_codes_to_the_registry( monkeypatch, requested, expected ): _routes_uk_area_regions(monkeypatch, routed=True) - _forbid_weight_matrix_regions(monkeypatch) country_module = _routed_uk_country_module() resolution = _resolve_region( @@ -1966,7 +1954,6 @@ def test_resolve_region_rejects_uk_areas_outside_the_routed_registry( monkeypatch, requested ): _routes_uk_area_regions(monkeypatch, routed=True) - _forbid_weight_matrix_regions(monkeypatch) with pytest.raises(ValueError, match="Unsupported UK region"): _resolve_region( @@ -1976,22 +1963,22 @@ def test_resolve_region_rejects_uk_areas_outside_the_routed_registry( ) -def test_resolve_region_keeps_weight_matrix_uk_areas_until_areas_are_routed( - monkeypatch, +@pytest.mark.parametrize( + "requested", ["constituency/Aldershot", "local_authority/E07000110"] +) +def test_resolve_region_rejects_uk_area_regions_on_an_unrouted_bundle( + monkeypatch, requested ): + # Without routing, policyengine.py's registry lists no area regions, and + # no weight-matrix region stands in for one: older bundles keep their own + # frozen apps. _routes_uk_area_regions(monkeypatch, routed=False) - # No region registry: an unrouted bundle must not read one. country_module = SimpleNamespace(model=SimpleNamespace(get_region=lambda _: None)) - resolution = _resolve_region( - country_module=country_module, - country="uk", - params={"region": "constituency/Aldershot"}, - ) - - assert resolution.code == "constituency/Aldershot" - assert isinstance(resolution.scoping_strategy, WeightReplacementStrategy) - assert resolution.scoping_strategy.region_code == "Aldershot" + with pytest.raises(ValueError, match=f"Unsupported UK region: {requested}"): + _resolve_region( + country_module=country_module, country="uk", params={"region": requested} + ) def test_resolve_dataset_reference_uses_bundle_default(): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 6c9388d2d..d46de4cb3 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -669,9 +669,6 @@ def stop_after_loading(*_args, **_kwargs): monkeypatch.setattr( worker, "detect_uk_local_authority_metadata", lambda *_a, **_k: None ) - monkeypatch.setattr( - simulation_runtime, "_require_certified_uk_weight_matrix", lambda *_: None - ) monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) planned = planned or _planned_simulation(SimulationRole.BASELINE) 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 deleted file mode 100644 index 7242afaa2..000000000 --- a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py +++ /dev/null @@ -1,174 +0,0 @@ -"""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 -import numpy as np -import pandas as pd -import pytest -from policyengine.core.scoping_strategy import ( - RowFilterStrategy, - 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=_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: - household = pd.DataFrame({"household_id": range(1, households + 1)}) - return SimpleNamespace( - year=year, - data=SimpleNamespace(entity_data={"household": household}), - ) - - -def test_fallback_regions_never_download_the_bucket_matrix() -> None: - region = sr._build_uk_weight_replacement_region("constituency/E14001063") - - assert region is not None - assert region.scoping_strategy.download_missing_assets is False - - -@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.provenance.manifest.get_release_manifest", - lambda _country: pytest.fail("no weight matrix should be resolved"), - ) - - sr._require_certified_uk_weight_matrix(strategy, _dataset(3), _SELECTION) - - -def test_guard_reads_the_matrix_certified_with_the_selected_dataset( - monkeypatch, tmp_path -) -> None: - _bundle(monkeypatch) - calls = _assets(monkeypatch, tmp_path) - - sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3), _SELECTION) - - assert calls == [("uk", "parliamentary_constituency_weights", tmp_path)] - - -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"has no parliamentary_constituency_weights\.h5 from the release of " - r"'enhanced_frs_2024_25'" - ), - ): - 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=r"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: - _bundle(monkeypatch) - _assets(monkeypatch, tmp_path) - - with pytest.raises(ValueError, match=r"has no weights for 2026 \(it covers 2025\)"): - sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3, 2026), _SELECTION) diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt index 35fed4bd6..50db66e6b 100644 --- a/rules/baselines/pyright-executor.txt +++ b/rules/baselines/pyright-executor.txt @@ -58,10 +58,7 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "model_dump" for class "object" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "region_type" of type "RegionType" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAssignmentType :: Type "dict[str, Any] | Policy | None" is not assignable to declared type "dict[str, Any] | None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "items" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column" diff --git a/rules/registry.toml b/rules/registry.toml index b7a59bfb3..dac2bb1a7 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1143,7 +1143,7 @@ job = "pyright" invocation = "python -m scripts.rules pyright" [rule.baseline] file = "rules/baselines/pyright-executor.txt" -count = 71 +count = 68 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good" From 17c73b1ebd9166f64d50eace08405f4981d1bb4a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:02:10 +0400 Subject: [PATCH 19/31] Count UK area regions as routed only to one local-area release dataset uk_area_regions_routed() was true whenever the bundle manifest's region_datasets named both area region types. That decides whether a national file may lack area codes, whether area regions resolve through .py's registry, whether a breakdown may run, and whether the routing manifest declares local_area_breakdown. Two bundles .py can certify were read wrongly: - regional_release_kind "area_coded_national" routes both area types to an area-coded national dataset, but .py's region registry lists no constituencies or local authorities for it. With the weight-matrix fallback removed, every area request would fail, yet the app would still declare local_area_breakdown: 1. - Both types routed to different datasets. routed_uk_area_dataset_name() then fails every breakdown, which again would be declared. uk_area_regions_routed() now requires regional_release_kind "local_area" (the condition .py's area_regions_routed_to_local_release uses) and routed_uk_area_dataset_name() to name one certified dataset. The routing manifest's capability follows from it, so it is exact too. Tests, written first, through the real check with a faked manifest: - test_area_regions_count_as_routed_only_to_one_local_area_dataset: the area_coded_national, missing-kind and different-datasets cases failed with "assert True is False". - test_bundle_manifest_declares_local_area_breakdowns_only_for_one_local_area_dataset (replacing the test that patched the routing check): the same three cases failed with "'local_area_breakdown' not in {... 'local_area_breakdown': 1 ...}". Suites: executor 1831 passed, 22 skipped; Pyright stays at 68; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../uk_local_authority_metadata.py | 41 +++++-- .../tests/test_uk_local_authority_metadata.py | 109 +++++++++++++----- .../tests/test_update_version_registry.py | 53 +++++++-- 3 files changed, 155 insertions(+), 48 deletions(-) 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 d30b4e095..f770c5efd 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 @@ -298,20 +298,37 @@ def detect_uk_local_authority_boundary_version( ) -def uk_area_regions_routed() -> bool: - """Whether the bundle routes UK constituency and local-authority regions. +# policyengine.py's ``regional_release_kind`` for a bundle that routes UK area +# regions to a separate local-area release. +_LOCAL_AREA_RELEASE_KIND = "local_area" + - 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. +def uk_area_regions_routed() -> bool: + """Whether the bundle routes UK area regions to one local-area dataset. + + Two things must hold. The manifest's ``regional_release_kind`` must be + ``local_area``: only then does policyengine.py's region registry list the + constituencies and local authorities, which the runtime resolves as row + filters on that dataset. Another kind, ``area_coded_national``, routes + both area types to a national dataset with area codes without + registering them. And both area types must route to one certified + dataset (``routed_uk_area_dataset_name``), on which a local-area + breakdown builds both levels. Only then may the UK default, a national + file, carry no area codes. """ 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) + manifest = get_release_manifest("uk") + release_kind = getattr(manifest, "regional_release_kind", None) + if release_kind != _LOCAL_AREA_RELEASE_KIND: + return False + try: + routed_uk_area_dataset_name() + except ValueError: + return False + return True def routed_uk_area_dataset_name() -> str: @@ -319,8 +336,8 @@ def routed_uk_area_dataset_name() -> str: Like policyengine.py's ``region_dataset_name``, this is the bundle dataset whose path the region type's template names. Both area region types must - share it: a national run builds both of its area breakdowns from one - baseline and reform pair on that dataset. + share it: a local-area breakdown builds both levels from one baseline and + reform pair on that dataset. """ from policyengine.provenance.manifest import ( get_release_manifest, @@ -342,8 +359,8 @@ def routed_uk_area_dataset_name() -> str: if len(names) > 1: msg = ( "The UK bundle routes constituency and local-authority regions to " - f"different datasets ({', '.join(sorted(names))}); a national run " - "builds both area breakdowns from one dataset" + f"different datasets ({', '.join(sorted(names))}); a local-area " + "breakdown builds both levels from one dataset" ) raise ValueError(msg) return names.pop() 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 4a2418bb2..f9e59e093 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 @@ -189,31 +189,72 @@ def test_dataset_detector_requires_codes_for_area_runs_even_when_routed( ) +_NATIONAL_FILE = "microcosm_uk_2024_25.h5" +_LOCAL_AREA_FILE = "microcosm_uk_2024_25_dense.h5" + + @pytest.mark.parametrize( - ("region_types", "routed"), + ("release_kind", "region_datasets", "routed"), [ - (("national",), False), - (("national", "constituency"), False), - (("national", "constituency", "local_authority"), True), + ("local_area", {"national": _NATIONAL_FILE}, False), + ( + "local_area", + {"national": _NATIONAL_FILE, "constituency": _LOCAL_AREA_FILE}, + False, + ), + ( + "local_area", + { + "national": _NATIONAL_FILE, + "constituency": _LOCAL_AREA_FILE, + "local_authority": _LOCAL_AREA_FILE, + }, + True, + ), + # policyengine.py registers no area regions for an area-coded national + # release, although it routes both area types. + ( + "area_coded_national", + { + "national": _NATIONAL_FILE, + "constituency": _NATIONAL_FILE, + "local_authority": _NATIONAL_FILE, + }, + False, + ), + ( + None, + { + "national": _NATIONAL_FILE, + "constituency": _LOCAL_AREA_FILE, + "local_authority": _LOCAL_AREA_FILE, + }, + False, + ), + # A breakdown builds both levels from one dataset. + ( + "local_area", + { + "national": _NATIONAL_FILE, + "constituency": _LOCAL_AREA_FILE, + "local_authority": _NATIONAL_FILE, + }, + False, + ), ], ) -def test_area_regions_count_as_routed_only_when_both_have_a_dataset( - monkeypatch, region_types: tuple[str, ...], *, routed: bool +def test_area_regions_count_as_routed_only_to_one_local_area_dataset( + monkeypatch, + release_kind: str | None, + region_datasets: dict[str, 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} - ), - ) + _bundle_routes(monkeypatch, region_datasets, release_kind=release_kind) assert uk_area_regions_routed() is routed -_NATIONAL_FILE = "microcosm_uk_2024_25.h5" -_LOCAL_AREA_FILE = "microcosm_uk_2024_25_dense.h5" - - @pytest.mark.parametrize( ("country", "region", "routed", "omitted"), [ @@ -239,20 +280,34 @@ def test_area_breakdowns_are_omitted_only_from_routed_uk_wide_and_nation_runs( ) -def _bundle_routes(monkeypatch, region_datasets: dict[str, str]) -> None: +def uk_manifest( + region_datasets: dict[str, str], *, release_kind: str | None = "local_area" +) -> SimpleNamespace: + """Return a UK manifest routing each region type to a bundle file.""" + return SimpleNamespace( + regional_release_kind=release_kind, + region_datasets={ + region_type: SimpleNamespace(path_template=path) + for region_type, path in region_datasets.items() + }, + datasets={ + "microcosm_uk_2024_25": SimpleNamespace(path=_NATIONAL_FILE), + "microcosm_uk_2024_25_dense": SimpleNamespace(path=_LOCAL_AREA_FILE), + }, + ) + + +def _bundle_routes( + monkeypatch, + region_datasets: dict[str, str], + *, + release_kind: str | None = "local_area", +) -> None: """Install a UK manifest routing each region type to a bundle file.""" + manifest = uk_manifest(region_datasets, release_kind=release_kind) monkeypatch.setattr( "policyengine.provenance.manifest.get_release_manifest", - lambda _country: SimpleNamespace( - region_datasets={ - region_type: SimpleNamespace(path_template=path) - for region_type, path in region_datasets.items() - }, - datasets={ - "microcosm_uk_2024_25": SimpleNamespace(path=_NATIONAL_FILE), - "microcosm_uk_2024_25_dense": SimpleNamespace(path=_LOCAL_AREA_FILE), - }, - ), + lambda _country: manifest, ) diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index cbe8fd3a5..4e8a6fdcb 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -190,10 +190,41 @@ def test_bundle_manifest_declares_job_progress_capability(monkeypatch): assert metadata["job_progress"] == JOB_PROGRESS_CAPABILITY_VERSION +_NATIONAL_FILE = "microcosm_uk_2024_25.h5" +_LOCAL_AREA_FILE = "microcosm_uk_2024_25_dense.h5" + + +def _uk_release_manifest(*, release_kind, constituency_file, local_authority_file): + from types import SimpleNamespace + + return SimpleNamespace( + regional_release_kind=release_kind, + region_datasets={ + "national": SimpleNamespace(path_template=_NATIONAL_FILE), + "constituency": SimpleNamespace(path_template=constituency_file), + "local_authority": SimpleNamespace(path_template=local_authority_file), + }, + datasets={ + "microcosm_uk_2024_25": SimpleNamespace(path=_NATIONAL_FILE), + "microcosm_uk_2024_25_dense": SimpleNamespace(path=_LOCAL_AREA_FILE), + }, + ) + + @pytest.mark.usefixtures("fake_bundle_metadata") -@pytest.mark.parametrize("routed", [True, False]) -def test_bundle_manifest_declares_local_area_breakdowns_only_for_routed_area_regions( - monkeypatch, routed +@pytest.mark.parametrize( + ("release_kind", "constituency_file", "local_authority_file", "declared"), + [ + ("local_area", _LOCAL_AREA_FILE, _LOCAL_AREA_FILE, True), + # Routes both area types but registers no area regions in .py. + ("area_coded_national", _NATIONAL_FILE, _NATIONAL_FILE, False), + # A breakdown needs both levels on one dataset. + ("local_area", _LOCAL_AREA_FILE, _NATIONAL_FILE, False), + (None, _LOCAL_AREA_FILE, _LOCAL_AREA_FILE, False), + ], +) +def test_bundle_manifest_declares_local_area_breakdowns_only_for_one_local_area_dataset( + monkeypatch, release_kind, constituency_file, local_authority_file, *, declared ): from policyengine_simulation_contract.local_area_breakdown import ( LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, @@ -202,10 +233,14 @@ def test_bundle_manifest_declares_local_area_breakdowns_only_for_routed_area_reg from policyengine_simulation_executor import spm as executor_spm monkeypatch.setattr(executor_spm, "runtime_spm_capability", lambda: None) + manifest = _uk_release_manifest( + release_kind=release_kind, + constituency_file=constituency_file, + local_authority_file=local_authority_file, + ) monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "uk_area_regions_routed", - lambda: routed, + "policyengine.provenance.manifest.get_release_manifest", + lambda _country: manifest, ) metadata = registry.build_bundle_manifest_metadata( @@ -213,9 +248,9 @@ def test_bundle_manifest_declares_local_area_breakdowns_only_for_routed_area_reg policyengine_version="6.10.0", ) - # A breakdown runs on the local-area dataset the bundle routes area - # regions to; the gateway sends breakdowns only to workers declaring it. - if routed: + # A breakdown runs on the one local-area dataset the bundle routes both + # area types to; the gateway sends breakdowns only to workers declaring it. + if declared: assert metadata[LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY] == ( LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION ) From 5ecf8d4760d06a6bb3211f528b5deca18f05a07a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:06:32 +0400 Subject: [PATCH 20/31] Check every area-output column before a routed area run simulates A constituency or local-authority run on the routed local-area dataset was checked by detect_uk_local_authority_metadata, which reads only la_code_oa. A release missing constituency_code_oa, region_code_oa or source_household_id ran both simulations and the impact analysis before the disclosure output raised. The runtime and the Stage 12 worker had the same gap. detect_uk_run_dataset_metadata() in uk_local_authority_metadata.py now runs the full routed-dataset check (detect_uk_area_dataset_metadata, all four UK_AREA_DATASET_COLUMNS, naming the dataset and the column) when the run needs area codes (uk_area_codes_required), the bundle routes area regions, and the loaded dataset is the routed one. Every other run keeps detect_uk_local_authority_metadata. Both _run_simulation_impl_core and the Stage 12 worker call it right after the dataset loads. Tests, written first, over the four columns for a constituency and a local-authority region: - runtime (test_a_routed_area_run_needs_every_column_its_outputs_read_before_simulating): three columns failed with "DID NOT RAISE " after both simulations were built; la_code_oa raised the generic "contains no la_code_oa column" error, which does not name the dataset. - Stage 12 (test_stage12_area_run_needs_every_column_its_outputs_read_before_simulating): three columns reached the simulation build ("simulated a dataset it cannot use"); la_code_oa raised the same generic error. A further runtime test pins that a complete area dataset still yields its LAD23 boundaries. The Stage 12 load helper now fakes the new function. Suites: executor 1848 passed, 22 skipped; Pyright stays at 68; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../simulation_runtime.py | 9 +- .../stage12_runtime/simulation.py | 5 +- .../uk_local_authority_metadata.py | 26 +++++ .../tests/test_simulation_output_builder.py | 55 ++++++++++- .../tests/test_stage12_runtime.py | 99 ++++++++++++++++++- 5 files changed, 187 insertions(+), 7 deletions(-) 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 569cceec8..2261097dc 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 @@ -980,12 +980,15 @@ def _run_simulation_impl_core( data_folder=baked_folder, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call - detect_uk_local_authority_metadata, + detect_uk_run_dataset_metadata, uk_area_breakdowns_omitted, ) - uk_local_authority_metadata = detect_uk_local_authority_metadata( - country, dataset, region_code=region_resolution.code + uk_local_authority_metadata = detect_uk_run_dataset_metadata( + country, + dataset, + region_code=region_resolution.code, + dataset_name=dataset_selection.name, ) with scopes.span(runtime, 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_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index d94f2d503..728331273 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 @@ -37,7 +37,7 @@ ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle from policyengine_simulation_executor.uk_local_authority_metadata import ( - detect_uk_local_authority_metadata, + detect_uk_run_dataset_metadata, ) from .cached_simulation import Stage12CachedSimulation, qualifying_cache_identity @@ -319,10 +319,11 @@ def calculate_simulation_frames( data_folder=resolve_data_folder(), ), ) - uk_local_authority_metadata = detect_uk_local_authority_metadata( + uk_local_authority_metadata = detect_uk_run_dataset_metadata( country, dataset, region_code=region.code, + dataset_name=dataset_selection.name, ) with stage_span(Stage.POLICY_NORMALIZATION): policy = _normalise_policy(simulation.policy) 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 f770c5efd..0d00d5640 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 @@ -417,6 +417,32 @@ def detect_uk_local_authority_metadata( ) +def detect_uk_run_dataset_metadata( + country: str, + dataset: object, + *, + region_code: str | None, + dataset_name: str, +) -> UKLocalAuthorityMetadata | None: + """Check a run's loaded dataset before anything is simulated on it. + + A UK run that needs area codes, on the dataset the bundle routes area + regions to, must carry every column its area outputs read + (``detect_uk_area_dataset_metadata``): a constituency or local-authority + run on a release missing one then fails here instead of after both + simulations and the impact analysis. Every other run keeps + ``detect_uk_local_authority_metadata``. + """ + if ( + country == "uk" + and uk_area_codes_required(region_code) + and uk_area_regions_routed() + and dataset_name == routed_uk_area_dataset_name() + ): + return detect_uk_area_dataset_metadata(dataset, dataset_name=dataset_name) + return detect_uk_local_authority_metadata(country, dataset, region_code=region_code) + + def detect_uk_area_dataset_metadata( dataset: object, *, dataset_name: str ) -> UKLocalAuthorityMetadata: 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 624e9222d..958b0c77a 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1616,10 +1616,22 @@ def country_module(self, country): return SimpleNamespace(model=SimpleNamespace(version="2.132.6")) def resolve_region(self, *, params, **_kwargs): - if params.get("region") == "country/wales": + region = params.get("region") + if region == "country/wales": return RegionResolution( code="country/wales", dataset_reference=None, scoping_strategy=_WALES ) + region_type = str(region).partition("/")[0] + if region_type in {"constituency", "local_authority"}: + # An area region names the routed local-area dataset itself. + return RegionResolution( + code=region, + dataset_reference=_LOCAL_AREA_SELECTION.name, + scoping_strategy=RowFilterStrategy( + variable_name=f"{region_type}_code_oa", + variable_value=region.partition("/")[2], + ), + ) return RegionResolution(code="uk", dataset_reference=None) def resolve_selection(self, _params, *, region_resolution): @@ -1737,6 +1749,47 @@ def test_a_routed_uk_wide_or_nation_run_builds_no_area_pair(monkeypatch, region) assert {"dataset_role": "uk_area"} not in _span_attributes(run.runtime.roots) +@pytest.mark.parametrize( + "region", ["constituency/E14001063", "local_authority/E06000063"] +) +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_a_routed_area_run_needs_every_column_its_outputs_read_before_simulating( + monkeypatch, region, missing +): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset(without=missing) + ) + + with pytest.raises( + ValueError, + match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " + rf"{missing} column", + ): + _run_simulation_impl_core( + {**_WALES_NATIONAL_PARAMS, "region": region}, runtime=run.runtime + ) + + assert run.loads == [_LOCAL_AREA_SELECTION] + assert run.builds == [] + + +def test_a_routed_area_run_with_every_column_detects_its_boundaries(monkeypatch): + run = _run_routed_uk( + monkeypatch, + params={**_WALES_NATIONAL_PARAMS, "region": "local_authority/E06000063"}, + local_area_dataset=_routed_area_dataset(), + ) + + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["uk_local_authority_metadata"] == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + assert builder_kwargs["omit_uk_area_breakdowns"] is False + + def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch): local_area_dataset = _routed_area_dataset() monkeypatch.setattr( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index d46de4cb3..eca8612bf 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -667,7 +667,7 @@ def stop_after_loading(*_args, **_kwargs): monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) monkeypatch.setattr(state_datasets, "baked_state_data_folder", lookup) monkeypatch.setattr( - worker, "detect_uk_local_authority_metadata", lambda *_a, **_k: None + worker, "detect_uk_run_dataset_metadata", lambda *_a, **_k: None ) monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) @@ -677,6 +677,103 @@ def stop_after_loading(*_args, **_kwargs): return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) +_UK_LOCAL_AREA_DATASET = "microcosm_uk_2024_25_dense" +_UK_LOCAL_AREA_URI = "hf://policyengine/uk/microcosm_uk_2024_25_dense.h5@revision" + + +def _uk_area_planned_simulation(region_code: str) -> PlannedSimulationExecutionInput: + """Return a planned UK area simulation whose provenance names the dense file.""" + payload = _planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") + payload["geography"] = {"country": "uk", "region": region_code} + payload["output_plan"]["country"] = "uk" + payload["bundle"]["country_package_name"] = "policyengine-uk" + payload["bundle"]["country_package_version"] = "2.132.6" + payload["bundle"]["dataset"]["identity"] = _UK_LOCAL_AREA_DATASET + payload["bundle"]["dataset"]["uri"] = _UK_LOCAL_AREA_URI + payload["population"]["artifact"]["uri"] = _UK_LOCAL_AREA_URI + return PlannedSimulationExecutionInput.model_validate(payload) + + +def _uk_area_household(*, without: str) -> pd.DataFrame: + household = { + "constituency_code_oa": ["E14001063", "W07000112"], + "la_code_oa": ["E06000063", "W06000001"], + "region_code_oa": ["SOUTH_EAST", "WALES"], + "source_household_id": [1, 2], + } + del household[without] + return pd.DataFrame(household) + + +@pytest.mark.parametrize( + "region_code", ["constituency/E14001063", "local_authority/E06000063"] +) +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_stage12_area_run_needs_every_column_its_outputs_read_before_simulating( + monkeypatch, region_code, missing +) -> None: + from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + ) + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + selection = DatasetSelection( + name=_UK_LOCAL_AREA_DATASET, uri=_UK_LOCAL_AREA_URI, is_default=False + ) + dataset = SimpleNamespace( + data=SimpleNamespace( + entity_data={"household": _uk_area_household(without=missing)} + ) + ) + for target, value in { + "uk_area_regions_routed": lambda: True, + "routed_uk_area_dataset_name": lambda: _UK_LOCAL_AREA_DATASET, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.uk_local_authority_metadata.{target}", + value, + ) + monkeypatch.setattr(worker, "_require_installed_bundle", lambda _: None) + monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: object() + ) + monkeypatch.setattr( + simulation_runtime, + "_resolve_region", + lambda **_: RegionResolution( + code=region_code, dataset_reference=_UK_LOCAL_AREA_DATASET + ), + ) + monkeypatch.setattr( + simulation_runtime, + "_resolve_dataset_selection", + lambda _params, *, region_resolution: selection, + ) + monkeypatch.setattr( + simulation_runtime, + "_load_dataset", + lambda _params, *, selection, country_module, data_folder=None: dataset, + ) + monkeypatch.setattr( + worker, + "build_stage12_simulation", + lambda *_args, **_kwargs: pytest.fail("simulated a dataset it cannot use"), + ) + + with pytest.raises( + ValueError, + match=rf"'{_UK_LOCAL_AREA_DATASET}', which the bundle routes .* has no " + rf"{missing} column", + ): + worker.calculate_simulation_frames(_uk_area_planned_simulation(region_code)) + + def _regional_selection(): from policyengine_simulation_executor.simulation_runtime import DatasetSelection From f2b7685e7f307c5dbb5968a3a5a851dc0ae7bf6a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:10:13 +0400 Subject: [PATCH 21/31] Accept only the UK nations .py registers in a breakdown request local_area_breakdown_error() accepted any country/, so a request such as country/ni (the Public API's spelling) or country/atlantis passed the contract's 422 check and failed only later, as a job. The contract now accepts only uk and the four nation regions policyengine.py 6.10.0 registers (UK_NATION_REGIONS: country/england, country/scotland, country/wales and country/northern_ireland), in any case. Any other region gets a 422 whose message lists the accepted values. The gateway installs no country package, so the contract cannot read .py's registry. An executor test checks UK_NATION_REGIONS against the nation regions of the installed registry instead. Tests, written first: - the request test's country/ni and country/atlantis cases failed with "DID NOT RAISE ValidationError"; - test_an_unknown_nation_is_refused_with_the_nations_a_breakdown_takes failed with "assert None is not None"; - test_breakdown_nations_are_the_uk_nations_policyengine_registers failed with AttributeError: no attribute 'UK_NATION_REGIONS'. The accepted-region test now covers all four nations and a mixed-case spelling. Suites: contract 319, gateway 267, entry 668 and executor 1849 passed; the gateway OpenAPI golden is unchanged; contract Pyright has no errors; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../local_area_breakdown.py | 27 ++++++++++++------- .../tests/test_gateway_models.py | 17 +++++++++++- .../tests/test_local_area_breakdown.py | 15 +++++++++++ .../tests/test_uk_local_authority_metadata.py | 12 +++++++++ 4 files changed, 61 insertions(+), 10 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py index cee1eac49..20dfbab5a 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py @@ -16,6 +16,17 @@ LOCAL_AREAS_BREAKDOWN = "local_areas" LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY = "local_area_breakdown" LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION = 1 +# The nation regions policyengine.py registers for the UK (its UK_COUNTRIES); +# an executor test checks this against the installed registry, since the +# gateway installs no country package. +UK_NATION_REGIONS = frozenset( + { + "country/england", + "country/scotland", + "country/wales", + "country/northern_ireland", + } +) def supported_local_area_breakdown_capability(value: object) -> int | None: @@ -38,8 +49,8 @@ def local_area_breakdown_error( """Return why a request cannot have a local-area breakdown, or None. A breakdown runs one UK-wide or nation region on the local-area dataset, - so it needs a UK request whose region is absent, ``uk`` or - ``country/``. It computes no cliff impacts and no region + so it needs a UK request whose region is absent, ``uk`` or one of + ``UK_NATION_REGIONS``, in any case. It computes no cliff impacts and no region segments, and takes no region group. ``include_cliffs`` and ``segmented`` may still be sent as false: the Public API always sends ``include_cliffs: false``. The gateway contract and the worker both apply @@ -47,13 +58,11 @@ def local_area_breakdown_error( """ if country.lower() != "uk": return "breakdown is supported only for UK requests" - if region is not None and region.lower() != "uk": - region_type, _, nation = region.partition("/") - if region_type.lower() != "country" or not nation.strip(): - return ( - "breakdown needs a UK-wide or nation region (uk or " - f"country/), not {region!r}" - ) + if region is not None and region.lower() not in {"uk", *UK_NATION_REGIONS}: + accepted = ", ".join(["uk", *sorted(UK_NATION_REGIONS)]) + return ( + f"breakdown needs a UK-wide or nation region ({accepted}), not {region!r}" + ) for flag, value in ( ("include_cliffs", include_cliffs), ("segmented", segmented), diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index a7f58ade3..266a6dafd 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -350,7 +350,19 @@ def test_breakdown_defaults_to_none(self): assert request.breakdown is None assert "breakdown" not in request.model_dump(exclude_none=True) - @pytest.mark.parametrize("region", [None, "uk", "UK", "country/wales"]) + @pytest.mark.parametrize( + "region", + [ + None, + "uk", + "UK", + "country/england", + "country/scotland", + "country/wales", + "country/northern_ireland", + "country/Wales", + ], + ) def test_local_areas_are_accepted_for_a_uk_wide_or_nation_region(self, region): # The Public API always sends include_cliffs: false. request = SimulationRequest( @@ -372,6 +384,9 @@ def test_local_areas_are_accepted_for_a_uk_wide_or_nation_region(self, region): ({"region": "local_authority/E06000063"}, "'local_authority/E06000063'"), ({"region": "wales"}, "'wales'"), ({"region": "country/"}, "'country/'"), + # The Public API's spelling; .py names it country/northern_ireland. + ({"region": "country/ni"}, "'country/ni'"), + ({"region": "country/atlantis"}, "country/northern_ireland"), ({"include_cliffs": True}, "include_cliffs"), ({"segmented": True}, "segmented"), ({"region_group": ["country/wales"]}, "region_group"), diff --git a/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py index 7c3421594..5f422a24b 100644 --- a/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py +++ b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py @@ -9,6 +9,7 @@ from policyengine_simulation_contract.local_area_breakdown import ( LOCAL_AREA_BREAKDOWN_CAPABILITY_KEY, LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + local_area_breakdown_error, supported_local_area_breakdown_capability, ) from policyengine_simulation_contract.macro_output import ( @@ -161,3 +162,17 @@ def test_a_breakdown_names_both_area_levels(field): def test_a_breakdown_has_its_own_kind_and_no_other_fields(change): with pytest.raises(ValidationError): LocalAreaBreakdownOutput.model_validate({**BREAKDOWN, **change}) + + +def test_an_unknown_nation_is_refused_with_the_nations_a_breakdown_takes(): + error = local_area_breakdown_error( + country="uk", + region="country/ni", + region_group=None, + include_cliffs=False, + segmented=None, + ) + + assert error is not None + for nation in ("england", "scotland", "wales", "northern_ireland"): + assert f"country/{nation}" in error 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 f9e59e093..650da0f48 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 @@ -500,3 +500,15 @@ def test_detector_rejects_an_unidentifiable_configuration() -> None: def test_detector_rejects_missing_or_unsupported_codes(value: object) -> None: with pytest.raises((TypeError, ValueError), match="local-authority code"): detect_uk_local_authority_boundary_version([value]) + + +def test_breakdown_nations_are_the_uk_nations_policyengine_registers() -> None: + from policyengine.countries.uk.regions import build_uk_region_registry + from policyengine_simulation_contract import local_area_breakdown + + registry = build_uk_region_registry( + include_constituencies=False, include_local_authorities=False + ) + + registered = frozenset(region.code for region in registry.get_by_type("country")) + assert registered == local_area_breakdown.UK_NATION_REGIONS From 1fd5c00449c16248d83f0f377c484904c2848ab8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:15:42 +0400 Subject: [PATCH 22/31] Type each breakdown level's records and the area pair precisely - LocalAreaBreakdownOutput.constituency_impact is now list[ConstituencyImpactRecord] and local_authority_impact list[LocalAuthorityImpactRecord], instead of GeographicImpactOutput, whose list takes either record type. A local-authority record can no longer validate in the constituency list, or the reverse. The JSON shape is unchanged: both fields are still arrays of the same records. The executor's builder converts each level's output with _area_records(), which refuses a record of the other level. - UKAreaSimulations.baseline and .reform are typed as the StageReportingSimulation that _build_simulation returns, and .dataset as Dataset. - SimulationOutputBuilder._uk_geography_disclosure is typed as UKGeographyDisclosureOutput | None. The gateway OpenAPI golden is regenerated: the two fields now reference their record schemas as arrays, and the model description gains one sentence. The generated client is regenerated by make check. Tests, written first: test_each_area_level_takes_only_its_own_records failed for both levels with "DID NOT RAISE ValidationError", and test_each_area_level_reads_its_records_as_its_own_type failed because the lists were not of the level's record type. Suites: contract 322, gateway 267, entry 668 and executor 1849 passed; make check passes with the executor at its 68-error Pyright baseline; the generated-client tests that need no services pass; the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- .../macro_output.py | 10 +++---- .../tests/test_local_area_breakdown.py | 29 +++++++++++++++++-- .../simulation_output_builder.py | 27 +++++++++++++++-- .../simulation_runtime.py | 8 ++--- .../tests/golden/openapi.json | 14 +++++++-- 5 files changed, 71 insertions(+), 17 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index bc2f7e020..674992ace 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -361,9 +361,9 @@ class LocalAreaBreakdownOutput(MacroOutputModel): the national run. ``kind`` is always ``local_areas``. ``dataset`` names the local-area dataset, ``data_version`` its release revision, and ``model_version`` and ``policyengine_version`` the country model and the - policyengine.py release that computed it. Areas whose income change - disclosure control withholds are left out of each list, so a list may be - empty. + policyengine.py release that computed it. Each list holds only its own + level's records. Areas whose income change disclosure control withholds + are left out of each list, so a list may be empty. """ kind: Literal["local_areas"] @@ -371,5 +371,5 @@ class LocalAreaBreakdownOutput(MacroOutputModel): policyengine_version: str dataset: str data_version: str - constituency_impact: GeographicImpactOutput - local_authority_impact: GeographicImpactOutput + constituency_impact: list[ConstituencyImpactRecord] + local_authority_impact: list[LocalAuthorityImpactRecord] diff --git a/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py index 5f422a24b..cdd80eaf1 100644 --- a/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py +++ b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py @@ -13,7 +13,9 @@ supported_local_area_breakdown_capability, ) from policyengine_simulation_contract.macro_output import ( + ConstituencyImpactRecord, LocalAreaBreakdownOutput, + LocalAuthorityImpactRecord, SingleYearMacroOutput, ) from pydantic import ValidationError @@ -143,8 +145,8 @@ def test_a_breakdown_may_publish_no_area_of_either_level(): output = LocalAreaBreakdownOutput.model_validate(payload) - assert output.constituency_impact.root == [] - assert output.local_authority_impact.root == [] + assert output.constituency_impact == [] + assert output.local_authority_impact == [] @pytest.mark.parametrize("field", ["constituency_impact", "local_authority_impact"]) @@ -176,3 +178,26 @@ def test_an_unknown_nation_is_refused_with_the_nations_a_breakdown_takes(): assert error is not None for nation in ("england", "scotland", "wales", "northern_ireland"): assert f"country/{nation}" in error + + +@pytest.mark.parametrize( + ("field", "record"), + [ + ("constituency_impact", LOCAL_AUTHORITY), + ("local_authority_impact", CONSTITUENCY), + ], +) +def test_each_area_level_takes_only_its_own_records(field, record): + with pytest.raises(ValidationError): + LocalAreaBreakdownOutput.model_validate({**BREAKDOWN, field: [record]}) + + +def test_each_area_level_reads_its_records_as_its_own_type(): + output = LocalAreaBreakdownOutput.model_validate(BREAKDOWN) + + assert [type(record) for record in output.constituency_impact] == [ + ConstituencyImpactRecord + ] + assert [type(record) for record in output.local_authority_impact] == [ + LocalAuthorityImpactRecord + ] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 7dd225c90..4900df636 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -30,6 +30,7 @@ BudgetaryImpact, CliffImpactOutput, CongressionalDistrictImpactOutput, + ConstituencyImpactRecord, DecileOutput, DetailedBudgetOutput, GeographicImpactOutput, @@ -37,6 +38,7 @@ IntraDecileOutput, LaborSupplyResponseOutput, LocalAreaBreakdownOutput, + LocalAuthorityImpactRecord, PovertyByGenderOutput, PovertyByRaceOutput, PovertyModuleOutputs, @@ -49,6 +51,19 @@ from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata +def _area_records[AreaRecord: (ConstituencyImpactRecord, LocalAuthorityImpactRecord)]( + output: GeographicImpactOutput, record_type: type[AreaRecord] +) -> list[AreaRecord]: + """Return one area level's records, refusing a record of the other level.""" + records: list[AreaRecord] = [] + for record in output.root: + if not isinstance(record, record_type): + msg = f"A {record_type.__name__} list holds a {type(record).__name__}" + raise TypeError(msg) + records.append(record) + return records + + @dataclass class SimulationOutputBuilder: """Assemble one year's macro output from a computed baseline and reform. @@ -80,7 +95,9 @@ class SimulationOutputBuilder: omit_uk_area_breakdowns: bool = False progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: Any = field(default=None, init=False) - _uk_geography_disclosure: Any = field(default=None, init=False) + _uk_geography_disclosure: ( + simulation_output_geographic.UKGeographyDisclosureOutput | None + ) = field(default=None, init=False) def __post_init__(self) -> None: """Normalize ``country`` to lower case.""" @@ -183,8 +200,12 @@ def build_local_area_breakdown(self) -> LocalAreaBreakdownOutput: policyengine_version=self._policyengine_version(), dataset=self.resolved_dataset_name, data_version=self._data_version(), - constituency_impact=constituency_impact, - local_authority_impact=local_authority_impact, + constituency_impact=_area_records( + constituency_impact, ConstituencyImpactRecord + ), + local_authority_impact=_area_records( + local_authority_impact, LocalAuthorityImpactRecord + ), ) def serialize_local_area_breakdown(self) -> dict[str, Any]: 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 2261097dc..ad2064e48 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 @@ -110,10 +110,10 @@ class UKAreaSimulations: breakdown reports. """ - baseline: Any - reform: Any + baseline: StageReportingSimulation + reform: StageReportingSimulation local_authority_metadata: UKLocalAuthorityMetadata - dataset: Any + dataset: Dataset dataset_name: str @@ -788,7 +788,7 @@ def _uk_area_simulations( local_authority_metadata = detect_uk_area_dataset_metadata( dataset, dataset_name=selection.name ) - simulations = [] + simulations: list[StageReportingSimulation] = [] for kind, policy in zip( ("uk_area_baseline", "uk_area_reform"), policies, strict=True ): diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 8df638f92..715cd2c5a 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1767,10 +1767,18 @@ "title": "Data Version" }, "constituency_impact": { - "$ref": "#/components/schemas/GeographicImpactOutput" + "items": { + "$ref": "#/components/schemas/ConstituencyImpactRecord" + }, + "type": "array", + "title": "Constituency Impact" }, "local_authority_impact": { - "$ref": "#/components/schemas/GeographicImpactOutput" + "items": { + "$ref": "#/components/schemas/LocalAuthorityImpactRecord" + }, + "type": "array", + "title": "Local Authority Impact" } }, "additionalProperties": false, @@ -1785,7 +1793,7 @@ "local_authority_impact" ], "title": "LocalAreaBreakdownOutput", - "description": "Completed result of a UK local-area breakdown run.\n\nA request with ``breakdown: \"local_areas\"`` returns this instead of a\n``SingleYearMacroOutput``. The run covers the requested UK-wide or nation\nregion on the local-area dataset alone, and computes only the\nconstituency and local-authority impacts; every other figure comes from\nthe national run. ``kind`` is always ``local_areas``. ``dataset`` names the\nlocal-area dataset, ``data_version`` its release revision, and\n``model_version`` and ``policyengine_version`` the country model and the\npolicyengine.py release that computed it. Areas whose income change\ndisclosure control withholds are left out of each list, so a list may be\nempty." + "description": "Completed result of a UK local-area breakdown run.\n\nA request with ``breakdown: \"local_areas\"`` returns this instead of a\n``SingleYearMacroOutput``. The run covers the requested UK-wide or nation\nregion on the local-area dataset alone, and computes only the\nconstituency and local-authority impacts; every other figure comes from\nthe national run. ``kind`` is always ``local_areas``. ``dataset`` names the\nlocal-area dataset, ``data_version`` its release revision, and\n``model_version`` and ``policyengine_version`` the country model and the\npolicyengine.py release that computed it. Each list holds only its own\nlevel's records. Areas whose income change disclosure control withholds\nare left out of each list, so a list may be empty." }, "LocalAuthorityImpactRecord": { "properties": { From dad344986c85d413720c252380b645bee37e3939 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:19:25 +0400 Subject: [PATCH 23/31] Build local-area breakdowns in their own builder SimulationOutputBuilder had two modes: serialize() for the macro output and serialize_local_area_breakdown() for a breakdown, with fields that applied to only one of them (resolved_dataset_name to the breakdown, omit_uk_area_breakdowns to the macro output). Nothing stopped a caller building a breakdown without a dataset name, or a macro output with one. The breakdown now has its own frozen dataclass, LocalAreaBreakdownBuilder in the new simulation_output_local_areas.py. Its required fields are exactly what a breakdown needs: the local-area baseline and reform, the authority boundary metadata, the dataset name and revision, the year that labels the disclosure release, and the country module for the model version. _run_local_area_breakdown_core builds it and calls serialize(). SimulationOutputBuilder loses resolved_dataset_name, both breakdown methods and _policyengine_version. No behaviour changes: the output, the span names and their nesting, the stage order (baseline, reform, then outputs) and the data-version fallback to the release's data version are the same. The year is the request's parsed year, which is also the year of the dataset loaded for it. Tests: the existing breakdown tests are unchanged in what they check. Only their seams moved: the builder tests construct LocalAreaBreakdownBuilder, and the runtime fakes replace it and assert its arguments. Suites: executor 1849 passed, 22 skipped; Pyright stays at 68; ruff and the quality guards pass. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_builder.py | 76 +-------- .../simulation_output_local_areas.py | 151 ++++++++++++++++++ .../simulation_runtime.py | 20 +-- .../tests/test_simulation_output_builder.py | 79 ++++----- 4 files changed, 203 insertions(+), 123 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 4900df636..66e69fbf5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -30,15 +30,12 @@ BudgetaryImpact, CliffImpactOutput, CongressionalDistrictImpactOutput, - ConstituencyImpactRecord, DecileOutput, DetailedBudgetOutput, GeographicImpactOutput, InequalityOutput, IntraDecileOutput, LaborSupplyResponseOutput, - LocalAreaBreakdownOutput, - LocalAuthorityImpactRecord, PovertyByGenderOutput, PovertyByRaceOutput, PovertyModuleOutputs, @@ -51,19 +48,6 @@ from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata -def _area_records[AreaRecord: (ConstituencyImpactRecord, LocalAuthorityImpactRecord)]( - output: GeographicImpactOutput, record_type: type[AreaRecord] -) -> list[AreaRecord]: - """Return one area level's records, refusing a record of the other level.""" - records: list[AreaRecord] = [] - for record in output.root: - if not isinstance(record, record_type): - msg = f"A {record_type.__name__} list holds a {type(record).__name__}" - raise TypeError(msg) - records.append(record) - return records - - @dataclass class SimulationOutputBuilder: """Assemble one year's macro output from a computed baseline and reform. @@ -73,12 +57,7 @@ class SimulationOutputBuilder: constituency and local-authority breakdowns out, for a UK-wide or nation-level run on a bundle that routes area regions to a local-area dataset (``uk_area_breakdowns_omitted``); those results come from a - separate local-area breakdown run. - - ``serialize()`` builds a ``SingleYearMacroOutput``. A local-area breakdown - run calls ``serialize_local_area_breakdown()`` instead: its ``baseline`` - and ``reform`` are the local-area pair, ``resolved_dataset_name`` names - that dataset and ``resolved_data_version`` its revision. + separate local-area breakdown run (``LocalAreaBreakdownBuilder``). """ country: str @@ -88,7 +67,6 @@ class SimulationOutputBuilder: baseline: Any reform: Any resolved_data_version: str | None = None - resolved_dataset_name: str | None = None resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None @@ -170,54 +148,6 @@ def serialize(self) -> dict[str, Any]: ): return output.model_dump(mode="json") - def build_local_area_breakdown(self) -> LocalAreaBreakdownOutput: - """Build the UK constituency and local-authority impacts alone. - - No national impact analysis runs. Both simulations are computed first, - baseline then reform, so a status poll moves through those stages - before the outputs stage, and both levels then come from one - disclosure release over the pair. - """ - self.baseline.ensure() - self.reform.ensure() - self.progress.enter(JobStage.OUTPUTS) - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD)): - constituency_impact = self._build_uk_constituency_impact() - local_authority_impact = self._build_uk_local_authority_impact() - if ( - constituency_impact is None - or local_authority_impact is None - or not self.resolved_dataset_name - ): - msg = ( - "A local-area breakdown is a UK output built on a named " - "local-area dataset" - ) - raise ValueError(msg) - return LocalAreaBreakdownOutput( - kind="local_areas", - model_version=self._model_version(), - policyengine_version=self._policyengine_version(), - dataset=self.resolved_dataset_name, - data_version=self._data_version(), - constituency_impact=_area_records( - constituency_impact, ConstituencyImpactRecord - ), - local_authority_impact=_area_records( - local_authority_impact, LocalAuthorityImpactRecord - ), - ) - - def serialize_local_area_breakdown(self) -> dict[str, Any]: - """Build a local-area breakdown and dump it to JSON-compatible data.""" - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): - output = self.build_local_area_breakdown() - with ( - self._span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)), - self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_MODEL_DUMP)), - ): - return output.model_dump(mode="json") - def _build_detailed_budget(self) -> DetailedBudgetOutput: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DETAILED_BUDGET)): return simulation_output_budget.build_detailed_budget(self.analysis) @@ -393,10 +323,6 @@ def _model_version(self) -> str: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): return str(getattr(self.country_module.model, "version", "")) - def _policyengine_version(self) -> str: - """Return the policyengine.py release whose bundle this run used.""" - return get_country_release_bundle(self.country).policyengine_version - def _data_version(self) -> str: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): if self.resolved_data_version: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py new file mode 100644 index 000000000..53f25f97e --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py @@ -0,0 +1,151 @@ +"""Build the result of a UK local-area breakdown run.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_observability.stages import ( + ANNUAL_IMPACT_STAGES, + Stage, +) + +from policyengine_simulation_executor import simulation_output_geographic +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + JobProgressReporter, +) +from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_executor.simulation_macro_output import ( + ConstituencyImpactRecord, + GeographicImpactOutput, + LocalAreaBreakdownOutput, + LocalAuthorityImpactRecord, +) + +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) + +# A local-area breakdown is a UK output only. +_COUNTRY = "uk" + + +def _area_records[AreaRecord: (ConstituencyImpactRecord, LocalAuthorityImpactRecord)]( + output: GeographicImpactOutput, record_type: type[AreaRecord] +) -> list[AreaRecord]: + """Return one area level's records, refusing a record of the other level.""" + records: list[AreaRecord] = [] + for record in output.root: + if not isinstance(record, record_type): + msg = f"A {record_type.__name__} list holds a {type(record).__name__}" + raise TypeError(msg) + records.append(record) + return records + + +@dataclass(frozen=True) +class LocalAreaBreakdownBuilder: + """Build a UK local-area breakdown from the local-area dataset's pair. + + ``baseline`` and ``reform`` are built on the dataset the bundle routes UK + area regions to, ``local_authority_metadata`` is the authority boundary + version detected on it, and ``dataset_name`` and ``data_version`` are its + bundle name and the revision the release manifest records for it (None + falls back to the release's data version, as the macro output does). + ``year`` labels the disclosure release, and ``country_module`` supplies + the model version. Each segment is built inside its own observability + span when a ``runtime`` is given. + """ + + baseline: StageReportingSimulation + reform: StageReportingSimulation + local_authority_metadata: UKLocalAuthorityMetadata + dataset_name: str + data_version: str | None + year: int + country_module: Any + runtime: ObservabilityRuntime | None = None + progress: JobProgressReporter = NULL_JOB_PROGRESS + + def _span(self, name: str) -> AbstractContextManager[object]: + return self.runtime.span(name) if self.runtime is not None else nullcontext() + + def build(self) -> LocalAreaBreakdownOutput: + """Build the UK constituency and local-authority impacts alone. + + No national impact analysis runs. Both simulations are computed first, + baseline then reform, so a status poll moves through those stages + before the outputs stage, and both levels then come from one + disclosure release over the pair. + """ + self.baseline.ensure() + self.reform.ensure() + self.progress.enter(JobStage.OUTPUTS) + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD)): + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY)): + disclosure = ( + simulation_output_geographic.compute_uk_geography_disclosure_output( + _COUNTRY, self.baseline, self.reform, year=self.year + ) + ) + constituency_impact = ( + simulation_output_geographic.build_uk_constituency_impact( + _COUNTRY, disclosure + ) + ) + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY)): + local_authority_impact = ( + simulation_output_geographic.build_uk_local_authority_impact( + _COUNTRY, + disclosure, + uk_local_authority_metadata=self.local_authority_metadata, + ) + ) + if constituency_impact is None or local_authority_impact is None: + msg = "UK area outputs returned no constituency or local-authority list" + raise ValueError(msg) + return LocalAreaBreakdownOutput( + kind="local_areas", + model_version=self._model_version(), + policyengine_version=( + get_country_release_bundle(_COUNTRY).policyengine_version + ), + dataset=self.dataset_name, + data_version=self._data_version(), + constituency_impact=_area_records( + constituency_impact, ConstituencyImpactRecord + ), + local_authority_impact=_area_records( + local_authority_impact, LocalAuthorityImpactRecord + ), + ) + + def serialize(self) -> dict[str, Any]: + """Build the breakdown and dump it to JSON-compatible data.""" + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): + output = self.build() + with ( + self._span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)), + self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_MODEL_DUMP)), + ): + return output.model_dump(mode="json") + + def _model_version(self) -> str: + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): + return str(getattr(self.country_module.model, "version", "")) + + def _data_version(self) -> str: + with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): + if self.data_version: + return self.data_version + try: + return get_country_release_bundle(_COUNTRY).data_version + except ValueError: + return "" 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 ad2064e48..232dc42b6 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 @@ -45,6 +45,9 @@ from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, ) +from policyengine_simulation_executor.simulation_output_local_areas import ( + LocalAreaBreakdownBuilder, +) if TYPE_CHECKING: from collections.abc import Iterator @@ -892,21 +895,18 @@ def _run_local_area_breakdown_core( _report_calculation_start(area.reform, progress, JobStage.REFORM) logger.info("Calculating the local-area breakdown") - builder = SimulationOutputBuilder( - country=country, - simulation_params=simulation_params, - country_module=country_module, - dataset=area.dataset, + builder = LocalAreaBreakdownBuilder( baseline=area.baseline, reform=area.reform, - resolved_data_version=_manifest_dataset_revision(country, area.dataset_name), - resolved_dataset_name=area.dataset_name, - resolved_region_code=region_resolution.code, + local_authority_metadata=area.local_authority_metadata, + dataset_name=area.dataset_name, + data_version=_manifest_dataset_revision(country, area.dataset_name), + year=_parse_year(simulation_params), + country_module=country_module, runtime=runtime, - uk_local_authority_metadata=area.local_authority_metadata, progress=progress, ) - output = builder.serialize_local_area_breakdown() + output = builder.serialize() logger.info("Local-area breakdown complete") return output 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 958b0c77a..aa7ae3a9e 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -68,6 +68,9 @@ build_congressional_district_impact, build_congressional_district_impact_output, ) +from policyengine_simulation_executor.simulation_output_local_areas import ( + LocalAreaBreakdownBuilder, +) from policyengine_simulation_executor.simulation_runtime import ( DatasetSelection, RegionResolution, @@ -1660,20 +1663,28 @@ def build_simulation(self, _params, *, dataset, dataset_selection, policy, regio def output_builder(self, **kwargs): self.run.builder_calls.append(kwargs) - return _RecordingOutputBuilder(self.run) + return _RecordingBuilder( + self.run, "macro", dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + ) + + def breakdown_builder(self, **kwargs): + self.run.builder_calls.append(kwargs) + return _RecordingBuilder( + self.run, "local_area_breakdown", {"kind": "local_areas"} + ) + +class _RecordingBuilder: + """Records which builder serialized the run's output.""" -class _RecordingOutputBuilder: - def __init__(self, run): + def __init__(self, run, kind, output): self.run = run + self.kind = kind + self.output = output def serialize(self): - self.run.serialized.append("serialize") - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - def serialize_local_area_breakdown(self): - self.run.serialized.append("serialize_local_area_breakdown") - return {"kind": "local_areas"} + self.run.serialized.append(self.kind) + return self.output def _install_routed_uk_fakes( @@ -1702,6 +1713,7 @@ def _install_routed_uk_fakes( "_load_dataset": seams.load_dataset, "_build_simulation": seams.build_simulation, "SimulationOutputBuilder": seams.output_builder, + "LocalAreaBreakdownBuilder": seams.breakdown_builder, }.items(): monkeypatch.setattr( f"policyengine_simulation_executor.simulation_runtime.{name}", value @@ -1745,7 +1757,7 @@ def test_a_routed_uk_wide_or_nation_run_builds_no_area_pair(monkeypatch, region) uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) ) assert "uk_area_simulations" not in builder_kwargs - assert run.serialized == ["serialize"] + assert run.serialized == ["macro"] assert {"dataset_role": "uk_area"} not in _span_attributes(run.runtime.roots) @@ -1818,15 +1830,13 @@ def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch (builder_kwargs,) = run.builder_calls assert builder_kwargs["baseline"] is run.builds[0]["simulation"] assert builder_kwargs["reform"] is run.builds[1]["simulation"] - assert builder_kwargs["dataset"] is local_area_dataset - assert builder_kwargs["uk_local_authority_metadata"] == UKLocalAuthorityMetadata( + assert builder_kwargs["local_authority_metadata"] == UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ) - assert builder_kwargs["resolved_dataset_name"] == _LOCAL_AREA_SELECTION.name - assert builder_kwargs["resolved_data_version"] == _LOCAL_AREA_REVISION - assert builder_kwargs["resolved_region_code"] == "country/wales" - assert builder_kwargs.get("uk_area_simulations") is None - assert run.serialized == ["serialize_local_area_breakdown"] + assert builder_kwargs["dataset_name"] == _LOCAL_AREA_SELECTION.name + assert builder_kwargs["data_version"] == _LOCAL_AREA_REVISION + assert builder_kwargs["year"] == 2026 + assert run.serialized == ["local_area_breakdown"] assert run.output == {"kind": "local_areas"} assert run.dataset_load_attributes() == [{"dataset_role": "uk_area"}] @@ -1863,7 +1873,7 @@ def test_a_uk_wide_local_area_breakdown_simulates_the_whole_dataset(monkeypatch) assert run.loads == [_LOCAL_AREA_SELECTION] assert [build["scoping_strategy"] for build in run.builds] == [None, None] - assert run.builder_calls[0]["resolved_region_code"] == "uk" + assert run.serialized == ["local_area_breakdown"] @pytest.mark.parametrize( @@ -2879,27 +2889,20 @@ def local_authorities(country, value, **kwargs): def no_impact_analysis(*_args, **_kwargs): pytest.fail("a local-area breakdown computes no national impact analysis") - builder = SimulationOutputBuilder( - country="uk", - simulation_params={ - "country": "uk", - "time_period": "2026", - "breakdown": "local_areas", - }, + builder = LocalAreaBreakdownBuilder( + baseline=_BreakdownSimulation("baseline", events), + reform=_BreakdownSimulation("reform", events), + local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), + dataset_name="microcosm_uk_2024_25_dense", + data_version=_LOCAL_AREA_REVISION, + year=2026, country_module=SimpleNamespace( model=SimpleNamespace(version="2.132.6"), economic_impact_analysis=no_impact_analysis, ), - dataset=SimpleNamespace(metadata={}), - baseline=_BreakdownSimulation("baseline", events), - reform=_BreakdownSimulation("reform", events), - resolved_data_version=_LOCAL_AREA_REVISION, - resolved_dataset_name="microcosm_uk_2024_25_dense", - resolved_region_code="uk", runtime=runtime, - uk_local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ), progress=_RecordingProgress(events), ) return builder, disclosure @@ -2909,13 +2912,13 @@ def test_a_local_area_breakdown_publishes_both_levels_from_one_release(monkeypat events: list[object] = [] builder, disclosure = _breakdown_builder(monkeypatch, events=events) - output = builder.serialize_local_area_breakdown() + output = builder.serialize() releases = [event for event in events if event[0] == "release"] assert releases == [("release", "uk", builder.baseline, builder.reform, 2026)] assert [event for event in events if event[0] != "release"][-2:] == [ ("constituency", disclosure), - ("local_authority", disclosure, builder.uk_local_authority_metadata), + ("local_authority", disclosure, builder.local_authority_metadata), ] assert output == { "kind": "local_areas", @@ -2932,7 +2935,7 @@ def test_a_local_area_breakdown_enters_outputs_after_both_simulations(monkeypatc events: list[object] = [] builder, _ = _breakdown_builder(monkeypatch, events=events) - builder.serialize_local_area_breakdown() + builder.serialize() simulations_and_stages = [event for event in events if not isinstance(event, tuple)] assert simulations_and_stages == ["baseline", "reform", JobStage.OUTPUTS] @@ -2946,7 +2949,7 @@ def test_a_local_area_breakdown_opens_only_registered_spans(monkeypatch): runtime = _TrackingRuntime() builder, _ = _breakdown_builder(monkeypatch, events=events, runtime=runtime) - builder.serialize_local_area_breakdown() + builder.serialize() registered = {stage.value for stage in ANNUAL_IMPACT_STAGES.stages} assert set(runtime.names) <= registered From f77a8beee342fa94e449ecd5241dc3347c580aa1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:21:13 +0400 Subject: [PATCH 24/31] Test the installed bundle's routing and a version-pinned refusal Two tests of behaviour the previous commits already have, so both passed when first run: - test_the_installed_bundle_routes_area_regions_to_the_dense_dataset reads only the manifest packaged with the installed policyengine.py 6.10.0, with no network access. It checks that routed_uk_area_dataset_name() is microcosm_uk_2024_25_dense, that uk_area_regions_routed() is true and agrees with .py's area_regions_routed_to_local_release(), and that _manifest_dataset_revision() returns that dataset's microcosm-uk-2024-25-dense-... revision from the manifest. The other routing tests fake the manifest; this one fails if a .py update changes the real routing. - test_a_worker_pinned_by_uk_model_version_without_it_is_refused: a breakdown pinned with version (a UK model version) to an older app whose manifest declares no local_area_breakdown gets 400 naming policyengine.py 3.9.0, with no spawn. The same pinned national request still routes to that app. Suites: gateway 268 and executor 1850 passed; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_uk_local_authority_metadata.py | 19 ++++++++++++ .../tests/test_local_area_breakdown.py | 29 +++++++++++++++++++ 2 files changed, 48 insertions(+) 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 650da0f48..8b9c4df2a 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 @@ -512,3 +512,22 @@ def test_breakdown_nations_are_the_uk_nations_policyengine_registers() -> None: registered = frozenset(region.code for region in registry.get_by_type("country")) assert registered == local_area_breakdown.UK_NATION_REGIONS + + +def test_the_installed_bundle_routes_area_regions_to_the_dense_dataset() -> None: + # Reads only the manifest packaged with the installed policyengine.py. + from policyengine.countries.uk.regions import area_regions_routed_to_local_release + from policyengine.provenance.manifest import get_release_manifest + from policyengine_simulation_executor.simulation_runtime import ( + _manifest_dataset_revision, + ) + + name = routed_uk_area_dataset_name() + revision = _manifest_dataset_revision("uk", name) + + assert name == "microcosm_uk_2024_25_dense" + assert uk_area_regions_routed() is True + assert uk_area_regions_routed() is area_regions_routed_to_local_release() + assert isinstance(revision, str) + assert revision.startswith("microcosm-uk-2024-25-dense-") + assert revision == get_release_manifest("uk").datasets[name].revision diff --git a/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py b/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py index 96aed95c4..3a5c61f2b 100644 --- a/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py +++ b/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py @@ -71,6 +71,35 @@ def test_an_older_pinned_worker_is_refused_before_spawning( assert breakdown_worker["func"].last_payload is None +def test_a_worker_pinned_by_uk_model_version_without_it_is_refused( + breakdown_worker, client: TestClient +): + # An older app serves UK model 2.60.0 and declares no breakdowns. + state = breakdown_worker["dicts"][ROUTING_STATE_DICT_NAME]["active"] + state["routes"]["uk"]["2.60.0"] = "policyengine-simulation-py3-9-0" + older = state["bundles"]["3.9.0"] + state["bundles"]["3.9.0"] = { + **older, + "uk": {**older["uk"], "model_version": "2.60.0"}, + } + pinned = {**BREAKDOWN_REQUEST, "version": "2.60.0"} + national = {key: value for key, value in pinned.items() if key != "breakdown"} + + refused = client.post("/simulate/economy/comparison", json=pinned) + + assert refused.status_code == 400 + assert "policyengine.py 3.9.0" in refused.json()["detail"] + assert breakdown_worker["func"].last_payload is None + # The same pinned route still serves a national request. + assert client.post("/simulate/economy/comparison", json=national).status_code == ( + 200 + ) + assert breakdown_worker["func"].last_from_name_call == ( + "policyengine-simulation-py3-9-0", + "run_simulation", + ) + + def test_legacy_only_routes_are_refused_before_spawning(mock_modal, client: TestClient): # Without a routing state no bundle manifest can declare the capability. del mock_modal["dicts"][ROUTING_STATE_DICT_NAME] From 2e3c1bf3e83e89638c8ea8c6c73cbb12797ebdd2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:22:18 +0400 Subject: [PATCH 25/31] Correct the changelog and gateway README for the review's points changelog_entry.yaml: - Removes the entry saying the release bound UK weight matrices to the certified data release; the weight-matrix fallback is removed in this release. - The weighted-microdata entry no longer says one recorded exception remains: the list is empty. - The entry presenting the end of national-run area computation as a fix is replaced. That computation was this PR's own unreleased code, not shipped behaviour. The new "changed" entry describes the shipped behaviour instead: on a bundle that routes UK area regions to a local-area release dataset, UK-wide and nation-level runs return null area fields and a breakdown run computes them; on other bundles they still compute both. - The weight-matrix removal entry says what a deployment rebuilds: only the app for the current .py version, now 6.10.0. - The breakdown entries name the four accepted nations, the exact routing condition, the per-level record types, version-pinned refusals, and the column check for constituency and local-authority runs. The gateway README's breakdown section lists the accepted nations and the exact condition under which the routing publisher declares the capability. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 13 ++++++------ .../policyengine-simulation-gateway/README.md | 21 ++++++++++++------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 23d4d55c0..4231810a0 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,21 +1,22 @@ - bump: patch changes: added: - - Added breakdown "local_areas" to the simulation request for UK requests whose region is absent, uk or country/, refused with 422 for other regions, other countries, include_cliffs=true, segmented=true or a region_group (include_cliffs=false stays accepted); the job runs one baseline and reform on the local-area dataset alone, with the region's filter, and returns a LocalAreaBreakdownOutput (kind local_areas, model_version, policyengine_version, dataset, data_version, constituency_impact and local_authority_impact), which the job status result now accepts alongside SingleYearMacroOutput - - Declared local_area_breakdown version 1 in the routing bundle manifest of apps whose policyengine.py bundle routes UK area regions, and made the gateway refuse breakdown with 400, before spawning, for any route whose manifest does not declare it, including legacy routes, since older apps ignore the field and would return a national result - - Made a routed local-area dataset fail as it loads, before any simulation, when it lacks any of the four household columns the area outputs read (constituency_code_oa, la_code_oa, region_code_oa and source_household_id); Stage 12 copies the same columns + - Added breakdown "local_areas" to the simulation request for UK requests whose region is absent, uk or one of country/england, country/scotland, country/wales and country/northern_ireland, refused with 422 for other regions (including country/ni), other countries, include_cliffs=true, segmented=true or a region_group (include_cliffs=false stays accepted); the job runs one baseline and reform on the local-area dataset alone, with the region's filter, and returns a LocalAreaBreakdownOutput (kind local_areas, model_version, policyengine_version, dataset, data_version, and constituency_impact and local_authority_impact as lists of their own record types), which the job status result now accepts alongside SingleYearMacroOutput + - Declared local_area_breakdown version 1 in the routing bundle manifest of apps whose policyengine.py bundle routes both UK area region types to one local-area release dataset (regional_release_kind local_area; an area_coded_national bundle, or one routing the two types to different datasets, does not count as routed), and made the gateway refuse breakdown with 400, before spawning, for any route whose manifest does not declare it, including legacy routes and routes pinned to an older version, since older apps ignore the field and would return a national result + - Made a routed local-area dataset fail as it loads, before any simulation, when it lacks any of the four household columns the area outputs read (constituency_code_oa, la_code_oa, region_code_oa and source_household_id), for breakdown, constituency and local-authority runs alike, in the runtime and in the Stage 12 worker; Stage 12 copies the same columns - Kept breakdown requests out of the Stage 12 shadow comparison, which skips them as unsupported_options, and added beta integration checks of a UK local-area breakdown and of a UK national result without area data - Added four registered dependency rules, checked by the locked-installs quality guard with proof fixtures; every project commits a uv.lock that PR CI checks with uv lock --check, every install in workflows, scripts, Makefiles, Dockerfiles, Compose files and Modal images installs the locked versions, every test command runs through uv run --frozen, and the daily update-dependencies workflow relocks every project; every uv sync and uv run now passes --frozen, and the five sites that cannot install from a lock are registered exceptions - Added optional disclosure fields to UK constituency and local-authority records (disclosure_tier, reported or approximate, and the 90% interval bounds of the average and relative income change); records stored before this change still validate - Added the rule sim.boundaries.sqlalchemy-and-alembic-only, with proof from the psycopg adapter of 08ebe1ca; code connects to databases only through SQLAlchemy, never by importing a database driver, and runtime code never changes the schema, which changes only through Alembic migrations; a guard parses every file under projects/*/src, libs/*/src and .github/scripts and rejects driver imports, alembic imports outside a migrations directory, create_all and drop_all, create and drop on tables, SQLAlchemy DDL statements and DDL text in string literals, which extends the Stage 12 DDL check's pattern to the gateway and every library - Made the budget's reform-minus-baseline totals log a warning when the baseline and reform household frames carry different weights, since the change is then weighted by the reform's weights alone - - Added weighted-microdata rules to the registry, with proof from the #767 household count; the executor's code never aggregates a weight column as a value, reads raw values from weighted frames only in simulation_output_common.py (the UK local-authority breakdown's plain household frame is its one recorded exception there, in a list that may only shrink), and its output-builder tests build entity frames through one shared factory with non-uniform weights, while a doc-only rule records that weights, values, labels and counts each allow different operations + - Added weighted-microdata rules to the registry, with proof from the #767 household count; the executor's code never aggregates a weight column as a value, reads raw values from weighted frames only in simulation_output_common.py, where any other read needs a recorded exception in a list that may only shrink and is now empty, and its output-builder tests build entity frames through one shared factory with non-uniform weights, while a doc-only rule records that weights, values, labels and counts each allow different operations - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: - - Removed the UK constituency and local-authority weight-matrix region fallback (_build_uk_weight_replacement_region, _require_certified_uk_weight_matrix and its Stage 12 call); executor apps built from now on bundle policyengine.py 6.10.0, which routes those regions to the local-area dataset, while apps already deployed for older releases are never rebuilt and keep their own code; on a bundle that does not route them, an area region is now an unsupported region + - Removed the UK constituency and local-authority weight-matrix region fallback (_build_uk_weight_replacement_region, _require_certified_uk_weight_matrix and its Stage 12 call); each deployment rebuilds only the app for the current policyengine.py version, now 6.10.0, which routes those regions to the local-area dataset, and apps deployed for older versions are not rebuilt and keep their own code; on a bundle that does not route them, an area region is now an unsupported region + - Made UK-wide and nation-level runs on a policyengine.py bundle that routes UK area regions to a local-area release dataset (6.10.0) load only the national dataset and return null for constituency_impact and local_authority_impact, whose results come from a breakdown=local_areas run instead; on a bundle that does not route them, those runs still compute both from the national dataset, and Stage 12 leaves the same breakdowns out of the same reports - Enabled every ruff rule (select ALL), ignoring only COM812 and ISC001, which conflict with ruff format, and the D203 and D213 halves of the mutually exclusive docstring pairs; complex functions, long parameter lists and magic numbers were split into smaller helpers and grouped parameters, each behind new characterization tests, across the contract and persistence validation, entry settings, app factory and Stage 12 adapter, gateway route resolution, Stage 12 coordination, workers, parity and qualification, the precompute and national reduce, the routing publisher and its command line, the release sentinel, the package update helper and the shared rules tool; Modal function, FastAPI route, Pydantic validator and command-line signatures and every OpenAPI document are unchanged; the contract library's batch-state builder and SPM selection resolver take their route values as one grouped argument, the simulation observability initializer drops a service_version parameter no caller passed, and no PLR0913 suppression is needed - Enabled ruff's argument, idiom and pytest-style rules (ARG, ASYNC, B, C4, FBT, PD, PERF, PGH, PLW, PT, PTH, RUF, SIM); unused arguments are removed, underscored or explained, FastAPI auth dependencies use Annotated, re-raised errors keep their cause, and boolean options of in-repo helpers are keyword-only; OpenAPI documents are unchanged, and the shared rules tool's unittest tests are exempt from pytest-style rules; tests take the unused-argument rules (ARG001, ARG002, ARG005) as file-level exceptions, since their fakes and stubs mirror real signatures; best-effort telemetry and cleanup calls that swallow any error use contextlib.suppress, with the reason as a comment - Enabled ruff's structure and import rules (INP, PLC, SLF, N, A, FAST, E); long lines are wrapped, deferred imports that Modal images, test seams, import cycles or heavy dependencies need carry a reason, script directories are declared to ruff as namespace packages, four internal exceptions gained an Error suffix, and published field names, OpenAPI documents and logged error types are unchanged; tests are also exempt from SLF001, INP001 and PLC0415; heavy or optional packages that load only inside functions are listed in ruff's banned-module-level-imports, which keeps them out of module-level imports (TID253) and lets their deferred imports go without a suppression, and the Modal utility programs take INP001 file-wide @@ -42,10 +43,8 @@ - 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-wide and nation-level runs on a policyengine.py bundle that routes constituency and local-authority regions to a local-area dataset from computing those breakdowns, which they had built from a second baseline and reform run on that dataset (about 70-90 s and 12 GiB per request, in every budget-window year and the UK release sentinel); they now return null for constituency_impact and local_authority_impact, Stage 12 leaves the same breakdowns out of the same reports, and constituency and local-authority runs are unchanged - Resolved UK constituency and local-authority regions by name or by GSS code in any case once policyengine.py routes them, as row filters from its region registry, so policyengine-api's constituency/{name} and local_authority/{name} requests keep working without weight matrices, and the 17 districts replaced in April 2023 are unsupported regions - Made the executor and Stage 12 image smokes check the routed dataset's area codes, downloading it with the Hugging Face credential, when the installed UK default carries none, and stopped Stage 12 copying area codes for a routed UK national or nation-level report, whose aggregate leaves the two area breakdowns out - - 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 - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated - Recorded a redacted, bounded Stage 12 failure cause naming the exception, failing stage, simulation role and message instead of only the exception class, and kept a failed worker's recorded cause on its simulation and report diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 82270c8a4..033836a95 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -107,17 +107,22 @@ instead of a `SingleYearMacroOutput`: Rules: -- The request must be for the UK, with `region` absent, `uk` or - `country/`. A constituency, local-authority or bare-name region, - `include_cliffs: true`, `segmented: true` or a `region_group` is refused - with 422, and so is any other `breakdown` value. `include_cliffs: false` +- The request must be for the UK, with `region` absent, `uk`, + `country/england`, `country/scotland`, `country/wales` or + `country/northern_ireland` (the nations policyengine.py registers). Any + other region, including a constituency, a local authority, a bare name or + `country/ni`, and `include_cliffs: true`, `segmented: true` or a + `region_group`, is refused with 422, and so is any other `breakdown` + value. `include_cliffs: false` is accepted, because the Public API always sends it. Budget-window requests do not take the field. - The gateway sends a breakdown only to a worker app whose routing bundle - manifest declares `local_area_breakdown: 1`; the routing publisher writes - it for apps whose bundle routes UK area regions. Any other route, including - one from the legacy version dicts, gets 400 before a job is spawned, and - the message names the policyengine.py version. Older apps ignore fields + manifest declares `local_area_breakdown: 1`. The routing publisher writes + it for apps whose bundle routes both UK area region types to one + local-area release dataset (`regional_release_kind` `local_area`). Any + other route, including one pinned to an older version and one from the + legacy version dicts, gets 400 before a job is spawned, and the message + names the policyengine.py version. Older apps ignore fields they do not know, so they would otherwise return a national result. - On those bundles, a UK-wide or nation-level request without `breakdown` computes no area results: `constituency_impact` and From ff0e3d97d6f72d4f1989bfb1f67f33459cb747b9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:22:32 +0400 Subject: [PATCH 26/31] Name the breakdown builder in the microdata-weights guidance docs/engineering/skills/microdata-weights.md named serialize_local_area_breakdown, which the builder refactor replaced with LocalAreaBreakdownBuilder. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/microdata-weights.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/engineering/skills/microdata-weights.md b/docs/engineering/skills/microdata-weights.md index f1bd2fe03..6d50283dd 100644 --- a/docs/engineering/skills/microdata-weights.md +++ b/docs/engineering/skills/microdata-weights.md @@ -89,7 +89,7 @@ disclosure-controlled geography output, which takes the simulations. The UK constituency and local-authority breakdowns are built only from the pair they are given: a constituency or local-authority run's own pair, or, for a UK-wide or nation-level region, the pair of a separate local-area -breakdown run (`serialize_local_area_breakdown`) on the local-area dataset. +breakdown run (`LocalAreaBreakdownBuilder`) on the local-area dataset. That pair's frames carry the local-area dataset's own weights, so the breakdowns aggregate them like any other run's. A UK-wide or nation-level run on a bundle that routes area regions builds no breakdown and never mixes the From 1eab140dc585c98088e2bdcf74dcc0786f2addbf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:51:59 +0400 Subject: [PATCH 27/31] Check every area-output column of the routed file in the image smokes The executor image smoke (smoke_app._validate_installed_uk_local_authority_dataset) and the Stage 12 static-files smoke (validate_installed_uk_local_authority_dataset) read only la_code_oa from the dataset that serves UK area results, which on .py 6.10.0 is the downloaded microcosm_uk_2024_25_dense file. A release of that file missing constituency_code_oa, region_code_oa or source_household_id passed both smokes and failed only when a request ran. Both smokes now call detect_uk_area_metadata_from_hdf() in uk_local_authority_metadata.py. It requires every column in UK_AREA_DATASET_COLUMNS in the household table, then detects the authority boundary version as before. The column check selects zero rows, so it reads only the table's column names, not its data. uk_hdf_household_has_local_authority_codes() shares the same column reader. The PR image smoke therefore checks the real file's columns before merge. Tests, written first, with a real table-format HDF file missing one column: - test_uk_dataset_smoke_requires_every_area_output_column (tests/test_smoke_app.py); - test_v2_static_runtime_file_smoke_requires_every_area_output_column (tests/test_stage12_modal_app.py). Before the change, in both files, the constituency_code_oa, region_code_oa and source_household_id cases failed with "DID NOT RAISE ". The la_code_oa case raised the generic "contains no la_code_oa column" error, which does not name the file. A further test pins that a complete file still gives lad23. The existing smoke tests fake the detection under its new name. Suites: executor 1859 passed, 22 skipped; Pyright stays at 68; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/smoke_app.py | 9 ++- .../modal/v2_static_runtime_files_smoke.py | 9 +-- .../uk_local_authority_metadata.py | 34 +++++++++- .../tests/test_smoke_app.py | 62 ++++++++++++++++++- .../tests/test_stage12_modal_app.py | 40 +++++++++++- 5 files changed, 140 insertions(+), 14 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 3cb4f820e..a4467e464 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -24,7 +24,7 @@ resolve_local_bundle_dataset_path, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( - detect_uk_local_authority_metadata_from_hdf, + detect_uk_area_metadata_from_hdf, uk_area_hdf_dataset_path, ) from src.modal.app import build_runtime_simulation_image, hf_secret @@ -101,15 +101,14 @@ def _validate_installed_uk_local_authority_dataset() -> str: That is the installed default, unless it carries no area codes and the bundle routes area regions to another dataset; then the routed dataset is - downloaded and checked instead (``uk_area_hdf_dataset_path``). + downloaded and checked instead (``uk_area_hdf_dataset_path``). Its + household table must also have every column the area outputs read. """ dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: msg = "installed certified UK dataset is unavailable" raise RuntimeError(msg) - metadata = detect_uk_local_authority_metadata_from_hdf( - uk_area_hdf_dataset_path(dataset_path) - ) + metadata = detect_uk_area_metadata_from_hdf(uk_area_hdf_dataset_path(dataset_path)) return metadata.boundary_version.value diff --git a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py index 9a2c7efbd..791f99e15 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py @@ -8,7 +8,7 @@ resolve_local_bundle_dataset_path, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( - detect_uk_local_authority_metadata_from_hdf, + detect_uk_area_metadata_from_hdf, load_uk_local_authority_resources, uk_area_hdf_dataset_path, ) @@ -24,14 +24,15 @@ def validate_installed_uk_local_authority_dataset() -> str: That is the installed UK dataset, unless it carries no area codes and the bundle routes area regions to another dataset; then the routed dataset is - downloaded and checked instead (``uk_area_hdf_dataset_path``). Raises if - the image has no installed certified UK dataset. + downloaded and checked instead (``uk_area_hdf_dataset_path``). Its + household table must also have every column the area outputs read. + Raises if the image has no installed certified UK dataset. """ dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: msg = "installed certified UK dataset is unavailable" raise RuntimeError(msg) - return detect_uk_local_authority_metadata_from_hdf( + return detect_uk_area_metadata_from_hdf( uk_area_hdf_dataset_path(dataset_path) ).boundary_version.value 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 0d00d5640..18194145d 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 @@ -507,8 +507,11 @@ def uk_area_hdf_dataset_path(installed_default_path: str) -> str: return routed.path -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``.""" +def _uk_hdf_household_columns(dataset_path: str) -> frozenset[str]: + """Return the columns of an installed UK HDF dataset's household table. + + Reads no rows: selecting zero rows returns only the table's columns. + """ with pd.HDFStore(dataset_path, mode="r") as store: if "household" not in store: msg = "UK dataset contains no household table" @@ -517,7 +520,32 @@ def uk_hdf_household_has_local_authority_codes(dataset_path: str) -> bool: if not isinstance(household, pd.DataFrame): msg = "UK dataset household table is not a data frame" raise TypeError(msg) - return "la_code_oa" in household.columns + return frozenset(str(column) for column in household.columns) + + +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``.""" + return "la_code_oa" in _uk_hdf_household_columns(dataset_path) + + +def detect_uk_area_metadata_from_hdf(dataset_path: str) -> UKLocalAuthorityMetadata: + """Check an installed UK HDF dataset that serves area results. + + Its household table must have every column the constituency and + local-authority outputs read (``UK_AREA_DATASET_COLUMNS``), checked from + the table's columns without reading its rows. Its authority codes must + then match one packaged boundary version, which is returned. + """ + columns = _uk_hdf_household_columns(dataset_path) + for column in UK_AREA_DATASET_COLUMNS: + if column not in columns: + msg = ( + f"UK dataset {Path(dataset_path).name} household table has no " + f"{column} column, which its constituency and local-authority " + "outputs read" + ) + raise ValueError(msg) + return detect_uk_local_authority_metadata_from_hdf(dataset_path) def _require_local_authority_column(household_chunk: pd.DataFrame) -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_smoke_app.py b/projects/policyengine-simulation-executor/tests/test_smoke_app.py index aac34be58..740e12bf3 100644 --- a/projects/policyengine-simulation-executor/tests/test_smoke_app.py +++ b/projects/policyengine-simulation-executor/tests/test_smoke_app.py @@ -3,6 +3,7 @@ import importlib import sys +import pandas as pd import pytest from fixtures.fake_modal import install_fake_modal from policyengine_simulation_contract.uk_geography import ( @@ -113,7 +114,7 @@ def _uk_area_dataset_smoke( ) monkeypatch.setattr( smoke_module, - "detect_uk_local_authority_metadata_from_hdf", + "detect_uk_area_metadata_from_hdf", lambda path: ( observed.append(path) or UKLocalAuthorityMetadata(boundary_version=boundary_version) @@ -154,6 +155,65 @@ def test_uk_dataset_smoke_checks_the_routed_dataset_of_a_national_default( assert observed == ["/installed/microcosm_uk_2024_25_dense.h5"] +_AREA_OUTPUT_COLUMNS = ( + "constituency_code_oa", + "la_code_oa", + "region_code_oa", + "source_household_id", +) + + +def write_area_household_hdf(path, *, without: str | None = None) -> str: + """Write a household table with the area-output columns, less ``without``.""" + household = { + "household_id": [1, 2], + "constituency_code_oa": ["E14001063", "W07000112"], + "la_code_oa": ["E06000063", "W06000001"], + "region_code_oa": ["SOUTH_EAST", "WALES"], + "source_household_id": [1, 2], + } + household.pop(without or "", None) + pd.DataFrame(household).to_hdf( + path, key="household", format="table", data_columns=True + ) + return str(path) + + +def _routed_area_smoke(monkeypatch, smoke_module, tmp_path, *, without): + installed = str(tmp_path / "microcosm_uk_2024_25.h5") + routed = write_area_household_hdf( + tmp_path / "microcosm_uk_2024_25_dense.h5", without=without + ) + monkeypatch.setattr( + smoke_module, + "resolve_local_bundle_dataset_path", + lambda _country, _requested_data: installed, + ) + monkeypatch.setattr( + smoke_module, + "uk_area_hdf_dataset_path", + lambda path: routed if path == installed else pytest.fail(path), + ) + + +@pytest.mark.parametrize("missing", _AREA_OUTPUT_COLUMNS) +def test_uk_dataset_smoke_requires_every_area_output_column( + monkeypatch, smoke_module, tmp_path, missing +) -> None: + _routed_area_smoke(monkeypatch, smoke_module, tmp_path, without=missing) + + with pytest.raises(ValueError, match=f"has no {missing} column"): + smoke_module._validate_installed_uk_local_authority_dataset() + + +def test_uk_dataset_smoke_accepts_a_routed_dataset_with_every_column( + monkeypatch, smoke_module, tmp_path +) -> None: + _routed_area_smoke(monkeypatch, smoke_module, tmp_path, without=None) + + assert smoke_module._validate_installed_uk_local_authority_dataset() == "lad23" + + def test_executor_smoke_can_download_a_routed_private_dataset(smoke_module) -> None: functions = dict(smoke_module.app.function_calls) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 325205936..ccf4a1918 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -243,7 +243,7 @@ def test_v2_static_runtime_file_smoke_checks_the_routed_area_dataset( ) monkeypatch.setattr( smoke, - "detect_uk_local_authority_metadata_from_hdf", + "detect_uk_area_metadata_from_hdf", lambda path: ( checked.append(path) or UKLocalAuthorityMetadata( @@ -256,6 +256,44 @@ def test_v2_static_runtime_file_smoke_checks_the_routed_area_dataset( assert checked == [routed] +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_v2_static_runtime_file_smoke_requires_every_area_output_column( + monkeypatch, tmp_path, missing +) -> None: + import pandas as pd + + _load(monkeypatch) + sys.modules.pop("src.modal.v2_static_runtime_files_smoke", None) + smoke = importlib.import_module("src.modal.v2_static_runtime_files_smoke") + installed = str(tmp_path / "microcosm_uk_2024_25.h5") + routed = str(tmp_path / "microcosm_uk_2024_25_dense.h5") + household = { + "household_id": [1, 2], + "constituency_code_oa": ["E14001063", "W07000112"], + "la_code_oa": ["E06000063", "W06000001"], + "region_code_oa": ["SOUTH_EAST", "WALES"], + "source_household_id": [1, 2], + } + del household[missing] + pd.DataFrame(household).to_hdf( + routed, key="household", format="table", data_columns=True + ) + monkeypatch.setattr( + smoke, "resolve_local_bundle_dataset_path", lambda _country, _data: installed + ) + monkeypatch.setattr( + smoke, + "uk_area_hdf_dataset_path", + lambda path: routed if path == installed else pytest.fail(path), + ) + + with pytest.raises(ValueError, match=f"has no {missing} column"): + smoke.validate_installed_uk_local_authority_dataset() + + def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: module = _load(monkeypatch) From 64fdfed74b66b69315812d01b269c5fcb21934a7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:54:00 +0400 Subject: [PATCH 28/31] Share the output builders' span and version helpers LocalAreaBreakdownBuilder carried copies of SimulationOutputBuilder's _span, _model_version and _data_version, including the data-version fallback. Both builders now use one module, simulation_output_provenance.py: - output_span(runtime, name) opens the span with scopes.span(runtime, name), as docs/engineering/skills/type-checking.md asks, or does nothing when no runtime is given; - output_model_version(country_module, runtime=...) reads the country model's version inside the output_model_version span; - output_data_version(country, runtime=..., resolved=..., simulation_params=..., dataset=...) returns the resolved revision, then a data_version request parameter, then the dataset's metadata, then the release's data version, or an empty string, inside the output_data_version span. SimulationOutputBuilder keeps its _span, _model_version and _data_version methods as one-line calls to these, so its many span sites are unchanged. The breakdown builder calls the helpers directly. No behaviour changes: the span names and nesting, the version values and the fallback order are the same, and the existing tests pass unchanged. Suites: executor 1859 passed, 22 skipped; Pyright stays at 68; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_builder.py | 36 +++++----- .../simulation_output_local_areas.py | 53 ++++++++------- .../simulation_output_provenance.py | 67 +++++++++++++++++++ 3 files changed, 111 insertions(+), 45 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 66e69fbf5..7afbd7252 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -2,7 +2,6 @@ from __future__ import annotations -from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -25,7 +24,6 @@ NULL_JOB_PROGRESS, JobProgressReporter, ) -from policyengine_simulation_executor.release_bundle import get_country_release_bundle from policyengine_simulation_executor.simulation_macro_output import ( BudgetaryImpact, CliffImpactOutput, @@ -42,8 +40,15 @@ PovertyOutput, SingleYearMacroOutput, ) +from policyengine_simulation_executor.simulation_output_provenance import ( + output_data_version, + output_model_version, + output_span, +) if TYPE_CHECKING: + from contextlib import AbstractContextManager + from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata @@ -82,7 +87,7 @@ def __post_init__(self) -> None: self.country = self.country.lower() def _span(self, name: str) -> AbstractContextManager[object]: - return self.runtime.span(name) if self.runtime is not None else nullcontext() + return output_span(self.runtime, name) @property def analysis(self) -> object: @@ -320,22 +325,13 @@ def _build_uk_local_authority_impact(self) -> GeographicImpactOutput | None: ) def _model_version(self) -> str: - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): - return str(getattr(self.country_module.model, "version", "")) + return output_model_version(self.country_module, runtime=self.runtime) def _data_version(self) -> str: - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): - if self.resolved_data_version: - return str(self.resolved_data_version) - if self.simulation_params.get("data_version"): - return str(self.simulation_params["data_version"]) - metadata = getattr(self.dataset, "metadata", {}) or {} - for key in ("data_version", "version"): - value = metadata.get(key) - if value is not None: - return str(value) - try: - return get_country_release_bundle(self.country).data_version - except ValueError: - pass - return "" + return output_data_version( + self.country, + runtime=self.runtime, + resolved=self.resolved_data_version, + simulation_params=self.simulation_params, + dataset=self.dataset, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py index 53f25f97e..206bb4595 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py @@ -2,7 +2,6 @@ from __future__ import annotations -from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -24,6 +23,11 @@ LocalAreaBreakdownOutput, LocalAuthorityImpactRecord, ) +from policyengine_simulation_executor.simulation_output_provenance import ( + output_data_version, + output_model_version, + output_span, +) if TYPE_CHECKING: from policyengine_observability import ObservabilityRuntime @@ -74,9 +78,6 @@ class LocalAreaBreakdownBuilder: runtime: ObservabilityRuntime | None = None progress: JobProgressReporter = NULL_JOB_PROGRESS - def _span(self, name: str) -> AbstractContextManager[object]: - return self.runtime.span(name) if self.runtime is not None else nullcontext() - def build(self) -> LocalAreaBreakdownOutput: """Build the UK constituency and local-authority impacts alone. @@ -88,8 +89,12 @@ def build(self) -> LocalAreaBreakdownOutput: self.baseline.ensure() self.reform.ensure() self.progress.enter(JobStage.OUTPUTS) - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD)): - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY)): + with output_span( + self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD) + ): + with output_span( + self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY) + ): disclosure = ( simulation_output_geographic.compute_uk_geography_disclosure_output( _COUNTRY, self.baseline, self.reform, year=self.year @@ -100,7 +105,9 @@ def build(self) -> LocalAreaBreakdownOutput: _COUNTRY, disclosure ) ) - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY)): + with output_span( + self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY) + ): local_authority_impact = ( simulation_output_geographic.build_uk_local_authority_impact( _COUNTRY, @@ -113,12 +120,16 @@ def build(self) -> LocalAreaBreakdownOutput: raise ValueError(msg) return LocalAreaBreakdownOutput( kind="local_areas", - model_version=self._model_version(), + model_version=output_model_version( + self.country_module, runtime=self.runtime + ), policyengine_version=( get_country_release_bundle(_COUNTRY).policyengine_version ), dataset=self.dataset_name, - data_version=self._data_version(), + data_version=output_data_version( + _COUNTRY, runtime=self.runtime, resolved=self.data_version + ), constituency_impact=_area_records( constituency_impact, ConstituencyImpactRecord ), @@ -129,23 +140,15 @@ def build(self) -> LocalAreaBreakdownOutput: def serialize(self) -> dict[str, Any]: """Build the breakdown and dump it to JSON-compatible data.""" - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): + with output_span(self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): output = self.build() with ( - self._span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)), - self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_MODEL_DUMP)), + output_span( + self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION) + ), + output_span( + self.runtime, + ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_MODEL_DUMP), + ), ): return output.model_dump(mode="json") - - def _model_version(self) -> str: - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): - return str(getattr(self.country_module.model, "version", "")) - - def _data_version(self) -> str: - with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): - if self.data_version: - return self.data_version - try: - return get_country_release_bundle(_COUNTRY).data_version - except ValueError: - return "" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py new file mode 100644 index 000000000..2316cf119 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py @@ -0,0 +1,67 @@ +"""Spans and version provenance shared by the simulation output builders.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, nullcontext +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_observability import scopes +from policyengine_simulation_observability.stages import ( + ANNUAL_IMPACT_STAGES, + Stage, +) + +from policyengine_simulation_executor.release_bundle import get_country_release_bundle + +if TYPE_CHECKING: + from collections.abc import Mapping + from types import ModuleType + + from policyengine_observability import ObservabilityRuntime + + +def output_span( + runtime: ObservabilityRuntime | None, name: str +) -> AbstractContextManager[object]: + """Open a span named ``name``, or nothing when no runtime is given.""" + if runtime is None: + return nullcontext() + return scopes.span(runtime, name) + + +def output_model_version( + country_module: ModuleType, *, runtime: ObservabilityRuntime | None +) -> str: + """Return the country model version that computed an output.""" + with output_span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_MODEL_VERSION)): + return str(getattr(country_module.model, "version", "")) + + +def output_data_version( + country: str, + *, + runtime: ObservabilityRuntime | None, + resolved: str | None, + simulation_params: Mapping[str, Any] | None = None, + dataset: object = None, +) -> str: + """Return the data version an output reports. + + In order: the ``resolved`` revision of the dataset the run loaded, a + ``data_version`` request parameter, the dataset's own metadata, and the + release's data version; an empty string when none is available. + """ + with output_span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_DATA_VERSION)): + if resolved: + return str(resolved) + if simulation_params is not None and simulation_params.get("data_version"): + return str(simulation_params["data_version"]) + metadata = getattr(dataset, "metadata", {}) or {} + for key in ("data_version", "version"): + value = metadata.get(key) + if value is not None: + return str(value) + try: + return get_country_release_bundle(country).data_version + except ValueError: + return "" From b34a705275e78f8cad4409590154e68c10b1dd10 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:55:38 +0400 Subject: [PATCH 29/31] Type the breakdown builder's country module and drop an unused field - LocalAreaBreakdownBuilder.country_module is typed as ModuleType, what _country_module returns, instead of Any. - UKAreaSimulations loses its dataset field. Since the breakdown moved to its own builder, nothing reads it: the builder needs the dataset's name, which the pair still carries, not the loaded dataset. No behaviour changes; the existing tests pass unchanged. Suites: executor 1859 passed, 22 skipped; Pyright stays at 68; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_local_areas.py | 4 +++- .../simulation_runtime.py | 8 +++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py index 206bb4595..9647ce091 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py @@ -30,6 +30,8 @@ ) if TYPE_CHECKING: + from types import ModuleType + from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata @@ -74,7 +76,7 @@ class LocalAreaBreakdownBuilder: dataset_name: str data_version: str | None year: int - country_module: Any + country_module: ModuleType runtime: ObservabilityRuntime | None = None progress: JobProgressReporter = NULL_JOB_PROGRESS 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 232dc42b6..257151403 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 @@ -108,15 +108,14 @@ class UKAreaSimulations: A local-area breakdown run builds its constituency and local-authority impacts from this pair, with the request's policies and nation filter, - and from the authority boundaries detected on that dataset. ``dataset`` - is the loaded dataset and ``dataset_name`` its bundle name, which the - breakdown reports. + and from the authority boundaries detected on that dataset. + ``dataset_name`` is that dataset's bundle name, which the breakdown + reports. """ baseline: StageReportingSimulation reform: StageReportingSimulation local_authority_metadata: UKLocalAuthorityMetadata - dataset: Dataset dataset_name: str @@ -814,7 +813,6 @@ def _uk_area_simulations( baseline=baseline, reform=reform, local_authority_metadata=local_authority_metadata, - dataset=dataset, dataset_name=selection.name, ) From dbf75b34d5973f647b6630e22065953d9f5626e3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:59:24 +0400 Subject: [PATCH 30/31] Test a Scotland breakdown on the installed bundle and the check fallback Two tests of behaviour the earlier commits already have, so both passed when first run: - test_a_scotland_breakdown_resolves_against_the_installed_bundle runs a country/scotland breakdown through _run_simulation_impl_core with the real country module, region resolution, dataset selection, routing and revision lookup of the installed .py 6.10.0 bundle. Only the dataset load, the simulations and the breakdown builder are faked, so no data or network is needed. It checks: - the one selection loaded is the non-default microcosm_uk_2024_25_dense; - both simulations use the nation's row filter (region == SCOTLAND); - the builder gets that dataset's manifest revision. - test_an_area_run_on_another_dataset_keeps_the_authority_code_check covers detect_uk_run_dataset_metadata's other branch. An area run on a dataset other than the routed one still gets only the la_code_oa check: it passes without the other three columns, and it fails without la_code_oa. A companion test shows the same dataset named as the routed one is refused for its missing constituency_code_oa. Suites: executor 1862 passed, contract 322, gateway 268 and entry 668; ruff passes. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_simulation_output_builder.py | 59 +++++++++++++++++++ .../tests/test_uk_local_authority_metadata.py | 55 +++++++++++++++++ 2 files changed, 114 insertions(+) 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 aa7ae3a9e..94723a931 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1841,6 +1841,65 @@ def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch assert run.dataset_load_attributes() == [{"dataset_role": "uk_area"}] +def test_a_scotland_breakdown_resolves_against_the_installed_bundle(monkeypatch): + # Only the dataset load, the simulations and the builder are faked: the + # country module, region resolution, dataset selection, routing and the + # revision come from the installed .py bundle, with no data or network. + from policyengine.provenance.manifest import get_release_manifest + from policyengine_simulation_executor import simulation_runtime + + loads: list[DatasetSelection] = [] + builds: list[dict] = [] + builder_calls: list[dict] = [] + + def load_dataset(_params, *, selection, country_module, data_folder=None): + loads.append(selection) + return _routed_area_dataset() + + def build_simulation(_params, *, dataset, dataset_selection, policy, region): + builds.append( + { + "selection": dataset_selection, + "region_code": region.code, + "scoping_strategy": region.scoping_strategy, + } + ) + return SimpleNamespace(on_first_ensure=lambda _callback: None) + + class RecordingBreakdownBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return {"kind": "local_areas"} + + for name, value in { + "_load_dataset": load_dataset, + "_build_simulation": build_simulation, + "LocalAreaBreakdownBuilder": RecordingBreakdownBuilder, + }.items(): + monkeypatch.setattr(simulation_runtime, name, value) + + output = _run_simulation_impl_core( + {**_WALES_BREAKDOWN_PARAMS, "region": "country/scotland"}, + runtime=_TrackingRuntime(), + ) + + (selection,) = loads + assert selection.name == "microcosm_uk_2024_25_dense" + assert selection.is_default is False + scotland = RowFilterStrategy(variable_name="region", variable_value="SCOTLAND") + assert [build["scoping_strategy"] for build in builds] == [scotland, scotland] + assert [build["selection"] for build in builds] == [selection, selection] + assert {build["region_code"] for build in builds} == {"country/scotland"} + (builder_kwargs,) = builder_calls + assert builder_kwargs["dataset_name"] == selection.name + assert builder_kwargs["data_version"] == ( + get_release_manifest("uk").datasets[selection.name].revision + ) + assert output == {"kind": "local_areas"} + + def test_a_dataset_revision_is_read_from_the_release_manifest(monkeypatch): from policyengine_simulation_executor import simulation_runtime 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 8b9c4df2a..20552792b 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 @@ -14,6 +14,7 @@ detect_uk_local_authority_boundary_version, detect_uk_local_authority_metadata, detect_uk_local_authority_metadata_from_hdf, + detect_uk_run_dataset_metadata, load_uk_local_authority_resources, routed_uk_area_dataset_name, uk_area_hdf_dataset_path, @@ -531,3 +532,57 @@ def test_the_installed_bundle_routes_area_regions_to_the_dense_dataset() -> None assert isinstance(revision, str) assert revision.startswith("microcosm-uk-2024-25-dense-") assert revision == get_release_manifest("uk").datasets[name].revision + + +def _authority_codes_only() -> SimpleNamespace: + return SimpleNamespace( + data=SimpleNamespace( + entity_data={"household": pd.DataFrame({"la_code_oa": ["E06000063"]})} + ) + ) + + +def _routes_area_regions_to_the_dense_dataset(monkeypatch) -> None: + _routes_area_regions(monkeypatch, routed=True) + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "routed_uk_area_dataset_name", + lambda: "microcosm_uk_2024_25_dense", + ) + + +def test_an_area_run_on_another_dataset_keeps_the_authority_code_check( + monkeypatch, +) -> None: + _routes_area_regions_to_the_dense_dataset(monkeypatch) + + metadata = detect_uk_run_dataset_metadata( + "uk", + _authority_codes_only(), + region_code="local_authority/E06000063", + dataset_name="enhanced_frs_2024_25", + ) + + assert metadata is not None + assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD23 + with pytest.raises(ValueError, match="contains no la_code_oa column"): + detect_uk_run_dataset_metadata( + "uk", + _national_dataset(), + region_code="constituency/E14001063", + dataset_name="enhanced_frs_2024_25", + ) + + +def test_an_area_run_on_the_routed_dataset_needs_every_area_output_column( + monkeypatch, +) -> None: + _routes_area_regions_to_the_dense_dataset(monkeypatch) + + with pytest.raises(ValueError, match="has no constituency_code_oa column"): + detect_uk_run_dataset_metadata( + "uk", + _authority_codes_only(), + region_code="local_authority/E06000063", + dataset_name="microcosm_uk_2024_25_dense", + ) From 9503823ddb310d7afc7826a3667dc6740b061fbd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:00:21 +0400 Subject: [PATCH 31/31] Record the image smokes' column check in the changelog The changelog entry for the image smokes now says that they also require the routed dataset's household table to have all four columns the area outputs read, checked from the column names only. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 4231810a0..a5e0b8bf6 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -44,7 +44,7 @@ - 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 - Resolved UK constituency and local-authority regions by name or by GSS code in any case once policyengine.py routes them, as row filters from its region registry, so policyengine-api's constituency/{name} and local_authority/{name} requests keep working without weight matrices, and the 17 districts replaced in April 2023 are unsupported regions - - Made the executor and Stage 12 image smokes check the routed dataset's area codes, downloading it with the Hugging Face credential, when the installed UK default carries none, and stopped Stage 12 copying area codes for a routed UK national or nation-level report, whose aggregate leaves the two area breakdowns out + - Made the executor and Stage 12 image smokes check the routed dataset's area codes, and that its household table has all four columns the area outputs read (reading only the column names), downloading it with the Hugging Face credential, when the installed UK default carries none, and stopped Stage 12 copying area codes for a routed UK national or nation-level report, whose aggregate leaves the two area breakdowns out - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated - Recorded a redacted, bounded Stage 12 failure cause naming the exception, failing stage, simulation role and message instead of only the exception class, and kept a failed worker's recorded cause on its simulation and report