Preserve supplied marital units and infer missing explicit couples

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.

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

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,

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()
+        )

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: