Add RegionGroupStrategy for union-of-regions scoping

changelog.d/region-group-scoping-strategy.added.md

@@ -0,0 +1 @@
+Add `RegionGroupStrategy`, a scoping strategy that runs one simulation over the union of several row-filter regions (e.g. multiple whole states), and refactor the household filter into a reusable `matching_household_ids` / `filter_dataset_by_household_ids` pair keyed on the stable `household_id`. This enables map-reduce national simulations that reconstruct exact whole-population results from parallel region-group runs.

src/policyengine/core/scoping_strategy.py

@@ -1,13 +1,17 @@
 """Region scoping strategies for geographic simulations.
 
-Provides two concrete strategies for scoping datasets to sub-national regions:
+Provides three concrete strategies for scoping datasets to sub-national regions:
 
 1. RowFilterStrategy: Filters dataset rows where a household variable matches
    a specific value (e.g., US states by 'state_fips', US congressional districts
    by 'congressional_district_geoid').
 
 2. WeightReplacementStrategy: Legacy strategy that replaces household weights from
    a pre-computed weight matrix resolved locally or from GCS.
+
+3. RegionGroupStrategy: Scopes to the union of several RowFilterStrategy regions
+   (e.g. multiple whole states) so one simulation covers the whole group. Used to
+   segment a national run into parallel region-group runs.
 """
 
 import logging
@@ -20,7 +24,9 @@
 from pydantic import BaseModel, Discriminator, Field
 
 from policyengine.utils.entity_utils import (
+    filter_dataset_by_household_ids,
     filter_dataset_by_household_variable,
+    matching_household_ids,
 )
 
 logger = logging.getLogger(__name__)
@@ -233,7 +239,45 @@ def cache_key(self) -> str:
         return f"weight_replacement:{self.weight_matrix_key}:{self.region_code}"
 
 
+class RegionGroupStrategy(RegionScopingStrategy):
+    """Scope to the UNION of several row-filter regions in one simulation.
+
+    Members are ordinary ``RowFilterStrategy`` regions (e.g. several whole
+    states); their households are unioned at the household level and the sim
+    runs once over that union. Because member regions never share households,
+    the union is a disjoint concatenation — no household is counted twice.
+    """
+
+    strategy_type: Literal["region_group"] = "region_group"
+    members: list[RowFilterStrategy] = Field(min_length=1)
+
+    def apply(
+        self,
+        entity_data: dict[str, MicroDataFrame],
+        group_entities: list[str],
+        year: int,
+    ) -> dict[str, MicroDataFrame]:
+        keep_household_ids: set = set()
+        for member in self.members:
+            keep_household_ids |= matching_household_ids(
+                entity_data,
+                member.variable_name,
+                member.variable_value,
+                member.additional_filters,
+            )
+        return filter_dataset_by_household_ids(
+            entity_data, group_entities, keep_household_ids
+        )
+
+    @property
+    def cache_key(self) -> str:
+        # Sorted so the key is independent of member order (deterministic
+        # simulation-ID hashing).
+        member_keys = sorted(member.cache_key for member in self.members)
+        return "region_group:" + "|".join(member_keys)
+
+
 ScopingStrategy = Annotated[
-    Union[RowFilterStrategy, WeightReplacementStrategy],
+    Union[RowFilterStrategy, WeightReplacementStrategy, RegionGroupStrategy],
     Discriminator("strategy_type"),
 ]

src/policyengine/utils/entity_utils.py

@@ -1,7 +1,7 @@
 """Shared utilities for entity relationship building and dataset filtering."""
 
 import logging
-from typing import Optional, Union
+from typing import Iterable, Optional, Union
 
 import pandas as pd
 from microdf import MicroDataFrame
@@ -52,37 +52,17 @@ def build_entity_relationships(
     return pd.DataFrame(columns)
 
 
-def filter_dataset_by_household_variable(
-    entity_data: dict[str, MicroDataFrame],
-    group_entities: list[str],
+def _household_mask(
+    household_data: pd.DataFrame,
     variable_name: str,
     variable_value: Union[str, int, float],
     additional_filters: Optional[dict[str, Union[str, int, float]]] = None,
-) -> dict[str, MicroDataFrame]:
-    """Filter dataset entities to only include households matching variables.
-
-    Uses an entity relationship approach: builds an explicit map of all
-    entity relationships, filters at the household level, and keeps all
-    persons in matching households to preserve entity integrity.
-
-    Args:
-        entity_data: Dict mapping entity names to their MicroDataFrames
-                     (from YearData.entity_data).
-        group_entities: List of group entity names for this country.
-        variable_name: The household-level variable to filter on.
-        variable_value: The value to match. Handles both str and bytes encoding.
-        additional_filters: Optional household-level filters that must also
-                            match, keyed by variable name.
+):
+    """Boolean mask over the household table for a single-value filter.
 
-    Returns:
-        A dict mapping entity names to filtered MicroDataFrames.
-
-    Raises:
-        ValueError: If variable_name is not found or no households match.
+    Local intermediate only — never crosses a function boundary — so no
+    positional-alignment invariant leaks out (callers key on household_id).
     """
-    person_data = pd.DataFrame(entity_data["person"])
-    household_data = pd.DataFrame(entity_data["household"])
-
     if variable_name not in household_data.columns:
         raise ValueError(
             f"Variable '{variable_name}' not found in household data. "
@@ -96,33 +76,63 @@ def filter_dataset_by_household_variable(
                 f"Available columns: {list(household_data.columns)}"
             )
 
-    # Build entity relationships
-    entity_rel = build_entity_relationships(person_data, group_entities)
+    mask = _values_match(household_data[variable_name].values, variable_value)
+    for extra_variable, extra_value in additional_filters.items():
+        mask &= _values_match(household_data[extra_variable].values, extra_value)
+    return mask
 
-    # Find matching household IDs
-    hh_ids = household_data["household_id"].values
 
-    hh_mask = _values_match(household_data[variable_name].values, variable_value)
-    for extra_variable, extra_value in additional_filters.items():
-        hh_mask &= _values_match(household_data[extra_variable].values, extra_value)
+def matching_household_ids(
+    entity_data: dict[str, MicroDataFrame],
+    variable_name: str,
+    variable_value: Union[str, int, float],
+    additional_filters: Optional[dict[str, Union[str, int, float]]] = None,
+) -> set:
+    """Return the set of ``household_id`` values matching a household filter.
 
-    matching_hh_ids = set(hh_ids[hh_mask])
+    This is the stable, order-proof boundary between "which households" and the
+    entity cascade: the household selection is expressed as a set of household
+    IDs, never a positional mask.
+    """
+    household_data = pd.DataFrame(entity_data["household"])
+    mask = _household_mask(
+        household_data, variable_name, variable_value, additional_filters
+    )
+    return set(household_data["household_id"].values[mask])
 
-    if len(matching_hh_ids) == 0:
-        raise ValueError(
-            f"No households found matching {variable_name}={variable_value}"
-        )
 
-    # Filter persons to those in matching households
-    person_mask = entity_rel["household_id"].isin(matching_hh_ids)
+def filter_dataset_by_household_ids(
+    entity_data: dict[str, MicroDataFrame],
+    group_entities: list[str],
+    keep_household_ids: Iterable,
+) -> dict[str, MicroDataFrame]:
+    """Filter every entity to the given households, cascading to persons and
+    all group entities.
+
+    The selection is keyed on the stable ``household_id`` value (a household
+    table primary key), so it is independent of row order. Both the single-value
+    filter and the multi-region (union) group strategy funnel through here, so
+    the cascade lives in exactly one place.
+    """
+    keep_household_ids = set(keep_household_ids)
+    # Guard the id-set entry point (direct callers and RegionGroupStrategy). The
+    # single-variable wrapper guards earlier with a variable-specific message, so
+    # this generic message is only surfaced when a caller passes an empty set.
+    if len(keep_household_ids) == 0:
+        raise ValueError("No households match the requested household id set.")
+
+    person_data = pd.DataFrame(entity_data["person"])
+    entity_rel = build_entity_relationships(person_data, group_entities)
+
+    # Keep persons whose household is selected, then collect the entity IDs
+    # those persons carry (household included, since it is a group entity).
+    person_mask = entity_rel["household_id"].isin(keep_household_ids)
     filtered_rel = entity_rel[person_mask]
 
-    # Collect filtered IDs for each entity
     filtered_ids = {"person": set(filtered_rel["person_id"])}
     for entity in group_entities:
         filtered_ids[entity] = set(filtered_rel[f"{entity}_id"])
 
-    # Filter each entity DataFrame
     result = {}
     for entity_name, mdf in entity_data.items():
         df = pd.DataFrame(mdf)
@@ -149,6 +159,48 @@ def filter_dataset_by_household_variable(
     return result
 
 
+def filter_dataset_by_household_variable(
+    entity_data: dict[str, MicroDataFrame],
+    group_entities: list[str],
+    variable_name: str,
+    variable_value: Union[str, int, float],
+    additional_filters: Optional[dict[str, Union[str, int, float]]] = None,
+) -> dict[str, MicroDataFrame]:
+    """Filter dataset entities to only include households matching variables.
+
+    Thin wrapper composing :func:`matching_household_ids` (which households) with
+    :func:`filter_dataset_by_household_ids` (the entity cascade). Public behavior
+    is unchanged.
+
+    Args:
+        entity_data: Dict mapping entity names to their MicroDataFrames
+                     (from YearData.entity_data).
+        group_entities: List of group entity names for this country.
+        variable_name: The household-level variable to filter on.
+        variable_value: The value to match. Handles both str and bytes encoding.
+        additional_filters: Optional household-level filters that must also
+                            match, keyed by variable name.
+
+    Returns:
+        A dict mapping entity names to filtered MicroDataFrames.
+
+    Raises:
+        ValueError: If variable_name is not found or no households match.
+    """
+    keep_household_ids = matching_household_ids(
+        entity_data, variable_name, variable_value, additional_filters
+    )
+    # Variable-specific message (asserted by tests and referenced by a documented
+    # incident); raised here before the generic id-set guard is ever reached.
+    if len(keep_household_ids) == 0:
+        raise ValueError(
+            f"No households found matching {variable_name}={variable_value}"
+        )
+    return filter_dataset_by_household_ids(
+        entity_data, group_entities, keep_household_ids
+    )
+
+
 def _values_match(values, expected: Union[str, int, float]):
     if isinstance(expected, str):
         return (values == expected) | (values == expected.encode())

tests/conftest.py

@@ -1,7 +1,11 @@
 """Pytest configuration and shared fixtures."""
 
+import pytest
+
 # Import fixtures from fixtures module so pytest can discover them
 from tests.fixtures.filtering_fixtures import (  # noqa: F401
+    create_uk_test_dataset,
+    create_us_test_dataset,
     uk_test_dataset,
     us_test_dataset,
 )
@@ -22,3 +26,34 @@
     high_income_single_filer,
     married_couple_with_kids,
 )
+
+# Group (non-person) entities per country, and a helper to extract the
+# entity_data dict a scoping strategy / filter operates on. Shared so the filter
+# and region-group tests don't each redefine them.
+US_GROUP_ENTITIES = ["household", "tax_unit", "spm_unit", "family", "marital_unit"]
+UK_GROUP_ENTITIES = ["benunit", "household"]
+
+
+def entity_data_of(dataset, group_entities):
+    data = dataset.data
+    return {entity: getattr(data, entity) for entity in ["person", *group_entities]}
+
+
+@pytest.fixture
+def us_group_entities():
+    return list(US_GROUP_ENTITIES)
+
+
+@pytest.fixture
+def uk_group_entities():
+    return list(UK_GROUP_ENTITIES)
+
+
+@pytest.fixture
+def us_entity_data():
+    return entity_data_of(create_us_test_dataset(), US_GROUP_ENTITIES)
+
+
+@pytest.fixture
+def uk_entity_data():
+    return entity_data_of(create_uk_test_dataset(), UK_GROUP_ENTITIES)