From ee518c91a4d4b7b9a72a8afb7f0bc46449d6386c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 19 Jun 2026 10:56:40 -0400 Subject: [PATCH] Add UK local candidate runner --- packages/populace-build/README.md | 10 + packages/populace-build/pyproject.toml | 4 + .../src/populace/build/uk/__init__.py | 26 + .../src/populace/build/uk/local_geography.py | 56 +- .../src/populace/build/uk/local_runner.py | 590 ++++++++++++++++++ .../tests/test_uk_local_geography.py | 51 ++ .../tests/test_uk_local_runner.py | 296 +++++++++ uv.lock | 8 +- 8 files changed, 1037 insertions(+), 4 deletions(-) create mode 100644 packages/populace-build/src/populace/build/uk/local_runner.py create mode 100644 packages/populace-build/tests/test_uk_local_runner.py diff --git a/packages/populace-build/README.md b/packages/populace-build/README.md index 3d678f8a..786cf6d7 100644 --- a/packages/populace-build/README.md +++ b/packages/populace-build/README.md @@ -67,6 +67,16 @@ wraps the Populace calibrator's log-weight optimizer for stacked local weights and records per-area/per-metric diagnostics before the solved weights are exported with `stacked_weights_to_long`. +`populace.build.uk.local_runner` is the Populace-owned candidate build path. It +loads explicit area and target tables, aligns a sorted household frame with +source-year/source-household/clone lineage, optionally computes household +metrics once per UK country by setting the PolicyEngine-UK `region` input, then +solves and writes `local_geography_weights.csv.gz`, +`solve_diagnostics.csv`, `area_support_summary.csv`, and `solve_summary.json`. +It accepts already-pooled or already-cloned household pools, so the compact UK +artifact can remain the fast national default while a separate `local` variant +scales up with pooled FRS years, cloned records, and L0 budget control. + ## US plan status `populace.build.us` declares the US build: stage order, donor graph with diff --git a/packages/populace-build/pyproject.toml b/packages/populace-build/pyproject.toml index 5fd4c295..72ed0d1a 100644 --- a/packages/populace-build/pyproject.toml +++ b/packages/populace-build/pyproject.toml @@ -20,6 +20,10 @@ dependencies = [ # content is declared in packaged manifests and interpreted by shared Populace # runtimes; country source loaders must not depend on incumbent data packages. us = ["policyengine-us>=1.729,<2", "h5py>=3"] +# The UK extra adds the rules engine for local metric generation from a +# Populace UK H5. Target tables remain explicit inputs, and the base package +# still does not import policyengine-uk at import time. +uk = ["policyengine-uk>=2.88", "h5py>=3"] [project.urls] Homepage = "https://populace.dev" diff --git a/packages/populace-build/src/populace/build/uk/__init__.py b/packages/populace-build/src/populace/build/uk/__init__.py index 68b887f5..29e94598 100644 --- a/packages/populace-build/src/populace/build/uk/__init__.py +++ b/packages/populace-build/src/populace/build/uk/__init__.py @@ -11,6 +11,20 @@ stacked_weights_to_long, write_long_geography_weights, ) +from populace.build.uk.local_runner import ( + UKLocalCandidateResult, + build_local_candidate, + build_local_candidate_from_dataset, + build_metric_tables_from_dataset, + load_metric_tables, + load_uk_dataset, + prepare_area_frame, + prepare_household_frame, + read_local_table, + set_simulation_area_group, + summarize_local_candidate, + write_local_candidate_outputs, +) from populace.build.uk.local_solver import ( StackedLocalSolveResult, solve_stacked_local_weights, @@ -36,16 +50,28 @@ "LONG_GEOGRAPHY_COLUMNS", "StackedLocalMatrix", "StackedLocalSolveResult", + "UKLocalCandidateResult", "align_area_targets", "area_support_summary", "area_groups_from_codes", + "build_local_candidate", + "build_local_candidate_from_dataset", + "build_metric_tables_from_dataset", "build_stacked_local_matrix", "compute_household_metrics", + "load_metric_tables", + "load_uk_dataset", "metric_names", "metric_tables_by_area_group", + "prepare_area_frame", + "prepare_household_frame", + "read_local_table", + "set_simulation_area_group", "solve_stacked_local_weights", "sort_households_by_id", "stacked_design_weights", "stacked_weights_to_long", + "summarize_local_candidate", + "write_local_candidate_outputs", "write_long_geography_weights", ] diff --git a/packages/populace-build/src/populace/build/uk/local_geography.py b/packages/populace-build/src/populace/build/uk/local_geography.py index d0c793f0..f87c1647 100644 --- a/packages/populace-build/src/populace/build/uk/local_geography.py +++ b/packages/populace-build/src/populace/build/uk/local_geography.py @@ -362,22 +362,64 @@ def stacked_weights_to_long( return out.loc[:, LONG_GEOGRAPHY_COLUMNS] -def area_support_summary(long_weights: pd.DataFrame) -> pd.DataFrame: - """Summarize non-zero household support by local area.""" +def area_support_summary( + long_weights: pd.DataFrame, + *, + area_codes: Sequence[str] | None = None, + area_type: str | None = None, +) -> pd.DataFrame: + """Summarize non-zero household support by local area. + + Passing ``area_codes`` includes requested areas with no positive assigned + households, which is important for sparse/L0 local solves. + """ missing = sorted(set(LONG_GEOGRAPHY_COLUMNS) - set(long_weights.columns)) if missing: raise ValueError(f"long weight frame is missing column(s): {missing}.") positive = long_weights[long_weights["weight"] > 0] - return ( + summary = ( positive.groupby(["area_type", "area_code"], sort=True) .agg( nonzero_households=("household_id", "nunique"), + nonzero_source_households=("source_household_key", "nunique"), weight_sum=("weight", "sum"), max_weight=("weight", "max"), + effective_sample_size=("weight", _effective_sample_size), ) .reset_index() ) + if area_codes is None: + return summary + + codes = _area_code_tuple(area_codes) + if area_type is None: + area_types = long_weights["area_type"].dropna().unique() + if len(area_types) != 1: + raise ValueError( + "area_type must be supplied when area_codes are supplied and " + "long_weights does not contain exactly one area_type." + ) + area_type = str(area_types[0]) + full = pd.DataFrame( + { + "area_type": area_type, + "area_code": list(codes), + } + ) + completed = full.merge(summary, on=["area_type", "area_code"], how="left") + completed["nonzero_households"] = ( + completed["nonzero_households"].fillna(0).astype(int) + ) + completed["nonzero_source_households"] = ( + completed["nonzero_source_households"].fillna(0).astype(int) + ) + completed["weight_sum"] = completed["weight_sum"].fillna(0.0).astype(float) + completed["max_weight"] = completed["max_weight"].fillna(0.0).astype(float) + completed["effective_sample_size"] = ( + completed["effective_sample_size"].fillna(0.0).astype(float) + ) + return completed def write_long_geography_weights( @@ -551,3 +593,11 @@ def _source_keys( else: keys.append(f"{year}:{household_id}") return np.asarray(keys, dtype=object) + + +def _effective_sample_size(weights: pd.Series) -> float: + values = weights.to_numpy(dtype=np.float64) + square_sum = float(np.square(values).sum()) + if square_sum == 0: + return 0.0 + return float(values.sum() ** 2 / square_sum) diff --git a/packages/populace-build/src/populace/build/uk/local_runner.py b/packages/populace-build/src/populace/build/uk/local_runner.py new file mode 100644 index 00000000..04b2e1ad --- /dev/null +++ b/packages/populace-build/src/populace/build/uk/local_runner.py @@ -0,0 +1,590 @@ +"""Reusable runner pieces for Populace-owned UK local builds. + +The core local-geography modules deliberately avoid data-package imports: +target providers hand Populace explicit area tables, and engine runners hand it +household metric tables. This module is the thin orchestration layer that ties +those pieces together for pilot and full UK local candidate builds. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd + +from populace.build.uk.local_geography import ( + StackedLocalMatrix, + area_support_summary, + build_stacked_local_matrix, + sort_households_by_id, + stacked_weights_to_long, + write_long_geography_weights, +) +from populace.build.uk.local_solver import ( + StackedLocalSolveResult, + solve_stacked_local_weights, +) +from populace.build.uk.local_targets import ( + COUNTRY_TO_REGION, + area_groups_from_codes, + compute_household_metrics, +) + + +@dataclass(frozen=True) +class UKLocalCandidateResult: + """Solved UK local candidate outputs and diagnostics.""" + + problem: StackedLocalMatrix + solve_result: StackedLocalSolveResult + long_weights: pd.DataFrame + support_summary: pd.DataFrame + + +def read_local_table(path: str | Path) -> pd.DataFrame: + """Read a local-build input table from CSV, CSV.GZ, or Parquet.""" + + table_path = Path(path) + suffixes = [suffix.lower() for suffix in table_path.suffixes] + if not suffixes: + raise ValueError(f"Cannot infer table format for {table_path}.") + if suffixes[-1] == ".csv" or suffixes[-2:] == [".csv", ".gz"]: + return pd.read_csv(table_path) + if suffixes[-1] in {".parquet", ".pq"}: + return pd.read_parquet(table_path) + raise ValueError( + f"Unsupported local table format for {table_path}; expected CSV, " + "CSV.GZ, or Parquet." + ) + + +def prepare_area_frame( + area_frame: pd.DataFrame | str | Path, + *, + code_column: str = "code", + group_column: str = "country", + sort_by_code: bool = True, + max_areas: int | None = None, +) -> pd.DataFrame: + """Return a validated, canonical area frame. + + ``code_column`` supplies the output area order. If ``group_column`` is + present it is authoritative for country/devolution grouping; otherwise the + ONS area-code prefix is used later by :func:`area_groups_from_codes`. + """ + + frame = _as_frame(area_frame).copy() + if code_column not in frame.columns: + raise ValueError(f"area frame is missing {code_column!r}.") + frame[code_column] = _normalise_nonblank_strings( + frame[code_column], + column=code_column, + ) + if frame[code_column].duplicated().any(): + duplicates = frame.loc[frame[code_column].duplicated(), code_column].unique() + raise ValueError( + f"area codes must be unique; duplicate value(s): " + f"{list(map(str, duplicates[:5]))}." + ) + + if group_column in frame.columns: + frame[group_column] = _normalise_nonblank_strings( + frame[group_column], + column=group_column, + ) + + if sort_by_code: + frame = frame.sort_values(code_column, kind="mergesort") + if max_areas is not None: + if not isinstance(max_areas, int) or max_areas <= 0: + raise ValueError("max_areas must be a positive integer when supplied.") + frame = frame.head(max_areas) + return frame.reset_index(drop=True) + + +def prepare_household_frame( + household_frame: pd.DataFrame | str | Path, + *, + id_column: str = "household_id", + weight_column: str = "household_weight", + source_year: int | None = None, +) -> pd.DataFrame: + """Sort households by ID, validate weights, and attach lineage columns.""" + + frame = sort_households_by_id(_as_frame(household_frame), id_column=id_column) + if id_column != "household_id": + frame = frame.rename(columns={id_column: "household_id"}) + if weight_column not in frame.columns: + raise ValueError(f"household frame is missing {weight_column!r}.") + weights = frame[weight_column].to_numpy(dtype=np.float64) + if not np.isfinite(weights).all() or (weights < 0).any(): + raise ValueError(f"{weight_column!r} must be finite and non-negative.") + + if weight_column != "household_weight": + frame = frame.rename(columns={weight_column: "household_weight"}) + if "source_household_id" not in frame.columns: + frame["source_household_id"] = frame["household_id"] + if "source_year" not in frame.columns and source_year is not None: + frame["source_year"] = source_year + if "clone_index" not in frame.columns: + frame["clone_index"] = 0 + if "source_household_key" not in frame.columns: + frame["source_household_key"] = _source_household_keys( + frame, + source_year=source_year, + ) + return frame.reset_index(drop=True) + + +def load_metric_tables( + paths: Mapping[str, str | Path], + *, + household_id_column: str = "household_id", +) -> dict[str, pd.DataFrame]: + """Load explicit household metric tables keyed by area group/country.""" + + if not paths: + raise ValueError("metric table paths must not be empty.") + return { + str(group): _metric_table_from_frame( + read_local_table(path), + household_id_column=household_id_column, + group=str(group), + ) + for group, path in paths.items() + } + + +def load_uk_dataset(path: str | Path) -> Any: + """Load a PolicyEngine-UK single-year H5 dataset lazily.""" + + try: + from policyengine_uk.data import UKSingleYearDataset + except ImportError as exc: # pragma: no cover - exercised only with engine absent + raise ImportError( + "Loading a UK H5 dataset requires policyengine-uk. Install the " + "UK engine before calling load_uk_dataset()." + ) from exc + return UKSingleYearDataset(file_path=str(path)) + + +def set_simulation_area_group( + sim: Any, + group: str, + *, + period: int | str, + n_households: int | None = None, + region_variable: str = "region", +) -> Any: + """Set a PolicyEngine-UK-like simulation to the area's devolution group.""" + + if group not in COUNTRY_TO_REGION: + raise ValueError( + f"Unknown UK area group {group!r}; expected one of " + f"{tuple(COUNTRY_TO_REGION)}." + ) + if n_households is None: + household_ids = _values( + sim.calculate("household_id", period=period, map_to="household") + ) + n_households = len(household_ids) + sim.set_input( + region_variable, + period, + np.asarray([COUNTRY_TO_REGION[group]] * n_households, dtype=object), + ) + return sim + + +def build_metric_tables_from_dataset( + dataset: Any, + area_groups: Mapping[str, str], + area_type: str, + *, + period: int | str | None = None, + household_ids: Sequence[Any] | None = None, + simulation_factory: Callable[[Any], Any] | None = None, +) -> dict[str, pd.DataFrame]: + """Compute household metric tables once per country/devolution group.""" + + if not area_groups: + raise ValueError("area_groups must not be empty.") + run_period = _infer_period(dataset, period) + if simulation_factory is None: + simulation_factory = _default_uk_simulation_factory + tables: dict[str, pd.DataFrame] = {} + for group in sorted(set(area_groups.values())): + sim = simulation_factory(dataset) + simulation_household_ids = _simulation_household_ids(sim, run_period) + set_simulation_area_group( + sim, + group, + period=run_period, + n_households=len(simulation_household_ids), + ) + table = compute_household_metrics( + sim, + area_type, + period=run_period, + ) + if not table.index.equals(pd.Index(simulation_household_ids)): + raise ValueError( + f"metric table {group!r} index must match simulation " + "household_id order." + ) + if household_ids is not None: + table = _align_metric_table_to_households( + table, + household_ids, + group=group, + ) + tables[group] = table + return tables + + +def build_local_candidate( + *, + area_type: str, + area_frame: pd.DataFrame | str | Path, + targets: pd.DataFrame | str | Path, + metrics: pd.DataFrame | Mapping[str, pd.DataFrame], + household_frame: pd.DataFrame | str | Path, + code_column: str = "code", + group_column: str = "country", + sort_areas_by_code: bool = True, + max_areas: int | None = None, + source_year: int | None = None, + weight_source: str = "populace_uk_local", + solver_options: Mapping[str, Any] | None = None, +) -> UKLocalCandidateResult: + """Build, solve, and export a UK local candidate in longwise form.""" + + areas = prepare_area_frame( + area_frame, + code_column=code_column, + group_column=group_column, + sort_by_code=sort_areas_by_code, + max_areas=max_areas, + ) + households = prepare_household_frame( + household_frame, + source_year=source_year, + ) + area_codes = tuple(areas[code_column].astype(str)) + area_groups = area_groups_from_codes( + areas, + code_column=code_column, + group_column=group_column, + ) + household_ids = households["household_id"].to_numpy() + base_weights = households["household_weight"].to_numpy(dtype=np.float64) + target_frame = _as_frame(targets) + problem = build_stacked_local_matrix( + metrics, + target_frame, + area_codes=area_codes, + area_groups=area_groups, + household_ids=household_ids, + area_type=area_type, + code_column=code_column, + ) + solve_result = solve_stacked_local_weights( + problem, + base_weights, + **dict(solver_options or {}), + ) + long_weights = stacked_weights_to_long( + solve_result.weights, + area_codes, + household_ids, + area_type=area_type, + household_frame=households, + source_year=source_year, + weight_source=weight_source, + ) + return UKLocalCandidateResult( + problem=problem, + solve_result=solve_result, + long_weights=long_weights, + support_summary=area_support_summary( + long_weights, + area_codes=area_codes, + area_type=area_type, + ), + ) + + +def build_local_candidate_from_dataset( + dataset: Any | str | Path, + *, + area_type: str, + area_frame: pd.DataFrame | str | Path, + targets: pd.DataFrame | str | Path, + household_frame: pd.DataFrame | str | Path, + period: int | str | None = None, + code_column: str = "code", + group_column: str = "country", + sort_areas_by_code: bool = True, + max_areas: int | None = None, + source_year: int | None = None, + weight_source: str = "populace_uk_local", + simulation_factory: Callable[[Any], Any] | None = None, + solver_options: Mapping[str, Any] | None = None, +) -> UKLocalCandidateResult: + """Build a UK local candidate from a Populace UK H5 or dataset object.""" + + dataset_obj = load_uk_dataset(dataset) if isinstance(dataset, str | Path) else dataset + areas = prepare_area_frame( + area_frame, + code_column=code_column, + group_column=group_column, + sort_by_code=sort_areas_by_code, + max_areas=max_areas, + ) + households = prepare_household_frame( + household_frame, + source_year=source_year, + ) + area_groups = area_groups_from_codes( + areas, + code_column=code_column, + group_column=group_column, + ) + household_ids = households["household_id"].to_numpy() + metrics = build_metric_tables_from_dataset( + dataset_obj, + area_groups, + area_type, + period=period, + household_ids=household_ids, + simulation_factory=simulation_factory, + ) + return build_local_candidate( + area_type=area_type, + area_frame=areas, + targets=targets, + metrics=metrics, + household_frame=households, + code_column=code_column, + group_column=group_column, + sort_areas_by_code=False, + source_year=source_year, + weight_source=weight_source, + solver_options=solver_options, + ) + + +def summarize_local_candidate(result: UKLocalCandidateResult) -> dict[str, Any]: + """Return a compact JSON-serializable summary for candidate run logs.""" + + support = result.support_summary + return { + "area_type": ( + None + if result.solve_result.diagnostics.empty + else str(result.solve_result.diagnostics["area_type"].iloc[0]) + ), + "n_areas": int(result.problem.n_areas), + "n_households": int(result.problem.n_households), + "n_targets": int(len(result.problem.targets)), + "n_long_rows": int(len(result.long_weights)), + "n_nonzero": int(result.solve_result.n_nonzero), + "initial_loss": float(result.solve_result.initial_loss), + "final_loss": float(result.solve_result.final_loss), + "weight_sum": float(result.long_weights["weight"].sum()), + "min_area_support": ( + 0 if support.empty else int(support["nonzero_households"].min()) + ), + "median_area_support": ( + 0.0 if support.empty else float(support["nonzero_households"].median()) + ), + "max_area_support": ( + 0 if support.empty else int(support["nonzero_households"].max()) + ), + "min_area_source_support": ( + 0 if support.empty else int(support["nonzero_source_households"].min()) + ), + "median_area_source_support": ( + 0.0 + if support.empty + else float(support["nonzero_source_households"].median()) + ), + "max_area_source_support": ( + 0 if support.empty else int(support["nonzero_source_households"].max()) + ), + "min_area_effective_sample_size": ( + 0.0 + if support.empty + else float(support["effective_sample_size"].min()) + ), + "median_area_effective_sample_size": ( + 0.0 + if support.empty + else float(support["effective_sample_size"].median()) + ), + } + + +def write_local_candidate_outputs( + result: UKLocalCandidateResult, + output_dir: str | Path, + *, + weights_filename: str = "local_geography_weights.csv.gz", +) -> dict[str, Any]: + """Write long weights, diagnostics, support, and a JSON summary.""" + + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + write_long_geography_weights(result.long_weights, out / weights_filename) + result.solve_result.diagnostics.to_csv(out / "solve_diagnostics.csv", index=False) + result.support_summary.to_csv(out / "area_support_summary.csv", index=False) + summary = summarize_local_candidate(result) + (out / "solve_summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + ) + return summary + + +def _as_frame(frame_or_path: pd.DataFrame | str | Path) -> pd.DataFrame: + if isinstance(frame_or_path, pd.DataFrame): + return frame_or_path.copy() + return read_local_table(frame_or_path) + + +def _normalise_nonblank_strings(values: pd.Series, *, column: str) -> pd.Series: + if values.isna().any(): + raise ValueError(f"{column!r} must not contain missing values.") + strings = values.astype(str).str.strip() + if (strings == "").any(): + raise ValueError(f"{column!r} must not contain blank values.") + return strings + + +def _source_household_keys( + household_frame: pd.DataFrame, + *, + source_year: int | None, +) -> list[str]: + if "source_year" in household_frame.columns: + years = household_frame["source_year"].tolist() + else: + years = [source_year] * len(household_frame) + keys = [] + for year, source_id in zip( + years, + household_frame["source_household_id"], + strict=True, + ): + if year is None or pd.isna(year): + keys.append(str(source_id)) + else: + keys.append(f"{year}:{source_id}") + return keys + + +def _metric_table_from_frame( + frame: pd.DataFrame, + *, + household_id_column: str, + group: str, +) -> pd.DataFrame: + if household_id_column not in frame.columns: + raise ValueError( + f"metric table {group!r} is missing {household_id_column!r}." + ) + table = frame.copy() + if table[household_id_column].isna().any(): + raise ValueError( + f"metric table {group!r} {household_id_column!r} must not contain " + "missing values." + ) + if pd.api.types.is_string_dtype(table[household_id_column]): + strings = table[household_id_column].astype(str).str.strip() + if (strings == "").any(): + raise ValueError( + f"metric table {group!r} {household_id_column!r} must not " + "contain blank values." + ) + table[household_id_column] = strings + if table[household_id_column].duplicated().any(): + duplicates = table.loc[ + table[household_id_column].duplicated(), + household_id_column, + ].unique() + raise ValueError( + f"metric table {group!r} household IDs must be unique; " + f"duplicate value(s): {list(map(str, duplicates[:5]))}." + ) + return table.set_index(household_id_column, drop=True) + + +def _simulation_household_ids(sim: Any, period: int | str) -> np.ndarray: + return _values(sim.calculate("household_id", period=period, map_to="household")) + + +def _align_metric_table_to_households( + table: pd.DataFrame, + household_ids: Sequence[Any], + *, + group: str, +) -> pd.DataFrame: + expected = pd.Index(household_ids) + if expected.has_duplicates: + duplicates = expected[expected.duplicated()].unique() + raise ValueError( + "household_ids must be unique before metric alignment; duplicate " + f"value(s): {list(map(str, duplicates[:5]))}." + ) + if table.index.has_duplicates: + duplicates = table.index[table.index.duplicated()].unique() + raise ValueError( + f"metric table {group!r} household index must be unique; " + f"duplicate value(s): {list(map(str, duplicates[:5]))}." + ) + missing = expected.difference(table.index) + if len(missing): + raise ValueError( + f"metric table {group!r} is missing household_id value(s): " + f"{list(map(str, missing[:5]))}." + ) + extra = table.index.difference(expected) + if len(extra): + raise ValueError( + f"metric table {group!r} has unexpected household_id value(s): " + f"{list(map(str, extra[:5]))}." + ) + return table.reindex(expected) + + +def _infer_period(dataset: Any, period: int | str | None) -> int | str: + if period is not None: + return period + for attr in ("time_period", "fiscal_year", "default_calculation_period"): + value = getattr(dataset, attr, None) + if value is not None: + return value + raise ValueError( + "period is required when it cannot be inferred from the dataset." + ) + + +def _default_uk_simulation_factory(dataset: Any) -> Any: + try: + from policyengine_uk import Microsimulation + except ImportError as exc: # pragma: no cover - exercised only with engine absent + raise ImportError( + "Computing UK local metrics requires policyengine-uk. Install the " + "UK engine or pass explicit metric tables." + ) from exc + return Microsimulation(dataset=dataset) + + +def _values(result: Any) -> np.ndarray: + if hasattr(result, "values"): + return np.asarray(result.values) + return np.asarray(result) diff --git a/packages/populace-build/tests/test_uk_local_geography.py b/packages/populace-build/tests/test_uk_local_geography.py index 736b9ebc..9ba714f9 100644 --- a/packages/populace-build/tests/test_uk_local_geography.py +++ b/packages/populace-build/tests/test_uk_local_geography.py @@ -182,4 +182,55 @@ def test_area_support_summary_counts_nonzero_households() -> None: assert summary["area_code"].tolist() == ["E001", "S001"] assert summary["nonzero_households"].tolist() == [1, 2] + assert summary["nonzero_source_households"].tolist() == [1, 2] assert summary["weight_sum"].tolist() == [0.5, 3.5] + np.testing.assert_allclose( + summary["effective_sample_size"], + [1.0, (1.5 + 2.0) ** 2 / (1.5**2 + 2.0**2)], + ) + + +def test_area_support_summary_includes_zero_support_areas() -> None: + long = stacked_weights_to_long( + [1.0, 0.0], + ["E001", "S001"], + [101], + area_type="constituency", + ) + + summary = area_support_summary( + long, + area_codes=["E001", "S001"], + area_type="constituency", + ) + + assert summary["area_code"].tolist() == ["E001", "S001"] + assert summary["nonzero_households"].tolist() == [1, 0] + assert summary["nonzero_source_households"].tolist() == [1, 0] + assert summary["weight_sum"].tolist() == [1.0, 0.0] + assert summary["effective_sample_size"].tolist() == [1.0, 0.0] + + +def test_area_support_summary_counts_cloned_source_support() -> None: + household_frame = pd.DataFrame( + { + "household_id": [101, 102], + "source_year": [2023, 2023], + "source_household_id": ["a", "a"], + "source_household_key": ["2023:a", "2023:a"], + "clone_index": [0, 1], + } + ) + long = stacked_weights_to_long( + [1.0, 1.0], + ["E001"], + [101, 102], + area_type="constituency", + household_frame=household_frame, + ) + + summary = area_support_summary(long) + + assert summary["nonzero_households"].tolist() == [2] + assert summary["nonzero_source_households"].tolist() == [1] + assert summary["effective_sample_size"].tolist() == [2.0] diff --git a/packages/populace-build/tests/test_uk_local_runner.py b/packages/populace-build/tests/test_uk_local_runner.py new file mode 100644 index 00000000..0ae45b74 --- /dev/null +++ b/packages/populace-build/tests/test_uk_local_runner.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +import populace.build.uk.local_runner as local_runner +from populace.build.uk import ( + build_local_candidate, + build_local_candidate_from_dataset, + build_metric_tables_from_dataset, + load_metric_tables, + prepare_area_frame, + prepare_household_frame, + read_local_table, + set_simulation_area_group, + write_local_candidate_outputs, +) + + +class Result: + def __init__(self, values): + self.values = np.asarray(values) + + +class FakeSimulation: + def __init__(self, _dataset): + self.inputs = {} + + def calculate(self, variable, **_kwargs): + assert variable == "household_id" + return Result([101, 102]) + + def set_input(self, variable, period, values): + self.inputs[(variable, period)] = list(values) + + +class SingleHouseholdSimulation(FakeSimulation): + def calculate(self, variable, **_kwargs): + assert variable == "household_id" + return Result([101]) + + +def test_prepare_area_frame_sorts_and_validates_codes() -> None: + areas = pd.DataFrame( + {"code": ["S001", "E001"], "country": ["Scotland", "England"]} + ) + + prepared = prepare_area_frame(areas) + + assert prepared["code"].tolist() == ["E001", "S001"] + assert prepared["country"].tolist() == ["England", "Scotland"] + + +def test_prepare_area_frame_rejects_duplicate_codes() -> None: + areas = pd.DataFrame({"code": ["E001", "E001"]}) + + with pytest.raises(ValueError, match="unique"): + prepare_area_frame(areas) + + +def test_prepare_household_frame_sorts_weights_and_fills_lineage() -> None: + households = pd.DataFrame( + {"household_id": [102, 101], "household_weight": [2.0, 1.0]} + ) + + prepared = prepare_household_frame(households, source_year=2023) + + assert prepared["household_id"].tolist() == [101, 102] + assert prepared["household_weight"].tolist() == [1.0, 2.0] + assert prepared["source_household_id"].tolist() == [101, 102] + assert prepared["source_year"].tolist() == [2023, 2023] + assert prepared["clone_index"].tolist() == [0, 0] + assert prepared["source_household_key"].tolist() == ["2023:101", "2023:102"] + + +def test_read_local_table_and_load_metric_tables(tmp_path: Path) -> None: + path = tmp_path / "metrics.csv" + pd.DataFrame( + { + "household_id": [101, 102], + "population": [1.0, 2.0], + } + ).to_csv(path, index=False) + + table = read_local_table(path) + metrics = load_metric_tables({"England": path}) + + assert table["population"].tolist() == [1.0, 2.0] + assert metrics["England"].index.tolist() == [101, 102] + assert metrics["England"]["population"].tolist() == [1.0, 2.0] + + +def test_set_simulation_area_group_sets_country_region() -> None: + sim = FakeSimulation(None) + + set_simulation_area_group(sim, "Scotland", period=2023) + + assert sim.inputs[("region", 2023)] == ["SCOTLAND", "SCOTLAND"] + + +def test_build_metric_tables_from_dataset_sets_each_country(monkeypatch) -> None: + calls = [] + + def factory(dataset): + sim = FakeSimulation(dataset) + calls.append(sim) + return sim + + def fake_compute(sim, area_type, *, period=None, household_ids=None): + assert household_ids is None + return pd.DataFrame( + {"population": [1.0, 2.0]}, + index=pd.Index([101, 102]), + ) + + monkeypatch.setattr(local_runner, "compute_household_metrics", fake_compute) + + tables = build_metric_tables_from_dataset( + dataset=type("Dataset", (), {"time_period": 2023})(), + area_groups={"E001": "England", "S001": "Scotland"}, + area_type="constituency", + household_ids=[101, 102], + simulation_factory=factory, + ) + + assert set(tables) == {"England", "Scotland"} + assert tables["England"].index.tolist() == [101, 102] + regions = [ + sim.inputs[("region", 2023)][0] + for sim in calls + ] + assert regions == ["SOUTH_EAST", "SCOTLAND"] + + +def test_build_metric_tables_from_dataset_reindexes_simulation_household_order( + monkeypatch, +) -> None: + class ReversedSimulation(FakeSimulation): + def calculate(self, variable, **_kwargs): + assert variable == "household_id" + return Result([102, 101]) + + def fake_compute(sim, area_type, *, period=None, household_ids=None): + assert household_ids is None + return pd.DataFrame( + {"population": [20.0, 10.0]}, + index=pd.Index([102, 101]), + ) + + monkeypatch.setattr(local_runner, "compute_household_metrics", fake_compute) + + tables = build_metric_tables_from_dataset( + dataset=type("Dataset", (), {"time_period": 2023})(), + area_groups={"E001": "England"}, + area_type="constituency", + household_ids=[101, 102], + simulation_factory=ReversedSimulation, + ) + + assert tables["England"].index.tolist() == [101, 102] + assert tables["England"]["population"].tolist() == [10.0, 20.0] + + +def test_build_local_candidate_solves_and_exports_long_weights() -> None: + areas = pd.DataFrame( + {"code": ["S001", "E001"], "country": ["Scotland", "England"]} + ) + targets = pd.DataFrame( + { + "code": ["E001", "S001"], + "population": [2.0, 2.0], + } + ) + metrics = { + "England": pd.DataFrame({"population": [1.0, 0.0]}, index=[101, 102]), + "Scotland": pd.DataFrame({"population": [0.0, 1.0]}, index=[101, 102]), + } + households = pd.DataFrame( + { + "household_id": [102, 101], + "household_weight": [2.0, 1.0], + "source_year": [2022, 2023], + "source_household_id": ["b", "a"], + "clone_index": [1, 0], + } + ) + + result = build_local_candidate( + area_type="constituency", + area_frame=areas, + targets=targets, + metrics=metrics, + household_frame=households, + solver_options={"epochs": 60, "learning_rate": 0.2, "seed": 1}, + ) + + assert result.problem.area_codes == ("E001", "S001") + assert result.solve_result.final_loss < result.solve_result.initial_loss + assert set(result.long_weights["area_code"]) == {"E001", "S001"} + assert set(result.long_weights["source_household_key"]) == { + "2023:a", + "2022:b", + } + assert result.support_summary["area_code"].tolist() == ["E001", "S001"] + assert "nonzero_source_households" in result.support_summary.columns + assert "effective_sample_size" in result.support_summary.columns + + +def test_build_local_candidate_can_limit_pilot_areas() -> None: + areas = pd.DataFrame( + { + "code": ["S001", "E001"], + "country": ["Scotland", "England"], + } + ) + targets = pd.DataFrame( + { + "code": ["E001", "S001"], + "population": [1.0, 1.0], + } + ) + metrics = { + "England": pd.DataFrame({"population": [1.0]}, index=[101]), + "Scotland": pd.DataFrame({"population": [1.0]}, index=[101]), + } + households = pd.DataFrame({"household_id": [101], "household_weight": [1.0]}) + + result = build_local_candidate( + area_type="constituency", + area_frame=areas, + targets=targets, + metrics=metrics, + household_frame=households, + max_areas=1, + solver_options={"epochs": 2}, + ) + + assert result.problem.area_codes == ("E001",) + assert result.long_weights["area_code"].unique().tolist() == ["E001"] + + +def test_build_local_candidate_from_dataset_computes_metrics(monkeypatch) -> None: + areas = pd.DataFrame({"code": ["E001"], "country": ["England"]}) + targets = pd.DataFrame({"code": ["E001"], "population": [1.0]}) + households = pd.DataFrame({"household_id": [101], "household_weight": [1.0]}) + + def fake_compute(sim, area_type, *, period=None, household_ids=None): + assert area_type == "constituency" + assert period == 2023 + assert sim.inputs[("region", 2023)] == ["SOUTH_EAST"] + assert household_ids is None + return pd.DataFrame({"population": [1.0]}, index=pd.Index([101])) + + monkeypatch.setattr(local_runner, "compute_household_metrics", fake_compute) + + result = build_local_candidate_from_dataset( + dataset=type("Dataset", (), {"time_period": 2023})(), + area_type="constituency", + area_frame=areas, + targets=targets, + household_frame=households, + simulation_factory=SingleHouseholdSimulation, + solver_options={"epochs": 2}, + ) + + assert result.problem.area_codes == ("E001",) + assert result.support_summary["nonzero_households"].tolist() == [1] + + +def test_write_local_candidate_outputs(tmp_path: Path) -> None: + areas = pd.DataFrame({"code": ["E001"], "country": ["England"]}) + targets = pd.DataFrame({"code": ["E001"], "population": [1.0]}) + metrics = pd.DataFrame({"population": [1.0]}, index=[101]) + households = pd.DataFrame({"household_id": [101], "household_weight": [1.0]}) + result = build_local_candidate( + area_type="la", + area_frame=areas, + targets=targets, + metrics=metrics, + household_frame=households, + solver_options={"epochs": 2}, + ) + + summary = write_local_candidate_outputs(result, tmp_path) + + assert (tmp_path / "local_geography_weights.csv.gz").exists() + assert (tmp_path / "solve_diagnostics.csv").exists() + assert (tmp_path / "area_support_summary.csv").exists() + saved = json.loads((tmp_path / "solve_summary.json").read_text()) + assert saved == summary + assert summary["area_type"] == "la" + assert summary["n_areas"] == 1 diff --git a/uv.lock b/uv.lock index 6243d5a2..347464b3 100644 --- a/uv.lock +++ b/uv.lock @@ -1172,6 +1172,10 @@ dependencies = [ ] [package.optional-dependencies] +uk = [ + { name = "h5py" }, + { name = "policyengine-uk" }, +] us = [ { name = "h5py" }, { name = "policyengine-us" }, @@ -1184,16 +1188,18 @@ dev = [ [package.metadata] requires-dist = [ + { name = "h5py", marker = "extra == 'uk'", specifier = ">=3" }, { name = "h5py", marker = "extra == 'us'", specifier = ">=3" }, { name = "numpy", specifier = ">=1.26" }, { name = "pandas", specifier = ">=2" }, + { name = "policyengine-uk", marker = "extra == 'uk'", specifier = ">=2.88" }, { name = "policyengine-us", marker = "extra == 'us'", specifier = ">=1.729,<2" }, { name = "populace-calibrate", editable = "packages/populace-calibrate" }, { name = "populace-fit", editable = "packages/populace-fit" }, { name = "populace-frame", editable = "packages/populace-frame" }, { name = "scipy", specifier = ">=1.13" }, ] -provides-extras = ["us"] +provides-extras = ["us", "uk"] [package.metadata.requires-dev] dev = [{ name = "pytest", specifier = ">=8" }]