From a8015a8765c23854b94b05d04bebf215804ab87c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 17 Aug 2026 12:15:58 -0400 Subject: [PATCH] Fix SPI prior allocation by replacement stratum --- .../630-spi-prior-stratum-allocation.fixed.md | 1 + .../microcosm/build/uk_runtime/spi_support.py | 64 ++++++++--- .../tests/test_uk_spi_support.py | 102 ++++++++++++++++++ 3 files changed, 151 insertions(+), 16 deletions(-) create mode 100644 changelog.d/630-spi-prior-stratum-allocation.fixed.md diff --git a/changelog.d/630-spi-prior-stratum-allocation.fixed.md b/changelog.d/630-spi-prior-stratum-allocation.fixed.md new file mode 100644 index 00000000..88e58f88 --- /dev/null +++ b/changelog.d/630-spi-prior-stratum-allocation.fixed.md @@ -0,0 +1 @@ +Allocate the UK SPI prior from each reviewed replacement stratum's incoming mass and spread its SPI share evenly over that stratum's sampled quota, preventing extreme calibrated donor lineages from propagating once per selected clone while preserving national and cell mass (microcosm#630). diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_support.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_support.py index 04d44d6f..3501d179 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_support.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_support.py @@ -337,9 +337,13 @@ def replace_uk_spi_support_tables( plus a zero-weight SPI-synthetic channel. This transform removes that dead channel, samples replacement source households within the exact reviewed clone/capital-gains/region quotas, rebuilds linked rows once, and allocates - a fixed share of national household mass to the new channel. The allocation - advances weights to ``IMPORTANCE`` and records a deliberate, factor-one - :class:`MassChangeRecord`; national mass never doubles or disappears. + a fixed share of every stratum's household mass to the new channel. Because + the source draw is uniform within each stratum, its SPI design weight is the + stratum's allocated mass divided by its sampled quota; selected donors do + not propagate their incoming calibrated weights. The allocation advances + weights to ``IMPORTANCE`` and records a deliberate, factor-one + :class:`MassChangeRecord`; national and stratum mass never double or + disappear. """ if not isinstance(input_weight_kind, WeightKind): @@ -420,6 +424,7 @@ def replace_uk_spi_support_tables( spi_prior_mass_share=share, input_weight_kind=input_weight_kind, mass_log=mass_log, + strata_columns=stratum_columns, ) return UKSPISupportResult( person=allocated.person, @@ -691,6 +696,7 @@ def _allocate_spi_prior_mass( spi_prior_mass_share: float, input_weight_kind: WeightKind, mass_log: tuple[MassChangeRecord, ...], + strata_columns: tuple[str, ...], ) -> UKSPISupportResult: household = result.household.copy() channels = household[support_channel_column("household")] @@ -705,24 +711,50 @@ def _allocate_spi_prior_mass( household["household_weight"], errors="coerce" ).to_numpy(dtype=float, na_value=np.nan) old_total = float(pre_weights.sum()) - source_weights = pd.Series( - pre_weights[base_mask], - index=household.loc[base_mask, "household_id"].to_numpy(), + _require_columns(household, strata_columns, label="rebuilt household") + allocation = household[list(strata_columns)].copy() + allocation["_base_mass"] = np.where(base_mask, pre_weights, 0.0) + allocation["_spi_count"] = spi_mask.astype(np.int64) + grouped = allocation.groupby( + list(strata_columns), + sort=False, + dropna=False, ) - spi_source_ids = household.loc[ - spi_mask, - support_source_id_column("household"), - ] - spi_raw = spi_source_ids.map(source_weights).to_numpy(dtype=np.float64) - if not np.isfinite(spi_raw).all() or not (spi_raw > 0.0).all(): + stratum_base_mass = grouped["_base_mass"].transform("sum").to_numpy(dtype=float) + stratum_spi_count = grouped["_spi_count"].transform("sum").to_numpy( + dtype=np.int64 + ) + unrepresented_base = base_mask & (stratum_base_mass > 0.0) & ( + stratum_spi_count == 0 + ) + if unrepresented_base.any(): + examples = ( + household.loc[unrepresented_base, list(strata_columns)] + .drop_duplicates() + .head(5) + .to_dict(orient="records") + ) + raise ValueError( + "Every positive-mass UK base stratum must have rebuilt SPI support; " + f"missing stratum example(s): {examples}." + ) + invalid_spi = spi_mask & ( + ~np.isfinite(stratum_base_mass) + | (stratum_base_mass <= 0.0) + | (stratum_spi_count <= 0) + ) + if invalid_spi.any(): raise ValueError( - "Every rebuilt SPI household must inherit a strictly positive " - "source-household prior." + "Every rebuilt SPI stratum must resolve to positive finite base " + "mass and a positive support quota." ) + final_weights = np.zeros_like(pre_weights) final_weights[base_mask] = pre_weights[base_mask] * (1.0 - spi_prior_mass_share) - final_weights[spi_mask] = spi_raw * ( - old_total * spi_prior_mass_share / float(spi_raw.sum()) + final_weights[spi_mask] = ( + stratum_base_mass[spi_mask] + * spi_prior_mass_share + / stratum_spi_count[spi_mask] ) # This is a newly assembled two-channel population, not an in-place diff --git a/packages/microcosm-build/tests/test_uk_spi_support.py b/packages/microcosm-build/tests/test_uk_spi_support.py index 20256e2c..062bba2e 100644 --- a/packages/microcosm-build/tests/test_uk_spi_support.py +++ b/packages/microcosm-build/tests/test_uk_spi_support.py @@ -303,6 +303,14 @@ def test_replace_spi_support_preserves_quotas_and_allocates_real_mass() -> None: dead_spi.groupby(strata, dropna=False).size(), spi.groupby(strata, dropna=False).size(), ) + incoming_base = dead.household[ + ~dead.household[HOUSEHOLD_IS_SPI_SYNTHETIC_COLUMN] + ] + pd.testing.assert_series_equal( + result.household.groupby(strata, dropna=False)["household_weight"].sum(), + incoming_base.groupby(strata, dropna=False)["household_weight"].sum(), + ) + assert spi.groupby(strata, dropna=False)["household_weight"].nunique().eq(1).all() assert len(result.mass_log) == 1 record = result.mass_log[-1] assert record.entity == "household" @@ -341,6 +349,100 @@ def test_replace_spi_support_builds_importance_pool_from_calibrated_base() -> No assert result.mass_log[-1].new_total == original_total +def _repeated_outlier_lineage_support(): + """Two extreme donor lineages repeated across ten real-shaped clone strata.""" + + household_rows: list[dict[str, object]] = [] + dead_source_ids: list[int] = [] + person_rows: list[dict[str, int]] = [] + benunit_rows: list[dict[str, int]] = [] + next_household_id = 1 + for clone_index in range(10): + for region, lineage, outlier_weight in ( + ("LONDON", 11_852, 1_000.0), + ("SCOTLAND", 14_876, 800.0), + ): + for position in range(100): + household_id = next_household_id + next_household_id += 1 + source_lineage = lineage if position == 0 else household_id + household_rows.append( + { + "household_id": household_id, + "household_weight": outlier_weight if position == 0 else 1.0, + "region": region, + "clone_index": clone_index, + "household_is_capital_gains_clone": False, + "source_household_id": source_lineage, + } + ) + if position < 60: + dead_source_ids.append(household_id) + person_rows.append( + { + "person_id": household_id, + "person_household_id": household_id, + "person_benunit_id": household_id, + } + ) + benunit_rows.append({"benunit_id": household_id}) + return create_uk_spi_support_tables( + person=pd.DataFrame(person_rows), + benunit=pd.DataFrame(benunit_rows), + household=pd.DataFrame(household_rows), + selected_household_ids=dead_source_ids, + source_year=2023, + ) + + +def test_spi_prior_spreads_stratum_mass_instead_of_propagating_outliers() -> None: + """#630: selected descendants must not each inherit an extreme donor weight.""" + + dead = _repeated_outlier_lineage_support() + result = replace_uk_spi_support_tables( + person=dead.person, + benunit=dead.benunit, + household=dead.household, + seed=42, + source_year=2023, + ) + channel = support_channel_column("household") + source_id = support_source_id_column("household") + spi = result.household[result.household[channel] == SPI_SYNTHETIC_SUPPORT_CHANNEL] + incoming_base = dead.household[ + ~dead.household[HOUSEHOLD_IS_SPI_SYNTHETIC_COLUMN] + ] + incoming_weights = incoming_base.set_index("household_id")["household_weight"] + selected_source_weights = spi[source_id].map(incoming_weights).to_numpy(dtype=float) + old_total = float(incoming_weights.sum()) + + # Reconstruct the removed behavior: copy selected calibrated weights and + # normalize them once across the whole SPI channel. Seed 42 selects twelve + # descendants of the two extreme lineages, matching the real defect's + # 11-12-row shape, and the legacy prior breaches the certified June fence. + legacy = np.concatenate( + ( + incoming_weights.to_numpy(dtype=float) * 0.5, + selected_source_weights + * (old_total * 0.5 / float(selected_source_weights.sum())), + ) + ) + legacy_median = float(np.median(legacy[legacy > 0.0])) + old_fence = 1151.2542195939373 + assert int((legacy > old_fence * legacy_median).sum()) == 12 + assert float(legacy.max() / legacy_median) > old_fence + + fixed = result.household["household_weight"].to_numpy(dtype=float) + fixed_median = float(np.median(fixed[fixed > 0.0])) + assert float(fixed.max() / fixed_median) < old_fence + strata = ["clone_index", "household_is_capital_gains_clone", "region"] + assert spi.groupby(strata, dropna=False)["household_weight"].nunique().eq(1).all() + pd.testing.assert_series_equal( + result.household.groupby(strata, dropna=False)["household_weight"].sum(), + incoming_base.groupby(strata, dropna=False)["household_weight"].sum(), + ) + + def test_replace_spi_support_corrects_rounding_residue_to_exact_mass() -> None: dead = _certified_like_dead_spi_support() household = dead.household.copy()