Skip to content

Let US calculate_household take lists of groups - #739

Draft
MaxGhenis wants to merge 3 commits into
mainfrom
claude/household-group-lists
Draft

MaxGhenis wants to merge 3 commits into
mainfrom
claude/household-group-lists

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #738

Summary

pe.us.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 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, so a single parent with one child got the couple's limit. 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 teenager's earnings counted toward their aged parents' SSI (#738 has the reproductions).

marital_unit, family, spm_unit and tax_unit now also accept a list of groups. Each group is a mapping with members, a list of indexes into people, plus that group's own inputs:

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
  • An entity passed as a list comes back as a list of EntityResults in the same order; under axes, each group's variables are axis series, as each person's are.
  • A mapping, or nothing, still builds one group holding everyone and returns one EntityResult, so existing calls return the same results.
  • household stays a single mapping: calculate_household computes one household.
  • A list that leaves a person out, names a person twice, names an index outside people, passes members in the mapping form, or passes a list for household raises ValueError before any calculation, naming the entity and group.
  • members takes a list, tuple, range or one-dimensional NumPy array of integer indexes. Booleans (NumPy's included), floats, strings, bytes, bytearray, memoryview and mappings raise.
  • Errors name a listed entity's groups by position, even in a list of one group, and the hint for a variable that belongs on a listed entity says to add it to one of its groups.
  • An empty list, or a list of key-value pairs, which the mapping form used to turn into a dict, now raises (changelog.d/738.changed.md).

docs/households.md documents the list form with the SSI and adult-child tax-unit examples, and its "Accessing the result" section now describes the result as a dict with attribute access (it said Pydantic model). It also says what each group needs (policyengine-us makes the oldest person aged 18 or over in each tax unit its head, and SPM measurement needs an adult in each SPM unit) and how to see each person's group: pass {entity}_id as a group input.

Invariants

Stated in tests/test_us_household_groups.py and checked there:

  1. Partition: calculate_household accepts a list of groups exactly when it assigns every person to exactly one group, and the situation it builds has memberships equal to that partition (Hypothesis, 200 random partitions). Dropping a person raises "belong to no", and repeating one raises "appears" (200 more).
  2. Backward compatibility: for mapping or omitted entities, the builder emits the same JSON, key order included, as origin/main's builder before list form, which the test carries verbatim. Hypothesis draws inputs on every entity, with and without axes (300 examples). The existing household, snapshot and SPM household tests pass unchanged.
  3. Differential: list-form results equal a policyengine_us.Simulation built by hand with the same groups:
    • for a marital-unit list: SSI, its joint-claim flag and marital earned income, Medicaid and its senior-or-disabled income limit, tax unit and SPM unit sizes, income tax, EITC, SNAP, SPM net income, household net income and household benefits;
    • for a household with every splittable entity as a list (three marital units, two families, two SPM units, two tax units): the same kinds of outputs on every entity, without axes, with one earnings axis, and with perpendicular axes on earnings and on an SPM unit's pre-subsidy child care expenses;
    • under one axis with two tax units, each tax unit's income tax series equals the hand-built [0::2] and [1::2] slices.
  4. Membership: ids passed as group inputs come back on each group and on each of its members.
  5. Semantics:
    • with its own marital unit, a child's earnings leave its aged parents' SSI unchanged, and one shared marital unit lowers it;
    • a single parent with their own marital unit gets the same Medicaid senior-or-disabled income limit as when living alone;
    • a couple with their own marital unit gets a higher limit than a single person, and a couple sharing one with their child gets the single person's limit.
  6. Composition: an SPM unit of children raises SPM_COMPOSITION_REQUIRED, and a tax unit of children has no head.

Review round 1 (5 November 2026)

An independent review (Opus, read-only) asked for these changes, and this push makes them:

  • The SSI wording was wrong. The summary, docs, docstring and changelog said SSI reads any two-person marital unit as a couple. Only Medicaid's limit does that; SSI's couple test counts aged, blind or disabled heads and spouses. All four now say so, and a test shows one shared marital unit lowering the parents' SSI.
  • The id check proved nothing. The differential compared marital_unit_id, which stays 0 for every group unless the caller passes it. It now compares SSI's joint-claim flag and marital earned income, group sizes and outputs on every entity, and a new test passes ids as group inputs and finds them on each group and member.
  • Composition: the docs now say what each tax unit and SPM unit needs, with a test for each.
  • More tests: end-to-end family and SPM unit lists, three marital units, perpendicular axes, an axis on a group variable, the builder differential against origin/main, and the input types below. The drop-a-person property no longer falls back to an out-of-range index.
  • Input types: the review expected np.True_ to fail operator.index; on NumPy 2.1.3 it returns 1, so members=[np.True_, 0] passed as [1, 0]. NumPy booleans now raise, a one-dimensional NumPy array of indexes works, and bytearray and memoryview raise. An empty members gets its own message.
  • Smaller fixes: errors name a one-group list's position, the misplaced-variable hint points at a listed entity's groups, HouseholdResult's docstring covers lists, the changelog notes the empty-list change, and the docs no longer read income_tax off a person.

Review rounds 2 and 3 (5 November 2026)

Round 2 (GPT-6.1 Sol, read-only) requested two changes, and 393a569 makes both:

  • NumPy members: tolist() turned datetime64 and timedelta64 arrays into integers before the index check ran, so np.array([0, 1], dtype="datetime64[ns]") resolved to members [0, 1]. The resolver now keeps NumPy's scalars until the check, and the tests cover datetime, timedelta and float arrays.
  • SSI wording: the docs and docstring now say SSI adds up income across the unit "unless exactly one member counts as living in a medical treatment facility", following ssi_couple_computation_applies.

Round 3 (GPT-6.1 Sol, the delta from 7c995cf to 393a569) approved with no findings. It ran the validators on all 33 resolver cases and NumPy dtypes, and ran the 2.2.1 SSI formulas across 64 facility and payment combinations.

Checks

  • tests/test_us_household_groups.py, test_household_impact.py, test_household_calculator_snapshot.py and test_spm_household.py: 141 passed.
  • tests/test_us_household_groups.py at 393a569: 55 passed.
  • The other suites that call the changed validators (test_uk_cgt_default.py, test_uk_cgt_responses.py, test_uk_dated_reforms.py, test_uk_budget_reform_contract.py, test_dict_reforms_on_simulation.py, test_spm_selection.py, test_us_reform_application.py): 243 passed.
  • ruff format --check . and ruff check . pass with ruff 0.17.3, which CI installs. The venv's older ruff 0.12.11 reports UP038 in 66 places, 64 of them already on main; this PR keeps the repository's tuple style in isinstance.
  • mypy src/policyengine reports nothing in us/household.py (CI runs mypy as informational).

No country-package pins change.

axiom: n/a: wrapper API for household group structure; no policy encoding change

🤖 Generated with Claude Code

MaxGhenis and others added 3 commits November 4, 2026 21:21
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

US calculate_household puts every person in one marital unit

1 participant