diff --git a/changelog.d/fix-lsr-branch-baseline-measurements.fixed.md b/changelog.d/fix-lsr-branch-baseline-measurements.fixed.md new file mode 100644 index 00000000000..a9f2ee4d6be --- /dev/null +++ b/changelog.d/fix-lsr-branch-baseline-measurements.fixed.md @@ -0,0 +1,3 @@ +Measure a branch's or clone's labor-supply and capital-gains responses against baseline policy applied to its own inputs, and give each simulation its own measurement cache. A branch of a reform simulation made before anything was calculated (and, on policyengine-core before #587, such a clone) measured against a new branch under the reform policy and reported no response, and a clone given other inputs reused its source's measurements. +Preserve zero responses without a reform, make capital-gains marginal tax rates independent of calculation order while reusing household measurements, and discard incompatible holders when applying baseline policy to measurement branches. +Forget deleted pre-response input records on Core 3.32.8 so fully calculated clones remeasure changed earnings, while retaining explicit overrides and the baseline household's inputs. diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index d5c540e0a96..7ff89faceba 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -83,6 +83,12 @@ COUNTY_FIPS_PATTERN = re.compile(r"[0-9]{5}") +# Where a simulation keeps its behavioural-response measurements, by period +# (``variables/gov/simulation/behavioral_response_measurements.py``). They are +# cached formula output kept outside the holders, so the simulation's clone, +# branch and invalidation paths below handle them explicitly. +BEHAVIORAL_RESPONSE_CACHE_ATTR = "_behavioral_response_measurements" + COUNTY_INPUT_FIX = ( 'send county_fips as a five-digit string (for example "06037"), or select ' 'geography_kind="national" in the spm configuration' @@ -711,7 +717,16 @@ def _isolate_parameter_tracing(self): branch._isolate_parameter_tracing() def get_branch(self, name="branch", clone_system=False): + creating = name != self.branch_name and name not in self.branches branch = super().get_branch(name, clone_system) + if creating: + # A branch starts from what this simulation has calculated, its + # behavioural-response measurements included, but keeps its own + # copy (``clone`` emptied it): what either measures afterwards, the + # other does not see. + measurements = self.__dict__.get(BEHAVIORAL_RESPONSE_CACHE_ATTR) + if measurements is not None: + branch.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] = dict(measurements) # Core names the branch and hands it this simulation's tracer after # cloning, so re-prime the branch's root with what it ended up holding. branch._isolate_parameter_tracing() @@ -908,8 +923,20 @@ def clone(self, debug=False, trace=False, clone_tax_benefit_system=True): cloned._rebind_holders() self._rebind_method_aliases(cloned) cloned._isolate_parameter_tracing() + # Core copies the instance dictionary by reference, which shared this + # simulation's measurements dictionary with the clone: a clone given + # other inputs reused them, and what it measured reached this + # simulation. A clone measures its own (``get_branch`` seeds a branch). + cloned.__dict__.pop(BEHAVIORAL_RESPONSE_CACHE_ATTR, None) return cloned + def _invalidate_all_caches(self): + # Core purges cached formula output here (``apply_reform``, + # ``subsample``); the measurements are formula output kept outside the + # holders. + self.__dict__.pop(BEHAVIORAL_RESPONSE_CACHE_ATTR, None) + super()._invalidate_all_caches() + def _rebind_method_aliases(self, cloned): """Point copied bound-method aliases at the clone, not the original. diff --git a/policyengine_us/system.py b/policyengine_us/system.py index c4c90688ea2..1b8a070ec6b 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -35,6 +35,7 @@ uprate_parameters, ) from .tools.default_uprating import add_default_uprating +from .tools.period_branch import delete_input_values from .tools.section_911 import validate_section_911_batch_inputs from .tools.per_capita_uprating import ( add_per_capita_parameters_for_parameter_uprating, @@ -264,13 +265,13 @@ def __init__(self, *args, **kwargs): for known_period in employment_income.get_known_periods(): array = employment_income.get_array(known_period) self.set_input("employment_income_before_lsr", known_period, array) - employment_income.delete_arrays(known_period) + delete_input_values(employment_income, known_period) self_employment_income = self.get_holder("self_employment_income") for known_period in self_employment_income.get_known_periods(): array = self_employment_income.get_array(known_period) self.set_input("self_employment_income_before_lsr", known_period, array) - self_employment_income.delete_arrays(known_period) + delete_input_values(self_employment_income, known_period) sstb_self_employment_income = self.get_holder("sstb_self_employment_income") for known_period in sstb_self_employment_income.get_known_periods(): @@ -278,13 +279,13 @@ def __init__(self, *args, **kwargs): self.set_input( "sstb_self_employment_income_before_lsr", known_period, array ) - sstb_self_employment_income.delete_arrays(known_period) + delete_input_values(sstb_self_employment_income, known_period) weekly_hours = self.get_holder("weekly_hours_worked") for known_period in weekly_hours.get_known_periods(): array = weekly_hours.get_array(known_period) self.set_input("weekly_hours_worked_before_lsr", known_period, array) - weekly_hours.delete_arrays(known_period) + delete_input_values(weekly_hours, known_period) # Capital gains responses @@ -294,7 +295,7 @@ def __init__(self, *args, **kwargs): self.set_input( "long_term_capital_gains_before_response", known_period, array ) - cg_holder.delete_arrays(known_period) + delete_input_values(cg_holder, known_period) # Geography backfill: state_code_str-only input -> state_code enum. _backfill_state_code_from_str(self) @@ -524,13 +525,13 @@ def __init__(self, *args, **kwargs): for known_period in employment_income.get_known_periods(): array = employment_income.get_array(known_period) self.set_input("employment_income_before_lsr", known_period, array) - employment_income.delete_arrays(known_period) + delete_input_values(employment_income, known_period) self_employment_income = self.get_holder("self_employment_income") for known_period in self_employment_income.get_known_periods(): array = self_employment_income.get_array(known_period) self.set_input("self_employment_income_before_lsr", known_period, array) - self_employment_income.delete_arrays(known_period) + delete_input_values(self_employment_income, known_period) sstb_self_employment_income = self.get_holder("sstb_self_employment_income") for known_period in sstb_self_employment_income.get_known_periods(): @@ -538,13 +539,13 @@ def __init__(self, *args, **kwargs): self.set_input( "sstb_self_employment_income_before_lsr", known_period, array ) - sstb_self_employment_income.delete_arrays(known_period) + delete_input_values(sstb_self_employment_income, known_period) weekly_hours = self.get_holder("weekly_hours_worked") for known_period in weekly_hours.get_known_periods(): array = weekly_hours.get_array(known_period) self.set_input("weekly_hours_worked_before_lsr", known_period, array) - weekly_hours.delete_arrays(known_period) + delete_input_values(weekly_hours, known_period) # Capital gains responses @@ -554,7 +555,7 @@ def __init__(self, *args, **kwargs): self.set_input( "long_term_capital_gains_before_response", known_period, array ) - cg_holder.delete_arrays(known_period) + delete_input_values(cg_holder, known_period) # Geography backfill: state_code_str-only input -> state_code enum. # Datasets carry state_fips, so this only fires for situations that diff --git a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py new file mode 100644 index 00000000000..58d6c5f2bf7 --- /dev/null +++ b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py @@ -0,0 +1,466 @@ +"""Behavioural responses of clones and branches of a reform simulation. + +A reform simulation measures its labour-supply and capital-gains responses +once per period, comparing household net income and marginal tax rates under +the reform and under baseline policy, with the responses neutralized +(``get_behavioral_response_measurements``). These tests cover the copies of +such a simulation, which a single simulation's tests never reach: + +- A branch made before anything was calculated used to measure against + ``branch.get_branch("baseline")``, a new branch under the *reform* policy, so + its response came out zero. A clone made that early did the same on + policyengine-core before #587. +- The measurements dictionary lived in the instance ``__dict__``, which core's + ``clone`` copies by reference, so a clone given other inputs reused its + source's measurements, and a branch's measurement for a new period landed in + its parent's dictionary. + +Invariants: + +1. A clone, at any point, and a branch made before its source measured, + report the response a new reform simulation with their inputs reports. +2. No two simulations share a measurements dictionary, and a measurement made + by one never reaches another. +3. A branch made after its source measured starts from a copy of its source's + measurements. Under d1021 (ruled 2026-10-08), marginal-tax-rate and + perturbation branches hold responses at the household's measured values. +4. Invalidating a simulation's cached formula output drops its measurements. +""" + +import numpy as np +import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st +from policyengine_core.periods import ETERNITY, YEAR as YEAR_PERIOD +from policyengine_core.periods import period as to_period +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable + +from policyengine_us import Simulation +from policyengine_us.entities import Household, Person +from policyengine_us.spm import BEHAVIORAL_RESPONSE_CACHE_ATTR +from policyengine_us.system import system +from policyengine_us.tools.period_branch import drop_inherited_values +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + _baseline_measurement_branch, +) + +YEAR = 2026 +FOREVER = f"{YEAR}-01-01.2100-12-31" +RATE_CUT = {"gov.irs.income.bracket.rates.2": {FOREVER: 0.10}} +LSR_REFORM = { + **RATE_CUT, + "gov.simulation.labor_supply_responses.elasticities.income": {FOREVER: -0.05}, + "gov.simulation.labor_supply_responses.elasticities.substitution.all": { + FOREVER: 0.25 + }, +} +WAGES = (20_000, 50_000, 100_000) + + +def household(wages): + return { + "people": {"you": {"age": {YEAR: 35}, "employment_income": {YEAR: wages}}}, + "tax_units": {"tu": {"members": ["you"]}}, + "spm_units": {"spm": {"members": ["you"]}}, + "families": {"f": {"members": ["you"]}}, + "marital_units": {"mu": {"members": ["you"]}}, + "households": {"h": {"members": ["you"], "state_name": {YEAR: "CA"}}}, + } + + +def make_simulation(**kwargs): + """Reuse read-only policy; each reform gets a private model and inputs.""" + kwargs.setdefault("tax_benefit_system", system) + return Simulation(**kwargs) + + +def response(simulation, period=YEAR): + return np.asarray( + simulation.calculate("labor_supply_behavioral_response", period), + dtype=float, + ) + + +_fresh = {} + + +def fresh_response(wages, period=YEAR): + """The response a new reform simulation with these wages reports. + + Values only, so sharing them across tests shares no simulation state. + """ + key = (wages, period) + if key not in _fresh: + _fresh[key] = response( + make_simulation(situation=household(wages), reform=LSR_REFORM), period + ) + return _fresh[key] + + +def set_wages(simulation, wages, period=YEAR): + """Give a copy other wages, as a marginal-tax-rate branch does its inputs. + + A copy made after its source calculated starts from the source's values, + and ``set_input`` clears nothing calculated from the old input, so drop + everything but inputs first. + """ + drop_inherited_values(simulation) + simulation.set_input("employment_income_before_lsr", period, np.array([wages])) + + +@settings(max_examples=3, deadline=None, derandomize=True) +@given( + kind=st.sampled_from(["clone", "branch"]), + after_source_measured=st.booleans(), + source_wages=st.sampled_from(WAGES), + copy_wages=st.sampled_from(WAGES), +) +# Each of these failed before: zero response, or the source's relative changes. +@example( + kind="branch", after_source_measured=False, source_wages=50_000, copy_wages=50_000 +) +@example( + kind="clone", after_source_measured=False, source_wages=50_000, copy_wages=50_000 +) +@example( + kind="clone", after_source_measured=True, source_wages=50_000, copy_wages=100_000 +) +@example( + kind="branch", after_source_measured=True, source_wages=50_000, copy_wages=100_000 +) +def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( + kind, after_source_measured, source_wages, copy_wages +): + source = make_simulation(situation=household(source_wages), reform=LSR_REFORM) + if after_source_measured: + source_response = response(source) + measured = { + name: values.copy() + for name, values in source.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR][ + str(YEAR) + ].items() + } + copy = source.clone() if kind == "clone" else source.get_branch("copy") + if copy_wages != source_wages: + set_wages(copy, copy_wages) + copy_response = response(copy) + if kind == "branch" and after_source_measured: + # d1021: perturbed wages use the household's measured relative changes. + # With this reform's uniform elasticities, the response scales with wages. + np.testing.assert_allclose( + copy_response, source_response * copy_wages / source_wages, atol=1e-3 + ) + copied_measurements = copy.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR][str(YEAR)] + assert copied_measurements.keys() == measured.keys() + for name, values in measured.items(): + np.testing.assert_array_equal(copied_measurements[name], values) + else: + np.testing.assert_allclose(copy_response, fresh_response(copy_wages), atol=1e-3) + if after_source_measured: + np.testing.assert_allclose( + response(source), fresh_response(source_wages), atol=1e-3 + ) + + +def test_migrated_input_records_are_private_and_explicit_overrides_survive(): + migrated_inputs = { + "employment_income": ("employment_income_before_lsr", 50_000), + "self_employment_income": ("self_employment_income_before_lsr", 5_000), + "sstb_self_employment_income": ( + "sstb_self_employment_income_before_lsr", + 2_000, + ), + "weekly_hours_worked": ("weekly_hours_worked_before_lsr", 40), + "long_term_capital_gains": ( + "long_term_capital_gains_before_response", + 1_000, + ), + } + situation = household(50_000) + situation["people"]["you"].update( + {name: {YEAR: value} for name, (_, value) in migrated_inputs.items()} + ) + source = make_simulation(situation=situation, reform=LSR_REFORM) + period = to_period(YEAR) + for original, (migrated, value) in migrated_inputs.items(): + assert (migrated, source.branch_name, period) in source._user_input_keys + assert (original, source.branch_name, period) not in source._user_input_keys + # Baseline was copied before migration and must retain its raw inputs. + assert ( + original, + source.branch_name, + period, + ) in source.baseline._user_input_keys + np.testing.assert_array_equal( + source.baseline.calculate(original, YEAR), [value] + ) + source_input_keys = set(source._user_input_keys) + baseline_input_keys = set(source.baseline._user_input_keys) + + clone = source.clone() + override_period = to_period(YEAR + 1) + override = np.array([123_456.0]) + clone.set_input("employment_income", override_period, override) + override_key = ("employment_income", clone.branch_name, override_period) + drop_inherited_values(clone) + assert override_key in clone._user_input_keys + np.testing.assert_array_equal( + clone.get_holder("employment_income").get_array( + override_period, clone.branch_name + ), + override, + ) + assert source._user_input_keys == source_input_keys + assert source.baseline._user_input_keys == baseline_input_keys + assert ( + source.get_holder("employment_income").get_array( + override_period, source.branch_name + ) + is None + ) + + +def test_a_branch_measures_a_new_period_without_touching_its_parent(): + parent = make_simulation(situation=household(50_000), reform=LSR_REFORM) + response(parent) + branch = parent.get_branch("next_year") + set_wages(branch, 100_000, YEAR + 1) + # The branch's inputs: this year's wages from the household, next year's + # set exactly (a household's own next-year wages would be uprated). + fresh = make_simulation(situation=household(50_000), reform=LSR_REFORM) + fresh.set_input("employment_income_before_lsr", YEAR + 1, np.array([100_000])) + np.testing.assert_allclose( + response(branch, YEAR + 1), response(fresh, YEAR + 1), atol=1e-3 + ) + assert str(YEAR + 1) not in parent.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] + np.testing.assert_allclose( + response(parent, YEAR + 1), fresh_response(50_000, YEAR + 1), atol=1e-3 + ) + + +def test_copies_own_their_measurements(): + source = make_simulation(situation=household(50_000), reform=LSR_REFORM) + response(source) + measurements = source.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] + assert str(YEAR) in measurements + + clone = source.clone() + assert BEHAVIORAL_RESPONSE_CACHE_ATTR not in clone.__dict__ + + branch = source.get_branch("after") + copied = branch.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] + assert copied is not measurements + assert copied.keys() == measurements.keys() + # Asking for an existing branch hands it back as it is. + copied["marker"] = None + assert source.get_branch("after") is branch + assert "marker" in branch.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] + assert "marker" not in measurements + + +class _no_change(Reform): + def apply(self): + pass + + +def test_invalidating_cached_output_drops_the_measurements(): + simulation = make_simulation(situation=household(50_000), reform=RATE_CUT) + branch = simulation.get_branch("child") + for member in (simulation, branch): + member.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] = {str(YEAR): "stale"} + # Core's apply_reform (and subsample) purge formula output through + # _invalidate_all_caches, which reaches every branch. + simulation.apply_reform(_no_change) + assert BEHAVIORAL_RESPONSE_CACHE_ATTR not in simulation.__dict__ + assert BEHAVIORAL_RESPONSE_CACHE_ATTR not in branch.__dict__ + + +def test_baseline_measurement_branch_is_baseline_policy_over_the_inputs(): + simulation = make_simulation(situation=household(50_000), reform=RATE_CUT) + reform_tax = simulation.calculate("income_tax", YEAR) + expected_tax = simulation.baseline.calculate("income_tax", YEAR) + assert not np.allclose(reform_tax, expected_tax) + + branch = _baseline_measurement_branch(simulation) + assert simulation.branches[BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH] is ( + branch + ) + assert branch.baseline is None + policy = branch.tax_benefit_system + assert policy is not simulation.baseline.tax_benefit_system + assert policy.simulation is branch + assert policy.parameters(YEAR).gov.irs.income.bracket.rates["2"] == 0.12 + # Nothing calculated under the reform survives, and the inputs do. + assert not branch.get_holder("income_tax").get_known_periods() + np.testing.assert_array_equal( + branch.calculate("employment_income_before_lsr", YEAR), [50_000] + ) + np.testing.assert_allclose(branch.calculate("income_tax", YEAR), expected_tax) + np.testing.assert_allclose(simulation.calculate("income_tax", YEAR), reform_tax) + # The policy is the branch's own: neutralizing in it reaches no one else. + policy.neutralize_variable("income_tax") + assert not simulation.baseline.tax_benefit_system.variables[ + "income_tax" + ].is_neutralized + simulation.branches.pop(BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH) + + +def test_baseline_measurement_branch_reads_a_branchs_own_inputs(): + # A reform household's baseline keeps the situation's employment_income as + # an override from construction, so a branch with other earnings cannot be + # measured from it. + simulation = make_simulation(situation=household(50_000), reform=RATE_CUT) + branch = simulation.get_branch("other_earnings") + branch.set_input("employment_income_before_lsr", YEAR, np.array([100_000.0])) + measurement = _baseline_measurement_branch(branch) + expected = make_simulation(situation=household(100_000)).calculate( + "household_net_income", YEAR + ) + np.testing.assert_allclose( + measurement.calculate("household_net_income", YEAR), expected + ) + np.testing.assert_allclose( + simulation.baseline.calculate("household_net_income", YEAR), + make_simulation(situation=household(50_000)).calculate( + "household_net_income", YEAR + ), + ) + + +def test_baseline_measurement_branch_of_a_traced_simulation(): + simulation = make_simulation(situation=household(50_000), reform=RATE_CUT) + simulation.trace = True + branch = _baseline_measurement_branch(simulation) + assert branch.trace + assert branch.tracer is simulation.tracer + # The swapped-in policy gets its own root node, primed for this branch, + # so its parameter reads are recorded under it, not on the reform's root. + root = branch.tax_benefit_system.parameters + assert root is not simulation.tax_benefit_system.parameters + assert root.trace and root.tracer is simulation.tracer + assert root.branch_name == BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH + assert simulation.tax_benefit_system.parameters.branch_name == "default" + + +@pytest.mark.parametrize("reform", [None, RATE_CUT]) +def test_without_a_baseline_the_baseline_measurement_is_the_simulations_policy( + reform, +): + # A simulation built without a reform, or a reform simulation's own + # baseline, has no baseline: there is no reform, so it is measured against + # its own policy, and every change is zero. + built = make_simulation(situation=household(50_000), reform=reform) + simulation = built if reform is None else built.baseline + assert simulation.baseline is None + branch = _baseline_measurement_branch(simulation) + assert branch.baseline is None + assert branch.tax_benefit_system is not simulation.tax_benefit_system + assert branch.tax_benefit_system.parameters(YEAR).gov.irs.income.bracket.rates[ + "2" + ] == pytest.approx(0.12) + np.testing.assert_allclose( + branch.calculate("income_tax", YEAR), simulation.calculate("income_tax", YEAR) + ) + + +@pytest.mark.parametrize("reform", [None, RATE_CUT]) +def test_response_variables_are_zero_without_a_reform(reform): + # policyengine.py reads these from its baseline simulation, which has no + # reform; they measure without the labour-supply response's own check. + built = make_simulation(situation=household(50_000), reform=reform) + simulation = built if reform is None else built.baseline + assert response(simulation) == 0 + for variable in ( + "income_elasticity_lsr", + "substitution_elasticity_lsr", + "relative_income_change", + "relative_wage_change", + "relative_capital_gains_mtr_change", + ): + np.testing.assert_array_equal(simulation.calculate(variable, YEAR), [0]) + + +CG_REFORM = { + **LSR_REFORM, + "gov.irs.capital_gains.rates.2": {FOREVER: 0.25}, + "gov.simulation.capital_gains_responses.elasticity": {FOREVER: -0.62}, +} + + +def test_capital_gains_marginal_tax_rate_does_not_depend_on_calculation_order(): + # Its branches start from the household's measurements however it is + # reached: first, or after the household's net income. + situation = household(60_000) + situation["people"]["you"]["long_term_capital_gains"] = {YEAR: 200_000} + first = make_simulation(situation=situation, reform=CG_REFORM) + rate_first = first.calculate("marginal_tax_rate_on_capital_gains", YEAR) + after = make_simulation(situation=situation, reform=CG_REFORM) + after.calculate("household_net_income", YEAR) + rate_after = after.calculate("marginal_tax_rate_on_capital_gains", YEAR) + np.testing.assert_allclose(rate_first, rate_after, atol=1e-5) + np.testing.assert_allclose( + first.calculate("household_net_income", YEAR), + after.calculate("household_net_income", YEAR), + ) + + +@pytest.mark.parametrize("change", ["period", "entity"]) +def test_baseline_measurement_drops_incompatible_holders(change): + # A policy swap must discard holders whose storage or population cannot + # serve the baseline definition, even when the variable name is the same. + class is_household_head(Variable): + definition_period = YEAR_PERIOD if change == "period" else ETERNITY + entity = Person if change == "period" else Household + + class ChangeHead(Reform): + def apply(self): + self.update_variable(is_household_head) + + simulation = make_simulation( + situation=household(50_000), tax_benefit_system=system, reform=ChangeHead + ) + # Core creates holders before applying the reform. Recreate this one so + # its populated storage fits the reform, then becomes unsuitable for + # baseline policy; an eternal source holder would already fit baseline. + population = simulation.get_variable_population("is_household_head") + population._holders.pop("is_household_head", None) + assert simulation.get_holder("is_household_head")._memory_storage.is_eternal == ( + change == "entity" + ) + simulation.set_input("is_household_head", YEAR, np.array([True])) + simulation.calculate("is_household_head", YEAR) + original_holder = simulation.get_holder("is_household_head") + branch = _baseline_measurement_branch(simulation) + + for population in branch.populations.values(): + for holder in population._holders.values(): + assert holder.variable.entity.key == population.entity.key + assert holder._memory_storage.is_eternal == ( + holder.variable.definition_period == ETERNITY + ) + np.testing.assert_array_equal( + branch.calculate("is_household_head", YEAR), + simulation.baseline.calculate("is_household_head", YEAR), + ) + # Pruning a branch cannot remove the source's holder or change its policy. + assert simulation.get_holder("is_household_head") is original_holder + assert ( + original_holder.variable.entity + is simulation.tax_benefit_system.variables["is_household_head"].entity + ) + assert original_holder.variable.definition_period == ( + YEAR_PERIOD if change == "period" else ETERNITY + ) + + +def test_a_fully_calculated_clone_remeasures_changed_wages(): + source = make_simulation(situation=household(50_000), reform=LSR_REFORM) + original_net_income = source.calculate("household_net_income", YEAR).copy() + cloned = source.clone() + set_wages(cloned, 100_000) + np.testing.assert_allclose(response(cloned), fresh_response(100_000), atol=1e-3) + np.testing.assert_array_equal( + source.calculate("household_net_income", YEAR), original_net_income + ) diff --git a/policyengine_us/tests/core/test_behavioral_response_measurements.py b/policyengine_us/tests/core/test_behavioral_response_measurements.py index dd488b51a7d..c220af8877b 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurements.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurements.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import policyengine_us.variables.gov.simulation.behavioral_response_measurements as measurements_module import policyengine_us.variables.gov.simulation.capital_gains_responses as capital_gains_module import policyengine_us.variables.gov.simulation.labor_supply_response.income_elasticity_lsr as income_lsr_module import policyengine_us.variables.gov.simulation.labor_supply_response.labor_supply_behavioral_response as labor_supply_module @@ -11,6 +12,7 @@ import policyengine_us.variables.gov.simulation.labor_supply_response.substitution_elasticity_lsr as substitution_lsr_module from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + BEHAVIORAL_RESPONSE_CACHE_ATTR, BEHAVIORAL_RESPONSE_INPUT_VARIABLES, BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES, @@ -100,7 +102,7 @@ def get_branch(self, name, clone_system=False): if name not in self.branches: if name == BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH: self.branches[name] = self.measurement_branch - elif name == "baseline": + elif name == BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH: self.branches[name] = self.baseline_branch else: raise AssertionError(f"Unexpected branch lookup: {name}") @@ -157,7 +159,7 @@ def make_parameters( ) -def test_behavioral_response_measurements_use_one_neutralized_pass(): +def test_behavioral_response_measurements_use_one_neutralized_pass(monkeypatch): period = 2026 reform_measurement = FakeBranch( FakePopulation( @@ -173,13 +175,21 @@ def test_behavioral_response_measurements_use_one_neutralized_pass(): capital_gains_mtr=np.array([0.15, 0.2]), ) ) - baseline_parent = FakeBranch( - child_branches={ - BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH: baseline_measurement - } - ) - simulation = FakeSimulation(reform_measurement, baseline_parent) + simulation = FakeSimulation(reform_measurement, baseline_measurement) person = FakePerson(simulation) + # The baseline branch's construction needs a real simulation; it is + # tested in test_behavioral_response_measurement_copies.py. + built_from = [] + + def baseline_measurement_branch(simulation): + built_from.append(simulation) + return simulation.get_branch(BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH) + + monkeypatch.setattr( + measurements_module, + "_baseline_measurement_branch", + baseline_measurement_branch, + ) measurements = get_behavioral_response_measurements(person, period) @@ -187,13 +197,12 @@ def test_behavioral_response_measurements_use_one_neutralized_pass(): assert np.array_equal( measurements["baseline_capital_gains_mtr"], np.array([0.15, 0.2]) ) + # Both measurement branches are branches of the measured simulation. assert simulation.get_branch_calls == [ (BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, True), - ("baseline", False), - ] - assert baseline_parent.get_branch_calls == [ - (BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, True) + (BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, False), ] + assert built_from == [simulation] for branch in (reform_measurement, baseline_measurement): assert branch.tax_benefit_system.neutralized_variables == list( NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES @@ -201,20 +210,40 @@ def test_behavioral_response_measurements_use_one_neutralized_pass(): assert [call[0] for call in branch.input_calls] == list( BEHAVIORAL_RESPONSE_INPUT_VARIABLES ) + assert ( + baseline_measurement.tax_benefit_system.parameters.simulation + is reform_measurement.tax_benefit_system.parameters.simulation + ) assert simulation.macro_cache_read is False assert simulation.macro_cache_write is False - assert BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH not in simulation.branches - assert ( - BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH not in baseline_parent.branches - ) + assert simulation.branches == {} cached = get_behavioral_response_measurements(person, period) assert cached is measurements - assert simulation.get_branch_calls == [ - (BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, True), - ("baseline", False), - ] + assert simulation.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] == {"2026": measurements} + assert len(simulation.get_branch_calls) == 2 + + +def test_behavioral_response_measurement_branches_are_removed_on_failure( + monkeypatch, +): + simulation = FakeSimulation(FakeBranch(), FakeBranch()) + person = FakePerson(simulation) + + def failing_baseline_measurement_branch(simulation): + simulation.get_branch(BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH) + raise RuntimeError("baseline measurement failed") + + monkeypatch.setattr( + measurements_module, + "_baseline_measurement_branch", + failing_baseline_measurement_branch, + ) + with pytest.raises(RuntimeError, match="baseline measurement failed"): + get_behavioral_response_measurements(person, 2026) + assert simulation.branches == {} + assert simulation.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] == {} def test_measurement_helper_calculations_clip_expected_values(): diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index f746b651c34..dd3b80f96da 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -24,6 +24,7 @@ from policyengine_us.data.dataset_schema import USMultiYearDataset, USSingleYearDataset from policyengine_us.entities import Person from policyengine_us.spm import ( + BEHAVIORAL_RESPONSE_CACHE_ATTR, DERIVED_POVERTY_OUTPUTS, REJECTED_DATASET_INPUTS, create_spm_provider, @@ -861,9 +862,14 @@ def test_subsampling_never_changes_which_policy_a_holder_belongs_to( simulation.calculate(variable, 2024) for n, seed, quantize_weights in draws: + previous_baseline = changed.baseline + for simulation in (changed, previous_baseline): + simulation.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] = {"2024": "stale"} changed.subsample(n, seed=seed, quantize_weights=quantize_weights) baseline = changed.branches["baseline"] + for simulation in (changed, previous_baseline, baseline): + assert BEHAVIORAL_RESPONSE_CACHE_ATTR not in simulation.__dict__ assert changed.baseline is baseline assert baseline.baseline is None assert changed.tax_benefit_system is policies["reform"] diff --git a/policyengine_us/tools/period_branch.py b/policyengine_us/tools/period_branch.py index 5e47797cf20..1270a9490ce 100644 --- a/policyengine_us/tools/period_branch.py +++ b/policyengine_us/tools/period_branch.py @@ -61,9 +61,8 @@ def _input_keys(branch: Simulation) -> Set[Tuple[str, str, Period]]: """The (variable, branch name, period) keys ``branch`` reads as inputs. policyengine-core records each key that ``set_input`` stores, whether - from the dataset, the situation or a branch, in one set shared by a - simulation and all its branches. ``branch`` reads its own keys and its - ancestors'. + from the dataset, the situation or a branch. ``branch`` reads its own + keys and its ancestors'. """ visible_branches = set(branch._get_visible_branch_names()) return { @@ -73,6 +72,36 @@ def _input_keys(branch: Simulation) -> Set[Tuple[str, str, Period]]: } +def delete_input_values(holder, period=None, branch_name="default") -> None: + """Delete arrays and forget the inputs whose storage no longer exists. + + Core 3.32.8 keeps deleted inputs in its replay record, making a later + formula result at that key look like supplied input. Forget the missing + slots before a clone can inherit that stale classification. + """ + holder.delete_arrays(period, branch_name) + simulation = holder.simulation + input_keys = getattr(simulation, "_user_input_keys", None) + if input_keys is None: + return + eternal_period = to_period(ETERNITY) + + def still_stored(key): + name, input_branch, input_period = key + if name != holder.variable.name: + return True + storage_period = ( + eternal_period if holder._memory_storage.is_eternal else input_period + ) + return f"{input_branch}:{storage_period}" in holder._memory_storage._arrays or ( + holder._disk_storage is not None + and f"{input_branch}_{storage_period}" in holder._disk_storage._files + ) + + # Older core clones share this set; leave the source's record untouched. + simulation._user_input_keys = {key for key in input_keys if still_stored(key)} + + def drop_inherited_values(branch: Simulation) -> None: """Delete every array ``branch`` holds except inputs. diff --git a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py index 34503291a64..f8daa8b8c89 100644 --- a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py +++ b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py @@ -1,10 +1,11 @@ from policyengine_us.model_api import * +from policyengine_us.spm import BEHAVIORAL_RESPONSE_CACHE_ATTR, clone_spm_system +from policyengine_us.tools.period_branch import drop_inherited_values BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH = "behavioral_response_measurement" BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH = ( "baseline_behavioral_response_measurement" ) -BEHAVIORAL_RESPONSE_CACHE_ATTR = "_behavioral_response_measurements" NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES = ( "employment_income_behavioral_response", "self_employment_income_behavioral_response", @@ -30,13 +31,67 @@ def _copy_behavioral_response_inputs(branch, person, period): def _behavioral_response_cache(simulation): - cache = getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR, None) + """This simulation's measurements, by period. + + Each simulation has its own dictionary: ``SPMSimulationMixin`` gives a + clone an empty one and a branch a copy of its parent's, and drops it with + the rest of the cached formula output (``_invalidate_all_caches``). + """ + cache = simulation.__dict__.get(BEHAVIORAL_RESPONSE_CACHE_ATTR) if cache is None: cache = {} - setattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR, cache) + simulation.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR] = cache return cache +def _baseline_measurement_branch(simulation): + """Branch ``simulation`` under baseline policy, keeping only its inputs. + + The baseline measurement applies baseline policy to the inputs of the + simulation being measured, as the reform measurement applies the + simulation's own policy to them. ``simulation.baseline`` has the baseline + policy, but its inputs need not be the simulation's: + + - A branch shares its parent's baseline, and on policyengine-core before + #587 so does a clone, so their baseline holds the parent's inputs. + - A reform simulation's baseline is a branch taken while it was built, + before this package moved ``employment_income`` and the other + pre-response inputs to their ``_before_lsr`` variables, so a household + baseline keeps them as overrides; later inputs never reach it. + + ``simulation.get_branch("baseline")``, used before, was worse: for a branch + or an early clone it made a new branch under the *reform* policy, so the + response was measured against the reform itself and came out zero. + + So branch the simulation itself and give the branch a private copy of the + baseline policy, as core does for a reform simulation's own baseline. Then + drop every array the branch copied except inputs: the rest was calculated + under the reform. Like that baseline, the branch has no baseline of its own, + so its behavioural responses are zero. + + A simulation with no baseline (one built without a reform, or a reform + simulation's own baseline) has no reform to respond to: its own policy is + the baseline, so both measurements match and every change is zero, as + before. The labour-supply response returns zero before measuring, but + ``income_elasticity_lsr``, ``substitution_elasticity_lsr`` and the + ``relative_*_change`` variables measure without that check. + """ + baseline = simulation.baseline if simulation.baseline is not None else simulation + branch = simulation.get_branch(BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH) + branch.tax_benefit_system = clone_spm_system(baseline.tax_benefit_system) + branch.baseline = None + # The copied holders were built for reform policy; its entities or + # ETERNITY storage may differ from the baseline's, as after subsampling. + branch._drop_holders_built_for_another_policy() + branch._rebind_holders() + drop_inherited_values(branch) + branch._isolate_parameter_tracing() + # The policy copy's receipts list the counties the baseline read; record + # the ones this branch reads, which are the simulation's. + branch._record_own_county_input_types() + return branch + + def get_behavioral_response_measurements(person, period): # pragma: no cover # Requires reform scenario with simulation branching - tested via microsim simulation = person.simulation @@ -46,18 +101,14 @@ def get_behavioral_response_measurements(person, period): # pragma: no cover if period_key in cache: return cache[period_key] - measurement_branch = simulation.get_branch( - BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, clone_system=True - ) - baseline_parent = simulation.get_branch("baseline") - baseline_branch = baseline_parent.get_branch( - BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, clone_system=True - ) - baseline_branch.tax_benefit_system.parameters.simulation = ( - measurement_branch.tax_benefit_system.parameters.simulation - ) - try: + measurement_branch = simulation.get_branch( + BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, clone_system=True + ) + baseline_branch = _baseline_measurement_branch(simulation) + baseline_branch.tax_benefit_system.parameters.simulation = ( + measurement_branch.tax_benefit_system.parameters.simulation + ) for branch in (measurement_branch, baseline_branch): _neutralize_behavioral_responses(branch) _copy_behavioral_response_inputs(branch, person, period) @@ -85,9 +136,7 @@ def get_behavioral_response_measurements(person, period): # pragma: no cover simulation.macro_cache_write = False return measurements finally: - baseline_parent.branches.pop( - BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, None - ) + simulation.branches.pop(BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, None) simulation.branches.pop(BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, None) diff --git a/policyengine_us/variables/gov/simulation/capital_gains_responses.py b/policyengine_us/variables/gov/simulation/capital_gains_responses.py index bc9bfc51d2c..bf6b0373c2a 100644 --- a/policyengine_us/variables/gov/simulation/capital_gains_responses.py +++ b/policyengine_us/variables/gov/simulation/capital_gains_responses.py @@ -104,6 +104,10 @@ def formula(person, period, parameters): # pragma: no cover simulation = person.simulation DELTA = 1_000 adult_index_values = person("adult_index_cg", period) + # Before branching, so a reform simulation has measured its behavioural + # responses and each branch starts from those measurements, as + # marginal_tax_rate's do; a branch made first measured its own. + household_net_income = person.household("household_net_income", period) inputs = user_set_variables(simulation) for adult_index in [1, 2]: alt_simulation = simulation.get_branch(f"adult_{adult_index}_cg_rise") @@ -118,7 +122,6 @@ def formula(person, period, parameters): # pragma: no cover person("capital_gains", period) + mask * DELTA, ) alt_person = alt_simulation.person - household_net_income = person.household("household_net_income", period) household_net_income_higher_earnings = alt_person.household( "household_net_income", period )