Add OA calibration pipeline: Phases 1-6 (crosswalk, clone-and-assign, L0 engine, sparse matrix, target DB, local H5 publishing)

[vahid-ahmadi]

Background

This PR implements Phases 1-6 of the OA calibration pipeline — porting the US-side clone-and-prune methodology to the UK at Output Area level (~235K OAs).

What this PR does

Phase 1: OA Crosswalk & Geographic Assignment

• Unified UK Output Area crosswalk from ONS/NRS/NISRA: OA → LSOA → MSOA → LA → constituency → region → country
• Population-weighted random OA assignment with country constraints and constituency collision avoidance
• Pre-built crosswalk: storage/oa_crosswalk.csv.gz (235K areas, 1.4MB)

Phase 2: Clone-and-Assign

• Clones each FRS household N times (10 production, 2 testing) with unique IDs across all entity tables
• Assigns each clone a different Output Area (population-weighted, country-constrained)
• Divides household_weight by N so aggregate population totals are preserved
• Wired into create_datasets.py after imputations, before uprating/calibration
• Pure pandas/numpy operations — no simulation overhead

Phase 3: L0 Calibration Engine

• Wraps l0-python's SparseCalibrationWeights (HardConcrete gates) with the existing target matrix interface
• Builds sparse (n_targets, n_records) calibration matrix with country masking baked into sparsity pattern
• Relative squared error loss with target group weighting for balanced metric contribution
• Existing calibrate.py preserved as fallback

Phase 4: Sparse Matrix Builder

• build_assignment_matrix(): sparse (n_areas, n_households) binary matrix from OA geography columns
• create_cloned_target_matrix(): backward-compatible (metrics, targets, country_mask) interface
• build_sparse_calibration_matrix(): direct sparse path producing (M_csr, y, group_ids)
• Consolidates metric computation and target loading duplicated between constituency and LA loss files

Phase 5: SQLite Target Database

• Hierarchical target storage with two parallel geographic branches:
  • Administrative: country → region → LA → MSOA → LSOA → OA
  • Parliamentary: country → constituency
• LA and constituency are parallel — a constituency can span multiple LAs and vice versa
• Schema: areas (hierarchy via parent_code), targets (definitions), target_values (year-indexed)
• ETL loads areas from OA crosswalk + area code CSVs, targets from registry + local CSV/XLSX sources
• Query API: get_targets(), get_area_targets(), get_area_children(), get_area_hierarchy()

Phase 6: Local Area H5 Publishing (new)

• publish_local_h5s(): extracts per-area H5 subsets from sparse L0-calibrated weight vector
• Each H5 contains only active (non-zero weight) households with linked person and benunit rows
• Supports both constituency (650) and LA (360) area types
• validate_local_h5s(): post-publish validation checking file existence, HDF5 structure, cross-area HH ID uniqueness
• Wired into create_datasets.py after calibration, before downrating
• Summary CSV with per-area statistics (n_households, n_active, total_weight)

Performance

Phase 2 clone step is pure pandas/numpy — seconds for ~20K households × 10 clones. Phase 4 sparse matrix builder avoids materialising dense (n_areas, n_cloned_households) matrices. Phase 6 publishing iterates sequentially over areas — for 650 constituencies this takes seconds; future Modal integration will parallelise for ~180K OAs.

Tests

• 25 crosswalk/assignment tests (Phase 1)
• 14 clone-and-assign tests (Phase 2)
• 6 L0 calibration tests (Phase 3)
• 10 sparse matrix builder tests (Phase 4)
• 12 target database tests (Phase 5)
• 13 local H5 publishing tests (Phase 6)

File summary

File	Purpose
calibration/oa_crosswalk.py	Downloads & builds unified UK OA crosswalk
calibration/oa_assignment.py	Population-weighted OA assignment with constraints
calibration/clone_and_assign.py	Clones FRS entities, remaps IDs, assigns geography
calibration/matrix_builder.py	Sparse assignment matrix, consolidated metrics & targets
calibration/publish_local_h5s.py	Per-area H5 extraction from sparse weights
utils/calibrate_l0.py	L0-regularised calibration with sparse matrices
db/schema.py	SQLite schema (areas, targets, target_values)
db/etl.py	ETL loading areas + targets from all sources
db/query.py	Query API for target retrieval
datasets/create_datasets.py	Pipeline integration (clone, calibrate, publish)
storage/oa_crosswalk.csv.gz	Pre-built crosswalk (235K areas)
docs/oa_calibration_pipeline.md	6-phase roadmap (all complete)

🤖 Generated with Claude Code

[baogorek]

Hi Vahid,

Most of this is from our boy Claude, as usual. This looks like a great setup! Can't wait to see HHs getting donated to the OAs! I'll approve, but please see the issues Claude found below.

Here's the code I used to poke around:

  from policyengine_uk_data.calibration.oa_crosswalk import load_oa_crosswalk
  xw = load_oa_crosswalk()                                                                                                                                    
  xw                                                                                                                                                          
                                                                                                                                                              
  # Population-weighted sampling demo                                 
  import numpy as np

  xw["population"] = xw["population"].astype(float)

  eng = xw[xw["country"] == "England"].copy()
  eng["prob"] = eng["population"] / eng["population"].sum()

  rng = np.random.default_rng(42)
  idx = rng.choice(len(eng), size=10_000, p=eng["prob"].values)
  sampled = eng.iloc[idx]

  sampled.groupby("oa_code")["population"].agg(["count", "first"]).rename(
      columns={"count": "times_sampled", "first": "population"}
  ).sort_values("times_sampled", ascending=False).head(20)

leads to:

Out[1]: 
           times_sampled  population
oa_code                             
E00179944              5      3354.0
E00035641              3       279.0
E00039569              3       263.0
E00066618              3       331.0
E00115325              2       319.0
E00136307              2       301.0
E00089585              2       333.0
E00167257              2       472.0
E00130843              2       406.0
E00021422              2       190.0
E00004742              2       313.0
E00044937              2       294.0
E00089725              2       240.0
E00044974              2       400.0
E00160095              2       401.0
E00016512              2       305.0
E00016490              2       380.0
E00089915              2       514.0
E00021502              2       396.0
E00105618              2       305.0

Interesting: "E00179944 with population 3,354 is a massive outlier (most OAs are 100–300 people)"

Bugs

1. load_oa_crosswalk loads population as string

load_oa_crosswalk() passes dtype=str for all columns (line 753 of oa_crosswalk.py), so population comes back as a string. This means any downstream arithmetic (e.g. computing probabilities) fails with TypeError: unsupported operand type(s) for /: 'str' and 'str'. Should either drop dtype=str or explicitly cast population to int on load.

2. NI households silently get no assignment

The crosswalk has 0 NI rows (NISRA 404), which is acknowledged, but assign_random_geography will silently produce None entries for NI households (country code 4). Worth either raising an error or logging a warning when a household's country has no distribution.

Code quality

3. Dead code in _assign_regions

Lines 602–606 of oa_crosswalk.py:

for k, v in la_to_region.items():
    if k[:3] == la_code[:3]:
        # Same LA type prefix
        pass

This loop does nothing — should be removed or finished.

4. Assignment inner loop should be vectorised

In oa_assignment.py lines 236–245, the for i, pos in enumerate(positions) loop storing results can be replaced with vectorised numpy indexing:

oa_codes[start + positions] = dist["oa_codes"][indices]

Same for all the other arrays. Will matter when n_clones * n_records gets large.

Worth noting

5. Scotland population weighting is effectively uniform

The fallback of ~117 per OA for all 46k Scottish OAs means population-weighted sampling is actually uniform for Scotland. This undermines the premise for ~20% of UK OAs. Might be worth a louder warning or a TODO to revisit once NRS fixes the 403.

[baogorek]

Approving Phase 1 — the crosswalk and assignment engine look good. Please see my comment above for a few things to address before merge.

[nwoodruff-co]

Putting a req changes here- due to importance of data here I'm going to say don't approve unless the PR is ready to merge at time of approval.

Aiming to block the least but these are the minimum:

• The constituency impacts (all 650) currently take less than 5 seconds to run after a completed national simulation. This probably increases that by several orders of magnitude to 10 minutes plus. Can you both confirm/reject, and argue in favour of your argument here? I agree yours is a theoretically better solution but we do need to consider this.

• This would be a major data change- need to run microsimulation regression tests to understand if outputs significantly change. At bare minimum this should include these examples:

a) the living standards outlook (rel change in real hbai household net income BHC from 2024 to 2029, broken down by age group

b) raising the higher rate to 41p (broken down by equiv hbai household net income bhc decile)

If you can say these don't change by 0.1p/0.1bn respectively, we can skip digging further

[MaxGhenis]

Ran the requested microsimulation regression checks locally on March 18, 2026.

Method:

• Held the model constant at policyengine_uk 2.74.0 / policyengine-core 3.23.6.
• Used one interpreter: /Users/maxghenis/worktrees/policyengine-uk-data-pr291/.venv/bin/python.
• Used the same dataset in both runs: enhanced_frs_2023_24.h5.
• Swapped only policyengine_uk_data between main and this PR worktree.

This is important because the latest PyPI policyengine-uk is newer (2.75.1 as of March 18, 2026), but upgrading the model while testing this data PR would confound the comparison.

Result: for the two examples below, main and this PR produced identical outputs at the precision shown.

1. Living standards outlook
   Relative change in real_hbai_household_net_income from 2024 to 2029, by age group:

• All: +0.156439%
• Children: +0.449643%
• Working-age adults: -0.119088%
• Seniors: +1.028031%

2. Raise higher rate to 41p in 2029
   Fiscal impact:

• +£2.697369bn

Relative change in household net income by household_income_decile:

• 1: -0.005543%
• 2: -0.001419%
• 3: -0.003603%
• 4: -0.007614%
• 5: -0.021598%
• 6: -0.035981%
• 7: -0.051815%
• 8: -0.091866%
• 9: -0.217914%
• 10: -0.514674%

So for these examples, the PR changes are 0 relative to main, which is well within the requested 0.1pp / 0.1bn thresholds.

This also matches the scope of the diff: Phase 1 adds OA crosswalk / assignment code and oa_crosswalk.csv.gz, but does not yet wire that path into calibration or modify the enhanced FRS dataset used by these runs.

[vahid-ahmadi]

@nwoodruff-co Re your performance concern about constituency impacts going from <5s to 10+ minutes:

Phase 1 has zero performance impact. This PR adds only new standalone files — zero existing files are modified. The new calibration/ package is not imported or called by anything in the existing pipeline:

• create_datasets.py — unchanged
• utils/calibrate.py — unchanged
• local_areas/constituencies/ — unchanged

The current <5s constituency impact calculation (weights @ metrics matrix multiply using pre-computed weights from parliamentary_constituency_weights.h5) is completely untouched. Max's regression tests confirmed this — identical outputs, because no existing code paths are affected.

The performance question is valid but applies to future phases (Phase 2: clone-and-assign, Phase 3: L0 calibration), where the weight matrix would grow from 650 × 100K to potentially 650 × 1M+. That's worth addressing when those PRs come, not here.

Between this and Max's regression results (zero change on both requested examples), both concerns from your changes-requested review should be resolved for Phase 1.

[nikhilwoodruff]

right so this pr doesn't actually change the production data? sure, but then why not just keep iterating within the pr. we don't need to merge yet if it's not actually changing package behaviour

[vahid-ahmadi]

Phase 2: Clone-and-Assign added

Following Nikhil's suggestion to keep iterating in this PR rather than merging Phase 1 alone, I've added Phase 2.

What's new

• calibration/clone_and_assign.py — clones each FRS household N times (10 production, 2 testing), remaps all entity IDs (household, person, benunit), divides weights by N, and attaches OA geography columns
• create_datasets.py — clone step wired in after imputations, before uprating/calibration
• tests/test_clone_and_assign.py — 14 tests (all passing) covering dimensions, weight preservation, ID uniqueness, FK integrity, country constraints

Re: runtime concern

The clone step is pure pandas/numpy (DataFrame copies + ID arithmetic + OA sampling). No microsimulation is run. For ~20K households × 10 clones this should take seconds. The existing constituency impact path (pre-computed weights matrix multiply from parliamentary_constituency_weights.h5) is completely untouched.

[vahid-ahmadi]

Phase 3: L0 Calibration Engine added

What's new

• utils/calibrate_l0.py — wraps l0-python's SparseCalibrationWeights (HardConcrete gates for continuous L0 relaxation). Uses the same matrix_fn/national_matrix_fn interface as the existing calibrate.py, so it's a drop-in alternative. Builds a sparse (n_targets × n_records) matrix with country masking baked into the sparsity pattern — avoids the dense (areas × households) memory overhead.
• pyproject.toml — added l0-python>=0.4.0 to dev dependencies (already PolicyEngine's own package)
• tests/test_calibrate_l0.py — 6 tests covering sparse matrix construction, country masking, zero-target filtering, group IDs, error reduction, and sparsity behaviour
• Existing calibrate.py preserved as fallback

Key design decisions

• Sparse matrix with baked-in country masking: Instead of a dense country mask r[i,j], the sparsity pattern only includes entries where a household belongs to an area's country. This is critical for the 10x-cloned dataset.
• Target group weighting: Each metric type (e.g. income, age, UC) contributes equally to the loss regardless of how many areas it spans. Prevents age targets (650 areas × 18 bands) from dominating income targets.
• Same interface as existing calibration: calibrate_l0() takes the same matrix_fn, national_matrix_fn, area_count, weight_file arguments — can be swapped in when ready.

All 45 tests passing (Phase 1: 25, Phase 2: 14, Phase 3: 6).

[vahid-ahmadi]

Phase 4: Sparse Matrix Builder added

Bridges Phase 2 (clone-and-assign) and Phase 3 (L0 calibration) — the missing piece that wires existing target sources into the sparse format the L0 engine consumes.

What's new

calibration/matrix_builder.py with three public functions:

• build_assignment_matrix() — builds a sparse (n_areas, n_households) binary matrix from the constituency_code_oa / la_code_oa columns that clone-and-assign attaches. Each household is in exactly one area. This replaces the dense country_mask that the existing loss.py files produce.

• create_cloned_target_matrix() — backward-compatible (metrics, targets, country_mask) interface, usable as matrix_fn for both calibrate_local_areas() (dense Adam) and calibrate_l0(). Densifies the sparse assignment for backward compat.

• build_sparse_calibration_matrix() — direct sparse path producing (M_csr, y, group_ids) without ever materialising a dense (n_areas, n_cloned_households) matrix. For 650 constituencies × 200K cloned households this avoids a 130M-entry dense array.

Also consolidates the metric computation and target loading that was copy-pasted between constituencies/loss.py and local_authorities/loss.py into shared helpers (_compute_household_metrics, _load_area_targets). Supports both constituency (650) and LA (360) area types.

Tests

10 tests covering assignment matrix shape, sparsity, binary values, unassigned households, area type switching, and unknown code handling. All passing.

[vahid-ahmadi]

Phase 5: SQLite Target Database added

What's new

policyengine_uk_data/db/ — new package with three modules:

• schema.py — SQLite schema with three tables:

  • areas: geographic hierarchy via parent_code — two parallel branches:
    • Administrative: country → region → LA → MSOA → LSOA → OA
    • Parliamentary: country → constituency
  • targets: one row per calibration target definition (name, variable, source, unit, geographic level, geo code, etc.)
  • target_values: year-indexed values for each target

• etl.py — ETL script that loads:

  • Areas from OA crosswalk (~235K OAs) + constituency/LA code CSVs
  • Registry targets (national/country/region from all 18 source modules)
  • Local targets from CSVs: constituency + LA age bands, HMRC SPI income, DWP UC households, LA extras (ONS income, tenure, private rent)
  • Run via python -m policyengine_uk_data.db.etl

• query.py — query API:

  • get_targets(geographic_level=, geo_code=, variable=, source=, year=) — flexible filtered queries
  • get_area_targets(geo_code, year) — all targets for a specific area
  • get_area_children(parent_code) — child areas in the hierarchy
  • get_area_hierarchy(code) — walk up from OA to country

Key design decision

LA and constituency are parallel branches, not parent-child. A constituency can span multiple LAs and vice versa. The parent_code chain follows the administrative branch (country → region → LA → MSOA → LSOA → OA), while constituencies parent directly to country.

Tests

12 tests covering schema creation, area hierarchy walks, LA→region→country chain, constituency→country chain, target queries by level/year/source/area. All passing.

[vahid-ahmadi]

Phase 6: Local Area H5 Publishing added

Completes the 6-phase pipeline. After L0 calibration produces a sparse weight vector, Phase 6 extracts per-area H5 files — each containing only the active (non-zero weight) households for that area.

What's new

calibration/publish_local_h5s.py with four public functions:

• _get_area_household_indices() — maps each area code to household row indices via the OA geography columns from clone-and-assign (constituency_code_oa / la_code_oa). O(n_households) scan.

• publish_area_h5() — writes a single per-area H5. Filters to active households (weight > 0 after L0 pruning), extracts linked persons via person_household_id FK and benunits via benunit_id // 100 FK. Stores as HDF5 groups (household/person/benunit) with metadata attributes (area_code, n_households, total_weight).

• publish_local_h5s() — orchestrates the full publish cycle: loads the sparse weight vector from the L0 output H5, iterates over all areas, writes one H5 per area to storage/local_h5s/{area_type}/, produces _summary.csv with per-area statistics.

• validate_local_h5s() — post-publish validation: checks all expected area files exist, verifies HDF5 structure (household/person/benunit groups), checks for cross-area household ID uniqueness (a household should only appear in one area after L0 pruning).

Pipeline integration: wired into create_datasets.py after constituency and LA calibration, before downrating. Publishes both constituency (650 files) and LA (360 files).

Tests

13 tests covering:

• Area-household index mapping (constituency, LA, unknown codes, full coverage)
• H5 file structure and metadata
• Zero-weight household exclusion
• Weight correctness
• Person/benunit FK integrity
• Full publish cycle with mock weight file
• Summary statistics
• Validation of published files
• Detection of missing files

All 80 tests passing (Phase 1: 25, Phase 2: 14, Phase 3: 6, Phase 4: 10, Phase 5: 12, Phase 6: 13).

Design note on Modal

The current implementation is sequential — fine for 650 constituencies or 360 LAs (seconds). For ~180K OA files, Modal parallelisation would be the next step: each OA publish is independent and embarrassingly parallel. The publish_area_h5() function is designed to be callable as a Modal remote function with no shared state.