Add changelog entry and format code

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,8 @@
+- bump: minor
+  changes:
+    added:
+    - Entity relationship approach for simulation filtering that preserves household integrity
+    - Reusable variable validation functions (`get_variable`, `validate_variable_entity`, `validate_household_variable`)
+    changed:
+    - Refactored `_filter_simulation_by_household_variable` to use explicit entity relationship mapping
+    - Place-level filtering now builds entity_rel DataFrame for cleaner filtering logic

policyengine/simulation.py

@@ -477,7 +477,9 @@ def _build_entity_relationships(
         Returns:
             A DataFrame indexed by person with columns for each entity ID.
         """
-        entity_rel = pd.DataFrame({"person_id": simulation.calculate("person_id").values})
+        entity_rel = pd.DataFrame(
+            {"person_id": simulation.calculate("person_id").values}
+        )
 
         # Add household relationship (required for all countries)
         entity_rel["household_id"] = simulation.calculate(
@@ -486,7 +488,12 @@ def _build_entity_relationships(
 
         # Add country-specific entity relationships
         tbs = simulation.tax_benefit_system
-        optional_entities = ["tax_unit_id", "spm_unit_id", "family_id", "marital_unit_id"]
+        optional_entities = [
+            "tax_unit_id",
+            "spm_unit_id",
+            "family_id",
+            "marital_unit_id",
+        ]
 
         for entity_id in optional_entities:
             if entity_id in tbs.variables:
@@ -525,7 +532,9 @@ def _filter_simulation_by_household_variable(
         Raises:
             ValueError: If the variable is not a household-level variable.
         """
-        validate_household_variable(simulation.tax_benefit_system, variable_name)
+        validate_household_variable(
+            simulation.tax_benefit_system, variable_name
+        )
 
         # Build entity relationships
         entity_rel = self._build_entity_relationships(simulation)
@@ -536,7 +545,9 @@ def _filter_simulation_by_household_variable(
 
         # Create mask for matching households, handling bytes encoding
         if isinstance(variable_value, str):
-            hh_mask = (hh_values == variable_value) | (hh_values == variable_value.encode())
+            hh_mask = (hh_values == variable_value) | (
+                hh_values == variable_value.encode()
+            )
         else:
             hh_mask = hh_values == variable_value
 

tests/fixtures/country/us_places.py

@@ -129,7 +129,9 @@
 # =============================================================================
 
 
-def create_mock_tax_benefit_system(household_variables: list[str] | None = None) -> Mock:
+def create_mock_tax_benefit_system(
+    household_variables: list[str] | None = None,
+) -> Mock:
     """Create a mock tax benefit system with variable entity information.
 
     Args:
@@ -152,12 +154,21 @@ def create_mock_tax_benefit_system(household_variables: list[str] | None = None)
         mock_tbs.variables[var_name] = mock_var
 
     # Add standard entity ID variables
-    for entity_id in ["person_id", "household_id", "tax_unit_id", "spm_unit_id", "family_id", "marital_unit_id"]:
+    for entity_id in [
+        "person_id",
+        "household_id",
+        "tax_unit_id",
+        "spm_unit_id",
+        "family_id",
+        "marital_unit_id",
+    ]:
         mock_var = Mock()
         mock_var.entity = Mock()
         # Entity IDs belong to their respective entities
         entity_name = entity_id.replace("_id", "")
-        mock_var.entity.key = entity_name if entity_name != "person" else "person"
+        mock_var.entity.key = (
+            entity_name if entity_name != "person" else "person"
+        )
         mock_tbs.variables[entity_id] = mock_var
 
     return mock_tbs
@@ -218,7 +229,12 @@ def mock_calculate(variable_name, map_to=None, period=None):
                 result.values = np.array(person_household_ids)
             else:
                 result.values = np.array(household_ids)
-        elif variable_name in ["tax_unit_id", "spm_unit_id", "family_id", "marital_unit_id"]:
+        elif variable_name in [
+            "tax_unit_id",
+            "spm_unit_id",
+            "family_id",
+            "marital_unit_id",
+        ]:
             # For simplicity, use household_id as proxy for other entity IDs
             if map_to == "person":
                 result.values = np.array(person_household_ids)
@@ -231,11 +247,13 @@ def mock_calculate(variable_name, map_to=None, period=None):
     mock_sim.calculate = mock_calculate
 
     # Mock to_input_dataframe to return person-level DataFrame
-    df = pd.DataFrame({
-        "person_id__2024": person_ids,
-        "household_id__2024": person_household_ids,
-        "place_fips__2024": person_place_fips,
-    })
+    df = pd.DataFrame(
+        {
+            "person_id__2024": person_ids,
+            "household_id__2024": person_household_ids,
+            "place_fips__2024": person_place_fips,
+        }
+    )
     mock_sim.to_input_dataframe.return_value = df
 
     return mock_sim
@@ -287,7 +305,12 @@ def mock_calculate(variable_name, map_to=None, period=None):
                 result.values = np.array(person_household_ids)
             else:
                 result.values = np.array(household_ids)
-        elif variable_name in ["tax_unit_id", "spm_unit_id", "family_id", "marital_unit_id"]:
+        elif variable_name in [
+            "tax_unit_id",
+            "spm_unit_id",
+            "family_id",
+            "marital_unit_id",
+        ]:
             if map_to == "person":
                 result.values = np.array(person_household_ids)
             else:
@@ -298,11 +321,13 @@ def mock_calculate(variable_name, map_to=None, period=None):
 
     mock_sim.calculate = mock_calculate
 
-    df = pd.DataFrame({
-        "person_id__2024": person_ids,
-        "household_id__2024": person_household_ids,
-        "place_fips__2024": person_place_fips,
-    })
+    df = pd.DataFrame(
+        {
+            "person_id__2024": person_ids,
+            "household_id__2024": person_household_ids,
+            "place_fips__2024": person_place_fips,
+        }
+    )
     mock_sim.to_input_dataframe.return_value = df
 
     return mock_sim