Admit UK datasets without area codes only when local-area runs are routed, and guard the weight-matrix fallback

changelog_entry.yaml

@@ -6,3 +6,6 @@
       - Ensured Stage 12-only deployments run authenticated candidate tests and verify every required deployment phase before reporting success
       - Allowed Cloud Run stable endpoint checks to wait for documented traffic propagation before restoring the previous revision
       - Added a dedicated read-only Hugging Face credential for private UK runtime datasets, validated it before synchronization, and tested its worker configuration
+      - 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

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py

@@ -210,11 +210,36 @@ def compute_and_format() -> CongressionalDistrictImpactOutput | None:
     return _try_compute_output("congressional district impacts", compute_and_format)
 
 
+def _output_household(simulation) -> pd.DataFrame:
+    return pd.DataFrame(simulation.output_dataset.data.household)
+
+
+def _require_uk_area_column(
+    household: pd.DataFrame, column: str, breakdown: str
+) -> None:
+    """Refuse a UK breakdown over a dataset that lacks its area codes.
+
+    A breakdown is never dropped silently: a dataset without area codes has
+    none of its own, and one built from a routed local-area dataset is a
+    separate run.
+    """
+
+    if column not in household.columns:
+        raise ValueError(
+            f"UK {breakdown} breakdowns need {column} on the simulated dataset, "
+            "which carries no area codes; for such a dataset they need a run on "
+            "the local-area dataset its bundle routes these regions to."
+        )
+
+
 def build_uk_constituency_impact(
     country: str, baseline, reform
 ) -> GeographicImpactOutput | None:
     if country != "uk":
         return None
+    _require_uk_area_column(
+        _output_household(baseline), "constituency_code_oa", "constituency"
+    )
 
     lookup_csv_path = _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC)
     impact = _output_module_function(
@@ -242,6 +267,8 @@ def build_uk_local_authority_impact(
     if country != "uk":
         return None
 
+    baseline_household = _output_household(baseline)
+    _require_uk_area_column(baseline_household, "la_code_oa", "local-authority")
     if uk_local_authority_metadata is None:
         raise ValueError("UK local-authority boundary metadata is required")
 
@@ -252,8 +279,7 @@ def build_uk_local_authority_impact(
         load_uk_local_authority_resources,
     )
 
-    baseline_household = pd.DataFrame(baseline.output_dataset.data.household)
-    reform_household = pd.DataFrame(reform.output_dataset.data.household)
+    reform_household = _output_household(reform)
     numeric_records = compute_longwise_uk_geography_impacts(
         baseline_household=baseline_household,
         reform_household=reform_household,

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py

@@ -269,10 +269,123 @@ def _build_uk_weight_replacement_region(region_code: str):
             lookup_csv_bucket=asset_spec.bucket,
             lookup_csv_key=asset_spec.lookup_csv_filename,
             region_code=value,
+            # The bucket copy is replaced by every data release; the matrix
+            # comes from the certified bundle instead, placed locally by
+            # ``_require_certified_uk_weight_matrix`` before the run.
+            download_missing_assets=False,
         ),
     )
 
 
+def _require_certified_uk_weight_matrix(
+    scoping_strategy, dataset, dataset_selection
+) -> None:
+    """Bind a weight-matrix region to the release of the dataset it reweights.
+
+    ``WeightReplacementStrategy`` gives each household the weight in its
+    column of the matrix, by position. Only the matrix certified with the
+    selected dataset (same repository and revision in the policyengine.py
+    bundle) lines up with its households: another release can have the same
+    shape and a different household order. This places that certified,
+    digest-checked matrix where the strategy looks first, confirms the
+    strategy will read exactly that file, and checks the run year and the
+    household dimension. The unversioned bucket copy, which every data
+    release replaces, is never downloaded (see
+    ``_build_uk_weight_replacement_region``).
+    """
+
+    from policyengine.core.scoping_strategy import WeightReplacementStrategy
+
+    if not isinstance(scoping_strategy, WeightReplacementStrategy):
+        return
+
+    from pathlib import Path
+
+    import h5py
+    import pandas as pd
+    from policyengine.data.uk_geography_assets import (
+        UKGeographyAssetSpec,
+        default_download_dir,
+        resolve_uk_geography_asset_paths,
+    )
+    from policyengine.provenance.dataset_materialization import materialize_dataset
+    from policyengine.provenance.manifest import get_release_manifest
+
+    from policyengine_simulation_executor import simulation_output_geographic
+
+    region = scoping_strategy.region_code
+    matrix_name = scoping_strategy.weight_matrix_key
+    manifest = get_release_manifest("uk")
+    package = manifest.data_package
+
+    def release(reference) -> tuple[str, str]:
+        return (
+            reference.repo_id or package.repo_id,
+            reference.revision or package.release_manifest_revision or package.version,
+        )
+
+    selected = manifest.datasets.get(dataset_selection.name)
+    if selected is None:
+        raise ValueError(
+            f"UK dataset {dataset_selection.name!r} is not in the certified bundle"
+        )
+    certified_name = next(
+        (
+            name
+            for name, reference in manifest.datasets.items()
+            if reference.path == matrix_name and release(reference) == release(selected)
+        ),
+        None,
+    )
+    if certified_name is None:
+        raise ValueError(
+            f"UK region {region!r} reweights households with {matrix_name}, but the "
+            f"certified bundle has no {matrix_name} from the release of "
+            f"{dataset_selection.name!r}; a matrix from another release does not "
+            "line up with its households."
+        )
+    certified = materialize_dataset(
+        "uk", certified_name, data_dir=default_download_dir()
+    )
+    spec = UKGeographyAssetSpec(
+        geography_type="weight replacement",
+        weight_matrix_filename=matrix_name,
+        lookup_csv_filename=scoping_strategy.lookup_csv_key,
+        bucket=scoping_strategy.weight_matrix_bucket,
+        weight_matrix_bucket=scoping_strategy.weight_matrix_bucket,
+        lookup_csv_bucket=scoping_strategy.lookup_csv_bucket,
+    )
+    simulation_output_geographic._required_uk_geography_lookup_csv_path(spec)
+    paths = resolve_uk_geography_asset_paths(spec, download_missing_assets=False)
+    if Path(paths.weight_matrix_path).resolve() != Path(certified.path).resolve():
+        raise ValueError(
+            f"UK region {region!r} would read {paths.weight_matrix_path}, not the "
+            f"{matrix_name} certified with {dataset_selection.name!r} "
+            f"({certified.path})."
+        )
+    year = str(dataset.year)
+    with h5py.File(paths.weight_matrix_path, "r") as matrix:
+        if year not in matrix:
+            covered = ", ".join(sorted(matrix))
+            raise ValueError(
+                f"UK region {region!r} reweights households with {matrix_name}, "
+                f"which has no weights for {year} (it covers {covered})."
+            )
+        weights = matrix[year]
+        if not isinstance(weights, h5py.Dataset):
+            raise TypeError(f"{matrix_name} entry {year} is not a weight matrix")
+        matrix_households = weights.shape[-1]
+    households = len(pd.DataFrame(dataset.data.entity_data["household"]))
+    if households != matrix_households:
+        raise ValueError(
+            f"UK region {region!r} reweights households with {matrix_name}, "
+            f"which was built for {matrix_households} households; the selected "
+            f"UK dataset has {households}. Constituency and local-authority "
+            "runs on this dataset need a local-area dataset that carries "
+            "constituency and local-authority codes."
+        )
+
+
 def _region_parent_dataset_reference(
     country_module,
     country: str,
@@ -605,7 +718,12 @@ def _run_simulation_impl_core(
         detect_uk_local_authority_metadata,
     )
 
-    uk_local_authority_metadata = detect_uk_local_authority_metadata(country, dataset)
+    uk_local_authority_metadata = detect_uk_local_authority_metadata(
+        country, dataset, region_code=region_resolution.code
+    )
+    _require_certified_uk_weight_matrix(
+        region_resolution.scoping_strategy, dataset, dataset_selection
+    )
     with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)):
         baseline_policy = _normalise_policy(simulation_params.get("baseline"))
         reform_policy = _normalise_policy(simulation_params.get("reform"))

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py

@@ -308,6 +308,7 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]:
             report.baseline.geography.country,
             baseline_calculation.uk_local_authority_metadata,
             reform_calculation.uk_local_authority_metadata,
+            region_code=report.baseline.geography.region,
         ),
     )
     v2_result = v2_report.get("result")

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py

@@ -21,12 +21,28 @@ def validate_uk_local_authority_metadata(
     country: CountryId,
     baseline: UKLocalAuthorityMetadata | None,
     reform: UKLocalAuthorityMetadata | None,
+    *,
+    region_code: str | None = None,
 ) -> UKLocalAuthorityMetadata | None:
-    """Require matching UK authority metadata and reject it elsewhere."""
+    """Require matching UK authority metadata and reject it elsewhere.
+
+    Retained artifacts are read without rerunning dataset detection, so a UK
+    pair without metadata passes only with the same evidence detection needs:
+    the report's region is national or nation-level and the bundle routes
+    both constituency and local-authority regions to another dataset.
+    """
 
     if country == "uk":
+        if baseline is None and reform is None:
+            from policyengine_simulation_executor.uk_local_authority_metadata import (
+                uk_area_codes_required,
+            )
+
+            if region_code is None or uk_area_codes_required(region_code):
+                raise ValueError("UK simulation artifact metadata is missing")
+            return None
         if baseline is None or reform is None:
-            raise ValueError("UK simulation artifact metadata is missing")
+            raise ValueError("UK simulation artifact metadata is missing on one side")
         if baseline != reform:
             raise ValueError("UK simulation artifact boundary versions do not match")
         return baseline

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py

@@ -420,6 +420,7 @@ def coordinate_report(
             report.baseline.geography.country,
             deserialize_uk_local_authority_metadata(baseline_payload),
             deserialize_uk_local_authority_metadata(reform_payload),
+            region_code=report.baseline.geography.region,
         )
         validate_output_frames(baseline_frames, output_plan)
         validate_output_frames(reform_frames, output_plan)

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py

@@ -130,6 +130,7 @@ def calculate_simulation_frames(
         _country_module,
         _load_dataset,
         _normalise_policy,
+        _require_certified_uk_weight_matrix,
         _resolve_dataset_selection,
         _resolve_region,
         setup_gcp_credentials,
@@ -183,6 +184,10 @@ def calculate_simulation_frames(
         uk_local_authority_metadata = detect_uk_local_authority_metadata(
             country,
             dataset,
+            region_code=region.code,
+        )
+        _require_certified_uk_weight_matrix(
+            region.scoping_strategy, dataset, dataset_selection
         )
         policy_span = (
             runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.POLICY_NORMALIZATION))

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py

@@ -80,9 +80,18 @@ def _check_uk_local_authority_resources() -> None:
 def _check_uk_local_authority_dataset(dataset_path: str) -> None:
     from policyengine_simulation_executor.uk_local_authority_metadata import (
         detect_uk_local_authority_metadata_from_hdf,
+        uk_area_regions_routed,
+        uk_hdf_household_has_local_authority_codes,
     )
 
     _check_uk_local_authority_resources()
+    # A national default may carry no area codes only when the bundle routes
+    # constituency and local-authority regions to another dataset; anything
+    # else is checked exactly as before.
+    if uk_area_regions_routed() and not uk_hdf_household_has_local_authority_codes(
+        dataset_path
+    ):
+        return
     detect_uk_local_authority_metadata_from_hdf(dataset_path)
 
 

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py

@@ -277,11 +277,55 @@ def detect_uk_local_authority_boundary_version(
     return UKLocalAuthorityMetadata(boundary_version=matches[0].boundary_version)
 
 
+UK_AREA_REGION_TYPES = ("constituency", "local_authority")
+_MISSING_AREA_CODES = (
+    "UK dataset household table contains no la_code_oa column; a UK dataset "
+    "may omit area codes only when the policyengine.py bundle routes "
+    "constituency and local-authority regions to another dataset"
+)
+
+
+def uk_area_regions_routed() -> bool:
+    """Whether the bundle routes UK constituency and local-authority regions.
+
+    policyengine.py lists a region type in ``region_datasets`` when its
+    ``regional_dataset_defaults`` names a dataset for it. Only when both area
+    region types have one may the UK default, a national file, carry no area
+    codes: their runs and breakdowns then belong to that local-area dataset.
+    """
+
+    from policyengine.provenance.manifest import get_release_manifest
+
+    region_datasets = get_release_manifest("uk").region_datasets
+    return all(region_type in region_datasets for region_type in UK_AREA_REGION_TYPES)
+
+
+def uk_area_codes_required(region_code: str | None) -> bool:
+    """Whether a UK run's dataset must carry area codes.
+
+    A constituency or local-authority run always needs them. A national or
+    nation-level run may use a dataset without them only when the bundle
+    routes both area region types to another dataset.
+    """
+
+    region_type = (region_code or "").split("/", maxsplit=1)[0]
+    if region_type in UK_AREA_REGION_TYPES:
+        return True
+    return not uk_area_regions_routed()
+
+
 def detect_uk_local_authority_metadata(
     country: str,
     dataset: object,
+    *,
+    region_code: str | None = None,
 ) -> UKLocalAuthorityMetadata | None:
-    """Inspect a complete dataset before any requested regional scoping."""
+    """Inspect a complete dataset before any requested regional scoping.
+
+    A UK dataset without ``la_code_oa`` passes, as ``None``, only for a run
+    that does not need area codes (see ``uk_area_codes_required``); a missing
+    column is never read as "this dataset is national-only" on its own.
+    """
 
     if country != "uk":
         return None
@@ -294,12 +338,26 @@ def detect_uk_local_authority_metadata(
         raise ValueError("UK dataset contains no household table")
     household_frame = pd.DataFrame(household)
     if "la_code_oa" not in household_frame:
-        raise ValueError("UK dataset household table contains no la_code_oa column")
+        if uk_area_codes_required(region_code):
+            raise ValueError(_MISSING_AREA_CODES)
+        return None
     return detect_uk_local_authority_boundary_version(
         household_frame["la_code_oa"].tolist()
     )
 
 
+def uk_hdf_household_has_local_authority_codes(dataset_path: str) -> bool:
+    """Whether an installed UK HDF dataset's household table has ``la_code_oa``."""
+
+    with pd.HDFStore(dataset_path, mode="r") as store:
+        if "household" not in store:
+            raise ValueError("UK dataset contains no household table")
+        household = store.select("household", stop=0)
+    if not isinstance(household, pd.DataFrame):
+        raise TypeError("UK dataset household table is not a data frame")
+    return "la_code_oa" in household.columns
+
+
 def detect_uk_local_authority_metadata_from_hdf(
     dataset_path: str,
 ) -> UKLocalAuthorityMetadata:

projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py

@@ -1752,8 +1752,14 @@ def fail_change_output_variable(*args, **kwargs):
 
 
 def test_uk_constituency_impact_uses_policyengine_output_function(monkeypatch):
-    baseline = object()
-    reform = object()
+    def coded_simulation() -> SimpleNamespace:
+        household = pd.DataFrame({"constituency_code_oa": ["E14001063"]})
+        return SimpleNamespace(
+            output_dataset=SimpleNamespace(data=SimpleNamespace(household=household))
+        )
+
+    baseline = coded_simulation()
+    reform = coded_simulation()
     expected = [_constituency_impact_record().model_dump(mode="json")]
 
     def fake_output_module_function(module_name, name):