Add maintenance loan calibration targets

Author: MaxGhenis

changelog.d/maintenance-loan-targets.added.md

@@ -0,0 +1 @@
+Add Student Loans Company maintenance-loan recipient-count and spend targets for England full-time undergraduates.

policyengine_uk_data/datasets/imputations/wealth.py

@@ -6,6 +6,7 @@
 Survey (WAS) data.
 """
 
+import numpy as np
 import pandas as pd
 from policyengine_uk_data.storage import STORAGE_FOLDER
 from policyengine_uk.data import UKSingleYearDataset
@@ -160,6 +161,93 @@ def _wealth_model_outputs_are_current(model: QRF) -> bool:
     return list(trained_outputs) == IMPUTE_VARIABLES
 
 
+def _person_column(person: pd.DataFrame, name: str, default) -> pd.Series:
+    if name in person:
+        return person[name]
+    return pd.Series(default, index=person.index)
+
+
+def _allocate_student_loan_balance_to_people(
+    household_balances: pd.Series,
+    person: pd.DataFrame,
+) -> np.ndarray:
+    """
+    Allocate household-imputed student loan balances to plausible holders.
+
+    The WAS target is household-level, but `student_loan_balance` is a person-
+    level input in `policyengine-uk`. We therefore allocate each household's
+    imputed balance to the most plausible holder set in priority order:
+    current repayers, reported borrowers, tertiary-qualified adults, current
+    tertiary students, then working-age adults as a final fallback.
+    """
+    balances = np.zeros(len(person), dtype=float)
+    if len(person) == 0:
+        return balances
+
+    age = (
+        pd.to_numeric(_person_column(person, "age", 0), errors="coerce")
+        .fillna(0)
+        .to_numpy()
+    )
+    repayments = (
+        pd.to_numeric(
+            _person_column(person, "student_loan_repayments", 0), errors="coerce"
+        )
+        .fillna(0)
+        .to_numpy()
+    )
+    reported_loans = (
+        pd.to_numeric(_person_column(person, "student_loans", 0), errors="coerce")
+        .fillna(0)
+        .to_numpy()
+    )
+    current_education = (
+        _person_column(person, "current_education", "NOT_IN_EDUCATION")
+        .fillna("NOT_IN_EDUCATION")
+        .astype(str)
+        .to_numpy()
+    )
+    highest_education = (
+        _person_column(person, "highest_education", "UPPER_SECONDARY")
+        .fillna("UPPER_SECONDARY")
+        .astype(str)
+        .to_numpy()
+    )
+
+    group_indices = person.groupby("person_household_id").indices
+
+    for household_id, household_balance in household_balances.items():
+        if household_balance <= 0 or household_id not in group_indices:
+            continue
+
+        idx = np.asarray(group_indices[household_id], dtype=int)
+        repayer_mask = repayments[idx] > 0
+        borrower_mask = reported_loans[idx] > 0
+        tertiary_grad_mask = highest_education[idx] == "TERTIARY"
+        current_student_mask = current_education[idx] == "TERTIARY"
+        working_age_mask = (age[idx] >= 18) & (age[idx] <= 55)
+
+        for mask in (
+            repayer_mask,
+            borrower_mask,
+            tertiary_grad_mask,
+            current_student_mask,
+            working_age_mask,
+            np.ones(len(idx), dtype=bool),
+        ):
+            if mask.any():
+                chosen = idx[mask]
+                break
+
+        if repayer_mask.any() and np.sum(repayments[idx][repayer_mask]) > 0:
+            weights = repayments[idx][repayer_mask]
+            balances[idx[repayer_mask]] += household_balance * (weights / weights.sum())
+        else:
+            balances[chosen] += household_balance / len(chosen)
+
+    return balances
+
+
 def save_imputation_models():
     """
     Train and save wealth imputation model.
@@ -213,7 +301,8 @@ def impute_wealth(dataset: UKSingleYearDataset) -> UKSingleYearDataset:
         dataset: PolicyEngine UK dataset to augment with wealth data.
 
     Returns:
-        Dataset with imputed wealth variables added to household table.
+        Dataset with household wealth variables added to the household table and
+        `student_loan_balance` allocated to people.
     """
     dataset = dataset.copy()
 
@@ -229,6 +318,12 @@ def impute_wealth(dataset: UKSingleYearDataset) -> UKSingleYearDataset:
     output_df = model.predict(input_df)
 
     for column in output_df.columns:
+        if column == "student_loan_balance":
+            dataset.person[column] = _allocate_student_loan_balance_to_people(
+                household_balances=output_df[column].clip(lower=0),
+                person=dataset.person,
+            )
+            continue
         dataset.household[column] = output_df[column].values
 
     dataset.validate()

policyengine_uk_data/targets/sources/slc.py

@@ -5,13 +5,19 @@
 Two target types are exposed:
 - `above_threshold`: borrowers liable to repay and earning above threshold
 - `liable`: all borrowers liable to repay, including below-threshold holders
+- `maintenance_loan`: full-time undergraduate England maintenance-loan
+  recipient counts and total amount paid
 
 Source: Explore Education Statistics — Student loan forecasts for England,
 Table 6a: Forecast number of student borrowers liable to repay and number
 earning above repayment threshold, by product. We use the "Higher education
 total" row which sums HE full-time, HE part-time, and Advanced Learner loans.
 Academic year 20XX-YY maps to calendar year 20XX+1 (e.g., 2024-25 → 2025).
 
+Maintenance-loan targets come from Student support for higher education in
+England 2025, Table 3A: Maintenance Loans paid to full-time undergraduate
+students. Academic year 20XX/YY maps to calendar year 20XX+1.
+
 Data permalink:
 https://explore-education-statistics.service.gov.uk/data-tables/permalink/6ff75517-7124-487c-cb4e-08de6eccf22d
 """
@@ -21,6 +27,7 @@
 import re
 from functools import lru_cache
 
+import pandas as pd
 import requests
 
 from policyengine_uk_data.targets.schema import Target, Unit
@@ -30,6 +37,10 @@
     f"https://explore-education-statistics.service.gov.uk"
     f"/data-tables/permalink/{_PERMALINK_ID}"
 )
+_MAINTENANCE_LOAN_URL = (
+    "https://assets.publishing.service.gov.uk/media/"
+    "691d9e662c6b98ecdbc5003f/slcsp052025.xlsx"
+)
 _TESTING_DATA = {
     "plan_2": {
         "above_threshold": {
@@ -68,6 +79,36 @@
         },
     },
 }
+_MAINTENANCE_LOAN_TESTING_DATA = {
+    "recipients": {
+        2014: 972_830,
+        2015: 963_084,
+        2016: 986_323,
+        2017: 1_013_354,
+        2018: 1_028_438,
+        2019: 1_044_973,
+        2020: 1_055_702,
+        2021: 1_117_591,
+        2022: 1_145_289,
+        2023: 1_151_607,
+        2024: 1_154_427,
+        2025: 1_159_761,
+    },
+    "amount_paid": {
+        2014: 3_783_626_551,
+        2015: 3_784_628_482,
+        2016: 3_996_708_360,
+        2017: 4_870_158_274,
+        2018: 5_746_431_691,
+        2019: 6_555_506_426,
+        2020: 7_113_141_652,
+        2021: 7_914_340_039,
+        2022: 8_332_837_845,
+        2023: 8_594_103_415,
+        2024: 8_881_701_387,
+        2025: 8_591_659_718,
+    },
+}
 
 
 def get_snapshot_data() -> dict:
@@ -80,6 +121,13 @@ def get_snapshot_data() -> dict:
     }
 
 
+def get_maintenance_loan_snapshot_data() -> dict:
+    """Return the checked-in maintenance-loan snapshot."""
+    return {
+        key: values.copy() for key, values in _MAINTENANCE_LOAN_TESTING_DATA.items()
+    }
+
+
 @lru_cache(maxsize=1)
 def _fetch_slc_data() -> dict:
     """Fetch and parse SLC Table 6a data from Explore Education Statistics.
@@ -166,9 +214,62 @@ def parse_values(row, start_index, years):
     }
 
 
+def _row_contains_text(df: pd.DataFrame, row_index: int, text: str) -> bool:
+    row = df.iloc[row_index].dropna()
+    return any(str(value).strip() == text for value in row)
+
+
+def _find_row(df: pd.DataFrame, text: str, start: int = 0) -> int:
+    for row_index in range(start, len(df)):
+        if _row_contains_text(df, row_index, text):
+            return row_index
+    raise ValueError(f"Could not find row containing {text!r}")
+
+
+@lru_cache(maxsize=1)
+def _fetch_maintenance_loan_data() -> dict:
+    """Fetch full-time England maintenance-loan recipient counts and spend."""
+    if os.environ.get("TESTING", "0") == "1":
+        return get_maintenance_loan_snapshot_data()
+
+    df = pd.read_excel(_MAINTENANCE_LOAN_URL, sheet_name="Table 3A", header=None)
+
+    count_header_row = _find_row(df, "Number of students paid (000s) [27]")
+    count_year_row = count_header_row + 1
+    count_total_row = _find_row(df, "Grand total", start=count_year_row + 1)
+
+    amount_header_row = _find_row(df, "Amount paid (£m)")
+    amount_year_row = amount_header_row + 1
+    amount_total_row = _find_row(df, "Grand total", start=amount_year_row + 1)
+
+    year_columns = {}
+    for column, value in df.iloc[count_year_row].items():
+        if isinstance(value, str) and re.fullmatch(r"\d{4}/\d{2}", value):
+            year_columns[column] = int(value[:4]) + 1
+
+    if not year_columns:
+        raise ValueError("Could not find maintenance-loan year columns")
+
+    recipients = {}
+    amount_paid = {}
+    for column, year in year_columns.items():
+        count_value = df.iloc[count_total_row, column]
+        amount_value = df.iloc[amount_total_row, column]
+        if pd.notna(count_value):
+            recipients[year] = int(round(float(count_value) * 1_000))
+        if pd.notna(amount_value):
+            amount_paid[year] = int(round(float(amount_value) * 1_000_000))
+
+    return {
+        "recipients": recipients,
+        "amount_paid": amount_paid,
+    }
+
+
 def get_targets() -> list[Target]:
     """Generate SLC calibration targets by fetching live data."""
     slc_data = _fetch_slc_data()
+    maintenance_loan_data = _fetch_maintenance_loan_data()
 
     targets = []
 
@@ -189,4 +290,26 @@ def get_targets() -> list[Target]:
                 )
             )
 
+    targets.extend(
+        [
+            Target(
+                name="slc/maintenance_loan_recipients",
+                variable="maintenance_loan",
+                source="slc",
+                unit=Unit.COUNT,
+                is_count=True,
+                values=maintenance_loan_data["recipients"],
+                reference_url=_MAINTENANCE_LOAN_URL,
+            ),
+            Target(
+                name="slc/maintenance_loan_spend",
+                variable="maintenance_loan",
+                source="slc",
+                unit=Unit.GBP,
+                values=maintenance_loan_data["amount_paid"],
+                reference_url=_MAINTENANCE_LOAN_URL,
+            ),
+        ]
+    )
+
     return targets