Count only inputs as already set when a set_input helper splits a longer period

changelog.d/set-input-helper-order.fixed.md

@@ -0,0 +1 @@
+Period-splitting input helpers now preserve supplied inputs in memory or on disk, replace calculated sub-period values, and invalidate overlapping caches consistently with the supplied-input tier record.

policyengine_core/holders/helpers.py

@@ -1,4 +1,5 @@
 import logging
+from typing import List
 
 import numpy
 from numpy.typing import ArrayLike
@@ -22,29 +23,178 @@ def get_stored_array(holder: Holder, period: Period, branch_name: str) -> ArrayL
     return holder._get_array_from_storage(period, branch_name)
 
 
+def _is_input(
+    holder: Holder, period: Period, branch_name: str, storage_name: str
+) -> bool:
+    """Whether the value stored for ``period`` under ``branch_name`` was set
+    as an input, rather than calculated by the simulation (a formula result,
+    or a default, carried-over or uprated value).
+
+    Use the same supplied-input tier record as reform replay. A cache write
+    can leave a registry entry behind or hide an input in another tier, and
+    its carry-over mark (``derived=False``) does not make it a supplied input.
+    A holder with no simulation, or a simulation with no record, keeps the
+    legacy behavior: every stored value counts as an input there.
+    """
+    simulation = getattr(holder, "simulation", None)
+    input_keys = getattr(simulation, "_user_input_keys", None)
+    if input_keys is None:
+        return True
+    period = holder._storage_period(period)
+    if (holder.variable.name, branch_name, period) not in input_keys:
+        return False
+    holder._seed_input_storage(period, branch_name)
+    return holder._stores_user_input(period, branch_name, storage_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.
+    """The supplied input under this exact branch, preferring memory to disk.
+
+    A surviving disk input is still known when a non-input memory cache
+    hides it; overlap cleanup removes that cache after the helper succeeds.
     """
-    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
+    # Checked before reading, so that a calculated value a branch still
+    # shares with its parent is not copied only to be replaced.
+    for storage_name, storage in (
+        ("memory", holder._memory_storage),
+        ("disk", holder._disk_storage),
     ):
-        storage = holder._disk_storage
-    else:
-        return None
-    if storage.is_derived(period, branch_name):
-        return None
-    return storage.get(period, branch_name)
+        if storage is not None and _is_input(holder, period, branch_name, storage_name):
+            value = storage.get(period, branch_name)
+            if value is not None:
+                return value
+    return None
+
+
+def _store_input(
+    holder: Holder, period: Period, array: ArrayLike, branch_name: str
+) -> None:
+    """Store ``array`` as the input for ``period``, replacing any value the
+    simulation calculated for it.
+
+    The value is added to the simulation's record of inputs even when the
+    helper is called directly instead of through ``Holder.set_input``, and
+    ``calculate``'s fast cache drops the value it held for the period.
+    """
+    simulation = getattr(holder, "simulation", None)
+    input_keys = getattr(simulation, "_user_input_keys", None)
+    # Let the holder normalize and record the supplied-input tier, including
+    # direct helper calls. Keep a legacy holder without an input record in
+    # that mode for the whole call, rather than creating a partial record.
+    holder._set(period, array, branch_name, is_input=input_keys is not None)
+
+
+def _evict_fast_cache(holder: Holder, changed_periods, branch_name: str) -> None:
+    """Drop ``calculate``'s fast-cache entries for the variable's
+    ``changed_periods``, if the simulation reads ``branch_name``.
+
+    The fast cache holds what this simulation's ``calculate`` returned; a
+    value stored under a branch it does not read changes none of that.
+    """
+    simulation = getattr(holder, "simulation", None)
+    fast_cache = getattr(simulation, "_fast_cache", None)
+    if not fast_cache:
+        return
+    get_visible_branch_names = getattr(simulation, "_get_visible_branch_names", None)
+    if (
+        get_visible_branch_names is not None
+        and branch_name not in get_visible_branch_names()
+    ):
+        return
+    for changed_period in changed_periods:
+        fast_cache.pop((holder.variable.name, changed_period), None)
+
+
+def _branches_read_with(holder: Holder, branch_name: str) -> List[str]:
+    """``branch_name`` and the branches whose values the holder's simulation
+    reads when ``branch_name`` stores none: its ancestors, then ``default``."""
+    simulation = getattr(holder, "simulation", None)
+    get_visible_branch_names = getattr(simulation, "_get_visible_branch_names", None)
+    if get_visible_branch_names is not None:
+        visible_branch_names = get_visible_branch_names()
+        if branch_name in visible_branch_names:
+            return visible_branch_names
+    return [branch_name]
+
+
+def _drop_calculated_overlapping(
+    holder: Holder, period: Period, branch_name: str
+) -> None:
+    """Drop what the simulation calculated for the variable over ``period``,
+    once a helper has set inputs for its sub-periods.
+
+    Besides the sub-periods the helper replaced, the simulation may hold
+    values it calculated from them: ``calculate_add`` caches the sum of a
+    variable over a longer period (a monthly variable's ``2013``), and
+    ``calculate_divide`` a twelfth of a yearly variable at a month. Any such
+    value for a period that overlaps ``period`` is out of date, and
+    ``calculate`` would return it instead of adding up or dividing the
+    input. Inputs are kept.
+
+    This holder's storage is searched under ``branch_name`` and under the
+    branches read with it (see ``_branches_read_with``): a branch starts with
+    its own index of what its parent calculated, and reads those values
+    until they are dropped. The simulation the branch was created from keeps
+    its values.
+    """
+    branch_names = set(_branches_read_with(holder, branch_name))
+    dropped_periods = set()
+    # Memory keys are "{branch}:{period}"; branch names cannot contain ":".
+    # Disk keys are "{branch}_{period}"; branch names can contain "_", but
+    # period strings cannot, so the period follows the last "_".
+    memory = holder._memory_storage
+    for key in list(memory._arrays):
+        stored_branch_name, period_string = key.split(":", 1)
+        if _is_calculated_overlapping(
+            holder, stored_branch_name, period_string, branch_names, period, "memory"
+        ):
+            del memory._arrays[key]
+            dropped_periods.add(_period_from_storage_key(period_string))
+    memory._stop_sharing_dropped_keys()
+    memory._unmark_dropped_keys()
+    disk = holder._disk_storage
+    if disk is not None:
+        for key in list(disk._files):
+            stored_branch_name, period_string = key.rsplit("_", 1)
+            if _is_calculated_overlapping(
+                holder, stored_branch_name, period_string, branch_names, period, "disk"
+            ):
+                del disk._files[key]
+                dropped_periods.add(_period_from_storage_key(period_string))
+        disk._derived.intersection_update(disk._files)
+    _evict_fast_cache(holder, dropped_periods, branch_name)
+
+
+def _period_from_storage_key(period_string: str) -> Period:
+    """Parse a stored period whose early year may have lost leading zeros."""
+    components = period_string.split(":")
+    date_index = 1 if len(components) > 1 else 0
+    date = components[date_index].split("-")
+    date[0] = date[0].zfill(4)
+    components[date_index] = "-".join(date)
+    return periods.period(":".join(components))
+
+
+def _is_calculated_overlapping(
+    holder: Holder,
+    stored_branch_name: str,
+    period_string: str,
+    branch_names: set,
+    period: Period,
+    storage_name: str,
+) -> bool:
+    if stored_branch_name not in branch_names:
+        return False
+    try:
+        stored_period = _period_from_storage_key(period_string)
+    except ValueError:
+        # Not a key ``put`` wrote (say, a file ``restore`` found).
+        return False
+    if stored_period.unit == periods.ETERNITY:
+        return False
+    if stored_period.start > period.stop or stored_period.stop < period.start:
+        return False
+    return not _is_input(holder, stored_period, stored_branch_name, storage_name)
 
 
 def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLike):
@@ -53,6 +203,20 @@ def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLik
 
     In this case, the variable will accept inputs on larger periods that its definition period, and the value for the larger period will be applied to all its subperiods.
 
+    Only inputs count as already known. A sub-period whose value the
+    simulation calculated (a formula result, or a default, carried-over or
+    uprated value) takes the new input like a sub-period with no value, so
+    the result does not depend on what was calculated before the input was
+    set.
+
+    A sub-period that already has an input keeps it, and that input, not the
+    value given for the larger period, is applied to the sub-periods after
+    it that have none: with ``3`` set for March, setting ``7`` for the year
+    gives ``7`` in January and February and ``3`` from March to December.
+    The value of a stock known at some month holds for the rest of the
+    period. (Setting the two inputs in the other order gives ``7`` in every
+    month but March.)
+
     To read more about ``set_input`` attributes, check the `documentation <https://openfisca.org/doc/coding-the-legislation/35_periods.html#set-input-automatically-process-variable-inputs-defined-for-periods-not-matching-the-definition-period>`_.
     """
     array = holder._to_array(array, validate_nan=True)
@@ -71,18 +235,19 @@ 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 months that already have an input
+    # Store the input data, skipping the sub-periods 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_input(holder, sub_period, branch_name)
-        if existing_array is None:
-            holder._set(sub_period, array, branch_name)
+        existing_input = get_stored_input(holder, sub_period, branch_name)
+        if existing_input is None:
+            _store_input(holder, sub_period, array, branch_name)
         else:
-            # The array of the current sub-period is reused for the next ones.
-            # TODO: refactor or document this behavior
-            array = existing_array
+            # The input of the current sub-period is applied to the next
+            # ones (see the docstring).
+            array = existing_input
         sub_period = sub_period.offset(1)
+    _drop_calculated_overlapping(holder, period, branch_name)
 
 
 def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike):
@@ -91,6 +256,13 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike)
 
     In this case, the variable will accept inputs on larger periods that its definition period, and the value for the larger period will be divided between its subperiods.
 
+    Sub-periods that already have an input keep it, and what is left of the
+    value is divided between the others. Only inputs count as already known.
+    A sub-period whose value the simulation calculated (a formula result, or
+    a default, carried-over or uprated value) takes its share of the new
+    input like a sub-period with no value, so the result does not depend on
+    what was calculated before the input was set.
+
     To read more about ``set_input`` attributes, check the `documentation <https://openfisca.org/doc/coding-the-legislation/35_periods.html#set-input-automatically-process-variable-inputs-defined-for-periods-not-matching-the-definition-period>`_.
     """
     if not isinstance(array, numpy.ndarray):
@@ -110,30 +282,28 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike)
 
     after_instant = period.start.offset(period_size, period_unit)
 
-    # Count the number of elementary periods to change, and the difference with what is already known.
+    # Find the elementary periods to change, and the difference with the inputs already known.
     branch_name = get_input_branch(holder)
     remaining_array = array.copy()
     sub_period = period.start.period(cached_period_unit)
-    sub_periods_count = 0
+    sub_periods_to_set = []
     while sub_period.start < after_instant:
-        existing_array = get_stored_input(holder, sub_period, branch_name)
-        if existing_array is not None:
-            remaining_array -= existing_array
+        existing_input = get_stored_input(holder, sub_period, branch_name)
+        if existing_input is not None:
+            remaining_array -= existing_input
         else:
-            sub_periods_count += 1
+            sub_periods_to_set.append(sub_period)
         sub_period = sub_period.offset(1)
 
-    # Cache the input data
-    if sub_periods_count > 0:
-        divided_array = remaining_array / sub_periods_count
-        sub_period = period.start.period(cached_period_unit)
-        while sub_period.start < after_instant:
-            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)
+    # Store the input data
+    if sub_periods_to_set:
+        divided_array = remaining_array / len(sub_periods_to_set)
+        for sub_period in sub_periods_to_set:
+            _store_input(holder, sub_period, divided_array, branch_name)
     elif not (remaining_array == 0).all():
         raise ValueError(
             "Inconsistent input: variable {0} has already been set for all months contained in period {1}, and value {2} provided for {1} doesn't match the total ({3}). This error may also be thrown if you try to call set_input twice for the same variable and period.".format(
                 holder.variable.name, period, array, array - remaining_array
             )
         )
+    _drop_calculated_overlapping(holder, period, branch_name)

policyengine_core/holders/holder.py

@@ -259,13 +259,18 @@ def _record_input_storage(
         locations = locations | {storage} if is_input else locations - {storage}
         recorded[slot] = (period, locations)
 
-    def _stores_user_input(self, period: Period, branch_name: str) -> bool:
+    def _stores_user_input(
+        self, period: Period, branch_name: str, storage_name: Optional[str] = None
+    ) -> bool:
         """Whether an actual input survives in either tier, without reading it.
 
         Called only for a slot already in the simulation's replay record.
         Register-only inputs (including restored inputs) may have no location
         metadata yet; their existing storage marks identify the surviving
         tier. Marks never add an unrecorded slot to the replay record.
+
+        With ``storage_name``, check only that tier so input helpers can
+        discard a cache in one tier while keeping an input in the other.
         """
         recorded = getattr(self, "_user_input_storage", {}).get(
             (branch_name, str(period))
@@ -276,7 +281,8 @@ def _stores_user_input(self, period: Period, branch_name: str) -> bool:
             ("disk", self._disk_storage),
         ):
             if (
-                location in locations
+                (storage_name is None or location == storage_name)
+                and location in locations
                 and storage is not None
                 and storage.has(period, branch_name)
                 and not storage.is_derived(period, branch_name)