From 3db0e60d04ede0ba6fa04ac0edff7d813a706904 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 22:55:16 -0400 Subject: [PATCH 1/5] Draw SPI incomes at each donor's FRS rank within their cell The SPI-synthetic copy drew each income at a random quantile of its earnings group's forest, conditioned on age, gender and region only, so a part-timer drew a full-timer's pay as often as a part-timer's. Each output is now drawn at the donor's weighted rank of the same income among FRS people in the same cell (earnings group, SPI age band, gender, region): the cells the forests condition on. Ties, mostly zeros, are put in random order and each person takes a uniform point in their slice of the cell's weight, so the quantiles are exactly uniform in every cell and the SPI's distribution within each cell is kept. The draw reimplements microimpute's chained predict at given quantiles, on 1,000 midpoints instead of microimpute 1.8's ten points between 1/11 and 10/11, which never drew from the top or bottom 9% of a cell. At microimpute's own random draws and grid it reproduces predict exactly (tested). Ranks come from the full FRS at grossing weights and also drive the FRS rows' dividend draw. Co-Authored-By: Claude Opus 5.5 --- .../datasets/imputations/income.py | 214 ++++++- .../tests/test_imputation_source_flags.py | 3 +- .../tests/test_spi_income_rank_draw.py | 551 ++++++++++++++++++ 3 files changed, 763 insertions(+), 5 deletions(-) create mode 100644 policyengine_uk_data/tests/test_spi_income_rank_draw.py diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index 799b5f35..363a1ad4 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -8,6 +8,13 @@ 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. + +Within that cell a person does not draw at a random quantile: each income is +drawn at the person's rank of the same income among FRS people in the same +cell (``draw_quantiles``). A part-time employee with low FRS pay draws low +SPI pay, a donor at the top of their cell draws from the top of the SPI, and +the SPI's distribution within each cell is kept, because the ranks are +uniform within every cell. """ import pandas as pd @@ -220,6 +227,77 @@ def generate_spi_table( "region", ] +# The SPI gives age only as a band, so the forests learn nothing within a +# band. A person's rank is taken within (earnings group, SPI age band, +# gender, region): the cells the forests condition on. Ranking over a wider +# pool would bias the draw: a London donor's national rank is higher than +# their rank in London, and London's SPI distribution is already higher. +SPI_AGE_BAND_STARTS = sorted({low for code, (low, _) in AGE_RANGES.items() if code > 0}) + + +def spi_age_band(age) -> np.ndarray: + """SPI age band (1 = 16-24, ..., 7 = 74 and over; 0 = under 16).""" + return np.searchsorted( + SPI_AGE_BAND_STARTS, np.asarray(age, dtype=float), side="right" + ) + + +def rank_cells(inputs: pd.DataFrame) -> np.ndarray: + """Integer cell of each person: earnings group, SPI age band, gender, region.""" + keys = pd.DataFrame( + { + "earnings_group": np.asarray(inputs["earnings_group"], dtype=object), + "band": spi_age_band(inputs["age"]), + "gender": np.asarray(inputs["gender"], dtype=object), + "region": np.asarray(inputs["region"], dtype=object), + } + ) + return keys.groupby(list(keys.columns), sort=True, dropna=False).ngroup().to_numpy() + + +def rank_quantiles(values, cells, weights, ids, seed: int = 0) -> np.ndarray: + """Each row's quantile, in [0, 1), in its cell's weighted distribution of + ``values``. + + Within a cell, rows are ordered by value, rows with equal values (most + often zero) in random order. Each row then holds the interval [weight of + the rows before it, that plus its own weight) as a share of the cell's + weight, and takes a uniform point in it. The intervals tile [0, 1), so + the quantiles are uniform within every cell, and a lower value always + gets a lower quantile than a higher one in the same cell. A cell of one + row gets a uniform random quantile. A cell with no weight is ranked with + equal weights. Random numbers are drawn in (cell, value, id) order, so + the result does not depend on the order of the rows. + """ + values = np.asarray(values, dtype=float) + weights = np.asarray(weights, dtype=float) + cells = pd.factorize(np.asarray(cells), sort=True)[0] + ids = np.asarray(ids) + if np.isnan(values).any() or np.isnan(weights).any(): + raise ValueError("Ranked values and weights must not be missing") + if (weights < 0).any(): + raise ValueError("Rank weights must not be negative") + n = len(values) + rng = np.random.default_rng(seed) + canonical = np.lexsort((ids, values, cells)) + tie_order, offset = np.empty(n), np.empty(n) + tie_order[canonical] = rng.random(n) + offset[canonical] = rng.random(n) + order = np.lexsort((tie_order, values, cells)) + sorted_rows = pd.DataFrame({"cell": cells[order], "weight": weights[order]}) + total = sorted_rows.groupby("cell").weight.transform("sum").to_numpy() + if (total <= 0).any(): + sorted_rows.loc[total <= 0, "weight"] = 1.0 + total = sorted_rows.groupby("cell").weight.transform("sum").to_numpy() + weight = sorted_rows.weight.to_numpy() + before = sorted_rows.groupby("cell").weight.cumsum().to_numpy() - weight + quantiles = np.empty(n) + quantiles[order] = np.minimum( + (before + offset[order] * weight) / total, np.nextafter(1.0, 0.0) + ) + return quantiles + + INCOME_COMPONENTS = [ "employment_income", "self_employment_income", @@ -245,6 +323,63 @@ def generate_spi_table( # its own policyengine-uk variable. IMPUTATIONS = INCOME_COMPONENTS + ["gift_aid", "charitable_investment_gifts"] +# Each output is drawn at the forest's conditional quantile nearest the +# person's rank, on this grid of 1,000 midpoints (0.0005 to 0.9995). The +# draw used to be microimpute 1.8's: a random pick from ten quantiles between +# 1/11 and 10/11, so no one drew from the top or bottom 9% of their cell. +DRAW_QUANTILE_GRID = (np.arange(1_000) + 0.5) / 1_000 +# Seed for the rank tie-breaks (and for the random quantiles of a draw made +# without ranks). +DRAW_SEED = 0 + + +def draw_quantile_column(variable: str) -> str: + return f"{variable}_draw_quantile" + + +def draw_at_quantiles(results, X: pd.DataFrame, quantiles: dict, grid) -> pd.DataFrame: + """Draw every output of a fitted microimpute QRF at given quantiles. + + Outputs are drawn in the model's order, each conditioned on the ones + drawn before it, as microimpute's own ``predict`` does; but where that + picks a random quantile, row ``i`` of ``variable`` takes the forest's + conditional quantile ``grid[int(quantiles[variable][i] * len(grid))]``. + With microimpute 1.8's grid and its random draws this reproduces its + ``predict`` exactly (tested). + """ + from microimpute.models.imputer import _ConstantValueModel + from microimpute.models.qrf import _get_sequential_predictors, _QRFModel + + grid = np.asarray(grid, dtype=float) + k = len(grid) + augmented, _ = results.preprocess_data_types( + X, results.original_predictors, getattr(results, "dummy_processor", None) + ) + output = pd.DataFrame(index=X.index) + for i, variable in enumerate(results.imputed_variables): + model = results.models[variable] + columns = results._get_encoded_predictors( + _get_sequential_predictors(results.predictors, results.imputed_variables, i) + ) + if isinstance(model, _ConstantValueModel): + values = np.asarray(model.predict(augmented)) + elif isinstance(model, _QRFModel): + features = augmented[columns] + if hasattr(model, "_align_features"): # microimpute >= 2 + features = model._align_features(features) + pred = np.asarray(model.qrf.predict(features, quantiles=list(grid))) + pred = pred.reshape(len(features), k) + index = np.clip( + (np.asarray(quantiles[variable], dtype=float) * k).astype(int), 0, k - 1 + ) + values = pred[np.arange(len(pred)), index] + else: + values = np.asarray(model.predict(augmented[columns], return_probs=False)) + output[variable] = values + augmented[variable] = values + augmented = results._encode_imputed_variable(augmented, variable) + return output + INCOME_MODEL_METADATA = { "spi_release_name": SPI_RELEASE_NAME, @@ -277,7 +412,10 @@ class EarningsGroupIncomeModel: 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``). + Each output is drawn at the quantile in the input's + ``draw_quantile_column(output)`` (see ``draw_quantiles``); without those + columns, at independent uniform random quantiles. ``predict`` returns NaN + for rows in no group (``NOT_IMPUTED``). """ def __init__(self, models: dict, metadata: dict | None = None): @@ -290,6 +428,15 @@ def imputed_variables(self) -> list[str]: def predict(self, X: pd.DataFrame) -> pd.DataFrame: groups = np.asarray(X["earnings_group"], dtype=object) + rng = np.random.default_rng(DRAW_SEED) + quantiles = { + variable: ( + np.asarray(X[draw_quantile_column(variable)], dtype=float) + if draw_quantile_column(variable) in X + else rng.random(len(X)) + ) + for variable in IMPUTATIONS + } output = pd.DataFrame(np.nan, index=X.index, columns=IMPUTATIONS) for group, model in self.models.items(): in_group = groups == group @@ -297,7 +444,12 @@ def predict(self, X: pd.DataFrame) -> pd.DataFrame: inputs = pd.DataFrame( {column: np.asarray(X[column])[in_group] for column in PREDICTORS} ) - draws = model.predict(inputs) + draws = draw_at_quantiles( + model.model, + inputs, + {v: q[in_group] for v, q in quantiles.items()}, + DRAW_QUANTILE_GRID, + ) output.loc[in_group, IMPUTATIONS] = draws[IMPUTATIONS].to_numpy() return output @@ -405,6 +557,43 @@ def income_model_inputs(dataset: UKSingleYearDataset) -> pd.DataFrame: return inputs +def draw_quantiles( + dataset: UKSingleYearDataset, inputs: pd.DataFrame | None = None +) -> pd.DataFrame: + """Each person's draw quantile for every output, indexed by person ID. + + The quantile is the person's rank of the same income in ``dataset``, + among people in the same cell (``rank_cells``), at household weights: + pay sets the pay draw, profit the profit draw, a private pension the + pension draw. Ties, most often at zero, are broken at random, so an + employee with no recorded pay draws from the bottom of their cell and + an output the FRS does not record (gift aid) is drawn at a uniform + random quantile. + """ + if inputs is None: + inputs = income_model_inputs(dataset) + person = dataset.person + weights = person.person_household_id.map( + dataset.household.set_index("household_id").household_weight + ) + cells = rank_cells(inputs) + return pd.DataFrame( + { + draw_quantile_column(variable): rank_quantiles( + person[variable] + if variable in person.columns + else np.zeros(len(person)), + cells, + weights, + person.person_id, + seed=DRAW_SEED + i, + ) + for i, variable in enumerate(IMPUTATIONS) + }, + index=person.person_id.to_numpy(), + ) + + def apply_income_draws( person: pd.DataFrame, draws: pd.DataFrame, groups, output_variables ) -> pd.DataFrame: @@ -425,24 +614,36 @@ def apply_income_draws( def impute_over_incomes( - dataset: UKSingleYearDataset, model, output_variables: list[str] + dataset: UKSingleYearDataset, + model, + output_variables: list[str], + quantiles: pd.DataFrame | None = None, ) -> 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. + (``frs_earnings_group``), at their draw quantiles; 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. + quantiles: Draw quantiles by person ID (``draw_quantiles``), + covering every person in ``dataset``. Defaults to ranks within + ``dataset`` itself. Returns: DataFrame with imputed income components. """ dataset = dataset.copy() input_df = income_model_inputs(dataset) + if quantiles is None: + quantiles = draw_quantiles(dataset, input_df) + person_quantiles = quantiles.loc[dataset.person.person_id.to_numpy()] + for column in person_quantiles.columns: + input_df[column] = person_quantiles[column].to_numpy() output_df = model.predict(input_df) dataset.person = apply_income_draws( dataset.person, output_df, input_df.earnings_group, output_variables @@ -506,6 +707,9 @@ def impute_income(dataset: UKSingleYearDataset) -> UKSingleYearDataset: zero_weight_copy = subsample_dataset(zero_weight_copy, 10_000) model = create_income_model() + # Ranks come from the full, weighted FRS: the copy is an unweighted + # subsample of it. + quantiles = draw_quantiles(dataset) # Impute just dividends on the original, full variable set on the copy @@ -513,6 +717,7 @@ def impute_income(dataset: UKSingleYearDataset) -> UKSingleYearDataset: zero_weight_copy, model, IMPUTATIONS, + quantiles, ) # Second-stage QRF: rewrite FRS-only variables (benefit `_reported` @@ -554,6 +759,7 @@ def impute_income(dataset: UKSingleYearDataset) -> UKSingleYearDataset: dataset, model, ["dividend_income"], + quantiles, ) zero_weight_copy.validate() diff --git a/policyengine_uk_data/tests/test_imputation_source_flags.py b/policyengine_uk_data/tests/test_imputation_source_flags.py index 2df1f9d8..5ef3f0d5 100644 --- a/policyengine_uk_data/tests/test_imputation_source_flags.py +++ b/policyengine_uk_data/tests/test_imputation_source_flags.py @@ -74,6 +74,7 @@ def test_impute_income_marks_spi_synthetic_households(monkeypatch): from policyengine_uk_data.datasets.imputations import frs_only monkeypatch.setattr(income_module, "create_income_model", lambda: object()) + monkeypatch.setattr(income_module, "draw_quantiles", lambda dataset: None) monkeypatch.setattr( income_module, "subsample_dataset", @@ -82,7 +83,7 @@ def test_impute_income_marks_spi_synthetic_households(monkeypatch): monkeypatch.setattr( income_module, "impute_over_incomes", - lambda dataset, _model, _output_variables: dataset, + lambda dataset, _model, _output_variables, _quantiles=None: dataset, ) monkeypatch.setattr( frs_only, diff --git a/policyengine_uk_data/tests/test_spi_income_rank_draw.py b/policyengine_uk_data/tests/test_spi_income_rank_draw.py new file mode 100644 index 00000000..a245cd7d --- /dev/null +++ b/policyengine_uk_data/tests/test_spi_income_rank_draw.py @@ -0,0 +1,551 @@ +"""SPI incomes are drawn at each person's FRS rank within their cell. + +Invariants (Hypothesis properties unless noted): + +1. ``rank_quantiles`` returns quantiles in [0, 1). Within a cell a lower + value never gets a higher quantile (and gets a lower one if it has + weight), and the rows' intervals (the weight + before a row, and that plus its own, over the cell's weight) tile [0, 1) + with each quantile inside its row's interval: the quantiles are uniform + within every cell, whatever the ties and weights. +2. With equal weights and m rows per grid point in a cell, every point of + ``DRAW_QUANTILE_GRID`` is drawn by exactly m rows (no tail is cut off). +3. The quantiles do not depend on row order, on any strictly increasing + transform of the values, on the values in other cells, or on the scale of + a cell's weights; equal seeds give equal results; a row whose value is + unique in its cell stays in its own interval whatever the seed. +4. Missing values or weights, and negative weights, are rejected. +5. ``draw_at_quantiles`` reproduces microimpute 1.8's ``predict`` exactly + at microimpute's own random draws and grid (differential test). +6. The first output drawn is non-decreasing in the quantile for a fixed + person, and every draw lies within the range of the training data. +7. ``EarningsGroupIncomeModel.predict`` draws each group at the quantiles in + the inputs, gives no draw to ``NOT_IMPUTED`` rows, and without quantiles + is deterministic. +8. ``impute_over_incomes`` passes each person's quantiles through by person + ID, so a subsample drawn with the full data's quantiles gets exactly + those quantiles, and a higher FRS income never draws lower than a lower + one in the same cell. +9. On a built enhanced FRS, SPI-synthetic employees' pay rises with their + hours as it does in the FRS (skipped when no build is present). +""" + +from __future__ import annotations + +import importlib.metadata + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, assume, 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 ( + DRAW_QUANTILE_GRID, + EMPLOYEE, + EMPLOYEE_STATUSES, + IMPUTATIONS, + NO_EARNINGS, + NOT_IMPUTED, + PREDICTORS, + SELF_EMPLOYED, + EarningsGroupIncomeModel, + draw_at_quantiles, + draw_quantile_column, + rank_cells, + rank_quantiles, + spi_age_band, +) + +RELAXED = settings(deadline=None, suppress_health_check=[HealthCheck.too_slow]) + +values_strategy = st.one_of( + st.just(0.0), + st.sampled_from([100.0, 5_000.0, 20_000.0]), # heaped values, as in the FRS + st.floats(-1e4, 1e7, allow_nan=False, allow_infinity=False), +) + + +@st.composite +def ranked_rows(draw, max_rows=60): + n = draw(st.integers(1, max_rows)) + return pd.DataFrame( + { + "value": draw(st.lists(values_strategy, min_size=n, max_size=n)), + "cell": draw(st.lists(st.integers(0, 4), min_size=n, max_size=n)), + "weight": draw( + st.lists( + st.one_of(st.just(0.0), st.floats(1e-3, 1e4)), + min_size=n, + max_size=n, + ) + ), + "id": draw( + st.lists(st.integers(0, 10**9), min_size=n, max_size=n, unique=True) + ), + } + ) + + +def _rank(rows: pd.DataFrame, seed: int = 0) -> np.ndarray: + return rank_quantiles(rows.value, rows.cell, rows.weight, rows.id, seed) + + +@RELAXED +@given(ranked_rows(), st.integers(0, 2**32 - 1)) +def test_rank_quantiles_tile_each_cell(rows, seed): + u = _rank(rows, seed) + assert ((u >= 0) & (u < 1)).all() + for _, cell in rows.assign(u=u).groupby("cell"): + weight = cell.weight.to_numpy() + if weight.sum() <= 0: + weight = np.ones(len(cell)) + # Rows without weight can share a quantile with the next value up. + order = np.lexsort((cell.value.to_numpy(), cell.u.to_numpy())) + w = weight[order] + total = w.sum() + before = np.cumsum(w) - w + q = cell.u.to_numpy()[order] + assert (q >= before / total - 1e-12).all() + assert (q <= (before + w) / total + 1e-12).all() + # Lower values get lower quantiles. + values = cell.value.to_numpy()[order] + assert (np.diff(values) >= 0).all() + + +@RELAXED +@given(ranked_rows()) +def test_rank_quantiles_increase_with_value(rows): + u = _rank(rows) + for _, cell in rows.assign(u=u).groupby("cell"): + v, q = cell.value.to_numpy(), cell.u.to_numpy() + # A cell with no weight is ranked with equal weights. + weighted = ( + cell.weight.to_numpy() > 0 + if cell.weight.sum() > 0 + else np.ones(len(cell), bool) + ) + lower = v[:, None] < v[None, :] + assert (q[:, None] <= q[None, :])[lower].all() + assert (q[:, None] < q[None, :])[lower & weighted[:, None]].all() + + +@RELAXED +@given(st.integers(1, 50), st.integers(0, 2**32 - 1), st.booleans()) +def test_equal_weights_fill_every_grid_point(m, seed, all_tied): + k = len(DRAW_QUANTILE_GRID) + n = m * k + rng = np.random.default_rng(seed) + values = np.zeros(n) if all_tied else rng.integers(0, 50, n).astype(float) + u = rank_quantiles(values, np.zeros(n), np.ones(n), np.arange(n), seed) + index = np.clip((u * k).astype(int), 0, k - 1) + assert (np.bincount(index, minlength=k) == m).all() + + +@RELAXED +@given(ranked_rows(), st.randoms(use_true_random=False), st.integers(0, 1000)) +def test_rank_quantiles_ignore_row_order(rows, random, seed): + order = list(range(len(rows))) + random.shuffle(order) + shuffled = rows.iloc[order] + a = pd.Series(_rank(rows, seed), index=rows.id) + b = pd.Series(_rank(shuffled, seed), index=shuffled.id) + pd.testing.assert_series_equal(a.sort_index(), b.sort_index()) + + +@RELAXED +@given(ranked_rows(), st.floats(1e-3, 1e3), st.floats(-1e6, 1e6)) +def test_rank_quantiles_ignore_increasing_transforms(rows, scale, shift): + transformed = rows.assign(value=rows.value * scale + shift) + # A transform that merges distinct values (float rounding) changes ties. + assume( + transformed.groupby("cell") + .value.nunique() + .equals(rows.groupby("cell").value.nunique()) + ) + np.testing.assert_array_equal(_rank(rows), _rank(transformed)) + + +@RELAXED +@given(ranked_rows(), st.integers(0, 4), st.data()) +def test_rank_quantiles_cells_are_independent(rows, cell, data): + changed = rows.copy() + in_cell = changed.cell == cell + changed.loc[in_cell, "value"] = data.draw( + st.lists( + values_strategy, min_size=int(in_cell.sum()), max_size=int(in_cell.sum()) + ) + ) + np.testing.assert_array_equal(_rank(rows)[~in_cell], _rank(changed)[~in_cell]) + + +@RELAXED +@given(ranked_rows(), st.integers(0, 4), st.floats(1e-2, 1e2)) +def test_rank_quantiles_ignore_weight_scale(rows, cell, factor): + scaled = rows.copy() + scaled.loc[scaled.cell == cell, "weight"] *= factor + np.testing.assert_allclose(_rank(rows), _rank(scaled), rtol=1e-9, atol=1e-12) + + +@RELAXED +@given(ranked_rows(), st.integers(0, 2**32 - 1), st.integers(0, 2**32 - 1)) +def test_rank_quantiles_seeds(rows, seed_a, seed_b): + np.testing.assert_array_equal(_rank(rows, seed_a), _rank(rows, seed_a)) + a, b = _rank(rows, seed_a), _rank(rows, seed_b) + weight = rows.weight.where( + rows.groupby("cell").weight.transform("sum") > 0, 1.0 + ).to_numpy() + total = pd.Series(weight).groupby(rows.cell.to_numpy()).transform("sum").to_numpy() + unique_in_cell = ~rows.duplicated(["cell", "value"], keep=False).to_numpy() + # A row with a value of its own keeps its interval of width weight/total. + width = weight / total + assert (np.abs(a - b)[unique_in_cell] <= width[unique_in_cell] + 1e-12).all() + + +@pytest.mark.parametrize( + "values, weights", + [ + ([1.0, np.nan], [1.0, 1.0]), + ([1.0, 2.0], [1.0, np.nan]), + ([1.0, 2.0], [1.0, -1.0]), + ], +) +def test_rank_quantiles_reject_bad_inputs(values, weights): + with pytest.raises(ValueError): + rank_quantiles(values, [0, 0], weights, [1, 2]) + + +def test_spi_age_bands_match_spi_codes(): + from policyengine_uk_data.datasets.spi import AGE_RANGES + + for code, (low, high) in AGE_RANGES.items(): + if code > 0: + assert (spi_age_band([low, (low + high) / 2, high - 0.01]) == code).all() + assert spi_age_band([0, 15.9])[0] == 0 and spi_age_band([95])[0] == 7 + + +def test_rank_cells_separate_every_key(): + inputs = pd.DataFrame( + { + "earnings_group": [EMPLOYEE, EMPLOYEE, SELF_EMPLOYED, EMPLOYEE, EMPLOYEE], + "age": [30, 31, 30, 50, 30], + "gender": ["MALE", "MALE", "MALE", "MALE", "FEMALE"], + "region": ["LONDON"] * 5, + } + ) + cells = rank_cells(inputs) + assert cells[0] == cells[1] + assert len(set(cells[[0, 2, 3, 4]])) == 4 + + +# ---- the draw --------------------------------------------------------------- + +OUTPUTS = ["employment_income", "dividend_income", "gift_aid"] + + +@pytest.fixture(scope="module") +def fitted_results(): + """A microimpute QRF on synthetic, SPI-like data with zeros and a tail.""" + from policyengine_uk_data.utils.qrf import QRF + + rng = np.random.default_rng(3) + n = 600 + X = pd.DataFrame( + { + "age": rng.uniform(16, 90, n), + "gender": rng.choice(["MALE", "FEMALE"], n), + "region": rng.choice(["LONDON", "WALES", "SCOTLAND"], n), + } + ) + pay = rng.lognormal(10, 0.8, n) * np.where(X.region == "LONDON", 1.5, 1.0) + y = pd.DataFrame( + { + "employment_income": pay, + "dividend_income": (rng.random(n) < 0.3) * rng.exponential(pay / 10), + "gift_aid": (rng.random(n) < 0.1) * rng.exponential(200, n), + } + ) + model = QRF() + model.fit(X, y) + return model.model, X, y + + +people = st.builds( + lambda age, gender, region: (age, gender, region), + st.floats(16, 90), + st.sampled_from(["MALE", "FEMALE"]), + st.sampled_from(["LONDON", "WALES", "SCOTLAND"]), +) + + +def _frame(rows): + return pd.DataFrame(rows, columns=PREDICTORS) + + +@pytest.mark.skipif( + importlib.metadata.version("microimpute").split(".")[0] != "1", + reason="microimpute 2+ draws on another grid; the locked version is 1.x", +) +@settings( + deadline=None, + max_examples=30, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given(st.lists(people, min_size=1, max_size=40)) +def test_draw_reproduces_microimpute_predict(fitted_results, rows): + results, _, _ = fitted_results + X = _frame(rows) + expected = results.predict(X) + u = np.random.default_rng(results.seed).beta(1, 1, size=len(X)) + grid = np.linspace(1 / 11, 1 - 1 / 11, 10) + drawn = draw_at_quantiles(results, X, {v: u for v in OUTPUTS}, grid) + for v in OUTPUTS: + np.testing.assert_array_equal(drawn[v].to_numpy(), expected[v].to_numpy()) + + +@settings( + deadline=None, + max_examples=30, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given(people, st.lists(st.floats(0, 1, exclude_max=True), min_size=2, max_size=12)) +def test_first_output_rises_with_the_quantile(fitted_results, person, us): + results, _, y = fitted_results + us = np.sort(np.array(us)) + X = _frame([person] * len(us)) + drawn = draw_at_quantiles(results, X, {v: us for v in OUTPUTS}, DRAW_QUANTILE_GRID) + assert (np.diff(drawn[OUTPUTS[0]].to_numpy()) >= 0).all() + for v in OUTPUTS: + assert (drawn[v] >= y[v].min() - 1e-9).all() + assert (drawn[v] <= y[v].max() + 1e-9).all() + + +def test_draw_reaches_the_tails(fitted_results): + """The grid's ends are the forest's 0.05th and 99.95th percentiles.""" + results, X, _ = fitted_results + row = X.iloc[[0, 0]].reset_index(drop=True) + u = np.array([0.0, np.nextafter(1.0, 0.0)]) + drawn = draw_at_quantiles(results, row, {v: u for v in OUTPUTS}, DRAW_QUANTILE_GRID) + model = results.models[OUTPUTS[0]] + features, _ = results.preprocess_data_types( + row, results.original_predictors, getattr(results, "dummy_processor", None) + ) + columns = results._get_encoded_predictors(results.predictors) + ends = np.asarray( + model.qrf.predict( + features[columns], quantiles=[DRAW_QUANTILE_GRID[0], DRAW_QUANTILE_GRID[-1]] + ) + )[0] + np.testing.assert_array_equal(drawn[OUTPUTS[0]].to_numpy(), ends) + + +# ---- the group model and the wiring ------------------------------------------ + + +@pytest.fixture(scope="module") +def group_model(fitted_results): + from policyengine_uk_data.utils.qrf import QRF + + results, X, y = fitted_results + full = y.reindex(columns=IMPUTATIONS, fill_value=0.0) + models = {} + for group in (EMPLOYEE, NO_EARNINGS): + model = QRF() + model.fit(X, full) + models[group] = model + return EarningsGroupIncomeModel(models, income_module.get_income_model_metadata()) + + +@settings( + deadline=None, + max_examples=25, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given( + st.lists( + st.tuples(people, st.sampled_from([EMPLOYEE, NO_EARNINGS, NOT_IMPUTED])), + min_size=1, + max_size=30, + ), + st.integers(0, 2**32 - 1), +) +def test_group_model_draws_at_the_given_quantiles(group_model, rows, seed): + X = _frame([r[0] for r in rows]) + X["earnings_group"] = [r[1] for r in rows] + rng = np.random.default_rng(seed) + for v in IMPUTATIONS: + X[draw_quantile_column(v)] = rng.random(len(X)) + drawn = group_model.predict(X) + groups = X.earnings_group.to_numpy() + assert drawn[groups == NOT_IMPUTED].isna().all().all() + for group, model in group_model.models.items(): + m = groups == group + if m.any(): + expected = draw_at_quantiles( + model.model, + X.loc[m, PREDICTORS].reset_index(drop=True), + {v: X[draw_quantile_column(v)].to_numpy()[m] for v in IMPUTATIONS}, + DRAW_QUANTILE_GRID, + ) + np.testing.assert_array_equal( + drawn.loc[m, IMPUTATIONS].to_numpy(), expected[IMPUTATIONS].to_numpy() + ) + # Without quantiles the draw is random but reproducible. + bare = X[PREDICTORS + ["earnings_group"]] + pd.testing.assert_frame_equal(group_model.predict(bare), group_model.predict(bare)) + + +class _EchoModel: + """Draws each output as the person's quantile for it.""" + + def predict(self, X): + drawn = pd.DataFrame( + {v: X[draw_quantile_column(v)].to_numpy() for v in IMPUTATIONS}, + index=X.index, + ) + drawn[X.earnings_group.to_numpy() == NOT_IMPUTED] = np.nan + return drawn + + +class _Dataset: + def __init__(self, person, household): + self.person, self.household = person, household + + def copy(self): + return _Dataset(self.person.copy(), self.household.copy()) + + +@st.composite +def frs_datasets(draw): + n = draw(st.integers(2, 40)) + households = draw(st.integers(1, n)) + person = pd.DataFrame( + { + "person_id": np.arange(n) * 7 + 3, + "person_household_id": draw( + st.lists(st.integers(0, households - 1), min_size=n, max_size=n) + ), + "age": draw(st.lists(st.integers(10, 80), min_size=n, max_size=n)), + "gender": draw( + st.lists(st.sampled_from(["MALE", "FEMALE"]), min_size=n, max_size=n) + ), + "employment_status": draw( + st.lists( + st.sampled_from(["FT_EMPLOYED", "PT_EMPLOYED", "RETIRED"]), + min_size=n, + max_size=n, + ) + ), + } + ) + for v in IMPUTATIONS[:6]: + person[v] = draw( + st.lists(st.one_of(st.just(0.0), st.floats(1, 1e6)), min_size=n, max_size=n) + ) + household = pd.DataFrame( + { + "household_id": np.arange(households), + "household_weight": draw( + st.lists(st.floats(1, 3_000), min_size=households, max_size=households) + ), + "region": draw( + st.lists( + st.sampled_from(["LONDON", "WALES"]), + min_size=households, + max_size=households, + ) + ), + } + ) + return _Dataset(person, household) + + +def _inputs_from_columns(dataset): + person = dataset.person + region = person.person_household_id.map( + dataset.household.set_index("household_id").region + ) + inputs = pd.DataFrame( + { + "age": person.age.to_numpy(), + "gender": person.gender.to_numpy(), + "region": region.to_numpy(), + } + ) + inputs["earnings_group"] = income_module.frs_earnings_group( + person.employment_status, + inputs.age, + person.employment_income, + person.self_employment_income, + ) + return inputs + + +@settings(deadline=None, max_examples=40, suppress_health_check=[HealthCheck.too_slow]) +@given(frs_datasets(), st.data()) +def test_impute_over_incomes_passes_quantiles_by_person(dataset, data): + with pytest.MonkeyPatch.context() as mp: + mp.setattr(income_module, "income_model_inputs", _inputs_from_columns) + quantiles = income_module.draw_quantiles(dataset) + keep = data.draw( + st.lists( + st.sampled_from(list(dataset.household.household_id)), + min_size=1, + unique=True, + ) + ) + sub = _Dataset( + dataset.person[dataset.person.person_household_id.isin(keep)].reset_index( + drop=True + ), + dataset.household[dataset.household.household_id.isin(keep)].reset_index( + drop=True + ), + ) + result = income_module.impute_over_incomes( + sub, _EchoModel(), IMPUTATIONS, quantiles + ) + groups = _inputs_from_columns(sub).earnings_group.to_numpy() + drawn = groups != NOT_IMPUTED + expected = quantiles.loc[sub.person.person_id.to_numpy()] + for v in IMPUTATIONS: + got = result.person[v].to_numpy() + np.testing.assert_array_equal( + got[drawn], expected[draw_quantile_column(v)].to_numpy()[drawn] + ) + if v in sub.person.columns: + np.testing.assert_array_equal(got[~drawn], sub.person[v].to_numpy()[~drawn]) + # Within a cell, more FRS pay never draws lower. + inputs = _inputs_from_columns(dataset) + cells = rank_cells(inputs) + pay = dataset.person.employment_income.to_numpy() + u = quantiles[draw_quantile_column("employment_income")].to_numpy() + same = cells[:, None] == cells[None, :] + assert (u[:, None] < u[None, :])[same & (pay[:, None] < pay[None, :])].all() + + +def test_built_enhanced_frs_spi_pay_rises_with_hours(enhanced_frs): + from scipy.stats import spearmanr + + person = enhanced_frs.person + household = enhanced_frs.household.set_index("household_id") + spi = person.person_household_id.map(household.household_is_spi_synthetic).to_numpy( + dtype=bool + ) + if not spi.any(): + pytest.skip("No SPI-synthetic rows in this build") + status = person.employment_status.astype(str).to_numpy() + working = ( + np.isin(status, EMPLOYEE_STATUSES) + & (person.employment_income > 0).to_numpy() + & (person.hours_worked > 0).to_numpy() + ) + frs_rho = spearmanr( + person.hours_worked[working & ~spi], person.employment_income[working & ~spi] + )[0] + spi_rho = spearmanr( + person.hours_worked[working & spi], person.employment_income[working & spi] + )[0] + # Random draws gave SPI rows almost no link between hours and pay. + assert spi_rho > 0.5 * frs_rho From df3a633063cee0b89e888273b8a9f9ae5be86f7e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 22:57:53 -0400 Subject: [PATCH 2/5] Split SPI earners with both pay and a trade by main source People with both pay and a trade drew from one forest, so an FRS self-employed person with a small PAYE job drew like a salaried person with a side trade. They are now split into employee-main and self-employed-main groups. In the SPI the split follows MAINSRCE: 1 means pay is the main source, 3 (sole trader) or 4 (partnership) a trade, otherwise the larger income. MAINSRCE is the self-assessment main source HMRC stratifies the SPI by, not the larger income. In the FRS the split follows the main job's status, otherwise the larger income. The FRS has only about 300 such people, one or two per rank cell, so the main source is most of what links their draw to their own jobs. Co-Authored-By: Claude Opus 5.5 --- .../spi-rank-preserving-draw.changed.md | 1 + .../datasets/imputations/income.py | 90 ++++++++++++++----- policyengine_uk_data/tests/test_spi_build.py | 1 + .../tests/test_spi_income_earnings_groups.py | 83 +++++++++++++---- 4 files changed, 135 insertions(+), 40 deletions(-) create mode 100644 changelog.d/spi-rank-preserving-draw.changed.md diff --git a/changelog.d/spi-rank-preserving-draw.changed.md b/changelog.d/spi-rank-preserving-draw.changed.md new file mode 100644 index 00000000..225b9f4d --- /dev/null +++ b/changelog.d/spi-rank-preserving-draw.changed.md @@ -0,0 +1 @@ +Draw each SPI-synthetic person's incomes at their FRS rank of that income within their cell (earnings group, SPI age band, gender, region) instead of at a random quantile, so part-timers and low earners draw low SPI pay and the top of each cell draws the SPI's top tail; split people with both pay and a trade by main source (SPI MAINSRCE, FRS main job); and draw on a 1,000-point quantile grid instead of microimpute 1.8's ten points between the 9th and 91st percentiles. The FRS rows' imputed dividends use the same ranks. diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index 363a1ad4..a8cd2a5e 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -75,16 +75,41 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: # 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. +# +# People with both pay and a trade are split by which is their main source. +# The SPI's MAINSRCE is the self-assessment "main source of income" HMRC +# stratifies its sample by (1 pay, 3 sole trader, 4 partnership), not the +# larger income: in 2022-23, 57% (weighted) of those whose main source is a +# trade had more pay than profit. The FRS equivalent is the main job's +# status. Where neither says (MAINSRCE -1 not classified, 2 occupational +# pension, 5 other; an FRS main job that is neither), the larger of pay and +# profit decides. FRS people with both are too few (about 300) for their rank +# within a cell to say much, so the main source is most of what links their +# draw to their own jobs. NO_EARNINGS = "NO_EARNINGS" EMPLOYEE = "EMPLOYEE" SELF_EMPLOYED = "SELF_EMPLOYED" -EMPLOYEE_AND_SELF_EMPLOYED = "EMPLOYEE_AND_SELF_EMPLOYED" +EMPLOYEE_MAIN_AND_SELF_EMPLOYED = "EMPLOYEE_MAIN_AND_SELF_EMPLOYED" +SELF_EMPLOYED_MAIN_AND_EMPLOYEE = "SELF_EMPLOYED_MAIN_AND_EMPLOYEE" EARNINGS_GROUPS = ( NO_EARNINGS, EMPLOYEE, SELF_EMPLOYED, - EMPLOYEE_AND_SELF_EMPLOYED, + EMPLOYEE_MAIN_AND_SELF_EMPLOYED, + SELF_EMPLOYED_MAIN_AND_EMPLOYEE, ) +PAY_GROUPS = ( + EMPLOYEE, + EMPLOYEE_MAIN_AND_SELF_EMPLOYED, + SELF_EMPLOYED_MAIN_AND_EMPLOYEE, +) +TRADE_GROUPS = ( + SELF_EMPLOYED, + EMPLOYEE_MAIN_AND_SELF_EMPLOYED, + SELF_EMPLOYED_MAIN_AND_EMPLOYEE, +) +SPI_PAY_MAIN_SOURCE = 1 +SPI_TRADE_MAIN_SOURCES = (3, 4) # 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 this FRS build does not record, and anyone under 16. With no @@ -99,29 +124,43 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: 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.""" +def earnings_group(has_pay, has_trade, pay_is_main) -> np.ndarray: + """Earnings group from whether a person has pay and a trade, and, for + people with both, whether pay is their main source.""" has_pay = np.asarray(has_pay, dtype=bool) has_trade = np.asarray(has_trade, dtype=bool) + both = has_pay & has_trade + pay_is_main = np.asarray(pay_is_main, dtype=bool) return np.select( - [has_pay & has_trade, has_trade, has_pay], - [EMPLOYEE_AND_SELF_EMPLOYED, SELF_EMPLOYED, EMPLOYEE], + [both & pay_is_main, both, has_trade, has_pay], + [ + EMPLOYEE_MAIN_AND_SELF_EMPLOYED, + SELF_EMPLOYED_MAIN_AND_EMPLOYEE, + SELF_EMPLOYED, + EMPLOYEE, + ], NO_EARNINGS, ).astype(object) def spi_earnings_group( - employment_income, self_employment_income, self_employed_indicator + employment_income, self_employment_income, self_employed_indicator, main_source ) -> 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. + pages filed, whatever the profit) or any assessable profit. Pay is the + main source if MAINSRCE says pay, or if it names neither pay nor a trade + and pay is at least the profit. """ - has_trade = (np.asarray(self_employed_indicator) == 1) | ( - np.asarray(self_employment_income, dtype=float) > 0 + pay = np.asarray(employment_income, dtype=float) + profit = np.asarray(self_employment_income, dtype=float) + main_source = np.asarray(main_source) + has_trade = (np.asarray(self_employed_indicator) == 1) | (profit > 0) + pay_is_main = (main_source == SPI_PAY_MAIN_SOURCE) | ( + ~np.isin(main_source, SPI_TRADE_MAIN_SOURCES) & (pay >= profit) ) - return earnings_group(np.asarray(employment_income, dtype=float) > 0, has_trade) + return earnings_group(pay > 0, has_trade, pay_is_main) def frs_earnings_group( @@ -132,20 +171,22 @@ def frs_earnings_group( 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. + source, and the main job's source is the main one. Everyone else, retired, + unemployed or inactive, draws from SPI records with neither, unless the + FRS records earnings for them. """ 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 - ) + pay = np.asarray(employment_income, dtype=float) + profit = np.asarray(self_employment_income, dtype=float) + employee = np.isin(status, EMPLOYEE_STATUSES) + self_employed = np.isin(status, SELF_EMPLOYED_STATUSES) + has_pay = employee | (pay > 0) + has_trade = self_employed | (profit > 0) + pay_is_main = employee | (~self_employed & (pay >= profit)) 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 - ) + return np.where( + is_child, NOT_IMPUTED, earnings_group(has_pay, has_trade, pay_is_main) + ).astype(object) def earnings_group_sample_sizes( @@ -197,7 +238,10 @@ def generate_spi_table( 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 + spi.employment_income, + spi.self_employment_income, + spi.SEINC_NUM, + spi.MAINSRCE, ) if sample_size is not None: diff --git a/policyengine_uk_data/tests/test_spi_build.py b/policyengine_uk_data/tests/test_spi_build.py index eff98ef7..04f29b16 100644 --- a/policyengine_uk_data/tests/test_spi_build.py +++ b/policyengine_uk_data/tests/test_spi_build.py @@ -67,6 +67,7 @@ "BPADUE", "MAIND", "SEINC_NUM", + "MAINSRCE", ] 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 44cf55de..8fb841b5 100644 --- a/policyengine_uk_data/tests/test_spi_income_earnings_groups.py +++ b/policyengine_uk_data/tests/test_spi_income_earnings_groups.py @@ -6,8 +6,12 @@ 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. + People with both are in the employee-main group if and only if the main + job is as an employee, or is neither and pay is at least the 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. + a trade if and only if SEINC_NUM = 1 or PROFITS > 0. People with both are + in the employee-main group if and only if MAINSRCE is 1 (pay), or is not 3 + or 4 (a trade) and pay is at least the profit. 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 @@ -39,7 +43,7 @@ CHILD_STATUS, EARNINGS_GROUPS, EMPLOYEE, - EMPLOYEE_AND_SELF_EMPLOYED, + EMPLOYEE_MAIN_AND_SELF_EMPLOYED, EMPLOYEE_STATUSES, IMPUTATIONS, MIN_GROUP_SAMPLE_SHARE, @@ -47,6 +51,7 @@ NOT_IMPUTED, PREDICTORS, SELF_EMPLOYED, + SELF_EMPLOYED_MAIN_AND_EMPLOYEE, SELF_EMPLOYED_STATUSES, EarningsGroupIncomeModel, apply_income_draws, @@ -56,8 +61,10 @@ spi_earnings_group, ) -PAY_GROUPS = {EMPLOYEE, EMPLOYEE_AND_SELF_EMPLOYED} -TRADE_GROUPS = {SELF_EMPLOYED, EMPLOYEE_AND_SELF_EMPLOYED} +PAY_GROUPS = set(income_module.PAY_GROUPS) +TRADE_GROUPS = set(income_module.TRADE_GROUPS) +BOTH_GROUPS = {EMPLOYEE_MAIN_AND_SELF_EMPLOYED, SELF_EMPLOYED_MAIN_AND_EMPLOYEE} +MAIN_SOURCES = (-1, 1, 2, 3, 4, 5, 6) NON_WORKING_STATUSES = ( "UNEMPLOYED", "RETIRED", @@ -179,27 +186,55 @@ def test_frs_group_elementwise_equals_rowwise(people): @RELAXED @given( st.lists( - st.tuples(amounts, amounts, st.sampled_from([-1, 0, 1])), + st.tuples( + amounts, amounts, st.sampled_from([-1, 0, 1]), st.sampled_from(MAIN_SOURCES) + ), 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) + pay, profit, indicator, source = (np.array(column) for column in zip(*records)) + groups = spi_earnings_group(pay, profit, indicator, source) 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] + # People with both: MAINSRCE decides, then the larger income. + both = np.isin(groups, list(BOTH_GROUPS)) + pay_main = (source == 1) | (~np.isin(source, [3, 4]) & (pay >= profit)) + np.testing.assert_array_equal( + groups[both] == EMPLOYEE_MAIN_AND_SELF_EMPLOYED, pay_main[both] + ) + rowwise = [spi_earnings_group([p], [q], [i], [s])[0] for p, q, i, s in records] assert list(groups) == rowwise +@RELAXED +@given(frs_people()) +def test_frs_both_group_follows_main_job(people): + groups = _frs_groups(people) + status = people.employment_status.to_numpy() + both = np.isin(groups, list(BOTH_GROUPS)) + pay, profit = ( + people.employment_income.to_numpy(), + people.self_employment_income.to_numpy(), + ) + pay_main = np.isin(status, EMPLOYEE_STATUSES) | ( + ~np.isin(status, SELF_EMPLOYED_STATUSES) & (pay >= profit) + ) + np.testing.assert_array_equal( + groups[both] == EMPLOYEE_MAIN_AND_SELF_EMPLOYED, pay_main[both] + ) + + @RELAXED @given( st.lists( - st.tuples(amounts, amounts, st.sampled_from([-1, 0, 1])), + st.tuples( + amounts, amounts, st.sampled_from([-1, 0, 1]), st.sampled_from(MAIN_SOURCES) + ), min_size=1, max_size=40, ), @@ -207,13 +242,15 @@ def test_spi_group_follows_pay_and_self_employment_pages(records): 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) + pay, profit, indicator, main = ( + np.array(column, dtype=float) for column in zip(*records) + ) + before = spi_earnings_group(pay, profit, indicator, main) if source == "pay": - after = spi_earnings_group(pay + extra, profit, indicator) + after = spi_earnings_group(pay + extra, profit, indicator, main) gained, other = PAY_GROUPS, TRADE_GROUPS else: - after = spi_earnings_group(pay, profit + extra, indicator) + after = spi_earnings_group(pay, profit + extra, indicator, main) gained, other = TRADE_GROUPS, PAY_GROUPS assert np.isin(after, list(gained)).all() np.testing.assert_array_equal( @@ -266,6 +303,15 @@ def _raw_spi(rng: np.random.Generator, n: int) -> pd.DataFrame: has_trade & (rng.random(n) < 0.7), rng.lognormal(9, 1, n), 0.0 ) raw["SEINC_NUM"] = has_trade.astype(int) + # MAINSRCE agrees with the group for people with both; anything for others. + raw["MAINSRCE"] = np.select( + [ + group == EMPLOYEE_MAIN_AND_SELF_EMPLOYED, + group == SELF_EMPLOYED_MAIN_AND_EMPLOYEE, + ], + [1, rng.choice([3, 4], n)], + rng.choice(MAIN_SOURCES, n), + ) 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) @@ -283,11 +329,14 @@ def test_generate_spi_table_resamples_within_groups(seed, 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 + table.employment_income, + table.self_employment_income, + table.SEINC_NUM, + table.MAINSRCE, ), ) raw_groups = spi_earnings_group( - raw.PAY + raw.EPB + raw.TAXTERM, raw.PROFITS, raw.SEINC_NUM + raw.PAY + raw.EPB + raw.TAXTERM, raw.PROFITS, raw.SEINC_NUM, raw.MAINSRCE ) weights = raw.FACT.groupby(raw_groups).sum().to_dict() assert table.earnings_group.value_counts().to_dict() == ( @@ -363,9 +412,9 @@ def test_model_draws_agree_with_group(fitted_model, inputs): assert (draws[drawn].to_numpy() >= 0).all() -@pytest.mark.parametrize("group", [SELF_EMPLOYED, EMPLOYEE_AND_SELF_EMPLOYED]) +@pytest.mark.parametrize("group", income_module.TRADE_GROUPS) def test_model_draws_some_profit_for_traders(fitted_model, group): - """Both trade groups draw zero and positive profits, as the SPI has both.""" + """Every trade group draws zero and positive profits, as the SPI has both.""" n = 400 inputs = pd.DataFrame( { From 6c2664465ca5f531bc8ff4acb6bbfad2c15a0ecb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:53:29 -0400 Subject: [PATCH 3/5] Stub the rank draw in the SPI-copy gainful self-employment test #525's test runs impute_income on a fake dataset with the imputation stubbed. impute_income now also computes the draw quantiles from the FRS, which builds a Microsimulation the fake dataset cannot support, so the test stubs draw_quantiles too and its impute_over_incomes stub takes the quantiles argument, as test_imputation_source_flags does. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/tests/test_uc_gainful_self_employment.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/policyengine_uk_data/tests/test_uc_gainful_self_employment.py b/policyengine_uk_data/tests/test_uc_gainful_self_employment.py index 48766685..5e7f6a5f 100644 --- a/policyengine_uk_data/tests/test_uc_gainful_self_employment.py +++ b/policyengine_uk_data/tests/test_uc_gainful_self_employment.py @@ -152,7 +152,7 @@ def test_spi_copy_flag_follows_its_own_imputed_incomes(rows): imputed_profit = [r[3] for r in rows] imputed_pay = [r[4] for r in rows] - def impute_over_incomes(dataset, _model, output_variables): + def impute_over_incomes(dataset, _model, output_variables, _quantiles=None): dataset = dataset.copy() if "self_employment_income" in output_variables: dataset.person["self_employment_income"] = imputed_profit @@ -161,6 +161,7 @@ def impute_over_incomes(dataset, _model, output_variables): with pytest.MonkeyPatch.context() as m: m.setattr(income_module, "create_income_model", lambda: object()) + m.setattr(income_module, "draw_quantiles", lambda dataset: None) m.setattr(income_module, "subsample_dataset", lambda d, _n: d.copy()) m.setattr(income_module, "impute_over_incomes", impute_over_incomes) m.setattr( From c134a6bcd609c56214f1b05279874a63778185ca Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 21:50:02 -0400 Subject: [PATCH 4/5] Test the rank draw's design invariants and qualify its claims Review r1 found 7 of 8 design-breaking mutants survived the tests. New tests: - a row's quantile is uniform on its slice across seeds (KS), so a lone donor does not always draw the median; - draw_quantiles tiles each rank cell at household weights; - outputs everyone ties on get independent quantiles; - each output is drawn at its own quantile; - impute_income passes one set of full-FRS quantiles to both draws. All eight mutants now fail. draw_at_quantiles raises on a model type it does not know, and rank_quantiles raises on duplicate ids. The module docstring no longer says the SPI distribution is kept: each group's first output keeps the forest's distribution for the cell in expectation, and later outputs carry extra dependence. The MAINSRCE comment cites the SPI 2022-23 public use tape documentation and lists every code. Co-Authored-By: Claude Opus 5.5 --- .../spi-rank-preserving-draw.changed.md | 2 +- .../datasets/imputations/income.py | 51 +++++--- .../tests/test_imputation_source_flags.py | 49 +++++++ .../tests/test_spi_income_rank_draw.py | 123 ++++++++++++++++-- 4 files changed, 198 insertions(+), 27 deletions(-) diff --git a/changelog.d/spi-rank-preserving-draw.changed.md b/changelog.d/spi-rank-preserving-draw.changed.md index 225b9f4d..bc9a64f8 100644 --- a/changelog.d/spi-rank-preserving-draw.changed.md +++ b/changelog.d/spi-rank-preserving-draw.changed.md @@ -1 +1 @@ -Draw each SPI-synthetic person's incomes at their FRS rank of that income within their cell (earnings group, SPI age band, gender, region) instead of at a random quantile, so part-timers and low earners draw low SPI pay and the top of each cell draws the SPI's top tail; split people with both pay and a trade by main source (SPI MAINSRCE, FRS main job); and draw on a 1,000-point quantile grid instead of microimpute 1.8's ten points between the 9th and 91st percentiles. The FRS rows' imputed dividends use the same ranks. +Draw each SPI-synthetic person's incomes at their FRS rank of that income within their cell (earnings group, SPI age band, gender, region) instead of at a random quantile, so part-timers and low earners draw low SPI pay and the top of each cell draws from the top of the SPI forest's distribution for that cell; split people with both pay and a trade by main source (SPI MAINSRCE, FRS main job); and draw on a 1,000-point quantile grid instead of microimpute 1.8's ten points between the 9th and 91st percentiles. The FRS rows' imputed dividends use the same ranks. diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index a8cd2a5e..b2cf5b36 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -12,9 +12,13 @@ Within that cell a person does not draw at a random quantile: each income is drawn at the person's rank of the same income among FRS people in the same cell (``draw_quantiles``). A part-time employee with low FRS pay draws low -SPI pay, a donor at the top of their cell draws from the top of the SPI, and -the SPI's distribution within each cell is kept, because the ranks are -uniform within every cell. +SPI pay, and a donor at the top of their cell draws from the top of the SPI. +The ranks are uniform within every cell, so each group's first output keeps +the forest's distribution for the cell, in expectation over the forest's +within-band age noise. Later outputs are conditioned on those drawn before +at correlated FRS ranks, so they carry more dependence between incomes than +the SPI does, though less than microimpute's single random quantile per +person did. """ import pandas as pd @@ -77,15 +81,18 @@ def _spi_age_bounds(age_code) -> tuple[int, int]: # (EMPSTATI) and any second-job earnings. # # People with both pay and a trade are split by which is their main source. -# The SPI's MAINSRCE is the self-assessment "main source of income" HMRC -# stratifies its sample by (1 pay, 3 sole trader, 4 partnership), not the -# larger income: in 2022-23, 57% (weighted) of those whose main source is a -# trade had more pay than profit. The FRS equivalent is the main job's -# status. Where neither says (MAINSRCE -1 not classified, 2 occupational -# pension, 5 other; an FRS main job that is neither), the larger of pay and -# profit decides. FRS people with both are too few (about 300) for their rank -# within a cell to say much, so the main source is most of what links their -# draw to their own jobs. +# The SPI's MAINSRCE is the main source indicator (1 pay, 2 occupational +# pension, 3 sole trader, 4 partnership, 5 other, 6 claims case), and the +# main source of income is one of the variables HMRC stratifies the +# self-assessment part of the sample by (HMRC, "Survey of Personal Incomes +# 2022-23: public use tape documentation", sample design, p. 3, and variable +# list, p. 15). It is not the larger income: in 2022-23, 57% (weighted) of +# those whose main source is a trade had more pay than profit. The FRS +# equivalent is the main job's status. Where neither names pay or a trade +# (any other MAINSRCE code, including -1 not classified; an FRS main job that +# is neither), the larger of pay and profit decides. FRS people with both are +# too few (about 300) for their rank within a cell to say much, so the main +# source is most of what links their draw to their own jobs. NO_EARNINGS = "NO_EARNINGS" EMPLOYEE = "EMPLOYEE" SELF_EMPLOYED = "SELF_EMPLOYED" @@ -311,7 +318,7 @@ def rank_quantiles(values, cells, weights, ids, seed: int = 0) -> np.ndarray: gets a lower quantile than a higher one in the same cell. A cell of one row gets a uniform random quantile. A cell with no weight is ranked with equal weights. Random numbers are drawn in (cell, value, id) order, so - the result does not depend on the order of the rows. + the result does not depend on the order of the rows; ids must be unique. """ values = np.asarray(values, dtype=float) weights = np.asarray(weights, dtype=float) @@ -321,6 +328,8 @@ def rank_quantiles(values, cells, weights, ids, seed: int = 0) -> np.ndarray: raise ValueError("Ranked values and weights must not be missing") if (weights < 0).any(): raise ValueError("Rank weights must not be negative") + if pd.Series(ids).duplicated().any(): + raise ValueError("Ranked ids must be unique") n = len(values) rng = np.random.default_rng(seed) canonical = np.lexsort((ids, values, cells)) @@ -370,7 +379,8 @@ def rank_quantiles(values, cells, weights, ids, seed: int = 0) -> np.ndarray: # Each output is drawn at the forest's conditional quantile nearest the # person's rank, on this grid of 1,000 midpoints (0.0005 to 0.9995). The # draw used to be microimpute 1.8's: a random pick from ten quantiles between -# 1/11 and 10/11, so no one drew from the top or bottom 9% of their cell. +# 1/11 and 10/11, so no one drew from the top or bottom 9% of the forest's +# conditional distribution at their predictors. DRAW_QUANTILE_GRID = (np.arange(1_000) + 0.5) / 1_000 # Seed for the rank tie-breaks (and for the random quantiles of a draw made # without ranks). @@ -392,7 +402,11 @@ def draw_at_quantiles(results, X: pd.DataFrame, quantiles: dict, grid) -> pd.Dat ``predict`` exactly (tested). """ from microimpute.models.imputer import _ConstantValueModel - from microimpute.models.qrf import _get_sequential_predictors, _QRFModel + from microimpute.models.qrf import ( + _get_sequential_predictors, + _QRFModel, + _RandomForestClassifierModel, + ) grid = np.asarray(grid, dtype=float) k = len(grid) @@ -417,8 +431,13 @@ def draw_at_quantiles(results, X: pd.DataFrame, quantiles: dict, grid) -> pd.Dat (np.asarray(quantiles[variable], dtype=float) * k).astype(int), 0, k - 1 ) values = pred[np.arange(len(pred)), index] - else: + elif isinstance(model, _RandomForestClassifierModel): values = np.asarray(model.predict(augmented[columns], return_probs=False)) + else: + raise TypeError( + f"Cannot draw {variable} at given quantiles from a " + f"{type(model).__name__}" + ) output[variable] = values augmented[variable] = values augmented = results._encode_imputed_variable(augmented, variable) diff --git a/policyengine_uk_data/tests/test_imputation_source_flags.py b/policyengine_uk_data/tests/test_imputation_source_flags.py index 5ef3f0d5..d1929c38 100644 --- a/policyengine_uk_data/tests/test_imputation_source_flags.py +++ b/policyengine_uk_data/tests/test_imputation_source_flags.py @@ -108,6 +108,55 @@ def test_impute_income_marks_spi_synthetic_households(monkeypatch): assert result.household.loc[2:, "household_weight"].eq(0).all() +def test_impute_income_draws_both_halves_at_full_frs_ranks(monkeypatch): + """Ranks are computed once, on the full, weighted FRS, and both draws use + them: the SPI copy is an unweighted subsample, so ranks taken within it + would be wrong.""" + from policyengine_uk_data.datasets.imputations import income as income_module + from policyengine_uk_data.datasets import disability_benefits + from policyengine_uk_data.datasets.imputations import frs_only + + full_frs_quantiles = pd.DataFrame({"q": [0.25, 0.75]}, index=[1, 2]) + ranked_weights, calls = [], [] + + def draw_quantiles(dataset): + ranked_weights.append(dataset.household.household_weight.tolist()) + return full_frs_quantiles + + def impute_over_incomes(dataset, _model, output_variables, quantiles=None): + calls.append((list(output_variables), quantiles)) + return dataset + + monkeypatch.setattr(income_module, "create_income_model", lambda: object()) + monkeypatch.setattr(income_module, "draw_quantiles", draw_quantiles) + monkeypatch.setattr( + income_module, + "subsample_dataset", + lambda dataset, _sample_size: dataset.copy(), + ) + monkeypatch.setattr(income_module, "impute_over_incomes", impute_over_incomes) + monkeypatch.setattr( + frs_only, + "impute_frs_only_variables", + lambda train_dataset, target_dataset: target_dataset, + ) + monkeypatch.setattr( + disability_benefits, + "strip_internal_disability_reported_amounts", + lambda dataset: dataset, + ) + monkeypatch.setattr(income_module, "stack_datasets", _stack_without_remapping) + + income_module.impute_income(_fake_dataset()) + + assert ranked_weights == [[1.0, 2.0]] + assert [variables for variables, _ in calls] == [ + income_module.IMPUTATIONS, + ["dividend_income"], + ] + assert all(quantiles is full_frs_quantiles for _, quantiles in calls) + + def test_impute_capital_gains_marks_capital_gains_clone_households(monkeypatch): cg_module = importlib.import_module( "policyengine_uk_data.datasets.imputations.capital_gains" diff --git a/policyengine_uk_data/tests/test_spi_income_rank_draw.py b/policyengine_uk_data/tests/test_spi_income_rank_draw.py index a245cd7d..444e3f06 100644 --- a/policyengine_uk_data/tests/test_spi_income_rank_draw.py +++ b/policyengine_uk_data/tests/test_spi_income_rank_draw.py @@ -13,19 +13,29 @@ 3. The quantiles do not depend on row order, on any strictly increasing transform of the values, on the values in other cells, or on the scale of a cell's weights; equal seeds give equal results; a row whose value is - unique in its cell stays in its own interval whatever the seed. -4. Missing values or weights, and negative weights, are rejected. + unique in its cell stays in its own interval whatever the seed. Across + seeds, a row's quantile is uniform on its interval (not a fixed point in + it). +4. Missing values or weights, negative weights and duplicate ids are + rejected. 5. ``draw_at_quantiles`` reproduces microimpute 1.8's ``predict`` exactly at microimpute's own random draws and grid (differential test). 6. The first output drawn is non-decreasing in the quantile for a fixed - person, and every draw lies within the range of the training data. + person, and every draw lies within the range of the training data. Each + output is drawn at its own quantile: raising one output's quantile never + lowers it and never changes an output drawn before it. A model type the + draw does not know is an error. 7. ``EarningsGroupIncomeModel.predict`` draws each group at the quantiles in the inputs, gives no draw to ``NOT_IMPUTED`` rows, and without quantiles is deterministic. 8. ``impute_over_incomes`` passes each person's quantiles through by person ID, so a subsample drawn with the full data's quantiles gets exactly those quantiles, and a higher FRS income never draws lower than a lower - one in the same cell. + one in the same cell. ``draw_quantiles`` ranks within ``rank_cells`` at + household weights: in every cell each person's quantile lies in their + slice of the cell's household weight. Outputs everyone ties on get + independent quantiles. (``impute_income`` uses one set of full-FRS + quantiles for both of its draws: test_imputation_source_flags.py.) 9. On a built enhanced FRS, SPI-synthetic employees' pay rises with their hours as it does in the FRS (skipped when no build is present). """ @@ -203,17 +213,33 @@ def test_rank_quantiles_seeds(rows, seed_a, seed_b): assert (np.abs(a - b)[unique_in_cell] <= width[unique_in_cell] + 1e-12).all() +def test_rank_quantiles_are_uniform_within_each_slice(): + """Across seeds a row takes a uniform point in its slice, not a fixed + one: a lone donor does not always draw its cell's median.""" + from scipy.stats import kstest + + seeds = range(2_000) + lone = [rank_quantiles([5.0], [0], [1.0], [7], seed)[0] for seed in seeds] + assert kstest(lone, "uniform").pvalue > 1e-3 + pair = np.array( + [rank_quantiles([1.0, 2.0], [0, 0], [1.0, 3.0], [1, 2], seed) for seed in seeds] + ) + assert kstest(pair[:, 0], "uniform", args=(0, 0.25)).pvalue > 1e-3 + assert kstest(pair[:, 1], "uniform", args=(0.25, 0.75)).pvalue > 1e-3 + + @pytest.mark.parametrize( - "values, weights", + "values, weights, ids", [ - ([1.0, np.nan], [1.0, 1.0]), - ([1.0, 2.0], [1.0, np.nan]), - ([1.0, 2.0], [1.0, -1.0]), + ([1.0, np.nan], [1.0, 1.0], [1, 2]), + ([1.0, 2.0], [1.0, np.nan], [1, 2]), + ([1.0, 2.0], [1.0, -1.0], [1, 2]), + ([1.0, 2.0], [1.0, 1.0], [1, 1]), ], ) -def test_rank_quantiles_reject_bad_inputs(values, weights): +def test_rank_quantiles_reject_bad_inputs(values, weights, ids): with pytest.raises(ValueError): - rank_quantiles(values, [0, 0], weights, [1, 2]) + rank_quantiles(values, [0, 0], weights, ids) def test_spi_age_bands_match_spi_codes(): @@ -340,6 +366,33 @@ def test_draw_reaches_the_tails(fitted_results): np.testing.assert_array_equal(drawn[OUTPUTS[0]].to_numpy(), ends) +@pytest.mark.parametrize("j", range(len(OUTPUTS))) +def test_each_output_is_drawn_at_its_own_quantile(fitted_results, j): + """Raising one output's quantile raises that output and leaves every + output drawn before it unchanged.""" + results, X, _ = fitted_results + X = X.iloc[:50].reset_index(drop=True) + low = {v: np.zeros(len(X)) for v in OUTPUTS} + high = {**low, OUTPUTS[j]: np.full(len(X), np.nextafter(1.0, 0.0))} + a = draw_at_quantiles(results, X, low, DRAW_QUANTILE_GRID) + b = draw_at_quantiles(results, X, high, DRAW_QUANTILE_GRID) + for v in OUTPUTS[:j]: + np.testing.assert_array_equal(a[v].to_numpy(), b[v].to_numpy()) + assert (b[OUTPUTS[j]] >= a[OUTPUTS[j]]).all() + assert (b[OUTPUTS[j]] > a[OUTPUTS[j]]).mean() > 0.5 + + +def test_draw_rejects_unknown_model_types(fitted_results): + import copy + + results, X, _ = fitted_results + results = copy.copy(results) + results.models = {**results.models, OUTPUTS[1]: object()} + u = {v: np.zeros(3) for v in OUTPUTS} + with pytest.raises(TypeError): + draw_at_quantiles(results, X.iloc[:3], u, DRAW_QUANTILE_GRID) + + # ---- the group model and the wiring ------------------------------------------ @@ -525,6 +578,56 @@ def test_impute_over_incomes_passes_quantiles_by_person(dataset, data): assert (u[:, None] < u[None, :])[same & (pay[:, None] < pay[None, :])].all() +@settings(deadline=None, max_examples=60, suppress_health_check=[HealthCheck.too_slow]) +@given(frs_datasets()) +def test_draw_quantiles_tile_each_cell_at_household_weights(dataset): + """Ranks are taken within ``rank_cells`` at household weights, so in + every cell each person's quantile lies in their slice of the cell's + household weight, for every output.""" + with pytest.MonkeyPatch.context() as mp: + mp.setattr(income_module, "income_model_inputs", _inputs_from_columns) + quantiles = income_module.draw_quantiles(dataset) + cells = rank_cells(_inputs_from_columns(dataset)) + weight = dataset.person.person_household_id.map( + dataset.household.set_index("household_id").household_weight + ).to_numpy() + for v in IMPUTATIONS: + q = quantiles[draw_quantile_column(v)].to_numpy() + for cell in np.unique(cells): + order = np.argsort(q[cells == cell], kind="stable") + w = weight[cells == cell][order] + before = np.cumsum(w) - w + u = q[cells == cell][order] + assert (u >= before / w.sum() - 1e-12).all() + assert (u <= (before + w) / w.sum() + 1e-12).all() + + +def test_tied_outputs_get_independent_quantiles(): + """Outputs everyone ties on (gift aid, which the FRS does not record) + are ranked in independent random orders, so they do not share one + quantile per person, as microimpute 1.8's draw did.""" + n = 2_000 + person = pd.DataFrame( + { + "person_id": np.arange(n), + "person_household_id": np.arange(n), + "age": 70, + "gender": "FEMALE", + "employment_status": "RETIRED", + } + ) + for v in IMPUTATIONS[:6]: + person[v] = 0.0 + household = pd.DataFrame( + {"household_id": np.arange(n), "household_weight": 1.0, "region": "WALES"} + ) + with pytest.MonkeyPatch.context() as mp: + mp.setattr(income_module, "income_model_inputs", _inputs_from_columns) + quantiles = income_module.draw_quantiles(_Dataset(person, household)) + rho = quantiles.corr(method="spearman").to_numpy() + assert np.abs(rho[~np.eye(len(rho), dtype=bool)]).max() < 0.1 + + def test_built_enhanced_frs_spi_pay_rises_with_hours(enhanced_frs): from scipy.stats import spearmanr From 5079ea6374a5f56358e498a894891c43785a67a2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 10:34:08 -0400 Subject: [PATCH 5/5] Describe rank cells as the forests' conditioning at SPI resolution The forests are queried at exact age (trained on uniform noise within the SPI age band) and, after each group's first output, at the outputs drawn before it, so the rank cells match their conditioning only up to that age noise and the chain. Comment only (review r2 nit). Co-Authored-By: Claude Opus 5.5 --- .../datasets/imputations/income.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/policyengine_uk_data/datasets/imputations/income.py b/policyengine_uk_data/datasets/imputations/income.py index b2cf5b36..8f7f0cb9 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -278,11 +278,15 @@ def generate_spi_table( "region", ] -# The SPI gives age only as a band, so the forests learn nothing within a -# band. A person's rank is taken within (earnings group, SPI age band, -# gender, region): the cells the forests condition on. Ranking over a wider -# pool would bias the draw: a London donor's national rank is higher than -# their rank in London, and London's SPI distribution is already higher. +# The SPI gives age only as a band, so the forests' training age is uniform +# noise within it and they learn no real age pattern within a band. A +# person's rank is taken within (earnings group, SPI age band, gender, +# region): the finest cells the SPI can tell apart. The forests are queried +# at the person's exact age and, after each group's first output, at the +# outputs drawn before it, so these cells are the forests' conditioning only +# up to that age noise and the chain. Ranking over a wider pool would bias +# the draw: a London donor's national rank is higher than their rank in +# London, and London's SPI distribution is already higher. SPI_AGE_BAND_STARTS = sorted({low for code, (low, _) in AGE_RANGES.items() if code > 0})