Give UK simulations a working subsample

[MaxGhenis]

Fixes #2254

What was wrong

Simulation.subsample() was inherited from policyengine-core and raised TypeError: Simulation.build_from_dataset() missing 1 required positional argument: 'dataset'. I reproduced it on main (2.125.2) with a six-household in-memory dataset, on policyengine-core 3.32.27 (the locked version) and 3.33.1.

Why the signature was not the whole problem

I read core's subsample and flat-file loader (master 8fdf59e9, 3.33.2) against the UK builders, and ran the issue's option 3:

1. Option 3 fails. With build_from_dataset(self, dataset=None) falling back to self.dataset, core's subsample hands the UK builder the flat-file Dataset it rebuilt with Dataset.from_dataframe. The UK builder calls dataset.load_dataset(), which reads only the table, time-period-array and array formats, and raises UnboundLocalError for a flat file. I ran this on core 3.32.27 and 3.33.2. The UK builder would also rebuild every year from the first one, projected forward, rather than keep each year's inputs.
2. Memberships. Core reads each person's household and benefit unit from the stored person_household_id and person_benunit_id inputs. On main the UK model stores both as 32-bit floats, which hold integers exactly only up to 2**24 (16,777,216). In the full-size UK dataset the issue's reproduction used, household identifiers reach about 900 million: its 535,080 distinct household identifiers are only 76,551 distinct values as float32, and its 618,980 benefit unit identifiers are 544,894. A rebuild from those columns would merge households and benefit units. (Store person_household_id and person_benunit_id as integers #2273 proposes storing them as integers. This PR does not depend on it.)
3. Baseline. Core rebuilds branches["baseline"]. A UK simulation's baseline is a separate Simulation, which core's method does not touch, so a reform and its baseline would hold different households.
4. Seed. With no seed, core reads self.dataset.name. UK datasets have no name.
5. Weights. Core rescales every column whose name contains weight__. In the UK that would include household_bus_fare_age_weight if it were supplied as an input; it shares out bus fare spending and is not a sampling weight.

So this is option 1 in the issue: a UK subsample. It does not call core's method and does not change build_from_dataset. Nothing needs to change in policyengine-core.

What the UK method does

Same signature as core's: subsample(n=None, frac=None, seed=None, time_period=None, quantize_weights=True).

• Draws n households with replacement and keeps every household drawn at least once, with its benefit units and people. Only households with a positive weight in time_period can be drawn: with quantize_weights=True (the default), each with probability proportional to that weight; without it, each with equal probability. A period before the first one the simulation was given household weights for has none, so the first given weights are used.
• Takes households, benefit units and people from the populations (ids and members_entity_id), which are exact, not from the stored identifier inputs.
• Keeps every supplied input (from the dataset or set_input) for the sampled records, for every period. It reads each input from the storage tier (memory or disk) the holder records it in, so a value calculated later and cached over the input is not kept. Calculated values are dropped and recalculated on the sample, and the tracer starts again.
• Replaces each supplied weight so that its total over its own entity is unchanged in every period. The weights are listed (household_weight, benunit_weight, person_weight, original_weight, raw_person_weight), and a test fails if a new variable ending in weight is not classified.
• Replaces the baseline, and its baseline, with copies subsampled with the same draw. It does not change a baseline in place: core's clone shares a separate baseline between a simulation and its clone, so subsampling a clone would otherwise cut down the source's baseline. The copy is made as get_branch makes a branch, so it shares the old baseline's arrays as read-only views instead of copying them, and it rebuilds its own at once. (A plain clone copied every cached array first, which took the full-size run below past a 20 GB memory limit.) The old baseline stays in memory until nothing holds it.
• Replaces dataset with the rows of the dataset the simulation was built from for the sampled records, weights rescaled the same way. The original object is not changed (it can be the shared entry of the URL cache). Like the dataset before subsampling, it holds whole years only, and not values set with set_input after the build.
• Works out every subsample before changing anything, so a refusal leaves every simulation as it was.

Differences from core's method, all deliberate:

• The draw uses a local numpy.random.RandomState seeded from a SHA-256 digest of the seed, so a seed draws the same households in every process, and the global random state is not touched. A seed therefore does not pick the same households as it would in core.
• With no seed, a fixed seed is used.
• A simulation built from a situation, one given no household weights or a negative weight, a branch, a benefit unit spread over two households, and a baseline holding other households are refused with a ValueError.
• The baseline object is replaced, not changed in place (see above).

Invariants

Stated in policyengine_uk/tests/test_subsample.py and tested with Hypothesis for generated datasets of one year (which the model projects) and of several years whose weights change differently from household to household (two explicit examples make sure the shares change), for every n, seed, sampling period and quantize_weights:

1. Size: between 1 and n households are kept, each with a positive weight in the sampling period.
2. Weight totals: the total household weight of every year is unchanged, to float32 precision, and so is the total of every other supplied weight over its own entity.
3. Quantized weights: a household's weight is a whole multiple of total / n, the multiples sum to n, and the same draws set every year's weights. With quantize_weights=False, it is proportional to its old weight times a whole number of draws.
4. Links: a kept household or benefit unit has exactly the people it had, a kept person keeps their household and benefit unit, and no other record is kept. This holds for neighbouring identifiers near 900 million, which float32 cannot tell apart, and whatever the stored membership inputs hold.
5. Inputs: every supplied value other than a weight is unchanged for the kept records in every period, and the subsample supplies the same variables and periods as before.
6. Determinism: the same seed keeps the same households with the same weights.
7. Baseline: the baseline is a copy holding the same households and weights, with its own policy. The baseline held before is unchanged, so subsampling a clone leaves the source and its baseline as they were.
8. Dataset (differential): dataset holds the same records as the populations. For a simulation whose inputs are all in its dataset, a simulation built from that dataset gives the same results as the subsampled simulation. This compares the two implementations of the subsample: the one that cuts the stored inputs and the one that cuts the dataset tables.
9. Unbiasedness: over 3,000 seeds, a household's mean quantized weight in the sampling period is its old weight (within 3%).
10. Trace and tiers: the tracer starts again, and an input stored on disk is kept rather than a value cached over it in memory. On core 3.33 and later, a differential test checks the tier rule against core's own Simulation._get_recorded_input_array (it skips on 3.32, which has no such method).

Intended exceptions, each pinned by a test:

• In a year other than the sampling one, a household's mean new weight is its share of the sampling year's total times that year's total. That equals its old weight only if its share is the same in both years.
• quantize_weights=False keeps the total exactly, which makes the new weights a ratio estimate. I measured the bias: about 6% at 40 draws and under 1% at 400, for a seven-household example. This is the same definition as core's.
• An input the dataset does not hold stays in the simulation and out of dataset: one set with set_input after the build, or one for part of a year, which a DataFrame can give and a dataset of whole years cannot hold. A weight set with set_input on the simulation alone is rescaled to its own total there, while the baseline and the dataset keep the dataset's weights.

Tests

• policyengine_uk/tests/test_subsample.py: 52 tests. They cover the reported call, each failure mode above, each refusal, the findings of two review rounds, and the properties.
• Mutation check at 01f5d051: I made 22 single changes to utils/subsample.py and utils/supplied_inputs.py and ran the file against each. 21 were caught. The survivor removes the reset of the input record, which _invalidate_all_caches repeats later in the same function, so it changes nothing observable.
• Local runs at 82e0c499 (the branch with main merged in), one process each, of the new file plus test_input_records.py, test_simulation_dataset_sources.py and test_dataframe_dataset.py:
  • core 3.32.27 (locked), Hypothesis 6.168.2, pandas 2.3.1, NumPy 2.1.3: 140 passed, and the core 3.33 differential test skipped.
  • core 3.33.2, Hypothesis 6.169.0 (what CI installs), pandas 3.0.6, NumPy 2.5.3: 141 passed.
• Full-size check under core 3.33.2 at 5dd27677 (the code has not changed since; later commits add a test and merge main), on the full-size UK dataset the issue's reproduction used (535,080 households): subsample(n=10_000, seed=0, time_period=2026) took 6.7 seconds and kept 8,660 households. From the weights alone, the expected number of distinct households in such a draw is 8,634. Memberships matched the dataset's 64-bit identifiers exactly, the 1,260 supplied variable-periods were the same as before, the baseline was a new copy with the same records and the old baseline still held every household, every year of dataset held the same records, and each year's household weight total was unchanged to 3 parts in 100 million. The same held for quantize_weights=False (9,864 households, 4.0 seconds).
• Memory: that run stayed under the 20 GB process limit that the plain-clone version exceeded. On a synthetic dataset of 60,000 households with 922 supplied variable-periods, the subsample raised the process's peak resident memory from 2,755 to 2,786 MiB (about 1%).
• Not run locally: the full suite. CI runs it. Earlier CI runs of this PR failed only test_dataframe_dataset.py::test_an_exported_frame_loads_back_and_exports_itself, which was failing on main (Stop deciding a household's UC health element rate by its position #2251 with Load a DataFrame simulation input as the dataset it describes #2270) until Export the identifiers for every year an input frame names, so it round-trips #2349; this branch now includes that fix.

Found on the way

• Simulation(dataset=DataFrame) fails when a household or benefit unit has more than one person #2266: Simulation(dataset=<DataFrame>) failed for households of more than one person. Fixed by Load a DataFrame simulation input as the dataset it describes #2270, merged; this branch includes it.
• person_household_id and person_benunit_id are stored as float32, which rounds large identifiers #2267: person_household_id and person_benunit_id are stored as 32-bit floats. Store person_household_id and person_benunit_id as integers #2273 is open. This PR reads memberships from the populations and does not depend on it.
• Elasticity-based labour supply responses compare a reform with a copy of itself #2268: the elasticity-based labour supply responses compared a reform with a copy of itself. Measure labour supply responses against the reform's baseline #2277 is open; this PR does not touch that code.
• supplied_input and supplied_input_periods count a value cached over a supplied input as supplied #2307: supplied_input and supplied_input_periods report a value calculated and cached over an input as supplied. Count only supplied inputs that are still stored in their tier #2308 is open and stacked on this branch: it builds on stored_supplied_value from this PR.

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

🤖 Generated with Claude Code

Reproducer: policyengine_uk/tests/test_subsample.py::test_subsample_runs_on_a_dataset_simulation, policyengine_uk/tests/test_subsample.py::test_subsample_is_the_uk_method_not_policyengine_cores, policyengine_uk/tests/test_subsample.py::test_identifiers_beyond_float32_precision_keep_their_households, policyengine_uk/tests/test_subsample.py::test_memberships_come_from_the_populations_not_the_stored_identifiers, policyengine_uk/tests/test_subsample.py::test_zero_weight_households_are_never_kept, policyengine_uk/tests/test_subsample.py::test_quantized_weights_count_draws_and_keep_every_years_total

[MaxGhenis]

Merge-order note from #2273 (fixes #2267, stores person_household_id and person_benunit_id as int32).

On a local merge of this branch (fa7b0f0a) with #2273, test_identifiers_beyond_float32_precision_keep_their_households fails on one line, its precondition:

assert len(set(stored.tolist())) < len(households)

With #2273 the stored values tell the 900-million households apart (6 < 6 fails). The test's own checks, that subsample keeps whole households with their benefit units and people, all pass on the merge. So did the rest of test_subsample.py.

Whichever PR lands second should:

• change that line to assert the stored memberships are exact, for example assert len(set(stored.tolist())) == len(households), and reword the comment above it;
• update the sentence in policyengine_uk/utils/subsample.py's module docstring that says the UK model stores person_household_id and person_benunit_id as 32-bit floats (and invariant 4's "identifiers that 32-bit floats cannot tell apart" in the test docstring, if wanted). Reading the populations stays the right approach either way.

🤖 Generated with Claude Code

[MaxGhenis]

UK hub, for this PR's owner: test interaction with #2270.

#2270 (DataFrame inputs load as the dataset they describe) is armed to land once its Test passes. After it, a DataFrame-built simulation has a dataset, so test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample in this PR's test_subsample.py, which asserts dataset is None, will fail.

• If Load a DataFrame simulation input as the dataset it describes #2270 lands first: merge main here and apply the replacement in Load a DataFrame simulation input as the dataset it describes #2270's body, also at ~/reviews/uk-2266/pr2265-test-update.patch. Load a DataFrame simulation input as the dataset it describes #2270's author reports that test_subsample.py passes (38 tests) on a local merge of both branches with the patch applied.
• If this PR lands first: Load a DataFrame simulation input as the dataset it describes #2270 carries the same patch.

This PR's last Test run was cancelled, so it has no green run yet. Merging main also brings the 120-minute Test limit (#2258).

[MaxGhenis]

For coordination: #2308 (fixes #2307) is stacked on this branch and builds on stored_supplied_value.

• It makes supplied_input and supplied_input_periods use the same tier rule, so they no longer count a value cached over a supplied input as supplied.
• It moves the tier lookup in stored_supplied_value into a private _supplied_value_storage, so that supplied_input_periods can check a tier without reading a value that may be on disk. What stored_supplied_value returns is unchanged.

Nothing in this PR needs to change for #2308. If stored_supplied_value changes here before merging, #2308 will follow it. #2308 moves to main when this PR merges.

[MaxGhenis]

CI run cancelled to free runners (UK hub, 2026-10-11 23:30Z). The PolicyEngine org is at GitHub's limit of 60 concurrent jobs, with about 360 queued. Fleet ops asked the hub to let the Budget-critical #2305 take runners first. This PR's queued run had not started, so I cancelled it. Nothing is wrong with the PR.

Please don't re-run it or push just to retrigger CI. The hub re-runs this run itself once #2305 has runners. The record is ~/reviews/uk-hub/ci-yield-2026-10-11.json; see notice 17 in ~/reviews/uk-hub/notices.md.