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 00000000..b9654be4 --- /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 2652d972..0adb39be 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 00000000..ecdaefdc --- /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 00000000..88dbc8ad --- /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 00000000..3a516160 --- /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