diff --git a/packages/populace-build/src/populace/build/us_runtime/voluntary_filing.py b/packages/populace-build/src/populace/build/us_runtime/voluntary_filing.py index 2e5854c3..ed5c4d39 100644 --- a/packages/populace-build/src/populace/build/us_runtime/voluntary_filing.py +++ b/packages/populace-build/src/populace/build/us_runtime/voluntary_filing.py @@ -745,15 +745,30 @@ def _source_receiver_rows( .groupby(rows["_source_id"]) .sum() ) - if not asec_counts.eq(1).all(): - bad = asec_counts.index[~asec_counts.eq(1)].tolist() + if asec_counts.gt(1).any(): + bad = asec_counts.index[asec_counts.gt(1)].tolist() raise ValueError( - "US voluntary-filing support source units require exactly one " - f"ASEC row; invalid source unit(s) {bad[:5]}." + "US voluntary-filing support source units carry duplicated " + f"ASEC rows; invalid source unit(s) {bad[:5]}." ) - source_rows = rows.loc[ - rows["_support_channel"].eq(_BASE_ASEC_SUPPORT_CHANNEL) - ].copy() + # Prefer each unit's ASEC row, but a frozen-support selection may + # legitimately keep only a unit's PUF clone (the L0-survivor case the + # SSI reporter lineage also handles — Build M's certified 57,240 + # selection does exactly this). Clones carry the unit's source + # predictors, so the surviving row predicts identically; pick it + # deterministically by channel then tax-unit id. + ordered_rows = rows.copy() + ordered_rows["_asec_rank"] = ( + ~ordered_rows["_support_channel"].eq(_BASE_ASEC_SUPPORT_CHANNEL) + ).astype(int) + source_rows = ( + ordered_rows.sort_values( + ["_source_id", "_asec_rank", "_support_channel", "_tax_unit_id"], + kind="stable", + ) + .drop_duplicates("_source_id", keep="first") + .drop(columns="_asec_rank") + ) else: if rows["_source_id"].duplicated().any(): duplicates = rows.loc[ diff --git a/packages/populace-build/tests/test_us_voluntary_filing.py b/packages/populace-build/tests/test_us_voluntary_filing.py index 1311bb72..a5ec9041 100644 --- a/packages/populace-build/tests/test_us_voluntary_filing.py +++ b/packages/populace-build/tests/test_us_voluntary_filing.py @@ -593,6 +593,53 @@ def fit( assert (by_source.nunique() == 1).all() +def test_puf_only_survivor_units_predict_from_the_surviving_clone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A unit whose ASEC row was dropped by selection still predicts once. + + Build M's sparse run died here: the certified frozen-support selection + keeps only the PUF clone for some source units (the L0-survivor case the + SSI reporter lineage already handles), and the receiver demanded exactly + one ASEC row per unit. The surviving clone carries the unit's source + predictors, so it serves as the prediction row; duplicated ASEC rows + remain a hard error. + """ + + expanded = clone_us_frame_for_puf_support(_frame(10)) + person = expanded.table("person") + tax_unit = expanded.table("tax_unit") + dropped_source = tax_unit["tax_unit_source_id"].iloc[0] + dropped_units = tax_unit.loc[ + tax_unit["tax_unit_source_id"].eq(dropped_source) + & tax_unit["tax_unit_support_channel"].eq("asec"), + "tax_unit_id", + ] + dropped = person["person_tax_unit_id"].isin(dropped_units) + sparse = expanded.select(~dropped.to_numpy()) + + class FakeFitted: + def predict(self, receiver: pd.DataFrame) -> pd.DataFrame: + return pd.DataFrame( + {_OUTPUT: np.ones(len(receiver), dtype=bool)}, + index=receiver.index, + ) + + class FakeQRF: + def __init__(self, **_kwargs: object) -> None: + pass + + def fit(self, *_args: object, **_kwargs: object) -> FakeFitted: + return FakeFitted() + + monkeypatch.setattr(module, "QRF", FakeQRF) + predicted = impute_us_voluntary_filing(sparse, _donor(), seed=17) + survivors = sparse.table("tax_unit")["tax_unit_source_id"].eq(dropped_source) + assert survivors.any() + assert len(predicted) == len(sparse.table("tax_unit")) + assert predicted[survivors.to_numpy()].all() + + def test_real_qrf_recomputation_is_deterministic() -> None: frame = _frame(14) donor = _donor(120)