Export the identifiers for every year an input frame names, so it round-trips

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,13 @@ 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.
 
 ```{important}
 The model stores `person_household_id` and `person_benunit_id` as 32-bit

policyengine_uk/data/input_dataframe.py

@@ -15,6 +15,12 @@
 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.
 """
 
 import warnings
@@ -267,6 +273,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 +379,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 +438,79 @@ 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 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; and in a year before its data, where inputs do not
+    carry back, its identifiers are the default, zero. 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, in order of period, which is
+    where the default export (sorted by period) puts an input set for that
+    year.
+
+    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.
+    """
+    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,10 @@
     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,
+    with_identifiers_in_each_year,
+)
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
@@ -426,6 +429,29 @@ 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 frame
+        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.