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..a2b4a3f48 100644 --- a/docs/python_api/simulations.md +++ b/docs/python_api/simulations.md @@ -24,6 +24,29 @@ 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: @@ -31,4 +54,3 @@ The `policyengine_core.simulations` module contains the definition of `Simulatio :inherited-members: :show-inheritance: ``` - diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index afcee0358..82f352046 100644 --- a/policyengine_core/reforms/reform.py +++ b/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 diff --git a/policyengine_core/simulations/individual_sim.py b/policyengine_core/simulations/individual_sim.py index c00933c87..4d896343f 100644 --- a/policyengine_core/simulations/individual_sim.py +++ b/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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index de5824c59..7490c1d35 100644 --- a/policyengine_core/simulations/simulation.py +++ b/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" diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index c16a767d3..c58acfa6c 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/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) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index ccadd2a39..b6be2543f 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/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 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_individual_sim.py b/tests/core/test_individual_sim.py new file mode 100644 index 000000000..ec5fe8b2d --- /dev/null +++ b/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"}) diff --git a/tests/core/test_join_with_persons_identity.py b/tests/core/test_join_with_persons_identity.py index 6ff739a88..d75c3e119 100644 --- a/tests/core/test_join_with_persons_identity.py +++ b/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(): diff --git a/tests/core/test_simulation_builder.py b/tests/core/test_simulation_builder.py index 049bedf78..bfeb7e177 100644 --- a/tests/core/test_simulation_builder.py +++ b/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()