Give a simulation its identifiers for a year before its data

changelog.d/2353.fixed.md

@@ -0,0 +1 @@
+- Fixed a simulation given an input for a year before its data holding no person, benefit unit or household identifiers in that year, so that formulas reading them, such as `months_since_last_birthday` and the Universal Credit deduction draws, gave different values there from the simulation its exported frame loads back to.

changelog.d/fix-roundtrip-recorded-inputs.fixed.md

@@ -0,0 +1 @@
+- Fixed `Simulation.to_input_dataframe()` leaving out the person, benefit unit and household identifiers for years in which a simulation was given inputs after it was built, such as the Universal Credit health element status for 2026 to 2029, so that an exported frame loads back to a simulation that exports the same frame.

docs/book/usage/simulations.md

@@ -200,6 +200,19 @@ print(sim.calculate("council_tax", 2025))  # One value per household
   period that is not one year are set on the simulation directly. A variable
   given for both months and the year is set column by column in the frame's
   order, since a year's value is spread over the months not already set.
+- `to_input_dataframe()` gives the identifier columns for every year that one
+  of its columns gives a whole year's value for, since a simulation built from
+  the frame holds the identifiers in each of those years. So an exported frame
+  loads back to a simulation that exports the same frame, even when the
+  simulation was given inputs for years outside its data, such as the
+  Universal Credit health element status the rebalancing records for 2026 to
+  2029.
+- Inputs carry forward to later years but not back to earlier ones. So when
+  `set_input` gives a simulation an input for a year before every year it
+  holds the identifiers for, it also gives the identifiers for that year.
+  Formulas that read identifiers, such as the draws that place each person in
+  their year of age, then read the same people and groups there as in the
+  simulation its exported frame loads back to.
 
 ```{important}
 The model stores `person_household_id` and `person_benunit_id` as 32-bit

policyengine_uk/data/input_dataframe.py

@@ -15,14 +15,24 @@
 Identifiers are read as whole numbers, value by value, and a float that may
 stand for more than one whole number is refused rather than allowed to merge
 two people or groups.
+
+The identifiers hold for every period, and a simulation built from a frame
+holds them all as inputs in each whole year a column names.
+``Simulation.to_input_dataframe`` gives them for each such year too
+(``with_identifiers_in_each_year``), so that an exported frame loads back to
+a simulation that exports it again. Inputs carry forward but not back, so
+``Simulation.set_input`` gives a simulation its identifiers for a year
+before every year it holds them for when it is given an input there
+(``hold_identifiers_in``). Formulas that read them then read the same people
+and groups in that year as the simulation the frame loads back to.
 """
 
 import warnings
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 import numpy as np
 import pandas as pd
-from policyengine_core.periods import period as period_
+from policyengine_core.periods import ETERNITY, period as period_
 
 from policyengine_uk.data.dataset_schema import (
     UKMultiYearDataset,
@@ -267,6 +277,17 @@ def _holds_floats(identifiers: pd.Series) -> bool:
     return values.dtype.kind == "f"
 
 
+def _is_whole_year(period) -> bool:
+    """Whether ``period`` is one calendar year, the period of a dataset
+    year."""
+    start = period.start
+    return (
+        period.unit == "year"
+        and period.size == 1
+        and (start.month, start.day) == (1, 1)
+    )
+
+
 def _same(left: np.ndarray, right: np.ndarray) -> np.ndarray:
     """Whether each pair of values is equal, counting two missing values as
     equal."""
@@ -362,11 +383,7 @@ def no_tables() -> Dict[str, Dict[str, np.ndarray]]:
             continue
         period = period_(period_text)
         start = period.start
-        whole_year = (
-            period.unit == "year"
-            and period.size == 1
-            and (start.month, start.day) == (1, 1)
-        )
+        whole_year = _is_whole_year(period)
         if variable in IDENTIFIER_VARIABLES:
             if whole_year:
                 identifier_years.add(start.year)
@@ -425,3 +442,137 @@ def no_tables() -> Dict[str, Dict[str, np.ndarray]]:
         for fiscal_year, tables in sorted(years.items())
     ]
     return UKMultiYearDataset(datasets=datasets), replayed, identifier_inputs
+
+
+def _ignored(variable, period) -> bool:
+    """Whether policyengine-core's ``Simulation.set_input`` stores nothing
+    for ``variable`` in ``period``, as it does for a period that starts after
+    the variable ends."""
+    return variable.end is not None and period.start.date > variable.end
+
+
+def hold_identifiers_in(simulation, variable_name: str, period) -> None:
+    """Give ``simulation`` its identifiers for the year an input of
+    ``variable_name`` for ``period`` starts in, if they start later.
+
+    policyengine-core carries an input forward to later periods but never
+    back (``Simulation._calculate``). So a simulation given an input for a
+    year before its data, such as 2023 for a frame of 2025, would hold no
+    identifiers there. Formulas would read every person, benefit unit and
+    household identifier there as zero, the default. Yet once the frame it
+    exports gives a whole year's value for that year, the frame loads back to
+    a simulation that holds the people and groups there: each whole year a
+    column of another variable names is a year of the dataset, which holds
+    them, and each other year an identifier column names gets the identifiers
+    as inputs (``dataset_from_input_dataframe``).
+
+    So each identifier that ``simulation`` holds as an input for some period,
+    but for none that starts by the first day of the year ``period`` starts
+    in, is set for that year to its value in the period it is held for that
+    starts first. A simulation holds one set of people and groups, so that
+    value names the same ones. From that year on the identifier carries
+    forward, so the simulation holds its identifiers in each year it holds an
+    input for. An identifier the simulation does not hold as an input, as in
+    one built from a situation, is left alone.
+
+    An input core does not store, for a period after its variable ends, gives
+    none. Each identifier is set through ``simulation.set_input``, which
+    calls this again for the identifier. That call finds the identifier held
+    for the year, or returns here if core did not store it, so the calls end.
+    """
+    period = period_(period)
+    variables = simulation.tax_benefit_system.variables
+    if period.unit == ETERNITY or _ignored(variables[variable_name], period):
+        return
+    year = period_(str(period.start.year))
+    branch = simulation.branch_name
+    for identifier in sorted(IDENTIFIER_VARIABLES):
+        holder = simulation.get_holder(identifier)
+        held = holder.get_input_periods(branch)
+        if not held:
+            continue
+        first = min(held, key=lambda held_period: held_period.start)
+        if first.start > year.start:
+            simulation.set_input(
+                identifier, year, np.array(holder.get_array(first, branch))
+            )
+
+
+def with_identifiers_in_each_year(
+    df: pd.DataFrame, tax_benefit_system: "TaxBenefitSystem"
+) -> pd.DataFrame:
+    """``df``, an exported frame, with each identifier it gives for some year
+    given for every whole calendar year a column of the frame names.
+
+    ``dataset_from_input_dataframe`` holds every identifier in each such year:
+    a year a column of another variable names is a year of the dataset, which
+    holds the people and groups, and each other year an identifier column
+    names gets them all as inputs. So a simulation built from the frame holds
+    the identifiers as inputs in those years, and exports them there.
+
+    The simulation that exported the frame may not have. One given an input
+    for a year after its data once it was built, as the Universal Credit
+    rebalancing records each benefit unit's health element status for 2026
+    to 2029, holds its identifiers in that year only by carrying them forward
+    from an earlier one. (In a year before its data, where inputs do not
+    carry back, ``hold_identifiers_in`` gives it them as inputs.)
+    policyengine-core exports an input only for the periods it was set for.
+    Giving every identifier for each year makes the frame load back to a
+    simulation that exports the same frame.
+
+    A simulation holds one set of people and groups, so every column of an
+    identifier names the same ones (the loader refuses a frame whose columns
+    differ). An identifier added for a year copies its column for the latest
+    year before it the frame gives, or else the earliest year it gives. Each
+    goes beside the identifier's other columns, before the first whose period
+    sorts after its own. The default export sorts each variable's periods, so
+    that is where it puts an input set for that year. (With
+    ``include_computed_variables=True``, policyengine-core exports each
+    variable's periods in the order they are stored, and the frame is not
+    expected to load back to the same frame.)
+
+    Only columns that name a variable of ``tax_benefit_system`` count, and an
+    identifier the frame does not give for any year is not added. ``df`` is
+    not changed: if every identifier it gives is given for each such year,
+    ``df`` itself is returned, and otherwise a new frame.
+    """
+    named: Dict[str, Dict[int, object]] = {}
+    years = set()
+    for column in df.columns:
+        variable, period_text = split_column(column)
+        if variable not in tax_benefit_system.variables:
+            continue
+        period = period_(period_text)
+        if not _is_whole_year(period):
+            continue
+        if variable in IDENTIFIER_VARIABLES:
+            named.setdefault(variable, {})[period.start.year] = column
+        years.add(period.start.year)
+    added = {
+        variable: sorted(years - set(columns))
+        for variable, columns in named.items()
+        if years - set(columns)
+    }
+    if not added:
+        return df
+    result = df.copy(deep=False)
+    for variable, missing in added.items():
+        columns = named[variable]
+        for year in missing:
+            earlier = [named_year for named_year in columns if named_year < year]
+            source = columns[max(earlier) if earlier else min(columns)]
+            name = f"{variable}{COLUMN_SEPARATOR}{year}"
+            # The identifier's columns, by period text, as the exporter sorts
+            # them; the new one goes before the first that sorts after it.
+            own = [
+                position
+                for position, column in enumerate(result.columns)
+                if split_column(column)[0] == variable
+            ]
+            later = [
+                position
+                for position in own
+                if split_column(result.columns[position])[1] > str(year)
+            ]
+            result.insert(later[0] if later else own[-1] + 1, name, df[source].copy())
+    return result

policyengine_uk/simulation.py

@@ -32,7 +32,11 @@
     reset_growthfactor_uprating,
 )
 from policyengine_uk.data.dataset_sources import materialize_gcs_dataset_url
-from policyengine_uk.data.input_dataframe import dataset_from_input_dataframe
+from policyengine_uk.data.input_dataframe import (
+    dataset_from_input_dataframe,
+    hold_identifiers_in,
+    with_identifiers_in_each_year,
+)
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
@@ -247,6 +251,14 @@ def delete_arrays(self, variable: str, period: Period = None) -> None:
         # and the current branch, ancestor branches and default branch.
         drop_missing_supplied_inputs(self, variable)
 
+    def set_input(self, variable_name: str, period, value) -> None:
+        """Set an input as policyengine-core does, then give the simulation
+        its identifiers for the year the input starts in if it holds them
+        only for later years, since inputs do not carry back. See
+        ``policyengine_uk.data.input_dataframe.hold_identifiers_in``."""
+        super().set_input(variable_name, period, value)
+        hold_identifiers_in(self, variable_name, period)
+
     def reset_calculations(self):
         for variable in self.tax_benefit_system.variables:
             if variable not in self.input_variables:
@@ -426,6 +438,30 @@ def build_from_dataframe(self, df: pd.DataFrame) -> None:
         for variable, time_period, values in replayed:
             self.set_input(variable, time_period, values)
 
+    def to_input_dataframe(
+        self, include_computed_variables: bool = False
+    ) -> pd.DataFrame:
+        """Export the simulation's inputs as a DataFrame that
+        ``Simulation(dataset=...)`` loads back.
+
+        policyengine-core exports each input for the periods it was set for.
+        The identifiers are also given for every whole year a column names,
+        where a simulation built from the frame holds them, so that the
+        default export loads back to a simulation that exports the same
+        frame. See
+        ``policyengine_uk.data.input_dataframe.with_identifiers_in_each_year``.
+
+        Args:
+            include_computed_variables: If True, export every stored value,
+                calculated ones included, as policyengine-core does.
+        """
+        return with_identifiers_in_each_year(
+            super().to_input_dataframe(
+                include_computed_variables=include_computed_variables
+            ),
+            self.tax_benefit_system,
+        )
+
     def build_from_dataset(self, dataset: Dataset) -> None:
         """Build simulation from a Dataset object.
 

policyengine_uk/tests/code_health/test_simulation_methods_defined_once.py

@@ -0,0 +1,77 @@
+"""Each method of the simulation classes is defined once.
+
+Pull requests often override methods of ``policyengine_uk.simulation.Simulation``
+(``set_input``, ``delete_arrays``, ``clone`` and others). Two overrides of one
+method, added at different places in the class, merge without a git conflict,
+and the later definition replaces the earlier one: whatever the earlier one
+did is lost without an error. Combine them into one method instead.
+
+``ruff check`` reports this as F811 when it runs. This test guards the merged
+class with the tests themselves, wherever lint is not run. Property setters
+and deleters, and ``typing.overload`` stubs, name a method again on purpose
+and are left out.
+"""
+
+import ast
+from collections import Counter
+from pathlib import Path
+
+import pytest
+
+PACKAGE = Path(__file__).resolve().parents[2]
+MODULES = ["simulation.py", "microsimulation.py"]
+
+
+def _names_again(decorator: ast.expr) -> bool:
+    """Whether ``decorator`` defines a method under a name already used."""
+    if isinstance(decorator, ast.Attribute):
+        return decorator.attr in ("setter", "deleter", "getter")
+    return isinstance(decorator, ast.Name) and decorator.id == "overload"
+
+
+def methods_defined_more_than_once(source: str) -> dict:
+    """The methods each class in ``source`` defines more than once, by
+    class."""
+    repeated = {}
+    for node in ast.walk(ast.parse(source)):
+        if not isinstance(node, ast.ClassDef):
+            continue
+        names = Counter(
+            method.name
+            for method in node.body
+            if isinstance(method, (ast.FunctionDef, ast.AsyncFunctionDef))
+            and not any(map(_names_again, method.decorator_list))
+        )
+        twice = sorted(name for name, count in names.items() if count > 1)
+        if twice:
+            repeated[node.name] = twice
+    return repeated
+
+
+@pytest.mark.parametrize("module", MODULES)
+def test_each_method_of_the_simulation_classes_is_defined_once(module):
+    repeated = methods_defined_more_than_once((PACKAGE / module).read_text())
+    assert not repeated, (
+        f"{module} defines these methods more than once, so only the last "
+        f"definition of each takes effect: {repeated}. Combine them."
+    )
+
+
+def test_a_method_defined_twice_is_found():
+    source = """
+class Simulation:
+    def set_input(self):
+        pass
+
+    @property
+    def size(self):
+        pass
+
+    @size.setter
+    def size(self, value):
+        pass
+
+    def set_input(self):
+        pass
+"""
+    assert methods_defined_more_than_once(source) == {"Simulation": ["set_input"]}