From a80a3b9641c5db1aef007de94ffff3d2b6dcb6ba Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:19:15 -0400 Subject: [PATCH 1/8] Carry over only inputs, the latest at or before the requested period Auto-carry-over took the latest-starting stored period of any kind and returned the default if it started after the requested period. Any later stored period therefore hid an earlier input: a later input (inputs for 2012 and 2014 gave the default for 2013), or a later period the simulation had already calculated (2014 calculated first made 2013 the default). And values the simulation calculated carried like inputs: a twelfth cached at a month by calculate_divide (the policyengine-us monthly_age bug #557 fixes), a value masked by defined_for, a default cached where defined_for was false everywhere, or a formula result from before the formula's end. Results depended on which periods were calculated first. The rule now: a period takes the input stored for the latest-starting period that starts no later than it (on a tie, the one ending last), preferring the variable's own definition-period unit; with none, the default. Values the simulation calculates are stored with derived=True: the mark lives in the storage with the value (InMemoryStorage and OnDiskStorage keep it per key, every put sets or clears it, it counts only while the key is stored, and clones copy it), so deletions, direct writes and #556's shared arrays cannot leave it stale. Holder.is_derived(period, branch_name) reports the mark of the value get_array reads, so a value one branch calculated never hides another branch's input. A derived value never replaces an input the branch reads (calculate_add/calculate_divide used to overwrite inputs stored in another unit), calculate_add caches only sums over several sub-periods, and dump/restore keeps the marks. Storages gain has(), which neither reads nor copies an array. Supersedes #557's carry-over change: its 253 tests pass here. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-carry-over-order.fixed.md | 1 + .../data_storage/in_memory_storage.py | 32 +- .../data_storage/on_disk_storage.py | 31 +- policyengine_core/holders/holder.py | 73 +++- policyengine_core/simulations/simulation.py | 74 +++- .../taxbenefitsystems/tax_benefit_system.py | 9 +- policyengine_core/tools/simulation_dumper.py | 26 +- tests/core/test_carry_over_order.py | 411 ++++++++++++++++++ tests/core/test_carry_over_order_property.py | 256 +++++++++++ tests/fixtures/carry_over.py | 171 ++++++++ 10 files changed, 1057 insertions(+), 27 deletions(-) create mode 100644 changelog.d/fix-carry-over-order.fixed.md create mode 100644 tests/core/test_carry_over_order.py create mode 100644 tests/core/test_carry_over_order_property.py create mode 100644 tests/fixtures/carry_over.py diff --git a/changelog.d/fix-carry-over-order.fixed.md b/changelog.d/fix-carry-over-order.fixed.md new file mode 100644 index 00000000..a2548855 --- /dev/null +++ b/changelog.d/fix-carry-over-order.fixed.md @@ -0,0 +1 @@ +Auto-carry-over now carries only inputs, taking the latest one stored for a period that starts no later than the requested period, so a carried value no longer depends on which periods were calculated first or on a later input. diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index a1bddf76..61c3b4c3 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -47,6 +47,10 @@ def __init__(self, is_eternal: bool): # with a copy the first time it is read. A key left here after code # outside this class empties ``_arrays`` costs one extra copy at most. self._shared = set() + # Keys whose value was stored with ``put(..., derived=True)``: calculated + # by the simulation rather than taken as input. A key counts only + # while it is stored, and every ``put`` sets or clears its mark. + self._derived = set() self.is_eternal = is_eternal def clone(self, share_arrays: bool = False) -> "InMemoryStorage": @@ -83,6 +87,7 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage": clone._arrays[key] = array.copy() else: clone._arrays = {key: array.copy() for key, array in self._arrays.items()} + clone._derived = set(self._derived) return clone def get(self, period: Period, branch_name: str = "default") -> ArrayLike: @@ -101,8 +106,29 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: self._shared.discard(key) return values + def has(self, period: Period, branch_name: str = "default") -> bool: + """Whether a value is stored for ``period`` under ``branch_name``. + + Unlike ``get``, this never copies an array shared by ``clone``. + """ + if self.is_eternal: + period = periods.period(periods.ETERNITY) + return f"{branch_name}:{periods.period(period)}" in self._arrays + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value stored for ``period`` under ``branch_name`` was + stored with ``derived=True``; ``False`` if none is stored.""" + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}:{periods.period(period)}" + return key in self._derived and key in self._arrays + def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -126,6 +152,10 @@ def put( key = f"{branch_name}:{period}" self._arrays[key] = value self._shared.discard(key) + if derived: + self._derived.add(key) + else: + self._derived.discard(key) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805..1f835f1c 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -22,6 +22,9 @@ def __init__( ): self._files = {} self._enums = {} + # File keys stored with ``put(..., derived=True)``; see + # ``InMemoryStorage``. + self._derived = set() self.is_eternal = is_eternal self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir @@ -43,6 +46,7 @@ def clone(self) -> "OnDiskStorage": ) clone._files = self._files.copy() clone._enums = self._enums.copy() + clone._derived = set(self._derived) clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) return clone @@ -63,8 +67,29 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: return None return self._decode_file(values) + def has(self, period: Period, branch_name: str = "default") -> bool: + """Whether a value is stored for ``period`` under ``branch_name``. + + Unlike ``get``, this reads no file. + """ + if self.is_eternal: + period = periods.period(periods.ETERNITY) + return f"{branch_name}_{periods.period(period)}" in self._files + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value stored for ``period`` under ``branch_name`` was + stored with ``derived=True``; ``False`` if none is stored.""" + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}_{periods.period(period)}" + return key in self._derived and key in self._files + def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -77,6 +102,10 @@ def put( value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path + if derived: + self._derived.add(filename) + else: + self._derived.discard(filename) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50..12a679c8 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -347,6 +347,7 @@ def _set( value: ArrayLike, branch_name: str = "default", validate_nan: bool = False, + derived: bool = False, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = getattr(simulation, "_user_input_contexts", None) @@ -367,17 +368,34 @@ def _set( ) if should_store_on_disk: - self._disk_storage.put(value, period, branch_name) + self._disk_storage.put(value, period, branch_name, derived=derived) else: - self._memory_storage.put(value, period, branch_name) + self._memory_storage.put(value, period, branch_name, derived=derived) if user_input_contexts: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() simulation._user_input_keys.add((self.variable.name, branch_name, period)) def put_in_cache( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: + """Cache ``value`` for ``period``. + + ``derived`` marks a value the simulation calculated rather than took + as input: a formula result, a carried, uprated or default value, a + twelfth of a yearly flow cached at a month by ``calculate_divide``, + or a sum over several sub-periods cached by ``calculate_add``. + Auto-carry-over never carries such a value into another period (see + ``is_derived``). The mark is stored with the value, for the (branch, + period) key written, and any later write to that key replaces it. + + A derived value never replaces an input that ``get_array(period, + branch_name)`` reads: the input is kept and nothing is stored. + """ if self._do_not_store: return @@ -388,7 +406,14 @@ def put_in_cache( ): return - self._set(period, value, branch_name) + if ( + derived + and self._branch_storing(period, branch_name) is not None + and not self.is_derived(period, branch_name) + ): + return + + self._set(period, value, branch_name, derived=derived) def default_array(self) -> ArrayLike: """ @@ -396,3 +421,43 @@ def default_array(self) -> ArrayLike: """ return self.variable.default_array(self.population.count) + + def _stores(self, period: Period, branch_name: str) -> bool: + """Whether a value is stored for ``period`` under ``branch_name``, + without reading or copying it.""" + return self._memory_storage.has(period, branch_name) or ( + self._disk_storage is not None + and self._disk_storage.has(period, branch_name) + ) + + def _branch_storing(self, period: Period, branch_name: str = "default") -> str: + """The branch whose stored value ``get_array(period, branch_name)`` + reads, in ``get_array``'s lookup order, or ``None`` if none stores one. + """ + if self._stores(period, branch_name): + return branch_name + if branch_name == "default": + return None + parent = ( + getattr(self.simulation, "parent_branch", None) if self.simulation else None + ) + while parent is not None: + if self._stores(period, parent.branch_name): + return parent.branch_name + parent = getattr(parent, "parent_branch", None) + return "default" if self._stores(period, "default") else None + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value ``get_array(period, branch_name)`` reads was + calculated by the simulation rather than set as an input. + + The answer comes from the branch that stores the value read: the + branch itself, else its ``parent_branch`` ancestors, else + ``default``. ``False`` if no value is stored for ``period``. + """ + storing = self._branch_storing(period, branch_name) + if storing is None: + return False + if self._memory_storage.has(period, storing): + return self._memory_storage.is_derived(period, storing) + return self._disk_storage.is_derived(period, storing) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d972..7f49e0ea 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -881,7 +881,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if np.all(~mask): array = holder.default_array() array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + holder.put_in_cache(array, period, self.branch_name, derived=True) return array array = None @@ -942,18 +942,49 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: and variable.calculate_output is None and len(known_periods) > 0 ): - # Variables with a calculate-output property specify - # Sort by period.start (temporal order). Sorting Period - # tuples lexicographically puts "year" before "month" - # alphabetically, so a known "2023" annual value would - # win over a later "2024-06" monthly value (bug H1). - last_known_period = max(known_periods, key=lambda p: p.start) - if last_known_period.start > period.start: - return holder.default_array() - # Pass branch_name through so auto-carry-over respects the - # active branch instead of reaching for the "default" - # branch's cache (bug H2). - array = holder.get_array(last_known_period, self.branch_name) + # Carry over the latest input: of the stored periods that + # start no later than ``period``, the one that starts last + # (on a tie, the one that ends last), preferring periods at + # the variable's own definition-period unit and using + # another unit only when there is none, as for an input to + # a variable with no ``set_input`` helper. Compare + # period.start (temporal order): sorting Period tuples + # lexicographically puts "year" before "month" + # alphabetically, so a known "2023" annual value would win + # over a later "2024-06" monthly value (bug H1). + # + # Only inputs carry. Every value this simulation + # calculated is marked derived when cached (formula + # results, carried, uprated and default values, a twelfth + # cached by ``calculate_divide``, a sum cached by + # ``calculate_add``), and carrying one would make the + # result depend on what was calculated first: a later + # period's carried value would hide an earlier input, and + # a value already masked by ``defined_for``, or given by a + # formula that has since ended, would carry forward. + # A later input does not carry backwards. + carry_over_periods = [ + known_period + for known_period in known_periods + if known_period.start <= period.start + and not holder.is_derived(known_period, self.branch_name) + ] + own_unit_periods = [ + known_period + for known_period in carry_over_periods + if known_period.unit == variable.definition_period + ] + if carry_over_periods: + last_known_period = max( + own_unit_periods or carry_over_periods, + key=lambda p: (p.start, p.stop), + ) + # Pass branch_name through so auto-carry-over respects + # the active branch instead of reaching for the + # "default" branch's cache (bug H2). + array = holder.get_array(last_known_period, self.branch_name) + else: + array = holder.default_array() else: array = holder.default_array() @@ -969,7 +1000,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: array = EnumArray(array, variable.possible_values) array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + # Calculated, not input: auto-carry-over never carries it. + holder.put_in_cache(array, period, self.branch_name, derived=True) except SpiralError: array = holder.default_array() @@ -1044,12 +1076,16 @@ def calculate_add( ) ) + sub_periods = list(period.get_subperiods(variable.definition_period)) result = sum( - self.calculate(variable_name, sub_period) - for sub_period in period.get_subperiods(variable.definition_period) + self.calculate(variable_name, sub_period) for sub_period in sub_periods ) - holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + # Cache only a sum over several sub-periods, as derived from them. A + # single sub-period's value is already stored, as an input or + # derived, by ``calculate``. + if len(sub_periods) > 1: + holder = self.get_holder(variable.name) + holder.put_in_cache(result, period, self.branch_name, derived=True) return result def calculate_divide( @@ -1082,7 +1118,7 @@ def calculate_divide( computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + holder.put_in_cache(result, period, self.branch_name, derived=True) return result elif period.unit == periods.YEAR: return self.calculate(variable_name, period) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index c381b11c..1fae168f 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -87,7 +87,14 @@ class TaxBenefitSystem: parameters_dir: str = None """Directory containing the YAML parameter tree.""" auto_carry_over_input_variables: bool = False - """Whether to automatically carry over input variables when calculating a variable for a period different from the period of the input variables.""" + """Whether a variable with no formula result for a period (and no + ``uprating`` path from an earlier period) takes an input from another + period. If so, the period takes the input stored for the latest-starting + period that starts no later than it (on a tie, the one that ends last), + masked by the variable's ``defined_for``; with no such input, the + default. An input never carries backwards, and values the simulation + calculated never carry (see ``Holder.is_derived``), so the result does + not depend on which periods were calculated first.""" basic_inputs: List[str] = None """Short list of basic inputs to get medium accuracy.""" modelled_policies: str = None diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index c3db0c4f..e2376e1e 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -9,6 +9,10 @@ from policyengine_core.periods import ETERNITY from policyengine_core.simulations import Simulation +# Periods, one per line, whose dumped value the simulation calculated (see +# ``Holder.is_derived``), so a restored simulation does not carry them over. +DERIVED_PERIODS_FILE = "derived_periods.txt" + def dump_simulation(simulation, directory): """ @@ -69,6 +73,20 @@ def _dump_holder(holder, directory): for period in holder.get_known_periods(): value = holder.get_array(period) disk_storage.put(value, period) + _dump_derived_periods(holder, disk_storage.storage_dir) + + +def _dump_derived_periods(holder, storage_dir, branch_name="default"): + derived = sorted( + { + str(period) + for period in holder.get_known_periods() + if holder.is_derived(period, branch_name) + } + ) + if derived: + with open(os.path.join(storage_dir, DERIVED_PERIODS_FILE), "w") as file: + file.write("\n".join(derived) + "\n") def _dump_entity(population, directory): @@ -135,6 +153,12 @@ def _restore_holder(simulation, variable, directory): holder = simulation.get_holder(variable) + derived_periods_path = os.path.join(storage_dir, DERIVED_PERIODS_FILE) + derived_periods = set() + if os.path.exists(derived_periods_path): + with open(derived_periods_path) as file: + derived_periods = set(file.read().split()) + for period in disk_storage.get_known_periods(): value = disk_storage.get(period) - holder.put_in_cache(value, period) + holder.put_in_cache(value, period, derived=str(period) in derived_periods) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py new file mode 100644 index 00000000..7e76a414 --- /dev/null +++ b/tests/core/test_carry_over_order.py @@ -0,0 +1,411 @@ +"""Auto-carry-over depends only on the inputs, not on what was calculated first. + +With ``auto_carry_over_input_variables``, a variable with no formula result +for a period takes its value from another period. Before this fix, +``Simulation._calculate`` took the latest-starting stored period other than +a ``calculate_add``/``calculate_divide`` cache (#557), and returned the +default if that period started after the requested one. Two things made the +result depend on calculation order: + +* Any later stored period blocked carry-over. A later input did (inputs for + 2012 and 2014 gave the default for 2013), and so did a value the simulation + had itself calculated: once 2014 had been calculated, carrying the 2012 + input into it, 2013 came out as the default instead of the 2012 input. +* Values the simulation calculated were carried like inputs. A value masked by + ``defined_for`` carried its mask into later periods, a default cached for a + period where ``defined_for`` was false everywhere replaced the input, and a + formula result carried past the formula's ``end``. + +The rule now is the reference rule in ``tests/fixtures/carry_over.py``. Every +value the simulation caches is marked derived for the (branch, period) key it +is stored under, ``Holder.is_derived`` reports the mark of the value a branch +reads, and only values that are not derived carry. A derived value never +replaces an input a branch can read. + +Each regression below checks a simulation that calculated other periods first +against a fresh simulation that calculates only the period in question; the +properties are in ``test_carry_over_order_property.py``. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.reforms import Reform +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.carry_over import alone, build_system, simulation + + +@pytest.fixture(scope="module") +def system(): + return build_system() + + +def _after(system, inputs, requests, variable, period): + built = simulation(system, inputs) + for requested_variable, requested_period in requests: + built.calculate(requested_variable, requested_period) + return built.calculate(variable, period) + + +@pytest.mark.parametrize( + "variable,earlier_input,later_period,between", + [ + ("carried", "2012", "2014", "2013"), + ("carried_count", "2012", "2014", "2013"), + ("carried_monthly", "2012-01", "2012-06", "2012-03"), + ], +) +def test_calculated_later_period_does_not_hide_earlier_input( + system, variable, earlier_input, later_period, between +): + inputs = {variable: {earlier_input: [7, 8]}} + result = _after(system, inputs, [(variable, later_period)], variable, between) + np.testing.assert_array_equal(result, [7, 8]) + np.testing.assert_array_equal(result, alone(system, inputs, variable, between)) + + +@pytest.mark.parametrize( + "variable,earlier,later,between", + [ + ("carried", "2012", "2014", "2013"), + ("carried_monthly", "2012-01", "2012-06", "2012-03"), + ], +) +def test_later_input_does_not_hide_earlier_input( + system, variable, earlier, later, between +): + inputs = {variable: {earlier: [7, 8], later: [9, 10]}} + np.testing.assert_array_equal(alone(system, inputs, variable, between), [7, 8]) + + +def test_inputs_are_stored_in_any_order(system): + inputs = {"carried": {"2014": [9, 10], "2012": [7, 8]}} + for year, expected in [("2013", [7, 8]), ("2014", [9, 10]), ("2016", [9, 10])]: + np.testing.assert_array_equal(alone(system, inputs, "carried", year), expected) + + +@pytest.mark.parametrize( + "first,second", [("2013", "2013-01"), ("2013-01", "2013")], ids=["year", "month"] +) +def test_inputs_starting_together_carry_the_one_ending_last(system, first, second): + values = {"2013": [1, 1], "2013-01": [2, 2]} + inputs = { + "year_input_without_helper": {first: values[first], second: values[second]} + } + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2014"), [1, 1] + ) + + +def test_later_input_does_not_carry_backwards(system): + built = simulation(system, {"carried": {"2012": [7, 8], "2014": [9, 10]}}) + np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) + np.testing.assert_array_equal(built.calculate("carried", "2015"), [9, 10]) + np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) + + +def test_value_masked_by_defined_for_does_not_carry_its_mask(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2013"), [0, 8] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +def test_default_cached_where_defined_for_is_false_does_not_carry(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, False], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2013"), [0, 0] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +@pytest.mark.parametrize("own_input", [None, "2012"]) +def test_formula_result_does_not_carry_past_the_formula_end(system, own_input): + inputs = {"carried": {"2012": [1, 2]}} + if own_input is not None: + inputs["formula_until_2013"] = {own_input: [40, 50]} + expected = alone(system, inputs, "formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [0, 0] if own_input is None else [40, 50]) + result = _after( + system, inputs, [("formula_until_2013", "2013")], "formula_until_2013", "2014" + ) + np.testing.assert_array_equal(result, expected) + + +def test_input_set_after_a_period_defaulted_still_carries(system): + built = simulation(system, {}) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [0, 0]) + built.set_input("carried", "2012", np.array([7.0, 8.0])) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) + + +def test_periods_with_no_earlier_input_are_cached(system): + built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) + for month in range(1, 13): + for _ in range(2): + built.calculate("carried_monthly", f"2012-{month:02d}") + holder = built.get_holder("carried_monthly") + for month in range(1, 13): + assert holder.get_array(periods.period(f"2012-{month:02d}")) is not None + + +def test_branch_input_is_not_hidden_by_a_value_the_parent_carried(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + branch = built.get_branch("reform") + branch.set_input("carried", "2012", np.array([70.0, 80.0])) + np.testing.assert_array_equal(branch.calculate("carried", "2014"), [70, 80]) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) + + +def test_derived_mark_belongs_to_the_branch_that_stored_the_value(system): + """A value calculated on one branch does not make another branch's input + for the same period derived.""" + + def build(): + built = simulation(system, {"carried": {"2011": [7, 7]}}) + built.get_holder("formula_until_2013").set_input( + periods.period("2012"), np.array([30.0, 30.0]), "other" + ) + return built + + expected = build().get_branch("other").calculate("formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [30, 30]) + built = build() + np.testing.assert_array_equal( + built.calculate("formula_until_2013", "2012"), [12, 12] + ) + other = built.get_branch("other") + holder = other.get_holder("formula_until_2013") + assert holder.is_derived(periods.period("2012"), "default") + assert not holder.is_derived(periods.period("2012"), "other") + np.testing.assert_array_equal( + other.calculate("formula_until_2013", "2014"), expected + ) + + +def test_add_over_a_calculated_period_keeps_it_derived(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate_add("carried_if_eligible", "2013"), [0, 8] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +def test_add_over_an_input_period_keeps_it_an_input(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + np.testing.assert_array_equal(built.calculate_add("carried", "2012"), [7, 8]) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + + +@pytest.mark.parametrize( + "variable,input_period,kind,request_period,later_period", + [ + # The sum over the two years would replace the two-year input. + ("year_input_without_helper", "year:2012:2", "add", "year:2012:2", "2014"), + # The twelfth of 2012 would replace the January input. + ("year_input_without_helper", "2012-01", "divide", "2012-01", "2014"), + # The sum of the months would replace the year input. + ("month_input_without_helper", "2013", "add", "2013", "2014-01"), + ], +) +@pytest.mark.parametrize("on_branch", [False, True], ids=["simulation", "branch"]) +def test_add_or_divide_never_replaces_an_input( + system, variable, input_period, kind, request_period, later_period, on_branch +): + inputs = {variable: {input_period: [120, 24]}} + expected = alone(system, inputs, variable, later_period) + np.testing.assert_array_equal(expected, [120, 24]) + built = simulation(system, inputs) + if on_branch: + built = built.get_branch("reform") + if kind == "add": + built.calculate_add(variable, request_period) + else: + built.calculate_divide(variable, request_period) + holder = built.get_holder(variable) + np.testing.assert_array_equal( + holder.get_array(periods.period(input_period), built.branch_name), [120, 24] + ) + np.testing.assert_array_equal(built.calculate(variable, later_period), expected) + + +def test_add_over_an_anchored_year_keeps_the_anchored_input(system): + """``set_input`` stores a month of a yearly flow at the year starting + then; ``ADD`` over that year is a single calendar-year sub-period.""" + inputs = {"carried": {"2012": [1, 2], "2013-06": [5, 6]}} + expected = alone(system, inputs, "carried", "2014") + np.testing.assert_array_equal(expected, [5, 6]) + built = simulation(system, inputs) + built.calculate_add("carried", "year:2013-06") + np.testing.assert_array_equal(built.calculate("carried", "2014"), expected) + + +def test_month_cache_of_a_year_input_does_not_carry_into_the_next_year(system): + """policyengine-us computes ``monthly_age`` from ``age`` one month at a + time; each month of a yearly flow caches a twelfth of the year (#557).""" + inputs = {"carried": {"2024": [40, 6]}} + built = simulation(system, inputs) + for month in range(1, 13): + built.calculate("carried", f"2024-{month:02d}") + np.testing.assert_allclose(built.calculate("carried", "2024-12"), [40 / 12, 6 / 12]) + np.testing.assert_array_equal(built.calculate("carried", "2025"), [40, 6]) + + +def test_own_unit_input_wins_over_an_input_at_another_unit(system): + """As before 3.24.0, and as in #557: the year input, not the later month + input of a variable with no ``set_input`` helper.""" + inputs = {"year_input_without_helper": {"2024": [120, 12], "2024-12": [9, 9]}} + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2025"), [120, 12] + ) + + +# Lifecycle paths that write or drop stored values outside ``Holder._set`` +# (the cases #557's review found to leave per-period marks stale). + +FLOW = "year_input_without_helper" + + +def test_reform_replay_restores_an_input_at_another_unit(system): + built = simulation(system, {FLOW: {"2024-06": [600, 60]}}) + branch = built.get_branch("child") + branch.calculate_divide(FLOW, "2024-06") + branch.apply_reform(_noop) + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [600, 60]) + + +def test_dump_and_restore_keep_carrying_the_input(system, tmp_path): + built = simulation(system, {FLOW: {"2024": [120, 12]}}) + built.calculate_divide(FLOW, "2024-06") + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), system) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [120, 12]) + np.testing.assert_array_equal(restored.calculate(FLOW, "2025"), [120, 12]) + + +def test_deleting_a_branch_cache_exposes_the_ancestor_input(system): + built = simulation(system, {FLOW: {"2024-06": [600, 60]}}) + branch = built.get_branch("child") + branch.calculate_divide(FLOW, "2024-06") + branch.get_holder(FLOW).delete_arrays(periods.period("2024-06"), "child") + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [600, 60]) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [600, 60]) + + +def test_deleting_a_branch_input_does_not_expose_the_ancestor_twelfth(system): + built = simulation(system, {FLOW: {"2024": [120, 12]}}) + built.calculate_divide(FLOW, "2024-06") + branch = built.get_branch("child") + branch.set_input(FLOW, "2024-06", np.array([600.0, 60.0])) + branch.get_holder(FLOW).delete_arrays(periods.period("2024-06"), "child") + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [120, 12]) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [120, 12]) + + +@pytest.mark.parametrize("delete_first", [True, False], ids=["deleted", "replaced"]) +def test_a_value_written_straight_to_storage_is_not_derived(system, delete_first): + """The mark is stored with the value: a value written later at the same + key, even straight to the storage as ``apply_reform``'s input replay + does, is not taken for a derived one.""" + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + holder = built.get_holder("carried") + assert holder.is_derived(periods.period("2013")) + if delete_first: + holder.delete_arrays(periods.period("2013")) + assert not holder.is_derived(periods.period("2013")) + holder._memory_storage.put(np.array([5.0, 5.0]), periods.period("2013")) + assert not holder.is_derived(periods.period("2013")) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [5, 5]) + + +def test_carry_over_in_a_branch_copies_only_the_array_it_reads(system): + """Deciding which periods are inputs reads no shared array (#556): the + branch copies only the input it carries.""" + built = simulation( + system, {"carried": {"2012": [1, 1], "2014": [2, 2], "2016": [3, 3]}} + ) + built.calculate("carried", "2017") + branch = built.get_branch("reform") + storage = branch.get_holder("carried")._memory_storage + shared = len(storage._shared) + np.testing.assert_array_equal(branch.calculate("carried", "2020"), [3, 3]) + assert len(storage._shared) == shared - 1 + + +def test_input_starting_with_the_requested_period_carries(system): + """An input in another unit that starts when the period does carries.""" + built = simulation(system, {"year_input_without_helper": {"2013-01": [3, 4]}}) + np.testing.assert_array_equal( + built.calculate("year_input_without_helper", "2013"), [3, 4] + ) + + +def test_carried_value_is_derived_and_the_input_is_not(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + holder = built.get_holder("carried") + assert holder.is_derived(periods.period("2013")) + assert not holder.is_derived(periods.period("2012")) + assert not holder.is_derived(periods.period("2015")) + branch = built.get_branch("reform") + assert branch.get_holder("carried").is_derived(periods.period("2013"), "reform") + + +def test_restored_simulation_keeps_derived_marks(system, tmp_path): + inputs = {"carried": {"2012": [1, 2]}} + built = simulation(system, inputs) + np.testing.assert_array_equal(built.calculate("formula_until_2013", "2012"), [6, 7]) + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), system) + expected = alone(system, inputs, "formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [0, 0]) + np.testing.assert_array_equal( + built.calculate("formula_until_2013", "2014"), expected + ) + np.testing.assert_array_equal( + restored.calculate("formula_until_2013", "2014"), expected + ) + np.testing.assert_array_equal(restored.calculate("carried", "2014"), [1, 2]) + + +class _noop(Reform): + def apply(self): + pass + + +def test_inputs_still_carry_after_apply_reform(system): + inputs = {"carried": {"2012": [7, 8]}} + built = simulation(system, inputs) + built.calculate("carried", "2014") + built.apply_reform(_noop) + holder = built.get_holder("carried") + assert not any(holder.is_derived(p) for p in holder.get_known_periods()) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) diff --git a/tests/core/test_carry_over_order_property.py b/tests/core/test_carry_over_order_property.py new file mode 100644 index 00000000..d343d830 --- /dev/null +++ b/tests/core/test_carry_over_order_property.py @@ -0,0 +1,256 @@ +"""Properties of auto-carry-over: order independence and the reference rule. + +For any inputs and any earlier requests (plain calculations in either unit, +``ADD`` and ``DIVIDE``), what a simulation then returns for a period, read +from the simulation or from a branch or nested branch forked after those +requests, equals what a fresh simulation given the same inputs returns for +that period alone, and equals the reference rule in +``tests/fixtures/carry_over.py``. Regressions are in +``test_carry_over_order.py``. +""" + +from __future__ import annotations + +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 import periods # noqa: E402 +from tests.fixtures.carry_over import ( # noqa: E402 + COUNT, + alone, + build_system, + reference, + request, + simulation, +) + +SYSTEM = build_system() + +YEARS = [str(year) for year in range(2010, 2017)] +MONTHS = [f"{year}-{month:02d}" for year in (2012, 2013) for month in (1, 2, 6, 12)] +MULTI_YEAR = ["year:2012:2"] + +YEAR_VARIABLES = [ + "carried", + "carried_count", + "eligible", + "carried_if_eligible", + "formula_until_2013", + "year_input_without_helper", +] +MONTH_VARIABLES = ["carried_monthly", "month_input_without_helper"] +# What each variable's result reads: itself and its formula and defined_for. +READS = { + "carried_if_eligible": ["carried_if_eligible", "eligible"], + "formula_until_2013": ["formula_until_2013", "carried"], +} + +FLOAT = st.floats(min_value=-100, max_value=100, allow_nan=False, width=32) +VALUES = { + "carried_count": st.integers(min_value=-100, max_value=100), + "eligible": st.booleans(), +} +INPUT_PERIODS = {variable: YEARS for variable in YEAR_VARIABLES} +INPUT_PERIODS["carried_monthly"] = MONTHS +# Variables with no set_input helper store inputs at any period. +INPUT_PERIODS["year_input_without_helper"] = YEARS + MONTHS + MULTI_YEAR +INPUT_PERIODS["month_input_without_helper"] = MONTHS + YEARS + + +def _shift(period, step): + period = periods.period(period) + return str(period.offset(step, period.unit)) + + +@st.composite +def scenarios(draw): + inputs = {} + for variable, input_periods in INPUT_PERIODS.items(): + values = draw( + st.dictionaries( + st.sampled_from(input_periods), + st.lists(VALUES.get(variable, FLOAT), min_size=COUNT, max_size=COUNT), + max_size=3, + ) + ) + if values: + inputs[variable] = values + variable = draw(st.sampled_from(YEAR_VARIABLES + MONTH_VARIABLES)) + unit_periods = YEARS if variable in YEAR_VARIABLES else MONTHS + period = draw(st.sampled_from(unit_periods)) + # Bias requests toward what the target reads and toward periods next to + # the inputs and the target: each bug needs a short chain of those. + nearby = {period} + for name in READS.get(variable, [variable]): + nearby.update(inputs.get(name, {})) + nearby.update( + _shift(known, step) + for known in list(nearby) + if not known.startswith("year:") + for step in (-1, 1) + ) + nearby = sorted(nearby) + request_variables = st.one_of( + st.sampled_from(READS.get(variable, [variable])), + st.sampled_from(YEAR_VARIABLES + MONTH_VARIABLES), + ) + + @st.composite + def one_request(draw): + name = draw(request_variables) + kind = draw(st.sampled_from(["calculate", "calculate", "add", "divide"])) + candidates = nearby + YEARS + MONTHS + if name in YEAR_VARIABLES: + candidates += MULTI_YEAR + requested = draw(st.sampled_from(candidates)) + unit = periods.period(requested).unit + if kind == "divide" and ( + name not in YEAR_VARIABLES or periods.period(requested).size != 1 + ): + kind = "calculate" + if name in MONTH_VARIABLES and requested.startswith("year:"): + kind = "add" + if kind == "add" and name in YEAR_VARIABLES and unit == periods.MONTH: + kind = "calculate" + if kind == "calculate" and requested.startswith("year:"): + kind = "add" + return kind, name, requested + + requests = draw(st.lists(one_request(), max_size=8)) + branch = draw(st.sampled_from([None, "reform", "nested"])) + branch_requests = draw(st.lists(one_request(), max_size=4)) if branch else [] + return inputs, requests, branch, branch_requests, (variable, period) + + +@settings( + max_examples=400, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given(scenario=scenarios()) +# One example per regression shape in test_carry_over_order.py. +@example( + scenario=( + {"carried": {"2012": [7, 8]}}, + [("calculate", "carried", "2014")], + None, + [], + ("carried", "2013"), + ) +) +@example( + scenario=( + { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + }, + [("calculate", "carried_if_eligible", "2013")], + None, + [], + ("carried_if_eligible", "2014"), + ) +) +@example( + scenario=( + { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, False], "2014": [True, True]}, + }, + [("add", "carried_if_eligible", "2013")], + "reform", + [], + ("carried_if_eligible", "2014"), + ) +) +@example( + scenario=( + {"carried": {"2012": [1, 2]}}, + [("calculate", "formula_until_2013", "2013")], + None, + [], + ("formula_until_2013", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"year:2012:2": [120, 24]}}, + [("add", "year_input_without_helper", "year:2012:2")], + None, + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2012-01": [120, 24]}}, + [("divide", "year_input_without_helper", "2012-01")], + "nested", + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"month_input_without_helper": {"2013": [120, 24]}}, + [("add", "month_input_without_helper", "2013")], + None, + [], + ("month_input_without_helper", "2013-06"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2013-01": [2, 2], "2013": [1, 1]}}, + [], + None, + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2013-01": [3, 4]}}, + [], + None, + [], + ("year_input_without_helper", "2013"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"year:2012:2": [120, 24]}}, + [], + "reform", + [("add", "year_input_without_helper", "year:2012:2")], + ("year_input_without_helper", "2014"), + ) +) +def test_carry_over_depends_only_on_the_inputs(scenario): + inputs, requests, branch, branch_requests, (variable, period) = scenario + built = simulation(SYSTEM, inputs) + for kind, name, requested in requests: + request(built, kind, name, requested) + if branch == "reform": + built = built.get_branch("reform") + elif branch == "nested": + built = built.get_branch("reform").get_branch("nested") + for kind, name, requested in branch_requests: + request(built, kind, name, requested) + result = built.calculate(variable, period) + + message = ( + f"{variable} {period} after {requests}, then {branch_requests} on {branch}" + ) + np.testing.assert_array_equal( + result, alone(SYSTEM, inputs, variable, period), err_msg=message + ) + np.testing.assert_array_equal( + result, reference(SYSTEM, inputs, variable, period), err_msg=message + ) diff --git a/tests/fixtures/carry_over.py b/tests/fixtures/carry_over.py new file mode 100644 index 00000000..ad535e3c --- /dev/null +++ b/tests/fixtures/carry_over.py @@ -0,0 +1,171 @@ +"""Variables, simulations and a reference rule for the auto-carry-over tests. + +Shared by ``tests/core/test_carry_over_order.py`` (regressions) and +``tests/core/test_carry_over_order_property.py`` (properties). + +The reference rule, for a variable with no formula result for a period: an +input stored for the period itself is read back as stored; otherwise the +period takes the input stored for the latest-starting period that starts no +later than it (on a tie, the one that ends last), preferring inputs at the +variable's own definition-period unit, masked by the variable's +``defined_for`` at the period; with no such input, the default. +""" + +from __future__ import annotations + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + +COUNT = 2 +FORMULA_END_YEAR = 2013 + + +class carried(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly flow input with no formula" + + +class carried_count(Variable): + value_type = int + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly integer input with no formula" + + +class eligible(Variable): + value_type = bool + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly boolean input that carried_if_eligible is defined for" + + +class carried_if_eligible(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + defined_for = "eligible" + label = "Yearly flow input with no formula, defined for eligible people" + + +class formula_until_2013(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + end = f"{FORMULA_END_YEAR}-12-31" + label = "Yearly variable whose formula ends after 2013" + + def formula(person, period, parameters): + return person("carried", period) + 5 + + +class carried_monthly(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly flow input with no formula" + + +class year_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = None + label = "Yearly input with no set_input helper, stored at any period" + + +class month_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + set_input = None + label = "Monthly input with no set_input helper, stored at any period" + + +VARIABLES = ( + carried, + carried_count, + eligible, + carried_if_eligible, + formula_until_2013, + carried_monthly, + year_input_without_helper, + month_input_without_helper, +) + + +def build_system() -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variables(*VARIABLES) + return system + + +def simulation(system, inputs): + """A simulation of ``COUNT`` people given ``{variable: {period: values}}``.""" + built = SimulationBuilder().build_default_simulation(system, count=COUNT) + for variable, values in inputs.items(): + for period, array in values.items(): + built.set_input(variable, period, np.array(array)) + return built + + +def alone(system, inputs, variable, period): + """What a fresh simulation given ``inputs`` returns for one request.""" + return simulation(system, inputs).calculate(variable, period) + + +def request(built, kind, variable, period): + if kind == "add": + built.calculate_add(variable, period) + elif kind == "divide": + built.calculate_divide(variable, period) + else: + built.calculate(variable, period) + + +def _stored_inputs(system, inputs, variable): + """The inputs ``set_input`` stores: it ignores those after the ``end``.""" + end = system.get_variable(variable).end + return { + periods.period(stored): values + for stored, values in inputs.get(variable, {}).items() + if end is None or periods.period(stored).start.date <= end + } + + +def reference(system, inputs, variable, period): + """The carry-over rule, computed from the inputs alone. + + ``period`` is in the variable's definition period. + """ + period = periods.period(period) + default = system.get_variable(variable).default_array(COUNT) + stored_inputs = _stored_inputs(system, inputs, variable) + if period in stored_inputs: + # An input for the period itself is read back as stored: neither a + # formula nor ``defined_for`` applies to it. + return np.array(stored_inputs[period], dtype=default.dtype) + if variable == "formula_until_2013" and period.start.year <= FORMULA_END_YEAR: + return (reference(system, inputs, "carried", period) + 5).astype(default.dtype) + earlier = [stored for stored in stored_inputs if stored.start <= period.start] + own_unit = [ + stored + for stored in earlier + if stored.unit == system.get_variable(variable).definition_period + ] + if earlier: + latest = max( + own_unit or earlier, key=lambda stored: (stored.start, stored.stop) + ) + value = np.array(stored_inputs[latest]).astype(default.dtype) + else: + value = default + if variable == "carried_if_eligible": + value = np.where(reference(system, inputs, "eligible", period), value, default) + return value From 3c83380b781b5d13bfb6bc7a5b0eb6b96df49c7e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:54:27 -0400 Subject: [PATCH 2/8] Keep uprating, input helpers, dumps and deletions consistent with input-only carry-over Round-2 review findings: - A default cached for a period with no earlier input became the base the uprating path uprated later periods from; the default is not cached for variables with uprating (as before). - A value calculated inside a set_input helper was written to the input's branch and recorded as a user input, bypassing the guard that keeps inputs; derived writes now stay on the branch they were calculated on. - dump_simulation reads each value's mark from the same branch and period as the value it dumps, so #552's branch dumps keep matching marks. - Carry-over checks candidates latest first and stops at the first input, instead of resolving the storing branch for every known period. - Storages drop the marks of deleted keys. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 2 + .../data_storage/on_disk_storage.py | 2 + policyengine_core/holders/holder.py | 7 +- policyengine_core/simulations/simulation.py | 43 ++++--- policyengine_core/tools/simulation_dumper.py | 23 ++-- tests/core/test_carry_over_order.py | 118 ++++++++++++++++++ tests/fixtures/carry_over.py | 32 +++++ 7 files changed, 196 insertions(+), 31 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 61c3b4c3..0a7ed4b4 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -168,6 +168,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: if not period_item.startswith(branch_prefix) } self._shared.intersection_update(self._arrays) + self._derived.intersection_update(self._arrays) return if self.is_eternal: @@ -186,6 +187,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: ) } self._shared.intersection_update(self._arrays) + self._derived.intersection_update(self._arrays) def get_known_periods(self) -> list: # Split on the first colon only: an anchored period's string form diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 1f835f1c..ed9d41b7 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -118,6 +118,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._files.items() if not period_item.startswith(branch_prefix) } + self._derived.intersection_update(self._files) return if self.is_eternal: @@ -130,6 +131,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._files.items() if not period_item == f"{branch_name}_{period}" } + self._derived.intersection_update(self._files) def get_known_periods(self) -> list: return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 12a679c8..872b7ed3 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -350,7 +350,12 @@ def _set( derived: bool = False, ) -> None: simulation = getattr(self, "simulation", None) - user_input_contexts = getattr(simulation, "_user_input_contexts", None) + # A value calculated while an input is being set (say, by a + # ``set_input`` helper that calculates) is not part of that input: it + # belongs to the branch it was calculated on. + user_input_contexts = ( + None if derived else getattr(simulation, "_user_input_contexts", None) + ) if user_input_contexts and branch_name == "default": branch_name = user_input_contexts[-1] value = self._to_array(value, validate_nan=validate_nan) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 7f49e0ea..a7ab065d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -963,26 +963,37 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # a value already masked by ``defined_for``, or given by a # formula that has since ended, would carry forward. # A later input does not carry backwards. - carry_over_periods = [ - known_period - for known_period in known_periods - if known_period.start <= period.start - and not holder.is_derived(known_period, self.branch_name) - ] - own_unit_periods = [ - known_period - for known_period in carry_over_periods - if known_period.unit == variable.definition_period - ] - if carry_over_periods: - last_known_period = max( - own_unit_periods or carry_over_periods, - key=lambda p: (p.start, p.stop), - ) + candidates = sorted( + dict.fromkeys( + known_period + for known_period in known_periods + if known_period.start <= period.start + ), + key=lambda p: ( + p.unit == variable.definition_period, + p.start, + p.stop, + ), + reverse=True, + ) + last_known_period = next( + ( + candidate + for candidate in candidates + if not holder.is_derived(candidate, self.branch_name) + ), + None, + ) + if last_known_period is not None: # Pass branch_name through so auto-carry-over respects # the active branch instead of reaching for the # "default" branch's cache (bug H2). array = holder.get_array(last_known_period, self.branch_name) + elif variable.uprating is not None: + # Not cached: the uprating path above would take a + # cached default as the value to uprate later periods + # from. + return holder.default_array() else: array = holder.default_array() else: diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index e2376e1e..65b3c664 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -70,23 +70,18 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + derived_periods = set() for period in holder.get_known_periods(): value = holder.get_array(period) disk_storage.put(value, period) - _dump_derived_periods(holder, disk_storage.storage_dir) - - -def _dump_derived_periods(holder, storage_dir, branch_name="default"): - derived = sorted( - { - str(period) - for period in holder.get_known_periods() - if holder.is_derived(period, branch_name) - } - ) - if derived: - with open(os.path.join(storage_dir, DERIVED_PERIODS_FILE), "w") as file: - file.write("\n".join(derived) + "\n") + # Read the mark of exactly the value dumped: the same period on the + # same branch as ``get_array``. + if holder.is_derived(period): + derived_periods.add(str(period)) + if derived_periods: + path = os.path.join(disk_storage.storage_dir, DERIVED_PERIODS_FILE) + with open(path, "w") as file: + file.write("\n".join(sorted(derived_periods)) + "\n") def _dump_entity(population, directory): diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 7e76a414..aa481582 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -33,6 +33,8 @@ import pytest from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.experimental import MemoryConfig from policyengine_core.reforms import Reform from policyengine_core.tools.simulation_dumper import ( dump_simulation, @@ -103,6 +105,21 @@ def test_inputs_starting_together_carry_the_one_ending_last(system, first, secon ) +@pytest.mark.parametrize( + "first,second", [("2013", "year:2013:2"), ("year:2013:2", "2013")] +) +def test_inputs_at_one_unit_starting_together_carry_the_one_ending_last( + system, first, second +): + values = {"2013": [1, 1], "year:2013:2": [2, 2]} + inputs = { + "year_input_without_helper": {first: values[first], second: values[second]} + } + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2015"), [2, 2] + ) + + def test_later_input_does_not_carry_backwards(system): built = simulation(system, {"carried": {"2012": [7, 8], "2014": [9, 10]}}) np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) @@ -359,6 +376,107 @@ def test_carry_over_in_a_branch_copies_only_the_array_it_reads(system): assert len(storage._shared) == shared - 1 +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_storages_keep_the_derived_mark_with_the_value(on_disk, tmp_path): + storage = ( + OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + if on_disk + else InMemoryStorage(is_eternal=False) + ) + year = periods.period("2013") + value = np.array([1.0, 2.0]) + assert not storage.has(year) and not storage.is_derived(year) + storage.put(value, year, derived=True) + assert storage.has(year) and storage.is_derived(year) + assert not storage.has(year, "reform") and not storage.is_derived(year, "reform") + clone = storage.clone() + assert clone.is_derived(year) + storage.put(value, year) + assert storage.has(year) and not storage.is_derived(year) + assert clone.is_derived(year) + storage.put(value, year, derived=True) + storage.delete(year) + assert not storage.has(year) and not storage.is_derived(year) + assert not storage._derived + storage.put(value, year) + assert not storage.is_derived(year) + + +def test_carry_over_reads_derived_marks_from_disk_storage(system): + built = simulation(system, {}) + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + built.set_input("carried", "2012", np.array([7.0, 8.0])) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) + assert holder._disk_storage.is_derived(periods.period("2014")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + holder.delete_arrays(periods.period("2014")) + built.set_input("carried", "2014", np.array([9.0, 10.0])) + assert not holder.is_derived(periods.period("2014")) + np.testing.assert_array_equal(built.calculate("carried", "2015"), [9, 10]) + + +def test_a_default_cached_before_an_input_does_not_become_an_uprating_base(system): + """With no input at or before 2012 the default is not cached for an + uprated variable, so 2014 still carries the 2013-01 input.""" + inputs = {"uprated_input_without_helper": {"2013-01": [10, 20]}} + expected = alone(system, inputs, "uprated_input_without_helper", "2014") + np.testing.assert_array_equal(expected, [10, 20]) + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("uprated_input_without_helper", "2012"), [0, 0] + ) + np.testing.assert_array_equal( + built.calculate("uprated_input_without_helper", "2014"), expected + ) + + +def test_a_value_calculated_while_an_input_is_set_stays_on_its_own_branch(system): + """A ``set_input`` helper that calculates writes the calculated value to + the simulation's own branch, as a derived value, not to the input's.""" + built = simulation(system, {"carried": {"2012": [1, 2]}}) + holder = built.get_holder("formula_until_2013") + holder.set_input(periods.period("2012"), np.array([20.0, 20.0]), "other") + built.get_holder("input_with_calculating_helper").set_input( + periods.period("2013-06"), np.array([1.0, 1.0]), "other" + ) + np.testing.assert_array_equal( + holder.get_array(periods.period("2012"), "other"), [20, 20] + ) + np.testing.assert_array_equal( + holder.get_array(periods.period("2012"), "default"), [6, 7] + ) + assert holder.is_derived(periods.period("2012"), "default") + assert not holder.is_derived(periods.period("2012"), "other") + assert ("formula_until_2013", "default", periods.period("2012")) not in ( + built._user_input_keys + ) + other = built.get_branch("other") + np.testing.assert_array_equal( + other.calculate("formula_until_2013", "2014"), [20, 20] + ) + + +def test_carry_over_checks_only_the_latest_candidates(system, monkeypatch): + from policyengine_core.holders import Holder + + years = {str(year): [year, year] for year in range(1900, 2010)} + built = simulation(system, {"carried": years}) + branch = built.get_branch("a").get_branch("b").get_branch("c") + calls = [] + is_derived = Holder.is_derived + + def counting(self, period, branch_name="default"): + calls.append(period) + return is_derived(self, period, branch_name) + + monkeypatch.setattr(Holder, "is_derived", counting) + np.testing.assert_array_equal(branch.calculate("carried", "2020"), [2009, 2009]) + assert len(calls) == 1 + + def test_input_starting_with_the_requested_period_carries(system): """An input in another unit that starts when the period does carries.""" built = simulation(system, {"year_input_without_helper": {"2013-01": [3, 4]}}) diff --git a/tests/fixtures/carry_over.py b/tests/fixtures/carry_over.py index ad535e3c..28ece23c 100644 --- a/tests/fixtures/carry_over.py +++ b/tests/fixtures/carry_over.py @@ -17,6 +17,7 @@ from policyengine_core import periods from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.parameters import ParameterNode from policyengine_core.simulations import SimulationBuilder from policyengine_core.variables import Variable @@ -87,6 +88,29 @@ class month_input_without_helper(Variable): label = "Monthly input with no set_input helper, stored at any period" +class uprated_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = None + uprating = "carry_over_test.index" + label = "Yearly uprated input with no set_input helper" + + +def _calculate_while_setting(holder, period, array): + """A ``set_input`` helper that calculates another variable first.""" + holder.simulation.calculate("formula_until_2013", "2012") + holder._set(period.this_year, array) + + +class input_with_calculating_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = _calculate_while_setting + label = "Yearly input whose set_input helper calculates" + + VARIABLES = ( carried, carried_count, @@ -96,12 +120,20 @@ class month_input_without_helper(Variable): carried_monthly, year_input_without_helper, month_input_without_helper, + uprated_input_without_helper, + input_with_calculating_helper, ) def build_system() -> CountryTaxBenefitSystem: system = CountryTaxBenefitSystem() system.auto_carry_over_input_variables = True + system.parameters.add_child( + "carry_over_test", + ParameterNode( + "carry_over_test", data={"index": {"values": {"2010-01-01": 1.0}}} + ), + ) system.add_variables(*VARIABLES) return system From 51318dd88302a876cd7aaf8a71e95cf5f687f35b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 03:02:54 -0400 Subject: [PATCH 3/8] Test that deleting a whole branch drops its derived marks Co-Authored-By: Claude Opus 5.5 --- tests/core/test_carry_over_order.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index aa481582..f3a398ab 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -398,6 +398,12 @@ def test_storages_keep_the_derived_mark_with_the_value(on_disk, tmp_path): storage.delete(year) assert not storage.has(year) and not storage.is_derived(year) assert not storage._derived + storage.put(value, year, derived=True) + storage.put(value, year, "reform", derived=True) + storage.delete(branch_name="reform") + assert storage.is_derived(year) and not storage.has(year, "reform") + storage.delete() + assert not storage._derived storage.put(value, year) assert not storage.is_derived(year) From 27d1bc353ad6a73215f0766eec215ac8b98a3dce Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 03:04:59 -0400 Subject: [PATCH 4/8] Cache the default only where master cached a carried value With no input at or before the period but a later period stored, master returned the default without caching it; caching it changed what formulas that test whether a value is stored see (policyengine-uk's maintenance loan and current_education formulas check get_array(period.last_year)), and gave the uprating path a calculated base. Return it uncached there, as master did; cache it (derived) otherwise, as master cached the value it carried. This replaces the uprating-only special case. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 13 +++++++++---- tests/core/test_carry_over_order.py | 13 ++++++++++--- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index a7ab065d..58190aab 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -989,10 +989,15 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # the active branch instead of reaching for the # "default" branch's cache (bug H2). array = holder.get_array(last_known_period, self.branch_name) - elif variable.uprating is not None: - # Not cached: the uprating path above would take a - # cached default as the value to uprate later periods - # from. + elif any( + known_period.start > period.start + for known_period in known_periods + ): + # No input to carry, but a later period is stored: as + # before, return the default without caching it. A + # cached default would change what a formula testing + # whether a value is stored sees, and the uprating path + # above would uprate later periods from it. return holder.default_array() else: array = holder.default_array() diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index f3a398ab..481932bc 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -175,14 +175,21 @@ def test_input_set_after_a_period_defaulted_still_carries(system): np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) -def test_periods_with_no_earlier_input_are_cached(system): - built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) +def test_defaults_are_cached_as_before(system): + """With nothing stored, or only earlier calculated values, the default is + cached (as master cached the value it carried); before a later stored + period it is not (as master returned it).""" + built = simulation(system, {}) + holder = built.get_holder("carried_monthly") for month in range(1, 13): for _ in range(2): built.calculate("carried_monthly", f"2012-{month:02d}") - holder = built.get_holder("carried_monthly") for month in range(1, 13): assert holder.get_array(periods.period(f"2012-{month:02d}")) is not None + built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) + holder = built.get_holder("carried_monthly") + np.testing.assert_array_equal(built.calculate("carried_monthly", "2012-03"), [0, 0]) + assert holder.get_array(periods.period("2012-03")) is None def test_branch_input_is_not_hidden_by_a_value_the_parent_carried(system): From 35e20d9fff4b34212b57f4551b1e3e3e5de28247 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 03:58:49 -0400 Subject: [PATCH 5/8] Read input periods in one pass over readable keys; clear marks with the caches Round-3 review findings: - A period stored only under a branch this one cannot read was taken for an input (own-unit preference ranked it first) and read back as NaN. Holder.get_input_periods(branch_name) now lists, in one pass over the stored keys, the periods whose readable value is an input, taking for each period the key get_array reads first; carry-over picks from those. This also replaces the per-period walk up the branch chain. - apply_reform's cache wipe kept the storages' marks, and rebuilding a disk index brought a stale mark back onto a new input; the wipe now clears the marks and OnDiskStorage.restore starts without any. - Storages pickled before the marks existed failed to clone, put or delete; __setstate__ now defaults them (and #556's shared-key set). - The public carry-over rule now states the own-unit preference and that an input for the period itself is read back as stored. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 7 ++ .../data_storage/on_disk_storage.py | 8 ++ policyengine_core/holders/holder.py | 65 ++++++++--- policyengine_core/simulations/simulation.py | 24 ++-- .../taxbenefitsystems/tax_benefit_system.py | 5 +- tests/core/test_carry_over_order.py | 105 ++++++++++++++++-- 6 files changed, 172 insertions(+), 42 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 0a7ed4b4..9feffb59 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -53,6 +53,13 @@ def __init__(self, is_eternal: bool): self._derived = set() self.is_eternal = is_eternal + def __setstate__(self, state: dict) -> None: + # A storage pickled before derived marks or shared arrays existed has + # neither: its values count as inputs and as its own. + state.setdefault("_derived", set()) + state.setdefault("_shared", set()) + self.__dict__.update(state) + def clone(self, share_arrays: bool = False) -> "InMemoryStorage": """Copy this storage. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index ed9d41b7..db3e2cad 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -29,6 +29,12 @@ def __init__( self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + def __setstate__(self, state: dict) -> None: + # A storage pickled before derived marks existed has none: its values + # count as inputs. + state.setdefault("_derived", set()) + self.__dict__.update(state) + def clone(self) -> "OnDiskStorage": """Create a private metadata view over this storage directory. @@ -144,6 +150,8 @@ def get_known_branch_periods(self) -> list: def restore(self) -> None: self._files = files = {} + # Files read back from a directory carry no derived marks. + self._derived = set() # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 872b7ed3..a0220889 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -435,22 +435,29 @@ def _stores(self, period: Period, branch_name: str) -> bool: and self._disk_storage.has(period, branch_name) ) + def _readable_branches(self, branch_name: str = "default") -> List[str]: + """``get_array``'s lookup order: the branch, its ``parent_branch`` + ancestors, then ``default``.""" + names = [branch_name] + if branch_name != "default": + parent = ( + getattr(self.simulation, "parent_branch", None) + if self.simulation + else None + ) + while parent is not None: + names.append(parent.branch_name) + parent = getattr(parent, "parent_branch", None) + names.append("default") + return list(dict.fromkeys(names)) + def _branch_storing(self, period: Period, branch_name: str = "default") -> str: """The branch whose stored value ``get_array(period, branch_name)`` - reads, in ``get_array``'s lookup order, or ``None`` if none stores one. - """ - if self._stores(period, branch_name): - return branch_name - if branch_name == "default": - return None - parent = ( - getattr(self.simulation, "parent_branch", None) if self.simulation else None - ) - while parent is not None: - if self._stores(period, parent.branch_name): - return parent.branch_name - parent = getattr(parent, "parent_branch", None) - return "default" if self._stores(period, "default") else None + reads, or ``None`` if none stores one.""" + for name in self._readable_branches(branch_name): + if self._stores(period, name): + return name + return None def is_derived(self, period: Period, branch_name: str = "default") -> bool: """Whether the value ``get_array(period, branch_name)`` reads was @@ -466,3 +473,33 @@ def is_derived(self, period: Period, branch_name: str = "default") -> bool: if self._memory_storage.has(period, storing): return self._memory_storage.is_derived(period, storing) return self._disk_storage.is_derived(period, storing) + + def get_input_periods(self, branch_name: str = "default") -> List[Period]: + """The periods for which the value ``get_array(period, branch_name)`` + reads is an input rather than a value the simulation calculated (see + ``put_in_cache``). Periods stored only under branches this one cannot + read are left out. + + One pass over the stored keys: for each period, the key ``get_array`` + reads first (the branch before its ancestors, memory before disk). + """ + rank = { + name: index + for index, name in enumerate(self._readable_branches(branch_name)) + } + storages = [self._memory_storage] + if self._disk_storage is not None: + storages.append(self._disk_storage) + read = {} + for order, storage in enumerate(storages): + for stored_branch, period in storage.get_known_branch_periods(): + if stored_branch not in rank: + continue + key = (rank[stored_branch], order) + if period not in read or key < read[period][0]: + read[period] = (key, storage, stored_branch) + return [ + period + for period, (_, storage, stored_branch) in read.items() + if not storage.is_derived(period, stored_branch) + ] diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 58190aab..ed818655 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -394,8 +394,10 @@ def _invalidate_all_caches(self) -> None: for population in self.populations.values(): for holder in population._holders.values(): holder._memory_storage._arrays = {} + holder._memory_storage._derived = set() if holder._disk_storage is not None: holder._disk_storage._files = {} + holder._disk_storage._derived = set() # Replay preserved user inputs so ``calculate`` still sees them. for user_input in preserved: holder = self.get_holder(user_input.variable_name) @@ -963,26 +965,20 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # a value already masked by ``defined_for``, or given by a # formula that has since ended, would carry forward. # A later input does not carry backwards. - candidates = sorted( - dict.fromkeys( - known_period - for known_period in known_periods - if known_period.start <= period.start + last_known_period = max( + ( + input_period + for input_period in holder.get_input_periods( + self.branch_name + ) + if input_period.start <= period.start ), key=lambda p: ( p.unit == variable.definition_period, p.start, p.stop, ), - reverse=True, - ) - last_known_period = next( - ( - candidate - for candidate in candidates - if not holder.is_derived(candidate, self.branch_name) - ), - None, + default=None, ) if last_known_period is not None: # Pass branch_name through so auto-carry-over respects diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 1fae168f..b7073081 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -91,8 +91,9 @@ class TaxBenefitSystem: ``uprating`` path from an earlier period) takes an input from another period. If so, the period takes the input stored for the latest-starting period that starts no later than it (on a tie, the one that ends last), - masked by the variable's ``defined_for``; with no such input, the - default. An input never carries backwards, and values the simulation + preferring inputs at the variable's own definition-period unit, masked by + the variable's ``defined_for``; with no such input, the default. An input + stored for the period itself is read back as stored. An input never carries backwards, and values the simulation calculated never carry (see ``Holder.is_derived``), so the result does not depend on which periods were calculated first.""" basic_inputs: List[str] = None diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 481932bc..42c23808 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -472,22 +472,103 @@ def test_a_value_calculated_while_an_input_is_set_stays_on_its_own_branch(system ) -def test_carry_over_checks_only_the_latest_candidates(system, monkeypatch): +def test_carry_over_reads_provenance_in_one_pass(system, monkeypatch): + """With many calculated periods and deep branches, a decision makes one + pass over the stored keys instead of walking the branches per period.""" from policyengine_core.holders import Holder - years = {str(year): [year, year] for year in range(1900, 2010)} - built = simulation(system, {"carried": years}) - branch = built.get_branch("a").get_branch("b").get_branch("c") - calls = [] - is_derived = Holder.is_derived + built = simulation(system, {"carried": {"1900": [1, 1]}}) + for year in range(1901, 2010): + built.calculate("carried", str(year)) + branch = built + for depth in range(20): + branch = branch.get_branch(f"b{depth}") + calls = {"stores": 0, "passes": 0} + stores, passes = Holder._stores, Holder.get_input_periods + + def counting_stores(self, *args, **kwargs): + calls["stores"] += 1 + return stores(self, *args, **kwargs) + + def counting_passes(self, *args, **kwargs): + calls["passes"] += 1 + return passes(self, *args, **kwargs) + + monkeypatch.setattr(Holder, "_stores", counting_stores) + monkeypatch.setattr(Holder, "get_input_periods", counting_passes) + np.testing.assert_array_equal(branch.calculate("carried", "2020"), [1, 1]) + # One pass for the decision; caching the result checks the period on each + # branch in the chain once (21), not every stored period on each (2,310). + assert calls["passes"] == 1 + assert calls["stores"] <= 21 + + +def test_a_period_stored_only_on_another_branch_does_not_carry(system): + """A branch never takes a period it cannot read for an input.""" + built = simulation(system, {"year_input_without_helper": {"2013-01": [20, 20]}}) + built.get_holder("year_input_without_helper").set_input( + periods.period("2012"), np.array([10.0, 10.0]), "other" + ) + other = built.get_branch("other") + np.testing.assert_array_equal( + built.calculate("year_input_without_helper", "2014"), [20, 20] + ) + np.testing.assert_array_equal( + other.calculate("year_input_without_helper", "2014"), [10, 10] + ) - def counting(self, period, branch_name="default"): - calls.append(period) - return is_derived(self, period, branch_name) - monkeypatch.setattr(Holder, "is_derived", counting) - np.testing.assert_array_equal(branch.calculate("carried", "2020"), [2009, 2009]) - assert len(calls) == 1 +def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): + built = simulation(system, {}) + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + for year in range(2010, 2020): + built.calculate("carried", str(year)) + built.apply_reform(_noop) + assert not holder._memory_storage._derived + assert not holder._disk_storage._derived + # A file written for a period whose value was derived, then read back + # by rebuilding the index, is an input. + built = simulation(system, {}) + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + writer = OnDiskStorage(holder._disk_storage.storage_dir, preserve_storage_dir=True) + built.calculate("carried", "2012") + assert holder.is_derived(periods.period("2012")) + built.apply_reform(_noop) + writer.put(np.array([10.0, 10.0]), periods.period("2012")) + holder._disk_storage.restore() + assert not holder.is_derived(periods.period("2012")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [10, 10]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_storages_pickled_without_marks_still_work(on_disk, tmp_path): + import pickle + + storage = ( + OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + if on_disk + else InMemoryStorage(is_eternal=False) + ) + year = periods.period("2012") + storage.put(np.array([10.0]), year) + state = dict(storage.__dict__) + del state["_derived"] + state.pop("_shared", None) + old = type(storage).__new__(type(storage)) + old.__setstate__(state) + restored = pickle.loads(pickle.dumps(old)) + assert restored.has(year) and not restored.is_derived(year) + restored.clone() + restored.put(np.array([20.0]), periods.period("2013"), derived=True) + assert restored.is_derived(periods.period("2013")) + restored.delete(periods.period("2013")) + np.testing.assert_array_equal(restored.get(year), [10.0]) def test_input_starting_with_the_requested_period_carries(system): From 1d3cb6a27b0dc531ad55c8eba710ade3ab60a34e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 04:04:54 -0400 Subject: [PATCH 6/8] Test memory-over-disk precedence, in-memory mark clearing and index rebuilds Closes three mutants the suite did not kill. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_carry_over_order.py | 38 +++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 42c23808..04801aca 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -546,6 +546,44 @@ def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): np.testing.assert_array_equal(built.calculate("carried", "2013"), [10, 10]) +def test_marks_are_cleared_with_in_memory_values_apply_reform_wipes(system): + built = simulation(system, {}) + holder = built.get_holder("carried") + for year in range(2010, 2020): + built.calculate("carried", str(year)) + built.apply_reform(_noop) + assert not holder._memory_storage._derived + + +def test_rebuilding_a_disk_index_reads_files_as_inputs(system): + """A derived file replaced by an input, then re-indexed, is an input.""" + built = simulation(system, {}) + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + built.calculate("carried", "2012") + assert holder.is_derived(periods.period("2012")) + writer = OnDiskStorage(holder._disk_storage.storage_dir, preserve_storage_dir=True) + writer.put(np.array([10.0, 10.0]), periods.period("2012")) + holder._disk_storage.restore() + assert not holder.is_derived(periods.period("2012")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [10, 10]) + + +def test_a_value_in_memory_takes_precedence_over_one_on_disk(system): + """``get_array`` reads memory before disk for the same key, and so does + the input test.""" + built = simulation(system, {}) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + year = periods.period("2012") + holder._disk_storage.put(np.array([0.0, 0.0]), year, derived=True) + holder._memory_storage.put(np.array([7.0, 8.0]), year) + assert year in holder.get_input_periods() + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + + @pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) def test_storages_pickled_without_marks_still_work(on_disk, tmp_path): import pickle From 912a4ac55efb6aab2f9c56dccf4dc6359c2c3c98 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 04:40:46 -0400 Subject: [PATCH 7/8] Break equal-extent carry-over ties by unit, then period Two inputs covering the same days in units other than the variable's (month:2013-01:2 and day:2013-01-01:59) resolved to whichever was stored first. Prefer the larger unit, then the period's string form, so the choice depends only on the inputs (round-4 review). Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 6 +++++- tests/core/test_carry_over_order.py | 12 ++++++++++++ tests/fixtures/carry_over.py | 11 +++++++++-- 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ed818655..ce22f3a0 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -946,7 +946,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: ): # Carry over the latest input: of the stored periods that # start no later than ``period``, the one that starts last - # (on a tie, the one that ends last), preferring periods at + # (on a tie, the one that ends last, then the larger unit, + # so the choice never depends on the order inputs were + # stored), preferring periods at # the variable's own definition-period unit and using # another unit only when there is none, as for an input to # a variable with no ``set_input`` helper. Compare @@ -977,6 +979,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: p.unit == variable.definition_period, p.start, p.stop, + periods.unit_weight(p.unit), + str(p), ), default=None, ) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 04801aca..9c077364 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -120,6 +120,18 @@ def test_inputs_at_one_unit_starting_together_carry_the_one_ending_last( ) +@pytest.mark.parametrize("reverse", [False, True], ids=["month-first", "day-first"]) +def test_inputs_with_the_same_extent_carry_the_larger_unit(system, reverse): + """Two inputs covering the same days in other units than the variable's + resolve the same way whichever was stored first.""" + values = {"month:2013-01:2": [10, 10], "day:2013-01-01:59": [20, 20]} + order = list(values)[::-1] if reverse else list(values) + inputs = {"year_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2014"), [10, 10] + ) + + def test_later_input_does_not_carry_backwards(system): built = simulation(system, {"carried": {"2012": [7, 8], "2014": [9, 10]}}) np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) diff --git a/tests/fixtures/carry_over.py b/tests/fixtures/carry_over.py index 28ece23c..83f7e3f3 100644 --- a/tests/fixtures/carry_over.py +++ b/tests/fixtures/carry_over.py @@ -6,7 +6,8 @@ The reference rule, for a variable with no formula result for a period: an input stored for the period itself is read back as stored; otherwise the period takes the input stored for the latest-starting period that starts no -later than it (on a tie, the one that ends last), preferring inputs at the +later than it (on a tie, the one that ends last, then the larger unit), +preferring inputs at the variable's own definition-period unit, masked by the variable's ``defined_for`` at the period; with no such input, the default. """ @@ -193,7 +194,13 @@ def reference(system, inputs, variable, period): ] if earlier: latest = max( - own_unit or earlier, key=lambda stored: (stored.start, stored.stop) + own_unit or earlier, + key=lambda stored: ( + stored.start, + stored.stop, + periods.unit_weight(stored.unit), + str(stored), + ), ) value = np.array(stored_inputs[latest]).astype(default.dtype) else: From 57c0d36b7aefb2426ea477758cd82a89b7d4eb5b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:15:36 -0400 Subject: [PATCH 8/8] Uprate only from inputs, never from calculated values Simulation._calculate uprated a variable from its latest earlier stored period in its own unit, including periods the simulation had calculated itself, so the result depended on calculation order. Executed on fix-carry-over-order (index +3.7%/yr, input at 2012, 2015 asked alone vs after 2013 and 2014): - int: [1116, 85] alone, [1115, 83] after (truncation compounds) - float: float32 rounding compounds - defined_for false for one person in 2013 only: [1115.16, 1115.16] alone, [0.0, 1115.16] after (the 2013 mask carries into 2015) Also: a cached default, a monthly input carried into a calculated year, and a reform formula's result before its end each became an uprating source. For variables with uprating, keep only the earlier same-unit periods in Holder.get_input_periods(branch_name) (the carry-over PR's branch-aware provenance): periods whose value this branch reads is an input. That also drops periods stored only under a branch this one cannot read, which read back as None and raised TypeError. With no earlier input, the carry-over path or the default applies, as when the period is asked alone. Tests: tests/core/test_uprating_order.py (regressions, no hypothesis), tests/core/test_uprating_order_property.py (derandomized Hypothesis: any requests, branches and branch inputs, carry-over on/off == alone, byte for byte, and == reference rule), shared fixtures in tests/fixtures/uprating_order.py. #551's path-independence test is now exact instead of rel=1e-5. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-uprating-order.fixed.md | 1 + policyengine_core/simulations/simulation.py | 27 +- tests/core/test_uprating_order.py | 285 +++++++++++++++++ tests/core/test_uprating_order_property.py | 299 ++++++++++++++++++ .../core/variables/test_variable_uprating.py | 11 +- tests/fixtures/uprating_order.py | 215 +++++++++++++ 6 files changed, 828 insertions(+), 10 deletions(-) create mode 100644 changelog.d/fix-uprating-order.fixed.md create mode 100644 tests/core/test_uprating_order.py create mode 100644 tests/core/test_uprating_order_property.py create mode 100644 tests/fixtures/uprating_order.py diff --git a/changelog.d/fix-uprating-order.fixed.md b/changelog.d/fix-uprating-order.fixed.md new file mode 100644 index 00000000..ae873047 --- /dev/null +++ b/changelog.d/fix-uprating-order.fixed.md @@ -0,0 +1 @@ +Uprating now starts from the latest earlier input in the variable's own unit and skips periods the simulation calculated, so calculating intermediate periods first no longer compounds rounding or truncation, or carries an eligibility mask or a default into later uprated values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ce22f3a0..e75864ce 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -903,13 +903,31 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if known_period.unit == variable.definition_period and known_period.start < period.start ] - if variable.uprating is not None and len(earlier_known_periods) > 0: + # Uprate only from an input this branch reads (see + # ``Holder.get_input_periods``). Every value this simulation + # calculated is marked derived when cached, and uprating from + # one would make the result depend on which periods were + # calculated first: an integer truncated, or a float32 + # rounded, at an intermediate period would compound, and a + # value masked by ``defined_for``, a default, or a value + # carried from another unit, cached there, would replace the + # input. A period stored only under a branch this one cannot + # read would read back as ``None``. + earlier_input_periods = [] + if variable.uprating is not None: + input_periods = set(holder.get_input_periods(self.branch_name)) + earlier_input_periods = [ + known_period + for known_period in earlier_known_periods + if known_period in input_periods + ] + if earlier_input_periods: # Take the latest period from the filtered list itself. # Indexing ``known_periods`` with a position in the # filtered list picked the wrong period whenever a later # one was stored first. latest_known_period = max( - earlier_known_periods, key=lambda p: p.start + earlier_input_periods, key=lambda p: p.start ) try: uprating_parameter = get_parameter( @@ -996,8 +1014,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # No input to carry, but a later period is stored: as # before, return the default without caching it. A # cached default would change what a formula testing - # whether a value is stored sees, and the uprating path - # above would uprate later periods from it. + # whether a value is stored sees. (The uprating path + # above skips derived periods, so it would not uprate + # from one.) return holder.default_array() else: array = holder.default_array() diff --git a/tests/core/test_uprating_order.py b/tests/core/test_uprating_order.py new file mode 100644 index 00000000..d67ffebe --- /dev/null +++ b/tests/core/test_uprating_order.py @@ -0,0 +1,285 @@ +"""Uprating depends only on the inputs, not on what was calculated first. + +A variable with ``uprating`` and no formula result for a period takes the +latest earlier value stored in its own unit and multiplies it by the ratio of +the uprating index at the two period starts. Until this fix, +``Simulation._calculate`` took that value from any stored period, including +one the simulation had itself calculated, so calculating an intermediate +period first changed the result: + +* An integer value truncated at each intermediate year compounded: an input of + 1001 for 2012 gave 1116 for 2015 asked alone, and 1115 after 2013 and 2014. +* A float32 value rounded at each intermediate year compounded the rounding. +* A value masked by ``defined_for``, or a default cached at an intermediate + year, replaced the input from then on: someone ineligible in 2013 only got + 0 for 2015 after 2013 was calculated, and the input uprated otherwise. +* A default, an input stored in another unit and carried into an + intermediate year, or the result of a formula a reform gave the variable up + to an ``end`` became a value to uprate, where the period asked alone gave + the default, the carried input or the uprated input. + +Uprating now starts only from an input this branch reads +(``Holder.get_input_periods``): values the simulation calculated are marked +derived when cached and skipped, as auto-carry-over skips them, and so are +periods stored only under branches this one cannot read. The rule is in +``tests/fixtures/uprating_order.py`` +and the order-independence property in ``test_uprating_order_property.py``. +Each regression here compares with a fresh simulation that calculates only +the period in question, byte for byte. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import entities +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable +from tests.fixtures.uprating_order import ( + UPRATING, + alone, + assert_bitwise_equal, + build_system, + index, + simulation, +) + +SYSTEMS = {True: build_system(True), False: build_system(False)} + + +@pytest.fixture(params=[True, False], ids=["carry_over", "no_carry_over"]) +def system(request): + return SYSTEMS[request.param] + + +def _after(system, inputs, requests, variable, period): + built = simulation(system, inputs) + for requested_variable, requested_period in requests: + built.calculate(requested_variable, requested_period) + return built.calculate(variable, period) + + +def test_integer_truncation_does_not_compound(system): + inputs = {"uprated_count": {"2012": [1001, 77]}} + requests = [("uprated_count", "2013"), ("uprated_count", "2014")] + result = _after(system, inputs, requests, "uprated_count", "2015") + # Before the fix: [1115, 83], truncated at 2013 and again at 2014. + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +def test_float32_rounding_does_not_compound(system): + inputs = {"uprated": {"2012": [1001.3, 77.7]}} + requests = [("uprated", "2013"), ("uprated", "2014")] + result = _after(system, inputs, requests, "uprated", "2015") + assert_bitwise_equal(result, alone(system, inputs, "uprated", "2015")) + np.testing.assert_allclose( + result, np.array([1001.3, 77.7]) * index(2015) / index(2012), rtol=1e-6 + ) + + +def test_value_masked_by_defined_for_is_not_uprated(system): + inputs = { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": { + "2012": [True, True], + "2013": [False, True], + "2014": [True, True], + "2015": [True, True], + }, + } + requests = [("uprated_if_eligible", "2013"), ("uprated_if_eligible", "2014")] + result = _after(system, inputs, requests, "uprated_if_eligible", "2015") + # Before the fix: [0, 1115.16]; the 2013 mask carried into 2015. + np.testing.assert_allclose( + result, [1000 * index(2015) / index(2012)] * 2, rtol=1e-6 + ) + assert_bitwise_equal(result, alone(system, inputs, "uprated_if_eligible", "2015")) + + +def test_default_cached_where_defined_for_is_false_is_not_uprated(system): + inputs = { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": { + "2012": [True, True], + "2013": [False, False], + "2014": [True, True], + }, + } + requests = [("uprated_if_eligible", "2013")] + result = _after(system, inputs, requests, "uprated_if_eligible", "2014") + # Before the fix: [0, 0], the default cached for 2013. + np.testing.assert_allclose( + result, [1000 * index(2014) / index(2012)] * 2, rtol=1e-6 + ) + assert_bitwise_equal(result, alone(system, inputs, "uprated_if_eligible", "2014")) + + +def test_cached_default_value_is_not_uprated(system): + requests = [("uprated_with_default", "2013")] + result = _after(system, {}, requests, "uprated_with_default", "2015") + # Before the fix: 10 * index(2015) / index(2013). + np.testing.assert_array_equal(result, [10, 10]) + assert_bitwise_equal(result, alone(system, {}, "uprated_with_default", "2015")) + + +def test_input_in_another_unit_is_not_uprated_through_a_calculated_year(system): + inputs = {"uprated_any_unit": {"2012-03": [10, 20]}} + requests = [("uprated_any_unit", "2013")] + result = _after(system, inputs, requests, "uprated_any_unit", "2015") + # Asked alone, 2015 has no earlier yearly input to uprate: auto-carry-over + # carries the monthly input unchanged, or it gets the default. Before the + # fix, calculating 2013 first cached that value at 2013, and 2015 uprated + # it. + expected = [10, 20] if system.auto_carry_over_input_variables else [0, 0] + np.testing.assert_array_equal(result, expected) + assert_bitwise_equal(result, alone(system, inputs, "uprated_any_unit", "2015")) + + +def test_reform_formula_result_is_not_uprated_past_the_formula_end(system): + """Core rejects a variable declaring both a formula and ``uprating``, but + a reform can give an uprated variable a formula, which keeps the + inherited ``uprating``. Past the formula's ``end`` the variable is uprated + again, from its input, not from a formula result.""" + + class formula_until_2013(Reform): + def apply(self): + class uprated(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + end = "2013-12-31" + label = "Uprated yearly input given a formula up to 2013" + + def formula(person, period, parameters): + return person.filled_array(42) + + self.update_variable(uprated) + + reformed = formula_until_2013(system) + assert reformed.get_variable("uprated").uprating == UPRATING + inputs = {"uprated": {"2012": [1, 2]}} + result = _after(reformed, inputs, [("uprated", "2013")], "uprated", "2015") + # Before the fix: the 2013 formula result, 42, uprated to 2015. + np.testing.assert_allclose( + result, np.array([1, 2]) * index(2015) / index(2012), rtol=1e-6 + ) + assert_bitwise_equal(result, alone(reformed, inputs, "uprated", "2015")) + + +def test_monthly_rounding_does_not_compound(system): + inputs = {"uprated_monthly": {"2011-06": [1001.3, 77.7]}} + requests = [("uprated_monthly", f"2012-{month:02d}") for month in range(1, 13)] + requests += [("uprated_monthly", "2013-01")] + result = _after(system, inputs, requests, "uprated_monthly", "2014-03") + assert_bitwise_equal(result, alone(system, inputs, "uprated_monthly", "2014-03")) + np.testing.assert_allclose( + result, np.array([1001.3, 77.7]) * index(2014) / index(2011), rtol=1e-6 + ) + + +@pytest.mark.parametrize("nested", [False, True]) +def test_branch_forked_after_intermediate_years_uprates_the_input(system, nested): + inputs = {"uprated_count": {"2012": [1001, 77]}} + built = simulation(system, inputs) + built.calculate("uprated_count", "2013") + built.calculate("uprated_count", "2014") + branch = built.get_branch("reform") + if nested: + branch = branch.get_branch("nested") + result = branch.calculate("uprated_count", "2015") + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +def test_parent_calculation_does_not_hide_a_branch_input(): + """An input stored only under a branch keeps its provenance when the + parent calculates a value for the same period.""" + system = SYSTEMS[True] + built = simulation(system, {"uprated": {"2011": [1, 2]}}) + built.get_holder("uprated").set_input( + periods.period("2012"), np.array([300.0, 400.0]), "other" + ) + built.calculate("uprated", "2012") + branch = built.get_branch("other") + np.testing.assert_allclose( + branch.calculate("uprated", "2014"), + np.array([300, 400]) * index(2014) / index(2012), + rtol=1e-6, + ) + + +def test_period_stored_only_under_an_unreadable_branch_is_not_uprated(system): + """A holder can store a period under a branch this simulation cannot + read (here through ``Holder.set_input`` with another branch name). It is + not an input this simulation can uprate from.""" + inputs = {"uprated_count": {"2012": [1001, 77]}} + built = simulation(system, inputs) + built.get_holder("uprated_count").set_input( + periods.period("2013"), np.array([2001, 177]), "sibling" + ) + # Before the fix: TypeError, the sibling's 2013 read back as None here. + result = built.calculate("uprated_count", "2015") + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +# Guards against over-correction: inputs at intermediate periods still count. + + +def test_latest_earlier_input_is_uprated(system): + inputs = {"uprated": {"2012": [1, 2], "2014": [300, 400]}} + for requests in ( + [], + [("uprated", "2013")], + [("uprated", "2013"), ("uprated", "2016")], + ): + result = _after(system, inputs, requests, "uprated", "2015") + np.testing.assert_allclose( + result, np.array([300, 400]) * index(2015) / index(2014), rtol=1e-6 + ) + + +def test_input_set_over_a_calculated_period_is_uprated(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + built.set_input("uprated", "2013", np.array([300.0, 400.0])) + np.testing.assert_allclose( + built.calculate("uprated", "2015"), + np.array([300, 400]) * index(2015) / index(2013), + rtol=1e-6, + ) + + +def test_branch_input_over_a_period_the_parent_calculated_is_uprated(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + branch = built.get_branch("reform") + branch.set_input("uprated", "2013", np.array([300.0, 400.0])) + np.testing.assert_allclose( + branch.calculate("uprated", "2015"), + np.array([300, 400]) * index(2015) / index(2013), + rtol=1e-6, + ) + np.testing.assert_allclose( + built.calculate("uprated", "2015"), + np.array([1, 2]) * index(2015) / index(2012), + rtol=1e-6, + ) + + +def test_add_over_an_input_period_keeps_it_an_input(system): + built = simulation(system, {"uprated_count": {"2012": [1001, 77]}}) + built.calculate_add("uprated_count", "2012") + built.calculate("uprated_count", "2013") + np.testing.assert_array_equal(built.calculate("uprated_count", "2015"), [1116, 85]) + + +def test_uprated_value_is_cached_as_derived(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + holder = built.get_holder("uprated") + assert holder.is_derived(periods.period("2013")) + assert not holder.is_derived(periods.period("2012")) diff --git a/tests/core/test_uprating_order_property.py b/tests/core/test_uprating_order_property.py new file mode 100644 index 00000000..5f015cc3 --- /dev/null +++ b/tests/core/test_uprating_order_property.py @@ -0,0 +1,299 @@ +"""Properties of uprating: order independence and the reference rule. + +Over the domain generated here, what a simulation returns for an uprated +variable, after earlier requests, equals byte for byte (dtype, shape and +bytes, so ``-0.0`` differs from ``0.0``) what a fresh simulation given the +same inputs returns for that period alone. It also equals the reference rule +in ``tests/fixtures/uprating_order.py``. Checked with and without +auto-carry-over. Regressions are in ``test_uprating_order.py``. + +The domain: + +* Inputs are stored in each variable's own unit. ``uprated_any_unit`` has no + ``set_input`` helper, so its inputs are years or months. +* Earlier requests are plain calculations in either unit, ``ADD`` and + ``DIVIDE``. +* The value may be read from the simulation, or from a branch or nested + branch forked after those requests. The branch may set its own input for + the target variable. The earlier requests then avoid the target period of + the target variable: a value the parent calculated before the fork stays + cached for the branch whatever its inputs are, which is not what this + module tests. +* Targets are in the variable's definition unit. +""" + +from __future__ import annotations + +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 import periods # noqa: E402 +from tests.fixtures.uprating_order import ( # noqa: E402 + COUNT, + assert_bitwise_equal, + build_system, + reference, + request, + simulation, +) + +SYSTEMS = {True: build_system(True), False: build_system(False)} + +YEARS = [str(year) for year in range(2009, 2019)] +INPUT_YEARS = [str(year) for year in range(2010, 2017)] +MONTHS = [f"{year}-{month:02d}" for year in range(2011, 2016) for month in range(1, 13)] +INPUT_MONTHS = [month for month in MONTHS if "2011-06" <= month <= "2013-12"] +REQUEST_MONTHS = [month for month in MONTHS if "2011-11" <= month <= "2015-06"] + +YEAR_VARIABLES = [ + "uprated", + "uprated_count", + "eligible", + "uprated_if_eligible", + "uprated_with_default", + "uprated_any_unit", +] +ALL_VARIABLES = YEAR_VARIABLES + ["uprated_monthly"] +FLOATS = st.floats(min_value=-2000, max_value=2000, allow_nan=False, width=32) +VALUES = {variable: FLOATS for variable in ALL_VARIABLES} +VALUES["uprated_count"] = st.integers(min_value=-2000, max_value=2000) +VALUES["eligible"] = st.booleans() +INPUT_PERIODS = {variable: INPUT_YEARS for variable in YEAR_VARIABLES} +INPUT_PERIODS["uprated_monthly"] = INPUT_MONTHS +INPUT_PERIODS["uprated_any_unit"] = INPUT_YEARS + INPUT_MONTHS +# Other variables read ``eligible`` (``uprated_if_eligible`` is defined for +# it), so a parent request for them caches it at their period: no branch +# input for it. +BRANCH_INPUT_VARIABLES = [name for name in ALL_VARIABLES if name != "eligible"] + + +def _variable_inputs(variable, min_size=0): + return st.dictionaries( + st.sampled_from(INPUT_PERIODS[variable]), + st.lists(VALUES[variable], min_size=COUNT, max_size=COUNT), + min_size=min_size, + max_size=3, + ) + + +def _requests_on(variable): + """Requests for ``variable``: plain calculations in either unit (a yearly + flow asked for a month caches a twelfth there, a monthly flow asked for a + year caches the sum of its months), and the ADD and DIVIDE options.""" + years = st.sampled_from(YEARS) + any_periods = st.sampled_from(YEARS + REQUEST_MONTHS) + if variable == "uprated_monthly": + return st.tuples( + st.sampled_from(["calculate", "add"]), st.just(variable), any_periods + ) + return st.one_of( + st.tuples(st.just("calculate"), st.just(variable), any_periods), + st.tuples(st.just("add"), st.just(variable), years), + st.tuples(st.just("divide"), st.just(variable), any_periods), + ) + + +def _overlaps(first, second): + first, second = periods.period(first), periods.period(second) + return first.start <= second.stop and second.start <= first.stop + + +target_strategy = st.one_of( + st.tuples(st.sampled_from(YEAR_VARIABLES), st.sampled_from(YEARS)), + st.tuples(st.just("uprated_monthly"), st.sampled_from(REQUEST_MONTHS)), +) + + +@st.composite +def scenarios(draw): + """A target, the inputs, the requests made before the target, the branch + it is read from and that branch's own input, if any. The target variable + always has an input, and about half the requests are for it: those are + the requests that cache values it could be uprated from.""" + variable, period = draw(target_strategy) + inputs = { + name: draw(_variable_inputs(name, min_size=int(name == variable))) + for name in ALL_VARIABLES + } + inputs = {name: values for name, values in inputs.items() if values} + any_request = st.one_of(*(_requests_on(name) for name in ALL_VARIABLES)) + requests = draw( + st.lists(st.one_of(_requests_on(variable), any_request), max_size=8) + ) + branch = draw(st.sampled_from([None, "reform", "nested"])) + branch_input = None + if branch is not None and variable in BRANCH_INPUT_VARIABLES: + branch_input = draw( + st.none() + | st.tuples( + st.sampled_from(INPUT_PERIODS[variable]), + st.lists(VALUES[variable], min_size=COUNT, max_size=COUNT), + ) + ) + if branch_input is not None: + requests = [ + (kind, name, requested) + for kind, name, requested in requests + if not (name == variable and _overlaps(requested, period)) + ] + return inputs, requests, variable, period, branch, branch_input + + +def _read_from(built, branch, variable, branch_input): + if branch == "reform": + built = built.get_branch("reform") + elif branch == "nested": + built = built.get_branch("reform").get_branch("nested") + if branch_input is not None: + built.set_input(variable, branch_input[0], branch_input[1]) + return built + + +@settings( + max_examples=500, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given(scenario=scenarios(), auto_carry_over=st.booleans()) +# The regressions in test_uprating_order.py, as examples of the property. +@example( + scenario=( + {"uprated_count": {"2012": [1001, 77]}}, + [ + ("calculate", "uprated_count", "2013"), + ("calculate", "uprated_count", "2014"), + ], + "uprated_count", + "2015", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated": {"2012": [1001.3, 77.7]}}, + [("calculate", "uprated", "2013"), ("divide", "uprated", "2014")], + "uprated", + "2015", + "nested", + None, + ), + auto_carry_over=False, +) +@example( + scenario=( + { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": {"2012": [True, True], "2013": [False, True]}, + }, + [("calculate", "uprated_if_eligible", "2013")], + "uprated_if_eligible", + "2015", + "reform", + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": {"2012": [True, True], "2013": [False, False]}, + }, + [("add", "uprated_if_eligible", "2013")], + "uprated_if_eligible", + "2014", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated_with_default": {"2016": [1, 2]}}, + [("calculate", "uprated_with_default", "2013")], + "uprated_with_default", + "2015", + None, + None, + ), + auto_carry_over=False, +) +@example( + scenario=( + {"uprated_any_unit": {"2012-03": [10, 20]}}, + [ + ("divide", "uprated_any_unit", "2012-03"), + ("calculate", "uprated_any_unit", "2013"), + ], + "uprated_any_unit", + "2015", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated_monthly": {"2011-06": [1001.3, 77.7]}}, + [ + ("calculate", "uprated_monthly", "2012-02"), + ("add", "uprated_monthly", "2013"), + ], + "uprated_monthly", + "2014-03", + None, + None, + ), + auto_carry_over=True, +) +# A branch input over a period the parent calculated before the fork. +@example( + scenario=( + {"uprated": {"2012": [1, 2]}}, + [("calculate", "uprated", "2013")], + "uprated", + "2015", + "reform", + ("2013", [300.0, 400.0]), + ), + auto_carry_over=True, +) +# Signed zero survives uprating from an input. +@example( + scenario=( + {"uprated": {"2012": [-0.0, 0.0]}}, + [("calculate", "uprated", "2013"), ("calculate", "uprated", "2014")], + "uprated", + "2015", + None, + None, + ), + auto_carry_over=False, +) +def test_uprating_depends_only_on_the_inputs(scenario, auto_carry_over): + inputs, requests, variable, period, branch, branch_input = scenario + system = SYSTEMS[auto_carry_over] + + built = simulation(system, inputs) + for kind, requested_variable, requested_period in requests: + request(built, kind, requested_variable, requested_period) + result = _read_from(built, branch, variable, branch_input).calculate( + variable, period + ) + + fresh = _read_from(simulation(system, inputs), branch, variable, branch_input) + message = f"{variable} {period} after {requests}, branch {branch} {branch_input}" + assert_bitwise_equal(result, fresh.calculate(variable, period), message) + + seen = {name: dict(values) for name, values in inputs.items()} + if branch_input is not None: + seen.setdefault(variable, {})[branch_input[0]] = branch_input[1] + assert_bitwise_equal(result, reference(system, seen, variable, period), message) diff --git a/tests/core/variables/test_variable_uprating.py b/tests/core/variables/test_variable_uprating.py index eb634a92..7cf6cc05 100644 --- a/tests/core/variables/test_variable_uprating.py +++ b/tests/core/variables/test_variable_uprating.py @@ -211,16 +211,15 @@ def test_matches_explicitly_backdated_index(system, backdated_system, known, req @pytest.mark.parametrize("known", range(2008, 2016)) def test_result_does_not_depend_on_intermediate_years_computed(system, known): - """Path independence: computing every year in turn (each uprating from the - one before, which is then known) gives the same final value as uprating - straight from the known year.""" + """Path independence: computing every year in turn first gives exactly + the value uprated straight from the known year. Each year uprates from + the input, not from the year before (see test_uprating_order.py); chained + from the year before, float32 rounding compounded.""" stepwise = simulate(system, {known: VALUE}) for year in range(known + 1, 2022): stepwise.calculate("uprated_income", year) stepwise_result = float(stepwise.calculate("uprated_income", 2022)[0]) - assert stepwise_result == pytest.approx( - calculate(system, {known: VALUE}, 2022), rel=1e-5 - ) + assert stepwise_result == calculate(system, {known: VALUE}, 2022) @pytest.mark.parametrize( diff --git a/tests/fixtures/uprating_order.py b/tests/fixtures/uprating_order.py new file mode 100644 index 00000000..826fcc8b --- /dev/null +++ b/tests/fixtures/uprating_order.py @@ -0,0 +1,215 @@ +"""Variables, systems and a reference rule for the uprating order tests. + +Shared by ``tests/core/test_uprating_order.py`` (regressions) and +``tests/core/test_uprating_order_property.py`` (properties). The simulation +helpers are the auto-carry-over tests' (``tests/fixtures/carry_over.py``). + +The reference rule, for a variable with ``uprating`` and no formula result +for a period: an input stored for the period itself is read back as stored; +otherwise the latest earlier input stored in the variable's own unit, times +the ratio of the uprating index at the two period starts, cast to the +variable's type. With no such input, the period gets what auto-carry-over +gives it (the input stored for the latest-starting period that starts no +later than it, on a tie the one that ends last, preferring the variable's own +unit) or, without auto-carry-over or such an input, the default. Either way the value is masked +by the variable's ``defined_for`` at the period. +""" + +from __future__ import annotations + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.parameters import ParameterNode, get_parameter +from policyengine_core.simulations.simulation import _uprating_index_value +from policyengine_core.variables import Variable +from tests.fixtures.carry_over import COUNT, alone, request, simulation + +__all__ = [ + "COUNT", + "INDEX", + "UPRATING", + "alone", + "assert_bitwise_equal", + "build_system", + "index", + "reference", + "request", + "simulation", +] + +# 3.7% a year: ratios that are not exact in float32, so rounding at an +# intermediate year shows. +INDEX = {f"{year}-01-01": 100 * 1.037 ** (year - 2010) for year in range(2010, 2021)} +UPRATING = "uprating_order.index" + + +class uprated(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + label = "Uprated yearly input" + + +class uprated_count(Variable): + value_type = int + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + label = "Uprated yearly integer input" + + +class eligible(Variable): + value_type = bool + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly boolean input that uprated_if_eligible is defined for" + + +class uprated_if_eligible(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + defined_for = "eligible" + label = "Uprated yearly input, defined for eligible people" + + +class uprated_with_default(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + default_value = 10 + label = "Uprated yearly input with a non-zero default" + + +class uprated_any_unit(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + set_input = None + label = "Uprated yearly input with no set_input helper, stored at any period" + + +class uprated_monthly(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + uprating = UPRATING + label = "Uprated monthly input" + + +VARIABLES = ( + uprated, + uprated_count, + eligible, + uprated_if_eligible, + uprated_with_default, + uprated_any_unit, + uprated_monthly, +) + + +def build_system(auto_carry_over: bool = True) -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = auto_carry_over + system.parameters.add_child( + "uprating_order", + ParameterNode("uprating_order", data={"index": {"values": INDEX}}), + ) + system.add_variables(*VARIABLES) + return system + + +def index(year: int) -> float: + return INDEX[f"{year}-01-01"] + + +def assert_bitwise_equal(actual, expected, message=""): + """Same dtype, shape and bytes: ``-0.0`` and ``0.0`` differ.""" + actual, expected = np.asarray(actual), np.asarray(expected) + assert actual.dtype == expected.dtype, ( + f"{message}: {actual.dtype} != {expected.dtype}" + ) + assert actual.shape == expected.shape, ( + f"{message}: {actual.shape} != {expected.shape}" + ) + assert actual.tobytes() == expected.tobytes(), f"{message}: {actual} != {expected}" + + +def _stored_inputs(system, inputs, variable): + """The inputs as ``set_input`` stores them, cast to the variable's type. + + The reference rule covers inputs in the variable's own unit, and inputs + in any unit to a variable with no ``set_input`` helper (stored as given). + A helper would split or copy an input in another unit across the + variable's own periods first, which the rule does not model. + """ + definition = system.get_variable(variable) + stored = {} + for period, values in inputs.get(variable, {}).items(): + period = periods.period(period) + if period.unit != definition.definition_period and definition.set_input: + raise ValueError( + f"{variable}: an input for {period} goes through set_input, " + "which the reference rule does not model" + ) + stored[period] = np.asarray(values, dtype=definition.dtype) + return stored + + +def _uprating_factor(system, definition, earlier, period): + parameter = get_parameter(system.parameters, definition.uprating) + value_then = _uprating_index_value(parameter, earlier.start) + value_now = _uprating_index_value(parameter, period.start) + if value_then is None or value_now is None or value_then == 0: + return 1 + return value_now / value_then + + +def reference(system, inputs, variable, period): + """The uprating and carry-over rules, computed from the inputs alone. + + ``period`` is in the variable's definition period. + """ + definition = system.get_variable(variable) + period = periods.period(period) + default = definition.default_array(COUNT) + stored = _stored_inputs(system, inputs, variable) + if period in stored: + # An input for the period itself is read back as stored: + # ``defined_for`` does not apply to it. + return stored[period] + earlier = [ + stored_period + for stored_period in stored + if stored_period.unit == definition.definition_period + and stored_period.start < period.start + ] + carried = [ + stored_period for stored_period in stored if stored_period.start <= period.start + ] + if definition.uprating is not None and earlier: + latest = max(earlier, key=lambda stored_period: stored_period.start) + value = stored[latest] * _uprating_factor(system, definition, latest, period) + elif system.auto_carry_over_input_variables and carried: + own_unit = [ + stored_period + for stored_period in carried + if stored_period.unit == definition.definition_period + ] + latest = max( + own_unit or carried, + key=lambda stored_period: (stored_period.start, stored_period.stop), + ) + value = stored[latest] + else: + value = default + if definition.defined_for is not None: + mask = reference(system, inputs, definition.defined_for, period) + value = np.where(mask, value, default) + return np.asarray(value).astype(definition.dtype)