Add federal vs. state budgetary impact to economic_impact_analysis

changelog.d/289.added.md

@@ -0,0 +1 @@
+Partition US policy reform budgetary impact into federal, state, and unattributed shares via `BudgetaryImpact` on `PolicyReformAnalysis` and the standalone `calculate_budgetary_impact` helper. `total` combines the change in `household_tax` and `household_benefits` with the change in the shared-funding health-program government cost (Medicaid FMAP, CHIP eFMAP, and Medicare Savings Programs), which `household_benefits` excludes by default — guarded against double-counting when `gov.simulation.include_health_benefits_in_net_income` is enabled. `federal` and `state` attribute the cleanly-assignable pieces (`income_tax` + `employee_payroll_tax` − `federal_benefit_cost`; `state_income_tax` − `state_benefit_cost`), and `unattributed` carries the residual so reforms to 100%-federal programs such as SSI or SNAP and to shared-funding health programs both surface instead of silently reading as near-zero totals. `economic_impact_analysis` materializes the `federal_benefit_cost` / `state_benefit_cost` aggregates automatically.

examples/us_budgetary_impact.py

@@ -117,6 +117,13 @@ def main():
     print("\nRunning full economic impact analysis...")
     analysis = economic_impact_analysis(baseline_sim, reform_sim)
 
+    print("\n=== Budgetary Impact (Federal / State / Unattributed) ===")
+    b = analysis.budgetary_impact
+    print(f"  Federal:      ${b.federal / 1e9:+8.1f}B")
+    print(f"  State:        ${b.state / 1e9:+8.1f}B")
+    print(f"  Unattributed: ${b.unattributed / 1e9:+8.1f}B")
+    print(f"  Total:        ${b.total / 1e9:+8.1f}B")
+
     print("\n=== Program-by-Program Impact ===")
     for prog in analysis.program_statistics.outputs:
         print(

src/policyengine/tax_benefit_models/us/__init__.py

@@ -30,7 +30,11 @@
     from policyengine.core import Dataset
     from policyengine.outputs import LaborSupplyResponse, ProgramStatistics
 
-    from .analysis import economic_impact_analysis
+    from .analysis import (
+        BudgetaryImpact,
+        calculate_budgetary_impact,
+        economic_impact_analysis,
+    )
     from .datasets import (
         PolicyEngineUSDataset,
         USYearData,
@@ -75,6 +79,8 @@
         "us_latest",
         "calculate_household",
         "economic_impact_analysis",
+        "calculate_budgetary_impact",
+        "BudgetaryImpact",
         "ProgramStatistics",
         "LaborSupplyResponse",
     ]

src/policyengine/tax_benefit_models/us/analysis.py

@@ -8,7 +8,7 @@
 
 from typing import Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, computed_field
 
 from policyengine.core import OutputCollection, Simulation
 from policyengine.outputs import (
@@ -22,10 +22,15 @@
     configure_labor_supply_response_variables,
     validate_program_statistics_config,
 )
+from policyengine.outputs.change_aggregate import (
+    ChangeAggregate,
+    ChangeAggregateType,
+)
 from policyengine.outputs.decile_impact import (
     DecileImpact,
     calculate_decile_impacts,
 )
+from policyengine.outputs.extra_variables import add_extra_variables
 from policyengine.outputs.inequality import (
     Inequality,
     USInequalityPreset,
@@ -56,11 +61,197 @@
 }
 
 
+class BudgetaryImpact(BaseModel):
+    """Federal/state/unattributed partition of a US reform's budgetary impact.
+
+    Sign convention: positive means the government is better off (more tax
+    revenue or lower benefit spending); negative means revenue lost or
+    spending incurred.
+
+    ``total`` combines the reform-minus-baseline change in ``household_tax`` and
+    ``household_benefits`` with the change in the government cost of the
+    shared-funding health programs (Medicaid FMAP, CHIP eFMAP, and Medicare
+    Savings Programs), which ``household_benefits`` excludes by default.
+    ``federal`` and ``state`` hold only the pieces that are cleanly
+    attributable to a level of government today; ``unattributed`` is the
+    residual (``total - federal - state``) that carries everything not yet
+    split. ``total`` is a computed field, so ``federal + state + unattributed``
+    always equals it by construction. See :func:`calculate_budgetary_impact`
+    for the exact variables behind each share.
+    """
+
+    federal: float = Field(
+        ...,
+        description="Federal budgetary impact, USD (positive = government gain).",
+    )
+    state: float = Field(
+        ...,
+        description="State budgetary impact, USD (positive = government gain).",
+    )
+    unattributed: float = Field(
+        ...,
+        description=(
+            "Budgetary impact not yet attributed to a level of government, "
+            "USD. Residual of total minus federal minus state."
+        ),
+    )
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def total(self) -> float:
+        """Total budgetary impact, USD: federal + state + unattributed."""
+        return self.federal + self.state + self.unattributed
+
+
+def _sum_change(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+    variable: str,
+) -> float:
+    """Reform minus baseline total for a variable."""
+    agg = ChangeAggregate(
+        baseline_simulation=baseline_simulation,
+        reform_simulation=reform_simulation,
+        variable=variable,
+        aggregate_type=ChangeAggregateType.SUM,
+    )
+    agg.run()
+    return float(agg.result)
+
+
+# Budgetary-impact variables that are not in the default US output
+# (household_tax, household_benefits, and the three tax variables are). They
+# must be materialized before the reform simulations run.
+_BUDGETARY_IMPACT_EXTRA_VARIABLES: dict[str, list[str]] = {
+    "person": ["federal_benefit_cost", "state_benefit_cost"],
+}
+
+
+def configure_budgetary_impact_variables(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> None:
+    """Materialize the budgetary-impact aggregates in the reform outputs.
+
+    ``federal_benefit_cost`` / ``state_benefit_cost`` are not in the default US
+    output, so they are added as extra output variables before the simulations
+    run. Call this before :meth:`Simulation.ensure` so the aggregates are
+    present when :func:`calculate_budgetary_impact` reads them;
+    :func:`economic_impact_analysis` does this automatically.
+    """
+    add_extra_variables(baseline_simulation, _BUDGETARY_IMPACT_EXTRA_VARIABLES)
+    add_extra_variables(reform_simulation, _BUDGETARY_IMPACT_EXTRA_VARIABLES)
+
+
+def _include_health_benefits_in_net_income(simulation: Simulation) -> bool:
+    """Whether ``household_benefits`` already includes Medicaid/CHIP/MSP.
+
+    ``gov.simulation.include_health_benefits_in_net_income`` (default False in
+    policyengine-us) controls whether the shared-funding health-program values
+    are folded into ``household_benefits``. When False, those programs are
+    excluded from ``household_benefits``, so :func:`calculate_budgetary_impact`
+    adds their government cost from ``federal_benefit_cost`` /
+    ``state_benefit_cost`` instead. The value is read from the baseline
+    simulation at its dataset year; a missing parameter or year falls back to
+    False (the default).
+    """
+    year = getattr(simulation.dataset, "year", None)
+    if year is None:
+        return False
+    try:
+        parameter = simulation.tax_benefit_model_version.get_parameter(
+            "gov.simulation.include_health_benefits_in_net_income"
+        )
+    except ValueError:
+        return False
+    return bool(parameter._core_param(f"{year}-01-01"))
+
+
+def calculate_budgetary_impact(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> BudgetaryImpact:
+    """Partition a US reform's budgetary impact into federal, state, and
+    unattributed shares.
+
+    All figures are reform-minus-baseline weighted totals. Sign convention:
+    positive means the government is better off (more tax revenue or lower
+    benefit spending).
+
+    ``total`` is the change in ``household_tax`` minus the change in
+    ``household_benefits``, plus the change in the shared-funding health-program
+    government cost (Medicaid FMAP, CHIP eFMAP, and Medicare Savings Programs)
+    read from ``federal_benefit_cost`` / ``state_benefit_cost``. Those health
+    programs are excluded from ``household_benefits`` by default
+    (``gov.simulation.include_health_benefits_in_net_income`` is False), so
+    their cost is added separately; when that parameter is True the values are
+    already in ``household_benefits`` and the cost is not added again, avoiding
+    a double count. ``total`` therefore covers taxes and every benefit the
+    model counts, including shared-funding health-program cost.
+
+    ``federal`` and ``state`` attribute only the cleanly-assignable pieces::
+
+        federal = Δincome_tax + Δemployee_payroll_tax - Δfederal_benefit_cost
+        state   = Δstate_income_tax                   - Δstate_benefit_cost
+
+    where ``federal_benefit_cost`` / ``state_benefit_cost`` split the
+    shared-funding health programs into their federal and state portions.
+
+    ``unattributed`` is the residual, ``total - federal - state``. It carries
+    everything not yet split by level of government — 100%-federal programs such
+    as SSI and SNAP, 100%-state supplements, and any tax outside the three tax
+    variables above — until finer attribution exists. Because ``total`` is
+    measured from the household aggregates and health cost rather than summed
+    from the attributed pieces, a reform to a shared-funding health program
+    surfaces in ``total`` and the federal/state split (residual zero), while a
+    reform to an unattributed program such as SSI surfaces in ``total`` and
+    ``unattributed`` — neither silently reads as zero.
+
+    ``federal_benefit_cost`` / ``state_benefit_cost`` are not in the default US
+    output; :func:`economic_impact_analysis` calls
+    :func:`configure_budgetary_impact_variables` to materialize them before the
+    simulations run. Callers using this helper directly must do the same.
+    """
+    household_tax_change = _sum_change(
+        baseline_simulation, reform_simulation, "household_tax"
+    )
+    household_benefits_change = _sum_change(
+        baseline_simulation, reform_simulation, "household_benefits"
+    )
+    federal_benefit_cost_change = _sum_change(
+        baseline_simulation, reform_simulation, "federal_benefit_cost"
+    )
+    state_benefit_cost_change = _sum_change(
+        baseline_simulation, reform_simulation, "state_benefit_cost"
+    )
+
+    total = household_tax_change - household_benefits_change
+    if not _include_health_benefits_in_net_income(baseline_simulation):
+        # Medicaid/CHIP/MSP are excluded from household_benefits by default, so
+        # add their government cost; when the parameter is True they are already
+        # in household_benefits and adding the cost would double-count.
+        total -= federal_benefit_cost_change + state_benefit_cost_change
+
+    federal = (
+        _sum_change(baseline_simulation, reform_simulation, "income_tax")
+        + _sum_change(baseline_simulation, reform_simulation, "employee_payroll_tax")
+        - federal_benefit_cost_change
+    )
+    state = (
+        _sum_change(baseline_simulation, reform_simulation, "state_income_tax")
+        - state_benefit_cost_change
+    )
+
+    unattributed = total - federal - state
+    return BudgetaryImpact(federal=federal, state=state, unattributed=unattributed)
+
+
 class PolicyReformAnalysis(BaseModel):
     """Complete policy reform analysis result."""
 
     decile_impacts: OutputCollection[DecileImpact]
     program_statistics: OutputCollection[ProgramStatistics]
+    budgetary_impact: BudgetaryImpact
     baseline_poverty: OutputCollection[Poverty]
     reform_poverty: OutputCollection[Poverty]
     baseline_inequality: Inequality
@@ -106,6 +297,7 @@ def economic_impact_analysis(
     )
     if include_cliff_impacts:
         configure_cliff_impact_variables(baseline_simulation, reform_simulation)
+    configure_budgetary_impact_variables(baseline_simulation, reform_simulation)
     _validate_program_statistics_config(baseline_simulation, reform_simulation)
 
     baseline_simulation.ensure()
@@ -149,9 +341,14 @@ def economic_impact_analysis(
         else None
     )
 
+    budgetary_impact = calculate_budgetary_impact(
+        baseline_simulation, reform_simulation
+    )
+
     return PolicyReformAnalysis(
         decile_impacts=decile_impacts,
         program_statistics=program_collection,
+        budgetary_impact=budgetary_impact,
         baseline_poverty=baseline_poverty,
         reform_poverty=reform_poverty,
         baseline_inequality=baseline_inequality,