From 9f1b5e8e83f47ca2129549fcda1511e0e13a3d96 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 16:34:23 -0400 Subject: [PATCH] Rebuild a subsampled simulation from its inputs, not its calculated values subsample exported to_input_dataframe(include_computed_variables=True), every stored value, and loaded it back as inputs of the rebuilt population. The baseline arm of a reform simulation is rebuilt from that population, so anything calculated under the reform before subsampling became an input of the baseline arm too: after calculating taxable_income on a policyengine-us reform simulation, the baseline arm read the reform's taxable income. subsample now rebuilds from _to_subsample_dataframe: for each variable, the periods Holder.get_input_periods reports as inputs (the storage's own input/derived mark, so a stale _user_input_keys entry cannot pass a calculated value off as an input), including inputs of formula-backed variables such as policyengine-us's dataset-supplied person_id and is_spm_independent_minor_role. A formula-backed id, membership or role variable that has no input is calculated at the dataset period, as the rebuild needs it. Tests: examples and a Hypothesis property (history independence, baseline arm equals an unreformed subsample, restriction, export depends only on inputs). Also carries the tests from #573, which this supersedes, and which pass here. Co-Authored-By: Claude Opus 5.5 --- changelog.d/subsample-inputs-only.fixed.md | 1 + policyengine_core/simulations/simulation.py | 77 +++++- tests/core/test_subsample_computed_values.py | 217 ++++++++++++++++ ...test_subsample_computed_values_property.py | 233 ++++++++++++++++++ tests/core/test_subsample_inputs_only.py | 106 ++++++++ .../test_subsample_inputs_only_property.py | 53 ++++ tests/fixtures/subsample_inputs.py | 109 ++++++++ 7 files changed, 793 insertions(+), 3 deletions(-) create mode 100644 changelog.d/subsample-inputs-only.fixed.md create mode 100644 tests/core/test_subsample_computed_values.py create mode 100644 tests/core/test_subsample_computed_values_property.py 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/subsample-inputs-only.fixed.md b/changelog.d/subsample-inputs-only.fixed.md new file mode 100644 index 000000000..e7c568aab --- /dev/null +++ b/changelog.d/subsample-inputs-only.fixed.md @@ -0,0 +1 @@ +Rebuild `Simulation.subsample` from the simulation's inputs only, so values calculated before subsampling (under a reform, say) no longer become inputs of the rebuilt arms and a reform simulation's baseline no longer reads the reform's values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 123a5f1ae..603370782 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2247,6 +2247,68 @@ def to_input_dict(self, include_computed_variables: bool = False) -> dict: return data + def _structural_variable_names(self) -> List[str]: + """The variables that place people in entities, under the names + ``build_from_dataset`` reads: the person's id, and each group + entity's id and each person's id and role in it.""" + person = self.tax_benefit_system.person_entity.key + names = [f"{person}_id"] + for group in self.tax_benefit_system.group_entities: + names += [ + f"{group.key}_id", + f"{person}_{group.key}_id", + f"{person}_{group.key}_role", + ] + return [name for name in names if name in self.tax_benefit_system.variables] + + def _to_subsample_dataframe(self) -> pd.DataFrame: + """The values ``subsample`` rebuilds the population from. + + Every value set as an input that this branch reads (see + ``Holder.get_input_periods``), whether or not its variable has a + formula: a dataset can supply a formula-backed variable (such as + policyengine-us's ``person_id``), and the rebuilt population must + read what this one read. A value the simulation calculated is never + kept. Reloaded as an input it would override the formula in every + arm of the rebuilt simulation, so after calculating under a reform, + the baseline arm would read the reform's values. + + A structural variable (see ``_structural_variable_names``) that has + no input but has a formula is calculated at the dataset's period and + kept, as the rebuild needs it to place people in entities. + """ + base_period = periods.period( + self.dataset.time_period or self.default_calculation_period + ) + structural = set(self._structural_variable_names()) + df = pd.DataFrame() + + for variable_name, variable in self.tax_benefit_system.variables.items(): + input_periods = self.get_holder(variable_name).get_input_periods( + self.branch_name + ) + if ( + not input_periods + and variable_name in structural + and not variable.is_input_variable() + ): + if variable.definition_period == ETERNITY: + input_periods = [periods.period(ETERNITY)] + elif variable.definition_period == YEAR: + input_periods = [base_period.this_year] + elif variable.definition_period == MONTH: + input_periods = [base_period.first_month] + else: + input_periods = [base_period.first_day] + for period in sorted(input_periods, key=str): + if variable.definition_period != period.unit: + continue + values = self.calculate(variable_name, period, map_to="person") + if values is not None: + df[f"{variable_name}__{period}"] = values + + return df + def subsample( self, n=None, @@ -2257,6 +2319,12 @@ def subsample( ) -> "Simulation": """Quantize the simulation to a smaller size by sampling households. + The sampled households are rebuilt from this simulation's inputs, + plus any entity id or membership a formula supplies where no input + does (see ``_to_subsample_dataframe``). No other calculated value is + carried over, so each arm of a reform simulation recalculates under + its own policy, whatever had been calculated before. + 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. @@ -2275,9 +2343,12 @@ 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) + # The rebuilt population takes this one's inputs, including those of + # formula-backed variables such as dataset-supplied IDs, and nothing + # calculated. Exporting calculated values as well (the previous + # ``include_computed_variables=True``) made them inputs of every arm, + # so a reform simulation's baseline read the reform's values. + df = self._to_subsample_dataframe() # Extract time period from DataFrame columns df_time_period = ( diff --git a/tests/core/test_subsample_computed_values.py b/tests/core/test_subsample_computed_values.py new file mode 100644 index 000000000..6466fccfe --- /dev/null +++ b/tests/core/test_subsample_computed_values.py @@ -0,0 +1,217 @@ +"""Regression tests: ``subsample`` rebuilds from inputs, never calculated values. + +``subsample`` exports the simulation, rebuilds the population from the sampled +rows, and rebuilds a reform simulation's baseline arm from the result. It +exported every value the simulation had stored +(``to_input_dataframe(include_computed_variables=True)``), so whatever had +been calculated before subsampling came back as an input of both arms: after +calculating income tax under a reform, the baseline arm read the reform's +income tax. It now exports the values ``Holder.get_input_periods`` reports as +inputs, whether or not their variable has a formula, plus any formula-backed +id or membership variable the rebuild needs and the inputs lack. + +The properties these examples come from are in +``test_subsample_computed_values_property.py``. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.country_template import Microsimulation +from policyengine_core.country_template.entities import Household +from policyengine_core.data import Dataset +from policyengine_core.model_api import Reform, Variable +from policyengine_core.periods import ETERNITY + +PERIOD = "2022-01" +LATER_PERIOD = "2022-02" +SEED = "subsample-computed-values" +N = 3 +REFORM_RATE = 0.42 +REFORM_DEFAULT_SALARY = 1000.0 +# A subsample of N households draws from 5; quantized weights lose nothing +# else, so a household's values after subsampling are its values before. +SALARIES = [1000.0, 2000.0, 0.0, 3500.0, 4000.0, 500.0, 6000.0, 0.0, 800.0, 2500.0] + + +class income_tax(Variable): + def formula(person, period, parameters): + return person("salary", period) * 0.5 + + +class halve_salaries_in_tax(Reform): + """A structural reform: income tax becomes half of salary.""" + + def apply(self): + self.update_variable(income_tax) + + +class salary(Variable): + default_value = REFORM_DEFAULT_SALARY + + +class change_salary_default(Reform): + def apply(self): + self.update_variable(salary) + + +RATE_REFORM = {"taxes.income_tax_rate": {"2022-01-01": REFORM_RATE}} + + +def _dataframe(**columns) -> pd.DataFrame: + """Five households of two people each, with salaries for ``PERIOD``.""" + df = pd.DataFrame( + { + "person_id__2022": list(range(1, 11)), + "household_id__2022": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + "person_household_id__2022": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + "household_weight__2022": [100.0] * 2 + + [200.0] * 2 + + [300.0] * 2 + + [400.0] * 2 + + [500.0] * 2, + f"salary__{PERIOD}": SALARIES, + } + ) + for name, values in columns.items(): + df[name] = values + return df + + +def _simulation(reform=None, df=None, tax_benefit_system=None) -> Microsimulation: + return Microsimulation( + dataset=Dataset.from_dataframe(_dataframe() if df is None else df, "2022"), + reform=reform, + tax_benefit_system=tax_benefit_system, + ) + + +def _subsample(simulation: Microsimulation) -> Microsimulation: + return simulation.subsample(n=N, seed=SEED, time_period="2022") + + +def _values(simulation, variable, period=PERIOD) -> np.ndarray: + return np.asarray(simulation.calculate(variable, period)) + + +@pytest.mark.parametrize( + "reform", + [ + pytest.param(RATE_REFORM, id="parametric"), + pytest.param(halve_salaries_in_tax, id="structural"), + ], +) +@pytest.mark.parametrize("variable", ["income_tax", "disposable_income"]) +def test_a_value_calculated_under_a_reform_is_not_the_baselines_after_subsample( + reform, variable +): + """The baseline arm equals an unreformed simulation of the same households.""" + simulation = _simulation(reform) + reform_before = _values(simulation, variable) + baseline_before = _values(simulation.baseline, variable) + assert not np.allclose(reform_before, baseline_before) + + _subsample(simulation) + unreformed = _subsample(_simulation()) + uncalculated = _subsample(_simulation(reform)) + + np.testing.assert_array_equal( + _values(simulation, "person_id", ETERNITY), + _values(unreformed, "person_id", ETERNITY), + ) + np.testing.assert_allclose( + _values(simulation.baseline, variable), _values(unreformed, variable) + ) + np.testing.assert_allclose( + _values(simulation, variable), _values(uncalculated, variable) + ) + + +def test_a_default_the_reform_changes_is_not_the_baselines_after_subsample(): + """A default read under the reform is calculated, not an input.""" + simulation = _simulation(change_salary_default) + assert (_values(simulation, "salary", LATER_PERIOD) == REFORM_DEFAULT_SALARY).all() + + _subsample(simulation) + + assert (_values(simulation.baseline, "salary", LATER_PERIOD) == 0).all() + assert (_values(simulation, "salary", LATER_PERIOD) == REFORM_DEFAULT_SALARY).all() + + +def test_a_value_calculated_where_an_input_was_deleted_is_not_kept(): + """An input that was set and then deleted leaves its record behind. + + policyengine-us moves dataset inputs such as ``employment_income`` to + another variable and deletes them, so ``_user_input_keys`` still lists + them. What is calculated there later is a calculated value all the same. + """ + simulation = _simulation(RATE_REFORM) + simulation.set_input("income_tax", PERIOD, np.zeros(10)) + simulation.get_holder("income_tax").delete_arrays(PERIOD) + assert any(key[0] == "income_tax" for key in simulation._user_input_keys) + reform_tax = _values(simulation, "income_tax") + assert reform_tax.sum() > 0 + + _subsample(simulation) + unreformed = _subsample(_simulation()) + + np.testing.assert_allclose( + _values(simulation.baseline, "income_tax"), _values(unreformed, "income_tax") + ) + + +def test_a_dataset_value_of_a_formula_backed_variable_is_kept(): + """A dataset can supply a variable that has a formula, overriding it. + + Subsampling keeps such a value for the sampled households: the rebuilt + population reads what the original read, in both arms. + """ + supplied = np.arange(10, dtype=float) * 7 + df = _dataframe(**{f"income_tax__{PERIOD}": supplied}) + simulation = _simulation(RATE_REFORM, df=df) + by_person = dict(zip(_values(simulation, "person_id", ETERNITY), supplied)) + + _subsample(simulation) + + ids = _values(simulation, "person_id", ETERNITY) + expected = np.array([by_person[person_id] for person_id in ids]) + np.testing.assert_array_equal(_values(simulation, "income_tax"), expected) + np.testing.assert_array_equal(_values(simulation.baseline, "income_tax"), expected) + + +class household_id(Variable): + value_type = int + entity = Household + definition_period = ETERNITY + label = "Formula-backed household ID for subsample regression tests." + + def formula(household, period): + return np.arange(household.count) + 1 + + +@pytest.mark.parametrize("calculate_first", [False, True]) +def test_a_formula_backed_id_the_dataset_lacks_is_calculated_for_the_rebuild( + isolated_tax_benefit_system, calculate_first +): + """The rebuild needs each person's household. When no input gives the + household ids but a formula does, subsample calculates them, whether or + not anything calculated them before.""" + isolated_tax_benefit_system.replace_variable(household_id) + df = _dataframe().drop(columns=["household_id__2022"]) + simulation = _simulation(df=df, tax_benefit_system=isolated_tax_benefit_system) + if calculate_first: + simulation.calculate("household_id", ETERNITY) + + _subsample(simulation) + + household_ids = _values(simulation, "household_id", ETERNITY) + assert simulation.populations["household"].count == len(household_ids) + assert len(household_ids) <= N + # Each sampled household keeps its people and their salaries. + members = _values(simulation, "person_household_id", ETERNITY) + for household in household_ids: + assert (members == household).sum() == 2 + assert set(_values(simulation, "salary")) <= set(SALARIES) diff --git a/tests/core/test_subsample_computed_values_property.py b/tests/core/test_subsample_computed_values_property.py new file mode 100644 index 000000000..d8ea56fb2 --- /dev/null +++ b/tests/core/test_subsample_computed_values_property.py @@ -0,0 +1,233 @@ +"""Properties of ``subsample`` on a reform simulation. + +For any population of one to four households of one to three people, any +salaries and household weights, with or without a dataset value for a +formula-backed variable (income tax), under any of five reforms or none, +after any calculations on either arm, and for any sample size and seed: + +* The export is a function of the inputs: calculating anything leaves + ``_to_subsample_dataframe`` unchanged. +* History independence (differential): subsampling after the calculations + rebuilds the same inputs, and both arms then return the same values, as + subsampling without them. +* Baseline fidelity (differential): the baseline arm returns what a + simulation without the reform, subsampled the same way, returns. +* Restriction: every sampled person reads, in each arm, what they read in a + simulation that was not subsampled. + +Examples, including the ones that failed before, are in +``test_subsample_computed_values.py``. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import HealthCheck, example, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.country_template import Microsimulation # noqa: E402 +from policyengine_core.data import Dataset # noqa: E402 +from policyengine_core.periods import ETERNITY # noqa: E402 +from tests.core.test_reform_baseline_system import add_doubled_salary # noqa: E402 +from tests.core.test_subsample_computed_values import ( # noqa: E402 + RATE_REFORM, + change_salary_default, + halve_salaries_in_tax, +) + +PERIOD = "2022-01" +REFORMS = { + "none": None, + "rate": RATE_REFORM, + "structural": halve_salaries_in_tax, + "salary default": change_salary_default, + "added variable": add_doubled_salary, +} +# Values each arm returns, by variable and period. None depends on weights. +OUTPUTS = [ + ("salary", PERIOD), + ("salary", "2022-02"), + ("income_tax", PERIOD), + ("income_tax", "2022-02"), + ("social_security_contribution", PERIOD), + ("disposable_income", PERIOD), + ("total_taxes", PERIOD), + ("household_income", PERIOD), + ("doubled_salary", PERIOD), +] +# What may be calculated before subsampling: the outputs and weights, +# including a household weight for a year the dataset has none for. +CALCULATIONS = OUTPUTS + [ + ("person_weight", "2022"), + ("household_weight", "2022"), + ("household_weight", "2023"), +] + + +@st.composite +def populations(draw): + sizes = draw(st.lists(st.integers(1, 3), min_size=1, max_size=4)) + people = sum(sizes) + amounts = st.integers(min_value=0, max_value=10_000).map(float) + salaries = draw(st.lists(amounts, min_size=people, max_size=people)) + weights = draw( + st.lists( + st.integers(min_value=1, max_value=1_000).map(float), + min_size=len(sizes), + max_size=len(sizes), + ) + ) + income_tax = draw(st.none() | st.lists(amounts, min_size=people, max_size=people)) + return sizes, salaries, weights, income_tax + + +def _dataset(population) -> Dataset: + sizes, salaries, weights, income_tax = population + # Ids that are not row positions, so a mix-up between the two shows. + household_ids = [10 * (index + 1) for index, size in enumerate(sizes)] + membership = np.repeat(household_ids, sizes) + df = pd.DataFrame( + { + "person_id__2022": 100 + np.arange(len(salaries)), + "household_id__2022": membership, + "person_household_id__2022": membership, + "household_weight__2022": np.repeat(weights, sizes), + f"salary__{PERIOD}": salaries, + } + ) + if income_tax is not None: + df[f"income_tax__{PERIOD}"] = income_tax + return Dataset.from_dataframe(df, "2022") + + +def _simulation(reform, population) -> Microsimulation: + return Microsimulation(dataset=_dataset(population), reform=REFORMS[reform]) + + +def _arms(simulation): + """The simulation, and its baseline arm if it has a reform.""" + baseline = getattr(simulation, "baseline", None) + return [simulation] if baseline is None else [simulation, baseline] + + +def _by_person(simulation, variable, period): + """What each person reads, keyed by person id.""" + ids = np.asarray(simulation.calculate("person_id", ETERNITY)) + values = np.asarray(simulation.calculate(variable, period, map_to="person")) + return dict(zip(ids.tolist(), values.tolist())) + + +def _check(reform, population, calculations, n, seed): + def subsample(simulation): + return simulation.subsample(n=n, seed=seed, time_period="2022") + + reference = _simulation(reform, population) + calculated = _simulation(reform, population) + export = calculated._to_subsample_dataframe() + for (variable, period), on_baseline in calculations: + arm = _arms(calculated)[-1] if on_baseline else calculated + if variable in arm.tax_benefit_system.variables: + arm.calculate(variable, period) + pd.testing.assert_frame_equal(calculated._to_subsample_dataframe(), export) + + subsample(calculated) + uncalculated = subsample(_simulation(reform, population)) + unreformed = subsample(_simulation("none", population)) + + pd.testing.assert_frame_equal( + calculated._to_subsample_dataframe(), + uncalculated._to_subsample_dataframe(), + ) + assert len(_arms(calculated)) == len(_arms(uncalculated)) == len(_arms(reference)) + for variable, period in OUTPUTS: + for arm, uncalculated_arm, reference_arm in zip( + _arms(calculated), _arms(uncalculated), _arms(reference) + ): + if variable not in arm.tax_benefit_system.variables: + continue + label = f"{variable}@{period} in {arm.branch_name}" + values = np.asarray(arm.calculate(variable, period)) + np.testing.assert_array_equal( + values, + np.asarray(uncalculated_arm.calculate(variable, period)), + err_msg=label, + ) + before = _by_person(reference_arm, variable, period) + after = _by_person(arm, variable, period) + assert after == {key: before[key] for key in after}, label + if variable in unreformed.tax_benefit_system.variables: + np.testing.assert_array_equal( + np.asarray(_arms(calculated)[-1].calculate(variable, period)), + np.asarray(unreformed.calculate(variable, period)), + err_msg=f"{variable}@{period} in the baseline arm", + ) + + +TWO_HOUSEHOLDS = ([1, 2], [1000.0, 2000.0, 3000.0], [1.0, 3.0], None) + + +@settings( + max_examples=200, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given( + reform=st.sampled_from(list(REFORMS)), + population=populations(), + calculations=st.lists( + st.tuples(st.sampled_from(CALCULATIONS), st.booleans()), max_size=6 + ), + n=st.integers(min_value=1, max_value=5), + seed=st.integers(min_value=0, max_value=2**16), +) +# Income tax calculated under the reform became the baseline's. +@example( + reform="rate", + population=TWO_HOUSEHOLDS, + calculations=[(("income_tax", PERIOD), False)], + n=2, + seed=0, +) +# So did a formula the reform replaces, and what it feeds. +@example( + reform="structural", + population=TWO_HOUSEHOLDS, + calculations=[(("disposable_income", PERIOD), False)], + n=2, + seed=0, +) +# A default the reform changes, read where no input is. +@example( + reform="salary default", + population=TWO_HOUSEHOLDS, + calculations=[(("salary", "2022-02"), False)], + n=2, + seed=0, +) +# A household weight read for a year the dataset has none for: its default +# came back as an input for that year. +@example( + reform="none", + population=TWO_HOUSEHOLDS, + calculations=[(("household_weight", "2023"), False)], + n=2, + seed=0, +) +# A dataset value for a formula-backed variable stays an input of both arms. +@example( + reform="rate", + population=([2, 1], [1000.0, 2000.0, 3000.0], [1.0, 3.0], [5.0, 6.0, 7.0]), + calculations=[(("income_tax", "2022-02"), True)], + n=3, + seed=1, +) +def test_subsample_after_calculating_is_subsample_before( + reform, population, calculations, n, seed +): + _check(reform, population, calculations, n, seed) 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