Map the stored WIC take-up draw onto takes_up_wic_if_eligible when loading US data

[MaxGhenis]

Fixes #530

The defect

policyengine-us 2.x renamed the WIC take-up input from would_claim_wic to takes_up_wic_if_eligible (Person, MONTH, default_value = True), and wic is defined_for it. Published US releases store the seeded draw under the old name only. That includes populace-us-2024-spm-20260915, this package's certified default (pinned in src/policyengine/data/bundle/manifest.json). policyengine-us 2.2.1 defines no would_claim_wic, so every load path skips the stored draw, and every WIC-eligible person takes WIC up.

PolicyEngine/microcosm#1026 measured this with full-population managed_microsimulation() runs. In 2024, WIC is $11.52B with the draw ignored and $6.63B with it restored; the published runtime gives $6.68B. Recipients are 12.65M and 6.66M. I did not reproduce those full-population figures here (see "What I verified").

This implements option 2 of microcosm#1026, as ruled in d271: a load-time rename in policyengine.py.

What changes

policyengine.tax_benefit_models.us.legacy_inputs holds a module-level register, LEGACY_INPUT_RENAMES = {"would_claim_wic": "takes_up_wic_if_eligible"}. One function, apply_legacy_input_renames(simulation, stored_tables), applies it. A rename applies when all of these hold:

• the stored data has the legacy column;
• the engine does not define the legacy name;
• the engine does define the live name;
• the data does not already store the live name.

The live input is then set from the stored draw, as a bool, for every month of every dataset year. The mapping switches itself off once the data stores the live name or the engine defines the legacy name again. No load path mentions WIC.

The mapping covers every place in src/ that builds a US simulation from stored data:

1. Simulation.run() (PolicyEngineUSLatest._build_simulation_from_dataset). Before this change, its input loop called set_input only for columns in system.variables. It now passes the entity tables, already reindexed to the simulation's IDs, to the mapping. This runs for the reform simulation and for its baseline (test_run_maps_the_baseline_of_a_reform_too). Region-scoped copies keep the column. For policyengine-core variable/period H5 files, the loader used to guess a column's entity from its length and dropped a column whose length matched two entities. It now places a pending legacy column on its live input's entity. It also refuses a legacy column stored for part of a year, as the managed path does, instead of reading the first stored month for the whole year. A file that also stores the live name is loaded natively on both paths, whatever period its legacy column is stored for.
2. managed_microsimulation(). policyengine-us loads an entity-table H5, the certified default's layout, into a USMultiYearDataset. Its per-year DataFrames keep the columns the engine skipped. For a policyengine-core variable/period file, the mapping reads only the ID and legacy columns again from dataset.file_path. Before setting anything, each stored table's person_id must equal the simulation's population IDs in order; otherwise the mapping raises ValueError. It maps every year policyengine-us extended the data to. It also maps branches that already exist, since a reform's baseline is branched off during construction.
3. create_datasets(), and so ensure_datasets() when it creates year files. This path builds a policyengine_us.Microsimulation from the source file, so it gets the same treatment. Each year file now stores takes_up_wic_if_eligible and the renames record.

Other simulation constructions in src/ do not load stored US data. household.py builds policyengine_us.Simulation from a situation dict, and the UK paths use policyengine-uk.

Where the record lives

Each run records the renames it applied as {legacy: live}, or {} when none applied. All of these are plain JSON-serialisable dicts:

• simulation.output_dataset.metadata["legacy_input_renames"] and simulation.release_bundle["legacy_input_renames"]. A run over a create_datasets() year file includes the renames applied when that file was cut, so the chain shows the rename even though the year file already stores the live name;
• files written by PolicyEngineUSDataset.save() (saved US outputs and create_datasets() year files), in an H5 dataset policyengine_legacy_input_renames, which PolicyEngineUSDataset.load() restores into metadata;
• run records: results.json gains legacy_input_renames;
• managed_microsimulation().policyengine_bundle["legacy_input_renames"].

{} means no stored column was mapped. It does not certify that the data carried a take-up draw: data that stores neither name runs with the live input's default, which is full take-up. The docs say so.

Behaviour change: files written before this fix are not reused (changelog.d/530.changed.md). Their missing record marks them.

• A saved US output has no record and may have been calculated without the draw. Simulation.load() now raises ValueError for it, so Simulation.ensure() runs the simulation again and saves the result. Direct Simulation.load() callers get the error. This follows the precedent of saved outputs that lack an SPM receipt.
• A year file that ensure_datasets() or create_datasets() wrote under pe.py 6.0.0–6.1.1 stores neither name, so the draw is lost. ensure_datasets() now treats such a file as missing and creates the year files again, and load_datasets() raises for it. A year file opened directly with PolicyEngineUSDataset(filepath=...) is not checked; the docs say so.

Invariants

These hold for every input. Each is tested with Hypothesis property tests and example tests:

Invariant	Property tests	Example and integration tests
Draw preserved. For any stored boolean draw and any set of dataset years, the live input equals the stored draw in every month of every year.	test_live_input_equals_the_stored_draw_every_month_of_every_year (200 examples: bool, int8, int64 and float64 0/1 draws, 1–4 years, 1–30 people); test_microsimulation_draw_is_preserved_on_every_branch; test_variable_centric_draw_is_preserved_every_month_of_every_year (core H5 files); test_real_engine_takes_up_any_stored_draw_every_month (real policyengine-us 2.2.1 engine, 25 examples)	test_run_keeps_a_stored_false_draw, test_run_over_a_core_h5_keeps_a_stored_false_draw, test_run_over_a_region_keeps_a_stored_false_draw, test_managed_entity_table_file_keeps_a_stored_false_draw, test_managed_variable_centric_file_keeps_a_stored_false_draw, test_managed_reform_maps_its_baseline_too, test_create_datasets_extracts_the_draw_under_the_live_name
No-op when it does not apply. Nothing is set when the engine defines the legacy name, when the engine lacks the live name, or when the data already stores the live name.	test_nothing_is_set_when_the_rename_does_not_apply (all three cases); test_only_years_whose_data_lacks_the_live_name_are_mapped	test_pending_renames_follow_what_the_engine_defines, test_variable_centric_file_storing_the_live_name_is_not_read, test_run_leaves_data_that_stores_the_live_name_alone, test_a_core_h5_storing_the_live_name_ignores_a_part_year_legacy_draw (both load paths)
Order check. A stored table that is not in the simulation's person order is refused, never silently misaligned, and nothing is set.	test_a_table_out_of_simulation_order_is_refused_and_nothing_is_set (a permutation, a foreign ID, a missing person, an extra person)	test_a_table_without_ids_is_refused, test_managed_person_table_out_of_order_is_refused (real engine), test_a_year_storing_its_own_ids_is_checked_against_them (a later year of a core H5 that stores its own person_id in another order), test_variable_centric_file_with_a_mislengthed_draw_is_refused
Idempotence. Applying the mapping twice gives the same inputs as applying it once.	test_applying_twice_gives_the_same_inputs_as_once; test_real_engine_takes_up_any_stored_draw_every_month (applies twice on the real engine)	test_managed_mapping_is_idempotent
Record round trip. A file's renames record reads back exactly as written, and writing again replaces it.	test_a_stored_record_reads_back_as_written (50 examples)	test_a_saved_output_keeps_the_renames_it_applied, test_create_datasets_extracts_the_draw_under_the_live_name (year file record survives a reload and reaches the run's record)
Stale files are not reused. A US file without the record is refused (load_datasets, Simulation.load()) or recomputed (ensure_datasets, Simulation.ensure()).		test_ensure_datasets_regenerates_a_year_file_cut_before_the_mapping, test_an_output_saved_before_the_mapping_is_not_reused
Whole years only. A legacy column stored for part of a year, in data that does not store the live name, is refused on both load paths.		test_a_part_year_period_is_refused, test_a_core_h5_with_a_part_year_draw_is_refused_by_run, test_managed_core_h5_with_a_part_year_draw_is_refused

Differential test. test_both_load_paths_give_the_same_take_up_and_wic checks that Simulation.run() and managed_microsimulation() give the same take-up and wic for the same household.

Fix versus no fix. The fixture is one adult with no income in Los Angeles County, an infant whose stored draw is False, and a two-year-old whose draw is True. Each test first asserts that both children are WIC-eligible, with is_wic_eligible true and wic_if_takes_up above zero. With the mapping, the infant gets wic == 0 and the toddler gets wic == wic_if_takes_up. test_run_without_the_mapping_gives_every_eligible_person_wic and test_managed_without_the_mapping_gives_every_eligible_person_wic empty the register, and then the infant gets a positive wic. All of this runs in CI's make test: the integration module needs only the [dev] extras and no downloads.

Mutation check. I ran five mutants of legacy_inputs.py against the unit and property tests: no order check, 11 months instead of 12, ignoring an engine that defines the legacy name, ignoring data that stores the live name, and setting inputs before every table has been checked. The property tests killed all five. After review, eight more mutants ran against both legacy test modules: skipping the baseline mapping in run(), ensure_datasets ignoring the record, load_datasets accepting a file without it, save() not writing it, load() not restoring it, run() ignoring the input dataset's record, run() reading a part-year draw from a core H5, and checking every year of a core H5 against the first period's person IDs. Each was killed by the new test aimed at it.

What I verified

• Certified data. I read the certified default in place from the local Hugging Face cache. Its SHA-256 is 6496cc43…aee84, the hash the manifest pins. It is an entity-table HDFStore with _time_period = 2024. The person table stores would_claim_wic as bool, True for 7,470 of 166,321 rows, and has no takes_up_wic_if_eligible.
• Real-data subset, not a population estimate. I picked 600 households from the certified file with np.random.default_rng(0).choice(household_ids, 600, replace=False), where household_ids is the household table's household_id column in stored order. That gives 1,700 people. I wrote them to an entity-table file in the certified layout, loaded it with PolicyEngineUSDataset (which projects the household weights onto the other entities), and ran Simulation.run() for 2024. With the mapping, person-weighted WIC was 0.606 of the unmapped figure and weighted recipients 0.620; unweighted, the ratios were 0.535 and 0.492. Of the 65 people with wic_if_takes_up > 0, 32 took WIC up with the mapping and 65 without it. For comparison, microcosm#1026's full-population ratios are 0.576 and 0.526. The sample is small, so it confirms only the direction of the fix and that the fix reaches real data. Peak RSS was 6.0 GiB.
• Stale year files on real data. On the same subset, create_datasets(years=[2026]) wrote a year file whose takes_up_wic_if_eligible equals the stored draw for every person, and which records the rename. I then removed the record and the live column, as a 6.1.1-era file would lack them. load_datasets refused that file, and ensure_datasets regenerated it with the record. A 2026 run over the regenerated file gave 32 of the 65 eligible people WIC, and its legacy_input_renames was {"would_claim_wic": "takes_up_wic_if_eligible"}.
• Engine behaviour. The mechanism statements above come from reading policyengine-us 2.2.1 (Microsimulation.__init__, USSingleYearDataset, USMultiYearDataset, extend_single_year_dataset, takes_up_wic_if_eligible) and policyengine-core 3.32.5 (build_from_dataset, Dataset.from_file, Holder.set_input) in this environment.
• Lint, format and tests. Results are in the next section.

Not verified

• I did not run a full-population managed_microsimulation() on the certified default: it needs about 60 GiB, and the machine is shared. The full-scale check, which should reproduce microcosm#1026's run C ($6.63B, 6.66M recipients in 2024), is left to a separate run.
• The Linux, Python 3.11–3.13 and CI-only jobs (bundle verification, changelog check) run only in CI. Locally I used Python 3.14.4 on macOS.
• mypy is informational in CI, and main already has errors in files this PR does not touch. The count is under "Commands and results".
• The renames record marks a file as written with this mapping, not with a particular register. If a later release adds a register entry, files written before that entry will carry a record and pass these checks, so that change would need to extend them. A comment on RENAMES_H5_DATASET says so.
• A year file opened directly with PolicyEngineUSDataset(filepath=...) is not checked for the record. Only ensure_datasets and load_datasets check it.

Commands and results

Local runs used Python 3.14.4 on macOS, with HF_HUB_OFFLINE=1 and no Hugging Face token. Each pytest process ran under a watchdog that caps RSS at 7.5 GB.

At 51065c6 (a follow-up: Simulation.run()'s part-year refusal now skips a core H5 that also stores the live name, as the managed path already did):

• Format and lint. With CI's ruff (0.16.9), ruff format --check . reports "267 files already formatted" (untracked files excluded), and ruff check . reports "All checks passed!".
• Tests. One process over test_us_legacy_inputs.py (45), test_us_legacy_inputs_integration.py (23), test_us_native_hdf_weights.py, test_dataset_persistence.py, test_dataset_runtime.py and test_us_long_term_datasets.py: 123 passed (peak 3.4 GiB).
• New test against the old code. With datasets.py restored to 7e0ae30, test_a_core_h5_storing_the_live_name_ignores_a_part_year_legacy_draw fails with "only yearly periods are supported"; with the fix it passes.
• Types. mypy --python-version 3.14 src/policyengine: 121 errors, unchanged.
• CI at 51065c6. Every check passed. Each of Test (3.11), (3.12), (3.13) and (3.14) ran the full suite: 1,188 passed, 9 skipped.

At 7e0ae30 (the review fixes):

• Format and lint. With CI's ruff (0.16.9), ruff format --check . reports "267 files already formatted", and ruff check . reports "All checks passed!". ruff 0.16.9 also checks Markdown, so the new changelog fragment raised the count from the 266 that CI's Lint job reported at d4bf360. The earlier body's "267" also counted an untracked agent notes file.
• Legacy tests. tests/test_us_legacy_inputs.py gave 45 passed. tests/test_us_legacy_inputs_integration.py gave 22 passed.
• Affected tests. Four batches, all passed:
  • legacy modules plus test_run_record.py, test_spm_model.py and test_us_native_alignment.py: 104 passed (peak 5.3 GiB);
  • test_be_axiom_pilot.py, test_budgetary_impact.py, test_cache.py, test_dataset_persistence.py, test_dataset_runtime.py and test_extra_variables.py: 45 passed, 4 skipped;
  • test_labor_supply_response.py, test_release_manifests.py, test_small_follow_ups.py and test_us_long_term_datasets.py: 85 passed;
  • test_us_microsim_structural_reforms.py, test_us_native_hdf_weights.py, test_us_program_statistics.py and test_dict_reforms_on_simulation.py: 45 passed.
• Full suite. I did not re-run the full suite locally at 7e0ae30; CI's Test jobs run it. At d4bf360, the local full suite, run in batches of eight files, gave 1,174 passed, 9 skipped and 0 failed. One of those batches (test_us_long_term_datasets.py through test_variable_labels.py) reached the cap in test_us_reform_application.py. Split into three processes, its files passed (61 + 6 + 37).
• Types. mypy --python-version 3.14 src/policyengine reports 121 errors on this branch and 124 on origin/main (4dc5959), both measured at this head. None are in legacy_inputs.py.
• Lockfile. At d4bf360, uv lock --check --offline was consistent and python scripts/release_lock.py passed. The review fixes do not touch pyproject.toml or uv.lock.

🤖 Generated with Claude Code

[MaxGhenis]

Merged by the session that opened it, on Max's go (decision d424, 2026-09-26), with head 51065c6ab pinned.

• Review: two independent Opus 5.5 reviewers (correctness lens and tests/claims lens). Round 1 requested changes (stale year files; the untested reform-baseline claim). Both approved round 2 at 51065c6ab.
• CI: 13 checks passed, including tests on Python 3.11–3.14.
• Full-population check (launchd, on the certified default populace-us-2024-spm-20260915, through this PR's own managed_microsimulation):
  • 2024: WIC $6.633B, 6.661M recipients;
  • 2030: WIC $7.114B, 6.42M recipients;
  • take-up share 3.8%.
    These match microcosm#1026's measured run C (draw restored: $6.63B / 6.66M) against run B (draw ignored: $11.52B / 12.65M).

🤖 Generated with Claude Code