From 6d839a148cee47981e8289de37da276f3e7beeec Mon Sep 17 00:00:00 2001 From: ZimingHua Date: Wed, 7 Oct 2026 15:51:11 -0400 Subject: [PATCH] Add country situation preprocessing and membership provenance --- changelog.d/situation-preprocessing.added.md | 1 + docs/python_api/simulations.md | 18 ++++- policyengine_core/simulations/simulation.py | 15 +++++ .../simulations/simulation_builder.py | 10 +++ .../taxbenefitsystems/tax_benefit_system.py | 14 ++++ .../test_dataset_unknown_column_warning.py | 1 + tests/core/test_simulation_builder.py | 66 +++++++++++++++++++ 7 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 changelog.d/situation-preprocessing.added.md diff --git a/changelog.d/situation-preprocessing.added.md b/changelog.d/situation-preprocessing.added.md new file mode 100644 index 000000000..a8dfee87d --- /dev/null +++ b/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. diff --git a/docs/python_api/simulations.md b/docs/python_api/simulations.md index 0b65c1dd9..ab300a966 100644 --- a/docs/python_api/simulations.md +++ b/docs/python_api/simulations.md @@ -24,6 +24,23 @@ 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: @@ -31,4 +48,3 @@ The `policyengine_core.simulations` module contains the definition of `Simulatio :inherited-members: :show-inheritance: ``` - diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 123a5f1ae..9637f4f67 100644 --- a/policyengine_core/simulations/simulation.py +++ b/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: diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index e3dee0cc5..f728f99dc 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/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)) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index b7073081d..7c45dfdd3 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/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 diff --git a/tests/core/test_dataset_unknown_column_warning.py b/tests/core/test_dataset_unknown_column_warning.py index 7450b92dc..d621cfa70 100644 --- a/tests/core/test_dataset_unknown_column_warning.py +++ b/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): diff --git a/tests/core/test_simulation_builder.py b/tests/core/test_simulation_builder.py index f838d6b3c..f11f34dc8 100644 --- a/tests/core/test_simulation_builder.py +++ b/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()