From d81c80a97553b258c1be9584766886f966893387 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 2 Jun 2026 22:43:52 -0400 Subject: [PATCH] Fix CBO income source PE-native loss units --- src/microplex_us/pipelines/pe_native_loss.py | 7 ++++++ tests/pipelines/test_pe_native_loss.py | 23 ++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/microplex_us/pipelines/pe_native_loss.py b/src/microplex_us/pipelines/pe_native_loss.py index 0bf81e5..75d33c9 100644 --- a/src/microplex_us/pipelines/pe_native_loss.py +++ b/src/microplex_us/pipelines/pe_native_loss.py @@ -218,6 +218,13 @@ def infer_pe_native_target_unit(target_name: str) -> str: return "households" if any(part in {"amount", "total"} for part in parts): return "dollars" + if ( + len(parts) >= 3 + and parts[0] == "nation" + and parts[1] == "cbo" + and parts[2] == "income_by_source" + ): + return "dollars" if any(part in {"count", "returns", "filers"} for part in parts): return "returns" if "spending" in normalized or "cost" in normalized or "tax" in normalized: diff --git a/tests/pipelines/test_pe_native_loss.py b/tests/pipelines/test_pe_native_loss.py index 925b381..43fa6de 100644 --- a/tests/pipelines/test_pe_native_loss.py +++ b/tests/pipelines/test_pe_native_loss.py @@ -7,6 +7,7 @@ from microplex_us.pipelines.pe_native_loss import ( build_pe_native_loss_arrays, + infer_pe_native_target_unit, pe_native_huber_loss, pe_native_huber_loss_terms, ) @@ -32,6 +33,28 @@ def test_bucketed_loss_downweights_tiny_baseline_outlier() -> None: assert terms[1] > terms[0] +def test_cbo_income_by_source_filers_targets_are_dollars() -> None: + assert ( + infer_pe_native_target_unit( + "nation/cbo/income_by_source/self_employment_income/filers" + ) + == "dollars" + ) + assert ( + infer_pe_native_target_unit( + "nation/cbo/income_by_source/taxable_interest_income+" + "non_qualified_dividend_income/filers" + ) + == "dollars" + ) + assert ( + infer_pe_native_target_unit( + "nation/irs/adjusted gross income/count/AGI in 0-25k/taxable/All" + ) + == "returns" + ) + + def test_robust_pe_native_optimizer_uses_huber_objective() -> None: matrix = np.asarray([[1.0, 0.0], [0.0, 1.0]]) target = np.asarray([1.0, 1.0])