diff --git a/changelog.d/reform-baseline-situation-defaults.fixed.md b/changelog.d/reform-baseline-situation-defaults.fixed.md new file mode 100644 index 000000000..32d1c92bb --- /dev/null +++ b/changelog.d/reform-baseline-situation-defaults.fixed.md @@ -0,0 +1 @@ +A reform simulation built from a situation now gives its baseline the inputs the baseline system makes of that situation: people the situation gives no value take the baseline's default instead of the reform's, and valid inputs for a variable the reform neutralizes reach the baseline. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ddda2ab94..ec5f10fcc 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -918,7 +918,16 @@ def __init__( # looked variables up there, and its holders kept the reform's # variables, so a variable the reform neutralized read as the # default in the baseline too. - self.baseline._bind_to_tax_benefit_system() + if ( + situation is not None + and populations is None + and self.baseline.tax_benefit_system is not None + ): + # So were its inputs: people the situation gives no value + # held the reform's default. + self.baseline._build_from_situation(situation) + else: + self.baseline._bind_to_tax_benefit_system() else: self.baseline = None @@ -1303,6 +1312,47 @@ def kept(name): if getattr(self, "input_variables", None) is not None: self.input_variables = [name for name in self.input_variables if kept(name)] + def _build_from_situation(self, situation: dict) -> None: + """Give this simulation the inputs its own system makes of ``situation``. + + For a reform simulation's baseline branch, which starts as a copy of + the reform simulation: the situation builder made that simulation's + inputs under the reform's system. A stored input is a whole array, so + a person the situation gives no value for a variable it gives others + held the reform's default for it, and the reform's system ignored + inputs for the variables it neutralizes. Here the simulation gets new + populations from its own system, and the situation is built into + them as the builder builds it under that system, except that inputs + the system cannot take are left out (see + ``SimulationBuilder.skip_rejected_inputs``), as + ``_bind_to_tax_benefit_system`` drops holders built for the reform's + definitions. The records of what the old populations held go with + them. The inputs are stored under ``"default"``, as in a simulation + built from the situation, which a branch reads through. + """ + from policyengine_core.simulations.simulation_builder import ( + SimulationBuilder, + ) # Import here to avoid circular dependency + + self.build_from_populations(self.tax_benefit_system.instantiate_entities()) + self._user_input_keys = set() + self._store_history = StoreHistory() + self._requested_variables = set() + self.invalidated_caches = set() + self._fast_cache = {} + builder = SimulationBuilder() + builder.default_period = self.default_input_period + builder.skip_rejected_inputs = True + # A situation without entities is set through ``set_input``, which + # stores under the simulation's branch name; one with entities is + # stored under ``"default"`` whatever the branch. + branch_name, self.branch_name = self.branch_name, "default" + try: + builder.build_from_dict(self.tax_benefit_system, situation, self) + finally: + self.branch_name = branch_name + self.has_axes = builder.has_axes + def create_shortcuts(self) -> None: for _key, population in self.populations.items(): # create shortcut simulation.person and simulation.household (for instance) diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index e3dee0cc5..e4104112e 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -63,6 +63,20 @@ def __init__(self): self.axes_memberships: typing.Dict[Entity.plural, typing.List[int]] = {} self.axes_roles: typing.Dict[Entity.plural, typing.List[int]] = {} + # Leave out, rather than raise for, inputs the tax-benefit system + # cannot take: a variable it does not have or defines for another + # entity, a value it rejects, an axis over a variable it does not + # have (the axis still shapes the population), and an array its + # ``set_input`` refuses with a ``ValueError`` (for a period the + # variable's definition does not take, say, but also for a NaN or + # for yearly and monthly inputs that do not add up). Entities, + # roles and the syntax of periods are still checked. A reform + # simulation builds its baseline's inputs this way, from a situation + # the reform's system took (see ``Simulation._build_from_situation``). + # That system may define a variable differently, and does not check + # inputs for the variables it neutralizes. + self.skip_rejected_inputs = False + def build_from_dict( self, tax_benefit_system: "TaxBenefitSystem", @@ -219,10 +233,14 @@ def build_from_variables( [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) + value = {self.default_period: value} + for period_str, dated_value in value.items(): + period = periods.period(period_str) + try: + simulation.set_input(variable, period, dated_value) + except (VariableNotFoundError, ValueError): + if not self.skip_rejected_inputs: + raise return simulation def build_default_simulation( @@ -508,8 +526,12 @@ def init_variable_values(self, entity, instance_object, instance_id): try: entity.check_variable_defined_for_entity(variable_name) except ValueError as e: # The variable is defined for another entity + if self.skip_rejected_inputs: + continue raise SituationParsingError(path_in_json, e.args[0]) except VariableNotFoundError as e: # The variable doesn't exist + if self.skip_rejected_inputs: + continue raise SituationParsingError(path_in_json, str(e), code=404) instance_index = self.get_ids(entity.plural).index(instance_id) @@ -554,6 +576,9 @@ def add_variable_value( try: value = variable.check_set_value(value) except ValueError as error: + if self.skip_rejected_inputs: + # The entry is left as if the situation did not give it. + return raise SituationParsingError(path_in_json, *error.args) array[instance_index] = value @@ -591,7 +616,13 @@ def finalize_variables_init(self, population): # TODO - this duplicates the check in Simulation.set_input, but # fixing that requires improving Simulation's handling of entities if (variable.end is None) or (period_value.start.date <= variable.end): - holder.set_input(period_value, array) + try: + holder.set_input(period_value, array) + except ValueError: + # ``Holder.set_input`` and core's ``set_input`` + # handlers raise before storing any of the input. + if not self.skip_rejected_inputs: + raise def raise_period_mismatch(self, entity, json, e): # This error happens when we try to set a variable value for a period that doesn't match its definition period @@ -688,11 +719,11 @@ def expand_axes(self): # Now generate input values along the specified axes # TODO - factor out the common logic here if len(self.axes) == 1 and len(self.axes[0]): - parallel_axes = self.axes[0] - first_axis = parallel_axes[0] - axis_count: int = first_axis["count"] - axis_entity = self.get_variable_entity(first_axis["name"]) - axis_entity_step_size = self.entity_counts[axis_entity.plural] + axis_count: int = self.axes[0][0]["count"] + parallel_axes = self._axes_to_set(self.axes[0]) + if parallel_axes: + axis_entity = self.get_variable_entity(parallel_axes[0]["name"]) + axis_entity_step_size = self.entity_counts[axis_entity.plural] # Distribute values along axes for axis in parallel_axes: axis_index = axis.get("index", 0) @@ -721,8 +752,11 @@ def expand_axes(self): ] axes_meshes = np.meshgrid(*axes_linspaces) for parallel_axes, mesh in zip(self.axes, axes_meshes): + axis_count = parallel_axes[0]["count"] + parallel_axes = self._axes_to_set(parallel_axes) + if not parallel_axes: + continue first_axis = parallel_axes[0] - axis_count = first_axis["count"] axis_entity = self.get_variable_entity(first_axis["name"]) axis_entity_step_size = self.entity_counts[axis_entity.plural] # Distribute values along the grid @@ -753,6 +787,18 @@ def expand_axes(self): self.has_axes = True + def _axes_to_set(self, parallel_axes: list) -> list: + """The axes of one dimension whose variable's values are set. + + All of them, except, when skipping rejected inputs, those over a + variable the tax-benefit system does not have. + """ + if not self.skip_rejected_inputs: + return parallel_axes + return [ + axis for axis in parallel_axes if axis["name"] in self.variable_entities + ] + def get_variable_entity(self, variable_name): return self.variable_entities[variable_name] diff --git a/tests/core/test_reform_baseline_situation.py b/tests/core/test_reform_baseline_situation.py new file mode 100644 index 000000000..27df24cae --- /dev/null +++ b/tests/core/test_reform_baseline_situation.py @@ -0,0 +1,638 @@ +"""A reform simulation's baseline holds what its own system makes of the situation. + +``Simulation(situation=..., reform=...)`` builds the situation under the +reform's system, then makes its baseline a branch of that simulation under the +baseline system. A stored input is a whole array: the situation builder gives +the people a situation gives no value for a variable it gives others (or sets +along axes) that variable's default, from the system it builds under. The +branch kept those arrays, so the baseline read the reform's default for those +people (PolicyEngine/policyengine-core#612). It also had no input for a +variable the reform neutralizes, whose input the reform's system ignores. The +baseline's inputs are now built from the situation under the baseline system, +leaving out the inputs that system cannot take. + +The property over any reform and situation is in +``test_reform_baseline_situation_property.py``. +""" + +from __future__ import annotations + +import copy +import datetime + +import numpy as np +import pytest + +from policyengine_core.country_template import Simulation +from policyengine_core.enums import Enum +from policyengine_core.errors import SituationParsingError, VariableNotFoundError +from policyengine_core.periods import ETERNITY, period +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable +from tests.core.test_reform_baseline_system import ( + REFORM_DEFAULT_SALARY, + add_doubled_salary, + change_salary_default, + move_salary_to_households, + neutralize_income_tax, +) + +PERIOD = "2022-01" +RATE = 0.15 +SYSTEM = Simulation.default_tax_benefit_system_instance +OUTPUTS = [ + "salary", + "income_tax", + "disposable_income", + "rent", + "housing_occupancy_status", + "total_taxes", + "household_income", +] +PARTIAL = { + "persons": {"a": {"salary": {PERIOD: 3000.0}}, "b": {}}, + "households": {"h": {"parents": ["a", "b"]}}, +} +REFORM_RENT = 600.0 +REFORM_BIRTH = datetime.date(1990, 6, 1) + + +class rent(Variable): + default_value = REFORM_RENT + + +class birth(Variable): + default_value = REFORM_BIRTH + + +class change_household_defaults(Reform): + def apply(self): + self.update_variable(rent) + # Each system loads its own enum classes, so the default is a member + # of this system's. + status = self.get_variable("housing_occupancy_status").possible_values + self.update_variable( + type( + "housing_occupancy_status", (Variable,), {"default_value": status.owner} + ) + ) + + +class change_birth_default(Reform): + def apply(self): + self.update_variable(birth) + + +class neutralize_salary(Reform): + def apply(self): + self.neutralize_variable("salary") + + +class ExtendedStatus(Enum): + owner = "Owner" + tenant = "Tenant" + free_lodger = "Free lodger" + homeless = "Homeless" + student = "Student accommodation" + + +class extend_housing_occupancy_status(Reform): + def apply(self): + # ``possible_values`` and ``default_value`` are read from the class + # itself, not from its bases. + self.update_variable( + type( + "housing_occupancy_status", + (Variable,), + { + "possible_values": ExtendedStatus, + "default_value": ExtendedStatus.tenant, + }, + ) + ) + + +class eternal_salary(Variable): + definition_period = ETERNITY + + +class make_salary_eternal(Reform): + def apply(self): + self.update_variable(type("salary", (eternal_salary,), {})) + + +def _built(situation, system=SYSTEM): + """What the situation builder alone makes of ``situation`` under ``system``.""" + builder = SimulationBuilder() + builder.default_period = PERIOD + return builder.build_from_dict(system, copy.deepcopy(situation)) + + +def _values(simulation, variable, period=PERIOD): + return np.asarray(simulation.calculate(variable, period)) + + +def assert_same(simulation, expected, outputs=OUTPUTS, periods=(PERIOD,)): + for variable in outputs: + for period in periods: + np.testing.assert_array_equal( + _values(simulation, variable, period), + _values(expected, variable, period), + err_msg=f"{variable} at {period}", + ) + + +def test_people_given_no_salary_take_the_baselines_default(): + reform = Simulation(situation=PARTIAL, reform=change_salary_default) + + np.testing.assert_array_equal(_values(reform.baseline, "salary"), [3000, 0]) + np.testing.assert_allclose( + _values(reform.baseline, "income_tax"), [3000 * RATE, 0], rtol=1e-6 + ) + # The reform simulation keeps the reform's default. + np.testing.assert_array_equal( + _values(reform, "salary"), [3000, REFORM_DEFAULT_SALARY] + ) + np.testing.assert_allclose( + _values(reform, "income_tax"), + [3000 * RATE, REFORM_DEFAULT_SALARY * RATE], + rtol=1e-6, + ) + assert_same(reform.baseline, _built(PARTIAL)) + reform_system = Simulation.default_tax_benefit_system(reform=change_salary_default) + assert_same(reform, _built(PARTIAL, reform_system)) + + +def test_people_off_an_axis_take_the_baselines_default(): + situation = { + "persons": {"a": {}, "b": {}}, + "households": {"h": {"parents": ["a", "b"]}}, + "axes": [ + [{"name": "salary", "count": 2, "min": 0, "max": 5000, "period": PERIOD}] + ], + } + reform = Simulation(situation=situation, reform=change_salary_default) + + np.testing.assert_array_equal(_values(reform.baseline, "salary"), [0, 0, 5000, 0]) + np.testing.assert_array_equal( + _values(reform, "salary"), + [0, REFORM_DEFAULT_SALARY, 5000, REFORM_DEFAULT_SALARY], + ) + assert reform.baseline.has_axes + assert_same(reform.baseline, _built(situation)) + + +def test_a_default_given_for_one_month_changes_how_a_year_is_divided(): + """Why the baseline is built again rather than given other defaults. + + ``salary`` divides a yearly input between the months that have no input. + ``b`` is given January's salary and ``a`` the year's, so January's input + gives ``a`` the default and the year's gives ``b`` the default. What is + left of each yearly value for February to December depends on the + default given for January. + """ + situation = { + "persons": { + "a": {"salary": {"2022": 12_000.0}}, + "b": {"salary": {PERIOD: 3000.0}}, + }, + "households": {"h": {"parents": ["a", "b"]}}, + } + reform = Simulation(situation=situation, reform=change_salary_default) + default = REFORM_DEFAULT_SALARY + + np.testing.assert_allclose( + _values(reform.baseline, "salary", "2022-02"), + [12_000 / 11, (0 - 3000) / 11], + rtol=1e-6, + ) + np.testing.assert_allclose( + _values(reform, "salary", "2022-02"), + [(12_000 - default) / 11, (default - 3000) / 11], + rtol=1e-6, + ) + assert_same(reform.baseline, _built(situation), periods=(PERIOD, "2022-02")) + + +def test_households_given_nothing_take_the_baselines_defaults(): + situation = { + "persons": {"a": {}, "b": {}}, + "households": { + "h": { + "parents": ["a"], + "rent": {PERIOD: 900.0}, + "housing_occupancy_status": {PERIOD: "free_lodger"}, + }, + "g": {"parents": ["b"]}, + }, + } + reform = Simulation(situation=situation, reform=change_household_defaults) + + np.testing.assert_array_equal(_values(reform.baseline, "rent"), [900, 0]) + assert list( + reform.baseline.calculate("housing_occupancy_status", PERIOD).decode_to_str() + ) == ["free_lodger", "tenant"] + np.testing.assert_array_equal(_values(reform, "rent"), [900, REFORM_RENT]) + assert list( + reform.calculate("housing_occupancy_status", PERIOD).decode_to_str() + ) == ["free_lodger", "owner"] + assert_same(reform.baseline, _built(situation)) + + +def test_people_given_no_birth_date_take_the_baselines_default(): + situation = { + "persons": {"a": {"birth": {"ETERNITY": "1980-01-01"}}, "b": {}}, + "households": {"h": {"parents": ["a", "b"]}}, + } + reform = Simulation(situation=situation, reform=change_birth_default) + + np.testing.assert_array_equal( + _values(reform.baseline, "birth"), + np.array(["1980-01-01", "1970-01-01"], dtype="datetime64[D]"), + ) + np.testing.assert_array_equal( + _values(reform, "birth"), + np.array(["1980-01-01", REFORM_BIRTH], dtype="datetime64[D]"), + ) + assert_same(reform.baseline, _built(situation), outputs=["birth", "age"]) + + +@pytest.mark.filterwarnings("ignore:You cannot set a value") +@pytest.mark.parametrize( + "situation", + [ + pytest.param( + { + "persons": { + "a": {"salary": {PERIOD: 3000.0}, "income_tax": {PERIOD: 7.0}} + }, + "households": {"h": {"parents": ["a"]}}, + }, + id="entities", + ), + # Without entities, the builder sets each variable through + # ``Simulation.set_input`` for one person. + pytest.param( + {"salary": {PERIOD: 3000.0}, "income_tax": {PERIOD: 7.0}}, + id="variables", + ), + ], +) +def test_an_input_the_reform_neutralizes_is_the_baselines(situation): + reform = Simulation(situation=situation, reform=neutralize_income_tax) + + np.testing.assert_array_equal(_values(reform.baseline, "income_tax"), [7]) + np.testing.assert_array_equal(_values(reform, "income_tax"), [0]) + assert "income_tax" in reform.baseline.input_variables + assert_same(reform.baseline, _built(situation), outputs=["salary", "income_tax"]) + + +def test_a_situation_without_entities_is_stored_as_in_a_simulation_built_from_it(): + """The baseline stores the inputs under ``"default"``, not its own name.""" + situation = {"salary": {PERIOD: [3000.0, 100.0]}} + baseline = Simulation(situation=situation, reform=change_salary_default).baseline + expected = Simulation(situation=situation) + + assert baseline.branch_name == "baseline" + assert baseline._user_input_keys == expected._user_input_keys + np.testing.assert_array_equal( + baseline.get_holder("salary").get_array(period(PERIOD)), [3000, 100] + ) + + +@pytest.mark.filterwarnings("ignore:You cannot set a value") +def test_a_neutralized_input_variable_is_the_baselines(): + reform = Simulation(situation=PARTIAL, reform=neutralize_salary) + + np.testing.assert_array_equal(_values(reform.baseline, "salary"), [3000, 0]) + np.testing.assert_array_equal(_values(reform, "salary"), [0, 0]) + assert_same(reform.baseline, _built(PARTIAL)) + + +def _without(situation, path): + """``situation`` without the entry at ``path``.""" + stripped = copy.deepcopy(situation) + entry = stripped + for key in path[:-1]: + entry = entry[key] + del entry[path[-1]] + return stripped + + +@pytest.mark.parametrize( + "reform, situation, rejected", + [ + pytest.param( + add_doubled_salary, + { + "persons": { + "a": {"salary": {PERIOD: 3000.0}, "doubled_salary": {PERIOD: 5.0}}, + "b": {"salary": {PERIOD: 100.0}}, + }, + "households": {"h": {"parents": ["a", "b"]}}, + }, + ["persons", "a", "doubled_salary"], + id="variable-only-the-reform-has", + ), + pytest.param( + move_salary_to_households, + { + "persons": {"a": {}, "b": {}}, + "households": {"h": {"parents": ["a", "b"], "salary": {PERIOD: 100.0}}}, + }, + ["households", "h", "salary"], + id="variable-the-reform-moves-to-another-entity", + ), + pytest.param( + extend_housing_occupancy_status, + { + "persons": {"a": {}, "b": {}}, + "households": { + "h": { + "parents": ["a"], + "housing_occupancy_status": {PERIOD: "student"}, + }, + "g": { + "parents": ["b"], + "housing_occupancy_status": {PERIOD: "owner"}, + }, + }, + }, + ["households", "h", "housing_occupancy_status"], + id="value-only-the-reform-has", + ), + pytest.param( + make_salary_eternal, + { + "persons": {"a": {"salary": {"ETERNITY": 3000.0}}, "b": {}}, + "households": {"h": {"parents": ["a", "b"]}}, + }, + ["persons", "a", "salary"], + id="period-only-the-reform-takes", + ), + ], +) +def test_inputs_the_baseline_system_cannot_take_are_left_out( + reform, situation, rejected +): + # Under the baseline system, the situation does not build at all. + with pytest.raises(SituationParsingError): + _built(situation) + + simulation = Simulation(situation=situation, reform=reform) + + assert_same(simulation.baseline, _built(_without(situation, rejected))) + assert sorted(simulation.baseline.input_variables) == sorted( + _built(_without(situation, rejected)).input_variables + ) + + +def test_an_axis_over_a_variable_only_the_reform_has_still_shapes_the_baseline(): + base = { + "persons": {"a": {"salary": {PERIOD: 3000.0}}, "b": {}}, + "households": {"h": {"parents": ["a", "b"]}}, + } + situation = { + **base, + "axes": [ + [ + { + "name": "doubled_salary", + "count": 3, + "min": 0, + "max": 10, + "period": PERIOD, + } + ], + [ + { + "name": "salary", + "count": 2, + "min": 0, + "max": 500, + "period": PERIOD, + "index": 1, + } + ], + ], + } + reform = Simulation(situation=situation, reform=add_doubled_salary) + baseline = reform.baseline + + assert baseline.persons.count == reform.persons.count == 12 + assert baseline.household.count == reform.household.count == 6 + np.testing.assert_array_equal(baseline.persons.ids, reform.persons.ids) + np.testing.assert_array_equal( + baseline.household.members_entity_id, reform.household.members_entity_id + ) + # ``salary``'s axis is set as it is in the reform simulation, whose + # ``salary`` is the baseline's. + np.testing.assert_array_equal( + _values(baseline, "salary"), _values(reform, "salary") + ) + assert "doubled_salary" not in baseline.input_variables + assert all(key[0] != "doubled_salary" for key in baseline._user_input_keys) + + +def test_the_baseline_records_and_exports_the_inputs_its_system_builds(): + # The axis sets ``a``'s salary in each copy; ``b``, given none, takes the + # default in each. + situation = { + **PARTIAL, + "axes": [ + [ + { + "name": "salary", + "count": 2, + "min": 0, + "max": 5000, + "period": PERIOD, + } + ] + ], + } + baseline = Simulation(situation=situation, reform=change_salary_default).baseline + expected = _built(situation) + + assert baseline._user_input_keys == expected._user_input_keys + assert sorted(baseline.input_variables) == sorted(expected.input_variables) + assert baseline.to_input_dict() == expected.to_input_dict() + assert baseline.to_input_dataframe().equals(expected.to_input_dataframe()) + assert baseline.has_axes == Simulation(situation=situation).has_axes + + +@pytest.mark.parametrize("trace", [False, True]) +def test_the_rebuilt_baseline_is_still_the_reform_simulations_baseline_branch(trace): + reform = Simulation(situation=PARTIAL, reform=change_salary_default, trace=trace) + baseline = reform.baseline + + assert reform.branches["baseline"] is baseline + assert reform.get_branch("baseline") is baseline + assert baseline.parent_branch is reform + assert baseline.branch_name == "baseline" + assert baseline.baseline is None + assert baseline.tracer is reform.tracer + assert baseline.tax_benefit_system is SYSTEM + entities = { + entity.key: entity for entity in [SYSTEM.person_entity, *SYSTEM.group_entities] + } + for key, population in baseline.populations.items(): + assert population.simulation is baseline + assert population.entity is entities[key] + assert getattr(baseline, key) is population + # A clone's copy of the baseline is a copy of the rebuilt one. + np.testing.assert_array_equal(_values(reform.clone().baseline, "salary"), [3000, 0]) + + +def test_the_baseline_keeps_its_inputs_through_the_reform_simulations_reform(): + """policyengine-us applies structural reforms after ``__init__``.""" + reform = Simulation(situation=PARTIAL, reform=change_salary_default) + reform.apply_reform(neutralize_income_tax) + + np.testing.assert_array_equal(_values(reform.baseline, "salary"), [3000, 0]) + np.testing.assert_allclose( + _values(reform.baseline, "income_tax"), [3000 * RATE, 0], rtol=1e-6 + ) + np.testing.assert_array_equal(_values(reform, "income_tax"), [0, 0]) + + +def test_inputs_set_on_the_reform_simulation_afterwards_stay_out_of_the_baseline(): + reform = Simulation(situation=PARTIAL, reform=change_salary_default) + reform.set_input("salary", PERIOD, [10.0, 20.0]) + + np.testing.assert_array_equal(_values(reform, "salary"), [10, 20]) + np.testing.assert_array_equal(_values(reform.baseline, "salary"), [3000, 0]) + + +def _strict_and_skipping(): + strict = SimulationBuilder() + strict.default_period = PERIOD + skipping = SimulationBuilder() + skipping.default_period = PERIOD + skipping.skip_rejected_inputs = True + return strict, skipping + + +@pytest.mark.parametrize( + "situation, rejected", + [ + pytest.param( + { + "persons": { + "a": {"salary": {PERIOD: 1.0}, "no_such_variable": {PERIOD: 1}} + }, + "households": {"h": {"parents": ["a"]}}, + }, + ["persons", "a", "no_such_variable"], + id="unknown-variable", + ), + pytest.param( + { + "persons": {"a": {"salary": {PERIOD: 1.0}, "rent": {PERIOD: 1.0}}}, + "households": {"h": {"parents": ["a"]}}, + }, + ["persons", "a", "rent"], + id="variable-of-another-entity", + ), + pytest.param( + { + "persons": {"a": {}, "b": {}}, + "households": { + "h": { + "parents": ["a"], + "housing_occupancy_status": {PERIOD: "castle"}, + }, + "g": { + "parents": ["b"], + "housing_occupancy_status": {PERIOD: "owner"}, + }, + }, + }, + ["households", "h", "housing_occupancy_status"], + id="rejected-value", + ), + pytest.param( + { + "persons": { + "a": {"salary": {"ETERNITY": 1.0}}, + "b": {"salary": {PERIOD: 2.0}}, + }, + "households": {"h": {"parents": ["a", "b"]}}, + }, + ["persons", "a", "salary", "ETERNITY"], + id="refused-period", + ), + pytest.param( + {"salary": {PERIOD: 1.0}, "no_such_variable": {PERIOD: 1.0}}, + ["no_such_variable"], + id="unknown-variable-without-entities", + ), + ], +) +def test_skipping_rejected_inputs_builds_the_situation_without_them( + situation, rejected +): + strict, skipping = _strict_and_skipping() + + with pytest.raises((SituationParsingError, VariableNotFoundError)): + strict.build_from_dict(SYSTEM, copy.deepcopy(situation)) + simulation = skipping.build_from_dict(SYSTEM, copy.deepcopy(situation)) + expected = _built(_without(situation, rejected)) + + assert simulation._user_input_keys == expected._user_input_keys + assert simulation.to_input_dict() == expected.to_input_dict() + assert_same( + simulation, expected, outputs=["salary", "rent", "housing_occupancy_status"] + ) + + +def test_skipping_rejected_inputs_keeps_an_axis_over_an_unknown_variable(): + situation = { + "persons": {"a": {"salary": {PERIOD: 1.0}}}, + "households": {"h": {"parents": ["a"]}}, + "axes": [ + [ + {"name": "no_such_variable", "count": 3, "min": 0, "max": 2}, + {"name": "salary", "count": 3, "min": 10, "max": 30, "period": PERIOD}, + ] + ], + } + _, skipping = _strict_and_skipping() + simulation = skipping.build_from_dict(SYSTEM, copy.deepcopy(situation)) + + assert simulation.persons.count == 3 + np.testing.assert_array_equal(_values(simulation, "salary"), [10, 20, 30]) + + +def test_the_builder_still_checks_entities_when_skipping_rejected_inputs(): + situation = { + "persons": {"a": {}}, + "households": {"h": {"parents": ["a", "nobody"]}}, + } + _, skipping = _strict_and_skipping() + + with pytest.raises(SituationParsingError): + skipping.build_from_dict(SYSTEM, situation) + + +def test_building_a_situation_replaces_what_the_simulation_held(): + """What the old populations held, calculated or given, goes with them.""" + simulation = Simulation(situation=PARTIAL) + simulation.calculate("income_tax", PERIOD) + simulation.set_input("rent", PERIOD, [450.0]) + situation = { + "persons": {"a": {}, "b": {}, "c": {"salary": {PERIOD: 100.0}}}, + "households": {"h": {"parents": ["a", "b"], "children": ["c"]}}, + "axes": [ + [{"name": "salary", "count": 2, "min": 0, "max": 2000, "period": PERIOD}] + ], + } + + simulation._build_from_situation(situation) + expected = Simulation(situation=situation) + + assert_same(simulation, expected) + assert simulation._user_input_keys == expected._user_input_keys + assert sorted(simulation.input_variables) == sorted(expected.input_variables) + assert simulation.has_axes == expected.has_axes diff --git a/tests/core/test_reform_baseline_situation_property.py b/tests/core/test_reform_baseline_situation_property.py new file mode 100644 index 000000000..8eeec5fbf --- /dev/null +++ b/tests/core/test_reform_baseline_situation_property.py @@ -0,0 +1,360 @@ +"""A reform simulation's baseline is what the builder makes of the situation. + +For any reform made of changes to input variables' defaults (person and +household variables; float, int, enum, date and string; monthly variables +whose yearly inputs are divided or dispatched between months, and an eternal +one), neutralized variables (inputs and formulas), a variable only the reform +has, and a changed tax rate; and for any situation of one to three people in +one or two households that gives any of those variables, for any of several +periods, to some people or households and not to others, with no axes, one +dimension of parallel axes, or two perpendicular dimensions (or that gives +each variable's values for everyone, without entities): + +* Differential: the baseline calculates, for each output and period, what the + simulation ``SimulationBuilder`` makes of the situation under the baseline + system alone calculates. The variable only the reform has is left out of + the situation first: the baseline system cannot take it. +* The baseline has the same population and records the same inputs + (``_user_input_keys``, ``input_variables``) as that simulation, and exports + the same ``to_input_dict()``. +* It is the reform simulation's ``"baseline"`` branch, under the baseline + system. +* The reform simulation calculates what the builder makes of the situation + under another copy of the reform's system: the fix leaves it as it was. + +Examples, including the cases that failed before (PolicyEngine/ +policyengine-core#612), are in ``test_reform_baseline_situation.py``. +""" + +from __future__ import annotations + +import copy +import datetime +import functools + +import numpy as np +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.periods import instant # noqa: E402 +from policyengine_core.reforms import Reform # noqa: E402 +from policyengine_core.simulations import SimulationBuilder # noqa: E402 +from policyengine_core.variables import Variable # noqa: E402 +from tests.core.test_reform_baseline_system import doubled_salary # noqa: E402 + +PERIOD = "2022-01" +SYSTEM = Simulation.default_tax_benefit_system_instance +ADDED = "doubled_salary" +MONTHS = ["2022-01", "2022-03"] +INPUT_PERIODS = [*MONTHS, "2022"] +STATUSES = ["owner", "tenant", "free_lodger", "homeless"] +DATES = st.dates(datetime.date(1930, 1, 1), datetime.date(2021, 12, 31)) +POSTAL_CODES = st.text(alphabet="0123456789", min_size=5, max_size=5) + +# Each input variable: the values a situation (or a reform's default) gives +# it, and the periods a situation gives it for. +PERSON_INPUTS = { + "salary": (st.integers(0, 10_000).map(float), INPUT_PERIODS), + "age": (st.integers(0, 100), INPUT_PERIODS), + "birth": (DATES.map(str), ["ETERNITY"]), +} +HOUSEHOLD_INPUTS = { + "rent": (st.integers(0, 2000).map(float), INPUT_PERIODS), + "accommodation_size": (st.integers(0, 300).map(float), INPUT_PERIODS), + "housing_occupancy_status": (st.sampled_from(STATUSES), INPUT_PERIODS), + "postal_code": (POSTAL_CODES, INPUT_PERIODS), +} +DEFAULTS = { + "salary": st.integers(1, 5000).map(float), + "age": st.integers(1, 90), + "birth": DATES, + "rent": st.integers(1, 2000).map(float), + "accommodation_size": st.integers(1, 300).map(float), + "housing_occupancy_status": st.sampled_from(STATUSES), + "postal_code": POSTAL_CODES, +} +NEUTRALIZABLE = ["salary", "rent", "income_tax", "housing_tax"] +AXES = { + "person": ["salary", "age"], + "household": ["rent", "accommodation_size"], +} +# What is compared: each variable at the periods its definition takes. +OUTPUT_MONTHS = [*MONTHS, "2022-07"] +OUTPUTS = { + "salary": OUTPUT_MONTHS, + "age": OUTPUT_MONTHS, + "birth": [PERIOD], + "income_tax": OUTPUT_MONTHS, + "social_security_contribution": OUTPUT_MONTHS, + "disposable_income": OUTPUT_MONTHS, + "rent": OUTPUT_MONTHS, + "accommodation_size": OUTPUT_MONTHS, + "housing_occupancy_status": OUTPUT_MONTHS, + "postal_code": OUTPUT_MONTHS, + "household_income": OUTPUT_MONTHS, + "total_taxes": OUTPUT_MONTHS, + "housing_tax": ["2022"], +} + + +@functools.lru_cache(maxsize=None) +def _reform(spec: tuple) -> type: + """A reform class from ``spec``: (defaults, neutralized, added, rate).""" + defaults, neutralized, added, rate = spec + + def apply(self): + for name, value in defaults: + if name == "housing_occupancy_status": + # A member of this system's own enum class. + value = self.get_variable(name).possible_values[value] + self.update_variable(type(name, (Variable,), {"default_value": value})) + for name in neutralized: + self.neutralize_variable(name) + if added: + self.update_variable(doubled_salary) + if rate is not None: + + def modify(parameters): + parameters.taxes.income_tax_rate.update( + start=instant("2015-01-01"), value=rate + ) + return parameters + + self.modify_parameters(modify) + + return type("generated_reform", (Reform,), {"apply": apply}) + + +@st.composite +def reform_specs(draw): + names = draw(st.lists(st.sampled_from(sorted(DEFAULTS)), unique=True)) + defaults = tuple((name, draw(DEFAULTS[name])) for name in names) + neutralized = tuple(draw(st.lists(st.sampled_from(NEUTRALIZABLE), unique=True))) + added = draw(st.booleans()) + rate = draw(st.none() | st.sampled_from([0.0, 0.3])) + return defaults, neutralized, added, rate + + +def _given(draw, inputs, extra=None): + """Variable values a person or household is given: some, none or all.""" + inputs = {**inputs, **(extra or {})} + entry = {} + for name in sorted(inputs): + values, periods = inputs[name] + dated = draw(st.dictionaries(st.sampled_from(periods), values, max_size=2)) + if dated: + entry[name] = dated + return entry + + +@st.composite +def axes(draw, counts): + dimensions = draw(st.sampled_from([0, 0, 1, 2])) + result = [] + for _ in range(dimensions): + entity = draw(st.sampled_from(sorted(AXES))) + count = draw(st.integers(1, 3)) + parallel = [] + for _ in range(draw(st.integers(1, 2))): + low = draw(st.integers(0, 1000)) + parallel.append( + { + "name": draw(st.sampled_from(AXES[entity])), + "count": count, + "index": draw(st.integers(0, counts[entity] - 1)), + "min": low, + "max": low + draw(st.integers(0, 5000)), + "period": draw(st.sampled_from([PERIOD, "2022"])), + } + ) + result.append(parallel) + return result + + +@st.composite +def situations(draw): + people = [f"p{index}" for index in range(draw(st.integers(1, 3)))] + household_count = draw(st.integers(1, min(2, len(people)))) + # The first people found the households; the rest join any of them. + homes = list(range(household_count)) + draw( + st.lists( + st.integers(0, household_count - 1), + min_size=len(people) - household_count, + max_size=len(people) - household_count, + ) + ) + added = {ADDED: (st.integers(0, 1000).map(float), MONTHS)} + situation = { + "persons": {person: _given(draw, PERSON_INPUTS, added) for person in people}, + "households": {}, + } + for home in range(household_count): + members = [person for person, at in zip(people, homes) if at == home] + household = {"parents": members[:2], **_given(draw, HOUSEHOLD_INPUTS)} + if members[2:]: + household["children"] = members[2:] + situation["households"][f"h{home}"] = household + dimensions = draw(axes({"person": len(people), "household": household_count})) + if dimensions: + situation["axes"] = dimensions + return situation + + +@st.composite +def situations_without_entities(draw): + """Variables' values for each of one to three people, each in a household. + + The builder takes the number of people from the first value. + """ + count = draw(st.integers(1, 3)) + inputs = { + **PERSON_INPUTS, + **HOUSEHOLD_INPUTS, + ADDED: (st.integers(0, 1000).map(float), MONTHS), + } + situation = {} + for name in draw(st.permutations(sorted(inputs))): + values, periods = inputs[name] + dated = draw( + st.dictionaries( + st.sampled_from(periods), + st.lists(values, min_size=count, max_size=count), + max_size=2, + ) + ) + if dated: + situation[name] = dated + if not set(situation) - {ADDED}: + # Left without the added variable, the situation keeps its people. + situation["salary"] = {PERIOD: [0.0] * count} + return situation + + +def _built(situation, system): + builder = SimulationBuilder() + builder.default_period = PERIOD + return builder.build_from_dict(system, copy.deepcopy(situation)) + + +def _without_added(situation): + stripped = copy.deepcopy(situation) + stripped.pop(ADDED, None) + for person in stripped.get("persons", {}).values(): + person.pop(ADDED, None) + return stripped + + +def _assert_calculates_the_same(simulation, expected): + for variable, periods in OUTPUTS.items(): + for period in periods: + np.testing.assert_array_equal( + np.asarray(simulation.calculate(variable, period)), + np.asarray(expected.calculate(variable, period)), + err_msg=f"{variable} at {period}", + ) + + +def _assert_same_population(simulation, expected): + for key, population in expected.populations.items(): + other = simulation.populations[key] + assert other.count == population.count + np.testing.assert_array_equal(other.ids, population.ids) + if not population.entity.is_person: + np.testing.assert_array_equal( + other.members_entity_id, population.members_entity_id + ) + np.testing.assert_array_equal(other.members_role, population.members_role) + + +PARTIAL_SALARY = { + "persons": {"p0": {"salary": {PERIOD: 3000.0}}, "p1": {}}, + "households": {"h0": {"parents": ["p0", "p1"]}}, +} + + +@pytest.mark.filterwarnings("ignore:You cannot set a value") +@settings( + max_examples=200, + deadline=None, + suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], +) +@given( + spec=reform_specs(), + situation=situations() | situations_without_entities(), +) +# The reproduction in #612: a salary given to one person of two. +@example(spec=((("salary", 1000.0),), (), False, None), situation=PARTIAL_SALARY) +# The same along an axis over the other person. +@example( + spec=((("salary", 1000.0),), (), False, None), + situation={ + **PARTIAL_SALARY, + "axes": [ + [ + { + "name": "salary", + "count": 2, + "index": 1, + "min": 0, + "max": 5000, + "period": PERIOD, + } + ] + ], + }, +) +# A household's enum given to one household of two, a yearly rent to the other. +@example( + spec=((("housing_occupancy_status", "owner"), ("rent", 500.0)), (), False, None), + situation={ + "persons": {"p0": {}, "p1": {}}, + "households": { + "h0": {"parents": ["p0"], "housing_occupancy_status": {PERIOD: "homeless"}}, + "h1": {"parents": ["p1"], "rent": {"2022": 6000.0}}, + }, + }, +) +# Without entities, the builder sets each variable through ``set_input``. +@example( + spec=((("salary", 1000.0),), ("income_tax",), False, None), + situation={"salary": {PERIOD: [3000.0, 0.0]}, "income_tax": {PERIOD: [7.0, 9.0]}}, +) +# An input the reform neutralizes, and one only the reform has. +@example( + spec=((), ("salary",), True, None), + situation={ + "persons": { + "p0": {"salary": {PERIOD: 3000.0}, ADDED: {PERIOD: 9.0}}, + "p1": {"salary": {"2022": 1200.0}}, + }, + "households": {"h0": {"parents": ["p0", "p1"]}}, + }, +) +def test_the_baseline_is_what_the_builder_makes_of_the_situation(spec, situation): + added = spec[2] + if not added: + # Only a reform that adds the variable takes it. + situation = _without_added(situation) + reform = _reform(spec) + simulation = Simulation(situation=copy.deepcopy(situation), reform=reform) + baseline = simulation.baseline + expected = _built(_without_added(situation), SYSTEM) + + _assert_calculates_the_same(baseline, expected) + _assert_same_population(baseline, expected) + assert baseline._user_input_keys == expected._user_input_keys + assert sorted(baseline.input_variables) == sorted(expected.input_variables) + assert baseline.to_input_dict() == expected.to_input_dict() + + assert simulation.branches["baseline"] is baseline + assert baseline.parent_branch is simulation + assert baseline.tax_benefit_system is SYSTEM + + reform_system = Simulation.default_tax_benefit_system(reform=reform) + _assert_calculates_the_same(simulation, _built(situation, reform_system))