git clone https://github.com/PolicyEngine/policyengine-uk-data
cd policyengine-uk-data

policyengine_uk_data/
├── datasets/          # Dataset definitions
│   └── frs/          # FRS enhancement
│       ├── raw_frs.py           # Raw FRS loader
│       ├── calibration.py       # Weight calibration
│       └── imputations/         # Variable imputation
│           ├── wealth.py        # WAS wealth imputation
│           ├── student_loans.py # Student loan balances
│           └── ...
└── storage/          # Data storage utilities

uv pip install policyengine-uk-data

uv pip install -e .

# In policyengine_uk_data/datasets/frs/imputations/my_variable.py

from policyengine_uk_data.datasets.frs.imputations.imputation_utils import (
    impute_from_was
)

def add_my_variable(frs, was):
    """
    Impute my_variable from WAS to FRS.

    Args:
        frs: Enhanced FRS DataFrame
        was: WAS DataFrame with target variable

    Returns:
        Enhanced FRS with imputed variable
    """
    return impute_from_was(
        donor=was,
        recipient=frs,
        target_variable='my_variable',
        common_variables=[
            'age',
            'region',
            'employment_status',
            # Add relevant predictors
        ],
        method='quantile_forest'  # Or other microimpute method
    )

# In the relevant module
RENAMES = {
    "was_variable_name": "standardized_name",
    "frs_variable_name": "standardized_name",
}

# policyengine_uk_data/datasets/frs/imputations/student_loans.py

def add_student_loan_balance(frs, was):
    """
    Impute student loan balances from WAS to FRS.

    WAS contains:
    - total_loans: All loan balances
    - total_loans_exc_slc: Loans excluding student loans

    Derived variable:
    - student_loan_balance = total_loans - total_loans_exc_slc
    """
    return impute_from_was(
        donor=was,
        recipient=frs,
        target_variable='student_loan_balance',
        common_variables=[
            'age',
            'highest_qualification',
            'region',
            'employment_status',
            'income'
        ],
        method='quantile_forest'
    )

make test

# Or pytest directly
pytest policyengine_uk_data/tests/ -v

# Check if imputation was added
pytest policyengine_uk_data/tests/test_imputations.py::test_student_loan_imputation

# Compare imputed FRS distribution to WAS source
import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.hist(was['my_variable'], bins=50)
ax1.set_title('WAS (source)')
ax2.hist(frs_imputed['my_variable'], bins=50)
ax2.set_title('FRS (imputed)')

# Check population-weighted totals match administrative data
weighted_total = (frs_imputed['my_variable'] * frs_imputed['weight']).sum()
print(f"Imputed total: {weighted_total:,.0f}")
# Compare to known UK aggregate

# Verify relationships are preserved
# E.g., student loan balance by age and qualification
frs_imputed.groupby(['age_band', 'qualification'])['student_loan_balance'].mean()

# Most common: direct variable imputation
def add_variable(frs, was):
    return impute_from_was(
        donor=was,
        recipient=frs,
        target_variable='my_var',
        common_variables=['age', 'income', 'region']
    )

# When WAS has components but not the exact variable
def add_derived_variable(frs, was):
    # First derive the variable in WAS
    was['net_wealth'] = was['total_assets'] - was['total_debts']

    # Then impute
    return impute_from_was(
        donor=was,
        recipient=frs,
        target_variable='net_wealth',
        common_variables=['age', 'income', 'region']
    )

# Impute several related variables together
def add_wealth_components(frs, was):
    variables = [
        'property_wealth',
        'financial_wealth',
        'pension_wealth',
        'debt'
    ]

    for var in variables:
        frs = impute_from_was(
            donor=was,
            recipient=frs,
            target_variable=var,
            common_variables=['age', 'income', 'region']
        )

    return frs

1. Load raw FRS
   ↓
2. Add WAS imputations (wealth, student loans, etc.)
   ↓
3. Calibrate weights to administrative benchmarks
   ↓
4. Validate against known UK totals
   ↓
5. Package for policyengine-uk
   ↓
6. Use for UK policy simulations

from policyengine_uk import Microsimulation

# Uses enhanced FRS under the hood
sim = Microsimulation()
sim.calculate('student_loan_repayment', period='2026')
# Uses imputed student_loan_balance variable