Protect stored draws during a calculation

changelog.d/2874.added.md

@@ -0,0 +1 @@
+Simulations read three new stored inputs, `uc_new_claimant_health_element_draw`, `pip_phase_in_draw` and `would_claim_student_loan`, in Universal Credit rebalancing, the PIP phase-in scenario and student loan uprating when a dataset supplies them, and keep their own seeded draws when it does not.

changelog.d/2874.changed.md

@@ -0,0 +1 @@
+The test suite fails on any new random draw, sampling call or hash of entity ids in the package, and lists today's draws in a ratchet.

policyengine_uk/data/core_route_checks.py

@@ -0,0 +1,154 @@
+"""Checks the policyengine-core ``Dataset`` route makes before it exports.
+
+``Simulation.build_from_dataset`` stores each variable of a core ``Dataset``
+through ``set_input``, which casts identifiers to int32, exports the
+simulation at the dataset's first period with
+``UKSingleYearDataset.from_simulation`` and rebuilds from that export. Two
+things a stored draw or decision needs can be lost on the way:
+
+* Identities. The rebuild declares benefit units and households in sorted
+  order of the cast identifiers, and a stored decision keyed on a raw
+  identifier no longer matches one keyed on the cast identifier. So a dataset
+  carrying a stored column must give identifiers that survive the int32 cast,
+  for people, benefit units and households.
+* Periods. Only the first period is exported. A column supplied only for
+  later periods would be exported as its default and then stored as a
+  supplied input. So a stored column, or an outcome that exempts a draw, that
+  the dataset supplies for any period must be supplied for the first; and a
+  draw's later periods must repeat the first, since only the first is kept.
+"""
+
+import warnings
+
+import numpy as np
+from policyengine_core.periods import period as period_
+
+# Stored columns the data build writes: the UC deduction draws, the
+# birthday position, the Warm Home Discount draw, the UC new-claimant draw,
+# the student loan decision and the PIP phase-in draw.
+STORED_COLUMNS = (
+    "uc_deduction_random_draw",
+    "uc_deduction_type_random_draw",
+    "months_since_last_birthday",
+    "warm_home_discount_random_draw",
+    "uc_new_claimant_health_element_draw",
+    "would_claim_student_loan",
+    "pip_phase_in_draw",
+)
+# Draws, whose later periods must repeat the first.
+DRAW_COLUMNS = (
+    "uc_deduction_random_draw",
+    "uc_deduction_type_random_draw",
+    "warm_home_discount_random_draw",
+    "uc_new_claimant_health_element_draw",
+    "pip_phase_in_draw",
+    "attends_private_school_random_draw",
+)
+# Columns that must be supplied for the first period if supplied at all: the
+# stored columns, the private school draw, and the outcomes that exempt a
+# draw.
+FIRST_PERIOD_COLUMNS = STORED_COLUMNS + (
+    "attends_private_school_random_draw",
+    "uc_receives_new_claimant_health_element",
+    "attends_private_school",
+)
+# The identifier arrays of each entity: its own and its members' memberships.
+ENTITY_IDENTIFIERS = {
+    "person": ("person_id",),
+    "benunit": ("benunit_id", "person_benunit_id"),
+    "household": ("household_id", "person_household_id"),
+}
+
+
+class CoreRouteError(ValueError):
+    """A core ``Dataset`` that the export and rebuild would change."""
+
+
+def _fails_int32_cast(ids) -> np.ndarray:
+    """Which identifiers change when cast to int32."""
+    values = np.asarray(ids)
+    if values.dtype.kind == "O":
+        try:
+            values = np.array(values.tolist(), dtype=np.float64)
+        except (TypeError, ValueError):
+            return np.ones(values.shape, dtype=bool)
+    if values.dtype.kind not in "iub":
+        values = values.astype(np.float64)
+        finite = np.isfinite(values)
+        safe = np.where(finite, values, 0.0)
+        in_range = (safe >= np.iinfo(np.int32).min) & (safe <= np.iinfo(np.int32).max)
+        return ~(finite & in_range & (safe == np.trunc(safe)))
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        return values.astype(np.int32).astype(values.dtype) != values
+
+
+def check_identifiers(data: dict) -> None:
+    """Refuse a dataset carrying a stored column whose identifiers do not
+    survive the int32 cast, naming the entity and the identifiers."""
+    carried = [column for column in STORED_COLUMNS if column in data]
+    if not carried:
+        return
+    for entity, columns in ENTITY_IDENTIFIERS.items():
+        failing = []
+        for column in columns:
+            for values in (data.get(column) or {}).values():
+                values = np.asarray(values)
+                bad = _fails_int32_cast(values)
+                failing.extend(values[bad].reshape(-1).tolist())
+        if failing:
+            unique = list(dict.fromkeys(failing))
+            raise CoreRouteError(
+                f"This core Dataset carries {', '.join(carried)}, and "
+                f"{len(unique)} {entity} identifier(s) do not survive the "
+                f"int32 cast its route stores identifiers with (the first: "
+                f"{unique[:5]}). The rebuild would key or order the stored "
+                "values on other identifiers. Renumber the identifiers within "
+                "the int32 range, or load the data as a UKSingleYearDataset."
+            )
+
+
+def check_supplied_periods(
+    data: dict, system, first_time_period, group_orders: dict
+) -> None:
+    """Refuse a dataset whose stored columns the first-period export would
+    change, naming the column and its periods."""
+    first = str(period_(first_time_period))
+    for column in FIRST_PERIOD_COLUMNS:
+        if column not in data or column not in system.variables:
+            continue
+        by_period = data[column]
+        if not len(by_period):
+            continue
+        supplied = {str(period_(key)): key for key in by_period}
+        if first not in supplied:
+            raise CoreRouteError(
+                f"This core Dataset supplies {column} for "
+                f"{', '.join(sorted(supplied))} but not for {first}, the only "
+                "period its route exports. The export would store the "
+                "variable's default as a supplied value."
+            )
+        if column not in DRAW_COLUMNS or len(supplied) == 1:
+            continue
+        variable = system.variables[column]
+        orders = group_orders.get(variable.entity.key)
+
+        def aligned(key):
+            values = np.asarray(by_period[key])
+            if orders is not None and values.ndim > 0:
+                values = values[orders.get(key, orders[supplied[first]])]
+            return values.astype(variable.dtype)
+
+        first_values = aligned(supplied[first])
+        differing = [
+            name
+            for name, key in supplied.items()
+            if name != first and not np.array_equal(aligned(key), first_values)
+        ]
+        if differing:
+            raise CoreRouteError(
+                f"This core Dataset supplies {column} for {', '.join(sorted(supplied))} "
+                f"with values in {', '.join(sorted(differing))} that differ "
+                f"from {first}'s. Its route keeps only {first}, so a draw "
+                "must not change between periods."
+            )

policyengine_uk/data/economic_assumptions.py

@@ -31,6 +31,10 @@
 import logging
 
 from policyengine_uk.utils.stochastic import splitmix64_uniform
+from policyengine_uk.utils.stored_draws import (
+    STUDENT_LOAN_TAKE_UP,
+    check_boolean_input,
+)
 
 # Approximate take-up rate for assigning loans to tertiary-educated NONE people.
 # This represents P(has loan AND earning above threshold | tertiary educated).
@@ -824,6 +828,20 @@ def uprate_rent(
     return current_year
 
 
+def stored_student_loan_take_up(person: pd.DataFrame) -> np.ndarray | None:
+    """The dataset's stored ``would_claim_student_loan`` for each row of
+    ``person``, or None where the dataset does not supply the column.
+
+    Raises if a value is missing or is not a boolean or the integer 0 or 1:
+    cast to a boolean, a missing value would read as True.
+    """
+    if STUDENT_LOAN_TAKE_UP not in person.columns:
+        return None
+    values = person[STUDENT_LOAN_TAKE_UP]
+    check_boolean_input(STUDENT_LOAN_TAKE_UP, values)
+    return np.asarray(values, dtype=bool)
+
+
 def uprate_student_loan_plans(
     current_year: UKSingleYearDataset,
     previous_year: UKSingleYearDataset,
@@ -925,13 +943,20 @@ def uprate_student_loan_plans(
     new_plan = base_plan.copy()
     repayments = person["student_loan_repayments"].values.copy()
 
-    # One draw per person, the same in every year.
-    draws = splitmix64_uniform(person["person_id"].values, salt=_STUDENT_LOAN_DRAW_SALT)
+    # Whether each person takes up a loan as a new borrower: the dataset's
+    # stored decision where it supplies one, else one draw per person, the
+    # same in every year, below the take-up rate.
+    would_claim = stored_student_loan_take_up(person)
+    if would_claim is None:
+        draws = splitmix64_uniform(
+            person["person_id"].values, salt=_STUDENT_LOAN_DRAW_SALT
+        )
+        would_claim = draws < _GRADUATE_LOAN_TAKE_UP
 
     # Helper to assign plans to eligible people
-    def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP):
-        """Assign the funding-nation plan to sampled new borrowers."""
-        sampled = mask & (draws < prob) & (new_borrower_plan != "NONE")
+    def assign_new_borrowers(mask):
+        """Assign the funding-nation plan to new borrowers who take it up."""
+        sampled = mask & would_claim & (new_borrower_plan != "NONE")
         new_plan[sampled] = new_borrower_plan[sampled]
         repayments[sampled] = 0.0
 
@@ -962,7 +987,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP):
         & (uni_start_year < _PLAN2_FIRST_UNI_START)
         & ~plan1_written_off
     )
-    assign_with_probability(plan1_new_cohort)
+    assign_new_borrowers(plan1_new_cohort)
 
     # === Step 3: Add holders in the extended 2012-2022 age bands ===
     # In the base year Plan 2 holders are at most the age of the 2012 intake
@@ -977,7 +1002,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP):
         & (uni_start_year >= _PLAN2_FIRST_UNI_START)
         & (uni_start_year < _PLAN5_FIRST_UNI_START)
     )
-    assign_with_probability(plan2_new_cohort)
+    assign_new_borrowers(plan2_new_cohort)
 
     # === Step 4: Add Plan 5 holders (new plan from 2023) ===
     # Plan 5 is England's, and did not exist in the base year. The other
@@ -993,7 +1018,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP):
         & (uni_start_year >= _PLAN5_FIRST_UNI_START)
         & (age >= 21)
     )
-    assign_with_probability(plan5_eligible)
+    assign_new_borrowers(plan5_eligible)
 
     # Write back to the person table (without the merged region column).
     person_out = current_year.person.copy()

policyengine_uk/scenarios/pip_reform.py

@@ -1,20 +1,39 @@
 from policyengine_uk.model_api import Scenario
 from policyengine_uk import Simulation
+from policyengine_uk.utils.stored_draws import input_source
 import numpy as np
 
+PHASE_IN_DRAW = "pip_phase_in_draw"
+
+
+def phased_out_by_draw(draw, share) -> np.ndarray:
+    """People whose draw falls below the year's phase-out share, compared in
+    float64: a stored draw is a float32 value, and comparing it with a Python
+    float in float32 would round the share first."""
+    return np.asarray(draw, dtype=np.float64) < share
+
 
 def modify_simulation(sim: Simulation):
-    np.random.seed(42)  # For reproducibility
-    pip_seed = np.random.random(len(sim.calculate("person_id")))
+    # The seeded stream, drawn only for a year that no stored draw reaches,
+    # and then once: each person takes the value at their row.
+    pip_seed = None
+    draw_population = sim.get_variable_population(PHASE_IN_DRAW)
 
     start_year = 2025
     end_year = 2029
 
     for year in range(start_year, end_year + 1):
+        if input_source(draw_population, PHASE_IN_DRAW, year) is not None:
+            draw = sim.calculate(PHASE_IN_DRAW, year)
+        else:
+            if pip_seed is None:
+                np.random.seed(42)  # For reproducibility
+                pip_seed = np.random.random(len(sim.calculate("person_id")))
+            draw = pip_seed
         # Copy: calculate returns the simulation's cached array.
         current_pip = np.array(sim.calculate("pip", year))
         percent_along_phase_in = (year - start_year) / (end_year - start_year)
-        current_pip[pip_seed < 0.25 * percent_along_phase_in] = 0
+        current_pip[phased_out_by_draw(draw, 0.25 * percent_along_phase_in)] = 0
         sim.set_input("pip", year, current_pip)
     return sim
 

policyengine_uk/scenarios/uc_reform.py

@@ -3,6 +3,7 @@
 from policyengine_uk.utils.data_source import built_from_data
 from policyengine_uk.data.economic_assumptions import PROJECTION_END_YEAR
 from policyengine_uk.utils.supplied_inputs import set_model_value, supplied_input
+from policyengine_uk.utils.stored_draws import input_source
 import numpy as np
 import warnings
 from weakref import WeakKeyDictionary
@@ -131,6 +132,7 @@ def _protected_existing_health_element_monthly(
 
 
 NEW_CLAIMANT_STATUS = "uc_receives_new_claimant_health_element"
+NEW_CLAIMANT_DRAW = "uc_new_claimant_health_element_draw"
 LCWRA_AMOUNT = "gov.dwp.universal_credit.elements.disabled.amount"
 
 
@@ -192,15 +194,15 @@ def warn_shadowed_lcwra_reform(sim) -> None:
 
 
 def _new_claimant_health_element_status(
-    sim: Microsimulation, year: int, drawn: np.ndarray, built_from_data: bool
+    sim: Microsimulation, year: int, drawn, built_from_data: bool
 ) -> np.ndarray:
     """Which benefit units get the new-claimant LCWRA rate in a year.
 
     A value the situation or dataset supplies is used as given. Data without
-    one takes the seeded draw, recorded as an input. A household situation
-    without one gets the protected amount (the variable's default): a draw
-    there would make a household's award depend on its position among the
-    other benefit units in the simulation.
+    one takes the draw (``drawn``, each benefit unit's decision), recorded
+    as an input. A household situation without one gets the protected amount
+    (the variable's default): a draw there would make a household's award
+    depend on its position among the other benefit units in the simulation.
     """
     # sim.get_array reads the simulation's own branch (and the branches it
     # inherits from); the holder's get_array reads only the default branch.
@@ -210,11 +212,20 @@ def _new_claimant_health_element_status(
     if built_from_data:
         status = np.asarray(drawn, dtype=bool)
     else:
-        status = np.zeros(len(drawn), dtype=bool)
+        status = np.zeros(
+            sim.get_variable_population(NEW_CLAIMANT_STATUS).count, dtype=bool
+        )
     sim.set_input(NEW_CLAIMANT_STATUS, year, status)
     return status
 
 
+def new_claimants_by_draw(draw, share) -> np.ndarray:
+    """Benefit units whose draw falls below the new-claimant share, compared
+    in float64: a stored draw is a float32 value, and comparing it with a
+    Python float in float32 would round the share first."""
+    return np.asarray(draw, dtype=np.float64) < share
+
+
 def _lcwra_amount_is_overridden(sim, year: int) -> bool:
     """An explicit LCWRA amount keeps precedence over the inactive path."""
     from policyengine_uk.utils.dated_reforms import legacy_override_bounds
@@ -308,9 +319,9 @@ def add_universal_credit_reform(sim: Microsimulation):
     _add_inactive_health_element_formula(sim)
     rebalancing = sim.tax_benefit_system.parameters.gov.dwp.universal_credit.rebalancing
 
-    generator = np.random.default_rng(43)
-
-    uc_seed = generator.random(len(sim.calculate("benunit_id")))
+    # The seeded stream, drawn only for a year whose dataset supplies no
+    # stored draw, and then once: each unit takes its position's value.
+    uc_seed = None
     # rebalancing.new_claimant_share: WPI Economics' estimates for 2026-27 to
     # 2029-30, held after that (an assumption; see the parameter).
     new_claimant_share = rebalancing.new_claimant_share
@@ -327,8 +338,23 @@ def add_universal_credit_reform(sim: Microsimulation):
             # Recalculate through the formula, preserving user-supplied awards.
             sim.delete_arrays("uc_LCWRA_element", year)
             continue
+        # Data without a status for the year takes the dataset's stored draw
+        # where one reaches the year, else the seeded stream.
+        drawn = None
+        if is_data and sim.get_array(NEW_CLAIMANT_STATUS, year) is None:
+            share = new_claimant_share(year)
+            draw_population = sim.get_variable_population(NEW_CLAIMANT_DRAW)
+            if input_source(draw_population, NEW_CLAIMANT_DRAW, year) is not None:
+                drawn = new_claimants_by_draw(
+                    sim.calculate(NEW_CLAIMANT_DRAW, year), share
+                )
+            else:
+                if uc_seed is None:
+                    generator = np.random.default_rng(43)
+                    uc_seed = generator.random(len(sim.calculate("benunit_id")))
+                drawn = new_claimants_by_draw(uc_seed, share)
         is_post_2025_claimant = _new_claimant_health_element_status(
-            sim, year, uc_seed < new_claimant_share(year), is_data
+            sim, year, drawn, is_data
         )
         # Copy: calculate returns the simulation's cached array, and the writes
         # below must reach the cache only through set_input.