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..bc9a64f8 --- /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 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 799b5f35..8f7f0cb9 100644 --- a/policyengine_uk_data/datasets/imputations/income.py +++ b/policyengine_uk_data/datasets/imputations/income.py @@ -8,6 +8,17 @@ 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, 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 @@ -68,16 +79,44 @@ 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 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" -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 @@ -92,29 +131,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( @@ -125,20 +178,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( @@ -190,7 +245,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: @@ -220,6 +278,83 @@ def generate_spi_table( "region", ] +# 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}) + + +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; ids must be unique. + """ + 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") + 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)) + 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 +380,73 @@ 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 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). +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, + _RandomForestClassifierModel, + ) + + 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] + 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) + return output + INCOME_MODEL_METADATA = { "spi_release_name": SPI_RELEASE_NAME, @@ -277,7 +479,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 +495,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 +511,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 +624,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 +681,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 +774,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 +784,7 @@ def impute_income(dataset: UKSingleYearDataset) -> UKSingleYearDataset: zero_weight_copy, model, IMPUTATIONS, + quantiles, ) # Second-stage QRF: rewrite FRS-only variables (benefit `_reported` @@ -554,6 +826,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..d1929c38 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, @@ -107,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_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( { 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..444e3f06 --- /dev/null +++ b/policyengine_uk_data/tests/test_spi_income_rank_draw.py @@ -0,0 +1,654 @@ +"""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. 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. 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. ``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). +""" + +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() + + +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, ids", + [ + ([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, ids): + with pytest.raises(ValueError): + rank_quantiles(values, [0, 0], weights, ids) + + +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) + + +@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 ------------------------------------------ + + +@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() + + +@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 + + 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 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(