SPI prior allocation: spread stratum mass over the sampled quota, not donor weights

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).

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

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()