diff --git a/src/microplex_us/pipelines/us.py b/src/microplex_us/pipelines/us.py index ad1736e..441b53a 100644 --- a/src/microplex_us/pipelines/us.py +++ b/src/microplex_us/pipelines/us.py @@ -281,7 +281,11 @@ PUF_SUPPORT_CLONE_TOP_TAIL_ROUGH_AGI_CAP = 78_999_999.0 PUF_SUPPORT_CLONE_TOP_TAIL_ROUGH_AGI_VARIABLES: tuple[str, ...] = ( "employment_income", + "employment_income_before_lsr", + "tip_income", + "fsla_overtime_premium", "self_employment_income", + "self_employment_income_before_lsr", "taxable_interest_income", "tax_exempt_interest_income", "capital_gains", @@ -296,11 +300,22 @@ "partnership_s_corp_income", "rental_income", "farm_income", + "farm_operations_income", + "farm_rent_income", "ira_distributions", "taxable_pension_income", + "taxable_private_pension_income", + "taxable_ira_distributions", + "taxable_401k_distributions", + "taxable_403b_distributions", + "taxable_sep_distributions", "total_pension_income", "taxable_social_security", "social_security", + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", ) PUF_SUPPORT_CLONE_TOP_TAIL_SCALE_VARIABLES: tuple[str, ...] = ( "capital_gains", @@ -5801,18 +5816,32 @@ def add(variable: str) -> bool: variables.append(variable) return True + def add_first(*variables: str) -> bool: + return any(add(variable) for variable in variables) + + def add_all(*variables: str) -> bool: + added = False + for variable in variables: + added = add(variable) or added + return added + + add_first("employment_income", "employment_income_before_lsr") + if "employment_income" not in variables: + add_all("tip_income", "fsla_overtime_premium") + + add_first("self_employment_income", "self_employment_income_before_lsr") + for variable in ( - "employment_income", - "self_employment_income", "taxable_interest_income", "tax_exempt_interest_income", "partnership_s_corp_income", "rental_income", - "farm_income", - "ira_distributions", ): add(variable) + if not add("farm_income"): + add_all("farm_operations_income", "farm_rent_income") + added_capital_gain_components = False if add("long_term_capital_gains_before_response"): added_capital_gain_components = True @@ -5828,10 +5857,22 @@ def add(variable: str) -> bool: add("qualified_dividend_income") add("non_qualified_dividend_income") - if not add("taxable_pension_income"): - add("total_pension_income") - if not add("taxable_social_security"): - add("social_security") + if not add("taxable_pension_income") and not add("total_pension_income"): + add_all( + "ira_distributions", + "taxable_private_pension_income", + "taxable_ira_distributions", + "taxable_401k_distributions", + "taxable_403b_distributions", + "taxable_sep_distributions", + ) + if not add("taxable_social_security") and not add("social_security"): + add_all( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", + ) if not components: return pd.Series(0.0, index=clone.index, dtype=float), [] diff --git a/tests/pipelines/test_us.py b/tests/pipelines/test_us.py index a9dc303..2fcb831 100644 --- a/tests/pipelines/test_us.py +++ b/tests/pipelines/test_us.py @@ -4193,6 +4193,49 @@ def test_puf_support_clone_top_tail_guard_scales_exported_ltcg_alias(self): ] assert "long_term_capital_gains_before_response" in summary["scaled_variables"] + def test_puf_support_clone_top_tail_guard_counts_exported_income_inputs(self): + pipeline = USMicroplexPipeline( + USMicroplexBuildConfig( + synthesis_backend="seed", + puf_support_clone_enabled=True, + puf_support_clone_top_tail_rough_agi_cap=78_999_999.0, + ) + ) + clone = pd.DataFrame( + { + "employment_income_before_lsr": [178_427.0], + "long_term_capital_gains_before_response": [95_000_000.0], + "taxable_private_pension_income": [127_000.0], + "social_security_retirement": [18_000.0], + } + ) + + guarded, summary = pipeline._apply_puf_support_clone_top_tail_guard( + clone, + integrated_variables=["long_term_capital_gains"], + ) + rough_agi, rough_agi_variables = pipeline._puf_support_clone_top_tail_rough_agi( + guarded + ) + + assert rough_agi.iloc[0] == pytest.approx(78_999_999.0) + assert rough_agi_variables == [ + "employment_income_before_lsr", + "long_term_capital_gains_before_response", + "taxable_private_pension_income", + "social_security_retirement", + ] + assert guarded["employment_income_before_lsr"].iloc[0] == pytest.approx( + 178_427.0 + ) + assert guarded["long_term_capital_gains_before_response"].iloc[ + 0 + ] == pytest.approx(78_999_999.0 - 178_427.0 - 127_000.0 - 18_000.0) + assert summary["affected_rows"] == 1 + assert summary["scale_basis_variables"] == [ + "long_term_capital_gains_before_response" + ] + def test_puf_support_clone_top_tail_guard_can_be_disabled(self): pipeline = USMicroplexPipeline( USMicroplexBuildConfig(