Skip to content

Give UK simulations a working subsample - #2265

Open
MaxGhenis wants to merge 22 commits into
mainfrom
fix-subsample-2254
Open

MaxGhenis wants to merge 22 commits into
mainfrom
fix-subsample-2254

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Oct 10, 2026 •

Copy link
Copy Markdown
Collaborator

Fixes #2254
Reproducer: policyengine_uk/tests/test_subsample.py::test_subsample_runs_on_a_dataset_simulation
axiom: n/a: simulation infrastructure, no policy rule change

Status: waiting for policyengine-core holder resync (policyengine-core#644). Four supplied/model restoration regressions stay strict expected failures until core releases the fix. main now locks policyengine-core 3.34.1, whose Holder still keeps the Variable it was built with, so the four cases still fail at 44f6deb7 (see Tests).

What was wrong

UK simulations inherited Simulation.subsample() from policyengine-core, and it 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 (then 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. Those hold whatever a dataset or set_input stores there, not necessarily the populations' memberships. On main at 2.125.2 the UK model stored 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, the 535,080 distinct household identifiers came to only 76,551 distinct float32 values, and the 618,980 benefit unit identifiers to 544,894. Store person_household_id and person_benunit_id as integers #2273 has since moved both to 32-bit integers. The UK method reads memberships from the populations either way.
  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 a caller supplied it as an input; it shares out bus fare spending and does not weight a sample.

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. Separately, core's holders keep the variable they were built with after a runtime reform replaces its definition. A UK workaround that rebound holders when core invalidated caches made core reject an existing neutralised CGT input replacement and came too late for writes inside a restoring reform, so this branch does not carry it. The four supplied/model restoration regressions stay strict expected failures with reason needs policyengine-core holder resync, with every assertion kept.

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 holds household weights for has none, so the method samples by the first weights it holds.
  • Takes households, benefit units and people from the populations (ids and members_entity_id), which hold them exactly, 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 it does not keep a value the model calculated later and cached over the input. It drops calculated values, which the model recalculates on the sample, and starts the tracer again.
  • Replaces each supplied weight so that its total over its own entity stays the same in every period. WEIGHT_VARIABLES lists the weights (household_weight, benunit_weight, person_weight, original_weight, raw_person_weight), and a test fails if a new variable ending in weight appears in neither list.
  • 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 method makes the copy as get_branch makes a branch, so the copy shares the old baseline's arrays as read-only views instead of copying them, and 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. It leaves the original object as it was (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. Each year keeps the dataset's energy price basis, ISA income contract and savings income series.
  • Keeps what the simulation records beyond its inputs, which a rebuild of the populations forgets:
    • values the model stored as its own with set_model_value (the Universal Credit health element reform stores them for every simulation, for 2026 to 2032), which it stores back with set_model_value, so supplied_input still does not read them as supplied;
    • stored inputs and model values of neutralised variables, which it writes back to the holder directly because set_input rejects a neutralised variable;
    • where each capital gains input came from (supplied, projected or default, by period), and the dataset's annotation of projected capital gains inputs, which it carries to the subsampled dataset;
    • the dataset's ISA income contract for each year and the series its savings income grows by, which set how the model represents and projects savings interest.
  • 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 method leaves the global random state alone. A seed therefore does not pick the same households as it would in core.
  • With no seed, the method uses a fixed seed.
  • The method refuses, with a ValueError, 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.
  • The method replaces the baseline object rather than changing it in place (see above).

Invariants

policyengine_uk/tests/test_subsample.py states these, and tests them 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: the sample keeps between 1 and n households, each with a positive weight in the sampling period.
  2. Weight totals: the total household weight of every year stays the same, to float32 precision, and so does the total of every other supplied weight over its own entity.
  3. Quantized weights: a household's weight equals its number of draws times the year's total over the draws of households with a positive weight that year: in the sampling year, a whole multiple of total / n, with the multiples summing to n. The same draws set every year's weights. With quantize_weights=False, it stays proportional to its old weight times a whole number of draws. Either way, a zero weight stays zero in every year, and the method rescales each supplied weight to its own entity's total (so a supplied person weight need not equal the new weight of the person's household).
  4. Links: a kept household or benefit unit has exactly the people it had, a kept person keeps their household and benefit unit, and the sample holds no other record. 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 stays the same 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 becomes a copy holding the same households and weights, with its own policy. The baseline held before stays as it was, 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 all come from 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 comes within 3% of its old weight.
  10. Trace and tiers: the tracer starts again, and the sample keeps an input stored on disk 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).
  11. Model values and dataset records: a value the model stored as its own stays the model's, a supplied value stays supplied, and the capital gains input origins reported by capital_gains_projection_metadata stay the same, including the default origin of a neutralised input. A simulation built from the subsampled dataset reports the same origins as one built from the full dataset. The ISA income contract of every year, the savings income series and each year's energy price basis stay the same on the simulation, its baseline copy and the dataset, and savings interest past the dataset's last year matches the full simulation for the same people.

Intended exceptions, each pinned by a test:

  • In a year other than the sampling one, a household's mean new weight equals its share of the sampling year's total times that year's total. That matches its old weight only if its share stays the same in both years.
  • quantize_weights=False keeps the total exactly, which makes the new weights a ratio estimate. I measured the bias at ab843c80: about 6% at 40 draws and under 1% at 400, for a seven-household example. Core uses the same definition.
  • 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. The method rescales a weight set with set_input on the simulation alone to its own total there, while the baseline and the dataset keep the dataset's weights.

Merge with main

29504123 merges main at 47a2817e (2.175.2, policyengine-core 3.34.1). Three files conflicted:

  • policyengine_uk/simulation.py: both sides added an import at the same place. The merge keeps main's measurement_branches import and this branch's subsample_simulation import.
  • policyengine_uk/utils/supplied_inputs.py: both sides edited the module docstring's last paragraph. The merge keeps main's wording for what a rebuild forgets ("when a rebuild installs populations that do not hold it") and this branch's paragraph on supplied_input_arrays.
  • policyengine_uk/tests/test_input_records.py: both sides added imports and tests. The merge keeps main's ten rebuilt-population tests and this branch's restoring-reform test, joins the imports, and renames the test's local period to input_period so that it can use main's period import.

main does not do anything this branch does, so the merge drops nothing. main added two dataset records that the rebuild in subsample resets: each year's ISA income contract and the savings income series. Without them, a subsample of a dataset under the legacy ISA contract projected savings interest in 2035 from the additive default (−110.29 instead of −100.00 for one person in the test). 44f6deb7 restores both on the simulation, and carries them and the energy price basis to each year of the subsampled dataset.

Tests

  • Merge with main at 47a2817e: at 44f6deb7, with the locked environment (policyengine-core 3.34.1) and pytest -n 4, 25 files gave 874 passed, 4 strict expected failures, 0 failed, 0 skipped in 43 minutes: test_subsample.py (57 passed / 2 xfailed), test_input_records.py (13 passed / 2 xfailed), every other file that calls subsample or the supplied-input helpers, the CGT projection and response-default files, the four UC health element files, dataset sources, the DataFrame round trip, input DataFrame export, reform-input lifecycle, reform representations, the ISA income contract and the energy price basis files. The four expected failures are the restoration cases above. policyengine_uk/tests/code_health gave 3,163 passed at 44f6deb7. At e804e8bc (which only adds a test), the two new tests pass. At e804e8bc, ruff format --check . (2,229 files), ruff check ., ruff check --select F811,F821 policyengine_uk, git diff --check against main, .github/check_changelog_fragments.py and .github/check_fix_pr.py on this title and body pass.
  • Revert checks at 44f6deb7, each in a scratch copy, running test_how_the_dataset_says_its_amounts_grow_is_kept: without the restore of the ISA contracts and savings series, it fails on the contracts; without the dataset metadata, it fails on the subsampled dataset's ISA income contract; the unchanged copy passes.
  • Mutation checks at 44f6deb7 of the post-merge code, each in a scratch copy, running test_a_model_value_stays_the_models and test_where_capital_gains_inputs_came_from_is_kept: dropping the capital gains origins, dropping the dataset's capital gains annotation and storing model values with set_input each fail one of them. Dropping the defaulted capital gains periods passed both, so e804e8bc adds test_a_neutralised_capital_gains_input_still_reads_as_its_default, which fails with that mutation and passes without it.
  • Earlier, before this merge (2026-10-23): the restoration regressions check writes inside restoration both directly and after subsampling; direct cases and the unchanged CGT test pass, while the r2 restoration cases and sampled write cases need core holder resync. At e4c54d7fa, the locked pytest -n 2 run passed 244 tests with 4 strict expected failures, 0 unexpected failures and 0 skips: subsample (56 passed / 2 xfailed), input records (3 passed / 2 xfailed), the whole CGT projection file (115 passed), input DataFrame export (25 passed), reform-input lifecycle (17 passed), and reform representations (28 passed).
  • Mutation check at 01f5d051: 22 single changes to utils/subsample.py and utils/supplied_inputs.py, each run against the test file. The tests caught 21. The survivor removes the reset of the input record, which _invalidate_all_caches repeats later in the same function, so it changes nothing observable.
  • At 7b4c3c28 (main at 7918fc2a merged in), with 4 workers, the subsample, input records, dataset sources, DataFrame round trip, four UC health element, CGT projection and CGT response default files gave 377 passed on core 3.33.2 and on core 3.33.12.
  • Full-size check under core 3.33.2 at 5dd27677, before any merge with main (not rerun since), 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 comes to 8,634. Memberships matched the dataset's 64-bit identifiers exactly, the 1,260 supplied variable-periods matched those before, the baseline became a new copy with the same records while the old baseline still held every household, every year of dataset held the same records, and each year's household weight total stayed the same to 3 parts in 100 million. The same held for quantize_weights=False (9,864 households, 4.0 seconds).
  • Memory at 5dd27677: 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, which CI runs, and tests that need survey data.

Aggregate impact

CI's aggregate impact report at e804e8bc, merged into main (040e9a5f07 against 4d3118a0ca, enhanced_frs_2024_25.h5, policyengine-uk-data 1.58.1, fiscal years 2025-26 to 2031-32), found:

Checks: 0 broken by this change, 1 failing on both sides, 0 fixed. Moves flagged for review: 0

No household's net income changes in any year.

The one check failing on both sides, sign:vat (VAT never goes negative), concerns code this PR does not touch. A full microsimulation does not call subsample, and the changes to supplied_inputs.py only add helpers.

Found on the way

🤖 Generated with Claude Code

MaxGhenis and others added 3 commits October 10, 2026 15:36
Simulation.subsample was inherited from policyengine-core and raised
TypeError: it rebuilds through build_from_dataset() with no arguments,
and the UK builder takes the dataset to load.

Making that argument optional would not have been enough. Core's method
reads each person's household and benefit unit back from the stored
person_household_id and person_benunit_id inputs, which the UK model
holds as 32-bit floats, so identifiers above 2**24 would merge
households. It also leaves Simulation.dataset as a core Dataset, while
UK formulas read dataset.years, and it rebuilds a baseline branch, while
a UK baseline is a simulation of its own.

The UK method takes households, benefit units and people from the
populations, keeps every supplied input for the sampled records,
rescales each supplied weight to keep its total, and subsamples the
baseline and the dataset with the same draw.

Fixes #2254

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A result calculated before subsampling is calculated again on the sample,
and the record of an input that a calculated value replaced is not kept.
Found by mutating the subsample module: removing the cache invalidation
went unnoticed until these assertions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The simulations usage page gains a short section on subsampling, since
the method now works for UK simulations. A simulation built from a
DataFrame keeps no dataset object, so test that subsampling it leaves
dataset as None and still cuts the baseline to the same households.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
CI formats Python blocks in Markdown with the latest ruff, which the
locked ruff does not check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

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

Copy link
Copy Markdown
Collaborator Author

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.

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 and others added 4 commits October 11, 2026 09:17
Address the independent review of #2265:

- Subsampling a clone changed the source simulation's baseline, because
  policyengine-core's clone shares a separate baseline. Replace each
  baseline with a subsampled copy instead of changing it in place.
- A calculated value cached in memory over an input stored on disk was
  kept as the input. Read each supplied input from the tier the holder
  records it in, as Holder._stores_user_input decides.
- A trace kept values calculated on the replaced populations. Give each
  simulation a new tracer, shared as before.
- Say that dataset holds the rows of the dataset the simulation was
  built from, without later set_input values, and that weights are
  unbiased in the sampling period only. Generate multi-year datasets
  whose weights change differently by household in the property test.

Also adapt to main: a DataFrame-built simulation now has a dataset, and
test memberships without assuming how the identifiers are stored.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Address the second independent review of #2265:

- Copying a baseline with clone copied every cached array before the
  subsample discarded them, which took a full-size run past a 20 GB
  memory limit. Announce the copy as get_branch does, so it shares the
  arrays as views and copies none.
- A negative household weight could leave a negative weight in the
  sample. Refuse negative sampling weights, and negative values in any
  weight being rescaled.
- An input for part of a year, which a DataFrame can give, stays in the
  simulation and out of the dataset of whole years. Say so, and test it.
- Give the property test two explicit cases whose weight shares change
  from year to year, which the generated cases had not reached.
- Qualify the docstring's sampling and unbiasedness claims by
  quantize_weights, and describe the first-weights fallback.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Mutation testing showed the rescaling check alone also caught a negative
household weight, so removing the draw's own check went unnoticed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

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

Copy link
Copy Markdown
Collaborator Author

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.

@vahid-ahmadi vahid-ahmadi left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fine once it is merged with main and keeps the model-value record across the rebuild.

@MaxGhenis, the approach is right. The memberships come from the populations, the baseline and the dataset are cut with the same draw, and the draw is deterministic and independent of process. The quantized weights are unbiased in the sampling year: E[count] = n·w/W, so E[count·W/n] = w. Two things need doing before merge, both from main moving on since 82e0c4995.

  1. must: merge main. policyengine_uk/utils/supplied_inputs.py conflicts in two places, the module docstring and the end of the file. Both sides add text there, so keeping both works: I did that locally and the import is fine. Run ruff format afterwards, because the joined file gets an extra blank line.
  2. must, after the merge: subsample turns a scenario's model value into a supplied input. Main now has set_model_value (supplied_inputs.py), and scenarios/uc_reform.py:163 uses it for uc_LCWRA_element. _apply (utils/subsample.py:336) calls build_from_ids, which on main calls build_from_populations, which calls forget_model_values. The loop at subsample.py:353 then stores every kept value with set_input, which main treats as supplied. On a local merge, with the six-household dataset() from test_subsample.py, I called set_model_value(sim, "uc_LCWRA_element", 2026, ...) and then subsample(n=4, seed=0, time_period=2026). supplied_input(benunit, "uc_LCWRA_element", 2026) was None before the subsample and returned the stored array after it. _model_values went from four keys to none. So a subsampled scenario simulation reads its own adjustment as a user's entry. That breaks the "same results as the same households in the full simulation" property for formulas that ask supplied_input. Suggested fix: in _plan, cut each _model_values array by the same entity positions; in _apply, write the cut arrays back into simulation._model_values after the set_input loop. Add a test that supplied_input is still None after the subsample.
  3. nit: rescale_weights (subsample.py:139). With quantize_weights=True, a record whose own weight is 0 gets a positive weight if its household was drawn: rescale_weights([0, 100, 100], [1, 1, 0]) gives [100, 100, 0]. In the sampling year this can't happen for households. It can for a household with zero weight in another year, or for a supplied person or benefit unit weight of 0. If that is intended, say so in the docstring. If not, use counts * (weights > 0).
  4. nit: each supplied weight is rescaled to its own entity's total separately (entity_counts in _plan). If a dataset ever supplies person_weight alongside household_weight, a person's new weight will not equal their household's, though the model defines it that way. None of the dataset code here supplies it today, so a line in the docstring is enough.

What I checked

  • Fresh uv venv at 82e0c4995. I then merged origin/main locally (the merge is not pushed), resolved the conflict as in point 1 and re-synced to the merged lock.
  • On that merge: test_subsample.py, test_input_records.py, test_simulation_dataset_sources.py and test_dataframe_dataset.py gave 141 passed. They don't cover point 2.
  • The model-value probe in point 2, and rescale_weights called directly for point 3.
  • ruff check and ruff format --check clean on utils/subsample.py, simulation.py and tests/test_subsample.py.
  • CI at 82e0c4995: Test, Lint, Fix evidence and the smoke imports passed. There is no run on a merge with current main, because GitHub reports the PR as conflicting.
  • Not checked: the full-size dataset run and the memory figures in the PR body. My #2193 adds apply_reform to simulation.py, in a different place from subsample, so there is no textual clash.

🤖 Generated with Claude Code

@juaristi22

Copy link
Copy Markdown
Collaborator

A note from #2419, the CGT retention-rate default, tracked in the checklist on PolicyEngine/policyengine.py#602.

Once #2419 lands, every reform simulation measures marginal rates on gains against its baseline, at the reform's own gains before the response. So the baseline has to hold the same people. A subsample() that leaves self.baseline at full size would make reform calculations fail. #2419 raises a clear error in that case ("The simulation has N people and its baseline M…"), but subsample() should subsample the baseline too, or rebuild it from the subsampled data.

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

UK hub takeover for the Budget fix sprint (Max, 21 Oct). The owner session local_eb7a5ce3 has been idle since mid-October. The hub woke it on 21 Oct at 23:58Z, and it did not start. To land this before the Budget model locks (the first pe-uk release cut by about 12:00Z Thu 22 Oct), the hub has taken it over as build 20261021-211252-sprint-2265, which merges main, re-tests, pushes as a fast-forward and marks it ready. The hub then reviews it and lands it through the merge queue.

If the owner session wakes, please don't push to this branch; comment here instead, so there's only one writer.

🤖 Generated with Claude Code

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

UK hub hand-off for review P2 at comment 6335788910: fix committed locally; validation and push blocked by host limits.

  • Local fix head: 9a91736f39990b87317116a2f59ebdc8c9054749, on top of conflict-free merge fc537e4e3 of origin/main at fe0b99a36. Existing PR author and history preserved; remote head remains db3a38e07.
  • Neutralised holders retain canonical stored input arrays, input records and model markers. Added supplied/model ANI regressions covering unchanged tax results and restoration, plus a changelog update. Ordinary replay is unchanged.
  • Passed: touched-file Ruff formatting/lint, git diff --check, and uvx ruff format --check . (1940 files).
  • Not run: locked environment setup and foreground pytest -n 2. After waiting, load1 was 172.65 and disk 54.98 GB, outside the requested limits (110 and 60 GB). No test suite or microsimulation started; the foreground resource wait was stopped before hand-off. The one permitted push remains unused.
  • Recovery: .job/HANDOFF.md, .job/PROGRESS.md, .job/commits.bundle and PR-body draft in /Users/maxghenis/.subfleet/worktrees/20261022-122605-sprint-2265-p2. Shared Git metadata is read-only, so local commits use git --git-dir=.job/git --work-tree=.. Bundle verified successfully.
  • Open items: restore host capacity; run locked targeted tests; push once to fix-subsample-2254; dispatch the hub's independent review at that final head. No independent review, aggregate-impact run, data build or merge performed. This is not Budget-material.

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

UK hub (Claude session, for Max): independent review r2 (GPT-6.1 Sol) at 9a91736f3: REQUEST_CHANGES. A hub fix build is running (fix-2265-r3).

[P2] Restoring a neutralized variable still returns zero after subsampling — policyengine_uk/utils/subsample.py:385-392; policyengine_uk/tests/test_subsample.py:584-586. Rebuilding creates holders bound to the neutralized variable instance. RestoreIncome replaces the system's variable, but the holder keeps its old reference: the system reports is_neutralized=False, the holder reports True, and Holder.get_array() returns zero despite retaining the £10,000 input. Both new regression variants fail. A separate probe confirms restoration works without subsampling; rebinding the sampled holder restores £10,000 in both supplied/model cases. Rebind holders to restored variable definitions and retain these assertions.

Everything else in the round-one fix holds: canonical stored inputs, supplied/model ownership and ordinary replay. CI at this head fails exactly the two new regressions.

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

UK hub (Claude session, for Max): independent review r3 (GPT-6.1 Sol) at 74b0c65d7: REQUEST_CHANGES. r2's read-after-restore case is fixed, but rebinding holders only at cache invalidation is too late:

  • Regression: test_capital_gains_projection.py::test_neutralized_supplied_gains_report_default_without_unused_warning[neutralize_variable] now returns £120,000 instead of £555. It passes once the new override is removed.
  • Write inside a restoring reform: a reform that calls replace_variable(original) and then set_input(..., [20_000]) hits a holder that is still neutralised. Core rejects the write and keeps £10,000.

The underlying bug is in policyengine-core, where a Holder caches a Variable that runtime reforms replace. A core fix is planned. The hub is running one more bounded pe-uk build (fix-2265-r4). If a correct pe-uk fix would need patching core internals, the build stops, marks the new regressions strict-xfail, and this PR waits for the core release.

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

UK hub (Claude session, for Max): waiting on policyengine-core. e4c54d7fa removes r3's _invalidate_all_caches override, so simulation.py and utils/subsample.py again match 9a91736f3. It also adds supplied/model × direct/subsampled regressions for writes inside a restoring reform.

A holder that tracks the system on every read and write would also reject the unchanged CGT test's £555 write made while the variable is neutralised. So a correct fix needs core's Holder to follow runtime variable replacement. Four cases are marked xfail(strict=True, reason="needs policyengine-core holder resync"): the two sampled writes and the two r2 restorations. Locally, 244 passed and 4 xfailed.

A core fix is planned. Once it ships, bump core, remove the four markers and their two allowlist entries, and re-review. This isn't Budget work, so it stays out of the merge queue until then.

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

Picking this up for the UK hub. Reply here if you're mid-change and I'll stop.


Generated by Claude Code

Conflicts:
- simulation.py: keep both imports (measurement_branches from main,
  subsample from this branch).
- utils/supplied_inputs.py: keep main's wording for what a rebuild
  forgets and this branch's paragraph on supplied_input_arrays.
- tests/test_input_records.py: keep main's rebuilt-population tests and
  this branch's restoring-reform test; join the imports and use main's
  `period` import.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LwQq6HhDW9nQVQhdZKvrUS
main records each dataset year's ISA income contract and the series its
savings income grows by on the simulation and the dataset. A rebuild of
the populations resets the simulation's copies, and the subsampled
dataset was built without them, so a legacy-contract dataset projected
savings interest past its last year differently after a subsample.
Restore both records after the rebuild, alongside the capital gains
origins, and carry each year's energy price basis, ISA income contract
and savings income series to the subsampled dataset.

Also reword the subsample module docstring for integer membership ids,
and rename the changelog fragment to 2254.fixed.md without list markers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LwQq6HhDW9nQVQhdZKvrUS
Removing the restore of the defaulted capital gains periods after a
subsample changed the reported origin of a neutralised input from
default to projected, and no test caught it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LwQq6HhDW9nQVQhdZKvrUS
@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

PR_TITLE.txt

Give UK simulations a working subsample

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

PR_BODY.md

Fixes #2254
Reproducer: `policyengine_uk/tests/test_subsample.py::test_subsample_runs_on_a_dataset_simulation`
axiom: n/a: simulation infrastructure, no policy rule change

**Status: waiting for policyengine-core holder resync.** Four supplied/model restoration regressions stay strict expected failures until core releases the fix. `main` now locks policyengine-core 3.34.1, whose `Holder` still keeps the `Variable` it was built with, so the four cases still fail at `44f6deb7` (see Tests).

## What was wrong

UK simulations inherited `Simulation.subsample()` from policyengine-core, and it 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 (then 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. Those hold whatever a dataset or `set_input` stores there, not necessarily the populations' memberships. On `main` at 2.125.2 the UK model stored 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, the 535,080 distinct household identifiers came to only 76,551 distinct float32 values, and the 618,980 benefit unit identifiers to 544,894. #2273 has since moved both to 32-bit integers. The UK method reads memberships from the populations either way.
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 a caller supplied it as an input; it shares out bus fare spending and does not weight a sample.

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`. Separately, core's holders keep the variable they were built with after a runtime reform replaces its definition. A UK workaround that rebound holders when core invalidated caches made core reject an existing neutralised CGT input replacement and came too late for writes inside a restoring reform, so this branch does not carry it. The four supplied/model restoration regressions stay strict expected failures with reason `needs policyengine-core holder resync`, with every assertion kept.

## 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 holds household weights for has none, so the method samples by the first weights it holds.
- Takes households, benefit units and people from the populations (`ids` and `members_entity_id`), which hold them exactly, 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 it does not keep a value the model calculated later and cached over the input. It drops calculated values, which the model recalculates on the sample, and starts the tracer again.
- Replaces each supplied weight so that its total over its own entity stays the same in every period. `WEIGHT_VARIABLES` lists the weights (`household_weight`, `benunit_weight`, `person_weight`, `original_weight`, `raw_person_weight`), and a test fails if a new variable ending in `weight` appears in neither list.
- 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 method makes the copy as `get_branch` makes a branch, so the copy shares the old baseline's arrays as read-only views instead of copying them, and 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. It leaves the original object as it was (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. Each year keeps the dataset's energy price basis, ISA income contract and savings income series.
- Keeps what the simulation records beyond its inputs, which a rebuild of the populations forgets:
  - values the model stored as its own with `set_model_value` (the Universal Credit health element reform stores them for every simulation, for 2026 to 2032), which it stores back with `set_model_value`, so `supplied_input` still does not read them as supplied;
  - stored inputs and model values of neutralised variables, which it writes back to the holder directly because `set_input` rejects a neutralised variable;
  - where each capital gains input came from (supplied, projected or default, by period), and the dataset's annotation of projected capital gains inputs, which it carries to the subsampled dataset;
  - the dataset's ISA income contract for each year and the series its savings income grows by, which set how the model represents and projects savings interest.
- 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 method leaves the global random state alone. A seed therefore does not pick the same households as it would in core.
- With no seed, the method uses a fixed seed.
- The method refuses, with a `ValueError`, 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.
- The method replaces the baseline object rather than changing it in place (see above).

## Invariants

`policyengine_uk/tests/test_subsample.py` states these, and tests them 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**: the sample keeps between 1 and `n` households, each with a positive weight in the sampling period.
2. **Weight totals**: the total household weight of every year stays the same, to float32 precision, and so does the total of every other supplied weight over its own entity.
3. **Quantized weights**: a household's weight equals its number of draws times the year's total over the draws of households with a positive weight that year: in the sampling year, a whole multiple of total / `n`, with the multiples summing to `n`. The same draws set every year's weights. With `quantize_weights=False`, it stays proportional to its old weight times a whole number of draws. Either way, a zero weight stays zero in every year, and the method rescales each supplied weight to its own entity's total (so a supplied person weight need not equal the new weight of the person's household).
4. **Links**: a kept household or benefit unit has exactly the people it had, a kept person keeps their household and benefit unit, and the sample holds no other record. 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 stays the same 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 becomes a copy holding the same households and weights, with its own policy. The baseline held before stays as it was, 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 all come from 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 comes within 3% of its old weight.
10. **Trace and tiers**: the tracer starts again, and the sample keeps an input stored on disk 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).
11. **Model values and dataset records**: a value the model stored as its own stays the model's, a supplied value stays supplied, and the capital gains input origins reported by `capital_gains_projection_metadata` stay the same, including the default origin of a neutralised input. A simulation built from the subsampled dataset reports the same origins as one built from the full dataset. The ISA income contract of every year, the savings income series and each year's energy price basis stay the same on the simulation, its baseline copy and the dataset, and savings interest past the dataset's last year matches the full simulation for the same people.

Intended exceptions, each pinned by a test:

- In a year other than the sampling one, a household's mean new weight equals its share of the sampling year's total times that year's total. That matches its old weight only if its share stays the same in both years.
- `quantize_weights=False` keeps the total exactly, which makes the new weights a ratio estimate. I measured the bias at `ab843c80`: about 6% at 40 draws and under 1% at 400, for a seven-household example. Core uses the same definition.
- 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. The method rescales a weight set with `set_input` on the simulation alone to its own total there, while the baseline and the dataset keep the dataset's weights.

## Merge with `main`

`29504123` merges `main` at `47a2817e` (2.175.2, policyengine-core 3.34.1). Three files conflicted:

- `policyengine_uk/simulation.py`: both sides added an import at the same place. The merge keeps `main`'s `measurement_branches` import and this branch's `subsample_simulation` import.
- `policyengine_uk/utils/supplied_inputs.py`: both sides edited the module docstring's last paragraph. The merge keeps `main`'s wording for what a rebuild forgets ("when a rebuild installs populations that do not hold it") and this branch's paragraph on `supplied_input_arrays`.
- `policyengine_uk/tests/test_input_records.py`: both sides added imports and tests. The merge keeps `main`'s ten rebuilt-population tests and this branch's restoring-reform test, joins the imports, and renames the test's local `period` to `input_period` so that it can use `main`'s `period` import.

`main` does not do anything this branch does, so the merge drops nothing. `main` added two dataset records that the rebuild in `subsample` resets: each year's ISA income contract and the savings income series. Without them, a subsample of a dataset under the legacy ISA contract projected savings interest in 2035 from the additive default (−110.29 instead of −100.00 for one person in the test). `44f6deb7` restores both on the simulation, and carries them and the energy price basis to each year of the subsampled dataset.

## Tests

- **Merge with `main` at `47a2817e`:** at `44f6deb7`, with the locked environment (policyengine-core 3.34.1) and `pytest -n 4`, 25 files gave **874 passed, 4 strict expected failures, 0 failed, 0 skipped** in 43 minutes: `test_subsample.py` (57 passed / 2 xfailed), `test_input_records.py` (13 passed / 2 xfailed), every other file that calls `subsample` or the supplied-input helpers, the CGT projection and response-default files, the four UC health element files, dataset sources, the DataFrame round trip, input DataFrame export, reform-input lifecycle, reform representations, the ISA income contract and the energy price basis files. The four expected failures are the restoration cases above. `policyengine_uk/tests/code_health` gave 3,163 passed at `44f6deb7`. At `e804e8bc` (which only adds a test), the two new tests pass. At `e804e8bc`, `ruff format --check .` (2,229 files), `ruff check .`, `ruff check --select F811,F821 policyengine_uk`, `git diff --check` against `main`, `.github/check_changelog_fragments.py` and `.github/check_fix_pr.py` on this title and body pass.
- Revert checks at `44f6deb7`, each in a scratch copy, running `test_how_the_dataset_says_its_amounts_grow_is_kept`: without the restore of the ISA contracts and savings series, it fails on the contracts; without the dataset metadata, it fails on the subsampled dataset's ISA income contract; the unchanged copy passes.
- Mutation checks at `44f6deb7` of the post-merge code, each in a scratch copy, running `test_a_model_value_stays_the_models` and `test_where_capital_gains_inputs_came_from_is_kept`: dropping the capital gains origins, dropping the dataset's capital gains annotation and storing model values with `set_input` each fail one of them. Dropping the defaulted capital gains periods passed both, so `e804e8bc` adds `test_a_neutralised_capital_gains_input_still_reads_as_its_default`, which fails with that mutation and passes without it.
- **Earlier, before this merge (2026-10-23):** the restoration regressions check writes inside restoration both directly and after subsampling; direct cases and the unchanged CGT test pass, while the r2 restoration cases and sampled write cases need core holder resync. At `e4c54d7fa`, the locked `pytest -n 2` run passed **244 tests with 4 strict expected failures, 0 unexpected failures and 0 skips**: subsample (56 passed / 2 xfailed), input records (3 passed / 2 xfailed), the whole CGT projection file (115 passed), input DataFrame export (25 passed), reform-input lifecycle (17 passed), and reform representations (28 passed).
- Mutation check at `01f5d051`: 22 single changes to `utils/subsample.py` and `utils/supplied_inputs.py`, each run against the test file. The tests caught 21. The survivor removes the reset of the input record, which `_invalidate_all_caches` repeats later in the same function, so it changes nothing observable.
- At `7b4c3c28` (`main` at `7918fc2a` merged in), with 4 workers, the subsample, input records, dataset sources, DataFrame round trip, four UC health element, CGT projection and CGT response default files gave 377 passed on core 3.33.2 and on core 3.33.12.
- Full-size check under core 3.33.2 at `5dd27677`, before any merge with `main` (not rerun since), 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 comes to 8,634. Memberships matched the dataset's 64-bit identifiers exactly, the 1,260 supplied variable-periods matched those before, the baseline became a new copy with the same records while the old baseline still held every household, every year of `dataset` held the same records, and each year's household weight total stayed the same to 3 parts in 100 million. The same held for `quantize_weights=False` (9,864 households, 4.0 seconds).
- Memory at `5dd27677`: 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, which CI runs, and tests that need survey data.

## Aggregate impact

Pending at `e804e8bc`: CI's aggregate impact check runs again when the branch moves. The last report, [on `7b4c3c28`](https://github.com/PolicyEngine/policyengine-uk/pull/2265#issuecomment-6331578994) merged into `main` (`b4495e5cb1` against `a8b2848c13`, `enhanced_frs_2024_25.h5`, policyengine-uk-data 1.58.1, fiscal years 2025-26 to 2031-32), found:

> **Checks:** 0 broken by this change, 1 failing on both sides, 0 fixed. **Moves flagged for review:** 0

> No household's net income changes in any year.

The one check failing on both sides, `sign:vat` (VAT never goes negative), concerns code this PR does not touch. A full microsimulation does not call `subsample`, and the changes to `supplied_inputs.py` only add helpers.

## Found on the way

- #2266: `Simulation(dataset=<DataFrame>)` failed for households of more than one person. #2270 fixed it.
- #2267: `person_household_id` and `person_benunit_id` used 32-bit floats. #2273 fixed it. This PR reads memberships from the populations and does not depend on it.
- #2268: the elasticity-based labour supply responses compared a reform with a copy of itself. #2277 fixed it; this PR does not touch that code.
- #2307: `supplied_input` and `supplied_input_periods` report a value the model calculated and cached over an input as supplied. #2308, still open, builds on `stored_supplied_value` from this PR.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Collaborator Author

Ready for hub review

  • Head: e804e8bc37110ba69d37a61abc42f5962752c2c6 on claude/merge-main-pr-2265-afvyrr. The harness only lets me push to this claude/... branch. It holds e4c54d7f (this PR's head) plus three commits. I did not push to fix-subsample-2254, and I did not open a PR.
    • 29504123: merges main.
    • 44f6deb7: the fix the merge needed.
    • e804e8bc: a test.
  • main moved. The brief named 95dfc62f. By the time I merged, origin/main had reached 47a2817e, which adds Count Maternity Allowance in the benefit cap #2811 and a version bump. I merged 47a2817e, which conflicts in the same three files as 95dfc62f.
  • No one mid-change: the PR's last commit dates from 2026-10-23. I found no comment saying someone was mid-change. I posted the pickup notice before pushing.

Merge resolutions (29504123)

  • policyengine_uk/simulation.py:87-92: both sides added an import at the same place. I kept both: main's measurement_branches import and this branch's subsample_simulation.
  • policyengine_uk/utils/supplied_inputs.py:28-39 (module docstring): I kept main's wording for what a rebuild forgets and this branch's paragraph on supplied_input_arrays. The code merged without conflict.
  • policyengine_uk/tests/test_input_records.py:
    • kept main's ten rebuilt-population tests and this branch's test_restoring_reform_can_replace_a_neutralized_input;
    • joined the imports;
    • in that test, renamed the local period to input_period so it can use main's from policyengine_core.periods import period. Its assertions are unchanged.
  • Nothing dropped: main has no subsample, and none of its code does what this branch does.

Findings

  1. Review 5514224679 (COMMENTED at 82e0c499), four points.
    • Point 1 (merge main): done again in 29504123.
    • Point 2 (model values): stays fixed. utils/subsample.py:414-417 stores model values back with set_model_value. With that line changed to set_input, test_a_model_value_stays_the_models fails.
    • Point 3 (zero weights): rescale_weights at utils/subsample.py:166-168 is unchanged and test_a_zero_weight_stays_zero_in_every_year passes.
    • Point 4 (separate rescaling): the comment on WEIGHT_VARIABLES (utils/subsample.py:61-75) still documents it.
  2. Merge-order note 6102408989 (Store person_household_id and person_benunit_id as integers #2273, integer ids). Store person_household_id and person_benunit_id as integers #2273 has merged. The test precondition the note named was already gone at e4c54d7f. I reworded the module docstring's float32 sentence (utils/subsample.py:12-16) and the PR body's item 2. The ids in the test are still ones float32 cannot tell apart, so I left the test name and invariant 4 alone.
  3. Note 6104724035 (Load a DataFrame simulation input as the dataset it describes #2270): already applied before this run. test_a_simulation_built_from_a_dataframe_subsamples_its_dataset passes.
  4. Note 6177826914 (Make a retention-rate elasticity of 1 the default CGT realisation response #2419, the baseline must hold the same people): stays resolved. Make a retention-rate elasticity of 1 the default CGT realisation response #2419 is on main, and test_a_reform_and_its_baseline_keep_the_same_households passes.
  5. r1 P2 in 6335788910 (neutralised stored inputs dropped): stays fixed at utils/subsample.py:403-412. The test of record is test_neutralized_inputs_and_model_values_survive_subsampling, which is strict xfail because of finding 6.
  6. r2 P2 in 6357819601 and r3 in 6359206473 (restoration and writes after a subsample need core holder resync): still open. main now locks policyengine-core 3.34.1. Its Holder still stores self.variable once in __init__ (holders/holder.py:31) and nothing rebinds it. All four strict xfails still fail at 44f6deb7, so their markers and allowlist entries stay.
  7. New, found during the merge: main added two simulation records that _start_new_population resets: _isa_income_contracts and _savings_income_projection. The subsampled dataset was also built without each year's isa_income_contract, savings_income_projection and energy_price_basis.
    • Effect: on a legacy-contract dataset, savings_interest_income in 2035 came out at −110.29 after a subsample against −100.00 in the full simulation, for the same person.
    • Fix (44f6deb7): DATASET_RECORDS (utils/subsample.py:92-97) now lists the two CGT records and the two new ones. _plan captures them and _apply restores them. _subsample_single_year_dataset (utils/subsample.py:230-241) passes the three dataset declarations through.
    • Test: test_how_the_dataset_says_its_amounts_grow_is_kept.

Tests (locked environment, Python 3.13, policyengine-core 3.34.1)

  • At 44f6deb7, pytest -n 4 on 25 files: 874 passed, 4 xfailed (strict, the four core-resync cases), 0 failed, 0 skipped, in 43 minutes. Per file:
    • test_subsample.py: 57 passed, 2 xfailed
    • test_input_records.py: 13 passed, 2 xfailed
    • test_capital_gains_projection.py: 115
    • test_capital_gains_response_default.py: 37
    • test_dataframe_dataset.py: 90
    • test_isa_income_contract.py: 53
    • test_energy_price_basis.py: 30
    • test_input_dataframe_export.py: 25
    • test_reform_inputs_survive_reset.py: 17
    • test_reform_representations.py: 28
    • test_simulation_dataset_sources.py: 11
    • the four test_uc_health_* files: 81
    • files that import the supplied-input helpers: aggregate-impact childcare provenance 42, aggregate-impact guard 63, benefit cap properties 46, childcare birthday terms 53, current education order 11, gift aid 12, household subset 9, HB accommodation dates 38, NI class 4 properties 25, trading allowance data year 18
  • policyengine_uk/tests/code_health at 44f6deb7: 3,163 passed.
  • At e804e8bc (adds one test): the two new tests pass. These also pass:
    • uvx ruff@0.17.2 format --check . (2,229 files), ruff check ., ruff check --select F811,F821 policyengine_uk
    • git diff --check origin/main HEAD
    • .github/check_changelog_fragments.py --base origin/main --head HEAD
    • .github/check_fix_pr.py on the posted title and body: PASS, 2 changed executable tests
  • YAML tests: the PR changes none and changes no variable, so there were none to run.

Revert and mutation checks (scratch copies, never committed)

  • Finding 7, at 44f6deb7:
    • Without the ISA and savings records in DATASET_RECORDS, the new test fails on the contracts.
    • Without the three dataset declarations, it fails on the subsampled dataset's isa_income_contract.
    • The unchanged copy passes.
  • The PR's three post-merge mutations, repeated because 44f6deb7 rewrote the code they cover. These three each fail test_where_capital_gains_inputs_came_from_is_kept or test_a_model_value_stays_the_models:
    • drop the CGT origins;
    • drop the dataset's _cgt_extension_provenance;
    • store model values with set_input.
  • A fourth mutation survived: dropping the restore of _cgt_defaulted_periods passed both of those tests. A probe showed a neutralised capital_gains_before_response's 2026 origin flipping from default to projected. e804e8bc adds test_a_neutralised_capital_gains_input_still_reads_as_its_default, which fails under that mutation and passes without it.

Overlap with #2308 (head 5b4dc214, base fix-subsample-2254)

Skipped

  • Survey-data tests and the full-size run: no access to private data. The full-size and memory figures in the body come from 5dd27677 and are labelled with it.
  • The full suite: CI runs it.

Judgment calls

Sentences that carry a CI figure (rerun when the branch moves)

  • In Aggregate impact:
    • "Pending at e804e8bc: CI's aggregate impact check runs again when the branch moves. The last report, on 7b4c3c28 …, found:"
    • the two quoted lines ("Checks: 0 broken …" and "No household's net income changes in any year.")
    • "The one check failing on both sides, sign:vat …"
  • In Tests: "Not run locally: the full suite, which CI runs, and tests that need survey data."

Hub documents: PR_TITLE.txt (comment 6652481878) and PR_BODY.md (comment 6652482158), above.


Generated by Claude Code

MaxGhenis added a commit that referenced this pull request Nov 5, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis
MaxGhenis marked this pull request as ready for review November 5, 2026 13:10
@MaxGhenis
MaxGhenis requested a review from juaristi22 November 5, 2026 13:10

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Simulation.subsample() raises TypeError: build_from_dataset() needs a dataset

4 participants