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,27 @@ 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, and returns a situation for the builder's normal validation. Both
+`Simulation(situation=...)` and the YAML runner use this path. 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. For datasets, it identifies groups with person
+membership columns, before loader defaults are inserted. An explicitly empty
+group mapping counts as supplied. Synthesized groups and the singleton groups
+created for variable-only input do not. 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/simulations/simulation.py

@@ -220,6 +220,15 @@ 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 with person membership columns before defaults
+    are inserted. 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."""
 
@@ -577,6 +586,12 @@ def get_eternity_array(name):
         entity_ids = get_eternity_array(entity_id_field)
         builder.declare_person_entity(person_entity.key, entity_ids)
 
+        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
+        )
         for group_entity in self.tax_benefit_system.group_entities:
             entity_id_field = f"{group_entity.key}_id"
             if self.dataset.data_format != Dataset.FLAT_FILE:

policyengine_core/simulations/simulation_builder.py

@@ -126,6 +126,15 @@ def build_from_entities(
                 simulation.get_variable_population(variable_name).entity,
             )
 
+        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
+        )
         check_type(input_dict, dict, ["error"])
         axes = input_dict.pop("axes", None)
 
@@ -243,6 +252,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.array(range(count))

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -101,6 +101,20 @@ class TaxBenefitSystem:
     modelled_policies: str = None
     """A YAML filepath containing metadata describing the modelled policies."""
 
+    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 the default input period (possibly None).
+        Return an entity-form situation for the builder's usual validation.
+        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_simulation_builder.py

@@ -1,4 +1,5 @@
 import datetime
+import copy
 from typing import Iterable
 
 import pytest
@@ -95,6 +96,71 @@ 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_add_person_entity(persons):
     persons_json = {"Alicia": {"salary": {}}, "Javier": {}}
     simulation_builder = SimulationBuilder()