Rebuild a subsampled simulation from its inputs, not its calculated values

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.

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 = (

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()

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)

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