Draw SPI incomes at each donor's FRS rank within their cell

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.

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"

policyengine_uk_data/tests/test_spi_build.py

@@ -67,6 +67,7 @@
     "BPADUE",
     "MAIND",
     "SEINC_NUM",
+    "MAINSRCE",
 ]
 
 

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,14 +43,15 @@
     CHILD_STATUS,
     EARNINGS_GROUPS,
     EMPLOYEE,
-    EMPLOYEE_AND_SELF_EMPLOYED,
+    EMPLOYEE_MAIN_AND_SELF_EMPLOYED,
     EMPLOYEE_STATUSES,
     IMPUTATIONS,
     MIN_GROUP_SAMPLE_SHARE,
     NO_EARNINGS,
     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,41 +186,71 @@ 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,
     ),
     st.floats(0.01, 1e6),
     st.sampled_from(["pay", "profit"]),
 )
 def test_spi_group_monotone_in_income(records, extra, source):
-    pay, profit, indicator = (np.array(column, dtype=float) for column in zip(*records))
-    before = spi_earnings_group(pay, profit, indicator)
+    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(
         {