diff --git a/changelog.d/us-acs-local-state-cd-sparse.added.md b/changelog.d/us-acs-local-state-cd-sparse.added.md new file mode 100644 index 000000000..3c2287aef --- /dev/null +++ b/changelog.d/us-acs-local-state-cd-sparse.added.md @@ -0,0 +1 @@ +Calibrate the US ACS local-area release (`tools/build_us_acs_local_release.py`) on a sparse target matrix and add the `--soi-mode state_cd` surface. The materialize stage now writes a (targets x households) float32 CSR matrix (`target_matrix.npz`, row-aligned with `target_registry.json`) and a row-aligned `target_roles.json` beside a structure-only lean H5; no dense households x targets matrix exists at any point, and calibration builds each training target from its registry spec with a CSR-row measure, so the unchanged calibrate kernel compiles it. District SOI rows are materialized as one geography-free carrier per concept restricted to each district's households, which equals the direct materialization (one row per carrier checked on the first chunk; every chunk rebuilds each `state_cd` state parent from its stored district rows, and in other modes checks each district row that has a same-concept state row against it). `state_cd` binds the TY2022 SOI congressional-district file's 21,743 district rows (427 districts; at-large states, the district file's mislabeled SALT and PTC columns, and four (state, concept) blocks outside a 1.25x factor band excluded) on top of the `state` surface, keeping one vintage per state concept: Historic Table 2 sets each state level, district rows keep only their within-state shares, and the six district-file-only measures are lifted onto the Historic Table 2 basis by a sibling concept's two state levels; a rebase or bridge ratio more than 1.25x from its measure's median across states is dropped and recorded, and a sign conflict or a district file whose state row disagrees with its own districts is refused. A hash-assigned 10% of (state x concept family) district blocks is held out of calibration and scored against a population pro-rata baseline; `calibration_summary.json` also records Kish ESS nationally, per state, per district and over distinct households, the top-1% weight share, and household-weight share by spine, and it, `weights_latest.npz`, `consumer_export.json` and `spine_qa.json` carry the run-identity digest that resume, finalize and package require, with the solver settings and a calibrated-weights digest binding the H5 to the summary that describes it; the run identity also binds the lean H5. `--sample-fraction` draws f001/f004/f010/f025 development rungs, which package refuses. `microcosm.build.holdout` gains `hash_holdout_unit`. The ladder population targets' district hierarchy ids now carry the 119th-plan prefix (`5001900US`) their ladder uses. diff --git a/docs/us-acs-local-soi-target-surface.md b/docs/us-acs-local-soi-target-surface.md index 02e71bad3..dbca25b16 100644 --- a/docs/us-acs-local-soi-target-surface.md +++ b/docs/us-acs-local-soi-target-surface.md @@ -1,15 +1,20 @@ # US ACS local-area SOI target surface `tools/build_us_acs_local_release.py --stage materialize` calibrates the ACS -local-area artifact to a state-level administrative surface: USDA SNAP, -CMS Medicaid enrollment and IRS SOI. `--soi-mode` decides which SOI specs it -keeps. +local-area artifact to an administrative surface: USDA SNAP, CMS Medicaid +enrollment and IRS SOI, plus PUMA-ladder population marginals. `--soi-mode` +decides which SOI specs it keeps. | `--soi-mode` | What it keeps | Status | |---|---|---| | `state` | State-level `irs_soi` specs at `ledger_geography_level == "state"` whose `ledger_layout_record_set_spec_id` is not a congressional-district file, of any `target_role` | Default | | `totals` | State-level `irs_soi` specs whose `target_role` is not `soi_fiscal_distribution` | Explicit opt-in | | `full` | Every state-level `irs_soi` spec | Explicit opt-in | +| `state_cd` | `state`, plus the congressional-district file's district rows reconciled to one vintage per state concept, and the district file's state rows for concepts Historic Table 2 lacks (see [The `state_cd` surface](#the-state_cd-surface)) | Explicit opt-in | + +Every mode now materializes into a sparse target matrix (see +[Target matrix storage](#target-matrix-storage)), so the surface size is no +longer bounded by a dense households x targets matrix. "State-level" means the spec's metadata carries `state_fips`, which the congressional-district SOI rows also carry. The rule lives in @@ -110,14 +115,13 @@ are aged to 2024. Its district rows reconcile to their state parents in the same file. Using them needs one vintage per state concept and a sparse target matrix. -### Matrix size +### Matrix size before the sparse checkpoint -This is arithmetic from the code, not a memory measurement. -`materialize_chunked` allocates one float32 column per admin spec for every -household (`np.memmap(..., dtype=np.float32, shape=(n_households, -len(names)))`), and `write_lean_checkpoint` copies that memmap into memory -(`np.array(admin_matrix)`). At 1,588,854 households, 4 bytes x households x -specs gives: +Until 2026-09-27 `materialize_chunked` allocated one float32 column per admin +spec for every household (`np.memmap(..., dtype=np.float32, +shape=(n_households, len(names)))`), and `write_lean_checkpoint` copied that +memmap into memory. At 1,588,854 households, 4 bytes x households x specs +gave (arithmetic from the code): | Surface | Admin specs | Dense float32 admin matrix | |---|---:|---:| @@ -125,10 +129,350 @@ specs gives: | `state` (default) | 3,972 | 23.5 GiB | | `full` | 31,066 | 183.9 GiB | -Measured peaks for comparison: Build P's ACS local release, on this surface, -recorded a 93.9 GB materialize peak (`build_manifest.json` → -`materialize.peak_rss_gb`); the 2026-09-22 hours rebuild on `totals` peaked at -75.6 GB in materialize. `full` does not fit a 128 GB machine as a dense matrix. +Measured peaks for comparison: Build P's ACS local release, on the `state` +surface, recorded a 93.9 GB materialize peak (`build_manifest.json` → +`materialize.peak_rss_gb`); the 2026-09-23 `state` run peaked at 77.9 GB, of +which the checkpoint write was the jump from 55.4 GB; the 2026-09-22 hours +rebuild on `totals` peaked at 75.6 GB. [Target matrix +storage](#target-matrix-storage) replaces this. + +## The `state_cd` surface + +`state_cd` binds district-level SOI targets. It is the `state` surface plus +the TY2022 SOI congressional-district file (`22incd.csv`, record-set spec +`irs_soi.congressional_district_2022.all_returns.v1`), reconciled so that +every state concept has one vintage. The rule is +`us_acs_local_cd_surface.state_cd_soi_surface`. + +### One vintage per state concept + +Both files carry state totals for 48 of the district file's measures: +Historic Table 2 (TY2022) and the district file's own `_total` rows. +After aging they disagree by a few percent. `state_cd` keeps **Historic Table +2 as the single vintage of every state concept it carries**, and uses the +district file only for each district's **share** of its state: + +``` +district target = Historic Table 2 state total + x district-file district value / sum of the district file's districts in that state +``` + +Why Historic Table 2: + +- **It is the level `state` already calibrates to** (Build O, Build P and the + 2026-09-23 run), so a `state_cd` build is comparable to them state by state. +- **The district file's state levels are off in ways the shares are not.** The + feed labels the file "Congressional District Data 2022" but stamps it tax + year 2023 (PR #1040, #1030), so it is aged one year less than Historic Table + 2; its taxable-interest rows are not rebased to Table 4.3 (Historic Table + 2's are; the rebase factor here is about 2.47). A within-state share is + unaffected by a level factor common to the state. +- **Two vintages of one concept are contradictory constraints**; the solve can + only split the difference. + +The rebase factors on the pinned feed, by measure across the 43 states with +district rows, are mostly 1.00 to 1.10. + +- **Counts cluster near 1.017** (median `return_count` factor), from 1.00 to + 1.03. Counts are never aged in either file (`aging_factor` 1, + `not_dollar_amount`), so this is purely a level difference between the two + publications: the district file's state totals count about 1.7% fewer + returns than Historic Table 2 (California: 18,242,570 against 18,487,690). +- **Amounts cluster near 1.05.** That is the same level difference times the + aging gap: the district file is aged 2023→2024 by one CBO growth factor + (about 1.087 for AGI), while Historic Table 2 is aged 2022→2024 on the + chained SOI and CBO series (about 1.120). The smallest amount factor, + 1.03, is that aging ratio alone. Correcting the district file's stamp + (#1030) would remove the aging part only. +- **Two factors reflect known level differences:** taxable interest (2.35 to + 2.70, the Table 4.3 rebase) and capital-gains amounts (0.77 to 1.00). +- **The largest single-state factors** were on qualified and ordinary + dividends (Hawaii, 2.68 and 2.17), tax-exempt interest (Utah, 1.43) and + rental income (New York, 0.72), where the two tables disagree about one + state's level. The factor band below drops those blocks. + +`materialize_rss.json` records the full table under +`soi_surface.rebase_factor_by_measure`. + +### The factor band + +A factor mixes a part common to every state (the coverage and aging gap, or +a measure-wide rebase such as taxable interest's Table 4.3 factor) with a +part specific to one state. The median across states absorbs the common +part. A state whose factor sits more than 1.25x from its concept's median +(`STATE_CD_FACTOR_BAND`, either direction) is one where the two publications +disagree about that state, so neither the district file's shares nor a +bridge built on that sibling is trusted there: + +- a rebase block out of band loses its district rows + (`rebase_out_of_band:`); its Historic Table 2 state row stays; +- a bridge is distrusted if either verdict on its sibling fails: the + sibling's two state levels (Historic Table 2 over the district file's state + row, against the median over every state carrying both), or, where the + sibling has district rows, the rebase of its own district block (against + the median over the states with district blocks). The two medians can + differ, because at-large states count only in the first. A distrusted + bridge loses the bridged state row and its district rows + (`level_bridge_out_of_band:`). + +Every failing verdict is recorded in the receipt under +`factor_band.out_of_band` with the verdict, its ratio, the median and the +relative gap. The band only judges finite positive ratios: a sibling whose +two levels differ in sign is a data defect and is refused outright. So is a +district file that disagrees with itself: wherever the file carries a state +row beside its district rows, the two must agree to 1e-6 +(`STATE_CD_INTERNAL_RTOL`), since no median absorbs a column or crosswalk +defect inside the file. On the pinned feed all 2,193 blocks agree to float +precision (largest relative gap 4e-16; the receipt records it under +`internal_consistency`, and the feed-gated test pins it). + +On the pinned feed the band drops eight: + +| State | Measure | Basis | Ratio | Median | Relative | +|---|---|---|---:|---:|---:| +| Hawaii | `qualified_dividends_amount` | rebase | 2.68 | 1.13 | 2.37 | +| Hawaii | `ordinary_dividends_amount` | rebase | 2.17 | 1.07 | 2.03 | +| New York | `rental_royalty_income_amount` | rebase | 0.72 | 1.05 | 0.68 | +| Utah | `tax_exempt_interest_amount` | rebase | 1.43 | 1.09 | 1.32 | +| Wyoming | `charitable_amount`, `interest_paid_deduction_amount` | bridge (itemized) | 1.75 | 1.06 | 1.65 | +| South Dakota | `charitable_amount`, `interest_paid_deduction_amount` | bridge (itemized) | 1.44 | 1.06 | 1.36 | + +That removes 34 district rows (Hawaii 2 x 2, New York 26, Utah 4) and four +bridged state rows (Wyoming and South Dakota are at-large, so they had no +district rows). The largest gaps the band keeps are 1.22x (Wisconsin +partnership and S-corporation income, West Virginia capital gains) and 1.20x +(Mississippi capital gains); every other kept block is within 1.17x. So on +this feed any tolerance above the largest kept gap (1.220) and below Utah's +(1.316) drops the same eight. The band +is a reviewed tolerance, not an estimate; changing it moves the contract +counts below, which a feed-gated test pins. + +The six measures only the district file has (`charitable_*`, +`interest_paid_deduction_*`, `qualified_business_income_deduction_*`) have no +Historic Table 2 parent. Their parent is the district file's own state row, +**lifted onto the Historic Table 2 basis** by a sibling concept both files +carry in the same state (`STATE_CD_LEVEL_BRIDGES`): + +- counts use `return_count`; +- charitable and interest-paid amounts use `itemized_deductions_amount`; +- QBI deduction amounts use `adjusted_gross_income`. + +The bridge assumes the measure shares its sibling's Historic Table 2 / +district-file ratio in that state: the coverage and aging gap above, plus +whatever the two publications disagree about for the sibling. It is an +estimate, not an identity, which is why the factor band applies to it too. +Without it these state rows and their 2,562 district rows would sit 2–6% +below every related concept. On the pinned feed the bridge factors are +1.00–1.03 for counts (median 1.016), 1.03–1.75 for the itemized-deduction +amounts (median 1.062; the two above 1.25x the median are dropped) and +1.03–1.08 for QBI amounts (median 1.049). The bridge is stamp-invariant: if #1030 corrects the +stamp, the sibling ratio shrinks with it. Its district rows keep their +within-state shares, as every other district row does. + +Every district row records its parent (`state_cd_parent_target_name`), the +parent's basis, the district file's own value and the rebase factor; a +bridged state row records its sibling and bridge factor. + +### What stays off the surface + +| Rows | Specs on the pinned feed | Why | +|---|---:|---| +| District-file `limited_state_local_taxes_*`, both geographies | 974 | Column A18425/N18425 is state and local income taxes, not the limited SALT deduction (microcosm#1038) | +| District-file `premium_tax_credit_returns`, both geographies | 487 | Column N85530 is the additional Medicare tax, not the premium tax credit (microcosm#1038) | +| District `tax_filer_individual_count` | 436 | No state parent in either vintage, so it could not nest in a bound state target | +| District rows of states with one district on the 117th plan (AK, DE, DC, MT, ND, SD, VT, WY) | 459 | The SOI file has no sub-state rows there. These rows are the state total copied, or for Montana split by population, so they carry no district information | +| Historic Table 2 rows copied to at-large districts | 360 | The same copies from the other vintage | +| District-file state rows for concepts Historic Table 2 carries | 2,295 | Second vintage of a state concept | +| District rows of the four out-of-band rebase blocks | 34 | The factor band (above) | +| Bridged state rows whose sibling ratio is out of band | 4 | The factor band (above) | + +PR #1040 (awaiting a ruling) excludes the same SALT and PTC columns in the +compiler and rescales the district file's capital-gains rows by one national +factor. `state_cd` already takes district capital gains as within-state +shares of Historic Table 2, which a national factor does not move, so their +level changes only when Historic Table 2's own capital-gains rows do +(#1036). + +North Carolina's district rows are bound. Before #1043 the packaged +117th→119th crosswalk was built from NC's 2016 plan rather than the 2019 plan +the 117th Congress used, and 37.0% of NC's population mapped to a different +119th district than the block plan registry (#1041) puts it in. #1043 rebuilt +the crosswalk from the registry. `STATE_CD_EXCLUDED_CD_STATES` is empty, and a +test pins the crosswalk digest it was reviewed against, so a later crosswalk +change forces the same review. + +### Counts on the pinned feed + +| Family | `state_cd` | +|---|---:| +| `usda_snap` | 102 | +| `cms_medicaid` (enrollment) | 51 | +| `irs_soi` state, Historic Table 2 | 3,819 | +| `irs_soi` state, district file (district-file-only measures) | 302 | +| `irs_soi` district (427 districts x 51 measures, less 34 banded) | 21,743 | +| **Admin specs** | **26,017** | + +Of the 21,743 district rows, 19,181 are rebased to a Historic Table 2 parent +and 2,562 keep a district-file parent. The 2,189 (state, concept) district +blocks each sum to their parent within 1e-9. Adding the 487 population +marginals gives 26,504 targets, before the holdout. The feed-gated test +`test_pinned_feed_state_cd_surface_matches_its_contract` pins these counts +and the reconciliation. + +### District plans + +The SOI district file is tabulated on the 117th-Congress plan and the +households carry 119th-plan districts (drawn within their PUMA from the +ladder). The compiler maps the 117th rows onto the 119th plan with the +packaged 2020-block population crosswalk, which #1043 builds from the block +plan registry (#1041); `state_cd` uses that mapping as is and records the +crosswalk's sha256. The mapping assumes returns spread with population inside +each 117th/119th intersection. A household column on the 117th plan +(`congressional_district_geoid__117th_congress`, written by the location v1 +block draw once the registry is attached) would let these targets bind as +exact block sums with no crosswalk; the district rows are materialized +against one named household column, so that change is a parameter, not a +rewrite. + +### Sigma + +No fact in the pinned feed carries an uncertainty field, and IRS SOI tables +are administrative. `target_roles.json` records `sigma` for every target (`null`, +`sigma_basis: "not_provided_by_feed"`), so a feed that supplies standard +errors surfaces them. The calibration loss is unchanged: fixed-scale capped +relative error. + +## CD holdout and the pro-rata baseline + +`state_cd` holds a hash-assigned subset of its district targets out of +calibration and scores it afterwards (`--cd-holdout-fraction`, default 0.1 in +`state_cd`, 0 elsewhere). + +**The held unit is a (state, SOI concept family) block** of district targets, +not a single district or target. With a state total and its sibling districts +trained, one held district is pinned by adding up, so it would score +perfectly for free. A concept family also groups measures that add up to each +other: every EITC measure is one family, because the per-child rows sum to +the EITC total. The state totals stay trained; the question a held block +answers is how well the calibrated file splits a known state total across +districts. + +**Assignment** is `microcosm.build.holdout.hash_holdout_unit`: SHA-256 of a +salt and the unit key `"|"`, read as a uniform on +[0, 1), held iff below the fraction. It is deterministic, does not depend on +which other units exist or their order, and is nested in the fraction. Held +targets are never built into the calibrator's `TargetSet` +(`calibration_target_set`); a property test spies on the solve to check it. + +**The baseline** allocates each held target's state parent by population: +`parent x district population / state population`, both from the PUMA +ladder's 119th-plan district overlap populations, so a held block's baseline +sums to its parent exactly. `calibration_summary.json` → `cd_holdout` +scores the held targets under the design weights, the calibrated weights and +the baseline (mean, median and p90 absolute relative error, share within 10%, +the capped loss the solve minimizes, per family, and the share of targets +where the calibration beats the baseline). It is report-only. + +## Effective sample size and weight by origin + +`calibration_summary.json` → `weight_origin` records, at the design and at +the calibrated weights: + +- Kish ESS over rows: nationally, per spine, per state and per district + (each with its distribution); +- Kish ESS over **distinct households**: rows summed by (`household_spine`, + `household_source_id`). ACS source ids are pre-offset and collide with donor + ids, so the spine is part of the key; a donor household's native row and its + PUF-detail clone count once; +- the share of total weight on the heaviest 1% of records; +- household-weight share by spine. + +The weight cap and `l2_lambda` are unchanged; the concentration they allow is +a known issue (`low_effective_sample_size_lambda_zero`) under review. This +reports it rather than tuning it. + +## Development rungs + +`--sample-fraction` draws a development rung, one of the stacked pool's +f001, f004, f010 or f025 (`tools/build_us_multispine_pool.py`; DESIGN.md +"Production US stacked spine" names f001, f010 and f100). + +- **The draw.** It samples whole households with + `microcosm.build.frame_sampling.sample_frame_households`, taking + `floor(fraction x n_h)` in every (spine, district) stratum `h`. +- **The weights.** Each drawn household's weight is scaled by its stratum's + inverse sampling rate `n_h / k_h`, so a drawn stratum keeps its household + mass in expectation (exactly, when a stratum's weights are equal). Each + spine is then scaled to its full household mass. +- **Strata that draw nothing.** A stratum that floors to zero draws loses its + households, and its weight is spread over the rest of its spine. The + receipt records how many strata that is and their weight share. At f001 + this is material for the donor spine, whose district strata are small, so + read district-level evidence from f010 or above. +- **Refusals.** `run_identity.json` records the rung, seed and selected-id + digest. The calibrate stage re-draws the same households before attaching + weights and refuses a mismatch, and `--stage package` refuses any rung but + f100. +- **Memory.** The staging frame is still loaded in full before sampling, so a + rung lowers the engine pass and the solve, not the load. + +## Target matrix storage + +The materialize stage writes four files: + +- `target_frame_lean.h5`: structure only (household id, geography, spine, + source id, design weight; person memberships; group ids). +- `target_registry.json`: every target as a `TargetSpec` with its + calibration hierarchy, held-out targets included. Its spec count and + digest therefore identify the materialized surface, not the trained set; + the trained count is `calibration_summary.json`'s `n_targets`, and the + roles file says which rows trained. +- `target_matrix.npz`: a (targets x households) CSR matrix with float32 + values, row *i* of which is registry spec *i*. +- `target_roles.json`, row-aligned with both: role (train or holdout), + holdout unit, geography, sigma, and for district rows their state parent + and pro-rata populations. + +No dense households x targets matrix exists at any point. Each engine chunk's +columns go straight into the CSR. District SOI rows are not materialized one +column each: the engine pass materializes one geography-free **carrier** +column per distinct SOI concept, and each district row is its carrier +restricted to the district's households. That equals the direct +materialization, because the SOI slice masks a tax unit by its household's +state and district; the first chunk of every run also materializes one +district row per carrier directly and compares it with the row the assembler +actually stored for it, refusing any difference. That check covers one row +per carrier in one chunk, which could miss the households a concept touches, +so every chunk also checks the stored district rows against a state row +materialized directly in the same pass (`cd_surface.district_row_parents`): + +- In `state_cd`, each district row names its state parent, and a block's + rows partition the parent's state. Laid side by side, they must equal the + parent's column on every household of the chunk, bit for bit, and every + block is checked in every chunk. A `state_cd` row without a compiled parent + is refused, so the check cannot silently skip it. +- In the other modes, a district row is checked against a state row of the + same materializer semantics in the same state (in `full`, the district + file's own state total), on the row's own households, since those blocks + need not cover the state. A row with no such state row is counted as + unchecked. This check and the carriers share one key + (`soi_materializer_semantics`), so a field the materializer reads that the + key omitted would slip past both; the `state_cd` check compares against a + named parent and does not share that blind spot. + +`materialize_rss.json` → `carrier_check` records the blocks and rows checked, +the nonzero households they covered, and the unchecked rows. The calibrate +stage builds +each training target from its registry spec (value, metadata, hierarchy) +with a callable measure that reads its CSR row, so the calibrate kernel's own +`build_constraint_matrix` compiles them one row at a time, unchanged, and the +schema-8 `calibration_diagnostics.json` names each measure by its matrix row +(`target_matrix_row[i]`). + +A differential test materializes the fixture both ways: dense float32 +columns compiled by the kernel from frame columns, and the carrier-split CSR +through the checkpoint. It requires the identical constraint matrix and the +identical calibrated weights. ## Where the chosen mode is recorded @@ -143,16 +487,64 @@ recorded a 93.9 GB materialize peak (`build_manifest.json` → reproduces the recorded surface even if the default changes again. `--stage package` refuses a checkpoint whose `materialize_rss.json` does not -record one of the three modes, before any release directory exists. +record one of the four modes, before any release directory exists. + +`state_cd` runs also record: + +- `run_identity.json`: the sha256 of `target_registry.json`, + `target_roles.json`, `target_matrix.npz` (with its shape and nnz) and the + lean H5; the CD holdout (unit, salt, fraction, held units and targets); the + sampling rung. Calibrate and finalize refuse any of the four files whose + bytes changed. +- `weights_latest.npz` and `calibration_summary.json` carry the digest of + that run identity (`run_identity_sha256`) and the solver settings (weight + cap, loss cap, `l2_lambda`, seed, epoch batch). `--resume` and the + calibrate stage's "already complete" shortcut refuse weights or a summary + from another materialization (another staging file, surface, holdout or + sample) or other settings; finalize and package refuse a summary from + another materialization. +- `consumer_export.json` (written with the calibrated H5) records the H5's + sha256, the run-identity digest, the solver settings and a digest of the + calibrated weights, which `calibration_summary.json` also records; + `spine_qa.json` records the run-identity digest and the sha256 of the bytes + it loaded. Before any release directory exists, finalize and package + refuse an H5 whose bytes are not the export's, an export and summary that + describe different weights (a recalibration that stopped part way), and QA + evidence from another materialization or of other bytes. The H5's path is + not compared; its sha binds the bytes wherever they are reached from. +- A new materialize deletes the previous calibration outputs, consumer + export, spine QA and gate report. A solve deletes the previous summary, + diagnostics, consumer export, null-fill manifest and spine QA (keeping the + resume weights) before it starts; a calibrated H5 or gate report left from + before fails the sha checks above. The "already complete" shortcut also + requires the summary to record the saved weights' digest. +- Materialize refuses a district row without a positive ladder district and + state population, since the pro-rata baseline could not score it. +- `calibration_summary.json`: `cd_holdout`, `weight_origin` (ESS nationally, + per spine, per state and per district, over distinct households, and the + top-1% weight share, at the design and the calibrated weights), + `n_holdout_targets` and the checkpoint matrix's shape and nnz, beside the + fit summary. +- `materialize_rss.json`: the full SOI surface receipt (`soi_surface`: counts, + every drop reason, rebase factors, crosswalk digest, sigma), the holdout + receipt with every held unit, and the carrier check. +- `gate_summary.json`: `cd_holdout` and `weight_origin` (report-only), and + the calibrated surface's trained and held-out counts. +- The release refresh recipe, which names `--cd-holdout-fraction`. ## Operating notes -- The default command calibrates to `state`. Raise the supervisor's RSS cap - above Build P's 93.9 GB peak before a full-scale run, and leave room on disk - for the 23.5 GiB memmap. +- The default command calibrates to `state`. The dense admin matrix is gone, + so the materialize peak is the staging frame plus one engine chunk; the + calibrate stage's peak is the artifact write. +- A checkpoint written before the sparse matrix (measures as dense H5 + columns) is refused by `--stage calibrate`; re-run `--stage materialize`. - To reproduce the 2026-09-22 hours rebuild (#974), pass `--soi-mode totals`. To reproduce a build that ran as `full` after `b7922b089`, pass `--soi-mode full`. - A checkpoint materialized under an earlier default records its own mode and packages as that mode. -- `--soi-mode` only matters when `soi` is in `--families`. +- `--soi-mode` only matters when `soi` is in `--families`. `state_cd` is + meant to run with `cd` in `--geographies` (the default), so district + population marginals are bound too; `--cd-holdout-fraction` above 0 is + refused without `state_cd`. diff --git a/experiments/us-acs-local-state-cd-20260927/concentration_report.py b/experiments/us-acs-local-state-cd-20260927/concentration_report.py new file mode 100644 index 000000000..4ae47da5a --- /dev/null +++ b/experiments/us-acs-local-state-cd-20260927/concentration_report.py @@ -0,0 +1,135 @@ +"""Weight concentration of the evaluation's weight vectors, arm by arm. + +Kish ESS nationally, per state and per congressional district, ESS over +distinct households, and the top-1% weight share, for each named weight +vector over one lean checkpoint's households. Uses the tool's own +``weight_origin_summary``. + + uv run python experiments/us-acs-local-state-cd-20260927/concentration_report.py \\ + --lean /target_frame_lean.h5 \\ + --arm design=:design --arm state=:weights ... \\ + --out concentration.json +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +REPO = Path(__file__).resolve().parents[2] + + +def load_tool(): + path = REPO / "tools" / "build_us_acs_local_release.py" + spec = importlib.util.spec_from_file_location("build_us_acs_local_release", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lean", type=Path, required=True) + parser.add_argument( + "--arm", + action="append", + required=True, + help="label=path.npz:key (weights aligned to the lean household table)", + ) + parser.add_argument("--focus-state", default="25") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + tool = load_tool() + with pd.HDFStore(args.lean, mode="r") as store: + households = ( + store.select( + "household", + columns=[ + "household_id", + "state_fips", + "congressional_district_geoid", + "household_spine", + "household_source_id", + ], + ) + if store.get_storer("household").is_table + else store["household"] + ) + origin = { + "spine": households["household_spine"].to_numpy(), + "source_id": households["household_source_id"].to_numpy(), + "state": pd.to_numeric(households["state_fips"]) + .astype(int) + .map("{:02d}".format) + .to_numpy(), + "district": pd.to_numeric(households["congressional_district_geoid"]) + .astype(int) + .map("{:04d}".format) + .to_numpy(), + } + report: dict[str, dict] = {} + for arm in args.arm: + label, rest = arm.split("=", 1) + path, key = rest.rsplit(":", 1) + weights = np.load(path)[key] + if len(weights) != len(households): + raise SystemExit(f"{label}: {len(weights)} weights, {len(households)} rows") + summary = tool.cd_surface.weight_origin_summary(weights, **origin) + focus_districts = { + code: value + for code, value in summary["effective_sample_size_by_district"].items() + if code.startswith(args.focus_state) + } + report[label] = { + "national_ess_rows": summary["effective_sample_size_rows"], + "national_ess_distinct_households": summary.get( + "effective_sample_size_distinct_households" + ), + "top_1pct_weight_share": summary["top_1pct_weight_share"], + "weight_share_by_spine": summary.get("weight_share_by_spine"), + "state_ess": summary["effective_sample_size_by_state_distribution"], + "district_ess": summary["effective_sample_size_by_district_distribution"], + f"state_{args.focus_state}_ess": summary[ + "effective_sample_size_by_state" + ].get(args.focus_state), + f"state_{args.focus_state}_district_ess": focus_districts, + "by_state": summary["effective_sample_size_by_state"], + "by_district": summary["effective_sample_size_by_district"], + } + labels = list(report) + if len(labels) >= 2: + base, *others = labels + for other in others: + for level in ("by_state", "by_district"): + a, b = report[base][level], report[other][level] + ratios = np.asarray([b[k] / a[k] for k in a if a[k] > 0 and k in b]) + report[other][f"{level}_ess_ratio_to_{base}"] = { + "min": float(ratios.min()), + "median": float(np.median(ratios)), + "max": float(ratios.max()), + "share_lower": float((ratios < 1).mean()), + } + args.out.write_text(json.dumps(report, indent=1)) + for label, entry in report.items(): + print( + label, + json.dumps( + { + key: value + for key, value in entry.items() + if key not in ("by_state", "by_district", "weight_share_by_spine") + }, + indent=None, + default=str, + )[:1500], + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experiments/us-acs-local-state-cd-20260927/engine_free_eval.py b/experiments/us-acs-local-state-cd-20260927/engine_free_eval.py new file mode 100644 index 000000000..df92d9f16 --- /dev/null +++ b/experiments/us-acs-local-state-cd-20260927/engine_free_eval.py @@ -0,0 +1,448 @@ +"""Engine-free full-scale evaluation of the sparse ACS local path and state_cd. + +Inputs (read-only): the 2026-09-23 ACS local release run +(``populace-us-2024-buildo-acs-local-767312d60``) -- its dense ``state``-mode +lean checkpoint (1,588,854 households x 4,459 float32 measures), its +calibrated weights, and its staging H5 -- plus the pinned Chronicle feed and +the PUMA ladder. + +Stages (each writes under ``--out`` and can be re-run alone): + +``convert`` + Stream the dense checkpoint into the sparse checkpoint format the tool + now writes (structure H5, target registry, CSR matrix and target roles). +``state`` + Re-calibrate the ``state`` surface from the sparse checkpoint with the + 09-23 settings through the tool's ``calibrate_surface`` and compare the + weights with the ones the dense path produced on 09-23. +``state_cd`` + Add the ``state_cd`` district rows whose state parent is a Historic + Table 2 row. Each such row is exactly its parent's CSR row restricted to + the district's households, because the parent and the district row + materialize identically apart from geography + (``state_cd_soi_surface`` refuses otherwise). Rows whose parent is a + district-file state total (charitable, interest paid, QBI deduction) have + no engine column in the 09-23 checkpoint and are left out. Draw the CD + holdout, calibrate, and score the held-out district targets under design, + ``state``-calibrated and ``state_cd``-calibrated weights against the + pro-rata baseline. + +This is development evidence, not a release: the engine columns are the +09-23 engine pass, and the surface omits the three district-file-only +concept families. +""" + +from __future__ import annotations + +import argparse +import gc +import importlib.util +import json +import os +import time +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse + +REPO = Path(__file__).resolve().parents[2] +RUN = Path( + "/Users/maxghenis/PolicyEngine/_recovered/scratch-backup/893/overnight-20260923/run" +) +DENSE_CHECKPOINT = RUN / "release" / "checkpoints" +STAGING_H5 = RUN / "staging" / "acs_multispine_staging.h5" +FEED = Path( + "/Users/maxghenis/PolicyEngine/_buildh-runtime/inputs/consumer_facts_us_c5e5bf8.jsonl" +) +LADDER = Path( + "/Users/maxghenis/PolicyEngine/_worktrees/populace-acs-clone/build/us/us_puma_ladder_2020.npz" +) +SETTINGS = dict( + epochs=800, + epoch_batch=400, + max_weight_ratio=5.0, + target_loss_cap=1.0, + l2_lambda=0.0, + seed=0, +) + + +def load_tool(): + path = REPO / "tools" / "build_us_acs_local_release.py" + spec = importlib.util.spec_from_file_location("build_us_acs_local_release", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def dense_to_csr(path: Path, n_rows_block: int = 25_000): + """Stream the dense float32 household block into a CSR (targets x hh).""" + + import h5py + + with h5py.File(path, "r") as handle: + group = handle["household"] + names = [item.decode() for item in group["block2_items"][:]] + values = group["block2_values"] + n_households, n_targets = values.shape + rows, cols, data = [], [], [] + for low in range(0, n_households, n_rows_block): + block = values[low : low + n_rows_block] + household, target = np.nonzero(block) + rows.append(target.astype(np.int32)) + cols.append((household + low).astype(np.int32)) + data.append(block[household, target].astype(np.float32)) + struct = pd.DataFrame( + group["block0_values"][:], + columns=[item.decode() for item in group["block0_items"][:]], + ) + weights = group["block3_values"][:, 0].astype(np.float64) + matrix = sparse.csr_array( + (np.concatenate(data), (np.concatenate(rows), np.concatenate(cols))), + shape=(n_targets, n_households), + ) + matrix.sort_indices() + return names, matrix, struct, weights + + +def staging_origin(household_ids: np.ndarray) -> pd.DataFrame: + with pd.HDFStore(STAGING_H5, mode="r") as store: + households = ( + store.select( + "household", + columns=["household_id", "household_spine", "household_source_id"], + ) + if store.get_storer("household").is_table + else store["household"][ + ["household_id", "household_spine", "household_source_id"] + ] + ) + if not np.array_equal(households["household_id"].to_numpy(), household_ids): + raise SystemExit("staging household ids differ from the lean checkpoint's") + return households + + +def roles_and_specs_for(tool, names, values, specs_by_name): + """The 09-23 targets as registry specs (the recompiled feed's) and roles.""" + + roles, specs, pop_names, pop_values = [], [], [], [] + for name, value in zip(names, values, strict=True): + spec = specs_by_name.get(name) + if spec is not None: + if abs(spec.value - value) > 1e-9 * max(1.0, abs(value)): + raise SystemExit(f"{name}: 09-23 value {value} != compile {spec.value}") + roles.append(tool.cd_surface.target_record(spec)) + specs.append(spec) + elif name.startswith(("pop_state_", "pop_cd_")): + pop_names.append(name) + pop_values.append(float(value)) + is_cd = name.startswith("pop_cd_") + roles.append( + { + "name": name, + "value": float(value), + "source": "us_puma_ladder_2020", + "family": "census_population_ladder", + "geography_level": "congressional_district" if is_cd else "state", + "state_fips": name[-4:-2] if is_cd else name[-2:], + "congressional_district_geoid": name[-4:] if is_cd else None, + } + ) + specs.append(None) + else: + raise SystemExit(f"09-23 target {name} is not on the recompiled surface") + pop_specs = iter(tool.population_target_specs(pop_names, pop_values)) + specs = [spec if spec is not None else next(pop_specs) for spec in specs] + return roles, specs + + +def stage_convert(tool, out: Path) -> None: + started = time.time() + dense_targets = json.loads((DENSE_CHECKPOINT / "targets.json").read_text()) + names, matrix, struct, weights = dense_to_csr( + DENSE_CHECKPOINT / "target_frame_lean.h5" + ) + if names != [target["measure"] for target in dense_targets]: + raise SystemExit("dense block order differs from targets.json") + print( + f"dense -> CSR: {matrix.shape} nnz {matrix.nnz:,} ({time.time() - started:.0f}s)" + ) + origin = staging_origin(struct["household_id"].to_numpy()) + struct = struct.assign( + household_spine=origin["household_spine"].astype(str).to_numpy(), + household_source_id=origin["household_source_id"].to_numpy(), + ) + with pd.HDFStore(DENSE_CHECKPOINT / "target_frame_lean.h5", mode="r") as store: + person = store["person"] + groups = {group: store[group] for group in tool.GROUP_IDS} + surface = tool.state_admin_surface( + FEED, ["snap", "medicaid", "soi"], soi_mode="state_cd" + ) + specs_by_name = {spec.name: spec for spec in surface.registry.specs} + roles, specs = roles_and_specs_for( + tool, + [target["name"] for target in dense_targets], + [target["value"] for target in dense_targets], + specs_by_name, + ) + tool.cd_surface.assign_target_roles(roles, fraction=0.0) + struct_tables = { + "household_struct": struct, + "person": person, + "groups": groups, + "weights": weights, + } + _path, _registry, digests = tool.write_lean_checkpoint( + struct_tables, matrix, specs, roles, out / "state" + ) + import pickle + + with open(out / "state_cd_surface.pkl", "wb") as handle: + pickle.dump( + {"specs": list(surface.registry.specs), "receipt": surface.soi_receipt}, + handle, + ) + (out / "state" / "identity.json").write_text( + json.dumps( + { + "source_checkpoint": str(DENSE_CHECKPOINT), + "dense_matrix_nnz": int(matrix.nnz), + "published_matrix_nnz": 24_773_532, + **digests, + "households": int(len(struct)), + }, + indent=2, + ) + ) + print(f"convert done ({time.time() - started:.0f}s, peak {tool.rss():.1f} GB)") + + +def calibrate_and_save(tool, checkpoint: Path, label: str): + """Calibrate a checkpoint through the tool's own functions and record it.""" + + from microcosm.calibrate import write_calibration_diagnostics + + frame, design, registry, roles, matrix = tool.load_checkpoint_surface(checkpoint) + target_set = tool.cd_surface.calibration_target_set( + roles, matrix, frame.n("household"), specs=registry.specs + ) + started = time.time() + result, _done = tool.calibrate_surface(frame, target_set, **SETTINGS) + weights = np.asarray(result.weights, dtype=np.float64) + summary = { + "soi_mode": label, + **SETTINGS, + "n_targets": result.problem.n_targets, + "matrix_format": result.options["matrix_format"], + "matrix_nnz": int(result.problem.matrix.nnz), + "initial_loss": round(result.initial_loss, 6), + "final_loss": round(result.final_loss, 6), + "fraction_within_10pct": round(result.fraction_within_10pct, 4), + "effective_sample_size": round(result.effective_sample_size, 1), + "realized_max_weight_ratio": round(result.realized_max_weight_ratio, 4), + "mass_conserved_ratio": round(float(weights.sum()) / float(design.sum()), 6), + **tool.calibration_evidence( + frame=frame, + roles=roles, + matrix=matrix, + design_weights=design, + weights=weights, + target_loss_cap=SETTINGS["target_loss_cap"], + ), + "total_wall_seconds": round(time.time() - started, 1), + "peak_rss_gb": round(tool.rss(), 3), + } + outcome = write_calibration_diagnostics( + result, + checkpoint / "calibration_diagnostics.json", + target_registry=registry, + build={"soi_mode": label, "experiment": "engine_free_eval"}, + ) + summary["calibration_diagnostics"] = { + "status": outcome.status, + "schema_version": getattr(outcome, "schema_version", None), + "message": getattr(outcome, "message", None), + } + np.savez(checkpoint / "weights.npz", weights=weights, design=design) + (checkpoint / "calibration_summary.json").write_text(json.dumps(summary, indent=1)) + return summary, weights, design, roles, matrix, frame + + +def stage_state(tool, out: Path) -> None: + diagnostics, weights, _design, _records, _matrix, _frame = calibrate_and_save( + tool, out / "state", "state" + ) + dense = np.load(DENSE_CHECKPOINT / "weights_latest.npz")["weights"] + difference = np.abs(weights - dense) / np.maximum(np.abs(dense), 1e-12) + comparison = { + "identical": bool(np.array_equal(weights, dense)), + "max_abs_diff": float(np.abs(weights - dense).max()), + "max_rel_diff": float(difference.max()), + "median_rel_diff": float(np.median(difference)), + "dense_final_loss_0923": 0.015491, + "sparse_final_loss": diagnostics["final_loss"], + "dense_ess_0923": 13631.3, + "sparse_ess": diagnostics["effective_sample_size"], + "torch_threads": int(os.environ.get("OMP_NUM_THREADS", "0") or 0), + } + (out / "state" / "dense_comparison.json").write_text( + json.dumps(comparison, indent=2) + ) + print(json.dumps(comparison, indent=2)) + + +def stage_state_cd( + tool, out: Path, fraction: float, state_weights_path: Path | None = None +) -> None: + import pickle + + cd = tool.cd_surface + frame, design, state_registry, state_records, state_matrix = ( + tool.load_checkpoint_surface(out / "state") + ) + households = frame.table("household") + hh_state = pd.to_numeric(households["state_fips"]).to_numpy(np.int64) + hh_cd = pd.to_numeric(households["congressional_district_geoid"]).to_numpy(np.int64) + with open(out / "state_cd_surface.pkl", "rb") as handle: + surface = pickle.load(handle) + row_of = {record["name"]: index for index, record in enumerate(state_records)} + district_specs, skipped = [], [] + for spec in surface["specs"]: + if spec.metadata.get("ledger_geography_level") != "congressional_district": + continue + if spec.metadata.get("state_cd_parent_basis") != "historic_table_2": + skipped.append(spec.name) + continue + district_specs.append(spec) + rows, cols, data = [], [], [] + records = list(state_records) + for offset, spec in enumerate(district_specs): + parent_row = row_of[spec.metadata["state_cd_parent_target_name"]] + start, stop = ( + state_matrix.indptr[parent_row], + state_matrix.indptr[parent_row + 1], + ) + indices = state_matrix.indices[start:stop] + values = state_matrix.data[start:stop] + keep = ( + hh_cd[indices] == int(spec.metadata["congressional_district_geoid"]) + ) & (hh_state[indices] == int(spec.metadata["state_fips"])) + row = len(state_records) + offset + rows.append(np.full(int(keep.sum()), row, dtype=np.int32)) + cols.append(indices[keep].astype(np.int32)) + data.append(values[keep]) + records.append(cd.target_record(spec)) + district_matrix = sparse.csr_array( + ( + np.concatenate(data), + (np.concatenate(rows) - len(state_records), np.concatenate(cols)), + ), + shape=(len(district_specs), state_matrix.shape[1]), + ) + matrix = sparse.vstack([state_matrix, district_matrix], format="csr") + matrix.sort_indices() + populations = tool.ladder_population(LADDER, ["state", "cd"]) + tool._attach_pro_rata_populations(records, populations["cd"]) + holdout = cd.assign_target_roles(records, fraction=fraction) + struct_tables = tool.extract_struct_tables(frame) + _path, _registry, digests = tool.write_lean_checkpoint( + struct_tables, + matrix, + (*state_registry.specs, *district_specs), + records, + out / "state_cd", + ) + (out / "state_cd" / "surface.json").write_text( + json.dumps( + { + "district_rows": len(district_specs), + "district_rows_nnz": int(district_matrix.nnz), + "state_rows_nnz": int(state_matrix.nnz), + "total_nnz": int(matrix.nnz), + "skipped_district_rows_without_engine_column": len(skipped), + "skipped_examples": skipped[:5], + "holdout": holdout, + "surface_receipt_counts": surface["receipt"]["counts"], + **digests, + }, + indent=1, + ) + ) + del frame, state_matrix, district_matrix, matrix + gc.collect() + diagnostics, weights, design, records, matrix, _frame = calibrate_and_save( + tool, out / "state_cd", "state_cd" + ) + state_weights = np.load(state_weights_path or out / "state" / "weights.npz")[ + "weights" + ] + same_held = cd.score_cd_holdout( + records, + matrix, + design_weights=design, + final_weights=state_weights, + cap=SETTINGS["target_loss_cap"], + ) + (out / "state_cd" / "holdout_under_state_weights.json").write_text( + json.dumps({k: v for k, v in same_held.items() if k != "targets"}, indent=1) + ) + summary = { + "state_cd": { + k: diagnostics[k] + for k in ( + "n_targets", + "n_holdout_targets", + "final_loss", + "fraction_within_10pct", + "effective_sample_size", + "matrix_nnz", + ) + }, + "holdout_state_cd_weights": { + k: v + for k, v in diagnostics["cd_holdout"].items() + if k not in ("targets", "by_family") + }, + "holdout_state_weights": { + k: v for k, v in same_held.items() if k not in ("targets", "by_family") + }, + "weight_origin": diagnostics["weight_origin"], + } + print(json.dumps(summary, indent=1, default=str)[:6000]) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("stage", choices=["convert", "state", "state_cd", "all"]) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--holdout-fraction", type=float, default=0.1) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument( + "--state-weights", + type=Path, + default=None, + help="state-only calibrated weights to score the holdout under " + "(default /state/weights.npz)", + ) + args = parser.parse_args() + import torch + + torch.set_num_threads(args.threads) + args.out.mkdir(parents=True, exist_ok=True) + tool = load_tool() + stages = ["convert", "state", "state_cd"] if args.stage == "all" else [args.stage] + for stage in stages: + print(f"=== {stage}", flush=True) + if stage == "convert": + stage_convert(tool, args.out) + elif stage == "state": + stage_state(tool, args.out) + else: + stage_state_cd(tool, args.out, args.holdout_fraction, args.state_weights) + print(f"peak RSS {tool.rss():.2f} GB", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/microcosm-build/src/microcosm/build/__init__.py b/packages/microcosm-build/src/microcosm/build/__init__.py index 3a77cc96d..1a883d3db 100644 --- a/packages/microcosm-build/src/microcosm/build/__init__.py +++ b/packages/microcosm-build/src/microcosm/build/__init__.py @@ -124,6 +124,8 @@ def _assert_frame_compatible(version: str, required: tuple[int, int]) -> None: weights_audit_gate, ) from microcosm.build.holdout import ( # noqa: E402 - after the compat gate + hash_holdout_uniform, + hash_holdout_unit, rotated_folds, summarize_rotations, ) @@ -270,6 +272,8 @@ def _assert_frame_compatible(version: str, required: tuple[int, int]) -> None: "export_surface_gate", "exported_nonzero_gate", "formula_owned_export_gate", + "hash_holdout_uniform", + "hash_holdout_unit", "input_column_coverage_gate", "input_mass_parity_gate", "ledger_compile_parity_gate", diff --git a/packages/microcosm-build/src/microcosm/build/holdout.py b/packages/microcosm-build/src/microcosm/build/holdout.py index 0de9a4a38..48da44667 100644 --- a/packages/microcosm-build/src/microcosm/build/holdout.py +++ b/packages/microcosm-build/src/microcosm/build/holdout.py @@ -10,16 +10,33 @@ The fold structure is pure and deterministic (seeded permutation), so a scorer on any machine reproduces the same rotation from ``(n_targets, n_folds, seed)``. + +:func:`hash_holdout_unit` is the other shape: a sealed holdout whose +membership is a pure function of each unit's own key. A seeded permutation +reshuffles every fold when one target is added or the surface is reordered; +a keyed hash does not, so a holdout assigned this way stays the same set of +units across surface revisions, and raising the fraction only adds units. """ from __future__ import annotations +import hashlib +import math from collections.abc import Iterable, Sequence from dataclasses import dataclass import numpy as np -__all__ = ["rotated_folds", "summarize_rotations", "RotationSummary"] +__all__ = [ + "rotated_folds", + "summarize_rotations", + "RotationSummary", + "hash_holdout_uniform", + "hash_holdout_unit", +] + +#: Units per draw: the first 8 bytes of a SHA-256 digest, read big-endian. +_HASH_UNIFORM_SCALE = float(2**64) def rotated_folds( @@ -94,3 +111,45 @@ def summarize_rotations(fold_losses: Iterable[float]) -> RotationSummary: worst_holdout_loss=float(np.max(losses)), fold_losses=tuple(losses), ) + + +def hash_holdout_uniform(key: str, *, salt: str) -> float: + """A unit's deterministic uniform draw on ``[0, 1)``. + + ``sha256(salt + "\x1f" + key)``, first 8 bytes big-endian, divided by + ``2**64``. The unit separator keeps ``("ab", "c")`` and ``("a", "bc")`` + apart. The draw depends only on the unit's own key and the salt, never on + which other units exist or their order. + + Raises: + ValueError: If ``key`` or ``salt`` is not a non-empty string. + """ + if not isinstance(key, str) or not key: + raise ValueError(f"holdout key must be a non-empty string, got {key!r}.") + if not isinstance(salt, str) or not salt: + raise ValueError(f"holdout salt must be a non-empty string, got {salt!r}.") + digest = hashlib.sha256(f"{salt}\x1f{key}".encode()).digest() + return int.from_bytes(digest[:8], "big") / _HASH_UNIFORM_SCALE + + +def hash_holdout_unit(key: str, *, fraction: float, salt: str) -> bool: + """Whether the unit named ``key`` is held out at ``fraction``. + + A unit is held out iff :func:`hash_holdout_uniform` is below + ``fraction``. Membership is therefore deterministic, independent of the + other units and their order, and nested in the fraction: every unit held + out at ``f1`` is held out at any ``f2 >= f1``. ``fraction=0`` holds out + nothing and ``fraction=1`` everything. + + Raises: + ValueError: If ``fraction`` is not a finite number in ``[0, 1]``, or + from :func:`hash_holdout_uniform`. + """ + if ( + isinstance(fraction, bool) + or not isinstance(fraction, int | float) + or not math.isfinite(fraction) + or not 0.0 <= float(fraction) <= 1.0 + ): + raise ValueError(f"holdout fraction must be in [0, 1], got {fraction!r}.") + return hash_holdout_uniform(key, salt=salt) < float(fraction) diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_cd_properties.py b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_cd_properties.py new file mode 100644 index 000000000..243568e5e --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_cd_properties.py @@ -0,0 +1,852 @@ +"""Property tests for the ACS local sparse surface, ``state_cd`` and holdout. + +Invariants, each for every input Hypothesis draws: + +1. Sparse and dense compile agree on every target: the kernel compiles the + identical constraint row from a CSR row (callable measure) as from the + same values stored as a dense float32 household column. +2. District targets add up to their state target: after ``state_cd`` + reconciliation every district block sums to its bound state parent within + ``RECONCILIATION_RTOL``, each state concept is bound once, and a rebase + only rescales a block (district shares are kept). +3. Held-out targets never reach the calibrator, and the hash holdout is + deterministic, order-free and nested in its fraction. +4. ESS over distinct households never exceeds row ESS; a pro-rata block sums + to its state parent. +""" + +from __future__ import annotations + +import dataclasses +import importlib.util +import math + +import numpy as np +import pandas as pd +import pytest +from scipy import sparse + +from microcosm.build.holdout import hash_holdout_uniform, hash_holdout_unit +from microcosm.calibrate import TargetSpec +from microcosm.calibrate.matrix import build_constraint_matrix +from microcosm.calibrate.target import Target, TargetSet +from microcosm.frame import US_SCHEMA, Frame, WeightKind, Weights +from test_support.paths import paths_for + +hypothesis = pytest.importorskip("hypothesis") +from hypothesis import HealthCheck, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +_TEST_PATHS = paths_for("microcosm-build") +_SETTINGS = settings( + max_examples=60, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) + + +def _load_tool_module(): + path = _TEST_PATHS.repository / "tools" / "build_us_acs_local_release.py" + spec = importlib.util.spec_from_file_location( + "build_us_acs_local_release_properties", path + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +_TOOL = _load_tool_module() +_CD = _TOOL.cd_surface + + +def _toy_frame(weights: np.ndarray, columns: dict | None = None) -> Frame: + """One person per household, every group one per person.""" + + n = len(weights) + ids = np.arange(1, n + 1, dtype=np.int64) + person = pd.DataFrame( + { + "person_id": ids, + "person_household_id": ids, + **{f"person_{group}_id": ids for group in US_SCHEMA.group_entities}, + } + ) + tables = {"person": person} + for group in US_SCHEMA.group_entities: + tables[group] = pd.DataFrame({f"{group}_id": ids}) + tables["household"] = pd.DataFrame({"household_id": ids, **(columns or {})}) + return Frame( + tables, + US_SCHEMA, + {"household": Weights(np.asarray(weights, float), WeightKind.DESIGN)}, + ) + + +@st.composite +def _sparse_surfaces(draw): + n_targets = draw(st.integers(1, 10)) + n_households = draw(st.integers(1, 30)) + density = draw(st.floats(0.0, 1.0)) + seed = draw(st.integers(0, 2**31 - 1)) + rng = np.random.default_rng(seed) + dense = rng.normal(0.0, 1e6, size=(n_targets, n_households)).astype(np.float32) + dense[rng.random(dense.shape) >= density] = 0.0 + # float32 values that are tiny, huge or negative all round-trip. + if dense.size: + dense.flat[0] = np.float32(draw(st.sampled_from([0.0, 1e-30, -3.5, 3e30]))) + weights = rng.uniform(0.5, 50.0, size=n_households) + return dense, weights + + +@_SETTINGS +@given(_sparse_surfaces()) +def test_sparse_and_dense_compile_agree_on_every_target(surface) -> None: + dense, weights = surface + n_targets, n_households = dense.shape + records = [ + { + "name": f"t{i}", + "value": float(i + 1), + "period": 2024, + "source": "property", + "role": "train", + } + for i in range(n_targets) + ] + sparse_problem = build_constraint_matrix( + _toy_frame(weights), + _CD.calibration_target_set(records, sparse.csr_array(dense), n_households), + ) + dense_frame = _toy_frame(weights, {f"m{i}": dense[i] for i in range(n_targets)}) + dense_problem = build_constraint_matrix( + dense_frame, + TargetSet( + [ + Target( + name=f"t{i}", + entity="household", + measure=f"m{i}", + value=float(i + 1), + period=2024, + source="property", + ) + for i in range(n_targets) + ] + ), + ) + for row in range(n_targets): + sparse_row = sparse_problem.matrix[[row]] + dense_row = dense_problem.matrix[[row]] + np.testing.assert_array_equal(sparse_row.indices, dense_row.indices) + np.testing.assert_array_equal(sparse_row.data, dense_row.data) + np.testing.assert_array_equal( + sparse_problem.estimates(weights), dense_problem.estimates(weights) + ) + + +_HT2 = "irs_soi.historic_table_2.state_broad_totals.v1" +_CD_FILE = "irs_soi.congressional_district_2022.all_returns.v1" +_MEASURES = { + "adjusted_gross_income": ("adjusted_gross_income", "sum"), + "return_count": ("count", "sum"), + "wages_salaries_amount": ("employment_income", "sum"), + "charitable_amount": ("charitable_deduction", "sum"), +} + + +def _soi_spec(name, measure, value, *, level, record_set, state, district=None): + variable, mode = _MEASURES[measure] + metadata = { + "ledger_geography_level": level, + "ledger_layout_record_set_spec_id": record_set, + "state_fips": state, + "source_measure_id": measure, + "variable": variable, + "source_variable": variable, + "measure_mode": mode, + "agi_lower_bound": "-inf", + "agi_upper_bound": "inf", + "filing_status": "All", + } + if district is not None: + metadata["congressional_district_geoid"] = district + return TargetSpec( + name=name, + entity="household", + measure=name, + value=value, + source="property", + family="irs_soi", + signed=value < 0, + metadata=metadata, + ) + + +@st.composite +def _state_cd_inputs(draw): + states = [f"{10 + index:02d}" for index in range(draw(st.integers(1, 4)))] + # return_count is always present in both files: it is the level bridge + # of every concept Historic Table 2 lacks here (_BRIDGES). + measures = ["return_count"] + draw( + st.lists( + st.sampled_from(sorted(set(_MEASURES) - {"return_count"})), + min_size=0, + max_size=3, + unique=True, + ) + ) + positive = st.floats(1.0, 1e10, allow_nan=False, allow_infinity=False) + specs, crosswalk, layout = [], [], {} + for state in states: + n_districts = draw(st.integers(1, 5)) + at_large = n_districts <= 2 and draw(st.booleans()) + districts = [f"{state}{index + 1:02d}" for index in range(n_districts)] + source = [f"{state}00"] if at_large else districts + for target in districts: + for origin in source if at_large else [target]: + crosswalk.append( + { + "source_geography_id": f"5001700US{origin}", + "target_geography_id": f"5001900US{target}", + } + ) + # A state with one source-plan district is at-large there: the SOI + # file has no sub-state rows for it, whatever the current plan says. + layout[state] = (districts, len(source) == 1) + for measure in measures: + if measure == "return_count" or draw(st.booleans()): + specs.append( + _soi_spec( + f"ht2.{state}.{measure}", + measure, + draw(positive), + level="state", + record_set=_HT2, + state=state, + ) + ) + district_values = [draw(positive) for _district in districts] + # The district file's own state total is the sum of its districts + # (it is on the pinned feed); an at-large state has no sub-state + # rows, so its total is free. + specs.append( + _soi_spec( + f"cdfile.{state}_total.{measure}", + measure, + draw(positive) if len(source) == 1 else math.fsum(district_values), + level="state", + record_set=_CD_FILE, + state=state, + ) + ) + for district, value in zip(districts, district_values, strict=True): + specs.append( + _soi_spec( + f"cdfile.{district}.{measure}", + measure, + value, + level="congressional_district", + record_set=_CD_FILE, + state=state, + district=district, + ) + ) + order = draw(st.permutations(range(len(specs)))) + return [specs[i] for i in order], pd.DataFrame(crosswalk), layout + + +_BRIDGES = {measure: "return_count" for measure in _MEASURES} + + +@_SETTINGS +@given(_state_cd_inputs()) +def test_district_targets_add_up_to_one_state_vintage(inputs) -> None: + specs, crosswalk, layout = inputs + surface = _CD.state_cd_soi_surface( + specs, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges=_BRIDGES, + # Random values make random vintage ratios; the band has its own test. + factor_band=math.inf, + ) + kept = {spec.name: spec for spec in surface.specs} + original = {spec.name: spec for spec in specs} + + # Every district block adds up to its bound state parent. + report = _CD.state_parent_reconciliation(surface.specs) + assert all(block["ok"] and block["parent_bound"] for block in report) + + # One vintage per state concept. + state_keys = [ + (spec.metadata["state_fips"], _CD.soi_concept_identity(spec)) + for spec in surface.specs + if spec.metadata["ledger_geography_level"] == "state" + ] + assert len(state_keys) == len(set(state_keys)) + + districts = [ + spec + for spec in surface.specs + if spec.metadata["ledger_geography_level"] == "congressional_district" + ] + by_parent: dict[str, list] = {} + for spec in districts: + by_parent.setdefault(spec.metadata["state_cd_parent_target_name"], []).append( + spec + ) + state = spec.metadata["state_fips"] + measure = spec.metadata["source_measure_id"] + ht2 = f"ht2.{state}.{measure}" + expected_parent = ht2 if ht2 in original else f"cdfile.{state}_total.{measure}" + assert spec.metadata["state_cd_parent_target_name"] == expected_parent + assert expected_parent in kept + if expected_parent != ht2: + # A district-file-only state level is lifted onto the Historic + # Table 2 basis by its bridge sibling's two levels in that state. + bridge = ( + original[f"ht2.{state}.return_count"].value + / original[f"cdfile.{state}_total.return_count"].value + ) + assert kept[expected_parent].value == pytest.approx( + original[expected_parent].value * bridge, rel=1e-12 + ) + # A rebase rescales a block and keeps each district's share of it. + for children in by_parent.values(): + ratios = [spec.value / original[spec.name].value for spec in children] + np.testing.assert_allclose(ratios, ratios[0], rtol=1e-12) + + # States with one source-plan district bind no district rows (they would + # duplicate the state row); the others bind every district of every + # concept. + for state, (district_ids, at_large) in layout.items(): + bound = { + spec.metadata["congressional_district_geoid"] + for spec in districts + if spec.metadata["state_fips"] == state + } + assert bound == (set() if at_large else set(district_ids)) + # A district-file state total survives only as a single-vintage parent. + for spec in surface.specs: + if spec.metadata["ledger_layout_record_set_spec_id"] == _CD_FILE and ( + spec.metadata["ledger_geography_level"] == "state" + ): + state = spec.metadata["state_fips"] + measure = spec.metadata["source_measure_id"] + assert f"ht2.{state}.{measure}" not in original + + +def test_an_incomplete_or_unabsorbable_district_block_is_refused() -> None: + crosswalk = pd.DataFrame( + { + "source_geography_id": ["5001700US1001", "5001700US1002"], + "target_geography_id": ["5001900US1001", "5001900US1002"], + } + ) + parent = _soi_spec( + "ht2.10.adjusted_gross_income", + "adjusted_gross_income", + 100.0, + level="state", + record_set=_HT2, + state="10", + ) + + def district(index, value): + return _soi_spec( + f"cdfile.100{index}.adjusted_gross_income", + "adjusted_gross_income", + value, + level="congressional_district", + record_set=_CD_FILE, + state="10", + district=f"100{index}", + ) + + predicate = _TOOL.soi_surface_predicate("state") + with pytest.raises(ValueError, match="expected the current plan"): + _CD.state_cd_soi_surface( + [parent, district(1, 5.0)], + state_surface_predicate=predicate, + crosswalk=crosswalk, + ) + with pytest.raises(ValueError, match="sum to zero"): + _CD.state_cd_soi_surface( + [parent, district(1, 0.0), district(2, 0.0)], + state_surface_predicate=predicate, + crosswalk=crosswalk, + ) + with pytest.raises(ValueError, match="opposite in sign"): + _CD.state_cd_soi_surface( + [parent, district(1, -5.0), district(2, -1.0)], + state_surface_predicate=predicate, + crosswalk=crosswalk, + ) + # The builder re-checks its own output: a block that failed to add up + # (here reported so by a stubbed reconciliation) is refused, not bound. + surface = _CD.state_cd_soi_surface( + [parent, district(1, 5.0), district(2, 15.0)], + state_surface_predicate=predicate, + crosswalk=crosswalk, + ) + assert [ + block["ok"] for block in _CD.state_parent_reconciliation(surface.specs) + ] == [True] + real = _CD.state_parent_reconciliation + try: + _CD.state_parent_reconciliation = lambda specs: [ + {**block, "ok": False} for block in real(specs) + ] + with pytest.raises(ValueError, match="do not add up"): + _CD.state_cd_soi_surface( + [parent, district(1, 5.0), district(2, 15.0)], + state_surface_predicate=predicate, + crosswalk=crosswalk, + ) + finally: + _CD.state_parent_reconciliation = real + + +@st.composite +def _records(draw): + n = draw(st.integers(1, 60)) + records = [] + for index in range(n): + eligible = draw(st.booleans()) + record = { + "name": f"target_{index}", + "value": float(index + 1), + "period": 2024, + "source": "property", + "family": "irs_soi" + if eligible + else draw( + st.sampled_from(["usda_snap", "census_population_ladder", "irs_soi"]) + ), + } + if eligible: + record.update( + { + "geography_level": "congressional_district", + "state_fips": draw(st.sampled_from(["06", "36", "48"])), + "source_measure_id": draw( + st.sampled_from( + [ + "adjusted_gross_income", + "eitc_one_child_amount", + "eitc_claims", + "return_count", + ] + ) + ), + "state_cd_parent_target_name": "parent", + } + ) + records.append(record) + return records + + +@_SETTINGS +@given(_records(), st.floats(0.0, 0.5), st.integers(0, 2**31 - 1)) +def test_held_out_targets_never_reach_the_calibrator(records, fraction, seed) -> None: + _CD.assign_target_roles(records, fraction=fraction) + rng = np.random.default_rng(seed) + n_households = 3 + matrix = sparse.csr_array( + rng.integers(0, 3, size=(len(records), n_households)).astype(np.float32) + ) + target_set = _CD.calibration_target_set(records, matrix, n_households) + seen = {target.name for target in target_set} + held = {record["name"] for record in records if record["role"] == "holdout"} + trained = {record["name"] for record in records if record["role"] == "train"} + assert seen == trained + assert not seen & held + # Only district SOI targets with a state parent are ever held out, and a + # held unit is held whole: every target of it shares one role. + roles_by_unit: dict[str, set[str]] = {} + for record in records: + if record["role"] == "holdout": + assert record["holdout_unit"] is not None + if record["holdout_unit"] is not None: + roles_by_unit.setdefault(record["holdout_unit"], set()).add(record["role"]) + assert all(len(roles) == 1 for roles in roles_by_unit.values()) + # EITC measures share one family, so per-child rows cannot be pinned by + # the trained EITC total of the same state. + for record in records: + if record["holdout_unit"] is not None and str( + record["source_measure_id"] + ).startswith("eitc"): + assert record["holdout_unit"].endswith("|eitc") + + +@_SETTINGS +@given( + st.lists(st.text(min_size=1, max_size=12), min_size=1, max_size=40, unique=True), + st.floats(0.0, 1.0), + st.floats(0.0, 1.0), + st.randoms(), +) +def test_hash_holdout_is_deterministic_order_free_and_nested( + keys, first, second, random +) -> None: + low, high = sorted((first, second)) + salt = "property-salt" + held_low = {key for key in keys if hash_holdout_unit(key, fraction=low, salt=salt)} + held_high = { + key for key in keys if hash_holdout_unit(key, fraction=high, salt=salt) + } + assert held_low <= held_high + shuffled = list(keys) + random.shuffle(shuffled) + assert { + key for key in shuffled if hash_holdout_unit(key, fraction=low, salt=salt) + } == held_low + for key in keys: + uniform = hash_holdout_uniform(key, salt=salt) + assert 0.0 <= uniform < 1.0 + assert uniform == hash_holdout_uniform(key, salt=salt) + assert not {k for k in keys if hash_holdout_unit(k, fraction=0.0, salt=salt)} + assert {k for k in keys if hash_holdout_unit(k, fraction=1.0, salt=salt)} == set( + keys + ) + + +def test_hash_holdout_rate_matches_the_fraction() -> None: + keys = [f"unit-{index}" for index in range(20_000)] + for fraction in (0.05, 0.1, 0.25): + held = sum( + hash_holdout_unit(key, fraction=fraction, salt="rate") for key in keys + ) + # Binomial sd at 20,000 draws is at most 0.0035; 0.02 is > 5 sd. + assert abs(held / len(keys) - fraction) < 0.02 + + +def test_hash_holdout_rejects_bad_inputs() -> None: + with pytest.raises(ValueError): + hash_holdout_unit("key", fraction=1.5, salt="salt") + with pytest.raises(ValueError): + hash_holdout_unit("key", fraction=float("nan"), salt="salt") + with pytest.raises(ValueError): + hash_holdout_unit("", fraction=0.1, salt="salt") + with pytest.raises(ValueError): + hash_holdout_unit("key", fraction=0.1, salt="") + + +@_SETTINGS +@given( + st.lists( + st.tuples( + st.sampled_from(["acs_2024_1yr", "asec_puf"]), + st.integers(1, 6), + st.floats(0.01, 1e4, allow_nan=False), + ), + min_size=1, + max_size=40, + ) +) +def test_distinct_household_ess_never_exceeds_row_ess(rows) -> None: + spine = np.asarray([row[0] for row in rows]) + source = np.asarray([row[1] for row in rows]) + weights = np.asarray([row[2] for row in rows]) + summary = _CD.weight_origin_summary(weights, spine=spine, source_id=source) + distinct = summary["effective_sample_size_distinct_households"] + assert distinct <= summary["effective_sample_size_rows"] * (1 + 1e-12) + assert distinct <= summary["distinct_households"] * (1 + 1e-12) + assert sum(summary["weight_share_by_spine"].values()) == pytest.approx(1.0) + keys = set(zip(spine, source, strict=True)) + if len(keys) == len(rows): + assert distinct == pytest.approx(summary["effective_sample_size_rows"]) + + +@_SETTINGS +@given( + st.floats(-1e9, 1e9, allow_nan=False).filter(lambda value: value != 0), + st.lists(st.floats(1.0, 1e6), min_size=1, max_size=12), +) +def test_a_pro_rata_block_sums_to_its_state_parent(parent_value, populations) -> None: + state_population = float(sum(populations)) + records = [{"name": "parent", "value": parent_value}] + for index, population in enumerate(populations): + records.append( + { + "name": f"cd_{index}", + "value": 0.0, + "state_cd_parent_target_name": "parent", + "cd_population": population, + "state_population": state_population, + } + ) + baseline = _CD.pro_rata_baseline(records, range(1, len(records))) + assert baseline.sum() == pytest.approx(parent_value, rel=1e-9, abs=1e-6) + + +@st.composite +def _banded_inputs(draw): + """Two-district states whose vintage ratio is the median times a jitter.""" + + n_states = draw(st.integers(3, 8)) + jitters = draw( + st.lists( + st.floats(0.3, 3.0, allow_nan=False), + min_size=n_states, + max_size=n_states, + ) + ) + median = draw(st.floats(0.5, 3.0, allow_nan=False)) + specs, crosswalk = [], [] + for index, jitter in enumerate(jitters): + state = f"{10 + index:02d}" + districts = [f"{state}01", f"{state}02"] + for district in districts: + crosswalk.append( + { + "source_geography_id": f"5001700US{district}", + "target_geography_id": f"5001900US{district}", + } + ) + cd_total = 100.0 + specs.append( + _soi_spec( + f"ht2.{state}.adjusted_gross_income", + "adjusted_gross_income", + cd_total * median * jitter, + level="state", + record_set=_HT2, + state=state, + ) + ) + specs.append( + _soi_spec( + f"cdfile.{state}_total.adjusted_gross_income", + "adjusted_gross_income", + cd_total, + level="state", + record_set=_CD_FILE, + state=state, + ) + ) + for district, share in zip(districts, (0.4, 0.6), strict=True): + specs.append( + _soi_spec( + f"cdfile.{district}.adjusted_gross_income", + "adjusted_gross_income", + cd_total * share, + level="congressional_district", + record_set=_CD_FILE, + state=state, + district=district, + ) + ) + return specs, pd.DataFrame(crosswalk), jitters, median + + +@_SETTINGS +@given(_banded_inputs()) +def test_blocks_off_their_concepts_median_ratio_are_dropped_and_recorded( + inputs, +) -> None: + """A block is kept iff its factor is within the band of the median factor.""" + + specs, crosswalk, jitters, median = inputs + band = _CD.STATE_CD_FACTOR_BAND + surface = _CD.state_cd_soi_surface( + specs, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges=_BRIDGES, + ) + factors = np.asarray([median * jitter for jitter in jitters]) + typical = float(np.median(factors)) + bound_states = { + spec.metadata["state_fips"] + for spec in surface.specs + if spec.metadata["ledger_geography_level"] == "congressional_district" + } + recorded = { + row["state_fips"] + for row in surface.receipt["factor_band"]["out_of_band"] + if row["basis"] == "rebase" + } + for index, factor in enumerate(factors): + state = f"{10 + index:02d}" + relative = factor / typical + if 1 / band * (1 + 1e-9) < relative < band * (1 - 1e-9): + assert state in bound_states and state not in recorded + elif not 1 / band * (1 - 1e-9) <= relative <= band * (1 + 1e-9): + assert state not in bound_states and state in recorded + # The Historic Table 2 state rows stay whatever happens to their districts. + assert { + spec.metadata["state_fips"] + for spec in surface.specs + if spec.metadata["ledger_geography_level"] == "state" + } == {f"{10 + index:02d}" for index in range(len(jitters))} + + +def _bridge_world(ratios: dict[str, float], at_large: dict[str, float]): + """States whose Historic Table 2 return count is ``ratio`` x the file's. + + Every state carries ``return_count`` in both files and ``charitable_amount`` + only in the district file (bridged by ``return_count``). States in + ``ratios`` have two districts; states in ``at_large`` one source district. + """ + + specs, crosswalk = [], [] + for state, ratio in {**ratios, **at_large}.items(): + split = state in ratios + districts = [f"{state}01", f"{state}02"] if split else [f"{state}01"] + for district in districts: + crosswalk.append( + { + "source_geography_id": f"5001700US{district if split else state + '00'}", + "target_geography_id": f"5001900US{district}", + } + ) + for measure, total, ht2 in ( + ("return_count", 100.0, 100.0 * ratio), + ("charitable_amount", 10.0, None), + ): + if ht2 is not None: + specs.append( + _soi_spec( + f"ht2.{state}.{measure}", + measure, + ht2, + level="state", + record_set=_HT2, + state=state, + ) + ) + specs.append( + _soi_spec( + f"cdfile.{state}_total.{measure}", + measure, + total, + level="state", + record_set=_CD_FILE, + state=state, + ) + ) + if split: + for district, share in zip(districts, (0.4, 0.6), strict=True): + specs.append( + _soi_spec( + f"cdfile.{district}.{measure}", + measure, + total * share, + level="congressional_district", + record_set=_CD_FILE, + state=state, + district=district, + ) + ) + return specs, pd.DataFrame(crosswalk) + + +def _bridge_surface(ratios, at_large=None): + specs, crosswalk = _bridge_world(ratios, at_large or {}) + return _CD.state_cd_soi_surface( + specs, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges={"charitable_amount": "return_count"}, + ) + + +def _bound(surface, state, measure, level): + return [ + spec.name + for spec in surface.specs + if spec.metadata["state_fips"] == state + and spec.metadata["source_measure_id"] == measure + and spec.metadata["ledger_geography_level"] == level + ] + + +def test_a_bridge_through_an_out_of_band_sibling_drops_its_state_and_districts() -> ( + None +): + """Either band verdict on the sibling distrusts the bridge in that state.""" + + # 1. State 14's two levels disagree 2x: both verdicts fail there. + surface = _bridge_surface({"10": 1.0, "11": 1.0, "12": 1.0, "13": 1.0, "14": 2.0}) + verdicts = { + (row["state_fips"], row["measure"], row["verdict"]) + for row in surface.receipt["factor_band"]["out_of_band"] + } + assert verdicts == { + ("14", "return_count", "rebase_factor"), + ("14", "charitable_amount", "sibling_vintage_ratio"), + ("14", "charitable_amount", "sibling_rebase_factor"), + } + assert not _bound(surface, "14", "charitable_amount", "state") + assert not _bound(surface, "14", "charitable_amount", "congressional_district") + assert not _bound(surface, "14", "return_count", "congressional_district") + assert _bound(surface, "14", "return_count", "state") # Historic Table 2 stays + for state in ("10", "11", "12", "13"): + assert ( + len(_bound(surface, state, "charitable_amount", "congressional_district")) + == 2 + ) + assert surface.receipt["dropped"]["level_bridge_out_of_band:charitable_amount"] == 3 + + # 2. The two verdicts can split: at-large states pull the vintage median + # to 1.2, so state 14's ratio of 1.3 passes it, but among the states with + # district blocks the median is 1.0 and its block fails. The bridge is + # distrusted by the sibling's rebase verdict alone. + surface = _bridge_surface( + {"10": 1.0, "11": 1.0, "12": 1.0, "13": 1.0, "14": 1.3}, + {f"{20 + index}": 1.2 for index in range(8)}, + ) + verdicts = { + (row["state_fips"], row["measure"], row["verdict"]) + for row in surface.receipt["factor_band"]["out_of_band"] + } + assert verdicts == { + ("14", "return_count", "rebase_factor"), + ("14", "charitable_amount", "sibling_rebase_factor"), + } + assert not _bound(surface, "14", "charitable_amount", "state") + assert not _bound(surface, "14", "charitable_amount", "congressional_district") + # At-large states keep their bridged state rows; they have no districts. + assert _bound(surface, "20", "charitable_amount", "state") + + +def test_a_sign_conflict_or_a_self_contradicting_district_file_is_refused() -> None: + """Defects inside the data are refused, not dropped as level gaps.""" + + specs, crosswalk = _bridge_world({"10": 1.0, "11": 1.0, "12": -1.0}, {}) + with pytest.raises(ValueError, match="has factor"): + _CD.state_cd_soi_surface( + specs, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges={"charitable_amount": "return_count"}, + ) + specs, crosswalk = _bridge_world({"10": 1.0, "11": 1.0, "12": 1.0}, {}) + broken = [ + dataclasses.replace(spec, value=spec.value * (1 + 1e-4)) + if spec.name == "cdfile.1101.return_count" + else spec + for spec in specs + ] + with pytest.raises(ValueError, match="differs from the sum of its district"): + _CD.state_cd_soi_surface( + broken, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges={"charitable_amount": "return_count"}, + ) + # Inside the tolerance (float noise) it is accepted. + nudged = [ + dataclasses.replace(spec, value=spec.value * (1 + 1e-9)) + if spec.name == "cdfile.1101.return_count" + else spec + for spec in specs + ] + _CD.state_cd_soi_surface( + nudged, + state_surface_predicate=_TOOL.soi_surface_predicate("state"), + crosswalk=crosswalk, + level_bridges={"charitable_amount": "return_count"}, + ) diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_release_tool.py b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_release_tool.py index be3d3a87d..a8c4eea70 100644 --- a/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_release_tool.py +++ b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_release_tool.py @@ -19,6 +19,9 @@ _TEST_PATHS = paths_for("microcosm-build") +#: The sampling block a full-scale materialize records in run_identity.json. +_FULL_RUNG = {"sample_fraction": 1.0, "rung": "f100", "sampled": False} + # Tests that write real H5 bytes go through pandas' HDFStore, which needs # pytables; the base wheel gate installs the shards without it. requires_pytables = pytest.mark.skipif( @@ -407,6 +410,7 @@ def _package_evidence_args(module, tmp_path: Path, monkeypatch, *, hours_report) "run_identity.json": { "staging_sha256": module._sha256(staging), "population_cells_dropped": [], + "sampling": _FULL_RUNG, }, "spine_qa.json": { "plain_consumption": True, @@ -587,6 +591,7 @@ def _package_args_before_evidence( { "staging_sha256": module._sha256(staging), "population_cells_dropped": [], + "sampling": _FULL_RUNG, } ) ) @@ -676,6 +681,7 @@ def test_do_package_requires_qa_and_consumer_evidence(tmp_path: Path) -> None: { "staging_sha256": module._sha256(staging), "population_cells_dropped": [], + "sampling": _FULL_RUNG, } ) ) @@ -739,10 +745,10 @@ def test_soi_mode_defaults_to_state_and_totals_and_full_are_explicit_opt_ins( asking for them.""" module = _load_tool_module() - assert module.SOI_MODES == ("state", "totals", "full") + assert module.SOI_MODES == ("state", "totals", "full", "state_cd") assert module.DEFAULT_SOI_MODE == module.SOI_MODE_STATE == "state" assert module._parse_args(_materialize_argv(tmp_path)).soi_mode == "state" - for mode in ("totals", "full"): + for mode in ("totals", "full", "state_cd"): assert ( module._parse_args(_materialize_argv(tmp_path, "--soi-mode", mode)).soi_mode == mode @@ -1214,6 +1220,7 @@ def test_finalize_report_round_trips_into_package(tmp_path, monkeypatch) -> None "staging_sha256": module._sha256(args.staging_h5), "ladder_sha256": module._sha256(args.ladder), "population_cells_dropped": [], + "sampling": _FULL_RUNG, }, "spine_qa.json": { "plain_consumption": True, @@ -1246,6 +1253,96 @@ def test_finalize_report_round_trips_into_package(tmp_path, monkeypatch) -> None ) +def test_a_sparse_era_calibration_is_bound_through_finalize_and_package( + tmp_path, monkeypatch +) -> None: + """The run-identity, weights and QA bindings at their real call sites. + + A sparse-era checkpoint (its identity records the roles digest) round-trips + when every piece of evidence names the same materialization, weights and + bytes; an interrupted recalibration (the export's weights differ from the + summary's) or QA of other bytes is refused before any gate report or + release directory is written. + """ + + module = _load_tool_module() + args = _finalize_args(module, tmp_path) + _write_frame_h5(args.out_h5, _plausible_hours_frame()) + artifact_sha = module._sha256(args.out_h5) + ckpt = args.checkpoint_dir + for name, content in ( + ("target_registry.json", "{}"), + (module.TARGET_ROLES_FILENAME, "[]"), + (module.TARGET_MATRIX_FILENAME, "matrix"), + ): + (ckpt / name).write_text(content) + identity = { + "staging_sha256": module._sha256(args.staging_h5), + "ladder_sha256": module._sha256(args.ladder), + "population_cells_dropped": [], + "sampling": _FULL_RUNG, + "target_registry_sha256": module._sha256(ckpt / "target_registry.json"), + "target_roles_sha256": module._sha256(ckpt / module.TARGET_ROLES_FILENAME), + "target_matrix": { + "sha256": module._sha256(ckpt / module.TARGET_MATRIX_FILENAME) + }, + } + stamp = module._run_identity_digest(identity) + summary = { + **json.loads((ckpt / "calibration_summary.json").read_text()), + "run_identity_sha256": stamp, + "weights_sha256": "w", + } + export = { + "staging_sha256": artifact_sha, + "run_identity_sha256": stamp, + "out_h5_sha256": artifact_sha, + "weights_sha256": "w", + } + qa = { + "plain_consumption": True, + "artifact_sha256": artifact_sha, + "per_spine": {}, + "run_identity_sha256": stamp, + } + evidence = { + "run_identity.json": identity, + "calibration_summary.json": summary, + "consumer_export.json": export, + "spine_qa.json": qa, + "consumer_reviewed_null_fills.json": {"columns_filled": []}, + } + for name, value in evidence.items(): + (ckpt / name).write_text(json.dumps(value)) + _patch_finalize_collaborators(module, monkeypatch, identity=False) + + # An interrupted recalibration: a new H5 and export, the old summary. + (ckpt / "consumer_export.json").write_text( + json.dumps({**export, "weights_sha256": "new"}) + ) + with pytest.raises(SystemExit, match="describe different weights"): + module.do_finalize(args) + assert not args.gate_report.exists() + (ckpt / "consumer_export.json").write_text(json.dumps(export)) + + module.do_finalize(args) + assert json.loads(args.out_summary.read_text())["simulation_ready"] is True + args.out = tmp_path / "release" + args.allow_dirty = True + + # QA of the wrong materialization: package refuses before any release dir. + (ckpt / "spine_qa.json").write_text( + json.dumps({**qa, "run_identity_sha256": "old"}) + ) + with pytest.raises(SystemExit, match="another materialization"): + module.do_package(args) + assert not (args.out / "releases").exists() + (ckpt / "spine_qa.json").write_text(json.dumps(qa)) + + result = module.do_package(args) + assert result["root_artifact"]["sha256"] == artifact_sha + + def _package_args_with_hours(module, tmp_path, monkeypatch, *, gate_state): """Real tiny H5 bytes; gate edits model stale separately resumed finalize.""" @@ -1281,6 +1378,7 @@ def _package_args_with_hours(module, tmp_path, monkeypatch, *, gate_state): "run_identity.json": { "staging_sha256": module._sha256(args.staging_h5), "population_cells_dropped": [], + "sampling": _FULL_RUNG, }, "spine_qa.json": { "plain_consumption": True, diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_sparse_cd.py b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_sparse_cd.py new file mode 100644 index 000000000..e8a647496 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/us/test_us_acs_local_sparse_cd.py @@ -0,0 +1,1280 @@ +"""The ACS local release's sparse target matrix and ``state_cd`` surface. + +The central test is a differential: the fake-engine fixture is materialized +once the old way (every target a dense float32 household column, compiled by +the calibrate kernel from frame columns) and once through the sparse path +(chunked carrier split -> CSR checkpoint -> callable rows). Both must give +the kernel the identical constraint matrix and the identical calibrated +weights. The dense construction lives only here. +""" + +from __future__ import annotations + +import dataclasses +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from scipy import sparse + +from microcosm.calibrate import TargetSpec +from microcosm.calibrate.matrix import build_constraint_matrix +from microcosm.calibrate.target import TargetSet +from microcosm.frame import US_SCHEMA, Frame +from test_support.paths import paths_for + +_TEST_PATHS = paths_for("microcosm-build") +requires_pytables = pytest.mark.skipif( + importlib.util.find_spec("tables") is None, + reason="requires pytables (the build environment)", +) + + +def _load_tool_module(): + path = _TEST_PATHS.repository / "tools" / "build_us_acs_local_release.py" + spec = importlib.util.spec_from_file_location("build_us_acs_local_release", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _load_fixtures(): + spec = importlib.util.spec_from_file_location( + "batched_materialization_fixtures", + Path(__file__).with_name("test_us_batched_target_materialization.py"), + ) + fixtures = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(fixtures) + return fixtures + + +def _district_hierarchy(name: str, district: str): + """A production-shaped hierarchy: its target id is the spec name.""" + + from microcosm.calibrate import ( + CalibrationHierarchy, + HierarchyCategory, + HierarchyGeography, + HierarchyNode, + ) + + return CalibrationHierarchy( + provider=HierarchyNode(id="irs_soi", label="IRS SOI"), + category=HierarchyCategory( + id="irs_soi.fixture", label="Fixture", provider_id="irs_soi" + ), + geography=HierarchyGeography( + id=f"5001900US{district}", + label=f"District {district}", + level="congressional_district", + ), + dimensions=(), + target=HierarchyNode(id=name, label=name), + ) + + +def _cd_soi(fixtures, name, variable, state, district, **metadata): + return dataclasses.replace( + fixtures._soi( + name, + variable, + state_fips=state, + congressional_district_geoid=district, + **metadata, + ), + hierarchy=_district_hierarchy(name, district), + ) + + +def _cd_surface_specs(fixtures) -> tuple: + """The fixture's targets plus district SOI rows of several concepts. + + The fixture households sit in districts 0601 and 0602 (CA), 3601 and + 3602 (NY) and 2401 (MD). The concepts cover an amount, a count, an AGI + band, a filing-status slice, an itemizer slice and the positive-EITC + domain, so every mask the SOI slice builds is exercised. + """ + + rows = [] + for state, district in ( + ("06", "0601"), + ("06", "0602"), + ("36", "3601"), + ("36", "3602"), + ("24", "2401"), + ): + rows += [ + _cd_soi( + fixtures, f"cd_{district}_agi", "adjusted_gross_income", state, district + ), + _cd_soi(fixtures, f"cd_{district}_returns", "count", state, district), + _cd_soi( + fixtures, + f"cd_{district}_agi_band_count", + "count", + state, + district, + agi_lower_bound="-100", + agi_upper_bound="900", + ), + _cd_soi( + fixtures, + f"cd_{district}_joint_wages", + "employment_income", + state, + district, + filing_status="Married Filing Jointly/Surviving Spouse", + ), + _cd_soi( + fixtures, + f"cd_{district}_itemized", + "itemized_taxable_income_deductions", + state, + district, + itemized_only="true", + ), + _cd_soi( + fixtures, + f"cd_{district}_eitc_returns", + "count", + state, + district, + ledger_domain=( + "individual_income_tax_returns_with_earned_income_credit" + ), + ), + ] + # A district row of a state no household lives in: an empty CSR row. + rows.append(_cd_soi(fixtures, "cd_4801_agi", "adjusted_gross_income", "48", "4801")) + return tuple(fixtures._TARGETS) + tuple(rows) + + +def _identity(digests: dict, *, households: int) -> dict: + return { + "households": households, + "target_registry_sha256": digests["target_registry_sha256"], + "target_roles_sha256": digests["target_roles_sha256"], + "target_matrix": {"sha256": digests["target_matrix_sha256"]}, + } + + +def _materialize(module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk): + fixtures._install_fake_engine(release, monkeypatch, reform_specs=fixtures._REFORMS) + monkeypatch.setattr(module, "project_input_only", lambda frame, **kw: (frame, {})) + monkeypatch.setattr(module, "fill_reviewed_nulls", lambda *args, **kw: None) + return module.materialize_chunked( + fixtures._nested_frame(), + specs, + hh_chunk=hh_chunk, + batch=2, + summary_path=tmp_path / "summary.json", + ) + + +def _dense_reference(fixtures, release, specs): + """The pre-sparse construction: one float32 household column per target.""" + + frame = fixtures._nested_frame() + target_frame, registry, _ = release._materialize_target_frame( + frame, tuple(specs), maximum_microsim_batch_size=2 + ) + households = target_frame.table("household") + columns = { + f"m{index:04d}": households[spec.measure].to_numpy(dtype=np.float32) + for index, spec in enumerate(registry.specs) + } + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + tables["household"] = pd.concat( + [ + tables["household"][["household_id"]].reset_index(drop=True), + pd.DataFrame(columns), + ], + axis=1, + ) + dense_frame = Frame( + tables, + US_SCHEMA, + {"household": frame.weights_for("household")}, + ) + # Exactly what TargetSpec.to_target() compiled on the dense checkpoint, + # reading each target's float32 column. + targets = TargetSet( + [ + dataclasses.replace(spec.to_target(), measure=f"m{index:04d}") + for index, spec in enumerate(registry.specs) + ] + ) + return dense_frame, targets, list(registry.specs) + + +def _sparse_records(module, materialized, fraction=0.0): + records = [ + { + "entity": "household", + "measure": spec.name, + "period": 2024, + **module.cd_surface.target_record(spec), + } + for spec in materialized.compiled_specs + ] + module.cd_surface.assign_target_roles(records, fraction=fraction) + return records + + +def _assert_same_problem(dense_problem, sparse_problem) -> None: + assert dense_problem.names == sparse_problem.names + assert dense_problem.matrix.shape == sparse_problem.matrix.shape + np.testing.assert_array_equal( + dense_problem.target_vector, sparse_problem.target_vector + ) + dense_matrix = sparse.csr_array(dense_problem.matrix) + sparse_matrix = sparse.csr_array(sparse_problem.matrix) + np.testing.assert_array_equal(dense_matrix.indptr, sparse_matrix.indptr) + np.testing.assert_array_equal(dense_matrix.indices, sparse_matrix.indices) + np.testing.assert_array_equal(dense_matrix.data, sparse_matrix.data) + + +@pytest.mark.parametrize("hh_chunk", [1, 2, 5]) +def test_sparse_materialization_equals_the_dense_columns_row_for_row( + monkeypatch, tmp_path, hh_chunk +) -> None: + """Every target's CSR row is its dense float32 column, zeros removed.""" + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + specs = _cd_surface_specs(fixtures) + materialized = _materialize( + module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk + ) + dense_frame, dense_targets, dense_specs = _dense_reference(fixtures, release, specs) + assert [spec.name for spec in materialized.compiled_specs] == [ + spec.name for spec in dense_specs + ] + assert materialized.matrix.data.dtype == np.float32 + dense_matrix = np.column_stack( + [ + dense_frame.table("household")[f"m{index:04d}"].to_numpy() + for index in range(len(dense_specs)) + ] + ).T + np.testing.assert_array_equal(materialized.matrix.toarray(), dense_matrix) + # The district rows came from carriers, and the first chunk proved one + # per carrier equal to its direct materialization. + check = materialized.carrier_check + assert check["district_rows"] == 32 + assert check["all_equal"] is True + assert check["checked_district_rows"] == check["carriers"] >= 6 + + +@requires_pytables +def test_dense_and_sparse_paths_give_identical_problems_and_weights( + monkeypatch, tmp_path +) -> None: + """The differential test of the sparse rewrite, end to end. + + Dense: float32 household columns compiled by the kernel from frame + columns (the pre-sparse checkpoint). Sparse: the carrier-split CSR, + written to and read back from the checkpoint, compiled by the kernel + from callable rows. Same constraint values, same calibrated weights. + """ + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + from microcosm.calibrate import calibrate + + fixtures = _load_fixtures() + specs = _cd_surface_specs(fixtures) + materialized = _materialize( + module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk=2 + ) + records = _sparse_records(module, materialized) + struct = module.extract_struct_tables(fixtures._nested_frame()) + checkpoint = tmp_path / "checkpoint" + _path, _registry, digests = module.write_lean_checkpoint( + struct, materialized.matrix, materialized.compiled_specs, records, checkpoint + ) + identity = _identity(digests, households=5) + lean_frame, design, registry, loaded_records, loaded_matrix = ( + module.load_checkpoint_surface(checkpoint, identity) + ) + dense_frame, dense_targets, _ = _dense_reference(fixtures, release, specs) + np.testing.assert_array_equal(design, dense_frame.weights_for("household").values) + + sparse_targets = module.cd_surface.calibration_target_set( + loaded_records, loaded_matrix, lean_frame.n("household"), specs=registry.specs + ) + _assert_same_problem( + build_constraint_matrix(dense_frame, dense_targets), + build_constraint_matrix(lean_frame, sparse_targets), + ) + + settings = dict(max_weight_ratio=5.0, target_loss_cap=1.0, l2_lambda=0.0, seed=0) + sparse_result, done = module.calibrate_surface( + lean_frame, + sparse_targets, + epochs=40, + epoch_batch=20, + **settings, + ) + assert done == 40 + warm = None + for _batch in range(2): + dense_result = calibrate( + dense_frame, + dense_targets, + weight_entity="household", + method="adam", + epochs=20, + learning_rate=0.02, + mass="conserve", + warm_start_weights=warm, + **settings, + ) + warm = dense_result.weights.copy() + np.testing.assert_array_equal(sparse_result.weights, dense_result.weights) + _assert_same_problem(dense_result.problem, sparse_result.problem) + + +def test_carriers_drop_the_district_hierarchy() -> None: + """A carrier is renamed, so it cannot keep a hierarchy naming its row.""" + + module = _load_tool_module() + fixtures = _load_fixtures() + plan = module.cd_surface.plan_carriers(_cd_surface_specs(fixtures)) + carriers = [ + spec + for spec in plan.engine_specs + if spec.name.startswith(module.cd_surface.CARRIER_PREFIX) + ] + assert len(carriers) == len(set(plan.carrier_of.values())) >= 6 + assert all(carrier.hierarchy is None for carrier in carriers) + assert all( + spec.hierarchy is not None + for spec in plan.declared + if module.cd_surface.is_cd_soi_spec(spec) and spec.name != "cd_0601_wages" + ) + + +def test_a_stored_district_row_that_disagrees_with_its_direct_row_is_refused( + monkeypatch, tmp_path +) -> None: + """The first-chunk check compares the rows the assembler stored.""" + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + assembler = module.cd_surface.SparseTargetAssembler + real_add_masked = assembler.add_masked + + def corrupting_add_masked(self, row, low, carrier, positions, *, name): + corrupted = np.asarray(carrier, dtype=np.float64).copy() + corrupted[positions] += 1.0 + return real_add_masked(self, row, low, corrupted, positions, name=name) + + monkeypatch.setattr(assembler, "add_masked", corrupting_add_masked) + with pytest.raises(RuntimeError, match="carrier split is not exact"): + _materialize( + module, + fixtures, + release, + monkeypatch, + tmp_path, + _cd_surface_specs(fixtures), + hh_chunk=5, + ) + + +def _parented_surface_specs(fixtures) -> tuple: + """The district surface with each state's AGI block under a state row.""" + + parents = { + state: fixtures._soi( + f"state_{state}_agi", "adjusted_gross_income", state_fips=state + ) + for state in ("06", "36", "24") + } + specs = [] + for spec in _cd_surface_specs(fixtures): + state = spec.metadata.get("state_fips") + if spec.name.endswith("_agi") and state in parents: + spec = dataclasses.replace( + spec, + metadata={ + **spec.metadata, + "state_cd_parent_target_name": parents[state].name, + }, + ) + specs.append(spec) + return (*specs, *parents.values()) + + +def test_every_chunk_rebuilds_each_state_parent_from_its_district_rows( + monkeypatch, tmp_path +) -> None: + """The per-chunk check runs on every chunk and covers real households.""" + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + materialized = _materialize( + module, + fixtures, + release, + monkeypatch, + tmp_path, + _parented_surface_specs(fixtures), + hh_chunk=2, + ) + # Three chunks, three state parents each; every household carries AGI. + assert materialized.carrier_check["parent_blocks_checked"] == 9 + assert materialized.carrier_check["parent_block_nonzero_households"] >= 4 + + +def test_households_outside_their_states_block_are_named(monkeypatch, tmp_path) -> None: + """A household whose district is in no row of its state's block.""" + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + specs = [ + spec for spec in _parented_surface_specs(fixtures) if spec.name != "cd_0602_agi" + ] + # Household 3 is the CA household in district 0602, whose row is missing. + with pytest.raises(RuntimeError, match=r"household_id \[3\]\) carry state_06_agi"): + _materialize( + module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk=5 + ) + + +def test_a_later_chunk_that_breaks_a_parent_is_refused(monkeypatch, tmp_path) -> None: + """Corruption past the first chunk escapes the row check, not this one.""" + + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + assembler = module.cd_surface.SparseTargetAssembler + real_add_masked = assembler.add_masked + + def corrupting_later_chunks(self, row, low, carrier, positions, *, name): + corrupted = np.asarray(carrier, dtype=np.float64).copy() + if low > 0: + corrupted[positions] += 1.0 + return real_add_masked(self, row, low, corrupted, positions, name=name) + + monkeypatch.setattr(assembler, "add_masked", corrupting_later_chunks) + with pytest.raises(RuntimeError, match="do not rebuild its directly"): + _materialize( + module, + fixtures, + release, + monkeypatch, + tmp_path, + _parented_surface_specs(fixtures), + hh_chunk=2, + ) + + +def _unparented_surface_specs(fixtures) -> tuple: + """Another mode's shape: district rows name no parent, and state rows of + the same concept sit beside them (as the district file's totals do).""" + + states = tuple( + fixtures._soi(f"state_{state}_agi", "adjusted_gross_income", state_fips=state) + for state in ("06", "36") + ) + return (*_cd_surface_specs(fixtures), *states) + + +def test_other_modes_check_each_district_row_against_a_same_concept_state_row( + monkeypatch, tmp_path +) -> None: + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + fixtures = _load_fixtures() + specs = _unparented_surface_specs(fixtures) + materialized = _materialize( + module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk=2 + ) + check = materialized.carrier_check + # The four CA and NY AGI district rows, in each of three chunks. + assert check["fallback_rows_checked"] == 4 * 3 + assert check["fallback_row_nonzero_households"] >= 4 + assert check["parent_blocks_checked"] == 0 + # Every other district row has no same-concept state row to check against. + assert check["rows_without_parent_check"] == check["district_rows"] - 4 + + assembler = module.cd_surface.SparseTargetAssembler + real_add_masked = assembler.add_masked + + def corrupting_later_chunks(self, row, low, carrier, positions, *, name): + corrupted = np.asarray(carrier, dtype=np.float64).copy() + if low > 0: + corrupted[positions] += 1.0 + return real_add_masked(self, row, low, corrupted, positions, name=name) + + monkeypatch.setattr(assembler, "add_masked", corrupting_later_chunks) + with pytest.raises(RuntimeError, match="differs from its same-concept state"): + _materialize( + module, fixtures, release, monkeypatch, tmp_path, specs, hh_chunk=2 + ) + + +def test_a_state_cd_row_without_a_compiled_parent_is_refused() -> None: + """Otherwise its per-chunk check would silently check nothing.""" + + module = _load_tool_module() + fixtures = _load_fixtures() + specs = [] + for spec in _parented_surface_specs(fixtures): + if spec.name == "cd_0601_agi": + spec = dataclasses.replace( + spec, + metadata={**spec.metadata, "state_cd_vintage_rule": "fixture"}, + ) + specs.append(spec) + plan = module.cd_surface.plan_carriers(specs) + compiled = {spec.name for spec in plan.engine_specs} + parents, _unchecked = module.cd_surface.district_row_parents(plan, compiled) + index = next( + i for i, spec in enumerate(plan.declared) if spec.name == "cd_0601_agi" + ) + assert parents[index] == ("state_06_agi", True) + with pytest.raises(ValueError, match="no compiled state parent"): + module.cd_surface.district_row_parents(plan, compiled - {"state_06_agi"}) + + +def test_district_rows_without_baseline_populations_are_refused() -> None: + module = _load_tool_module() + targets = [ + { + "name": "cd_0601", + "family": "irs_soi", + "congressional_district_geoid": "0601", + }, + {"name": "state_06", "family": "irs_soi", "state_fips": "06"}, + {"name": "pop", "family": "census_population"}, + ] + module._attach_pro_rata_populations(targets, {601: 10.0, 602: 30.0}) + module._require_pro_rata_populations(targets) + assert targets[0]["cd_population"] == 10.0 + assert targets[0]["state_population"] == 40.0 + module._attach_pro_rata_populations(targets, {602: 30.0}) + with pytest.raises(SystemExit, match="1 district SOI target"): + module._require_pro_rata_populations(targets) + + +def test_calibration_never_sees_a_held_out_target(monkeypatch, tmp_path) -> None: + module = _load_tool_module() + import build_us_fiscal_refresh_release as release + + import microcosm.calibrate as calibrate_package + + fixtures = _load_fixtures() + materialized = _materialize( + module, + fixtures, + release, + monkeypatch, + tmp_path, + _cd_surface_specs(fixtures), + hh_chunk=5, + ) + records = _sparse_records(module, materialized) + for record in records: + if record.get("geography_level") is None and record["family"] == "irs_soi": + if record.get("congressional_district_geoid"): + record["geography_level"] = "congressional_district" + record["state_cd_parent_target_name"] = "agi_amount" + module.cd_surface.assign_target_roles(records, fraction=0.5) + held = {record["name"] for record in records if record["role"] == "holdout"} + assert held, "the fixture must hold something out at 0.5" + seen: list[set[str]] = [] + real_calibrate = calibrate_package.calibrate + + def spy(frame, targets, **kwargs): + seen.append({target.name for target in targets}) + return real_calibrate(frame, targets, **kwargs) + + monkeypatch.setattr(calibrate_package, "calibrate", spy) + frame = fixtures._nested_frame() + module.calibrate_surface( + frame, + module.cd_surface.calibration_target_set(records, materialized.matrix, 5), + epochs=4, + epoch_batch=2, + max_weight_ratio=5.0, + target_loss_cap=1.0, + l2_lambda=0.0, + seed=0, + ) + assert len(seen) == 2 + trained = {record["name"] for record in records if record["role"] == "train"} + for names in seen: + assert names == trained + assert not names & held + + +def test_population_rows_match_the_dense_household_size_indicator() -> None: + module = _load_tool_module() + households = pd.DataFrame( + { + "household_id": [1, 2, 3, 4], + "state_fips": [6, 6, 36, 6], + "congressional_district_geoid": [601, 602, 3601, 601], + } + ) + size = np.asarray([2.0, 0.0, 3.0, 1.0], dtype=np.float32) + populations = { + "state": {6: 100.0, 36: 50.0, 24: 10.0}, + "cd": {601: 60.0, 602: 40.0, 3601: 50.0}, + } + records, matrix, dropped = module.cd_surface.population_rows( + households, size, populations, ["state", "cd"] + ) + assert dropped == ["pop_state_24"] + assert [record["name"] for record in records] == [ + "pop_state_06", + "pop_state_36", + "pop_cd_0601", + "pop_cd_0602", + "pop_cd_3601", + ] + # pop_cd_0602's only household has no persons: the row stays, empty. + expected = np.asarray( + [ + [2, 0, 0, 1], + [0, 0, 3, 0], + [2, 0, 0, 1], + [0, 0, 0, 0], + [0, 0, 3, 0], + ], + dtype=np.float32, + ) + np.testing.assert_array_equal(matrix.toarray(), expected) + assert records[2]["state_fips"] == "06" + assert records[2]["congressional_district_geoid"] == "0601" + + +@requires_pytables +def test_checkpoint_refuses_changed_bytes_and_dense_predecessors(tmp_path) -> None: + module = _load_tool_module() + frame = _load_fixtures()._nested_frame() + struct = module.extract_struct_tables(frame) + matrix = sparse.csr_array(np.eye(2, 5, dtype=np.float32)) + specs = [ + TargetSpec( + name=f"t{i}", entity="household", measure=f"t{i}", value=1.0, source="x" + ) + for i in range(2) + ] + records = [ + {"name": f"t{i}", "value": 1.0, "source": "x", "family": "f", "role": "train"} + for i in range(2) + ] + checkpoint = tmp_path / "ckpt" + _path, _registry, digests = module.write_lean_checkpoint( + struct, matrix, specs, records, checkpoint + ) + identity = _identity(digests, households=5) + module.load_checkpoint_surface(checkpoint, identity) + for key, file_name in ( + ("target_registry_sha256", "target_registry.json"), + ("target_roles_sha256", "target_roles.json"), + ): + with pytest.raises(SystemExit, match=f"{file_name} changed"): + module.load_checkpoint_surface(checkpoint, {**identity, key: "0" * 64}) + with pytest.raises(SystemExit, match="target_matrix.npz changed"): + module.load_checkpoint_surface( + checkpoint, {**identity, "target_matrix": {"sha256": "0" * 64}} + ) + # The lean H5 the calibrator reads is bound too, once the identity has it. + with_h5 = {**identity, "lean_h5_sha256": digests["lean_h5_sha256"]} + module.load_checkpoint_surface(checkpoint, with_h5) + with pytest.raises(SystemExit, match="target_frame_lean.h5 changed"): + module.load_checkpoint_surface( + checkpoint, {**with_h5, "lean_h5_sha256": "0" * 64} + ) + with pytest.raises(ValueError, match="row-aligned"): + module.write_lean_checkpoint(struct, matrix, specs, records[::-1], checkpoint) + (checkpoint / module.TARGET_MATRIX_FILENAME).unlink() + with pytest.raises(SystemExit, match="predates the sparse target matrix"): + module.load_checkpoint_surface(checkpoint, identity) + + +def test_holdout_scoring_against_the_pro_rata_baseline() -> None: + module = _load_tool_module() + cd_surface = module.cd_surface + records = [ + {"name": "state_agi", "value": 100.0, "family": "irs_soi", "role": "train"}, + *( + { + "name": f"cd_{d}", + "value": value, + "family": "irs_soi", + "geography_level": "congressional_district", + "state_fips": "06", + "congressional_district_geoid": d, + "source_measure_id": "adjusted_gross_income", + "state_cd_parent_target_name": "state_agi", + "cd_population": population, + "state_population": 100.0, + } + for d, value, population in (("0601", 70.0, 50.0), ("0602", 30.0, 50.0)) + ), + ] + cd_surface.assign_target_roles(records, fraction=0.5) + for record in records[1:]: + record["role"] = "holdout" + # Household 0 is in 0601, household 1 in 0602; each carries AGI 1. + matrix = sparse.csr_array(np.asarray([[1, 1], [1, 0], [0, 1]], dtype=np.float32)) + baseline = cd_surface.pro_rata_baseline(records, [1, 2]) + np.testing.assert_allclose(baseline, [50.0, 50.0]) + assert baseline.sum() == pytest.approx(records[0]["value"]) + scored = cd_surface.score_cd_holdout( + records, + matrix, + design_weights=np.asarray([50.0, 50.0]), + final_weights=np.asarray([70.0, 30.0]), + cap=1.0, + ) + assert scored["n_targets"] == 2 + assert scored["calibrated"]["mean_abs_rel_error"] == pytest.approx(0.0) + assert scored["pro_rata_baseline"]["mean_abs_rel_error"] == pytest.approx( + (20 / 70 + 20 / 30) / 2 + ) + assert scored["calibrated_beats_pro_rata_share"] == 1.0 + assert scored["by_family"]["adjusted_gross_income"]["n_targets"] == 2 + + +def test_weight_origin_counts_copies_of_one_household_once() -> None: + cd_surface = _load_tool_module().cd_surface + weights = np.asarray([1.0, 1.0, 2.0, 4.0]) + spine = np.asarray(["asec_puf", "asec_puf", "acs_2024_1yr", "asec_puf"]) + # Rows 0 and 1 are the native and PUF-detail copies of donor household 7; + # ACS source id 7 is a different household. + source = np.asarray([7, 7, 7, 9]) + summary = cd_surface.weight_origin_summary(weights, spine=spine, source_id=source) + assert summary["distinct_households"] == 3 + assert summary["effective_sample_size_distinct_households"] == pytest.approx( + 8.0**2 / (2.0**2 + 2.0**2 + 4.0**2) + ) + assert summary["effective_sample_size_rows"] == pytest.approx( + 8.0**2 / (1 + 1 + 4 + 16) + ) + assert summary["weight_share_by_spine"] == { + "acs_2024_1yr": 0.25, + "asec_puf": 0.75, + } + + +def test_sampling_is_a_rung_and_package_refuses_anything_below_full( + tmp_path, +) -> None: + module = _load_tool_module() + argv = [ + "--stage", + "materialize", + "--staging-h5", + str(tmp_path / "staging.h5"), + "--feed", + str(tmp_path / "feed.jsonl"), + "--checkpoint-dir", + str(tmp_path / "ckpt"), + ] + assert module._parse_args(argv).sample_fraction == 1.0 + assert module._parse_args([*argv, "--sample-fraction", "0.1"]).sample_seed == 578 + with pytest.raises(SystemExit): + module._parse_args([*argv, "--sample-fraction", "0.3"]) + with pytest.raises(SystemExit): + module._parse_args([*argv, "--cd-holdout-fraction", "0.6"]) + module._require_full_rung({"sampling": {"sampled": False, "rung": "f100"}}) + with pytest.raises(SystemExit, match="f010 development rung"): + module._require_full_rung({"sampling": {"sampled": True, "rung": "f010"}}) + with pytest.raises(SystemExit, match="no sampling block"): + module._require_full_rung({}) + + +def test_full_rung_returns_the_frame_unchanged() -> None: + cd_surface = _load_tool_module().cd_surface + frame = _load_fixtures()._nested_frame() + same, receipt = cd_surface.sample_staging_frame(frame, fraction=1.0, seed=578) + assert same is frame + assert receipt["rung"] == "f100" + assert receipt["sampled"] is False + + +def test_the_refresh_recipe_reproduces_the_recorded_cd_holdout() -> None: + """Re-running a recorded recipe draws the recorded holdout, 0 included.""" + + import shlex + + module = _load_tool_module() + assert "--cd-holdout-fraction" not in module.release_refresh_recipe("state", 0.0) + for fraction in (0.0, 0.1, 0.1234567): + recipe = module.release_refresh_recipe("state_cd", fraction) + assert f"--soi-mode state_cd --cd-holdout-fraction {fraction!r} " in recipe + argv = [ + token.replace("<", "").replace(">", "") for token in shlex.split(recipe)[3:] + ] + args = module._parse_args(argv) + assert module._cd_holdout_fraction(args) == fraction + + +def test_holdout_needs_the_state_cd_surface(tmp_path) -> None: + module = _load_tool_module() + args = SimpleNamespace( + families="snap,medicaid,soi", + geographies="state,cd", + soi_mode="state", + cd_holdout_fraction=0.1, + ) + with pytest.raises(SystemExit, match="needs --soi-mode state_cd"): + module.do_materialize(args) + assert module._cd_holdout_fraction( + SimpleNamespace(soi_mode="state_cd", cd_holdout_fraction=None) + ) == pytest.approx(0.1) + assert ( + module._cd_holdout_fraction( + SimpleNamespace(soi_mode="state", cd_holdout_fraction=None) + ) + == 0.0 + ) + + +def test_district_exclusions_were_reviewed_against_the_packaged_crosswalk() -> None: + """Regenerating the 117th->119th crosswalk must revisit the exclusions. + + North Carolina's district rows came back with #1043's registry-built + crosswalk; a later crosswalk change needs the same review. + """ + + from microcosm.build.us_runtime import ( + default_congressional_district_vintage_crosswalk_path, + ) + + module = _load_tool_module() + digest = module._sha256(default_congressional_district_vintage_crosswalk_path()) + assert digest == module.cd_surface.STATE_CD_REVIEWED_CROSSWALK_SHA256, ( + "The packaged CD crosswalk changed. Re-check which states' district " + "rows it maps correctly, update STATE_CD_EXCLUDED_CD_STATES, and re-pin " + "the reviewed digest." + ) + assert module.cd_surface.STATE_CD_EXCLUDED_CD_STATES == {} + + +def _pinned_feed() -> Path: + import os + + feed = os.environ.get("MICROCOSM_US_CHRONICLE_FACTS") + if not feed or not Path(feed).exists(): + pytest.skip("MICROCOSM_US_CHRONICLE_FACTS is not set to the pinned feed") + return Path(feed) + + +def test_pinned_feed_state_cd_surface_matches_its_contract() -> None: + """On the pinned feed the ``state_cd`` surface has the documented shape. + + Counts are docs/us-acs-local-soi-target-surface.md's; every district + block adds up to its state parent; each state concept has one vintage; + every district of North Carolina is bound and no defective district-file + column is. + """ + + feed = _pinned_feed() + module = _load_tool_module() + surface = module.state_admin_surface( + feed, ["snap", "medicaid", "soi"], soi_mode="state_cd" + ) + specs = list(surface.registry.specs) + receipt = surface.soi_receipt + soi = [spec for spec in specs if spec.family == "irs_soi"] + assert receipt["counts"] == { + "historic_table_2_state": 3819, + "cd_file_state": 302, + "congressional_district": 21743, + "congressional_district_by_parent_basis": { + "historic_table_2": 19181, + "cd_file_state_total_bridged": 2562, + }, + "total": 25864, + } + assert len(soi) == 25864 + assert len(specs) == 25864 + 102 + 51 + reconciliation = module.cd_surface.state_parent_reconciliation(soi) + assert len(reconciliation) == 2189 + assert all(block["ok"] for block in reconciliation) + state_rows = [ + spec for spec in soi if spec.metadata.get("ledger_geography_level") == "state" + ] + keys = [ + (spec.metadata["state_fips"], module.cd_surface.soi_concept_identity(spec)) + for spec in state_rows + ] + assert len(keys) == len(set(keys)), "a state concept is bound twice" + districts = [ + spec + for spec in soi + if spec.metadata.get("ledger_geography_level") == "congressional_district" + ] + # North Carolina's 14 districts are bound since #1043's crosswalk. + assert ( + len([spec for spec in districts if spec.metadata["state_fips"] == "37"]) + == 14 * 51 + ) + defective = set(module.cd_surface.STATE_CD_DEFECTIVE_CD_FILE_MEASURES) + assert not [ + spec.name + for spec in soi + if module.cd_surface.is_cd_file_spec(spec) + and spec.metadata["source_measure_id"] in defective + ] + assert all(spec.metadata.get("state_cd_parent_target_name") for spec in districts) + assert receipt["sigma"]["targets_with_sigma"] == 0 + # The factor band drops these (state, measure) pairs, and only these: + # four district blocks whose two publications disagree about the state, + # and the itemized bridge of two at-large states (no district rows). + # The district file agrees with itself in every block, to float precision. + internal = receipt["internal_consistency"] + assert internal["blocks_checked"] == 2193 + assert internal["max_relative_gap"] < 1e-12 + band = receipt["factor_band"] + assert band["tolerance"] == module.cd_surface.STATE_CD_FACTOR_BAND == 1.25 + assert { + (row["state_fips"], row["measure"], row["basis"].split(":")[0]) + for row in band["out_of_band"] + } == { + ("15", "ordinary_dividends_amount", "rebase"), + ("15", "qualified_dividends_amount", "rebase"), + ("36", "rental_royalty_income_amount", "rebase"), + ("49", "tax_exempt_interest_amount", "rebase"), + ("46", "charitable_amount", "level_bridge"), + ("46", "interest_paid_deduction_amount", "level_bridge"), + ("56", "charitable_amount", "level_bridge"), + ("56", "interest_paid_deduction_amount", "level_bridge"), + } + # 43 states carry district rows (51 minus the 8 at-large on the 117th + # plan), less the band's four blocks. + rebase_n = { + measure: entry["n"] + for measure, entry in receipt["rebase_factor_by_measure"].items() + } + banded = { + "ordinary_dividends_amount", + "qualified_dividends_amount", + "rental_royalty_income_amount", + "tax_exempt_interest_amount", + } + assert {measure for measure, n in rebase_n.items() if n == 42} == banded + assert set(rebase_n.values()) == {42, 43} + # Every district-file-only state level is bridged, in all 51 states but + # the band's two itemized outliers. + bridges = receipt["level_bridge_factor_by_measure"] + assert set(bridges) == set(module.cd_surface.STATE_CD_LEVEL_BRIDGES) + assert {measure: entry["n"] for measure, entry in bridges.items()} == { + "charitable_amount": 49, + "charitable_returns": 51, + "interest_paid_deduction_amount": 49, + "interest_paid_deduction_returns": 51, + "qualified_business_income_deduction_amount": 51, + "qualified_business_income_deduction_returns": 51, + } + assert json.dumps(receipt) # the receipt is manifest-serializable + + +def test_the_cd_holdout_hash_is_pinned() -> None: + """Golden values: a changed hash formula would silently redraw the holdout. + + The US holdout port reuses the same function, so both lines move together. + """ + + from microcosm.build.holdout import hash_holdout_uniform + + salt = _load_tool_module().cd_surface.CD_HOLDOUT_SALT + assert hash_holdout_uniform("06|eitc", salt=salt) == 0.7980954941465199 + assert ( + hash_holdout_uniform("36|adjusted_gross_income", salt=salt) + == 0.037229772944579555 + ) + + +def _stratified_frame(donor_sizes: tuple[int, ...]) -> Frame: + """Two spines x five districts; one weight per stratum (so sums are exact).""" + + from microcosm.frame import WeightKind, Weights + + districts = ["0601", "0602", "3601", "3602", "2401"] + rows = [] + for spine, sizes in (("acs_2024_1yr", (40,) * 5), ("asec_puf", donor_sizes)): + for index, (district, size) in enumerate(zip(districts, sizes, strict=True)): + rows += [ + (spine, district, 10.0 + index + (50 if spine == "asec_puf" else 0)) + ] * size + n = len(rows) + ids = np.arange(1, n + 1, dtype=np.int64) + person = pd.DataFrame( + { + "person_id": ids, + "person_household_id": ids, + **{f"person_{group}_id": ids for group in US_SCHEMA.group_entities}, + } + ) + tables = {"person": person} + for group in US_SCHEMA.group_entities: + tables[group] = pd.DataFrame({f"{group}_id": ids}) + tables["household"] = pd.DataFrame( + { + "household_id": ids, + "household_spine": [row[0] for row in rows], + "congressional_district_geoid": [row[1] for row in rows], + } + ) + weights = np.asarray([row[2] for row in rows]) + return Frame(tables, US_SCHEMA, {"household": Weights(weights, WeightKind.DESIGN)}) + + +def _stratum_mass(frame) -> pd.Series: + table = frame.table("household").assign( + weight=frame.weights_for("household").values + ) + return table.groupby( + ["household_spine", "congressional_district_geoid"] + ).weight.sum() + + +def test_a_rung_keeps_every_drawn_stratum_at_full_mass() -> None: + """Inverse-rate stratum weights: with one weight per stratum, exact.""" + + cd_surface = _load_tool_module().cd_surface + frame = _stratified_frame(donor_sizes=(8, 8, 8, 8, 8)) + sampled, receipt = cd_surface.sample_staging_frame(frame, fraction=0.25, seed=578) + assert receipt["rung"] == "f025" + assert receipt["zero_draw_strata"] == 0 + assert all( + factor == pytest.approx(1.0, rel=1e-12) + for factor in receipt["per_spine_normalization_factor"].values() + ) + pd.testing.assert_series_equal(_stratum_mass(sampled), _stratum_mass(frame)) + + +def test_zero_draw_strata_are_recorded_and_spread_over_their_spine() -> None: + cd_surface = _load_tool_module().cd_surface + # The first donor district (3 households) floors to zero draws at f025. + frame = _stratified_frame(donor_sizes=(3, 8, 8, 8, 8)) + full = _stratum_mass(frame) + sampled, receipt = cd_surface.sample_staging_frame(frame, fraction=0.25, seed=578) + assert receipt["zero_draw_strata"] == 1 + lost = full[("asec_puf", "0601")] + assert receipt["zero_draw_weight_share"] == pytest.approx(lost / full.sum()) + mass = _stratum_mass(sampled) + assert ("asec_puf", "0601") not in mass.index + donor_full = full.xs("asec_puf").sum() + factor = donor_full / (donor_full - lost) + for district in ("0602", "3601", "3602", "2401"): + assert mass[("asec_puf", district)] == pytest.approx( + full[("asec_puf", district)] * factor, rel=1e-12 + ) + assert mass.xs("asec_puf").sum() == pytest.approx(donor_full, rel=1e-12) + # A spine none of whose strata draws cannot be represented at the rung. + with pytest.raises(ValueError, match="drew no weight on spine 'asec_puf'"): + cd_surface.sample_staging_frame( + _stratified_frame(donor_sizes=(3, 3, 3, 3, 3)), fraction=0.25, seed=578 + ) + + +@requires_pytables +def test_calibration_outputs_are_bound_to_the_materialization(tmp_path) -> None: + """Resume, the complete shortcut, finalize and package refuse stale evidence.""" + + module = _load_tool_module() + identity = { + "staging_sha256": "s", + "target_roles_sha256": "r", + "target_registry_sha256": "g", + "target_matrix": {"sha256": "m"}, + "sampling": {"rung": "f100", "sampled": False}, + } + stamp = module._run_identity_digest(identity) + assert stamp == module._run_identity_digest(dict(reversed(identity.items()))) + changed = {**identity, "sampling": {"rung": "f010", "sampled": True}} + assert module._run_identity_digest(changed) != stamp + module._require_current_calibration( + identity, {"run_identity_sha256": stamp}, stage="finalize" + ) + for summary in ({}, {"run_identity_sha256": module._run_identity_digest(changed)}): + with pytest.raises(SystemExit, match="does not belong to this checkpoint"): + module._require_current_calibration(identity, summary, stage="package") + # A pre-sparse identity (no roles digest) stays readable. + module._require_current_calibration({"staging_sha256": "s"}, {}, stage="finalize") + assert set(module.CALIBRATION_OUTPUT_FILENAMES) == { + "weights_latest.npz", + "calibration_summary.json", + "calibration_diagnostics.json", + "consumer_export.json", + "consumer_reviewed_null_fills.json", + "spine_qa.json", + } + # The calibrated H5 and its evidence are bound too: to the materialization, + # to the summary's weights, and QA to the H5's bytes. Not to a path. + h5 = tmp_path / "out.h5" + h5.write_bytes(b"artifact") + h5_sha = module._sha256(h5) + weights = np.array([1.0, 2.5, 3.0]) + weights_sha = module._weights_digest(weights) + assert weights_sha == module._weights_digest(weights.astype(np.float64).copy()) + assert weights_sha != module._weights_digest(weights * (1 + 1e-12)) + export = { + "run_identity_sha256": stamp, + "out_h5": str(tmp_path / "elsewhere" / "out.h5"), + "out_h5_sha256": h5_sha, + "weights_sha256": weights_sha, + } + summary = {"run_identity_sha256": stamp, "weights_sha256": weights_sha} + qa = {"run_identity_sha256": stamp, "artifact_sha256": h5_sha} + module._require_current_artifact( + identity, + consumer_export=export, + summary=summary, + spine_qa=qa, + out_h5=h5, + out_h5_sha256=h5_sha, + stage="package", + ) + for bad_export, bad_summary, bad_qa, message in ( + ({**export, "run_identity_sha256": "old"}, summary, None, "consumer_export"), + ({**export, "out_h5_sha256": "0" * 64}, summary, None, "is not the calibrated"), + # An interrupted recalibration: new H5 and export, the old summary. + (export, {**summary, "weights_sha256": "old"}, None, "different weights"), + ({**export, "weights_sha256": None}, summary, None, "different weights"), + (export, summary, {**qa, "run_identity_sha256": "old"}, "another material"), + (export, summary, {**qa, "artifact_sha256": "0" * 64}, "certifies other"), + ): + with pytest.raises(SystemExit, match=message): + module._require_current_artifact( + identity, + consumer_export=bad_export, + summary=bad_summary, + spine_qa=bad_qa, + out_h5=h5, + out_h5_sha256=h5_sha, + stage="package", + ) + + +def test_resume_refuses_weights_without_the_run_identity_stamp( + tmp_path, monkeypatch +) -> None: + module = _load_tool_module() + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + identity = {"staging_sha256": "s", "target_roles_sha256": "r"} + monkeypatch.setattr(module, "_verify_run_identity", lambda args: identity) + frame = _load_fixtures()._nested_frame() + registry = SimpleNamespace(specs=()) + monkeypatch.setattr( + module, + "load_checkpoint_surface", + lambda *a, **k: ( + frame, + np.ones(5), + registry, + [], + sparse.csr_array((0, 5), dtype=np.float32), + ), + ) + monkeypatch.setattr(module.cd_surface, "calibration_target_set", lambda *a, **k: []) + args = SimpleNamespace( + checkpoint_dir=checkpoint, + resume=True, + epochs=10, + epoch_batch=5, + max_weight_ratio=5.0, + target_loss_cap=1.0, + l2_lambda=0.0, + seed=0, + ) + settings = module._solver_settings(args) + for saved in ( + {"weights": np.ones(5), "epochs_done": 5, "staging_sha256": "s"}, + { + "weights": np.ones(5), + "epochs_done": 5, + "run_identity_sha256": module._run_identity_digest( + {**identity, "target_roles_sha256": "other"} + ), + }, + ): + np.savez(checkpoint / "weights_latest.npz", **saved) + with pytest.raises(SystemExit, match="different materialization"): + module.do_calibrate(args) + # A matching stamp under other solver settings is refused too. + np.savez( + checkpoint / "weights_latest.npz", + weights=np.ones(5), + epochs_done=5, + run_identity_sha256=module._run_identity_digest(identity), + solver_settings=json.dumps({**settings, "l2_lambda": 0.1}, sort_keys=True), + ) + with pytest.raises(SystemExit, match="different solver settings"): + module.do_calibrate(args) + # A matching stamp with every epoch done still needs a matching summary. + np.savez( + checkpoint / "weights_latest.npz", + weights=np.ones(5), + epochs_done=10, + run_identity_sha256=module._run_identity_digest(identity), + solver_settings=json.dumps(settings, sort_keys=True), + ) + (checkpoint / "calibration_summary.json").write_text( + json.dumps({"run_identity_sha256": "stale"}) + ) + with pytest.raises(SystemExit, match="belongs to another materialization"): + module.do_calibrate(args) + # ... including its solver settings. + (checkpoint / "calibration_summary.json").write_text( + json.dumps( + { + "run_identity_sha256": module._run_identity_digest(identity), + "solver_settings": {**settings, "seed": 1}, + } + ) + ) + with pytest.raises(SystemExit, match="solver settings or weights"): + module.do_calibrate(args) + # ... and its weights: a summary from before the weights digest would + # otherwise send --resume round the shortcut forever. + (checkpoint / "calibration_summary.json").write_text( + json.dumps( + { + "run_identity_sha256": module._run_identity_digest(identity), + "solver_settings": settings, + } + ) + ) + with pytest.raises(SystemExit, match="run without --resume"): + module.do_calibrate(args) + # A solve deletes every previous calibration output before it starts, so + # one that stops part way leaves no evidence describing other weights. + for name in module.CALIBRATION_OUTPUT_FILENAMES: + (checkpoint / name).write_text("{}") + + def interrupted(*_args, **_kwargs): + raise RuntimeError("solver interrupted") + + monkeypatch.setattr(module, "calibrate_surface", interrupted) + with pytest.raises(RuntimeError, match="interrupted"): + module.do_calibrate(SimpleNamespace(**{**vars(args), "resume": False})) + assert sorted(path.name for path in checkpoint.iterdir()) == ["weights_latest.npz"] diff --git a/test_support/microcosm_build/frame_serializer_registry.py b/test_support/microcosm_build/frame_serializer_registry.py index c5ac485a6..6e5bec73a 100644 --- a/test_support/microcosm_build/frame_serializer_registry.py +++ b/test_support/microcosm_build/frame_serializer_registry.py @@ -11,6 +11,7 @@ import numpy as np import pandas as pd import pytest +from scipy import sparse from microcosm.build.frame_checkpoint import ( load_frame_checkpoint, @@ -307,13 +308,10 @@ def _round_trip_acs_lean(tmp_path: Path, nullable_case: str) -> BooleanRoundTrip }, "weights": np.asarray([1.0, 2.0, 3.0]), } - path, _targets = tool.write_lean_checkpoint( + path, _registry, _digests = tool.write_lean_checkpoint( struct, - np.empty((3, 0), dtype=np.float64), - [], - [], - [], - [], + sparse.csr_array((0, 3), dtype=np.float32), + (), [], tmp_path / "acs-lean", ) diff --git a/tools/build_us_acs_local_release.py b/tools/build_us_acs_local_release.py index e58e12734..7c3cf446c 100644 --- a/tools/build_us_acs_local_release.py +++ b/tools/build_us_acs_local_release.py @@ -10,22 +10,29 @@ Stages (``--stage all`` runs materialize -> calibrate -> qa -> finalize -> package; each is separately resumable): - materialize : compile the state-level administrative surface from the - Ledger feed exactly like the production path + materialize : compile the administrative surface from the Ledger feed + exactly like the production path (``compile_us_fiscal_target_registry`` -> RI Medicaid substitution -> state {usda_snap, cms_medicaid[enrollment], irs_soi}; ``--soi-mode state`` by default -- Build O's - state-geography SOI contract -- with ``totals`` and ``full`` - as explicit opt-ins), run the household-chunked engine pass under the - nullable-artifact contract (input-schema projection + - reviewed-null fill), add PUMA-ladder population marginals - (state + congressional district), and write a lean float32 - target-frame checkpoint + targets.json. Heavy stage; a crash - in calibrate never re-runs the microsim. - calibrate : epoch-batched warm-start calibrate on the lean checkpoint - (adam, mass conserved, hard weight-ratio cap; resumable with - --resume), write diagnostics and the calibrated weights onto - a copy of the staging H5. + state-geography SOI contract -- with ``totals``, ``full`` and + ``state_cd`` (district SOI rows, one vintage per state + concept) as explicit opt-ins), run the household-chunked + engine pass under the nullable-artifact contract (input-schema + projection + reviewed-null fill), add PUMA-ladder population + marginals (state + congressional district), assign the CD + holdout, and write the lean checkpoint: a structure-only H5, + target_registry.json, a sparse (targets x households) float32 + CSR target matrix row-aligned with it (target_matrix.npz) and + target_roles.json. ``--sample-fraction`` draws a development + rung. + Heavy stage; a crash in calibrate never re-runs the microsim. + calibrate : epoch-batched warm-start calibrate on the checkpoint's + training targets (adam, mass conserved, hard weight-ratio cap; + resumable with --resume), score the held-out district targets + against the pro-rata baseline, record ESS over rows and + distinct households, and write the calibrated weights onto a + copy of the staging H5. qa : chunked engine probe of the calibrated artifact recording per-spine SSI incidence and intensity (the microcosm#403 signature, measured rather than assumed). @@ -59,6 +66,7 @@ import sys import threading import time +from collections import defaultdict from datetime import UTC, datetime from pathlib import Path @@ -74,6 +82,8 @@ # the sibling release tool rather than forking it. sys.path.insert(0, str(_TOOLS_DIR)) +import us_acs_local_cd_surface as cd_surface # noqa: E402 - needs tools/ on sys.path + PERIOD = 2024 RELEASE_NAMESPACE = "buildo_acs_local" RELEASE_ID_PREFIX = "populace-us-2024-buildo-acs-local" @@ -104,13 +114,21 @@ #: ``totals`` keeps only the specs whose ``target_role`` is not #: ``soi_fiscal_distribution`` -- on the pinned feed, ACA premium tax credit #: rows and no state AGI, income-tax or EITC total. ``full`` keeps every -#: state-bearing spec, including the TY2023 congressional-district file, which -#: needs a dense matrix too large for one 128 GB machine. Both are explicit -#: opt-ins. docs/us-acs-local-soi-target-surface.md has the measured surfaces. +#: state-bearing spec, including the congressional-district file, with both +#: vintages of each state concept. Both are explicit opt-ins. +#: +#: ``state_cd`` (explicit opt-in) binds the congressional-district file's +#: district rows on top of ``state``, keeping one vintage per state concept: +#: Historic Table 2 gives a state concept's level, the district file only +#: each district's share of it (``us_acs_local_cd_surface.state_cd_soi_surface``). +#: A hash-assigned block of its district targets is held out of calibration +#: and scored against a pro-rata baseline. The target matrix is sparse in +#: every mode. docs/us-acs-local-soi-target-surface.md has the surfaces. SOI_MODE_STATE = "state" SOI_MODE_TOTALS = "totals" SOI_MODE_FULL = "full" -SOI_MODES = (SOI_MODE_STATE, SOI_MODE_TOTALS, SOI_MODE_FULL) +SOI_MODE_STATE_CD = "state_cd" +SOI_MODES = (SOI_MODE_STATE, SOI_MODE_TOTALS, SOI_MODE_FULL, SOI_MODE_STATE_CD) DEFAULT_SOI_MODE = SOI_MODE_STATE #: Ledger record-set specs from a congressional-district file start with this; #: ``state`` mode excludes them, which is what Build O's switch did. @@ -126,19 +144,30 @@ def _require_soi_mode(soi_mode: str) -> str: return soi_mode -def release_refresh_recipe(soi_mode: str) -> str: +def release_refresh_recipe( + soi_mode: str, cd_holdout_fraction: float | None = None +) -> str: """The one-command release refresh, pinned to the SOI surface it built. The recipe names ``--soi-mode`` explicitly so re-running it reproduces - the recorded surface even if the parser default changes again. + the recorded surface even if the parser default changes again, and, for + every ``state_cd`` build, ``--cd-holdout-fraction`` at the recorded value + (0 included, full precision), since ``state_cd`` otherwise defaults to a + 10% holdout. """ _require_soi_mode(soi_mode) + holdout = ( + f"--cd-holdout-fraction {float(cd_holdout_fraction)!r} " + if soi_mode == SOI_MODE_STATE_CD and cd_holdout_fraction is not None + else "" + ) return ( "uv run tools/build_us_acs_local_release.py --stage all " "--staging-h5 /acs_multispine_staging.h5 " "--feed --feed-sha256 " f"--soi-mode {soi_mode} " + f"{holdout}" "--ladder build/us/us_puma_ladder_2020.npz " "--checkpoint-dir /checkpoints " "--out-h5 /populace_us_2024_acs_local.h5 " @@ -214,9 +243,13 @@ def soi_surface_predicate(soi_mode: str): AGI-band slices: every SOI fact without a named role gets it, which at state level includes the all-income-range state and district rows. - ``full`` keeps them all. + - ``state_cd`` is not a filter: this returns the ``full`` candidates it + starts from, and :func:`state_admin_surface` reconciles them. """ _require_soi_mode(soi_mode) + if soi_mode == SOI_MODE_STATE_CD: + soi_mode = SOI_MODE_FULL def selected(spec) -> bool: metadata = spec.metadata @@ -245,11 +278,34 @@ def state_admin_specs( ): """Select the state-level admin surface from the production compile path. + Returns ``(registry, ri_substitutions)``; :func:`state_admin_surface` + also returns the SOI surface receipt. + """ + + surface = state_admin_surface(feed, families, soi_mode=soi_mode) + return surface.registry, surface.ri_substitutions + + +class AdminSurface: + """The compiled admin registry plus how its SOI slice was chosen.""" + + def __init__(self, registry, ri_substitutions, soi_receipt: dict) -> None: + self.registry = registry + self.ri_substitutions = ri_substitutions + self.soi_receipt = soi_receipt + + +def state_admin_surface( + feed: str | Path, families: list[str], soi_mode: str = DEFAULT_SOI_MODE +) -> AdminSurface: + """Select the admin surface from the production compile path. + feed -> ``compile_us_fiscal_target_registry(age_targets=True)`` -> ``apply_us_medicaid_enrollment_substitutions`` (RI FIPS-44) -> state-level {usda_snap, cms_medicaid[enrollment], irs_soi}. The SOI slice follows - :func:`soi_surface_predicate`: ``state`` (the default), ``totals`` or - ``full``. + :func:`soi_surface_predicate` for ``state`` (the default), ``totals`` and + ``full``; ``state_cd`` adds the reconciled congressional-district rows + (:func:`us_acs_local_cd_surface.state_cd_soi_surface`). """ # Refuse an unknown mode before loading the feed and compiling the registry. @@ -269,14 +325,12 @@ def state_admin_specs( from microcosm.calibrate.registry import TargetRegistry artifact = load_ledger_consumer_artifact(str(feed)) + crosswalk_path = default_congressional_district_vintage_crosswalk_path() + crosswalk = load_congressional_district_vintage_crosswalk(crosswalk_path) registry = compile_us_fiscal_target_registry( artifact.facts, target_period=PERIOD, - congressional_district_vintage_crosswalk=( - load_congressional_district_vintage_crosswalk( - default_congressional_district_vintage_crosswalk_path() - ) - ), + congressional_district_vintage_crosswalk=crosswalk, age_targets=True, ) registry, ri_substitutions = apply_us_medicaid_enrollment_substitutions(registry) @@ -295,11 +349,26 @@ def state_level(spec) -> bool: and spec.metadata.get("target_role") == "medicaid_enrollment" ), ).specs + soi_receipt: dict = {"soi_mode": soi_mode} if "soi" in families: - picked += registry.select( + candidates = registry.select( family="irs_soi", predicate=soi_surface_predicate(soi_mode) ).specs - return TargetRegistry(list(picked), country="us"), ri_substitutions + if soi_mode == SOI_MODE_STATE_CD: + surface = cd_surface.state_cd_soi_surface( + candidates, + state_surface_predicate=soi_surface_predicate(SOI_MODE_STATE), + crosswalk=crosswalk, + crosswalk_sha256=_sha256(crosswalk_path), + ) + picked += list(surface.specs) + soi_receipt.update(surface.receipt) + else: + picked += candidates + soi_receipt["counts"] = {"total": len(candidates)} + return AdminSurface( + TargetRegistry(list(picked), country="us"), ri_substitutions, soi_receipt + ) # --------------------------------------------------------------------------- @@ -499,6 +568,106 @@ def fill_reviewed_nulls( return fills, drift_warnings +class MaterializedSurface: + """The sparse admin surface one engine pass produced.""" + + def __init__(self, matrix, compiled_specs, chunk_stats, carrier_check) -> None: + #: (compiled specs x households) CSR, float32 data, declared order. + self.matrix = matrix + self.compiled_specs = compiled_specs + self.chunk_stats = chunk_stats + #: Direct-vs-carrier evidence from the first chunk (district rows). + self.carrier_check = carrier_check + + @property + def names(self) -> list[str]: + return [spec.measure for spec in self.compiled_specs] + + +def _check_district_rows_against_parents( + parents, captured, target_households, measure_of, *, n_chunk: int, tally: dict +) -> None: + """Every chunk: each stored district row against its state parent's column. + + ``parents`` comes from ``cd_surface.district_row_parents``. A strict + (``state_cd``) block's district rows partition the parent's state, so the + float32 values stored for them, laid side by side, must equal the parent's + directly materialized column on every household of the chunk, bit for bit. + A fallback parent (another mode's same-concept state row) is compared on + the row's own households. Unlike the first-chunk row check, this covers + every household each concept touches, so it cannot be vacuous. + """ + + strict_blocks: dict[str, list[int]] = defaultdict(list) + fallback_rows = fallback_nonzero = 0 + for index in sorted(captured): + if index not in parents: + continue + parent_name, strict = parents[index] + if strict: + strict_blocks[parent_name].append(index) + continue + positions, values = captured[index] + direct = target_households[measure_of[parent_name]].to_numpy(dtype=np.float32) + if not np.array_equal(values, direct[positions]): + differing = int((values != direct[positions]).sum()) + raise RuntimeError( + f"A stored district row differs from its same-concept state " + f"row {parent_name} on {differing} of its household(s); the " + "carrier split is not exact for this surface." + ) + fallback_rows += 1 + fallback_nonzero += int(np.count_nonzero(values)) + checked = nonzero = 0 + for parent_name, rows in strict_blocks.items(): + rebuilt = np.zeros(n_chunk, dtype=np.float32) + written = np.zeros(n_chunk, dtype=bool) + for index in rows: + positions, values = captured[index] + if written[positions].any(): + raise RuntimeError( + f"District rows of {parent_name} overlap on a household." + ) + written[positions] = True + rebuilt[positions] = values + direct = target_households[measure_of[parent_name]].to_numpy(dtype=np.float32) + if not np.array_equal(rebuilt, direct): + outside = np.flatnonzero(~written & (direct != 0)) + if len(outside): + household_ids = target_households["household_id"].to_numpy() + raise RuntimeError( + f"{len(outside)} household(s) (household_id " + f"{household_ids[outside[:5]].tolist()}) carry {parent_name} " + "but sit in no district of its block: their congressional " + "district is not one of their state's current-plan " + "districts, or the block is missing that district's row." + ) + differing = int((rebuilt != direct).sum()) + raise RuntimeError( + f"The stored district rows of {parent_name} do not rebuild its " + f"directly materialized column ({differing} household(s)); the " + "carrier split is not exact for this surface." + ) + checked += 1 + nonzero += int(np.count_nonzero(direct)) + for key, value in ( + ("parent_blocks_checked", checked), + ("parent_block_nonzero_households", nonzero), + ("fallback_rows_checked", fallback_rows), + ("fallback_row_nonzero_households", fallback_nonzero), + ): + tally[key] = tally.get(key, 0) + value + + +def _household_codes(release_tool, households, column: str) -> np.ndarray: + return np.asarray( + release_tool._integer_geography_codes( + households[column].to_numpy(), column=column + ), + dtype=np.int64, + ) + + def materialize_chunked( base_frame, specs, @@ -509,15 +678,24 @@ def materialize_chunked( dropped_manifest_path: Path | None = None, summary_path: Path | None = None, fills_manifest_path: Path | None = None, - matrix_path: Path | None = None, -): - """Household-chunked production materialization (memory-bounded). +) -> MaterializedSurface: + """Household-chunked production materialization into a sparse matrix. One full-frame Microsimulation caches the entire dependency closure of the SOI components (the 1.6M-row jetsam class), so the unchanged production ``_materialize_target_frame`` runs on whole-household - sub-frames; each chunk's calc cache dies with its simulation and the - float32 measure matrix accumulates on disk when ``matrix_path`` is set. + sub-frames; each chunk's calc cache dies with its simulation. Each + chunk's measure columns go straight into a (targets x households) CSR + matrix with float32 values (``SparseTargetAssembler``): no dense + households x targets matrix exists at any point. + + District SOI rows are not materialized one column each. The engine pass + materializes one geography-free carrier per distinct SOI concept, and + each district row is its carrier restricted to the district's households + (``us_acs_local_cd_surface.plan_carriers``). That equals the direct + materialization because the SOI slice masks a tax unit by its household's + state and district. The first chunk also materializes one district row + per carrier directly and refuses any difference. """ import build_us_fiscal_refresh_release as release_tool @@ -561,9 +739,12 @@ def materialize_chunked( "table; frame is invalid." ) - measure_names = None - compiled_specs = None - matrix = None + plan = cd_surface.plan_carriers(specs) + by_name = {spec.name: spec for spec in plan.declared} + assembler = cd_surface.SparseTargetAssembler(plan.n_rows, n_households) + compiled_names: set[str] | None = None + carrier_check: dict[str, object] = {"district_rows": len(plan.split)} + parents: dict[int, tuple[str, bool]] | None = None chunk_stats = [] n_chunks = (n_households + hh_chunk - 1) // hh_chunk if n_chunks > 1: @@ -574,32 +755,36 @@ def materialize_chunked( started = time.time() mask = (person_position >= low) & (person_position < high) sub_frame = projected.select(mask) + chunk_households = sub_frame.table("household") + state_codes = district_codes = None + if plan.split: + state_codes = _household_codes(release_tool, chunk_households, "state_fips") + district_codes = _household_codes( + release_tool, chunk_households, "congressional_district_geoid" + ) + check_specs = ( + cd_surface.carrier_check_specs(plan, district_codes) + if chunk_index == 0 and plan.split + else () + ) target_frame, compiled_registry, _ = release_tool._materialize_target_frame( sub_frame, - tuple(specs), + plan.engine_specs + check_specs, maximum_microsim_batch_size=batch, refuse_population_aggregates=True if n_chunks > 1 else None, ) - names = [spec.measure for spec in compiled_registry.specs] - if measure_names is None: - measure_names = names - compiled_specs = list(compiled_registry.specs) - if matrix_path is not None: - matrix = np.memmap( - matrix_path, - dtype=np.float32, - mode="w+", - shape=(n_households, len(names)), - ) - else: - matrix = np.zeros((n_households, len(names)), dtype=np.float32) - elif names != measure_names: + compiled = {spec.name for spec in compiled_registry.specs} + engine_compiled = compiled - {spec.name for spec in check_specs} + if compiled_names is None: + compiled_names = engine_compiled + elif engine_compiled != compiled_names: raise RuntimeError( f"chunk {chunk_index} compiled a different measure set " - f"({len(names)} vs {len(measure_names)}); refusing to assemble." + f"({len(engine_compiled)} vs {len(compiled_names)}); refusing " + "to assemble." ) - chunk_households = target_frame.table("household") - got_ids = chunk_households["household_id"].to_numpy() + target_households = target_frame.table("household") + got_ids = target_households["household_id"].to_numpy() if len(got_ids) != high - low or not np.array_equal( got_ids, household_ids[low:high] ): @@ -607,11 +792,91 @@ def materialize_chunked( f"chunk {chunk_index} household order/id mismatch; refusing " "to assemble." ) - for j, measure in enumerate(measure_names): - matrix[low:high, j] = chunk_households[measure].to_numpy(dtype=np.float32) - if matrix_path is not None: - matrix.flush() - del target_frame, compiled_registry, chunk_households, sub_frame, got_ids + for index, measure in plan.direct_rows: + if plan.declared[index].name in compiled_names: + assembler.add_column( + index, + low, + target_households[measure].to_numpy(dtype=np.float64), + name=plan.declared[index].name, + ) + if plan.split: + carriers = { + carrier: target_households[carrier].to_numpy(dtype=np.float64) + for carrier in sorted(set(plan.carrier_of.values())) + if carrier in compiled_names + } + split = tuple(row for row in plan.split if row[3] in carriers) + index_of = {spec.name: index for index, spec in enumerate(plan.declared)} + captured = cd_surface.split_carriers_into( + assembler, + cd_surface.CarrierPlan( + declared=plan.declared, + engine_specs=plan.engine_specs, + direct_rows=plan.direct_rows, + split=split, + carrier_of=plan.carrier_of, + ), + carriers, + low=low, + state_codes=state_codes, + district_codes=district_codes, + capture=frozenset(index for index, *_rest in split), + ) + if parents is None: + parents, unchecked = cd_surface.district_row_parents( + plan, compiled_names + ) + carrier_check["rows_without_parent_check"] = len(unchecked) + _check_district_rows_against_parents( + parents, + captured, + target_households, + {name: by_name[name].measure for name, _strict in parents.values()}, + n_chunk=high - low, + tally=carrier_check, + ) + if check_specs: + checked = [] + for spec in check_specs: + carrier = plan.carrier_of[spec.name] + if spec.name not in compiled or carrier not in carriers: + raise RuntimeError( + f"carrier check spec {spec.name} or its carrier " + f"{carrier} did not materialize." + ) + direct = target_households[spec.measure].to_numpy(dtype=np.float64) + positions, values = captured[index_of[spec.name]] + stored = np.zeros(high - low, dtype=np.float32) + stored[positions] = values + if not np.array_equal(stored, direct.astype(np.float32)): + differing = int((stored != direct.astype(np.float32)).sum()) + raise RuntimeError( + f"District row {spec.name} materialized directly " + "differs from the row stored from its carrier " + f"({differing} household(s)); the carrier split is " + "not exact for this surface." + ) + checked.append( + { + "target": spec.name, + "carrier": carrier, + "nonzero_households": int(np.count_nonzero(direct)), + } + ) + carrier_check.update( + { + "carriers": len(carriers), + "checked_district_rows": len(checked), + "checked_nonzero_households": sum( + entry["nonzero_households"] for entry in checked + ), + "compared": "stored CSR row vs direct materialization", + "all_equal": True, + "checked": checked, + } + ) + del target_frame, compiled_registry, target_households, sub_frame, got_ids gc.collect() stat = { "chunk": chunk_index + 1, @@ -624,7 +889,27 @@ def materialize_chunked( log(f"materialize chunk {chunk_index + 1}/{n_chunks}: {stat['wall_s']}s") del projected gc.collect() - return matrix, measure_names, compiled_specs, chunk_stats + matrix = assembler.to_csr() + keep = [ + index + for index, spec in enumerate(plan.declared) + if spec.name in compiled_names + or plan.carrier_of.get(spec.name) in (compiled_names or set()) + ] + compiled_specs = [by_name[plan.declared[index].name] for index in keep] + if len(keep) != plan.n_rows: + matrix = sparse_rows(matrix, keep) + return MaterializedSurface(matrix, compiled_specs, chunk_stats, carrier_check) + + +def sparse_rows(matrix, rows): + """The CSR rows ``rows`` of ``matrix`` (float32 data preserved).""" + + from scipy import sparse + + selected = sparse.csr_array(matrix[rows]) + selected.sort_indices() + return selected # --------------------------------------------------------------------------- @@ -668,45 +953,40 @@ def ladder_population(ladder_path: Path, geographies: list[str]): return populations -def population_measure_arrays(frame, ladder_populations, geographies: list[str]): - """Household-grain population measures as float32 arrays. +def household_sizes(frame) -> np.ndarray: + """Persons per household, aligned to the household table (float32).""" + + households = frame.table("household") + size = frame.table("person").groupby("person_household_id").size() + return households["household_id"].map(size).fillna(0).to_numpy(np.float32) + + +def population_targets(frame, ladder_populations, geographies: list[str]): + """Household-grain population marginals as sparse target rows. Population in a geography = sum over persons-in-geo of household weight; every person shares its household's geography, so the household-grain measure is household_size x 1[household geo == value]. No engine involved. A ladder cell with no supporting household is DROPPED from the - surface and returned in the fourth element — the caller decides whether - a shrunken surface is acceptable (a capped smoke) or a defect (release). + surface and returned in the third element: the caller decides whether a + shrunken surface is acceptable (a capped smoke) or a defect (release). + Returns ``(target records, CSR rows, dropped cell names)``. """ - households = frame.table("household") - persons = frame.table("person") - size = persons.groupby("person_household_id").size() - household_size = households["household_id"].map(size).fillna(0).to_numpy(np.float32) - names, arrays, values, dropped = [], [], [], [] - column_map = { - "state": ("state_fips", "state"), - "cd": ("congressional_district_geoid", "cd"), - } - for geography in geographies: - column, key = column_map[geography] - geo_values = pd.to_numeric(households[column]).to_numpy() - for value, population in sorted(ladder_populations[key].items()): - present = geo_values == value - width = 2 if geography == "state" else 4 - name = f"pop_{geography}_{value:0{width}d}" - if not present.any(): - dropped.append(name) - continue - names.append(name) - arrays.append((household_size * present).astype(np.float32)) - values.append(float(population)) - return names, arrays, values, dropped + return cd_surface.population_rows( + frame.table("household"), + household_sizes(frame), + ladder_populations, + geographies, + ) def population_target_specs(names, values): """Declare ladder population targets with the shared diagnostics hierarchy.""" + from microcosm.build.us_runtime.congressional_district_vintage import ( + CURRENT_CONGRESSIONAL_DISTRICT_PREFIX, + ) from microcosm.calibrate import TargetSpec from microcosm.calibrate.geography_constants import US_STATE_FIPS_TO_POSTAL from microcosm.calibrate.hierarchy import ( @@ -736,8 +1016,11 @@ def population_target_specs(names, values): "at-large" if district == "00" else f"district {int(district)}" ) label = f"{state} congressional {district_label}" + # The ladder's districts are the 119th Congress plan. geography = HierarchyGeography( - f"5001800US{geoid}", label, "congressional_district" + f"{CURRENT_CONGRESSIONAL_DISTRICT_PREFIX}{geoid}", + label, + "congressional_district", ) else: # pragma: no cover - names originate in population_measure_arrays raise ValueError(f"Unknown population target name {name!r}.") @@ -763,19 +1046,30 @@ def population_target_specs(names, values): return tuple(specs) +#: Household columns the lean checkpoint keeps: identity, geography, and the +#: origin tags the distinct-household ESS and spine shares need. +LEAN_HOUSEHOLD_COLUMNS = ( + "household_id", + "state_fips", + "congressional_district_geoid", + "county_fips", + "household_spine", + "household_source_id", +) +#: The checkpoint's sparse target matrix: (targets x households) CSR with +#: float32 values, row ``i`` of which is ``target_registry.json`` spec ``i``. +TARGET_MATRIX_FILENAME = "target_matrix.npz" +#: Per-target build roles, row-aligned with the registry: train or holdout, +#: geography, sigma, and a district row's state parent and populations. +TARGET_ROLES_FILENAME = "target_roles.json" + + def extract_struct_tables(frame): """Small structural/geography copies so the big frame can be freed early.""" households = frame.table("household") struct_columns = [ - column - for column in ( - "household_id", - "state_fips", - "congressional_district_geoid", - "county_fips", - ) - if column in households.columns + column for column in LEAN_HOUSEHOLD_COLUMNS if column in households.columns ] person_columns = [ column for column in PERSON_STRUCT if column in frame.table("person").columns @@ -793,33 +1087,41 @@ def extract_struct_tables(frame): def write_lean_checkpoint( struct, - admin_matrix, - admin_names, - admin_specs, - pop_names, - pop_arrays, - pop_values, + matrix, + specs, + roles: list[dict], checkpoint_dir: Path, ): - """Assemble the lean target frame and its versioned target registry.""" + """Write the lean checkpoint: structure H5, registry, sparse matrix, roles. + + ``target_frame_lean.h5`` holds only structure (household identity, + geography, origin tags and design weights; person memberships; group + ids). ``target_registry.json`` holds every target, held-out ones + included; ``target_matrix.npz`` is a (targets x households) CSR with + float32 values whose row ``i`` is registry spec ``i``; + ``target_roles.json`` is row-aligned with both. Returns + ``(h5 path, registry, digests)``. + """ + + from scipy import sparse from microcosm.calibrate import TargetRegistry from microcosm.frame import put_frame_table checkpoint_dir.mkdir(parents=True, exist_ok=True) - if isinstance(admin_matrix, np.memmap): - admin_matrix = np.array(admin_matrix) - parts = { - column: struct["household_struct"][column] - for column in struct["household_struct"].columns - } - for j, name in enumerate(admin_names): - parts[name] = admin_matrix[:, j] - for name, array in zip(pop_names, pop_arrays, strict=True): - parts[name] = array - lean_households = pd.DataFrame(parts) - del parts - gc.collect() + specs = tuple(specs) + matrix = sparse.csr_array(matrix) + n_households = len(struct["household_struct"]) + if matrix.shape != (len(specs), n_households): + raise ValueError( + f"target matrix shape {matrix.shape} does not pair with " + f"{len(specs)} targets x {n_households} households." + ) + if [role["name"] for role in roles] != [spec.name for spec in specs]: + raise ValueError("target roles are not row-aligned with the target specs.") + if matrix.data.dtype != np.float32: + matrix = sparse.csr_array(matrix, dtype=np.float32) + lean_households = struct["household_struct"].copy() lean_households["household_weight"] = struct["weights"] checkpoint_h5 = checkpoint_dir / "target_frame_lean.h5" with pd.HDFStore(checkpoint_h5, mode="w") as store: @@ -843,17 +1145,28 @@ def write_lean_checkpoint( preferred_format="fixed", ) store.put("_time_period", pd.Series([PERIOD]), format="table") - registry = TargetRegistry( - (*admin_specs, *population_target_specs(pop_names, pop_values)), - country="us", - ) + registry = TargetRegistry(specs, country="us") registry.to_json(checkpoint_dir / "target_registry.json") + matrix_sha = cd_surface.save_target_matrix( + checkpoint_dir / TARGET_MATRIX_FILENAME, matrix + ) + (checkpoint_dir / TARGET_ROLES_FILENAME).write_text(json.dumps(roles, indent=1)) log( - f"checkpoint: {checkpoint_h5.name} ({len(lean_households)} hh, " - f"{len(admin_names) + len(pop_names)} measures), target_registry.json " - f"({len(registry)} targets)" + f"checkpoint: {checkpoint_h5.name} ({n_households} hh), " + f"target_registry.json ({len(registry)} targets), " + f"{TARGET_MATRIX_FILENAME} ({matrix.shape[0]} x {matrix.shape[1]}, " + f"nnz {matrix.nnz:,})" + ) + return ( + checkpoint_h5, + registry, + { + "target_registry_sha256": _sha256(checkpoint_dir / "target_registry.json"), + "target_matrix_sha256": matrix_sha, + "target_roles_sha256": _sha256(checkpoint_dir / TARGET_ROLES_FILENAME), + "lean_h5_sha256": _sha256(checkpoint_h5), + }, ) - return checkpoint_h5, registry def load_lean_frame(checkpoint_h5: Path): @@ -892,9 +1205,71 @@ def _load_staging_frame(path: Path): return _load_base_frame(Path(path)) +def _cd_holdout_fraction(args) -> float: + """The CD holdout fraction this run draws (0 when there is no CD surface).""" + + fraction = getattr(args, "cd_holdout_fraction", None) + if fraction is None: + return ( + cd_surface.DEFAULT_STATE_CD_HOLDOUT_FRACTION + if args.soi_mode == SOI_MODE_STATE_CD + else 0.0 + ) + return float(fraction) + + +def _attach_pro_rata_populations(targets: list[dict], cd_populations: dict) -> None: + """Give each district SOI target its district and state populations. + + The pro-rata baseline allocates the state parent by these; both come + from the PUMA ladder's 119th-plan district overlap populations, so a + state's district shares sum to one. + """ + + state_population: dict[int, float] = {} + for district, population in cd_populations.items(): + state = int(district) // 100 + state_population[state] = state_population.get(state, 0.0) + float(population) + for target in targets: + geoid = target.get("congressional_district_geoid") + if target.get("family") != "irs_soi" or not geoid: + continue + district = int(geoid) + target["cd_population"] = float(cd_populations.get(district, 0.0)) + target["state_population"] = float(state_population.get(district // 100, 0.0)) + + +def _require_pro_rata_populations(targets: list[dict]) -> None: + """Refuse, before any solve, a district row the baseline cannot score.""" + + unpopulated = [ + target["name"] + for target in targets + if target.get("family") == "irs_soi" + and target.get("congressional_district_geoid") + and not ( + target.get("cd_population", 0) > 0 and target.get("state_population", 0) > 0 + ) + ] + if unpopulated: + raise SystemExit( + f"{len(unpopulated)} district SOI target(s) have no positive " + "ladder district or state population for the pro-rata " + f"baseline (e.g. {unpopulated[:3]}); refusing before the solve." + ) + + def do_materialize(args) -> None: + from scipy import sparse + families = [item.strip() for item in args.families.split(",") if item.strip()] geographies = [item.strip() for item in args.geographies.split(",") if item.strip()] + holdout_fraction = _cd_holdout_fraction(args) + if holdout_fraction and args.soi_mode != SOI_MODE_STATE_CD: + raise SystemExit( + f"--cd-holdout-fraction {holdout_fraction} needs --soi-mode " + f"{SOI_MODE_STATE_CD}; no other surface binds district SOI targets." + ) started = time.time() if args.feed_sha256: actual = _sha256(args.feed) @@ -903,12 +1278,11 @@ def do_materialize(args) -> None: f"Ledger feed sha256 mismatch: {args.feed} is {actual}, " f"expected {args.feed_sha256}." ) - registry, ri_substitutions = state_admin_specs( - args.feed, families, soi_mode=args.soi_mode - ) + surface = state_admin_surface(args.feed, families, soi_mode=args.soi_mode) + registry = surface.registry log( f"admin specs: {len(registry)} ({families}, soi_mode={args.soi_mode}); " - f"RI substitution records={len(ri_substitutions)}" + f"RI substitution records={len(surface.ri_substitutions)}" ) summary_path = _staging_summary_path(args) if not summary_path.exists(): @@ -922,15 +1296,23 @@ def do_materialize(args) -> None: ladder_sha = _sha256(args.ladder) frame = _load_staging_frame(args.staging_h5) _require_local_hours(frame, _load_json(summary_path)) + frame, sampling = cd_surface.sample_staging_frame( + frame, fraction=args.sample_fraction, seed=args.sample_seed + ) + gc.collect() log( f"loaded staging frame households={frame.n('household')} " - f"({time.time() - started:.1f}s)" + f"(rung {sampling['rung']}; {time.time() - started:.1f}s)" ) started = time.time() - matrix_path = args.checkpoint_dir / "measures_f32.mmap" args.checkpoint_dir.mkdir(parents=True, exist_ok=True) - matrix, admin_names, compiled_specs, chunk_stats = materialize_chunked( + (args.checkpoint_dir / "measures_f32.mmap").unlink(missing_ok=True) + for name in CALIBRATION_OUTPUT_FILENAMES: + (args.checkpoint_dir / name).unlink(missing_ok=True) + if getattr(args, "gate_report", None) is not None: + Path(args.gate_report).unlink(missing_ok=True) + materialized = materialize_chunked( frame, registry.specs, hh_chunk=args.hh_chunk, @@ -939,20 +1321,28 @@ def do_materialize(args) -> None: dropped_manifest_path=args.checkpoint_dir / "held_back_columns.json", summary_path=summary_path, fills_manifest_path=args.checkpoint_dir / "reviewed_null_fills.json", - matrix_path=matrix_path, ) + admin_matrix = materialized.matrix + compiled_specs = materialized.compiled_specs + chunk_stats = materialized.chunk_stats log( - f"materialized admin: {len(admin_names)} measures over " - f"{len(chunk_stats)} chunks ({time.time() - started:.1f}s)" + f"materialized admin: {len(compiled_specs)} measures over " + f"{len(chunk_stats)} chunks, nnz {admin_matrix.nnz:,} " + f"({time.time() - started:.1f}s)" ) - if len(admin_names) != len(registry): + if len(compiled_specs) != len(registry): raise SystemExit( - f"Compiled admin surface has {len(admin_names)} measures but " + f"Compiled admin surface has {len(compiled_specs)} measures but " f"{len(registry)} specs were declared; admin targets must never " "disappear silently between compile and materialization." ) - populations = ladder_population(args.ladder, geographies) - pop_names, pop_arrays, pop_values, pop_dropped = population_measure_arrays( + admin_roles = [cd_surface.target_record(spec) for spec in compiled_specs] + + populations = ladder_population(args.ladder, sorted(set(geographies) | {"cd"})) + if args.soi_mode == SOI_MODE_STATE_CD: + _attach_pro_rata_populations(admin_roles, populations["cd"]) + _require_pro_rata_populations(admin_roles) + pop_roles, pop_matrix, pop_dropped = population_targets( frame, populations, geographies ) if pop_dropped: @@ -967,26 +1357,39 @@ def do_materialize(args) -> None: "--allow-partial-geography only for capped smokes." ) log("WARNING " + message + " Continuing (--allow-partial-geography).") - log(f"population measures: {len(pop_names)} ({geographies})") + pop_specs = population_target_specs( + [record["name"] for record in pop_roles], + [record["value"] for record in pop_roles], + ) + log(f"population measures: {len(pop_specs)} ({geographies})") + roles = admin_roles + pop_roles + holdout = cd_surface.assign_target_roles(roles, fraction=holdout_fraction) + log( + f"CD holdout: {holdout['held_units']}/{holdout['eligible_units']} units, " + f"{holdout['held_targets']} district targets held out of calibration" + ) + matrix = sparse.vstack([admin_matrix, pop_matrix], format="csr") + del admin_matrix, pop_matrix struct = extract_struct_tables(frame) n_households = frame.n("household") del frame gc.collect() - write_lean_checkpoint( + _checkpoint_h5, _registry, digests = write_lean_checkpoint( struct, matrix, - admin_names, - compiled_specs, - pop_names, - pop_arrays, - pop_values, + (*compiled_specs, *pop_specs), + roles, args.checkpoint_dir, ) - del matrix, struct, pop_arrays + matrix_shape = [int(value) for value in matrix.shape] + matrix_nnz = int(matrix.nnz) + del matrix, struct gc.collect() - matrix_path.unlink(missing_ok=True) - registry_digest = _sha256(args.checkpoint_dir / "target_registry.json") + holdout_identity = { + key: holdout[key] + for key in ("unit", "salt", "fraction", "held_units", "held_targets") + } (args.checkpoint_dir / "run_identity.json").write_text( json.dumps( { @@ -994,11 +1397,23 @@ def do_materialize(args) -> None: "staging_sha256": staging_sha, "ladder_sha256": ladder_sha, "households": n_households, - "n_targets": len(admin_names) + len(pop_names), - "target_registry_sha256": registry_digest, + "n_targets": len(roles), + "target_registry_sha256": digests["target_registry_sha256"], + "target_roles_sha256": digests["target_roles_sha256"], + "lean_h5_sha256": digests["lean_h5_sha256"], + "target_matrix": { + "file": TARGET_MATRIX_FILENAME, + "sha256": digests["target_matrix_sha256"], + "format": "csr_float32", + "shape": matrix_shape, + "nnz": matrix_nnz, + }, "declared_admin_specs": len(registry), - "compiled_admin_specs": len(admin_names), + "compiled_admin_specs": len(compiled_specs), "population_cells_dropped": pop_dropped, + "soi_mode": args.soi_mode, + "cd_holdout": holdout_identity, + "sampling": sampling, }, indent=2, ) @@ -1008,11 +1423,16 @@ def do_materialize(args) -> None: { "soi_mode": args.soi_mode, "families": families, - "n_admin": len(admin_names), - "n_population": len(pop_names), + "n_admin": len(compiled_specs), + "n_population": len(pop_specs), "population_cells_dropped": pop_dropped, "hh_chunk": args.hh_chunk, "chunk_stats": chunk_stats, + "target_matrix": {"shape": matrix_shape, "nnz": matrix_nnz}, + "carrier_check": materialized.carrier_check, + "soi_surface": surface.soi_receipt, + "cd_holdout": holdout, + "sampling": sampling, "materialize_peak_rss_gb": round(rss(), 3), }, indent=2, @@ -1043,49 +1463,352 @@ def _verify_run_identity(args, *, require: bool = True) -> dict: return identity -def do_calibrate(args) -> None: - from microcosm.calibrate import ( - TargetRegistry, - calibrate, - write_calibration_diagnostics, +#: Outputs of the calibrate stage that describe one materialization; a new +#: materialize removes them so none can outlive the surface it described. +CALIBRATION_OUTPUT_FILENAMES = ( + "weights_latest.npz", + "calibration_summary.json", + "calibration_diagnostics.json", + "consumer_export.json", + "consumer_reviewed_null_fills.json", + "spine_qa.json", +) + + +#: The fixed solver settings, shared by the solve and the stamp describing it. +SOLVER_METHOD = "adam" +SOLVER_LEARNING_RATE = 0.02 +SOLVER_MASS = "conserve" + + +def _solver_settings(args) -> dict: + """The calibrate-stage settings a resume or reuse must share.""" + + return { + "method": SOLVER_METHOD, + "learning_rate": SOLVER_LEARNING_RATE, + "mass": SOLVER_MASS, + "max_weight_ratio": args.max_weight_ratio, + "target_loss_cap": args.target_loss_cap, + "l2_lambda": args.l2_lambda, + "seed": args.seed, + "epoch_batch": args.epoch_batch, + } + + +def _weights_digest(weights) -> str: + """Digest of a calibrated household weight vector (float64 bytes).""" + + return hashlib.sha256( + np.ascontiguousarray(weights, dtype=np.float64).tobytes() + ).hexdigest() + + +def _run_identity_digest(identity: dict) -> str: + """Content digest of a run identity (canonical JSON).""" + + return hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + + +def _verify_checkpoint_digests(checkpoint_dir: Path, identity: dict) -> None: + """Refuse a registry, roles file or matrix whose bytes changed.""" + + checks = ( + ( + checkpoint_dir / "target_registry.json", + identity.get("target_registry_sha256"), + ), + (checkpoint_dir / TARGET_ROLES_FILENAME, identity.get("target_roles_sha256")), + ( + checkpoint_dir / TARGET_MATRIX_FILENAME, + (identity.get("target_matrix") or {}).get("sha256"), + ), ) + if "lean_h5_sha256" in identity: + checks += ( + (checkpoint_dir / "target_frame_lean.h5", identity["lean_h5_sha256"]), + ) + for path, recorded in checks: + if not path.exists() or _sha256(path) != recorded: + raise SystemExit( + f"{path.name} changed since materialize (or the run identity " + "records none); the checkpoint and surface no longer agree. " + "Re-run --stage materialize." + ) - identity = _verify_run_identity(args) - checkpoint_h5 = args.checkpoint_dir / "target_frame_lean.h5" - registry_path = args.checkpoint_dir / "target_registry.json" - registry_sha = _sha256(registry_path) - if registry_sha != identity.get("target_registry_sha256"): + +def _require_current_artifact( + identity: dict, + *, + consumer_export: dict, + summary: dict, + spine_qa: dict | None, + out_h5: Path, + out_h5_sha256: str, + stage: str, +) -> None: + """Refuse a calibrated H5, or evidence about it, from another calibration. + + The H5's bytes must be the ones ``consumer_export.json`` records; that + export and the calibration summary must name the same materialization and + the same weights; and QA evidence, when given, must be of these bytes and + this materialization. The path is not compared: the sha binds the bytes + wherever they are reached from. + """ + + if "target_roles_sha256" not in identity: + return + digest = _run_identity_digest(identity) + if consumer_export.get("run_identity_sha256") != digest: + raise SystemExit( + "consumer_export.json was written for another materialization " + f"(or records none); re-run --stage calibrate before --stage {stage}." + ) + if consumer_export.get("out_h5_sha256") != out_h5_sha256: + raise SystemExit( + f"{out_h5} is not the calibrated H5 this checkpoint's calibrate stage " + f"wrote; re-run --stage calibrate before --stage {stage}." + ) + weights_sha = consumer_export.get("weights_sha256") + if weights_sha is None or summary.get("weights_sha256") != weights_sha: + raise SystemExit( + "The calibrated H5 and calibration_summary.json describe different " + "weights (an interrupted recalibration?); re-run --stage calibrate " + f"before --stage {stage}." + ) + if spine_qa is not None: + if spine_qa.get("run_identity_sha256") != digest: + raise SystemExit( + "spine_qa.json was written for another materialization; re-run " + f"--stage qa before --stage {stage}." + ) + if spine_qa.get("artifact_sha256") != out_h5_sha256: + raise SystemExit( + "spine_qa.json certifies other bytes than the calibrated H5 " + f"({str(spine_qa.get('artifact_sha256'))[:12]}… vs " + f"{out_h5_sha256[:12]}…); re-run --stage qa before --stage " + f"{stage}." + ) + + +def _require_current_calibration(identity: dict, summary: dict, *, stage: str) -> None: + """Refuse a calibration summary from another materialization. + + A sparse-era run identity (one that records ``target_roles_sha256``) + binds its calibration: the summary must carry that identity's digest. + Pre-sparse checkpoints stay readable (finalize's read-only path); package + refuses them separately (no sampling block). + """ + + if "target_roles_sha256" not in identity: + return + if summary.get("run_identity_sha256") != _run_identity_digest(identity): + raise SystemExit( + "calibration_summary.json does not belong to this checkpoint's " + "materialization (its run identity differs or is not recorded); " + f"re-run --stage calibrate before --stage {stage}." + ) + + +def load_checkpoint_surface(checkpoint_dir: Path, identity: dict | None = None): + """The lean frame, design weights, registry, roles and matrix of a checkpoint. + + With ``identity`` (the materialize-time run identity) the registry, the + roles and the matrix must still be the bytes materialize wrote. A + checkpoint from before the sparse matrix (dense measure columns in the + lean H5) is refused: re-run ``--stage materialize``. + Returns ``(frame, design_weights, registry, roles, matrix)``. + """ + + from microcosm.calibrate import TargetRegistry + + matrix_path = checkpoint_dir / TARGET_MATRIX_FILENAME + roles_path = checkpoint_dir / TARGET_ROLES_FILENAME + registry_path = checkpoint_dir / "target_registry.json" + if not matrix_path.exists() or not roles_path.exists(): raise SystemExit( - "target_registry.json changed since materialize; the checkpoint and " - "surface no longer agree. Re-run --stage materialize." + f"{checkpoint_dir} has no {TARGET_MATRIX_FILENAME} or " + f"{TARGET_ROLES_FILENAME}: it predates the sparse target matrix " + "(its measures are dense H5 columns). Re-run --stage materialize " + "with the current tool." ) + if identity is not None: + _verify_checkpoint_digests(checkpoint_dir, identity) registry = TargetRegistry.from_json(registry_path) - frame, design_weights = load_lean_frame(checkpoint_h5) + roles = json.loads(roles_path.read_text()) + matrix = cd_surface.load_target_matrix(matrix_path) + frame, design_weights = load_lean_frame(checkpoint_dir / "target_frame_lean.h5") n_households = frame.n("household") - if n_households != identity.get("households"): + if identity is not None and n_households != identity.get("households"): raise SystemExit( f"Lean checkpoint has {n_households} households but the run " f"identity pins {identity.get('households')}." ) - target_set = registry.to_target_set() + if [role["name"] for role in roles] != [spec.name for spec in registry.specs]: + raise SystemExit( + f"{TARGET_ROLES_FILENAME} is not row-aligned with target_registry.json." + ) + if matrix.shape != (len(registry), n_households): + raise SystemExit( + f"{TARGET_MATRIX_FILENAME} is {matrix.shape}, not " + f"{len(registry)} targets x {n_households} households." + ) + return frame, design_weights, registry, roles, matrix + + +def calibrate_surface( + frame, + target_set, + *, + epochs: int, + epoch_batch: int, + max_weight_ratio: float, + target_loss_cap: float, + l2_lambda: float, + seed: int, + warm: np.ndarray | None = None, + done: int = 0, + on_batch=None, +): + """Epoch-batched warm-start calibration of ``target_set``. + + ``target_set`` comes from ``cd_surface.calibration_target_set``, which + builds only the training targets, each a callable row of the checkpoint + CSR. Each batch calls the kernel's ``calibrate``, which compiles those + rows into its own CSR constraint matrix. Returns ``(result, epochs_done)``. + """ + + from microcosm.calibrate import calibrate + + batch = epoch_batch if epoch_batch > 0 else epochs + result = None + while done < epochs: + this_batch = min(batch, epochs - done) + batch_started = time.time() + result = calibrate( + frame, + target_set, + weight_entity="household", + method=SOLVER_METHOD, + epochs=this_batch, + learning_rate=SOLVER_LEARNING_RATE, + mass=SOLVER_MASS, + max_weight_ratio=max_weight_ratio, + target_loss_cap=target_loss_cap, + l2_lambda=l2_lambda, + seed=seed, + warm_start_weights=warm, + ) + done += this_batch + warm = result.weights.copy() + if on_batch is not None: + on_batch(warm, done) + log( + f"batch -> {done}/{epochs} ep, " + f"{time.time() - batch_started:.1f}s, " + f"loss={result.final_loss:.5f}, " + f"within10%={result.fraction_within_10pct:.2%}, " + f"ESS={result.effective_sample_size:,.0f}" + ) + return result, done + + +def _origin_columns(frame) -> dict: + """The household columns ``weight_origin_summary`` groups by, if present.""" + + households = frame.table("household") + columns = { + "spine": "household_spine", + "source_id": "household_source_id", + "state": "state_fips", + "district": "congressional_district_geoid", + } + return { + key: households[column].to_numpy() if column in households.columns else None + for key, column in columns.items() + } + + +def calibration_evidence( + *, + frame, + roles: list[dict], + matrix, + design_weights: np.ndarray, + weights: np.ndarray, + target_loss_cap: float, +) -> dict: + """The sparse-surface evidence the calibration summary adds. + + The CD holdout scored against the pro-rata baseline, and ESS over rows + and distinct households with household-weight share by spine, at the + design and the calibrated weights. + """ + + origin = _origin_columns(frame) + return { + "n_targets_on_surface": len(roles), + "n_holdout_targets": len(cd_surface.holdout_rows(roles)), + "checkpoint_matrix": { + "format": "csr_float32", + "shape": [int(x) for x in matrix.shape], + "nnz": int(matrix.nnz), + }, + "weight_origin": { + "design": cd_surface.weight_origin_summary(design_weights, **origin), + "calibrated": cd_surface.weight_origin_summary(weights, **origin), + }, + "cd_holdout": cd_surface.score_cd_holdout( + roles, + matrix, + design_weights=design_weights, + final_weights=weights, + cap=target_loss_cap, + ), + } + + +def do_calibrate(args) -> None: + from microcosm.calibrate import write_calibration_diagnostics + + identity = _verify_run_identity(args) + frame, design_weights, registry, roles, matrix = load_checkpoint_surface( + args.checkpoint_dir, identity + ) + n_households = frame.n("household") + target_set = cd_surface.calibration_target_set( + roles, matrix, n_households, specs=registry.specs + ) log( - f"calibrate: households={n_households}, targets={len(target_set)}, " - f"design_total={design_weights.sum():,.0f}" + f"calibrate: households={n_households}, targets={len(target_set)} " + f"trained + {len(roles) - len(target_set)} held out, " + f"nnz={matrix.nnz:,}, design_total={design_weights.sum():,.0f}" ) + stamp = _run_identity_digest(identity) + settings = _solver_settings(args) resume_npz = args.checkpoint_dir / "weights_latest.npz" warm, done = None, 0 if args.resume and resume_npz.exists(): saved = np.load(resume_npz) - saved_identity = ( - str(saved["staging_sha256"]) if "staging_sha256" in saved else None + saved_stamp = ( + str(saved["run_identity_sha256"]) + if "run_identity_sha256" in saved + else None ) - if saved_identity is not None and saved_identity != identity.get( - "staging_sha256" - ): + saved_settings = ( + json.loads(str(saved["solver_settings"])) + if "solver_settings" in saved + else None + ) + if saved_stamp != stamp or saved_settings != settings: raise SystemExit( - "weights_latest.npz was produced against a different staging " - "H5; refusing to warm-start from a foreign checkpoint." + "weights_latest.npz was calibrated for a different " + "materialization (staging, surface, holdout or sample) or " + "under different solver settings, or predates the stamp; " + "refusing to warm-start from it. Delete it to recalibrate." ) warm, done = saved["weights"], int(saved["epochs_done"]) if len(warm) != n_households: @@ -1096,57 +1819,61 @@ def do_calibrate(args) -> None: log(f"RESUME from {done} epochs") if done >= args.epochs: summary_path = args.checkpoint_dir / "calibration_summary.json" - if summary_path.exists(): + previous = _load_json(summary_path) + if ( + previous.get("run_identity_sha256") == stamp + and previous.get("solver_settings") == settings + and previous.get("weights_sha256") == _weights_digest(warm) + ): log( f"calibration already complete at {done} epochs and " "the calibration summary exists; nothing to do (delete " "weights_latest.npz to recalibrate)." ) _write_calibrated_artifact( - args, np.asarray(warm, dtype=np.float64), identity + args, np.asarray(warm, dtype=np.float64), identity, settings ) return raise SystemExit( f"weights_latest.npz reports {done} epochs (>= --epochs " - f"{args.epochs}) but calibration_summary.json is missing. " - "Delete the checkpoint to recalibrate, or raise --epochs." - ) - batch = args.epoch_batch if args.epoch_batch > 0 else args.epochs - result = None - started = time.time() - while done < args.epochs: - this_batch = min(batch, args.epochs - done) - batch_started = time.time() - result = calibrate( - frame, - target_set, - weight_entity="household", - method="adam", - epochs=this_batch, - learning_rate=0.02, - mass="conserve", - max_weight_ratio=args.max_weight_ratio, - target_loss_cap=args.target_loss_cap, - l2_lambda=args.l2_lambda, - seed=args.seed, - warm_start_weights=warm, + f"{args.epochs}) but calibration_summary.json is missing or " + "belongs to another materialization, solver settings or weights " + "(a summary written before the weights digest). Delete " + "weights_latest.npz or run without --resume to recalibrate, or " + "raise --epochs." ) - done += this_batch - warm = result.weights.copy() + + # A solve replaces every output describing the calibration: none of the + # previous ones may outlive it if this run stops part way. + for name in CALIBRATION_OUTPUT_FILENAMES: + if name != "weights_latest.npz": + (args.checkpoint_dir / name).unlink(missing_ok=True) + + def save(weights: np.ndarray, epochs_done: int) -> None: np.savez( resume_npz, - weights=warm, - epochs_done=done, + weights=weights, + epochs_done=epochs_done, initial_weights=design_weights, staging_sha256=np.str_(identity["staging_sha256"]), + run_identity_sha256=np.str_(stamp), + solver_settings=np.str_(json.dumps(settings, sort_keys=True)), ) - log( - f"batch -> {done}/{args.epochs} ep, " - f"{time.time() - batch_started:.1f}s, " - f"loss={result.final_loss:.5f}, " - f"within10%={result.fraction_within_10pct:.2%}, " - f"ESS={result.effective_sample_size:,.0f}" - ) + + started = time.time() + result, done = calibrate_surface( + frame, + target_set, + epochs=args.epochs, + epoch_batch=args.epoch_batch, + max_weight_ratio=args.max_weight_ratio, + target_loss_cap=args.target_loss_cap, + l2_lambda=args.l2_lambda, + seed=args.seed, + warm=warm, + done=done, + on_batch=save, + ) if result.problem.skipped: skipped = [getattr(item, "name", str(item)) for item in result.problem.skipped] @@ -1156,16 +1883,21 @@ def do_calibrate(args) -> None: "measures or the targets before shipping." ) summary = { + "run_identity_sha256": stamp, + "solver_settings": settings, + "weights_sha256": _weights_digest(result.weights), "households": n_households, "n_targets": result.problem.n_targets, "families": args.families, "geographies": args.geographies, + "soi_mode": identity.get("soi_mode"), "matrix_format": result.options["matrix_format"], "matrix_shape": [int(x) for x in result.problem.matrix.shape], "matrix_nnz": int(result.problem.matrix.nnz), "epochs": args.epochs, "epoch_batch": args.epoch_batch, "max_weight_ratio": args.max_weight_ratio, + "target_loss_cap": args.target_loss_cap, "l2_lambda": args.l2_lambda, "seed": args.seed, "initial_loss": round(result.initial_loss, 6), @@ -1177,13 +1909,23 @@ def do_calibrate(args) -> None: "mass_conserved_ratio": round( float(result.weights.sum()) / float(design_weights.sum()), 6 ), + "sampling": identity.get("sampling"), + **calibration_evidence( + frame=frame, + roles=roles, + matrix=matrix, + design_weights=design_weights, + weights=np.asarray(result.weights, dtype=np.float64), + target_loss_cap=args.target_loss_cap, + ), "total_wall_seconds": round(time.time() - started, 1), "peak_rss_gb": round(rss(), 3), } _write_calibrated_artifact( - args, np.asarray(result.weights, dtype=np.float64), identity + args, np.asarray(result.weights, dtype=np.float64), identity, settings ) + holdout = summary["cd_holdout"] outcome = write_calibration_diagnostics( result, args.checkpoint_dir / "calibration_diagnostics.json", @@ -1192,12 +1934,15 @@ def do_calibrate(args) -> None: "dataset_role": "non_default_local_area", "families": args.families, "geographies": args.geographies, + "soi_mode": identity.get("soi_mode"), "epochs": args.epochs, "epoch_batch": args.epoch_batch, "total_wall_seconds": summary["total_wall_seconds"], "peak_rss_gb": summary["peak_rss_gb"], "ess_fraction": summary["ess_fraction"], "mass_conserved_ratio": summary["mass_conserved_ratio"], + "n_holdout_targets": summary["n_holdout_targets"], + "sampling_rung": (identity.get("sampling") or {}).get("rung"), }, ) summary["calibration_diagnostics"] = ( @@ -1221,10 +1966,45 @@ def do_calibrate(args) -> None: f"calibrate stage complete: loss={summary['final_loss']}, " f"within10%={summary['fraction_within_10pct']:.2%}, " f"diagnostics={outcome.status}" + + ( + f"; CD holdout {holdout['n_targets']} targets: calibrated " + f"{holdout['calibrated']['mean_abs_rel_error']:.2%} vs pro-rata " + f"{holdout['pro_rata_baseline']['mean_abs_rel_error']:.2%} mean " + "abs rel error" + if holdout.get("n_targets") + else "" + ) + ) + + +def _resample_like_materialize(frame, identity: dict): + """Re-draw the development rung materialize drew, or refuse. + + A full-rung (or pre-sampling) identity returns the frame unchanged. A + sampled one re-draws with the recorded fraction and seed and must select + exactly the recorded households. + """ + + sampling = identity.get("sampling") or {} + if not sampling.get("sampled"): + return frame + sampled, receipt = cd_surface.sample_staging_frame( + frame, + fraction=float(sampling["sample_fraction"]), + seed=int(sampling["sample_seed"]), ) + recorded = sampling.get("selected_household_ids_sha256") + if receipt.get("selected_household_ids_sha256") != recorded: + raise SystemExit( + "Re-drawing the recorded development rung selected different " + "households than materialize did; refusing to attach weights." + ) + return sampled -def _write_calibrated_artifact(args, weights: np.ndarray, identity: dict) -> None: +def _write_calibrated_artifact( + args, weights: np.ndarray, identity: dict, settings: dict +) -> None: """Write the consumer-ready calibrated H5 with verified attachment. The weights attach by verified household-id vector equality between the @@ -1245,6 +2025,7 @@ def _write_calibrated_artifact(args, weights: np.ndarray, identity: dict) -> Non frame = _load_staging_frame(args.staging_h5) _require_local_hours(frame, _load_json(_staging_summary_path(args))) + frame = _resample_like_materialize(frame, identity) staging_ids = frame.table("household")["household_id"].to_numpy() with pd.HDFStore(args.checkpoint_dir / "target_frame_lean.h5", mode="r") as store: lean_ids = store["household"]["household_id"].to_numpy() @@ -1286,6 +2067,10 @@ def _write_calibrated_artifact(args, weights: np.ndarray, identity: dict) -> Non json.dumps( { "out_h5": str(Path(args.out_h5).resolve()), + "out_h5_sha256": _sha256(args.out_h5), + "run_identity_sha256": _run_identity_digest(identity), + "weights_sha256": _weights_digest(weights), + "solver_settings": settings, "staging_sha256": identity.get("staging_sha256"), "held_back_formula_owned": dropped, "held_back_total": sum(len(v) for v in dropped.values()), @@ -1412,7 +2197,12 @@ def do_qa(args) -> None: if entry["person_weight"] else 0.0 ) + # Recorded, not re-verified: finalize and package verify the identity. + qa_identity = _load_json(args.checkpoint_dir / "run_identity.json") payload = { + "run_identity_sha256": ( + _run_identity_digest(qa_identity) if qa_identity else None + ), "period": PERIOD, "variable": "ssi", "artifact": str(Path(args.out_h5).resolve()), @@ -1581,6 +2371,93 @@ def finalize_reviewed_limitations( return list(deduped.values()) +def state_cd_reviewed_limitations(materialize_rss: dict) -> list[dict]: + """Reviewed limitations a ``state_cd`` surface adds to the register.""" + + if materialize_rss.get("soi_mode") != SOI_MODE_STATE_CD: + return [] + surface = materialize_rss.get("soi_surface") or {} + holdout = materialize_rss.get("cd_holdout") or {} + excluded = surface.get("excluded_cd_states") or {} + exclusion = ( + " District rows of state(s) " + + ", ".join(sorted(excluded)) + + " are excluded: " + + "; ".join(f"{state}: {why}" for state, why in sorted(excluded.items())) + if excluded + else " No state's district rows are excluded for the mapping." + ) + return [ + { + "id": "cd_soi_117th_plan_population_crosswalk", + "status": "reviewed_construction", + "reason": ( + "The SOI congressional-district file (22incd.csv, TY2022) is " + "tabulated on the 117th-Congress plan. Its district rows are " + "mapped onto the households' 119th-plan districts by the " + "packaged 2020-block population crosswalk (built from the " + "block plan registry since #1043), so each 119th district " + "target assumes returns spread with population inside every " + "117th/119th intersection." + exclusion + ), + "treatment": ( + "A household 117th-plan district column " + "(congressional_district_geoid__117th_congress, from the " + "location v1 block draw with the plan registry attached) lets " + "these targets bind as exact block sums with no crosswalk." + ), + "excluded_states": surface.get("excluded_cd_states"), + "crosswalk": surface.get("crosswalk"), + "calibration_blocker": False, + }, + { + "id": "cd_soi_one_vintage_per_state_concept", + "status": "reviewed_construction", + "reason": surface.get("vintage_rule_description"), + "rebase_factor_by_measure": surface.get("rebase_factor_by_measure"), + "level_bridges": surface.get("level_bridges"), + "level_bridge_factor_by_measure": surface.get( + "level_bridge_factor_by_measure" + ), + "dropped": surface.get("dropped"), + "calibration_blocker": False, + }, + { + "id": "cd_soi_defective_district_columns_excluded", + "status": "reviewed_exclusion", + "reason": ( + "District-file measures that read the wrong IRS column are " + "off the surface at both geographies (microcosm#1038)." + ), + "measures": surface.get("defective_cd_file_measures"), + "calibration_blocker": False, + }, + { + "id": "cd_holdout_sealed", + "status": "reviewed_construction", + "reason": ( + f"{holdout.get('held_targets')} district SOI targets in " + f"{holdout.get('held_units')} (state x concept family) units " + f"(fraction {holdout.get('fraction')}, salt " + f"{holdout.get('salt')}) never reach the calibrator; they are " + "scored against a pro-rata baseline in calibration_summary.json " + "(cd_holdout) and gate_summary.json (gates.cd_holdout)." + ), + "calibration_blocker": False, + }, + { + "id": "cd_soi_sigma_absent", + "status": "reviewed_data_gap", + "reason": ( + "The pinned facts feed carries no uncertainty for any IRS SOI " + "fact, so no target carries sigma and the loss is unchanged " + "(fixed-scale capped relative error)." + ), + "calibration_blocker": False, + }, + ] + + def _local_hours_gate(frame, staging_summary: dict): audit = staging_summary.get("reviewed_engine_input_nulls") if not isinstance(audit, list) or not all(isinstance(item, dict) for item in audit): @@ -1613,6 +2490,9 @@ def do_finalize(args) -> None: "--stage calibrate first." ) identity = _verify_run_identity(args) + _require_current_calibration(identity, diagnostics, stage="finalize") + if "target_roles_sha256" in identity: + _verify_checkpoint_digests(args.checkpoint_dir, identity) ladder_sha = _sha256(args.ladder) if ladder_sha != identity.get("ladder_sha256"): raise SystemExit( @@ -1622,7 +2502,10 @@ def do_finalize(args) -> None: ) materialize_rss = _load_json(args.checkpoint_dir / "materialize_rss.json") registry_path = args.checkpoint_dir / "target_registry.json" - if registry_path.is_file(): + roles_path = args.checkpoint_dir / TARGET_ROLES_FILENAME + if roles_path.is_file(): + targets = json.loads(roles_path.read_text()) + elif registry_path.is_file(): registry = TargetRegistry.from_json(registry_path) targets = [ {"name": spec.name, "family": spec.family} for spec in registry.specs @@ -1641,6 +2524,15 @@ def do_finalize(args) -> None: if not args.out_h5.exists(): raise SystemExit(f"Calibrated H5 not found: {args.out_h5}.") hours_artifact_sha = _sha256(args.out_h5) + _require_current_artifact( + identity, + consumer_export=consumer_export, + summary=diagnostics, + spine_qa=spine_qa or None, + out_h5=args.out_h5, + out_h5_sha256=hours_artifact_sha, + stage="finalize", + ) frame = _load_staging_frame(args.out_h5) local_hours_gate = _local_hours_gate(frame, staging_summary) households = frame.table("household") @@ -1667,7 +2559,13 @@ def do_finalize(args) -> None: breakdown: dict[str, int] = {} for target in targets: name = target["name"] - if name.startswith("pop_state"): + if target.get("role") == cd_surface.ROLE_HOLDOUT: + key = "cd_holdout" + elif name.startswith("irs_soi") and ( + target.get("geography_level") == cd_surface.GEOGRAPHY_CD + ): + key = "soi_congressional_district" + elif name.startswith("pop_state"): key = "population_state" elif name.startswith("pop_cd"): key = "population_cd" @@ -1710,7 +2608,8 @@ def do_finalize(args) -> None: # per-target error bars) are maintainer-adjudicated surface # policy (#398-class), recorded here rather than invented. "passed": bool( - diagnostics.get("final_loss", 1.0) < args.target_loss_cap + diagnostics.get("final_loss", 1.0) + < diagnostics.get("target_loss_cap", args.target_loss_cap) and diagnostics.get("final_loss", 1.0) < diagnostics.get("initial_loss", 0.0) and abs(mass - 1.0) < 1e-3 @@ -1728,9 +2627,35 @@ def do_finalize(args) -> None: "n_targets": len(targets), "n_admin": materialize_rss.get("n_admin"), "n_population": materialize_rss.get("n_population"), + "n_trained": diagnostics.get("n_targets"), + "n_holdout": diagnostics.get("n_holdout_targets"), "breakdown": breakdown, + "target_matrix": materialize_rss.get("target_matrix"), + "rung": (identity.get("sampling") or {}).get("rung"), }, }, + "cd_holdout": { + # Report-only: held-out district targets never reach the solve; + # the comparison with the pro-rata baseline is evidence, not a + # bound (d487 asks whether district fidelity beats pro-rata). + "passed": True, + "report_only": True, + "detail": { + key: value + for key, value in (diagnostics.get("cd_holdout") or {}).items() + if key != "targets" + }, + }, + "weight_origin": { + "passed": True, + "report_only": True, + "note": ( + "Household-weight share by spine and Kish ESS over rows and " + "over distinct households (household_spine, " + "household_source_id), at design and calibrated weights." + ), + "detail": diagnostics.get("weight_origin"), + }, "input_coverage": { # Enforced upstream: the staging driver's donor-coverage gate # hard-fails before transfer, so reaching finalize means it held. @@ -1783,6 +2708,7 @@ def do_finalize(args) -> None: } limitations = finalize_reviewed_limitations(staging_summary, diagnostics, spine_qa) + limitations += state_cd_reviewed_limitations(materialize_rss) hard_failures = [ name for name in ( @@ -1884,6 +2810,29 @@ def _require_recorded_soi_mode(materialize_rss: dict) -> str: return soi_mode +def _require_full_rung(identity: dict) -> dict: + """Refuse to package a development rung, or a run that cannot show its rung. + + ``--sample-fraction`` below 1 is for development runs (DESIGN.md + "Production US stacked spine"); a release is always full scale. + """ + + sampling = identity.get("sampling") + if not isinstance(sampling, dict) or "sampled" not in sampling: + raise SystemExit( + "run_identity.json records no sampling block, so packaging cannot " + "establish that the surface was materialized at full scale. " + "Re-run --stage materialize with the current tool." + ) + if sampling.get("sampled") or sampling.get("rung") != "f100": + raise SystemExit( + f"The checkpoint was materialized on the {sampling.get('rung')} " + "development rung; a release is packaged only at full scale " + "(--sample-fraction 1)." + ) + return sampling + + def _require_stored_inputs(calibrated_h5: Path) -> dict[str, object]: """Refuse an artifact that stores a model input the installed engine lacks. @@ -1946,6 +2895,8 @@ def do_package(args) -> dict: # Before any release directory exists: a refused smoke leaves nothing. staging_orchestration = _require_uncapped_staging(staging_summary) soi_mode = _require_recorded_soi_mode(materialize_rss) + _require_full_rung(identity) + _require_current_calibration(identity, diagnostics, stage="package") calibrated_h5 = Path(args.out_h5) if not calibrated_h5.exists(): raise SystemExit(f"Calibrated H5 not found: {calibrated_h5}.") @@ -1953,15 +2904,10 @@ def do_package(args) -> dict: # as build.built_with_model_package, so the artifact may store no model # input that engine does not define. A refused artifact leaves nothing. stored_inputs_gate = _require_stored_inputs(calibrated_h5) - - code = _repo_code_identity(args.allow_dirty) - timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") - release_id = f"{RELEASE_ID_PREFIX}-{code['sha']}-{timestamp}" - release_dir = args.out / "releases" / release_id - release_dir.mkdir(parents=True, exist_ok=True) - log("hashing calibrated H5 …") h5_sha = _sha256(calibrated_h5) + # The H5 and its evidence must come from this checkpoint's calibration; + # a refused artifact leaves no release directory behind. # The gate report certifies specific artifact bytes: the QA probe # recorded the sha it loaded plain. Packaging different bytes (a # recalibrate without re-running qa+finalize) is refused. @@ -1992,6 +2938,22 @@ def do_package(args) -> dict: f"certified ({h5_sha[:12]}… vs {str(qa_sha)[:12]}…). Re-run " "--stage qa and --stage finalize against the current artifact." ) + _require_current_artifact( + identity, + consumer_export=consumer_export, + summary=diagnostics, + spine_qa=spine_qa, + out_h5=calibrated_h5, + out_h5_sha256=h5_sha, + stage="package", + ) + + code = _repo_code_identity(args.allow_dirty) + timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + release_id = f"{RELEASE_ID_PREFIX}-{code['sha']}-{timestamp}" + release_dir = args.out / "releases" / release_id + release_dir.mkdir(parents=True, exist_ok=True) + gates = gate_report.get("gates") hours_gate = gates.get("hours_worked_signal") if isinstance(gates, dict) else None if not isinstance(hours_gate, dict) or hours_gate.get("passed") is not True: @@ -2069,7 +3031,9 @@ def _version(package: str) -> str: "unchanged." ), "staging": LEGACY_STAGING_REFRESH_RECIPE, - "release": release_refresh_recipe(soi_mode), + "release": release_refresh_recipe( + soi_mode, (identity.get("cd_holdout") or {}).get("fraction") + ), "publish": ( "tools/publish_release.sh --no-latest " f"--artifact-root --repo-id {HF_REPO_ID}" @@ -2118,6 +3082,12 @@ def _version(package: str) -> str: }, "materialize": { "soi_mode": soi_mode, + "soi_surface_counts": (materialize_rss.get("soi_surface") or {}).get( + "counts" + ), + "target_matrix": identity.get("target_matrix"), + "cd_holdout": identity.get("cd_holdout"), + "sampling": identity.get("sampling"), "peak_rss_gb": materialize_rss.get("materialize_peak_rss_gb"), "hh_chunk": materialize_rss.get("hh_chunk"), "engine_pass": ( @@ -2325,16 +3295,43 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: choices=SOI_MODES, default=DEFAULT_SOI_MODE, help=( - "State SOI target surface for --stage materialize. 'state' " - "(default) is Build O's contract: every state-geography SOI spec " - "outside the congressional-district file. 'totals' drops every " + "SOI target surface for --stage materialize. 'state' (default) " + "is Build O's contract: every state-geography SOI spec outside " + "the congressional-district file. 'totals' drops every " "soi_fiscal_distribution spec (no state AGI, income-tax or EITC " - "total); 'full' keeps every state-bearing spec, including the " - "district file, and needs a much larger dense admin matrix " - "(contents and sizes in docs/us-acs-local-soi-target-surface.md). " - "Later stages use the mode the checkpoint recorded." + "total); 'full' keeps every state-bearing spec, both vintages " + "of each state concept included. 'state_cd' adds the district " + "file's district rows to 'state' with one vintage per state " + "concept, and holds a hash-assigned block of them out " + "(docs/us-acs-local-soi-target-surface.md). Later stages use the " + "mode the checkpoint recorded." + ), + ) + parser.add_argument( + "--cd-holdout-fraction", + type=float, + default=None, + help=( + "Share of (state x SOI concept family) district-target units held " + "out of calibration and scored against a pro-rata baseline " + f"(state_cd only; default {cd_surface.DEFAULT_STATE_CD_HOLDOUT_FRACTION}" + f" there, max {cd_surface.MAX_CD_HOLDOUT_FRACTION}; 0 disables)." + ), + ) + parser.add_argument( + "--sample-fraction", + type=float, + default=1.0, + help=( + "Development rung: sample whole staging households at this " + f"fraction ({sorted(cd_surface.SAMPLE_RUNG_TOKENS)}), stratified " + "by spine x district and normalized to each spine's full mass. " + "Recorded in the run identity; package refuses anything below 1." ), ) + parser.add_argument( + "--sample-seed", type=int, default=cd_surface.DEFAULT_SAMPLE_SEED + ) parser.add_argument("--epochs", type=int, default=800) parser.add_argument("--epoch-batch", type=int, default=400) parser.add_argument("--max-weight-ratio", type=float, default=5.0) @@ -2404,7 +3401,19 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: args.out_summary = args.out_h5.with_suffix(".summary.json") if args.gate_report is None: args.gate_report = args.checkpoint_dir / "gate_summary.json" + try: + cd_surface.rung_token(args.sample_fraction) + except ValueError as error: + parser.error(str(error)) + if args.cd_holdout_fraction is not None and not ( + 0.0 <= args.cd_holdout_fraction <= cd_surface.MAX_CD_HOLDOUT_FRACTION + ): + parser.error( + "--cd-holdout-fraction must be in " + f"[0, {cd_surface.MAX_CD_HOLDOUT_FRACTION}]." + ) args.stages = stages + return args diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index a3b28053a..9a5cefeb2 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -536,6 +536,9 @@ def _acs_local_release_argv( "families": OptionFlag("--families", str), "geographies": OptionFlag("--geographies", str), "soi_mode": OptionFlag("--soi-mode", str), + "cd_holdout_fraction": OptionFlag("--cd-holdout-fraction", float), + "sample_fraction": OptionFlag("--sample-fraction", float), + "sample_seed": OptionFlag("--sample-seed", int), "epochs": OptionFlag("--epochs", int), "epoch_batch": OptionFlag("--epoch-batch", int), "max_weight_ratio": OptionFlag("--max-weight-ratio", float), diff --git a/tools/us_acs_local_cd_surface.py b/tools/us_acs_local_cd_surface.py new file mode 100644 index 000000000..edd37055f --- /dev/null +++ b/tools/us_acs_local_cd_surface.py @@ -0,0 +1,1767 @@ +"""Congressional-district calibration pieces of the ACS local-area release. + +``tools/build_us_acs_local_release.py`` imports this module for everything +that makes a district-level SOI surface affordable and checkable: + +- **The ``state_cd`` SOI surface** (:func:`state_cd_soi_surface`): the TY2022 + congressional-district SOI file's district rows plus the state surface, + keeping one vintage per state concept. Historic Table 2 supplies a state + concept's level wherever it has that concept; the district file then + supplies only each district's share of its state, and its district rows + are rebased to sum to the Historic Table 2 parent. Concepts only the + district file has keep the district file's own state row as their parent. +- **The sparse target matrix** (:class:`SparseTargetAssembler`): one CSR row + per target over the household weight vector, assembled chunk by chunk. + District SOI rows never exist as dense columns: each distinct SOI concept + is materialized once without geography (a *carrier* column) and every + district row is that carrier restricted to the district's households, + which is exactly what the materializer's district mask computes, because a + tax unit's geography is its household's. +- **The CD holdout** (:func:`assign_target_roles`): hash-assigned blocks of + district targets that never reach the calibrator, scored against a naive + pro-rata baseline (:func:`score_cd_holdout`). +- **Effective sample size over distinct households** and weight share by + spine (:func:`weight_origin_summary`). +- **Sampling rungs** for development runs (:func:`sample_staging_frame`). + +Everything here is engine-free; the tool supplies the engine pass. +""" + +from __future__ import annotations + +import hashlib +import math +from collections import Counter, defaultdict +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse + +# --------------------------------------------------------------------------- +# Surface constants +# --------------------------------------------------------------------------- + +#: Ledger record-set specs from the congressional-district SOI file. +CD_FILE_RECORD_SET_SPEC_PREFIX = "irs_soi.congressional_district_" +#: How the ``state_cd`` surface resolves the two vintages of a state concept. +STATE_CD_VINTAGE_RULE = "state_cd.historic_table_2_level_cd_file_shares.v1" + +#: District-file measures that read the wrong IRS column (microcosm#1038). +#: PR #1040 excludes the same columns in the compiler; until it lands this +#: register keeps them off the ``state_cd`` surface at both geographies. +STATE_CD_DEFECTIVE_CD_FILE_MEASURES: Mapping[str, str] = { + "limited_state_local_taxes_amount": ( + "22incd.csv column A18425 is state and local income taxes (Schedule A " + "line 5a), not the limited SALT deduction (A18460); the district " + "file sums to 1.98x Historic Table 2 (microcosm#1038)." + ), + "limited_state_local_taxes_returns": ( + "22incd.csv column N18425 counts state and local income-tax " + "deductions, not the limited SALT deduction (microcosm#1038)." + ), + "premium_tax_credit_returns": ( + "22incd.csv column N85530 is the additional Medicare tax (Form 8959 " + "line 24), not the premium tax credit (N85770) (microcosm#1038)." + ), +} +#: District-file measures with no state parent in either vintage, so their +#: district rows could not nest in a bound state target. +STATE_CD_UNPARENTED_CD_MEASURES: Mapping[str, str] = { + "tax_filer_individual_count": ( + "The compiler drops this measure at state and national geography " + "(fiscal_targets._soi_reference_from_fact), so no state target exists " + "for its district rows to reconcile to." + ), +} +#: District-file-only concepts (Historic Table 2 lacks them) and the sibling +#: concept whose two state levels bridge theirs onto the Historic Table 2 +#: basis. The district file's own state levels sit below Historic Table 2's +#: for two reasons: it covers only the returns processed in its window (about +#: 1.7% fewer), and the feed stamps it one tax year late, so its amounts are +#: aged one year less. A sibling both files carry, in the same state, estimates +#: that gap, on the assumption that the measure shares the sibling's ratio: +#: counts bridge by ``return_count`` (counts are not aged), amounts by the +#: closest amount Historic Table 2 carries. The factor band below applies to +#: the sibling. The bridge is stamp-invariant: if the stamp is corrected +#: (#1030) both levels move together. +STATE_CD_LEVEL_BRIDGES: Mapping[str, str] = { + "charitable_amount": "itemized_deductions_amount", + "charitable_returns": "return_count", + "interest_paid_deduction_amount": "itemized_deductions_amount", + "interest_paid_deduction_returns": "return_count", + "qualified_business_income_deduction_amount": "adjusted_gross_income", + "qualified_business_income_deduction_returns": "return_count", +} +#: States whose district SOI rows stay off the surface, with the evidence. +#: Empty since #1043 rebuilt the 117th->119th crosswalk from the block plan +#: registry: the packaged crosswalk before it (sha256 c7cb040b...) carried +#: North Carolina's 2016 plan instead of the 2019 plan the 117th Congress used, +#: and 37.0% of NC's population mapped to a different 119th district, so NC's +#: district rows were held off until then. +STATE_CD_EXCLUDED_CD_STATES: Mapping[str, str] = {} +#: How far a state's Historic Table 2 / district-file ratio may sit from its +#: measure's median across states before the district file's data for that +#: (state, concept) is refused. The ratio mixes a common part (the file's +#: coverage and aging gap, or a measure-wide rebase such as taxable interest's +#: x2.47 onto Table 4.3), which the median absorbs, with a state-specific part. +#: A state-specific part beyond 25% means the two publications disagree about +#: that state (Hawaii's dividends differ 2.4x from the typical gap on the +#: pinned feed), so the district file's shares, or a bridge built on that +#: sibling, are not trusted: the block's district rows and any bridged state +#: row are dropped and recorded; the Historic Table 2 state row stays. +STATE_CD_FACTOR_BAND = 1.25 +#: The district file's own state row must equal the sum of its district rows +#: to this relative tolerance wherever both are present (the receipt records +#: the largest gap; the feed-gated contract test pins it on the pinned feed). +#: A gap means a column or crosswalk defect inside the file (#1038 class), +#: which no median absorbs. +STATE_CD_INTERNAL_RTOL = 1e-6 +#: The packaged crosswalk the exclusions above were reviewed against. A test +#: pins it, so regenerating the crosswalk forces a second look at them. +STATE_CD_REVIEWED_CROSSWALK_SHA256 = ( + "a347303fff3fea145f758488c43cb94355df5a8cb5e512552b2ee3791d2aa233" +) + +#: Relative tolerance for "district rows sum to their state parent". The +#: rebase makes it exact up to float64 rounding. +RECONCILIATION_RTOL = 1e-9 + +SIGMA_BASIS_FEED = "feed_standard_error" +SIGMA_BASIS_NONE = "not_provided_by_feed" + +# --------------------------------------------------------------------------- +# Holdout constants +# --------------------------------------------------------------------------- + +#: Salt of the CD holdout hash. Changing it re-draws every held unit. +CD_HOLDOUT_SALT = "microcosm.us_acs_local.cd_holdout.v1" +#: The held unit: every district target of one SOI concept family in one +#: state. A single district (or a single district target) is not a fair +#: holdout: with its state total and sibling districts trained it is pinned +#: by adding up. Concept families also group measures that add up to each +#: other (EITC by number of children sums to the EITC total). +CD_HOLDOUT_UNIT = "state_x_soi_concept_family" +DEFAULT_STATE_CD_HOLDOUT_FRACTION = 0.1 +MAX_CD_HOLDOUT_FRACTION = 0.5 +ROLE_TRAIN = "train" +ROLE_HOLDOUT = "holdout" + +# --------------------------------------------------------------------------- +# Sampling constants: the rung tokens of tools/build_us_multispine_pool.py +# (DESIGN.md "Production US stacked spine" names f001, f010 and f100). +# --------------------------------------------------------------------------- + +SAMPLE_RUNG_TOKENS: Mapping[float, str] = { + 0.01: "f001", + 0.04: "f004", + 0.10: "f010", + 0.25: "f025", + 1.0: "f100", +} +DEFAULT_SAMPLE_SEED = 578 + +GEOGRAPHY_STATE = "state" +GEOGRAPHY_CD = "congressional_district" + + +def _release_tool(): + """The sibling release tool (``tools/`` is on ``sys.path`` via the ACS tool).""" + + import build_us_fiscal_refresh_release as release_tool + + return release_tool + + +# --------------------------------------------------------------------------- +# SOI concept identity +# --------------------------------------------------------------------------- + + +def _metadata(spec) -> Mapping[str, str]: + return spec.metadata + + +def geography_level(spec) -> str | None: + return _metadata(spec).get("ledger_geography_level") + + +def record_set_spec_id(spec) -> str: + return str(_metadata(spec).get("ledger_layout_record_set_spec_id") or "") + + +def is_cd_file_spec(spec) -> bool: + return record_set_spec_id(spec).startswith(CD_FILE_RECORD_SET_SPEC_PREFIX) + + +def soi_materializer_semantics(spec) -> tuple: + """Everything the SOI slice reads from a spec, except its geography. + + ``_base_simulation_household_columns`` builds an ``irs_soi`` column from + exactly these fields plus ``state_fips`` and + ``congressional_district_geoid``. Two specs with equal semantics + therefore materialize the same household column up to their geography + masks. + """ + + tool = _release_tool() + metadata = _metadata(spec) + return ( + spec.family, + spec.entity, + spec.filter, + metadata.get("source_variable", metadata.get("variable")), + repr(tool._as_bound(metadata["agi_lower_bound"])), + repr(tool._as_bound(metadata["agi_upper_bound"])), + metadata.get("taxable_only") == "true", + metadata.get("filing_status"), + tool._soi_eitc_child_count_filter(metadata), + tool._soi_requires_positive_eitc_filter(metadata), + metadata.get("itemized_only") == "true", + metadata.get("measure_mode") == "indicator_sum", + tool._unsupported_soi_ledger_filters(metadata), + ) + + +def soi_concept_identity(spec) -> tuple[str, ...]: + """The published concept a spec measures, independent of vintage. + + Historic Table 2 and the district file name the same concept with the + same ``source_measure_id``; the AGI band, filing status and EITC child + count separate the Historic Table 2 band rows from the all-income rows. + """ + + tool = _release_tool() + metadata = _metadata(spec) + return ( + str(metadata.get("source_measure_id")), + repr(tool._as_bound(metadata["agi_lower_bound"])), + repr(tool._as_bound(metadata["agi_upper_bound"])), + str(metadata.get("filing_status", "")).lower(), + str(tool._soi_eitc_child_count_filter(metadata)), + ) + + +def soi_concept_family(source_measure_id: str) -> str: + """The holdout family of an SOI measure: its concept without count/amount. + + Every EITC measure is one family, because the per-child-count rows add up + to the EITC total and holding one out would leave it pinned by the rest. + """ + + measure = str(source_measure_id) + if measure.startswith("eitc"): + return "eitc" + for suffix in ("_amount", "_returns", "_claims", "_count"): + if measure.endswith(suffix): + return measure[: -len(suffix)] + return measure + + +# --------------------------------------------------------------------------- +# The state_cd surface +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class StateCdSurface: + """The ``state_cd`` SOI specs and the receipt that explains them.""" + + specs: tuple + receipt: dict + + +def _crosswalk_district_sets(crosswalk: pd.DataFrame): + """(source districts per state, target districts per state) as SSDD.""" + + source: dict[str, set[str]] = defaultdict(set) + target: dict[str, set[str]] = defaultdict(set) + for source_id, target_id in zip( + crosswalk["source_geography_id"].astype(str), + crosswalk["target_geography_id"].astype(str), + strict=True, + ): + source[source_id[-4:-2]].add(source_id[-4:]) + target[target_id[-4:-2]].add(target_id[-4:]) + return dict(source), dict(target) + + +def _factor_summary(values: Sequence[float]) -> dict[str, float | int]: + array = np.asarray(values, dtype=np.float64) + return { + "n": int(array.size), + "min": float(array.min()), + "median": float(np.median(array)), + "max": float(array.max()), + } + + +def state_cd_soi_surface( + soi_specs: Sequence, + *, + state_surface_predicate: Callable[[object], bool], + crosswalk: pd.DataFrame, + crosswalk_sha256: str | None = None, + level_bridges: Mapping[str, str] = STATE_CD_LEVEL_BRIDGES, + factor_band: float = STATE_CD_FACTOR_BAND, +) -> StateCdSurface: + """Select and reconcile the ``state_cd`` SOI surface. + + Args: + soi_specs: Every compiled ``irs_soi`` spec (the ``full`` surface). + state_surface_predicate: The ``state`` mode's predicate; its specs + (Historic Table 2) are the state surface. + crosswalk: The 117th->119th CD crosswalk the compiler used; it names + each state's source-plan and current-plan districts. + crosswalk_sha256: Recorded in the receipt. + level_bridges: District-file-only measure -> sibling measure whose + Historic Table 2 / district-file ratio in the same state lifts + the district file's state level onto the Historic Table 2 basis. + factor_band: Largest allowed ratio between a state's Historic Table 2 + / district-file ratio and its measure's median across states + (either direction); see ``STATE_CD_FACTOR_BAND``. + + Returns: + The surface in a fixed order (Historic Table 2 state rows, district + file state rows, district rows) and its receipt. + + Raises: + ValueError: If a district concept is incomplete in a state, if its + district rows disagree on what they materialize, if a parent's + materialization differs from its children's, or if a parent + cannot absorb its children (zero sum under a nonzero parent, or + a sign conflict). + """ + + source_districts, target_districts = _crosswalk_district_sets(crosswalk) + at_large_source_states = sorted( + state for state, districts in source_districts.items() if len(districts) == 1 + ) + dropped: Counter = Counter() + dropped_examples: dict[str, str] = {} + + def drop(spec, reason: str) -> None: + dropped[reason] += 1 + dropped_examples.setdefault(reason, spec.name) + + ht2_state: list = [] + cd_file_state: list = [] + cd_rows: list = [] + for spec in soi_specs: + level = geography_level(spec) + if state_surface_predicate(spec): + ht2_state.append(spec) + elif is_cd_file_spec(spec) and level == GEOGRAPHY_STATE: + cd_file_state.append(spec) + elif is_cd_file_spec(spec) and level == GEOGRAPHY_CD: + cd_rows.append(spec) + elif level == GEOGRAPHY_CD: + drop(spec, "historic_table_2_at_large_proxy_district_row") + else: + drop(spec, f"not_state_cd_geography:{level}") + + ht2_by_key = {} + for spec in ht2_state: + key = (spec.metadata["state_fips"], soi_concept_identity(spec)) + if key in ht2_by_key: + raise ValueError( + f"Two Historic Table 2 state specs measure {key}: " + f"{ht2_by_key[key].name} and {spec.name}." + ) + ht2_by_key[key] = spec + cd_file_state_by_key = {} + for spec in cd_file_state: + key = (spec.metadata["state_fips"], soi_concept_identity(spec)) + if key in cd_file_state_by_key: + raise ValueError( + f"Two district-file state specs measure {key}: " + f"{cd_file_state_by_key[key].name} and {spec.name}." + ) + cd_file_state_by_key[key] = spec + + # Each state concept both files carry: its Historic Table 2 / district-file + # ratio, judged against the measure's median ratio across states. + vintage_ratio: dict[tuple, float] = {} + for key in ht2_by_key.keys() & cd_file_state_by_key.keys(): + ht2_value = float(ht2_by_key[key].value) + cd_value = float(cd_file_state_by_key[key].value) + vintage_ratio[key] = ( + ht2_value / cd_value + if cd_value != 0.0 and (ht2_value > 0) == (cd_value > 0) + else math.nan + ) + ratios_by_identity: dict[tuple, list[float]] = defaultdict(list) + for (_state, identity), ratio in vintage_ratio.items(): + if math.isfinite(ratio): + ratios_by_identity[identity].append(ratio) + median_ratio = { + identity: float(np.median(values)) + for identity, values in ratios_by_identity.items() + } + out_of_band: list[dict] = [] + + def within_band( + key: tuple, + ratio: float, + median: float, + basis: str, + measure: str, + *, + verdict: str, + ) -> bool: + relative = ratio / median if math.isfinite(ratio) and median else math.nan + ok = math.isfinite(relative) and 1 / factor_band <= relative <= factor_band + if not ok: + out_of_band.append( + { + "state_fips": key[0], + "measure": measure, + "concept": key[1][0], + "basis": basis, + "verdict": verdict, + "ratio": ratio, + "median_ratio": median, + "relative": relative, + } + ) + return ok + + kept_cd_by_key: dict[tuple, list] = defaultdict(list) + for spec in cd_rows: + metadata = spec.metadata + measure = str(metadata.get("source_measure_id")) + state = str(metadata["state_fips"]) + if measure in STATE_CD_DEFECTIVE_CD_FILE_MEASURES: + drop(spec, f"defective_cd_file_column:{measure}") + elif measure in STATE_CD_UNPARENTED_CD_MEASURES: + drop(spec, f"no_state_parent:{measure}") + elif state in at_large_source_states: + drop(spec, "at_large_on_source_plan") + elif state in STATE_CD_EXCLUDED_CD_STATES: + drop(spec, f"excluded_state:{state}") + else: + kept_cd_by_key[(state, soi_concept_identity(spec))].append(spec) + + # Each Historic-Table-2-parented block's rebase factor (parent over its + # district rows' sum), judged against its concept's median across states. + block_factor: dict[tuple, float] = {} + for key, block in kept_cd_by_key.items(): + parent = ht2_by_key.get(key) + child_sum = math.fsum(float(spec.value) for spec in block) + if parent is not None and child_sum != 0.0: + block_factor[key] = float(parent.value) / child_sum + factors_by_identity: dict[tuple, list[float]] = defaultdict(list) + for (_state, identity), factor in block_factor.items(): + if math.isfinite(factor) and factor > 0: + factors_by_identity[identity].append(factor) + median_factor = { + identity: float(np.median(values)) + for identity, values in factors_by_identity.items() + } + + # One basis for every state level: a concept Historic Table 2 lacks keeps + # the district file's state row, scaled by a sibling's two levels. + bridged: dict[str, object] = {} + bridge_factors: dict[str, list[float]] = defaultdict(list) + for key, spec in sorted(cd_file_state_by_key.items()): + state, identity = key + measure = identity[0] + if key in ht2_by_key or measure in STATE_CD_DEFECTIVE_CD_FILE_MEASURES: + continue + sibling = level_bridges.get(measure) + if sibling is None: + raise ValueError( + f"District-file-only concept {measure!r} has no level bridge; " + "register a sibling measure in STATE_CD_LEVEL_BRIDGES before " + "binding it." + ) + sibling_key = (state, (sibling, *identity[1:])) + ht2_sibling = ht2_by_key.get(sibling_key) + cd_sibling = cd_file_state_by_key.get(sibling_key) + if ht2_sibling is None or cd_sibling is None or float(cd_sibling.value) == 0: + raise ValueError( + f"Level bridge {measure} -> {sibling} in state {state} needs " + "the sibling's state total in both files." + ) + factor = float(ht2_sibling.value) / float(cd_sibling.value) + # A sign conflict is a data defect, not a level disagreement: refuse + # it before the band, which only judges finite positive ratios. + if not (math.isfinite(factor) and factor > 0.0): + raise ValueError( + f"Level bridge {measure} -> {sibling} in state {state} has " + f"factor {factor}." + ) + # The sibling is trusted for bridging only if both of the band's + # verdicts on it pass: its two state levels, and (where it has district + # rows) the rebase of its own district block. + trusted = within_band( + sibling_key, + factor, + median_ratio.get(sibling_key[1], math.nan), + f"level_bridge:{measure}", + measure, + verdict="sibling_vintage_ratio", + ) + if sibling_key in block_factor: + trusted = ( + within_band( + sibling_key, + block_factor[sibling_key], + median_factor.get(sibling_key[1], math.nan), + f"level_bridge:{measure}", + measure, + verdict="sibling_rebase_factor", + ) + and trusted + ) + if not trusted: + continue + metadata = dict(spec.metadata) + metadata.update( + { + "state_cd_vintage_rule": STATE_CD_VINTAGE_RULE, + "state_cd_cd_file_value": repr(float(spec.value)), + "state_cd_level_bridge_sibling": ht2_sibling.name, + "state_cd_level_bridge_cd_file_sibling": cd_sibling.name, + "state_cd_level_bridge_factor": repr(factor), + } + ) + value = float(spec.value) * factor + bridged[spec.name] = replace( + spec, value=value, signed=value < 0.0, metadata=metadata + ) + bridge_factors[measure].append(factor) + + rebased: dict[str, object] = {} + internal_gaps: list[float] = [] + factors_by_measure: dict[str, list[float]] = defaultdict(list) + parent_basis_counts: Counter = Counter() + kept_cd_file_parents: set[str] = set() + for key in sorted(kept_cd_by_key): + state, identity = key + children = sorted( + kept_cd_by_key[key], + key=lambda spec: spec.metadata["congressional_district_geoid"], + ) + districts = [spec.metadata["congressional_district_geoid"] for spec in children] + expected = sorted(target_districts.get(state, ())) + if districts != expected: + raise ValueError( + f"District concept {identity} in state {state} covers " + f"{districts}, expected the current plan's {expected}; " + "rebasing an incomplete set would misstate every district." + ) + semantics = {soi_materializer_semantics(spec) for spec in children} + if len(semantics) != 1: + raise ValueError( + f"District rows of {identity} in state {state} materialize " + f"differently: {sorted(map(repr, semantics))}." + ) + (child_semantics,) = semantics + parent = ht2_by_key.get(key) + basis = "historic_table_2" + if parent is None: + parent = cd_file_state_by_key.get(key) + basis = "cd_file_state_total_bridged" + if parent is None: + raise ValueError( + f"District concept {identity} in state {state} has no " + "state parent in either vintage." + ) + if parent.name not in bridged: + for child in children: + drop(child, f"level_bridge_out_of_band:{identity[0]}") + continue + parent = bridged[parent.name] + kept_cd_file_parents.add(parent.name) + if soi_materializer_semantics(parent) != child_semantics: + raise ValueError( + f"State parent {parent.name} materializes differently from " + f"its district rows ({identity}, state {state})." + ) + child_sum = math.fsum(float(spec.value) for spec in children) + parent_value = float(parent.value) + if child_sum == 0.0: + if parent_value != 0.0: + raise ValueError( + f"District rows of {identity} in state {state} sum to zero " + f"under a nonzero parent {parent.name}={parent_value}." + ) + factor = 1.0 + else: + factor = parent_value / child_sum + if factor < 0.0: + raise ValueError( + f"District rows of {identity} in state {state} sum to " + f"{child_sum}, opposite in sign to parent {parent.name}=" + f"{parent_value}." + ) + # The district file must agree with itself before its shares are used. + cd_state = cd_file_state_by_key.get(key) + if cd_state is not None: + cd_state_value = float(cd_state.value) + gap = abs(cd_state_value - child_sum) + internal_gaps.append(gap / max(abs(cd_state_value), 1.0)) + if gap > STATE_CD_INTERNAL_RTOL * max(abs(cd_state_value), 1.0): + raise ValueError( + f"District-file state row {cd_state.name}={cd_state_value} " + f"differs from the sum of its district rows ({child_sum}) " + f"for {identity} in state {state}: a column or crosswalk " + "defect inside the district file." + ) + # An all-zero block agrees with a zero parent whatever the median. + if ( + basis == "historic_table_2" + and child_sum != 0.0 + and not within_band( + key, + factor, + median_factor.get(identity, math.nan), + "rebase", + identity[0], + verdict="rebase_factor", + ) + ): + for child in children: + drop(child, f"rebase_out_of_band:{identity[0]}") + continue + factors_by_measure[identity[0]].append(factor) + parent_basis_counts[basis] += len(children) + for child in children: + metadata = dict(child.metadata) + metadata.update( + { + "state_cd_vintage_rule": STATE_CD_VINTAGE_RULE, + "state_cd_parent_target_name": parent.name, + "state_cd_parent_basis": basis, + "state_cd_parent_value": repr(parent_value), + "state_cd_cd_file_value": repr(float(child.value)), + "state_cd_cd_file_state_sum": repr(child_sum), + "state_cd_rebase_factor": repr(factor), + } + ) + value = float(child.value) * factor + rebased[child.name] = replace( + child, + value=value, + signed=value < 0.0, + metadata=metadata, + ) + + kept_cd_file_state = [] + for spec in cd_file_state: + measure = str(spec.metadata.get("source_measure_id")) + key = (spec.metadata["state_fips"], soi_concept_identity(spec)) + if measure in STATE_CD_DEFECTIVE_CD_FILE_MEASURES: + drop(spec, f"defective_cd_file_column:{measure}") + elif key in ht2_by_key: + drop(spec, "second_vintage_of_state_concept") + elif spec.name not in bridged: + drop(spec, f"level_bridge_out_of_band:{measure}") + else: + kept_cd_file_state.append(bridged[spec.name]) + missing_parents = kept_cd_file_parents - {spec.name for spec in kept_cd_file_state} + if missing_parents: + raise ValueError( + f"District-file state parents {sorted(missing_parents)[:5]} were " + "dropped while their district rows were kept." + ) + + cd_specs = tuple(rebased[spec.name] for spec in cd_rows if spec.name in rebased) + specs = (*ht2_state, *kept_cd_file_state, *cd_specs) + # Fail closed on the invariant the whole surface rests on. + unreconciled = [ + block for block in state_parent_reconciliation(specs) if not block["ok"] + ] + if unreconciled: + raise ValueError( + f"{len(unreconciled)} district block(s) do not add up to a bound " + f"state parent, e.g. {unreconciled[:2]}." + ) + with_sigma = sum(1 for spec in specs if getattr(spec, "se", None) is not None) + receipt = { + "vintage_rule": STATE_CD_VINTAGE_RULE, + "vintage_rule_description": ( + "Historic Table 2 (TY2022) is the single vintage of every state " + "concept it carries. District-file (22incd.csv) district rows " + "keep only their within-state shares and are rebased to sum to " + "that parent. Concepts only the district file carries keep its " + "own state row as their parent, lifted onto the Historic Table 2 " + "basis by a sibling concept's ratio of the two files' state " + "totals in the same state (STATE_CD_LEVEL_BRIDGES)." + ), + "counts": { + "historic_table_2_state": len(ht2_state), + "cd_file_state": len(kept_cd_file_state), + "congressional_district": len(cd_specs), + "congressional_district_by_parent_basis": dict(parent_basis_counts), + "total": len(specs), + }, + "dropped": dict(sorted(dropped.items())), + "dropped_examples": dict(sorted(dropped_examples.items())), + "defective_cd_file_measures": dict(STATE_CD_DEFECTIVE_CD_FILE_MEASURES), + "unparented_cd_measures": dict(STATE_CD_UNPARENTED_CD_MEASURES), + "excluded_cd_states": dict(STATE_CD_EXCLUDED_CD_STATES), + "at_large_on_source_plan": at_large_source_states, + "internal_consistency": { + "rtol": STATE_CD_INTERNAL_RTOL, + "blocks_checked": len(internal_gaps), + "max_relative_gap": max(internal_gaps, default=0.0), + }, + "factor_band": { + "tolerance": factor_band, + "relative_to": "the measure's median Historic Table 2 / " + "district-file ratio across states", + "out_of_band": sorted( + out_of_band, key=lambda row: (row["concept"], row["state_fips"]) + ), + }, + "level_bridges": dict(level_bridges), + "level_bridge_factor_by_measure": { + measure: _factor_summary(values) + for measure, values in sorted(bridge_factors.items()) + }, + "rebase_factor_by_measure": { + measure: _factor_summary(values) + for measure, values in sorted(factors_by_measure.items()) + }, + "crosswalk": { + "source_plan": "117th_congress", + "target_plan": "119th_congress", + "method": "packaged 2020-block population crosswalk", + "sha256": crosswalk_sha256, + }, + "sigma": { + "targets_with_sigma": with_sigma, + "targets_without_sigma": len(specs) - with_sigma, + "note": ( + "The pinned feed carries no uncertainty field for any fact; " + "IRS SOI tables are administrative. target_roles.json records " + "sigma and sigma_basis for every target, null wherever a spec " + "carries none." + ), + }, + } + return StateCdSurface(specs=specs, receipt=receipt) + + +def state_parent_reconciliation( + specs: Sequence, *, rtol: float = RECONCILIATION_RTOL +) -> list[dict]: + """Every (state parent, district rows) block and whether it adds up. + + A block is the district specs naming one ``state_cd_parent_target_name``. + Returns one entry per block with the parent value, the district sum and + ``ok`` (relative difference within ``rtol``). + """ + + by_name = {spec.name: spec for spec in specs} + blocks: dict[str, list] = defaultdict(list) + for spec in specs: + parent = spec.metadata.get("state_cd_parent_target_name") + if parent is not None: + blocks[parent].append(spec) + report = [] + for parent_name, children in sorted(blocks.items()): + parent = by_name.get(parent_name) + child_sum = math.fsum(float(spec.value) for spec in children) + parent_value = None if parent is None else float(parent.value) + scale = max(abs(parent_value or 0.0), 1.0) + report.append( + { + "parent": parent_name, + "parent_bound": parent is not None, + "parent_value": parent_value, + "district_sum": child_sum, + "n_districts": len(children), + "ok": parent is not None + and abs(child_sum - parent_value) <= rtol * scale, + } + ) + return report + + +# --------------------------------------------------------------------------- +# Sparse target matrix assembly +# --------------------------------------------------------------------------- + + +def is_cd_soi_spec(spec) -> bool: + return spec.family == "irs_soi" and bool( + spec.metadata.get("congressional_district_geoid") + ) + + +def _without_geography(metadata: Mapping[str, str]) -> dict[str, str]: + return { + key: value + for key, value in metadata.items() + if key not in ("state_fips", "congressional_district_geoid") + } + + +@dataclass(frozen=True) +class CarrierPlan: + """How declared specs map onto what the engine pass materializes. + + ``engine_specs`` are handed to the materializer: every declared spec + that is not a district SOI row, then one carrier per distinct district + SOI concept. ``split`` lists, per district SOI row, its declared row + index, state (``NO_STATE_MASK`` when the spec carries no ``state_fips``, + which the materializer then does not mask on), district and carrier + measure. + """ + + declared: tuple + engine_specs: tuple + direct_rows: tuple[tuple[int, str], ...] + split: tuple[tuple[int, int, int, str], ...] + carrier_of: Mapping[str, str] + + @property + def n_rows(self) -> int: + return len(self.declared) + + +CARRIER_PREFIX = "__state_cd_carrier__" +#: ``split`` state value of a district row whose spec has no ``state_fips``. +NO_STATE_MASK = -1 + + +def plan_carriers(specs: Sequence) -> CarrierPlan: + """Factor district SOI rows into per-concept carriers. + + Each carrier is the first district row of its concept with the state and + district keys removed, so it materializes the concept for every + household; a district row is the carrier masked to the district's + households (``SparseTargetAssembler`` applies the mask). + """ + + declared = tuple(specs) + names = [spec.name for spec in declared] + if len(set(names)) != len(names): + raise ValueError("Declared target specs must have unique names.") + direct_rows: list[tuple[int, str]] = [] + split: list[tuple[int, int, int, str]] = [] + carriers: dict[tuple, object] = {} + carrier_of: dict[str, str] = {} + for index, spec in enumerate(declared): + if not is_cd_soi_spec(spec): + direct_rows.append((index, spec.measure)) + continue + semantics = soi_materializer_semantics(spec) + carrier = carriers.get(semantics) + if carrier is None: + carrier_name = f"{CARRIER_PREFIX}{len(carriers):05d}" + # A carrier is an engine-pass column, never a calibration target: + # it drops the hierarchy, whose target id names the district row. + carrier = replace( + spec, + name=carrier_name, + measure=carrier_name, + metadata=_without_geography(spec.metadata), + hierarchy=None, + ) + carriers[semantics] = carrier + district = str(spec.metadata["congressional_district_geoid"]) + state = spec.metadata.get("state_fips") + split.append( + ( + index, + NO_STATE_MASK if state is None else int(state), + int(district), + carrier.measure, + ) + ) + carrier_of[spec.name] = carrier.measure + engine_specs = tuple(spec for spec in declared if not is_cd_soi_spec(spec)) + tuple( + carriers.values() + ) + return CarrierPlan( + declared=declared, + engine_specs=engine_specs, + direct_rows=tuple(direct_rows), + split=tuple(split), + carrier_of=carrier_of, + ) + + +def carrier_check_specs(plan: CarrierPlan, district_codes: np.ndarray) -> tuple: + """District rows the first chunk also materializes directly, as a check. + + One per carrier: the row whose district has the most households in the + chunk (ties to the lower row). ``materialize_chunked`` compares each + directly materialized row with the row the assembler stored for it and + refuses any difference. + """ + + counts = Counter(int(code) for code in district_codes) + best: dict[str, tuple[int, int]] = {} + for index, _state, district, carrier in plan.split: + score = counts.get(district, 0) + current = best.get(carrier) + if current is None or score > current[0]: + best[carrier] = (score, index) + return tuple(plan.declared[index] for _score, index in best.values()) + + +class SparseTargetAssembler: + """Accumulate a (targets x households) CSR matrix one chunk at a time. + + Values are stored as float32, exactly as the dense checkpoint stored its + columns, and a zero after float32 rounding is not stored, so the matrix + is the dense float32 matrix with its zeros removed. + """ + + def __init__(self, n_rows: int, n_households: int) -> None: + self.n_rows = int(n_rows) + self.n_households = int(n_households) + self._rows: list[np.ndarray] = [] + self._cols: list[np.ndarray] = [] + self._data: list[np.ndarray] = [] + + def add_column(self, row: int, low: int, values: np.ndarray, *, name: str) -> None: + """Add one target's values for households ``low .. low+len(values)``.""" + + values32 = np.asarray(values, dtype=np.float32) + if not np.isfinite(values32).all(): + bad = int((~np.isfinite(values32)).sum()) + raise ValueError( + f"Target {name!r} materialized {bad} non-finite household " + "value(s); refusing to assemble the target matrix." + ) + nonzero = np.flatnonzero(values32) + if not len(nonzero): + return + self._rows.append(np.full(len(nonzero), row, dtype=np.int32)) + self._cols.append((nonzero + low).astype(np.int32)) + self._data.append(values32[nonzero]) + + def add_masked( + self, + row: int, + low: int, + carrier: np.ndarray, + positions: np.ndarray, + *, + name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Add a carrier restricted to chunk ``positions`` (the district's). + + Returns the chunk positions and float32 values actually stored. + """ + + values32 = np.asarray(carrier[positions], dtype=np.float32) + if not np.isfinite(values32).all(): + raise ValueError( + f"Target {name!r} derives non-finite values from its carrier." + ) + keep = values32 != 0 + stored = (positions[keep], values32[keep]) + if not keep.any(): + return stored + self._rows.append(np.full(int(keep.sum()), row, dtype=np.int32)) + self._cols.append((positions[keep] + low).astype(np.int32)) + self._data.append(values32[keep]) + return stored + + def extend(self, other_rows: sparse.csr_array, row_offset: int) -> None: + """Append already-assembled CSR rows at ``row_offset``.""" + + coo = sparse.coo_array(other_rows) + self._rows.append((coo.row + row_offset).astype(np.int32)) + self._cols.append(coo.col.astype(np.int32)) + self._data.append(coo.data.astype(np.float32)) + + def to_csr(self) -> sparse.csr_array: + if self._data: + rows = np.concatenate(self._rows) + cols = np.concatenate(self._cols) + data = np.concatenate(self._data) + else: + rows = np.empty(0, dtype=np.int32) + cols = np.empty(0, dtype=np.int32) + data = np.empty(0, dtype=np.float32) + self._rows, self._cols, self._data = [], [], [] + coo = sparse.coo_array( + (data, (rows, cols)), shape=(self.n_rows, self.n_households) + ) + stored = coo.nnz + coo.sum_duplicates() + if coo.nnz != stored: + raise ValueError( + f"{stored - coo.nnz} (target, household) cell(s) were written " + "twice; a chunk or a district split overlapped." + ) + matrix = sparse.csr_array(coo) + matrix.sort_indices() + if matrix.data.dtype != np.float32: + raise TypeError("target matrix data must stay float32.") + return matrix + + +def district_positions(district_codes: np.ndarray) -> dict[int, np.ndarray]: + """Chunk positions of each district's households, in household order.""" + + codes = np.asarray(district_codes, dtype=np.int64) + order = np.argsort(codes, kind="stable") + sorted_codes = codes[order] + boundaries = np.flatnonzero(np.diff(sorted_codes)) + 1 + starts = np.concatenate(([0], boundaries)) + stops = np.concatenate((boundaries, [len(codes)])) + return { + int(sorted_codes[start]): order[start:stop] + for start, stop in zip(starts, stops, strict=True) + if stop > start + } + + +def split_carriers_into( + assembler: SparseTargetAssembler, + plan: CarrierPlan, + carrier_columns: Mapping[str, np.ndarray], + *, + low: int, + state_codes: np.ndarray, + district_codes: np.ndarray, + capture: set[int] | frozenset[int] = frozenset(), +) -> dict[int, tuple[np.ndarray, np.ndarray]]: + """Add every district SOI row of one chunk from its carrier column. + + Returns, for each declared row index in ``capture``, the chunk positions + and float32 values the assembler stored for it. + """ + + positions_by_district = district_positions(district_codes) + states = np.asarray(state_codes, dtype=np.int64) + empty = np.empty(0, dtype=np.int64) + captured: dict[int, tuple[np.ndarray, np.ndarray]] = {} + for index, state, district, carrier in plan.split: + positions = positions_by_district.get(district, empty) + if len(positions) and state != NO_STATE_MASK: + positions = positions[states[positions] == state] + stored = assembler.add_masked( + index, + low, + carrier_columns[carrier], + positions, + name=plan.declared[index].name, + ) + if index in capture: + captured[index] = stored + return captured + + +def district_row_parents( + plan: CarrierPlan, compiled_names: Iterable[str] +) -> tuple[dict[int, tuple[str, bool]], list[int]]: + """Each district row's directly materialized state parent, for the check. + + Returns ``{row index: (parent name, strict)}`` and the rows with no parent, + over the rows whose carrier compiled. + + - A ``state_cd`` row names its parent (``state_cd_parent_target_name``), + and its block partitions that parent's state, so the check is strict: + the stored district rows, laid side by side, must equal the parent's + column on every household. A ``state_cd`` row without a compiled parent + is refused, since the check would silently skip it. + - Any other district row falls back to a compiled state row of the same + materializer semantics in the same state (in ``full``, the district + file's own state total). Its block need not cover the state, so the + check compares the row's stored values with that column on the row's + own households only. + """ + + compiled = set(compiled_names) + same_concept_state_row: dict[tuple, str] = {} + for index, _measure in plan.direct_rows: + spec = plan.declared[index] + metadata = spec.metadata + if ( + spec.family == "irs_soi" + and spec.name in compiled + and metadata.get("state_fips") + and not metadata.get("congressional_district_geoid") + ): + same_concept_state_row.setdefault( + (soi_materializer_semantics(spec), str(metadata["state_fips"])), + spec.name, + ) + parents: dict[int, tuple[str, bool]] = {} + unchecked: list[int] = [] + for index, _state, _district, carrier in plan.split: + if carrier not in compiled: + continue # the row is not materialized, so not on the surface + spec = plan.declared[index] + metadata = spec.metadata + named = metadata.get("state_cd_parent_target_name") + if "state_cd_vintage_rule" in metadata and named not in compiled: + raise ValueError( + f"state_cd district row {spec.name} has no compiled state " + f"parent ({named!r}); its per-chunk check would be vacuous." + ) + if named in compiled: + parents[index] = (named, True) + continue + fallback = same_concept_state_row.get( + (soi_materializer_semantics(spec), str(metadata.get("state_fips"))) + ) + if fallback is None: + unchecked.append(index) + else: + parents[index] = (fallback, False) + return parents, unchecked + + +def population_rows( + households: pd.DataFrame, + household_size: np.ndarray, + ladder_populations: Mapping[str, Mapping[int, float]], + geographies: Sequence[str], +) -> tuple[list[dict], sparse.csr_array, list[str]]: + """Ladder population marginals as sparse rows (household size x 1[geo]). + + Returns (target records, CSR rows, dropped cell names). A ladder cell + with no supporting household is dropped and named, as before. + """ + + column_map = { + "state": ("state_fips", "state", 2, GEOGRAPHY_STATE), + "cd": ("congressional_district_geoid", "cd", 4, GEOGRAPHY_CD), + } + size32 = np.asarray(household_size, dtype=np.float32) + records: list[dict] = [] + rows: list[np.ndarray] = [] + cols: list[np.ndarray] = [] + data: list[np.ndarray] = [] + dropped: list[str] = [] + for geography in geographies: + column, key, width, level = column_map[geography] + codes = pd.to_numeric(households[column]).to_numpy() + for value, population in sorted(ladder_populations[key].items()): + name = f"pop_{geography}_{value:0{width}d}" + present = np.flatnonzero(codes == value) + if not len(present): + dropped.append(name) + continue + present = present[size32[present] != 0] + row = len(records) + rows.append(np.full(len(present), row, dtype=np.int32)) + cols.append(present.astype(np.int32)) + data.append(size32[present]) + record = { + "name": name, + "value": float(population), + "source": "us_puma_ladder_2020", + "family": "census_population_ladder", + "geography_level": level, + "state_fips": f"{int(value) // (100 if geography == 'cd' else 1):02d}", + "congressional_district_geoid": ( + f"{int(value):04d}" if geography == "cd" else None + ), + } + records.append(record) + matrix = sparse.csr_array( + ( + np.concatenate(data) if data else np.empty(0, np.float32), + ( + np.concatenate(rows) if rows else np.empty(0, np.int32), + np.concatenate(cols) if cols else np.empty(0, np.int32), + ), + ), + shape=(len(records), len(households)), + ) + matrix.sort_indices() + return records, matrix, dropped + + +def save_target_matrix(path: Path, matrix: sparse.csr_array) -> str: + """Write the CSR (float32 data) uncompressed; return its sha256.""" + + matrix = sparse.csr_array(matrix) + if matrix.data.dtype != np.float32: + raise TypeError("target matrix data must be float32.") + sparse.save_npz(path, matrix, compressed=False) + 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_target_matrix(path: Path) -> sparse.csr_array: + matrix = sparse.csr_array(sparse.load_npz(path)) + matrix.sort_indices() + return matrix + + +def csr_row_measure(matrix: sparse.csr_array, row: int, n: int): + """A callable measure returning one CSR row as a dense float64 vector. + + The calibrate kernel compiles targets through ``build_constraint_matrix``, + which reads each target's row and keeps its nonzeros; a callable measure + lets it read the row straight from the checkpoint CSR (the UK rowwise + precedent, ``uk_runtime.local_rowwise._rowwise_target_set``). Only one + dense row exists at a time. float32 values widen exactly to float64, as + the dense checkpoint's float32 columns did. + """ + + start, stop = int(matrix.indptr[row]), int(matrix.indptr[row + 1]) + indices = matrix.indices[start:stop] + data = matrix.data[start:stop] + + def measure(frame) -> np.ndarray: + values = np.zeros(n, dtype=np.float64) + values[indices] = data + return values + + # Calibration diagnostics name a callable measure by its qualified name; + # name the matrix row rather than this closure. + measure.__qualname__ = f"target_matrix_row[{row}]" + return measure + + +# --------------------------------------------------------------------------- +# Target records, roles and the calibration target set +# --------------------------------------------------------------------------- + + +def target_record(spec) -> dict: + """The target_roles.json record for one admin spec (metadata subset).""" + + metadata = spec.metadata + level = metadata.get("ledger_geography_level") + se = getattr(spec, "se", None) + record = { + "name": spec.name, + "entity": spec.entity, + "period": spec.period, + "value": float(spec.value), + "source": spec.source or "ledger_feed", + "family": spec.family, + "geography_level": level, + "state_fips": metadata.get("state_fips"), + "congressional_district_geoid": metadata.get("congressional_district_geoid"), + "source_measure_id": metadata.get("source_measure_id"), + "sigma": None if se is None else float(se), + "sigma_basis": SIGMA_BASIS_NONE if se is None else SIGMA_BASIS_FEED, + } + for key in ( + "state_cd_parent_target_name", + "state_cd_parent_basis", + "state_cd_rebase_factor", + "state_cd_cd_file_value", + ): + if key in metadata: + record[key] = metadata[key] + return record + + +def cd_holdout_unit_key(record: Mapping) -> str | None: + """The holdout unit of a district SOI target, or None if not eligible.""" + + if record.get("family") != "irs_soi": + return None + if record.get("geography_level") != GEOGRAPHY_CD: + return None + if not record.get("state_cd_parent_target_name"): + return None + return ( + f"{record['state_fips']}|" + f"{soi_concept_family(str(record.get('source_measure_id')))}" + ) + + +def assign_target_roles( + records: Sequence[dict], + *, + fraction: float, + salt: str = CD_HOLDOUT_SALT, +) -> dict: + """Set ``role`` and ``holdout_unit`` on every record; return the receipt. + + Only district SOI targets with a bound state parent are eligible; every + other target trains. A unit is held out by + :func:`microcosm.build.holdout.hash_holdout_unit` on its key. + """ + + from microcosm.build.holdout import hash_holdout_unit + + if not 0.0 <= float(fraction) <= MAX_CD_HOLDOUT_FRACTION: + raise ValueError( + f"CD holdout fraction must be in [0, {MAX_CD_HOLDOUT_FRACTION}], " + f"got {fraction!r}." + ) + units: dict[str, bool] = {} + for record in records: + unit = cd_holdout_unit_key(record) + record["holdout_unit"] = unit + if unit is None: + record["role"] = ROLE_TRAIN + continue + if unit not in units: + units[unit] = hash_holdout_unit(unit, fraction=fraction, salt=salt) + record["role"] = ROLE_HOLDOUT if units[unit] else ROLE_TRAIN + held_units = sorted(unit for unit, held in units.items() if held) + return { + "unit": CD_HOLDOUT_UNIT, + "salt": salt, + "fraction": float(fraction), + "eligible_units": len(units), + "held_units": len(held_units), + "held_unit_keys": held_units, + "eligible_targets": sum( + 1 for record in records if record["holdout_unit"] is not None + ), + "held_targets": sum(1 for record in records if record["role"] == ROLE_HOLDOUT), + "hash": "sha256(salt + 0x1f + unit)[:8] / 2**64 < fraction", + } + + +def train_rows(records: Sequence[Mapping]) -> list[int]: + return [ + index for index, record in enumerate(records) if record["role"] == ROLE_TRAIN + ] + + +def holdout_rows(records: Sequence[Mapping]) -> list[int]: + return [ + index for index, record in enumerate(records) if record["role"] == ROLE_HOLDOUT + ] + + +def calibration_target_set( + records: Sequence[Mapping], + matrix: sparse.csr_array, + n_households: int, + *, + specs: Sequence | None = None, +): + """The TargetSet the calibrator sees: training targets only. + + Held-out targets are never constructed, so they cannot reach the solve. + With ``specs`` (the checkpoint's registry specs, row-aligned with + ``records``) each target keeps its spec's value, period, source, + metadata and calibration hierarchy, as ``TargetSpec.to_target`` would, + and only its measure becomes the CSR row. + """ + + from microcosm.calibrate.target import Target, TargetSet + + if matrix.shape != (len(records), n_households): + raise ValueError( + f"target matrix shape {matrix.shape} does not match " + f"{len(records)} targets x {n_households} households." + ) + if specs is not None and [spec.name for spec in specs] != [ + record["name"] for record in records + ]: + raise ValueError("target specs are not row-aligned with the roles.") + targets = [] + for index in train_rows(records): + record = records[index] + if record["role"] != ROLE_TRAIN: # pragma: no cover - train_rows filters + raise AssertionError("a held-out target reached the calibration set") + measure = csr_row_measure(matrix, index, n_households) + if specs is None: + targets.append( + Target( + name=record["name"], + entity=record.get("entity", "household"), + measure=measure, + value=float(record["value"]), + period=record["period"], + source=record["source"], + ) + ) + continue + spec = specs[index] + if spec.filter is not None: + raise ValueError( + f"{spec.name}: a filtered spec cannot take a matrix-row measure; " + "the row already carries the filter." + ) + targets.append( + Target( + name=spec.name, + entity=spec.entity, + measure=measure, + value=spec.value, + period=spec.period, + tolerance=spec.tolerance, + source=spec.source, + metadata=spec.metadata, + hierarchy=spec.hierarchy, + ) + ) + held = {records[index]["name"] for index in holdout_rows(records)} + leaked = held & {target.name for target in targets} + if leaked: + raise AssertionError(f"held-out targets reached the calibrator: {leaked}") + return TargetSet(targets) + + +# --------------------------------------------------------------------------- +# Holdout scoring against the pro-rata baseline +# --------------------------------------------------------------------------- + + +def _error_summary(estimates, targets, *, cap: float) -> dict: + from microcosm.calibrate import relative_error_loss + + estimates = np.asarray(estimates, dtype=np.float64) + targets = np.asarray(targets, dtype=np.float64) + scale = np.maximum(np.abs(targets), 1.0) + errors = np.abs(estimates - targets) / scale + return { + "mean_abs_rel_error": float(errors.mean()), + "median_abs_rel_error": float(np.median(errors)), + "p90_abs_rel_error": float(np.quantile(errors, 0.9)), + "max_abs_rel_error": float(errors.max()), + "fraction_within_10pct": float((errors <= 0.10).mean()), + "capped_loss": float( + relative_error_loss(estimates, targets, target_loss_cap=cap) + ), + } + + +def pro_rata_baseline(records: Sequence[Mapping], rows: Sequence[int]) -> np.ndarray: + """State parent value x district population share, for each row.""" + + by_name = {record["name"]: record for record in records} + values = [] + for index in rows: + record = records[index] + parent = by_name.get(record.get("state_cd_parent_target_name")) + if parent is None: + raise ValueError( + f"{record['name']} has no bound state parent for the baseline." + ) + cd_population = record.get("cd_population") + state_population = record.get("state_population") + if not cd_population or not state_population: + raise ValueError( + f"{record['name']} carries no district/state population for " + "the pro-rata baseline." + ) + values.append( + float(parent["value"]) * float(cd_population) / float(state_population) + ) + return np.asarray(values, dtype=np.float64) + + +def score_cd_holdout( + records: Sequence[Mapping], + matrix: sparse.csr_array, + *, + design_weights: np.ndarray, + final_weights: np.ndarray, + cap: float, + receipt: Mapping | None = None, +) -> dict: + """Held-out district targets: calibrated vs design vs pro-rata baseline.""" + + rows = holdout_rows(records) + base = {"report_only": True, **dict(receipt or {})} + if not rows: + return {**base, "n_targets": 0} + held = sparse.csr_array(matrix[rows]).astype(np.float64) + targets = np.asarray([records[i]["value"] for i in rows], dtype=np.float64) + final = held @ np.asarray(final_weights, dtype=np.float64) + design = held @ np.asarray(design_weights, dtype=np.float64) + baseline = pro_rata_baseline(records, rows) + scale = np.maximum(np.abs(targets), 1.0) + final_error = np.abs(final - targets) / scale + baseline_error = np.abs(baseline - targets) / scale + families: dict[str, list[int]] = defaultdict(list) + for position, index in enumerate(rows): + families[str(records[index]["holdout_unit"]).split("|", 1)[1]].append(position) + by_family = {} + for family, positions in sorted(families.items()): + idx = np.asarray(positions) + by_family[family] = { + "n_targets": len(positions), + "calibrated_mean_abs_rel_error": float(final_error[idx].mean()), + "pro_rata_mean_abs_rel_error": float(baseline_error[idx].mean()), + "calibrated_beats_pro_rata": float( + (final_error[idx] < baseline_error[idx]).mean() + ), + } + return { + **base, + "n_targets": len(rows), + "calibrated": _error_summary(final, targets, cap=cap), + "design_weights": _error_summary(design, targets, cap=cap), + "pro_rata_baseline": _error_summary(baseline, targets, cap=cap), + "calibrated_beats_pro_rata_share": float((final_error < baseline_error).mean()), + "baseline": ( + "state parent target x (district population / state population), " + "populations from the PUMA ladder's 119th-plan district overlap" + ), + "by_family": by_family, + "targets": [ + { + "name": records[index]["name"], + "unit": records[index]["holdout_unit"], + "target": float(targets[position]), + "calibrated": float(final[position]), + "design": float(design[position]), + "pro_rata": float(baseline[position]), + } + for position, index in enumerate(rows) + ], + } + + +# --------------------------------------------------------------------------- +# Effective sample size over distinct households, weight share by spine +# --------------------------------------------------------------------------- + + +def kish_ess(weights: np.ndarray) -> float: + weights = np.asarray(weights, dtype=np.float64) + denominator = float((weights**2).sum()) + return float(weights.sum() ** 2 / denominator) if denominator > 0 else 0.0 + + +def distinct_household_weights( + weights: np.ndarray, spine: np.ndarray, source_id: np.ndarray +) -> np.ndarray: + """Weights summed over rows that are copies of one source household. + + A household is ``(household_spine, household_source_id)``: ACS source ids + are the pre-offset ACS ids and collide with donor ids, and a donor + household can appear as its native row and its PUF-detail clone. + """ + + frame = pd.DataFrame( + { + "spine": np.asarray(spine).astype(str), + "source": np.asarray(source_id), + "weight": np.asarray(weights, dtype=np.float64), + } + ) + return frame.groupby(["spine", "source"], sort=True)["weight"].sum().to_numpy() + + +def top_weight_share(weights: np.ndarray, fraction: float = 0.01) -> float: + """Share of total weight on the heaviest ``ceil(fraction * n)`` records.""" + + weights = np.asarray(weights, dtype=np.float64) + if not len(weights) or float(weights.sum()) <= 0.0: + return 0.0 + k = max(1, math.ceil(fraction * len(weights))) + return float(np.sort(weights)[-k:].sum() / weights.sum()) + + +def ess_by_group(weights: np.ndarray, codes: np.ndarray) -> dict[str, float]: + """Kish ESS of the records in each group, keyed by the group code.""" + + frame = pd.DataFrame( + {"code": np.asarray(codes).astype(str), "w": np.asarray(weights, float)} + ) + frame["w2"] = frame["w"] ** 2 + sums = frame.groupby("code", sort=True)[["w", "w2"]].sum() + return { + str(code): float(row.w**2 / row.w2) if row.w2 > 0 else 0.0 + for code, row in sums.iterrows() + } + + +def _distribution(values: Mapping[str, float]) -> dict[str, float | int]: + array = np.asarray(list(values.values()), dtype=np.float64) + if not len(array): + return {"n": 0} + return { + "n": int(len(array)), + "min": float(array.min()), + "p10": float(np.quantile(array, 0.1)), + "median": float(np.median(array)), + "p90": float(np.quantile(array, 0.9)), + "max": float(array.max()), + } + + +def weight_origin_summary( + weights: np.ndarray, + *, + spine: np.ndarray | None, + source_id: np.ndarray | None, + state: np.ndarray | None = None, + district: np.ndarray | None = None, +) -> dict: + """Concentration and origin of a household weight vector. + + Kish ESS over rows (national, per spine, per state, per district), over + distinct households, the top-1% weight share, and household-weight share + by spine. + """ + + weights = np.asarray(weights, dtype=np.float64) + summary: dict[str, object] = { + "rows": int(len(weights)), + "effective_sample_size_rows": kish_ess(weights), + "top_1pct_weight_share": top_weight_share(weights, 0.01), + } + for label, codes in (("state", state), ("district", district)): + if codes is None: + continue + by_group = ess_by_group(weights, codes) + summary[f"effective_sample_size_by_{label}"] = by_group + summary[f"effective_sample_size_by_{label}_distribution"] = _distribution( + by_group + ) + if spine is None: + summary["note"] = "no household_spine column; distinct ESS unavailable" + return summary + spine = np.asarray(spine).astype(str) + total = float(weights.sum()) + summary["weight_share_by_spine"] = { + str(value): float(weights[spine == value].sum() / total) if total else 0.0 + for value in sorted(set(spine)) + } + summary["effective_sample_size_rows_by_spine"] = { + str(value): kish_ess(weights[spine == value]) for value in sorted(set(spine)) + } + if source_id is not None: + distinct = distinct_household_weights(weights, spine, source_id) + summary["distinct_households"] = int(len(distinct)) + summary["effective_sample_size_distinct_households"] = kish_ess(distinct) + summary["distinct_key"] = ["household_spine", "household_source_id"] + return summary + + +# --------------------------------------------------------------------------- +# Sampling rungs for development runs +# --------------------------------------------------------------------------- + + +def rung_token(fraction: float) -> str: + for value, token in SAMPLE_RUNG_TOKENS.items(): + if math.isclose(float(fraction), value, rel_tol=0.0, abs_tol=1e-12): + return token + raise ValueError( + f"sample fraction {fraction!r} is not a rung; use one of " + f"{sorted(SAMPLE_RUNG_TOKENS)} (the stacked pool's rungs)." + ) + + +def _staging_strata(households: pd.DataFrame, tag: str) -> pd.Series: + return ( + households[tag] + .astype(str) + .str.cat(households["congressional_district_geoid"].astype(str), sep="|cd=") + ) + + +def sample_staging_frame(frame, *, fraction: float, seed: int): + """Sample whole households at a development rung. + + ``fraction == 1`` returns the frame unchanged. Otherwise + :func:`microcosm.build.frame_sampling.sample_frame_households` draws + ``floor(fraction * n_h)`` households in every (spine, district) stratum + ``h``, each drawn household's weight is scaled by ``n_h / k_h`` (the + stratum's inverse sampling rate, so a drawn stratum keeps its full + household mass), and each spine is then scaled to its full household + mass. A stratum that floors to zero draws loses its households; the + receipt records how many strata and how much weight that is, and the + per-spine scaling spreads that weight over the spine's drawn strata. At + f001 that is material for the donor spine (its district strata are + small), so read district-level evidence from f010 or above. + """ + + from microcosm.build.frame_sampling import ( + sample_frame_households, + validate_sample_seed, + ) + from microcosm.build.us_runtime.base_pool import spine_column + from microcosm.frame import MassChange + + token = rung_token(fraction) + validate_sample_seed(seed) + households = frame.table("household") + if token == "f100": + return frame, { + "sample_fraction": 1.0, + "rung": token, + "sample_seed": int(seed), + "sampled": False, + "households": int(len(households)), + } + tag = spine_column("household") + strata = _staging_strata(households, tag) + spine = households[tag].astype(str).to_numpy() + full_weights = np.asarray(frame.weights_for("household").values, dtype=np.float64) + full_mass = { + value: float(full_weights[spine == value].sum()) for value in sorted(set(spine)) + } + stratum_size = strata.value_counts() + stratum_mass = pd.Series(full_weights).groupby(strata.to_numpy()).sum() + sampled, receipt = sample_frame_households( + frame, + fraction=float(fraction), + seed=int(seed), + source_name="ACS local staging", + unit_strata=strata.to_numpy(), + floor_context="the ACS local development rung", + ) + sampled_households = sampled.table("household") + sampled_strata = _staging_strata(sampled_households, tag) + drawn = sampled_strata.value_counts() + inverse_rate = sampled_strata.map(stratum_size).to_numpy( + dtype=np.float64 + ) / sampled_strata.map(drawn).to_numpy(dtype=np.float64) + sampled_spine = sampled_households[tag].astype(str).to_numpy() + weights = sampled.weights_for("household") + values = np.asarray(weights.values, dtype=np.float64) * inverse_rate + factors = {} + for value, mass in full_mass.items(): + mask = sampled_spine == value + sampled_mass = float(values[mask].sum()) + if sampled_mass <= 0.0: + raise ValueError( + f"The {token} rung drew no weight on spine {value!r}; the " + "sample cannot stand for it." + ) + factors[value] = mass / sampled_mass + values[mask] *= factors[value] + zero_draw = sorted(set(stratum_size.index) - set(drawn.index)) + new_total = float(values.sum()) + normalized = sampled.with_weights( + "household", + weights.with_values(values, weights.kind), + mass=MassChange( + factor=new_total / float(weights.total), + reason=( + f"ACS local {token} development rung: inverse stratum sampling " + "rate, then per-spine normalization to the full staging " + "household mass" + ), + ), + ) + receipt = { + **{key: value for key, value in receipt.items() if key != "strata"}, + "rung": token, + "sample_fraction": float(fraction), + "sample_seed": int(seed), + "sampled": True, + "strata": "household_spine x congressional_district_geoid", + "n_strata": int(len(stratum_size)), + "zero_draw_strata": len(zero_draw), + "zero_draw_weight_share": float( + stratum_mass.reindex(zero_draw).sum() / full_weights.sum() + ) + if zero_draw + else 0.0, + "per_spine_normalization_factor": factors, + "full_household_mass_by_spine": full_mass, + "households": int(len(sampled_households)), + } + return normalized, receipt + + +def iter_chunks(n: int, size: int) -> Iterable[tuple[int, int]]: + for low in range(0, n, size): + yield low, min(low + size, n)