diff --git a/changelog.d/subsample-reform-baseline-tests.changed.md b/changelog.d/subsample-reform-baseline-tests.changed.md new file mode 100644 index 000000000..1ec203c06 --- /dev/null +++ b/changelog.d/subsample-reform-baseline-tests.changed.md @@ -0,0 +1 @@ +Add regression coverage for reform and baseline fidelity when subsampling simulations after calculations. diff --git a/tests/core/test_subsample_computed_values.py b/tests/core/test_subsample_computed_values.py new file mode 100644 index 000000000..4139f63c7 --- /dev/null +++ b/tests/core/test_subsample_computed_values.py @@ -0,0 +1,173 @@ +"""Regression tests: ``subsample`` rebuilds from inputs, never calculated values. + +``subsample`` exports surviving supplied inputs from their recorded storage +tiers, including overrides of formula-backed variables, and rebuilds both +the reform and baseline arms. Prior calculations must not become inputs: +each rebuilt arm must continue to use its own formulas and defaults. + +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.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. These outputs do not depend on +# weights, so each retained person's values should stay the same. +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) -> Microsimulation: + return Microsimulation( + dataset=Dataset.from_dataframe(_dataframe() if df is None else df, "2022"), + reform=reform, + ) + + +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(): + """Deleting an override restores the formula without creating a new input. + + A subsequent calculation under the reform must not become an override + of the baseline formula when subsampling rebuilds both arms. + """ + simulation = _simulation(RATE_REFORM) + simulation.set_input("income_tax", PERIOD, np.zeros(10)) + simulation.get_holder("income_tax").delete_arrays(PERIOD) + 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) 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..188c38e86 --- /dev/null +++ b/tests/core/test_subsample_computed_values_property.py @@ -0,0 +1,253 @@ +"""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 four reforms or none, +after any calculations on either arm, and for any sample size and seed: + +* The recorded-input export is a function of the inputs: calculating + anything on either arm leaves it 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 _recorded_inputs(simulation): + """The input export ``subsample`` uses, including formula-backed inputs.""" + return simulation._to_person_dataframe(simulation._get_set_input_periods) + + +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) + exports = [_recorded_inputs(arm) for arm in _arms(calculated)] + 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) + for arm, export in zip(_arms(calculated), exports): + pd.testing.assert_frame_equal(_recorded_inputs(arm), export) + + subsample(calculated) + uncalculated = subsample(_simulation(reform, population)) + unreformed = subsample(_simulation("none", population)) + + assert len(_arms(calculated)) == len(_arms(uncalculated)) == len(_arms(reference)) + for arm, uncalculated_arm in zip(_arms(calculated), _arms(uncalculated)): + pd.testing.assert_frame_equal( + _recorded_inputs(arm), + _recorded_inputs(uncalculated_arm), + ) + 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, +) +# Combined calculations on both arms cannot affect either rebuild. +@example( + reform="structural", + population=TWO_HOUSEHOLDS, + calculations=[ + (("income_tax", PERIOD), False), + (("disposable_income", PERIOD), True), + (("household_weight", "2023"), False), + (("salary", "2022-02"), True), + ], + n=2, + seed=0, +) +def test_subsample_after_calculating_is_subsample_before( + reform, population, calculations, n, seed +): + _check(reform, population, calculations, n, seed)