Load a DataFrame simulation input as the dataset it describes

[MaxGhenis]

Fixes #2266

axiom: n/a: simulation infrastructure, no policy rule change

What was wrong

Simulation.build_from_dataframe set every column of a frame (one row per person, columns variable__period, as to_input_dataframe() exports) as one value per person. So:

• any benefit unit or household with more than one person raised ValueError: Unable to set value ... as its length is 2 while there are 1 benunits (the issue's reproduction, confirmed on a3752304);
• single-person groups whose rows were not in identifier order got each other's values without any error: households 30, 10, 20 with council tax 3,000, 1,000, 2,000 came out as 3,000, 1,000, 2,000 for households 10, 20, 30;
• identifier columns were found by substring, so household_id__ matched person_household_id__2025;
• a frame with no group columns (the example in docs/book/usage/simulations.md) raised IndexError.

What changed

• policyengine_uk/data/input_dataframe.py (new): dataset_from_input_dataframe turns the frame into the UKMultiYearDataset it describes. It has a year for each whole calendar-year period the columns name. Each benefit unit and household takes the value on its members' rows, and those rows must agree. Periods that are not one whole year, such as months, are returned to be set directly.
• Simulation.build_from_dataframe loads that dataset with build_from_multi_year_dataset, the same path as a multi-year dataset, then sets the other periods.
• Columns match by exact name: column.split("__")[0] == variable (via str.partition).
• Identifiers are read as exact 64-bit integers. Columns that mix integers and floats (object columns) are read value by value.
• A variable given for both months and the year has all its columns set again after the dataset loads, in the frame's column order. The previous loader set every column in turn, so this keeps its results.
• nhs_budget_change, education_budget_change and other_public_spending_budget_change skip income deciles with no weighted households. Before, they raised KeyError for any dataset-built simulation whose data left a decile empty, such as a small sample (confirmed for a seven-household dataset). Setting dataset on frame-built simulations sends them down this path. A full dataset fills every decile, so its results do not change.
• Docs: the DataFrame section of docs/book/usage/simulations.md now shows a multi-person household and says how groups are named.

Decisions

Simulation.dataset is now the multi-year dataset, not None. Formulas branch on simulation.dataset is not None or read dataset.years. These include basic_state_pension / new_state_pension / additional_state_pension (data year), consumption_shareholding, shareholding, and nhs_budget_change and the other budget-incidence variables. I took a dataset-built simulation of seven households, exported it with to_input_dataframe(include_computed_variables=True) and reloaded it:

• with dataset set, every one of the model's 1,117 variables matched in 2025 and 2027;
• with dataset forced to None, 72 variables differed.

The test test_an_exported_frame_gives_the_results_of_its_dataset now checks the first result. A frame-built simulation also behaves like a dataset-built one under no_economic_assumptions, which resets uprating on dataset.

Identifiers that 32-bit floats round (#2267, not fixed here). person_household_id and person_benunit_id are stored as float32, so an exported frame holds them rounded above 2**24. The loader:

• takes groups from the exact integer household_id / benunit_id columns first, then the person_ columns, then one group per person;
• compares person_<group>_id with <group>_id exactly when the person_ value is an integer (ready for person_household_id and person_benunit_id are stored as float32, which rounds large identifiers #2267 making them integers). A float value agrees if it is the group's identifier exactly, or the identifier rounded to a 32-bit float, either as held or as CSV writes it. That covers the exported float32 column, and the same frame saved as CSV: CSV writes the float32 900000064 as 900000060.0, a different 64-bit float. Any other float that merely rounds to the same 32-bit float names another group and is refused, so two exact 64-bit floats naming each other's households do not pass;
• warns when groups are named only by float person_ columns of 2**24 or more. Those may have been rounded before the frame was made, for example an export saved as CSV with the <group>_id columns dropped, and no check can see that;
• refuses, rather than merges, a frame that names its groups or people only with float32 identifiers of 2**24 or more. 2**24 + 1 rounds to 2**24, so the bound is >=. The same rule applies to any float, at its own limit (2**53 for 64-bit floats), and to each value of a mixed object column. In such a column, np.float32(2**24) compares equal to 2**24 + 1, a case review round 1 found.

It also refuses:

• a repeated person identifier;
• missing identifiers, identifiers that are not whole numbers, or identifiers beyond 64 bits;
• identifier columns that differ between periods;
• <group>_id and person_<group>_id columns that name different groups;
• group members whose rows hold different values for a group variable;
• columns not named variable__period (these crashed before too).

Months and years. Core spreads a year's value over the months not yet set. It divides what is left for floats, and fills the rest for other types. So when a frame gives one variable for months and for the year, the order of its columns decides the months' values, as it did when the previous loader set each column in turn. The loader keeps that order. A frame to_input_dataframe() exports never gives one variable for both: it exports each variable only for periods of its own definition period.

Years. The dataset's years are the years the frame gives values for. Identifier columns add none, because identifiers hold for every period. This matters because the State Pension formulas read the dataset's first year as the data year. An identifier column named for 2025 beside data for 2027 used to add an empty 2025 and make the pension zero. Review round 2 found that; it is now tested against the previous loader. A frame with values in several years has its first such year as the data year, as a multi-year dataset does.

An identifier column named for a year with no data still has its identifiers set for that year, as the previous loader set them. Inputs carry forward but not back, so they would otherwise read as zero before the data.

A frame is loaded with exactly the years it gives values for. A one-year frame is not projected forward the way a UKSingleYearDataset is: that would change later-year results for frames that load today. A frame that names no whole year is held in default_input_period.

Not checked (as for the dataset loaders): a benefit unit whose members are in different households.

Not changed: shareholding, consumption_shareholding and consumer_incident_tax_revenue_change divide by a weighted total. That total can be zero while the unweighted total is not, for example when every household holding wealth has zero weight. Dataset-built simulations already behave this way, and frame-built ones now share it. Review round 2 raised it as a nit. It needs its own change to those formulas: #2296.

Invariants (property-tested with Hypothesis in policyengine_uk/tests/test_dataframe_dataset.py)

For every frame of people in benefit units in households, in any row order, whichever identifier columns name the groups (<group>_id, person_<group>_id, both, or as exported with float32 person_ columns), with identifiers near 1, near 2**24 and near 9e8:

1. People and groups. People are in row order. Benefit units and households are the distinct identifiers in ascending order, and each person is in the ones their row names, also for identifiers float32 cannot tell apart.

2. Values. Every person, benefit unit and household variable has the value its rows give. Projected back to people, each variable gives its column again.

3. Dataset. dataset holds the same people, groups and values with exact identifiers.

4. Row order. Reordering rows changes no entity's values.

5. Round trip. to_input_dataframe() loads back to the same people, groups and inputs, and exports the same frame again. So does that frame saved as CSV.

6. Differential. A simulation built from a dataset and one built from its exported frame agree on every variable of the model. Checked with the default export and with include_computed_variables=True, and for a reform and its baseline.

7. Exact names. Columns whose names only contain a variable's name change nothing.

8. No mutation. The frame is not changed.

9. Previous loader (differential). On every frame the previous loader read correctly (single-person groups, rows in identifier order), both loaders give every input the same value in every period they set. This includes monthly inputs given for some months and the year, in any column order.

10. Mixed identifier types. With identifiers written as Python or NumPy integers or floats in object columns, a frame loads its people and groups exactly, or is refused when an identifier the groups are read from is a float32 of 2**24 or more.

11. Small-sample spending. For the three budget reforms on a seven-household frame and on its dataset, allocations are equal, finite and positive. Households in the same decile get the same amount, and the total is under the budget.

Intended exceptions are the refusals above. Each has an example test, including the minimal counterexamples found while writing this (2**24 + 1 vs 2**24, and repeated person ids), which merged households before those checks were added.

A limit of the default export, not of the loader: to_input_dataframe() leaves out inputs a dataset supplies for variables with formulas, such as state_pension_reported. include_computed_variables=True keeps them. Filed as #2297.

Interaction with #2265

#2265 (open) adds Simulation.subsample. I merged both branches locally and ran test_subsample.py: every test passed but one. The failure is test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample, which asserts that a frame-built simulation has dataset is None. That stops holding here, by design. With this change, subsampling a frame-built simulation also cuts its dataset down to the sampled households, which is the behaviour #2265 already has for dataset-built simulations.

Whichever of the two PRs lands second needs that test replaced. The replacement, test_a_simulation_built_from_a_dataframe_subsamples_its_dataset, checks:

• the dataset is a UKMultiYearDataset;
• the dataset's households, people and weight total match the subsample.

It passes on the merge. The patch against #2265's test file:

Patch for policyengine_uk/tests/test_subsample.py

diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py
index f116507d5..5f58dcb4e 100644
--- a/policyengine_uk/tests/test_subsample.py
+++ b/policyengine_uk/tests/test_subsample.py
@@ -499,7 +499,9 @@ def test_subsampling_a_single_year_dataset_keeps_its_type():
     assert len(data.household) == 6
 
 
-def test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample():
+def test_a_simulation_built_from_a_dataframe_subsamples_its_dataset():
+    # A frame loads as the multi-year dataset it describes (#2266), which the
+    # subsample cuts down with the populations.
     households = np.arange(4)
     frame = pd.DataFrame(
         {
@@ -513,12 +515,11 @@ def test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample():
         }
     )
     simulation = Microsimulation(dataset=frame)
-    assert simulation.dataset is None
+    assert isinstance(simulation.dataset, UKMultiYearDataset)
     ages = by_id(simulation, "age", 2025)
 
     simulation.subsample(n=3, seed=1)
 
-    assert simulation.dataset is None
     kept = ids(simulation, "household")
     assert 1 <= len(kept) <= 3
     assert ids(simulation.baseline, "household") == kept
@@ -526,6 +527,12 @@ def test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample():
     assert by_id(simulation, "age", 2025) == {
         person: ages[person] for person in ids(simulation, "person")
     }
+    household = simulation.dataset[2025].household
+    assert household.household_id.tolist() == kept
+    assert household.household_weight.sum() == pytest.approx(100.0)
+    assert simulation.dataset[2025].person.person_id.tolist() == ids(
+        simulation, "person"
+    )
 
 
 # Refusals.

#2265's last Test run was cancelled, so it has no green run that predates this change.

Review

• Round 1 (GPT-6.1 Sol, at 620c2483): REQUEST_CHANGES with three findings, all addressed in 8b2baaa1:

  • Blocker: mixed-type identifiers merged groups.
  • Should-fix: monthly and yearly values were set out of column order.
  • Should-fix: the budget allocations raised KeyError on small frames.

• Round 2 (Opus 5.5, at cc126d88, by reading only: its lane had no shell): REQUEST_CHANGES. It confirmed the three round-1 fixes and raised two should-fix findings and four nits. I executed the findings and confirmed A, B and D. Addressed in 6d689341:

  • A, should-fix: an exported frame saved as CSV was refused.
  • B, should-fix: identifier columns added an empty data year.
  • D, nit: a Python int beyond 64 bits raised OverflowError.
  • E, nit: rounding in float person_ columns given alone cannot be detected, so the loader now warns.
  • F, nit: 2**24 rendered as bold.
  • C, nit: left alone; see "Not changed".

• Round 3 (Opus 5.5, at 6d689341, by reading only): APPROVE with three nits. Addressed in 58af8c52:

  • Nit 1: identifiers were not stored for a year their column names when that year has no data.
  • Nit 2: the float comparison accepted exact floats that name another group.
  • Nit 3 (the warning's stack level): cosmetic, left as it is.

• Round 4 (Opus 5.5, at 58af8c52, by reading only): APPROVE, no blocker or should-fix. It checked that pandas 2.3.1 to_csv writes a float32 column as str(np.float32(v)), which is the form the loader accepts. Two nits are left as they are:

  • a float64 person_ value is compared exactly only below 2**53 (the integer column still decides the group);
  • one docstring line is longer than its neighbours.

• Round 5 (delta review of 4892db98, which adds two tests and rewraps a docstring): APPROVE at 4892db98, the current head (Opus 5.5, by reading). I confirmed with git diff --stat 58af8c52 4892db98 that the commit touches only the test file and that docstring.

• CI on 58af8c52 failed one test of 3,697. test_an_exported_frame_gives_the_results_of_its_dataset[default export] failed at its check that the loader leaves the exported frame unchanged: the column's array class differed (EnumArray against ndarray), not its values. From policyengine-core 3.33, which CI resolves, to_input_dataframe() exports an enum input as its stored EnumArray. pandas 3 keeps that class until a copy, and assert_frame_equal compares classes. The whole-model differential on the line before passed. 3b8f54b2 compares a copy. Round 6 (delta review of 3b8f54b2): APPROVE at 3b8f54b2, the current head (Opus 5.5, by reading). It traced the class change to pandas 3's deep copy merging same-dtype blocks with np.vstack. It confirmed that comparing a copy still catches any change to the frame's values, and that the loader reads enum columns as plain integer codes.

Rounds 2 to 6 ran on Opus 5.5 lanes without a shell: the GPT-6.1 Sol lanes that ran round 1 stayed queued. So the reviews are by reading, and the execution evidence is the local test and mutation runs below, plus CI.

Tests run

Locally, Python 3.13, policyengine-core 3.32.27, one process each, -p no:cacheprovider, fresh --basetemp deleted afterwards, at 4892db98 (the last code change; 3b8f54b2 changes one test assertion) unless noted:

• policyengine_uk/tests/test_dataframe_dataset.py: 71 passed in 160 s, peak RSS 0.7 GB, with the locked versions (pandas 2.3.1, policyengine-core 3.32.27).

• The same file with the versions CI resolves (pandas 3.0.6, numpy 2.5.4, policyengine-core 3.33.2), at 3b8f54b2: 71 passed in 158 s, peak RSS 0.8 GB.

• test_input_records.py, test_simulation_dataset_sources.py, test_dataset_sources.py, test_no_economic_assumptions.py, test_household_income_decile.py: 22 passed (at 58af8c52; only tests and a docstring changed after).

• Mutation check: 36 deliberately broken versions, each applied to a scratch copy of the package, and the test file must fail on each. All 36 are killed. 34 were killed at 58af8c52. The other two showed gaps in the tests, not in the code: a fraction inside a mixed identifier column, and a decile whose households all have zero weight. 4892db98 adds a test for each, and both mutants are now killed. Without the zero-weight guard the formula raises ZeroDivisionError. The mutants cover:

  • group order, and which identifier column wins;
  • the float guards and their bounds, in typed and in mixed columns;
  • agreement between the two group columns, at integer, float32 and CSV precision;
  • the member-agreement, missing-identifier, repeated-person, whole-number and 64-bit checks;
  • substring matching of identifier and value columns;
  • the replay of month columns: only other periods, no delete first, or nothing replayed;
  • months taken as years, only the first year loaded, identifier columns adding data years, identifiers not set for their own years;
  • dataset left as None, and the rounding warning removed;
  • the empty-decile and zero-weight-decile guards in the budget formula.

  Earlier runs on earlier heads found two more gaps the same way: a January column taken as the whole year, and a float in a mixed person_ column that names the wrong group.

• The docs example, run as written.

• test_subsample.py from Give UK simulations a working subsample #2265, on a local merge of the two branches at 58af8c52, with the replacement test (in "Interaction with Give UK simulations a working subsample #2265") applied: 38 passed.

Not run locally: the full suite and the docs build (CI runs both).

🤖 Generated with Claude Code

[MaxGhenis]

Hand-off to the UK hub (session local_1d38e8ba-7439-4c06-a8ef-b8bdee97c760).

• Head: 3b8f54b254562e5bf579412a5e5e12ed05c886cb.
• Independent review: APPROVE at this exact head. Round 6, Opus 5.5 (subfleet uk2266-review-r6-opus), saved at ~/reviews/uk-2266/review-r6.md.
  • Rounds 1 to 5 are in the same folder (review-r1.md … review-r5.md).
  • Round 1 (GPT-6.1 Sol) and round 2 requested changes, all addressed. Rounds 3 to 6 approved.
  • The Sol lanes were saturated after round 1, so rounds 2 to 6 ran on Opus, by reading. The execution evidence is local.
• Passed: Lint, docs build, smoke imports on 3.11 to 3.14. Locally:
  • test_dataframe_dataset.py (71 tests) with the locked versions, and with CI's versions (pandas 3.0.6, core 3.33.2);
  • the targeted files;
  • 36 of 36 mutants killed.
• Pending: CI Test on this head, about an hour. The run on 58af8c52 failed only the frame-unchanged assertion, from an EnumArray class difference under core 3.33. 3b8f54b2 fixes it, and the fix passes locally under CI's versions.
• Mergeable: yes. No CHANGES_REQUESTED review. Nothing in this PR needs Max: it is a bug fix plus a guard in three contrib budget formulas, which leaves a full dataset's results unchanged.
• Merge method: merge commit, per the brief for Simulation(dataset=DataFrame) fails when a household or benefit unit has more than one person #2266.
• Depends on: nothing.
• Interacts with Give UK simulations a working subsample #2265 (open):
  • Whichever lands second must replace test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample in test_subsample.py. It asserts dataset is None, which this PR changes.
  • The replacement patch is in this PR's body and at ~/reviews/uk-2266/pr2265-test-update.patch.
  • On a local merge of both branches with the patch applied, test_subsample.py passes (38 tests).
  • Give UK simulations a working subsample #2265's last Test run was cancelled, so it has no green run that predates this PR.
• Follow-ups filed: Wealth and consumption shares divide by a weighted total that can be zero #2296 (zero weighted totals in the share formulas) and to_input_dataframe() drops dataset inputs for variables with formulas, such as state_pension_reported #2297 (to_input_dataframe() drops dataset inputs for formula variables).
• Holds or decisions: none.

[MaxGhenis]

UK hub: landed on the gates at head 3b8f54b254562e5bf579412a5e5e12ed05c886cb (merge commit).

• Checks: gh pr checks exit 0 on this head, every check passing, including Test (the landing script merges only after a complete green run).
• Mergeable: MERGEABLE; not a draft; no CHANGES_REQUESTED review.
• Independent review: Opus 5.5 round 6, APPROVE at this exact head (~/reviews/uk-2266/review-r6.md). Rounds 1 (GPT-6.1 Sol) and 2 requested changes, which were addressed; rounds 3 to 6 approved. The Sol lanes were saturated after round 1, so the later rounds ran on Opus.
• Impact: no Enhanced FRS run. Full-dataset results are unchanged, which the hub checked by reading the diff.
  • The three contrib budget formulas (nhs_budget_change, education_budget_change, other_public_spending_budget_change) now skip a decile only when households_per_decile.get(i, 0) > 0 fails, which happens only for a decile with no weighted households. The Enhanced FRS populates all ten deciles.
  • The loader change applies only to simulations built from a DataFrame (build_from_dataframe).
• Axiom: n/a (simulation infrastructure).
• For Give UK simulations a working subsample #2265 (open): whichever of Give UK simulations a working subsample #2265 and this PR lands second must replace test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample, which asserts dataset is None. The replacement patch is in this PR's body.
• Follow-ups filed: Wealth and consumption shares divide by a weighted total that can be zero #2296, to_input_dataframe() drops dataset inputs for variables with formulas, such as state_pension_reported #2297.