diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 2da3c0513..a5e0b8bf6 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,15 +1,22 @@ - bump: patch changes: added: + - 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); 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 @@ -27,6 +34,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 @@ -35,7 +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 - - 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 + - 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, 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 diff --git a/docs/engineering/skills/microdata-weights.md b/docs/engineering/skills/microdata-weights.md index 9a84a2d87..6d50283dd 100644 --- a/docs/engineering/skills/microdata-weights.md +++ b/docs/engineering/skills/microdata-weights.md @@ -82,7 +82,15 @@ 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. + +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 (`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 +national file's weights with the local-area file's. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 686c1ea01..7fa416bbc 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`: 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/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 953ca0d70..4ecb5fc0e 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -212,10 +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 +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/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..20dfbab5a --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/local_area_breakdown.py @@ -0,0 +1,74 @@ +"""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 + +# 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 +# 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: + """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 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 + these rules. + """ + if country.lower() != "uk": + return "breakdown is supported only for UK requests" + 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), + ): + 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 a45cddec4..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 @@ -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): @@ -320,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. 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"] + model_version: str + policyengine_version: str + dataset: str + data_version: str + constituency_impact: list[ConstituencyImpactRecord] + local_authority_impact: list[LocalAuthorityImpactRecord] diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index 8fb126fe4..266a6dafd 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -341,6 +341,100 @@ 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/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( + 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/'"), + # 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"), + ], + ) + 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..cdd80eaf1 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_local_area_breakdown.py @@ -0,0 +1,203 @@ +"""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, + local_area_breakdown_error, + supported_local_area_breakdown_capability, +) +from policyengine_simulation_contract.macro_output import ( + ConstituencyImpactRecord, + LocalAreaBreakdownOutput, + LocalAuthorityImpactRecord, + 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 == [] + assert output.local_authority_impact == [] + + +@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}) + + +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 + + +@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/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-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) 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/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/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 1b655bdd5..a4467e464 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -24,9 +24,10 @@ 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 +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,22 @@ 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``). 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(dataset_path) + metadata = detect_uk_area_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/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/src/modal/v2_static_runtime_files_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py index 3c022295a..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,10 +8,11 @@ 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, ) -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,25 @@ 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. + 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``). 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( - dataset_path + return detect_uk_area_metadata_from_hdf( + 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/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 5a86a0482..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 @@ -53,7 +58,11 @@ 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. ``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 (``LocalAreaBreakdownBuilder``). """ country: str @@ -66,15 +75,19 @@ class SimulationOutputBuilder: resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + omit_uk_area_breakdowns: bool = False progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: 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.""" 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: @@ -257,38 +270,68 @@ def _build_congressional_district_impact( region_code=self.resolved_region_code, ) + 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. None for other countries. + """ + if self.country != "uk": + return None + if self._uk_geography_disclosure is None: + self._uk_geography_disclosure = ( + simulation_output_geographic.compute_uk_geography_disclosure_output( + self.country, + self.baseline, + self.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 return simulation_output_geographic.build_uk_constituency_impact( - self.country, self.baseline, self.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 return simulation_output_geographic.build_uk_local_authority_impact( self.country, - self.baseline, - self.reform, + self._uk_geography_disclosure_output(), uk_local_authority_metadata=self.uk_local_authority_metadata, ) 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_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..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 @@ -2,13 +2,30 @@ 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, + RECOVERABLE_PARENT, + RELATIVE_CHANGE, + REPORTED, + UNCHECKED_OVERLAP, +) from policyengine_simulation_executor.simulation_macro_output import ( CongressionalDistrictImpactOutput, @@ -16,19 +33,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 +82,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 +261,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 +350,382 @@ 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} + + +# 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, ...]]] + - Downloads the constituency lookup if needed and raises when the dataset - lacks constituency codes or a record lacks lookup name or coordinates. +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 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)) + ) + 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]], + *, + 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, ...], + 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. 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 []: + 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 + key = (outcome, finding.reform, finding.year) + if finding.suggested and finding.level in published_levels: + 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 + + +def _resolved_public_table( + impact: UKGeographyDisclosure, + area_codes: Mapping[str, Mapping[str, float]], + *, + year: int, + inside: _UKAreasInside, +) -> 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, inside) - 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. + + 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, inside=_uk_areas_inside(household) + ) + 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 +733,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 +771,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/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..9647ce091 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_local_areas.py @@ -0,0 +1,156 @@ +"""Build the result of a UK local-area breakdown run.""" + +from __future__ import annotations + +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, +) +from policyengine_simulation_executor.simulation_output_provenance import ( + output_data_version, + output_model_version, + output_span, +) + +if TYPE_CHECKING: + from types import ModuleType + + 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: ModuleType + runtime: ObservabilityRuntime | None = None + progress: JobProgressReporter = NULL_JOB_PROGRESS + + 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 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 + ) + ) + constituency_impact = ( + simulation_output_geographic.build_uk_constituency_impact( + _COUNTRY, disclosure + ) + ) + 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, + 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=output_model_version( + self.country_module, runtime=self.runtime + ), + policyengine_version=( + get_country_release_bundle(_COUNTRY).policyengine_version + ), + dataset=self.dataset_name, + data_version=output_data_version( + _COUNTRY, runtime=self.runtime, resolved=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 output_span(self.runtime, ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): + output = self.build() + with ( + 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") 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 "" 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 f21114f4b..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 @@ -17,11 +17,14 @@ 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 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 ( @@ -42,15 +45,19 @@ 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 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, + ) from policyengine_simulation_observability.telemetry import TelemetryEnvelope from policyengine_simulation_executor.stage_reporting_simulation import ( @@ -95,6 +102,23 @@ 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_name`` is that dataset's bundle name, which the breakdown + reports. + """ + + baseline: StageReportingSimulation + reform: StageReportingSimulation + local_authority_metadata: UKLocalAuthorityMetadata + dataset_name: str + + def _normalize_credentials_blob(creds_json: str) -> str: """Return the raw JSON blob, decoding the outer escape if present. @@ -358,167 +382,38 @@ def _normalise_region_code(country: str, region: object) -> str: return f"{prefix}/{value}" -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 _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 _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 +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. - 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, + 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 = 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) + 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 _region_parent_dataset_reference( @@ -581,9 +476,9 @@ 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": - region = _build_uk_weight_replacement_region(region_code) if region is None: msg = f"Unsupported {country.upper()} region: {region_code}" raise ValueError(msg) @@ -732,13 +627,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 @@ -848,6 +749,166 @@ 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 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, + routed_uk_area_dataset_name, + ) + + area_role = {"dataset_role": "uk_area"} + with scopes.span( + runtime, + 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 scopes.span( + runtime, 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: list[StageReportingSimulation] = [] + for kind, policy in zip( + ("uk_area_baseline", "uk_area_reform"), policies, strict=True + ): + with scopes.span( + runtime, + 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, + 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 = LocalAreaBreakdownBuilder( + baseline=area.baseline, + reform=area.reform, + 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, + progress=progress, + ) + output = builder.serialize() + logger.info("Local-area breakdown complete") + return output + + def _run_simulation_impl_core( params: dict, *, @@ -866,13 +927,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( @@ -882,6 +937,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) @@ -919,14 +978,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 - ) - _require_certified_uk_weight_matrix( - region_resolution.scoping_strategy, dataset, dataset_selection + 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")) @@ -973,6 +1033,11 @@ def _run_simulation_impl_core( resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, + # 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() @@ -997,6 +1062,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/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index ca4f2d8f5..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 @@ -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_breakdowns_omitted, + ) country = report.baseline.geography.country country_module = _country_module(country) @@ -232,6 +238,9 @@ 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=uk_area_breakdowns_omitted( + country, 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..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 @@ -23,12 +23,19 @@ ) from policyengine_simulation_observability.stages import Stage +from policyengine_simulation_executor.uk_local_authority_metadata import ( + UK_AREA_DATASET_COLUMNS, + uk_area_codes_required, +) + if TYPE_CHECKING: import pandas as pd from policyengine_simulation_contract.stage12_bundle import CountryId +# Input-dataset columns the UK constituency and local-authority breakdowns +# 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"), + "household": UK_AREA_DATASET_COLUMNS, } @@ -105,12 +112,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 +188,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/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 3b2bb5b56..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 @@ -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, @@ -320,13 +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, - ) - _require_certified_uk_weight_matrix( - region.scoping_strategy, dataset, dataset_selection + 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 0ce031c54..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 @@ -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,15 @@ def detect_uk_local_authority_boundary_version( UK_AREA_REGION_TYPES = ("constituency", "local_authority") +# 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 " @@ -288,20 +298,72 @@ def detect_uk_local_authority_boundary_version( ) +# 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" + + def uk_area_regions_routed() -> bool: - """Whether the bundle routes UK constituency and local-authority regions. + """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, + ) + + 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 + - 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 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 local-area breakdown builds both levels from one baseline and + reform pair on that dataset. """ from policyengine.provenance.manifest import ( get_release_manifest, ) - region_datasets = get_release_manifest("uk").region_datasets - return all(region_type in region_datasets for region_type in UK_AREA_REGION_TYPES) + 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 local-area " + "breakdown builds both levels from one dataset" + ) + raise ValueError(msg) + return names.pop() def uk_area_codes_required(region_code: str | None) -> bool: @@ -317,6 +379,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, @@ -331,6 +407,69 @@ 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_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: + """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 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_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, which its constituency and local-authority outputs read" + ) + 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,18 +479,39 @@ 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: - """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" @@ -360,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_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_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-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index f4abd82e8..94723a931 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,13 @@ fake_analysis, ) from fixtures.weighted_frames import weighted_entity_frame +from policyengine.core.scoping_strategy import RowFilterStrategy +from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.uk_geography import ( 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, @@ -65,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, @@ -80,9 +86,7 @@ from policyengine_simulation_executor.simulation_runtime import ( _load_dataset as _load_selected_dataset, ) -from policyengine_simulation_observability.stages import Stage - -CONSTITUENCY_CSV_PATH = "/tmp/constituencies_2024.csv" +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage def _params_with_spm(params, selection): @@ -217,6 +221,7 @@ class BuilderOptions: include_cliffs: bool | None = None resolved_region_code: str | None = "us" uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + omit_uk_area_breakdowns: bool = False def _simulation_output_builder( @@ -256,6 +261,7 @@ def economic_impact_analysis( reform=reform, resolved_region_code=options.resolved_region_code, uk_local_authority_metadata=options.uk_local_authority_metadata, + omit_uk_area_breakdowns=options.omit_uk_area_breakdowns, ) @@ -1150,6 +1156,7 @@ def serialize(self): "resolved_region_code": "us", "runtime": runtime, "uk_local_authority_metadata": None, + "omit_uk_area_breakdowns": False, "progress": NULL_JOB_PROGRESS, } ] @@ -1461,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={ @@ -1512,6 +1522,585 @@ def serialize(self): assert metadata == UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 ) + assert builder_calls[0]["omit_uk_area_breakdowns"] is False + + +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 +) + + +_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.""" + + 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) + ] + + +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_region(self, *, params, **_kwargs): + 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): + if region_resolution.dataset_reference == _LOCAL_AREA_SELECTION.name: + return _LOCAL_AREA_SELECTION + return _NATIONAL_SELECTION + + 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) + self.run.builds.append( + { + "dataset": dataset, + "selection": dataset_selection, + "policy": policy, + "scoping_strategy": region.scoping_strategy, + "simulation": simulation, + } + ) + return simulation + + def output_builder(self, **kwargs): + self.run.builder_calls.append(kwargs) + 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.""" + + def __init__(self, run, kind, output): + self.run = run + self.kind = kind + self.output = output + + def serialize(self): + self.run.serialized.append(self.kind) + return self.output + + +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": 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, + "LocalAreaBreakdownBuilder": seams.breakdown_builder, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_runtime.{name}", value + ) + return run + + +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 _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, "region": region}, + local_area_dataset=_routed_area_dataset(), + ) + + # 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 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) + ) + assert "uk_area_simulations" not in builder_kwargs + assert run.serialized == ["macro"] + 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( + "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["local_authority_metadata"] == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + 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"}] + + +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 + + 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.serialized == ["local_area_breakdown"] + + +@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 = ( + "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) + ) + + +@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) + 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) + + with pytest.raises(ValueError, match="Unsupported UK region"): + _resolve_region( + country_module=_routed_uk_country_module(), + country="uk", + params={"region": requested}, + ) + + +@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) + country_module = SimpleNamespace(model=SimpleNamespace(get_region=lambda _: None)) + + 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(): @@ -2189,77 +2778,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)) - ) +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]] = [] - baseline = coded_simulation() - reform = coded_simulation() - expected = [_constituency_impact_record().model_dump(mode="json")] + def compute(country, baseline_simulation, reform_simulation, *, year): + computed.append((country, baseline_simulation, reform_simulation, year)) + return disclosure - def fake_output_module_function(module_name, name): - assert module_name == "constituency_impact" - assert name == "compute_uk_constituency_impacts" + def build_constituencies(country, value): + built.append(("constituency", value)) + return GeographicImpactOutput.model_validate([]) - 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) - - 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." - "_required_uk_geography_lookup_csv_path", - lambda _spec: CONSTITUENCY_CSV_PATH, + "build_uk_constituency_impact", + build_constituencies, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "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", @@ -2268,16 +2862,166 @@ 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)] + + +def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): + 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"), + ) + 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 + + +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 = 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, + ), + runtime=runtime, + 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() + + 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.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() + + 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() + + 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"]} 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..84636b533 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,150 @@ 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, inside={} + ) + + 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, 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")} diff --git a/projects/policyengine-simulation-executor/tests/test_smoke_app.py b/projects/policyengine-simulation-executor/tests/test_smoke_app.py index 3e778e46e..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 ( @@ -91,27 +92,129 @@ 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, - "detect_uk_local_authority_metadata_from_hdf", + "uk_area_hdf_dataset_path", + lambda path: area_dataset if path == installed else pytest.fail(path), + ) + monkeypatch.setattr( + smoke_module, + "detect_uk_area_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"] + + +_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) + + 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 00ae66afc..ccf4a1918 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -202,6 +202,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", @@ -215,6 +219,81 @@ 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_area_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] + + +@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) 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..b1dd1b800 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} @@ -148,10 +161,55 @@ def test_uk_plan_uses_uk_defaults_without_us_budget_variables() -> None: 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")) +@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", + "region_code_oa", + "source_household_id", + } + if planned + else set() + ) + + +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)) @@ -256,12 +314,17 @@ 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, ) -> 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 83ee55cf0..eca8612bf 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -667,10 +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 - ) - monkeypatch.setattr( - simulation_runtime, "_require_certified_uk_weight_matrix", lambda *_: None + worker, "detect_uk_run_dataset_metadata", lambda *_a, **_k: None ) monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) @@ -680,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 @@ -1832,6 +1926,66 @@ 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 + # 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( monkeypatch, ) -> None: @@ -1894,6 +2048,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"] == { @@ -1913,6 +2068,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: 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..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 @@ -10,10 +10,14 @@ 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, + detect_uk_run_dataset_metadata, 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, ) @@ -186,25 +190,264 @@ 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( + ("release_kind", "region_datasets", "routed"), + [ + ("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_to_one_local_area_dataset( + monkeypatch, + release_kind: str | None, + region_datasets: dict[str, str], + *, + routed: bool, +) -> None: + _bundle_routes(monkeypatch, region_datasets, release_kind=release_kind) + + assert uk_area_regions_routed() is routed + + @pytest.mark.parametrize( - ("region_types", "routed"), + ("country", "region", "routed", "omitted"), [ - (("national",), False), - (("national", "constituency"), False), - (("national", "constituency", "local_authority"), True), + ("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_regions_count_as_routed_only_when_both_have_a_dataset( - monkeypatch, region_types: tuple[str, ...], *, routed: bool +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 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: object() for region_type in region_types} - ), + lambda _country: manifest, ) - assert uk_area_regions_routed() is routed + +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)}) + ) + + +# 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") + + assert metadata.boundary_version is UKLocalAuthorityBoundaryVersion.LAD23 + + +@pytest.mark.parametrize("missing", list(_ROUTED_AREA_HOUSEHOLD)) +def test_routed_area_dataset_needs_every_column_area_outputs_read( + missing: str, +) -> None: + columns = dict(_ROUTED_AREA_HOUSEHOLD) + 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: @@ -258,3 +501,88 @@ 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 + + +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 + + +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", + ) 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/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index e4b56a658..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,6 +190,74 @@ 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( + ("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, + LOCAL_AREA_BREAKDOWN_CAPABILITY_VERSION, + ) + 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.provenance.manifest.get_release_manifest", + lambda _country: manifest, + ) + + metadata = registry.build_bundle_manifest_metadata( + app_name="policyengine-simulation-py6-10-0", + policyengine_version="6.10.0", + ) + + # 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 + ) + 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( 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]] diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index adea361f5..033836a95 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -89,6 +89,50 @@ 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`, + `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 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 + `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 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/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 76d3f626b..715cd2c5a 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": { @@ -1309,10 +1368,14 @@ { "$ref": "#/components/schemas/SingleYearMacroOutput" }, + { + "$ref": "#/components/schemas/LocalAreaBreakdownOutput" + }, { "type": "null" } - ] + ], + "title": "Result" }, "errors": { "anyOf": [ @@ -1680,8 +1743,119 @@ "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": { + "items": { + "$ref": "#/components/schemas/ConstituencyImpactRecord" + }, + "type": "array", + "title": "Constituency Impact" + }, + "local_authority_impact": { + "items": { + "$ref": "#/components/schemas/LocalAuthorityImpactRecord" + }, + "type": "array", + "title": "Local Authority Impact" + } + }, + "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. 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": { + "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" @@ -2353,6 +2527,18 @@ } ], "title": "Region" + }, + "breakdown": { + "anyOf": [ + { + "type": "string", + "const": "local_areas" + }, + { + "type": "null" + } + ], + "title": "Breakdown" } }, "additionalProperties": false, @@ -2361,7 +2547,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 ): 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..3a5c61f2b --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_local_area_breakdown.py @@ -0,0 +1,126 @@ +"""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_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] + 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 diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt index 77651b94b..50db66e6b 100644 --- a/rules/baselines/pyright-executor.txt +++ b/rules/baselines/pyright-executor.txt @@ -56,18 +56,9 @@ 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" -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/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..dac2bb1a7 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 = 68 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"