diff --git a/changelog.d/region-group-scoping-strategy.added.md b/changelog.d/region-group-scoping-strategy.added.md new file mode 100644 index 00000000..58bc2698 --- /dev/null +++ b/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. diff --git a/src/policyengine/core/scoping_strategy.py b/src/policyengine/core/scoping_strategy.py index abd117f3..3cdd3f1e 100644 --- a/src/policyengine/core/scoping_strategy.py +++ b/src/policyengine/core/scoping_strategy.py @@ -1,6 +1,6 @@ """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 @@ -8,6 +8,10 @@ 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"), ] diff --git a/src/policyengine/utils/entity_utils.py b/src/policyengine/utils/entity_utils.py index 481ff62a..3f8b54e7 100644 --- a/src/policyengine/utils/entity_utils.py +++ b/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()) diff --git a/tests/conftest.py b/tests/conftest.py index 75e2b55c..0f730005 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/test_entity_utils_characterization.py b/tests/test_entity_utils_characterization.py new file mode 100644 index 00000000..88fb1272 --- /dev/null +++ b/tests/test_entity_utils_characterization.py @@ -0,0 +1,113 @@ +"""Characterization tests for the household filter cascade, plus direct tests of +``filter_dataset_by_household_ids`` (the id-set "Job B" reused by the group). + +These pin the behavior the household-ID-set refactor must preserve: filtering +keeps exactly the matching households and cascades to persons and every group +entity consistently. Uses the ``us_test_dataset`` fixture (households 1,2 = +state_fips 6 / CA; household 3 = state_fips 34 / NJ). +""" + +import pandas as pd +import pytest + +from policyengine.utils.entity_utils import ( + filter_dataset_by_household_ids, + filter_dataset_by_household_variable, +) + + +def _ids(result, entity): + return set(pd.DataFrame(result[entity])[f"{entity}_id"]) + + +class TestFilterCascadeCharacterization: + def test__filter_ca__keeps_ca_households_and_cascades( + self, us_entity_data, us_group_entities + ): + result = filter_dataset_by_household_variable( + entity_data=us_entity_data, + group_entities=us_group_entities, + variable_name="state_fips", + variable_value=6, + ) + assert _ids(result, "household") == {1, 2} + assert set(pd.DataFrame(result["person"])["person_id"]) == {1, 2, 3, 4} + assert _ids(result, "tax_unit") == {1, 2} + assert _ids(result, "spm_unit") == {1, 2} + assert _ids(result, "family") == {1, 2} + assert _ids(result, "marital_unit") == {1, 2} + + def test__filter_nj__keeps_nj_household_and_cascades( + self, us_entity_data, us_group_entities + ): + result = filter_dataset_by_household_variable( + entity_data=us_entity_data, + group_entities=us_group_entities, + variable_name="state_fips", + variable_value=34, + ) + assert _ids(result, "household") == {3} + assert set(pd.DataFrame(result["person"])["person_id"]) == {5, 6} + assert _ids(result, "tax_unit") == {3} + assert _ids(result, "marital_unit") == {3} + + def test__filter_preserves_national_weights( + self, us_entity_data, us_group_entities + ): + result = filter_dataset_by_household_variable( + entity_data=us_entity_data, + group_entities=us_group_entities, + variable_name="state_fips", + variable_value=6, + ) + household = pd.DataFrame(result["household"]) + # National weights preserved (not reweighted): CA = 1000 + 1000. + assert float(household["household_weight"].sum()) == 2000.0 + + def test__no_match_raises_value_error(self, us_entity_data, us_group_entities): + with pytest.raises(ValueError, match="No households found"): + filter_dataset_by_household_variable( + entity_data=us_entity_data, + group_entities=us_group_entities, + variable_name="state_fips", + variable_value=99, + ) + + def test__congressional_district_filter_cascades( + self, us_entity_data, us_group_entities + ): + # CD 601 is household 1 only (persons 1,2). + result = filter_dataset_by_household_variable( + entity_data=us_entity_data, + group_entities=us_group_entities, + variable_name="congressional_district_geoid", + variable_value=601, + ) + assert _ids(result, "household") == {1} + assert set(pd.DataFrame(result["person"])["person_id"]) == {1, 2} + + +class TestFilterDatasetByHouseholdIds: + """Direct tests of the id-set entry point reused by RegionGroupStrategy.""" + + def test__explicit_ids_cascade_to_all_entities( + self, us_entity_data, us_group_entities + ): + result = filter_dataset_by_household_ids( + us_entity_data, us_group_entities, {1, 3} + ) + assert _ids(result, "household") == {1, 3} + assert set(pd.DataFrame(result["person"])["person_id"]) == {1, 2, 5, 6} + assert _ids(result, "tax_unit") == {1, 3} + + def test__accepts_any_iterable_and_dedupes(self, us_entity_data, us_group_entities): + result = filter_dataset_by_household_ids( + us_entity_data, us_group_entities, [1, 1, 2] + ) + assert _ids(result, "household") == {1, 2} + + def test__empty_selection_raises(self, us_entity_data, us_group_entities): + with pytest.raises( + ValueError, match="No households match the requested household id set" + ): + filter_dataset_by_household_ids(us_entity_data, us_group_entities, set()) diff --git a/tests/test_region_group_strategy.py b/tests/test_region_group_strategy.py new file mode 100644 index 00000000..179f84f8 --- /dev/null +++ b/tests/test_region_group_strategy.py @@ -0,0 +1,123 @@ +"""Tests for RegionGroupStrategy — scoping to the union of several regions.""" + +import pandas as pd +import pytest +from pydantic import TypeAdapter + +from policyengine.core.scoping_strategy import ( + RegionGroupStrategy, + RowFilterStrategy, + ScopingStrategy, +) + + +def _hh_ids(result): + return set(pd.DataFrame(result["household"])["household_id"]) + + +def _state(value): + return RowFilterStrategy(variable_name="state_fips", variable_value=value) + + +class TestRegionGroupStrategy: + def test__has_region_group_strategy_type(self): + assert RegionGroupStrategy(members=[_state(6)]).strategy_type == "region_group" + + def test__apply_unions_members(self, us_entity_data, us_group_entities): + # CA (households 1,2) ∪ NJ (household 3) == all households. + strategy = RegionGroupStrategy(members=[_state(6), _state(34)]) + result = strategy.apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + assert _hh_ids(result) == {1, 2, 3} + assert set(pd.DataFrame(result["person"])["person_id"]) == {1, 2, 3, 4, 5, 6} + + def test__single_member_group_matches_row_filter( + self, us_entity_data, us_group_entities + ): + grouped = RegionGroupStrategy(members=[_state(34)]).apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + row = _state(34).apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + assert _hh_ids(grouped) == _hh_ids(row) == {3} + + def test__disjoint_union_has_no_duplicates(self, us_entity_data, us_group_entities): + strategy = RegionGroupStrategy(members=[_state(6), _state(34)]) + result = strategy.apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + household = pd.DataFrame(result["household"]) + assert len(household) == 3 + assert household["household_id"].is_unique + + def test__overlapping_members_are_not_double_counted( + self, us_entity_data, us_group_entities + ): + # CA (state_fips=6 -> households 1,2) overlaps CA-01 (geoid 601 -> hh 1). + strategy = RegionGroupStrategy( + members=[ + _state(6), + RowFilterStrategy( + variable_name="congressional_district_geoid", variable_value=601 + ), + ] + ) + result = strategy.apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + household = pd.DataFrame(result["household"]) + assert _hh_ids(result) == {1, 2} # household 1 counted once, not twice + assert household["household_id"].is_unique + + def test__uk_country_group(self, uk_entity_data, uk_group_entities): + # ENGLAND (households 1,2) ∪ SCOTLAND (household 3) == all UK households. + strategy = RegionGroupStrategy( + members=[ + RowFilterStrategy(variable_name="country", variable_value="ENGLAND"), + RowFilterStrategy(variable_name="country", variable_value="SCOTLAND"), + ] + ) + result = strategy.apply( + entity_data=uk_entity_data, group_entities=uk_group_entities, year=2024 + ) + assert _hh_ids(result) == {1, 2, 3} + + def test__member_additional_filters_are_applied( + self, us_entity_data, us_group_entities + ): + # CA households additionally filtered to CD 601 -> only household 1. + member = RowFilterStrategy( + variable_name="state_fips", + variable_value=6, + additional_filters={"congressional_district_geoid": 601}, + ) + result = RegionGroupStrategy(members=[member]).apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + assert _hh_ids(result) == {1} + + def test__empty_union_raises(self, us_entity_data, us_group_entities): + strategy = RegionGroupStrategy(members=[_state(99)]) # no such state + with pytest.raises(ValueError, match="No households match"): + strategy.apply( + entity_data=us_entity_data, group_entities=us_group_entities, year=2024 + ) + + def test__cache_key_is_member_order_independent(self): + a = RegionGroupStrategy(members=[_state(6), _state(34)]) + b = RegionGroupStrategy(members=[_state(34), _state(6)]) + assert a.cache_key == b.cache_key + assert "region_group" in a.cache_key + + def test__roundtrips_through_scoping_strategy_union(self): + strategy = RegionGroupStrategy(members=[_state(6), _state(34)]) + restored = TypeAdapter(ScopingStrategy).validate_python(strategy.model_dump()) + assert isinstance(restored, RegionGroupStrategy) + assert restored.strategy_type == "region_group" + assert len(restored.members) == 2 + + def test__requires_at_least_one_member(self): + with pytest.raises(ValueError): + RegionGroupStrategy(members=[])