Add country situation preprocessing and supplied membership provenance

changelog.d/situation-preprocessing.added.md

@@ -0,0 +1 @@
+Add country-specific entity-situation preprocessing and retain which group memberships were explicitly supplied by a situation or dataset.

docs/python_api/simulations.md

@@ -24,11 +24,33 @@ The `policyengine_core.simulations` module contains the definition of `Simulatio
 
 ## SimulationBuilder
 
+Country models can override `TaxBenefitSystem.preprocess_situation(situation,
+default_period)` to fill missing groups before an entity-form situation is
+built. The hook receives a deep copy after singular entity aliases have been
+expanded, with any `axes` still present, and must return a dict for the
+builder's normal validation. `default_period` is the builder's default period
+as set, so it may be `None`, a string, an integer or a `Period`. Both
+`Simulation(situation=...)` and the YAML runner use this path, and a `Reform`
+delegates the hook to its baseline. Dataset and variable-only construction do
+not call the hook. The default hook returns its input unchanged.
+
+`Simulation.input_group_entities` is an immutable set of singular group keys
+supplied before preprocessing. An explicitly empty group mapping counts as
+supplied. Synthesized groups, the singleton groups created for variable-only
+input, and the placeholder and default groups `IndividualSim` adds do not. For
+datasets, it records groups declared by person membership columns or, in flat
+files, by the group's ID column on each person row. Core's loaders require
+membership for every group, so a loaded dataset records all of them; a country
+that overrides `Simulation.build_from_dataset` must set the attribute itself.
+Country formulas can use this provenance to distinguish reported relationships
+from defaults without inspecting group sizes or identifier values.
+Population-only construction does not assert input provenance; simulation
+clones preserve their source's set.
+
 ```{eval-rst}
 .. autoclass:: policyengine_core.simulations.simulation_builder.SimulationBuilder
     :members:
     :undoc-members:
     :inherited-members:
     :show-inheritance:
 ```
-

policyengine_core/reforms/reform.py

@@ -147,6 +147,11 @@ def __init__(self, baseline: TaxBenefitSystem):
     def __getattr__(self, attribute):
         return getattr(self.baseline, attribute)
 
+    def preprocess_situation(self, situation: dict, default_period=None) -> dict:
+        # TaxBenefitSystem defines a no-op hook, so ``__getattr__`` never
+        # reaches the baseline's override; delegate to it explicitly.
+        return self.baseline.preprocess_situation(situation, default_period)
+
     @property
     def full_key(self) -> str:
         key = self.key

policyengine_core/simulations/individual_sim.py

@@ -54,6 +54,13 @@ def __init__(
             setattr(self, f"add_{entity}", partial(self.add_data, entity=entity))
 
     def build(self):
+        # The empty mappings seeded in __init__ and the required groups
+        # filled in below are not groups the caller supplied.
+        supplied_groups = frozenset(
+            entity.key
+            for entity in self.system.group_entities
+            if self.situation_data.get(entity.plural)
+        )
         if self.required_entities is not None:
             # Check for missing entities
             entities = {entity.key: entity for entity in self.system.entities}
@@ -80,6 +87,9 @@ def build(self):
             situation=self.situation_data,
             reform=self.reform,
         )
+        for simulation in (self.simulation, getattr(self.simulation, "baseline", None)):
+            if simulation is not None:
+                simulation.input_group_entities = supplied_groups
         self.simulation.trace = True
         self.sim = self.simulation
 

policyengine_core/simulations/simulation.py

@@ -725,6 +725,18 @@ class Simulation:
     is_over_dataset: bool = False
     """Whether this simulation is built over a dataset."""
 
+    input_group_entities: frozenset[str] = frozenset()
+    """Singular keys of group entities explicitly supplied by the caller.
+
+    Entity-form situations record groups before country preprocessing.
+    Datasets record groups declared by person membership columns or, in
+    flat files, by the group's ID column on each person row; Core's loaders
+    require membership for every group, so a loaded dataset records all of
+    them, and a country that overrides ``build_from_dataset`` must set this
+    itself. Variable-only and population-only construction do not declare
+    input groups. This immutable provenance survives cloning.
+    """
+
     macro_cache_read: bool = False
     """Whether to read from the macro cache."""
 
@@ -1134,6 +1146,18 @@ def get_eternity_array(name):
         entity_ids = get_eternity_array(entity_id_field)
         builder.declare_person_entity(person_entity.key, entity_ids)
 
+        # A dataset declares a group through its person membership column
+        # or, in a flat file, the group's ID column on each person row.
+        membership_columns = {column.split("__")[0] for column in data}
+        self.input_group_entities = frozenset(
+            group_entity.key
+            for group_entity in self.tax_benefit_system.group_entities
+            if f"{person_entity.key}_{group_entity.key}_id" in membership_columns
+            or (
+                self.dataset.data_format == Dataset.FLAT_FILE
+                and f"{group_entity.key}_id" in membership_columns
+            )
+        )
         for group_entity in self.tax_benefit_system.group_entities:
             entity_id_field = f"{group_entity.key}_id"
             person_membership_id_field = f"{person_entity.key}_{group_entity.key}_id"

policyengine_core/simulations/simulation_builder.py

@@ -226,6 +226,19 @@ def build_from_entities(
             )
 
         check_type(input_dict, dict, ["error"])
+        simulation.input_group_entities = frozenset(
+            entity.key
+            for entity in tax_benefit_system.group_entities
+            if input_dict.get(entity.plural) is not None
+        )
+        input_dict = tax_benefit_system.preprocess_situation(
+            input_dict, self.default_period
+        )
+        if not isinstance(input_dict, dict):
+            raise TypeError(
+                f"{type(tax_benefit_system).__name__}.preprocess_situation must "
+                f"return a dict, not {type(input_dict).__name__}."
+            )
         axes = input_dict.pop("axes", None)
 
         unexpected_entities = [
@@ -350,6 +363,7 @@ def build_default_simulation(
                 tax_benefit_system=tax_benefit_system,
                 populations=tax_benefit_system.instantiate_entities(),
             )
+        simulation.input_group_entities = frozenset()
         for population in simulation.populations.values():
             population.count = count
             population.ids = np.arange(id_count, dtype=np.int64)

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -150,6 +150,22 @@ class TaxBenefitSystem:
     _defined_for_checks_deferred: bool = False
     """Whether ``load_variable`` leaves ``defined_for`` checks to a later pass over every variable."""
 
+    def preprocess_situation(self, situation: dict, default_period=None) -> dict:
+        """Normalize entity-form input before the builder creates populations.
+
+        Country models can override this to fill missing entity definitions.
+        The builder passes a deep copy of the caller's situation, with singular
+        entity aliases expanded and any ``axes`` still present, and its default
+        period as set (None, a string, an integer or a ``Period``). Return a
+        dict holding an entity-form situation for the builder's usual
+        validation. A ``Reform`` delegates this hook to its baseline. Dataset
+        and variable-only inputs do not use this hook.
+
+        ``Simulation.input_group_entities`` records the groups supplied by
+        the caller before this hook runs, separately from inferred groups.
+        """
+        return situation
+
     def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None:
         if entities is None:
             entities = self.entities

tests/core/test_dataset_unknown_column_warning.py

@@ -38,6 +38,7 @@ def test__given_unknown_dataset_column__then_warns_and_still_loads(caplog):
 
     # And the known columns still load normally
     np.testing.assert_array_equal(salary.values, np.array([100.0, 200.0, 300.0]))
+    assert simulation.input_group_entities == frozenset({"household"})
 
 
 def test__given_only_known_columns__then_no_unknown_column_warning(caplog):

tests/core/test_individual_sim.py

@@ -0,0 +1,26 @@
+from policyengine_core.country_template import CountryTaxBenefitSystem
+from policyengine_core.simulations.individual_sim import IndividualSim
+
+
+class TemplateIndividualSim(IndividualSim):
+    tax_benefit_system = CountryTaxBenefitSystem
+    required_entities = ["household"]
+    default_roles = {"household": "parent"}
+
+
+def test_placeholder_and_default_groups_are_not_supplied():
+    # __init__ seeds an empty mapping for every group and build() fills the
+    # required household with every person; neither was supplied by the
+    # caller.
+    simulation = TemplateIndividualSim(year=2024)
+    simulation.add_person(name="Alicia")
+    simulation.build()
+    assert simulation.simulation.input_group_entities == frozenset()
+
+
+def test_added_groups_are_supplied():
+    simulation = TemplateIndividualSim(year=2024)
+    simulation.add_person(name="Alicia")
+    simulation.add_household(name="home", parents=["Alicia"])
+    simulation.build()
+    assert simulation.simulation.input_group_entities == frozenset({"household"})

tests/core/test_join_with_persons_identity.py

@@ -99,6 +99,7 @@ class InteriorEmptyHousehold(Dataset):
 
     household = simulation.populations["household"]
     assert household.members_entity_id.tolist() == [0, 2]
+    assert simulation.input_group_entities == frozenset({"household"})
     assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [
         100.0,
         300.0,
@@ -189,6 +190,8 @@ def test_flat_file_with_only_a_household_id_column_reads_it_as_membership(
     household = simulation.populations["household"]
     assert household.count == 2
     assert household.members_entity_id.tolist() == [0, 0, 1]
+    # The ID column declares the households, so they count as supplied.
+    assert simulation.input_group_entities == frozenset({"household"})
 
 
 def test_flat_file_default_roles_cover_every_person():

tests/core/test_simulation_builder.py

@@ -1,4 +1,5 @@
 import datetime
+import copy
 from decimal import Decimal
 from fractions import Fraction
 from typing import Iterable
@@ -7,12 +8,20 @@
 import pytest
 
 from policyengine_core import periods, tools
-from policyengine_core.country_template import entities, situation_examples
+from policyengine_core.country_template import (
+    CountryTaxBenefitSystem,
+    entities,
+    situation_examples,
+)
+from policyengine_core.entities import build_entity
 from policyengine_core.enums import Enum
 from policyengine_core.errors import SituationParsingError
 from policyengine_core.populations import Population
+from policyengine_core.reforms import Reform
 from policyengine_core.simulations import Simulation, SimulationBuilder
+from policyengine_core.taxbenefitsystems import TaxBenefitSystem
 from policyengine_core.tools import test_runner
+from policyengine_core.tools.test_runner import _get_tax_benefit_system
 from policyengine_core.variables import Variable
 
 
@@ -125,6 +134,124 @@ def test_explicit_singular_entities(tax_benefit_system):
     }
 
 
+def test_country_preprocessing_preserves_caller_input_and_group_provenance(
+    tax_benefit_system, monkeypatch
+):
+    seen_periods = []
+
+    def preprocess(situation, default_period):
+        seen_periods.append(default_period)
+        situation.setdefault(
+            "households", {"inferred": {"parents": ["Alicia", "Javier"]}}
+        )
+        return situation
+
+    monkeypatch.setattr(tax_benefit_system, "preprocess_situation", preprocess)
+    situation = {"persons": {"Alicia": {}, "Javier": {}}}
+    original = copy.deepcopy(situation)
+    builder = SimulationBuilder()
+    builder.set_default_period("2024-01")
+    inferred = builder.build_from_dict(tax_benefit_system, situation)
+    assert list(inferred.household.ids) == ["inferred"]
+    assert inferred.input_group_entities == frozenset()
+    assert situation == original
+    assert seen_periods == ["2024-01"]
+
+    # Singular aliases are normalized before the hook; explicit membership
+    # remains distinguishable even if it has the same shape as the default.
+    explicit = builder.build_from_dict(
+        tax_benefit_system,
+        {**situation, "household": {"parents": ["Alicia", "Javier"]}},
+    )
+    assert explicit.input_group_entities == frozenset({"household"})
+    assert explicit.clone().input_group_entities == explicit.input_group_entities
+
+    # Reusing a simulation for variable-only construction resets provenance
+    # and does not run entity preprocessing.
+    calls = len(seen_periods)
+    builder.build_default_simulation(tax_benefit_system, 2, explicit)
+    assert explicit.input_group_entities == frozenset()
+    assert len(seen_periods) == calls
+
+
+def test_country_preprocessing_applies_to_simulation_api(
+    tax_benefit_system, monkeypatch
+):
+    def preprocess(situation, default_period):
+        situation["households"] = {
+            name: {"parents": [name]} for name in situation["persons"]
+        }
+        return situation
+
+    monkeypatch.setattr(tax_benefit_system, "preprocess_situation", preprocess)
+    simulation = Simulation(
+        tax_benefit_system=tax_benefit_system,
+        situation={"persons": {"Alicia": {}, "Javier": {}}},
+    )
+    assert simulation.household.members_entity_id.tolist() == [0, 1]
+    assert simulation.input_group_entities == frozenset()
+
+
+def test_explicit_empty_group_mapping_is_recorded(tax_benefit_system):
+    simulation = SimulationBuilder().build_from_dict(
+        tax_benefit_system, {"persons": {"Alicia": {}}, "households": {}}
+    )
+    assert simulation.input_group_entities == frozenset({"household"})
+
+
+def test_only_supplied_groups_are_recorded():
+    person = build_entity("person", "persons", "", is_person=True)
+    household = build_entity(
+        "household", "households", "", roles=[{"key": "member", "plural": "members"}]
+    )
+    family = build_entity(
+        "family", "families", "", roles=[{"key": "member", "plural": "members"}]
+    )
+    simulation = SimulationBuilder().build_from_dict(
+        TaxBenefitSystem([person, household, family]),
+        {"persons": {"Alicia": {}}, "households": {"home": {"members": ["Alicia"]}}},
+    )
+    assert simulation.input_group_entities == frozenset({"household"})
+
+
+class _PreprocessingTaxBenefitSystem(CountryTaxBenefitSystem):
+    def preprocess_situation(self, situation, default_period=None):
+        situation.setdefault("households", {"inferred": {"parents": ["Alicia"]}})
+        return situation
+
+
+class _NoOpReform(Reform):
+    def apply(self):
+        pass
+
+
+def test_reform_delegates_preprocessing_to_its_baseline():
+    # Reform inherits TaxBenefitSystem's no-op hook, so a reform, including
+    # the one the YAML runner builds for inline parameters, must delegate.
+    baseline = _PreprocessingTaxBenefitSystem()
+    for system in (
+        _NoOpReform(baseline),
+        _get_tax_benefit_system(baseline, [_NoOpReform], []),
+    ):
+        simulation = SimulationBuilder().build_from_dict(
+            system, {"persons": {"Alicia": {}}}
+        )
+        assert list(simulation.household.ids) == ["inferred"]
+        assert simulation.input_group_entities == frozenset()
+
+
+def test_country_preprocessing_must_return_a_dict(tax_benefit_system, monkeypatch):
+    monkeypatch.setattr(
+        tax_benefit_system,
+        "preprocess_situation",
+        lambda situation, default_period: None,
+    )
+    with pytest.raises(TypeError, match="preprocess_situation must return a dict"):
+        SimulationBuilder().build_from_dict(
+            tax_benefit_system, {"persons": {"Alicia": {}}}
+        )
+
+
 def test_add_person_entity(persons):
     persons_json = {"Alicia": {"salary": {}}, "Javier": {}}
     simulation_builder = SimulationBuilder()