Impute UC deductions and private school attendance in constituency and local-authority runs

changelog.d/built-from-dataset-imputation-gates.fixed.md

@@ -0,0 +1 @@
+- UC deduction draws and private school attendance now apply to every simulation built from data, including a constituency or local authority filtered from the national data, instead of only those carrying over a million people of weight. A filtered region ranks private school attendance by its own income distribution. `filter_dataset` carries each person's private school attendance into the household it extracts. Simulations rebuilt in place, and policyengine-core simulations built over the UK system from data, are recognised as data-built too.

docs/book/validation/uc-deductions.md

@@ -81,4 +81,4 @@ A cap of 1 − *x* is equivalent to a protected minimum floor at *x* of the stan
 
 Statutory parameters — the cap, the protected floor, the minimum payable penny, the abolition switches — live under `gov.dwp.universal_credit.deductions`. The calibrated distributions live under `gov.simulation.uc_deductions`: they describe the world, not the law.
 
-The assignment formulas are an explicit fallback. The end state imputes `uc_latent_deduction_rate` and `uc_deduction_combination` at dataset build, at which point the model consumes them as plain inputs and the fallback retires; raw-FRS users keep working through the fallback until then. Assignment uses deterministic splitmix64 hashes of `benunit_id`, reproducible across runs and machines, and overridable by datasets or situations through `uc_deduction_random_draw` and `uc_deduction_type_random_draw`. Single-household simulations get no deductions unless set explicitly.
+The assignment formulas are an explicit fallback. The end state imputes `uc_latent_deduction_rate` and `uc_deduction_combination` at dataset build, at which point the model consumes them as plain inputs and the fallback retires; raw-FRS users keep working through the fallback until then. Assignment uses deterministic splitmix64 hashes of `benunit_id`, reproducible across runs and machines, and overridable by datasets or situations through `uc_deduction_random_draw` and `uc_deduction_type_random_draw`. Every simulation built from data gets the hashed draws, including a constituency or local authority filtered from the national data, and a benefit unit gets the same draw there as in the national run. Household situations get no deductions unless set explicitly.

policyengine_uk/data/filter_dataset.py

@@ -20,6 +20,12 @@ def filter_dataset(
     This function creates a new dataset containing only the specified household
     and the associated benefit units and people within that household.
 
+    Values imputed across the whole population (months_since_last_birthday and
+    attends_private_school) are taken from sim for the given year and carried
+    in as inputs, so the extract keeps them in later years too. Private school
+    attendance ranks incomes, so the first extract from a simulation computes
+    its taxes and benefits.
+
     Parameters
     ----------
     sim : Microsimulation
@@ -37,17 +43,17 @@ def filter_dataset(
     dataset: UKSingleYearDataset = sim.dataset[year]
     new_dataset = dataset.copy()
     person = new_dataset.person
-    if "months_since_last_birthday" not in person.columns:
-        # Birthdays are spread over the year across the whole population, so
-        # carry each person's place across rather than recompute it for one
-        # household.
-        months = pd.Series(
-            np.asarray(sim.calculate("months_since_last_birthday", year)),
-            index=np.asarray(sim.calculate("person_id", year)),
-        )
-        person = person.assign(
-            months_since_last_birthday=months.loc[person.person_id].values
-        )
+    # These are imputed across the whole population: birthdays are spread over
+    # the year, and private school attendance follows each household's income
+    # percentile (one household alone would rank at the 100th). Carry each
+    # person's value across rather than recompute it for one household.
+    for variable in ("months_since_last_birthday", "attends_private_school"):
+        if variable not in person.columns:
+            values = pd.Series(
+                np.asarray(sim.calculate(variable, year)),
+                index=np.asarray(sim.calculate("person_id", year)),
+            )
+            person = person.assign(**{variable: values.loc[person.person_id].values})
     new_dataset.person = person[person.person_household_id == household_id]
     new_dataset.household = new_dataset.household[
         new_dataset.household.household_id == household_id

policyengine_uk/simulation.py

@@ -107,8 +107,10 @@ class Simulation(CoreSimulation):
     dataset = None
     # True when built from survey or other microdata rather than a situation
     # dictionary. Variables that impute unobserved detail across a population
-    # (such as months_since_last_birthday) read it; unlike the sum of weights,
-    # it stays true for a region or constituency filtered from the data.
+    # (such as months_since_last_birthday) read it through
+    # utils.data_source.built_from_data; unlike the sum of weights, it stays
+    # true for a region or constituency filtered from the data. The builders
+    # set it, so rebuilding a simulation in place keeps it right.
     built_from_dataset: bool = False
 
     def __init__(
@@ -181,7 +183,6 @@ def __init__(
             self.build_from_dataset_source(get_default_dataset_url())
         else:
             raise ValueError(f"Unsupported dataset type: {dataset.__class__}")
-        self.built_from_dataset = situation is None
 
         # Universal Credit reform (July 2025). Needs closer integration in the baseline,
         # but adding here for ease of toggling on/off via the 'active' parameter.
@@ -270,6 +271,7 @@ def build_from_situation(self, situation: Dict) -> None:
         Args:
             situation: Dictionary describing household composition and characteristics
         """
+        self._start_new_population(built_from_dataset=False)
         self.build_from_populations(self.tax_benefit_system.instantiate_entities())
         from policyengine_core.simulations.simulation_builder import (
             SimulationBuilder,
@@ -519,6 +521,17 @@ def build_from_multi_year_dataset(self, dataset: UKMultiYearDataset) -> None:
 
         self.dataset = dataset
 
+    def _start_new_population(self, built_from_dataset: bool) -> None:
+        """Record the new population's source, and drop what the simulation
+        cached for the previous one, so rebuilding in place (e.g. a clone)
+        never reads arrays sized for the old population."""
+        self.built_from_dataset = built_from_dataset
+        if getattr(self, "_fast_cache", None) is not None:
+            self._fast_cache = {}
+        if getattr(self, "_user_input_keys", None) is not None:
+            # A clone shares this set with its original, so replace it.
+            self._user_input_keys = set()
+
     def build_from_ids(
         self,
         person_id: np.ndarray,
@@ -536,6 +549,8 @@ def build_from_ids(
             benunit_id: Array of benefit unit IDs
             household_id: Array of household IDs
         """
+        # Every data source (DataFrame, dataset, file, URL) builds through here.
+        self._start_new_population(built_from_dataset=True)
         from policyengine_core.simulations.simulation_builder import (
             SimulationBuilder,
         )  # Import here to avoid circular dependency

policyengine_uk/tests/policy/baseline/contrib/labour/attends_private_school.yaml

@@ -1,41 +1,32 @@
-- name: Attends private school returns False when attendance rate is 0%
+- name: A household situation attends no private school unless set, even at the top of the income distribution
   period: 2024
   input:
-    gov.simulation.private_school_vat.private_school_attendance_rate.100: 0
-    gov.simulation.private_school_vat.private_school_attendance_rate.95: 0
     people:
-      adult: 
+      adult:
         age: 25
       child:
         age: 10
+        attends_private_school_random_draw: 0
     households:
       household:
-        household_weight: 0.001
+        household_weight: 1_000_000_000
         household_market_income: 1_000_000_000
         household_benefits: 0
-        members: [
-          adult,
-          child
-        ]
+        members: [adult, child]
   output:
     attends_private_school: [False, False]
 
-- name: Attends private school successfully returns False when household weights are 0
+- name: A household situation keeps attendance it sets
   period: 2024
   input:
     people:
-      adult: 
+      adult:
         age: 25
       child:
         age: 10
+        attends_private_school: true
     households:
       household:
-        household_weight: 0
-        household_market_income: 1_000_000_000
-        household_benefits: 0
-        members: [
-          adult,
-          child
-        ]
+        members: [adult, child]
   output:
-    attends_private_school: [False, False]
+    attends_private_school: [False, True]