From bf75f4df64e91920edaee20eab87b1c2af85670c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 10 Oct 2026 15:36:24 -0400 Subject: [PATCH 1/7] Give UK simulations a working subsample 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 --- changelog.d/fix-subsample-2254.fixed.md | 1 + policyengine_uk/simulation.py | 51 ++ policyengine_uk/tests/test_subsample.py | 803 +++++++++++++++++++++++ policyengine_uk/utils/subsample.py | 404 ++++++++++++ policyengine_uk/utils/supplied_inputs.py | 29 +- 5 files changed, 1287 insertions(+), 1 deletion(-) create mode 100644 changelog.d/fix-subsample-2254.fixed.md create mode 100644 policyengine_uk/tests/test_subsample.py create mode 100644 policyengine_uk/utils/subsample.py diff --git a/changelog.d/fix-subsample-2254.fixed.md b/changelog.d/fix-subsample-2254.fixed.md new file mode 100644 index 0000000000..cd6e668f1d --- /dev/null +++ b/changelog.d/fix-subsample-2254.fixed.md @@ -0,0 +1 @@ +- Fixed `Simulation.subsample()` and `Microsimulation.subsample()`, which raised `TypeError: Simulation.build_from_dataset() missing 1 required positional argument: 'dataset'`. They now keep a weighted sample of households with every supplied weight total unchanged, and subsample the baseline and `dataset` to the same households. diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index bf21684c48..91e8f7e5dd 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -33,6 +33,7 @@ ) from policyengine_uk.data.dataset_sources import materialize_gcs_dataset_url from policyengine_uk.utils.dependencies import get_variable_dependencies +from policyengine_uk.utils.subsample import subsample_simulation from policyengine_uk.utils.supplied_inputs import ( SUPPLIED_INPUT_VARIABLES, drop_missing_supplied_inputs, @@ -580,6 +581,56 @@ def build_from_ids( self.build_from_populations(builder.populations) + def subsample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + seed=None, + time_period=None, + quantize_weights: bool = True, + ) -> "Simulation": + """Cut the simulation down to a weighted random sample of households. + + Draws ``n`` households with replacement, each with probability + proportional to its weight in ``time_period``, and keeps each + household drawn at least once, with its benefit units and people. So + the sample holds at most ``n`` households. + + Every value the simulation was given (from its dataset or through + ``set_input``) is kept for the sampled records, for every period. + Calculated values are dropped and recalculated on the sample. Each + supplied weight is replaced so that its total is unchanged in every + period: a household's new weight is proportional to the number of + times it was drawn. + + The baseline and ``dataset`` are subsampled to the same households, + and each simulation keeps its own policy. + + Args: + n: The number of households to draw. + frac: The number to draw as a share of all households, used when + ``n`` is not given. + seed: A number or text that fixes the draw. The same seed always + draws the same households; without one, a fixed seed is used. + time_period: Draw households by their weight in this period. + Defaults to the default calculation period. + quantize_weights: If false, draw households with a positive + weight with equal probability, and make a household's new + weight proportional to its old weight times the number of + times it was drawn. + + Returns: + This simulation, subsampled. + """ + return subsample_simulation( + self, + n=n, + frac=frac, + seed=seed, + time_period=time_period, + quantize_weights=quantize_weights, + ) + def move_values(self, variable_donor: str, variable_target: str) -> None: """Move values from one variable to another across all branches. diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py new file mode 100644 index 0000000000..9a50ac2bba --- /dev/null +++ b/policyengine_uk/tests/test_subsample.py @@ -0,0 +1,803 @@ +"""``Simulation.subsample`` keeps a weighted sample of a dataset's households. + +Before, the method was inherited from policyengine-core and raised +``TypeError``: it rebuilds through ``build_from_dataset()`` with no arguments, +and the UK builder needs the dataset. The UK method subsamples the +simulation's populations, its baseline and its dataset together. + +Invariants, for every dataset, sample size ``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: with ``quantize_weights`` (the default), a household's + weight is a whole multiple of total / ``n``, and the multiples sum to ``n``. + Without it, a household's weight 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 identifiers that 32-bit floats cannot tell apart. +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. Calculated values are not kept. +6. Determinism: the same seed keeps the same households with the same weights. +7. Baseline: the baseline holds the same households and weights, and each + simulation keeps its own policy. +8. Dataset (differential): ``dataset`` holds the same records and weights as + the populations, and a simulation built from it gives the same results as + the subsampled simulation. +9. Unbiasedness: over many seeds, a household's mean quantized weight is its + old weight. Intended exception: unquantized weights keep the total exactly + and are therefore a ratio estimate, whose mean is off by a few per cent in + a small sample and closes on the old weight as the sample grows. + +Intended exceptions: a simulation built from a situation, one given no +household weights, a branch, and a baseline holding other households raise +``ValueError`` and are left unchanged. +""" + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st +from policyengine_core.periods import period as period_ + +from policyengine_uk import CountryTaxBenefitSystem, Microsimulation, Simulation +from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset +from policyengine_uk.utils.subsample import ( + WEIGHT_VARIABLES, + draw_household_counts, + rescale_weights, + subsample_dataset, +) +from policyengine_uk.utils.supplied_inputs import supplied_input_arrays + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +YEARS = range(2025, 2031) +# Neighbouring identifiers of the size the UK datasets use. As 32-bit floats, +# which is how the model stores person_household_id and person_benunit_id, +# identifiers this large are exact only at multiples of 64. +FIRST_HOUSEHOLD_ID = 900_000_001 +FIRST_BENUNIT_ID = 906_000_001 +FIRST_PERSON_ID = 967_000_001 +PERSONAL_ALLOWANCE = "gov.hmrc.income_tax.allowances.personal_allowance.amount" +REGIONS = ["LONDON", "WALES", "SCOTLAND", "NORTH_EAST"] + +# Six households: (weight, people in each benefit unit). +HOUSEHOLDS = [ + (100.0, [1, 1]), + (200.0, [1]), + (300.0, [2, 2]), + (400.0, [1]), + (500.0, [1, 1]), + (600.0, [2]), +] + + +def dataset(households=HOUSEHOLDS, year=2025, weights=True) -> UKSingleYearDataset: + """A dataset of ``households``, each a weight and the number of people in + each of its benefit units.""" + people, benunits, rows = [], [], [] + for index, (weight, units) in enumerate(households): + household_id = FIRST_HOUSEHOLD_ID + index + rows.append( + { + "household_id": household_id, + "household_weight": weight, + "region": REGIONS[index % len(REGIONS)], + "rent": 0.0, + "tenure_type": "OWNED_OUTRIGHT", + "council_tax": 1_000.0 + index, + } + ) + for size in units: + benunit_id = FIRST_BENUNIT_ID + len(benunits) + benunits.append({"benunit_id": benunit_id}) + for _ in range(size): + people.append( + { + "person_id": FIRST_PERSON_ID + len(people), + "person_benunit_id": benunit_id, + "person_household_id": household_id, + "age": 25 + 3 * len(people), + "employment_income": 9_000.0 * (1 + len(people) % 7), + } + ) + household = pd.DataFrame(rows) + if not weights: + household = household.drop(columns="household_weight") + return UKSingleYearDataset( + person=pd.DataFrame(people), + benunit=pd.DataFrame(benunits), + household=household, + fiscal_year=year, + ) + + +def values(simulation, variable, period) -> np.ndarray: + return np.asarray(simulation.calculate(variable, period)) + + +def ids(simulation, entity) -> list: + return np.asarray(simulation.populations[entity].ids).tolist() + + +def members(simulation, entity) -> dict: + """The people in each ``entity``, by its identifier, from the populations.""" + population = simulation.populations[entity] + groups = np.asarray(population.ids)[population.members_entity_id] + members_by_group = {group: set() for group in ids(simulation, entity)} + for person, group in zip(ids(simulation, "person"), groups.tolist()): + members_by_group[group].add(person) + return members_by_group + + +def by_id(simulation, variable, period) -> dict: + entity = simulation.tax_benefit_system.get_variable(variable).entity.key + return dict(zip(ids(simulation, entity), values(simulation, variable, period))) + + +def weights_by_id(simulation, period) -> dict: + return by_id(simulation, "household_weight", period) + + +def supplied(simulation) -> set: + """The variables and periods the simulation holds supplied values for.""" + return { + (name, str(period)) for name, period, _ in supplied_input_arrays(simulation) + } + + +# The reported failure. + + +def test_subsample_runs_on_a_dataset_simulation(): + simulation = Microsimulation(dataset=dataset()) + total = values(simulation, "household_weight", 2026).sum() + + result = simulation.subsample(n=4, seed=0, time_period=2026) + + assert result is simulation + assert 1 <= len(ids(simulation, "household")) <= 4 + assert values(simulation, "household_weight", 2026).sum() == pytest.approx( + total, rel=1e-5 + ) + assert np.isfinite(values(simulation, "household_net_income", 2026)).all() + + +def test_subsample_is_the_uk_method_not_policyengine_cores(): + """policyengine-core's method ends in ``build_from_dataset()``, which the + UK builder cannot answer, so the UK classes must not inherit it.""" + assert "subsample" in vars(Simulation) + assert Microsimulation.subsample is Simulation.subsample + + +# Links. + + +def test_identifiers_beyond_float32_precision_keep_their_households(): + simulation = Microsimulation(dataset=dataset()) + households = members(simulation, "household") + benunits = members(simulation, "benunit") + # The stored membership inputs cannot tell these households apart. + stored = values(simulation, "person_household_id", 2025) + assert len(set(stored.tolist())) < len(households) + + simulation.subsample(n=5, seed=3, time_period=2025) + + kept = ids(simulation, "household") + assert 1 < len(kept) <= 5 + assert members(simulation, "household") == { + household: households[household] for household in kept + } + kept_people = set(ids(simulation, "person")) + assert members(simulation, "benunit") == { + benunit: people for benunit, people in benunits.items() if people <= kept_people + } + assert kept_people == set().union(*(households[household] for household in kept)) + assert values(simulation, "household_id", 2025).tolist() == kept + + +def test_zero_weight_households_are_never_kept(): + households = [(0.0, [1]), (50.0, [2]), (0.0, [1, 1]), (70.0, [1])] + for seed in range(5): + simulation = Microsimulation(dataset=dataset(households)) + simulation.subsample(n=6, seed=seed) + assert set(ids(simulation, "household")) <= { + FIRST_HOUSEHOLD_ID + 1, + FIRST_HOUSEHOLD_ID + 3, + } + + +# Weights. + + +def test_quantized_weights_count_draws_and_keep_every_years_total(): + simulation = Microsimulation(dataset=dataset()) + totals = { + year: values(simulation, "household_weight", year).sum() for year in YEARS + } + + simulation.subsample(n=7, seed=1, time_period=2027) + + draws = None + for year in YEARS: + weights = values(simulation, "household_weight", year) + assert weights.sum() == pytest.approx(totals[year], rel=1e-5) + multiples = weights * 7 / totals[year] + assert np.abs(multiples - multiples.round()).max() < 1e-3 + assert multiples.round().sum() == 7 + # The same draws set the weights of every year. + assert draws is None or multiples.round().tolist() == draws + draws = multiples.round().tolist() + + +def test_unquantized_weights_scale_the_old_weights_by_draws(): + simulation = Microsimulation(dataset=dataset()) + before = weights_by_id(simulation, 2025) + total = sum(before.values()) + + simulation.subsample(n=5, seed=2, quantize_weights=False) + + after = weights_by_id(simulation, 2025) + assert sum(after.values()) == pytest.approx(total, rel=1e-5) + ratios = np.array([after[household] / before[household] for household in after]) + draws = ratios * 5 / ratios.sum() + assert np.abs(draws - draws.round()).max() < 1e-3 + assert draws.round().sum() == 5 + + +def test_other_supplied_weights_keep_their_totals_over_their_own_entity(): + data = dataset() + data.household["original_weight"] = [7.0, 5.0, 3.0, 2.0, 1.0, 9.0] + data.person["raw_person_weight"] = np.arange(1.0, len(data.person) + 1) + simulation = Microsimulation(dataset=data) + benunit_weight = values(simulation, "benunit_weight", 2025) + simulation.set_input("benunit_weight", 2025, benunit_weight * 2) + totals = { + (name, year): values(simulation, name, year).sum() + for name in ("original_weight", "raw_person_weight") + for year in YEARS + } + totals["benunit_weight", 2025] = benunit_weight.sum() * 2 + + simulation.subsample(n=4, seed=5) + + for (name, year), total in totals.items(): + assert values(simulation, name, year).sum() == pytest.approx(total, rel=1e-5) + table = getattr( + simulation.dataset[year], + simulation.tax_benefit_system.get_variable(name).entity.key, + ) + if name in table: + assert table[name].sum() == pytest.approx(total, rel=1e-5) + + +def test_every_weight_variable_is_listed(): + """A new sampling weight must be rescaled by a subsample, so it has to be + listed. A variable that only ends in "weight" must be left out on purpose: + household_bus_fare_age_weight shares out bus fare spending by age.""" + named = { + name for name in CountryTaxBenefitSystem().variables if name.endswith("weight") + } + assert named == WEIGHT_VARIABLES | {"household_bus_fare_age_weight"} + + +# Inputs. + + +def test_supplied_values_are_kept_and_calculated_values_are_not(): + simulation = Microsimulation(dataset=dataset()) + hours = np.arange(1.0, len(ids(simulation, "person")) + 1) + simulation.set_input("hours_worked", 2027, hours) + energy = np.arange(10.0, 10.0 + len(ids(simulation, "household"))) + simulation.set_input("monthly_domestic_energy_consumption", "2026-03", energy) + hours_by_person = dict(zip(ids(simulation, "person"), hours)) + energy_by_household = dict(zip(ids(simulation, "household"), energy)) + ages = {year: by_id(simulation, "age", year) for year in YEARS} + regions = by_id(simulation, "region", 2025) + inputs = supplied(simulation) + # Calculate results before subsampling; they must not become inputs. + for variable in ("income_tax", "household_net_income", "universal_credit"): + simulation.calculate(variable, 2026) + simulation.calculate("age", 2035) + + simulation.subsample(n=4, seed=4) + + assert supplied(simulation) == inputs + for variable in ("income_tax", "household_net_income", "universal_credit"): + assert simulation.get_holder(variable).get_known_periods() == [] + assert period_(2035) not in simulation.get_holder("age").get_known_periods() + people = ids(simulation, "person") + households = ids(simulation, "household") + assert by_id(simulation, "hours_worked", 2027) == { + person: hours_by_person[person] for person in people + } + assert by_id(simulation, "monthly_domestic_energy_consumption", "2026-03") == { + household: energy_by_household[household] for household in households + } + for year in YEARS: + assert by_id(simulation, "age", year) == { + person: ages[year][person] for person in people + } + assert by_id(simulation, "region", 2025) == { + household: regions[household] for household in households + } + + +def test_results_match_the_same_households_in_the_full_simulation(): + """Subsampling changes which households are held and their weights, + not what the model calculates for a household.""" + full = Microsimulation(dataset=dataset()) + sample = Microsimulation(dataset=dataset()) + sample.subsample(n=4, seed=6, time_period=2026) + + for variable in ("income_tax", "household_net_income", "council_tax"): + expected = by_id(full, variable, 2026) + for record, value in by_id(sample, variable, 2026).items(): + assert value == pytest.approx(expected[record]) + + +# Determinism. + + +def test_the_same_seed_keeps_the_same_households(): + samples = [] + for seed in (0, 0, "text seed", "text seed", None, None): + simulation = Microsimulation(dataset=dataset()) + simulation.subsample(n=3, seed=seed, time_period=2026) + samples.append(weights_by_id(simulation, 2026)) + + assert samples[0] == samples[1] + assert samples[2] == samples[3] + assert samples[4] == samples[5] + + +def test_different_seeds_keep_different_households(): + weights = np.array([weight for weight, _ in HOUSEHOLDS]) + draws = {tuple(draw_household_counts(weights, 2, seed)) for seed in range(12)} + assert len(draws) > 1 + + +def test_frac_sets_the_number_of_draws(): + by_frac = Microsimulation(dataset=dataset()) + by_frac.subsample(frac=0.5, seed=8) + by_n = Microsimulation(dataset=dataset()) + by_n.subsample(n=3, seed=8) + + assert weights_by_id(by_frac, 2025) == weights_by_id(by_n, 2025) + + +def test_a_period_before_the_data_samples_by_the_first_weights(): + """The model has no weights before the data's first year: it would give + every household the default weight of one.""" + early = Microsimulation(dataset=dataset()) + assert set(values(early, "household_weight", 2020).tolist()) == {1.0} + early.subsample(n=3, seed=9, time_period=2020) + first = Microsimulation(dataset=dataset()) + first.subsample(n=3, seed=9, time_period=2025) + + assert weights_by_id(early, 2025) == weights_by_id(first, 2025) + + +# Baseline. + + +def test_a_reform_and_its_baseline_keep_the_same_households(): + reform = {PERSONAL_ALLOWANCE: 0} + simulation = Microsimulation(dataset=dataset(), reform=reform) + + simulation.subsample(n=4, seed=0, time_period=2026) + + baseline = simulation.baseline + assert 1 <= len(ids(simulation, "household")) <= 4 + for held in (baseline, baseline.baseline): + assert ids(held, "household") == ids(simulation, "household") + assert ids(held, "person") == ids(simulation, "person") + assert weights_by_id(held, 2026) == weights_by_id(simulation, 2026) + # Each simulation keeps its own policy, on the sample. + sample = simulation.dataset + reformed = Microsimulation(dataset=sample, reform=reform) + unreformed = Microsimulation(dataset=sample) + for held, expected in ((simulation, reformed), (baseline, unreformed)): + assert values(held, "income_tax", 2026).tolist() == pytest.approx( + values(expected, "income_tax", 2026).tolist() + ) + assert ( + values(simulation, "income_tax", 2026).sum() + > values(baseline, "income_tax", 2026).sum() + ) + + +def test_a_baseline_holding_other_households_is_refused_and_nothing_changes(): + simulation = Microsimulation(dataset=dataset()) + simulation.baseline.subsample(n=2, seed=1) + households = ids(simulation, "household") + data = simulation.dataset + + with pytest.raises(ValueError, match="same households"): + simulation.subsample(n=3, seed=0) + + assert ids(simulation, "household") == households + assert simulation.dataset is data + + +# Dataset. + + +def test_the_dataset_holds_the_sample(): + simulation = Microsimulation(dataset=dataset()) + full = simulation.dataset + full_households = len(full[2025].household) + + simulation.subsample(n=4, seed=0, time_period=2026) + + sample = simulation.dataset + assert isinstance(sample, UKMultiYearDataset) + assert sample is not full + assert sample.years == full.years + assert simulation.baseline.dataset is sample + # The dataset the simulation was built from is not changed. + assert len(full[2025].household) == full_households + for year in sample.years: + tables = sample[year] + assert tables.household.household_id.tolist() == ids(simulation, "household") + assert tables.benunit.benunit_id.tolist() == ids(simulation, "benunit") + assert tables.person.person_id.tolist() == ids(simulation, "person") + assert tables.household.household_weight.tolist() == pytest.approx( + values(simulation, "household_weight", year).tolist(), rel=1e-5 + ) + assert tables.household.region.tolist() == list( + values(simulation, "region", year) + ) + + +def test_subsampling_a_single_year_dataset_keeps_its_type(): + data = dataset() + counts = np.array([0, 2, 0, 1, 0, 0]) + household_ids = data.household.household_id.values + + sample = subsample_dataset(data, household_ids, counts) + + assert isinstance(sample, UKSingleYearDataset) + assert sample.time_period == data.time_period + assert sample.household.household_id.tolist() == household_ids[counts > 0].tolist() + assert sample.household.household_weight.tolist() == pytest.approx([1_400.0, 700.0]) + assert len(data.household) == 6 + + +# Refusals. + + +def test_subsample_needs_n_or_frac(): + simulation = Microsimulation(dataset=dataset()) + with pytest.raises(ValueError, match="n or frac"): + simulation.subsample(seed=0) + + +@pytest.mark.parametrize("n", [0, -3]) +def test_subsample_needs_at_least_one_draw(n): + simulation = Microsimulation(dataset=dataset()) + with pytest.raises(ValueError, match="at least one household"): + simulation.subsample(n=n) + assert len(ids(simulation, "household")) == len(HOUSEHOLDS) + + +def test_a_simulation_without_weights_is_refused(): + simulation = Microsimulation(dataset=dataset(weights=False)) + with pytest.raises(ValueError, match="household_weight"): + simulation.subsample(n=2) + + +def test_a_simulation_with_no_positive_weight_is_refused(): + households = [(0.0, [1]), (0.0, [2])] + simulation = Microsimulation(dataset=dataset(households)) + with pytest.raises(ValueError, match="positive weight"): + simulation.subsample(n=2) + + +def test_a_situation_is_refused(): + simulation = Simulation(situation={"age": 30}) + with pytest.raises(ValueError, match="built from a dataset"): + simulation.subsample(n=1) + + +def test_a_branch_is_refused(): + simulation = Microsimulation(dataset=dataset()) + branch = simulation.get_branch("other") + with pytest.raises(ValueError, match="branches"): + branch.subsample(n=2) + + +def test_branches_are_made_again_on_the_sample(): + simulation = Microsimulation(dataset=dataset()) + simulation.get_branch("other").calculate("income_tax", 2026) + + simulation.subsample(n=3, seed=0) + + assert simulation.branches == {} + branch = simulation.get_branch("other") + assert ids(branch, "household") == ids(simulation, "household") + assert len(branch.calculate("income_tax", 2026)) == len(ids(simulation, "person")) + + +def test_a_benefit_unit_across_households_is_refused(): + data = dataset() + # Move the second person of the third household's first benefit unit to + # the fourth household. + data.person.loc[3, "person_household_id"] = FIRST_HOUSEHOLD_ID + 3 + simulation = Microsimulation(dataset=data) + with pytest.raises(ValueError, match="more than one household"): + simulation.subsample(n=3) + + +def test_a_total_that_cannot_be_kept_names_the_weight(): + simulation = Microsimulation(dataset=dataset()) + simulation.set_input("original_weight", 2026, [np.inf, 1.0, 1.0, 1.0, 1.0, 1.0]) + households = ids(simulation, "household") + + with pytest.raises(ValueError, match="original_weight for 2026.*finite total"): + simulation.subsample(n=3) + + assert ids(simulation, "household") == households + + +def test_unquantized_weights_need_a_sampled_record_with_weight(): + """The third household holds all of the original weight and is never + drawn, so scaling the old weights cannot keep that total; counting draws + can.""" + data = dataset([(100.0, [1]), (200.0, [2]), (0.0, [1])]) + data.household["original_weight"] = [0.0, 0.0, 9.0] + simulation = Microsimulation(dataset=data) + + with pytest.raises(ValueError, match="original_weight for 2025.*no weight"): + simulation.subsample(n=2, quantize_weights=False) + + simulation.subsample(n=2) + assert values(simulation, "original_weight", 2025).sum() == pytest.approx(9.0) + + +# Properties of the draw and the weights. + +DRAW_SETTINGS = settings(max_examples=200, deadline=None, derandomize=True) +weight = st.one_of( + st.just(0.0), + st.floats(min_value=1e-3, max_value=1e6, allow_nan=False, allow_infinity=False), +) +weight_lists = st.lists(weight, min_size=1, max_size=40).filter( + lambda weights: max(weights) > 0 +) +seeds = st.one_of(st.integers(), st.text(max_size=8)) + + +@DRAW_SETTINGS +@given( + weights=weight_lists, + n=st.integers(1, 60), + seed=seeds, + quantize=st.booleans(), +) +def test_draws_sum_to_n_and_skip_weightless_households(weights, n, seed, quantize): + weights = np.array(weights) + + counts = draw_household_counts(weights, n, seed, quantize) + + assert counts.shape == weights.shape + assert counts.sum() == n + assert (counts >= 0).all() + assert (counts[weights <= 0] == 0).all() + assert 1 <= (counts > 0).sum() <= n + assert counts.tolist() == draw_household_counts(weights, n, seed, quantize).tolist() + + +@DRAW_SETTINGS +@given( + weights=weight_lists, + n=st.integers(1, 60), + seed=seeds, + quantize=st.booleans(), +) +def test_rescaled_weights_keep_the_total(weights, n, seed, quantize): + weights = np.array(weights) + counts = draw_household_counts(weights, n, seed, quantize) + + rescaled = rescale_weights(weights, counts, quantize) + + assert rescaled.sum() == pytest.approx(weights.sum(), rel=1e-9) + assert (rescaled[counts == 0] == 0).all() + assert (rescaled[counts > 0] > 0).all() + drawn = counts > 0 + shares = rescaled[drawn] / ( + counts[drawn] if quantize else weights[drawn] * counts[drawn] + ) + assert shares == pytest.approx(np.full(drawn.sum(), shares[0]), rel=1e-9) + + +@DRAW_SETTINGS +@given( + weights=weight_lists, + counts=st.data(), + quantize=st.booleans(), +) +def test_a_second_weight_keeps_its_own_total(weights, counts, quantize): + """A weight is rescaled by draws made from another weight, as when a + year's weights are rescaled by a sample drawn in another year.""" + weights = np.array(weights) + counts = np.array( + counts.draw( + st.lists(st.integers(0, 5), min_size=len(weights), max_size=len(weights)) + ) + ) + drawn_weight = (counts if quantize else weights * counts).sum() + if weights.sum() > 0 and drawn_weight == 0: + with pytest.raises(ValueError, match="no weight"): + rescale_weights(weights, counts, quantize) + return + + rescaled = rescale_weights(weights, counts, quantize) + + assert rescaled.sum() == pytest.approx(weights.sum(), rel=1e-9) + assert (rescaled[counts == 0] == 0).all() + + +def test_weights_without_a_total_are_refused_or_stay_zero(): + counts = np.array([1, 2]) + assert rescale_weights(np.zeros(2), counts).tolist() == [0.0, 0.0] + for weights in ([np.nan, 1.0], [np.inf, 1.0]): + with pytest.raises(ValueError, match="finite total"): + rescale_weights(np.array(weights), counts) + with pytest.raises(ValueError, match="finite"): + draw_household_counts(np.array([np.inf, 1.0]), 2, 0) + + +def mean_new_weights(weights, n, seeds, quantize): + return np.mean( + [ + rescale_weights( + weights, draw_household_counts(weights, n, seed, quantize), quantize + ) + for seed in range(seeds) + ], + axis=0, + ) + + +def test_a_households_mean_quantized_weight_is_its_old_weight(): + weights = np.array([1.0, 2.0, 3.0, 4.0, 0.0, 10.0, 30.0]) + + mean = mean_new_weights(weights, n=40, seeds=3_000, quantize=True) + + assert mean == pytest.approx(weights, rel=0.03) + + +def test_unquantized_weights_are_a_ratio_estimate(): + """Intended: keeping the total exactly makes the unquantized weights a + ratio estimate, so a household's mean new weight is off by a few per cent + in a small sample and closes on its old weight as the sample grows.""" + weights = np.array([1.0, 2.0, 3.0, 4.0, 0.0, 10.0, 30.0]) + positive = weights > 0 + + def error(n): + mean = mean_new_weights(weights, n=n, seeds=1_500, quantize=False) + return np.abs(mean[positive] / weights[positive] - 1).max() + + assert 0.04 < error(40) < 0.10 + assert error(400) < 0.02 + + +# Properties of a subsampled simulation. + +SIMULATION_SETTINGS = settings( + max_examples=10, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], +) +household_specs = st.lists( + st.tuples( + st.one_of(st.just(0.0), st.floats(min_value=0.5, max_value=5_000.0)), + st.lists(st.integers(1, 3), min_size=1, max_size=2), + ), + min_size=2, + max_size=7, +).filter(lambda households: any(weight > 0 for weight, _ in households)) + + +@SIMULATION_SETTINGS +@given( + households=household_specs, + n=st.integers(1, 10), + seed=st.integers(0, 10_000), + year=st.sampled_from(list(YEARS)), + quantize=st.booleans(), +) +def test_subsample_properties(households, n, seed, year, quantize): + def subsample(): + simulation = Microsimulation(dataset=dataset(households)) + return simulation.subsample( + n=n, seed=seed, time_period=year, quantize_weights=quantize + ) + + full = Microsimulation(dataset=dataset(households)) + sample = subsample() + kept = ids(sample, "household") + people = ids(sample, "person") + + # 1. Size. + assert 1 <= len(kept) <= n + full_weights = {period: weights_by_id(full, period) for period in YEARS} + assert all(full_weights[year][household] > 0 for household in kept) + + for period in YEARS: + weights = weights_by_id(sample, period) + total = sum(full_weights[period].values()) + # 2. Weight totals. + assert sum(weights.values()) == pytest.approx(total, rel=1e-5) + # 3. Quantized weights: whole numbers of draws, of the old weight + # when the weights are not quantized. + shares = np.array( + [ + weights[household] + / (1 if quantize else full_weights[period][household]) + for household in kept + ] + ) + draws = shares * n / shares.sum() + assert np.abs(draws - draws.round()).max() < 1e-2 + assert draws.round().sum() == n + + # 4. Links. + for entity in ("household", "benunit"): + full_members = members(full, entity) + sample_members = members(sample, entity) + assert all(sample_members.values()) + assert sample_members == { + group: full_members[group] for group in sample_members + } + assert set(people) == set().union(*members(sample, "household").values()) + assert set(people) == set().union(*members(sample, "benunit").values()) + assert len(set(people)) == len(people) + + # 5. Inputs. + assert supplied(sample) == supplied(full) + for period in YEARS: + for variable in ("age", "employment_income", "region", "council_tax"): + expected = by_id(full, variable, period) + assert by_id(sample, variable, period) == { + record: expected[record] + for record in ids( + sample, sample.tax_benefit_system.get_variable(variable).entity.key + ) + } + + # 6. Determinism. + again = subsample() + assert ids(again, "person") == people + for period in YEARS: + assert weights_by_id(again, period) == weights_by_id(sample, period) + + # 7. Baseline. + assert ids(sample.baseline, "household") == kept + assert ids(sample.baseline, "person") == people + assert weights_by_id(sample.baseline, year) == weights_by_id(sample, year) + + # 8. Dataset, and a simulation built from it. + rebuilt = Microsimulation(dataset=sample.dataset) + for entity in ("person", "benunit", "household"): + assert ids(rebuilt, entity) == ids(sample, entity) + assert getattr(sample.dataset[year], entity)[f"{entity}_id"].tolist() == ids( + sample, entity + ) + for period in YEARS: + assert list(weights_by_id(rebuilt, period).values()) == pytest.approx( + list(weights_by_id(sample, period).values()), rel=1e-5 + ) + for variable in ("household_net_income", "income_tax", "universal_credit"): + assert values(rebuilt, variable, year).tolist() == pytest.approx( + values(sample, variable, year).tolist() + ) diff --git a/policyengine_uk/utils/subsample.py b/policyengine_uk/utils/subsample.py new file mode 100644 index 0000000000..e636bc919c --- /dev/null +++ b/policyengine_uk/utils/subsample.py @@ -0,0 +1,404 @@ +"""Subsample a simulation built from a dataset. + +``Simulation.subsample`` keeps a weighted random sample of households and +gives the sample weights with the same totals as before. policyengine-core's +``subsample`` cannot do this for a UK simulation: + +- it rebuilds through ``Simulation.build_from_dataset()`` with no arguments, + and the UK builder takes the dataset to load, which it projects to later + years again; +- it reads each person's household and benefit unit back from the stored + ``person_household_id`` and ``person_benunit_id`` inputs. The UK model stores + those as 32-bit floats, which hold integers exactly only up to 2**24, so + larger identifiers would merge households and benefit units; +- it leaves ``Simulation.dataset`` as a policyengine-core ``Dataset``, while UK + formulas read ``dataset.years``; +- it rebuilds a ``baseline`` branch, while a UK simulation's baseline is a + simulation of its own. + +So the functions here take the households, benefit units and people from the +simulation's populations, whose identifiers and memberships are exact, and +subsample the baseline and the dataset with the same draw. +""" + +import hashlib +import operator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +from policyengine_core.periods import Period, period as period_ + +from policyengine_uk.data.dataset_schema import ( + UKMultiYearDataset, + UKSingleYearDataset, +) +from policyengine_uk.utils.supplied_inputs import supplied_input_arrays + +if TYPE_CHECKING: + from policyengine_uk.simulation import Simulation + +# Variables that count how many units of the population a record stands for. +# A subsample rescales each one that is supplied as an input, so that its +# total over its own entity is unchanged. +WEIGHT_VARIABLES = frozenset( + { + "household_weight", + "benunit_weight", + "person_weight", + "original_weight", + "raw_person_weight", + } +) + +# policyengine-core seeds a subsample with the dataset's name by default. UK +# datasets have no name, so every UK subsample without a seed uses this one. +DEFAULT_SEED = "policyengine-uk" + +UKDataset = Union[UKSingleYearDataset, UKMultiYearDataset] + + +def _random_state(seed) -> np.random.RandomState: + """A generator for ``seed``, which may be a number or text. + + ``RandomState`` keeps the same stream in every NumPy version, so a seed + chooses the same households wherever it runs. The digest gives the same + number in every process, which the built-in ``hash`` does not. + """ + digest = hashlib.sha256(str(seed).encode("utf-8")).digest() + return np.random.RandomState(int.from_bytes(digest[:4], "big")) + + +def draw_household_counts( + weights: np.ndarray, n: int, seed, quantize_weights: bool = True +) -> np.ndarray: + """Draw ``n`` households with replacement and count each one's draws. + + Only households with a positive weight can be drawn. Each draw picks a + household with probability proportional to its weight, or with equal + probability when ``quantize_weights`` is false. + + Returns: + The number of times each household was drawn; the counts sum to ``n``. + """ + n = operator.index(n) + if n < 1: + raise ValueError(f"A subsample needs at least one household, not {n}.") + weights = np.asarray(weights, dtype=float) + eligible = np.flatnonzero(weights > 0) + if eligible.size == 0: + raise ValueError("No household has a positive weight to sample by.") + if not np.isfinite(weights[eligible]).all(): + raise ValueError("Household weights must be finite to sample by them.") + probabilities = ( + weights[eligible] / weights[eligible].sum() if quantize_weights else None + ) + draws = _random_state(seed).choice( + eligible.size, size=n, replace=True, p=probabilities + ) + counts = np.zeros(weights.size, dtype=np.int64) + counts[eligible] = np.bincount(draws, minlength=eligible.size) + return counts + + +def rescale_weights( + weights: np.ndarray, counts: np.ndarray, quantize_weights: bool = True +) -> np.ndarray: + """Weights for the drawn records that keep the total of ``weights``. + + A drawn record's new weight is proportional to the number of times its + household was drawn (``counts``), and also to its old weight when + ``quantize_weights`` is false. A record that was not drawn gets zero. + """ + weights = np.asarray(weights, dtype=float) + total = weights.sum() + if not np.isfinite(total): + raise ValueError("the weights do not have a finite total") + if total == 0: + return np.zeros_like(weights) + drawn = counts.astype(float) if quantize_weights else weights * counts + drawn_total = drawn.sum() + if drawn_total == 0: + raise ValueError( + "the sampled records have no weight. Use quantize_weights=True or " + "a different sample" + ) + return drawn * (total / drawn_total) + + +def _rescaled(name: str, period, weights, counts, quantize_weights) -> np.ndarray: + try: + return rescale_weights(weights, counts, quantize_weights) + except ValueError as error: + raise ValueError( + f"Cannot preserve the total of {name} for {period}: {error}." + ) from error + + +def _subsample_table( + table: pd.DataFrame, counts: np.ndarray, period, quantize_weights: bool +) -> pd.DataFrame: + """The rows of ``table`` drawn at least once, with their new weights.""" + kept = counts > 0 + subsample = table.loc[kept].reset_index(drop=True) + for column in table.columns: + if column in WEIGHT_VARIABLES: + weights = _rescaled( + column, period, table[column].values, counts, quantize_weights + ) + subsample[column] = weights[kept] + return subsample + + +def _subsample_single_year_dataset( + dataset: UKSingleYearDataset, + household_ids: np.ndarray, + counts: np.ndarray, + quantize_weights: bool, +) -> UKSingleYearDataset: + # Read the stored tables, in which enum columns stay encoded; the public + # properties decode every enum column of the whole table. + person, benunit, household = dataset.tables + draws = pd.Series(counts, index=household_ids) + + def draws_of(households) -> np.ndarray: + return draws.reindex(np.asarray(households)).fillna(0).values + + # A benefit unit belongs to the household of its members. + benunit_household = person.groupby("person_benunit_id")[ + "person_household_id" + ].first() + tables = [ + _subsample_table(table, table_counts, dataset.time_period, quantize_weights) + for table, table_counts in ( + (person, draws_of(person["person_household_id"])), + ( + benunit, + draws_of(benunit_household.reindex(benunit["benunit_id"].values)), + ), + (household, draws_of(household["household_id"])), + ) + ] + subsample = UKSingleYearDataset( + person=tables[0], + benunit=tables[1], + household=tables[2], + fiscal_year=dataset.time_period, + ) + subsample._enum_columns = dict(dataset._enum_columns) + return subsample + + +def subsample_dataset( + dataset: UKDataset, + household_ids: np.ndarray, + counts: np.ndarray, + quantize_weights: bool = True, +) -> UKDataset: + """A copy of ``dataset`` holding the drawn households, for every year. + + Args: + dataset: The dataset to subsample. It is not changed. + household_ids: The identifier of each household that could be drawn. + counts: The number of times each of those households was drawn. + quantize_weights: As for ``Simulation.subsample``. + """ + if isinstance(dataset, UKMultiYearDataset): + return UKMultiYearDataset( + datasets=[ + _subsample_single_year_dataset( + dataset[year], household_ids, counts, quantize_weights + ) + for year in dataset.years + ] + ) + return _subsample_single_year_dataset( + dataset, household_ids, counts, quantize_weights + ) + + +@dataclass +class _Subsample: + """What one simulation holds once it is subsampled.""" + + simulation: "Simulation" + # The arguments of Simulation.build_from_ids. + ids: Tuple[np.ndarray, ...] + # Each supplied input, cut down to the sample. + inputs: List[Tuple[str, Period, np.ndarray]] + dataset: Optional[UKDataset] + + +def _sorted_positions(kept: np.ndarray, ids: np.ndarray) -> np.ndarray: + """Positions of the kept members of a group population, ordered by + identifier, which is the order ``Simulation.build_from_ids`` gives them.""" + positions = np.flatnonzero(kept) + return positions[np.argsort(ids[positions], kind="stable")] + + +def _plan( + simulation: "Simulation", + household_ids: np.ndarray, + counts: np.ndarray, + quantize_weights: bool, + datasets: Dict[int, UKDataset], +) -> _Subsample: + """Work out the subsample of ``simulation`` without changing it.""" + person = simulation.populations["person"] + benunit = simulation.populations["benunit"] + household = simulation.populations["household"] + if not np.array_equal(np.asarray(household.ids), household_ids): + raise ValueError( + "A simulation and its baseline must hold the same households to " + "be subsampled together." + ) + person_household = np.asarray(household.members_entity_id) + person_benunit = np.asarray(benunit.members_entity_id) + + # A benefit unit is kept with its household, so it must sit within one. + benunit_household = np.full(benunit.count, -1) + benunit_household[person_benunit] = person_household + if not np.array_equal(benunit_household[person_benunit], person_household): + raise ValueError( + "Cannot subsample: a benefit unit has members in more than one household." + ) + members = np.bincount(person_household, minlength=household.count) + if (members[counts > 0] == 0).any(): + raise ValueError("Cannot subsample: a sampled household has no members.") + + entity_counts = { + "person": counts[person_household], + "benunit": np.where(benunit_household >= 0, counts[benunit_household], 0), + "household": counts, + } + positions = { + "person": np.flatnonzero(entity_counts["person"] > 0), + "benunit": _sorted_positions( + entity_counts["benunit"] > 0, np.asarray(benunit.ids) + ), + "household": _sorted_positions(counts > 0, household_ids), + } + + inputs = [] + for name, period, values in supplied_input_arrays(simulation): + entity_key = simulation.tax_benefit_system.get_variable(name).entity.key + if name in WEIGHT_VARIABLES: + values = _rescaled( + name, period, values, entity_counts[entity_key], quantize_weights + ) + inputs.append((name, period, values[positions[entity_key]])) + + # A simulation and its baseline usually share one dataset object, so + # subsample each dataset once and share the result in the same way. + dataset = simulation.dataset + if dataset is not None: + if id(dataset) not in datasets: + datasets[id(dataset)] = subsample_dataset( + dataset, household_ids, counts, quantize_weights + ) + dataset = datasets[id(dataset)] + + kept_people = positions["person"] + return _Subsample( + simulation=simulation, + ids=( + np.asarray(person.ids)[kept_people], + np.asarray(benunit.ids)[person_benunit[kept_people]], + household_ids[person_household[kept_people]], + np.asarray(benunit.ids)[positions["benunit"]], + household_ids[positions["household"]], + ), + inputs=inputs, + dataset=dataset, + ) + + +def _apply(subsample: _Subsample) -> None: + """Rebuild a simulation on its subsample.""" + simulation = subsample.simulation + start_instant = simulation.start_instant + # Branches hold the populations the subsample replaces; formulas make + # them again when they next ask for them. + simulation.branches = {} + # The rebuild replaces every stored value, so the record of inputs, and + # of what was calculated from them, starts again with the inputs it sets. + simulation._user_input_keys = set() + simulation._store_history = None + simulation.build_from_ids(*subsample.ids) + for name, period, values in subsample.inputs: + simulation.set_input(name, period, values) + simulation.start_instant = start_instant + # Drop results calculated on the populations the subsample replaced. + simulation._invalidate_all_caches() + simulation.dataset = subsample.dataset + + +def _sampling_weights(simulation: "Simulation", time_period) -> np.ndarray: + """Household weights for ``time_period``, to draw the sample by. + + A period before the first supplied weights has none: the model would + return the default weight of one for every household. Use the first + supplied weights for such a period. + """ + periods = sorted( + ( + period + for name, period, _ in supplied_input_arrays(simulation) + if name == "household_weight" + ), + key=lambda period: period.start, + ) + if not periods: + raise ValueError( + "Cannot subsample a simulation that was given no household_weight values." + ) + requested = period_(time_period) + if requested.start < periods[0].start: + requested = periods[0] + return np.asarray(simulation.calculate("household_weight", requested)) + + +def subsample_simulation( + simulation: "Simulation", + n: Optional[int] = None, + frac: Optional[float] = None, + seed=None, + time_period=None, + quantize_weights: bool = True, +) -> "Simulation": + """Subsample ``simulation`` in place; see ``Simulation.subsample``.""" + if not simulation.built_from_dataset: + raise ValueError("Only a simulation built from a dataset can be subsampled.") + if simulation.branch_name != "default": + raise ValueError("Subsample a simulation, not one of its branches.") + if n is None and frac is None: + raise ValueError("Either n or frac must be provided.") + if time_period is None: + time_period = simulation.default_calculation_period + + weights = _sampling_weights(simulation, time_period) + if n is None: + n = int(len(weights) * frac) + counts = draw_household_counts( + weights, n, DEFAULT_SEED if seed is None else seed, quantize_weights + ) + + # The baseline is a simulation of its own, with a baseline of its own + # when this one has a scenario. Give each of them the same households. + simulations = [] + while simulation is not None and simulation not in simulations: + simulations.append(simulation) + simulation = getattr(simulation, "baseline", None) + + # Work out every subsample before applying any, so that a simulation + # that cannot be subsampled leaves them all as they were. + household_ids = np.asarray(simulations[0].populations["household"].ids) + datasets: Dict[int, UKDataset] = {} + subsamples = [ + _plan(each, household_ids, counts, quantize_weights, datasets) + for each in simulations + ] + for subsample in subsamples: + _apply(subsample) + return simulations[0] diff --git a/policyengine_uk/utils/supplied_inputs.py b/policyengine_uk/utils/supplied_inputs.py index 39371fa0cc..9363e5be20 100644 --- a/policyengine_uk/utils/supplied_inputs.py +++ b/policyengine_uk/utils/supplied_inputs.py @@ -19,9 +19,13 @@ refill it, and the helpers only answer for those variables. ``Simulation.delete_arrays`` forgets such records too, for any variable it deletes. + +``supplied_input_arrays`` answers for every variable. It reads only a record +whose supplied value is still stored, which the holder tells apart from a +value the engine carried over or calculated into the same period. """ -from typing import List, Optional +from typing import List, Optional, Tuple import numpy as np from policyengine_core.periods import Period @@ -129,3 +133,26 @@ def supplied_input( return holder.default_array() return value return None + + +def supplied_input_arrays(simulation) -> List[Tuple[str, Period, np.ndarray]]: + """Every value ``simulation`` was given as an input and still stores, on + its own branch, with its variable and the period it is stored for. + + Ordered by variable and period, so that the same simulation always gives + the same list. + """ + input_keys = getattr(simulation, "_user_input_keys", None) or () + arrays = [] + for variable_name, branch_name, period in input_keys: + if branch_name != simulation.branch_name: + continue + holder = simulation.get_holder(variable_name) + # A record can outlive its stored value (see above), and a carried + # over or calculated value can then take its place. + if not holder._stores_user_input(period, branch_name): + continue + value = holder._get_array_from_storage(period, branch_name) + if value is not None: + arrays.append((variable_name, period, value)) + return sorted(arrays, key=lambda item: (item[0], item[1].start, str(item[1]))) From 155c090075fb85badd48380aa1df70fb9eca49e7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 10 Oct 2026 15:59:41 -0400 Subject: [PATCH 2/7] Test that a subsample recalculates results and drops replaced inputs 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 --- policyengine_uk/tests/test_subsample.py | 28 +++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py index 9a50ac2bba..84ecc3ea4d 100644 --- a/policyengine_uk/tests/test_subsample.py +++ b/policyengine_uk/tests/test_subsample.py @@ -306,6 +306,7 @@ def test_supplied_values_are_kept_and_calculated_values_are_not(): for variable in ("income_tax", "household_net_income", "universal_credit"): simulation.calculate(variable, 2026) simulation.calculate("age", 2035) + income_tax = by_id(simulation, "income_tax", 2026) simulation.subsample(n=4, seed=4) @@ -315,6 +316,10 @@ def test_supplied_values_are_kept_and_calculated_values_are_not(): assert period_(2035) not in simulation.get_holder("age").get_known_periods() people = ids(simulation, "person") households = ids(simulation, "household") + # A result calculated before is calculated again, on the sample. + assert by_id(simulation, "income_tax", 2026) == { + person: income_tax[person] for person in people + } assert by_id(simulation, "hours_worked", 2027) == { person: hours_by_person[person] for person in people } @@ -330,6 +335,29 @@ def test_supplied_values_are_kept_and_calculated_values_are_not(): } +def test_a_calculated_value_that_replaced_an_input_is_not_kept(): + """The record of an input outlives a calculated value stored over it. + The subsample must not keep the record: the next calculation would fill + the period again, and the record would pass the result off as supplied.""" + simulation = Microsimulation(dataset=dataset()) + people = len(ids(simulation, "person")) + simulation.set_input("hours_worked", 2027, np.full(people, 1_000.0)) + simulation.get_holder("hours_worked").put_in_cache( + np.full(people, 500.0), period_(2027) + ) + assert ("hours_worked", "2027") not in supplied(simulation) + + simulation.subsample(n=4, seed=4) + simulation.calculate("hours_worked", 2027) + + assert ("hours_worked", "2027") not in supplied(simulation) + assert "2027" not in simulation.to_input_dict().get("hours_worked", {}) + assert not any( + name == "hours_worked" and str(period) == "2027" + for name, _, period in simulation._user_input_keys + ) + + def test_results_match_the_same_households_in_the_full_simulation(): """Subsampling changes which households are held and their weights, not what the model calculates for a household.""" From 21a87b2f6c9fdb3e50daa3f6de58d876891ede43 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 10 Oct 2026 16:03:23 -0400 Subject: [PATCH 3/7] Document subsample and test a simulation that has no dataset 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 --- docs/book/usage/simulations.md | 13 +++++++++++ policyengine_uk/tests/test_subsample.py | 29 +++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index 2ff8dc3f9c..c9835a9ea1 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -437,6 +437,19 @@ weighted_total = (unweighted_values * modified_weights).sum() print(f"Total with 10% higher weights: £{weighted_total / 1e9:.1f}bn") ``` +### Subsampling a microsimulation + +`subsample` cuts a simulation built from a dataset down to a weighted random sample of its households: + +```python +sim = Microsimulation(dataset="hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5") +sim.subsample(n=10_000, seed=0, time_period=2026) +``` + +It draws `n` households with replacement, each with probability proportional to its weight in `time_period`, and keeps every household drawn at least once, with its benefit units and people. So the sample holds at most `n` households. A household's new weight is proportional to the number of times it was drawn, and the total household weight of every year is unchanged. The values the simulation was given are kept for the sampled records; results are calculated again on the sample, so population estimates then carry sampling error. + +The simulation's `baseline` and `dataset` are cut down to the same households, and the same `seed` always draws the same households. + ## Debugging simulations When simulations don't behave as expected: diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py index 84ecc3ea4d..f116507d54 100644 --- a/policyengine_uk/tests/test_subsample.py +++ b/policyengine_uk/tests/test_subsample.py @@ -499,6 +499,35 @@ 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(): + households = np.arange(4) + frame = pd.DataFrame( + { + "person_id__2025": FIRST_PERSON_ID + households, + "person_benunit_id__2025": FIRST_BENUNIT_ID + households, + "person_household_id__2025": FIRST_HOUSEHOLD_ID + households, + "benunit_id__2025": FIRST_BENUNIT_ID + households, + "household_id__2025": FIRST_HOUSEHOLD_ID + households, + "household_weight__2025": [10.0, 20.0, 30.0, 40.0], + "age__2025": [30, 40, 50, 60], + } + ) + simulation = Microsimulation(dataset=frame) + assert simulation.dataset is None + 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 + assert sum(weights_by_id(simulation, 2025).values()) == pytest.approx(100.0) + assert by_id(simulation, "age", 2025) == { + person: ages[person] for person in ids(simulation, "person") + } + + # Refusals. From fa7b0f0a42bf845a4fc89bbedd84749673bf56c6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 10 Oct 2026 16:09:54 -0400 Subject: [PATCH 4/7] Wrap the subsample example in the docs to the line length CI formats Python blocks in Markdown with the latest ruff, which the locked ruff does not check. Co-Authored-By: Claude Opus 5.5 --- docs/book/usage/simulations.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index c9835a9ea1..1630e102af 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -442,7 +442,9 @@ print(f"Total with 10% higher weights: £{weighted_total / 1e9:.1f}bn") `subsample` cuts a simulation built from a dataset down to a weighted random sample of its households: ```python -sim = Microsimulation(dataset="hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5") +sim = Microsimulation( + dataset="hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5" +) sim.subsample(n=10_000, seed=0, time_period=2026) ``` From db55173b66c70bf63798b4d77eceb5c7b866d476 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 09:28:09 -0400 Subject: [PATCH 5/7] Copy baselines, read inputs from their tier, and restart the trace 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 --- docs/book/usage/simulations.md | 2 +- policyengine_uk/simulation.py | 23 +- policyengine_uk/tests/test_subsample.py | 306 ++++++++++++++++++++--- policyengine_uk/utils/subsample.py | 66 +++-- policyengine_uk/utils/supplied_inputs.py | 43 +++- 5 files changed, 381 insertions(+), 59 deletions(-) diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index 98a9091743..a6cf807d71 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -476,7 +476,7 @@ sim.subsample(n=10_000, seed=0, time_period=2026) It draws `n` households with replacement, each with probability proportional to its weight in `time_period`, and keeps every household drawn at least once, with its benefit units and people. So the sample holds at most `n` households. A household's new weight is proportional to the number of times it was drawn, and the total household weight of every year is unchanged. The values the simulation was given are kept for the sampled records; results are calculated again on the sample, so population estimates then carry sampling error. -The simulation's `baseline` and `dataset` are cut down to the same households, and the same `seed` always draws the same households. +The simulation's `baseline` and `dataset` are replaced with copies cut down to the same households; the objects they held before are left as they were. The `dataset` holds the rows of the dataset the simulation was built from, so it does not include values set with `set_input` after the simulation was built. The same `seed` always draws the same households. ## Debugging simulations diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index c5f5779977..dc8dbfdbdd 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -585,13 +585,22 @@ def subsample( Every value the simulation was given (from its dataset or through ``set_input``) is kept for the sampled records, for every period. - Calculated values are dropped and recalculated on the sample. Each - supplied weight is replaced so that its total is unchanged in every - period: a household's new weight is proportional to the number of - times it was drawn. - - The baseline and ``dataset`` are subsampled to the same households, - and each simulation keeps its own policy. + Calculated values are dropped and recalculated on the sample, and the + tracer starts again. Each supplied weight is replaced so that its + total is unchanged in every period: a household's new weight is + proportional to the number of times it was drawn. Averaged over + draws, a household's new weight in ``time_period`` is its old one; in + another period that holds only if the household's share of the total + weight is the same there. + + ``baseline`` is replaced with a copy subsampled to the same + households, which keeps its own policy and inputs. The baseline the + simulation held before is left as it was, since another simulation + (a clone of this one) can hold it too. ``dataset`` is replaced with + the rows of the dataset the simulation was built from for the + sampled records, with its weights rescaled in the same way. Like the + dataset before, it does not hold values set with ``set_input`` after + the simulation was built. Args: n: The number of households to draw. diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py index f116507d54..bc15f10040 100644 --- a/policyengine_uk/tests/test_subsample.py +++ b/policyengine_uk/tests/test_subsample.py @@ -19,32 +19,53 @@ 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 identifiers that 32-bit floats cannot tell apart. + kept. This holds for identifiers that 32-bit floats cannot tell apart, and + whatever the stored ``person_household_id`` and ``person_benunit_id`` + 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. Calculated values are not kept. 6. Determinism: the same seed keeps the same households with the same weights. -7. Baseline: the baseline holds the same households and weights, and each - simulation keeps its own policy. -8. Dataset (differential): ``dataset`` holds the same records and weights as - the populations, and a simulation built from it gives the same results as - the subsampled simulation. -9. Unbiasedness: over many seeds, a household's mean quantized weight is its - old weight. Intended exception: unquantized weights keep the total exactly - and are therefore a ratio estimate, whose mean is off by a few per cent in - a small sample and closes on the old weight as the sample grows. +7. Baseline: the baseline is replaced with a copy that holds the same + households and weights and keeps its own policy. The baseline held before + is unchanged, so subsampling a clone leaves the source simulation and its + baseline as they were. +8. Dataset (differential): ``dataset`` holds the same records as the + populations, from the rows of the dataset the simulation was built from, + with weights rescaled in the same way. For a simulation given no inputs + after it was built, a simulation built from that dataset gives the same + results as the subsampled simulation. Intended: an input set with + ``set_input`` after the build stays out of the dataset, as it was out of + the dataset before subsampling. +9. Unbiasedness: over many seeds, a household's mean quantized weight in the + sampling period is its old weight. Intended exceptions: in another period, + the mean is the household's share of the sampling period's total times + that period's total, which is its old weight only if its share is the same + in both periods; and unquantized weights keep the total exactly, so they + are a ratio estimate, whose mean is off by a few per cent in a small + sample and closes on the old weight as the sample grows. +10. Trace and inputs: the tracer starts again, and a supplied input stored on + disk is kept, not a value calculated later and cached over it in memory. + +The properties are generated for datasets of one year, which the model +projects, and of several years whose weights change differently from +household to household. Intended exceptions: a simulation built from a situation, one given no household weights, a branch, and a baseline holding other households raise ``ValueError`` and are left unchanged. """ +from types import SimpleNamespace + import numpy as np import pandas as pd import pytest from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st +from policyengine_core.data_storage import OnDiskStorage from policyengine_core.periods import period as period_ +from policyengine_core.tracers import FullTracer from policyengine_uk import CountryTaxBenefitSystem, Microsimulation, Simulation from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset @@ -54,13 +75,15 @@ rescale_weights, subsample_dataset, ) -from policyengine_uk.utils.supplied_inputs import supplied_input_arrays +from policyengine_uk.utils.supplied_inputs import ( + stored_supplied_value, + supplied_input_arrays, +) pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") YEARS = range(2025, 2031) # Neighbouring identifiers of the size the UK datasets use. As 32-bit floats, -# which is how the model stores person_household_id and person_benunit_id, # identifiers this large are exact only at multiples of 64. FIRST_HOUSEHOLD_ID = 900_000_001 FIRST_BENUNIT_ID = 906_000_001 @@ -184,9 +207,6 @@ def test_identifiers_beyond_float32_precision_keep_their_households(): simulation = Microsimulation(dataset=dataset()) households = members(simulation, "household") benunits = members(simulation, "benunit") - # The stored membership inputs cannot tell these households apart. - stored = values(simulation, "person_household_id", 2025) - assert len(set(stored.tolist())) < len(households) simulation.subsample(n=5, seed=3, time_period=2025) @@ -203,6 +223,24 @@ def test_identifiers_beyond_float32_precision_keep_their_households(): assert values(simulation, "household_id", 2025).tolist() == kept +def test_memberships_come_from_the_populations_not_the_stored_identifiers(): + simulation = Microsimulation(dataset=dataset()) + households = members(simulation, "household") + benunits = members(simulation, "benunit") + people = len(ids(simulation, "person")) + for variable in ("person_household_id", "person_benunit_id"): + simulation.set_input(variable, 2025, np.zeros(people)) + + simulation.subsample(n=5, seed=3) + + for entity, before in (("household", households), ("benunit", benunits)): + assert members(simulation, entity) == { + group: before[group] for group in ids(simulation, entity) + } + # Intended: the stored values are supplied inputs, so they are kept. + assert set(values(simulation, "person_household_id", 2025).tolist()) == {0} + + def test_zero_weight_households_are_never_kept(): households = [(0.0, [1]), (50.0, [2]), (0.0, [1, 1]), (70.0, [1])] for seed in range(5): @@ -358,6 +396,80 @@ def test_a_calculated_value_that_replaced_an_input_is_not_kept(): ) +def input_on_disk_under_a_cached_value(simulation, folder): + """Store household_net_income for 2025 as an input on disk, then cache a + calculated value over it in memory. Returns the input.""" + period = period_(2025) + holder = simulation.get_holder("household_net_income") + holder._disk_storage = OnDiskStorage(str(folder), preserve_storage_dir=True) + holder._on_disk_storable = True + # A memory limit no simulation can stay under sends the input to disk. + simulation.memory_config = SimpleNamespace(max_memory_occupation_pc=-1) + supplied_value = np.full(len(HOUSEHOLDS), 111.0) + simulation.set_input("household_net_income", period, supplied_value) + holder._on_disk_storable = False + simulation.memory_config = None + holder.put_in_cache(np.full(len(HOUSEHOLDS), 999.0), period) + return supplied_value + + +def test_an_input_on_disk_is_kept_not_a_value_cached_over_it(tmp_path): + simulation = Microsimulation(dataset=dataset()) + supplied_value = input_on_disk_under_a_cached_value(simulation, tmp_path) + holder = simulation.get_holder("household_net_income") + assert holder._memory_storage.get(period_(2025)).tolist() == [999.0] * 6 + + assert stored_supplied_value(holder, period_(2025), "default").tolist() == ( + supplied_value.tolist() + ) + simulation.subsample(n=2, seed=1) + assert set(values(simulation, "household_net_income", 2025).tolist()) == {111.0} + + +def test_supplied_values_match_policyengine_cores_recorded_inputs(tmp_path): + """Differential: policyengine-core 3.33 reads a recorded input from its + tier in ``Simulation._get_recorded_input_array``; earlier versions have + no such method, and ``stored_supplied_value`` follows the same rule.""" + simulation = Microsimulation(dataset=dataset()) + if not hasattr(simulation, "_get_recorded_input_array"): + pytest.skip("policyengine-core reads recorded inputs from 3.33") + input_on_disk_under_a_cached_value(simulation, tmp_path) + people = len(ids(simulation, "person")) + simulation.set_input("hours_worked", 2027, np.full(people, 1_000.0)) + simulation.get_holder("hours_worked").put_in_cache( + np.full(people, 500.0), period_(2027) + ) + simulation.calculate("income_tax", 2026) + + keys = sorted(simulation._user_input_keys, key=lambda key: (key[0], str(key[2]))) + assert keys + for variable, branch, period in keys: + ours = stored_supplied_value(simulation.get_holder(variable), period, branch) + cores = simulation._get_recorded_input_array(variable, period) + if cores is None: + assert ours is None, (variable, period) + else: + assert np.asarray(ours).tolist() == np.asarray(cores).tolist(), ( + variable, + period, + ) + + +def test_a_trace_starts_again_on_the_sample(): + simulation = Microsimulation(dataset=dataset(), trace=True) + simulation.calculate("household_net_income", 2025) + # A baseline traced in its simulation stays traced in it. + simulation.baseline.tracer = simulation.tracer + + simulation.subsample(n=1, seed=1) + simulation.calculate("household_net_income", 2025) + + assert isinstance(simulation.tracer, FullTracer) + assert simulation.baseline.tracer is simulation.tracer + flat = simulation.tracer.get_serialized_flat_trace() + assert len(flat["household_net_income<2025, (default)>"]["value"]) == 1 + + def test_results_match_the_same_households_in_the_full_simulation(): """Subsampling changes which households are held and their weights, not what the model calculates for a household.""" @@ -442,6 +554,47 @@ def test_a_reform_and_its_baseline_keep_the_same_households(): ) +@pytest.mark.parametrize("reform", [None, {PERSONAL_ALLOWANCE: 0}]) +def test_subsampling_a_clone_leaves_the_source_and_its_baseline_alone(reform): + """policyengine-core's ``clone`` shares a separate baseline between the + source and the clone, so the clone must not subsample it in place.""" + source = Microsimulation(dataset=dataset(), reform=reform) + baselines = [] + held = source + while getattr(held, "baseline", None) is not None: + held = held.baseline + baselines.append((held, held.dataset, weights_by_id(held, 2026))) + clone = source.clone() + assert clone.baseline is source.baseline + + clone.subsample(n=2, seed=1) + + assert len(ids(clone, "household")) <= 2 + assert ids(clone.baseline, "household") == ids(clone, "household") + assert len(ids(source, "household")) == len(HOUSEHOLDS) + held = source + for baseline, data, weights in baselines: + assert held.baseline is baseline + assert len(ids(baseline, "household")) == len(HOUSEHOLDS) + assert baseline.dataset is data + assert weights_by_id(baseline, 2026) == weights + held = baseline + # The source can still be subsampled on its own. + source.subsample(n=3, seed=2) + assert ids(source.baseline, "household") == ids(source, "household") + + +def test_the_baseline_held_before_is_left_as_it_was(): + simulation = Microsimulation(dataset=dataset()) + before = simulation.baseline + + simulation.subsample(n=2, seed=1) + + assert simulation.baseline is not before + assert len(ids(before, "household")) == len(HOUSEHOLDS) + assert ids(simulation.baseline, "household") == ids(simulation, "household") + + def test_a_baseline_holding_other_households_is_refused_and_nothing_changes(): simulation = Microsimulation(dataset=dataset()) simulation.baseline.subsample(n=2, seed=1) @@ -485,6 +638,35 @@ def test_the_dataset_holds_the_sample(): ) +def test_inputs_set_after_the_build_stay_out_of_the_dataset(): + """Intended: ``dataset`` is the subsample of the dataset the simulation + was built from, which held no values set afterwards either. So a later + input stays in the simulation only, and a weight set on the simulation + alone is rescaled to its own total there, while the baseline and the + dataset keep the dataset's weights.""" + simulation = Microsimulation(dataset=dataset()) + people = len(ids(simulation, "person")) + simulation.set_input("hours_worked", 2025, np.full(people, 1_234.0)) + simulation.set_input("household_weight", 2025, np.ones(len(HOUSEHOLDS))) + hours_in_data = Microsimulation(dataset=simulation.dataset).calculate( + "hours_worked", 2025 + ) + + simulation.subsample(n=2, seed=1) + + rebuilt = Microsimulation(dataset=simulation.dataset) + assert set(values(simulation, "hours_worked", 2025).tolist()) == {1_234.0} + assert set(values(rebuilt, "hours_worked", 2025).tolist()) <= set( + np.asarray(hours_in_data).tolist() + ) + assert values(simulation, "household_weight", 2025).sum() == pytest.approx( + len(HOUSEHOLDS) + ) + total = sum(weight for weight, _ in HOUSEHOLDS) + for held in (simulation.baseline, rebuilt): + assert values(held, "household_weight", 2025).sum() == pytest.approx(total) + + def test_subsampling_a_single_year_dataset_keeps_its_type(): data = dataset() counts = np.array([0, 2, 0, 1, 0, 0]) @@ -499,7 +681,7 @@ 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(): households = np.arange(4) frame = pd.DataFrame( { @@ -513,12 +695,10 @@ def test_a_simulation_built_from_a_dataframe_has_no_dataset_to_subsample(): } ) simulation = Microsimulation(dataset=frame) - assert simulation.dataset is None 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 +706,10 @@ 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") } + assert isinstance(simulation.dataset, UKMultiYearDataset) + tables = simulation.dataset[2025] + assert tables.household.household_id.tolist() == kept + assert tables.household.household_weight.sum() == pytest.approx(100.0) # Refusals. @@ -747,6 +931,25 @@ def error(n): assert error(400) < 0.02 +def test_another_periods_mean_weight_follows_the_sampling_periods_shares(): + """Intended: households are drawn by their weights in the sampling period + only, so in another period a household's mean new weight is its share of + the sampling period's total times that period's total. That is its old + weight there only if its share is the same in both periods.""" + sampling = np.array([1.0, 2.0, 3.0, 4.0, 10.0]) + other = np.array([5.0, 1.0, 3.0, 4.0, 2.0]) + expected = sampling / sampling.sum() * other.sum() + mean = np.mean( + [ + rescale_weights(other, draw_household_counts(sampling, 40, seed)) + for seed in range(3_000) + ], + axis=0, + ) + assert mean == pytest.approx(expected, rel=0.03) + assert not np.allclose(mean, other, rtol=0.2) + + # Properties of a subsampled simulation. SIMULATION_SETTINGS = settings( @@ -765,32 +968,76 @@ def error(n): ).filter(lambda households: any(weight > 0 for weight, _ in households)) +MULTI_YEARS = (2025, 2026, 2027) + + +def multi_year_dataset(households, factors) -> UKMultiYearDataset: + """A dataset of ``households`` for each of MULTI_YEARS, with each + household's weight multiplied by its own factor in each year.""" + return UKMultiYearDataset( + datasets=[ + dataset( + [ + (weight * factor, units) + for (weight, units), factor in zip(households, year_factors) + ], + year=year, + ) + for year, year_factors in zip(MULTI_YEARS, factors) + ] + ) + + @SIMULATION_SETTINGS @given( households=household_specs, n=st.integers(1, 10), seed=st.integers(0, 10_000), - year=st.sampled_from(list(YEARS)), quantize=st.booleans(), + multi_year=st.booleans(), + data=st.data(), ) -def test_subsample_properties(households, n, seed, year, quantize): +def test_subsample_properties(households, n, seed, quantize, multi_year, data): + if multi_year: + factors = data.draw( + st.lists( + st.lists( + st.floats(min_value=0.5, max_value=2.0), + min_size=len(households), + max_size=len(households), + ), + min_size=len(MULTI_YEARS), + max_size=len(MULTI_YEARS), + ), + label="factors", + ) + years = list(MULTI_YEARS) + else: + years = list(YEARS) + year = data.draw(st.sampled_from(years), label="year") + + def build(): + if multi_year: + return Microsimulation(dataset=multi_year_dataset(households, factors)) + return Microsimulation(dataset=dataset(households)) + def subsample(): - simulation = Microsimulation(dataset=dataset(households)) - return simulation.subsample( + return build().subsample( n=n, seed=seed, time_period=year, quantize_weights=quantize ) - full = Microsimulation(dataset=dataset(households)) + full = build() sample = subsample() kept = ids(sample, "household") people = ids(sample, "person") # 1. Size. assert 1 <= len(kept) <= n - full_weights = {period: weights_by_id(full, period) for period in YEARS} + full_weights = {period: weights_by_id(full, period) for period in years} assert all(full_weights[year][household] > 0 for household in kept) - for period in YEARS: + draws_by_period = {} + for period in years: weights = weights_by_id(sample, period) total = sum(full_weights[period].values()) # 2. Weight totals. @@ -807,6 +1054,9 @@ def subsample(): draws = shares * n / shares.sum() assert np.abs(draws - draws.round()).max() < 1e-2 assert draws.round().sum() == n + draws_by_period[period] = draws.round().tolist() + # The same draws set the weights of every period. + assert len({tuple(draws) for draws in draws_by_period.values()}) == 1 # 4. Links. for entity in ("household", "benunit"): @@ -822,7 +1072,7 @@ def subsample(): # 5. Inputs. assert supplied(sample) == supplied(full) - for period in YEARS: + for period in years: for variable in ("age", "employment_income", "region", "council_tax"): expected = by_id(full, variable, period) assert by_id(sample, variable, period) == { @@ -835,7 +1085,7 @@ def subsample(): # 6. Determinism. again = subsample() assert ids(again, "person") == people - for period in YEARS: + for period in years: assert weights_by_id(again, period) == weights_by_id(sample, period) # 7. Baseline. @@ -850,7 +1100,7 @@ def subsample(): assert getattr(sample.dataset[year], entity)[f"{entity}_id"].tolist() == ids( sample, entity ) - for period in YEARS: + for period in years: assert list(weights_by_id(rebuilt, period).values()) == pytest.approx( list(weights_by_id(sample, period).values()), rel=1e-5 ) diff --git a/policyengine_uk/utils/subsample.py b/policyengine_uk/utils/subsample.py index e636bc919c..06c5aeeb5a 100644 --- a/policyengine_uk/utils/subsample.py +++ b/policyengine_uk/utils/subsample.py @@ -5,14 +5,16 @@ ``subsample`` cannot do this for a UK simulation: - it rebuilds through ``Simulation.build_from_dataset()`` with no arguments, - and the UK builder takes the dataset to load, which it projects to later - years again; + and the UK builder takes the dataset to load; +- given that dataset, the UK builder still cannot load what core rebuilds + from: core's flat-file ``Dataset``, which ``Dataset.load_dataset`` does not + read. The UK builder would also rebuild every year from the first one, + projected forward, rather than keep each year's inputs; - it reads each person's household and benefit unit back from the stored - ``person_household_id`` and ``person_benunit_id`` inputs. The UK model stores - those as 32-bit floats, which hold integers exactly only up to 2**24, so - larger identifiers would merge households and benefit units; -- it leaves ``Simulation.dataset`` as a policyengine-core ``Dataset``, while UK - formulas read ``dataset.years``; + ``person_household_id`` and ``person_benunit_id`` inputs, which need not be + exact: the UK model has stored them as 32-bit floats, which hold integers + exactly only up to 2**24, so larger identifiers would merge households and + benefit units (policyengine-uk#2267); - it rebuilds a ``baseline`` branch, while a UK simulation's baseline is a simulation of its own. @@ -29,6 +31,7 @@ import numpy as np import pandas as pd from policyengine_core.periods import Period, period as period_ +from policyengine_core.tracers import SimpleTracer from policyengine_uk.data.dataset_schema import ( UKMultiYearDataset, @@ -314,9 +317,11 @@ def _plan( ) -def _apply(subsample: _Subsample) -> None: - """Rebuild a simulation on its subsample.""" - simulation = subsample.simulation +def _apply( + subsample: _Subsample, simulation: "Simulation", tracer: SimpleTracer +) -> None: + """Rebuild ``simulation`` on ``subsample``, which was worked out for it + or for the simulation it is a copy of.""" start_instant = simulation.start_instant # Branches hold the populations the subsample replaces; formulas make # them again when they next ask for them. @@ -325,6 +330,9 @@ def _apply(subsample: _Subsample) -> None: # of what was calculated from them, starts again with the inputs it sets. simulation._user_input_keys = set() simulation._store_history = None + # A trace of calculations on the replaced populations would otherwise be + # reported for the sample: a flat trace keeps a variable's first value. + simulation.tracer = tracer simulation.build_from_ids(*subsample.ids) for name, period, values in subsample.inputs: simulation.set_input(name, period, values) @@ -391,14 +399,42 @@ def subsample_simulation( simulations.append(simulation) simulation = getattr(simulation, "baseline", None) - # Work out every subsample before applying any, so that a simulation - # that cannot be subsampled leaves them all as they were. + # Work out every subsample before changing anything, so that a + # simulation that cannot be subsampled leaves them all as they were. household_ids = np.asarray(simulations[0].populations["household"].ids) datasets: Dict[int, UKDataset] = {} subsamples = [ _plan(each, household_ids, counts, quantize_weights, datasets) for each in simulations ] - for subsample in subsamples: - _apply(subsample) - return simulations[0] + + # Each simulation gets a new tracer of the same kind. Simulations that + # shared a tracer share the new one. The old tracer is kept with its + # replacement so that its id is not reused while this runs. + tracers: Dict[int, Tuple[SimpleTracer, SimpleTracer]] = {} + + def new_tracer(old: SimpleTracer) -> SimpleTracer: + if id(old) not in tracers: + tracers[id(old)] = (old, type(old)()) + return tracers[id(old)][1] + + first = simulations[0] + _apply(subsamples[0], first, new_tracer(first.tracer)) + + # Replace each baseline with a subsampled copy rather than subsample it + # in place: policyengine-core's ``clone`` shares a separate baseline, so + # another simulation can hold this one and still hold every household. + # Copying each baseline after the simulation that holds it is cut down + # keeps the memory in use to about what the simulations held before. + copies = {id(first): first} + for subsample in subsamples[1:]: + baseline = subsample.simulation + copy = baseline.clone(debug=baseline.debug, trace=baseline.trace) + _apply(subsample, copy, new_tracer(baseline.tracer)) + copies[id(baseline)] = copy + for each in simulations: + holder = copies[id(each)] + baseline = getattr(holder, "baseline", None) + if baseline is not None and id(baseline) in copies: + holder.baseline = copies[id(baseline)] + return first diff --git a/policyengine_uk/utils/supplied_inputs.py b/policyengine_uk/utils/supplied_inputs.py index 9363e5be20..e9b7a5cdcf 100644 --- a/policyengine_uk/utils/supplied_inputs.py +++ b/policyengine_uk/utils/supplied_inputs.py @@ -21,8 +21,9 @@ deletes. ``supplied_input_arrays`` answers for every variable. It reads only a record -whose supplied value is still stored, which the holder tells apart from a -value the engine carried over or calculated into the same period. +whose supplied value is still stored, from the storage tier (memory or disk) +the holder records it in. A value the engine carried over or calculated into +the same period, in either tier, is not read as supplied. """ from typing import List, Optional, Tuple @@ -147,12 +148,38 @@ def supplied_input_arrays(simulation) -> List[Tuple[str, Period, np.ndarray]]: for variable_name, branch_name, period in input_keys: if branch_name != simulation.branch_name: continue - holder = simulation.get_holder(variable_name) - # A record can outlive its stored value (see above), and a carried - # over or calculated value can then take its place. - if not holder._stores_user_input(period, branch_name): - continue - value = holder._get_array_from_storage(period, branch_name) + value = stored_supplied_value( + simulation.get_holder(variable_name), period, branch_name + ) if value is not None: arrays.append((variable_name, period, value)) return sorted(arrays, key=lambda item: (item[0], item[1].start, str(item[1]))) + + +def stored_supplied_value(holder, period: Period, branch_name: str): + """The supplied value ``holder`` still stores for ``period``, or None. + + A record can outlive its stored value (see above), and a value calculated + later can take its place in one tier while the input survives in the + other. So read only a tier that the holder records the input in (either + tier when it records none) and that holds a value it did not calculate, + as ``Holder._stores_user_input`` decides whether an input survives. + Reading memory before disk instead could return a calculated value cached + over an input stored on disk. + """ + recorded = (getattr(holder, "_user_input_storage", None) or {}).get( + (branch_name, str(period)) + ) + locations = recorded[1] if recorded is not None else {"memory", "disk"} + for location, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if ( + location in locations + and storage is not None + and storage.has(period, branch_name) + and not storage.is_derived(period, branch_name) + ): + return storage.get(period, branch_name) + return None From 5dd27677c564c4d13460932822b9247aec62a1a6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 11:23:51 -0400 Subject: [PATCH 6/7] Share arrays when copying baselines and refuse negative weights 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 --- docs/book/usage/simulations.md | 4 +- policyengine_uk/simulation.py | 25 +++-- policyengine_uk/tests/test_subsample.py | 135 +++++++++++++++++++----- policyengine_uk/utils/subsample.py | 50 +++++++-- 4 files changed, 172 insertions(+), 42 deletions(-) diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index a6cf807d71..748b2d93c6 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -474,9 +474,9 @@ sim = Microsimulation( sim.subsample(n=10_000, seed=0, time_period=2026) ``` -It draws `n` households with replacement, each with probability proportional to its weight in `time_period`, and keeps every household drawn at least once, with its benefit units and people. So the sample holds at most `n` households. A household's new weight is proportional to the number of times it was drawn, and the total household weight of every year is unchanged. The values the simulation was given are kept for the sampled records; results are calculated again on the sample, so population estimates then carry sampling error. +It draws `n` households with replacement, each with probability proportional to its weight in `time_period` (with the default `quantize_weights=True`), and keeps every household drawn at least once, with its benefit units and people. So the sample holds at most `n` households. A household's new weight is proportional to the number of times it was drawn, and the total household weight of every year is unchanged. The values the simulation was given are kept for the sampled records; results are calculated again on the sample, so population estimates then carry sampling error. -The simulation's `baseline` and `dataset` are replaced with copies cut down to the same households; the objects they held before are left as they were. The `dataset` holds the rows of the dataset the simulation was built from, so it does not include values set with `set_input` after the simulation was built. The same `seed` always draws the same households. +The simulation's `baseline` and `dataset` are replaced with copies cut down to the same households; the objects they held before are left as they were. The `dataset` holds the rows of the dataset the simulation was built from, so it holds whole years only and does not include values set with `set_input` after the simulation was built. The same `seed` always draws the same households. ## Debugging simulations diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index dc8dbfdbdd..485d074114 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -578,20 +578,25 @@ def subsample( ) -> "Simulation": """Cut the simulation down to a weighted random sample of households. - Draws ``n`` households with replacement, each with probability - proportional to its weight in ``time_period``, and keeps each - household drawn at least once, with its benefit units and people. So - the sample holds at most ``n`` households. + Draws ``n`` households with replacement and keeps each household + drawn at least once, with its benefit units and people. So the sample + holds at most ``n`` households. Only households with a positive + weight in ``time_period`` can be drawn: with ``quantize_weights`` + (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 instead. Every value the simulation was given (from its dataset or through ``set_input``) is kept for the sampled records, for every period. Calculated values are dropped and recalculated on the sample, and the tracer starts again. Each supplied weight is replaced so that its total is unchanged in every period: a household's new weight is - proportional to the number of times it was drawn. Averaged over - draws, a household's new weight in ``time_period`` is its old one; in - another period that holds only if the household's share of the total - weight is the same there. + proportional to the number of times it was drawn, and also to its old + weight without ``quantize_weights``. With ``quantize_weights``, a + household's new weight in ``time_period``, averaged over draws, is + its old one; in another period that holds only if the household's + share of the total weight is the same there. ``baseline`` is replaced with a copy subsampled to the same households, which keeps its own policy and inputs. The baseline the @@ -599,8 +604,8 @@ def subsample( (a clone of this one) can hold it too. ``dataset`` is replaced with the rows of the dataset the simulation was built from for the sampled records, with its weights rescaled in the same way. Like the - dataset before, it does not hold values set with ``set_input`` after - the simulation was built. + dataset before, it holds whole years only, and not values set with + ``set_input`` after the simulation was built. Args: n: The number of households to draw. diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py index bc15f10040..dffa3e9d11 100644 --- a/policyengine_uk/tests/test_subsample.py +++ b/policyengine_uk/tests/test_subsample.py @@ -32,11 +32,13 @@ baseline as they were. 8. Dataset (differential): ``dataset`` holds the same records as the populations, from the rows of the dataset the simulation was built from, - with weights rescaled in the same way. For a simulation given no inputs - after it was built, a simulation built from that dataset gives the same - results as the subsampled simulation. Intended: an input set with - ``set_input`` after the build stays out of the dataset, as it was out of - the dataset before subsampling. + with weights rescaled in the same way. For a simulation whose inputs are + all in its dataset, a simulation built from that dataset gives the same + results as the subsampled simulation. Intended: an input the dataset does + not hold stays out of it, as it was out of the dataset before + subsampling. That is an input set with ``set_input`` after the build, or + one for a part of a year, which a DataFrame can give and a dataset of + whole years cannot hold. 9. Unbiasedness: over many seeds, a household's mean quantized weight in the sampling period is its old weight. Intended exceptions: in another period, the mean is the household's share of the sampling period's total times @@ -52,8 +54,8 @@ household to household. Intended exceptions: a simulation built from a situation, one given no -household weights, a branch, and a baseline holding other households raise -``ValueError`` and are left unchanged. +household weights or a negative one, a branch, and a baseline holding other +households raise ``ValueError`` and are left unchanged. """ from types import SimpleNamespace @@ -61,7 +63,7 @@ import numpy as np import pandas as pd import pytest -from hypothesis import HealthCheck, given, settings +from hypothesis import HealthCheck, example, given, settings from hypothesis import strategies as st from policyengine_core.data_storage import OnDiskStorage from policyengine_core.periods import period as period_ @@ -595,6 +597,32 @@ def test_the_baseline_held_before_is_left_as_it_was(): assert ids(simulation.baseline, "household") == ids(simulation, "household") +def test_a_baseline_copy_does_not_copy_the_baselines_arrays(monkeypatch): + """The copy of a baseline rebuilds its arrays at once, so copying the + baseline's cached arrays first would only double its memory.""" + from policyengine_core.data_storage import InMemoryStorage + + from policyengine_uk.utils import subsample as module + + if module._branch_clone is None: + pytest.skip("this policyengine-core cannot share a clone's arrays") + simulation = Microsimulation(dataset=dataset(), reform={PERSONAL_ALLOWANCE: 0}) + simulation.baseline.calculate("household_net_income", 2026) + calls = [] + clone = InMemoryStorage.clone + + def recording_clone(self, share_arrays=False): + calls.append(share_arrays) + return clone(self, share_arrays=share_arrays) + + monkeypatch.setattr(InMemoryStorage, "clone", recording_clone) + + simulation.subsample(n=3, seed=1) + + assert calls + assert all(calls) + + def test_a_baseline_holding_other_households_is_refused_and_nothing_changes(): simulation = Microsimulation(dataset=dataset()) simulation.baseline.subsample(n=2, seed=1) @@ -667,6 +695,28 @@ def test_inputs_set_after_the_build_stay_out_of_the_dataset(): assert values(held, "household_weight", 2025).sum() == pytest.approx(total) +def test_an_input_for_part_of_a_year_stays_out_of_the_dataset(): + """Intended: a DataFrame can give a monthly input, which the simulation + keeps and a dataset of whole years cannot hold, so a simulation built + from the subsampled dataset does not have it.""" + frame = pd.DataFrame( + { + "person_id__2025": [1, 2], + "household_id__2025": [1, 2], + "household_weight__2025": [1.0, 1.0], + "ebr_energy_bills_credit_eligible__2025-01": [False, False], + } + ) + simulation = Microsimulation(dataset=frame) + + simulation.subsample(n=1, seed=1) + + monthly = ("ebr_energy_bills_credit_eligible", "2025-01") + assert monthly in supplied(simulation) + assert values(simulation, *monthly).tolist() == [False] + assert monthly not in supplied(Microsimulation(dataset=simulation.dataset)) + + def test_subsampling_a_single_year_dataset_keeps_its_type(): data = dataset() counts = np.array([0, 2, 0, 1, 0, 0]) @@ -735,6 +785,28 @@ def test_a_simulation_without_weights_is_refused(): simulation.subsample(n=2) +def test_a_negative_household_weight_is_refused_and_nothing_changes(): + simulation = Microsimulation(dataset=dataset([(1.0, [1]), (-2.0, [1])])) + weights = weights_by_id(simulation, 2025) + + with pytest.raises(ValueError, match="negative"): + simulation.subsample(n=1, seed=0) + + assert weights_by_id(simulation, 2025) == weights + + +def test_another_negative_weight_is_refused(): + data = dataset() + data.household["original_weight"] = [1.0, 2.0, 3.0, -4.0, 5.0, 6.0] + simulation = Microsimulation(dataset=data) + households = ids(simulation, "household") + + with pytest.raises(ValueError, match="original_weight for 2025.*negative"): + simulation.subsample(n=3) + + assert ids(simulation, "household") == households + + def test_a_simulation_with_no_positive_weight_is_refused(): households = [(0.0, [1]), (0.0, [2])] simulation = Microsimulation(dataset=dataset(households)) @@ -952,6 +1024,7 @@ def test_another_periods_mean_weight_follows_the_sampling_periods_shares(): # Properties of a subsampled simulation. +MAX_HOUSEHOLDS = 7 SIMULATION_SETTINGS = settings( max_examples=10, deadline=None, @@ -964,7 +1037,7 @@ def test_another_periods_mean_weight_follows_the_sampling_periods_shares(): st.lists(st.integers(1, 3), min_size=1, max_size=2), ), min_size=2, - max_size=7, + max_size=MAX_HOUSEHOLDS, ).filter(lambda households: any(weight > 0 for weight, _ in households)) @@ -988,6 +1061,26 @@ def multi_year_dataset(households, factors) -> UKMultiYearDataset: ) +FACTORS = st.lists( + st.lists( + st.sampled_from([0.5, 0.8, 1.0, 1.25, 2.0]), + min_size=MAX_HOUSEHOLDS, + max_size=MAX_HOUSEHOLDS, + ), + min_size=len(MULTI_YEARS), + max_size=len(MULTI_YEARS), +) +# Households whose shares of the total weight change from year to year. +CHANGING_SHARES = dict( + households=[(100.0, [1]), (300.0, [2]), (0.0, [1]), (50.0, [1, 2])], + n=6, + seed=11, + multi_year=True, + factors=[[1.0] * MAX_HOUSEHOLDS, [2.0, 0.5, 1.0, 1.25, 1, 1, 1], [0.5] + [2.0] * 6], + year_index=1, +) + + @SIMULATION_SETTINGS @given( households=household_specs, @@ -995,26 +1088,20 @@ def multi_year_dataset(households, factors) -> UKMultiYearDataset: seed=st.integers(0, 10_000), quantize=st.booleans(), multi_year=st.booleans(), - data=st.data(), + factors=FACTORS, + year_index=st.integers(0, len(YEARS) - 1), ) -def test_subsample_properties(households, n, seed, quantize, multi_year, data): +@example(quantize=True, **CHANGING_SHARES) +@example(quantize=False, **CHANGING_SHARES) +def test_subsample_properties( + households, n, seed, quantize, multi_year, factors, year_index +): if multi_year: - factors = data.draw( - st.lists( - st.lists( - st.floats(min_value=0.5, max_value=2.0), - min_size=len(households), - max_size=len(households), - ), - min_size=len(MULTI_YEARS), - max_size=len(MULTI_YEARS), - ), - label="factors", - ) + factors = [year_factors[: len(households)] for year_factors in factors] years = list(MULTI_YEARS) else: years = list(YEARS) - year = data.draw(st.sampled_from(years), label="year") + year = years[year_index % len(years)] def build(): if multi_year: diff --git a/policyengine_uk/utils/subsample.py b/policyengine_uk/utils/subsample.py index 06c5aeeb5a..704c182ca7 100644 --- a/policyengine_uk/utils/subsample.py +++ b/policyengine_uk/utils/subsample.py @@ -33,6 +33,17 @@ from policyengine_core.periods import Period, period as period_ from policyengine_core.tracers import SimpleTracer +try: + # How ``get_branch`` makes a clone that shares the cached arrays instead + # of copying them (policyengine-core 3.32.27 and later). See + # ``_copy_sharing_arrays``. + from policyengine_core.simulations.simulation import ( + _BranchClone, + _branch_clone, + ) +except ImportError: # pragma: no cover - a core without shared clones + _BranchClone = _branch_clone = None + from policyengine_uk.data.dataset_schema import ( UKMultiYearDataset, UKSingleYearDataset, @@ -65,9 +76,10 @@ def _random_state(seed) -> np.random.RandomState: """A generator for ``seed``, which may be a number or text. - ``RandomState`` keeps the same stream in every NumPy version, so a seed - chooses the same households wherever it runs. The digest gives the same - number in every process, which the built-in ``hash`` does not. + NumPy's compatibility guarantee keeps ``RandomState``'s results for a + fixed seed and calls the same, up to roundoff, so a seed chooses the same + households wherever it runs. The digest gives the same number in every + process, which the built-in ``hash`` does not. """ digest = hashlib.sha256(str(seed).encode("utf-8")).digest() return np.random.RandomState(int.from_bytes(digest[:4], "big")) @@ -89,6 +101,8 @@ def draw_household_counts( if n < 1: raise ValueError(f"A subsample needs at least one household, not {n}.") weights = np.asarray(weights, dtype=float) + if (weights < 0).any(): + raise ValueError("Household weights must not be negative to sample by them.") eligible = np.flatnonzero(weights > 0) if eligible.size == 0: raise ValueError("No household has a positive weight to sample by.") @@ -118,6 +132,8 @@ def rescale_weights( total = weights.sum() if not np.isfinite(total): raise ValueError("the weights do not have a finite total") + if (weights < 0).any(): + raise ValueError("the weights include negative values") if total == 0: return np.zeros_like(weights) drawn = counts.astype(float) if quantize_weights else weights * counts @@ -342,6 +358,27 @@ def _apply( simulation.dataset = subsample.dataset +def _copy_sharing_arrays(simulation: "Simulation") -> "Simulation": + """A clone of ``simulation`` that shares its cached arrays. + + ``clone`` copies every cached array, and the subsample discards them all + at once, so for a full-size baseline that doubles its memory for nothing. + Announce the clone as ``get_branch`` does: the copy then reads views of + the arrays and copies one only if it reads it, which it never does. The + simulation is not changed; with a core that cannot share, the arrays are + copied. + """ + if _branch_clone is None: # pragma: no cover - a core without shared clones + return simulation.clone(debug=simulation.debug, trace=simulation.trace) + request = _BranchClone(simulation) + token = _branch_clone.set(request) + try: + return simulation.clone(debug=simulation.debug, trace=simulation.trace) + finally: + request.pending = False + _branch_clone.reset(token) + + def _sampling_weights(simulation: "Simulation", time_period) -> np.ndarray: """Household weights for ``time_period``, to draw the sample by. @@ -424,12 +461,13 @@ def new_tracer(old: SimpleTracer) -> SimpleTracer: # Replace each baseline with a subsampled copy rather than subsample it # in place: policyengine-core's ``clone`` shares a separate baseline, so # another simulation can hold this one and still hold every household. - # Copying each baseline after the simulation that holds it is cut down - # keeps the memory in use to about what the simulations held before. + # The copy shares the baseline's arrays rather than copying them, so it + # costs little memory; the old baselines stay in memory until nothing + # holds them. copies = {id(first): first} for subsample in subsamples[1:]: baseline = subsample.simulation - copy = baseline.clone(debug=baseline.debug, trace=baseline.trace) + copy = _copy_sharing_arrays(baseline) _apply(subsample, copy, new_tracer(baseline.tracer)) copies[id(baseline)] = copy for each in simulations: From 01f5d051be089ba93429fdbf3268f247d4f25393 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 11:30:34 -0400 Subject: [PATCH 7/7] Test that a negative sampling weight is refused before the draw 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 --- policyengine_uk/tests/test_subsample.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/policyengine_uk/tests/test_subsample.py b/policyengine_uk/tests/test_subsample.py index dffa3e9d11..73426a93f1 100644 --- a/policyengine_uk/tests/test_subsample.py +++ b/policyengine_uk/tests/test_subsample.py @@ -958,6 +958,13 @@ def test_a_second_weight_keeps_its_own_total(weights, counts, quantize): assert (rescaled[counts == 0] == 0).all() +def test_negative_sampling_weights_are_refused_before_drawing(): + """A negative weight would otherwise just never be drawn, and the total + it lowers would be shared out among the drawn households.""" + with pytest.raises(ValueError, match="must not be negative to sample"): + draw_household_counts(np.array([1.0, -2.0, 3.0]), 2, 0) + + def test_weights_without_a_total_are_refused_or_stay_zero(): counts = np.array([1, 2]) assert rescale_weights(np.zeros(2), counts).tolist() == [0.0, 0.0]