From bbb5d65181d5b97f385c5512dc725632ca95bd52 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 4 Nov 2026 21:21:46 -0500 Subject: [PATCH 1/3] Let US calculate_household take lists of groups calculate_household put every person in one marital unit, family, SPM unit and tax unit. policyengine-us defines a marital unit as one unmarried person or one married couple, and its SSI and Medicaid rules read a two-person marital unit as a couple and sum some incomes across it. So a single parent with one child got the couple's Medicaid income limit, and a teenager's earnings counted as their married parents' SSI income. marital_unit, family, spm_unit and tax_unit now also accept a list of groups, each a mapping with members (indexes into people) plus the group's inputs. Every person must belong to exactly one group of the entity; a list that leaves someone out, names someone twice or names an index outside people raises before any calculation. An entity passed as a list comes back as a list of results in the same order, with one axis series per group under axes. A mapping still builds one group holding everyone, so existing calls return the same results. Fixes #738 Co-Authored-By: Claude Opus 5.5 --- changelog.d/738.added.md | 1 + docs/households.md | 61 ++- .../tax_benefit_models/us/household.py | 296 ++++++++---- tests/test_us_household_groups.py | 440 ++++++++++++++++++ 4 files changed, 711 insertions(+), 87 deletions(-) create mode 100644 changelog.d/738.added.md create mode 100644 tests/test_us_household_groups.py diff --git a/changelog.d/738.added.md b/changelog.d/738.added.md new file mode 100644 index 000000000..85a7c2cfd --- /dev/null +++ b/changelog.d/738.added.md @@ -0,0 +1 @@ +`pe.us.calculate_household` accepts a list of groups for `marital_unit`, `family`, `spm_unit` and `tax_unit`. Each group is a mapping with `members`, a list of indexes into `people`, plus its own inputs, and each person belongs to exactly one group of the entity. An entity passed as a list comes back as a list of results in the same order, with one axis series per group under axes. A mapping still puts every person in one group, so existing calls return the same results. Callers can now give a household with children or other adults its own marital units: policyengine-us reads a two-person marital unit as a married couple in its SSI and Medicaid rules, so one marital unit for everyone gave a single parent with one child the couple's Medicaid income limit and counted a child's earnings as their married parents' SSI income. diff --git a/docs/households.md b/docs/households.md index 789db687e..9bc7804fb 100644 --- a/docs/households.md +++ b/docs/households.md @@ -14,6 +14,8 @@ result = pe.us.calculate_household( {"age": 8}, {"age": 5}, ], + # The married couple, then each child: see "Several groups of one entity". + marital_unit=[{"members": [0, 1]}, {"members": [2]}, {"members": [3]}], tax_unit={"filing_status": "JOINT"}, household={"state_code": "TX", "county_fips": "48201"}, year=2026, @@ -25,13 +27,59 @@ result = pe.us.calculate_household( | Argument | Purpose | |---|---| | `people` | List of person dicts. Keys are any person-level variable on the model. | -| `tax_unit` | Tax-unit inputs (e.g. `filing_status`). | -| `spm_unit` | SPM-unit inputs. | +| `tax_unit` | Tax-unit inputs (e.g. `filing_status`), or a list of tax units. | +| `spm_unit` | SPM-unit inputs, or a list of SPM units. | | `household` | Household inputs, including `state_code` and observed five-digit `county_fips` for the default SPM geography selection. | -| `family` | Family-level inputs. | -| `marital_unit` | Marital-unit inputs. | +| `family` | Family-level inputs, or a list of families. | +| `marital_unit` | Marital-unit inputs, or a list of marital units. | -All adults default to one shared tax unit and household. For separate tax units (e.g. two adult roommates), construct the `Simulation` directly and set the entity-membership arrays. +A mapping puts every person in one group of that entity, and the household always holds everyone. + +### Several groups of one entity + +To split people into several marital units, families, SPM units or tax units, pass a list of groups. Each group is a mapping with `members`, a list of indexes into `people`, plus that group's own inputs. Each person belongs to exactly one group of the entity: `calculate_household` raises if a list leaves someone out, names someone twice or names an index outside `people`. + +policyengine-us defines a marital unit as one unmarried person or one married couple. Its SSI and Medicaid rules read a two-person marital unit as a couple and add up some incomes across it, so pass marital units whenever the household holds anyone besides one person or one couple: + +```python +result = pe.us.calculate_household( + people=[ + {"age": 68}, + {"age": 66}, + {"age": 17, "employment_income": 6_000}, + ], + marital_unit=[{"members": [0, 1]}, {"members": [2]}], + tax_unit={"filing_status": "JOINT"}, + household={"state_code": "CA"}, + year=2026, + spm={"geography_kind": "national"}, + extra_variables=["ssi"], +) +result.marital_unit[1] # the 17-year-old's marital unit +``` + +With one marital unit for all three, the couple's SSI would count the 17-year-old's earnings as their own income, and Medicaid's optional senior-or-disabled pathway would hold the couple to the individual income limit. + +A list of tax units covers, for example, an adult child who files their own return: + +```python +result = pe.us.calculate_household( + people=[ + {"age": 52, "employment_income": 60_000}, + {"age": 22, "employment_income": 18_000}, + ], + marital_unit=[{"members": [0]}, {"members": [1]}], + tax_unit=[ + {"members": [0], "filing_status": "SINGLE"}, + {"members": [1], "filing_status": "SINGLE"}, + ], + household={"state_code": "NY", "county_fips": "36061"}, + year=2026, +) +result.tax_unit[1].income_tax # the adult child's return +``` + +An entity passed as a list comes back as a list of results in the same order. Under axes, each group's variables are axis series, as each person's are. ### SPM geography and measurement selection @@ -232,10 +280,11 @@ each variable on that person is its own axis series. result.person[0].income_tax # first person result.person[2].age # third person result.tax_unit.income_tax # single tax unit +result.marital_unit[1] # second marital unit, when passed as a list result.household.household_net_income # single household ``` -The result is a Pydantic model — `.model_dump()` gives you a dict, individual sections are regular attribute lookups. +The result is a `dict` with attribute access: `.to_dict()` gives a plain dict, `.write(path)` writes it as JSON, and individual sections are regular attribute lookups. ## Errors diff --git a/src/policyengine/tax_benefit_models/us/household.py b/src/policyengine/tax_benefit_models/us/household.py index 17f828a42..ba224aa28 100644 --- a/src/policyengine/tax_benefit_models/us/household.py +++ b/src/policyengine/tax_benefit_models/us/household.py @@ -9,36 +9,35 @@ import policyengine as pe # Single parent with one child in New York, $45k wages. - result = pe.us.calculate_household( + inputs = dict( people=[ {"age": 32, "employment_income": 45000, "is_tax_unit_head": True}, {"age": 6, "is_tax_unit_dependent": True}, ], + # A marital unit is one unmarried person or one married couple. + marital_unit=[{"members": [0]}, {"members": [1]}], tax_unit={"filing_status": "HEAD_OF_HOUSEHOLD"}, household={"state_code": "NY", "county_fips": "36061"}, year=2026, - extra_variables=["adjusted_gross_income"], + ) + result = pe.us.calculate_household( + **inputs, extra_variables=["adjusted_gross_income"] ) print(result.tax_unit.income_tax) print(result.tax_unit.ctc, result.tax_unit.eitc) print(result.household.household_net_income) + print(result.marital_unit[0]) # one result per marital unit # Reform: zero out SNAP. reformed = pe.us.calculate_household( - people=[ - {"age": 32, "employment_income": 45000, "is_tax_unit_head": True}, - {"age": 6, "is_tax_unit_dependent": True}, - ], - tax_unit={"filing_status": "HEAD_OF_HOUSEHOLD"}, - household={"state_code": "NY", "county_fips": "36061"}, - year=2026, - reform={"gov.usda.snap.income.deductions.earned_income": 0}, + **inputs, reform={"gov.usda.snap.income.deductions.earned_income": 0} ) """ from __future__ import annotations -from collections.abc import Mapping -from typing import Any, Optional +import operator +from collections.abc import Mapping, Sequence +from typing import Any, Optional, Union from policyengine.tax_benefit_models.common import ( EntityResult, @@ -55,6 +54,19 @@ from .spm import SPMSelection, calculation_provenance, resolve_spm_selection _GROUP_ENTITIES = ("marital_unit", "family", "spm_unit", "tax_unit", "household") +# Group entities a caller may split into several groups. The household stays +# one group: ``calculate_household`` computes a single household. +_SPLITTABLE_GROUP_ENTITIES = ("marital_unit", "family", "spm_unit", "tax_unit") +_GROUP_PLURALS = { + "marital_unit": "marital_units", + "family": "families", + "spm_unit": "spm_units", + "tax_unit": "tax_units", + "household": "households", +} + +GroupInput = Union[Mapping[str, Any], Sequence[Mapping[str, Any]]] +"""One group's inputs, or a list of groups that each name their ``members``.""" def _raise_unexpected_kwargs(unexpected: Mapping[str, Any]) -> None: @@ -96,14 +108,123 @@ def _safe_convert(value: Any) -> Any: return str(value) if value is not None else None +def _person_index(value: Any, *, location: str, people_count: int) -> int: + if isinstance(value, bool): + index = None + else: + try: + index = operator.index(value) + except TypeError: + index = None + if index is None: + raise ValueError( + f"{location} must hold integer indexes into people; got {value!r}." + ) + if not 0 <= index < people_count: + raise ValueError( + f"{location} names person {index}, but people has " + f"{people_count} entries (indexes 0 to {people_count - 1})." + ) + return index + + +def _resolve_groups( + entity: str, + spec: Optional[GroupInput], + *, + people_count: int, +) -> tuple[list[dict[str, Any]], bool]: + """Return one entity's groups, each with ``members`` as person indexes. + + A mapping (or ``None``) is one group holding every person; the flag in the + return value is then ``False``. A list holds several groups, each naming + its ``members``; every person must belong to exactly one of them. + """ + everyone = list(range(people_count)) + if spec is None: + return [{"members": everyone}], False + if isinstance(spec, Mapping): + if "members" in spec: + raise ValueError( + f"{entity} is a single group holding every person, so it " + "takes no 'members'. To split people into several groups, " + f"pass a list: {entity}=[{{'members': [0, 1]}}, " + "{'members': [2]}]." + ) + return [{"members": everyone, **spec}], False + if entity not in _SPLITTABLE_GROUP_ENTITIES: + raise ValueError( + f"{entity} takes one mapping of inputs: calculate_household " + "computes a single household." + ) + if isinstance(spec, (str, bytes)) or not isinstance(spec, Sequence): + raise ValueError( + f"{entity} must be a mapping of inputs or a list of groups; " + f"got {type(spec).__name__}." + ) + groups = list(spec) + if not groups: + raise ValueError(f"{entity} must list at least one group.") + + owner: dict[int, int] = {} + resolved: list[dict[str, Any]] = [] + for position, group in enumerate(groups): + location = f"{entity}[{position}]" + if not isinstance(group, Mapping): + raise ValueError( + f"{location} must be a mapping with 'members' and the group's " + f"inputs; got {type(group).__name__}." + ) + members = group.get("members") + if ( + members is None + or isinstance(members, (str, bytes, Mapping)) + or not isinstance(members, Sequence) + or not members + ): + raise ValueError( + f"{location} must include 'members', a non-empty list of " + "indexes into people." + ) + indexes = [] + for value in members: + index = _person_index( + value, + location=f"{location}['members']", + people_count=people_count, + ) + if index in owner: + where = ( + "twice" + if owner[index] == position + else f"in both {entity}[{owner[index]}] and {location}" + ) + raise ValueError( + f"people[{index}] appears {where}; each person belongs to " + f"exactly one {entity}." + ) + owner[index] = position + indexes.append(index) + resolved.append({**group, "members": indexes}) + + unassigned = [index for index in everyone if index not in owner] + if unassigned: + listed = ", ".join(f"people[{index}]" for index in unassigned) + raise ValueError( + f"{listed} belong to no {entity}; each person belongs to exactly " + f"one {entity}." + ) + return resolved, True + + +def _group_inputs(group: Mapping[str, Any]) -> dict[str, Any]: + return {key: value for key, value in group.items() if key != "members"} + + def _build_situation( *, people: list[Mapping[str, Any]], - marital_unit: Mapping[str, Any], - family: Mapping[str, Any], - spm_unit: Mapping[str, Any], - tax_unit: Mapping[str, Any], - household: Mapping[str, Any], + groups: Mapping[str, list[dict[str, Any]]], year: int, axes: Optional[list[Any]] = None, ) -> dict[str, Any]: @@ -116,16 +237,17 @@ def _periodise(spec: Mapping[str, Any]) -> dict[str, dict[str, Any]]: persons = {pid: _periodise(person) for pid, person in zip(person_ids, people)} def _group(spec: Mapping[str, Any]) -> dict[str, Any]: - return {"members": list(person_ids), **_periodise(spec)} - - situation = { - "people": persons, - "marital_units": {"marital_unit_0": _group(marital_unit)}, - "families": {"family_0": _group(family)}, - "spm_units": {"spm_unit_0": _group(spm_unit)}, - "tax_units": {"tax_unit_0": _group(tax_unit)}, - "households": {"household_0": _group(household)}, - } + return { + "members": [person_ids[index] for index in spec["members"]], + **_periodise(_group_inputs(spec)), + } + + situation: dict[str, Any] = {"people": persons} + for entity in _GROUP_ENTITIES: + situation[_GROUP_PLURALS[entity]] = { + f"{entity}_{position}": _group(spec) + for position, spec in enumerate(groups[entity]) + } if axes is not None: situation["axes"] = axes return situation @@ -151,10 +273,10 @@ def _group(spec: Mapping[str, Any]) -> dict[str, Any]: def calculate_household( *, people: list[Mapping[str, Any]], - marital_unit: Optional[Mapping[str, Any]] = None, - family: Optional[Mapping[str, Any]] = None, - spm_unit: Optional[Mapping[str, Any]] = None, - tax_unit: Optional[Mapping[str, Any]] = None, + marital_unit: Optional[GroupInput] = None, + family: Optional[GroupInput] = None, + spm_unit: Optional[GroupInput] = None, + tax_unit: Optional[GroupInput] = None, household: Optional[Mapping[str, Any]] = None, year: int = 2026, reform: Optional[Mapping[str, Any]] = None, @@ -172,7 +294,17 @@ def calculate_household( marital_unit, family, spm_unit, tax_unit, household: Optional per-entity overrides, each keyed by variable name (e.g. ``tax_unit={"filing_status": "SINGLE"}``, - ``household={"state_code": "NY"}``). + ``household={"state_code": "NY"}``). A mapping applies to one + group holding every person. ``marital_unit``, ``family``, + ``spm_unit`` and ``tax_unit`` also accept a list of groups, each a + mapping with ``members`` (indexes into ``people``) plus that + group's inputs; each person must belong to exactly one group of + the entity. policyengine-us defines a marital unit as one + unmarried person or one married couple, and its SSI and Medicaid + rules read a two-person marital unit as a couple, so a household + with children or other adults needs the list form, e.g. + ``marital_unit=[{"members": [0, 1]}, {"members": [2]}]`` for a + married couple and their child. year: Calendar year to compute for. Defaults to 2026. reform: Optional reform as ``{parameter_path: value}`` or ``{parameter_path: {effective_date: value}}``. Scalar @@ -211,8 +343,10 @@ def calculate_household( Returns: :class:`HouseholdResult` with dot-accessible per-entity - variables. Singleton entities (``tax_unit``, ``household``, ...) - return :class:`EntityResult`; ``person`` returns a list of them. + variables. An entity passed as a mapping (or omitted) returns one + :class:`EntityResult`; an entity passed as a list of groups returns + a list of them in the same order, and ``person`` returns one per + person. Raises: ValueError: if any input dict uses an unknown variable name, @@ -221,39 +355,47 @@ def calculate_household( / ``reform`` names a variable or parameter path not defined on the US model. Raises if ``year`` is not an annual calendar year or if household input values are already periodized. + Raises if a list of groups leaves a person out, names a + person twice, or names an index outside ``people``. """ if unexpected: _raise_unexpected_kwargs(unexpected) people = list(people) - entities = { - "marital_unit": dict(marital_unit or {}), - "family": dict(family or {}), - "spm_unit": dict(spm_unit or {}), - "tax_unit": dict(tax_unit or {}), - "household": dict(household or {}), + if not people: + raise ValueError( + "people must be a non-empty list. At minimum pass people=[{'age': }]." + ) + specs: dict[str, Optional[GroupInput]] = { + "marital_unit": marital_unit, + "family": family, + "spm_unit": spm_unit, + "tax_unit": tax_unit, + "household": household, + } + groups: dict[str, list[dict[str, Any]]] = {} + listed: dict[str, bool] = {} + for name, spec in specs.items(): + groups[name], listed[name] = _resolve_groups( + name, spec, people_count=len(people) + ) + records = { + name: [_group_inputs(group) for group in entity_groups] + for name, entity_groups in groups.items() } year = validate_annual_household_inputs( year=year, - entities={ - "people": people, - **{name: [value] for name, value in entities.items()}, - }, + entities={"people": people, **records}, ) from policyengine_us import Simulation from spm_calculator.policyengine_adapter import validate_policyengine_inputs - validate_policyengine_inputs( - {"person": people, **{name: [value] for name, value in entities.items()}} - ) + validate_policyengine_inputs({"person": people, **records}) validate_household_input( model_version=us_latest, - entities={ - "person": people, - **{name: [value] for name, value in entities.items()}, - }, + entities={"person": people, **records}, ) extra_by_entity = dispatch_extra_variables( @@ -277,11 +419,7 @@ def calculate_household( simulation = Simulation( situation=_build_situation( people=people, - marital_unit=entities["marital_unit"], - family=entities["family"], - spm_unit=entities["spm_unit"], - tax_unit=entities["tax_unit"], - household=entities["household"], + groups=groups, year=year, axes=normalized_axes, ), @@ -291,43 +429,39 @@ def calculate_household( result = HouseholdResult() for entity, columns in output_columns.items(): - raw = { - variable: list( - simulation.calculate( + values = { + variable: [ + _safe_convert(value) + for value in simulation.calculate( variable, period=year, map_to=entity, decode_enums=True, ) - ) + ] for variable in columns } - if entity == "person": - result["person"] = [ - EntityResult( - { - variable: values_for_entity( - [_safe_convert(value) for value in raw[variable]], - entity_index=i, - entity_count=len(people), - axes_active=axes_active, - ) - for variable in columns - } - ) - for i in range(len(people)) - ] - else: - result[entity] = EntityResult( + count = len(people) if entity == "person" else len(groups[entity]) + # Core lays each entity's instances out copy by copy across the axis + # grid, so instance i of this entity sits at i, i + count, ... + entity_results = [ + EntityResult( { - variable: ( - [_safe_convert(value) for value in raw[variable]] - if axes_active - else _safe_convert(raw[variable][0]) + variable: values_for_entity( + values[variable], + entity_index=index, + entity_count=count, + axes_active=axes_active, ) for variable in columns } ) + for index in range(count) + ] + if entity == "person" or listed[entity]: + result[entity] = entity_results + else: + result[entity] = entity_results[0] result["provenance"] = { "spm_config": dict(simulation.spm_config), "spm": calculation_provenance(simulation), diff --git a/tests/test_us_household_groups.py b/tests/test_us_household_groups.py new file mode 100644 index 000000000..34d83abc4 --- /dev/null +++ b/tests/test_us_household_groups.py @@ -0,0 +1,440 @@ +"""Several groups of one entity in ``pe.us.calculate_household`` (#738). + +``marital_unit``, ``family``, ``spm_unit`` and ``tax_unit`` accept either one +mapping (one group holding every person) or a list of groups that each name +their ``members`` by index into ``people``. + +Invariants checked here: + +- Partition: a list of groups is accepted exactly when it assigns every person + to one group, and the situation's memberships equal that partition. +- Backward compatibility: a mapping (or nothing) builds the same single group + holding every person as before, under the same ids. +- Differential: list-form results equal a ``policyengine_us.Simulation`` + built by hand with the same groups, with and without axes. +- Semantics: with its own marital unit, a child's earnings leave its married + parents' SSI unchanged, and a single parent faces the individual Medicaid + income limit. + +US cases select national SPM geography: they test the calculator interface +without claiming a local SPM area. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import policyengine as pe +from policyengine.tax_benefit_models.common import EntityResult +from policyengine.tax_benefit_models.us.household import ( + _build_situation, + _resolve_groups, +) +from policyengine.tax_benefit_models.us.spm import resolve_spm_selection + +YEAR = 2026 +NATIONAL = {"geography_kind": "national"} +GROUP_PLURALS = { + "marital_unit": "marital_units", + "family": "families", + "spm_unit": "spm_units", + "tax_unit": "tax_units", + "household": "households", +} +SPLITTABLE = ["marital_unit", "family", "spm_unit", "tax_unit"] + + +def _direct_values(people, groups, variables, *, axes=None): + """Run policyengine_us directly on a hand-written situation. + + ``groups`` maps each group entity to a list of member-index lists plus + inputs, e.g. ``{"marital_unit": [([0, 1], {}), ([2], {})]}``. Entities not + listed hold everyone in one group. + """ + from policyengine_us import Simulation + + ids = [f"p{index}" for index in range(len(people))] + situation = { + "people": { + pid: {key: {YEAR: value} for key, value in person.items()} + for pid, person in zip(ids, people) + } + } + for entity, plural in GROUP_PLURALS.items(): + specs = groups.get(entity, [(list(range(len(people))), {})]) + situation[plural] = { + f"{entity}{position}": { + "members": [ids[index] for index in members], + **{key: {YEAR: value} for key, value in inputs.items()}, + } + for position, (members, inputs) in enumerate(specs) + } + if axes is not None: + situation["axes"] = axes + simulation = Simulation(situation=situation, spm=resolve_spm_selection(NATIONAL)) + out = {} + for entity, names in variables.items(): + for name in names: + values = simulation.calculate( + name, period=YEAR, map_to=entity, decode_enums=True + ) + out[(entity, name)] = [ + float(value) if not isinstance(value, str) else value + for value in values + ] + return out + + +# --- Resolving groups --------------------------------------------------------- + + +class TestResolveGroups: + def test__mapping__then_one_group_of_everyone(self): + groups, listed = _resolve_groups( + "tax_unit", {"filing_status": "JOINT"}, people_count=3 + ) + assert listed is False + assert groups == [{"members": [0, 1, 2], "filing_status": "JOINT"}] + + def test__none__then_one_group_of_everyone(self): + groups, listed = _resolve_groups("family", None, people_count=2) + assert listed is False + assert groups == [{"members": [0, 1]}] + + def test__list__then_members_and_inputs_kept_in_order(self): + groups, listed = _resolve_groups( + "tax_unit", + [ + {"members": [2, 0], "filing_status": "JOINT"}, + {"members": (1,), "filing_status": "SINGLE"}, + ], + people_count=3, + ) + assert listed is True + assert groups == [ + {"members": [2, 0], "filing_status": "JOINT"}, + {"members": [1], "filing_status": "SINGLE"}, + ] + + def test__numpy_integer_indexes__then_accepted(self): + groups, _ = _resolve_groups( + "marital_unit", + [{"members": [np.int64(0)]}, {"members": [np.int32(1)]}], + people_count=2, + ) + assert groups == [{"members": [0]}, {"members": [1]}] + + @pytest.mark.parametrize( + ("spec", "match"), + [ + ([{"members": [0]}], r"people\[1\] belong to no marital_unit"), + ([{"members": [0, 0, 1]}], r"people\[0\] appears twice"), + ( + [{"members": [0, 1]}, {"members": [1]}], + r"people\[1\] appears in both marital_unit\[0\] and " + r"marital_unit\[1\]", + ), + ([{"members": [0, 2]}], r"names person 2, but people has 2 entries"), + ([{"members": [-1, 0, 1]}], r"names person -1"), + ([{"members": [0.0, 1]}], r"integer indexes into people; got 0.0"), + ([{"members": ["0", 1]}], r"integer indexes into people; got '0'"), + ([{"members": [True, 0]}], r"integer indexes into people; got True"), + ([{"members": []}, {"members": [0, 1]}], r"marital_unit\[0\] must"), + ([{"filing_status": "JOINT"}], r"marital_unit\[0\] must include"), + ([{"members": "01"}], r"marital_unit\[0\] must include 'members'"), + ([{"members": {0: 1}}], r"marital_unit\[0\] must include 'members'"), + ([], r"marital_unit must list at least one group"), + ([[0, 1]], r"marital_unit\[0\] must be a mapping"), + ("01", r"marital_unit must be a mapping of inputs or a list"), + ({"members": [0, 1]}, r"takes no 'members'"), + ], + ) + def test__invalid_groups__then_raises(self, spec, match): + with pytest.raises(ValueError, match=match): + _resolve_groups("marital_unit", spec, people_count=2) + + def test__household_list__then_raises(self): + with pytest.raises(ValueError, match="computes a single household"): + _resolve_groups("household", [{"members": [0]}], people_count=1) + + +@st.composite +def _partitions(draw): + """People count, plus groups that partition them in a random order.""" + people_count = draw(st.integers(min_value=1, max_value=8)) + labels = draw( + st.lists( + st.integers(min_value=0, max_value=people_count - 1), + min_size=people_count, + max_size=people_count, + ) + ) + blocks: dict[int, list[int]] = {} + for index, label in enumerate(labels): + blocks.setdefault(label, []).append(index) + groups = [draw(st.permutations(members)) for members in blocks.values()] + return people_count, draw(st.permutations(groups)) + + +class TestPartitionProperties: + @settings(max_examples=200, deadline=None) + @given(case=_partitions(), entity=st.sampled_from(SPLITTABLE)) + def test__any_partition__then_memberships_equal_it(self, case, entity): + people_count, partition = case + spec = [{"members": list(members)} for members in partition] + groups = { + name: _resolve_groups(name, None, people_count=people_count)[0] + for name in GROUP_PLURALS + } + groups[entity], listed = _resolve_groups( + entity, spec, people_count=people_count + ) + assert listed is True + situation = _build_situation( + people=[{"age": 30}] * people_count, groups=groups, year=YEAR + ) + built = list(situation[GROUP_PLURALS[entity]].values()) + assert [group["members"] for group in built] == [ + [f"person_{index}" for index in members] for members in partition + ] + everyone = sorted(pid for group in built for pid in group["members"]) + assert everyone == sorted(f"person_{i}" for i in range(people_count)) + + @settings(max_examples=200, deadline=None) + @given(case=_partitions(), data=st.data()) + def test__dropped_or_repeated_person__then_raises(self, case, data): + people_count, partition = case + groups = [list(members) for members in partition] + position = data.draw(st.integers(0, len(groups) - 1)) + if data.draw(st.booleans()): + groups[position].pop(data.draw(st.integers(0, len(groups[position]) - 1))) + groups = [members for members in groups if members] + if not groups: + groups = [[people_count]] # every person dropped: out of range + else: + extra = data.draw(st.integers(0, people_count - 1)) + groups[position].append(extra) + with pytest.raises(ValueError): + _resolve_groups( + "tax_unit", + [{"members": members} for members in groups], + people_count=people_count, + ) + + @settings(max_examples=100, deadline=None) + @given( + people_count=st.integers(min_value=1, max_value=6), + state=st.sampled_from(["CA", "NY", "TX"]), + ) + def test__mapping_form__then_builds_the_single_group_structure( + self, people_count, state + ): + groups = { + name: _resolve_groups( + name, + {"state_code": state} if name == "household" else None, + people_count=people_count, + )[0] + for name in GROUP_PLURALS + } + situation = _build_situation( + people=[{"age": 40}] * people_count, groups=groups, year=YEAR + ) + everyone = [f"person_{i}" for i in range(people_count)] + assert situation == { + "people": {pid: {"age": {str(YEAR): 40}} for pid in everyone}, + "marital_units": {"marital_unit_0": {"members": everyone}}, + "families": {"family_0": {"members": everyone}}, + "spm_units": {"spm_unit_0": {"members": everyone}}, + "tax_units": {"tax_unit_0": {"members": everyone}}, + "households": { + "household_0": { + "members": everyone, + "state_code": {str(YEAR): state}, + } + }, + } + + +# --- Calculating -------------------------------------------------------------- + + +COUPLE_AND_TEEN = dict( + people=[{"age": 68}, {"age": 66}, {"age": 17, "employment_income": 6_000}], + tax_unit={"filing_status": "JOINT"}, + household={"state_code": "CA"}, + year=YEAR, + spm=NATIONAL, +) + + +class TestCalculateWithGroups: + def test__invalid_groups__then_raises_before_calculation(self): + with pytest.raises(ValueError, match=r"people\[2\] belong to no"): + pe.us.calculate_household( + **{**COUPLE_AND_TEEN, "marital_unit": [{"members": [0, 1]}]} + ) + + def test__empty_people__then_raises(self): + with pytest.raises(ValueError, match="people must be a non-empty list"): + pe.us.calculate_household( + people=[], marital_unit=[{"members": [0]}], spm=NATIONAL + ) + + def test__group_input_on_wrong_entity__then_raises_with_entity_hint(self): + with pytest.raises(ValueError, match="filing_status"): + pe.us.calculate_household( + **{ + **COUPLE_AND_TEEN, + "marital_unit": [ + {"members": [0, 1], "filing_status": "JOINT"}, + {"members": [2]}, + ], + } + ) + + def test__periodized_group_input__then_names_the_group(self): + with pytest.raises(ValueError, match=r"tax_unit\[1\]\.filing_status"): + pe.us.calculate_household( + **{ + **COUPLE_AND_TEEN, + "tax_unit": [ + {"members": [0, 1], "filing_status": "JOINT"}, + {"members": [2], "filing_status": {"2026": "SINGLE"}}, + ], + } + ) + + def test__list_form__then_list_of_results_and_singletons_elsewhere(self): + result = pe.us.calculate_household( + **COUPLE_AND_TEEN, + marital_unit=[{"members": [0, 1]}, {"members": [2]}], + ) + assert isinstance(result.marital_unit, list) + assert len(result.marital_unit) == 2 + assert all(isinstance(unit, EntityResult) for unit in result.marital_unit) + assert isinstance(result.tax_unit, EntityResult) + assert isinstance(result.household, EntityResult) + as_dict = result.to_dict() + assert isinstance(as_dict["marital_unit"], list) + assert set(as_dict["marital_unit"][0]) == set(result.marital_unit[0]) + + def test__list_form__then_matches_direct_simulation(self): + variables = { + "person": [ + "ssi", + "medicaid", + "medicaid_optional_senior_or_disabled_income_limit", + ], + "marital_unit": ["marital_unit_id"], + "tax_unit": ["income_tax", "eitc"], + "spm_unit": ["snap", "spm_unit_net_income"], + "household": ["household_net_income", "household_benefits"], + } + result = pe.us.calculate_household( + **COUPLE_AND_TEEN, + marital_unit=[{"members": [0, 1]}, {"members": [2]}], + extra_variables=variables["person"], + ) + direct = _direct_values( + COUPLE_AND_TEEN["people"], + { + "marital_unit": [([0, 1], {}), ([2], {})], + "tax_unit": [([0, 1, 2], {"filing_status": "JOINT"})], + "household": [([0, 1, 2], {"state_code": "CA"})], + }, + variables, + ) + for (entity, name), expected in direct.items(): + if entity == "person": + actual = [person[name] for person in result.person] + elif entity == "marital_unit": + actual = [unit[name] for unit in result.marital_unit] + else: + actual = [result[entity][name]] + assert actual == pytest.approx(expected), (entity, name) + + def test__own_marital_unit__then_child_earnings_leave_parents_ssi(self): + def adult_ssi(child_earnings, marital_unit): + people = [ + {"age": 68}, + {"age": 66}, + {"age": 17, "employment_income": child_earnings}, + ] + result = pe.us.calculate_household( + **{**COUPLE_AND_TEEN, "people": people}, + marital_unit=marital_unit, + extra_variables=["ssi"], + ) + return [person.ssi for person in result.person[:2]] + + own_units = [{"members": [0, 1]}, {"members": [2]}] + with_earnings = adult_ssi(6_000, own_units) + assert with_earnings[0] > 0 + assert with_earnings == adult_ssi(0, own_units) + + def test__single_parent_own_unit__then_individual_medicaid_limit(self): + parent = {"age": 63, "is_blind": True, "social_security_retirement": 12_000} + variable = "medicaid_optional_senior_or_disabled_income_limit" + common = dict(household={"state_code": "MI"}, year=YEAR, spm=NATIONAL) + with_child = pe.us.calculate_household( + people=[parent, {"age": 13}], + marital_unit=[{"members": [0]}, {"members": [1]}], + tax_unit={"filing_status": "HEAD_OF_HOUSEHOLD"}, + extra_variables=[variable], + **common, + ) + alone = pe.us.calculate_household( + people=[parent], + tax_unit={"filing_status": "SINGLE"}, + extra_variables=[variable], + **common, + ) + assert with_child.person[0][variable] > 0 + assert with_child.person[0][variable] == alone.person[0][variable] + + def test__list_form_with_axes__then_each_group_gets_its_own_series(self): + people = [ + {"age": 40, "employment_income": 0}, + {"age": 20, "employment_income": 15_000}, + ] + axes = [{"name": "employment_income", "min": 0, "max": 60_000, "count": 4}] + result = pe.us.calculate_household( + people=people, + marital_unit=[{"members": [0]}, {"members": [1]}], + tax_unit=[ + {"members": [0], "filing_status": "SINGLE"}, + {"members": [1], "filing_status": "SINGLE"}, + ], + household={"state_code": "TX"}, + year=YEAR, + spm=NATIONAL, + axes=axes, + ) + direct = _direct_values( + people, + { + "marital_unit": [([0], {}), ([1], {})], + "tax_unit": [ + ([0], {"filing_status": "SINGLE"}), + ([1], {"filing_status": "SINGLE"}), + ], + "household": [([0, 1], {"state_code": "TX"})], + }, + {"tax_unit": ["income_tax"], "person": ["employment_income"]}, + axes=[[{**axes[0], "period": YEAR}]], + ) + income_tax = direct[("tax_unit", "income_tax")] + assert [unit.income_tax for unit in result.tax_unit] == [ + pytest.approx(income_tax[0::2]), + pytest.approx(income_tax[1::2]), + ] + assert result.person[0].employment_income == [0, 20_000, 40_000, 60_000] + assert result.person[1].employment_income == [15_000] * 4 + parent_tax, child_tax = (unit.income_tax for unit in result.tax_unit) + assert all(a <= b for a, b in zip(parent_tax, parent_tax[1:])) + assert parent_tax[-1] > parent_tax[0] + assert child_tax == [child_tax[0]] * 4 From 7c995cf4bf8cec53f38f4f174f7215428610b7de Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 4 Nov 2026 22:58:48 -0500 Subject: [PATCH 2/3] Tighten calculate_household group lists after review Reject NumPy booleans and other non-index types in members, and give an empty members list its own error. Name a one-group list's position in validation errors and point the misplaced-variable hint at the listed entity's groups. Cover family and SPM unit lists, three marital units, perpendicular axes and an axis on a group variable. Add a differential test against the builder on origin/main, and correct the docs and changelog. Co-Authored-By: Claude Opus 5.5 --- changelog.d/738.added.md | 2 +- changelog.d/738.changed.md | 1 + docs/households.md | 10 +- .../tax_benefit_models/common/household.py | 28 +- .../tax_benefit_models/common/result.py | 7 +- .../tax_benefit_models/us/household.py | 47 +- .../utils/household_validation.py | 27 +- tests/test_us_household_groups.py | 443 ++++++++++++++++-- 8 files changed, 504 insertions(+), 61 deletions(-) create mode 100644 changelog.d/738.changed.md diff --git a/changelog.d/738.added.md b/changelog.d/738.added.md index 85a7c2cfd..b6fac343d 100644 --- a/changelog.d/738.added.md +++ b/changelog.d/738.added.md @@ -1 +1 @@ -`pe.us.calculate_household` accepts a list of groups for `marital_unit`, `family`, `spm_unit` and `tax_unit`. Each group is a mapping with `members`, a list of indexes into `people`, plus its own inputs, and each person belongs to exactly one group of the entity. An entity passed as a list comes back as a list of results in the same order, with one axis series per group under axes. A mapping still puts every person in one group, so existing calls return the same results. Callers can now give a household with children or other adults its own marital units: policyengine-us reads a two-person marital unit as a married couple in its SSI and Medicaid rules, so one marital unit for everyone gave a single parent with one child the couple's Medicaid income limit and counted a child's earnings as their married parents' SSI income. +`pe.us.calculate_household` accepts a list of groups for `marital_unit`, `family`, `spm_unit` and `tax_unit`. Each group is a mapping with `members`, a list of indexes into `people`, plus its own inputs, and each person belongs to exactly one group of the entity. An entity passed as a list comes back as a list of results in the same order, with one axis series per group under axes. A mapping still puts every person in one group, so existing calls return the same results. Callers can now give children and other adults their own marital units. With one marital unit for everyone, Medicaid's optional pathway for seniors and people with disabilities gave a single parent with one child the couple's income limit, and SSI added a child's earnings to the income of aged parents. diff --git a/changelog.d/738.changed.md b/changelog.d/738.changed.md new file mode 100644 index 000000000..95732e4ee --- /dev/null +++ b/changelog.d/738.changed.md @@ -0,0 +1 @@ +`calculate_household` now reads a list passed as `marital_unit`, `family`, `spm_unit` or `tax_unit` as a list of groups, so an empty list or a list of key-value pairs, which it used to turn into a mapping of inputs, now raises. diff --git a/docs/households.md b/docs/households.md index 9bc7804fb..585aafc7d 100644 --- a/docs/households.md +++ b/docs/households.md @@ -39,7 +39,7 @@ A mapping puts every person in one group of that entity, and the household alway To split people into several marital units, families, SPM units or tax units, pass a list of groups. Each group is a mapping with `members`, a list of indexes into `people`, plus that group's own inputs. Each person belongs to exactly one group of the entity: `calculate_household` raises if a list leaves someone out, names someone twice or names an index outside `people`. -policyengine-us defines a marital unit as one unmarried person or one married couple. Its SSI and Medicaid rules read a two-person marital unit as a couple and add up some incomes across it, so pass marital units whenever the household holds anyone besides one person or one couple: +policyengine-us defines a marital unit as one unmarried person or one married couple, and two of its rules depend on that. Medicaid's optional pathway for seniors and people with disabilities applies the couple income limit to a marital unit of exactly two people. SSI treats two aged, blind or disabled heads or spouses in one marital unit as a couple, and then adds up countable income across the whole unit. So pass marital units whenever the household holds anyone besides one person or one couple: ```python result = pe.us.calculate_household( @@ -58,7 +58,7 @@ result = pe.us.calculate_household( result.marital_unit[1] # the 17-year-old's marital unit ``` -With one marital unit for all three, the couple's SSI would count the 17-year-old's earnings as their own income, and Medicaid's optional senior-or-disabled pathway would hold the couple to the individual income limit. +With one marital unit for all three, SSI would add the 17-year-old's earnings to the couple's income, and Medicaid's optional senior-or-disabled pathway would hold the couple to the individual income limit. A list of tax units covers, for example, an adult child who files their own return: @@ -81,6 +81,10 @@ result.tax_unit[1].income_tax # the adult child's return An entity passed as a list comes back as a list of results in the same order. Under axes, each group's variables are axis series, as each person's are. +Each group needs the people its rules expect. policyengine-us makes the oldest person aged 18 or over in each tax unit its head, so a tax unit of children has no head. SPM measurement needs an adult in each SPM unit: a unit of children raises `SPM_COMPOSITION_REQUIRED` when its poverty status computes, as it does among the default outputs. + +Group ids such as `tax_unit_id` keep the values you pass, 0 by default. To see each person's group in `result.person`, pass ids as group inputs, for example `tax_unit=[{"members": [0], "tax_unit_id": 1}, {"members": [1], "tax_unit_id": 2}]`. + ### SPM geography and measurement selection US household results include SPM resources and poverty by default. Provide the @@ -277,7 +281,7 @@ each variable on that person is its own axis series. ## Accessing the result ```python -result.person[0].income_tax # first person +result.person[0].medicaid # first person result.person[2].age # third person result.tax_unit.income_tax # single tax unit result.marital_unit[1] # second marital unit, when passed as a list diff --git a/src/policyengine/tax_benefit_models/common/household.py b/src/policyengine/tax_benefit_models/common/household.py index e3aced331..6f7c60aae 100644 --- a/src/policyengine/tax_benefit_models/common/household.py +++ b/src/policyengine/tax_benefit_models/common/household.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Collection, Mapping, Sequence from typing import Any @@ -8,10 +8,16 @@ def validate_annual_household_inputs( *, year: Any, entities: Mapping[str, Sequence[Mapping[str, Any]]], + listed_entities: Collection[str] = (), ) -> int: - """Validate annual-only household calculator inputs.""" + """Validate annual-only household calculator inputs. + + Errors name an input of an entity in ``listed_entities``, which the caller + passed as a list of groups, by its group's position, even in a list of + one group. + """ validated_year = _validate_annual_year(year) - _validate_unperiodized_values(entities) + _validate_unperiodized_values(entities, listed_entities=listed_entities) return validated_year @@ -36,16 +42,24 @@ def _annual_period_error(year: Any) -> ValueError: def _validate_unperiodized_values( entities: Mapping[str, Sequence[Mapping[str, Any]]], + *, + listed_entities: Collection[str] = (), ) -> None: for entity, records in entities.items(): for index, record in enumerate(records): for variable, value in record.items(): if variable != "id" and isinstance(value, Mapping): + location = _input_location( + entity, + index, + len(records), + variable, + listed=entity in listed_entities, + ) raise ValueError( "Periodized household inputs are not supported by " "calculate_household. Pass annual scalar input values " - f"only; received a periodized value for " - f"{_input_location(entity, index, len(records), variable)}." + f"only; received a periodized value for {location}." ) @@ -54,7 +68,9 @@ def _input_location( index: int, record_count: int, variable: str, + *, + listed: bool = False, ) -> str: - if record_count == 1 and entity != "people": + if record_count == 1 and entity != "people" and not listed: return f"{entity}.{variable}" return f"{entity}[{index}].{variable}" diff --git a/src/policyengine/tax_benefit_models/common/result.py b/src/policyengine/tax_benefit_models/common/result.py index 78b5c2218..3c79c3a2b 100644 --- a/src/policyengine/tax_benefit_models/common/result.py +++ b/src/policyengine/tax_benefit_models/common/result.py @@ -41,9 +41,10 @@ def __setattr__(self, name: str, value: Any) -> None: # pragma: no cover class HouseholdResult(dict): """Full household calculation result; one key per entity. - Singleton entities (``household``, ``tax_unit``, ``benunit``, ...) - map to a single :class:`EntityResult`; multi-member entities (like - ``person``) map to a ``list[EntityResult]``. + ``person`` maps to a ``list[EntityResult]``, one per person. A group + entity (``household``, ``tax_unit``, ``benunit``, ...) maps to a single + :class:`EntityResult`, or to a ``list[EntityResult]`` in the caller's + order when the caller passed that entity as a list of groups. """ def __getattr__(self, name: str) -> Any: diff --git a/src/policyengine/tax_benefit_models/us/household.py b/src/policyengine/tax_benefit_models/us/household.py index ba224aa28..680bad4d9 100644 --- a/src/policyengine/tax_benefit_models/us/household.py +++ b/src/policyengine/tax_benefit_models/us/household.py @@ -39,6 +39,8 @@ from collections.abc import Mapping, Sequence from typing import Any, Optional, Union +import numpy as np + from policyengine.tax_benefit_models.common import ( EntityResult, HouseholdResult, @@ -109,7 +111,9 @@ def _safe_convert(value: Any) -> Any: def _person_index(value: Any, *, location: str, people_count: int) -> int: - if isinstance(value, bool): + # NumPy booleans support operator.index (np.True_ gives 1), so they need + # their own check, like Python's. + if isinstance(value, (bool, np.bool_)): index = None else: try: @@ -175,16 +179,23 @@ def _resolve_groups( f"{location} must be a mapping with 'members' and the group's " f"inputs; got {type(group).__name__}." ) - members = group.get("members") - if ( - members is None - or isinstance(members, (str, bytes, Mapping)) - or not isinstance(members, Sequence) - or not members - ): + if "members" not in group: + raise ValueError( + f"{location} must include 'members', a list of indexes into people." + ) + members = group["members"] + if isinstance(members, np.ndarray) and members.ndim == 1: + members = members.tolist() + if isinstance( + members, (str, bytes, bytearray, memoryview, Mapping) + ) or not isinstance(members, Sequence): + raise ValueError( + f"{location}['members'] must be a list of indexes into " + f"people; got {type(members).__name__}." + ) + if not members: raise ValueError( - f"{location} must include 'members', a non-empty list of " - "indexes into people." + f"{location}['members'] is empty; each group holds at least one person." ) indexes = [] for value in members: @@ -300,11 +311,14 @@ def calculate_household( mapping with ``members`` (indexes into ``people``) plus that group's inputs; each person must belong to exactly one group of the entity. policyengine-us defines a marital unit as one - unmarried person or one married couple, and its SSI and Medicaid - rules read a two-person marital unit as a couple, so a household - with children or other adults needs the list form, e.g. - ``marital_unit=[{"members": [0, 1]}, {"members": [2]}]`` for a - married couple and their child. + unmarried person or one married couple. Medicaid's optional + pathway for seniors and people with disabilities applies the + couple income limit to a marital unit of exactly two people, and + SSI treats two aged, blind or disabled heads or spouses in one + marital unit as a couple and then adds up income across the whole + unit. So a household with children or other adults needs the list + form, e.g. ``marital_unit=[{"members": [0, 1]}, {"members": [2]}]`` + for a married couple and their child. year: Calendar year to compute for. Defaults to 2026. reform: Optional reform as ``{parameter_path: value}`` or ``{parameter_path: {effective_date: value}}``. Scalar @@ -383,9 +397,11 @@ def calculate_household( name: [_group_inputs(group) for group in entity_groups] for name, entity_groups in groups.items() } + listed_entities = frozenset(name for name, is_list in listed.items() if is_list) year = validate_annual_household_inputs( year=year, entities={"people": people, **records}, + listed_entities=listed_entities, ) from policyengine_us import Simulation @@ -396,6 +412,7 @@ def calculate_household( validate_household_input( model_version=us_latest, entities={"person": people, **records}, + listed_entities=listed_entities, ) extra_by_entity = dispatch_extra_variables( diff --git a/src/policyengine/utils/household_validation.py b/src/policyengine/utils/household_validation.py index 6be90fb26..f8996aa01 100644 --- a/src/policyengine/utils/household_validation.py +++ b/src/policyengine/utils/household_validation.py @@ -14,7 +14,7 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping +from collections.abc import Collection, Iterable, Mapping from difflib import get_close_matches from typing import TYPE_CHECKING @@ -48,6 +48,7 @@ def validate_household_input( *, model_version: TaxBenefitModelVersion, entities: Mapping[str, Iterable[Mapping[str, object]]], + listed_entities: Collection[str] = (), ) -> None: """Raise ``ValueError`` on unknown or mis-placed entity variables. @@ -58,6 +59,10 @@ def validate_household_input( suggestion. - If the key is a known variable but defined on a different entity, the error names the correct entity and shows the kwarg swap. + + ``listed_entities`` names the entities the caller passed as a list of + groups: errors name their groups by position, and a variable that + belongs on one of them is moved into a group rather than a mapping. """ if "person" in entities and not list(entities["person"]): raise ValueError( @@ -70,15 +75,20 @@ def validate_household_input( misplaced: list[tuple[str, str, str]] = [] for entity_name, records in entities.items(): - for record in records: + for index, record in enumerate(records): + location = ( + f"{entity_name}[{index}]" + if entity_name in listed_entities + else entity_name + ) for key in record: if key in _STRUCTURAL_KEYS: continue variable = variables_by_name.get(key) if variable is None: - unknown.append((entity_name, key)) + unknown.append((location, key)) elif variable.entity != entity_name: - misplaced.append((entity_name, key, variable.entity)) + misplaced.append((location, key, variable.entity)) if not unknown and not misplaced: return @@ -104,10 +114,15 @@ def validate_household_input( lines.append("") lines.append("Variables passed on the wrong entity:") for wrong_entity, key, correct_entity in misplaced: + if correct_entity in listed_entities: + move = ( + f"add '{key}': to the {correct_entity} group it applies to" + ) + else: + move = f"pass {correct_entity}={{'{key}': }}" lines.append( f" - '{key}' was given on {wrong_entity}; it belongs on " - f"{correct_entity}. Move it: pass " - f"{correct_entity}={{'{key}': }}." + f"{correct_entity}. Move it: {move}." ) raise ValueError("\n".join(lines)) diff --git a/tests/test_us_household_groups.py b/tests/test_us_household_groups.py index 34d83abc4..8472b7bb0 100644 --- a/tests/test_us_household_groups.py +++ b/tests/test_us_household_groups.py @@ -8,13 +8,19 @@ - Partition: a list of groups is accepted exactly when it assigns every person to one group, and the situation's memberships equal that partition. -- Backward compatibility: a mapping (or nothing) builds the same single group - holding every person as before, under the same ids. +- Backward compatibility: for mapping (or omitted) entities, the builder emits + the same situation, key order included, as origin/main's builder before + list form, whatever the inputs on each entity and with or without axes. - Differential: list-form results equal a ``policyengine_us.Simulation`` - built by hand with the same groups, with and without axes. + built by hand with the same groups, for every splittable entity, without + axes, with one axis, and with perpendicular axes on a person variable and a + group variable. +- Membership: ids passed as group inputs come back on each group and on each + of its members. - Semantics: with its own marital unit, a child's earnings leave its married - parents' SSI unchanged, and a single parent faces the individual Medicaid - income limit. + parents' SSI unchanged, while one shared unit lowers it; a single parent + faces the individual Medicaid income limit, and so does a couple sharing a + unit with their child. US cases select national SPM geography: they test the calculator interface without claiming a local SPM area. @@ -22,10 +28,15 @@ from __future__ import annotations +import json +from collections.abc import Mapping +from typing import Any, Optional + import numpy as np import pytest from hypothesis import given, settings from hypothesis import strategies as st +from spm_calculator.errors import SPMInputError import policyengine as pe from policyengine.tax_benefit_models.common import EntityResult @@ -88,6 +99,38 @@ def _direct_values(people, groups, variables, *, axes=None): return out +def _result_values(result, entity, name): + """One variable of a ``calculate_household`` result in core's layout. + + Each group gives one value, or an axis series under axes; core lays copy c + of instance i out at ``c * count + i``. + """ + section = result[entity] + instances = section if isinstance(section, list) else [section] + series = [instance[name] for instance in instances] + if not isinstance(series[0], list): + return series + return [ + series[index][point] + for point in range(len(series[0])) + for index in range(len(instances)) + ] + + +def _list_kwargs(groups): + """``calculate_household`` keyword arguments for ``_direct_values`` groups.""" + kwargs = {} + for entity, specs in groups.items(): + if entity == "household": + ((_, inputs),) = specs + kwargs[entity] = dict(inputs) + else: + kwargs[entity] = [ + {"members": list(members), **inputs} for members, inputs in specs + ] + return kwargs + + # --- Resolving groups --------------------------------------------------------- @@ -127,6 +170,15 @@ def test__numpy_integer_indexes__then_accepted(self): ) assert groups == [{"members": [0]}, {"members": [1]}] + def test__numpy_array_and_range_members__then_accepted_as_ints(self): + groups, _ = _resolve_groups( + "spm_unit", + [{"members": np.array([2, 0])}, {"members": range(1, 2)}], + people_count=3, + ) + assert groups == [{"members": [2, 0]}, {"members": [1]}] + assert all(type(index) is int for group in groups for index in group["members"]) + @pytest.mark.parametrize( ("spec", "match"), [ @@ -142,13 +194,43 @@ def test__numpy_integer_indexes__then_accepted(self): ([{"members": [0.0, 1]}], r"integer indexes into people; got 0.0"), ([{"members": ["0", 1]}], r"integer indexes into people; got '0'"), ([{"members": [True, 0]}], r"integer indexes into people; got True"), - ([{"members": []}, {"members": [0, 1]}], r"marital_unit\[0\] must"), - ([{"filing_status": "JOINT"}], r"marital_unit\[0\] must include"), - ([{"members": "01"}], r"marital_unit\[0\] must include 'members'"), - ([{"members": {0: 1}}], r"marital_unit\[0\] must include 'members'"), + ( + [{"members": [np.True_, 0]}], + r"integer indexes into people; got np\.True_", + ), + ( + [{"members": np.array([True, False])}], + r"integer indexes into people; got True", + ), + ( + [{"members": []}, {"members": [0, 1]}], + r"marital_unit\[0\]\['members'\] is empty", + ), + ( + [{"filing_status": "JOINT"}], + r"marital_unit\[0\] must include 'members'", + ), + ([{"members": "01"}], r"marital_unit\[0\]\['members'\] .* got str"), + ([{"members": {0: 1}}], r"marital_unit\[0\]\['members'\] .* got dict"), + ([{"members": b"\x00\x01"}], r"\['members'\] .* got bytes"), + ( + [{"members": bytearray(b"\x00\x01")}], + r"\['members'\] .* got bytearray", + ), + ( + [{"members": memoryview(b"\x00\x01")}], + r"\['members'\] .* got memoryview", + ), + ([{"members": np.array([[0, 1]])}], r"\['members'\] .* got ndarray"), + ([{"members": 0}], r"marital_unit\[0\]\['members'\] .* got int"), ([], r"marital_unit must list at least one group"), ([[0, 1]], r"marital_unit\[0\] must be a mapping"), ("01", r"marital_unit must be a mapping of inputs or a list"), + ( + (group for group in [{"members": [0, 1]}]), + r"marital_unit must be a mapping of inputs or a list of groups; " + r"got generator", + ), ({"members": [0, 1]}, r"takes no 'members'"), ], ) @@ -209,15 +291,18 @@ def test__dropped_or_repeated_person__then_raises(self, case, data): people_count, partition = case groups = [list(members) for members in partition] position = data.draw(st.integers(0, len(groups) - 1)) - if data.draw(st.booleans()): - groups[position].pop(data.draw(st.integers(0, len(groups[position]) - 1))) + # Dropping the only person leaves no group at all, which is the + # empty-list error that TestResolveGroups covers. + if people_count > 1 and data.draw(st.booleans()): + members = groups[position] + dropped = members.pop(data.draw(st.integers(0, len(members) - 1))) groups = [members for members in groups if members] - if not groups: - groups = [[people_count]] # every person dropped: out of range + match = rf"people\[{dropped}\] belong to no tax_unit" else: extra = data.draw(st.integers(0, people_count - 1)) groups[position].append(extra) - with pytest.raises(ValueError): + match = rf"people\[{extra}\] appears" + with pytest.raises(ValueError, match=match): _resolve_groups( "tax_unit", [{"members": members} for members in groups], @@ -259,6 +344,99 @@ def test__mapping_form__then_builds_the_single_group_structure( } +# --- Backward compatibility with the builder before list form ---------------- + + +def _origin_main_build_situation( + *, + people: list[Mapping[str, Any]], + marital_unit: Mapping[str, Any], + family: Mapping[str, Any], + spm_unit: Mapping[str, Any], + tax_unit: Mapping[str, Any], + household: Mapping[str, Any], + year: int, + axes: Optional[list[Any]] = None, +) -> dict[str, Any]: + """``_build_situation`` as origin/main had it at 286fce481, verbatim.""" + year_str = str(year) + + def _periodise(spec: Mapping[str, Any]) -> dict[str, dict[str, Any]]: + return {key: {year_str: value} for key, value in spec.items() if key != "id"} + + person_ids = [f"person_{i}" for i in range(len(people))] + persons = {pid: _periodise(person) for pid, person in zip(person_ids, people)} + + def _group(spec: Mapping[str, Any]) -> dict[str, Any]: + return {"members": list(person_ids), **_periodise(spec)} + + situation = { + "people": persons, + "marital_units": {"marital_unit_0": _group(marital_unit)}, + "families": {"family_0": _group(family)}, + "spm_units": {"spm_unit_0": _group(spm_unit)}, + "tax_units": {"tax_unit_0": _group(tax_unit)}, + "households": {"household_0": _group(household)}, + } + if axes is not None: + situation["axes"] = axes + return situation + + +_RECORDS = st.dictionaries( + st.sampled_from(["id", "age", "employment_income", "filing_status", "x"]), + st.one_of( + st.integers(-3, 3), + st.booleans(), + st.sampled_from(["JOINT", "CA"]), + st.floats(allow_nan=False, allow_infinity=False, width=32), + ), + max_size=4, +) + + +class TestBackwardCompatibility: + @settings(max_examples=300, deadline=None) + @given( + people=st.lists(_RECORDS, min_size=1, max_size=6), + specs=st.fixed_dictionaries( + {name: st.one_of(st.none(), _RECORDS) for name in GROUP_PLURALS} + ), + axes=st.one_of( + st.none(), + st.just( + [ + [ + { + "name": "employment_income", + "min": 0, + "max": 10, + "count": 2, + "period": YEAR, + } + ] + ] + ), + ), + ) + def test__mapping_form__then_same_situation_as_origin_main( + self, people, specs, axes + ): + groups = { + name: _resolve_groups(name, spec, people_count=len(people))[0] + for name, spec in specs.items() + } + built = _build_situation(people=people, groups=groups, year=YEAR, axes=axes) + before = _origin_main_build_situation( + people=people, + year=YEAR, + axes=axes, + **{name: dict(spec or {}) for name, spec in specs.items()}, + ) + # Compare serialized forms, so key order has to match too. + assert json.dumps(built) == json.dumps(before) + + # --- Calculating -------------------------------------------------------------- @@ -270,6 +448,56 @@ def test__mapping_form__then_builds_the_single_group_structure( spm=NATIONAL, ) +# A parent and child, plus a grandparent who forms their own family, SPM unit +# and tax unit. +THREE_GENERATIONS_PEOPLE = [ + {"age": 34, "employment_income": 18_000}, + {"age": 9}, + {"age": 70, "social_security_retirement": 14_000}, +] +THREE_GENERATIONS_GROUPS = { + "marital_unit": [([0], {}), ([1], {}), ([2], {})], + "family": [([0, 1], {}), ([2], {})], + "spm_unit": [([0, 1], {}), ([2], {})], + "tax_unit": [ + ([0, 1], {"filing_status": "HEAD_OF_HOUSEHOLD"}), + ([2], {"filing_status": "SINGLE"}), + ], + "household": [([0, 1, 2], {"state_code": "PA"})], +} +THREE_GENERATIONS_VARIABLES = { + "person": [ + "employment_income", + "medicaid", + "ssi", + "ssi_claim_is_joint", + "ssi_marital_earned_income", + ], + "family": ["is_married"], + "spm_unit": [ + "spm_unit_pre_subsidy_childcare_expenses", + "spm_unit_size", + "snap", + "spm_unit_net_income", + ], + "tax_unit": ["tax_unit_size", "income_tax", "eitc"], + "household": ["household_size", "household_net_income"], +} +EARNINGS_AXIS = { + "name": "employment_income", + "min": 0, + "max": 50_000, + "count": 3, + "period": YEAR, +} +CHILDCARE_AXIS = { + "name": "spm_unit_pre_subsidy_childcare_expenses", + "min": 0, + "max": 8_000, + "count": 2, + "period": YEAR, +} + class TestCalculateWithGroups: def test__invalid_groups__then_raises_before_calculation(self): @@ -284,8 +512,8 @@ def test__empty_people__then_raises(self): people=[], marital_unit=[{"members": [0]}], spm=NATIONAL ) - def test__group_input_on_wrong_entity__then_raises_with_entity_hint(self): - with pytest.raises(ValueError, match="filing_status"): + def test__group_input_on_wrong_entity__then_names_the_group(self): + with pytest.raises(ValueError) as caught: pe.us.calculate_household( **{ **COUPLE_AND_TEEN, @@ -295,6 +523,26 @@ def test__group_input_on_wrong_entity__then_raises_with_entity_hint(self): ], } ) + message = str(caught.value) + assert "'filing_status' was given on marital_unit[0]" in message + assert "pass tax_unit={'filing_status': }" in message + + def test__input_belonging_on_a_listed_entity__then_hint_names_a_group(self): + with pytest.raises(ValueError) as caught: + pe.us.calculate_household( + **{ + **COUPLE_AND_TEEN, + "marital_unit": [ + {"members": [0, 1], "filing_status": "JOINT"}, + {"members": [2]}, + ], + "tax_unit": [{"members": [0, 1, 2]}], + } + ) + assert ( + "add 'filing_status': to the tax_unit group it applies to" + in str(caught.value) + ) def test__periodized_group_input__then_names_the_group(self): with pytest.raises(ValueError, match=r"tax_unit\[1\]\.filing_status"): @@ -308,6 +556,17 @@ def test__periodized_group_input__then_names_the_group(self): } ) + def test__periodized_input_in_one_group_list__then_names_its_position(self): + with pytest.raises(ValueError, match=r"tax_unit\[0\]\.filing_status"): + pe.us.calculate_household( + **{ + **COUPLE_AND_TEEN, + "tax_unit": [ + {"members": [0, 1, 2], "filing_status": {"2026": "JOINT"}} + ], + } + ) + def test__list_form__then_list_of_results_and_singletons_elsewhere(self): result = pe.us.calculate_household( **COUPLE_AND_TEEN, @@ -322,22 +581,27 @@ def test__list_form__then_list_of_results_and_singletons_elsewhere(self): assert isinstance(as_dict["marital_unit"], list) assert set(as_dict["marital_unit"][0]) == set(result.marital_unit[0]) - def test__list_form__then_matches_direct_simulation(self): + def test__marital_unit_list__then_matches_direct_simulation(self): variables = { "person": [ "ssi", + "ssi_claim_is_joint", + "ssi_marital_earned_income", "medicaid", "medicaid_optional_senior_or_disabled_income_limit", ], - "marital_unit": ["marital_unit_id"], - "tax_unit": ["income_tax", "eitc"], - "spm_unit": ["snap", "spm_unit_net_income"], + "tax_unit": ["tax_unit_size", "income_tax", "eitc"], + "spm_unit": ["spm_unit_size", "snap", "spm_unit_net_income"], "household": ["household_net_income", "household_benefits"], } result = pe.us.calculate_household( **COUPLE_AND_TEEN, marital_unit=[{"members": [0, 1]}, {"members": [2]}], - extra_variables=variables["person"], + extra_variables=[ + *variables["person"], + "tax_unit_size", + "spm_unit_size", + ], ) direct = _direct_values( COUPLE_AND_TEEN["people"], @@ -349,14 +613,112 @@ def test__list_form__then_matches_direct_simulation(self): variables, ) for (entity, name), expected in direct.items(): - if entity == "person": - actual = [person[name] for person in result.person] - elif entity == "marital_unit": - actual = [unit[name] for unit in result.marital_unit] - else: - actual = [result[entity][name]] + actual = ( + [person[name] for person in result.person] + if entity == "person" + else _result_values(result, entity, name) + ) assert actual == pytest.approx(expected), (entity, name) + @pytest.mark.parametrize( + "axes", + [None, [[EARNINGS_AXIS]], [[EARNINGS_AXIS], [CHILDCARE_AXIS]]], + ids=["no_axes", "one_axis", "perpendicular_axes"], + ) + def test__every_entity_as_a_list__then_matches_direct_simulation(self, axes): + variables = THREE_GENERATIONS_VARIABLES + # Default outputs among these come back once either way. + extra = sorted({name for names in variables.values() for name in names}) + result = pe.us.calculate_household( + people=THREE_GENERATIONS_PEOPLE, + **_list_kwargs(THREE_GENERATIONS_GROUPS), + year=YEAR, + spm=NATIONAL, + axes=axes, + extra_variables=extra, + ) + direct = _direct_values( + THREE_GENERATIONS_PEOPLE, + THREE_GENERATIONS_GROUPS, + variables, + axes=axes, + ) + for (entity, name), expected in direct.items(): + assert _result_values(result, entity, name) == pytest.approx(expected), ( + entity, + name, + ) + points = 1 if axes is None else 3 * (2 if len(axes) == 2 else 1) + assert len(direct[("person", "employment_income")]) == 3 * points + assert [len(unit) for unit in (result.family, result.spm_unit)] == [2, 2] + assert [unit.tax_unit_size for unit in result.tax_unit] == ( + [2, 1] if axes is None else [[2] * points, [1] * points] + ) + + def test__ids_passed_as_group_inputs__then_on_each_group_and_member(self): + result = pe.us.calculate_household( + people=[{"age": 40}, {"age": 38}, {"age": 12}, {"age": 19}], + marital_unit=[ + {"members": [0, 1], "marital_unit_id": 7}, + {"members": [2], "marital_unit_id": 8}, + {"members": [3], "marital_unit_id": 9}, + ], + family=[ + {"members": [0, 1, 2], "family_id": 4}, + {"members": [3], "family_id": 5}, + ], + spm_unit=[ + {"members": [0, 1, 2], "spm_unit_id": 2}, + {"members": [3], "spm_unit_id": 3}, + ], + tax_unit=[ + {"members": [0, 1, 2], "tax_unit_id": 1, "filing_status": "JOINT"}, + {"members": [3], "tax_unit_id": 6, "filing_status": "SINGLE"}, + ], + household={"state_code": "OH"}, + year=YEAR, + spm=NATIONAL, + ) + members = { + "marital_unit": [7, 7, 8, 9], + "family": [4, 4, 4, 5], + "spm_unit": [2, 2, 2, 3], + "tax_unit": [1, 1, 1, 6], + } + for entity, ids in members.items(): + variable = f"{entity}_id" + assert [person[variable] for person in result.person] == ids + assert [group[variable] for group in result[entity]] == list( + dict.fromkeys(ids) + ) + + def test__spm_unit_of_children__then_composition_error(self): + with pytest.raises(SPMInputError) as caught: + pe.us.calculate_household( + people=[{"age": 40, "employment_income": 20_000}, {"age": 13}], + marital_unit=[{"members": [0]}, {"members": [1]}], + spm_unit=[{"members": [0]}, {"members": [1]}], + household={"state_code": "TX"}, + year=YEAR, + spm=NATIONAL, + ) + assert caught.value.code == "SPM_COMPOSITION_REQUIRED" + + def test__tax_unit_of_children__then_it_has_no_head(self): + result = pe.us.calculate_household( + people=[{"age": 40, "employment_income": 20_000}, {"age": 13}], + marital_unit=[{"members": [0]}, {"members": [1]}], + tax_unit=[ + {"members": [0], "filing_status": "SINGLE"}, + {"members": [1], "filing_status": "SINGLE"}, + ], + household={"state_code": "TX"}, + year=YEAR, + spm=NATIONAL, + extra_variables=["is_tax_unit_head"], + ) + assert [person.is_tax_unit_head for person in result.person] == [1, 0] + def test__own_marital_unit__then_child_earnings_leave_parents_ssi(self): def adult_ssi(child_earnings, marital_unit): people = [ @@ -375,6 +737,10 @@ def adult_ssi(child_earnings, marital_unit): with_earnings = adult_ssi(6_000, own_units) assert with_earnings[0] > 0 assert with_earnings == adult_ssi(0, own_units) + # One marital unit for all three adds the child's earnings to the + # couple's income, as the docs warn. + shared = [{"members": [0, 1, 2]}] + assert sum(adult_ssi(6_000, shared)) < sum(adult_ssi(0, shared)) def test__single_parent_own_unit__then_individual_medicaid_limit(self): parent = {"age": 63, "is_blind": True, "social_security_retirement": 12_000} @@ -396,6 +762,29 @@ def test__single_parent_own_unit__then_individual_medicaid_limit(self): assert with_child.person[0][variable] > 0 assert with_child.person[0][variable] == alone.person[0][variable] + def test__couple_sharing_a_unit_with_their_child__then_individual_limit(self): + variable = "medicaid_optional_senior_or_disabled_income_limit" + + def parent_limits(marital_unit): + result = pe.us.calculate_household( + **{**COUPLE_AND_TEEN, "household": {"state_code": "MI"}}, + marital_unit=marital_unit, + extra_variables=[variable], + ) + return [person[variable] for person in result.person[:2]] + + alone = pe.us.calculate_household( + people=[{"age": 68}], + tax_unit={"filing_status": "SINGLE"}, + household={"state_code": "MI"}, + year=YEAR, + spm=NATIONAL, + extra_variables=[variable], + ).person[0][variable] + couple = parent_limits([{"members": [0, 1]}, {"members": [2]}]) + assert couple[0] > alone > 0 + assert parent_limits([{"members": [0, 1, 2]}]) == [alone, alone] + def test__list_form_with_axes__then_each_group_gets_its_own_series(self): people = [ {"age": 40, "employment_income": 0}, From 393a569cfea034891a124a7de7e74165a066dc18 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 4 Nov 2026 23:15:06 -0500 Subject: [PATCH 3/3] Refuse non-integer NumPy members and qualify the SSI note tolist() turned datetime64 and timedelta64 arrays into integers before _person_index could refuse them, so keep NumPy's scalars until then. SSI adds up income across the unit only when the couple computation applies, which excludes exactly one member in a medical treatment facility. Co-Authored-By: Claude Opus 5.5 --- docs/households.md | 2 +- .../tax_benefit_models/us/household.py | 14 +++++++++----- tests/test_us_household_groups.py | 14 +++++++++++++- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/docs/households.md b/docs/households.md index 585aafc7d..d0c740553 100644 --- a/docs/households.md +++ b/docs/households.md @@ -39,7 +39,7 @@ A mapping puts every person in one group of that entity, and the household alway To split people into several marital units, families, SPM units or tax units, pass a list of groups. Each group is a mapping with `members`, a list of indexes into `people`, plus that group's own inputs. Each person belongs to exactly one group of the entity: `calculate_household` raises if a list leaves someone out, names someone twice or names an index outside `people`. -policyengine-us defines a marital unit as one unmarried person or one married couple, and two of its rules depend on that. Medicaid's optional pathway for seniors and people with disabilities applies the couple income limit to a marital unit of exactly two people. SSI treats two aged, blind or disabled heads or spouses in one marital unit as a couple, and then adds up countable income across the whole unit. So pass marital units whenever the household holds anyone besides one person or one couple: +policyengine-us defines a marital unit as one unmarried person or one married couple, and two of its rules depend on that. Medicaid's optional pathway for seniors and people with disabilities applies the couple income limit to a marital unit of exactly two people. SSI treats two aged, blind or disabled heads or spouses in one marital unit as a joint claim and, unless exactly one member counts as living in a medical treatment facility, adds up countable income across the whole unit. So pass marital units whenever the household holds anyone besides one person or one couple: ```python result = pe.us.calculate_household( diff --git a/src/policyengine/tax_benefit_models/us/household.py b/src/policyengine/tax_benefit_models/us/household.py index 680bad4d9..f51b1d3bb 100644 --- a/src/policyengine/tax_benefit_models/us/household.py +++ b/src/policyengine/tax_benefit_models/us/household.py @@ -185,7 +185,9 @@ def _resolve_groups( ) members = group["members"] if isinstance(members, np.ndarray) and members.ndim == 1: - members = members.tolist() + # list(), not tolist(): tolist() turns dates and durations into + # integers before _person_index can refuse them. + members = list(members) if isinstance( members, (str, bytes, bytearray, memoryview, Mapping) ) or not isinstance(members, Sequence): @@ -315,10 +317,12 @@ def calculate_household( pathway for seniors and people with disabilities applies the couple income limit to a marital unit of exactly two people, and SSI treats two aged, blind or disabled heads or spouses in one - marital unit as a couple and then adds up income across the whole - unit. So a household with children or other adults needs the list - form, e.g. ``marital_unit=[{"members": [0, 1]}, {"members": [2]}]`` - for a married couple and their child. + marital unit as a joint claim and, unless exactly one member + counts as living in a medical treatment facility, adds up income + across the whole unit. So a household with children or other + adults needs the list form, e.g. + ``marital_unit=[{"members": [0, 1]}, {"members": [2]}]`` for a + married couple and their child. year: Calendar year to compute for. Defaults to 2026. reform: Optional reform as ``{parameter_path: value}`` or ``{parameter_path: {effective_date: value}}``. Scalar diff --git a/tests/test_us_household_groups.py b/tests/test_us_household_groups.py index 8472b7bb0..7a7a763a9 100644 --- a/tests/test_us_household_groups.py +++ b/tests/test_us_household_groups.py @@ -200,7 +200,19 @@ def test__numpy_array_and_range_members__then_accepted_as_ints(self): ), ( [{"members": np.array([True, False])}], - r"integer indexes into people; got True", + r"integer indexes into people; got np\.True_", + ), + ( + [{"members": np.array([0, 1], dtype="datetime64[ns]")}], + r"integer indexes into people; got np\.datetime64", + ), + ( + [{"members": np.array([0, 1], dtype="timedelta64[ns]")}], + r"integer indexes into people; got np\.timedelta64", + ), + ( + [{"members": np.array([0.0, 1.0])}], + r"integer indexes into people; got np\.float64", ), ( [{"members": []}, {"members": [0, 1]}],