From e95c2a60ec0ba23b2eb67195b2b3edd72e44c1ba Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:36:32 -0400 Subject: [PATCH 01/13] Rebuild a subsampled simulation from its inputs only subsample exported every stored value, calculated ones included, and loaded them all back as inputs. A formula result calculated before subsampling then replaced its formula for good: it was carried over past the formula's end and survived apply_reform and later set_input calls, so results after subsampling depended on what had been calculated before. It now exports the values the simulation was given (loaded from the dataset or passed to set_input), for variables with a formula too, so formula-backed structural IDs from the dataset are still kept. Co-Authored-By: Claude Opus 5.5 --- .../fix-subsample-inputs-only.fixed.md | 1 + policyengine_core/simulations/simulation.py | 32 ++++- tests/core/test_subsample_inputs_only.py | 106 +++++++++++++++++ .../test_subsample_inputs_only_property.py | 53 +++++++++ tests/fixtures/subsample_inputs.py | 109 ++++++++++++++++++ 5 files changed, 295 insertions(+), 6 deletions(-) create mode 100644 changelog.d/fix-subsample-inputs-only.fixed.md create mode 100644 tests/core/test_subsample_inputs_only.py create mode 100644 tests/core/test_subsample_inputs_only_property.py create mode 100644 tests/fixtures/subsample_inputs.py diff --git a/changelog.d/fix-subsample-inputs-only.fixed.md b/changelog.d/fix-subsample-inputs-only.fixed.md new file mode 100644 index 000000000..b9654be46 --- /dev/null +++ b/changelog.d/fix-subsample-inputs-only.fixed.md @@ -0,0 +1 @@ +Rebuild a simulation in `subsample` from the values it was given (the dataset and `set_input`), leaving calculated values out, so results after subsampling no longer depend on what was calculated before it. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..0adb39be2 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1789,6 +1789,15 @@ def _get_exportable_input_periods( if not self._is_exportable_input_variable(variable_name): return [] + return self._get_set_input_periods(variable_name) + + def _get_set_input_periods(self, variable_name: str) -> List[Period]: + """Periods of ``variable_name`` whose stored value the simulation was given. + + That is every value loaded from the dataset or passed to ``set_input`` + on a branch this simulation reads, including for a variable that has + a formula. Values the simulation calculated are left out. + """ user_input_periods = { period for input_variable_name, branch_name, period in getattr( @@ -1824,13 +1833,19 @@ def to_input_dataframe( pd.DataFrame: The DataFrame containing the input values. """ + return self._to_person_dataframe( + lambda variable: self._get_exportable_input_periods( + variable, include_computed_variables + ) + ) + + def _to_person_dataframe(self, get_periods) -> pd.DataFrame: + """Person-level DataFrame of each variable at ``get_periods(variable)``.""" df = pd.DataFrame() for variable in self.tax_benefit_system.variables: variable_meta = self.tax_benefit_system.variables[variable] - for period in self._get_exportable_input_periods( - variable, include_computed_variables - ): + for period in get_periods(variable): # Test if period matches entity definition period if variable_meta.definition_period != period.unit: continue @@ -1898,9 +1913,14 @@ def subsample( if time_period is None: time_period = self.default_calculation_period - # Subsampling rebuilds the complete dataset, so preserve computed - # structural variables such as formula-backed IDs. - df = self.to_input_dataframe(include_computed_variables=True) + # Subsampling rebuilds the simulation from what it was given: every + # value loaded from the dataset or set with ``set_input``, including + # for variables that have a formula (structural IDs, for instance). + # Calculated values are left out. Reloaded as inputs they would + # replace their formulas for good, so later results (after a reform, + # or carried over to another year) would depend on what happened to + # be calculated before subsampling. + df = self._to_person_dataframe(self._get_set_input_periods) # Extract time period from DataFrame columns df_time_period = ( diff --git a/tests/core/test_subsample_inputs_only.py b/tests/core/test_subsample_inputs_only.py new file mode 100644 index 000000000..ecdaefdc4 --- /dev/null +++ b/tests/core/test_subsample_inputs_only.py @@ -0,0 +1,106 @@ +"""``subsample`` rebuilds a simulation from its inputs, not its calculations. + +``subsample`` used to export every stored value, calculated ones included, +and load them all back as inputs. A formula result calculated before +subsampling then replaced its formula for good: it was carried over past the +formula's end, and survived ``apply_reform`` and later ``set_input`` calls, +so results after subsampling depended on what had been calculated before. +""" + +import numpy as np + +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + build, + doubled, + stored, +) + +SEED = "inputs-only" + + +def subsampled(prior=(), n=4): + """A simulation subsampled after calculating ``prior`` (variable, period).""" + simulation = build() + for variable, period in prior: + simulation.calculate(variable, period) + simulation.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + return simulation + + +def test_formula_result_is_not_carried_past_its_end_after_subsample(): + fresh = subsampled() + used = subsampled(prior=[("ended", DATASET_YEAR)]) + assert used.calculate("ended", "2023").tolist() == [0.0] * used.persons.count + assert ( + used.calculate("ended", "2023").tolist() + == fresh.calculate("ended", "2023").tolist() + ) + + +def test_subsample_stores_the_same_values_whatever_was_calculated_before(): + fresh = subsampled() + used = subsampled( + prior=[ + ("doubled", DATASET_YEAR), + ("household_total", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("ended", DATASET_YEAR), + ("base", "2024"), + ] + ) + assert stored(used) == stored(fresh) + assert not any(name == "doubled" for name, _, _ in stored(used)) + + +def test_calculated_value_follows_a_new_input_after_subsample(): + used = subsampled(prior=[("doubled", DATASET_YEAR)]) + count = used.persons.count + used.set_input("base", DATASET_YEAR, np.full(count, 5.0)) + assert used.calculate("doubled", DATASET_YEAR).tolist() == [10.0] * count + + +def test_calculated_value_follows_a_reform_after_subsample(): + class tripled(Variable): + value_type = float + entity = doubled.entity + definition_period = doubled.definition_period + label = "Three times the input amount" + + def formula(person, period): + return 3 * person("base", period) + + tripled.__name__ = "doubled" + + class reform(Reform): + def apply(self): + self.update_variable(tripled) + + used = subsampled(prior=[("doubled", DATASET_YEAR)]) + used.apply_reform(reform) + base = used.calculate("base", DATASET_YEAR) + assert used.calculate("doubled", DATASET_YEAR).tolist() == (3 * base).tolist() + + +def test_dataset_values_for_a_formula_variable_survive_subsample(): + # ``overridden`` has a formula returning -1, and dataset values that + # replace it: person ``i`` (1-based ID) has ``i - 1``. + simulation = subsampled() + ids = simulation.calculate("person_id", DATASET_YEAR) + assert ( + simulation.calculate("overridden", DATASET_YEAR).tolist() + == (ids - 1).astype(float).tolist() + ) + assert ("overridden", "default", DATASET_YEAR) in stored(simulation) + + +def test_formulas_run_on_the_subsample_after_subsample(): + simulation = subsampled(prior=[("household_total", DATASET_YEAR)]) + assert ( + simulation.calculate("doubled", DATASET_YEAR).tolist() + == (2 * simulation.calculate("base", DATASET_YEAR)).tolist() + ) + household_total = simulation.calculate("household_total", DATASET_YEAR) + assert household_total.sum() == simulation.calculate("doubled", DATASET_YEAR).sum() diff --git a/tests/core/test_subsample_inputs_only_property.py b/tests/core/test_subsample_inputs_only_property.py new file mode 100644 index 000000000..88dbc8ad8 --- /dev/null +++ b/tests/core/test_subsample_inputs_only_property.py @@ -0,0 +1,53 @@ +"""Property: what was calculated before ``subsample`` leaves no trace. + +For any calculations made before subsampling (formula variables, weights and +inputs, at the dataset year and other years, with auto-carry-over on or off), +the subsampled simulation stores exactly the values a simulation subsampled +straight after loading does, and every later result agrees. +""" + +import pytest + +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + FORMULA_VARIABLES, + build, + stored, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +VARIABLES = FORMULA_VARIABLES + ["base", "overridden", "household_weight"] +YEARS = ["2021", DATASET_YEAR, "2023", "2024"] +REQUESTS = st.tuples(st.sampled_from(VARIABLES), st.sampled_from(YEARS)) + + +def subsampled(carry_over, prior, n, seed): + simulation = build(carry_over) + for variable, period in prior: + simulation.calculate(variable, period) + simulation.subsample(n=n, seed=seed, time_period=DATASET_YEAR) + return simulation + + +@hypothesis.settings(max_examples=150, deadline=None) +@hypothesis.given( + carry_over=st.booleans(), + prior=st.lists(REQUESTS, min_size=1, max_size=6), + later=st.lists(REQUESTS, min_size=1, max_size=4), + n=st.integers(1, 8), + seed=st.sampled_from(["a", "b", "c"]), +) +def test_calculations_before_subsample_leave_no_trace( + carry_over, prior, later, n, seed +): + fresh = subsampled(carry_over, [], n, seed) + used = subsampled(carry_over, prior, n, seed) + assert stored(used) == stored(fresh) + for variable, period in later: + assert ( + used.calculate(variable, period).tolist() + == fresh.calculate(variable, period).tolist() + ) + assert stored(used) == stored(fresh) diff --git a/tests/fixtures/subsample_inputs.py b/tests/fixtures/subsample_inputs.py new file mode 100644 index 000000000..3a5161608 --- /dev/null +++ b/tests/fixtures/subsample_inputs.py @@ -0,0 +1,109 @@ +"""A small dataset-backed simulation for testing what ``subsample`` keeps. + +Shared by ``tests/core/test_subsample_inputs_only.py`` and its Hypothesis +property module, which pytest must be able to skip on its own when +Hypothesis is not installed. +""" + +import numpy as np +import pandas as pd + +from policyengine_core.country_template import CountryTaxBenefitSystem, Simulation +from policyengine_core.country_template.entities import Household, Person +from policyengine_core.data import Dataset +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable + +DATASET_YEAR = "2022" +HOUSEHOLD_OF_PERSON = [1, 1, 2, 2, 2, 3, 4, 4, 5, 6] +BASE = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0] +WEIGHTS = {1: 100.0, 2: 200.0, 3: 300.0, 4: 400.0, 5: 500.0, 6: 600.0} + + +class base(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Input amount" + + +class doubled(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Twice the input amount" + + def formula(person, period): + return 2 * person("base", period) + + +class ended(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Seven, until its formula ends with the dataset year" + end = f"{DATASET_YEAR}-12-31" + + def formula(person, period): + return person.filled_array(7.0) + + +class household_total(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household sum of the doubled amount" + + def formula(household, period): + return household.sum(household.members("doubled", period)) + + +class overridden(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "A formula variable the dataset gives values for" + + def formula(person, period): + return person.filled_array(-1.0) + + +FORMULA_VARIABLES = ["doubled", "ended", "household_total", "person_weight"] +OVERRIDDEN = [float(index) for index in range(len(BASE))] + + +def dataframe() -> pd.DataFrame: + return pd.DataFrame( + { + f"person_id__{DATASET_YEAR}": list(range(1, len(BASE) + 1)), + f"household_id__{DATASET_YEAR}": HOUSEHOLD_OF_PERSON, + f"person_household_id__{DATASET_YEAR}": HOUSEHOLD_OF_PERSON, + f"household_weight__{DATASET_YEAR}": [ + WEIGHTS[household] for household in HOUSEHOLD_OF_PERSON + ], + f"base__{DATASET_YEAR}": BASE, + f"overridden__{DATASET_YEAR}": OVERRIDDEN, + } + ) + + +def build(carry_over: bool = True) -> Simulation: + """A ten-person, six-household simulation loaded from ``dataframe()``.""" + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = carry_over + system.add_variables(base, doubled, ended, household_total, overridden) + return Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(dataframe(), DATASET_YEAR), + ) + + +def stored(simulation) -> dict: + """Every stored value: ``{(variable, branch, period): values}``.""" + values = {} + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + for branch, period in holder.get_known_branch_periods(): + array = holder.get_array(period, branch) + values[(name, branch, str(period))] = np.asarray(array).tolist() + return values From e0b41258f43179352bfa2f2d5f43e8d28d2dfb51 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 16:56:36 -0400 Subject: [PATCH 02/13] Discard stale formula branches when subsampling Preserve the baseline policy while recreating calculation branches on the sampled population. Extend the existing differential property to formulas that create branches. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 7 +++++ tests/core/test_subsample_inputs_only.py | 31 +++++++++++++++++++ .../test_subsample_inputs_only_property.py | 7 +++++ tests/fixtures/subsample_inputs.py | 22 +++++++++++-- 4 files changed, 65 insertions(+), 2 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index a5399bf85..19a39fd7f 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2387,6 +2387,13 @@ def subsample( # as "size X != Y = count" projection errors. self._invalidate_all_caches() + # Other branches still contain the pre-subsample populations. Forget + # them so a formula's next ``get_branch`` starts from the new sample. + # Keep the baseline long enough to preserve its policy below. + self.branches = { + name: branch for name, branch in self.branches.items() if name == "baseline" + } + # Ensure the baseline branch has the new data: rebuild it from the # subsampled simulation through ``get_branch`` (the same wiring # ``__init__`` uses), then restore the saved baseline tax-benefit diff --git a/tests/core/test_subsample_inputs_only.py b/tests/core/test_subsample_inputs_only.py index ecdaefdc4..cee15f00f 100644 --- a/tests/core/test_subsample_inputs_only.py +++ b/tests/core/test_subsample_inputs_only.py @@ -8,6 +8,7 @@ """ import numpy as np +import pytest from policyengine_core.reforms import Reform from policyengine_core.variables import Variable @@ -104,3 +105,33 @@ def test_formulas_run_on_the_subsample_after_subsample(): ) household_total = simulation.calculate("household_total", DATASET_YEAR) assert household_total.sum() == simulation.calculate("doubled", DATASET_YEAR).sum() + + +@pytest.mark.parametrize("with_baseline", [False, True]) +def test_subsample_recreates_formula_branches_on_the_new_population(with_baseline): + fresh = build() + used = build() + if with_baseline: + for simulation in (fresh, used): + simulation.baseline = simulation.get_branch("baseline", clone_system=True) + baseline_system = used.baseline.tax_benefit_system + + used.calculate("via_branch", DATASET_YEAR) + stale_branch = used.get_branch("probe") + for simulation in (fresh, used): + simulation.subsample(n=4, seed=SEED, time_period=DATASET_YEAR) + + assert "probe" not in used.branches + if with_baseline: + assert used.baseline is used.branches["baseline"] + assert used.baseline.tax_benefit_system is baseline_system + assert used.baseline.persons.count == used.persons.count + + for year in (DATASET_YEAR, "2023"): + np.testing.assert_array_equal( + used.calculate("via_branch", year), fresh.calculate("via_branch", year) + ) + branch = used.get_branch("probe") + assert branch is not stale_branch + assert branch.persons.count == used.persons.count + assert branch.household.count == used.household.count diff --git a/tests/core/test_subsample_inputs_only_property.py b/tests/core/test_subsample_inputs_only_property.py index 88dbc8ad8..e26dfddb1 100644 --- a/tests/core/test_subsample_inputs_only_property.py +++ b/tests/core/test_subsample_inputs_only_property.py @@ -32,6 +32,13 @@ def subsampled(carry_over, prior, n, seed): @hypothesis.settings(max_examples=150, deadline=None) +@hypothesis.example( + carry_over=True, + prior=[("via_branch", DATASET_YEAR)], + later=[("via_branch", DATASET_YEAR), ("via_branch", "2023")], + n=4, + seed="a", +) @hypothesis.given( carry_over=st.booleans(), prior=st.lists(REQUESTS, min_size=1, max_size=6), diff --git a/tests/fixtures/subsample_inputs.py b/tests/fixtures/subsample_inputs.py index 3a5161608..9d6f1ed41 100644 --- a/tests/fixtures/subsample_inputs.py +++ b/tests/fixtures/subsample_inputs.py @@ -68,7 +68,25 @@ def formula(person, period): return person.filled_array(-1.0) -FORMULA_VARIABLES = ["doubled", "ended", "household_total", "person_weight"] +class via_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "A formula that changes an input in its own branch" + + def formula(person, period): + branch = person.simulation.get_branch("probe") + branch.set_input("base", period, np.full(person.count, 1.0)) + return branch.calculate("doubled", period) + + +FORMULA_VARIABLES = [ + "doubled", + "ended", + "household_total", + "person_weight", + "via_branch", +] OVERRIDDEN = [float(index) for index in range(len(BASE))] @@ -91,7 +109,7 @@ def build(carry_over: bool = True) -> Simulation: """A ten-person, six-household simulation loaded from ``dataframe()``.""" system = CountryTaxBenefitSystem() system.auto_carry_over_input_variables = carry_over - system.add_variables(base, doubled, ended, household_total, overridden) + system.add_variables(base, doubled, ended, household_total, overridden, via_branch) return Simulation( tax_benefit_system=system, dataset=Dataset.from_dataframe(dataframe(), DATASET_YEAR), From 01eb84414aa04dc67a203cdf1e4b903eaa127506 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:03:28 -0400 Subject: [PATCH 03/13] Register restored source inputs before subsampling Keep restored input provenance in the input export registry while excluding values listed in derived_periods.txt. Add a differential restore and subsample regression without relying on PR #560. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/tools/simulation_dumper.py | 7 +- tests/core/test_subsample_restore_inputs.py | 106 +++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 tests/core/test_subsample_restore_inputs.py diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index 65b3c6644..24a4d5b5b 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -156,4 +156,9 @@ def _restore_holder(simulation, variable, directory): for period in disk_storage.get_known_periods(): value = disk_storage.get(period) - holder.put_in_cache(value, period, derived=str(period) in derived_periods) + derived = str(period) in derived_periods + holder.put_in_cache(value, period, derived=derived) + if not derived: + # Restored source values must also be available to input-only + # exports, including explicit values of formula-backed variables. + simulation._user_input_keys.add((variable, simulation.branch_name, period)) diff --git a/tests/core/test_subsample_restore_inputs.py b/tests/core/test_subsample_restore_inputs.py new file mode 100644 index 000000000..77fec9f8e --- /dev/null +++ b/tests/core/test_subsample_restore_inputs.py @@ -0,0 +1,106 @@ +"""Dump/restore preserves the source inputs used to rebuild a subsample. + +For the same inputs and seed, a restored simulation and one never dumped +produce the same subsample regardless of prior calculations. Cached derived +values stay derived after restore and do not become inputs to the subsample; +explicit values of formula-backed variables remain inputs. +""" + +import numpy as np +import pytest + +from policyengine_core.periods import period +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, stored + +SEED = "restored-inputs" + + +def build_with_later_input(carry_over): + simulation = build(carry_over) + simulation.set_input("overridden", "2023", np.full(simulation.persons.count, 7.0)) + return simulation + + +@pytest.mark.parametrize("carry_over", [False, True]) +@pytest.mark.parametrize("calculate_before", [False, True]) +@pytest.mark.parametrize("n", [1, 4]) +def test_restored_subsample_matches_never_dumped_inputs( + tmp_path, carry_over, calculate_before, n +): + fresh = build_with_later_input(carry_over) + source = build_with_later_input(carry_over) + if calculate_before: + for variable, requested_period in [ + ("doubled", DATASET_YEAR), + ("household_total", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("ended", DATASET_YEAR), + ("base", "2024"), + ("overridden", "2024"), + ]: + source.calculate(variable, requested_period) + + directory = str(tmp_path / "dump") + dump_simulation(source, directory) + restored = restore_simulation(directory, source.tax_benefit_system) + # The dump stores populations and holder values, not dataset metadata. + restored.dataset = source.dataset + restored.default_calculation_period = source.default_calculation_period + + for variable in ["base", "overridden", "household_weight"]: + assert (variable, "default", period(DATASET_YEAR)) in restored._user_input_keys + assert ("overridden", "default", period("2023")) in restored._user_input_keys + if calculate_before: + for variable in ["doubled", "household_total", "person_weight", "ended"]: + holder = restored.get_holder(variable) + assert holder.get_array(period(DATASET_YEAR)) is not None + assert holder.is_derived(period(DATASET_YEAR)) + assert ( + variable, + "default", + period(DATASET_YEAR), + ) not in restored._user_input_keys + assert ("base", "default", period("2024")) not in restored._user_input_keys + assert ( + "overridden", + "default", + period("2024"), + ) not in restored._user_input_keys + + fresh.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + restored.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + assert stored(restored) == stored(fresh) + assert not any( + variable in ["doubled", "household_total", "person_weight", "ended"] + for variable, _, _ in stored(restored) + ) + for variable, requested_period in [ + ("base", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("overridden", DATASET_YEAR), + ("overridden", "2023"), + ("ended", "2023"), + ]: + np.testing.assert_array_equal( + restored.calculate(variable, requested_period), + fresh.calculate(variable, requested_period), + ) + + # The old calculation cannot override the formula after new source data. + for simulation in [fresh, restored]: + simulation.set_input( + "base", DATASET_YEAR, np.full(simulation.persons.count, 5.0) + ) + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 10.0), + ) + np.testing.assert_array_equal( + restored.calculate("household_total", DATASET_YEAR), + fresh.calculate("household_total", DATASET_YEAR), + ) + assert stored(restored) == stored(fresh) From f26f86281079212b06fed4de97a6f559e1d0917e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:07:58 -0400 Subject: [PATCH 04/13] Use requested-period weights only for subsample selection Calculate sampling probabilities from the requested period independently of the input export. Keep existing input-column normalization and recompute derived weights on the sample. Add focused differential and property coverage. Co-Authored-By: Claude Opus 5.5 --- .../fix-subsample-inputs-only.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 24 +++-- tests/core/test_subsample_sampling_period.py | 100 ++++++++++++++++++ ...test_subsample_sampling_period_property.py | 24 +++++ 4 files changed, 140 insertions(+), 10 deletions(-) create mode 100644 tests/core/test_subsample_sampling_period.py create mode 100644 tests/core/test_subsample_sampling_period_property.py diff --git a/changelog.d/fix-subsample-inputs-only.fixed.md b/changelog.d/fix-subsample-inputs-only.fixed.md index b9654be46..b62ecfb0b 100644 --- a/changelog.d/fix-subsample-inputs-only.fixed.md +++ b/changelog.d/fix-subsample-inputs-only.fixed.md @@ -1 +1 @@ -Rebuild a simulation in `subsample` from the values it was given (the dataset and `set_input`), leaving calculated values out, so results after subsampling no longer depend on what was calculated before it. +Rebuild subsamples from source inputs, preserve restored inputs, recreate calculation branches and select households using requested-period weights without freezing calculated values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 19a39fd7f..ab341ec05 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2272,6 +2272,12 @@ def subsample( ) -> "Simulation": """Quantize the simulation to a smaller size by sampling households. + Rebuild from recorded inputs, including inputs of formula-backed + variables. Calculate household weights for ``time_period`` to choose + the sample, then rescale the input weight columns. Calculated weights + and other formula results are recomputed on the sample. Calculation + branches are recreated on its populations; baseline policy is kept. + Args: n (int, optional): The number of households to sample. Defaults to 10_000. frac (float, optional): The fraction of households to sample. Defaults to None. @@ -2308,20 +2314,20 @@ def subsample( df_household_id_column = f"household_id__{df_time_period}" df_person_id_column = f"person_id__{df_time_period}" - # Determine the appropriate household weight column - if f"household_weight__{time_period}" in df.columns: - household_weight_column = f"household_weight__{time_period}" - else: - household_weight_column = f"household_weight__{df_time_period}" - - # Group by household ID and get the first entry for each group - h_df = df.groupby(df_household_id_column).first() + # Sampling needs the requested year's weight even when it was + # calculated rather than supplied as an input. Use it for selection + # without adding it to the data that will be reloaded as inputs. + sampling_df = df[[df_household_id_column]].copy() + sampling_df["household_weight"] = np.asarray( + self.calculate("household_weight", time_period, map_to="person") + ) + h_df = sampling_df.groupby(df_household_id_column).first() h_ids = pd.Series(h_df.index) if n is None and frac is None: raise ValueError("Either n or frac must be provided.") if n is None: n = int(len(h_ids) * frac) - h_weights = pd.Series(h_df[household_weight_column].values) + h_weights = pd.Series(h_df["household_weight"].values) frac = n / len(h_ids) diff --git a/tests/core/test_subsample_sampling_period.py b/tests/core/test_subsample_sampling_period.py new file mode 100644 index 000000000..9eda3da6a --- /dev/null +++ b/tests/core/test_subsample_sampling_period.py @@ -0,0 +1,100 @@ +"""Requested-year weights select the sample without becoming source inputs.""" + +from unittest.mock import patch + +import numpy as np +import pytest + +from policyengine_core.country_template import Microsimulation, Simulation +from policyengine_core.country_template.entities import Household +from policyengine_core.data import Dataset +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +class household_weight(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Weights growing at a different rate in each household" + + def formula(household, period): + growth = household("household_id", period) % 3 + 1 + return household("household_weight", period.last_year) * growth + + +def make_simulation(simulation_class): + system = build().tax_benefit_system + system.update_variable(household_weight) + data = dataframe() + # Microsimulation's export calculates weights as a side effect of + # reading this input, reproducing the s17 variable-order case. + data["base__2023"] = data[f"base__{DATASET_YEAR}"] + 1 + return simulation_class( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + + +def check_sampling_period(simulation_class, year, n, seed, quantize_weights): + reference = make_simulation(simulation_class) + expected_ids = np.asarray(reference.calculate("household_id", DATASET_YEAR)) + expected_weights = np.asarray(reference.calculate("household_weight", year)) + expected_probabilities = expected_weights / expected_weights.sum() + fresh = make_simulation(simulation_class) + used = make_simulation(simulation_class) + used.calculate("household_weight", year) + used.calculate("person_weight", year) + + for simulation in (fresh, used): + original_weight_total = simulation.dataset.load()[ + f"household_weight__{DATASET_YEAR}" + ].sum() + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample( + n=n, seed=seed, time_period=year, quantize_weights=quantize_weights + ) + args, kwargs = choose.call_args + np.testing.assert_array_equal(args[0], expected_ids) + if quantize_weights: + np.testing.assert_allclose(kwargs["p"], expected_probabilities) + else: + assert kwargs["p"] is None + data = simulation.dataset.load() + assert data[f"household_weight__{DATASET_YEAR}"].sum() == pytest.approx( + original_weight_total + ) + if year != DATASET_YEAR: + assert f"household_weight__{year}" not in data + assert not any(name == "person_weight" for name, _, _ in stored(simulation)) + assert simulation.default_calculation_period == DATASET_YEAR + + assert stored(used) == stored(fresh) + for variable in ("household_weight", "person_weight"): + np.testing.assert_array_equal( + used.calculate(variable, year), fresh.calculate(variable, year) + ) + if year != DATASET_YEAR: + assert used.get_holder(variable).is_derived(year) + + +@pytest.mark.parametrize("simulation_class", [Simulation, Microsimulation]) +@pytest.mark.parametrize("year", [DATASET_YEAR, "2023", "2024"]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_sampling_uses_requested_year_weights_independently_of_calculations( + simulation_class, year, quantize_weights +): + check_sampling_period(simulation_class, year, 6, "probe", quantize_weights) + + +def test_sampling_period_is_preserved_after_repeated_subsamples(): + simulation = make_simulation(Simulation) + for n in (6, 3): + simulation.subsample(n=n, seed="again", time_period="2024") + assert simulation.default_calculation_period == DATASET_YEAR + assert "household_weight__2024" not in simulation.dataset.load() + assert not any( + name == "household_weight" and period == "2024" + for name, _, period in stored(simulation) + ) diff --git a/tests/core/test_subsample_sampling_period_property.py b/tests/core/test_subsample_sampling_period_property.py new file mode 100644 index 000000000..cc009b190 --- /dev/null +++ b/tests/core/test_subsample_sampling_period_property.py @@ -0,0 +1,24 @@ +"""Differential sampling coverage across requested years, seeds and sizes.""" + +import pytest + +from policyengine_core.country_template import Microsimulation, Simulation +from tests.core.test_subsample_sampling_period import check_sampling_period +from tests.fixtures.subsample_inputs import DATASET_YEAR + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + + +@hypothesis.settings(max_examples=35, deadline=None) +@hypothesis.given( + simulation_class=st.sampled_from([Simulation, Microsimulation]), + year=st.sampled_from([DATASET_YEAR, "2023", "2024"]), + n=st.integers(1, 8), + seed=st.sampled_from(["a", "b", "c"]), + quantize_weights=st.booleans(), +) +def test_requested_year_sampling_matches_direct_weights( + simulation_class, year, n, seed, quantize_weights +): + check_sampling_period(simulation_class, year, n, seed, quantize_weights) From ea1de0d203b3ff7e9293360e44fd8eb54b0e776f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 10:30:33 -0400 Subject: [PATCH 05/13] Sample missing input household IDs by recorded memberships Use the flat-file loader\x27s recorded membership labels when an input-only household ID is absent. Keep calculated default IDs out of the rebuilt dataset and verify complete household partitions with shuffled, nonconsecutive labels. Co-Authored-By: Claude Opus 5.5 --- .../fix-subsample-inputs-only.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 23 +++-- tests/core/test_subsample_membership_ids.py | 95 +++++++++++++++++++ 3 files changed, 111 insertions(+), 9 deletions(-) create mode 100644 tests/core/test_subsample_membership_ids.py diff --git a/changelog.d/fix-subsample-inputs-only.fixed.md b/changelog.d/fix-subsample-inputs-only.fixed.md index b62ecfb0b..971637318 100644 --- a/changelog.d/fix-subsample-inputs-only.fixed.md +++ b/changelog.d/fix-subsample-inputs-only.fixed.md @@ -1 +1 @@ -Rebuild subsamples from source inputs, preserve restored inputs, recreate calculation branches and select households using requested-period weights without freezing calculated values. +Rebuild subsamples from source inputs, preserve restored inputs and recorded household memberships, recreate calculation branches and select households using requested-period weights without freezing calculated values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ab341ec05..310a8717f 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2277,6 +2277,8 @@ def subsample( the sample, then rescale the input weight columns. Calculated weights and other formula results are recomputed on the sample. Calculation branches are recreated on its populations; baseline policy is kept. + If the flat-file data omits an input-only household ID, recorded + person-to-household IDs retain the partition used by the loader. Args: n (int, optional): The number of households to sample. Defaults to 10_000. @@ -2305,14 +2307,19 @@ def subsample( # be calculated before subsampling. df = self._to_person_dataframe(self._get_set_input_periods) - # Extract time period from DataFrame columns - df_time_period = ( - df.columns[df.columns.str.contains("household_id__")] - .values[0] - .split("__")[1] - ) - df_household_id_column = f"household_id__{df_time_period}" - df_person_id_column = f"person_id__{df_time_period}" + # Flat-file loading can infer households from recorded membership + # IDs when the input-only household ID is absent. Use those same + # labels for selection without materializing a default ID input. + household_id_columns = [ + column for column in df if column.startswith("household_id__") + ] + if not household_id_columns and self._is_exportable_input_variable( + "household_id" + ): + household_id_columns = [ + column for column in df if column.startswith("person_household_id__") + ] + df_household_id_column = household_id_columns[0] # Sampling needs the requested year's weight even when it was # calculated rather than supplied as an input. Use it for selection diff --git a/tests/core/test_subsample_membership_ids.py b/tests/core/test_subsample_membership_ids.py new file mode 100644 index 000000000..9c434f456 --- /dev/null +++ b/tests/core/test_subsample_membership_ids.py @@ -0,0 +1,95 @@ +"""Subsampling preserves a flat-file partition recorded by membership IDs.""" + +from unittest.mock import patch + +import numpy as np +import pytest + +from policyengine_core.country_template import Simulation +from policyengine_core.data import Dataset +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +def partition(simulation): + """Person IDs grouped by the population's actual household membership.""" + members = simulation.household.members_entity_id + return { + frozenset(simulation.persons.ids[members == index].tolist()) + for index in range(simulation.household.count) + } + + +@pytest.mark.parametrize("nonconsecutive_ids", [False, True]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("n", [1, 4]) +def test_missing_input_household_id_uses_recorded_membership( + nonconsecutive_ids, quantize_weights, n +): + data = dataframe().drop(columns=[f"household_id__{DATASET_YEAR}"]) + membership_column = f"person_household_id__{DATASET_YEAR}" + if nonconsecutive_ids: + data[membership_column] = data[membership_column].map( + {1: 901, 2: -20, 3: 33, 4: 88, 5: 205, 6: 411} + ) + data = data.iloc[[9, 0, 4, 6, 2, 8, 1, 5, 3, 7]].reset_index(drop=True) + + def make_simulation(): + return Simulation( + tax_benefit_system=build().tax_benefit_system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + + fresh, used = make_simulation(), make_simulation() + source_partition = partition(fresh) + assert fresh.household.count == 6 + used.calculate("household_id", DATASET_YEAR) + + choice = np.random.choice + for simulation in (fresh, used): + selected = {} + + def choose_households(*args, **kwargs): + selected["ids"] = choice(*args, **kwargs) + return selected["ids"] + + with patch("numpy.random.choice", side_effect=choose_households): + simulation.subsample( + n=n, + seed="recorded-memberships", + time_period=DATASET_YEAR, + quantize_weights=quantize_weights, + ) + # The resulting inputs select whole recorded groups, without adding + # the calculated default household ID. + sampled = simulation.dataset.load() + selected_ids = set(selected["ids"]) + assert set(sampled["person_household_id__ETERNITY"]) == selected_ids + expected_people = set( + data.loc[ + data[membership_column].isin(selected_ids), f"person_id__{DATASET_YEAR}" + ] + ) + assert set(simulation.persons.ids) == expected_people + assert simulation.household.count == len(selected_ids) + assert 1 <= len(selected_ids) <= n + assert partition(simulation) <= source_partition + assert not any(column.startswith("household_id__") for column in sampled) + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + assert simulation.default_calculation_period == DATASET_YEAR + expected_totals = 2 * data.groupby(membership_column)[ + f"base__{DATASET_YEAR}" + ].transform("sum") + totals_by_person = dict( + zip(data[f"person_id__{DATASET_YEAR}"], expected_totals) + ) + np.testing.assert_array_equal( + simulation.calculate("household_total", DATASET_YEAR, map_to="person"), + [totals_by_person[person_id] for person_id in simulation.persons.ids], + ) + + assert stored(used) == stored(fresh) + assert partition(used) == partition(fresh) + np.testing.assert_array_equal( + used.calculate("household_weight", DATASET_YEAR), + fresh.calculate("household_weight", DATASET_YEAR), + ) From dbc7a72d93f79912ca244b896063c9826cf3ca42 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:12:56 -0400 Subject: [PATCH 06/13] Pin monthly sampling to existing annual weight conversion Extend differential sampling examples and properties to monthly requests. Verify annual FLOW weights divide by twelve while normalized probabilities, source inputs and the default calculation period stay consistent. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_subsample_sampling_period.py | 14 +++++++++++--- .../test_subsample_sampling_period_property.py | 6 +++--- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/core/test_subsample_sampling_period.py b/tests/core/test_subsample_sampling_period.py index 9eda3da6a..e6e161972 100644 --- a/tests/core/test_subsample_sampling_period.py +++ b/tests/core/test_subsample_sampling_period.py @@ -1,4 +1,4 @@ -"""Requested-year weights select the sample without becoming source inputs.""" +"""Requested-period weights select the sample without becoming source inputs.""" from unittest.mock import patch @@ -80,9 +80,9 @@ def check_sampling_period(simulation_class, year, n, seed, quantize_weights): @pytest.mark.parametrize("simulation_class", [Simulation, Microsimulation]) -@pytest.mark.parametrize("year", [DATASET_YEAR, "2023", "2024"]) +@pytest.mark.parametrize("year", [DATASET_YEAR, "2023", "2024", "2022-07", "2023-07"]) @pytest.mark.parametrize("quantize_weights", [False, True]) -def test_sampling_uses_requested_year_weights_independently_of_calculations( +def test_sampling_uses_requested_period_weights_independently_of_calculations( simulation_class, year, quantize_weights ): check_sampling_period(simulation_class, year, 6, "probe", quantize_weights) @@ -98,3 +98,11 @@ def test_sampling_period_is_preserved_after_repeated_subsamples(): name == "household_weight" and period == "2024" for name, _, period in stored(simulation) ) + + +def test_monthly_sampling_uses_existing_annual_flow_conversion(): + simulation = make_simulation(Simulation) + annual = np.asarray(simulation.calculate("household_weight", "2023")) + monthly = np.asarray(simulation.calculate("household_weight", "2023-07")) + np.testing.assert_allclose(monthly, annual / 12) + np.testing.assert_allclose(monthly / monthly.sum(), annual / annual.sum()) diff --git a/tests/core/test_subsample_sampling_period_property.py b/tests/core/test_subsample_sampling_period_property.py index cc009b190..21489b91d 100644 --- a/tests/core/test_subsample_sampling_period_property.py +++ b/tests/core/test_subsample_sampling_period_property.py @@ -1,4 +1,4 @@ -"""Differential sampling coverage across requested years, seeds and sizes.""" +"""Differential sampling coverage across requested periods, seeds and sizes.""" import pytest @@ -13,12 +13,12 @@ @hypothesis.settings(max_examples=35, deadline=None) @hypothesis.given( simulation_class=st.sampled_from([Simulation, Microsimulation]), - year=st.sampled_from([DATASET_YEAR, "2023", "2024"]), + year=st.sampled_from([DATASET_YEAR, "2023", "2024", "2022-07", "2023-07"]), n=st.integers(1, 8), seed=st.sampled_from(["a", "b", "c"]), quantize_weights=st.booleans(), ) -def test_requested_year_sampling_matches_direct_weights( +def test_requested_period_sampling_matches_direct_weights( simulation_class, year, n, seed, quantize_weights ): check_sampling_period(simulation_class, year, n, seed, quantize_weights) From 04772070f29a78bdfcb804b78261c0c8d3095a46 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:49:41 -0400 Subject: [PATCH 07/13] Preserve subsample input provenance, structural leaves and entity weight totals --- policyengine_core/simulations/simulation.py | 141 ++++++++++-- tests/core/test_subsample_deleted_inputs.py | 148 +++++++++++++ .../test_subsample_entity_weight_totals.py | 179 ++++++++++++++++ tests/core/test_subsample_sampling_period.py | 12 +- tests/core/test_subsample_structural_ids.py | 200 ++++++++++++++++++ tests/core/test_subsample_weight_fallback.py | 107 ++++++++++ 6 files changed, 762 insertions(+), 25 deletions(-) create mode 100644 tests/core/test_subsample_deleted_inputs.py create mode 100644 tests/core/test_subsample_entity_weight_totals.py create mode 100644 tests/core/test_subsample_structural_ids.py create mode 100644 tests/core/test_subsample_weight_fallback.py diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 310a8717f..9376ff3d7 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -605,12 +605,12 @@ def get_eternity_array(name): person_membership_ids = get_eternity_array(person_membership_id_field) person_role_field = f"{person_entity.key}_{group_entity.key}_role" - if person_role_field in data: + if any(column.split("__")[0] == person_role_field for column in data): person_roles = get_eternity_array(person_role_field) elif "role" in data: person_roles = get_eternity_array("role") elif self.default_role is not None: - person_roles = np.full(len(entity_ids), self.default_role) + person_roles = np.full(self.persons.count, self.default_role) else: raise ValueError( f"Missing {person_role_field} column in the dataset. Each group entity must have a person role array defined for ETERNITY." @@ -2173,7 +2173,8 @@ def _get_set_input_periods(self, variable_name: str) -> List[Period]: That is every value loaded from the dataset or passed to ``set_input`` on a branch this simulation reads, including for a variable that has - a formula. Values the simulation calculated are left out. + a formula. Values the simulation calculated are left out, including + values calculated after a previously recorded input was deleted. """ user_input_periods = { period @@ -2187,10 +2188,10 @@ def _get_set_input_periods(self, variable_name: str) -> List[Period]: return [] variable = self.tax_benefit_system.get_variable(variable_name) holder = self.get_holder(variable_name) + input_periods = set(holder.get_input_periods(self.branch_name)) if variable.definition_period == ETERNITY: - return holder.get_known_periods() - known_periods = set(holder.get_known_periods()) - return sorted(user_input_periods & known_periods, key=str) + return sorted(input_periods, key=str) + return sorted(user_input_periods & input_periods, key=str) def to_input_dataframe( self, @@ -2262,6 +2263,54 @@ def to_input_dict(self, include_computed_variables: bool = False) -> dict: return data + def _subsample_weight_period(self, time_period) -> Period: + """Use dataset input weights if the requested period has no source. + + A calculated default is not evidence of available sampling weights. + Check the input/formula paths, rather than comparing values with the + default, so a formula intentionally returning uniform weights is kept. + """ + requested = periods.period(time_period) + variable = self.tax_benefit_system.get_variable("household_weight") + calculation_period = ( + requested.this_year + if variable.definition_period == YEAR and requested.unit == MONTH + else requested + ) + input_period = ( + periods.period(ETERNITY) + if variable.definition_period == ETERNITY + else calculation_period + ) + inputs = self.get_holder("household_weight").get_input_periods(self.branch_name) + has_formula = ( + variable.get_formula(calculation_period) is not None + or variable.adds + or variable.subtracts + or variable.is_neutralized + ) + can_uprate = variable.uprating is not None and any( + source.unit == variable.definition_period + and source.start < calculation_period.start + for source in inputs + ) + can_carry = ( + self.tax_benefit_system.auto_carry_over_input_variables + and variable.calculate_output is None + and any(source.start <= calculation_period.start for source in inputs) + ) + if input_period in inputs or has_formula or can_uprate or can_carry: + return requested + dataset_period = periods.period(self.dataset.time_period).start.period( + variable.definition_period + ) + if dataset_period in inputs: + return dataset_period + raise ValueError( + f"No household_weight source is available for {requested}, " + f"and no input weights exist for dataset period {dataset_period}." + ) + def subsample( self, n=None, @@ -2274,11 +2323,14 @@ def subsample( Rebuild from recorded inputs, including inputs of formula-backed variables. Calculate household weights for ``time_period`` to choose - the sample, then rescale the input weight columns. Calculated weights + the sample, falling back to dataset-period input weights when that + period has no input, formula, uprating or carry-over source. Rescale + each input weight column to preserve its own entity's total. Calculated weights and other formula results are recomputed on the sample. Calculation branches are recreated on its populations; baseline policy is kept. - If the flat-file data omits an input-only household ID, recorded - person-to-household IDs retain the partition used by the loader. + If structural inputs are absent, loaded population IDs, memberships + and roles supply the leaves needed to rebuild the same partition. + This also supports structural variables whose IDs have formulas. Args: n (int, optional): The number of households to sample. Defaults to 10_000. @@ -2307,15 +2359,46 @@ def subsample( # be calculated before subsampling. df = self._to_person_dataframe(self._get_set_input_periods) - # Flat-file loading can infer households from recorded membership - # IDs when the input-only household ID is absent. Use those same - # labels for selection without materializing a default ID input. + def retain_structure(name, values): + variable = self.tax_benefit_system.get_variable(name) + if variable is None or any(column.split("__")[0] == name for column in df): + return + structural_period = ( + ETERNITY + if variable.definition_period == ETERNITY + else periods.period(self.dataset.time_period).start.period( + variable.definition_period + ) + ) + df[f"{name}__{structural_period}"] = values + + # These are loaded population leaves, not calculated ID formulas. + # A simulation built from populations may have no stored ID inputs. + person_key = self.persons.entity.key + retain_structure(f"{person_key}_id", self.persons.ids) + for entity in self.tax_benefit_system.group_entities: + population = self.populations[entity.key] + retain_structure( + f"{person_key}_{entity.key}_id", + population.ids[population.members_entity_id], + ) + retain_structure( + f"{person_key}_{entity.key}_role", + [ + getattr(role, "key", "") + for role in np.broadcast_to( + population.members_role, self.persons.count + ) + ], + ) + + # Flat-file loading can infer households from membership IDs even + # when the household ID variable has a formula. Use the same leaves + # for selection without materializing a calculated ID input. household_id_columns = [ column for column in df if column.startswith("household_id__") ] - if not household_id_columns and self._is_exportable_input_variable( - "household_id" - ): + if not household_id_columns: household_id_columns = [ column for column in df if column.startswith("person_household_id__") ] @@ -2326,7 +2409,11 @@ def subsample( # without adding it to the data that will be reloaded as inputs. sampling_df = df[[df_household_id_column]].copy() sampling_df["household_weight"] = np.asarray( - self.calculate("household_weight", time_period, map_to="person") + self.calculate( + "household_weight", + self._subsample_weight_period(time_period), + map_to="person", + ) ) h_df = sampling_df.groupby(df_household_id_column).first() h_ids = pd.Series(h_df.index) @@ -2372,16 +2459,32 @@ def subsample( lambda x: household_id_to_count.get(x, 0) ) - # Adjust household weights to maintain the total weight + # Conserve each source column's native entity total. Group weights + # repeat on person rows, so summing those rows counts larger groups + # more often. The loader reads each retained group's first member; + # that member can differ from its original first member when another + # group partition crosses the sampled households. for col in subset_df.columns: if "weight__" in col: - target_total_weight = df[col].values.sum() + variable_name, weight_period = col.split("__") + entity = self.tax_benefit_system.get_variable(variable_name).entity + target_total_weight = np.asarray( + self.calculate(variable_name, weight_period) + ).sum() if not quantize_weights: subset_df[col] *= household_counts.values else: subset_df[col] = household_counts.values - subset_df[col] *= target_total_weight / subset_df[col].values.sum() + if entity.is_person: + selected_total = subset_df[col].values.sum() + else: + memberships = self.populations[entity.key].members_entity_id[ + subset_df.index + ] + _, first_retained = np.unique(memberships, return_index=True) + selected_total = subset_df[col].values[first_retained].sum() + subset_df[col] *= target_total_weight / selected_total df = subset_df diff --git a/tests/core/test_subsample_deleted_inputs.py b/tests/core/test_subsample_deleted_inputs.py new file mode 100644 index 000000000..6e524ec7f --- /dev/null +++ b/tests/core/test_subsample_deleted_inputs.py @@ -0,0 +1,148 @@ +"""Deleted source inputs must not turn later calculations into source data. + +Subsampling after deleting an override must match subsampling without reading +the deleted variable first. This holds for branches and both storage backends; +new overrides are still preserved. +""" + +import numpy as np +import pytest + +from policyengine_core.country_template.entities import Person +from policyengine_core.experimental import MemoryConfig +from policyengine_core.periods import ETERNITY, YEAR, period +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, stored + + +class deleted_leaf(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "An input-only variable whose deleted value becomes a default" + default_value = 17.0 + + +class deleted_eternal(Variable): + value_type = float + entity = Person + definition_period = ETERNITY + label = "An eternal formula variable with a deleted override" + + def formula(person, period): + return person("person_id", period) * 2 + + +def prepare(variable, on_disk=False, branch=False): + simulation = build() + simulation.tax_benefit_system.add_variables(deleted_leaf, deleted_eternal) + holder = simulation.get_holder(variable) + if on_disk: + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input(variable, DATASET_YEAR, np.full(simulation.persons.count, 99)) + if branch: + simulation = simulation.get_branch("child") + return simulation + + +def delete_override(simulation, variable, through_holder): + if through_holder: + holder = simulation.get_holder(variable) + for branch_name in simulation._get_visible_branch_names(): + holder.delete_arrays(DATASET_YEAR, branch_name) + else: + simulation.delete_arrays(variable, DATASET_YEAR) + + +@pytest.mark.parametrize("variable", ["doubled", "deleted_leaf", "deleted_eternal"]) +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("through_holder", [False, True], ids=["simulation", "holder"]) +def test_deleted_input_recalculated_as_derived_is_not_exportable( + variable, on_disk, through_holder +): + simulation = prepare(variable, on_disk) + delete_override(simulation, variable, through_holder) + simulation.calculate(variable, DATASET_YEAR) + + assert simulation.get_holder(variable).is_derived( + period(DATASET_YEAR), simulation.branch_name + ) + assert simulation._get_set_input_periods(variable) == [] + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("branch", [False, True], ids=["default", "child"]) +def test_deleted_formula_override_does_not_freeze_after_subsample(on_disk, branch): + simulation = prepare("doubled", on_disk, branch) + simulation.delete_arrays("doubled", DATASET_YEAR) + simulation.calculate("doubled", DATASET_YEAR) + + simulation.subsample(n=3, seed="deleted-input", time_period=DATASET_YEAR) + simulation.set_input("base", DATASET_YEAR, np.full(simulation.persons.count, 5)) + + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 10), + ) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("branch", [False, True], ids=["default", "child"]) +def test_new_override_after_deletion_is_still_exportable(on_disk, branch): + simulation = prepare("doubled", on_disk, branch) + simulation.delete_arrays("doubled", DATASET_YEAR) + simulation.calculate("doubled", DATASET_YEAR) + simulation.set_input("doubled", DATASET_YEAR, np.full(simulation.persons.count, 7)) + + assert simulation._get_set_input_periods("doubled") == [period(DATASET_YEAR)] + simulation.subsample(n=3, seed="deleted-input", time_period=DATASET_YEAR) + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 7), + ) + + +def test_reading_deleted_overrides_does_not_change_subsample(): + """Property: storage and all later formulas ignore pre-sample reads. + + Vary the deleted override, storage, branch, deletion API, sample size, + and subsequent leaf input. A never-read twin is the differential oracle. + """ + hypothesis = pytest.importorskip("hypothesis") + st = hypothesis.strategies + + @hypothesis.settings(max_examples=40, deadline=None) + @hypothesis.given( + override=st.integers(-100, 100), + on_disk=st.booleans(), + branch=st.booleans(), + through_holder=st.booleans(), + n=st.integers(1, 6), + later_input=st.integers(-100, 100), + ) + def check(override, on_disk, branch, through_holder, n, later_input): + fresh = prepare("doubled", on_disk, branch) + used = prepare("doubled", on_disk, branch) + for simulation in (fresh, used): + simulation.set_input( + "doubled", DATASET_YEAR, np.full(simulation.persons.count, override) + ) + delete_override(simulation, "doubled", through_holder) + used.calculate("doubled", DATASET_YEAR) + for simulation in (fresh, used): + simulation.subsample(n=n, seed="deleted-input", time_period=DATASET_YEAR) + + assert stored(used) == stored(fresh) + for simulation in (fresh, used): + simulation.set_input( + "base", DATASET_YEAR, np.full(simulation.persons.count, later_input) + ) + np.testing.assert_array_equal( + used.calculate("doubled", DATASET_YEAR), + fresh.calculate("doubled", DATASET_YEAR), + ) + assert stored(used) == stored(fresh) + + check() diff --git a/tests/core/test_subsample_entity_weight_totals.py b/tests/core/test_subsample_entity_weight_totals.py new file mode 100644 index 000000000..ecb6c23c5 --- /dev/null +++ b/tests/core/test_subsample_entity_weight_totals.py @@ -0,0 +1,179 @@ +"""Every recorded weight column conserves its own entity's total. + +The flat-file representation repeats group weights per person. Those repeated +values must not determine the total conserved by subsampling. The independent +reference below groups by the first surviving member, matching the loader even +when a second group partition crosses households. +""" + +from collections import Counter +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.data import Dataset +from policyengine_core.entities import build_entity +from policyengine_core.periods import ETERNITY, YEAR +from policyengine_core.simulations import Simulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.variables import Variable + +YEAR_INPUT = "2022" +Person = build_entity(key="person", plural="persons", label="Person", is_person=True) +Household = build_entity( + key="household", + plural="households", + label="Household", + roles=[{"key": "member", "plural": "members", "label": "Members"}], +) +Family = build_entity( + key="family", + plural="families", + label="Family", + roles=[{"key": "member", "plural": "members", "label": "Members"}], +) +ENTITIES = {"person": Person, "household": Household, "family": Family} +WEIGHT_YEARS = (YEAR_INPUT, "2023") + + +def weight_simulation(data): + """Load only leaf inputs into a small three-entity system.""" + system = TaxBenefitSystem(list(ENTITIES.values())) + for name, entity, period in ( + ("person_id", Person, ETERNITY), + ("household_id", Household, ETERNITY), + ("person_household_id", Person, ETERNITY), + ("family_id", Family, ETERNITY), + ("person_family_id", Person, ETERNITY), + ("person_weight", Person, YEAR), + ("household_weight", Household, YEAR), + ("family_weight", Family, YEAR), + ): + system.add_variable( + type( + name, + (Variable,), + { + "value_type": float if period == YEAR else int, + "entity": entity, + "definition_period": period, + "label": name.replace("_", " "), + }, + ) + ) + return Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, YEAR_INPUT), + ) + + +def weight_data( + households, families, household_weights, family_weights, person_weights +): + data = pd.DataFrame( + { + "person_id__ETERNITY": np.arange(len(households)) + 101, + "household_id__ETERNITY": households, + "person_household_id__ETERNITY": households, + "family_id__ETERNITY": families, + "person_family_id__ETERNITY": families, + } + ) + for year_index, year in enumerate(WEIGHT_YEARS): + # Distinct period distributions expose accidental reuse of one column's + # total or scale factor when rescaling all recorded weight inputs. + data[f"household_weight__{year}"] = data["household_id__ETERNITY"].map( + { + group: weight * (1 + year_index * (index % 3 + 1)) + for index, (group, weight) in enumerate( + sorted(household_weights.items()) + ) + } + ) + data[f"family_weight__{year}"] = data["family_id__ETERNITY"].map( + { + group: weight * (1 + year_index * (index % 2 + 1)) + for index, (group, weight) in enumerate(sorted(family_weights.items())) + } + ) + data[f"person_weight__{year}"] = np.asarray(person_weights) * ( + 1 + year_index * (np.arange(len(households)) % 4 + 1) + ) + return data + + +def assert_weight_totals_and_reference(data, chosen, quantize_weights): + simulation = weight_simulation(data) + originals = { + (entity, year): simulation.calculate(f"{entity}_weight", year).sum() + for entity in ENTITIES + for year in WEIGHT_YEARS + } + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="entity-weight-totals", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + retained = data[data["household_id__ETERNITY"].isin(chosen)].copy() + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + for entity in ENTITIES: + # Native entity order is independent of the flat-file person order. + entity_ids = simulation.calculate(f"{entity}_id", ETERNITY) + for year in WEIGHT_YEARS: + column = f"{entity}_weight__{year}" + actual = simulation.calculate(f"{entity}_weight", year) + target = originals[(entity, year)] + assert actual.sum() == pytest.approx(target, rel=2e-6) + + # This reference uses pandas group labels, rather than simulation + # population indices or member counts. A family can lose its first + # source member and retain members with unequal household counts. + retained["candidate_weight"] = ( + counts if quantize_weights else retained[column] * counts + ) + reference = retained.groupby(f"{entity}_id__ETERNITY", sort=True)[ + "candidate_weight" + ].first() + reference *= target / reference.sum() + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("chosen", [[20, 40, 20], [30, 20, 20]]) +def test_subsample_preserves_native_entity_weight_totals(chosen, quantize_weights): + data = weight_data( + households=[30, 10, 20, 20, 30, 40, 40, 20, 40], + families=[110, 110, 220, 110, 330, 330, 220, 220, 440], + household_weights={10: 100.0, 20: 200.0, 30: 300.0, 40: 400.0}, + family_weights={110: 5.0, 220: 11.0, 330: 17.0, 440: 23.0}, + person_weights=np.arange(2.0, 11.0), + ) + assert_weight_totals_and_reference(data, chosen, quantize_weights) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("seed", range(8)) +def test_entity_total_invariant_across_deterministic_random_partitions( + seed, quantize_weights +): + rng = np.random.default_rng(seed) + households = np.repeat([10, 20, 30, 40], rng.integers(1, 5, size=4)) + rng.shuffle(households) + families = rng.choice([110, 220, 330], size=len(households)) + family_ids = np.unique(families) + data = weight_data( + households=households, + families=families, + household_weights=dict(zip([10, 20, 30, 40], rng.integers(1, 100, 4))), + family_weights=dict(zip(family_ids, rng.integers(1, 100, len(family_ids)))), + person_weights=rng.integers(1, 100, len(households)), + ) + chosen = rng.choice([10, 20, 30, 40], size=5, replace=True) + assert_weight_totals_and_reference(data, chosen, quantize_weights) diff --git a/tests/core/test_subsample_sampling_period.py b/tests/core/test_subsample_sampling_period.py index e6e161972..ed71a4605 100644 --- a/tests/core/test_subsample_sampling_period.py +++ b/tests/core/test_subsample_sampling_period.py @@ -48,9 +48,9 @@ def check_sampling_period(simulation_class, year, n, seed, quantize_weights): used.calculate("person_weight", year) for simulation in (fresh, used): - original_weight_total = simulation.dataset.load()[ - f"household_weight__{DATASET_YEAR}" - ].sum() + original_weight_total = np.asarray( + simulation.calculate("household_weight", DATASET_YEAR) + ).sum() with patch("numpy.random.choice", wraps=np.random.choice) as choose: simulation.subsample( n=n, seed=seed, time_period=year, quantize_weights=quantize_weights @@ -62,9 +62,9 @@ def check_sampling_period(simulation_class, year, n, seed, quantize_weights): else: assert kwargs["p"] is None data = simulation.dataset.load() - assert data[f"household_weight__{DATASET_YEAR}"].sum() == pytest.approx( - original_weight_total - ) + assert np.asarray( + simulation.calculate("household_weight", DATASET_YEAR) + ).sum() == pytest.approx(original_weight_total) if year != DATASET_YEAR: assert f"household_weight__{year}" not in data assert not any(name == "person_weight" for name, _, _ in stored(simulation)) diff --git a/tests/core/test_subsample_structural_ids.py b/tests/core/test_subsample_structural_ids.py new file mode 100644 index 000000000..d6ada5250 --- /dev/null +++ b/tests/core/test_subsample_structural_ids.py @@ -0,0 +1,200 @@ +"""Subsampling uses the loaded population when structural IDs have formulas. + +The source partition and roles are leaves: a structural formula's cache must +neither replace them nor decide which people constitute a whole household. +""" + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem, Simulation +from policyengine_core.country_template.entities import Household, Person +from policyengine_core.data import Dataset +from policyengine_core.periods import ETERNITY +from policyengine_core.simulations.simulation_builder import SimulationBuilder +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + HOUSEHOLD_OF_PERSON, + base, + dataframe, + doubled, + household_total, + stored, +) + + +class household_id(Variable): + value_type = int + entity = Household + definition_period = ETERNITY + label = "Formula household identifier" + + def formula(household, period): + # Deliberately unsuitable for deciding population membership. + return household.filled_array(-999) + + +class person_id(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Formula person identifier" + + def formula(person, period): + return np.arange(person.count) + 700 + + +class person_household_id(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Formula person household identifier" + + def formula(person, period): + return person.filled_array(-999) + + +class person_household_role(Variable): + value_type = str + entity = Person + definition_period = ETERNITY + label = "Formula person household role" + + def formula(person, period): + return person.filled_array("child") + + +STRUCTURAL_VARIABLES = ( + household_id, + person_id, + person_household_id, + person_household_role, +) + + +def partition(simulation): + return { + frozenset( + simulation.persons.ids[ + simulation.household.members_entity_id == index + ].tolist() + ) + for index in range(simulation.household.count) + } + + +def roles_by_person(simulation): + return { + person: getattr(role, "key", role) + for person, role in zip( + simulation.persons.ids, simulation.household.members_role + ) + } + + +def make_population_simulation(recorded_structural_inputs, shuffled): + """A population with structural leaves and optional formula overrides.""" + system = CountryTaxBenefitSystem() + system.add_variables(base, doubled, household_total) + for variable in STRUCTURAL_VARIABLES: + system.replace_variable(variable) + + data = dataframe() + labels = {1: 901, 2: -20, 3: 33, 4: 88, 5: 205, 6: 411} + membership = np.array([labels[group] for group in HOUSEHOLD_OF_PERSON]) + # Explicit first/second parent and child roles make role loss observable. + roles = np.array([0, 2, 0, 1, 2, 0, 0, 2, 0, 0]) + order = np.array([9, 0, 4, 6, 2, 8, 1, 5, 3, 7]) if shuffled else np.arange(10) + data = data.iloc[order].reset_index(drop=True) + membership, roles = membership[order], roles[order] + + builder = SimulationBuilder() + builder.create_entities(system) + builder.declare_person_entity("person", data[f"person_id__{DATASET_YEAR}"]) + household = builder.declare_entity("household", list(labels.values())) + builder.join_with_persons(household, membership, roles) + simulation = Simulation( + tax_benefit_system=system, + populations=builder.populations, + default_calculation_period=DATASET_YEAR, + ) + # Attach period metadata without reloading the explicitly built population. + simulation.dataset = Dataset.from_dataframe(data, DATASET_YEAR) + simulation.is_over_dataset = True + simulation.set_input("base", DATASET_YEAR, data[f"base__{DATASET_YEAR}"].values) + simulation.set_input( + "household_weight", DATASET_YEAR, [100, 200, 300, 400, 500, 600] + ) + if recorded_structural_inputs: + simulation.set_input("person_id", ETERNITY, simulation.persons.ids) + simulation.set_input("person_household_id", ETERNITY, membership) + simulation.set_input( + "person_household_role", + ETERNITY, + [role.key for role in simulation.household.members_role], + ) + return simulation + + +def test_formula_only_household_id_uses_recorded_membership(): + """The reviewed failure: no household ID source gives an IndexError.""" + data = dataframe().drop(columns=[f"household_id__{DATASET_YEAR}"]) + system = CountryTaxBenefitSystem() + system.replace_variable(household_id) + system.add_variables(base, doubled, household_total) + simulation = Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + source_partition = partition(simulation) + simulation.calculate("household_id", ETERNITY) + + simulation.subsample(n=4, seed="formula-household-id", time_period=DATASET_YEAR) + + assert partition(simulation) <= source_partition + assert 1 <= simulation.household.count <= 4 + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + + +@pytest.mark.parametrize("recorded_structural_inputs", [False, True]) +@pytest.mark.parametrize("shuffled", [False, True]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("n,seed", [(1, "one"), (4, "four"), (8, "repeats")]) +def test_loaded_structural_leaves_survive_independent_of_formula_caches( + recorded_structural_inputs, shuffled, quantize_weights, n, seed +): + """Differential: formula caches cannot alter IDs, groups, roles or sums.""" + fresh = make_population_simulation(recorded_structural_inputs, shuffled) + used = make_population_simulation(recorded_structural_inputs, shuffled) + source_partition = partition(fresh) + source_roles = roles_by_person(fresh) + source_totals = dict( + zip( + fresh.persons.ids, + fresh.calculate("household_total", DATASET_YEAR, map_to="person"), + ) + ) + for variable in STRUCTURAL_VARIABLES: + used.calculate(variable.__name__, ETERNITY) + + for simulation in (fresh, used): + simulation.subsample( + n=n, + seed=seed, + time_period=DATASET_YEAR, + quantize_weights=quantize_weights, + ) + assert partition(simulation) <= source_partition + assert 1 <= simulation.household.count <= min(n, 6) + assert roles_by_person(simulation) == { + person: source_roles[person] for person in simulation.persons.ids + } + np.testing.assert_array_equal( + simulation.calculate("household_total", DATASET_YEAR, map_to="person"), + [source_totals[person] for person in simulation.persons.ids], + ) + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + + assert partition(used) == partition(fresh) + assert stored(used) == stored(fresh) diff --git a/tests/core/test_subsample_weight_fallback.py b/tests/core/test_subsample_weight_fallback.py new file mode 100644 index 000000000..ee1311080 --- /dev/null +++ b/tests/core/test_subsample_weight_fallback.py @@ -0,0 +1,107 @@ +"""Unavailable sampling periods use recorded weights rather than defaults.""" + +from unittest.mock import patch + +import numpy as np +import pytest +from hypothesis import given, settings, strategies as st + +from policyengine_core.country_template import Microsimulation, Simulation +from policyengine_core.country_template.entities import Household +from policyengine_core.data import Dataset +from policyengine_core.parameters import Parameter +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +class household_weight(Variable): + value_type = float + entity = Household + definition_period = YEAR + default_value = 1 + uprating = "weight_index" + + +def make_simulation(simulation_class, uprated, carry_over): + system = build(carry_over=carry_over).tax_benefit_system + if uprated: + system.update_variable(household_weight) + system.parameters.add_child( + "weight_index", + Parameter( + "weight_index", + data={"values": {"2021-01-01": 1, "2022-01-01": 2, "2023-01-01": 3}}, + ), + ) + return simulation_class( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(dataframe(), DATASET_YEAR), + ) + + +def check_fallback(simulation_class, uprated, carry_over, year, n, seed): + fresh = make_simulation(simulation_class, uprated, carry_over) + used = make_simulation(simulation_class, uprated, carry_over) + expected_weights = np.asarray(fresh.calculate("household_weight", DATASET_YEAR)) + expected_ids = fresh.calculate("household_id", DATASET_YEAR) + used.calculate("household_weight", year) + for simulation in (fresh, used): + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=n, seed=seed, time_period=year) + args, kwargs = choose.call_args + np.testing.assert_array_equal(args[0], expected_ids) + np.testing.assert_allclose( + kwargs["p"], expected_weights / expected_weights.sum() + ) + assert f"household_weight__{year}" not in simulation.dataset.load() + assert simulation.default_calculation_period == DATASET_YEAR + assert stored(fresh) == stored(used) + + +@pytest.mark.parametrize("simulation_class", [Simulation, Microsimulation]) +@pytest.mark.parametrize( + "uprated,carry_over,year", + [ + (False, False, "2023"), + (False, False, "2023-07"), + (False, True, "2021"), + (True, False, "2021"), + (True, True, "2021-07"), + ], +) +def test_missing_period_weights_use_dataset_input( + simulation_class, uprated, carry_over, year +): + check_fallback(simulation_class, uprated, carry_over, year, 3, "fallback") + + +@settings(max_examples=25, deadline=None) +@given( + simulation_class=st.sampled_from([Simulation, Microsimulation]), + case=st.sampled_from( + [(False, False, "2023"), (False, True, "2021"), (True, True, "2021-07")] + ), + n=st.integers(1, 8), + seed=st.integers(0, 100), +) +def test_weight_fallback_is_independent_of_prior_default_calculation( + simulation_class, case, n, seed +): + check_fallback(simulation_class, *case, n, seed) + + +def test_available_requested_input_takes_precedence_over_dataset_weights(): + simulation = make_simulation(Simulation, False, False) + weights = np.arange(1, simulation.household.count + 1, dtype=float)[::-1] + simulation.set_input("household_weight", "2023", weights) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="explicit", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], weights / weights.sum()) + + +def test_absent_weight_source_has_clear_error(): + simulation = make_simulation(Simulation, False, False) + simulation.delete_arrays("household_weight") + with pytest.raises(ValueError, match="household_weight.*2023.*2022"): + simulation.subsample(n=1, seed="missing", time_period="2023") From ba0fb376c6abc7c1152b0280ced5f59ce7861a69 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:50:45 -0400 Subject: [PATCH 08/13] Clarify legacy dump provenance and subsample branch disposal --- changelog.d/fix-subsample-inputs-only.fixed.md | 2 +- policyengine_core/tools/simulation_dumper.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/changelog.d/fix-subsample-inputs-only.fixed.md b/changelog.d/fix-subsample-inputs-only.fixed.md index 971637318..61ade9b4b 100644 --- a/changelog.d/fix-subsample-inputs-only.fixed.md +++ b/changelog.d/fix-subsample-inputs-only.fixed.md @@ -1 +1 @@ -Rebuild subsamples from source inputs, preserve restored inputs and recorded household memberships, recreate calculation branches and select households using requested-period weights without freezing calculated values. +Rebuild subsamples from current source inputs and loaded memberships, preserve each input weight column's entity total, select with available requested-period or dataset-period weights, and discard pre-sample calculation branches, including caller-created branches. diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index 24a4d5b5b..038dd5d19 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -150,6 +150,7 @@ def _restore_holder(simulation, variable, directory): derived_periods_path = os.path.join(storage_dir, DERIVED_PERIODS_FILE) derived_periods = set() + # Legacy dumps have no provenance record; their calculated values may restore as inputs. if os.path.exists(derived_periods_path): with open(derived_periods_path) as file: derived_periods = set(file.read().split()) From ca3faea9c46eadda8b52c4faf4ce22b38743decc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:07:52 -0400 Subject: [PATCH 09/13] Handle period labels, integer roles and zero weight totals when rebuilding samples --- policyengine_core/simulations/simulation.py | 43 ++++++++++++++----- .../test_subsample_entity_weight_totals.py | 18 ++++++++ tests/core/test_subsample_structural_ids.py | 21 +++++++++ tests/core/test_subsample_weight_fallback.py | 35 ++++++++++++++- 4 files changed, 105 insertions(+), 12 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9376ff3d7..1a797b389 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2277,6 +2277,11 @@ def _subsample_weight_period(self, time_period) -> Period: if variable.definition_period == YEAR and requested.unit == MONTH else requested ) + calculation_periods = ( + requested.get_subperiods(MONTH) + if variable.definition_period == MONTH and requested.unit == YEAR + else [calculation_period] + ) input_period = ( periods.period(ETERNITY) if variable.definition_period == ETERNITY @@ -2284,7 +2289,9 @@ def _subsample_weight_period(self, time_period) -> Period: ) inputs = self.get_holder("household_weight").get_input_periods(self.branch_name) has_formula = ( - variable.get_formula(calculation_period) is not None + any( + variable.get_formula(value) is not None for value in calculation_periods + ) or variable.adds or variable.subtracts or variable.is_neutralized @@ -2299,7 +2306,14 @@ def _subsample_weight_period(self, time_period) -> Period: and variable.calculate_output is None and any(source.start <= calculation_period.start for source in inputs) ) - if input_period in inputs or has_formula or can_uprate or can_carry: + if ( + requested in inputs + or input_period in inputs + or any(value in inputs for value in calculation_periods) + or has_formula + or can_uprate + or can_carry + ): return requested dataset_period = periods.period(self.dataset.time_period).start.period( variable.definition_period @@ -2382,15 +2396,19 @@ def retain_structure(name, values): f"{person_key}_{entity.key}_id", population.ids[population.members_entity_id], ) - retain_structure( - f"{person_key}_{entity.key}_role", - [ - getattr(role, "key", "") - for role in np.broadcast_to( - population.members_role, self.persons.count - ) - ], - ) + role_name = f"{person_key}_{entity.key}_role" + role_variable = self.tax_benefit_system.get_variable(role_name) + roles = np.broadcast_to(population.members_role, self.persons.count) + if role_variable is not None and role_variable.value_type == int: + role_values = [ + entity.flattened_roles.index(role) + if role in entity.flattened_roles + else 0 + for role in roles + ] + else: + role_values = [getattr(role, "key", "") for role in roles] + retain_structure(role_name, role_values) # Flat-file loading can infer households from membership IDs even # when the household ID variable has a formula. Use the same leaves @@ -2472,6 +2490,9 @@ def retain_structure(name, values): target_total_weight = np.asarray( self.calculate(variable_name, weight_period) ).sum() + if target_total_weight == 0: + subset_df[col] = 0.0 + continue if not quantize_weights: subset_df[col] *= household_counts.values else: diff --git a/tests/core/test_subsample_entity_weight_totals.py b/tests/core/test_subsample_entity_weight_totals.py index ecb6c23c5..5c35f79c3 100644 --- a/tests/core/test_subsample_entity_weight_totals.py +++ b/tests/core/test_subsample_entity_weight_totals.py @@ -177,3 +177,21 @@ def test_entity_total_invariant_across_deterministic_random_partitions( ) chosen = rng.choice([10, 20, 30, 40], size=5, replace=True) assert_weight_totals_and_reference(data, chosen, quantize_weights) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_zero_total_input_weight_column_remains_zero(quantize_weights): + data = weight_data( + households=[10, 10, 20], + families=[110, 110, 220], + household_weights={10: 100, 20: 200}, + family_weights={110: 0, 220: 0}, + person_weights=[0, 0, 0], + ) + simulation = weight_simulation(data) + simulation.subsample(n=2, seed="zero-column", quantize_weights=quantize_weights) + for entity in ("person", "family"): + for year in WEIGHT_YEARS: + np.testing.assert_array_equal( + simulation.calculate(f"{entity}_weight", year), 0 + ) diff --git a/tests/core/test_subsample_structural_ids.py b/tests/core/test_subsample_structural_ids.py index d6ada5250..2b070a0f5 100644 --- a/tests/core/test_subsample_structural_ids.py +++ b/tests/core/test_subsample_structural_ids.py @@ -198,3 +198,24 @@ def test_loaded_structural_leaves_survive_independent_of_formula_caches( assert partition(used) == partition(fresh) assert stored(used) == stored(fresh) + + +@pytest.mark.parametrize("shuffled", [False, True]) +def test_formula_only_integer_roles_preserve_loaded_subroles(shuffled): + simulation = make_population_simulation(False, shuffled) + source_roles = roles_by_person(simulation) + + class integer_role(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + + def formula(person, period): + return person.filled_array(2) + + integer_role.__name__ = "person_household_role" + simulation.tax_benefit_system.replace_variable(integer_role) + simulation.subsample(n=4, seed="integer-roles", time_period=DATASET_YEAR) + assert roles_by_person(simulation) == { + person: source_roles[person] for person in simulation.persons.ids + } diff --git a/tests/core/test_subsample_weight_fallback.py b/tests/core/test_subsample_weight_fallback.py index ee1311080..806edb7b5 100644 --- a/tests/core/test_subsample_weight_fallback.py +++ b/tests/core/test_subsample_weight_fallback.py @@ -10,7 +10,7 @@ from policyengine_core.country_template.entities import Household from policyengine_core.data import Dataset from policyengine_core.parameters import Parameter -from policyengine_core.periods import YEAR +from policyengine_core.periods import MONTH, YEAR from policyengine_core.variables import Variable from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored @@ -105,3 +105,36 @@ def test_absent_weight_source_has_clear_error(): simulation.delete_arrays("household_weight") with pytest.raises(ValueError, match="household_weight.*2023.*2022"): simulation.subsample(n=1, seed="missing", time_period="2023") + + +def test_annual_weight_fallback_normalizes_dataset_period(): + simulation = make_simulation(Simulation, False, False) + simulation.dataset.time_period = "2022-01" + weights = np.asarray(simulation.calculate("household_weight", DATASET_YEAR)) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="dataset-month", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], weights / weights.sum()) + + +def test_annual_request_uses_monthly_formula_starting_midyear(): + class monthly_weight(Variable): + value_type = float + entity = Household + definition_period = MONTH + + def formula_2023_07(household, period): + return household.filled_array(2.0) + + monthly_weight.__name__ = "household_weight" + system = build(carry_over=False).tax_benefit_system + system.replace_variable(monthly_weight) + data = dataframe().rename( + columns={"household_weight__2022": "household_weight__2022-01"} + ) + simulation = Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="monthly-formula", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], np.full(6, 1 / 6)) From b6069c9fc308d86f7fdb1a39f3de0a58b57599ce Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:13:05 -0400 Subject: [PATCH 10/13] Focus disk provenance coverage and label boundary-case variables --- tests/core/test_subsample_deleted_inputs.py | 13 ++++++++++++- tests/core/test_subsample_structural_ids.py | 1 + tests/core/test_subsample_weight_fallback.py | 1 + 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/core/test_subsample_deleted_inputs.py b/tests/core/test_subsample_deleted_inputs.py index 6e524ec7f..e6b276673 100644 --- a/tests/core/test_subsample_deleted_inputs.py +++ b/tests/core/test_subsample_deleted_inputs.py @@ -38,10 +38,21 @@ def prepare(variable, on_disk=False, branch=False): simulation.tax_benefit_system.add_variables(deleted_leaf, deleted_eternal) holder = simulation.get_holder(variable) if on_disk: - simulation.memory_config = MemoryConfig(max_memory_occupation=0) + # Exercise this variable's provenance on disk, including after the + # rebuild, without spilling unrelated object-valued role arrays. + simulation.memory_config = MemoryConfig( + max_memory_occupation=0, + priority_variables=[ + name + for name in simulation.tax_benefit_system.variables + if name != variable + ], + ) holder._disk_storage = holder.create_disk_storage() holder._on_disk_storable = True simulation.set_input(variable, DATASET_YEAR, np.full(simulation.persons.count, 99)) + if on_disk: + assert holder._disk_storage.has(period(DATASET_YEAR)) if branch: simulation = simulation.get_branch("child") return simulation diff --git a/tests/core/test_subsample_structural_ids.py b/tests/core/test_subsample_structural_ids.py index 2b070a0f5..9216904d8 100644 --- a/tests/core/test_subsample_structural_ids.py +++ b/tests/core/test_subsample_structural_ids.py @@ -209,6 +209,7 @@ class integer_role(Variable): value_type = int entity = Person definition_period = ETERNITY + label = "Formula integer role" def formula(person, period): return person.filled_array(2) diff --git a/tests/core/test_subsample_weight_fallback.py b/tests/core/test_subsample_weight_fallback.py index 806edb7b5..9e7905704 100644 --- a/tests/core/test_subsample_weight_fallback.py +++ b/tests/core/test_subsample_weight_fallback.py @@ -121,6 +121,7 @@ class monthly_weight(Variable): value_type = float entity = Household definition_period = MONTH + label = "Monthly household weight" def formula_2023_07(household, period): return household.filled_array(2.0) From 1c084485f8d1c68356661130053901b1def9d3c1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:15:26 -0400 Subject: [PATCH 11/13] Reject samples that cannot preserve a positive input weight total --- policyengine_core/simulations/simulation.py | 5 ++++ .../test_subsample_entity_weight_totals.py | 24 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 1a797b389..fe8b8ca28 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2505,6 +2505,11 @@ def retain_structure(name, values): ] _, first_retained = np.unique(memberships, return_index=True) selected_total = subset_df[col].values[first_retained].sum() + if selected_total == 0: + raise ValueError( + f"Cannot preserve {col} total: the retained entities have " + "zero weight. Use quantize_weights=True or a different sample." + ) subset_df[col] *= target_total_weight / selected_total df = subset_df diff --git a/tests/core/test_subsample_entity_weight_totals.py b/tests/core/test_subsample_entity_weight_totals.py index 5c35f79c3..645aa482e 100644 --- a/tests/core/test_subsample_entity_weight_totals.py +++ b/tests/core/test_subsample_entity_weight_totals.py @@ -195,3 +195,27 @@ def test_zero_total_input_weight_column_remains_zero(quantize_weights): np.testing.assert_array_equal( simulation.calculate(f"{entity}_weight", year), 0 ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_positive_total_with_only_zero_retained_weights(quantize_weights): + data = weight_data( + households=[10, 20], + families=[110, 220], + household_weights={10: 1, 20: 1}, + family_weights={110: 1, 220: 1}, + person_weights=[0, 5], + ) + simulation = weight_simulation(data) + original_dataset = simulation.dataset + with patch("numpy.random.choice", return_value=np.array([10])): + if quantize_weights: + simulation.subsample(n=1, seed="zero-retained", quantize_weights=True) + assert simulation.calculate("person_weight", YEAR_INPUT).sum() == 5 + else: + with pytest.raises(ValueError, match="person_weight__2022.*zero weight"): + simulation.subsample(n=1, seed="zero-retained", quantize_weights=False) + assert simulation.dataset is original_dataset + np.testing.assert_array_equal( + simulation.calculate("person_weight", YEAR_INPUT), [0, 5] + ) From d5e775675dd9fe5c3cc1fd268bfae07d4b5d7605 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 07:00:04 -0400 Subject: [PATCH 12/13] Normalize subsample weights by the exported membership --- policyengine_core/simulations/simulation.py | 12 +- .../test_subsample_entity_weight_totals.py | 118 +++++++++++++++++- 2 files changed, 125 insertions(+), 5 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 7c8b99f61..37957736d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2609,7 +2609,8 @@ def retain_structure(name, values): # repeat on person rows, so summing those rows counts larger groups # more often. The loader reads each retained group's first member; # that member can differ from its original first member when another - # group partition crosses the sampled households. + # group partition crosses the sampled households. Group by the same + # exported membership the rebuild reads, including later overrides. for col in subset_df.columns: if "weight__" in col: @@ -2630,9 +2631,12 @@ def retain_structure(name, values): if entity.is_person: selected_total = subset_df[col].values.sum() else: - memberships = self.populations[entity.key].members_entity_id[ - subset_df.index - ] + membership_column = next( + column + for column in subset_df + if column.split("__")[0] == f"{person_key}_{entity.key}_id" + ) + memberships = subset_df[membership_column].values _, first_retained = np.unique(memberships, return_index=True) selected_total = subset_df[col].values[first_retained].sum() if selected_total == 0: diff --git a/tests/core/test_subsample_entity_weight_totals.py b/tests/core/test_subsample_entity_weight_totals.py index 645aa482e..1f5d8f098 100644 --- a/tests/core/test_subsample_entity_weight_totals.py +++ b/tests/core/test_subsample_entity_weight_totals.py @@ -15,10 +15,11 @@ from policyengine_core.data import Dataset from policyengine_core.entities import build_entity -from policyengine_core.periods import ETERNITY, YEAR +from policyengine_core.periods import ETERNITY, YEAR, period from policyengine_core.simulations import Simulation from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.variables import Variable +from tests.core.test_user_input_keys import _store_on_disk YEAR_INPUT = "2022" Person = build_entity(key="person", plural="persons", label="Person", is_person=True) @@ -158,6 +159,121 @@ def test_subsample_preserves_native_entity_weight_totals(chosen, quantize_weight assert_weight_totals_and_reference(data, chosen, quantize_weights) +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_subsample_normalizes_using_exported_membership(quantize_weights): + data = weight_data( + households=[30, 10, 20, 20, 30, 40, 40, 20, 40], + families=[110, 110, 220, 110, 330, 330, 220, 220, 440], + household_weights={10: 100.0, 20: 200.0, 30: 300.0, 40: 400.0}, + family_weights={110: 5.0, 220: 11.0, 330: 17.0, 440: 23.0}, + person_weights=np.arange(2.0, 11.0), + ) + simulation = weight_simulation(data) + exported_memberships = np.array([110, 110, 220, 220, 330, 330, 330, 220, 440]) + simulation.set_input("person_family_id", ETERNITY, exported_memberships) + + # set_input changes the recorded column without changing loaded populations. + # The retained rows contain three exported families, but four old families. + chosen = [20, 40, 20] + retained = data[data["household_id__ETERNITY"].isin(chosen)].copy() + retained["person_family_id__ETERNITY"] = exported_memberships[retained.index] + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + np.testing.assert_array_equal( + simulation.to_input_dataframe()["person_family_id__ETERNITY"], + exported_memberships, + ) + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="exported-membership", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + for entity in ENTITIES: + entity_ids = simulation.calculate(f"{entity}_id", ETERNITY) + membership_column = ( + "person_id__ETERNITY" + if entity == "person" + else f"person_{entity}_id__ETERNITY" + ) + for year in WEIGHT_YEARS: + column = f"{entity}_weight__{year}" + # Each period's target comes independently from the source's native + # entity groups, before the membership override or subsampling. + target = data.groupby(f"{entity}_id__ETERNITY")[column].first().sum() + candidate = ( + counts.astype(float) if quantize_weights else retained[column] * counts + ) + reference = candidate.groupby(retained[membership_column]).first() + reference *= target / reference.sum() + actual = simulation.calculate(f"{entity}_weight", year) + assert actual.sum() == pytest.approx(target, rel=2e-6) + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_subsample_conserves_recorded_disk_weight_total_under_memory_cache( + quantize_weights, +): + data = weight_data( + households=[10, 10, 20, 30, 30], + families=[110, 110, 220, 330, 110], + household_weights={10: 100.0, 20: 200.0, 30: 300.0}, + family_weights={110: 7.0, 220: 11.0, 330: 13.0}, + person_weights=np.arange(2.0, 7.0), + ) + simulation = weight_simulation(data) + simulation.delete_arrays("family_weight") + _store_on_disk(simulation, "family_weight") + source_weights = data.groupby("family_id__ETERNITY").first() + holder = simulation.get_holder("family_weight") + for year in WEIGHT_YEARS: + values = source_weights[f"family_weight__{year}"].to_numpy() + simulation.set_input("family_weight", year, values) + np.testing.assert_array_equal(holder._disk_storage.get(period(year)), values) + + simulation.memory_config.max_memory_occupation_pc = 101 + for year in WEIGHT_YEARS: + # An unmarked memory cache can shadow a surviving disk input in reads. + # Normalization must preserve the recorded source total from disk. + shadow = np.full(len(source_weights), 999.0) + holder._set(period(year), shadow, derived=False, is_input=False) + assert not holder._stores_user_input(period(year), "default", "memory") + assert holder._stores_user_input(period(year), "default", "disk") + np.testing.assert_array_equal( + simulation.calculate("family_weight", year), shadow + ) + + chosen = [20, 30, 20] + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="recorded-disk-weight-total", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + retained = data[data["household_id__ETERNITY"].isin(chosen)] + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + entity_ids = simulation.calculate("family_id", ETERNITY) + for year in WEIGHT_YEARS: + column = f"family_weight__{year}" + target = source_weights[column].sum() + candidate = ( + counts.astype(float) if quantize_weights else retained[column] * counts + ) + reference = candidate.groupby(retained["person_family_id__ETERNITY"]).first() + reference *= target / reference.sum() + actual = simulation.calculate("family_weight", year) + assert actual.sum() == pytest.approx(target, rel=2e-6) + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + @pytest.mark.parametrize("quantize_weights", [False, True]) @pytest.mark.parametrize("seed", range(8)) def test_entity_total_invariant_across_deterministic_random_partitions( From 0b91121d844e602ef08ab54182b625c362f42989 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 17:11:54 -0400 Subject: [PATCH 13/13] Skip weight fallback tests when Hypothesis is unavailable --- tests/core/test_subsample_weight_fallback.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/core/test_subsample_weight_fallback.py b/tests/core/test_subsample_weight_fallback.py index 9e7905704..8acc27152 100644 --- a/tests/core/test_subsample_weight_fallback.py +++ b/tests/core/test_subsample_weight_fallback.py @@ -4,7 +4,9 @@ import numpy as np import pytest -from hypothesis import given, settings, strategies as st + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies from policyengine_core.country_template import Microsimulation, Simulation from policyengine_core.country_template.entities import Household @@ -76,8 +78,8 @@ def test_missing_period_weights_use_dataset_input( check_fallback(simulation_class, uprated, carry_over, year, 3, "fallback") -@settings(max_examples=25, deadline=None) -@given( +@hypothesis.settings(max_examples=25, deadline=None) +@hypothesis.given( simulation_class=st.sampled_from([Simulation, Microsimulation]), case=st.sampled_from( [(False, False, "2023"), (False, True, "2021"), (True, True, "2021-07")]