Make branch-measured marginal tax rates match fresh re-simulation

changelog.d/fix-mtr-branch-divergence.fixed.md

@@ -0,0 +1,9 @@
+Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise, and they hold a reform's behavioral responses at their simulated values. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. A reformed simulation's baseline branch now has the moved inputs too, so the baseline rates that labor supply responses compare against no longer erase pre-response wages.
+
+Restoring a rebuilt baseline's policy also rebinds its holders, removes reform-only inputs, and preserves tracing isolation, so neutralized reform variables cannot corrupt baseline comparisons.
+
+Behavioral-response measurements also neutralize hours responses, avoiding recursion when capital gains are queried before labor supply with both elasticities enabled.
+
+Inputs explicitly set after construction, including identification inputs on a measuring branch, now remain inputs in its raised branch. This prevents earnings raises from restoring child or ACA premium tax credits denied in the base simulation, and also preserves branch inputs in health-inclusive rates and behavioral-response measurements.
+- Recompute calculated values when the same variable is explicitly supplied in another period, preserving the supplied branch inputs.
+- Choose aggregate versus pre-response earnings inputs for the period being measured, so a future-year aggregate input cannot change the current year's marginal rate or hide its earnings raise from TANF.

policyengine_us/spm.py

@@ -1051,6 +1051,25 @@ def set_input(self, variable_name, period, value):
                 "Use primitive inputs and retain observed outputs under report-only names."
             )
         result = super().set_input(variable_name, period, value)
+        # Core snapshots input_variables during construction, but later
+        # explicit inputs must also survive counterfactual branches. Keep
+        # only stored inputs; neutralized and expired writes are ignored.
+        input_variables = getattr(self, "input_variables", None)
+        if input_variables is not None and variable_name not in input_variables:
+            holder = self.get_holder(variable_name)
+            variable = holder.variable
+            input_period = periods_.period(period)
+            if (
+                not variable.is_neutralized
+                and (variable.end is None or input_period.start.date <= variable.end)
+                and any(
+                    branch_name == self.branch_name
+                    for branch_name, _ in holder.get_known_branch_periods()
+                )
+            ):
+                # Core shallow-copies metadata when branching. Assign a new
+                # list so this input cannot change another branch's registry.
+                self.input_variables = [*input_variables, variable_name]
         if variable_name == "county_fips":
             self._record_county_input_types(period)
         return result

policyengine_us/system.py

@@ -22,6 +22,9 @@
     backdate_parameters,
 )
 from policyengine_us.reforms import create_structural_reforms_from_parameters
+from policyengine_us.variables.gov.simulation.behavioral_response_measurements import (
+    PRE_RESPONSE_INPUTS,
+)
 from policyengine_core.parameters.operations.homogenize_parameters import (
     homogenize_parameter_structures,
 )
@@ -189,6 +192,78 @@ def clone(self):
 system = CountryTaxBenefitSystem()
 
 
+def _move_inputs_to_pre_response_variables(simulation):
+    """Store aggregate inputs on the pre-response variables they derive from."""
+    for aggregate, pre_response in PRE_RESPONSE_INPUTS.items():
+        holder = simulation.get_holder(aggregate)
+        for known_period in holder.get_known_periods():
+            array = holder.get_array(known_period)
+            simulation.set_input(pre_response, known_period, array)
+            delete_input_values(holder, known_period)
+
+
+def _sync_input_variables(simulation):
+    """List the variables that hold inputs after the moves and backfill.
+
+    Core records ``input_variables`` before the pre-response moves and the
+    ``state_code`` backfill run, so the recorded list still names the emptied
+    aggregates and ``state_code_str`` and omits the variables that now hold
+    their values. Branch-based variables such as ``marginal_tax_rate`` delete
+    every variable outside this list, so a stale list erases those inputs in
+    the branch: a wage-earning parent's TANF then counts no earnings.
+    """
+    moved = (
+        *PRE_RESPONSE_INPUTS,
+        *PRE_RESPONSE_INPUTS.values(),
+        "state_code_str",
+        "state_code",
+    )
+
+    def holds_input(variable):
+        return len(simulation.get_holder(variable).get_known_periods()) > 0
+
+    input_variables = [
+        variable
+        for variable in simulation.input_variables
+        if variable not in moved or holds_input(variable)
+    ]
+    simulation.input_variables = input_variables + [
+        variable
+        for variable in moved
+        if variable not in input_variables and holds_input(variable)
+    ]
+
+
+def _rebuild_baseline_branch(simulation):
+    """Branch a reformed simulation's baseline again, after the moves.
+
+    Core branches the baseline while it constructs a reformed simulation,
+    before the wrappers move inputs and backfill ``state_code``. The branch
+    therefore kept ``employment_income`` and ``weekly_hours_worked`` as
+    inputs, left ``employment_income_before_lsr`` and
+    ``weekly_hours_worked_before_lsr`` empty, and kept the stale
+    ``input_variables``. Its TANF counted no earnings, and the baseline
+    marginal tax rate that labor supply responses compare against erased
+    pre-response wages. Branching again from the simulation, as core's
+    ``subsample`` does, gives the baseline the same inputs as the reform.
+    """
+    baseline = simulation.branches.pop("baseline", None)
+    if baseline is None:
+        return
+    # A branch copies its parent's attributes, and a baseline whose own
+    # baseline is set would compute behavioral responses.
+    simulation.baseline = None
+    rebuilt = simulation.get_branch("baseline")
+    rebuilt.trace = simulation.trace
+    rebuilt.tracer = simulation.tracer
+    rebuilt.tax_benefit_system = baseline.tax_benefit_system
+    # Cloning bound holders to reform variables. Restoring policy must also
+    # restore those bindings and remove inputs that exist only in the reform.
+    rebuilt._rebind_holders()
+    rebuilt._isolate_parameter_tracing()
+    simulation.baseline = rebuilt
+
+
 def _backfill_state_code_from_str(simulation):
     """Backfill the ``state_code`` enum from a ``state_code_str`` input.
 
@@ -259,46 +334,11 @@ def __init__(self, *args, **kwargs):
         if reform is not None:
             self.apply_reform(reform)
 
-        # Labor supply responses
-
-        employment_income = self.get_holder("employment_income")
-        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)
-            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)
-            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():
-            array = sstb_self_employment_income.get_array(known_period)
-            self.set_input(
-                "sstb_self_employment_income_before_lsr", known_period, array
-            )
-            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)
-            delete_input_values(weekly_hours, known_period)
-
-        # Capital gains responses
-
-        cg_holder = self.get_holder("long_term_capital_gains")
-        for known_period in cg_holder.get_known_periods():
-            array = cg_holder.get_array(known_period)
-            self.set_input(
-                "long_term_capital_gains_before_response", known_period, array
-            )
-            delete_input_values(cg_holder, known_period)
-
+        _move_inputs_to_pre_response_variables(self)
         # Geography backfill: state_code_str-only input -> state_code enum.
         _backfill_state_code_from_str(self)
+        _sync_input_variables(self)
+        _rebuild_baseline_branch(self)
 
 
 def _download_or_explain(dataset_str, download):
@@ -519,67 +559,13 @@ def __init__(self, *args, **kwargs):
         if reform is not None:
             self.apply_reform(reform)
 
-        # Labor supply responses
-
-        employment_income = self.get_holder("employment_income")
-        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)
-            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)
-            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():
-            array = sstb_self_employment_income.get_array(known_period)
-            self.set_input(
-                "sstb_self_employment_income_before_lsr", known_period, array
-            )
-            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)
-            delete_input_values(weekly_hours, known_period)
-
-        # Capital gains responses
-
-        cg_holder = self.get_holder("long_term_capital_gains")
-        for known_period in cg_holder.get_known_periods():
-            array = cg_holder.get_array(known_period)
-            self.set_input(
-                "long_term_capital_gains_before_response", known_period, array
-            )
-            delete_input_values(cg_holder, known_period)
-
+        _move_inputs_to_pre_response_variables(self)
         # Geography backfill: state_code_str-only input -> state_code enum.
         # Datasets carry state_fips, so this only fires for situations that
         # explicitly supply state_code_str.
         _backfill_state_code_from_str(self)
-
-        self.input_variables = [
-            variable
-            for variable in self.input_variables
-            if variable
-            not in [
-                "employment_income",
-                "self_employment_income",
-                "sstb_self_employment_income",
-                "weekly_hours_worked",
-                "capital_gains",
-            ]
-        ] + [
-            "employment_income_before_lsr",
-            "self_employment_income_before_lsr",
-            "sstb_self_employment_income_before_lsr",
-            "weekly_hours_worked_before_lsr",
-            "long_term_capital_gains_before_response",
-        ]
+        _sync_input_variables(self)
+        _rebuild_baseline_branch(self)
 
 
 class IndividualSim(CoreIndividualSim):  # Deprecated

policyengine_us/tests/core/test_behavioral_response_measurement_copies.py

@@ -187,12 +187,21 @@ def test_migrated_input_records_are_private_and_explicit_overrides_survive():
     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.
+        # The baseline is rebuilt after migration, so both arms retain the
+        # same pre-response inputs and forget the emptied aggregate inputs.
         assert (
-            original,
+            migrated,
             source.branch_name,
             period,
         ) in source.baseline._user_input_keys
+        assert (
+            original,
+            source.branch_name,
+            period,
+        ) not in source.baseline._user_input_keys
+        np.testing.assert_array_equal(
+            source.baseline.get_array(migrated, YEAR), [value]
+        )
         np.testing.assert_array_equal(
             source.baseline.calculate(original, YEAR), [value]
         )
@@ -308,9 +317,8 @@ def test_baseline_measurement_branch_is_baseline_policy_over_the_inputs():
 
 
 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.
+    # A reform household's baseline keeps the situation's original earnings,
+    # 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]))