From f93f62d9ae004832908786975e6d02b4680dffee Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 14:20:40 -0400 Subject: [PATCH 1/2] Replace a simulation's population when the builder writes a situation into it SimulationBuilder.build_from_dict (and build_from_entities, build_from_variables and build_default_simulation) can write a situation into an existing simulation. It wrote the situation into that simulation's populations as they were. A simulation built over a dataset and then given a situation: - kept is_over_dataset and dataset, so the macro cache and country models that read is_over_dataset treated the situation as data; - answered from its fast cache with values calculated for the dataset; - kept the dataset's arrays for every variable the situation did not set; - kept the dataset's input record, so to_input_dataframe failed; - kept its branches, the baseline among them, on the dataset. Once the situation parses, the builder now gives the simulation new populations (Simulation._replace_population), with no dataset, cache, input record or store history from the old ones. After writing it records the input variables and axes and rebuilds the baseline branch on the situation (Simulation._rebuild_branches, moved out of subsample, which rebuilt it the same way). subsample now says it needs a dataset instead of failing on the missing one. Co-Authored-By: Claude Opus 5.5 --- .../builder-replaces-population.fixed.md | 1 + policyengine_core/simulations/simulation.py | 105 +++++--- .../simulations/simulation_builder.py | 55 +++- .../core/test_builder_replaces_population.py | 239 ++++++++++++++++++ ...st_builder_replaces_population_property.py | 227 +++++++++++++++++ 5 files changed, 596 insertions(+), 31 deletions(-) create mode 100644 changelog.d/builder-replaces-population.fixed.md create mode 100644 tests/core/test_builder_replaces_population.py create mode 100644 tests/core/test_builder_replaces_population_property.py diff --git a/changelog.d/builder-replaces-population.fixed.md b/changelog.d/builder-replaces-population.fixed.md new file mode 100644 index 000000000..5308f16af --- /dev/null +++ b/changelog.d/builder-replaces-population.fixed.md @@ -0,0 +1 @@ +Writing a situation into an existing simulation with `SimulationBuilder` now replaces its population: it is no longer over a dataset (`is_over_dataset` is false and `dataset` is `None`), keeps nothing loaded or calculated for the old population, and rebuilds its baseline branch on the situation; `subsample` explains that it needs a dataset instead of failing. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ddda2ab94..86fb1fb61 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1077,6 +1077,68 @@ def build_from_populations(self, populations: Dict[str, Population]) -> None: self.link_to_entities_instances() self.create_shortcuts() + def _replace_population(self) -> None: + """Give this simulation new, empty populations, and forget the old ones. + + ``SimulationBuilder`` calls this before it writes a situation into a + simulation it was given, which may hold a population loaded from a + dataset and values calculated for it. Writing the situation into the + old populations kept the old population's arrays for every variable + the situation did not set, the fast cache and the input record kept + answering for it, and ``dataset`` and ``is_over_dataset`` still named + the dataset, so the macro cache, and country models that treat a + simulation over data differently, treated the situation as data. + + The simulation keeps its tax-benefit system and settings. Its branches + keep the old populations until :meth:`_rebuild_branches`. + """ + self.dataset = None + self.is_over_dataset = False + self.situation_input = None + self.has_axes = False + self.input_variables = [] + self._store_history = StoreHistory() + self._user_input_keys = set() + self.invalidated_caches = set() + self._fast_cache = {} + self.build_from_populations(self.tax_benefit_system.instantiate_entities()) + + def _rebuild_branches(self) -> None: + """Recreate this simulation's branches on its current populations. + + Called after the populations are replaced in place (``subsample``, or + ``SimulationBuilder`` writing a situation into this simulation). + """ + # Other branches still contain the old populations. Forget them so a + # formula's next ``get_branch`` starts from the new ones. Keep the + # baseline long enough to preserve its policy below. + self.branches = { + name: branch for name, branch in self.branches.items() if name == "baseline" + } + + # Ensure the baseline branch has the new data: rebuild it from this + # simulation through ``get_branch`` (the same wiring ``__init__`` + # uses), then restore the saved baseline tax-benefit system. + # Previously ``subsample`` assigned the saved system to + # ``self.branches["tax_benefit_system"]`` — a stray dict key — so the + # rebuilt baseline branch kept the reform system and reform-vs-baseline + # comparisons after ``subsample`` compared the reform against itself. + if "baseline" in self.branches: + baseline_tax_benefit_system = self.branches["baseline"].tax_benefit_system + del self.branches["baseline"] + baseline = self.get_branch("baseline") + # As in ``__init__``, the branch is traced in this simulation, + # uses the baseline system's entities and variables, and has no + # baseline of its own: ``get_branch`` gave it this simulation's, + # the branch it replaces, which kept the old population alive. + baseline.trace = self.trace + baseline.tracer = self.tracer + baseline.tax_benefit_system = baseline_tax_benefit_system + baseline._bind_to_tax_benefit_system() + baseline.baseline = None + if getattr(self, "baseline", None) is not None: + self.baseline = baseline + def build_from_dataset(self) -> None: """Build a simulation from a dataset.""" self.build_from_populations(self.tax_benefit_system.instantiate_entities()) @@ -3384,7 +3446,19 @@ def subsample( Returns: Simulation: The quantized simulation. + + Raises: + ValueError: if the simulation has no dataset, as one built from a + situation, or whose data a situation replaced, has none. """ + if getattr(self, "dataset", None) is None: + # The sample is rebuilt as a dataset in the dataset's period, and + # seeded with its name by default. + raise ValueError( + "Only a simulation with a dataset can be subsampled. This one " + "has none: it was built from a situation or populations, or a " + "situation replaced its data." + ) default_calculation_period = self.default_calculation_period # Set default key if not provided if seed is None: @@ -3567,35 +3641,8 @@ def retain_structure(name, values): # as "size X != Y = count" projection errors. self._invalidate_all_caches() - # Other branches still contain the pre-subsample populations. Forget - # them so a formula's next ``get_branch`` starts from the new sample. - # Keep the baseline long enough to preserve its policy below. - self.branches = { - name: branch for name, branch in self.branches.items() if name == "baseline" - } - - # Ensure the baseline branch has the new data: rebuild it from the - # subsampled simulation through ``get_branch`` (the same wiring - # ``__init__`` uses), then restore the saved baseline tax-benefit - # system. Previously the saved system was assigned to - # ``self.branches["tax_benefit_system"]`` — a stray dict key — so the - # rebuilt baseline branch kept the reform system and reform-vs-baseline - # comparisons after ``subsample`` compared the reform against itself. - if "baseline" in self.branches: - baseline_tax_benefit_system = self.branches["baseline"].tax_benefit_system - del self.branches["baseline"] - baseline = self.get_branch("baseline") - # As in ``__init__``, the branch is traced in this simulation, - # uses the baseline system's entities and variables, and has no - # baseline of its own: ``get_branch`` gave it this simulation's, - # the branch it replaces, which kept the old population alive. - baseline.trace = self.trace - baseline.tracer = self.tracer - baseline.tax_benefit_system = baseline_tax_benefit_system - baseline._bind_to_tax_benefit_system() - baseline.baseline = None - if getattr(self, "baseline", None) is not None: - self.baseline = baseline + # Branches still contain the pre-subsample populations. + self._rebuild_branches() self.default_calculation_period = default_calculation_period return self diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index e3dee0cc5..6afc49f6a 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -73,8 +73,15 @@ def build_from_dict( Build a simulation from ``input_dict``. Optionally overwrites an existing simulation. This method uses :any:`build_from_entities` if entities are fully specified, or :any:`build_from_variables` if not. + + Overwriting replaces the simulation's population with the situation's + (see :meth:`Simulation._replace_population`): it is no longer over a + dataset, and holds nothing loaded or calculated for its old population. + It keeps its tax-benefit system and settings, and its baseline branch + is rebuilt on the new population. """ + replacing = simulation is not None input_dict = self.explicit_singular_entities(tax_benefit_system, input_dict) if any( key in tax_benefit_system.entities_plural() for key in input_dict.keys() @@ -87,13 +94,28 @@ def build_from_dict( tax_benefit_system, input_dict, simulation ) + if not replacing: + # A replaced simulation's were recorded before its branches were + # rebuilt (see ``_finish_replacing``). + self._record_input_variables(simulation) + + return simulation + + def _record_input_variables(self, simulation: Simulation) -> None: simulation.input_variables = [ variable.name for variable in simulation.tax_benefit_system.variables.values() if len(simulation.get_holder(variable.name).get_known_periods()) > 0 ] - return simulation + def _finish_replacing(self, simulation: Simulation) -> None: + """Finish writing a situation into a simulation whose population it replaced.""" + simulation.has_axes = self.has_axes + self._record_input_variables(simulation) + # Its branches still hold the old population. Rebuilt after the + # inputs are recorded, the baseline records them too, as in + # ``Simulation.__init__``. + simulation._rebuild_branches() def build_from_entities( self, @@ -113,6 +135,7 @@ def build_from_entities( """ input_dict = copy.deepcopy(input_dict) + replacing = simulation is not None if simulation is None: simulation = Simulation( tax_benefit_system=tax_benefit_system, @@ -177,6 +200,10 @@ def build_from_entities( self.axes = axes self.expand_axes() + if replacing: + # The situation parsed: write it into new populations. + simulation._replace_population() + try: self.finalize_variables_init(simulation.persons) except PeriodMismatchError as e: @@ -189,6 +216,9 @@ def build_from_entities( except PeriodMismatchError as e: self.raise_period_mismatch(population.entity, instances_json, e) + if replacing: + self._finish_replacing(simulation) + return simulation def build_from_variables( @@ -209,7 +239,8 @@ def build_from_variables( ) """ count = _get_person_count(input_dict) - simulation = self.build_default_simulation( + replacing = simulation is not None + simulation = self._build_default_populations( tax_benefit_system, count, simulation ) for variable, value in input_dict.items(): @@ -223,6 +254,8 @@ def build_from_variables( else: for period_str, dated_value in value.items(): simulation.set_input(variable, period_str, dated_value) + if replacing: + self._finish_replacing(simulation) return simulation def build_default_simulation( @@ -236,13 +269,31 @@ def build_default_simulation( - There are ``count`` persons - There are ``count`` instances of each group entity, containing one person - Every person has, in each entity, the first role + + Given a simulation, replace its population with this one, as + :meth:`build_from_dict` does. """ + replacing = simulation is not None + simulation = self._build_default_populations( + tax_benefit_system, count, simulation + ) + if replacing: + self._finish_replacing(simulation) + return simulation + def _build_default_populations( + self, + tax_benefit_system: "TaxBenefitSystem", + count: int, + simulation: Simulation = None, + ) -> Simulation: if simulation is None: simulation = Simulation( tax_benefit_system=tax_benefit_system, populations=tax_benefit_system.instantiate_entities(), ) + else: + simulation._replace_population() for population in simulation.populations.values(): population.count = count population.ids = np.array(range(count)) diff --git a/tests/core/test_builder_replaces_population.py b/tests/core/test_builder_replaces_population.py new file mode 100644 index 000000000..85c65ed1b --- /dev/null +++ b/tests/core/test_builder_replaces_population.py @@ -0,0 +1,239 @@ +"""``SimulationBuilder`` replaces the population of a simulation it is given. + +``build_from_dict`` (like ``build_from_entities``, ``build_from_variables`` +and ``build_default_simulation``) writes a situation into a simulation passed +as ``simulation``. It wrote the situation into that simulation's populations +as they were, so a simulation built over a dataset and then given a +situation: + +* kept ``is_over_dataset`` and ``dataset``, so the macro cache, and country + models that read ``is_over_dataset``, treated the situation as data; +* answered from its fast cache with values calculated for the dataset; +* kept the dataset's arrays for every variable the situation did not set; +* kept the dataset's input record, so ``to_input_dataframe`` failed; +* kept its branches, the baseline among them, on the dataset. + +The property-based version, which compares against a simulation built from +the situation alone, is in ``test_builder_replaces_population_property.py``. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.country_template import Microsimulation +from policyengine_core.country_template import Simulation as CountryTemplateSimulation +from policyengine_core.errors import SituationParsingError +from policyengine_core.simulations import Simulation, SimulationBuilder + +PERIOD = "2022-01" +SALARY = 3000.0 +SITUATION = { + "persons": {"you": {"salary": {PERIOD: SALARY}}}, + "households": {"home": {"parents": ["you"]}}, +} +REFORM = {"taxes.income_tax_rate": {"2022-01-01": 0.4}} +# Variables the situation sets (salary) or leaves to their defaults and +# formulas, including ones the data sets. +OUTPUTS = ["salary", "income_tax", "disposable_income", "age", "household_weight"] + + +def _data() -> pd.DataFrame: + """Ten people in five households, each with an age and salary.""" + return pd.DataFrame( + { + "person_id__2022": list(range(1, 11)), + "household_id__2022": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + "person_household_id__2022": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + "household_weight__2022": [100.0] * 10, + "salary__2022-01": [1000.0] * 10, + "age__2022-01": [40] * 10, + } + ) + + +def _period(variable: str) -> str: + return "2022" if variable == "household_weight" else PERIOD + + +def _assert_same_values(simulation, reference, label=""): + for variable in OUTPUTS: + np.testing.assert_array_equal( + simulation.calculate(variable, _period(variable)), + reference.calculate(variable, _period(variable)), + err_msg=f"{label} {variable}", + ) + + +def test_a_situation_written_over_data_is_not_over_a_dataset(tax_benefit_system): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + assert simulation.is_over_dataset + assert simulation.dataset is not None + + SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=simulation + ) + + assert not simulation.is_over_dataset + assert simulation.dataset is None + + +def test_it_holds_nothing_from_the_data(tax_benefit_system): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + # Calculated for the data, these filled the fast cache and holders. + for variable in OUTPUTS: + simulation.calculate(variable, _period(variable)) + + SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=simulation + ) + + reference = SimulationBuilder().build_from_dict(tax_benefit_system, SITUATION) + # Before, income tax was the ten people's (150 each) and age and weight + # the data's. + np.testing.assert_allclose(simulation.calculate("income_tax", PERIOD), [450]) + _assert_same_values(simulation, reference) + assert simulation.input_variables == reference.input_variables + + +def test_its_inputs_are_the_situations(tax_benefit_system): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + + SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=simulation + ) + + reference = SimulationBuilder().build_from_dict(tax_benefit_system, SITUATION) + assert simulation.true_input_variables == ["salary"] + # Before, the export failed: the data's ten-person inputs were still + # recorded on a one-person population. + pd.testing.assert_frame_equal( + simulation.to_input_dataframe(), reference.to_input_dataframe() + ) + + +@pytest.mark.parametrize( + "build", + [ + pytest.param( + lambda builder, system, simulation: builder.build_from_entities( + system, SITUATION, simulation + ), + id="build_from_entities", + ), + pytest.param( + lambda builder, system, simulation: builder.build_from_variables( + system, {"salary": {PERIOD: [SALARY, 0.0]}}, simulation + ), + id="build_from_variables", + ), + pytest.param( + lambda builder, system, simulation: builder.build_default_simulation( + system, 2, simulation + ), + id="build_default_simulation", + ), + pytest.param( + lambda builder, system, simulation: builder.build_from_dict( + system, {"salary": {PERIOD: [SALARY, 0.0]}}, simulation + ), + id="build_from_dict without entities", + ), + ], +) +def test_every_builder_replaces_the_population(tax_benefit_system, build): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + simulation.calculate("income_tax", PERIOD) + simulation.get_branch("scenario") + + build(SimulationBuilder(), tax_benefit_system, simulation) + + reference = build(SimulationBuilder(), tax_benefit_system, None) + assert not simulation.is_over_dataset + assert simulation.dataset is None + # A branch of the data's population is forgotten. + assert simulation.branches == {} + _assert_same_values(simulation, reference) + + +def test_a_situation_that_fails_to_parse_leaves_the_data(tax_benefit_system): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + + with pytest.raises(SituationParsingError): + SimulationBuilder().build_from_dict( + tax_benefit_system, + {**SITUATION, "companies": {"acme": {}}}, + simulation=simulation, + ) + + assert simulation.is_over_dataset + assert simulation.dataset is not None + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [1000.0] * 10) + + +def test_a_reform_simulation_compares_against_the_situation(): + simulation = CountryTemplateSimulation(dataset=_data(), reform=REFORM) + simulation.calculate("income_tax", PERIOD) + simulation.baseline.calculate("income_tax", PERIOD) + + SimulationBuilder().build_from_dict( + simulation.tax_benefit_system, SITUATION, simulation=simulation + ) + + # The baseline branch is rebuilt on the situation, under the baseline + # policy, and is the simulation's baseline. + assert simulation.baseline is simulation.branches["baseline"] + assert list(simulation.branches) == ["baseline"] + assert simulation.baseline.baseline is None + assert not simulation.baseline.is_over_dataset + _assert_same_values( + simulation, CountryTemplateSimulation(situation=SITUATION, reform=REFORM) + ) + _assert_same_values( + simulation.baseline, CountryTemplateSimulation(situation=SITUATION) + ) + np.testing.assert_allclose( + simulation.calculate("income_tax", PERIOD), [SALARY * 0.4] + ) + np.testing.assert_allclose( + simulation.baseline.calculate("income_tax", PERIOD), [SALARY * 0.15] + ) + + +def test_the_macro_cache_is_not_read_for_a_situation(isolated_tax_benefit_system): + # The country template's dataset is stored as time-period arrays, which + # the macro cache can serve, and income tax depends on one parameter. + isolated_tax_benefit_system.variables[ + "income_tax" + ].exhaustive_parameter_dependencies = ["taxes.income_tax_rate"] + simulation = Microsimulation(tax_benefit_system=isolated_tax_benefit_system) + simulation.macro_cache_read = True + assert simulation.check_macro_cache("income_tax", PERIOD) + + SimulationBuilder().build_from_dict( + isolated_tax_benefit_system, SITUATION, simulation=simulation + ) + + assert not simulation.check_macro_cache("income_tax", PERIOD) + + +@pytest.mark.parametrize("data_first", [True, False]) +def test_only_a_simulation_with_a_dataset_can_be_subsampled( + tax_benefit_system, data_first +): + simulation = ( + Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + if data_first + else None + ) + simulation = SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=simulation + ) + + # Before, one written over data subsampled the data's leftover inputs and + # failed on their length, and one built from a situation alone raised an + # AttributeError reading the default seed (the dataset's name). + with pytest.raises(ValueError, match="Only a simulation with a dataset"): + simulation.subsample(n=1) diff --git a/tests/core/test_builder_replaces_population_property.py b/tests/core/test_builder_replaces_population_property.py new file mode 100644 index 000000000..22ce0fbb6 --- /dev/null +++ b/tests/core/test_builder_replaces_population_property.py @@ -0,0 +1,227 @@ +"""Writing a situation into a simulation is the same as building it anew. + +For any dataset, any reform (or none), any calculations, branches and +earlier situations before it, and any situation (with entities or without, +with axes or without), once ``SimulationBuilder.build_from_dict`` writes the +situation into the simulation: + +* the simulation is not over a dataset, and has none; +* it, and its baseline, calculate what a simulation the builder makes from + the situation alone, under the same system, calculates (differential); +* it records the same inputs, input variables and axes as + ``Simulation(situation=..., reform=...)``, and exports the same inputs; +* of its branches, only the baseline is left. + +Examples, including the ones that failed before, are in +``test_builder_replaces_population.py``. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import HealthCheck, example, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.country_template import ( # noqa: E402 + CountryTaxBenefitSystem, + Simulation, +) +from policyengine_core.simulations import SimulationBuilder # noqa: E402 +from tests.core.test_reform_baseline_system import ( # noqa: E402 + add_doubled_salary, + change_salary_default, + neutralize_income_tax, +) + +PERIOD = "2022-01" +REFORMS = { + None: None, + "rate": {"taxes.income_tax_rate": {"2022-01-01": 0.42}}, + "neutralize income tax": neutralize_income_tax, + "salary default": change_salary_default, + "added variable": add_doubled_salary, +} +OUTPUTS = { + "salary": PERIOD, + "income_tax": PERIOD, + "disposable_income": PERIOD, + "household_income": PERIOD, + "age": PERIOD, + "household_weight": "2022", + "doubled_salary": PERIOD, +} +AMOUNT = st.integers(min_value=0, max_value=10_000).map(float) +BASELINE_SYSTEM = CountryTaxBenefitSystem() + + +@st.composite +def data(draw): + """People in households of one or two, with salaries and weights.""" + sizes = draw(st.lists(st.integers(1, 2), min_size=1, max_size=4)) + household = [index for index, size in enumerate(sizes, 1) for _ in range(size)] + count = len(household) + columns = { + "person_id__2022": list(range(1, count + 1)), + "household_id__2022": household, + "person_household_id__2022": household, + "household_weight__2022": [100.0 * h for h in household], + "salary__2022-01": draw(st.lists(AMOUNT, min_size=count, max_size=count)), + } + if draw(st.booleans()): + columns["age__2022-01"] = [30 + h for h in household] + return pd.DataFrame(columns) + + +@st.composite +def situations(draw): + """One household of up to three people, or people without entities.""" + count = draw(st.integers(1, 3)) + salaries = draw( + st.lists(st.one_of(st.none(), AMOUNT), min_size=count, max_size=count) + ) + if draw(st.booleans()): + # No entities: every person given a salary. + return {"salary": {PERIOD: [amount or 0.0 for amount in salaries]}} + people = [f"person_{index}" for index in range(count)] + situation = { + "persons": { + person: {} if amount is None else {"salary": {PERIOD: amount}} + for person, amount in zip(people, salaries) + }, + "households": {"home": {"parents": people[:2], "children": people[2:]}}, + } + if draw(st.booleans()): + situation["axes"] = [ + [ + { + "name": "salary", + "count": draw(st.integers(2, 3)), + "min": 0, + "max": 5000, + "period": PERIOD, + } + ] + ] + return situation + + +STEPS = st.lists( + st.one_of( + st.tuples(st.just("calculate"), st.sampled_from(list(OUTPUTS))), + st.tuples(st.just("calculate on the baseline"), st.sampled_from(list(OUTPUTS))), + st.tuples(st.just("branch"), st.sampled_from(list(OUTPUTS))), + st.tuples(st.just("write a situation"), situations()), + ), + max_size=4, +) + + +def _write(simulation, situation): + builder = SimulationBuilder() + builder.default_period = simulation.default_input_period + return builder.build_from_dict( + simulation.tax_benefit_system, situation, simulation=simulation + ) + + +def _calculate(simulation, variable): + if variable in simulation.tax_benefit_system.variables: + simulation.calculate(variable, OUTPUTS[variable]) + + +def _assert_calculates_as(simulation, reference, label): + for variable, period in OUTPUTS.items(): + if variable not in reference.tax_benefit_system.variables: + continue + np.testing.assert_array_equal( + simulation.calculate(variable, period), + reference.calculate(variable, period), + err_msg=f"{label} {variable}", + ) + + +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given( + data_=data(), + reform=st.sampled_from(list(REFORMS)), + steps=STEPS, + situation=situations(), +) +# Values calculated for the data, on the simulation and its baseline. +@example( + data_=pd.DataFrame( + { + "person_id__2022": [1, 2], + "household_id__2022": [1, 2], + "person_household_id__2022": [1, 2], + "household_weight__2022": [100.0, 200.0], + "salary__2022-01": [1000.0, 2000.0], + } + ), + reform="rate", + steps=[("calculate", "income_tax"), ("calculate on the baseline", "income_tax")], + situation={ + "persons": {"person_0": {"salary": {PERIOD: 3000.0}}}, + "households": {"home": {"parents": ["person_0"]}}, + }, +) +def test_writing_a_situation_is_building_it_anew(data_, reform, steps, situation): + simulation = Simulation(dataset=data_, reform=REFORMS[reform]) + for index, (kind, argument) in enumerate(steps): + if kind == "calculate": + _calculate(simulation, argument) + elif kind == "calculate on the baseline" and simulation.baseline is not None: + _calculate(simulation.baseline, argument) + elif kind == "branch": + _calculate(simulation.get_branch(f"branch_{index}"), argument) + elif kind == "write a situation": + _write(simulation, argument) + _calculate(simulation, "income_tax") + + _write(simulation, situation) + + assert not simulation.is_over_dataset + assert simulation.dataset is None + assert set(simulation.branches) == ({"baseline"} if reform else set()) + + # Built by the builder alone, under the same systems. + builder = SimulationBuilder() + builder.default_period = PERIOD + _assert_calculates_as( + simulation, + builder.build_from_dict(simulation.tax_benefit_system, situation), + "simulation", + ) + if reform: + assert simulation.baseline is simulation.branches["baseline"] + assert not simulation.baseline.is_over_dataset + builder = SimulationBuilder() + builder.default_period = PERIOD + _assert_calculates_as( + simulation.baseline, + builder.build_from_dict(BASELINE_SYSTEM, situation), + "baseline", + ) + else: + assert simulation.baseline is None + + # Built from the situation by ``Simulation``. + reference = Simulation(situation=situation, reform=REFORMS[reform]) + assert simulation._user_input_keys == reference._user_input_keys + assert simulation.input_variables == reference.input_variables + assert simulation.has_axes == reference.has_axes + pd.testing.assert_frame_equal( + simulation.to_input_dataframe(), reference.to_input_dataframe() + ) + if reform: + assert simulation.baseline.input_variables == reference.baseline.input_variables From f6673aea2261d57610cdef60db44228437718a71 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 14:48:42 -0400 Subject: [PATCH 2/2] Refuse branches and running calculations; define failed writes Review of the first commit (GPT-6.1 Sol, static): - A branch shares its parent's population and baseline, so writing a situation into one replaced its population but not its baseline. _replace_population now refuses a branch. - A formula could replace the population of the simulation calculating it. _replace_population now refuses while a calculation runs, and advances _input_epoch so a calculation elsewhere that read the simulation does not keep what it calculated from the old population. - requires_computation_after accepted a prerequisite requested for the old population, and a FullTracer kept the old population's values. Both are reset. - build_from_variables resized the populations before checking the situation. It now checks variable names and periods first, so a situation that fails to parse leaves the simulation as it was, as build_from_entities did. A situation that fails while it is written leaves the simulation replaced with what was written; the builder now records it and rebuilds branches in that case too. - The property test compared the baseline with one built under the baseline system, which differs whenever a reform changes the default of a variable the situation gives for some people only (a known limitation of the situation builder). It now compares with Simulation(situation=..., reform=...)'s baseline. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 22 +++ .../simulations/simulation_builder.py | 64 ++++--- .../core/test_builder_replaces_population.py | 171 +++++++++++++++++- ...st_builder_replaces_population_property.py | 42 ++--- 4 files changed, 252 insertions(+), 47 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 86fb1fb61..229e31f27 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1091,7 +1091,22 @@ def _replace_population(self) -> None: The simulation keeps its tax-benefit system and settings. Its branches keep the old populations until :meth:`_rebuild_branches`. + + Raises: + ValueError: if the simulation is a branch, which shares its + parent's population and baseline. + RuntimeError: if a calculation is running in the simulation. """ + if getattr(self, "parent_branch", None) is not None: + raise ValueError( + f"Cannot replace the population of branch '{self.branch_name}': " + "it shares its parent's population and baseline. Write the " + "situation into the simulation it was made from, or a new one." + ) + if self._calculations_in_flight: + raise RuntimeError( + "Cannot replace a simulation's population while it is calculating." + ) self.dataset = None self.is_over_dataset = False self.situation_input = None @@ -1099,8 +1114,15 @@ def _replace_population(self) -> None: self.input_variables = [] self._store_history = StoreHistory() self._user_input_keys = set() + self._requested_variables = None self.invalidated_caches = set() self._fast_cache = {} + # A calculation elsewhere that read this simulation does not keep what + # it calculated from the old population (``_Frame.is_stale``). + self._input_epoch = self._input_epoch + 1 + if isinstance(getattr(self, "tracer", None), FullTracer): + # Its trees hold the old population's values. + self.tracer = FullTracer() self.build_from_populations(self.tax_benefit_system.instantiate_entities()) def _rebuild_branches(self) -> None: diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index 6afc49f6a..84a13ba0b 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -78,7 +78,10 @@ def build_from_dict( (see :meth:`Simulation._replace_population`): it is no longer over a dataset, and holds nothing loaded or calculated for its old population. It keeps its tax-benefit system and settings, and its baseline branch - is rebuilt on the new population. + is rebuilt on the new population. A situation that fails to parse + (unknown entities or variables, say) leaves the simulation as it was; + one that fails while it is written leaves the simulation replaced, + with what was written. A branch cannot be overwritten. """ replacing = simulation is not None @@ -109,7 +112,11 @@ def _record_input_variables(self, simulation: Simulation) -> None: ] def _finish_replacing(self, simulation: Simulation) -> None: - """Finish writing a situation into a simulation whose population it replaced.""" + """Finish writing a situation into a simulation whose population it replaced. + + Also when writing failed partway, so the simulation's records and + branches describe what was written. + """ simulation.has_axes = self.has_axes self._record_input_variables(simulation) # Its branches still hold the old population. Rebuilt after the @@ -205,19 +212,20 @@ def build_from_entities( simulation._replace_population() try: - self.finalize_variables_init(simulation.persons) - except PeriodMismatchError as e: - self.raise_period_mismatch(simulation.persons.entity, persons_json, e) - - for entity_class in tax_benefit_system.group_entities: try: - population = simulation.populations[entity_class.key] - self.finalize_variables_init(population) + self.finalize_variables_init(simulation.persons) except PeriodMismatchError as e: - self.raise_period_mismatch(population.entity, instances_json, e) + self.raise_period_mismatch(simulation.persons.entity, persons_json, e) - if replacing: - self._finish_replacing(simulation) + for entity_class in tax_benefit_system.group_entities: + try: + population = simulation.populations[entity_class.key] + self.finalize_variables_init(population) + except PeriodMismatchError as e: + self.raise_period_mismatch(population.entity, instances_json, e) + finally: + if replacing: + self._finish_replacing(simulation) return simulation @@ -239,23 +247,29 @@ def build_from_variables( ) """ count = _get_person_count(input_dict) + # Check the situation before building on it, so one that fails here + # leaves a given simulation as it was (as ``build_from_entities``). + for variable, value in input_dict.items(): + tax_benefit_system.get_variable(variable, check_existence=True) + if not isinstance(value, dict) and self.default_period is None: + raise SituationParsingError( + [variable], + "Can't deal with type: expected object. Input variables should be set for specific periods. For instance: {'salary': {'2017-01': 2000, '2017-02': 2500}}, or {'birth_date': {'ETERNITY': '1980-01-01'}}.", + ) replacing = simulation is not None simulation = self._build_default_populations( tax_benefit_system, count, simulation ) - for variable, value in input_dict.items(): - if not isinstance(value, dict): - if self.default_period is None: - raise SituationParsingError( - [variable], - "Can't deal with type: expected object. Input variables should be set for specific periods. For instance: {'salary': {'2017-01': 2000, '2017-02': 2500}}, or {'birth_date': {'ETERNITY': '1980-01-01'}}.", - ) - simulation.set_input(variable, self.default_period, value) - else: - for period_str, dated_value in value.items(): - simulation.set_input(variable, period_str, dated_value) - if replacing: - self._finish_replacing(simulation) + try: + for variable, value in input_dict.items(): + if not isinstance(value, dict): + simulation.set_input(variable, self.default_period, value) + else: + for period_str, dated_value in value.items(): + simulation.set_input(variable, period_str, dated_value) + finally: + if replacing: + self._finish_replacing(simulation) return simulation def build_default_simulation( diff --git a/tests/core/test_builder_replaces_population.py b/tests/core/test_builder_replaces_population.py index 85c65ed1b..e3be1c0db 100644 --- a/tests/core/test_builder_replaces_population.py +++ b/tests/core/test_builder_replaces_population.py @@ -25,8 +25,11 @@ from policyengine_core.country_template import Microsimulation from policyengine_core.country_template import Simulation as CountryTemplateSimulation -from policyengine_core.errors import SituationParsingError +from policyengine_core.country_template.entities import Person +from policyengine_core.errors import SituationParsingError, VariableNotFoundError +from policyengine_core.periods import MONTH from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.variables import Variable PERIOD = "2022-01" SALARY = 3000.0 @@ -173,6 +176,172 @@ def test_a_situation_that_fails_to_parse_leaves_the_data(tax_benefit_system): np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [1000.0] * 10) +@pytest.mark.parametrize( + "situation, error", + [ + # A value without a period, and no default period to give it. + ({"salary": SALARY}, SituationParsingError), + ({"wages": {PERIOD: [SALARY]}}, VariableNotFoundError), + ], +) +def test_variables_that_fail_to_parse_leave_the_data( + tax_benefit_system, situation, error +): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + + with pytest.raises(error): + SimulationBuilder().build_from_dict( + tax_benefit_system, situation, simulation=simulation + ) + + # Before, the data's populations were resized to the situation's first. + assert simulation.persons.count == 10 + assert simulation.is_over_dataset + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [1000.0] * 10) + + +def test_a_situation_that_fails_while_written_leaves_what_was_written( + tax_benefit_system, +): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + simulation.get_branch("scenario") + + with pytest.raises(ValueError): + SimulationBuilder().build_from_dict( + tax_benefit_system, + {"salary": {PERIOD: [SALARY]}, "age": {PERIOD: ["forty"]}}, + simulation=simulation, + ) + + # Replaced, with the salary written before the age failed, and records + # and branches that describe that. + assert not simulation.is_over_dataset + assert simulation.dataset is None + assert simulation.branches == {} + assert simulation.input_variables == ["salary"] + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [SALARY]) + + +def test_a_branch_cannot_be_overwritten(tax_benefit_system): + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=_data()) + branch = simulation.get_branch("scenario") + + # It shares its parent's population and baseline. + with pytest.raises(ValueError, match="branch 'scenario'"): + SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=branch + ) + + assert branch.persons.count == 10 + assert branch.is_over_dataset + + +class overwrites_its_simulation(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Salary, read before overwriting the simulation calculating it" + + def formula(person, period, parameters): + salary = person("salary", period) + simulation = person.simulation + SimulationBuilder().build_default_simulation( + simulation.tax_benefit_system, 1, simulation + ) + return salary + + +def test_a_simulation_cannot_be_overwritten_while_it_calculates( + isolated_tax_benefit_system, +): + isolated_tax_benefit_system.add_variable(overwrites_its_simulation) + simulation = Simulation( + tax_benefit_system=isolated_tax_benefit_system, dataset=_data() + ) + + # Before, it returned the ten people's salaries for a one-person + # simulation, which failed when stored. + with pytest.raises(RuntimeError, match="while it is calculating"): + simulation.calculate("overwrites_its_simulation", PERIOD) + + assert simulation.persons.count == 10 + assert simulation.is_over_dataset + + +class taxed_salary(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Salary, once income tax has been requested" + requires_computation_after = "income_tax" + + def formula(person, period, parameters): + return person("salary", period) + + +def test_a_prerequisite_requested_for_the_data_is_not_requested_for_the_situation( + isolated_tax_benefit_system, +): + isolated_tax_benefit_system.add_variable(taxed_salary) + simulation = Simulation( + tax_benefit_system=isolated_tax_benefit_system, dataset=_data() + ) + simulation.calculate("income_tax", PERIOD) + simulation.calculate("taxed_salary", PERIOD) + + SimulationBuilder().build_from_dict( + isolated_tax_benefit_system, SITUATION, simulation=simulation + ) + + # As in a simulation built from the situation alone. + reference = SimulationBuilder().build_from_dict( + isolated_tax_benefit_system, SITUATION + ) + for built in (reference, simulation): + with pytest.raises(ValueError, match="requires income_tax to be requested"): + built.calculate("taxed_salary", PERIOD) + + +def test_a_trace_shows_the_situations_values(tax_benefit_system): + simulation = Simulation( + tax_benefit_system=tax_benefit_system, dataset=_data(), trace=True + ) + simulation.calculate("income_tax", PERIOD) + + SimulationBuilder().build_from_dict( + tax_benefit_system, SITUATION, simulation=simulation + ) + simulation.calculate("income_tax", PERIOD) + + # Before, the trace kept the first calculation's (the data's) values. + trace = simulation.tracer.get_serialized_flat_trace() + np.testing.assert_allclose( + trace[f"income_tax<{PERIOD}, (default)>"]["value"], [450] + ) + + +def test_a_situation_written_over_another_replaces_it(): + first = { + **SITUATION, + "axes": [[{"name": "salary", "count": 3, "min": 0, "max": 6000}]], + } + simulation = CountryTemplateSimulation(situation=first) + simulation.calculate("income_tax", PERIOD) + assert simulation.has_axes + + builder = SimulationBuilder() + builder.default_period = PERIOD + builder.build_from_dict( + simulation.tax_benefit_system, SITUATION, simulation=simulation + ) + + # The first situation's axes and its record as the situation the + # simulation was built from describe it no longer. + assert not simulation.has_axes + assert simulation.situation_input is None + _assert_same_values(simulation, CountryTemplateSimulation(situation=SITUATION)) + + def test_a_reform_simulation_compares_against_the_situation(): simulation = CountryTemplateSimulation(dataset=_data(), reform=REFORM) simulation.calculate("income_tax", PERIOD) diff --git a/tests/core/test_builder_replaces_population_property.py b/tests/core/test_builder_replaces_population_property.py index 22ce0fbb6..86e208e6b 100644 --- a/tests/core/test_builder_replaces_population_property.py +++ b/tests/core/test_builder_replaces_population_property.py @@ -6,12 +6,19 @@ situation into the simulation: * the simulation is not over a dataset, and has none; -* it, and its baseline, calculate what a simulation the builder makes from - the situation alone, under the same system, calculates (differential); +* it calculates what a simulation the builder makes from the situation + alone, under the same system, calculates (differential); * it records the same inputs, input variables and axes as ``Simulation(situation=..., reform=...)``, and exports the same inputs; + its baseline calculates and records what that simulation's baseline does; * of its branches, only the baseline is left. +The baseline is compared with ``Simulation(situation=..., reform=...)``'s, +not with a simulation built under the baseline system: given a variable for +some people only (or along axes), the situation builder stores the reform's +default for the others as an input, in both (see +``test_reform_baseline_system.py``). + Examples, including the ones that failed before, are in ``test_builder_replaces_population.py``. """ @@ -28,10 +35,7 @@ from hypothesis import HealthCheck, example, given, settings # noqa: E402 from hypothesis import strategies as st # noqa: E402 -from policyengine_core.country_template import ( # noqa: E402 - CountryTaxBenefitSystem, - Simulation, -) +from policyengine_core.country_template import Simulation # noqa: E402 from policyengine_core.simulations import SimulationBuilder # noqa: E402 from tests.core.test_reform_baseline_system import ( # noqa: E402 add_doubled_salary, @@ -57,7 +61,6 @@ "doubled_salary": PERIOD, } AMOUNT = st.integers(min_value=0, max_value=10_000).map(float) -BASELINE_SYSTEM = CountryTaxBenefitSystem() @st.composite @@ -147,7 +150,7 @@ def _assert_calculates_as(simulation, reference, label): @settings( - max_examples=50, + max_examples=200, deadline=None, suppress_health_check=[HealthCheck.too_slow], ) @@ -194,7 +197,7 @@ def test_writing_a_situation_is_building_it_anew(data_, reform, steps, situation assert simulation.dataset is None assert set(simulation.branches) == ({"baseline"} if reform else set()) - # Built by the builder alone, under the same systems. + # Built by the builder alone, under the same system. builder = SimulationBuilder() builder.default_period = PERIOD _assert_calculates_as( @@ -202,18 +205,6 @@ def test_writing_a_situation_is_building_it_anew(data_, reform, steps, situation builder.build_from_dict(simulation.tax_benefit_system, situation), "simulation", ) - if reform: - assert simulation.baseline is simulation.branches["baseline"] - assert not simulation.baseline.is_over_dataset - builder = SimulationBuilder() - builder.default_period = PERIOD - _assert_calculates_as( - simulation.baseline, - builder.build_from_dict(BASELINE_SYSTEM, situation), - "baseline", - ) - else: - assert simulation.baseline is None # Built from the situation by ``Simulation``. reference = Simulation(situation=situation, reform=REFORMS[reform]) @@ -224,4 +215,13 @@ def test_writing_a_situation_is_building_it_anew(data_, reform, steps, situation simulation.to_input_dataframe(), reference.to_input_dataframe() ) if reform: + assert simulation.baseline is simulation.branches["baseline"] + assert simulation.baseline.baseline is None + assert not simulation.baseline.is_over_dataset + _assert_calculates_as(simulation.baseline, reference.baseline, "baseline") + assert simulation.baseline._user_input_keys == ( + reference.baseline._user_input_keys + ) assert simulation.baseline.input_variables == reference.baseline.input_variables + else: + assert simulation.baseline is None