diff --git a/docs/informed-l0-warm-start-design.md b/docs/informed-l0-warm-start-design.md new file mode 100644 index 00000000..ec0adb72 --- /dev/null +++ b/docs/informed-l0-warm-start-design.md @@ -0,0 +1,232 @@ +# Informed-L0 / warm-start selection — design (populace#328) + +Reconstructs the informed-L0 / warm-start selection behind the certified US +default as **committed, spec-declared machinery**, closing the #159 discipline gap +(no uncommitted certified steps). Grounded in the Phase 1 forensics +(`informed-l0-warm-start-forensics.md`), which established that the frozen +57,240-record support is recoverable from the published H5's source identities and +joins 100% cleanly onto a freshly-rebuilt base. + +## Problem restated + +The live default is a 1,500-epoch dense polish (`method = "dense_no_l0"`) of a +**frozen 57,240-record support** warm-started from an ad-hoc run-dir `.npz`. The +selection that produced that support was never committed. Consequently: + +- The **sparse** default cannot be reproduced from `main`: cold L0 on the rebuilt + 337,704 pool gives federal income tax **+19.9%**, loss 0.119, within-10% 50.2% + (vs certified 0.044 / 94.7%). +- The **dense** parent overshoots the JCT mortgage tax-expenditure **+59.7%** — + the same cold-vs-informed-selection gap. + +Both are fixed by making the informed selection a first-class, reproducible build +input. + +## Two modes + +**(a) Frozen-support mode** — select *exactly* the named record set, then calibrate +weights on that fixed support. This is the certified pattern: reduce the candidate +frame to the frozen 57,240 support, then run the dense polish (or an L0-refit) on +it. Requires no change to `populace-calibrate`. + +**(b) Informed-init mode** — initialize the L0 selection *probabilities* from the +named artifact's selection, then let L0 optimize. More robust to base drift (a +selected source record that no longer exists just starts cold instead of failing), +but requires a new per-record selection-probability-init parameter in +`populace-calibrate.calibrate_l0_refit` (which today exposes only a scalar +`init_mean`). + +### Which first, and why + +**Frozen-support mode ships first.** Phase 1 proved the join onto the current +reconstructed base is exact (57,240/57,240, 0 unmapped, 0 ambiguous), so there is +no base drift to tolerate today — mode (b)'s robustness buys nothing yet, and its +cost (a cross-package solver change) is real. Mode (a) reproduces the certified +artifact's construction faithfully and is fully testable now. Mode (b) is specified +here so the seam, spec schema, and manifest fields are forward-compatible, and is +implemented when a future base rebuild first produces a non-empty unmapped set. + +## The seam + +### Spec/CLI input (#159 discipline) + +A release declares a **selection source** naming an artifact plus the **join +contract**. Two surfaces, same underlying object: + +1. **CLI** on `tools/build_us_fiscal_refresh_release.py`: + - `--selection-source-h5 PATH` — a populace US H5 whose record set defines the + frozen support (e.g. the certified live default). + - `--selection-join-key KEY` — the identity key; default + `source_year,source_household_id,household_support_channel,household_support_clone_index` + (the Phase-1-verified unique key). Named, not hard-coded, so the contract is + auditable and overridable. + - `--selection-mode {frozen_support,informed_init}` — default `frozen_support`. + - `--selection-source-manifest PATH` (alternative to `--selection-source-h5`): + a committed JSON listing the selected identities directly (see below), for a + laptop reproduction that does not depend on downloading the 354 MB H5. + +2. **Release spec field** (forward-looking, per #159): the same three values as a + `selection_source` block the country release spec can carry, so a spec-compiled + build passes them without ad-hoc flags. The CLI flags remain the authoritative + runtime surface; the spec field is the declarative mirror. + +The join key vocabulary is restricted to a **whitelist of stable identity +columns** (`source_year`, `source_household_id`, `source_person_id`, +`household_support_channel`, `household_support_clone_index`, `household_source_id`) +— never an assigned row id (`household_id`), which is order-dependent and would +silently misjoin across a rebuild. + +### Selection-source manifest (committed, laptop-reproducible) + +To make the sparse default reproducible from `main` on a laptop **without** +re-downloading the source H5 every time, the selection can be frozen once into a +small committed JSON: + +```json +{ + "schema_version": 1, + "join_key": ["source_year", "source_household_id", + "household_support_channel", "household_support_clone_index"], + "source": {"repo_id": "policyengine/populace-us", + "revision": "populace-us-2024-sparse-l0-refit-57k-...-20260701", + "path": "populace_us_2024.h5", + "sha256": "c2065b64..."}, + "n_selected": 57240, + "identities_sha256": "", + "identities": [[2024, 13, "asec", 0], ...] +} +``` + +`identities_sha256` binds the exact set; the `source` block records provenance back +to the HF artifact it was distilled from. A tool +(`tools/build_us_selection_source_manifest.py`) produces this from any source H5. +This manifest is the committed, spec-declared embodiment of the previously +uncommitted `a7a77c4` selection. + +### Runtime support code (`us_runtime/`, never `build/us/`) + +New module `packages/populace-build/src/populace/build/us_runtime/warm_start_selection.py`: + +- `SelectionSource` — a frozen dataclass: the join key (validated against the + whitelist), the identity rows, and provenance (source sha / repo / revision, or + the source H5 path). +- `load_selection_source_from_h5(path, join_key)` — read the identity rows from a + source H5's person + household tables, building the per-household key exactly as + Phase 1 did (household-level components are constant within a household; the + channel/clone come from the household table). Returns a `SelectionSource`. +- `load_selection_source_from_manifest(path)` — read + validate the committed JSON, + including the `identities_sha256` integrity check. +- `select_frozen_support(base_frame, source) -> (Frame, SelectionReport)` — the + core join. Builds the base's per-household key, intersects with the source's + identity set, expands to a person mask, and calls `Frame.select`. Returns the + reduced frame and a report (`n_source`, `n_base_candidates`, `n_selected`, + `n_unmapped`, `n_ambiguous`, the join key, source provenance). + +**Refusal contract (no fabricated selections):** `select_frozen_support` raises if +- the join key is not fully unique in either frame (would make selection + order-dependent), or +- any source identity fails to map onto the base (`n_unmapped > 0`) in + `frozen_support` mode — a silently-unmappable artifact must **fail loudly**, not + select a smaller-than-declared or wrong support, or +- a source identity maps to more than one base row (`n_ambiguous > 0`). + +Informed-init mode relaxes only the unmapped rule (unmapped source records are +dropped from the init set, recorded in the report, and the L0 selector starts those +records cold) — but still refuses on a non-unique key or ambiguous match. + +### Where it slots into the build + +Immediately after the base frame is loaded and **before** target materialization +(`_materialize_target_frame`), when a selection source is supplied and the mode is +`frozen_support`: + +``` +base_frame = _load_frame(base_h5) +if selection_source is not None and mode == "frozen_support": + base_frame, selection_report = select_frozen_support(base_frame, selection_source) +# ... existing immigration/take-up input attach, then materialize on the reduced frame +target_frame, registry, compilation = _materialize_target_frame(base_frame, ...) +``` + +Reducing *before* materialization matches the certified run (`n_candidate = 57,240`, +`compiled_candidate_targets = 32,637`) and materializes PolicyEngine over 57k rows +instead of 337k. The existing calibration call then runs unchanged on the reduced +frame — dense (`--dense-default-dataset`, the certified polish) or L0-refit. + +For informed-init mode the base frame is *not* reduced; instead the per-record init +probabilities are passed into `calibrate_l0_refit` (new parameter) so L0 optimizes +from the informed start over the full pool. + +### Manifest provenance + +A `selection_source` block is added to both `build_manifest.json` and +`release_manifest.json` (alongside the existing `warm_start_calibration` and +`default_dataset` blocks): + +```json +"selection_source": { + "mode": "frozen_support", + "join_key": ["source_year", "source_household_id", + "household_support_channel", "household_support_clone_index"], + "source": {"kind": "h5"|"manifest", "path": "...", "sha256": "...", + "repo_id": "...", "revision": "..."}, + "n_source": 57240, "n_base_candidates": 337704, + "n_selected": 57240, "n_unmapped": 0, "n_ambiguous": 0 +} +``` + +This is the recorded provenance #328 requires — the selection is reproducible from +`main` given the (committed manifest or sha-pinned HF) source, and the build +records exactly which source it descended from and how the join resolved. + +### Relationship to the existing warm-start weights seam + +`--warm-start-calibration-npz` (positional weight start) is **complementary and +unchanged**. After frozen-support reduction, the reduced frame's household order is +deterministic (base `household_id` ascending), so an aligned warm-start weight +vector for the *reduced* frame can still be supplied the existing way. The +certified pattern is: (1) frozen-support reduction → 57,240 frame, (2) optional +positional warm-start of those 57,240 weights, (3) dense polish. This design adds +step (1) as committed machinery; steps (2)–(3) already exist. The `solver_config` +identity omission (#326) is orthogonal and not blocking here. + +## Testing (synthetic fixtures + refusal) + +Unit tests in `packages/populace-build/tests/test_us_warm_start_selection.py`, all +on tiny synthetic frames (no network, no large H5): + +1. **Join maps a clone-aware support** — a base with 2-year × 2-clone records; a + source naming a subset; assert the reduced frame is exactly the named + households, in the right order, with the source's identities. +2. **Order-independence** — shuffle the base row order; assert the same support is + recovered (proves identity-join, not positional). +3. **Refuse silently-unmappable (frozen)** — a source identity absent from the base + raises with a message naming the unmapped identity; assert **no** frame is + returned (no fabricated / truncated selection). +4. **Refuse non-unique key** — a key missing the channel/clone component (so a + source household's two clones collide) raises before selecting. +5. **Refuse ambiguous match** — a source identity mapping to >1 base row raises. +6. **Manifest round-trip + integrity** — write a selection-source manifest, reload, + assert identities match; corrupt `identities_sha256` and assert it raises. +7. **Informed-init drops-and-records** — in `informed_init` mode an unmapped source + record is dropped from the init set and appears in the report, without raising. + +A build-level test asserts the manifest carries the `selection_source` block with +the resolved counts when a selection source is passed. + +## Validation (Phase 4) + +Because Phase 1 recovered a usable selection, run the sparse path against the Build +F base (`18833fb6…`) with the recovered frozen support and report income tax / SS / +loss / within-10% versus the certified reference (0.044 / 94.7%) and cold-L0 +(0.119 / 50.2%). Detached compute (nohup + pidfile). If the run reproduces certified +quality, the sparse default is once again buildable from `main`. + +## Out of scope / deferred + +- Informed-init mode's calibrate-library parameter (implemented when a rebuild + first yields `n_unmapped > 0`). +- `solver_config_sha256` (#326) — orthogonal identity hardening. +- The export mass-parity gate reference fix (#327) — separate one-line reference + change; the mortgage residual it exposes is resolved by *this* work (the frozen + support fits the mortgage tax-expenditure tighter than cold dense). diff --git a/docs/informed-l0-warm-start-forensics.md b/docs/informed-l0-warm-start-forensics.md new file mode 100644 index 00000000..c98ca0f2 --- /dev/null +++ b/docs/informed-l0-warm-start-forensics.md @@ -0,0 +1,112 @@ +# Informed-L0 / warm-start reconstruction — Phase 1 forensics (populace#328) + +This records the Phase 1 forensic findings for populace#328 — whether the frozen +57,240-record support behind the certified live default can be recovered, and how +its record identities map onto a freshly-rebuilt base. It is the evidence base for +the design in `informed-l0-warm-start-design.md`. + +All Hugging Face reads here were **read-only** against `policyengine/populace-us`. +No prod write path, tag, or `latest.json` was touched. + +## Question + +The certified live default +`populace-us-2024-sparse-l0-refit-57k-71a0887-national-only-20260701` is, per its +manifest, a 1,500-epoch dense polish (`method = "dense_no_l0"`, final loss 0.0440) +of a **frozen 57,240-record support** warm-started from +`out/sparse-default-release-20260701/refit-selected-current-surface-a7a77c4-20260701/artifacts/a7a77c4_reconstructed_warm_start.npz`. +That selection artifact was never committed and descends from an ad-hoc run dir. +Can the frozen support be recovered, and does it map onto Build F's rebuilt base +(`18833fb6…`, rebuilt from raw ASEC, so row identities differ)? + +## Findings + +### F1 — The named warm-start artifact is NOT on Hugging Face + +`policyengine/populace-us` has one branch (`main`) and 15 tags. The certified tag +exists and carries `populace_us_2024.h5` (354 MB, 57,240 households), +`populace_us_2024_calibration.npz` (2.2 MB, the shipped weights), and the release +manifests. There is **no** `refit-selected-current-surface-a7a77c4` tag or file. +The manifest's `warm_start_calibration.path` is a local `out/…` run-dir path; the +`.npz` (sha `7ba46106…`) is not published. So the frozen support cannot be +retrieved as a named artifact. + +The build manifest records **no** `base_dataset_sha256` and no lineage back to the +337,704 pool or to the selection run — the same lineage gap #328 describes. + +### F2 — The frozen support IS recoverable from the published H5's record identities + +The published certified H5 carries, on its **person** table, the stable source +identity triple already used elsewhere in the build +(`us_runtime/take_up.py::_SOURCE_IDENTITY_COLUMNS`): + +- `source_year` ∈ {2024, 2023, 2022} (persons 57,185 / 54,665 / 54,452 — the + 3-year pool), `source_household_id`, `source_person_id` (stable ASEC person id + strings), `person_support_channel` ∈ {asec, puf_tax_detail}, + `person_support_clone_index` ∈ {0, 1}. + +and on its **household** table: `household_id` (assigned row id — order-dependent), +`household_source_id`, `household_support_channel`, `household_support_clone_index`. + +The Build F base carries the identical identity columns. So the *set of source +records* the certified default selected is fully readable from the published H5, +independent of the assigned `household_id`. + +### F3 — The identity join onto the rebuilt base is PERFECT and unambiguous + +Household-level identity key +`(source_year, source_household_id, household_support_channel, household_support_clone_index)`: + +- **Unique** in the Build F base (337,704 / 337,704 distinct) and in the certified + support (57,240 / 57,240 distinct). +- Joining the certified support onto the Build F base by this key: + **57,240 of 57,240 match (100.00%), 0 unmapped, 0 ambiguous** (no certified key + hits more than one base row). + +Two other keys were also verified fully unique in both frames and are viable +fallbacks: `(source_year, source_household_id, channel, min(source_person_id))` +and `(source_year, household_source_id, clone_index)`. The channel/clone component +is **required** — without it, a source ASEC household's two support-channel clones +collide. + +### F4 — Positional warm-start is unusable across a rebuild (why identity-join is mandatory) + +The selected base rows are **scattered**, not a prefix: in Build F base +`household_id` order they sit at positions 4, 7, 10, 15, 19, … up to 337,698. The +two existing warm-start seams are both strictly positional: + +- `tools/build_us_fiscal_refresh_release.py::_load_warm_start_calibration_npz` + requires `household_weight.shape == n(household)` **and** asserts the stored + `initial_household_weight` matches the current frame's initial weights + (`max_abs_initial_household_weight_delta` must be ~0 — it is exactly 0.0 in the + certified manifest, i.e. that run's warm start was positionally aligned to its + own frozen support). +- `us_runtime/l0_refit_export.py::load_l0_refit_npz` / `attach_l0_refit_weights` + require a full-length vector aligned to the base's own row order and apply the + selection mask positionally. + +Against a rebuilt base (57,240-length vector vs 337,704 rows; different order), +both would shape-mismatch or silently misalign. Recovery therefore **must** join on +the stable identity, which is exactly what this reconstruction adds. + +## Decision for Phase 2 + +Because the join is 100% clean and unambiguous, **frozen-support mode (select +exactly the named record set, then calibrate weights on that fixed support — the +certified pattern) is implementable now with no change to the calibrate library**, +and is the first mode to ship. Informed-init mode (initialize L0 selection +probabilities from the named artifact, then optimize) is designed as the drift- +robust successor but requires a new selection-probability-init parameter in +`populace-calibrate`; it is deferred behind mode (a) precisely because Phase 1 +shows the current base needs no drift tolerance. + +## Reproduction + +- HF reads: `huggingface_hub` against `policyengine/populace-us`, revision = + certified tag, `repo_type="dataset"` (read-only). Token from the shared + `HUGGING_FACE_TOKEN` agent secret. +- Join verification: `_buildg-runtime/forensics/join_check.py` (kept out of the + repo tree). Inputs: certified H5 (HF) + Build F base + `out/base-f-20260705/base_populace_us_2024_puf_support.h5` (sha `18833fb6…`). +- Recovered selection mask over the Build F base: + `_buildg-runtime/forensics/frozen_support_mask.npy` (337,704 bool, 57,240 True). diff --git a/packages/populace-build/src/populace/build/us_runtime/warm_start_selection.py b/packages/populace-build/src/populace/build/us_runtime/warm_start_selection.py new file mode 100644 index 00000000..a59a8f67 --- /dev/null +++ b/packages/populace-build/src/populace/build/us_runtime/warm_start_selection.py @@ -0,0 +1,552 @@ +"""Recover a frozen support onto a rebuilt base by stable source identity. + +The certified US default is a dense polish of a *frozen* 57,240-record support +whose selection descended from an uncommitted run-dir artifact (populace#328). +This module makes that selection first-class, committed machinery: it recovers +which source records a named selection artifact chose and reduces a freshly-built +candidate frame to exactly that support — keyed on the records' **stable source +identity**, never their assigned row id. + +Why identity, not row id: a base rebuilt from raw ASEC assigns fresh, order- +dependent ``household_id`` values, so a positional warm-start (the existing +``--warm-start-calibration-npz`` and ``load_l0_refit_npz`` seams) cannot recover +a support across a rebuild. The identity a record keeps across rebuilds is the +one already used for seeded draws in :mod:`populace.build.us_runtime.take_up`: +its origin ASEC ``source_year`` and ``source_household_id``, disambiguated across +the two PUF support-channel clones by ``household_support_channel`` / +``household_support_clone_index``. On the current reconstructed base this key is +unique in both the certified support (57,240) and the pool (337,704) and joins +100% cleanly (see ``docs/informed-l0-warm-start-forensics.md``). + +Two modes are supported (``docs/informed-l0-warm-start-design.md``): + +- ``frozen_support`` — select exactly the named record set, then calibrate on the + fixed support (the certified pattern). Any source identity that fails to map + onto the base is a hard error: a silently-unmappable artifact must fail loudly, + never select a truncated or wrong support. +- ``informed_init`` — a source identity that no longer exists is dropped from the + init set (recorded in the report) rather than raising, so a later L0 selection + can start those records cold. Reserved for a future rebuild that first drifts; + the runtime wiring for it lands with the calibrate-library parameter it needs. + +Both modes refuse on a non-unique join key (selection would be order-dependent) +and on an ambiguous match (one source identity hitting more than one base row). +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +import numpy as np +import pandas as pd + +from populace.frame import Frame, WeightKind, Weights +from populace.frame.units import US_SCHEMA + +__all__ = [ + "DEFAULT_SELECTION_JOIN_KEY", + "SELECTION_JOIN_KEY_COLUMNS", + "SELECTION_MODES", + "SelectionReport", + "SelectionSource", + "build_selection_source_manifest_from_h5", + "load_selection_source_from_h5", + "load_selection_source_from_manifest", + "main", + "select_frozen_support", + "write_selection_source_manifest", +] + +SelectionMode = Literal["frozen_support", "informed_init"] + +#: The selection modes this module supports. +SELECTION_MODES: tuple[SelectionMode, ...] = ("frozen_support", "informed_init") + +#: Whitelist of stable identity columns a join key may reference. Deliberately +#: excludes assigned row ids (``household_id``, ``person_id``): those are +#: order-dependent and would silently misjoin across a base rebuild. Each column +#: is tagged with the entity table it lives on. +SELECTION_JOIN_KEY_COLUMNS: dict[str, str] = { + "source_year": "person", + "source_household_id": "person", + "source_person_id": "person", + "household_support_channel": "household", + "household_support_clone_index": "household", + "household_source_id": "household", +} + +#: The Phase-1-verified unique, clone-aware household identity key. +DEFAULT_SELECTION_JOIN_KEY: tuple[str, ...] = ( + "source_year", + "source_household_id", + "household_support_channel", + "household_support_clone_index", +) + +_MANIFEST_SCHEMA_VERSION = 1 + + +def _validate_join_key(join_key: Sequence[str]) -> tuple[str, ...]: + key = tuple(join_key) + if not key: + raise ValueError("selection join key must name at least one column.") + unknown = [column for column in key if column not in SELECTION_JOIN_KEY_COLUMNS] + if unknown: + raise ValueError( + "selection join key may only reference stable identity columns " + f"{sorted(SELECTION_JOIN_KEY_COLUMNS)}; got unsupported {unknown}. " + "Assigned row ids are rejected because they are order-dependent." + ) + if len(set(key)) != len(key): + raise ValueError(f"selection join key has duplicate columns: {key}.") + return key + + +def _identities_digest(join_key: Sequence[str], identities: Sequence[Sequence]) -> str: + """Order-independent content hash binding the exact selected identity set.""" + canonical = sorted( + json.dumps(list(row), separators=(",", ":"), sort_keys=False) + for row in identities + ) + payload = json.dumps( + {"join_key": list(join_key), "identities": canonical}, + separators=(",", ":"), + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class SelectionReport: + """The outcome of resolving a selection source against a base frame.""" + + join_key: tuple[str, ...] + n_source: int + n_base_candidates: int + n_selected: int + n_unmapped: int + n_ambiguous: int + mode: SelectionMode + provenance: Mapping[str, object] = field(default_factory=dict) + + def as_manifest(self) -> dict[str, object]: + return { + "mode": self.mode, + "join_key": list(self.join_key), + "source": dict(self.provenance), + "n_source": int(self.n_source), + "n_base_candidates": int(self.n_base_candidates), + "n_selected": int(self.n_selected), + "n_unmapped": int(self.n_unmapped), + "n_ambiguous": int(self.n_ambiguous), + } + + +@dataclass(frozen=True) +class SelectionSource: + """A named record set to recover, plus its join contract and provenance. + + ``identities`` is a list of rows, each row a list of values aligned to + ``join_key`` (e.g. ``[2024, 13, "asec", 0]``). The rows are the source + artifact's selected records expressed in stable-identity space. + """ + + join_key: tuple[str, ...] + identities: list[list[object]] + provenance: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + key = _validate_join_key(self.join_key) + object.__setattr__(self, "join_key", key) + rows = [list(row) for row in self.identities] + for row in rows: + if len(row) != len(key): + raise ValueError( + f"each identity row must have {len(key)} values aligned to " + f"the join key {key}; got {row}." + ) + object.__setattr__(self, "identities", rows) + object.__setattr__(self, "provenance", dict(self.provenance)) + + @property + def n_identities(self) -> int: + return len(self.identities) + + @property + def identities_sha256(self) -> str: + return _identities_digest(self.join_key, self.identities) + + def _source_key_frame(self) -> pd.DataFrame: + frame = pd.DataFrame(self.identities, columns=list(self.join_key)) + return _canonicalize_key_frame(frame, self.join_key) + + def base_selection_mask( + self, base_frame: Frame, *, mode: SelectionMode = "frozen_support" + ) -> tuple[np.ndarray, SelectionReport]: + """Return a household-level boolean mask over ``base_frame`` and a report. + + The mask marks base households whose identity is in this source's set. + """ + if mode not in SELECTION_MODES: + raise ValueError( + f"unknown selection mode {mode!r}; expected one of {SELECTION_MODES}." + ) + base_key = _base_household_key_frame(base_frame, self.join_key) + _assert_key_unique(base_key, self.join_key, context="base frame") + + source_key = self._source_key_frame() + if source_key.duplicated().any(): + raise ValueError( + "selection source contains duplicate identities under the join " + f"key {self.join_key}." + ) + + # Left-join source identities onto base rows to find matches and count + # ambiguity (one source identity hitting >1 base row). The base key is + # already asserted unique above, so a clean base yields no ambiguity; + # this catches a malformed base that slipped a duplicate identity. + merged = source_key.merge( + base_key.assign(_base_row=np.arange(len(base_key))), + on=list(self.join_key), + how="left", + ) + match_counts = merged.groupby(list(self.join_key), sort=False)[ + "_base_row" + ].transform("count") + n_ambiguous = int((match_counts > 1).sum()) + if n_ambiguous: + raise ValueError( + f"{n_ambiguous} source identities map to more than one base row " + f"under the join key {self.join_key}; refusing an ambiguous " + "selection." + ) + + unmapped = merged["_base_row"].isna() + n_unmapped = int(unmapped.sum()) + if n_unmapped and mode == "frozen_support": + examples = merged.loc[unmapped, list(self.join_key)].head(5).values.tolist() + raise ValueError( + f"{n_unmapped} of {len(source_key)} source identities are " + "unmapped onto the base under the join key " + f"{self.join_key}; refusing to fabricate or truncate a frozen " + f"support. Example unmapped identities: {examples}." + ) + + selected_rows = merged.loc[~unmapped, "_base_row"].to_numpy(dtype="int64") + mask = np.zeros(base_frame.n("household"), dtype=bool) + mask[selected_rows] = True + report = SelectionReport( + join_key=self.join_key, + n_source=len(source_key), + n_base_candidates=int(base_frame.n("household")), + n_selected=int(mask.sum()), + n_unmapped=n_unmapped, + n_ambiguous=n_ambiguous, + mode=mode, + provenance=dict(self.provenance), + ) + return mask, report + + +def _canonicalize_key_frame( + frame: pd.DataFrame, join_key: Sequence[str] +) -> pd.DataFrame: + """Coerce key columns to comparable dtypes so joins match across sources. + + String identity columns are compared as strings; integer identity columns as + Python ints. This keeps a base read from H5 (numpy int64) joinable to a + manifest read from JSON (Python int) and to a source frame. + """ + out = pd.DataFrame(index=frame.index) + for column in join_key: + series = frame[column] + if column in ("household_support_channel", "source_person_id"): + out[column] = series.astype(str) + else: + out[column] = pd.to_numeric(series).astype("int64") + return out + + +def _base_household_key_frame( + base_frame: Frame, join_key: Sequence[str] +) -> pd.DataFrame: + """Build the per-household identity key frame in base household-row order. + + Person-level identity columns (``source_year``, ``source_household_id``, + ``source_person_id``) are read from the person table and reduced to one value + per household via the frame's declared person->household linkage. Household- + level columns are read directly. Order matches the household table so the + resulting boolean mask aligns to :meth:`Frame.select`'s expectations. + """ + key = _validate_join_key(join_key) + schema = base_frame.schema + household = base_frame.table("household") + hh_id_column = schema.id_column("household") + result = pd.DataFrame({hh_id_column: household[hh_id_column].to_numpy()}) + + person_columns = [c for c in key if SELECTION_JOIN_KEY_COLUMNS[c] == "person"] + if person_columns: + person = base_frame.table("person") + membership = schema.membership_column("household") + missing = [c for c in person_columns if c not in person.columns] + if missing: + raise ValueError( + f"base person table is missing identity columns {missing} " + f"required by the join key {key}." + ) + # source_person_id reduces by min (a household's smallest member id); + # the others are constant within a household, so first() is exact. + agg = {} + for column in person_columns: + agg[column] = (column, "min" if column == "source_person_id" else "first") + identity = person.groupby(membership, sort=False).agg(**agg).reset_index() + result = result.merge( + identity, left_on=hh_id_column, right_on=membership, how="left" + ) + if result[person_columns].isna().any().any(): + raise ValueError( + "base households without member persons cannot carry a source " + "identity; the join key requires person-level identity." + ) + + household_columns = [c for c in key if SELECTION_JOIN_KEY_COLUMNS[c] == "household"] + missing_hh = [c for c in household_columns if c not in household.columns] + if missing_hh: + raise ValueError( + f"base household table is missing identity columns {missing_hh} " + f"required by the join key {key}." + ) + for column in household_columns: + result[column] = household[column].to_numpy() + + return _canonicalize_key_frame(result[list(key)], key) + + +def _assert_key_unique( + key_frame: pd.DataFrame, join_key: Sequence[str], *, context: str +) -> None: + duplicated = key_frame.duplicated() + if duplicated.any(): + n_dup = int(duplicated.sum()) + raise ValueError( + f"selection join key {tuple(join_key)} is not unique in the " + f"{context} ({n_dup} duplicate rows); a non-unique key makes " + "selection order-dependent. Add a clone/channel component." + ) + + +def select_frozen_support( + base_frame: Frame, source: SelectionSource +) -> tuple[Frame, SelectionReport]: + """Reduce ``base_frame`` to exactly the support named by ``source``. + + Returns the reduced frame and a :class:`SelectionReport`. Raises on a + non-unique join key, any unmapped source identity, or an ambiguous match — + the frozen-support contract forbids a fabricated or truncated selection. + """ + mask, report = source.base_selection_mask(base_frame, mode="frozen_support") + person_mask = _household_mask_to_person_mask(base_frame, mask) + reduced = base_frame.select(person_mask) + return reduced, report + + +def _household_mask_to_person_mask( + base_frame: Frame, household_mask: np.ndarray +) -> np.ndarray: + schema = base_frame.schema + household = base_frame.table("household") + hh_id_column = schema.id_column("household") + selected_ids = household[hh_id_column].to_numpy()[household_mask] + membership = base_frame.table("person")[ + schema.membership_column("household") + ].to_numpy() + return np.isin(membership, selected_ids) + + +def load_selection_source_from_h5( + source: Frame, + *, + join_key: Sequence[str] = DEFAULT_SELECTION_JOIN_KEY, + provenance: Mapping[str, object] | None = None, +) -> SelectionSource: + """Distill a :class:`SelectionSource` from a loaded source frame. + + ``source`` is a populace US frame (e.g. the certified live default loaded via + the release tool's frame loader) whose record set defines the support. Its + per-household identities under ``join_key`` become the selection. + """ + key = _validate_join_key(join_key) + if source.schema != US_SCHEMA: + raise ValueError("selection source frame must use the US schema.") + key_frame = _base_household_key_frame(source, key) + _assert_key_unique(key_frame, key, context="selection source frame") + identities = [list(row) for row in key_frame.itertuples(index=False, name=None)] + return SelectionSource( + join_key=key, + identities=identities, + provenance=dict(provenance or {"kind": "h5"}), + ) + + +def write_selection_source_manifest(source: SelectionSource, path: str | Path) -> Path: + """Write ``source`` to a committed selection-source manifest JSON. + + The manifest embeds ``identities_sha256`` so a later load can verify the + exact set is intact, and the source ``provenance`` so a build records which + artifact the selection descended from. + """ + path = Path(path) + payload = { + "schema_version": _MANIFEST_SCHEMA_VERSION, + "join_key": list(source.join_key), + "source": dict(source.provenance), + "n_selected": source.n_identities, + "identities_sha256": source.identities_sha256, + "identities": source.identities, + } + path.write_text(json.dumps(payload, indent=1, sort_keys=False)) + return path + + +def load_selection_source_from_manifest(path: str | Path) -> SelectionSource: + """Load and integrity-check a committed selection-source manifest.""" + path = Path(path) + payload = json.loads(path.read_text()) + version = payload.get("schema_version") + if version != _MANIFEST_SCHEMA_VERSION: + raise ValueError( + f"unsupported selection-source manifest schema_version {version!r}; " + f"expected {_MANIFEST_SCHEMA_VERSION}." + ) + join_key = _validate_join_key(payload["join_key"]) + identities = [list(row) for row in payload["identities"]] + expected = payload.get("identities_sha256") + actual = _identities_digest(join_key, identities) + if expected != actual: + raise ValueError( + "selection-source manifest identities_sha256 integrity check failed: " + f"recorded {expected!r} != computed {actual!r}." + ) + declared = payload.get("n_selected") + if declared is not None and int(declared) != len(identities): + raise ValueError( + f"selection-source manifest n_selected {declared} does not match the " + f"{len(identities)} identities present." + ) + return SelectionSource( + join_key=join_key, + identities=identities, + provenance=dict(payload.get("source", {})), + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_us_frame_from_h5(path: Path) -> Frame: + """Load a populace US H5 into a Frame (lazy PE-US import, like the release tool).""" + from policyengine_us.data import USSingleYearDataset + + dataset = USSingleYearDataset(file_path=str(path)) + tables = { + "person": dataset.person.copy(), + "household": dataset.household.copy(), + "tax_unit": dataset.tax_unit.copy(), + "spm_unit": dataset.spm_unit.copy(), + "family": dataset.family.copy(), + "marital_unit": dataset.marital_unit.copy(), + } + weights = tables["household"].pop("household_weight").to_numpy(dtype=np.float64) + return Frame( + tables, + US_SCHEMA, + {"household": Weights(weights, WeightKind.CALIBRATED)}, + ) + + +def build_selection_source_manifest_from_h5( + source_h5: str | Path, + manifest_path: str | Path, + *, + join_key: Sequence[str] = DEFAULT_SELECTION_JOIN_KEY, + repo_id: str | None = None, + revision: str | None = None, +) -> SelectionSource: + """Distill a source H5's record set into a committed selection-source manifest. + + The manifest freezes exactly which source records the artifact selected, keyed + on their stable identity, plus provenance (the H5 sha256 and, when given, the + HF repo/revision) so a build records what the selection descended from. + """ + source_h5 = Path(source_h5) + frame = _load_us_frame_from_h5(source_h5) + provenance: dict[str, object] = { + "kind": "h5", + "path": str(source_h5), + "sha256": _sha256_file(source_h5), + } + if repo_id is not None: + provenance["repo_id"] = repo_id + if revision is not None: + provenance["revision"] = revision + source = load_selection_source_from_h5( + frame, join_key=join_key, provenance=provenance + ) + write_selection_source_manifest(source, manifest_path) + return source + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Distill a populace US H5's record set into a committed " + "selection-source manifest (populace#328). The manifest names the " + "frozen support's stable record identities so the sparse default is " + "reproducible without re-downloading the source H5." + ) + ) + parser.add_argument("--source-h5", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument( + "--join-key", + default=",".join(DEFAULT_SELECTION_JOIN_KEY), + help="Comma-separated stable-identity columns (default: %(default)s).", + ) + parser.add_argument("--repo-id", default=None) + parser.add_argument("--revision", default=None) + return parser + + +def main(argv: list[str] | None = None) -> None: + args = _parser().parse_args(argv) + join_key = tuple(p.strip() for p in str(args.join_key).split(",") if p.strip()) + source = build_selection_source_manifest_from_h5( + args.source_h5, + args.output, + join_key=join_key or DEFAULT_SELECTION_JOIN_KEY, + repo_id=args.repo_id, + revision=args.revision, + ) + print( + json.dumps( + { + "output": str(args.output), + "n_selected": source.n_identities, + "join_key": list(source.join_key), + "identities_sha256": source.identities_sha256, + "source": dict(source.provenance), + }, + indent=2, + sort_keys=True, + ) + ) diff --git a/packages/populace-build/tests/test_us_fiscal_refresh_builder.py b/packages/populace-build/tests/test_us_fiscal_refresh_builder.py index 55cdb569..de047c44 100644 --- a/packages/populace-build/tests/test_us_fiscal_refresh_builder.py +++ b/packages/populace-build/tests/test_us_fiscal_refresh_builder.py @@ -4096,6 +4096,167 @@ def __len__(self): assert artifact["sha256"] +def _minimal_manifest_kwargs(builder, release_id, release_dir, artifact_root): + result = SimpleNamespace( + skipped=(), + diagnostics=( + SimpleNamespace( + name=f"nation/cbo/individual_income_tax@{builder.PERIOD}", + target=1.0, + initial_estimate=1.0, + final_estimate=1.0, + ), + ), + initial_loss=2.0, + final_loss=1.0, + ) + + class FakeRegistry: + version = "registry-sha" + + def __len__(self): + return 1 + + return dict( + release_id=release_id, + release_dir=release_dir, + artifact_root=artifact_root, + result=result, + registry=FakeRegistry(), + dropped={"dropped_target_names": []}, + target_profile_gate=builder.GateResult( + name="target_profile_coverage", + passed=True, + details={"requirements_checked": 1}, + ), + default_dataset={"method": "dense_no_l0", "sparse": False}, + ) + + +def test_build_manifests_records_selection_source_provenance( + monkeypatch, tmp_path +) -> None: + # A frozen-support build records its selection provenance in both manifests + # so the informed-L0 step is reproducible from main (populace#328). + builder = _load_builder_module() + release_id = "populace-us-2024-sel-20260706" + release_dir = tmp_path / "release" / release_id + release_dir.mkdir(parents=True) + artifact_root = tmp_path / "artifacts" + artifact_root.mkdir() + (artifact_root / builder.DATASET_FILENAME).write_bytes(b"h5") + (artifact_root / builder.CALIBRATION_FILENAME).write_bytes(b"npz") + (release_dir / "calibration_diagnostics.json").write_text("{}") + (release_dir / "us_source_coverage.json").write_text("{}") + monkeypatch.setattr( + builder, + "_runtime_versions", + lambda: { + "python": "3.14.0", + "populace-data": "0.1.0", + "policyengine-core": "3.26.11", + "policyengine-us": "1.752.2", + }, + ) + monkeypatch.setattr( + builder, + "_git_output", + lambda *args: "a" * 40 if args == ("rev-parse", "HEAD") else "", + ) + monkeypatch.setattr( + builder, + "diagnostics_payload", + lambda result, target_registry: { + "initial_loss": 2.0, + "final_loss": 1.0, + "fraction_within_10pct": 1.0, + "target_surface": {"sha256": "b" * 64, "n_targets": 1}, + }, + ) + + selection_source = { + "mode": "frozen_support", + "join_key": [ + "source_year", + "source_household_id", + "household_support_channel", + "household_support_clone_index", + ], + "source": { + "kind": "h5", + "path": "certified.h5", + "sha256": "c" * 64, + }, + "n_source": 57_240, + "n_base_candidates": 337_704, + "n_selected": 57_240, + "n_unmapped": 0, + "n_ambiguous": 0, + } + + builder._build_manifests( + selection_source=selection_source, + **_minimal_manifest_kwargs(builder, release_id, release_dir, artifact_root), + ) + + build_manifest = json.loads((release_dir / "build_manifest.json").read_text()) + release_manifest = json.loads((release_dir / "release_manifest.json").read_text()) + assert build_manifest["calibration"]["selection_source"] == selection_source + assert release_manifest["build"]["selection_source"] == selection_source + assert build_manifest["calibration"]["selection_source"]["n_selected"] == 57_240 + assert build_manifest["calibration"]["selection_source"]["n_unmapped"] == 0 + + +def test_build_manifests_selection_source_absent_by_default( + monkeypatch, tmp_path +) -> None: + # A build with no selection source records the disabled sentinel, not None. + builder = _load_builder_module() + release_id = "populace-us-2024-nosel-20260706" + release_dir = tmp_path / "release" / release_id + release_dir.mkdir(parents=True) + artifact_root = tmp_path / "artifacts" + artifact_root.mkdir() + (artifact_root / builder.DATASET_FILENAME).write_bytes(b"h5") + (artifact_root / builder.CALIBRATION_FILENAME).write_bytes(b"npz") + (release_dir / "calibration_diagnostics.json").write_text("{}") + (release_dir / "us_source_coverage.json").write_text("{}") + monkeypatch.setattr( + builder, + "_runtime_versions", + lambda: { + "python": "3.14.0", + "populace-data": "0.1.0", + "policyengine-core": "3.26.11", + "policyengine-us": "1.752.2", + }, + ) + monkeypatch.setattr( + builder, + "_git_output", + lambda *args: "a" * 40 if args == ("rev-parse", "HEAD") else "", + ) + monkeypatch.setattr( + builder, + "diagnostics_payload", + lambda result, target_registry: { + "initial_loss": 2.0, + "final_loss": 1.0, + "fraction_within_10pct": 1.0, + "target_surface": {"sha256": "b" * 64, "n_targets": 1}, + }, + ) + + builder._build_manifests( + **_minimal_manifest_kwargs(builder, release_id, release_dir, artifact_root), + ) + + build_manifest = json.loads((release_dir / "build_manifest.json").read_text()) + release_manifest = json.loads((release_dir / "release_manifest.json").read_text()) + assert build_manifest["calibration"]["selection_source"] == {"enabled": False} + assert release_manifest["build"]["selection_source"] == {"enabled": False} + + def test_build_manifests_uses_incumbent_aware_calibration_gate( monkeypatch, tmp_path ) -> None: diff --git a/packages/populace-build/tests/test_us_warm_start_selection.py b/packages/populace-build/tests/test_us_warm_start_selection.py new file mode 100644 index 00000000..14320d83 --- /dev/null +++ b/packages/populace-build/tests/test_us_warm_start_selection.py @@ -0,0 +1,364 @@ +"""Frozen-support / warm-start selection recovery (populace#328). + +These tests exercise the identity join that recovers a frozen support onto a +freshly-rebuilt base, and the refusal contract that forbids fabricated or +truncated selections. All fixtures are tiny synthetic frames — no network, no +large H5. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from populace.build.us_runtime.warm_start_selection import ( + DEFAULT_SELECTION_JOIN_KEY, + SelectionSource, + load_selection_source_from_h5, + load_selection_source_from_manifest, + select_frozen_support, + write_selection_source_manifest, +) +from populace.frame import US_SCHEMA, Frame, WeightKind, Weights + +# A source ASEC household appears twice in the pool: clone 0 (asec channel) and +# clone 1 (puf_tax_detail channel). The identity that survives a base rebuild is +# (source_year, source_household_id, channel, clone) — never the assigned row id. + + +def _record( + *, + person_id: int, + household_id: int, + source_year: int, + source_household_id: int, + source_person_id: str, + channel: str, + clone: int, +) -> dict[str, object]: + return { + "person_id": person_id, + "person_household_id": household_id, + "person_tax_unit_id": household_id, + "person_spm_unit_id": household_id, + "person_family_id": household_id, + "person_marital_unit_id": person_id, + "source_year": source_year, + "source_household_id": source_household_id, + "source_person_id": source_person_id, + "person_support_channel": channel, + "person_support_clone_index": clone, + "household_id": household_id, + "household_support_channel": channel, + "household_support_clone_index": clone, + } + + +def _frame_from_records(records: list[dict[str, object]]) -> Frame: + """Build a minimal US-schema frame, one person per household, from records. + + One person per household keeps the household identity unambiguous while still + exercising the person->household expansion in the selector. + """ + df = pd.DataFrame(records) + person = pd.DataFrame( + { + "person_id": df["person_id"].to_numpy(dtype="int64"), + "person_household_id": df["person_household_id"].to_numpy(dtype="int64"), + "person_tax_unit_id": df["person_tax_unit_id"].to_numpy(dtype="int64"), + "person_spm_unit_id": df["person_spm_unit_id"].to_numpy(dtype="int64"), + "person_family_id": df["person_family_id"].to_numpy(dtype="int64"), + "person_marital_unit_id": df["person_marital_unit_id"].to_numpy( + dtype="int64" + ), + "source_year": df["source_year"].to_numpy(dtype="int64"), + "source_household_id": df["source_household_id"].to_numpy(dtype="int64"), + "source_person_id": df["source_person_id"].astype(str).to_numpy(), + "person_support_channel": df["person_support_channel"] + .astype(str) + .to_numpy(), + "person_support_clone_index": df["person_support_clone_index"].to_numpy( + dtype="int64" + ), + } + ) + household = pd.DataFrame( + { + "household_id": df["household_id"].to_numpy(dtype="int64"), + "household_support_channel": df["household_support_channel"] + .astype(str) + .to_numpy(), + "household_support_clone_index": df[ + "household_support_clone_index" + ].to_numpy(dtype="int64"), + } + ) + n = len(df) + return Frame( + { + "person": person, + "household": household, + "tax_unit": pd.DataFrame( + {"tax_unit_id": df["person_tax_unit_id"].to_numpy(dtype="int64")} + ), + "spm_unit": pd.DataFrame( + {"spm_unit_id": df["person_spm_unit_id"].to_numpy(dtype="int64")} + ), + "family": pd.DataFrame( + {"family_id": df["person_family_id"].to_numpy(dtype="int64")} + ), + "marital_unit": pd.DataFrame( + { + "marital_unit_id": df["person_marital_unit_id"].to_numpy( + dtype="int64" + ) + } + ), + }, + US_SCHEMA, + { + "household": Weights( + np.full(n, 100.0, dtype="float64"), WeightKind.CALIBRATED + ) + }, + ) + + +def _pool_records() -> list[dict[str, object]]: + """A 4-source-household × 2-clone base = 8 households across two years.""" + records: list[dict[str, object]] = [] + hid = 1 + pid = 1 + for source_year, source_hh in [(2024, 11), (2024, 22), (2023, 11), (2023, 44)]: + for clone, channel in [(0, "asec"), (1, "puf_tax_detail")]: + records.append( + _record( + person_id=pid, + household_id=hid, + source_year=source_year, + source_household_id=source_hh, + source_person_id=f"{source_year}{source_hh:04d}0101", + channel=channel, + clone=clone, + ) + ) + hid += 1 + pid += 1 + return records + + +def test__given_named_support__then_reduces_base_to_exactly_that_support() -> None: + # Given a base pool and a source naming 3 of its 8 households (clone-aware) + base = _frame_from_records(_pool_records()) + picks = [ + (2024, 11, "asec", 0), + (2024, 22, "puf_tax_detail", 1), + (2023, 44, "asec", 0), + ] + source = SelectionSource( + join_key=DEFAULT_SELECTION_JOIN_KEY, + identities=[list(p) for p in picks], + provenance={"kind": "test"}, + ) + + # When + reduced, report = select_frozen_support(base, source) + + # Then exactly those households remain, and the report is clean + assert reduced.n("household") == 3 + assert report.n_selected == 3 + assert report.n_unmapped == 0 + assert report.n_ambiguous == 0 + got = set( + zip( + reduced.table("household")["household_support_channel"].tolist(), + reduced.table("household")["household_support_clone_index"].tolist(), + strict=True, + ) + ) + assert ("asec", 0) in got and ("puf_tax_detail", 1) in got + + +def _reassigned_ids( + records: list[dict[str, object]], order: list[int] +) -> list[dict[str, object]]: + """Re-emit the same source identities in a different arrangement. + + ``order`` permutes which source record lands at each ascending household id, + so the assigned ``household_id``/``person_id`` differ from the natural build + while the stable source identity of each record is unchanged. The Frame + invariant (household ids sorted ascending) is preserved. + """ + out: list[dict[str, object]] = [] + for row_position, source_index in enumerate(order, start=1): + rec = dict(records[source_index]) + rec["person_id"] = row_position + rec["household_id"] = row_position + rec["person_household_id"] = row_position + rec["person_tax_unit_id"] = row_position + rec["person_spm_unit_id"] = row_position + rec["person_family_id"] = row_position + rec["person_marital_unit_id"] = row_position + out.append(rec) + return out + + +def test__given_reassigned_base_ids__then_recovers_the_same_support() -> None: + # Given the same source records but assigned to DIFFERENT household ids + records = _pool_records() + base = _frame_from_records(records) + reassigned = _frame_from_records(_reassigned_ids(records, [7, 5, 3, 1, 6, 4, 2, 0])) + picks = [(2024, 11, "asec", 0), (2023, 11, "puf_tax_detail", 1)] + source = SelectionSource( + join_key=DEFAULT_SELECTION_JOIN_KEY, + identities=[list(p) for p in picks], + provenance={"kind": "test"}, + ) + + # When + from_ordered, _ = select_frozen_support(base, source) + from_shuffled, _ = select_frozen_support(reassigned, source) + + # Then the identity join recovers the same support regardless of row order + def support_ids(frame: Frame) -> set[tuple[int, int, str, int]]: + p = frame.table("person") + return set( + zip( + p["source_year"].tolist(), + p["source_household_id"].tolist(), + p["person_support_channel"].tolist(), + p["person_support_clone_index"].tolist(), + strict=True, + ) + ) + + assert support_ids(from_ordered) == support_ids(from_shuffled) + assert from_ordered.n("household") == from_shuffled.n("household") == 2 + + +def test__given_unmappable_identity__then_frozen_mode_refuses() -> None: + # Given a source identity that does not exist in the base + base = _frame_from_records(_pool_records()) + source = SelectionSource( + join_key=DEFAULT_SELECTION_JOIN_KEY, + identities=[[2024, 11, "asec", 0], [1999, 99, "asec", 0]], + provenance={"kind": "test"}, + ) + + # When / Then: it refuses loudly and returns nothing (no truncated selection) + with pytest.raises(ValueError, match="unmapped"): + select_frozen_support(base, source) + + +def test__given_key_missing_channel__then_refuses_non_unique_key() -> None: + # Given a join key without the clone/channel component, a source household's + # two clones collide, so the key is not unique in the base. + base = _frame_from_records(_pool_records()) + source = SelectionSource( + join_key=("source_year", "source_household_id"), + identities=[[2024, 11], [2023, 44]], + provenance={"kind": "test"}, + ) + + # When / Then + with pytest.raises(ValueError, match="unique"): + select_frozen_support(base, source) + + +def test__given_ambiguous_source_row__then_refuses() -> None: + # Given a base where two rows share the full identity key (a malformed base), + # a matching source identity would be ambiguous. + records = _pool_records() + dup = dict(records[0]) + dup["person_id"] = 999 + dup["household_id"] = 999 + dup["person_household_id"] = 999 + dup["person_tax_unit_id"] = 999 + dup["person_spm_unit_id"] = 999 + dup["person_family_id"] = 999 + dup["person_marital_unit_id"] = 999 + base = _frame_from_records([*records, dup]) + source = SelectionSource( + join_key=DEFAULT_SELECTION_JOIN_KEY, + identities=[ + list(records[0][k] for k in ("source_year", "source_household_id")) + + [ + records[0]["household_support_channel"], + records[0]["household_support_clone_index"], + ] + ], + provenance={"kind": "test"}, + ) + + with pytest.raises(ValueError, match="unique|ambiguous"): + select_frozen_support(base, source) + + +def test__given_source_h5_frame__then_loads_identities_from_it() -> None: + # Given a "source" frame standing in for a published H5 + source_frame = _frame_from_records( + [r for r in _pool_records() if r["source_year"] == 2024] + ) + + # When + source = load_selection_source_from_h5( + source_frame, join_key=DEFAULT_SELECTION_JOIN_KEY + ) + + # Then it captures one identity per source household + assert source.n_identities == 4 + assert list(source.join_key) == list(DEFAULT_SELECTION_JOIN_KEY) + + +def test__given_manifest_roundtrip__then_identities_match_and_integrity_holds( + tmp_path: Path, +) -> None: + # Given a selection source distilled from a frame + source_frame = _frame_from_records(_pool_records()) + source = load_selection_source_from_h5( + source_frame, + join_key=DEFAULT_SELECTION_JOIN_KEY, + provenance={"kind": "h5", "path": "source.h5", "sha256": "deadbeef"}, + ) + path = tmp_path / "selection.json" + write_selection_source_manifest(source, path) + + # When + reloaded = load_selection_source_from_manifest(path) + + # Then the identities and key match + assert reloaded.n_identities == source.n_identities + assert list(reloaded.join_key) == list(source.join_key) + reloaded_reduced, report = select_frozen_support( + _frame_from_records(_pool_records()), reloaded + ) + assert report.n_selected == source.n_identities + + # And a corrupted integrity hash is rejected + payload = json.loads(path.read_text()) + payload["identities_sha256"] = "0" * 64 + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="identities_sha256|integrity"): + load_selection_source_from_manifest(path) + + +def test__given_informed_init_mode__then_drops_unmapped_without_raising() -> None: + # Given an unmappable identity but informed_init mode + base = _frame_from_records(_pool_records()) + source = SelectionSource( + join_key=DEFAULT_SELECTION_JOIN_KEY, + identities=[[2024, 11, "asec", 0], [1999, 99, "asec", 0]], + provenance={"kind": "test"}, + ) + + # When: informed_init tolerates unmapped records (they start cold) + mask, report = source.base_selection_mask(base, mode="informed_init") + + # Then: one mapped, one dropped-and-recorded, no raise + assert report.n_selected == 1 + assert report.n_unmapped == 1 + assert int(mask.sum()) == 1 diff --git a/tools/build_us_fiscal_refresh_release.py b/tools/build_us_fiscal_refresh_release.py index 6b668078..05b3c7c7 100644 --- a/tools/build_us_fiscal_refresh_release.py +++ b/tools/build_us_fiscal_refresh_release.py @@ -91,6 +91,13 @@ reform_validation_payload, write_reform_validation, ) +from populace.build.us_runtime.warm_start_selection import ( + DEFAULT_SELECTION_JOIN_KEY, + SELECTION_MODES, + load_selection_source_from_h5, + load_selection_source_from_manifest, + select_frozen_support, +) from populace.calibrate import TargetRegistry, calibrate, calibrate_l0_refit from populace.calibrate.diagnostics import ( diagnostics_payload, @@ -670,6 +677,52 @@ def _parse_args() -> argparse.Namespace: "frame before using household_weight as the optimizer start." ), ) + parser.add_argument( + "--selection-source-h5", + type=Path, + help=( + "Optional populace US H5 whose record set defines a frozen support " + "to recover onto the base pool (populace#328). The base frame is " + "reduced to exactly the source's records — matched by stable source " + "identity (see --selection-join-key), not row order — before target " + "materialization and calibration. This reconstructs the certified " + "default's informed-L0/warm-start step as committed machinery. " + "Mutually exclusive with --selection-source-manifest." + ), + ) + parser.add_argument( + "--selection-source-manifest", + type=Path, + help=( + "Optional committed selection-source manifest JSON (produced by " + "tools/build_us_selection_source_manifest.py) naming the frozen " + "support's record identities directly, so the sparse default is " + "reproducible on a laptop without downloading the source H5. " + "Mutually exclusive with --selection-source-h5." + ), + ) + parser.add_argument( + "--selection-join-key", + default=",".join(DEFAULT_SELECTION_JOIN_KEY), + help=( + "Comma-separated stable-identity columns used to match the selection " + "source's records onto the base pool. Defaults to " + f"{','.join(DEFAULT_SELECTION_JOIN_KEY)!r}. Only stable identity " + "columns are allowed; assigned row ids are rejected because they are " + "order-dependent across a base rebuild." + ), + ) + parser.add_argument( + "--selection-mode", + choices=SELECTION_MODES, + default="frozen_support", + help=( + "How the selection source is applied. 'frozen_support' (default) " + "reduces the base to exactly the named records and refuses any " + "unmapped identity. 'informed_init' is reserved for a future " + "drifting rebuild and is not yet wired into calibration." + ), + ) parser.add_argument("--seed", type=int, default=0) parser.add_argument( "--maximum-microsim-batch-size", @@ -1419,6 +1472,51 @@ def _load_frame(path: Path) -> Frame: ) +def _resolve_selection_source(args): + """Build the frozen-support selection source from CLI args, or None. + + Reconstructs the certified informed-L0/warm-start selection (populace#328) as + a committed input: either a source H5 whose record set defines the support, or + a committed selection-source manifest naming the identities directly. Returns + ``(source, join_key)`` or ``(None, join_key)`` when no selection is requested. + """ + join_key = tuple( + part.strip() for part in str(args.selection_join_key).split(",") if part.strip() + ) + if not join_key: + join_key = DEFAULT_SELECTION_JOIN_KEY + if ( + args.selection_source_h5 is not None + and args.selection_source_manifest is not None + ): + raise ValueError( + "Pass at most one of --selection-source-h5 / --selection-source-manifest." + ) + if args.selection_source_manifest is not None: + source = load_selection_source_from_manifest(args.selection_source_manifest) + if tuple(source.join_key) != join_key and ( + args.selection_join_key != ",".join(DEFAULT_SELECTION_JOIN_KEY) + ): + raise ValueError( + "--selection-join-key " + f"{join_key} conflicts with the manifest's own join key " + f"{tuple(source.join_key)}; omit the flag to use the manifest's." + ) + return source, tuple(source.join_key) + if args.selection_source_h5 is not None: + source_frame = _load_frame(args.selection_source_h5) + provenance = { + "kind": "h5", + "path": str(args.selection_source_h5), + "sha256": _sha256(args.selection_source_h5), + } + source = load_selection_source_from_h5( + source_frame, join_key=join_key, provenance=provenance + ) + return source, join_key + return None, join_key + + def _assert_cd_vintage_support_matches( h5_path: Path, crosswalk_metadata: Mapping[str, object] | None, @@ -4158,6 +4256,7 @@ def _write_release_calibration_diagnostics( gate_failures: Iterable[str], timing: Mapping[str, object] | None = None, warm_start_calibration: Mapping[str, object] | None = None, + selection_source: Mapping[str, object] | None = None, default_dataset: Mapping[str, object] | None = None, incumbent_diagnostics_path: Path | None = None, incumbent_diagnostics: Mapping[str, Mapping[str, object]] | None = None, @@ -4260,6 +4359,11 @@ def _write_release_calibration_diagnostics( if warm_start_calibration is not None else {"enabled": False} ), + "selection_source": ( + dict(selection_source) + if selection_source is not None + else {"enabled": False} + ), "default_dataset": ( dict(default_dataset) if default_dataset is not None else None ), @@ -4467,6 +4571,7 @@ def _build_manifests( ecps_parity_gate: GateResult | None = None, timing: Mapping[str, object] | None = None, warm_start_calibration: Mapping[str, object] | None = None, + selection_source: Mapping[str, object] | None = None, default_dataset: Mapping[str, object] | None = None, staging: Mapping[str, object] | None = None, ledger_artifact: Mapping[str, object] | None = None, @@ -4502,6 +4607,9 @@ def _build_manifests( if warm_start_calibration is not None else {"enabled": False} ) + selection_source_payload = ( + dict(selection_source) if selection_source is not None else {"enabled": False} + ) default_dataset_payload = ( dict(default_dataset) if default_dataset is not None else None ) @@ -4529,6 +4637,7 @@ def _build_manifests( "filename": CALIBRATION_FILENAME, "sha256": calibration_sha, "warm_start": warm_start_payload, + "selection_source": selection_source_payload, "target_surface": { "sha256": diag["target_surface"]["sha256"], "n_targets": diag["target_surface"]["n_targets"], @@ -4634,6 +4743,7 @@ def _build_manifests( "timing": timing_payload, "ledger_artifact": dict(ledger_artifact) if ledger_artifact else None, "warm_start_calibration": warm_start_payload, + "selection_source": selection_source_payload, "default_dataset": default_dataset_payload, **( { @@ -4951,6 +5061,41 @@ def main() -> None: if telemetry is not None: telemetry.stage("load_base_frame", message="Loading base population H5.") base_frame = _load_frame(base_h5) + + # Frozen-support recovery (populace#328): if a selection source is supplied, + # reduce the base pool to exactly that support by stable source identity, + # BEFORE the population mass repair — matching the certified sequence, where + # the 340.1M-person frozen support is rescaled to the 334.2M Census benchmark + # (not the full pool). Reducing here also materializes PolicyEngine over the + # ~57k support instead of the ~338k pool. + selection_source, _selection_join_key = _resolve_selection_source(args) + selection_source_payload: dict | None = None + if selection_source is not None: + if args.selection_mode != "frozen_support": + raise ValueError( + f"--selection-mode {args.selection_mode!r} is not yet wired into " + "calibration; only 'frozen_support' is implemented (populace#328)." + ) + if telemetry is not None: + telemetry.stage( + "frozen_support_selection", + message="Recovering frozen support by stable source identity.", + n_source=selection_source.n_identities, + join_key=list(selection_source.join_key), + ) + base_frame, selection_report = select_frozen_support( + base_frame, selection_source + ) + selection_source_payload = selection_report.as_manifest() + if telemetry is not None: + telemetry.stage( + "frozen_support_selection_done", + message="Reduced base pool to the frozen support.", + n_selected=selection_report.n_selected, + n_base_candidates=selection_report.n_base_candidates, + n_unmapped=selection_report.n_unmapped, + ) + base_frame, base_population_repair = _with_base_population_mass_repair(base_frame) base_frame, social_security_component_repair = ( _with_social_security_component_value_repair(base_frame, target_specs) @@ -5371,6 +5516,7 @@ def main() -> None: "social_security_components": social_security_component_repair }, warm_start_calibration=warm_start_calibration, + selection_source=selection_source_payload, audit_export_targets=bool(args.audit_export_targets), gate_failures=gate_failures, timing=timing, @@ -5572,6 +5718,7 @@ def main() -> None: ecps_parity_gate=enforced_ecps_parity_gate, timing=timing, warm_start_calibration=warm_start_calibration, + selection_source=selection_source_payload, ledger_artifact=ledger_artifact.provenance(), default_dataset=default_dataset, staging=( diff --git a/tools/build_us_selection_source_manifest.py b/tools/build_us_selection_source_manifest.py new file mode 100644 index 00000000..44eb3b2e --- /dev/null +++ b/tools/build_us_selection_source_manifest.py @@ -0,0 +1,11 @@ +"""CLI wrapper: distill a source H5 into a committed selection-source manifest. + +See populace#328. The manifest names a frozen support's stable record +identities so the sparse US default is reproducible without re-downloading the +source H5. +""" + +from populace.build.us_runtime.warm_start_selection import main + +if __name__ == "__main__": + main()