Export the identifiers for every year an input frame names, so it round-trips

[MaxGhenis]

Fixes #2333

Fixes the main CI failure from #2251 × #2270.

What and why

Main has been red since #2251 merged (cfa0c55). One test fails: test_dataframe_dataset.py::test_an_exported_frame_loads_back_and_exports_itself (main run 38143054188 at 7343f94: 1 failed, 3749 passed, 1 skipped, 2 xfailed, 1 xpassed). #2270 added that property: a frame to_input_dataframe() exports should load back to a simulation that exports the same frame.

Mechanism

1. Stop deciding a household's UC health element rate by its position #2251 records inputs for later years.
   • Simulation.__init__ applies the Universal Credit rebalancing modifier to every simulation (policyengine_uk/simulation.py:198).
   • For each of 2026 to 2029 in which the rebalancing is active, _new_claimant_health_element_status records uc_receives_new_claimant_health_element with sim.set_input when nothing supplies it (policyengine_uk/scenarios/uc_reform.py:114-121, called from :140-145).
   • The variable has no formula, so it is an exportable input.
2. policyengine-core exports an input only for the periods it was set for. It takes the _user_input_keys records (core 3.32.27 simulations/simulation.py:2208-2234, to_input_dataframe at :2236-2268).
   • A simulation built from a 2025 frame set its identifiers for 2025 only. Its dataset has one year.
   • It holds them in 2026 to 2029 only because core carries inputs forward (core :1198-1240; auto_carry_over_input_variables, policyengine_uk/tax_benefit_system.py:75).
   • So the first export has the status for 2026 to 2029 but person_id, benunit_id, household_id, person_benunit_id and person_household_id for 2025 alone.
3. The loader holds the identifiers in every year a frame gives values for.
   • dataset_from_input_dataframe makes a dataset year for each whole calendar year a non-identifier column names (policyengine_uk/data/input_dataframe.py:368-372, :401).
   • Every year's tables hold all five identifiers (:427-439).
   • It also sets them for each other year an identifier column names (:418-426).
   • build_from_multi_year_dataset then calls set_input for every column of every year (policyengine_uk/simulation.py:544-549).
4. Result. The reloaded simulation records the identifiers for 2026 to 2029 and exports them: 18 columns become 38.

The same gap was there before #2251 for any input set after building for a year outside the data, for example sim.set_input("hours_worked", 2027, ...). #2251 made every simulation do that.

Bisect

The hub ran repro_roundtrip.py (the minimal example: one person aged 0, benunit 1, household 1, LONDON, every other input 0 or false, id_columns="both"):

Commit	First export	Second export
83001d5 (main after #2270)	14 columns	14 columns
cfa0c55 (after #2251)	18 columns	38 columns
this branch (my run of the same script)	38 columns	38 columns

On this branch before the fix (6d89338), I reproduced 18 → 38. The extra columns are the five identifiers for each of 2026 to 2029.

Fix: the exporter

The export now gives every identifier for each whole calendar year one of its columns names. Those are exactly the years in which a simulation built from the frame holds them as inputs.

• with_identifiers_in_each_year (policyengine_uk/data/input_dataframe.py:443) adds a column for each identifier the export gives, for each such year it lacks.
  • Each added column copies the identifier's latest earlier column, or else its earliest. The loader already requires every column of an identifier to agree (_identifier).
  • It is inserted beside the identifier's other columns in period order, which is where core's default export puts an input set for that year. assert_frame_equal checks column order.
  • The loader and the helper share one "whole year" predicate, _is_whole_year (:276), extracted from the loader without changing it.
• Simulation.to_input_dataframe (policyengine_uk/simulation.py:432-453) calls core's export, then the helper.

Why the exporter, not the loader. The loader treats identifiers as holding in every period (input_dataframe.py:368-369) and holds them in every dataset year by design. Invariant 3 asserts that every dataset year holds the identifiers. With this fix, the loader and every simulation it builds are unchanged; the export now states the identifiers the loader would have set anyway.

A loader-side fix would make the loader forget identifier input records for years the frame does not name and rely on carry-forward. It needs special cases:

• inputs do not carry back to a year before the first named identifier;
• a frame can name its identifiers only for months.

Why not #2251's recording. #2251 records a seeded draw for data-built simulations, and a frame-built simulation is one, so not recording the variable's default would not help here. Any other later-year set_input breaks the round trip the same way.

The other set_input calls under policyengine_uk/scenarios set pip (pip_reform.py:18, which has adds) and uc_LCWRA_element (uc_reform.py:159, which has a formula). The default export leaves formula variables out, so neither reaches it; the fix covers them anyway.

File ownership. The simulation.py diff is the import and the new to_input_dataframe method beside build_from_dataframe, nothing else.

Invariants

• Fixed point. For every simulation S built from a frame, export(load(export(S))) == export(S), column order and dtypes included. This holds when the frame names identifiers and values for different years, and when S was given inputs (or a lone identifier) for years outside its data once built.
• Identifiers in every year. An export gives person_id and the benefit unit and household identifiers exactly for every whole year any of its columns names, and every identifier it gives for some year, for all of those years.
• Idempotence. with_identifiers_in_each_year is idempotent, and does not change its argument.

Tests

New, in policyengine_uk/tests/test_dataframe_dataset.py:

• The minimal counterexample as an explicit @example on test_an_exported_frame_loads_back_and_exports_itself.
• test_an_export_names_the_people_in_each_year_the_rebalancing_records: the Stop deciding a household's UC health element rate by its position #2251 case, deterministic.
• test_an_input_set_once_built_for_another_year_round_trips: an input set by hand for 2024, 2027 or 2032, and a lone person_id for 2032.
• test_an_export_gives_its_identifiers_for_each_year_it_gives_values_for: a unit test of the helper. It covers order, sources and dtypes, months, years only an identifier names, non-variables, idempotence, and that the argument is not changed.
• Property test_a_frame_over_several_years_exports_the_same_frame_again, with two explicit examples. Each frame names identifiers for 1–2 years and gives values for 1–3 years, both drawn from 2024–2030. The simulation is given hours_worked and person_id for other years once built. The property asserts that the second export equals the first and that the identifiers are exact in every year.

The invariant docstring (item 5) is extended. No test was weakened, skipped or marked xfail.

Results

All runs used policyengine-core 3.32.27 (the uv.lock pin), Python 3.13 and Hypothesis 6.168.2, on an overloaded host (load average 200–470).

The minimal example, through the hub's repro_roundtrip.py:

Code	First export	Second export	Extra columns
Before the fix (6d89338)	18	38	20
With the fix	38	38	none

pytest policyengine_uk/tests/test_dataframe_dataset.py policyengine_uk/tests/test_uc_rebalancing.py policyengine_uk/tests/test_input_records.py -n 3 --dist worksteal on 49e7ef3: 96 passed, 1 failed (1:54:39). The later head a966445 only reflows a docstring.

File	Passed	Failed
test_dataframe_dataset.py, including every new test	78	0
test_uc_rebalancing.py	17	1
test_input_records.py	1	0

The one failure is test_uc_rebalancing.py::test_section_4_floor_properties. It raised hypothesis.errors.FailedHealthCheck: Input generation is slow: Hypothesis only generated 2 valid inputs after 656.73 seconds, which is host starvation, not a property violation. That test checks pure arithmetic in uc_reform.protected_lcwra_floor, which this PR does not touch, and it passed on main CI. Run alone with the same seed (pytest policyengine_uk/tests/test_uc_rebalancing.py::test_section_4_floor_properties --hypothesis-seed=324857880990766857290461818381978056954), it passes: 1 passed.

Mutation check. I removed the UK to_input_dataframe override in-process, so core's export ran as on main, and ran the new regressions and the @example: all 6 fail (6 failed, 72 deselected).

• The @example fails at the round trip's assert_frame_equal.
• The regressions fail because the identifier columns for the later years are missing.

Independent review. An Opus 5.5 agent reviewed the diff, reading code only, and agreed with the exporter-side fix. I fixed its two minor findings in 49e7ef3:

• the helper now also covers years only an identifier names, as the loader does;
• a docstring claim about years before the data is corrected.

Not run locally: the full make test suite (CI runs it) and make documentation; the docs change is one Markdown bullet.

Impact

No calculated value or aggregate changes: the fix touches only the DataFrame export and a no-op refactor of the loader.

• A Microsimulation subclasses Simulation (policyengine_uk/microsimulation.py:14) and has no constructor or export of its own.
• Simulation.__init__ sends only a pd.DataFrame to build_from_dataframe (simulation.py:180-181), the only caller of dataset_from_input_dataframe (:415). An h5 path or URL goes through build_from_dataset_source → build_from_url / build_from_file → build_from_multi_year_dataset (:178, :309-398, :519-551), as do dataset objects and the default dataset (:182-189).
• Nothing in policyengine_uk outside tests calls to_input_dataframe. In core 3.32.27 its only caller is Simulation.subsample (core simulations/simulation.py:2330).
• I did not run a Microsimulation comparison. It isn't needed for an export-only change, and the host's compressor was at 39.5–52 GB, over the 30 GB limit.

Follow-ups (not in this PR)

• Differential gap (pre-existing). A simulation given an input for a year before its data holds its identifiers as zero in that year, because inputs do not carry back (core :1220, :1241-1250). The frame it exports loads back with the real identifiers there, as it did before this PR.
  • So formulas that read identifiers for the period can differ in that year between the two simulations: is_benunit_eldest_child.py:22, months_since_last_birthday.py:36 and uc_deduction_random_draw.py:25.
  • The round trip holds; invariant 6 (the differential) does not, for such simulations.
• Core's subsample, from reading, not run. It rebuilds with self.build_from_dataset() and no argument (core :2410), but UK's build_from_dataset(self, dataset) (simulation.py:455) requires one.
• to_input_dict() is unchanged and still gives the identifiers only for the periods they were set for.

Test plan

• For policy changes: not applicable (no policy change).
• ruff format and ruff check on the changed files (ruff 0.15.7): clean.
• Changelog fragment in changelog.d/fix-roundtrip-recorded-inputs.fixed.md.

axiom: n/a: DataFrame export round trip only; no policy, parameter or formula change

🤖 Generated with Claude Code

[MaxGhenis]

UK hub landing audit for #2349, head a9664455e12e415ab96dd5cce7f0fe03f14a75a1.

• Why: main has been red since Stop deciding a household's UC health element rate by its position #2251 (cfa0c55). Every PE-UK PR fails test_dataframe_dataset.py::test_an_exported_frame_loads_back_and_exports_itself. This fixes the exporter so an export is a fixed point of load-then-export.
• Independent review: GPT-6.1 Sol, APPROVE at this exact head (~/reviews/uk-hub/reviews/uk2349/review.md).
  • It verified each mechanism claim against the code and the locked policyengine-core 3.32.27.
  • It confirmed that the loader and build_from_dataframe are unchanged and that the h5 path is untouched.
  • It found that Lift the benefit cap only for the support component and the people HB reg 75F and UC reg 83 name, and cap only the family's benefits #2076, if it lands too, should still round-trip.
  • It found no blockers and no should-fix items, only four nits, which are follow-ups (helper docstring scope, one shared _whole_years rule, return-copy docstring, a test that disagreeing identifier columns still raise).
  • Its approval was conditional on the test runs; the CI Test job on this head covers that.
• Gates at merge: gh pr checks exits 0 on this head (Test included), MERGEABLE, not a draft, merged with a merge commit by --match-head-commit.
• Impact: no calculated value or aggregate changes. Only Simulation.to_input_dataframe output changes. h5, URL and dataset builds never reach it (code path in the PR body and the review).
• After merge: the hub tells Count 16-year-olds outside education as UC and Pension Credit qualifying young persons until the 1 September after their 16th birthday #2279, Relieve pension contributions to relevant earnings and charge the annual allowance excess once #2240, Stop charging property stocks as purchases when a household buys its home #2241, Apply the Universal Credit Act 2025 standard allowance uplifts to 2029-30 #2242, Fix benefit cap exemption defects (#1818) #1820, Lift the benefit cap only for the support component and the people HB reg 75F and UC reg 83 name, and cap only the family's benefits #2076, Split the CI Test job into shards and cut its runtime #2269 and Assert the 2023-24 Class 4 lower profits limit instead of expecting it to fail #2303 to merge main and push once.