diff --git a/packages/populace-build/src/populace/build/us_runtime/puf_support.py b/packages/populace-build/src/populace/build/us_runtime/puf_support.py index fc06ad0b..209ba353 100644 --- a/packages/populace-build/src/populace/build/us_runtime/puf_support.py +++ b/packages/populace-build/src/populace/build/us_runtime/puf_support.py @@ -191,6 +191,28 @@ def target_order(self) -> tuple[str, ...]: # half is unaffected by the ordinary all-channel sparsification loop. _PUF_TAX_DETAIL_PRESERVE_BASE_ASEC_OUTPUTS = frozenset({"alimony_income"}) +# Sparse, sign-mixed, heavy-tailed person outputs whose imputed *signed* mass +# must be pinned to the donor instrument. The regime-gated QRF imputes such a +# column as an independent sign gate (a HistGradientBoostingClassifier) times +# per-sign magnitude forests. On a rare loss-mixed column the gate regresses the +# positive/negative/zero shares toward balance (it over-predicts the rarer +# leg) and the magnitude forests inflate each leg, so nothing pins the aggregate +# signed total: the imputed net -- a small difference of two large legs -- +# regresses toward a fixed balance point that is nearly independent of the +# donor's true net. A loss-heavy source (SOI Schedule F, net-negative +# nationally) is then dragged toward or past zero, flipping the export sign +# (farm_operations_income) or manufacturing a spurious cancelling leg +# (partnership_self_employment_net_earnings, populace #432). These columns get +# a per-leg mass calibration in finalization -- the signed generalization of the +# donor-positive-rate sparsification -- so the imputed per-unit-weight positive +# and negative leg masses match the donor's and the net sign tracks the source. +_PUF_TAX_DETAIL_SIGNED_MASS_CALIBRATED_PERSON_OUTPUTS = frozenset( + { + "farm_operations_income", + "partnership_self_employment_net_earnings", + } +) + # Known formula-owned outputs the PUF tax-detail donor must never carry as # persistable leaves. This is a documented *seed* set, not the whole story: # :func:`resolve_formula_owned_outputs` unions it with the set derived live @@ -825,6 +847,17 @@ def finalize_us_puf_tax_detail_predictions( person_channel=person_channel, tax_unit_channel=tax_unit_channel, ) + for column in person_outputs: + if column in _PUF_TAX_DETAIL_SIGNED_MASS_CALIBRATED_PERSON_OUTPUTS: + _calibrate_tax_unit_person_output_signed_mass_to_donor( + tables, + column=column, + donor_values=donor[column], + donor_weights=donor["weight"], + household_weights=frame.weights_for("household").values, + person_channel=person_channel, + tax_unit_channel=tax_unit_channel, + ) return Frame( tables, frame.schema, @@ -1566,6 +1599,134 @@ def _sparsify_tax_unit_person_output_to_donor_positive_rate( ) +def _weighted_signed_leg_masses( + values: pd.Series, weights: pd.Series +) -> tuple[float, float, float]: + """Return donor-scale (positive, negative, total) weighted leg masses. + + The two legs are reported separately so a signed calibration can pin each + one; the total weight is returned so callers can work in donor-invariant + per-unit-weight leg masses rather than population-scaled totals. + """ + + numeric_values = ( + pd.to_numeric(values, errors="coerce").fillna(0.0).to_numpy(dtype=np.float64) + ) + numeric_weights = ( + pd.to_numeric(weights, errors="coerce") + .fillna(0.0) + .clip(lower=0.0) + .to_numpy(dtype=np.float64) + ) + positive_mass = float((np.maximum(numeric_values, 0.0) * numeric_weights).sum()) + negative_mass = float((np.minimum(numeric_values, 0.0) * numeric_weights).sum()) + return positive_mass, negative_mass, float(numeric_weights.sum()) + + +def _calibrate_tax_unit_person_output_signed_mass_to_donor( + tables: Mapping[str, pd.DataFrame], + *, + column: str, + donor_values: pd.Series, + donor_weights: pd.Series, + household_weights: np.ndarray, + person_channel: str, + tax_unit_channel: str, +) -> None: + """Pin a signed person output's per-leg mass to the donor instrument. + + Scales the imputed positive and negative legs on the PUF support channel by + separate scalars so each leg's per-unit-weight weighted mass equals the + donor's. The gate's sign assignment (which rows are positive, negative, or + zero) and each leg's relative shape are untouched; only the two leg totals + move, restoring the net signed mass that the regime-gated QRF regresses + toward balance on a sparse, sign-mixed, heavy-tailed column. This is the + signed generalization of + :func:`_sparsify_tax_unit_person_output_to_donor_positive_rate`. + + Applied to the PUF support channel only: the ASEC channel carries measured + source observations (Schedule F operations income is measured on ASEC as + ``FRSE_VAL``) that must never be rescaled. + """ + + person = tables["person"] + household = tables["household"] + tax_unit = tables["tax_unit"] + if len(household_weights) != len(household): + raise ValueError( + "household_weights must align with household rows, got " + f"{len(household_weights)} weights for {len(household)} households." + ) + + donor_positive, donor_negative, donor_weight_total = _weighted_signed_leg_masses( + donor_values, donor_weights + ) + if donor_weight_total <= 0.0: + return + donor_positive_per_weight = donor_positive / donor_weight_total + donor_negative_per_weight = donor_negative / donor_weight_total + + household_weight = pd.Series( + np.asarray(household_weights, dtype=np.float64), + index=household["household_id"], + ) + tax_unit_household_id = ( + person.groupby("person_tax_unit_id", sort=False)["person_household_id"] + .first() + .astype("int64") + ) + tax_unit_weight = tax_unit_household_id.map(household_weight).fillna(0.0) + tax_unit_amount = ( + pd.to_numeric(person[column], errors="coerce") + .fillna(0.0) + .groupby(person["person_tax_unit_id"], sort=False) + .sum() + ) + + channel_tax_unit_ids = tax_unit.loc[ + tax_unit[tax_unit_channel] == PUF_TAX_DETAIL_SUPPORT_CHANNEL, + "tax_unit_id", + ] + amounts = tax_unit_amount.reindex(channel_tax_unit_ids).fillna(0.0) + weights = tax_unit_weight.reindex(channel_tax_unit_ids).fillna(0.0) + channel_weight_total = float(weights.sum()) + if channel_weight_total <= 0.0: + return + + values_array = amounts.to_numpy(dtype=np.float64) + weights_array = weights.to_numpy(dtype=np.float64) + positive = values_array > 0.0 + negative = values_array < 0.0 + imputed_positive_per_weight = ( + float((values_array[positive] * weights_array[positive]).sum()) + / channel_weight_total + ) + imputed_negative_per_weight = ( + float((values_array[negative] * weights_array[negative]).sum()) + / channel_weight_total + ) + + calibrated = values_array.copy() + # A leg with imputed mass can be rescaled to the donor's; a leg the gate + # produced but the donor lacks is scaled to zero (a one-sided donor pins the + # sign). A donor leg the gate produced no rows for cannot be injected by + # scaling, so it is left untouched rather than fabricated. + if imputed_positive_per_weight > 0.0: + calibrated[positive] *= donor_positive_per_weight / imputed_positive_per_weight + if imputed_negative_per_weight < 0.0: + calibrated[negative] *= donor_negative_per_weight / imputed_negative_per_weight + + calibrated_totals = pd.Series(calibrated, index=amounts.index) + mask = person[person_channel] == PUF_TAX_DETAIL_SUPPORT_CHANNEL + _write_person_tax_unit_totals( + person, + mask=mask, + column=column, + totals=calibrated_totals, + nonnegative=False, + ) + + def _reconcile_puf_social_security_components( predictions: pd.DataFrame, person: pd.DataFrame, diff --git a/packages/populace-build/tests/test_us_farm_business_income.py b/packages/populace-build/tests/test_us_farm_business_income.py index 8e44b238..fea34cb6 100644 --- a/packages/populace-build/tests/test_us_farm_business_income.py +++ b/packages/populace-build/tests/test_us_farm_business_income.py @@ -294,7 +294,16 @@ def test_weighted_qrf_preserves_asec_and_writes_signed_puf_support( assert asec[_OPERATIONS].tolist()[:4] == [500.0, -200.0, 800.0, -300.0] assert not asec[_RENT].any() - assert puf[_OPERATIONS].tolist()[:4] == [700.0, -400.0, 900.0, -500.0] + # farm_operations_income now carries the signed-mass calibration: each + # PUF-channel leg is rescaled so its per-unit-weight mass equals the donor's + # (positive 40/4 = 10.0, negative -20/4 = -5.0), pinning the imputed net to + # the source. The raw draw [700, -400, 900, -500] at PUF weight 0.5 has + # per-weight legs 80.0 / -45.0, so the positive leg scales by 10/80 = 0.125 + # and the negative by 5/45 = 1/9; the signs and the measured ASEC leg are + # untouched. farm_rent_income is not signed-mass calibrated. + assert puf[_OPERATIONS].tolist()[:4] == pytest.approx( + [87.5, -400.0 / 9.0, 112.5, -500.0 / 9.0] + ) assert puf[_RENT].tolist()[:4] == [300.0, -100.0, 600.0, -200.0] gate = us_farm_business_income_signal_gate(result) assert gate.passed, gate.failures diff --git a/packages/populace-build/tests/test_us_puf_support.py b/packages/populace-build/tests/test_us_puf_support.py index dbb7a99d..c4e73a58 100644 --- a/packages/populace-build/tests/test_us_puf_support.py +++ b/packages/populace-build/tests/test_us_puf_support.py @@ -1,6 +1,7 @@ """US PUF support-channel expansion tests.""" import importlib +from collections.abc import Sequence import numpy as np import pandas as pd @@ -20,6 +21,7 @@ support_source_id_column, ) from populace.build.us_runtime.puf_support import ( + _PUF_TAX_DETAIL_SIGNED_MASS_CALIBRATED_PERSON_OUTPUTS, PUF_TAX_DETAIL_DEFAULT_PERSON_OUTPUTS, PUF_TAX_DETAIL_DEFAULT_TAX_UNIT_OUTPUTS, PUF_TAX_DETAIL_FORMULA_OWNED_OUTPUTS, @@ -1566,3 +1568,279 @@ def test_production_fit_records_design_kind_under_the_real_engine(self) -> None: assert result.details["resolved_weight_kinds"] == { US_PUF_SUPPORT_FIT_NAME: "design" } + + +# --------------------------------------------------------------------------- # +# Signed-mass calibration: pin the imputed net signed mass of a sparse, +# sign-mixed, heavy-tailed person output (farm_operations_income, populace +# farm-chain-sign-structure; partnership_self_employment_net_earnings, #432) to +# the donor instrument, so the regime-gated QRF's regression toward balance +# cannot flip or cancel the source's net sign. +# --------------------------------------------------------------------------- # + + +def _puf_recipient_frame( + employment_income: Sequence[float], + self_employment_income: Sequence[float], + household_weights: Sequence[float], +) -> Frame: + """One person per household/tax-unit, so person and tax-unit weights align.""" + + n = len(household_weights) + ids = np.arange(1, n + 1, dtype="int64") + person = pd.DataFrame( + { + "person_id": ids, + "person_household_id": ids, + "person_tax_unit_id": ids, + "person_spm_unit_id": ids, + "person_family_id": ids, + "person_marital_unit_id": ids, + "employment_income": np.asarray(employment_income, dtype=np.float64), + "self_employment_income": np.asarray( + self_employment_income, dtype=np.float64 + ), + } + ) + tables = { + "person": person, + "household": pd.DataFrame( + {"household_id": ids, "state_fips": np.full(n, 6, dtype="int64")} + ), + "tax_unit": pd.DataFrame( + {"tax_unit_id": ids, "filing_status_input": ["SINGLE"] * n} + ), + "spm_unit": pd.DataFrame({"spm_unit_id": ids}), + "family": pd.DataFrame({"family_id": ids}), + "marital_unit": pd.DataFrame({"marital_unit_id": ids}), + } + weights = { + "household": Weights( + np.asarray(household_weights, dtype=np.float64), WeightKind.DESIGN + ) + } + return Frame(tables, US_SCHEMA, weights) + + +def _per_weight_legs( + values: Sequence[float], weights: Sequence[float] +) -> tuple[float, float]: + """Return (positive, negative) per-unit-weight weighted leg masses.""" + + values = np.asarray(values, dtype=np.float64) + weights = np.asarray(weights, dtype=np.float64) + total = weights.sum() + positive = float((np.maximum(values, 0.0) * weights).sum() / total) + negative = float((np.minimum(values, 0.0) * weights).sum() / total) + return positive, negative + + +def _puf_channel_person_legs(frame: Frame, column: str) -> tuple[float, float]: + """Per-unit-weight positive/negative leg masses on the PUF person channel.""" + + person = frame.table("person") + household_weight = pd.Series( + frame.weights_for("household").values, + index=frame.table("household")["household_id"], + ) + person_weight = person["person_household_id"].map(household_weight).to_numpy() + puf = ( + person[support_channel_column("person")] == PUF_TAX_DETAIL_SUPPORT_CHANNEL + ).to_numpy() + values = pd.to_numeric(person[column], errors="coerce").fillna(0.0).to_numpy() + return _per_weight_legs(values[puf], person_weight[puf]) + + +@pytest.mark.parametrize( + "column", + [ + "farm_operations_income", + "partnership_self_employment_net_earnings", + ], +) +def test_finalize_pins_signed_person_leg_masses_to_loss_heavy_donor( + monkeypatch: pytest.MonkeyPatch, column: str +) -> None: + """A sign-flipped QRF draw is recalibrated to the donor's signed leg masses. + + Both columns run through the one signed-mass-calibration code path, so a + single fix restores the net sign for the farm defect and #432 alike. + """ + + assert column in _PUF_TAX_DETAIL_SIGNED_MASS_CALIBRATED_PERSON_OUTPUTS + + # The fake QRF returns a sign-flipped draw (net positive, both legs) for the + # six PUF recipients, reproducing the regime-gated QRF's regression toward a + # positive balance point even though the source instrument is loss-heavy. + raw_prediction = np.asarray([600.0, 600.0, 600.0, 600.0, -50.0, -50.0]) + + class FlippedQRF: + def __init__(self, *, n_estimators: int, seed: int) -> None: + pass + + def fit(self, frame, predictors, outputs, *, weights) -> "FlippedQRF": + assert outputs == [column] + assert weights == "design" + return self + + @property + def weight_kind(self) -> str: + return "design" + + def predict( + self, features: pd.DataFrame, *, release_models: bool = False + ) -> pd.DataFrame: + return pd.DataFrame({column: raw_prediction}, index=features.index) + + monkeypatch.setattr(puf_support_module, "QRF", FlippedQRF) + + frame = clone_us_frame_for_puf_support( + _puf_recipient_frame( + employment_income=[50_000.0] * 6, + self_employment_income=[0.0] * 6, + household_weights=[1.0, 1.0, 1.0, 3.0, 3.0, 3.0], + ) + ) + # Loss-heavy donor: small positive per-unit-weight mass, large negative. + donor = pd.DataFrame( + { + "employment_income": [50_000.0] * 6, + column: [200.0, 0.0, -400.0, -400.0, 0.0, 0.0], + "weight": [1.0, 1.0, 2.0, 2.0, 1.0, 1.0], + } + ) + donor_pos, donor_neg = _per_weight_legs( + donor[column].to_numpy(), donor["weight"].to_numpy() + ) + assert donor_pos + donor_neg < 0.0 # the source instrument is loss-heavy + + raw_pos, raw_neg = _per_weight_legs(raw_prediction, [0.5, 0.5, 0.5, 1.5, 1.5, 1.5]) + assert raw_pos + raw_neg > 0.0 # the raw QRF draw flips the net sign + + imputed = impute_us_puf_tax_detail_support( + frame, + donor, + predictors=("puf_predictor_employment_income",), + person_outputs=(column,), + tax_unit_outputs=(), + n_estimators=4, + seed=0, + ) + + final_pos, final_neg = _puf_channel_person_legs(imputed, column) + # Each leg's per-unit-weight mass now equals the donor's, so the net sign + # tracks the loss-heavy source instead of the QRF's flipped draw. + assert final_pos == pytest.approx(donor_pos) + assert final_neg == pytest.approx(donor_neg) + assert final_pos + final_neg == pytest.approx(donor_pos + donor_neg) + assert final_pos + final_neg < 0.0 + + +def test_regime_gated_qrf_farm_net_flips_and_calibration_restores_the_source_sign() -> ( + None +): + """End-to-end: fit the chain's real QRF on a faithful loss-heavy Schedule-F + donor, show the raw draw flips the national net sign, and confirm the + signed-mass calibration restores the source's net-negative mass.""" + + column = "farm_operations_income" + seed = 7 + + def _features(rng: np.random.Generator, n: int) -> tuple[np.ndarray, np.ndarray]: + return ( + np.abs(rng.normal(50_000, 40_000, n)), + np.abs(rng.normal(8_000, 30_000, n)), + ) + + recipient_rng = np.random.default_rng(seed) + wage, self_emp = _features(recipient_rng, 2_500) + frame = clone_us_frame_for_puf_support( + _puf_recipient_frame( + employment_income=wage, + self_employment_income=self_emp, + household_weights=np.clip( + np.exp(recipient_rng.normal(4.0, 1.5, 2_500)), 1.0, 30_000.0 + ), + ) + ) + + donor_rng = np.random.default_rng(seed) + donor_wage, donor_self_emp = _features(donor_rng, 8_000) + donor_weight = np.clip(np.exp(donor_rng.normal(4.0, 1.5, 8_000)), 1.0, 30_000.0) + # Sparse (~3%) signed Schedule-F: 40% gains, 60% heavier losses -> the + # weighted national mass is net-negative, like SOI sole-proprietor farm + # income. + nonzero = donor_rng.random(8_000) < 0.03 + indices = np.flatnonzero(nonzero) + draws = donor_rng.random(len(indices)) + gains = indices[draws < 0.40] + losses = indices[draws >= 0.40] + farm = np.zeros(8_000) + farm[gains] = donor_rng.lognormal(9.7, 1.0, len(gains)) + farm[losses] = -donor_rng.lognormal(10.0, 1.0, len(losses)) + donor = pd.DataFrame( + { + "employment_income": donor_wage, + "self_employment_income": donor_self_emp, + column: farm, + "weight": donor_weight, + } + ) + donor_pos, donor_neg = _per_weight_legs(farm, donor_weight) + donor_net = donor_pos + donor_neg + assert donor_net < 0.0 # loss-heavy source instrument + + raw_holder: dict[str, np.ndarray] = {} + imputed = impute_us_puf_tax_detail_support( + frame, + donor, + predictors=( + "puf_predictor_employment_income", + "puf_predictor_self_employment_income", + ), + person_outputs=(column,), + tax_unit_outputs=(), + n_estimators=60, + seed=0, + raw_predictions_callback=lambda predictions: raw_holder.__setitem__( + "raw", predictions[column].to_numpy().copy() + ), + ) + + # Recover the recipient tax-unit weights to score the raw draw's national net. + person = imputed.table("person") + tax_unit = imputed.table("tax_unit") + household_weight = pd.Series( + imputed.weights_for("household").values, + index=imputed.table("household")["household_id"], + ) + tax_unit_household = person.groupby("person_tax_unit_id", sort=False)[ + "person_household_id" + ].first() + tax_unit_weight = ( + tax_unit["tax_unit_id"].map(tax_unit_household).map(household_weight).to_numpy() + ) + puf_tax_unit = ( + tax_unit[support_channel_column("tax_unit")] == PUF_TAX_DETAIL_SUPPORT_CHANNEL + ).to_numpy() + raw_pos, raw_neg = _per_weight_legs( + raw_holder["raw"], tax_unit_weight[puf_tax_unit] + ) + raw_net = raw_pos + raw_neg + # The defect: the raw QRF draw regresses the net toward a balance point far + # from the loss-heavy source. The dominant loss leg collapses and the + # positive leg inflates, so the imputed net badly mis-tracks the donor -- + # it flips outright on base-m, and here it lands near or past zero. These + # are mechanism properties (robust across solver builds); the exact draw is + # not, which is why the fix pins the mass rather than trusting the draw. + assert raw_neg > donor_neg # the loss leg collapses (less negative) + assert raw_pos > donor_pos # the positive leg inflates + assert abs(raw_net) < 0.5 * abs(donor_net) # net regresses toward balance + + final_pos, final_neg = _puf_channel_person_legs(imputed, column) + # The fix: each PUF-channel leg is pinned to the donor's per-unit-weight + # mass, so the imputed national net tracks the source's net-negative sign. + assert final_pos == pytest.approx(donor_pos, rel=1e-6) + assert final_neg == pytest.approx(donor_neg, rel=1e-6) + assert final_pos + final_neg == pytest.approx(donor_net, rel=1e-6) + assert final_pos + final_neg < 0.0