From 56e0e89d15607dd44a781530dccdad118b711e03 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:37:24 -0400 Subject: [PATCH 01/11] Draw SPI incomes within earnings groups set by FRS employment status The SPI income model drew every SPI-synthetic row's six incomes from age, gender and region alone, while the row kept its FRS donor's employment status. Children drew pay, the unemployed and retired drew earnings, and self-employed rows rarely drew a profit. Fit one QRF per earnings group (employee, self-employed, both, neither) on the SPI records in that group, using PAY and the SPI self-employment indicator SEINC_NUM, and draw each FRS row from the group its employment status and recorded earnings put it in. Children keep their own incomes. Co-Authored-By: Claude Opus 5.5 --- .../spi-income-earnings-groups.fixed.md | 1 + .../datasets/imputations/income.py | 282 ++++++++++++++++-- 2 files changed, 253 insertions(+), 30 deletions(-) create mode 100644 changelog.d/spi-income-earnings-groups.fixed.md diff --git a/changelog.d/spi-income-earnings-groups.fixed.md b/changelog.d/spi-income-earnings-groups.fixed.md new file mode 100644 index 00000000..724e2fca --- /dev/null +++ b/changelog.d/spi-income-earnings-groups.fixed.md @@ -0,0 +1 @@ +Draw SPI incomes for the enhanced FRS's SPI-synthetic rows within earnings groups (employee, self-employed, both, neither) set by each row's FRS employment status, so employees draw pay, the self-employed draw a trade, people out of work draw no earnings, and children keep their own (zero) incomes. diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index feafcb27..7e1d6c71 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -4,11 +4,16 @@ This module imputes detailed income components (employment, self-employment, pensions, property, savings interest, dividends) using machine learning models trained on HMRC Survey of Personal Incomes (SPI) data. + +The draw is conditioned on each person's earnings group (see +``EARNINGS_GROUPS``) as well as their age, gender and region, so the incomes +agree with the employment status the FRS donor row keeps. """ import pandas as pd import numpy as np import os +import pickle from policyengine_uk_data.storage import STORAGE_FOLDER from policyengine_uk.data import UKSingleYearDataset from policyengine_uk import Microsimulation @@ -51,6 +56,103 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: return AGE_RANGES[-1] +# Earnings groups that both surveys identify. The SPI has no ILO employment +# status and no hours, but it records each taxpayer's income sources: pay, and +# whether they file the self-employment pages of a tax return for a trade or +# partnership (SEINC_NUM, "Indicator for self-employed cases"). PROFITS is +# floored at zero, so SEINC_NUM is what finds break-even and loss-making +# traders. The FRS gives the same split through the main job's ILO status +# (EMPSTATI) and any second-job earnings. +NO_EARNINGS = "NO_EARNINGS" +EMPLOYEE = "EMPLOYEE" +SELF_EMPLOYED = "SELF_EMPLOYED" +EMPLOYEE_AND_SELF_EMPLOYED = "EMPLOYEE_AND_SELF_EMPLOYED" +EARNINGS_GROUPS = ( + NO_EARNINGS, + EMPLOYEE, + SELF_EMPLOYED, + EMPLOYEE_AND_SELF_EMPLOYED, +) +# FRS rows that take no SPI draw and keep their own values: the FRS child +# table (dependent children, including 16-19-year-olds in education), whose +# earnings the FRS build does not record, and anyone under 16. The SPI has no +# record of them as earners. +NOT_IMPUTED = "NOT_IMPUTED" + +# policyengine-uk EmploymentStatus names. +EMPLOYEE_STATUSES = ("FT_EMPLOYED", "PT_EMPLOYED") +SELF_EMPLOYED_STATUSES = ("FT_SELF_EMPLOYED", "PT_SELF_EMPLOYED") +CHILD_STATUS = "CHILD" + +# Every earnings group gets at least this share of the training sample, so the +# small self-employed groups are not fitted on a few thousand records. +MIN_GROUP_SAMPLE_SHARE = 0.1 + + +def earnings_group(has_pay, has_trade) -> np.ndarray: + """Earnings group from whether a person has pay and a trade.""" + has_pay = np.asarray(has_pay, dtype=bool) + has_trade = np.asarray(has_trade, dtype=bool) + return np.select( + [has_pay & has_trade, has_trade, has_pay], + [EMPLOYEE_AND_SELF_EMPLOYED, SELF_EMPLOYED, EMPLOYEE], + NO_EARNINGS, + ).astype(object) + + +def spi_earnings_group( + employment_income, self_employment_income, self_employed_indicator +) -> np.ndarray: + """Earnings group of SPI records. + + Pay is PAY + EPB + TAXTERM. A trade is SEINC_NUM = 1 (self-employment + pages filed, whatever the profit) or any assessable profit. + """ + has_trade = (np.asarray(self_employed_indicator) == 1) | ( + np.asarray(self_employment_income, dtype=float) > 0 + ) + return earnings_group(np.asarray(employment_income, dtype=float) > 0, has_trade) + + +def frs_earnings_group( + employment_status, age, employment_income, self_employment_income +) -> np.ndarray: + """Earnings group of FRS rows, or ``NOT_IMPUTED`` for children. + + An employee main job always draws pay and a self-employed main job always + draws a trade, whatever the FRS recorded for the donor. Earnings the FRS + records outside the main job (a second job or a side trade) add the other + source. Everyone else, retired, unemployed or inactive, draws from SPI + records with neither. + """ + status = np.asarray(employment_status, dtype=object) + has_pay = np.isin(status, EMPLOYEE_STATUSES) | ( + np.asarray(employment_income, dtype=float) > 0 + ) + has_trade = np.isin(status, SELF_EMPLOYED_STATUSES) | ( + np.asarray(self_employment_income, dtype=float) > 0 + ) + is_child = (status == CHILD_STATUS) | (np.asarray(age, dtype=float) < 16) + return np.where(is_child, NOT_IMPUTED, earnings_group(has_pay, has_trade)).astype( + object + ) + + +def earnings_group_sample_sizes( + group_weights: dict[str, float], sample_size: int +) -> dict[str, int]: + """Training records per earnings group: the group's weighted share of + ``sample_size``, but never under ``MIN_GROUP_SAMPLE_SHARE`` of it. Groups + with no weight get no records.""" + weights = {group: float(w) for group, w in group_weights.items() if w > 0} + total = sum(weights.values()) + floor = int(np.ceil(MIN_GROUP_SAMPLE_SHARE * sample_size)) + return { + group: max(int(round(sample_size * w / total)), floor) + for group, w in weights.items() + } + + def generate_spi_table( spi: pd.DataFrame, seed: int = 0, @@ -61,9 +163,13 @@ def generate_spi_table( Args: spi: Raw SPI survey data DataFrame. + seed: Seed for the age draw and the resample. + sample_size: If set, resample records with replacement, in proportion + to their weight within each earnings group, with each group's + record count from ``earnings_group_sample_sizes``. Returns: - Cleaned DataFrame with age and region mappings applied. + Cleaned DataFrame with age, region and earnings group. """ rng = np.random.default_rng(seed) age_range = spi.AGERANGE @@ -78,16 +184,25 @@ def generate_spi_table( spi[rename] = spi[SPI_RENAMES[rename]] spi["employment_income"] = spi[["PAY", "EPB", "TAXTERM"]].sum(axis=1) + spi["earnings_group"] = spi_earnings_group( + spi.employment_income, spi.self_employment_income, spi.SEINC_NUM + ) if sample_size is not None: + sizes = earnings_group_sample_sizes( + spi.groupby("earnings_group").person_weight.sum().to_dict(), + sample_size, + ) spi = pd.concat( [ - spi.sample( - sample_size, - weights=spi.person_weight, + spi[spi.earnings_group == group].sample( + sizes[group], + weights="person_weight", replace=True, - random_state=seed, - ), + random_state=seed + i, + ) + for i, group in enumerate(EARNINGS_GROUPS) + if group in sizes ] ) @@ -130,6 +245,8 @@ def generate_spi_table( "spi_release_name": SPI_RELEASE_NAME, "spi_tab_filename": SPI_TAB_FILENAME, "imputations": tuple(IMPUTATIONS), + "earnings_groups": EARNINGS_GROUPS, + "min_group_sample_share": MIN_GROUP_SAMPLE_SHARE, } INCOME_MODEL_PATH = STORAGE_FOLDER / f"income_{SPI_RELEASE_NAME}.pkl" INCOME_MODEL_SAMPLE_SIZE = 100_000 @@ -149,29 +266,98 @@ def get_income_model_metadata() -> dict: } +class EarningsGroupIncomeModel: + """One QRF per earnings group, each fitted on that group's SPI records. + + A person is drawn only from SPI records in their own earnings group, so an + employee always draws pay, a self-employed person always draws a trade + (whose profit can be zero), and someone with neither draws neither. + ``predict`` returns NaN for rows in no group (``NOT_IMPUTED``). + """ + + def __init__(self, models: dict, metadata: dict | None = None): + self.models = models + self.metadata = metadata or {} + + @property + def imputed_variables(self) -> list[str]: + return list(IMPUTATIONS) + + def predict(self, X: pd.DataFrame) -> pd.DataFrame: + groups = np.asarray(X["earnings_group"], dtype=object) + output = pd.DataFrame(np.nan, index=X.index, columns=IMPUTATIONS) + for group, model in self.models.items(): + in_group = groups == group + if in_group.any(): + inputs = pd.DataFrame( + {column: np.asarray(X[column])[in_group] for column in PREDICTORS} + ) + draws = model.predict(inputs) + output.loc[in_group, IMPUTATIONS] = draws[IMPUTATIONS].to_numpy() + return output + + def save(self, file_path): + with open(file_path, "wb") as f: + pickle.dump( + { + "models": { + group: model.model for group, model in self.models.items() + }, + "input_columns": PREDICTORS, + "metadata": self.metadata, + }, + f, + ) + + @classmethod + def load(cls, file_path): + """The cached model, or None if the file holds another format.""" + from policyengine_uk_data.utils.qrf import QRF + + with open(file_path, "rb") as f: + data = pickle.load(f) + if not isinstance(data, dict) or "models" not in data: + return None + models = {} + for group, fitted in data["models"].items(): + model = QRF() + model.model = fitted + model.input_columns = data.get("input_columns", PREDICTORS) + models[group] = model + return cls(models, data.get("metadata", {})) + + def _income_model_matches_current_release(model) -> bool: - if getattr(model, "metadata", {}) != get_income_model_metadata(): + if model is None or getattr(model, "metadata", {}) != get_income_model_metadata(): return False - cached_outputs = set(getattr(model.model, "imputed_variables", [])) - return cached_outputs == set(IMPUTATIONS) + models = getattr(model, "models", {}) + if set(models) != set(EARNINGS_GROUPS): + return False + return all( + set(getattr(group_model.model, "imputed_variables", [])) == set(IMPUTATIONS) + for group_model in models.values() + ) def save_imputation_models(): """ - Train and save income imputation model. + Train and save the income imputation model: one QRF per earnings group. Returns: - Trained QRF model for income imputation. + Trained ``EarningsGroupIncomeModel``. """ from policyengine_uk_data.utils import QRF - income = QRF() - income.metadata = get_income_model_metadata() spi = pd.read_csv(SPI_TAB_FOLDER / SPI_TAB_FILENAME, delimiter="\t") spi = generate_spi_table(spi, sample_size=get_income_model_sample_size()) - spi = spi[PREDICTORS + IMPUTATIONS] - income.fit(spi[PREDICTORS], spi[IMPUTATIONS]) + models = {} + for group in EARNINGS_GROUPS: + training = spi[spi.earnings_group == group] + model = QRF() + model.fit(training[PREDICTORS], training[IMPUTATIONS]) + models[group] = model + income = EarningsGroupIncomeModel(models, get_income_model_metadata()) income.save(INCOME_MODEL_PATH) return income @@ -180,46 +366,82 @@ def create_income_model(overwrite_existing: bool = False): """ Create or load income imputation model. - If a cached model exists and its training metadata or output columns don't - match the current SPI release and ``IMPUTATIONS`` list, the cache is - discarded and the model is retrained. + If a cached model exists and its training metadata, earnings groups or + output columns don't match the current SPI release and ``IMPUTATIONS`` + list, the cache is discarded and the model is retrained. Args: overwrite_existing: Whether to retrain model if it exists. Returns: - QRF model for income imputation. + ``EarningsGroupIncomeModel`` for income imputation. """ - from policyengine_uk_data.utils.qrf import QRF - if INCOME_MODEL_PATH.exists() and not overwrite_existing: - cached = QRF(file_path=INCOME_MODEL_PATH) + cached = EarningsGroupIncomeModel.load(INCOME_MODEL_PATH) if _income_model_matches_current_release(cached): return cached - # Cached model was trained against a different SPI release or output set. + # Cached model was trained against a different SPI release, output + # set or grouping. return save_imputation_models() +def income_model_inputs(dataset: UKSingleYearDataset) -> pd.DataFrame: + """Predictors (age, gender, region) and earnings group of each person.""" + sim = Microsimulation(dataset=dataset) + frame = sim.calculate_dataframe(["age", "gender", "region"]) + inputs = pd.DataFrame({column: np.asarray(frame[column]) for column in PREDICTORS}) + person = dataset.person + inputs["earnings_group"] = frs_earnings_group( + person.employment_status, + inputs.age, + person.employment_income, + person.self_employment_income, + ) + return inputs + + +def apply_income_draws( + person: pd.DataFrame, draws: pd.DataFrame, groups, output_variables +) -> pd.DataFrame: + """Write ``draws`` over ``output_variables`` for every person in an + earnings group (missing draws become zero). ``NOT_IMPUTED`` rows keep + their own values.""" + person = person.copy() + drawn = np.asarray(groups, dtype=object) != NOT_IMPUTED + for column in output_variables: + draw = np.nan_to_num(np.asarray(draws[column], dtype=float), nan=0.0) + own = ( + np.asarray(person[column], dtype=float) + if column in person.columns + else np.zeros(len(person)) + ) + person[column] = np.where(drawn, draw, own) + return person + + def impute_over_incomes( dataset: UKSingleYearDataset, model, output_variables: list[str] ) -> pd.DataFrame: """ Impute specified income components using trained model. + Each person draws from SPI records in their earnings group + (``frs_earnings_group``); children keep their own values. + Args: dataset: PolicyEngine UK dataset to augment with income data. + model: Fitted ``EarningsGroupIncomeModel``. output_variables: List of income components to impute. Returns: DataFrame with imputed income components. """ dataset = dataset.copy() - sim = Microsimulation(dataset=dataset) - input_df = sim.calculate_dataframe(["age", "gender", "region"]) + input_df = income_model_inputs(dataset) output_df = model.predict(input_df) - - for column in output_variables: - dataset.person[column] = output_df[column].fillna(0).values + dataset.person = apply_income_draws( + dataset.person, output_df, input_df.earnings_group, output_variables + ) # Housing costs (rent, mortgage interest, mortgage capital) used to be # rescaled here by new_income_total / original_income_total across @@ -239,8 +461,8 @@ def impute_income(dataset: UKSingleYearDataset) -> UKSingleYearDataset: Impute detailed income components using trained model. Uses SPI-trained models to predict various income sources for individuals - based on age, gender, and region. Creates a synthetic population with - the imputed income data. + based on age, gender, region and earnings group. Creates a synthetic + population with the imputed income data. Args: dataset: PolicyEngine UK dataset to augment with income data. From 313468b4282a751d82ec85d054947bd9358ffe67 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:39:12 -0400 Subject: [PATCH 02/11] Add hypothesis to the dev dependencies Same lines as #514 and #524, so whichever lands second merges cleanly. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 5 ++-- uv.lock | 73 +++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index beff7f1a..7973343c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,8 +45,9 @@ dev = [ "yaml-changelog>=0.1.7", "itables", "quantile-forest", - "build", "towncrier>=24.8.0", - + "build", + "towncrier>=24.8.0", + "hypothesis>=6.168.3", ] [tool.setuptools] diff --git a/uv.lock b/uv.lock index a3e9f44d..672737ea 100644 --- a/uv.lock +++ b/uv.lock @@ -577,6 +577,75 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/35/f4/124858007ddf3c61e9b144107304c9152fa80b5b6c168da07d86fe583cc1/huggingface_hub-1.1.5-py3-none-any.whl", hash = "sha256:e88ecc129011f37b868586bbcfae6c56868cae80cd56a79d61575426a3aa0d7d", size = 516000, upload-time = "2025-11-20T15:49:30.926Z" }, ] +[[package]] +name = "hypothesis" +version = "6.168.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sortedcontainers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/b7/13118bbc45d6d8b9d04e2de779e2a4ff23145ea39efa692b994fb874ca72/hypothesis-6.168.3.tar.gz", hash = "sha256:a43388f9067678fef6e13bdff325b6cfa6961a590498bb37f7ff31589c83bc75", size = 511022, upload-time = "2026-09-28T05:20:58.499Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/cd/746a2fc1e5e5ba34f07ea6ee1ad07525a3dde5ff4836db3a7980af49a891/hypothesis-6.168.3-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d20972ca134e652e9928ecd200966a8a2857adc2812c1d74b12a872e5cd503eb", size = 791536, upload-time = "2026-09-28T05:18:57.162Z" }, + { url = "https://files.pythonhosted.org/packages/c3/0f/7a158e377b69556c8e12c25c8fd811d0103e54c24a12dab1f3b9819f202e/hypothesis-6.168.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:eafbec09d3e87d13d8242411d1f5868b5e879f1bd6d95e233e9ff80c527bea1c", size = 787313, upload-time = "2026-09-28T05:20:18.122Z" }, + { url = "https://files.pythonhosted.org/packages/45/c6/7df9104ee359e8fcd77791781c47e70794964ca69673665e0de2a6590f4c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:73c5627497968cc62d140e9fde1ea21e12cac7b64dc419fea8286cd34ff1ad4e", size = 1124036, upload-time = "2026-09-28T05:20:04.372Z" }, + { url = "https://files.pythonhosted.org/packages/92/0f/8a61715404a73b9a82cb1f976803428d603cdea976a9c30266c72ad6a7ce/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:29dc56e6dc6eeb0aeb08ac463f847279bde2336c4786faf7ee0af4b93cd031a7", size = 1147875, upload-time = "2026-09-28T05:19:15.495Z" }, + { url = "https://files.pythonhosted.org/packages/e4/3f/2b16e95cc3b9a069afa23830a012aefe377ce4a457a227e5a5026eb827ff/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:753bb501f8560d2e321ed3b62596b4496c56f15668b0a0669231ccc3e6c4e80d", size = 1149489, upload-time = "2026-09-28T05:19:18.886Z" }, + { url = "https://files.pythonhosted.org/packages/f6/e7/d7cf6dc068bb02732b2a6e38a35a0dcb7f2137608198fc79d740f69d197c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:71ab606c472449cb872ef2a7acaec679ee4f651b7f0a477a04ebc85d8933ef8c", size = 1191992, upload-time = "2026-09-28T05:19:21.83Z" }, + { url = "https://files.pythonhosted.org/packages/5f/26/f1e5b25dec15998e8375d722825ec3dae1075cbe7cb1f74221b2312a2e33/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:26928956c54748e4dfa587333741ab750246123b93484955646ff316cca27eef", size = 1169911, upload-time = "2026-09-28T05:19:27.996Z" }, + { url = "https://files.pythonhosted.org/packages/03/36/901234a49147e6bf8a5b9e54b775706a223c2b52befea1106ea2aaf15d48/hypothesis-6.168.3-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:0bdfc53c041b61c854fa3761735736b991bc2bddafee45a361cfe4d43027c1ef", size = 1129406, upload-time = "2026-09-28T05:18:42.632Z" }, + { url = "https://files.pythonhosted.org/packages/55/c1/5fc9ff91ec9fba6bb527833386c66812524b5cc75fe7d4a4f9122ee1c420/hypothesis-6.168.3-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:650528e2b1e2a45e95c2df624d4b9364b8ac5a027ba84e1eea2bc5398dd01dcd", size = 1160399, upload-time = "2026-09-28T05:20:23.911Z" }, + { url = "https://files.pythonhosted.org/packages/9b/d2/49ad1ef5c55547c8abfd0fc571b3493806dda530fd5cb3d8263dbad86a72/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:209dc54cdb1b4d6d7020ad8d09e44c49756361b7183adacea4d6f5705085b595", size = 1299906, upload-time = "2026-09-28T05:18:58.417Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5f/f98a26094c4f462c4215807bed0f9a6b5508b27a3c329bcba1bc8adc7ac7/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:af8ca98cd11dc7f9427bb90483abcd9688d4df2e8e63893a30a2423b027ebb11", size = 1425538, upload-time = "2026-09-28T05:20:49.074Z" }, + { url = "https://files.pythonhosted.org/packages/d0/b8/c5b4406e3a6591e51a42f85469351d1d824e0a54b84167aec9b2353759fe/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:f8122cfdd0bc0ba843063effa22ac310bdebdb3a1319bf1e63f14baa119c72ab", size = 1377084, upload-time = "2026-09-28T05:20:06.618Z" }, + { url = "https://files.pythonhosted.org/packages/e0/01/0d84a8ea469d024c15f602799c0b8410fbe03d99fdb0370449b33c3c916f/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4bedbb379eab34f792af7ee9a05aae04e9c08bbb52e0f90f5a2110d4fe4b2fbd", size = 1281162, upload-time = "2026-09-28T05:19:42.768Z" }, + { url = "https://files.pythonhosted.org/packages/22/b5/ea6435038de1a795f005cb603531e8ccc053febc50180f3ab55583c9c60e/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:5a953115b9f5c95133ab2d04efffeec96e5658c3207923dca7285f7db3e6bbef", size = 1300433, upload-time = "2026-09-28T05:18:55.634Z" }, + { url = "https://files.pythonhosted.org/packages/6c/78/fa6c77d9f64b6d592e69e582d83cbce7d7e56897faa4762885a2123b0166/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6173558e676ad25ed1e20507fa4024a0d816dd90f715f77006e5a28f19109026", size = 1336273, upload-time = "2026-09-28T05:19:08.273Z" }, + { url = "https://files.pythonhosted.org/packages/f1/56/4acd0d778bab818c9ea336fcba768bf43bc70db7f84cc373579298df4641/hypothesis-6.168.3-cp310-abi3-win32.whl", hash = "sha256:dc66390fb12d80585aa9222bf538ce8b7aa22cf5d118250647355a1c9e8f62f4", size = 678211, upload-time = "2026-09-28T05:19:23.355Z" }, + { url = "https://files.pythonhosted.org/packages/b0/db/0a02b146ad1c30f9716362f2e1a379dfd6b0f7be03ac1e22c9da24c9e2a9/hypothesis-6.168.3-cp310-abi3-win_amd64.whl", hash = "sha256:92325b276360fe86c5bf71a568c0d53a6d140b0de36dcf029f9164a17803bb24", size = 684906, upload-time = "2026-09-28T05:18:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c0/958deaf726848f96f52250740bf39f13f476b068e22f73e2b358eaa07532/hypothesis-6.168.3-cp310-abi3-win_arm64.whl", hash = "sha256:3cf6f1eeaf41cd8d60cf1f88fde905ca1dd77c906929a507d6ac7f66f2ccba2a", size = 683292, upload-time = "2026-09-28T05:19:57.404Z" }, + { url = "https://files.pythonhosted.org/packages/0c/a2/6787da846d929e52fc3344d803299c45782dbeac287528af08380e984bf8/hypothesis-6.168.3-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1b230a850de63334c16654a34a2d547e0179d36b9071d4439b3e7237f6d077e7", size = 793254, upload-time = "2026-09-28T05:19:36.294Z" }, + { url = "https://files.pythonhosted.org/packages/89/ba/2893f5ca42501f4d562cba3229c8994c3a3e5fac66ef60c1e0e0b5220a5d/hypothesis-6.168.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d63b0226cd3e0d8bdd97c3384b22a21934ed4d53246c1c93575dada616672499", size = 784699, upload-time = "2026-09-28T05:18:54.194Z" }, + { url = "https://files.pythonhosted.org/packages/b7/aa/7d7349daf75b71f6f35876f8de115e974c5c04d600e86df1b2779b0883ab/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0369f5df055f96e117ab12e5f249668ff144731ab5280bc7a205fdf81f990b89", size = 1123016, upload-time = "2026-09-28T05:20:36.372Z" }, + { url = "https://files.pythonhosted.org/packages/00/f0/7774e1ea072708ea46cb25c4aeee5f9978b6c271764809069e8a6789858e/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f076bcd0f77fdcdb7797826c879099573d03a02e228ebe649ea917081b962ac0", size = 1168989, upload-time = "2026-09-28T05:20:08.44Z" }, + { url = "https://files.pythonhosted.org/packages/a7/7d/113992abed9efbd7944e3a496a382da58152dce126fddf6419ff31a0a57d/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b823ba1fcec8da730f29316d010b06d3f7e0c3828dcf91020e24c55e7d24652a", size = 1298626, upload-time = "2026-09-28T05:19:40.986Z" }, + { url = "https://files.pythonhosted.org/packages/d5/65/3659fa5e733027e5b37a27486e57f4053535e40853d23501c422bb275043/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dd2849c269d674e4618590f3b48d443bd4c06b5aef3d8d869086c2e6d213d248", size = 1335184, upload-time = "2026-09-28T05:20:20.087Z" }, + { url = "https://files.pythonhosted.org/packages/b4/6c/35aab2221b5ea65125340f236e1a765a9e9f3b28d89a389ea0033fed29f7/hypothesis-6.168.3-cp313-cp313-win_amd64.whl", hash = "sha256:3ef7d26f5789e691401d5f87eafed9bd2763f0dbe47af6d6b66012a509404766", size = 682173, upload-time = "2026-09-28T05:19:12.762Z" }, + { url = "https://files.pythonhosted.org/packages/6e/79/27b0cb56ff5d2bc92458bf6fcebdb1f0dc01f57f043a178e9f9d74498169/hypothesis-6.168.3-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:e2d4c68729a13df9af4998d2652cfb5d541c5609b88c306880a5dfb284938ec2", size = 793287, upload-time = "2026-09-28T05:20:38.522Z" }, + { url = "https://files.pythonhosted.org/packages/e8/bd/5342f95c3bc36586777ef8cb8eba87ac7acf0184210327fb0e6d8654e295/hypothesis-6.168.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:b345f818083ec99966a43ca4f7b38feb62c6920ce28572bc2bde4948a8b7eaba", size = 784857, upload-time = "2026-09-28T05:20:42.609Z" }, + { url = "https://files.pythonhosted.org/packages/4a/5f/fc774241518a0588d362680a3084275058aab86d6bc02a2c61e1073142d0/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0608c610fc002978fc5de0f471770d8817e9455cb8649a47983c9f8bccfc1897", size = 1123330, upload-time = "2026-09-28T05:19:29.735Z" }, + { url = "https://files.pythonhosted.org/packages/b0/e6/131f16775a3dca5f4fe27f0d6ad9a9851600098a54a872d163f6ebf0b664/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9dcf6448b1ecc37f2b23f2d1b3ddfc9ff6b6910614c15a642dc82f419b96037b", size = 1169178, upload-time = "2026-09-28T05:19:53.436Z" }, + { url = "https://files.pythonhosted.org/packages/06/61/c9f5bff8b73321c9fd12f2d69666c6e4861b97b60f5f24fdd6ac293007fd/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ae4f9f094041dcce5119ebd7bab71062743ab02b6e654d056b370beda78c19e2", size = 1299127, upload-time = "2026-09-28T05:20:14.486Z" }, + { url = "https://files.pythonhosted.org/packages/e4/6d/d4618f8ab12dd76c4d58ea052252759aa2bf0c4e71f4b502ba5d577e0b9f/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d479985fe73af97badfb72dc6d20c6a353e486a36f6d065f600026e0cf954a86", size = 1335390, upload-time = "2026-09-28T05:20:10.326Z" }, + { url = "https://files.pythonhosted.org/packages/95/b6/727161e17cc297df1aeea945de05c532033b66751625bec192bc5b5ff707/hypothesis-6.168.3-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:785e2c45f8c08e274b4bf1ccae97f1a4e09407e9a68e80790cfab1d17a4a45fa", size = 624291, upload-time = "2026-09-28T05:18:36.376Z" }, + { url = "https://files.pythonhosted.org/packages/77/23/2f9b506392a0b8712440bcb781abc5713be9ffac1303b6d10f3669050882/hypothesis-6.168.3-cp314-cp314-win_amd64.whl", hash = "sha256:320920b1e3dae8611eee8a03d063cf2187446f2a17c38cfb8a7fc1466f71eee2", size = 682064, upload-time = "2026-09-28T05:20:16.336Z" }, + { url = "https://files.pythonhosted.org/packages/50/93/efabfd95eb2b69c0c9fa1e1c83c8290aa2715d46baeb5a7764faf09c133a/hypothesis-6.168.3-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:35380baa981108a7f60c4eab71e46acd8d8f58440520346a1a6aba06dca7e074", size = 791875, upload-time = "2026-09-28T05:19:37.789Z" }, + { url = "https://files.pythonhosted.org/packages/1b/fc/2a0ada1623a9a33048bbf9f00a7c92513398c6b8865cba5efb854262ce30/hypothesis-6.168.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f0aaaed00438fa6673d856aaee12a4a8afbde82a9f3c5ff487cacfb064239afe", size = 783412, upload-time = "2026-09-28T05:18:48.266Z" }, + { url = "https://files.pythonhosted.org/packages/41/31/73c615d37eeb12da0209555612c32ced628d860267c23cfd2691c6207aea/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f27e6df1576bf838e7f484d4ae5cba92114497ca30afb917056bf1ea33e95675", size = 1121662, upload-time = "2026-09-28T05:19:26.225Z" }, + { url = "https://files.pythonhosted.org/packages/92/54/1893344a3b8bbf83fa7c9bb7e96f6836580213854f46671cb749f7b4706c/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3bc85014577982ec6d2e266edc5cd6e7a7d3674c648791ca983c67a52f89a0c4", size = 1167761, upload-time = "2026-09-28T05:18:49.85Z" }, + { url = "https://files.pythonhosted.org/packages/26/97/a61d82968febd25a0f08ffbbfabfaeead66f68fc5fa73f3c9ece1c36fc8f/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ccfc29505aa1cdcc254cb9cd701fd0811fef5480addd97d4127a00df023410f7", size = 1297309, upload-time = "2026-09-28T05:19:51.711Z" }, + { url = "https://files.pythonhosted.org/packages/bc/01/fe8e4cf6d39d02f4efadaa351427fa13722086cfc54a2da3579741c5177a/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32d0699566aaa93f9e97a44705d7164386f1e91de78d277f53bada35e78bcc96", size = 1334260, upload-time = "2026-09-28T05:19:00.121Z" }, + { url = "https://files.pythonhosted.org/packages/39/6a/09177ebc62f4778f94379ecade6a51299d3d8c4f4b75ba636bf3a0202364/hypothesis-6.168.3-cp314-cp314t-win_amd64.whl", hash = "sha256:28d88fa174ecbd4ecbd7bb290f06d0db084a3971c3f511ae2830b65b5e25500f", size = 681992, upload-time = "2026-09-28T05:19:48.043Z" }, + { url = "https://files.pythonhosted.org/packages/f9/f6/890bf33d608cd63348d3146a3ca83362fbf09e589c5e5230a00c403070f4/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_10_12_x86_64.whl", hash = "sha256:61f5782d807b1e6aa5c9beef1054cb2037e7d1add48a64972cd6ee447778781f", size = 791231, upload-time = "2026-09-28T05:18:43.998Z" }, + { url = "https://files.pythonhosted.org/packages/03/dd/fc36b204f8aa7437c676576d9437f0422de46aa8a729e6cd5bbb0224e759/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_11_0_arm64.whl", hash = "sha256:7aacf3cf40c7ce8f9e4924d5b57bc0b068beafdfed2347bcf160a28b4cfbba7b", size = 783171, upload-time = "2026-09-28T05:19:20.353Z" }, + { url = "https://files.pythonhosted.org/packages/5c/d2/3cd087d577db67c404991d7723e64d7cd75570cdc82a0f7329aed6f2086c/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:01768a03a30dc54df7fe457c0b34016c84d598fa00d5965ededab93ba4eb3408", size = 1121152, upload-time = "2026-09-28T05:18:39.095Z" }, + { url = "https://files.pythonhosted.org/packages/42/6f/a05a68cc66e7a22701c50bf94fb9ddeb36f4f5d0de3cc130066709eda621/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1a8a4ffc6c6e6e577f2bfbcfebf7cffbb310283ba2532c729570a0a752cebaaa", size = 1144069, upload-time = "2026-09-28T05:20:44.705Z" }, + { url = "https://files.pythonhosted.org/packages/4a/f4/fdd7a093fb2fb3a0ce24256c92ad5bf67936f4669bd06d5ceae4298150ec/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b987d73eba95183a7e59cca6d1925c588aa4307922d9852cc2b4d282a2ae4128", size = 1146692, upload-time = "2026-09-28T05:20:25.83Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/6dbd4377e935505ae4fc8ee4c7b18c69994c4bbaca015447c470218bdbee/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4569c39bd97d9573e946429ed676f3b55a7c8ed80920addd67d384a47d59b38", size = 1189480, upload-time = "2026-09-28T05:20:46.734Z" }, + { url = "https://files.pythonhosted.org/packages/ba/bf/ff288b496b690000d2686dcfa7f67855e4c5a6dddb464ed8e4880be83bb2/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6042b8707a4b25bbbfe20b110258fb7549a68951e5525608a8ce06091039e9fc", size = 1167109, upload-time = "2026-09-28T05:19:24.758Z" }, + { url = "https://files.pythonhosted.org/packages/d9/d1/3fee2bc29fc747fabbb27e174515592e390809ca465d3bbbabbcfa3235a8/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_31_riscv64.whl", hash = "sha256:f413629de94d38a7a2ad259697a6752526143cb49c2e2d7bdba19381c80693f6", size = 1126823, upload-time = "2026-09-28T05:18:52.636Z" }, + { url = "https://files.pythonhosted.org/packages/58/6a/585294392fa6a6d9719446a042a777ba98ecca265d9e48cbedc10c88df0f/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f071737e4e775bebba07e319eb645a880d1e0186b4d24bad23255e849d86d483", size = 1155773, upload-time = "2026-09-28T05:20:51.268Z" }, + { url = "https://files.pythonhosted.org/packages/cd/4d/fa283ff79996debf1ff08593e01f8b773eb8211b19b3d6d5d491317a65bf/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_aarch64.whl", hash = "sha256:f59a3912858d0609c26054aa1c474847c797937f5e0e1c77f0d3f08032e50f09", size = 1296671, upload-time = "2026-09-28T05:18:46.937Z" }, + { url = "https://files.pythonhosted.org/packages/42/05/98f9c2f628da5afabb6c3e8df5d2d299d2d50bf306b14a6e26d3954c4951/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_armv7l.whl", hash = "sha256:f09c05a23a8025dd5cad07c2ed299a46778e72d49ff0dd5bf353bcc0f7dc1580", size = 1422001, upload-time = "2026-09-28T05:19:04.188Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c9/f2109ade29a7ec284d55bca328d784743e8b13e321576fd34ee7e65f99af/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_i686.whl", hash = "sha256:44ace770bda3a0301739fc5a413c764de790c739df1f1a7218dd049f8594d9f5", size = 1374182, upload-time = "2026-09-28T05:20:27.764Z" }, + { url = "https://files.pythonhosted.org/packages/ce/6d/e05d5f72441564a3bebc71fa155deafd0fd3b015d6014ec8a00edf6e42bb/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_ppc64le.whl", hash = "sha256:03b131043608f94a2578a2a896a1a72092acb7079815eef5b8513b08447e0e62", size = 1278402, upload-time = "2026-09-28T05:20:56.332Z" }, + { url = "https://files.pythonhosted.org/packages/96/d7/6988a7f1f69c5c530687dce03a32c157e2d878e3a2e32af9269ce016b78a/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_riscv64.whl", hash = "sha256:e9784aca26eddfe99b03a0292320db742cc8a74200ef864fdce949f526f973cd", size = 1297770, upload-time = "2026-09-28T05:19:11.387Z" }, + { url = "https://files.pythonhosted.org/packages/a2/28/5bb82b60b836bd2329e1fe01ad94efc2b14dfa806f8e50ed14b795d78770/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_x86_64.whl", hash = "sha256:2b52ac363096232bebc2add117e9178d91f4248f4cdb919fd1026b5f86a4bb16", size = 1333977, upload-time = "2026-09-28T05:20:02.483Z" }, + { url = "https://files.pythonhosted.org/packages/8e/7d/d419841b8f65481ea1a50c4ba36f2670d48b8c7a51e4e8586089564cf7f6/hypothesis-6.168.3-cp315-abi3.abi3t-win32.whl", hash = "sha256:d28e3a6b511a74ce37df5274b51f36c2b274fea7365e00b16f0c81e22acd5957", size = 675394, upload-time = "2026-09-28T05:19:55.237Z" }, + { url = "https://files.pythonhosted.org/packages/d8/d2/1de6a2ad100e44817f2e9a8e8ba3eeaf4769621f4f87f2d4966c4971fcc4/hypothesis-6.168.3-cp315-abi3.abi3t-win_amd64.whl", hash = "sha256:7b9638789548361a57d984f56619ac694a914c328181d911271618409261ff4a", size = 681699, upload-time = "2026-09-28T05:19:09.957Z" }, + { url = "https://files.pythonhosted.org/packages/c0/79/be3370fd02734d6b1d950183ae9224580d19fa8d356b95348d60f1e3eeb7/hypothesis-6.168.3-cp315-abi3.abi3t-win_arm64.whl", hash = "sha256:65d78e4357ec48ed2c67825f06740ee3599be4cfe770a6092bed07108d679ac5", size = 679849, upload-time = "2026-09-28T05:19:06.884Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -1364,7 +1433,7 @@ wheels = [ [[package]] name = "policyengine-uk-data" -version = "1.56.16" +version = "1.57.4" source = { editable = "." } dependencies = [ { name = "google-auth" }, @@ -1392,6 +1461,7 @@ dependencies = [ dev = [ { name = "build" }, { name = "furo" }, + { name = "hypothesis" }, { name = "itables" }, { name = "l0-python" }, { name = "pytest" }, @@ -1410,6 +1480,7 @@ requires-dist = [ { name = "google-auth" }, { name = "google-cloud-storage" }, { name = "huggingface-hub" }, + { name = "hypothesis", marker = "extra == 'dev'", specifier = ">=6.168.3" }, { name = "itables", marker = "extra == 'dev'" }, { name = "l0-python", marker = "extra == 'dev'", specifier = ">=0.4.0" }, { name = "microcalibrate", specifier = ">=0.18.0" }, From 708ab3e06521e3790d73121ffb98ec35a32f08c9 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:43:30 -0400 Subject: [PATCH 03/11] Test that SPI income draws follow earnings groups Hypothesis properties for the FRS and SPI group mappings, the training sample allocation, the per-group model's draws and the donor values kept for children, plus a check on a built enhanced FRS. The cache tests write the per-group format. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/tests/test_spi_build.py | 55 +- .../tests/test_spi_income_earnings_groups.py | 482 ++++++++++++++++++ 2 files changed, 504 insertions(+), 33 deletions(-) create mode 100644 policyengine_uk_data/tests/test_spi_income_earnings_groups.py diff --git a/policyengine_uk_data/tests/test_spi_build.py b/policyengine_uk_data/tests/test_spi_build.py index efb37f81..dc307442 100644 --- a/policyengine_uk_data/tests/test_spi_build.py +++ b/policyengine_uk_data/tests/test_spi_build.py @@ -66,6 +66,7 @@ "MCAS", "BPADUE", "MAIND", + "SEINC_NUM", ] @@ -386,6 +387,24 @@ def __init__(self, path): ] +def _write_income_model_cache(path, income_module, metadata): + """A cache in the per-earnings-group format, with stub fitted models.""" + with path.open("wb") as f: + pickle.dump( + { + "models": { + group: SimpleNamespace( + imputed_variables=list(income_module.IMPUTATIONS) + ) + for group in income_module.EARNINGS_GROUPS + }, + "input_columns": income_module.PREDICTORS, + "metadata": metadata, + }, + f, + ) + + def test_income_model_cache_rejects_stale_spi_release(tmp_path, monkeypatch): from policyengine_uk_data.datasets.imputations import income as income_module @@ -395,17 +414,7 @@ def test_income_model_cache_rejects_stale_spi_release(tmp_path, monkeypatch): "spi_release_name": "spi_2020_21", "spi_tab_filename": "put2021uk.tab", } - with cache.open("wb") as f: - pickle.dump( - { - "model": SimpleNamespace( - imputed_variables=list(income_module.IMPUTATIONS) - ), - "input_columns": income_module.PREDICTORS, - "metadata": stale_metadata, - }, - f, - ) + _write_income_model_cache(cache, income_module, stale_metadata) sentinel = object() monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) @@ -423,17 +432,7 @@ def test_income_model_cache_rejects_stale_sample_size(tmp_path, monkeypatch): **income_module.get_income_model_metadata(), "sample_size": income_module.TESTING_INCOME_MODEL_SAMPLE_SIZE, } - with cache.open("wb") as f: - pickle.dump( - { - "model": SimpleNamespace( - imputed_variables=list(income_module.IMPUTATIONS) - ), - "input_columns": income_module.PREDICTORS, - "metadata": stale_metadata, - }, - f, - ) + _write_income_model_cache(cache, income_module, stale_metadata) sentinel = object() monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) @@ -447,17 +446,7 @@ def test_income_model_cache_accepts_current_spi_release(tmp_path, monkeypatch): cache = tmp_path / "income_spi_2022_23.pkl" current_metadata = income_module.get_income_model_metadata() - with cache.open("wb") as f: - pickle.dump( - { - "model": SimpleNamespace( - imputed_variables=list(income_module.IMPUTATIONS) - ), - "input_columns": income_module.PREDICTORS, - "metadata": current_metadata, - }, - f, - ) + _write_income_model_cache(cache, income_module, current_metadata) monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) monkeypatch.setattr( diff --git a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py new file mode 100644 index 00000000..5225aadc --- /dev/null +++ b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py @@ -0,0 +1,482 @@ +"""SPI income draws are conditioned on earnings group. + +Invariants (Hypothesis properties unless noted): + +1. FRS rows: children (FRS child table or under 16) are ``NOT_IMPUTED``. + Everyone else is in exactly one earnings group, which has pay if and only + if the main job is as an employee or the FRS records pay, and a trade if + and only if the main job is self-employment or the FRS records a profit. +2. SPI records: the group has pay if and only if PAY + EPB + TAXTERM > 0, and + a trade if and only if SEINC_NUM = 1 or PROFITS > 0. +3. Both mappings are monotone (more income never removes a source) and give + the same answer elementwise as row by row. +4. Training sample: every group with weight gets at least + ``MIN_GROUP_SAMPLE_SHARE`` of the sample, groups without weight get none, + and groups above the floor get their weighted share. +5. ``generate_spi_table`` resamples each group only from its own records. +6. Model draws: pay is positive exactly in the groups with pay, profit is zero + in the groups without a trade, and ``NOT_IMPUTED`` rows get no draw. +7. ``apply_income_draws`` overwrites drawn rows (missing draws become zero), + leaves ``NOT_IMPUTED`` rows and other columns alone, and does not mutate + its input. +8. A cached model in another format or for another grouping is retrained. +9. On a built enhanced FRS, SPI-synthetic rows' incomes agree with their + employment status (skipped when no build is present). +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +from policyengine_uk_data.datasets.imputations import income as income_module +from policyengine_uk_data.datasets.imputations.income import ( + CHILD_STATUS, + EARNINGS_GROUPS, + EMPLOYEE, + EMPLOYEE_AND_SELF_EMPLOYED, + EMPLOYEE_STATUSES, + IMPUTATIONS, + MIN_GROUP_SAMPLE_SHARE, + NO_EARNINGS, + NOT_IMPUTED, + PREDICTORS, + SELF_EMPLOYED, + SELF_EMPLOYED_STATUSES, + EarningsGroupIncomeModel, + apply_income_draws, + earnings_group_sample_sizes, + frs_earnings_group, + generate_spi_table, + spi_earnings_group, +) + +PAY_GROUPS = {EMPLOYEE, EMPLOYEE_AND_SELF_EMPLOYED} +TRADE_GROUPS = {SELF_EMPLOYED, EMPLOYEE_AND_SELF_EMPLOYED} +NON_WORKING_STATUSES = ( + "UNEMPLOYED", + "RETIRED", + "STUDENT", + "CARER", + "LONG_TERM_DISABLED", + "SHORT_TERM_DISABLED", + "OTHER_INACTIVE", +) +STATUSES = ( + (CHILD_STATUS,) + EMPLOYEE_STATUSES + SELF_EMPLOYED_STATUSES + NON_WORKING_STATUSES +) +REGIONS = ("LONDON", "WALES", "SCOTLAND", "NORTH_EAST") + +amounts = st.one_of( + st.just(0.0), st.floats(0.01, 5e6, allow_nan=False, allow_infinity=False) +) + + +def test_statuses_are_policyengine_uk_enum_names(): + from policyengine_uk.variables.household.income.employment_status import ( + EmploymentStatus, + ) + + names = {status.name for status in EmploymentStatus} + # Every status the FRS build writes is covered: an employee, self-employed, + # child or out-of-work group. + assert set(STATUSES) == names + + +@st.composite +def frs_people(draw): + n = draw(st.integers(1, 40)) + return pd.DataFrame( + { + "employment_status": draw( + st.lists(st.sampled_from(STATUSES), min_size=n, max_size=n) + ), + "age": draw(st.lists(st.integers(0, 95), min_size=n, max_size=n)), + "employment_income": draw(st.lists(amounts, min_size=n, max_size=n)), + "self_employment_income": draw(st.lists(amounts, min_size=n, max_size=n)), + } + ) + + +def _frs_groups(people: pd.DataFrame) -> np.ndarray: + return frs_earnings_group( + people.employment_status, + people.age, + people.employment_income, + people.self_employment_income, + ) + + +@given(frs_people()) +def test_frs_group_follows_status_and_recorded_earnings(people): + groups = _frs_groups(people) + status = people.employment_status.to_numpy() + child = (status == CHILD_STATUS) | (people.age.to_numpy() < 16) + + assert (groups[child] == NOT_IMPUTED).all() + adult = groups[~child] + assert np.isin(adult, EARNINGS_GROUPS).all() + + has_pay = np.isin(status, EMPLOYEE_STATUSES) | ( + people.employment_income.to_numpy() > 0 + ) + has_trade = np.isin(status, SELF_EMPLOYED_STATUSES) | ( + people.self_employment_income.to_numpy() > 0 + ) + np.testing.assert_array_equal(np.isin(adult, list(PAY_GROUPS)), has_pay[~child]) + np.testing.assert_array_equal(np.isin(adult, list(TRADE_GROUPS)), has_trade[~child]) + + +@given(frs_people(), st.floats(0.01, 1e6), st.sampled_from(["pay", "profit"])) +def test_frs_group_monotone_in_income(people, extra, source): + before = _frs_groups(people) + column = "employment_income" if source == "pay" else "self_employment_income" + after = _frs_groups(people.assign(**{column: people[column] + extra})) + imputed = before != NOT_IMPUTED + np.testing.assert_array_equal(after == NOT_IMPUTED, ~imputed) + gained, other = ( + (PAY_GROUPS, TRADE_GROUPS) if source == "pay" else (TRADE_GROUPS, PAY_GROUPS) + ) + # The source that grew is now present; the other source is unchanged. + assert np.isin(after[imputed], list(gained)).all() + np.testing.assert_array_equal( + np.isin(after[imputed], list(other)), np.isin(before[imputed], list(other)) + ) + + +@given(frs_people()) +def test_frs_group_elementwise_equals_rowwise(people): + vectorised = _frs_groups(people) + rowwise = [ + frs_earnings_group( + [row.employment_status], + [row.age], + [row.employment_income], + [row.self_employment_income], + )[0] + for row in people.itertuples() + ] + as_lists = frs_earnings_group( + people.employment_status.tolist(), + people.age.tolist(), + people.employment_income.tolist(), + people.self_employment_income.tolist(), + ) + assert list(vectorised) == rowwise == list(as_lists) + + +@given( + st.lists( + st.tuples(amounts, amounts, st.sampled_from([-1, 0, 1])), + min_size=1, + max_size=40, + ) +) +def test_spi_group_follows_pay_and_self_employment_pages(records): + pay, profit, indicator = (np.array(column) for column in zip(*records)) + groups = spi_earnings_group(pay, profit, indicator) + assert np.isin(groups, EARNINGS_GROUPS).all() + np.testing.assert_array_equal(np.isin(groups, list(PAY_GROUPS)), pay > 0) + np.testing.assert_array_equal( + np.isin(groups, list(TRADE_GROUPS)), (indicator == 1) | (profit > 0) + ) + rowwise = [spi_earnings_group([p], [q], [i])[0] for p, q, i in records] + assert list(groups) == rowwise + + +@given( + st.dictionaries( + st.sampled_from(EARNINGS_GROUPS), + st.one_of(st.just(0.0), st.floats(1e-3, 1e8)), + min_size=1, + ), + st.integers(1, 200_000), +) +def test_group_sample_sizes(weights, sample_size): + sizes = earnings_group_sample_sizes(weights, sample_size) + positive = {group for group, w in weights.items() if w > 0} + assert set(sizes) == positive + if not positive: + return + floor = int(np.ceil(MIN_GROUP_SAMPLE_SHARE * sample_size)) + total = sum(weights[group] for group in positive) + for group, size in sizes.items(): + share = sample_size * weights[group] / total + assert size >= floor + if share > floor + 1: + assert abs(size - share) <= 0.5 + 1e-9 + assert sizes == earnings_group_sample_sizes(weights, sample_size) + + +def _raw_spi(rng: np.random.Generator, n: int) -> pd.DataFrame: + """A raw SPI-shaped frame with all four earnings groups.""" + group = rng.choice(EARNINGS_GROUPS, n) + has_pay = np.isin(group, list(PAY_GROUPS)) + has_trade = np.isin(group, list(TRADE_GROUPS)) + raw = { + column: np.zeros(n) + for column in set(income_module.SPI_RENAMES.values()) + | {"PAY", "EPB", "TAXTERM", "SEINC_NUM", "GIFTINV"} + } + raw["PAY"] = np.where(has_pay, rng.lognormal(10, 1, n), 0.0) + raw["EPB"] = np.where(has_pay & (rng.random(n) < 0.1), 500.0, 0.0) + # Three in ten traders make no assessable profit. + raw["PROFITS"] = np.where( + has_trade & (rng.random(n) < 0.7), rng.lognormal(9, 1, n), 0.0 + ) + raw["SEINC_NUM"] = has_trade.astype(int) + for column in ("INCBBS", "DIVIDENDS", "PENSION", "INCPROP", "GIFTAID", "GIFTINV"): + raw[column] = rng.exponential(1_000, n) * (rng.random(n) < 0.3) + raw["FACT"] = rng.uniform(1, 500, n) + raw["SEX"] = rng.choice([1, 2], n) + raw["GORCODE"] = rng.choice([1, 7, 10, 11], n) + raw["AGERANGE"] = rng.choice([1, 2, 3, 4, 5, 6, 7], n) + return pd.DataFrame(raw) + + +@settings(deadline=None, max_examples=25) +@given(st.integers(0, 2**31 - 1), st.integers(50, 2_000)) +def test_generate_spi_table_resamples_within_groups(seed, sample_size): + raw = _raw_spi(np.random.default_rng(seed), 600) + table = generate_spi_table(raw.copy(), seed=seed % 1000, sample_size=sample_size) + np.testing.assert_array_equal( + table.earnings_group.to_numpy(), + spi_earnings_group( + table.employment_income, table.self_employment_income, table.SEINC_NUM + ), + ) + raw_groups = spi_earnings_group( + raw.PAY + raw.EPB + raw.TAXTERM, raw.PROFITS, raw.SEINC_NUM + ) + weights = raw.FACT.groupby(raw_groups).sum().to_dict() + assert table.earnings_group.value_counts().to_dict() == ( + earnings_group_sample_sizes(weights, sample_size) + ) + + +@pytest.fixture(scope="module") +def fitted_model(): + """A real per-group QRF fitted through ``generate_spi_table``.""" + from policyengine_uk_data.utils.qrf import QRF + + table = generate_spi_table( + _raw_spi(np.random.default_rng(1), 4_000), seed=0, sample_size=2_000 + ) + models = {} + for group in EARNINGS_GROUPS: + training = table[table.earnings_group == group] + model = QRF() + model.fit(training[PREDICTORS], training[IMPUTATIONS]) + models[group] = model + return EarningsGroupIncomeModel(models, income_module.get_income_model_metadata()) + + +@st.composite +def model_inputs(draw): + n = draw(st.integers(1, 30)) + return pd.DataFrame( + { + "age": draw(st.lists(st.floats(16, 90), min_size=n, max_size=n)), + "gender": draw( + st.lists(st.sampled_from(["MALE", "FEMALE"]), min_size=n, max_size=n) + ), + "region": draw( + st.lists( + st.sampled_from(["NORTH_EAST", "LONDON", "WALES", "SCOTLAND"]), + min_size=n, + max_size=n, + ) + ), + "earnings_group": draw( + st.lists( + st.sampled_from(EARNINGS_GROUPS + (NOT_IMPUTED,)), + min_size=n, + max_size=n, + ) + ), + }, + index=draw( + st.lists(st.integers(0, 10**6), min_size=n, max_size=n, unique=True) + ), + ) + + +@settings( + deadline=None, + max_examples=40, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given(model_inputs()) +def test_model_draws_agree_with_group(fitted_model, inputs): + draws = fitted_model.predict(inputs) + assert list(draws.columns) == IMPUTATIONS + assert draws.index.equals(inputs.index) + groups = inputs.earnings_group.to_numpy() + drawn = groups != NOT_IMPUTED + assert draws[~drawn].isna().all().all() + assert draws[drawn].notna().all().all() + pay = draws.employment_income.to_numpy()[drawn] + profit = draws.self_employment_income.to_numpy()[drawn] + np.testing.assert_array_equal(pay > 0, np.isin(groups[drawn], list(PAY_GROUPS))) + assert (profit[~np.isin(groups[drawn], list(TRADE_GROUPS))] == 0).all() + assert (draws[drawn].to_numpy() >= 0).all() + + +def test_model_draws_some_profit_for_traders(fitted_model): + """The trade groups draw zero and positive profits, as the SPI has both.""" + n = 400 + inputs = pd.DataFrame( + { + "age": np.linspace(20, 80, n), + "gender": ["MALE", "FEMALE"] * (n // 2), + "region": ["LONDON"] * n, + "earnings_group": [SELF_EMPLOYED] * n, + } + ) + profit = fitted_model.predict(inputs).self_employment_income + assert 0.3 < (profit > 0).mean() < 1 + + +@st.composite +def person_and_draws(draw): + n = draw(st.integers(1, 30)) + columns = ["employment_income", "dividend_income", "rent"] + person = pd.DataFrame( + {c: draw(st.lists(amounts, min_size=n, max_size=n)) for c in columns} + ) + draws = pd.DataFrame( + { + c: draw( + st.lists(st.one_of(amounts, st.just(np.nan)), min_size=n, max_size=n) + ) + for c in IMPUTATIONS + } + ) + groups = draw( + st.lists( + st.sampled_from(EARNINGS_GROUPS + (NOT_IMPUTED,)), min_size=n, max_size=n + ) + ) + outputs = draw( + st.lists(st.sampled_from(IMPUTATIONS), min_size=1, max_size=4, unique=True) + ) + return person, draws, np.array(groups, dtype=object), outputs + + +@given(person_and_draws()) +def test_apply_income_draws(case): + person, draws, groups, outputs = case + before = person.copy() + result = apply_income_draws(person, draws, groups, outputs) + pd.testing.assert_frame_equal(person, before) + + kept = groups == NOT_IMPUTED + for column in outputs: + own = ( + before[column].to_numpy() + if column in before.columns + else np.zeros(len(before)) + ) + np.testing.assert_array_equal(result[column].to_numpy()[kept], own[kept]) + np.testing.assert_array_equal( + result[column].to_numpy()[~kept], + np.nan_to_num(draws[column].to_numpy(dtype=float)[~kept], nan=0.0), + ) + for column in set(before.columns) - set(outputs): + pd.testing.assert_series_equal(result[column], before[column]) + + +def test_old_single_model_cache_is_retrained(tmp_path, monkeypatch): + import pickle + from types import SimpleNamespace + + cache = tmp_path / "income_spi_2022_23.pkl" + old_metadata = { + key: value + for key, value in income_module.get_income_model_metadata().items() + if key not in ("earnings_groups", "min_group_sample_share") + } + with cache.open("wb") as f: + pickle.dump( + { + "model": SimpleNamespace(imputed_variables=list(IMPUTATIONS)), + "input_columns": PREDICTORS, + "metadata": old_metadata, + }, + f, + ) + sentinel = object() + monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) + monkeypatch.setattr(income_module, "save_imputation_models", lambda: sentinel) + assert income_module.create_income_model() is sentinel + + +def test_cache_missing_a_group_is_retrained(tmp_path, monkeypatch, fitted_model): + cache = tmp_path / "income_spi_2022_23.pkl" + partial = EarningsGroupIncomeModel( + {g: m for g, m in fitted_model.models.items() if g != SELF_EMPLOYED}, + income_module.get_income_model_metadata(), + ) + partial.save(cache) + sentinel = object() + monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) + monkeypatch.setattr(income_module, "save_imputation_models", lambda: sentinel) + assert income_module.create_income_model() is sentinel + + +def test_current_cache_round_trips(tmp_path, monkeypatch, fitted_model): + cache = tmp_path / "income_spi_2022_23.pkl" + fitted_model.save(cache) + monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache) + monkeypatch.setattr( + income_module, + "save_imputation_models", + lambda: pytest.fail("a current cache should be reused"), + ) + loaded = income_module.create_income_model() + inputs = pd.DataFrame( + { + "age": [30.0, 50.0, 70.0, 10.0], + "gender": ["MALE", "FEMALE", "MALE", "FEMALE"], + "region": ["LONDON", "WALES", "SCOTLAND", "LONDON"], + "earnings_group": [EMPLOYEE, SELF_EMPLOYED, NO_EARNINGS, NOT_IMPUTED], + } + ) + pd.testing.assert_frame_equal(loaded.predict(inputs), fitted_model.predict(inputs)) + + +def test_built_enhanced_frs_spi_rows_agree_with_status(enhanced_frs): + person = enhanced_frs.person + household = enhanced_frs.household.set_index("household_id") + spi = person.person_household_id.map(household.household_is_spi_synthetic) + if not spi.any(): + pytest.skip("No SPI-synthetic rows in this build") + status = person.employment_status.astype(str) + rows = person[spi.to_numpy(dtype=bool)] + rows_status = status[spi.to_numpy(dtype=bool)] + + employees = rows_status.isin(EMPLOYEE_STATUSES) + assert (rows.employment_income[employees] > 0).all() + + children = rows_status.eq(CHILD_STATUS) + assert (rows.employment_income[children] == 0).all() + assert (rows.self_employment_income[children] == 0).all() + + # The SPI self-employed group draws a profit for about nine in ten + # (zero for traders who break even or make a loss); before the groups + # it was under one in ten. + self_employed = rows_status.isin(SELF_EMPLOYED_STATUSES) + if self_employed.sum() >= 50: + assert (rows.self_employment_income[self_employed] > 0).mean() > 0.6 + + # Out of work: earnings only where the FRS donor recorded some, which + # the FRS does for almost no one. + out_of_work = rows_status.isin(NON_WORKING_STATUSES) + if out_of_work.sum() >= 50: + earning = (rows.employment_income[out_of_work] > 0) | ( + rows.self_employment_income[out_of_work] > 0 + ) + assert earning.mean() < 0.02 From 40fbfea5ba554ca583848237de1a020cf3bb650e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 14:33:33 -0400 Subject: [PATCH 04/11] Calibrate to LFS counts of employees and the self-employed With SPI incomes drawn by earnings group, the SPI-synthetic rows no longer give pay to children and people out of work. Calibration then met the HMRC counts of income-tax payers with employment income, which are annual and include part-year earners, by moving about 4m people's weight from out of work to employees (32.2m against the LFS's 29.1m for 2024). Nothing tied employment status to an official count. Add national targets for the ONS LFS employee and self-employed levels (MGRN, MGRQ annual averages, 2022-2025), counted from the FRS ILO main-job status. Move the status groups to utils/employment_status.py so the income imputation and the targets share them. Co-Authored-By: Claude Opus 5.5 --- .../datasets/imputations/income.py | 9 ++- .../targets/sources/ons_labour_market.py | 81 +++++++++++++++++++ .../utils/employment_status.py | 7 ++ 3 files changed, 93 insertions(+), 4 deletions(-) create mode 100644 policyengine_uk_data/targets/sources/ons_labour_market.py create mode 100644 policyengine_uk_data/utils/employment_status.py diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index 7e1d6c71..f1d9d130 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -25,6 +25,11 @@ ) from policyengine_uk_data.utils.stack import stack_datasets from policyengine_uk_data.utils.subsample import subsample_dataset +from policyengine_uk_data.utils.employment_status import ( + CHILD_STATUS, + EMPLOYEE_STATUSES, + SELF_EMPLOYED_STATUSES, +) SPI_TAB_FOLDER = STORAGE_FOLDER / SPI_RELEASE_NAME SPI_RENAMES = dict( @@ -79,10 +84,6 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: # record of them as earners. NOT_IMPUTED = "NOT_IMPUTED" -# policyengine-uk EmploymentStatus names. -EMPLOYEE_STATUSES = ("FT_EMPLOYED", "PT_EMPLOYED") -SELF_EMPLOYED_STATUSES = ("FT_SELF_EMPLOYED", "PT_SELF_EMPLOYED") -CHILD_STATUS = "CHILD" # Every earnings group gets at least this share of the training sample, so the # small self-employed groups are not fitted on a few thousand records. diff --git a/policyengine_uk_data/targets/sources/ons_labour_market.py b/policyengine_uk_data/targets/sources/ons_labour_market.py new file mode 100644 index 00000000..6a6f5063 --- /dev/null +++ b/policyengine_uk_data/targets/sources/ons_labour_market.py @@ -0,0 +1,81 @@ +"""ONS Labour Force Survey employment levels: employees and the self-employed. + +The FRS records each adult's ILO employment status for their main job +(EMPSTATI, `employment_status`), the same concept the LFS uses. Without these +targets nothing in the calibration ties the number of employees or +self-employed people to an official count. The HMRC counts of income-tax +payers with employment income (SPI tables 3.6 and 3.15) are annual: they +include people with pay for part of the year whose status at interview is +out of work. So calibration can meet them by moving weight from people out of +work to employees. + +The FRS child table (dependent children, including 16-19-year-olds in +non-advanced education) carries no employment status, so their jobs are not +counted here, while the LFS counts them. On the FRS 2024-25 grossing weights, +the FRS has 28.1m employees against the LFS's 29.1m for 2024. + +Source: ONS Labour market overview, series MGRN (LFS: Employees: UK: All, +aged 16 and over, seasonally adjusted) and MGRQ (LFS: Self-employed: UK: +All), annual four-quarter averages, release of 14 September 2026. +""" + +import numpy as np + +from policyengine_uk_data.targets.schema import ( + GeographicLevel, + Target, + Unit, +) +from policyengine_uk_data.utils.employment_status import ( + EMPLOYEE_STATUSES, + SELF_EMPLOYED_STATUSES, +) + +_REF = ( + "https://www.ons.gov.uk/employmentandlabourmarket/peopleinwork/" + "employmentandemployeetypes/timeseries/{series}/lms" +) + +# Annual four-quarter averages, people. +EMPLOYEES = { + 2022: 28_564_000.0, + 2023: 28_821_000.0, + 2024: 29_126_000.0, + 2025: 29_590_000.0, +} +SELF_EMPLOYED = { + 2022: 4_244_000.0, + 2023: 4_380_000.0, + 2024: 4_340_000.0, + 2025: 4_395_000.0, +} + + +def _status_count(statuses: tuple[str, ...]): + def compute(ctx, target, year) -> np.ndarray: + status = np.asarray(ctx.pe_person("employment_status")).astype(str) + return ctx.household_from_person(np.isin(status, statuses).astype(float)) + + return compute + + +def get_targets() -> list[Target]: + return [ + Target( + name=f"ons/lfs_{name}", + variable="employment_status", + source="ons", + unit=Unit.COUNT, + geographic_level=GeographicLevel.NATIONAL, + geo_code="K02000001", + geo_name="United Kingdom", + values=dict(values), + is_count=True, + reference_url=_REF.format(series=series), + custom_compute=_status_count(statuses), + ) + for name, values, series, statuses in ( + ("employees", EMPLOYEES, "mgrn", EMPLOYEE_STATUSES), + ("self_employed", SELF_EMPLOYED, "mgrq", SELF_EMPLOYED_STATUSES), + ) + ] diff --git a/policyengine_uk_data/utils/employment_status.py b/policyengine_uk_data/utils/employment_status.py new file mode 100644 index 00000000..950cbc49 --- /dev/null +++ b/policyengine_uk_data/utils/employment_status.py @@ -0,0 +1,7 @@ +"""Groups of policyengine-uk ``EmploymentStatus`` names (the FRS ILO main-job status).""" + +EMPLOYEE_STATUSES = ("FT_EMPLOYED", "PT_EMPLOYED") +SELF_EMPLOYED_STATUSES = ("FT_SELF_EMPLOYED", "PT_SELF_EMPLOYED") +# The FRS child table: dependent children, including 16-19-year-olds in +# non-advanced education. +CHILD_STATUS = "CHILD" From beccbad55d3ba71dbebcb392ac8b464f4ab87434 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 15:26:13 -0400 Subject: [PATCH 05/11] Don't fail the SPI earnings-group properties on slow input generation A loaded runner tripped Hypothesis's too_slow health check while the properties themselves held. Drop the deadline and that health check. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_spi_income_earnings_groups.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py index 5225aadc..f6233a5d 100644 --- a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py +++ b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py @@ -70,6 +70,9 @@ ) REGIONS = ("LONDON", "WALES", "SCOTLAND", "NORTH_EAST") +# Generation and QRF draws are slow on a loaded runner; that is not a failure. +RELAXED = settings(deadline=None, suppress_health_check=[HealthCheck.too_slow]) + amounts = st.one_of( st.just(0.0), st.floats(0.01, 5e6, allow_nan=False, allow_infinity=False) ) @@ -110,6 +113,7 @@ def _frs_groups(people: pd.DataFrame) -> np.ndarray: ) +@RELAXED @given(frs_people()) def test_frs_group_follows_status_and_recorded_earnings(people): groups = _frs_groups(people) @@ -130,6 +134,7 @@ def test_frs_group_follows_status_and_recorded_earnings(people): np.testing.assert_array_equal(np.isin(adult, list(TRADE_GROUPS)), has_trade[~child]) +@RELAXED @given(frs_people(), st.floats(0.01, 1e6), st.sampled_from(["pay", "profit"])) def test_frs_group_monotone_in_income(people, extra, source): before = _frs_groups(people) @@ -147,6 +152,7 @@ def test_frs_group_monotone_in_income(people, extra, source): ) +@RELAXED @given(frs_people()) def test_frs_group_elementwise_equals_rowwise(people): vectorised = _frs_groups(people) @@ -168,6 +174,7 @@ def test_frs_group_elementwise_equals_rowwise(people): assert list(vectorised) == rowwise == list(as_lists) +@RELAXED @given( st.lists( st.tuples(amounts, amounts, st.sampled_from([-1, 0, 1])), @@ -187,6 +194,7 @@ def test_spi_group_follows_pay_and_self_employment_pages(records): assert list(groups) == rowwise +@RELAXED @given( st.dictionaries( st.sampled_from(EARNINGS_GROUPS), @@ -237,7 +245,7 @@ def _raw_spi(rng: np.random.Generator, n: int) -> pd.DataFrame: return pd.DataFrame(raw) -@settings(deadline=None, max_examples=25) +@settings(deadline=None, max_examples=25, suppress_health_check=[HealthCheck.too_slow]) @given(st.integers(0, 2**31 - 1), st.integers(50, 2_000)) def test_generate_spi_table_resamples_within_groups(seed, sample_size): raw = _raw_spi(np.random.default_rng(seed), 600) @@ -307,7 +315,7 @@ def model_inputs(draw): @settings( deadline=None, max_examples=40, - suppress_health_check=[HealthCheck.function_scoped_fixture], + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], ) @given(model_inputs()) def test_model_draws_agree_with_group(fitted_model, inputs): @@ -366,6 +374,7 @@ def person_and_draws(draw): return person, draws, np.array(groups, dtype=object), outputs +@RELAXED @given(person_and_draws()) def test_apply_income_draws(case): person, draws, groups, outputs = case From fc46a560f303acbd48a3aaf26c41fc1158b9e0c6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 15:47:54 -0400 Subject: [PATCH 06/11] Report, not train on, the local HMRC employment-income counts HMRC's counts of income-tax payers with employment income by area are annual, so they include people with pay for part of the year whose FRS status at interview is out of work. Trained on, they held employees at 32m against the LFS target of 29.6m. The area amounts and the national counts by income band still train; the counts stay in the calibration logs as validation targets. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/datasets/create_datasets.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/policyengine_uk_data/datasets/create_datasets.py b/policyengine_uk_data/datasets/create_datasets.py index 29a73028..d3c1ef0b 100644 --- a/policyengine_uk_data/datasets/create_datasets.py +++ b/policyengine_uk_data/datasets/create_datasets.py @@ -7,6 +7,15 @@ logging.basicConfig(level=logging.INFO) +# Local targets the calibration logs but does not train on. HMRC's counts of +# income-tax payers with employment income by area (SPI table 3.15) are annual: +# they include people with pay for part of the year whose FRS status at +# interview is out of work, and who therefore have no pay in the FRS. Training +# on them moves weight from people out of work to employees, away from the LFS +# employee count (targets/sources/ons_labour_market.py). The area amounts of +# employment income, and the national counts by income band, still train. +VALIDATION_ONLY_LOCAL_TARGETS = ["hmrc/employment_income/count"] + def _get_positive_int_env(name: str, default: int) -> int: raw_value = os.environ.get(name) @@ -252,7 +261,7 @@ def main(): area_count=650, weight_file="parliamentary_constituency_weights.h5", dataset_key=str(frs_release.calibration_year), - excluded_training_targets=[], + excluded_training_targets=VALIDATION_ONLY_LOCAL_TARGETS, log_csv="constituency_calibration_log.csv", verbose=True, # Enable nested progress display area_name="Constituency", @@ -279,7 +288,7 @@ def main(): area_count=360, weight_file="local_authority_weights.h5", dataset_key=str(frs_release.calibration_year), - excluded_training_targets=[], + excluded_training_targets=VALIDATION_ONLY_LOCAL_TARGETS, log_csv="la_calibration_log.csv", verbose=True, # Enable nested progress display area_name="Local Authority", From 363d5f251a0236ec70981315562f46f0e800cd08 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 15:48:04 -0400 Subject: [PATCH 07/11] Describe the LFS targets and validation-only counts in the changelog Co-Authored-By: Claude Opus 5.5 --- changelog.d/spi-income-earnings-groups.fixed.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/spi-income-earnings-groups.fixed.md b/changelog.d/spi-income-earnings-groups.fixed.md index 724e2fca..af1530c1 100644 --- a/changelog.d/spi-income-earnings-groups.fixed.md +++ b/changelog.d/spi-income-earnings-groups.fixed.md @@ -1 +1 @@ -Draw SPI incomes for the enhanced FRS's SPI-synthetic rows within earnings groups (employee, self-employed, both, neither) set by each row's FRS employment status, so employees draw pay, the self-employed draw a trade, people out of work draw no earnings, and children keep their own (zero) incomes. +Draw SPI incomes for the enhanced FRS's SPI-synthetic rows within earnings groups (employee, self-employed, both, neither) set by each row's FRS employment status, so employees draw pay, the self-employed draw a trade, people out of work draw no earnings, and children keep their own (zero) incomes (#504). Calibrate to the ONS LFS counts of employees and the self-employed, and report rather than train on HMRC's local counts of taxpayers with employment income, which are annual and include part-year earners the FRS records as out of work. From 54f09741f499b4ad07ed00ebaf90b74b955fac27 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 15:56:35 -0400 Subject: [PATCH 08/11] Test the LFS employment targets Values pinned against the ONS release, year resolution, source discovery, a Hypothesis property for the household column, and a check that a built enhanced FRS lands near both counts. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_lfs_employment_targets.py | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 policyengine_uk_data/tests/test_lfs_employment_targets.py diff --git a/policyengine_uk_data/tests/test_lfs_employment_targets.py b/policyengine_uk_data/tests/test_lfs_employment_targets.py new file mode 100644 index 00000000..aaaf5f53 --- /dev/null +++ b/policyengine_uk_data/tests/test_lfs_employment_targets.py @@ -0,0 +1,125 @@ +"""ONS LFS employee and self-employed calibration targets. + +Invariants: + +1. Two national count targets, with the ONS values (held here independently + of the source module, so a wrong value is caught) for every year listed. +2. The calibration and base years resolve to a value. +3. The household column counts the household's members whose main-job status + is in the target's group, for any statuses and household layout + (Hypothesis), so the weighted column total is the weighted head count. +4. On a built enhanced FRS, the weighted counts are near the targets. + This also catches the local HMRC employment counts going back into + training (``VALIDATION_ONLY_LOCAL_TARGETS``), which moves employees up. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +from policyengine_uk_data.datasets.frs_release import CURRENT_FRS_RELEASE +from policyengine_uk_data.targets.build_loss_matrix import _resolve_value +from policyengine_uk_data.targets.registry import discover_source_modules +from policyengine_uk_data.targets.sources.ons_labour_market import get_targets +from policyengine_uk_data.utils.employment_status import ( + EMPLOYEE_STATUSES, + SELF_EMPLOYED_STATUSES, +) + +# ONS Labour market overview, 14 September 2026: MGRN and MGRQ, annual +# four-quarter averages, thousands. +ONS_THOUSANDS = { + "ons/lfs_employees": {2022: 28_564, 2023: 28_821, 2024: 29_126, 2025: 29_590}, + "ons/lfs_self_employed": {2022: 4_244, 2023: 4_380, 2024: 4_340, 2025: 4_395}, +} +STATUS_GROUPS = { + "ons/lfs_employees": EMPLOYEE_STATUSES, + "ons/lfs_self_employed": SELF_EMPLOYED_STATUSES, +} +ALL_STATUSES = ( + EMPLOYEE_STATUSES + + SELF_EMPLOYED_STATUSES + + ("CHILD", "UNEMPLOYED", "RETIRED", "STUDENT", "CARER", "OTHER_INACTIVE") +) +# Calibration only partly pulls national counts in (see the public sector +# employment test); the built-data check guards against the 4m drift seen +# without these targets, not against small misses. +BUILT_RELATIVE_TOLERANCE = 0.08 + + +def _by_name(): + return {target.name: target for target in get_targets()} + + +def test_targets_and_values(): + targets = _by_name() + assert set(targets) == set(ONS_THOUSANDS) + for name, values in ONS_THOUSANDS.items(): + target = targets[name] + assert target.is_count + assert target.source == "ons" + assert target.variable == "employment_status" + assert target.values == {year: v * 1e3 for year, v in values.items()} + + +def test_source_module_is_discovered(): + # Discovery only imports the source modules; collecting every target + # would download the other sources' tables. + modules = {module.__name__ for module in discover_source_modules()} + assert "policyengine_uk_data.targets.sources.ons_labour_market" in modules + + +@pytest.mark.parametrize( + "year", + sorted({CURRENT_FRS_RELEASE.base_year, CURRENT_FRS_RELEASE.calibration_year}), +) +def test_model_years_resolve(year): + for name, target in _by_name().items(): + assert _resolve_value(target, year) == ONS_THOUSANDS[name][year] * 1e3 + + +class _FakeContext: + def __init__(self, status, household): + self.status = np.array(status, dtype=object) + self.household = np.array(household) + + def pe_person(self, variable): + assert variable == "employment_status" + return self.status + + def household_from_person(self, values): + return np.bincount(self.household, weights=values, minlength=10) + + +@settings(deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given( + st.lists( + st.tuples(st.sampled_from(ALL_STATUSES), st.integers(0, 9)), + min_size=1, + max_size=60, + ) +) +def test_column_counts_household_members_in_group(people): + status, household = zip(*people) + ctx = _FakeContext(status, household) + for name, target in _by_name().items(): + column = target.custom_compute(ctx, target, 2025) + expected = np.zeros(10) + for s, h in people: + expected[h] += s in STATUS_GROUPS[name] + np.testing.assert_array_equal(column, expected) + + +def test_built_enhanced_frs_near_lfs(enhanced_frs, baseline): + year = CURRENT_FRS_RELEASE.calibration_year + status = baseline.calculate("employment_status", year).values.astype(str) + weight = baseline.calculate("person_weight", year).values + for name, statuses in STATUS_GROUPS.items(): + estimate = weight[np.isin(status, statuses)].sum() + target = ONS_THOUSANDS[name][year] * 1e3 + assert abs(estimate / target - 1) < BUILT_RELATIVE_TOLERANCE, ( + f"{name}: {estimate / 1e6:.2f}m against {target / 1e6:.2f}m" + ) From c3cb23753cd4b03b0de4d3ebdd061bcb98e9630d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 18:12:09 -0400 Subject: [PATCH 09/11] Keep the calibration validation loss finite when every validation target is local With VALIDATION_ONLY_LOCAL_TARGETS no national target is held out, so the national validation term was the mean of an empty set and the logged validation loss was NaN. An empty mask now contributes zero, as the local term already did. Training and the saved weights are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../test_calibrate_validation_targets.py | 95 +++++++++++++++++++ policyengine_uk_data/utils/calibrate.py | 7 +- 2 files changed, 100 insertions(+), 2 deletions(-) create mode 100644 policyengine_uk_data/tests/test_calibrate_validation_targets.py diff --git a/policyengine_uk_data/tests/test_calibrate_validation_targets.py b/policyengine_uk_data/tests/test_calibrate_validation_targets.py new file mode 100644 index 00000000..4dfba548 --- /dev/null +++ b/policyengine_uk_data/tests/test_calibrate_validation_targets.py @@ -0,0 +1,95 @@ +"""Validation-only targets in calibrate_local_areas. + +create_datasets passes VALIDATION_ONLY_LOCAL_TARGETS, which exist only in the +local matrices. The calibration must then train on the other targets, leave +the excluded one out of training, and log a finite validation loss (an empty +national validation set used to make it NaN). +""" + +from __future__ import annotations + +import importlib.util + +import numpy as np +import pandas as pd +import pytest + +if ( + importlib.util.find_spec("torch") is None + or importlib.util.find_spec("policyengine_uk") is None +): + pytest.skip( + "torch/policyengine_uk not available in test environment", + allow_module_level=True, + ) + +from policyengine_uk_data.tests.test_calibrate_save import _StubDataset + + +def _performance(weights, _m_c, _y_c, m_n, y_n, _excluded_targets): + estimate = float((weights.sum(axis=0) @ m_n).iloc[0]) + target = float(y_n.iloc[0]) + return pd.DataFrame( + { + "name": ["UK"], + "metric": ["national_total"], + "estimate": [estimate], + "target": [target], + "error": [estimate - target], + "abs_error": [abs(estimate - target)], + "rel_abs_error": [abs(estimate - target) / target], + "validation": [False], + } + ) + + +def _calibrate(tmp_path, monkeypatch, excluded, target_b): + import h5py + import torch + + from policyengine_uk_data.utils import calibrate as calibrate_module + from policyengine_uk_data.utils.calibrate import calibrate_local_areas + + monkeypatch.setattr(calibrate_module, "STORAGE_FOLDER", tmp_path) + # Household 0 feeds local target "a", household 1 feeds "b". + matrix = pd.DataFrame({"a": [1.0, 0.0], "b": [0.0, 1.0]}) + local = pd.DataFrame({"a": [100.0], "b": [target_b]}) + national = pd.DataFrame({"national_total": [1.0, 1.0]}) + + log_csv = tmp_path / f"log_{len(excluded)}_{target_b:.0f}.csv" + # Weight dropout is random; seed it so runs differ only by their inputs. + torch.manual_seed(0) + calibrate_local_areas( + dataset=_StubDataset(np.array([50.0, 50.0])), + matrix_fn=lambda _d: (matrix.copy(), local.copy(), np.ones((1, 2))), + national_matrix_fn=lambda _d: (national.copy(), pd.Series([200.0])), + area_count=1, + weight_file=f"weights_{len(excluded)}_{target_b:.0f}.h5", + dataset_key="2025", + epochs=21, + excluded_training_targets=excluded, + log_csv=log_csv, + get_performance=_performance, + verbose=True, + ) + with h5py.File(tmp_path / f"weights_{len(excluded)}_{target_b:.0f}.h5") as f: + weights = f["2025"][:] + return weights, pd.read_csv(log_csv) + + +def test_local_only_validation_target_logs_finite_validation_loss( + tmp_path, monkeypatch +): + _, log = _calibrate(tmp_path, monkeypatch, ["b"], 100.0) + assert np.isfinite(log["validation_loss"]).all() + + +def test_validation_only_target_does_not_train(tmp_path, monkeypatch): + # Changing an excluded target's value leaves the weights unchanged; + # changing it when it trains moves them. + excluded_low, _ = _calibrate(tmp_path, monkeypatch, ["b"], 100.0) + excluded_high, _ = _calibrate(tmp_path, monkeypatch, ["b"], 10_000.0) + np.testing.assert_allclose(excluded_low, excluded_high) + trained_low, _ = _calibrate(tmp_path, monkeypatch, [], 100.0) + trained_high, _ = _calibrate(tmp_path, monkeypatch, [], 10_000.0) + assert not np.allclose(trained_low, trained_high) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index a31dcc1f..f97c4e1e 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -349,8 +349,11 @@ def loss(w, validation: bool = False): mask = validation_targets_national else: mask = ~validation_targets_national - pred_national = pred_national[mask] - mse_national = torch.mean(sre(pred_national, y_national[mask])) + if mask.any(): + mse_national = torch.mean(sre(pred_national[mask], y_national[mask])) + else: + # Validation targets can all be local; an empty mean is NaN. + mse_national = pred_national.sum() * 0 else: mse_national = torch.mean(sre(pred_national, y_national)) From 88ef576a950d8a961b8e0a253c0ca7ddab9a3a5c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:32:42 -0400 Subject: [PATCH 10/11] Address review: reduced-build bounds, sample-floor wording, SPI monotonicity - Widen the reduced (TESTING) bounds on salary-sacrifice users (40% total, 45% below cap) and above-AEA gains (relative error 5). The 32-epoch build stops short of the OBR salary-sacrifice targets (4.5m at epoch 0, 5.5m at epoch 30, 7.9m at epoch 510 in the full build), and main's reduced build met the old bounds only through about 1.35m SPI-synthetic children with salary sacrifice. Seed-0 reduced builds of main and of this branch both exceed the old gains bound. Full builds keep the strict bounds and pass. - State the training-sample floor as a share of the nominal size, with the total's bounds tested. - Add an SPI monotonicity property and check mixed profits in both trade groups. - ONS release date is 15 September 2026; note how later years resolve. Co-Authored-By: Claude Opus 5.5 --- .../datasets/imputations/income.py | 16 ++++--- .../targets/sources/ons_labour_market.py | 4 +- .../tests/test_cgt_band_donors.py | 8 +++- .../tests/test_lfs_employment_targets.py | 2 +- .../tests/test_salary_sacrifice_headcount.py | 25 ++++++----- .../tests/test_spi_income_earnings_groups.py | 45 ++++++++++++++++--- 6 files changed, 73 insertions(+), 27 deletions(-) diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index f1d9d130..fdf466b7 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -80,13 +80,15 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: ) # FRS rows that take no SPI draw and keep their own values: the FRS child # table (dependent children, including 16-19-year-olds in education), whose -# earnings the FRS build does not record, and anyone under 16. The SPI has no -# record of them as earners. +# earnings this FRS build does not record, and anyone under 16. With no +# earnings or job information, there is nothing to tie a taxpayer's SPI +# incomes to. NOT_IMPUTED = "NOT_IMPUTED" -# Every earnings group gets at least this share of the training sample, so the -# small self-employed groups are not fitted on a few thousand records. +# Every earnings group gets at least this share of the nominal training sample +# size, so the small self-employed groups are not fitted on a few thousand +# records. Floors are added on top, so the sample can exceed the nominal size. MIN_GROUP_SAMPLE_SHARE = 0.1 @@ -143,8 +145,10 @@ def earnings_group_sample_sizes( group_weights: dict[str, float], sample_size: int ) -> dict[str, int]: """Training records per earnings group: the group's weighted share of - ``sample_size``, but never under ``MIN_GROUP_SAMPLE_SHARE`` of it. Groups - with no weight get no records.""" + ``sample_size``, but never under ``MIN_GROUP_SAMPLE_SHARE`` of + ``sample_size``. Groups with no weight get no records. Floored groups are + not offset elsewhere, so the total can exceed ``sample_size`` (by at most + one floor per group, plus rounding).""" weights = {group: float(w) for group, w in group_weights.items() if w > 0} total = sum(weights.values()) floor = int(np.ceil(MIN_GROUP_SAMPLE_SHARE * sample_size)) diff --git a/policyengine_uk_data/targets/sources/ons_labour_market.py b/policyengine_uk_data/targets/sources/ons_labour_market.py index 6a6f5063..bae191b8 100644 --- a/policyengine_uk_data/targets/sources/ons_labour_market.py +++ b/policyengine_uk_data/targets/sources/ons_labour_market.py @@ -16,7 +16,9 @@ Source: ONS Labour market overview, series MGRN (LFS: Employees: UK: All, aged 16 and over, seasonally adjusted) and MGRQ (LFS: Self-employed: UK: -All), annual four-quarter averages, release of 14 September 2026. +All), annual four-quarter averages, release of 15 September 2026. The loss +matrix holds the latest year's value for up to three later years and drops the +targets after that (``_resolve_value``), so add each new annual average. """ import numpy as np diff --git a/policyengine_uk_data/tests/test_cgt_band_donors.py b/policyengine_uk_data/tests/test_cgt_band_donors.py index a24e2e1d..a03b375a 100644 --- a/policyengine_uk_data/tests/test_cgt_band_donors.py +++ b/policyengine_uk_data/tests/test_cgt_band_donors.py @@ -158,6 +158,12 @@ def test_stack_cgt_band_donors(frs): # still catches order-of-magnitude pathologies without failing on # reduced-fidelity calibration noise. Full builds get the strict bounds. _REDUCED_BUILD_SLACK = 5.0 if os.environ.get("TESTING") == "1" else 1.0 +# Reduced builds sit far above HMRC's total gains: the seed-0 reduced builds +# of main and of #529 carry about £243bn and £267bn (relative errors 2.7 and +# 3.1, beyond the 0.5 x slack bound), while their full builds carry about +# £53bn. Under TESTING the total-gains check only guards against +# order-of-magnitude errors; the full build keeps the 50% bound. +_TOTAL_GAINS_RELATIVE_BOUND = 5.0 if _REDUCED_BUILD_SLACK > 1 else 0.5 def _built_with_band_donors(enhanced_frs): @@ -215,7 +221,7 @@ def test_built_total_gains(enhanced_frs): enhanced_frs = _built_with_band_donors(enhanced_frs) gains, weights = _person_gains_and_weights(enhanced_frs) total = float((gains * weights)[gains > _AEA].sum()) - assert abs(total / _HMRC_TOTAL_GAINS - 1) < 0.5 * _REDUCED_BUILD_SLACK, ( + assert abs(total / _HMRC_TOTAL_GAINS - 1) < _TOTAL_GAINS_RELATIVE_BOUND, ( f"£{total / 1e9:.1f}bn of above-AEA gains against HMRC's " f"£{_HMRC_TOTAL_GAINS / 1e9:.1f}bn " f"(relative error {abs(total / _HMRC_TOTAL_GAINS - 1):.0%})." diff --git a/policyengine_uk_data/tests/test_lfs_employment_targets.py b/policyengine_uk_data/tests/test_lfs_employment_targets.py index aaaf5f53..ae49b68a 100644 --- a/policyengine_uk_data/tests/test_lfs_employment_targets.py +++ b/policyengine_uk_data/tests/test_lfs_employment_targets.py @@ -29,7 +29,7 @@ SELF_EMPLOYED_STATUSES, ) -# ONS Labour market overview, 14 September 2026: MGRN and MGRQ, annual +# ONS Labour market overview, 15 September 2026: MGRN and MGRQ, annual # four-quarter averages, thousands. ONS_THOUSANDS = { "ons/lfs_employees": {2022: 28_564, 2023: 28_821, 2024: 29_126, 2025: 29_590}, diff --git a/policyengine_uk_data/tests/test_salary_sacrifice_headcount.py b/policyengine_uk_data/tests/test_salary_sacrifice_headcount.py index cbae0972..17a18a42 100644 --- a/policyengine_uk_data/tests/test_salary_sacrifice_headcount.py +++ b/policyengine_uk_data/tests/test_salary_sacrifice_headcount.py @@ -9,21 +9,24 @@ from policyengine_uk_data.datasets.frs_release import CURRENT_FRS_RELEASE -# The total combines below-cap and above-cap users and moves slightly with -# each generated FRS calibration refresh. Widened from 0.16 after the -# household-weight alignment fix (#436) shifted the calibration starting point -# under the reduced-epoch CI build (TESTING=1). -TOTAL_TOLERANCE = 0.20 -# The below-cap count sits right at the 15% boundary under the -# reduced-epoch CI build (observed 15.09% on one TESTING=1 run and passing -# the next), so the tolerance widens under TESTING like TOTAL_TOLERANCE -# did after #436. Full builds keep the strict 15%. -TOLERANCE = 0.25 if os.environ.get("TESTING") == "1" else 0.15 +REDUCED_BUILD = os.environ.get("TESTING") == "1" +# The 32-epoch reduced build stops well short of the OBR salary-sacrifice +# counts that calibration targets: in the seed-0 full build of #529 the users +# are 4.5m at epoch 0, 5.5m at epoch 30 and 7.9m at epoch 510. Before #529, +# reduced builds met the old bounds (20% total, 25% below cap) only because +# SPI-synthetic children were imputed salary sacrifice (about 1.35m users on +# the seed-0 reduced build of main, which held 5.0m users without them). +# #529 gives those children no pay, so the reduced bounds widen to 40% and +# 45%. They still catch a collapse; full builds keep the strict bounds. +# The total was earlier widened from 0.16 after the household-weight +# alignment fix (#436). +TOTAL_TOLERANCE = 0.40 if REDUCED_BUILD else 0.20 +TOLERANCE = 0.45 if REDUCED_BUILD else 0.15 # Widened under the reduced-epoch CI build after the benunit-table sort fix # (#462) shifted the calibration starting point (observed 20.4% on a # TESTING=1 run), following the precedent of TOTAL_TOLERANCE (#436) and # TOLERANCE above. Full builds keep the strict 20%. -ABOVE_CAP_TOLERANCE = 0.25 if os.environ.get("TESTING") == "1" else 0.20 +ABOVE_CAP_TOLERANCE = 0.25 if REDUCED_BUILD else 0.20 PERIOD = CURRENT_FRS_RELEASE.calibration_year diff --git a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py index f6233a5d..44cf55de 100644 --- a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py +++ b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py @@ -8,11 +8,13 @@ and only if the main job is self-employment or the FRS records a profit. 2. SPI records: the group has pay if and only if PAY + EPB + TAXTERM > 0, and a trade if and only if SEINC_NUM = 1 or PROFITS > 0. -3. Both mappings are monotone (more income never removes a source) and give - the same answer elementwise as row by row. +3. Both mappings are monotone (more of one income adds that source and leaves + the other alone) and give the same answer elementwise as row by row. 4. Training sample: every group with weight gets at least - ``MIN_GROUP_SAMPLE_SHARE`` of the sample, groups without weight get none, - and groups above the floor get their weighted share. + ``MIN_GROUP_SAMPLE_SHARE`` of the nominal sample size, groups without + weight get none, and groups above the floor get their weighted share. + Floors are not offset elsewhere, so the total lies between the nominal size + (less rounding) and the nominal size plus one floor per group. 5. ``generate_spi_table`` resamples each group only from its own records. 6. Model draws: pay is positive exactly in the groups with pay, profit is zero in the groups without a trade, and ``NOT_IMPUTED`` rows get no draw. @@ -194,6 +196,31 @@ def test_spi_group_follows_pay_and_self_employment_pages(records): assert list(groups) == rowwise +@RELAXED +@given( + st.lists( + st.tuples(amounts, amounts, st.sampled_from([-1, 0, 1])), + min_size=1, + max_size=40, + ), + st.floats(0.01, 1e6), + st.sampled_from(["pay", "profit"]), +) +def test_spi_group_monotone_in_income(records, extra, source): + pay, profit, indicator = (np.array(column, dtype=float) for column in zip(*records)) + before = spi_earnings_group(pay, profit, indicator) + if source == "pay": + after = spi_earnings_group(pay + extra, profit, indicator) + gained, other = PAY_GROUPS, TRADE_GROUPS + else: + after = spi_earnings_group(pay, profit + extra, indicator) + gained, other = TRADE_GROUPS, PAY_GROUPS + assert np.isin(after, list(gained)).all() + np.testing.assert_array_equal( + np.isin(after, list(other)), np.isin(before, list(other)) + ) + + @RELAXED @given( st.dictionaries( @@ -216,6 +243,9 @@ def test_group_sample_sizes(weights, sample_size): assert size >= floor if share > floor + 1: assert abs(size - share) <= 0.5 + 1e-9 + groups = len(positive) + assert sample_size - 0.5 * groups <= sum(sizes.values()) + assert sum(sizes.values()) <= sample_size + groups * (floor + 0.5) assert sizes == earnings_group_sample_sizes(weights, sample_size) @@ -333,15 +363,16 @@ def test_model_draws_agree_with_group(fitted_model, inputs): assert (draws[drawn].to_numpy() >= 0).all() -def test_model_draws_some_profit_for_traders(fitted_model): - """The trade groups draw zero and positive profits, as the SPI has both.""" +@pytest.mark.parametrize("group", [SELF_EMPLOYED, EMPLOYEE_AND_SELF_EMPLOYED]) +def test_model_draws_some_profit_for_traders(fitted_model, group): + """Both trade groups draw zero and positive profits, as the SPI has both.""" n = 400 inputs = pd.DataFrame( { "age": np.linspace(20, 80, n), "gender": ["MALE", "FEMALE"] * (n // 2), "region": ["LONDON"] * n, - "earnings_group": [SELF_EMPLOYED] * n, + "earnings_group": [group] * n, } ) profit = fitted_model.predict(inputs).self_employment_income From d39371d508c33065e636fd9dcc5f15e750177fa2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:42:18 -0400 Subject: [PATCH 11/11] Bound the reduced-build gains check on both sides A relative-error bound of 5 could never fail low: a total of zero has a relative error of 1. Reduced builds now require the total to lie within a factor of 6 of HMRC's in either direction; full builds keep the 50% bound. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_cgt_band_donors.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/policyengine_uk_data/tests/test_cgt_band_donors.py b/policyengine_uk_data/tests/test_cgt_band_donors.py index a03b375a..ec97d651 100644 --- a/policyengine_uk_data/tests/test_cgt_band_donors.py +++ b/policyengine_uk_data/tests/test_cgt_band_donors.py @@ -162,8 +162,9 @@ def test_stack_cgt_band_donors(frs): # of main and of #529 carry about £243bn and £267bn (relative errors 2.7 and # 3.1, beyond the 0.5 x slack bound), while their full builds carry about # £53bn. Under TESTING the total-gains check only guards against -# order-of-magnitude errors; the full build keeps the 50% bound. -_TOTAL_GAINS_RELATIVE_BOUND = 5.0 if _REDUCED_BUILD_SLACK > 1 else 0.5 +# order-of-magnitude errors, in either direction: the total must lie within a +# factor of 6 of HMRC's. The full build keeps the 50% bound. +_REDUCED_GAINS_FACTOR = 6.0 def _built_with_band_donors(enhanced_frs): @@ -221,11 +222,15 @@ def test_built_total_gains(enhanced_frs): enhanced_frs = _built_with_band_donors(enhanced_frs) gains, weights = _person_gains_and_weights(enhanced_frs) total = float((gains * weights)[gains > _AEA].sum()) - assert abs(total / _HMRC_TOTAL_GAINS - 1) < _TOTAL_GAINS_RELATIVE_BOUND, ( + ratio = total / _HMRC_TOTAL_GAINS + message = ( f"£{total / 1e9:.1f}bn of above-AEA gains against HMRC's " - f"£{_HMRC_TOTAL_GAINS / 1e9:.1f}bn " - f"(relative error {abs(total / _HMRC_TOTAL_GAINS - 1):.0%})." + f"£{_HMRC_TOTAL_GAINS / 1e9:.1f}bn (ratio {ratio:.2f})." ) + if _REDUCED_BUILD_SLACK > 1: + assert 1 / _REDUCED_GAINS_FACTOR < ratio < _REDUCED_GAINS_FACTOR, message + else: + assert abs(ratio - 1) < 0.5, message @pytest.mark.slow