diff --git a/changelog.d/fix-stock-option-caches.fixed.md b/changelog.d/fix-stock-option-caches.fixed.md new file mode 100644 index 000000000..3c6d80e77 --- /dev/null +++ b/changelog.d/fix-stock-option-caches.fixed.md @@ -0,0 +1 @@ +Cache `calculate_add` and `calculate_divide` results as derived only where a plain read returns them without changing dtype, preserving visible inputs and allowing explicit aggregate refresh after an input changes or is deleted. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 0b56d00d1..96d05624e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1342,12 +1342,7 @@ def calculate_add( result = sum( self.calculate(variable_name, sub_period) for sub_period in sub_periods ) - # 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) + self._cache_option_result(variable, period, result) return result def calculate_divide( @@ -1379,8 +1374,7 @@ def calculate_divide( if period.unit == periods.MONTH: 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, derived=True) + self._cache_option_result(variable, period, result) return result elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1391,6 +1385,55 @@ def calculate_divide( ) ) + def _cache_option_result( + self, variable: Variable, period: Period, result: ArrayLike + ) -> None: + """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. + + A value cached at ``period`` is what every later ``calculate`` of the + variable at ``period`` returns. ``_calculate`` computes a FLOW + variable over a period of another unit with these same options (a + monthly variable over a year with ``calculate_add``, a yearly one over + a month with ``calculate_divide``), so their result is the plain value + there and is cached. Anywhere else it is not: + + - A STOCK variable's plain value over a year is its last month's, and + over a month the year's, not the sum or the twelfth. + - Over several periods of the variable's own unit, a plain read + raises instead. + - A day variable's plain read over a month or a year does not sum. + - Over a single period of its own unit, the sum is the value + ``calculate`` has already stored. + + Caching there would make a later plain read depend on whether the + option ran first. Nor is an option result cached when storing it would + change it (the twelfth of an integer or a count of true months is + stored as the variable's own type, so a later read would return the + truncated value where the first returned the exact one). + + The result is cached as derived, so auto-carry-over and uprating + never take it for an input, and ``put_in_cache`` keeps an input a + plain read finds at ``period`` (in this branch, an ancestor or + ``default``) instead of storing it. It does replace a value + calculated there before, which may predate a change to the inputs. + Whether the value read is an input comes from the storage's mark for + the first visible value, which every write sets. The simulation's + supplied-input record follows inputs across storage tiers for replay + and export; it can name a disk input hidden by a derived memory value, + and raw holder inputs need not be registered there. Accepted cache + writes never register the aggregate as a supplied input. + """ + if variable.quantity_type == QuantityType.STOCK: + return + routed = (variable.definition_period == MONTH and period.unit == YEAR) or ( + variable.definition_period == YEAR and period.unit == MONTH + ) + if not routed or np.asarray(result).dtype != variable.dtype: + return + self.get_holder(variable.name).put_in_cache( + result, period, self.branch_name, derived=True + ) + def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ Calculate the value of a variable using the ``calculate_output`` attribute of the variable. diff --git a/tests/core/test_option_result_cache.py b/tests/core/test_option_result_cache.py new file mode 100644 index 000000000..921e7ba44 --- /dev/null +++ b/tests/core/test_option_result_cache.py @@ -0,0 +1,268 @@ +"""ADD and DIVIDE results are cached only where a plain read returns them. + +``calculate_add`` and ``calculate_divide`` used to store their result at the +requested period, where every later plain ``calculate`` of the variable at +that period found it. That is right for a FLOW variable, whose plain value +over a period of another unit is exactly that sum or twelfth, and wrong +elsewhere: a STOCK variable's plain value over a year is its last month's, +and over a month the year's. A plain read then returned 120 or 1 instead of +10 or 12, depending only on whether the option had run first. +""" + +import pytest + +from policyengine_core import periods +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.option_caches import ( + DAY, + FLOW, + MONTH, + PROBE, + STOCK, + YEAR, + build, + make_probe, + run, +) + + +def fresh_and_after(probe, inputs, target, prior, carry_over=False): + """``target`` in a new simulation, and after running ``prior`` first.""" + fresh = build(probe, inputs, carry_over) + used = build(probe, inputs, carry_over) + prior_result = run(used, *prior) + return run(fresh, *target), run(used, *target), prior_result + + +def test_add_over_a_monthly_stock_leaves_its_annual_value_alone(): + probe = make_probe(MONTH, STOCK) + fresh, after, prior = fresh_and_after( + probe, + {"2012-01": 10}, + ("calculate", "2012"), + ("add", "2012"), + carry_over=True, + ) + assert prior == [120.0] + assert fresh == after == [10.0] + + +def test_divide_over_a_yearly_stock_leaves_its_monthly_value_alone(): + probe = make_probe(YEAR, STOCK, value_type=int) + fresh, after, prior = fresh_and_after( + probe, + {"2012": 12}, + ("calculate", "2012-06"), + ("divide", "2012-06"), + ) + assert prior == [1.0] + assert fresh == after == [12.0] + + +def test_stock_with_a_formula_keeps_its_last_month_value(): + probe = make_probe(MONTH, STOCK, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012"), ("add", "2012") + ) + assert prior == [sum(1200.0 + month for month in range(1, 13))] + assert fresh == after == [1212.0] + + +def test_flow_add_is_still_cached_as_the_plain_annual_value(): + probe = make_probe(MONTH, FLOW, with_formula=True) + simulation = build(probe) + total = simulation.calculate_add(PROBE, "2012") + holder = simulation.get_holder(PROBE) + assert holder.get_array("2012").tolist() == total.tolist() + assert holder.is_derived(periods.period("2012")) + assert periods.period("2012") not in holder.get_input_periods() + assert simulation.calculate(PROBE, "2012").tolist() == total.tolist() + + +def test_flow_divide_is_still_cached_as_the_plain_monthly_value(): + probe = make_probe(YEAR, FLOW, with_formula=True) + simulation = build(probe) + twelfth = simulation.calculate_divide(PROBE, "2012-03") + holder = simulation.get_holder(PROBE) + assert holder.get_array("2012-03").tolist() == twelfth.tolist() + assert holder.is_derived(periods.period("2012-03")) + assert periods.period("2012-03") not in holder.get_input_periods() + assert twelfth.tolist() == [pytest.approx(1012.0 / 12)] + + +def test_add_does_not_replace_an_input_stored_at_the_year(): + # A monthly FLOW variable whose input is stored at the year itself. + probe = make_probe(MONTH, FLOW, with_formula=True, set_input=None) + fresh, after, prior = fresh_and_after( + probe, {"2012": 7}, ("calculate", "2012"), ("add", "2012") + ) + assert prior == [sum(1200.0 + month for month in range(1, 13))] + assert fresh == after == [7.0] + + +def test_divide_does_not_replace_an_input_stored_at_the_month(): + probe = make_probe(YEAR, FLOW, with_formula=True, set_input=None) + fresh, after, prior = fresh_and_after( + probe, {"2012-02": 24}, ("calculate", "2012-02"), ("divide", "2012-02") + ) + assert prior == [pytest.approx(1012.0 / 12)] + assert fresh == after == [24.0] + + +@pytest.mark.parametrize( + "definition_period, period", + [(YEAR, "year:2012:2"), (MONTH, "month:2012-01:3")], +) +def test_add_over_several_own_periods_does_not_make_a_plain_read_succeed( + definition_period, period +): + probe = make_probe(definition_period, FLOW, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", period), ("add", period) + ) + assert isinstance(prior, list) + assert fresh == after == "ValueError" + + +def test_add_over_a_day_variable_leaves_its_monthly_value_alone(): + # A plain read of a day variable over a month runs its formula for the + # month; it does not sum the days. + probe = make_probe(DAY, FLOW, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012-01"), ("add", "2012-01") + ) + assert prior == [31 * 1201.0] + assert fresh == after == [1201.0] + + +def test_stock_add_in_a_branch_leaves_the_branch_annual_value_alone(): + probe = make_probe(MONTH, STOCK, with_formula=True) + fresh = build(probe).get_branch("other") + used = build(probe).get_branch("other") + used.calculate_add(PROBE, "2012") + assert run(used, "calculate", "2012") == run(fresh, "calculate", "2012") + assert run(used, "calculate", "2012") == [1212.0] + + +def test_a_twelfth_of_an_integer_flow_is_not_cached_truncated(): + # Storing 1012 / 12 in an integer variable would truncate it to 84, so a + # second monthly read would differ from the first. + probe = make_probe(YEAR, FLOW, value_type=int, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012-01"), ("calculate", "2012-01") + ) + assert fresh == after == prior == [pytest.approx(1012 / 12)] + + +@pytest.mark.parametrize( + "definition_period, native, target, option, first, second", + [ + (MONTH, "2012-01", "2012", "add", [12.0], [24.0]), + (YEAR, "2012", "2012-01", "divide", [1.0], [2.0]), + ], +) +def test_an_option_refreshes_the_aggregate_it_cached_before( + definition_period, native, target, option, first, second +): + # The cached aggregate predates the new input; running the option again + # replaces it, as on master, so the plain read follows the input. + simulation = build(make_probe(definition_period, FLOW)) + simulation.set_input(PROBE, native, [12]) + assert run(simulation, option, target) == first + simulation.set_input(PROBE, native, [24]) + assert run(simulation, option, target) == second + assert run(simulation, "calculate", target) == second + + +@pytest.mark.parametrize( + "definition_period, target, option, result", + [ + (MONTH, "2012", "add", [sum(1200.0 + month for month in range(1, 13))]), + (YEAR, "2012-02", "divide", [pytest.approx(1012.0 / 12)]), + ], +) +@pytest.mark.parametrize("nested", [False, True]) +def test_an_option_in_a_branch_does_not_replace_a_parent_input( + definition_period, target, option, result, nested +): + probe = make_probe(definition_period, FLOW, with_formula=True, set_input=None) + simulation = build(probe, {target: 7}) + if nested: + simulation = simulation.get_branch("parent") + # A raw holder write has an input mark, without a set_input record. + # The nearest ancestor must take precedence over default's input. + simulation.get_holder(PROBE).put_in_cache( + [9], periods.period(target), simulation.branch_name + ) + branch = simulation.get_branch("other") + assert run(branch, option, target) == result + assert run(branch, "calculate", target) == [9.0 if nested else 7.0] + assert not branch.get_holder(PROBE).is_derived( + periods.period(target), branch.branch_name + ) + + +@pytest.mark.parametrize( + "definition_period, deleted, native, target, option, first, second", + [ + (MONTH, "2012", "2012-01", "2012", "add", [12.0], [24.0]), + (YEAR, "2012-01", "2012", "2012-01", "divide", [1.0], [2.0]), + ], +) +def test_an_option_refreshes_its_aggregate_where_a_deleted_input_was( + definition_period, deleted, native, target, option, first, second +): + # Deleting the target input also removes its supplied-input record. The + # aggregate subsequently cached there is derived, so running the option + # again replaces it with the result from the current native inputs. + probe = make_probe(definition_period, FLOW, set_input=None) + simulation = build(probe, {deleted: 7}) + simulation.delete_arrays(PROBE, deleted) + simulation.set_input(PROBE, native, [12]) + assert run(simulation, option, target) == first + simulation.set_input(PROBE, native, [24]) + assert run(simulation, option, target) == second + assert run(simulation, "calculate", target) == second + + +@pytest.mark.parametrize( + "definition_period, stored, option, result", + [ + (MONTH, "2012", "add", [sum(1200.0 + month for month in range(1, 13))]), + (YEAR, "2012-02", "divide", [pytest.approx(1012.0 / 12)]), + ], +) +def test_an_option_does_not_replace_an_input_restored_from_a_dump( + definition_period, stored, option, result, tmp_path +): + # ``restore_simulation`` restores storage's input mark and registers the + # supplied input for replay. The option must preserve the visible input. + probe = make_probe(definition_period, FLOW, with_formula=True, set_input=None) + built = build(probe, {stored: 7}) + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), built.tax_benefit_system) + assert run(restored, option, stored) == result + assert run(restored, "calculate", stored) == [7.0] + assert not restored.get_holder(PROBE).is_derived(periods.period(stored)) + + +@pytest.mark.parametrize( + "definition_period, option, target, later", + [ + (MONTH, "add", "2012", "2013-01"), + (YEAR, "divide", "2012-06", "2013"), + ], +) +def test_an_option_result_does_not_carry_past_a_formula_end( + definition_period, option, target, later +): + probe = make_probe(definition_period, FLOW, with_formula=True) + probe.end = "2012-12-31" + fresh = build(probe, carry_over=True) + used = build(probe, carry_over=True) + run(used, option, target) + assert used.get_holder(PROBE).is_derived(periods.period(target)) + assert run(used, "calculate", later) == run(fresh, "calculate", later) == [0.0] diff --git a/tests/core/test_option_result_cache_property.py b/tests/core/test_option_result_cache_property.py new file mode 100644 index 000000000..d8a71fda8 --- /dev/null +++ b/tests/core/test_option_result_cache_property.py @@ -0,0 +1,136 @@ +"""Property: running ADD or DIVIDE first never changes a later result. + +For a monthly or yearly variable of any quantity type and value type, with +or without a formula, and with inputs stored at any periods, every request +(plain ``calculate``, ``calculate_add`` or ``calculate_divide`` at a single +or multi-period period) returns the same values, or raises the same kind of +error, whatever requests ran before it. The reference is the same request in +a new simulation with the same inputs. + +Auto-carry-over is off here: on its own, carrying over calculated values is +a separate order dependence (policyengine-core#562). +""" + +import pytest + +from policyengine_core import periods +from tests.fixtures.option_caches import ( + FLOW, + MONTH, + PROBE, + STOCK, + YEAR, + build, + make_probe, + run, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +OPERATIONS = ["calculate", "add", "divide"] +PERIODS = [ + "2012", + "2013", + "2012-01", + "2012-06", + "2012-12", + "2013-01", + "year:2012:2", + "month:2012-01:3", +] +INPUT_PERIODS = ["2012", "2013", "2012-01", "2012-06", "2012-12"] +REQUESTS = st.tuples(st.sampled_from(OPERATIONS), st.sampled_from(PERIODS)) + + +@st.composite +def scenarios(draw): + probe = dict( + definition_period=draw(st.sampled_from([MONTH, YEAR])), + quantity_type=draw(st.sampled_from([STOCK, FLOW])), + value_type=draw(st.sampled_from([float, int, bool])), + with_formula=draw(st.booleans()), + set_input=draw(st.sampled_from(["default", None])), + ) + inputs = draw( + st.dictionaries(st.sampled_from(INPUT_PERIODS), st.integers(0, 40), max_size=3) + ) + prior = draw(st.lists(REQUESTS, max_size=5)) + target = draw(REQUESTS) + return probe, inputs, prior, target + + +def build_or_none(probe, inputs): + try: + return build(probe, inputs) + except (TypeError, ValueError): + # Inputs this variable cannot take at those periods (spreading a + # boolean over months, for instance, subtracts booleans). + return None + + +@hypothesis.settings(max_examples=400, deadline=None) +@hypothesis.given(scenarios()) +def test_earlier_requests_do_not_change_a_result(scenario): + attributes, inputs, prior, target = scenario + probe = make_probe(**attributes) + fresh = build_or_none(probe, inputs) + hypothesis.assume(fresh is not None) + used = build(probe, inputs) + for request in prior: + run(used, *request) + assert run(used, *target) == run(fresh, *target) + + +@hypothesis.settings(max_examples=100, deadline=None) +@hypothesis.example( + definition_period=MONTH, + values=[12, 24], + deleted=True, + depth=2, + trace=False, +) +@hypothesis.example( + definition_period=YEAR, + values=[12, 24], + deleted=True, + depth=2, + trace=True, +) +@hypothesis.given( + definition_period=st.sampled_from([MONTH, YEAR]), + values=st.lists( + st.integers(-100, 100).map(lambda value: value * 12), min_size=2, max_size=5 + ), + deleted=st.booleans(), + depth=st.integers(0, 2), + trace=st.booleans(), +) +def test_explicit_refresh_matches_a_fresh_plain_read( + definition_period, values, deleted, depth, trace +): + """Refreshing a derived aggregate follows current inputs, even after deletion.""" + native, target, option = ( + ("2012-01", "2012", "add") + if definition_period == MONTH + else ("2012", "2012-01", "divide") + ) + probe = make_probe(definition_period, FLOW, set_input=None) + used = build(probe, {target: 7} if deleted else {}) + if deleted: + used.delete_arrays(PROBE, target) + used.trace = trace + for index in range(depth): + used = used.get_branch(f"branch_{index}") + for value in values: + used.set_input(PROBE, native, [value]) + fresh = build(probe, {native: value}) + expected = run(fresh, "calculate", target) + assert run(used, option, target) == expected + assert run(used, "calculate", target) == expected + assert used.get_holder(PROBE).is_derived( + periods.period(target), used.branch_name + ) + assert periods.period(target) not in used.get_holder(PROBE).get_input_periods( + used.branch_name + ) diff --git a/tests/core/test_option_result_input_provenance.py b/tests/core/test_option_result_input_provenance.py new file mode 100644 index 000000000..7464c6c8a --- /dev/null +++ b/tests/core/test_option_result_input_provenance.py @@ -0,0 +1,120 @@ +"""Option caches remain derived alongside supplied-input tier records.""" + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.experimental import MemoryConfig +from tests.fixtures.option_caches import ( + FLOW, + MONTH, + PROBE, + YEAR, + build, + make_probe, + run, +) + + +def _store_on_disk(simulation): + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder(PROBE) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + + +@pytest.mark.parametrize( + "definition_period, native, target, option, expected", + [ + (MONTH, "2012-01", "2012", "add", [12.0]), + (YEAR, "2012", "2012-01", "divide", [1.0]), + ], +) +@pytest.mark.parametrize("on_disk", [False, True]) +@pytest.mark.parametrize("deleted_target", [False, True]) +def test_option_cache_is_never_registered_exported_or_replayed_as_input( + definition_period, native, target, option, expected, on_disk, deleted_target +): + simulation = build(make_probe(definition_period, FLOW, set_input=None)) + if on_disk: + _store_on_disk(simulation) + holder = simulation.get_holder(PROBE) + target_period = periods.period(target) + native_period = periods.period(native) + target_key = (PROBE, "default", target_period) + if deleted_target: + simulation.set_input(PROBE, target, [7]) + simulation.delete_arrays(PROBE, target) + assert target_key not in simulation._user_input_keys + simulation.set_input(PROBE, native, [12]) + + assert run(simulation, option, target) == expected + + storage = holder._disk_storage if on_disk else holder._memory_storage + assert storage.get(target_period).tolist() == expected + assert storage.is_derived(target_period) + assert not holder._stores_user_input(target_period, "default") + assert target_key not in simulation._user_input_keys + assert (PROBE, "default", native_period) in simulation._user_input_keys + frame = simulation.to_input_dataframe() + assert { + column: frame[column].tolist() + for column in frame.columns + if column.startswith(f"{PROBE}__") + } == {f"{PROBE}__{native}": [12.0]} + assert simulation.to_input_dict()[PROBE] == {native: [12.0]} + + simulation._invalidate_all_caches() + + assert storage.get(target_period) is None + assert storage.get(native_period).tolist() == [12.0] + assert target_key not in simulation._user_input_keys + assert holder._stores_user_input(native_period, "default") + assert simulation.to_input_dataframe()[f"{PROBE}__{native}"].tolist() == [12.0] + assert run(simulation, "calculate", target) == expected + assert storage.is_derived(target_period) + + +@pytest.mark.parametrize( + "definition_period, native, target, option, expected", + [ + (MONTH, "2012-01", "2012", "add", [12.0]), + (YEAR, "2012", "2012-01", "divide", [1.0]), + ], +) +def test_option_refresh_preserves_a_supplied_disk_input_hidden_by_memory_cache( + definition_period, native, target, option, expected +): + simulation = build(make_probe(definition_period, FLOW, set_input=None)) + _store_on_disk(simulation) + simulation.set_input(PROBE, native, [12]) + simulation.set_input(PROBE, target, [7]) + holder = simulation.get_holder(PROBE) + target_period = periods.period(target) + input_path = holder._disk_storage._files[f"default_{target}"] + # Model an existing derived value in the higher-priority memory tier. + # The supplied disk value remains independently recorded for replay. + holder._memory_storage.put( + np.array([0.0], dtype=holder.variable.dtype), target_period, derived=True + ) + + assert run(simulation, option, target) == expected + assert run(simulation, "calculate", target) == expected + assert holder._memory_storage.get(target_period).tolist() == expected + assert holder._memory_storage.is_derived(target_period) + assert holder._disk_storage.get(target_period).tolist() == [7.0] + assert holder._user_input_storage[("default", target)] == ( + target_period, + frozenset({"disk"}), + ) + assert holder._stores_user_input(target_period, "default") + assert (PROBE, "default", target_period) in simulation._user_input_keys + + simulation._invalidate_all_caches() + + assert holder._memory_storage.get(target_period) is None + assert holder._disk_storage._files[f"default_{target}"] == input_path + assert not holder._disk_storage.is_derived(target_period) + assert run(simulation, "calculate", target) == [7.0] + assert holder._stores_user_input(target_period, "default") + assert (PROBE, "default", target_period) in simulation._user_input_keys diff --git a/tests/fixtures/option_caches.py b/tests/fixtures/option_caches.py new file mode 100644 index 000000000..b28de4004 --- /dev/null +++ b/tests/fixtures/option_caches.py @@ -0,0 +1,98 @@ +"""One-variable simulations for testing what ADD and DIVIDE results cache. + +Shared by ``tests/core/test_option_result_cache.py`` and its Hypothesis +property module, which pytest must be able to skip on its own when +Hypothesis is not installed. +""" + +from typing import Any, Dict, Optional + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import QuantityType, Variable + +PROBE = "probe" + + +def period_code(period) -> float: + """A value that differs for every month and year, for probe formulas.""" + start = period.start + if period.unit == periods.YEAR: + return 1000.0 + start.year % 100 + return float(100 * (start.year % 100) + start.month) + + +def make_probe( + definition_period: str, + quantity_type: str, + value_type: type = float, + with_formula: bool = False, + set_input: Optional[Any] = "default", +) -> type: + """A Person variable named ``probe`` with the given time behaviour. + + With ``with_formula``, its formula returns ``period_code(period)``, so + sums over months and twelfths of years are easy to tell apart from + single-period values. ``set_input=None`` stores inputs at the period + given instead of spreading them over the variable's own periods. + """ + attributes = dict( + value_type=value_type, + entity=Person, + definition_period=definition_period, + quantity_type=quantity_type, + label="Option cache probe", + ) + if set_input is None: + attributes["set_input"] = None + if with_formula: + + def formula(person, period): + return person.filled_array(period_code(period)) + + attributes["formula"] = formula + return type(PROBE, (Variable,), attributes) + + +def build( + probe: type, + inputs: Optional[Dict[str, Any]] = None, + carry_over: bool = False, +): + """One-person simulation of ``probe`` with ``inputs`` ({period: value}).""" + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = carry_over + system.add_variable(probe) + simulation = SimulationBuilder().build_default_simulation(system, count=1) + for period, value in (inputs or {}).items(): + simulation.set_input(PROBE, period, np.array([value])) + return simulation + + +def run(simulation, operation: str, period: str): + """Run ``calculate``, ``add`` or ``divide`` on ``probe``. + + Returns the value as a list, or the error's type name, so that results of + fresh and reused simulations compare equal exactly when both return the + same values or both raise the same kind of error. + """ + method = { + "calculate": simulation.calculate, + "add": simulation.calculate_add, + "divide": simulation.calculate_divide, + }[operation] + try: + return np.asarray(method(PROBE, period)).astype(float).tolist() + except ValueError as error: + return type(error).__name__ + + +STOCK = QuantityType.STOCK +FLOW = QuantityType.FLOW +MONTH = periods.MONTH +YEAR = periods.YEAR +DAY = periods.DAY