Generate random seeds in dataset for reproducible stochastic simulations

This change moves random number generation from policyengine-uk into the
dataset generation, following the pattern established in policyengine-us-data.

Changes:
- Add random seed generation in FRS dataset for 11 independent random decisions
  (4 person-level, 4 benunit-level, 3 household-level seeds)
- Update SPI dataset to use seeded generator for age assignment
- Update income imputation to use seeded generator for age assignment
- Update capital gains imputation to use seeded generator for quantile sampling
- Update childcare assumptions to use seeded generator

All random generation now uses np.random.default_rng(seed=100) for full
reproducibility across dataset builds.

Each seed corresponds to a specific independent random decision to avoid
artificial correlations between unrelated stochastic processes.

Related: policyengine-uk PR (must be merged after this)

Author: MaxGhenis

policyengine_uk_data/datasets/frs.py

@@ -751,48 +751,41 @@ def determine_education_level(fted_val, typeed2_val, age_val):
         paragraph_3 | paragraph_4 | paragraph_5
     )
 
-    # Add random variables which are for now in policyengine-uk.
-
-    RANDOM_VARIABLES = [
-        "would_evade_tv_licence_fee",
-        "would_claim_pc",
-        "would_claim_uc",
-        "would_claim_child_benefit",
-        "main_residential_property_purchased_is_first_home",
-        "household_owns_tv",
-        "is_higher_earner",
-        "attends_private_school",
-    ]
-
-    for variable in RANDOM_VARIABLES:
-        value = sim.calculate(variable).values
-        entity = sim.tax_benefit_system.variables[variable].entity.key
-        if entity == "person":
-            pe_person[variable] = value
-        elif entity == "household":
-            pe_household[variable] = value
-        elif entity == "benunit":
-            pe_benunit[variable] = value
-
-    # Add Tax-Free Childcare assumptions
-
-    count_benunits = len(pe_benunit)
-
-    extended_would_claim = np.random.random(count_benunits) < 0.812
-    tfc_would_claim = np.random.random(count_benunits) < 0.586
-    universal_would_claim = np.random.random(count_benunits) < 0.563
-    targeted_would_claim = np.random.random(count_benunits) < 0.597
+    # Add random seed variables for stochastic simulation
+    # Each seed is for a specific independent random decision to avoid artificial correlations
+    # Random seeds are generated once during dataset creation and stored
+
+    generator = np.random.default_rng(seed=100)
+
+    # Person-level seeds
+    pe_person["is_disabled_for_benefits_seed"] = generator.random(len(pe_person))
+    pe_person["marriage_allowance_take_up_seed"] = generator.random(len(pe_person))
+    pe_person["is_higher_earner_seed"] = generator.random(len(pe_person))
+    pe_person["attends_private_school_seed"] = generator.random(len(pe_person))
+
+    # Benefit unit-level seeds
+    pe_benunit["child_benefit_take_up_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["child_benefit_opts_out_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["pension_credit_take_up_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["universal_credit_take_up_seed"] = generator.random(len(pe_benunit))
+
+    # Household-level seeds
+    pe_household["first_home_purchase_seed"] = generator.random(len(pe_household))
+    pe_household["household_owns_tv_seed"] = generator.random(len(pe_household))
+    pe_household["tv_licence_evasion_seed"] = generator.random(len(pe_household))
+
+    # Add childcare take-up seeds
+    # These will be used by the formulas in policyengine-uk with parameters
+    pe_benunit["tax_free_childcare_take_up_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["extended_childcare_take_up_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["universal_childcare_take_up_seed"] = generator.random(len(pe_benunit))
+    pe_benunit["targeted_childcare_take_up_seed"] = generator.random(len(pe_benunit))
 
     # Generate extended childcare hours usage values with mean 15.019 and sd 4.972
-    extended_hours_values = np.random.normal(15.019, 4.972, count_benunits)
+    extended_hours_values = generator.normal(15.019, 4.972, len(pe_benunit))
     # Clip values to be between 0 and 30 hours
     extended_hours_values = np.clip(extended_hours_values, 0, 30)
 
-    pe_benunit["would_claim_extended_childcare"] = extended_would_claim
-    pe_benunit["would_claim_tfc"] = tfc_would_claim
-    pe_benunit["would_claim_universal_childcare"] = universal_would_claim
-    pe_benunit["would_claim_targeted_childcare"] = targeted_would_claim
-
     # Add the maximum extended childcare hours usage
     pe_benunit["maximum_extended_childcare_hours_usage"] = (
         extended_hours_values

policyengine_uk_data/datasets/imputations/capital_gains.py

@@ -125,6 +125,9 @@ def loss(blend_factor):
 
     logging.info("Imputing capital gains among those with gains")
 
+    # Use seeded generator for reproducibility
+    generator = np.random.default_rng(seed=100)
+
     for i in range(len(capital_gains)):
         row = capital_gains.iloc[i]
         spline = UnivariateSpline(
@@ -136,7 +139,7 @@ def loss(blend_factor):
         upper = row.maximum_total_income
         ti_in_range = (ti >= lower) * (ti < upper)
         in_target_range = has_cg * ti_in_range > 0
-        quantiles = np.random.random(int(in_target_range.sum()))
+        quantiles = generator.random(int(in_target_range.sum()))
         pred_capital_gains = spline(quantiles)
         new_cg[in_target_range] = pred_capital_gains
 

policyengine_uk_data/datasets/imputations/income.py

@@ -51,7 +51,9 @@ def generate_spi_table(spi: pd.DataFrame):
     LOWER = np.array([0, 16, 25, 35, 45, 55, 65, 75])
     UPPER = np.array([16, 25, 35, 45, 55, 65, 75, 80])
     age_range = spi.AGERANGE
-    spi["age"] = LOWER[age_range] + np.random.rand(len(spi)) * (
+    # Use seeded generator for reproducibility
+    generator = np.random.default_rng(seed=100)
+    spi["age"] = LOWER[age_range] + generator.random(len(spi)) * (
         UPPER[age_range] - LOWER[age_range]
     )
 

policyengine_uk_data/datasets/spi.py

@@ -73,8 +73,10 @@ def create_spi(
     age_range = df.AGERANGE
 
     # Randomly assign ages in age ranges
+    # Use seeded generator for reproducibility
 
-    percent_along_age_range = np.random.rand(len(df))
+    generator = np.random.default_rng(seed=100)
+    percent_along_age_range = generator.random(len(df))
     min_age = np.array([AGE_RANGES[age][0] for age in age_range])
     max_age = np.array([AGE_RANGES[age][1] for age in age_range])
     person["age"] = (