From af63395fee652cafb0708e43f873b0f8bed6a5cc Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 7 Oct 2026 15:55:43 -0400 Subject: [PATCH] Preserve marital inputs and infer missing units from explicit pairs --- .../preserve-marital-unit-inputs.fixed.md | 1 + docs/usage/microsimulation.md | 28 ++++ policyengine_us/system.py | 4 + .../test_marital_unit_inputs.py | 151 ++++++++++++++++++ .../ssi/income/ssi_marital_earned_income.yaml | 73 +++++++++ .../tax_unit/is_tax_unit_spouse.yaml | 123 ++++++++++++++ policyengine_us/tools/marital_units.py | 121 ++++++++++++++ .../tax_unit/is_tax_unit_spouse.py | 8 + pyproject.toml | 6 +- uv.lock | 10 +- 10 files changed, 517 insertions(+), 8 deletions(-) create mode 100644 changelog.d/preserve-marital-unit-inputs.fixed.md create mode 100644 policyengine_us/tests/microsimulation/test_marital_unit_inputs.py create mode 100644 policyengine_us/tools/marital_units.py diff --git a/changelog.d/preserve-marital-unit-inputs.fixed.md b/changelog.d/preserve-marital-unit-inputs.fixed.md new file mode 100644 index 00000000000..107ebe746cf --- /dev/null +++ b/changelog.d/preserve-marital-unit-inputs.fixed.md @@ -0,0 +1 @@ +Preserve supplied marital membership when inferring tax spouses, and construct omitted household marital units from explicit head/spouse pairs with singletons for other people. diff --git a/docs/usage/microsimulation.md b/docs/usage/microsimulation.md index d1559310db7..189223890fd 100644 --- a/docs/usage/microsimulation.md +++ b/docs/usage/microsimulation.md @@ -1,5 +1,33 @@ # Microsimulation +## Marital membership and tax roles + +Dataset `person_marital_unit_id` and `marital_unit_id` columns define marital +membership. Loading the dataset and extending its years preserve those IDs. +Microcosm constructs these units from source spouse pointers; policyengine-us +does not reconstruct them from ages or tax-unit roles. + +For household situations, an explicit `marital_units` mapping is also +authoritative. An empty mapping assigns each person their own unit using +core's existing unallocated-person behavior. A married couple may share a +marital unit while belonging to separate tax units. + +When an entity-form situation omits `marital_units`, the model pairs only an +unambiguous, explicitly supplied `is_tax_unit_head` / `is_tax_unit_spouse` +pair in the same tax unit and household. Their true role inputs must overlap +in period, neither may be input as a dependent or separated, and changing +roles require an explicit marital mapping because population membership is +static. Other people receive singleton units. A singleton synthesized without +relationship inputs does not establish that the person is unmarried. + +The model never uses calculated tax roles to construct marital units. When +marital membership was explicitly supplied, inferred tax spouses must share +the head's marital unit as well as their tax unit. When membership was omitted, +the existing tax-role fallback is retained. Direct tax-role inputs still +override formulas. These construction rules apply to both Python household +simulations and entity-form YAML tests; variable-only tests retain core's +singleton defaults. + The `Microsimulation` class is the primary tool for population-level policy analysis in PolicyEngine US. It combines representative survey microdata with PolicyEngine's tax-benefit model to estimate how policies affect the entire US population. ## Getting started diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 158e6443f43..1ff7390d63c 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.marital_units import infer_missing_marital_units from .tools.per_capita_uprating import ( add_per_capita_parameters_for_parameter_uprating, add_per_capita_uprating, @@ -101,6 +102,9 @@ class CountryTaxBenefitSystem(TaxBenefitSystem): ] modelled_policies = COUNTRY_DIR / "programs.yaml" + def preprocess_situation(self, situation, default_period=None): + return infer_missing_marital_units(situation, default_period) + def __init__( self, reform: tuple | None = None, diff --git a/policyengine_us/tests/microsimulation/test_marital_unit_inputs.py b/policyengine_us/tests/microsimulation/test_marital_unit_inputs.py new file mode 100644 index 00000000000..6617bac9af3 --- /dev/null +++ b/policyengine_us/tests/microsimulation/test_marital_unit_inputs.py @@ -0,0 +1,151 @@ +"""Construction/data contracts that the policy YAML runner cannot assert. + +Policy outcomes live in is_tax_unit_spouse.yaml and +ssi_marital_earned_income.yaml. These tests use one six-person dataset and +one small household to check identifiers, year extension, caller isolation, +and clone provenance. Input-validation tests call the preprocessing helper +directly, without constructing simulations. +""" + +from copy import deepcopy + +import numpy as np +import pandas as pd +import pytest + +from policyengine_us import Microsimulation, Simulation +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.tools.marital_units import infer_missing_marital_units + + +def _explicit_pair(): + return { + "people": { + "head": {"is_tax_unit_head": {"2024": True}}, + "spouse": {"is_tax_unit_spouse": {"2024": True}}, + "child": {"age": {"2024": 10}}, + }, + "tax_units": {"tax_unit": {"members": ["head", "spouse", "child"]}}, + } + + +def test_household_construction_does_not_mutate_input_and_expands_axes(): + situation = _explicit_pair() + situation["axes"] = [ + [ + { + "name": "employment_income", + "period": "2024", + "min": 0, + "max": 10_000, + "count": 2, + } + ] + ] + before = deepcopy(situation) + sim = Simulation(situation=situation) + assert situation == before + assert sim.marital_unit.members_entity_id.tolist() == [0, 0, 1, 2, 2, 3] + assert "marital_unit" not in sim.input_group_entities + clone = sim.get_branch("copy") + assert clone.input_group_entities == sim.input_group_entities + np.testing.assert_array_equal( + clone.marital_unit.members_entity_id, sim.marital_unit.members_entity_id + ) + + +def test_dataset_marital_ids_survive_loading_and_year_extension(): + memberships = np.array([0, 17, 42, 42, 99, 99]) + tax_memberships = np.array([10, 10, 20, 20, 30, 40]) + person = pd.DataFrame( + { + "person_id": [101, 102, 201, 202, 301, 302], + "age": [45, 18, 60, 40, 65, 63], + "person_marital_unit_id": memberships, + "person_tax_unit_id": tax_memberships, + "person_household_id": [1] * 6, + "person_family_id": [1] * 6, + "person_spm_unit_id": [1] * 6, + } + ) + dataset = USSingleYearDataset( + time_period=2024, + person=person, + household=pd.DataFrame({"household_id": [1], "household_weight": [1.0]}), + family=pd.DataFrame({"family_id": [1]}), + spm_unit=pd.DataFrame({"spm_unit_id": [1]}), + tax_unit=pd.DataFrame({"tax_unit_id": [10, 20, 30, 40]}), + marital_unit=pd.DataFrame({"marital_unit_id": [0, 17, 42, 99]}), + ) + original_person = person.copy(deep=True) + sim = Microsimulation(dataset=dataset, dataset_end_year=2025) + assert "marital_unit" in sim.input_group_entities + np.testing.assert_array_equal( + np.asarray(sim.marital_unit.ids)[sim.marital_unit.members_entity_id], + memberships, + ) + for year in (2024, 2025): + np.testing.assert_array_equal( + sim.calculate("person_marital_unit_id", year), memberships + ) + np.testing.assert_array_equal( + sim.calculate("is_tax_unit_spouse", year), + [False, False, False, True, False, False], + ) + pd.testing.assert_frame_equal(person, original_person) + + +@pytest.mark.parametrize( + "marital_units", + [ + {}, + {"reported": {"members": ["head", "child"]}}, + ], +) +def test_explicit_marital_mapping_is_left_untouched(marital_units): + situation = _explicit_pair() + situation["marital_units"] = marital_units + original = deepcopy(situation) + assert infer_missing_marital_units(situation) == original + assert situation["marital_units"] is marital_units + + +def test_no_pairing_across_households_or_disjoint_role_periods(): + situation = _explicit_pair() + situation["households"] = { + "first": {"members": ["head", "child"]}, + "second": {"members": ["spouse"]}, + } + inferred = infer_missing_marital_units(situation) + assert all(len(unit["members"]) == 1 for unit in inferred["marital_units"].values()) + situation = _explicit_pair() + situation["people"]["spouse"]["is_tax_unit_spouse"] = {"2025": True} + inferred = infer_missing_marital_units(situation) + assert all(len(unit["members"]) == 1 for unit in inferred["marital_units"].values()) + + +@pytest.mark.parametrize( + "extra_inputs", + [ + {"is_tax_unit_spouse": {"2024": True, "2025": False}}, + {"is_tax_unit_head": {"2024": True}}, + {"is_tax_unit_dependent": {"2024": True}}, + ], +) +def test_ambiguous_or_changing_roles_require_explicit_membership(extra_inputs): + situation = _explicit_pair() + situation["people"]["spouse"].update(extra_inputs) + with pytest.raises(ValueError, match="Supply marital_units explicitly"): + infer_missing_marital_units(situation) + + +def test_separated_pair_and_partial_tax_units_stay_singletons(): + separated = _explicit_pair() + separated["people"]["spouse"]["is_separated"] = {"2024": True} + partial = _explicit_pair() + partial["tax_units"]["tax_unit"]["members"] = ["head", "child"] + for situation in (separated, partial): + inferred = infer_missing_marital_units(situation) + assert all( + len(unit["members"]) == 1 for unit in inferred["marital_units"].values() + ) diff --git a/policyengine_us/tests/policy/baseline/gov/ssa/ssi/income/ssi_marital_earned_income.yaml b/policyengine_us/tests/policy/baseline/gov/ssa/ssi/income/ssi_marital_earned_income.yaml index 9b45bfaa00d..1f9931697cf 100644 --- a/policyengine_us/tests/policy/baseline/gov/ssa/ssi/income/ssi_marital_earned_income.yaml +++ b/policyengine_us/tests/policy/baseline/gov/ssa/ssi/income/ssi_marital_earned_income.yaml @@ -30,6 +30,34 @@ output: ssi_marital_earned_income: [10, 20] +- name: Omitted marital units infer two explicit couples separately within one household. + period: 2023 + input: + people: + first_head: + is_tax_unit_head: true + ssi_couple_computation_applies: true + ssi_earned_income: 10 + first_spouse: + is_tax_unit_spouse: true + ssi_couple_computation_applies: true + ssi_earned_income: 20 + second_head: + is_tax_unit_head: true + ssi_couple_computation_applies: true + ssi_earned_income: 100 + second_spouse: + is_tax_unit_spouse: true + ssi_couple_computation_applies: true + ssi_earned_income: 200 + tax_units: + first_tax_unit: + members: [first_head, first_spouse] + second_tax_unit: + members: [second_head, second_spouse] + output: + ssi_marital_earned_income: [30, 30, 300, 300] + - name: Earned income is shared if head has income period: 2023 input: @@ -46,6 +74,51 @@ output: ssi_marital_earned_income: [10, 10] +- name: Omitted marital units pair explicit spouses and exclude their dependents. + period: 2023 + input: + people: + head: + age: 40 + is_tax_unit_head: true + ssi_couple_computation_applies: true + ssi_earned_income: 10 + spouse: + age: 38 + is_tax_unit_spouse: true + ssi_couple_computation_applies: true + ssi_earned_income: 20 + child: + age: 10 + is_tax_unit_dependent: true + ssi_couple_computation_applies: true + ssi_earned_income: 100 + parent: + age: 75 + is_tax_unit_dependent: true + ssi_couple_computation_applies: true + ssi_earned_income: 200 + tax_units: + tax_unit: + members: [head, spouse, child, parent] + output: + ssi_marital_earned_income: [30, 30, 100, 200] + +- name: Omitted marital units do not pair adults solely by age. + period: 2023 + input: + people: + older: + age: 45 + ssi_couple_computation_applies: true + ssi_earned_income: 10 + younger: + age: 20 + ssi_couple_computation_applies: true + ssi_earned_income: 20 + output: + ssi_marital_earned_income: [10, 20] + - name: Earned income is shared if spouse has income period: 2023 input: diff --git a/policyengine_us/tests/policy/baseline/household/demographic/tax_unit/is_tax_unit_spouse.yaml b/policyengine_us/tests/policy/baseline/household/demographic/tax_unit/is_tax_unit_spouse.yaml index 2bc08f8d0ee..d6eac9b119d 100644 --- a/policyengine_us/tests/policy/baseline/household/demographic/tax_unit/is_tax_unit_spouse.yaml +++ b/policyengine_us/tests/policy/baseline/household/demographic/tax_unit/is_tax_unit_spouse.yaml @@ -103,3 +103,126 @@ is_tax_unit_head: [true, false, false, false] is_tax_unit_spouse: [false, true, false, false] tax_unit_dependents: 2 + +- name: Supplied singleton marital units keep an adult child out of the spouse role. + period: 2025 + input: + people: + parent: + age: 45 + employment_income: 20_000 + adult_child: + age: 18 + tax_units: + tax_unit: + members: [parent, adult_child] + marital_units: + parent_unit: + members: [parent] + child_unit: + members: [adult_child] + output: + is_tax_unit_head: [true, false] + is_tax_unit_spouse: [false, false] + is_tax_unit_dependent: [false, true] + filing_status: HEAD_OF_HOUSEHOLD + eitc_child_count: 1 + eitc: 4_328 + +- name: The supplied marital partner takes precedence over an older adult in the same tax unit. + period: 2025 + input: + people: + head: + age: 60 + relative: + age: 55 + spouse: + age: 50 + tax_units: + tax_unit: + members: [head, relative, spouse] + marital_units: + couple: + members: [head, spouse] + relative_unit: + members: [relative] + output: + is_tax_unit_spouse: [false, false, true] + is_tax_unit_dependent: [false, true, false] + +- name: Multiple supplied couples remain separate within one household. + period: 2025 + input: + people: + first_head: + age: 60 + first_spouse: + age: 50 + second_head: + age: 40 + second_spouse: + age: 35 + tax_units: + first_tax_unit: + members: [first_head, first_spouse] + second_tax_unit: + members: [second_head, second_spouse] + marital_units: + first_couple: + members: [first_head, first_spouse] + second_couple: + members: [second_head, second_spouse] + output: + is_tax_unit_head: [true, false, true, false] + is_tax_unit_spouse: [false, true, false, true] + +- name: A supplied couple in separate tax units does not acquire a tax unit spouse. + period: 2025 + input: + people: + first: + age: 60 + second: + age: 50 + tax_units: + first_tax_unit: + members: [first] + second_tax_unit: + members: [second] + marital_units: + couple: + members: [first, second] + output: + is_tax_unit_head: [true, true] + is_tax_unit_spouse: [false, false] + +- name: An explicit empty marital mapping means supplied singleton memberships. + period: 2025 + input: + people: + parent: + age: 45 + adult_child: + age: 18 + tax_units: + tax_unit: + members: [parent, adult_child] + marital_units: {} + output: + is_tax_unit_spouse: [false, false] + +- name: Explicit tax role overrides remain authoritative even with supplied marital units. + period: 2025 + input: + people: + first: + is_tax_unit_head: true + second: + is_tax_unit_spouse: true + tax_units: + tax_unit: + members: [first, second] + marital_units: {} + output: + is_tax_unit_spouse: [false, true] diff --git a/policyengine_us/tools/marital_units.py b/policyengine_us/tools/marital_units.py new file mode 100644 index 00000000000..62ee134ab27 --- /dev/null +++ b/policyengine_us/tools/marital_units.py @@ -0,0 +1,121 @@ +"""Construct missing household marital units without replacing supplied data.""" + +from policyengine_core import periods + + +def _true_input_periods(person, variable, default_period, require_constant=False): + values = person.get(variable, {}) + if not isinstance(values, dict): + values = {default_period: values} + supplied = {key: value for key, value in values.items() if value is not None} + if require_constant and True in supplied.values() and False in supplied.values(): + raise ValueError( + f"Cannot infer static marital units from changing {variable} inputs. " + "Supply marital_units explicitly." + ) + return { + periods.period(key) + for key, value in supplied.items() + if value is True and key is not None + } + + +def _membership(situation, plural, people): + groups = situation.get(plural) + if groups is None: + return {name: None for name in people} + # Core assigns unlisted people their own group, including when the + # caller supplies an empty mapping rather than omitting the entity. + result = {name: ("singleton", name) for name in people} + for name, group in groups.items(): + members = group.get("members", []) + if not isinstance(members, list): + members = [members] + for member in members: + result[str(member)] = ("supplied", name) + return result + + +def infer_missing_marital_units(situation, default_period=None): + """Use only explicit head/spouse pairs when marital units are omitted. + + Supplied membership, including an empty mapping, is authoritative. With + no supplied membership, an unambiguous pair must have true head/spouse + inputs for a common period and share both a tax unit and a household. + Other people receive singleton units, without claiming that their + marital status is known. Calculated roles and ages are never consulted. + Core calls this hook on a private copy before building populations. + """ + if situation.get("marital_units") is not None: + return situation + people = situation.get("people", {}) + # Leave malformed entity objects to the builder's standard validation. + if not isinstance(people, dict) or any( + not isinstance(person, dict) for person in people.values() + ): + return situation + for plural in ("tax_units", "households"): + supplied = situation.get(plural) + if supplied is not None and ( + not isinstance(supplied, dict) + or any(not isinstance(group, dict) for group in supplied.values()) + ): + return situation + people = {str(name): person for name, person in people.items()} + tax_units = _membership(situation, "tax_units", people) + households = _membership(situation, "households", people) + groups = {} + for name in people: + groups.setdefault((tax_units[name], households[name]), []).append(name) + + units = {name: {"members": [name]} for name in people} + for members in groups.values(): + heads = { + name: dates + for name in members + if ( + dates := _true_input_periods( + people[name], "is_tax_unit_head", default_period + ) + ) + } + spouses = { + name: dates + for name in members + if ( + dates := _true_input_periods( + people[name], "is_tax_unit_spouse", default_period + ) + ) + } + if not heads or not spouses: + continue + if len(heads) != 1 or len(spouses) != 1 or heads.keys() == spouses.keys(): + raise ValueError( + "Cannot infer an unambiguous head/spouse pair. " + "Supply marital_units explicitly." + ) + head = next(iter(heads)) + spouse = next(iter(spouses)) + if not heads[head] & spouses[spouse]: + continue + _true_input_periods(people[head], "is_tax_unit_head", default_period, True) + _true_input_periods(people[spouse], "is_tax_unit_spouse", default_period, True) + for name in (head, spouse): + if _true_input_periods( + people[name], "is_tax_unit_dependent", default_period + ): + raise ValueError( + "Cannot infer a marital pair from an explicit dependent " + "also marked as head or spouse. Supply marital_units explicitly." + ) + if any( + _true_input_periods(people[name], "is_separated", default_period) + for name in (head, spouse) + ): + continue + units[head]["members"].append(spouse) + del units[spouse] + + situation["marital_units"] = units + return situation diff --git a/policyengine_us/variables/household/demographic/tax_unit/is_tax_unit_spouse.py b/policyengine_us/variables/household/demographic/tax_unit/is_tax_unit_spouse.py index 24ccbf62c4c..73f7d360a15 100644 --- a/policyengine_us/variables/household/demographic/tax_unit/is_tax_unit_spouse.py +++ b/policyengine_us/variables/household/demographic/tax_unit/is_tax_unit_spouse.py @@ -17,5 +17,13 @@ def formula(person, period, parameters): head = person("is_tax_unit_head", period) eligible = candidate & ~head & ~is_separated tax_unit = person.tax_unit + if "marital_unit" in person.simulation.input_group_entities: + # Compare the actual population memberships, not the ID variable: + # household situations can supply groups without numeric ID inputs. + marital_unit = person.simulation.populations[ + "marital_unit" + ].members_entity_id + head_marital_unit = tax_unit.max(where(head, marital_unit, -1)) + eligible = eligible & (marital_unit == head_marital_unit) age = person("age", period) return person.get_rank(tax_unit, -age, eligible) == 0 diff --git a/pyproject.toml b/pyproject.toml index b53e647d362..bde4caa262b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,8 @@ dependencies = [ # pandas 3 copy-on-write enables the shallow-copy fast path in dataset # extension and is available on every supported Python version. "pandas>=3.0", - "policyengine-core>=3.32.8", + # Draft prerequisite: replace with a released minimum before merging. + "policyengine-core @ git+https://github.com/PolicyEngine/policyengine-core.git@6d839a148cee47981e8289de37da276f3e7beeec", "spm-calculator>=1.0.0,<=1.0.0.post1", "tables>=3.9", "tqdm>=4.67.1", @@ -60,6 +61,9 @@ dev = [ [tool.hatch.build.targets.wheel] packages = ["policyengine_us"] +[tool.hatch.metadata] +allow-direct-references = true + [tool.towncrier] package = "policyengine_us" diff --git a/uv.lock b/uv.lock index d7bfef4956f..f517beed211 100644 --- a/uv.lock +++ b/uv.lock @@ -1791,8 +1791,8 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.8" -source = { registry = "https://pypi.org/simple" } +version = "3.32.21" +source = { git = "https://github.com/PolicyEngine/policyengine-core.git?rev=6d839a148cee47981e8289de37da276f3e7beeec#6d839a148cee47981e8289de37da276f3e7beeec" } dependencies = [ { name = "dpath" }, { name = "h5py" }, @@ -1811,10 +1811,6 @@ dependencies = [ { name = "standard-imghdr" }, { name = "wheel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/df/2fe428d690f02cb109e91d22b8af31ca6dcb047920d7177a2b28487a8ea1/policyengine_core-3.32.8.tar.gz", hash = "sha256:925a8990985ec22885f4d6c4ed129cd13dcc2eab02cf9a130939ca563394fcc3", size = 426369, upload-time = "2026-09-27T19:10:50.428Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1a/45/bbba25be7ddc52ad9bb6f2b8ea044eb3c670f4168c67044328f631ccb025/policyengine_core-3.32.8-py3-none-any.whl", hash = "sha256:84e470d1f03d531d1d5ee77c173c59444bbbdea50132faf9150a82aa00d54397", size = 246996, upload-time = "2026-09-27T19:10:48.871Z" }, -] [[package]] name = "policyengine-us" @@ -1850,7 +1846,7 @@ requires-dist = [ { name = "jupyter-book", marker = "extra == 'dev'", specifier = ">=1.0.4.post1" }, { name = "microdf-python", specifier = ">=1.0.0" }, { name = "pandas", specifier = ">=3.0" }, - { name = "policyengine-core", specifier = ">=3.32.8" }, + { name = "policyengine-core", git = "https://github.com/PolicyEngine/policyengine-core.git?rev=6d839a148cee47981e8289de37da276f3e7beeec" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, { name = "setuptools", marker = "extra == 'dev'", specifier = ">=80.9.0" }, { name = "spm-calculator", specifier = ">=1.0.0,<=1.0.0.post1" },