Uprate only from inputs, so uprated values don't depend on calculation order

changelog.d/fix-disk-storage-clone-writes.fixed.md

@@ -0,0 +1 @@
+Simulations and storages cloned from one another no longer write over each other's values in on-disk storage: a value stored for a key whose file a clone also reads goes to a new file of its own. Each cloned simulation also keeps its own record of user inputs, so `apply_reform` on one no longer keeps, as an input, a value it calculated for a period the other was given as input.

changelog.d/fix-uprating-order.fixed.md

@@ -0,0 +1 @@
+Uprating now starts from the latest earlier input in the variable's own unit and skips periods the simulation calculated, and a yearly input given through a `set_input` helper now replaces months already calculated instead of keeping them, so calculating intermediate periods first no longer compounds rounding or truncation, or carries an eligibility mask, a default or a calculated month into later uprated values.

policyengine_core/data_storage/on_disk_storage.py

@@ -1,5 +1,6 @@
 import os
 import shutil
+import tempfile
 
 import numpy
 from numpy.typing import ArrayLike
@@ -8,6 +9,12 @@
 from policyengine_core.enums import EnumArray
 from policyengine_core.periods import Period
 
+# Subdirectory of a storage directory for the files ``put`` writes when a
+# key's usual file is shared with a clone (see ``OnDiskStorage._path_to_write``).
+# ``restore`` reads only the files directly in the directory, so it never
+# reads these.
+REPLACEMENTS_DIR = "replaced"
+
 
 class OnDiskStorage:
     """
@@ -25,25 +32,44 @@ def __init__(
         # File keys stored with ``put(..., derived=True)``; see
         # ``InMemoryStorage``.
         self._derived = set()
+        # For each key, the file this storage last wrote for it, while it has
+        # not shared that file since: the only files of its family ``put``
+        # writes over (see ``_path_to_write``). Kept when the key is deleted,
+        # so that writing it again reuses the file.
+        self._own_paths = {}
+        # Paths of every file this storage, the storage it was cloned from or
+        # any other clone of either wrote or stored when cloned: one set,
+        # shared by all of them.
+        self._family_files = set()
         self.is_eternal = is_eternal
         self.preserve_storage_dir = preserve_storage_dir
         self.storage_dir = storage_dir
 
+    def __getstate__(self) -> dict:
+        # Whatever this state is read into reads the same files, so from now
+        # on this storage writes over none of them, as after ``clone``.
+        self._own_paths = {}
+        return self.__dict__.copy()
+
     def __setstate__(self, state: dict) -> None:
         # A storage pickled before derived marks existed has none: its values
         # count as inputs.
         state.setdefault("_derived", set())
+        # Nor did it record the files it wrote: it writes over none it stores.
+        state.setdefault("_own_paths", {})
+        if "_family_files" not in state:
+            state["_family_files"] = set(state.get("_files", {}).values())
         self.__dict__.update(state)
 
     def clone(self) -> "OnDiskStorage":
         """Create a private metadata view over this storage directory.
 
         The file and enum mappings are copied so deleting or rewiring entries
         through the clone does not mutate the source storage. The underlying
-        ``.npy`` files remain shared: writing the same ``{branch}_{period}``
-        key from two views targets the same path and can overwrite the file.
-        Clones retain the original cleanup owner so the shared directory stays
-        alive, but never own cleanup themselves.
+        ``.npy`` files are shared, so neither storage writes over them: a later
+        ``put`` of one of their keys, through either storage, writes a new
+        file (see ``_path_to_write``). Clones retain the original cleanup owner
+        so the shared directory stays alive, but never own cleanup themselves.
         """
         clone = OnDiskStorage(
             self.storage_dir,
@@ -54,6 +80,11 @@ def clone(self) -> "OnDiskStorage":
         clone._enums = self._enums.copy()
         clone._derived = set(self._derived)
         clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self)
+        # Both storages now read every file stored so far, including any this
+        # family did not write (read back by ``restore``, say).
+        self._family_files.update(self._files.values())
+        clone._family_files = self._family_files
+        self._own_paths = {}
         return clone
 
     def _decode_file(self, file: str) -> ArrayLike:
@@ -102,17 +133,45 @@ def put(
         period = periods.period(period)
 
         filename = f"{branch_name}_{period}"
-        path = os.path.join(self.storage_dir, filename) + ".npy"
+        path = self._path_to_write(filename)
         if isinstance(value, EnumArray):
             self._enums[path] = value.possible_values
             value = value.view(numpy.ndarray)
         numpy.save(path, value)
         self._files[filename] = path
+        self._own_paths[filename] = path
+        self._family_files.add(path)
         if derived:
             self._derived.add(filename)
         else:
             self._derived.discard(filename)
 
+    def _path_to_write(self, filename: str) -> str:
+        """The path ``put`` writes the value for key ``filename`` to.
+
+        Storages cloned from one another share a directory and the files
+        stored before cloning, and each can store the same key afterwards,
+        which names the same path. Writing over a file another of them reads
+        would change that storage's value. So ``put`` writes over one of
+        their files only if this storage wrote it and has not shared it since;
+        otherwise the value goes to a new file of its own, in
+        ``REPLACEMENTS_DIR``. A file only storages outside this family wrote
+        is written over, as before.
+        """
+        own = self._own_paths.get(filename)
+        if own is not None:
+            return own
+        path = os.path.join(self.storage_dir, filename) + ".npy"
+        if path not in self._family_files:
+            return path
+        directory = os.path.join(self.storage_dir, REPLACEMENTS_DIR)
+        os.makedirs(directory, exist_ok=True)
+        descriptor, path = tempfile.mkstemp(
+            prefix=f"{filename}.", suffix=".npy", dir=directory
+        )
+        os.close(descriptor)
+        return path
+
     def delete(self, period: Period = None, branch_name: str = "default") -> None:
         if period is None:
             # Only wipe files belonging to the requested branch (previously

policyengine_core/holders/helpers.py

@@ -22,6 +22,31 @@ def get_stored_array(holder: Holder, period: Period, branch_name: str) -> ArrayL
     return holder._get_array_from_storage(period, branch_name)
 
 
+def get_stored_input(holder: Holder, period: Period, branch_name: str) -> ArrayLike:
+    """The input stored for ``period`` under ``branch_name`` itself, or
+    ``None``.
+
+    A value the simulation calculated and cached there (see
+    ``Holder.put_in_cache``) is not an input, so the helpers below treat the
+    sub-period as empty and the new input replaces it, as an input set for
+    that sub-period alone would. Treating it as an input kept the calculated
+    value under the input, copied it into the later sub-periods (dispatch),
+    or took it out of the total (divide), so the result depended on what had
+    been calculated before the input was set.
+    """
+    if holder._memory_storage.has(period, branch_name):
+        storage = holder._memory_storage
+    elif holder._disk_storage is not None and holder._disk_storage.has(
+        period, branch_name
+    ):
+        storage = holder._disk_storage
+    else:
+        return None
+    if storage.is_derived(period, branch_name):
+        return None
+    return storage.get(period, branch_name)
+
+
 def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLike):
     """
     This function can be declared as a ``set_input`` attribute of a variable.
@@ -46,11 +71,11 @@ def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLik
 
     after_instant = period.start.offset(period_size, period_unit)
 
-    # Cache the input data, skipping the existing cached months
+    # Cache the input data, skipping the months that already have an input
     branch_name = get_input_branch(holder)
     sub_period = period.start.period(cached_period_unit)
     while sub_period.start < after_instant:
-        existing_array = get_stored_array(holder, sub_period, branch_name)
+        existing_array = get_stored_input(holder, sub_period, branch_name)
         if existing_array is None:
             holder._set(sub_period, array, branch_name)
         else:
@@ -91,7 +116,7 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike)
     sub_period = period.start.period(cached_period_unit)
     sub_periods_count = 0
     while sub_period.start < after_instant:
-        existing_array = get_stored_array(holder, sub_period, branch_name)
+        existing_array = get_stored_input(holder, sub_period, branch_name)
         if existing_array is not None:
             remaining_array -= existing_array
         else:
@@ -103,7 +128,7 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike)
         divided_array = remaining_array / sub_periods_count
         sub_period = period.start.period(cached_period_unit)
         while sub_period.start < after_instant:
-            if get_stored_array(holder, sub_period, branch_name) is None:
+            if get_stored_input(holder, sub_period, branch_name) is None:
                 holder._set(sub_period, divided_array, branch_name)
             sub_period = sub_period.offset(1)
     elif not (remaining_array == 0).all():

policyengine_core/simulations/simulation.py

@@ -935,13 +935,31 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
                     if known_period.unit == variable.definition_period
                     and known_period.start < period.start
                 ]
-                if variable.uprating is not None and len(earlier_known_periods) > 0:
+                # Uprate only from an input this branch reads (see
+                # ``Holder.get_input_periods``). Every value this simulation
+                # calculated is marked derived when cached, and uprating from
+                # one would make the result depend on which periods were
+                # calculated first: an integer truncated, or a float32
+                # rounded, at an intermediate period would compound, and a
+                # value masked by ``defined_for``, a default, or a value
+                # carried from another unit, cached there, would replace the
+                # input. A period stored only under a branch this one cannot
+                # read would read back as ``None``.
+                earlier_input_periods = []
+                if variable.uprating is not None:
+                    input_periods = set(holder.get_input_periods(self.branch_name))
+                    earlier_input_periods = [
+                        known_period
+                        for known_period in earlier_known_periods
+                        if known_period in input_periods
+                    ]
+                if earlier_input_periods:
                     # Take the latest period from the filtered list itself.
                     # Indexing ``known_periods`` with a position in the
                     # filtered list picked the wrong period whenever a later
                     # one was stored first.
                     latest_known_period = max(
-                        earlier_known_periods, key=lambda p: p.start
+                        earlier_input_periods, key=lambda p: p.start
                     )
                     try:
                         uprating_parameter = get_parameter(
@@ -1028,8 +1046,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
                         # No input to carry, but a later period is stored: as
                         # before, return the default without caching it. A
                         # cached default would change what a formula testing
-                        # whether a value is stored sees, and the uprating path
-                        # above would uprate later periods from it.
+                        # whether a value is stored sees. (The uprating path
+                        # above skips derived periods, so it would not uprate
+                        # from one.)
                         return holder.default_array()
                     else:
                         array = holder.default_array()
@@ -1525,6 +1544,21 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non
         _fast_cache = getattr(self, "_fast_cache", None)
         if _fast_cache is not None:
             _fast_cache.pop((variable_name, period), None)
+            if variable.set_input and period.unit != variable.definition_period:
+                # The helper wrote the input's sub-periods, replacing any
+                # value calculated there, so what ``calculate`` returned for
+                # them is stale too. (``_end_order``, not ``stop``: ``stop``
+                # raises for a period that ends after 9999-12-31.)
+                stale = [
+                    key
+                    for key in _fast_cache
+                    if key[0] == variable_name
+                    and isinstance(key[1], Period)
+                    and period.start <= key[1].start
+                    and _end_order(key[1]) <= _end_order(period)
+                ]
+                for key in stale:
+                    del _fast_cache[key]
 
     def get_variable_population(self, variable_name: str) -> Population:
         variable = self.tax_benefit_system.get_variable(
@@ -1588,6 +1622,11 @@ def clone(
             ):
                 new_dict[key] = value
         new._fast_cache = {}
+        # Each simulation records its own inputs. A shared record let an
+        # input set on one replay, in the other's ``apply_reform``, whatever
+        # the other had calculated for that period, as an input.
+        if hasattr(self, "_user_input_keys"):
+            new._user_input_keys = set(self._user_input_keys)
 
         # Only pass ``share_arrays`` when sharing, so a population or holder
         # ``clone`` override with the earlier signature still deep-copies.