Make two formulas that read stored values independent of calculation order

changelog.d/fix-order-dependent-stored-reads-medicaid.fixed.md

@@ -0,0 +1 @@
+The Medicaid enrollment that SSI state supplements read (medicaid_enrolled_for_ssi_state_supplement) now reads medicaid_enrolled directly only when it was supplied, so whether medicaid_enrolled had already been calculated no longer changes it.

changelog.d/fix-order-dependent-stored-reads.fixed.md

@@ -0,0 +1 @@
+Without Form 2555 leaf inputs, the foreign earned income exclusion now uprates only from the latest earlier entered amount, as Core does, so a later year no longer depends on whether earlier years were calculated first, and an amount stored only under a branch the simulation cannot read no longer raises a TypeError.

policyengine_us/tests/core/test_ssi_state_supplement_medicaid_dependency.py

@@ -2,6 +2,9 @@
 import pytest
 
 from policyengine_us import Simulation
+
+hypothesis = pytest.importorskip("hypothesis")
+st = pytest.importorskip("hypothesis.strategies")
 from policyengine_us.variables.household.demographic.geographic.state_code import (
     StateCode,
 )
@@ -77,6 +80,129 @@ def test_ssp_preserves_medicaid_inputs_set_after_initialization(
     assert simulation.calculate(variable, 2030).tolist() == [False, True]
 
 
+SSP = {"IN": ("in_ssp", "in_ssp_sapn_eligible"), "KS": ("ks_sspp", "ks_sspp_eligible")}
+SSP_MEDICAID = "medicaid_enrolled_for_ssi_state_supplement"
+PASS_THROUGH = "medicaid_community_engagement_pass_through_eligible"
+# Variables a caller may calculate before the SSP's Medicaid enrollment, each
+# with the period it is asked for.
+CALCULATED_FIRST = {
+    "medicaid_enrolled": 2027,
+    "is_medicaid_eligible": 2027,
+    PASS_THROUGH: "2027-01",
+    "snap": "2027-01",
+    "in_ssp": "2027-01",
+    "ks_sspp": "2027-01",
+}
+
+
+def supplied_situation(state, supplied, year=2027):
+    """``situation`` with one value on the SNAP-SSP-Medicaid cycle supplied.
+
+    Supplying SNAP or the pass-through means calculating medicaid_enrolled no
+    longer runs the SSP on the way, so medicaid_enrolled can be stored before
+    the SSP's Medicaid enrollment is first asked for.
+    """
+    data = situation(year, state)
+    if supplied == "snap":
+        data["spm_units"]["recipient"]["snap"] = {year: 0}
+        data["spm_units"]["adult"]["snap"] = {year: 1_200}
+    elif supplied == "pass_through":
+        data["people"]["adult"][PASS_THROUGH] = {f"{year}-01": True}
+    elif supplied == "medicaid_enrolled":
+        data["people"]["recipient"]["medicaid_enrolled"] = {year: True}
+        data["people"]["adult"]["medicaid_enrolled"] = {year: True}
+    return data
+
+
+@hypothesis.settings(max_examples=6, deadline=None, derandomize=True)
+# The reproducing case: with SNAP supplied, the SNAP pass-through adult's SSP
+# enrollment was False asked alone, but True once medicaid_enrolled, which
+# keeps the pass-through, had been calculated.
+@hypothesis.example(state="IN", supplied="snap", first=["medicaid_enrolled"])
+@hypothesis.given(
+    state=st.sampled_from(["IN", "KS"]),
+    supplied=st.sampled_from([None, "snap", "pass_through", "medicaid_enrolled"]),
+    first=st.lists(st.sampled_from(list(CALCULATED_FIRST)), unique=True, max_size=4),
+)
+def test_ssp_medicaid_enrollment_does_not_depend_on_calculation_order(
+    state, supplied, first
+):
+    benefit, eligibility = SSP[state]
+    fresh = Simulation(situation=supplied_situation(state, supplied))
+    simulation = Simulation(situation=supplied_situation(state, supplied))
+    for name in first:
+        simulation.calculate(name, CALCULATED_FIRST[name])
+    for name, year in ((SSP_MEDICAID, 2027), (eligibility, "2027-01"), (benefit, 2027)):
+        np.testing.assert_array_equal(
+            simulation.calculate(name, year), fresh.calculate(name, year)
+        )
+
+
+@pytest.mark.parametrize("route", ["snap", "tanf"])
+@pytest.mark.parametrize("first", [(), ("medicaid_enrolled",)])
+def test_pass_through_adult_drops_the_exclusion_in_either_order(route, first):
+    """The adult is enrolled in Medicaid only through the SNAP or TANF
+    exclusion from the community engagement requirement. SNAP is supplied, so
+    medicaid_enrolled can be stored first. The SSP's enrollment drops the
+    exclusion whether or not medicaid_enrolled was calculated first."""
+    data = supplied_situation("IN", "snap")
+    if route == "tanf":
+        data["people"]["adult"].pop("is_snap_workfare_participant")
+        data["spm_units"]["adult"]["snap"] = {2027: 0}
+        data["spm_units"]["adult"]["receives_tanf"] = {2027: True}
+        data["spm_units"]["adult"]["meets_tanf_work_requirements"] = {2027: True}
+    simulation = Simulation(situation=data)
+    for name in first:
+        simulation.calculate(name, CALCULATED_FIRST[name])
+    assert simulation.calculate(SSP_MEDICAID, 2027).tolist() == [True, False]
+    assert simulation.calculate("medicaid_enrolled", 2027).tolist() == [True, True]
+
+
+@pytest.mark.parametrize("first", [(), ("medicaid_enrolled",)])
+def test_ssi_recipient_in_the_adult_group_gets_the_same_sapn_in_either_order(first):
+    """A reform that ends Indiana's Medicaid coverage of SSI recipients puts
+    this 40-year-old, institutionalized SSI recipient with no disability
+    finding in the adult group, where the community engagement requirement
+    applies; ``is_ssi_recipient_for_medicaid: False`` does the same here. The
+    recipient is exempt only through the supplied SNAP pass-through, which
+    the SSP's enrollment drops. SAPN was $22 when medicaid_enrolled had been
+    calculated first and $0 otherwise; it is now $0 in either order."""
+    year = 2027
+    people = {
+        "recipient": {
+            "age": {year: 40},
+            "receives_ssi": {year: True},
+            "is_ssi_recipient_for_medicaid": {year: False},
+            "ssi_lives_in_medical_treatment_facility": {year: True},
+            "ssi_medicaid_pays_majority_of_care": {year: True},
+            PASS_THROUGH: {f"{year}-01": True},
+        }
+    }
+    unit = {"unit": {"members": ["recipient"]}}
+    simulation = Simulation(
+        situation={
+            "people": people,
+            **{e: unit for e in ("tax_units", "spm_units", "families")},
+            "households": {
+                "unit": {"members": ["recipient"], "state_code": {year: "IN"}}
+            },
+        }
+    )
+    for name in first:
+        simulation.calculate(name, CALCULATED_FIRST[name])
+    # Ask for the SSP before medicaid_enrolled, so that with nothing
+    # calculated first this is the order in which the SSP comes first.
+    assert simulation.calculate(SSP_MEDICAID, year).tolist() == [False]
+    assert simulation.calculate("in_ssp_sapn_eligible", f"{year}-01").tolist() == [
+        False
+    ]
+    assert simulation.calculate("in_ssp", year).tolist() == [0]
+    assert simulation.calculate("medicaid_category", year).decode_to_str().tolist() == [
+        "ADULT"
+    ]
+    assert simulation.calculate("medicaid_enrolled", year).tolist() == [True]
+
+
 @pytest.mark.parametrize("year", [2027, 2030, 2035])
 def test_future_poverty_in_all_states(year):
     """Exercise state benefit dependencies together, as a national dataset does."""

policyengine_us/tests/test_foreign_earned_income_exclusion_leaves.py

@@ -198,18 +198,25 @@ def test_set_input_marks_a_leaf_explicit_after_its_default_was_cached(entered):
     )
 
 
-def expected_legacy_carry(simulation, entered):
-    # The original pure dollar input used default uprating, then the system
-    # converted that national total to its per-capita series. Adding a
-    # formula must preserve the projected value.
+def uprated(simulation, amounts, entered, year):
+    """The amounts entered for one year, uprated straight to ``year``.
+
+    The original pure dollar input used default uprating, then the system
+    converted that national total to its per-capita series. Adding a formula
+    must preserve the projected value.
+    """
     parameters = simulation.tax_benefit_system.parameters
-    before = parameters(
-        str(YEAR)
-    ).calibration.gov.cbo.income_by_source.adjusted_gross_income_per_capita
-    after = parameters(
-        str(YEAR + 1)
-    ).calibration.gov.cbo.income_by_source.adjusted_gross_income_per_capita
-    return np.array([entered], dtype=np.float32) * (after / before)
+
+    def index(index_year):
+        return parameters(
+            str(index_year)
+        ).calibration.gov.cbo.income_by_source.adjusted_gross_income_per_capita
+
+    return np.array(amounts, dtype=np.float32) * (index(year) / index(entered))
+
+
+def expected_legacy_carry(simulation, entered):
+    return uprated(simulation, [entered], YEAR, YEAR + 1)
 
 
 @pytest.mark.parametrize("legacy", [0, 10_000, -5_000])
@@ -267,6 +274,103 @@ def test_future_leaf_inputs_do_not_replace_legacy_carry():
     )
 
 
+@st.composite
+def legacy_calculation_orders(draw):
+    """Amounts entered for one year, a later year to ask for, and the years
+    calculated first, in the order drawn. Some may follow the target year."""
+    entered = draw(st.integers(2015, 2030))
+    target = draw(st.integers(entered + 1, entered + 6))
+    first = draw(
+        st.lists(st.integers(entered + 1, target + 2), unique=True, max_size=5)
+    )
+    amounts = draw(
+        st.lists(
+            st.floats(-1e6, 1e7, allow_nan=False, width=32), min_size=1, max_size=4
+        )
+    )
+    return amounts, entered, target, first
+
+
+@hypothesis.settings(max_examples=6, deadline=None, derandomize=True)
+# The reproducing case: 10,000 entered for 2024 gave 11,090.326171875 for 2026
+# asked alone, but 11,090.3271484375 after 2025, which was uprated from the
+# float32 result calculated for 2025 rather than from the input.
+@hypothesis.example(([10_000], 2024, 2026, [2025]))
+@hypothesis.given(legacy_calculation_orders())
+def test_legacy_uprating_does_not_depend_on_calculation_order(case):
+    amounts, entered, target, first = case
+
+    def build():
+        return simulation_for_tax_units(
+            [
+                {"foreign_earned_income_exclusion": {entered: amount}}
+                for amount in amounts
+            ]
+        )
+
+    fresh = build().calculate("foreign_earned_income_exclusion", target)
+    simulation = build()
+    for year in first:
+        simulation.calculate("foreign_earned_income_exclusion", year)
+    after = simulation.calculate("foreign_earned_income_exclusion", target)
+    # Bit for bit: the drift was in the last bits of a float32.
+    np.testing.assert_array_equal(after.view(np.uint32), fresh.view(np.uint32))
+    np.testing.assert_array_equal(after, uprated(simulation, amounts, entered, target))
+
+
+@pytest.mark.parametrize("order", [(YEAR + 1, YEAR + 3), (YEAR + 3, YEAR + 1)])
+def test_legacy_uprating_starts_from_the_latest_earlier_entered_year(order):
+    # Each year uprates from the latest year entered before it, whichever of
+    # the other years were calculated first.
+    simulation = simulation_for_tax_units(
+        [{"foreign_earned_income_exclusion": {YEAR: 10_000, YEAR + 2: 7_000}}]
+    )
+    values = {
+        year: simulation.calculate("foreign_earned_income_exclusion", year)
+        for year in order
+    }
+    np.testing.assert_array_equal(
+        values[YEAR + 1], uprated(simulation, [10_000], YEAR, YEAR + 1)
+    )
+    np.testing.assert_array_equal(
+        values[YEAR + 3], uprated(simulation, [7_000], YEAR + 2, YEAR + 3)
+    )
+
+
+@pytest.mark.parametrize("first", [(), (2013, 2015)])
+def test_legacy_amount_is_held_flat_until_the_index_starts(first):
+    # The per-capita AGI index starts in 2015. An amount entered for 2012 is
+    # carried unchanged until then and uprated from the index's first value
+    # after it, whichever earlier years were calculated first.
+    simulation = simulation_for_tax_units(
+        [{"foreign_earned_income_exclusion": {2012: 12_345.67}}]
+    )
+    for year in first:
+        simulation.calculate("foreign_earned_income_exclusion", year)
+    np.testing.assert_array_equal(
+        simulation.calculate("foreign_earned_income_exclusion", 2014),
+        np.array([12_345.67], dtype=np.float32),
+    )
+    np.testing.assert_array_equal(
+        simulation.calculate("foreign_earned_income_exclusion", 2016),
+        uprated(simulation, [12_345.67], 2015, 2016),
+    )
+
+
+def test_legacy_uprating_ignores_an_amount_only_another_branch_holds():
+    # An amount stored only under a branch this simulation cannot read is not
+    # one of its inputs, so the year after it takes the legacy default. When
+    # any stored year could be the source, that amount was chosen, read back
+    # as None and raised TypeError.
+    simulation = simulation_for_tax_units([{}])
+    simulation.get_holder("foreign_earned_income_exclusion").set_input(
+        period(YEAR), np.array([10_000], dtype=np.float32), "sibling"
+    )
+    np.testing.assert_array_equal(
+        simulation.calculate("foreign_earned_income_exclusion", YEAR + 1), [0]
+    )
+
+
 @hypothesis.settings(max_examples=3, deadline=None, derandomize=True)
 @hypothesis.given(st.integers(1, 150_000), st.integers(0, 150_000))
 def test_form_2555_filers_skip_worksheet_b_even_when_stacking_is_zero(amount, excess):

policyengine_us/variables/gov/hhs/medicaid/medicaid_enrolled_for_ssi_state_supplement.py

@@ -26,9 +26,15 @@ def formula(person, period, parameters):
             return person("medicaid_enrolled", period)
 
         simulation = person.simulation
-        enrolled = simulation.get_array("medicaid_enrolled", period)
-        if enrolled is not None:
-            return enrolled
+        # Read enrollment directly only when it was supplied. A calculated
+        # value keeps the SNAP/TANF exclusion, and whether one is stored yet
+        # depends on what was calculated first, so reading it made this
+        # variable order-dependent: it then differed for anyone who relies on
+        # the exclusion, including an SSI recipient whom a reform moves into
+        # the adult group.
+        holder = simulation.get_holder("medicaid_enrolled")
+        if period in holder.get_input_periods(simulation.branch_name):
+            return holder.get_array(period, simulation.branch_name)
 
         # SNAP counts these state supplements as income, so asking for receipt
         # status here would recurse. SSP candidates qualify through SSI or

policyengine_us/variables/gov/irs/income/taxable_income/adjusted_gross_income/above_the_line_deductions/foreign_earned_income.py

@@ -1,6 +1,9 @@
 from policyengine_us.model_api import *
 from policyengine_core.parameters import get_parameter
-from policyengine_core.simulations.simulation import _uprating_index_value
+from policyengine_core.simulations.simulation import (
+    _latest_input_key,
+    _uprating_index_value,
+)
 from policyengine_us.tools.default_uprating import DEFAULT_DOLLAR_INPUT_UPRATING
 from policyengine_us.tools.section_911 import has_section_911_leaf_inputs
 
@@ -22,17 +25,26 @@ def formula(tax_unit, period, parameters):
             # A formula cannot also declare Core uprating. Preserve the
             # old dollar input's per-capita default uprating here instead,
             # including signed inputs and Core's held-flat index boundaries.
+            # As Core does, uprate only from an entered input this branch
+            # reads. A year this simulation calculated was itself uprated
+            # and rounded to float32, so uprating from it would make the
+            # result depend on which years were calculated first.
             simulation = tax_unit.simulation
             holder = simulation.get_holder("foreign_earned_income_exclusion")
-            earlier_periods = [
-                known_period
-                for known_period in holder.get_known_periods()
-                if known_period.unit == period.unit
-                and known_period.start < period.start
+            earlier_inputs = [
+                input_period
+                for input_period in holder.get_input_periods(simulation.branch_name)
+                if input_period.unit == period.unit
+                and input_period.start < period.start
             ]
-            if not earlier_periods:
+            if not earlier_inputs:
                 return None
-            latest = max(earlier_periods, key=lambda known_period: known_period.start)
+            # Core's tie-break: of two inputs starting the same year, the one
+            # that ends last, so the source never depends on storage order.
+            latest = max(
+                earlier_inputs,
+                key=lambda input_period: _latest_input_key(input_period, period.unit),
+            )
             index = get_parameter(
                 simulation.tax_benefit_system.parameters,
                 DEFAULT_DOLLAR_INPUT_UPRATING + "_per_capita",