diff --git a/changelog.d/fix-county-branch-carry-over.fixed.md b/changelog.d/fix-county-branch-carry-over.fixed.md new file mode 100644 index 00000000000..2bd90bcb572 --- /dev/null +++ b/changelog.d/fix-county-branch-carry-over.fixed.md @@ -0,0 +1 @@ +County and has_tin read stored values with the asking branch's name, so a branch reused for a later year keeps the household's county instead of falling to County.UNKNOWN, and a branch's own has_itin input reaches has_tin. diff --git a/policyengine_us/tests/core/test_has_tin_branch_inputs.py b/policyengine_us/tests/core/test_has_tin_branch_inputs.py new file mode 100644 index 00000000000..d847da9eee7 --- /dev/null +++ b/policyengine_us/tests/core/test_has_tin_branch_inputs.py @@ -0,0 +1,163 @@ +"""``has_tin`` reads its inputs as the simulation asking for it sees them. + +A branch stores its inputs under its own name. The formula read the default +branch's, so a ``has_tin`` or ``has_itin`` (the deprecated alias) input set +on a branch, or inherited by a nested branch from its parent, was ignored and +``has_tin`` came from the default branch's inputs instead. + +The contract: a canonical ``has_tin`` input wins over a legacy ``has_itin`` +input, which wins over the default ``True``. Each is read as the asking +simulation sees it: its own value, then its nearest ancestor's, then the +default branch's. A sibling branch's inputs are invisible. + +``calculate`` returns a canonical ``has_tin`` input the simulation can read +without running the formula, so the canonical-input tests call the formula +itself. +""" + +import numpy as np +import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st +from policyengine_core.periods import period as as_period + +from policyengine_us import Simulation +from policyengine_us.variables.household.demographic.person.has_tin import has_tin + +YEAR = 2026 + + +def _situation() -> dict: + return { + "people": {"adult": {"age": {YEAR: 30}}}, + "tax_units": {"tax_unit": {"members": ["adult"]}}, + "households": {"household": {"members": ["adult"], "state_code": {YEAR: "CA"}}}, + } + + +def test_a_branch_reads_its_own_legacy_has_itin_input(): + simulation = Simulation(situation=_situation()) + branch = simulation.get_branch("without_tin") + branch.set_input("has_itin", YEAR, np.array([False])) + + assert branch.calculate("has_tin", YEAR).tolist() == [False] + # The input stays in the branch. + assert simulation.calculate("has_tin", YEAR).tolist() == [True] + + +def test_a_nested_branch_reads_its_parents_legacy_has_itin_input(): + simulation = Simulation(situation=_situation()) + branch = simulation.get_branch("without_tin") + branch.set_input("has_itin", YEAR, np.array([False])) + nested = branch.get_branch("nested") + + assert nested.calculate("has_tin", YEAR).tolist() == [False] + + +def test_a_branch_still_reads_the_default_branchs_legacy_input(): + situation = _situation() + situation["people"]["adult"]["has_itin"] = {YEAR: False} + simulation = Simulation(situation=situation) + branch = simulation.get_branch("unchanged") + + assert branch.calculate("has_tin", YEAR).tolist() == [False] + assert simulation.calculate("has_tin", YEAR).tolist() == [False] + + +@pytest.mark.parametrize("nested", [False, True]) +def test_formula_reads_branch_or_ancestor_canonical_has_tin(nested): + simulation = Simulation(situation=_situation()) + simulation.set_input("has_tin", YEAR, np.array([True])) + simulation.set_input("has_itin", YEAR, np.array([True])) + branch = simulation.get_branch("without_tin") + branch.set_input("has_tin", YEAR, np.array([False])) + target = branch.get_branch("nested") if nested else branch + sibling = simulation.get_branch("unrelated") + + # The branch's canonical False wins over the legacy True it inherits, and + # a nested branch inherits it; the simulation and an unrelated sibling + # keep the default branch's True. + for current, expected in ( + (target, [False]), + (simulation, [True]), + (sibling, [True]), + ): + actual = has_tin.formula(current.persons, as_period(YEAR), None) + assert actual.tolist() == expected + + +# A simulation tree: the simulation's own (has_tin, has_itin) inputs, then +# each branch as (index of the simulation it branches from, has_tin input, +# has_itin input), where index 0 is the simulation and branch i is index +# i + 1; ``None`` sets no input. Then the order to calculate in. +INPUT = st.sampled_from([None, True, False]) + + +@st.composite +def simulation_trees(draw): + root = (draw(INPUT), draw(INPUT)) + branches = [ + (draw(st.integers(0, i)), draw(INPUT), draw(INPUT)) + for i in range(draw(st.integers(0, 5))) + ] + order = draw(st.permutations(range(len(branches) + 1))) + return root, branches, order + + +def _expected_has_tin(chain_inputs) -> bool: + """The contract, from (has_tin, has_itin) inputs ordered asker first.""" + for position in (0, 1): # has_tin, then has_itin + for inputs in chain_inputs: + if inputs[position] is not None: + return inputs[position] + return True + + +def _set_inputs(simulation, inputs): + for variable, value in zip(("has_tin", "has_itin"), inputs): + if value is not None: + simulation.set_input(variable, YEAR, np.array([value])) + + +@settings( + max_examples=12, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given(simulation_trees()) +# The canonical case above: a branch's has_tin input, a nested branch and an +# unrelated sibling. +@example( + ((True, True), [(0, False, None), (1, None, None), (0, None, None)], [2, 0, 3, 1]) +) +# A branch's legacy input, inherited by a nested branch. +@example(((None, None), [(0, None, False), (1, None, None)], [0, 1, 2])) +# A nearer legacy input loses to an ancestor's canonical one. +@example(((True, None), [(0, None, False), (1, None, None)], [2, 1, 0])) +def test_has_tin_follows_the_branch_read_contract(tree): + root_inputs, branches, order = tree + simulation = Simulation(situation=_situation()) + _set_inputs(simulation, root_inputs) + # Each simulation's inputs are set before anything branches from it, as a + # formula sets a branch's override right after creating the branch. + simulations, parents, inputs = [simulation], [None], [root_inputs] + for i, (parent, *branch_inputs) in enumerate(branches): + branch = simulations[parent].get_branch(f"branch_{i}") + _set_inputs(branch, branch_inputs) + simulations.append(branch) + parents.append(parent) + inputs.append(tuple(branch_inputs)) + + def expected(index): + chain = [] + while index is not None: + chain.append(inputs[index]) + index = parents[index] + return _expected_has_tin(chain) + + for index in order: + formula = has_tin.formula(simulations[index].persons, as_period(YEAR), None) + assert formula.tolist() == [expected(index)], ("formula", index) + for index in order: + calculated = simulations[index].calculate("has_tin", YEAR) + assert calculated.tolist() == [expected(index)], ("calculate", index) diff --git a/policyengine_us/tests/microsimulation/test_county_branch_reuse.py b/policyengine_us/tests/microsimulation/test_county_branch_reuse.py new file mode 100644 index 00000000000..db840fa72a7 --- /dev/null +++ b/policyengine_us/tests/microsimulation/test_county_branch_reuse.py @@ -0,0 +1,252 @@ +"""County must stay the household's county in branches reused across years. + +``tax_unit_itemizes`` branches the simulation into ``itemizing`` and +``not_itemizing`` (each with a nested ``no_salt``), and later years reuse +those branches. A branch stores what it calculates under its own name, so +the county formula's stored-county read has to use the branch name: it used +to read the default branch's value for the latest known period, found +nothing for a year only the branch had calculated, and returned ``None``. +Core then fell back to its own carry-over. Carrying over only inputs +(PolicyEngine/policyengine-core#562) left those branches at +``County.UNKNOWN``, which raised ``ParameterNotFoundError`` for Maryland +child care rates in a 2025-2027 run over the enhanced CPS; carrying over the +latest calculated value (core master) left a branch at ``UNKNOWN`` whenever +it had calculated a later year first. + +These pass on both core versions. The order tests enumerate every order of +three requests; the Hypothesis test draws longer request sequences on trees +that grow as they are asked. +""" + +from itertools import permutations + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st + +from policyengine_core.periods import period as as_period +from policyengine_us import Microsimulation +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.variables.household.demographic.geographic.county.county_enum import ( + County, +) + +DATASET_YEAR = 2024 +COUNTY_FIPS = ["24031", "36047", "06037"] +# None of these is its state's first county, so a fall back to +# first_county_in_state shows up as a wrong county. +EXPECTED = [ + "MONTGOMERY_COUNTY_MD", + "KINGS_COUNTY_NY", + "LOS_ANGELES_COUNTY_CA", +] +YEARS = (2025, 2026, 2027) + + +def _dataset(county_fips: list) -> USSingleYearDataset: + n = len(county_fips) + ids = np.arange(1, n + 1) + person = pd.DataFrame( + { + "person_id": ids, + "person_household_id": ids, + "person_tax_unit_id": ids, + "person_spm_unit_id": ids, + "person_family_id": ids, + "person_marital_unit_id": ids, + "age": np.full(n, 40.0), + "employment_income": np.full(n, 100_000.0), + } + ) + household = pd.DataFrame( + { + "household_id": ids, + "state_fips": np.array([int(fips[:2]) for fips in COUNTY_FIPS]), + "county_fips": county_fips, + "household_weight": np.full(n, 1.0), + } + ) + return USSingleYearDataset( + person=person, + household=household, + tax_unit=pd.DataFrame({"tax_unit_id": ids}), + spm_unit=pd.DataFrame({"spm_unit_id": ids}), + family=pd.DataFrame({"family_id": ids}), + marital_unit=pd.DataFrame({"marital_unit_id": ids}), + time_period=DATASET_YEAR, + ) + + +def _simulation(stored: str) -> Microsimulation: + """A dataset storing either ``county_fips`` or one year's ``county``. + + The loader extends a dataset's columns to later years, so a ``county`` + column would be stored for every year and never reach the formula. + Storing the county for the dataset's year only, with no county FIPS, + makes later years carry it forward (falling back would give each state's + first county instead). + """ + if stored == "county_fips": + return Microsimulation(dataset=_dataset(COUNTY_FIPS)) + simulation = Microsimulation(dataset=_dataset([""] * len(COUNTY_FIPS))) + simulation.set_input("county", DATASET_YEAR, County.encode(np.array(EXPECTED))) + return simulation + + +def _counties(simulation, year) -> list: + return simulation.calculate("county_str", year).tolist() + + +@pytest.mark.parametrize("stored", ["county_fips", "county"]) +@pytest.mark.parametrize("years", list(permutations(YEARS)), ids=str) +def test_branch_reused_across_years_keeps_the_county(stored, years): + simulation = _simulation(stored) + itemizing = simulation.get_branch("itemizing") + no_salt = itemizing.get_branch("no_salt") + + for year in years: + for label, branch in (("itemizing", itemizing), ("no_salt", no_salt)): + assert _counties(branch, year) == EXPECTED, (label, year) + for year in years: + assert _counties(simulation, year) == EXPECTED, ("default", year) + + +@pytest.mark.parametrize("stored", ["county_fips", "county"]) +def test_branch_created_in_one_year_keeps_the_county_in_later_years(stored): + """The enhanced CPS order: each year the simulation, then its branches. + + The ``itemizing`` branch is created after the simulation has calculated + the first year's county, calculates the second year itself, and then + holds no default-branch value for it when the third year asks. + """ + simulation = _simulation(stored) + branches = {} + for year in YEARS: + assert _counties(simulation, year) == EXPECTED, ("default", year) + if not branches: + branches["itemizing"] = simulation.get_branch("itemizing") + branches["no_salt"] = branches["itemizing"].get_branch("no_salt") + for label, branch in branches.items(): + assert _counties(branch, year) == EXPECTED, (label, year) + + +@pytest.mark.parametrize("stored", ["county_fips", "county"]) +@pytest.mark.parametrize( + "requests", + [ + # Every order of these three year-and-simulation requests, with the + # branches created before any of them. + list(order) + for order in permutations( + [("default", 2025), ("itemizing", 2026), ("no_salt", 2027)] + ) + ] + + [ + list(order) + for order in permutations( + [("itemizing", 2027), ("default", 2026), ("no_salt", 2025)] + ) + ], + ids=str, +) +def test_county_does_not_depend_on_the_order_of_requests(stored, requests): + simulation = _simulation(stored) + itemizing = simulation.get_branch("itemizing") + simulations = { + "default": simulation, + "itemizing": itemizing, + "no_salt": itemizing.get_branch("no_salt"), + } + for label, year in requests: + assert _counties(simulations[label], year) == EXPECTED, (label, year) + # Every simulation then reads the same county for every year. + for label, each in simulations.items(): + for year in YEARS: + assert _counties(each, year) == EXPECTED, (label, year) + + +# Each step either branches from a simulation already in the tree or asks one +# of them for a year's county; index 0 is the simulation and branch i is +# index i + 1, so a branch can be created after its parent has calculated. +REQUEST_YEARS = range(DATASET_YEAR, DATASET_YEAR + 5) +MAX_BRANCHES = 4 + + +@st.composite +def request_sequences(draw): + stored = draw(st.sampled_from(["county_fips", "county"])) + steps = [] + size = 1 + for _ in range(draw(st.integers(1, 10))): + if size <= MAX_BRANCHES and draw(st.booleans()): + steps.append(("branch", draw(st.integers(0, size - 1)))) + size += 1 + else: + index = draw(st.integers(0, size - 1)) + steps.append(("ask", index, draw(st.sampled_from(REQUEST_YEARS)))) + return stored, steps + + +@settings( + max_examples=8, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given(request_sequences()) +# The enhanced CPS order: the simulation's first year, then branches created +# from it that calculate the next year and are reused for the one after. +@example( + ( + "county", + [ + ("ask", 0, 2025), + ("branch", 0), + ("branch", 1), + ("ask", 1, 2026), + ("ask", 2, 2026), + ("ask", 0, 2027), + ("ask", 1, 2027), + ("ask", 2, 2027), + ], + ) +) +def test_county_stays_the_households_county_whatever_the_requests(case): + stored, steps = case + simulations = [_simulation(stored)] + for step in steps: + if step[0] == "branch": + parent = simulations[step[1]] + simulations.append(parent.get_branch(f"branch_{len(simulations)}")) + else: + _, index, year = step + assert _counties(simulations[index], year) == EXPECTED, step + for index, simulation in enumerate(simulations): + for year in REQUEST_YEARS: + assert _counties(simulation, year) == EXPECTED, (index, year) + + +def test_a_period_only_another_branch_can_read_is_skipped(): + """With nothing readable, county comes from ``county_fips``.""" + simulation = _simulation("county_fips") + holder = simulation.get_holder("county") + holder.put_in_cache( + County.encode(np.array(["ALBANY_COUNTY_NY"] * len(COUNTY_FIPS))), + as_period(2030), + "sibling", + ) + + assert _counties(simulation, 2026) == EXPECTED + + +def test_the_latest_period_the_branch_can_read_wins(): + simulation = _simulation("county") + branch = simulation.get_branch("moved") + moved = ["BALTIMORE_COUNTY_MD", "QUEENS_COUNTY_NY", "SAN_DIEGO_COUNTY_CA"] + branch.set_input("county", 2026, County.encode(np.array(moved))) + + # The branch reads its own 2026 county; the simulation it was branched + # from cannot, and keeps the dataset's. + assert _counties(branch, 2027) == moved + assert _counties(simulation, 2027) == EXPECTED diff --git a/policyengine_us/tests/policy/baseline/household/demographic/person/has_tin.yaml b/policyengine_us/tests/policy/baseline/household/demographic/person/has_tin.yaml index 363a9b836bd..df3ad2ab364 100644 --- a/policyengine_us/tests/policy/baseline/household/demographic/person/has_tin.yaml +++ b/policyengine_us/tests/policy/baseline/household/demographic/person/has_tin.yaml @@ -25,3 +25,15 @@ has_itin: true output: has_tin: true + +- name: Canonical TIN branch-precedence seed + period: 2026 + input: + people: + adult: {age: 30, has_tin: true, has_itin: true} + tax_units: + tax_unit: {members: [adult]} + households: + household: {members: [adult], state_code: CA} + output: + has_tin: [true] diff --git a/policyengine_us/variables/household/demographic/geographic/county/county.py b/policyengine_us/variables/household/demographic/geographic/county/county.py index c15394eceb5..cebabf37a9e 100644 --- a/policyengine_us/variables/household/demographic/geographic/county/county.py +++ b/policyengine_us/variables/household/demographic/geographic/county/county.py @@ -9,6 +9,33 @@ ) +def latest_readable_county(simulation: Simulation): + """The county this simulation can read for its latest known period. + + The holder lists the periods of every array in its storage, whichever + branch stored them. A branch reused for a later year holds its own values + for earlier years, which only a read with its branch name finds, so read + each period as this simulation would (its own value, then its ancestors', + then the default branch's), latest start first. Returns ``None`` when no + known period is readable. + """ + holder = simulation.get_holder("county") + known_periods = sorted( + set(holder.get_known_periods()), + key=lambda known_period: ( + known_period.start, + known_period.stop, + str(known_period), + ), + reverse=True, + ) + for known_period in known_periods: + array = holder.get_array(known_period, simulation.branch_name) + if array is not None: + return array + return None + + class county(Variable): value_type = Enum possible_values = County @@ -22,13 +49,10 @@ def formula(household, period, parameters): # When running over a dataset, use stored county data if available # (geographic variables like county are time-invariant for households) - if simulation.is_over_dataset: # pragma: no cover - # Microsimulation-specific path - tested via microsim - holder = simulation.get_holder("county") - known_periods = holder.get_known_periods() - if len(known_periods) > 0: - last_known_period = sorted(known_periods)[-1] - return holder.get_array(last_known_period) + if simulation.is_over_dataset: + stored_county = latest_readable_county(simulation) + if stored_county is not None: + return stored_county # No stored county: fall through to the county_fips mapping, so # datasets that store county_fips compute counties from it # instead of collapsing to first_county_in_state. diff --git a/policyengine_us/variables/household/demographic/person/has_tin.py b/policyengine_us/variables/household/demographic/person/has_tin.py index 2dd744b99d8..31689a58b42 100644 --- a/policyengine_us/variables/household/demographic/person/has_tin.py +++ b/policyengine_us/variables/household/demographic/person/has_tin.py @@ -11,17 +11,21 @@ class has_tin(Variable): def formula(person, period, parameters): simulation = person.simulation + # Read inputs as this simulation sees them: in a branch, its own + # value, then its ancestors', then the default branch's. + branch_name = simulation.branch_name + # Canonical path: allow direct `has_tin` inputs to override the formula. holder = simulation.get_holder("has_tin") if period in holder.get_known_periods(): - array = holder.get_array(period) + array = holder.get_array(period, branch_name) if array is not None: return array # Temporary migration path: honor legacy `has_itin` inputs until callers move. legacy_holder = simulation.get_holder("has_itin") if period in legacy_holder.get_known_periods(): - array = legacy_holder.get_array(period) + array = legacy_holder.get_array(period, branch_name) if array is not None: return array