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..229e31f27 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1077,6 +1077,90 @@ 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`. + + 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 + self.has_axes = False + 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: + """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 +3468,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 +3663,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..84a13ba0b 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -73,8 +73,18 @@ 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. 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 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 +97,32 @@ 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. + + 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 + # inputs are recorded, the baseline records them too, as in + # ``Simulation.__init__``. + simulation._rebuild_branches() def build_from_entities( self, @@ -113,6 +142,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,17 +207,25 @@ def build_from_entities( self.axes = axes self.expand_axes() - try: - self.finalize_variables_init(simulation.persons) - except PeriodMismatchError as e: - self.raise_period_mismatch(simulation.persons.entity, persons_json, e) + if replacing: + # The situation parsed: write it into new populations. + simulation._replace_population() - for entity_class in tax_benefit_system.group_entities: + try: 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) + + 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 @@ -209,20 +247,29 @@ def build_from_variables( ) """ count = _get_person_count(input_dict) - simulation = self.build_default_simulation( + # 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) + 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( @@ -236,13 +283,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..e3be1c0db --- /dev/null +++ b/tests/core/test_builder_replaces_population.py @@ -0,0 +1,408 @@ +"""``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.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 +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) + + +@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) + 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..86e208e6b --- /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 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``. +""" + +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 Simulation # noqa: E402 +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) + + +@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=200, + 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 system. + builder = SimulationBuilder() + builder.default_period = PERIOD + _assert_calculates_as( + simulation, + builder.build_from_dict(simulation.tax_benefit_system, situation), + "simulation", + ) + + # 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 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