Let US calculate_household take lists of groups

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 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.

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.

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,63 @@ 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, 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(
+    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, 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:
+
+```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.
+
+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
 
@@ -229,13 +281,14 @@ 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
 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
 

src/policyengine/tax_benefit_models/common/household.py

@@ -1,17 +1,23 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Collection, Mapping, Sequence
 from typing import Any
 
 
 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}"

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: