Restore inputs as inputs, evict the fast cache on holder writes, reject non-numeric uprating and defined_for

changelog.d/holder-write-fast-cache.fixed.md

@@ -0,0 +1 @@
+Writing to or deleting from a holder directly (`Holder.set_input`, `Holder.put_in_cache`, `Holder.delete_arrays`) now drops the `Simulation.calculate` fast-cache entries it replaces, so `calculate` no longer returns the value the holder held before.

changelog.d/non-numeric-uprating-defined-for.changed.md

@@ -0,0 +1 @@
+`uprating` on an Enum, `str` or date variable, and a `defined_for` that names one, are now rejected with a `ValueError` when the variables are registered (or when `uprating` is assigned), so a tax-benefit system containing either no longer loads; before, it loaded and raised `TypeError` when the variable was uprated or calculated, or, for a group variable `defined_for` a person Enum, masked on the sum of its members' Enum indices. `defined_for` must name a `bool`, `int` or `float` variable.

changelog.d/replace-variable-rollback.fixed.md

@@ -0,0 +1 @@
+`TaxBenefitSystem.replace_variable` now keeps the existing variable when the replacement is rejected, instead of leaving the system without it.

changelog.d/restore-simulation-input-record.fixed.md

@@ -0,0 +1 @@
+`restore_simulation` now restores as inputs the values the dumped simulation stored as inputs (`dump_simulation` lists their periods in an `inputs.txt` beside each variable's arrays), so `apply_reform` on a restored simulation no longer drops them; a dump written before this has no such list, so every value in it is restored as an input, with a warning.

policyengine_core/holders/holder.py

@@ -105,6 +105,7 @@ def delete_arrays(
         self._memory_storage.delete(period, branch_name)
         if self._disk_storage:
             self._disk_storage.delete(period, branch_name)
+        self._evict_fast_cache(period, branch_name, contained=True)
 
     def _get_array_from_storage(
         self, period: Period, branch_name: str = "default"
@@ -370,6 +371,7 @@ def _set(
             self._disk_storage.put(value, period, branch_name)
         else:
             self._memory_storage.put(value, period, branch_name)
+        self._evict_fast_cache(period, branch_name)
         if user_input_contexts:
             if not hasattr(simulation, "_user_input_keys"):
                 simulation._user_input_keys = set()
@@ -396,3 +398,52 @@ def default_array(self) -> ArrayLike:
         """
 
         return self.variable.default_array(self.population.count)
+
+    def _evict_fast_cache(
+        self, period: Period, branch_name: str, contained: bool = False
+    ) -> None:
+        """Drop the simulation's ``_fast_cache`` entries a storage write or delete makes stale.
+
+        ``Simulation.calculate`` answers a repeated request from
+        ``_fast_cache``, keyed by ``(variable name, requested period)``,
+        before it reads this holder. So every write into, or delete from,
+        this holder's storage drops the entries for the periods it changes;
+        otherwise ``calculate`` kept returning the value the storage no
+        longer held (for example after ``holder.set_input``).
+
+        The fast cache belongs to this holder's simulation: a branch has its
+        own holders and its own fast cache, and keeps the values it started
+        with, so nothing outside this simulation is touched. Nor is anything
+        here when ``branch_name`` is a branch this simulation does not read.
+
+        A write changes one storage key: ``period``, or, for an ETERNITY
+        variable, the one value every period reads. A delete (``contained``)
+        removes every period ``period`` contains, or every period when
+        ``period`` is ``None``.
+        """
+        simulation = self.simulation
+        fast_cache = getattr(simulation, "_fast_cache", None)
+        if not fast_cache:
+            return
+        name = self.variable.name
+        drop_all = period is None or self.variable.definition_period == periods.ETERNITY
+        if not drop_all:
+            period = periods.period(period)
+            if not contained and (name, period) not in fast_cache:
+                return
+        visible_branch_names = getattr(simulation, "_get_visible_branch_names", None)
+        if visible_branch_names is not None and branch_name not in (
+            visible_branch_names()
+        ):
+            return
+        if not drop_all and not contained:
+            del fast_cache[(name, period)]
+            return
+        stale_keys = [
+            key
+            for key in fast_cache
+            if key[0] == name
+            and (drop_all or not isinstance(key[1], Period) or period.contains(key[1]))
+        ]
+        for key in stale_keys:
+            del fast_cache[key]

policyengine_core/simulations/simulation.py

@@ -874,6 +874,17 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
         self._check_period_consistency(period, variable)
 
         if variable.defined_for is not None:
+            # Registration rejects a non-numeric defined_for variable (see
+            # ``TaxBenefitSystem._check_defined_for``). This catches one set,
+            # or a variable replaced, afterwards. It reads the variable's
+            # type, not the values: mapped to a group entity, an Enum's
+            # indices are summed into numbers, and str or date values fail
+            # inside the mapping.
+            defined_for_variable = self.tax_benefit_system.get_variable(
+                variable.defined_for
+            )
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
             mask = (
                 self.calculate(variable.defined_for, period, map_to=variable.entity.key)
                 > 0
@@ -902,6 +913,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
                     and known_period.start < period.start
                 ]
                 if variable.uprating is not None and len(earlier_known_periods) > 0:
+                    # Registration rejects these; an ``uprating`` assigned
+                    # past the setter gets the same message here.
+                    variable.check_uprating_value_type()
                     # Take the latest period from the filtered list itself.
                     # Indexing ``known_periods`` with a position in the
                     # filtered list picked the wrong period whenever a later

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -51,6 +51,7 @@
 from policyengine_core.periods import Instant, Period
 from policyengine_core.populations import GroupPopulation, Population
 from policyengine_core.variables import Variable
+from policyengine_core.variables.variable import NUMERIC_VALUE_TYPES
 
 log = logging.getLogger(__name__)
 
@@ -92,6 +93,8 @@ class TaxBenefitSystem:
     """Short list of basic inputs to get medium accuracy."""
     modelled_policies: str = None
     """A YAML filepath containing metadata describing the modelled policies."""
+    _defined_for_checks_deferred: bool = False
+    """Whether ``load_variable`` leaves ``defined_for`` checks to a later pass over every variable."""
 
     def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None:
         if entities is None:
@@ -119,7 +122,15 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None:
         self.variable_module_metadata = {}
 
         if self.variables_dir is not None:
-            self.add_variables_from_directory(self.variables_dir)
+            # Check every variable once, after the whole directory is loaded,
+            # instead of scanning the variables loaded so far each time a
+            # non-numeric one is added.
+            self._defined_for_checks_deferred = True
+            try:
+                self.add_variables_from_directory(self.variables_dir)
+            finally:
+                self._defined_for_checks_deferred = False
+            self._check_defined_for_variables()
         self.data_modified = False
 
         if self.parameters_dir is not None:
@@ -215,10 +226,44 @@ def load_variable(
             )
 
         variable = variable_class(baseline_variable=baseline_variable)
+        if not self._defined_for_checks_deferred:
+            self._check_defined_for(variable)
         self.variables[variable.name] = variable
 
         return variable
 
+    def _check_defined_for(self, variable: Variable) -> None:
+        """Check the ``defined_for`` links that ``variable`` takes part in.
+
+        Both directions: the variable ``variable`` is defined for, if it is
+        registered, and, when ``variable`` cannot be compared with zero, every
+        registered variable defined for it. See
+        :meth:`Variable.check_defined_for_variable`.
+        """
+        defined_for = variable.defined_for
+        if defined_for is not None:
+            defined_for_variable = (
+                variable
+                if defined_for == variable.name
+                else self.variables.get(defined_for)
+            )
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
+        if variable.value_type in NUMERIC_VALUE_TYPES:
+            return
+        for other in self.variables.values():
+            if other.defined_for == variable.name and other.name != variable.name:
+                other.check_defined_for_variable(variable)
+
+    def _check_defined_for_variables(self) -> None:
+        """Check every registered variable's ``defined_for`` variable."""
+        for variable in self.variables.values():
+            if variable.defined_for is None:
+                continue
+            defined_for_variable = self.variables.get(variable.defined_for)
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
+
     def add_variable(self, variable: Type[Variable]) -> Variable:
         """Adds an OpenFisca variable to the tax and benefit system.
 
@@ -243,9 +288,14 @@ def replace_variable(self, variable: Type[Variable]) -> None:
         :param Variable variable: New variable class to add. Must be a subclass of Variable.
         """
         name = variable.__name__
-        if self.variables.get(name) is not None:
-            del self.variables[name]
-        self.load_variable(variable, update=False)
+        replaced = self.variables.pop(name, None)
+        try:
+            self.load_variable(variable, update=False)
+        except Exception:
+            # A rejected replacement leaves the system as it was.
+            if replaced is not None:
+                self.variables[name] = replaced
+            raise
         self.data_modified = True
 
     def update_variable(self, variable: Type[Variable]) -> Variable:

policyengine_core/tools/simulation_dumper.py

@@ -2,13 +2,22 @@
 
 
 import os
+import warnings
 
 import numpy as np
 
 from policyengine_core.data_storage import OnDiskStorage
+from policyengine_core import periods
 from policyengine_core.periods import ETERNITY
 from policyengine_core.simulations import Simulation
 
+# Next to each variable's arrays: the periods, one per line, whose dumped
+# value was an input (stored through ``set_input``). ``restore_simulation``
+# registers exactly these as inputs, so ``apply_reform``, which keeps inputs
+# and drops calculated values, keeps the same values in the restored
+# simulation as in the dumped one.
+INPUT_PERIODS_FILE = "inputs.txt"
+
 
 def dump_simulation(simulation, directory):
     """
@@ -26,18 +35,27 @@ def dump_simulation(simulation, directory):
     entities_dump_dir = os.path.join(directory, "__entities__")
     os.mkdir(entities_dump_dir)
 
+    input_keys = _input_storage_keys(simulation)
     for entity in simulation.populations.values():
         # Dump entity structure
         _dump_entity(entity, entities_dump_dir)
 
         # Dump variable values
         for holder in entity._holders.values():
-            _dump_holder(holder, directory)
+            _dump_holder(holder, directory, input_keys)
 
 
 def restore_simulation(directory, tax_benefit_system, **kwargs):
     """
     Restore simulation from directory
+
+    Values the dumped simulation stored as inputs are restored as inputs
+    (recorded in ``_user_input_keys``, as ``set_input`` records them), and
+    every other value as a calculated one, so ``apply_reform`` keeps and drops
+    the same values it would have in the dumped simulation. A dump written
+    before inputs were recorded (no ``inputs.txt``) does not say which values
+    were inputs, so every value in it is restored as an input, with a
+    warning: ``apply_reform`` then keeps its calculated values as dumped.
     """
     simulation = Simulation(
         tax_benefit_system, tax_benefit_system.instantiate_entities()
@@ -58,17 +76,58 @@ def restore_simulation(directory, tax_benefit_system, **kwargs):
     variables_to_restore = (
         variable for variable in os.listdir(directory) if variable != "__entities__"
     )
-    for variable in variables_to_restore:
-        _restore_holder(simulation, variable, directory)
+    without_input_record = [
+        variable
+        for variable in variables_to_restore
+        if not _restore_holder(simulation, variable, directory)
+    ]
+    if without_input_record:
+        warnings.warn(
+            f"The simulation dump in {directory} does not record which values "
+            f"were inputs ({len(without_input_record)} variables have no "
+            f"{INPUT_PERIODS_FILE}; it was written by an earlier version of "
+            "policyengine-core). Every value in it is restored as an input, "
+            "so apply_reform keeps the calculated values as dumped instead of "
+            "recalculating them. Dump the simulation again to record its "
+            "inputs.",
+            stacklevel=2,
+        )
 
     return simulation
 
 
-def _dump_holder(holder, directory):
+def _dump_holder(holder, directory, input_keys=frozenset()):
     disk_storage = holder.create_disk_storage(directory, preserve=True)
+    input_periods = []
     for period in holder.get_known_periods():
         value = holder.get_array(period)
         disk_storage.put(value, period)
+        # The input record of exactly the value dumped: ``get_array`` above
+        # reads the default branch.
+        if (holder.variable.name, "default", str(period)) in input_keys:
+            input_periods.append(str(period))
+    path = os.path.join(disk_storage.storage_dir, INPUT_PERIODS_FILE)
+    with open(path, "w") as file:
+        file.write("".join(f"{period}\n" for period in dict.fromkeys(input_periods)))
+
+
+def _input_storage_keys(simulation):
+    """The storage keys ``_user_input_keys`` records as inputs.
+
+    Each record entry becomes ``(variable, branch, period)`` with the period
+    as storage writes it, as ``Simulation._invalidate_all_caches`` reads the
+    record back through the storage: an ETERNITY variable's one value is an
+    input whatever period its entry names.
+    """
+    input_keys = set()
+    for name, branch_name, period in getattr(simulation, "_user_input_keys", ()):
+        variable = simulation.tax_benefit_system.get_variable(name)
+        if variable is not None and variable.definition_period == ETERNITY:
+            period = ETERNITY
+        elif period is None:
+            continue
+        input_keys.add((name, branch_name, str(periods.period(period))))
+    return input_keys
 
 
 def _dump_entity(population, directory):
@@ -123,6 +182,7 @@ def _restore_entity(population, directory):
 
 
 def _restore_holder(simulation, variable, directory):
+    """Restore one variable's values; return whether its inputs were recorded."""
     storage_dir = os.path.join(directory, variable)
     is_variable_eternal = (
         simulation.tax_benefit_system.get_variable(variable).definition_period
@@ -135,6 +195,34 @@ def _restore_holder(simulation, variable, directory):
 
     holder = simulation.get_holder(variable)
 
+    input_periods_path = os.path.join(storage_dir, INPUT_PERIODS_FILE)
+    if os.path.exists(input_periods_path):
+        with open(input_periods_path) as file:
+            input_periods = set(file.read().split())
+    else:
+        # Dumped before inputs were recorded: nothing says which values were
+        # calculated, so keep every value as an input.
+        input_periods = None
+
     for period in disk_storage.get_known_periods():
         value = disk_storage.get(period)
-        holder.put_in_cache(value, period)
+        if input_periods is None or str(period) in input_periods:
+            _restore_input(simulation, holder, period, value)
+        else:
+            holder.put_in_cache(value, period)
+    return input_periods is not None
+
+
+def _restore_input(simulation, holder, period, value):
+    """Store ``value`` as an input, recorded as ``set_input`` records one.
+
+    ``Holder.set_input`` would also run the variable's ``set_input`` helper,
+    but a dump holds values already split into the variable's own periods.
+    """
+    if not hasattr(simulation, "_user_input_contexts"):
+        simulation._user_input_contexts = []
+    simulation._user_input_contexts.append(simulation.branch_name)
+    try:
+        holder._set(period, value, simulation.branch_name)
+    finally:
+        simulation._user_input_contexts.pop()