From a80a3b9641c5db1aef007de94ffff3d2b6dcb6ba Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:19:15 -0400 Subject: [PATCH 01/10] 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 000000000..a25488556 --- /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 a1bddf76b..61c3b4c38 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 3563c805d..1f835f1c2 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 11a8ae50b..12a679c8a 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 2652d9723..7f49e0eaf 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 c381b11c4..1fae168fe 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 c3db0c4fa..e2376e1e9 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 000000000..7e76a4149 --- /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 000000000..d343d830d --- /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 000000000..ad535e3c4 --- /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 02/10] 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 61c3b4c38..0a7ed4b4e 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 1f835f1c2..ed9d41b75 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 12a679c8a..872b7ed33 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 7f49e0eaf..a7ab065d5 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 e2376e1e9..65b3c6644 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 7e76a4149..aa481582a 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 ad535e3c4..28ece23c0 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 03/10] 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 aa481582a..f3a398abf 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 04/10] 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 a7ab065d5..58190aabd 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 f3a398abf..481932bcf 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 05/10] 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 0a7ed4b4e..9feffb59e 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 ed9d41b75..db3e2cad0 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 872b7ed33..a02208894 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 58190aabd..ed8186559 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 1fae168fe..b7073081d 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 481932bcf..42c238081 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 06/10] 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 42c238081..04801aca6 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 07/10] 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 ed8186559..ce22f3a08 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 04801aca6..9c077364b 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 28ece23c0..83f7e3f32 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 d89a2ea91b3678130c8e949dfa24cb392acba1e1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 10:45:42 -0400 Subject: [PATCH 08/10] Record a storage's inputs, not its derived values, and share them with clones Measured on a policyengine-us household (policyengine-us 2.18.3, one household_net_income calculation, 6,185 storages, 30,925 holder clones across 5 branches), master with #578 retains 20.94 MB per simulation. With derived marks as a set per storage that holds a derived value, the merge retained 24.35 MB: 1,514 sets in the simulation and a copy of its source's marks in every holder clone. That is 2.6 times the 1.3 MB per simulation that ran policyengine-us's CI out of memory (#577). 5,201 of the 5,204 stored values are derived, so record the inputs instead: is_derived is "stored and not an input". The 3 storages holding inputs have a set (648 bytes); the rest read _NO_INPUTS. A storage and its clones share one frozenset of inputs until one of them changes its own. Retained memory is now 21.25 MB per simulation, +0.31 MB on master: one more 8-byte attribute slot in each of the 37,110 storage objects. A pickle records that inputs were recorded; a state from before (no marker) counts every stored value as an input, as it did then. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 136 ++++++--- tests/core/test_carry_over_order.py | 22 +- tests/core/test_storage_derived_index.py | 170 ----------- tests/core/test_storage_input_index.py | 282 ++++++++++++++++++ .../test_storage_shared_index_differential.py | 44 +-- 5 files changed, 411 insertions(+), 243 deletions(-) delete mode 100644 tests/core/test_storage_derived_index.py create mode 100644 tests/core/test_storage_input_index.py diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index dcc609a2f..f1c3304f4 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -38,9 +38,14 @@ def _read_only_view(array: numpy.ndarray) -> numpy.ndarray: # share nothing, and an empty set is 216 bytes. _NOTHING_SHARED: FrozenSet[str] = frozenset() -# What ``InMemoryStorage._derived`` is while no stored value is derived, for -# the same reason: most storages in a simulation hold no calculated value. -_NOTHING_DERIVED: FrozenSet[str] = frozenset() +# What ``InMemoryStorage._inputs`` is while a storage holds no input, for the +# same reason: most storages in a simulation hold no input. +_NO_INPUTS: FrozenSet[str] = frozenset() + +# Key a pickled storage's state carries once it records its inputs; a state +# without it comes from a version that recorded neither inputs nor derived +# values (see ``__setstate__``). +_INPUTS_RECORDED = "_inputs_recorded" class InMemoryStorage: @@ -56,20 +61,47 @@ class InMemoryStorage: # this class attribute. A key left in the set after code outside this # class empties ``_arrays`` costs one extra copy at most. _shared: Union[Set[str], FrozenSet[str]] = _NOTHING_SHARED - # 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. As with - # ``_shared``, a storage has a set of its own only while at least one - # stored value is derived; otherwise it reads this class attribute. A - # storage pickled before derived marks existed reads it too, so its - # values count as inputs. - _derived: Union[Set[str], FrozenSet[str]] = _NOTHING_DERIVED + # Keys whose value was stored as an input: with ``put(..., derived=False)``, + # the default. Every other stored value was calculated by the simulation + # (``put(..., derived=True)``), which ``is_derived`` reports. A key counts + # only while it is stored, and every ``put`` sets or clears its mark. + # + # The storage records inputs rather than derived values because a + # simulation calculates far more values than it is given (in a + # policyengine-us household, 5,201 of 5,204 stored values are derived). + # As with ``_shared``, a storage has a set of its own only while it holds + # an input; otherwise it reads this class attribute. + # + # A storage's own inputs are a set it alone changes, or a frozenset it + # shares with storages cloned from it (or it from them): ``clone`` hands + # both the same frozenset, and whichever changes its inputs first + # replaces it with a set of its own (``_own_inputs``). A branch clones + # every holder, so copying at each clone would cost a set per input + # holder per branch for inputs most branches never change. + _inputs: Union[Set[str], FrozenSet[str]] = _NO_INPUTS is_eternal: bool def __init__(self, is_eternal: bool): self._arrays = {} self.is_eternal = is_eternal + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state[_INPUTS_RECORDED] = True + return state + + def __setstate__(self, state: dict) -> None: + if not state.pop(_INPUTS_RECORDED, False): + # Pickled before storages recorded their inputs, when nothing + # marked a calculated value: every stored value counts as an + # input, as it did then, except any a development version of this + # change marked derived. + derived = state.pop("_derived", ()) + inputs = set(state.get("_arrays", {})).difference(derived) + if inputs: + state["_inputs"] = inputs + self.__dict__.update(state) + def clone(self, share_arrays: bool = False) -> "InMemoryStorage": """Copy this storage. @@ -107,14 +139,19 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage": clone._shared = shared else: clone._arrays = {key: array.copy() for key, array in self._arrays.items()} - # Copy the marks of the keys the clone holds: a mark left behind by - # code outside this class emptying ``_arrays`` never counts (see - # ``is_derived``), so a clone has no use for it. - derived = self._derived - if derived: - marks = derived.intersection(clone._arrays) - if marks: - clone._derived = marks + # Share the inputs (see ``_inputs``). A mark left behind by code + # outside this class emptying ``_arrays`` never counts (see + # ``is_derived``), so the clone gets only those of keys it holds. + inputs = self._inputs + if inputs: + if inputs <= clone._arrays.keys(): + if not isinstance(inputs, frozenset): + inputs = self.__dict__["_inputs"] = frozenset(inputs) + clone._inputs = inputs + else: + held = inputs.intersection(clone._arrays) + if held: + clone._inputs = frozenset(held) return clone def get(self, period: Period, branch_name: str = "default") -> ArrayLike: @@ -148,7 +185,7 @@ def is_derived(self, period: Period, branch_name: str = "default") -> bool: 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 + return key in self._arrays and key not in self._inputs # Each method below reads ``self._shared`` once and works on that set, and # releases it with one ``dict.pop`` that never raises. A storage is still @@ -181,32 +218,41 @@ def _release_if_empty(self, shared: Set[str]) -> None: if not shared: self.__dict__.pop("_shared", None) - def _mark_derived(self, key: str) -> None: - """Record that the value stored for ``key`` was calculated.""" - derived = self.__dict__.get("_derived") - if derived is None: - derived = self.__dict__["_derived"] = set() - derived.add(key) - - def _unmark_derived(self, key: str) -> None: + def _own_inputs(self) -> Set[str]: + """This storage's inputs as a set that only it changes, made from the + inputs it shares, or from nothing, if need be.""" + inputs = self.__dict__.get("_inputs") + if inputs is None: + inputs = self.__dict__["_inputs"] = set() + elif isinstance(inputs, frozenset): + inputs = self.__dict__["_inputs"] = set(inputs) + return inputs + + def _mark_input(self, key: str) -> None: """Record that the value stored for ``key`` is an input.""" - derived = self._derived - if key in derived: - derived.discard(key) - self._release_derived_if_empty(derived) + if key not in self._inputs: + self._own_inputs().add(key) - def _unmark_dropped_keys(self) -> None: - """Forget the derived marks of keys ``_arrays`` no longer has.""" - derived = self._derived - if derived: - derived.intersection_update(self._arrays) - self._release_derived_if_empty(derived) + def _unmark_input(self, key: str) -> None: + """Record that the value stored for ``key`` was calculated.""" + if key in self._inputs: + inputs = self._own_inputs() + inputs.discard(key) + self._release_inputs_if_empty(inputs) - def _release_derived_if_empty(self, derived: Set[str]) -> None: - """Go back to the class's nothing-derived object once no stored value - is derived, as ``_release_if_empty`` does for ``_shared``.""" - if not derived: - self.__dict__.pop("_derived", None) + def _unmark_dropped_keys(self) -> None: + """Forget the input marks of keys ``_arrays`` no longer has.""" + inputs = self._inputs + if inputs and not inputs <= self._arrays.keys(): + inputs = self._own_inputs() + inputs.intersection_update(self._arrays) + self._release_inputs_if_empty(inputs) + + def _release_inputs_if_empty(self, inputs: Set[str]) -> None: + """Go back to the class's no-inputs object once no stored value is an + input, as ``_release_if_empty`` does for ``_shared``.""" + if not inputs: + self.__dict__.pop("_inputs", None) def put( self, @@ -238,9 +284,9 @@ def put( self._arrays[key] = value self._stop_sharing(key) if derived: - self._mark_derived(key) + self._unmark_input(key) else: - self._unmark_derived(key) + self._mark_input(key) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 879ef9a53..a82ea7417 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -395,6 +395,12 @@ def test_carry_over_in_a_branch_copies_only_the_array_it_reads(system): assert len(storage._shared) == shared - 1 +def _marks(storage): + """A storage's own record of which values are which: an in-memory + storage records its inputs, an on-disk one its derived values.""" + return storage._derived if isinstance(storage, OnDiskStorage) else storage._inputs + + @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 = ( @@ -416,13 +422,13 @@ def test_storages_keep_the_derived_mark_with_the_value(on_disk, tmp_path): storage.put(value, year, derived=True) storage.delete(year) assert not storage.has(year) and not storage.is_derived(year) - assert not storage._derived + assert not _marks(storage) 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 + assert not _marks(storage) storage.put(value, year) assert not storage.is_derived(year) @@ -539,7 +545,7 @@ def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): for year in range(2010, 2020): built.calculate("carried", str(year)) built.apply_reform(_noop) - assert not holder._memory_storage._derived + assert not holder._memory_storage._inputs 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. @@ -564,7 +570,7 @@ def test_marks_are_cleared_with_in_memory_values_apply_reform_wipes(system): for year in range(2010, 2020): built.calculate("carried", str(year)) built.apply_reform(_noop) - assert not holder._memory_storage._derived + assert not holder._memory_storage._inputs def test_rebuilding_a_disk_index_reads_files_as_inputs(system): @@ -607,12 +613,12 @@ def test_storages_pickled_without_marks_still_work(on_disk, tmp_path): ) year = periods.period("2012") storage.put(np.array([10.0]), year) + # The state an older version pickled: its arrays or files, and no marks. state = dict(storage.__dict__) - # An in-memory storage holding only inputs has no marks of its own. - state.pop("_derived", None) - state.pop("_shared", None) + for name in ("_derived", "_inputs", "_shared"): + state.pop(name, None) old = type(storage).__new__(type(storage)) - old.__dict__.update(state) + old.__setstate__(state) restored = pickle.loads(pickle.dumps(old)) assert restored.has(year) and not restored.is_derived(year) restored.clone() diff --git a/tests/core/test_storage_derived_index.py b/tests/core/test_storage_derived_index.py deleted file mode 100644 index 1e343b4c2..000000000 --- a/tests/core/test_storage_derived_index.py +++ /dev/null @@ -1,170 +0,0 @@ -"""A storage holds a set of derived keys only while it holds a derived value. - -``InMemoryStorage._derived`` records which stored values the simulation -calculated rather than took as input (``put(..., derived=True)``), so that -auto-carry-over and uprating start only from inputs. Like ``_shared`` -(``test_storage_shared_index.py``), it is one empty object for every storage -until a storage holds a derived value: a simulation has a storage for every -variable of the tax-benefit system, and an empty set in each is 216 bytes -per variable per simulation. ``test_storage_shared_index_differential.py`` -checks random sequences of operations against the storage that always had -both sets. -""" - -from __future__ import annotations - -import copy -import pickle - -import numpy as np -import pytest - -from policyengine_core import periods -from policyengine_core.data_storage import InMemoryStorage -from policyengine_core.data_storage.in_memory_storage import ( - _NOTHING_DERIVED, - _NOTHING_SHARED, -) -from policyengine_core.reforms import Reform -from tests.fixtures.branch_shared_arrays import build_simulation - -JANUARY = periods.period("2017-01") - - -def _storages(simulation): - return { - name: holder._memory_storage - for population in simulation.populations.values() - for name, holder in population._holders.items() - } - - -def test_the_nothing_derived_index_is_empty_and_cannot_change(): - assert isinstance(_NOTHING_DERIVED, frozenset) - assert not _NOTHING_DERIVED - - -def test_storing_reading_and_deleting_inputs_never_needs_a_set(): - storage = InMemoryStorage(is_eternal=False) - assert storage._derived is _NOTHING_DERIVED - storage.put(np.array([1.0]), "2017-01") - storage.get("2017-01") - storage.delete("2017-01") - storage.put(np.array([1.0]), "2017-02") - storage.delete() - assert storage._derived is _NOTHING_DERIVED - assert "_derived" not in vars(storage) - - -@pytest.mark.parametrize( - "drop", - [ - lambda storage: storage.put(np.array([2.0]), "2017-01"), - lambda storage: storage.delete("2017-01"), - lambda storage: storage.delete("2017"), - lambda storage: storage.delete(), - ], - ids=["replaced_by_an_input", "deleted", "deleted_by_a_containing_period", "wiped"], -) -def test_a_derived_value_gets_a_set_that_is_released_with_it(drop): - storage = InMemoryStorage(is_eternal=False) - storage.put(np.array([1.0]), "2017-01", derived=True) - assert storage._derived == {"default:2017-01"} - assert "_derived" in vars(storage) - assert storage.is_derived("2017-01") - - drop(storage) - assert not storage.is_derived("2017-01") - # Released, not replaced by another empty object. - assert "_derived" not in vars(storage) - assert storage._derived is _NOTHING_DERIVED - - -def test_the_set_is_kept_while_another_value_is_derived(): - storage = InMemoryStorage(is_eternal=False) - storage.put(np.array([1.0]), "2017-01", derived=True) - storage.put(np.array([1.0]), "2017-02", derived=True) - storage.delete("2017-01") - assert storage._derived == {"default:2017-02"} - storage.put(np.array([3.0]), "2017-02") - assert storage._derived is _NOTHING_DERIVED - - -@pytest.mark.parametrize("share_arrays", [False, True], ids=["copy", "share"]) -def test_clones_copy_the_marks_into_a_set_of_their_own(share_arrays): - inputs = InMemoryStorage(is_eternal=False) - inputs.put(np.array([1.0]), "2017-01") - assert inputs.clone(share_arrays)._derived is _NOTHING_DERIVED - - derived = InMemoryStorage(is_eternal=False) - derived.put(np.array([1.0]), "2017-01", derived=True) - clone = derived.clone(share_arrays) - assert clone._derived == derived._derived == {"default:2017-01"} - assert clone._derived is not derived._derived - clone.put(np.array([2.0]), "2017-01") - assert clone._derived is _NOTHING_DERIVED - assert derived.is_derived("2017-01") - - -@pytest.mark.parametrize( - "duplicate", - [copy.deepcopy, lambda storage: pickle.loads(pickle.dumps(storage))], - ids=["deepcopy", "pickle"], -) -def test_copied_or_unpickled_storage_keeps_the_marks(duplicate): - inputs = InMemoryStorage(is_eternal=False) - inputs.put(np.array([1.0]), "2017-01") - twin = duplicate(inputs) - assert twin._derived is _NOTHING_DERIVED - assert twin._shared is _NOTHING_SHARED - assert not twin.is_derived("2017-01") - - derived = InMemoryStorage(is_eternal=False) - derived.put(np.array([1.0]), "2017-01", derived=True) - twin = duplicate(derived) - assert twin.is_derived("2017-01") - assert twin._derived is not derived._derived - - -def test_storage_pickled_with_an_empty_set_still_works(): - # Before storages shared one empty object, each pickled its own sets. - storage = InMemoryStorage(is_eternal=False) - storage.put(np.array([1.0]), "2017-01") - storage.__dict__["_derived"] = set() - storage.__dict__["_shared"] = set() - restored = pickle.loads(pickle.dumps(storage)) - assert not restored.is_derived("2017-01") - restored.put(np.array([2.0]), "2017-02", derived=True) - assert restored.is_derived("2017-02") - restored.delete() - assert restored.get("2017-01") is None - - -def test_only_storages_holding_a_derived_value_have_a_set(tax_benefit_system): - simulation = build_simulation(tax_benefit_system) - simulation.calculate("disposable_income", JANUARY) - - storages = _storages(simulation) - assert len(storages) == len(tax_benefit_system.variables) - with_derived = {name for name, storage in storages.items() if storage._derived} - # Calculating disposable income derives some values and not others. - assert with_derived and with_derived < set(storages) - for name, storage in storages.items(): - assert ("_derived" in vars(storage)) == (name in with_derived), name - if name not in with_derived: - assert storage._derived is _NOTHING_DERIVED, name - assert storage._shared is _NOTHING_SHARED, name - - -def test_apply_reform_releases_the_marks(tax_benefit_system): - class _noop(Reform): - def apply(self): - pass - - simulation = build_simulation(tax_benefit_system) - simulation.calculate("disposable_income", JANUARY) - assert any(storage._derived for storage in _storages(simulation).values()) - simulation.apply_reform(_noop) - for name, storage in _storages(simulation).items(): - assert storage._derived is _NOTHING_DERIVED, name - assert "_derived" not in vars(storage), name diff --git a/tests/core/test_storage_input_index.py b/tests/core/test_storage_input_index.py new file mode 100644 index 000000000..b5ecb3c39 --- /dev/null +++ b/tests/core/test_storage_input_index.py @@ -0,0 +1,282 @@ +"""A storage records which of its values are inputs, at no cost when none is. + +``InMemoryStorage.is_derived`` tells a value the simulation calculated +(``put(..., derived=True)``) from an input, so that auto-carry-over and +uprating start only from inputs. The storage records the inputs, in +``_inputs``, because a simulation calculates far more values than it is +given. Like ``_shared`` (``test_storage_shared_index.py``), ``_inputs`` is +one empty object for every storage until a storage holds an input: a +simulation has a storage for every variable of the tax-benefit system, and +an empty set in each is 216 bytes per variable per simulation (#577). A +storage and its clones share one frozenset of inputs until one of them +changes its own. ``test_storage_shared_index_differential.py`` checks random +sequences of operations against a storage that always had its own sets. +""" + +from __future__ import annotations + +import copy +import pickle + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.in_memory_storage import ( + _INPUTS_RECORDED, + _NO_INPUTS, + _NOTHING_SHARED, +) +from policyengine_core.reforms import Reform +from tests.fixtures.branch_shared_arrays import build_simulation + +JANUARY = periods.period("2017-01") + + +def _storages(simulation): + return { + name: holder._memory_storage + for population in simulation.populations.values() + for name, holder in population._holders.items() + } + + +def test_the_no_inputs_index_is_empty_and_cannot_change(): + assert isinstance(_NO_INPUTS, frozenset) + assert not _NO_INPUTS + + +def test_storing_reading_and_deleting_derived_values_never_needs_a_set(): + storage = InMemoryStorage(is_eternal=False) + assert storage._inputs is _NO_INPUTS + storage.put(np.array([1.0]), "2017-01", derived=True) + assert storage.is_derived("2017-01") + storage.get("2017-01") + storage.delete("2017-01") + storage.put(np.array([1.0]), "2017-02", derived=True) + storage.delete() + assert storage._inputs is _NO_INPUTS + assert "_inputs" not in vars(storage) + + +@pytest.mark.parametrize( + "drop", + [ + lambda storage: storage.put(np.array([2.0]), "2017-01", derived=True), + lambda storage: storage.delete("2017-01"), + lambda storage: storage.delete("2017"), + lambda storage: storage.delete(), + ], + ids=[ + "replaced_by_a_derived_value", + "deleted", + "deleted_by_a_containing_period", + "wiped", + ], +) +def test_an_input_gets_a_set_that_is_released_with_it(drop): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01") + assert storage._inputs == {"default:2017-01"} + assert "_inputs" in vars(storage) + assert not storage.is_derived("2017-01") + + drop(storage) + # Released, not replaced by another empty object. + assert "_inputs" not in vars(storage) + assert storage._inputs is _NO_INPUTS + + +def test_the_set_is_kept_while_another_value_is_an_input(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01") + storage.put(np.array([1.0]), "2017-02") + storage.delete("2017-01") + assert storage._inputs == {"default:2017-02"} + storage.put(np.array([3.0]), "2017-02", derived=True) + assert storage._inputs is _NO_INPUTS + assert storage.is_derived("2017-02") + + +def test_only_stored_values_are_derived(): + storage = InMemoryStorage(is_eternal=False) + assert not storage.is_derived("2017-01") + storage.put(np.array([1.0]), "2017-01", "reform", derived=True) + assert not storage.is_derived("2017-01") + assert storage.is_derived("2017-01", "reform") + + +@pytest.mark.parametrize("share_arrays", [False, True], ids=["copy", "share"]) +def test_clones_share_the_inputs_until_one_of_them_changes_its_own(share_arrays): + derived_only = InMemoryStorage(is_eternal=False) + derived_only.put(np.array([1.0]), "2017-01", derived=True) + assert derived_only.clone(share_arrays)._inputs is _NO_INPUTS + + source = InMemoryStorage(is_eternal=False) + source.put(np.array([1.0]), "2017-01") + source.put(np.array([1.0]), "2017-02") + first = source.clone(share_arrays) + second = source.clone(share_arrays) + grandchild = first.clone(share_arrays) + # One immutable set for all four: a clone costs no set of its own. + inputs = source._inputs + assert isinstance(inputs, frozenset) + assert inputs == {"default:2017-01", "default:2017-02"} + assert first._inputs is inputs and second._inputs is inputs + assert grandchild._inputs is inputs + + # Storing an input over an input changes nothing. + first.put(np.array([5.0]), "2017-01") + assert first._inputs is inputs + + # The first storage to change its inputs gets a set of its own. + first.put(np.array([2.0]), "2017-01", derived=True) + assert first.is_derived("2017-01") and not first.is_derived("2017-02") + assert isinstance(first._inputs, set) + for other in (source, second, grandchild): + assert other._inputs is inputs and not other.is_derived("2017-01") + + source.put(np.array([3.0]), "2017-03") + assert source._inputs == inputs | {"default:2017-03"} + second.put(np.array([3.0]), "2017-03", derived=True) + assert second.is_derived("2017-03") and not source.is_derived("2017-03") + second.delete("2017-02") + assert second._inputs == {"default:2017-01"} + assert grandchild._inputs is inputs + grandchild.delete() + assert grandchild._inputs is _NO_INPUTS + assert not source.is_derived("2017-01") and not source.is_derived("2017-02") + + +def test_clone_of_an_emptied_storage_gets_only_the_inputs_it_holds(): + source = InMemoryStorage(is_eternal=False) + source.put(np.array([1.0]), "2017-01") + source._arrays.clear() # as policyengine-uk clears a dropped branch + assert source.clone()._inputs is _NO_INPUTS + source.put(np.array([1.0]), "2017-02") + clone = source.clone() + assert clone._inputs == {"default:2017-02"} + assert isinstance(clone._inputs, frozenset) + + +@pytest.mark.parametrize( + "duplicate", + [copy.deepcopy, lambda storage: pickle.loads(pickle.dumps(storage))], + ids=["deepcopy", "pickle"], +) +def test_copied_or_unpickled_storage_keeps_the_marks(duplicate): + derived = InMemoryStorage(is_eternal=False) + derived.put(np.array([1.0]), "2017-01", derived=True) + twin = duplicate(derived) + assert twin._inputs is _NO_INPUTS + assert twin._shared is _NOTHING_SHARED + assert twin.is_derived("2017-01") + + for cloned_first in (False, True): + inputs = InMemoryStorage(is_eternal=False) + inputs.put(np.array([1.0]), "2017-01") + if cloned_first: + inputs.clone() # its inputs are now a frozenset it shares + twin = duplicate(inputs) + assert not twin.is_derived("2017-01") + # Changing either one's inputs leaves the other's alone. + twin.put(np.array([2.0]), "2017-01", derived=True) + assert twin.is_derived("2017-01") + assert not inputs.is_derived("2017-01") + inputs.put(np.array([2.0]), "2017-02") + assert not twin.has("2017-02") and not twin.is_derived("2017-02") + + +def _older_state(storage, **extra): + """The state a version that recorded no inputs pickled.""" + state = {"_arrays": dict(storage._arrays), "is_eternal": storage.is_eternal} + state.update(extra) + return state + + +def test_values_from_an_older_pickle_count_as_inputs(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01", derived=True) + storage.put(np.array([1.0]), "2017-02", derived=True) + restored = InMemoryStorage.__new__(InMemoryStorage) + # 3.32.12 pickled an empty set of shared keys in every storage. + restored.__setstate__(_older_state(storage, _shared=set())) + assert not restored.is_derived("2017-01") + assert not restored.is_derived("2017-02") + assert restored._inputs == {"default:2017-01", "default:2017-02"} + assert _INPUTS_RECORDED not in vars(restored) + + # A development version of this change recorded derived values instead. + restored = InMemoryStorage.__new__(InMemoryStorage) + restored.__setstate__(_older_state(storage, _derived={"default:2017-02"})) + assert not restored.is_derived("2017-01") + assert restored.is_derived("2017-02") + assert "_derived" not in vars(restored) + + # An older storage with nothing stored needs no set. + restored = InMemoryStorage.__new__(InMemoryStorage) + restored.__setstate__(_older_state(InMemoryStorage(is_eternal=False))) + assert restored._inputs is _NO_INPUTS + restored.put(np.array([2.0]), "2017-03", derived=True) + assert restored.is_derived("2017-03") + + +def test_a_pickle_records_that_the_inputs_were_recorded(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01", derived=True) + state = storage.__getstate__() + assert state[_INPUTS_RECORDED] is True + assert _INPUTS_RECORDED not in vars(storage) + assert pickle.loads(pickle.dumps(storage)).is_derived("2017-01") + + +def test_only_storages_holding_an_input_have_a_set(tax_benefit_system): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + + storages = _storages(simulation) + assert len(storages) == len(tax_benefit_system.variables) + with_inputs = {name for name, storage in storages.items() if storage._inputs} + with_derived = { + name + for name, storage in storages.items() + if any( + storage.is_derived(key.split(":", 1)[1], key.split(":", 1)[0]) + for key in storage._arrays + ) + } + # The situation's inputs, and the values calculated from them. + assert with_inputs and with_derived + assert with_inputs | with_derived < set(storages) + for name, storage in storages.items(): + assert ("_inputs" in vars(storage)) == (name in with_inputs), name + if name not in with_inputs: + assert storage._inputs is _NO_INPUTS, name + assert storage._shared is _NOTHING_SHARED, name + + +def test_a_branch_shares_its_parents_inputs(tax_benefit_system): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + branch = simulation.get_branch("branch") + parents = _storages(simulation) + assert any(storage._inputs for storage in parents.values()) + for name, storage in _storages(branch).items(): + assert storage._inputs is parents[name]._inputs, name + + +def test_apply_reform_keeps_only_the_replayed_inputs(tax_benefit_system): + class _noop(Reform): + def apply(self): + pass + + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + before = { + name: set(storage._inputs) for name, storage in _storages(simulation).items() + } + simulation.apply_reform(_noop) + for name, storage in _storages(simulation).items(): + assert set(storage._inputs) == before.get(name, set()), name + assert set(storage._arrays) == set(storage._inputs), name diff --git a/tests/core/test_storage_shared_index_differential.py b/tests/core/test_storage_shared_index_differential.py index f3158d85c..343229466 100644 --- a/tests/core/test_storage_shared_index_differential.py +++ b/tests/core/test_storage_shared_index_differential.py @@ -1,15 +1,16 @@ """A storage without a set of its own behaves like one that always had one. ``_EagerIndexStorage`` below is ``InMemoryStorage`` as it was when every -storage owned a set of shared keys and a set of derived keys. Random -sequences of operations, storing inputs and derived values, run on both, and -after every step the two must hold the same arrays, the same shared keys and -the same derived keys; at the end every read and every ``is_derived`` must -agree. The storage under test must also keep its own invariants: it refers -to the one shared-nothing (nothing-derived) object exactly when it shares -nothing (holds no derived value), and no two storages hold the same set. -``test_storage_shared_index.py`` and ``test_storage_derived_index.py`` pin -the same behaviour with examples. +storage owned a set of shared keys and a set of derived keys. The storage +under test records its inputs instead, and has each set only while it is +not empty. Random sequences of operations, storing inputs and derived +values, run on both, and after every step the two must hold the same arrays, +the same shared keys, and the same stored values marked derived; at the end +every read and every ``is_derived`` must agree. The storage under test must +also keep its own invariants: it refers to the one shared-nothing (no-inputs) +object exactly when it shares nothing (holds no input), and no two storages +hold the same mutable set. ``test_storage_shared_index.py`` and +``test_storage_input_index.py`` pin the same behaviour with examples. """ from __future__ import annotations @@ -24,7 +25,7 @@ from policyengine_core import periods from policyengine_core.data_storage import InMemoryStorage from policyengine_core.data_storage.in_memory_storage import ( - _NOTHING_DERIVED, + _NO_INPUTS, _NOTHING_SHARED, _can_share, _read_only_view, @@ -201,9 +202,10 @@ def _assert_same_state(tested, reference): for key, array in storage._arrays.items(): assert _same_array(array, expected._arrays[key]), (index, key) assert storage._shared == expected._shared, index - # The marks that count: a storage emptied from outside keeps the marks - # of keys it no longer holds, which the reference copies into clones. - assert storage._derived & storage._arrays.keys() == ( + # The stored values marked derived. A storage emptied from outside + # keeps the marks of keys it no longer holds, which the reference + # copies into clones; they never count. + assert storage._arrays.keys() - storage._inputs == ( expected._derived & expected._arrays.keys() ), index @@ -223,17 +225,19 @@ def _assert_index_invariants(storages, cleared): for key in shared & storage._arrays.keys(): # A shared array cannot be written through. assert not storage._arrays[key].flags.writeable, (index, key) - derived = storage._derived - assert (derived is _NOTHING_DERIVED) == (not derived), index - assert ("_derived" in vars(storage)) == bool(derived), index - if derived is not _NOTHING_DERIVED: - own_sets.append(derived) + inputs = storage._inputs + assert (inputs is _NO_INPUTS) == (not inputs), index + assert ("_inputs" in vars(storage)) == bool(inputs), index + # Inputs are a set of the storage's own, or a frozenset it may share + # with storages cloned from it or it from them. + if isinstance(inputs, set): + own_sets.append(inputs) if index not in cleared: - assert derived <= storage._arrays.keys(), index + assert inputs <= storage._arrays.keys(), index # A set belongs to one storage, and to one of its two indexes. assert len({id(own) for own in own_sets}) == len(own_sets) assert not _NOTHING_SHARED - assert not _NOTHING_DERIVED + assert not _NO_INPUTS @hypothesis.settings( From 6ddbaf9fc382ed72d81765e1595632b387290d5a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 10:50:14 -0400 Subject: [PATCH 09/10] Carry over an input that ends after the last representable date The carry-over sort key evaluated period.stop for every candidate input. A period that ends after 9999-12-31 (day:9999-12-30:3, or day:2012-01-01:4000000) has no stop: computing it raises OverflowError. Master carried such inputs ([10, 20]); this branch raised. Sort by _end_order, which puts such a period after every period that has a stop, so it still ends last. Found by the review of #582, which uses the same key for uprating. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 13 +++++++++++- tests/core/test_carry_over_order.py | 23 +++++++++++++++++++++ tests/fixtures/carry_over.py | 9 ++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index fe7bb8484..ea13ab31c 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -103,6 +103,17 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: return defined[-1].value +def _end_order(period: Period) -> tuple: + """Sorts periods by when they end. A period that ends after the last + date ``datetime`` can represent (``day:9999-12-30:3``) has no ``stop``, + which raises ``OverflowError``; it sorts after every period that has + one.""" + try: + return (0, period.stop) + except OverflowError: + return (1,) + + if TYPE_CHECKING: from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -978,7 +989,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: key=lambda p: ( p.unit == variable.definition_period, p.start, - p.stop, + _end_order(p), periods.unit_weight(p.unit), str(p), ), diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index a82ea7417..87e10fff5 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -132,6 +132,29 @@ def test_inputs_with_the_same_extent_carry_the_larger_unit(system, reverse): ) +@pytest.mark.parametrize( + "stored,asked", + [("day:9999-12-30:3", "9999-12-31"), ("day:2012-01-01:4000000", "2013-01-01")], +) +def test_an_input_ending_after_the_last_date_carries(system, stored, asked): + """A period that ends after 9999-12-31 has no ``stop`` (it raises + ``OverflowError``); its input carries all the same.""" + inputs = {"day_input_without_helper": {stored: [10, 20]}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", asked), [10, 20] + ) + + +@pytest.mark.parametrize("longer_first", [True, False], ids=["longer", "shorter"]) +def test_an_input_ending_after_the_last_date_ends_last(system, longer_first): + values = {"day:2012-01-01:4000000": [10, 20], "day:2012-01-01:2": [1, 2]} + order = list(values) if longer_first else list(values)[::-1] + inputs = {"day_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", "2013-01-01"), [10, 20] + ) + + 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 83f7e3f32..c68c08047 100644 --- a/tests/fixtures/carry_over.py +++ b/tests/fixtures/carry_over.py @@ -89,6 +89,14 @@ class month_input_without_helper(Variable): label = "Monthly input with no set_input helper, stored at any period" +class day_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.DAY + set_input = None + label = "Daily input with no set_input helper, stored at any period" + + class uprated_input_without_helper(Variable): value_type = float entity = entities.Person @@ -121,6 +129,7 @@ class input_with_calculating_helper(Variable): carried_monthly, year_input_without_helper, month_input_without_helper, + day_input_without_helper, uprated_input_without_helper, input_with_calculating_helper, ) From 2580d43be1f68259cdb1e60fbe052ea83e0c1c2f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 14:33:43 -0400 Subject: [PATCH 10/10] Order period ends by day number, so ends after 9999-12-31 compare too 6ddbaf9f sorted every period whose stop cannot be computed after those that have one, but equal to each other, so of two inputs ending after 9999-12-31 the string form decided: day:9999-12-30:9 was carried over day:9999-12-30:10, and 24 months over 800 days from the same day (delta review, P2). Month and year periods past 9999 also raised ValueError, or returned a stop no date can hold, rather than OverflowError. _end_order is now the number of the day after the period's last day, counted as date.toordinal counts, by integer arithmetic on the same calendar. It equals stop + 1 day wherever stop has a value and orders every other period by its true end. Tests: the review's two cases in both set orders; a property that _end_order agrees with Period.stop where it has a value and with numpy's calendar everywhere (1,000 examples, sizes to 5,000,000); a longer period from the same day ends later. Also asserts that a pickled state without a set of shared keys gets none (a surviving mutant). Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 39 ++++++++--- tests/core/test_carry_over_order.py | 28 ++++++++ tests/core/test_carry_over_order_property.py | 73 ++++++++++++++++++++ tests/core/test_storage_input_index.py | 3 + 4 files changed, 134 insertions(+), 9 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ea13ab31c..fa3acfa2e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -103,15 +103,36 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: return defined[-1].value -def _end_order(period: Period) -> tuple: - """Sorts periods by when they end. A period that ends after the last - date ``datetime`` can represent (``day:9999-12-30:3``) has no ``stop``, - which raises ``OverflowError``; it sorts after every period that has - one.""" - try: - return (0, period.stop) - except OverflowError: - return (1,) +_DAYS_BEFORE_MONTH = (0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334) + + +def _end_order(period: Period) -> float: + """Sorts periods by when they end: the number of the day after the + period's last day, counting days as ``date.toordinal`` does. + + ``period.stop`` gives the same order wherever it has a value, but it + raises for a period that ends after 9999-12-31 (``day:9999-12-30:3``), + the last date ``datetime`` can represent. This is integer arithmetic on + the same (proleptic Gregorian) calendar, so every period has a value and + one that ends later always sorts later. + """ + unit, (year, month, day), size = period + if unit == ETERNITY: + return float("inf") + if unit == periods.DAY: + day += size + elif unit == MONTH: + year, month = divmod(year * 12 + month - 1 + size, 12) + month += 1 + else: + year += size + # The first day of that month, then ``day - 1`` days on: a day past the + # month's end runs into the next month, as it does in ``period.stop``. + leap = year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) + days_before_year = ( + (year - 1) * 365 + (year - 1) // 4 - (year - 1) // 100 + (year - 1) // 400 + ) + return days_before_year + _DAYS_BEFORE_MONTH[month - 1] + (leap and month > 2) + day if TYPE_CHECKING: diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 87e10fff5..0fc88939f 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -155,6 +155,34 @@ def test_an_input_ending_after_the_last_date_ends_last(system, longer_first): ) +@pytest.mark.parametrize("longer_first", [True, False], ids=["longer", "shorter"]) +def test_of_two_inputs_ending_after_the_last_date_the_later_carries( + system, longer_first +): + """Both end after 9999-12-31. As strings, ``:10`` sorts before ``:9``.""" + values = {"day:9999-12-30:10": [10, 10], "day:9999-12-30:9": [9, 9]} + order = list(values) if longer_first else list(values)[::-1] + inputs = {"day_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", "9999-12-31"), [10, 10] + ) + + +@pytest.mark.parametrize("months_first", [True, False], ids=["months", "days"]) +def test_inputs_in_two_units_ending_after_the_last_date_carry_the_later( + system, months_first +): + """800 days from 1 February 9999 end after 24 months from the same day. + The larger unit and the string order both favour the months, so only + the end picks the days.""" + values = {"day:9999-02-01:800": [8, 8], "month:9999-02:24": [24, 24]} + order = list(values)[::-1] if months_first 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", "year:9999-02"), [8, 8] + ) + + 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/core/test_carry_over_order_property.py b/tests/core/test_carry_over_order_property.py index d343d830d..23111e262 100644 --- a/tests/core/test_carry_over_order_property.py +++ b/tests/core/test_carry_over_order_property.py @@ -254,3 +254,76 @@ def test_carry_over_depends_only_on_the_inputs(scenario): np.testing.assert_array_equal( result, reference(SYSTEM, inputs, variable, period), err_msg=message ) + + +# ----- The order of period ends ----- # + +from datetime import date # noqa: E402 + +from policyengine_core.simulations.simulation import _end_order # noqa: E402 + +_starts = st.dates(min_value=date(1, 1, 1), max_value=date(9999, 12, 31)) +_units = st.sampled_from([periods.DAY, periods.MONTH, periods.YEAR]) + + +def _period(unit, start, size): + return periods.Period( + (unit, periods.Instant((start.year, start.month, start.day)), size) + ) + + +def _numpy_day_after(period): + """The day after the period's last day, on numpy's calendar, which runs + past year 9999: an implementation that shares nothing with + ``_end_order`` or ``Period.stop``.""" + unit, (year, month, day), size = period + if unit == periods.DAY: + return np.datetime64(f"{year:04d}-{month:02d}-{day:02d}") + np.timedelta64( + size, "D" + ) + months = size if unit == periods.MONTH else 12 * size + first = np.datetime64(f"{year:04d}-{month:02d}") + np.timedelta64(months, "M") + return first.astype("datetime64[D]") + np.timedelta64(day - 1, "D") + + +@settings(max_examples=1000, deadline=None, derandomize=True) +@given(unit=_units, start=_starts, size=st.integers(min_value=1, max_value=5_000_000)) +@example(unit=periods.DAY, start=date(9999, 12, 30), size=3) +@example(unit=periods.MONTH, start=date(2012, 1, 31), size=1) +@example(unit=periods.YEAR, start=date(2012, 2, 29), size=1) +@example(unit=periods.MONTH, start=date(9999, 2, 1), size=24) +def test_end_order_is_the_day_after_the_last_day(unit, start, size): + """Against two other implementations: ``Period.stop`` where it has a + value, and numpy's calendar everywhere.""" + period = _period(unit, start, size) + day_one = np.datetime64("0001-01-01") + expected = int((_numpy_day_after(period) - day_one) / np.timedelta64(1, "D")) + 1 + assert _end_order(period) == expected + try: + stop = date(*period.stop) + except (OverflowError, ValueError): + # Ends after 9999-12-31, the last date ``datetime`` has: ``stop`` + # raises, or gives a year no ``date`` can hold. + assert expected > date.max.toordinal() + 1 + else: + assert _end_order(period) == stop.toordinal() + 1 + + +@settings(max_examples=500, deadline=None, derandomize=True) +@given( + unit=_units, + start=_starts, + size=st.integers(min_value=1, max_value=5_000_000), + more=st.integers(min_value=1, max_value=1000), +) +def test_a_longer_period_from_the_same_day_ends_later(unit, start, size, more): + assert _end_order(_period(unit, start, size + more)) > _end_order( + _period(unit, start, size) + ) + + +def test_eternity_ends_after_every_other_period(): + eternity = periods.period(periods.ETERNITY) + assert _end_order(eternity) > _end_order( + _period(periods.YEAR, date(9999, 1, 1), 5_000_000) + ) diff --git a/tests/core/test_storage_input_index.py b/tests/core/test_storage_input_index.py index b5ecb3c39..ffbd75bcd 100644 --- a/tests/core/test_storage_input_index.py +++ b/tests/core/test_storage_input_index.py @@ -213,11 +213,14 @@ def test_values_from_an_older_pickle_count_as_inputs(): assert not restored.is_derived("2017-01") assert restored.is_derived("2017-02") assert "_derived" not in vars(restored) + # A state with no set of shared keys gets none. + assert "_shared" not in vars(restored) # An older storage with nothing stored needs no set. restored = InMemoryStorage.__new__(InMemoryStorage) restored.__setstate__(_older_state(InMemoryStorage(is_eternal=False))) assert restored._inputs is _NO_INPUTS + assert restored._shared is _NOTHING_SHARED restored.put(np.array([2.0]), "2017-03", derived=True) assert restored.is_derived("2017-03")