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

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.

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

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)