Consolidate economic impact analysis into shared compute functions

changelog.d/consolidate-economic-impact.added.md

@@ -0,0 +1 @@
+Shared compute functions for economic impact analysis: AnalysisStrategy protocol, economic_impact_analysis, BudgetSummaryItem, compute_program_statistics, compute_decile_impacts, and PolicyReformAnalysis

examples/us_budgetary_impact.py

@@ -0,0 +1,155 @@
+"""Example: US budgetary impact comparison between baseline and reform.
+
+Demonstrates the canonical policyengine.py workflow:
+1. Ensure datasets exist (download + compute or load from cache)
+2. Define a parametric reform
+3. Run baseline and reform simulations
+4. Use economic_impact_analysis() for the full analysis
+5. Use ChangeAggregate for targeted single-metric queries
+
+Run: python examples/us_budgetary_impact.py
+"""
+
+import datetime
+
+from policyengine.core import Parameter, ParameterValue, Policy, Simulation
+from policyengine.outputs.change_aggregate import (
+    ChangeAggregate,
+    ChangeAggregateType,
+)
+from policyengine.tax_benefit_models.us import (
+    economic_impact_analysis,
+    ensure_datasets,
+    us_latest,
+)
+
+
+def main():
+    year = 2026
+
+    # ── Step 1: Get dataset (downloads from HuggingFace on first run) ──
+    print("Ensuring datasets are available...")
+    datasets = ensure_datasets(
+        datasets=["hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"],
+        years=[year],
+        data_folder="./data",
+    )
+    dataset = datasets[f"enhanced_cps_2024_{year}"]
+    print(f"  Loaded: {dataset}")
+
+    # ── Step 2: Define a reform ──
+    # Example: double the standard deduction for single filers
+    param = Parameter(
+        name="gov.irs.deductions.standard.amount.SINGLE",
+        tax_benefit_model_version=us_latest,
+    )
+    reform = Policy(
+        name="Double standard deduction (single)",
+        parameter_values=[
+            ParameterValue(
+                parameter=param,
+                start_date=datetime.date(year, 1, 1),
+                end_date=datetime.date(year, 12, 31),
+                value=30_950,
+            ),
+        ],
+    )
+
+    # ── Step 3: Create simulations ──
+    baseline_sim = Simulation(
+        dataset=dataset,
+        tax_benefit_model_version=us_latest,
+    )
+    reform_sim = Simulation(
+        dataset=dataset,
+        tax_benefit_model_version=us_latest,
+        policy=reform,
+    )
+
+    # ── Step 4a: Quick budgetary number via ChangeAggregate ──
+    # This requires running the simulations first.
+    print("\nRunning simulations...")
+    baseline_sim.run()
+    reform_sim.run()
+
+    tax_change = ChangeAggregate(
+        baseline_simulation=baseline_sim,
+        reform_simulation=reform_sim,
+        variable="household_tax",
+        aggregate_type=ChangeAggregateType.SUM,
+    )
+    tax_change.run()
+    print("\nQuick budgetary result:")
+    print(f"  Tax revenue change: ${tax_change.result / 1e9:.2f}B")
+
+    # Count winners and losers
+    winners = ChangeAggregate(
+        baseline_simulation=baseline_sim,
+        reform_simulation=reform_sim,
+        variable="household_net_income",
+        aggregate_type=ChangeAggregateType.COUNT,
+        change_geq=1,
+    )
+    losers = ChangeAggregate(
+        baseline_simulation=baseline_sim,
+        reform_simulation=reform_sim,
+        variable="household_net_income",
+        aggregate_type=ChangeAggregateType.COUNT,
+        change_leq=-1,
+    )
+    winners.run()
+    losers.run()
+    print(f"  Winners: {winners.result / 1e6:.2f}M households")
+    print(f"  Losers: {losers.result / 1e6:.2f}M households")
+
+    # ── Step 4b: Full analysis via economic_impact_analysis ──
+    # Note: this calls .ensure() internally, which is a no-op here since
+    # we already ran the simulations above. If we hadn't called .run(),
+    # ensure() would run + cache them automatically.
+    print("\nRunning full economic impact analysis...")
+    analysis = economic_impact_analysis(baseline_sim, reform_sim)
+
+    print("\n=== Program-by-Program Impact ===")
+    for prog in analysis.program_statistics.outputs:
+        print(
+            f"  {prog.program_name:30s}  "
+            f"baseline=${prog.baseline_total / 1e9:8.1f}B  "
+            f"reform=${prog.reform_total / 1e9:8.1f}B  "
+            f"change=${prog.change / 1e9:+8.1f}B"
+        )
+
+    print("\n=== Decile Impacts ===")
+    for d in analysis.decile_impacts.outputs:
+        print(
+            f"  Decile {d.decile:2d}:  "
+            f"avg change=${d.absolute_change:+8.0f}  "
+            f"relative={d.relative_change:+.2%}"
+        )
+
+    print("\n=== Poverty ===")
+    for bp, rp in zip(
+        analysis.baseline_poverty.outputs,
+        analysis.reform_poverty.outputs,
+        strict=True,
+    ):
+        print(
+            f"  {bp.poverty_type:30s}  "
+            f"baseline={bp.rate:.4f}  "
+            f"reform={rp.rate:.4f}  "
+            f"change={rp.rate - bp.rate:+.4f}"
+        )
+
+    print("\n=== Inequality ===")
+    bi = analysis.baseline_inequality
+    ri = analysis.reform_inequality
+    print(f"  Gini:           baseline={bi.gini:.4f}  reform={ri.gini:.4f}")
+    print(
+        f"  Top 10% share:  baseline={bi.top_10_share:.4f}  reform={ri.top_10_share:.4f}"
+    )
+    print(
+        f"  Top 1% share:   baseline={bi.top_1_share:.4f}  reform={ri.top_1_share:.4f}"
+    )
+
+
+if __name__ == "__main__":
+    main()

src/policyengine/outputs/__init__.py

@@ -1,5 +1,15 @@
 from policyengine.core import Output, OutputCollection
 from policyengine.outputs.aggregate import Aggregate, AggregateType
+from policyengine.outputs.analysis_strategy import (
+    AnalysisStrategy,
+    InequalityResult,
+    PovertyResult,
+    ProgramDefinition,
+)
+from policyengine.outputs.budget_summary import (
+    BudgetSummaryItem,
+    compute_budget_summary,
+)
 from policyengine.outputs.change_aggregate import (
     ChangeAggregate,
     ChangeAggregateType,
@@ -15,6 +25,10 @@
 from policyengine.outputs.decile_impact import (
     DecileImpact,
     calculate_decile_impacts,
+    compute_decile_impacts,
+)
+from policyengine.outputs.economic_impact import (
+    economic_impact_analysis,
 )
 from policyengine.outputs.inequality import (
     UK_INEQUALITY_INCOME_VARIABLE,
@@ -31,6 +45,7 @@
     LocalAuthorityImpact,
     compute_uk_local_authority_impacts,
 )
+from policyengine.outputs.policy_reform_analysis import PolicyReformAnalysis
 from policyengine.outputs.poverty import (
     AGE_GROUPS,
     GENDER_GROUPS,
@@ -48,6 +63,7 @@
     calculate_us_poverty_by_race,
     calculate_us_poverty_rates,
 )
+from policyengine.outputs.program_statistics import compute_program_statistics
 
 __all__ = [
     "Output",
@@ -86,4 +102,14 @@
     "compute_uk_constituency_impacts",
     "LocalAuthorityImpact",
     "compute_uk_local_authority_impacts",
+    "BudgetSummaryItem",
+    "compute_budget_summary",
+    "compute_decile_impacts",
+    "compute_program_statistics",
+    "PolicyReformAnalysis",
+    "AnalysisStrategy",
+    "ProgramDefinition",
+    "PovertyResult",
+    "InequalityResult",
+    "economic_impact_analysis",
 ]

src/policyengine/outputs/aggregate.py

@@ -47,10 +47,15 @@ def run(self):
 
         # Get variable object
         var_obj = next(
-            v
-            for v in self.simulation.tax_benefit_model_version.variables
-            if v.name == self.variable
+            (
+                v
+                for v in self.simulation.tax_benefit_model_version.variables
+                if v.name == self.variable
+            ),
+            None,
         )
+        if var_obj is None:
+            raise ValueError(f"Variable '{self.variable}' not found in model")
 
         # Get the target entity data
         target_entity = self.entity or var_obj.entity
@@ -68,10 +73,17 @@ def run(self):
         # Apply filters
         if self.filter_variable is not None:
             filter_var_obj = next(
-                v
-                for v in self.simulation.tax_benefit_model_version.variables
-                if v.name == self.filter_variable
+                (
+                    v
+                    for v in self.simulation.tax_benefit_model_version.variables
+                    if v.name == self.filter_variable
+                ),
+                None,
             )
+            if filter_var_obj is None:
+                raise ValueError(
+                    f"Filter variable '{self.filter_variable}' not found in model"
+                )
 
             if filter_var_obj.entity != target_entity:
                 filter_mapped = self.simulation.output_dataset.data.map_to_entity(

src/policyengine/outputs/analysis_strategy.py

@@ -0,0 +1,88 @@
+"""Strategy protocol and result types for economic impact analysis."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, TypedDict, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict
+
+from policyengine.core import OutputCollection
+from policyengine.outputs.inequality import Inequality
+from policyengine.outputs.poverty import Poverty
+
+if TYPE_CHECKING:
+    from policyengine.core.simulation import Simulation
+
+
+class ProgramDefinition(TypedDict):
+    """Definition of a program for program statistics computation."""
+
+    entity: str
+    is_tax: bool
+
+
+class PovertyResult(BaseModel):
+    """Standardised poverty result returned by a country strategy."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    baseline_poverty: OutputCollection[Poverty]
+    reform_poverty: OutputCollection[Poverty]
+    baseline_poverty_by_age: OutputCollection[Poverty] | None = None
+    reform_poverty_by_age: OutputCollection[Poverty] | None = None
+    baseline_poverty_by_gender: OutputCollection[Poverty] | None = None
+    reform_poverty_by_gender: OutputCollection[Poverty] | None = None
+    baseline_poverty_by_race: OutputCollection[Poverty] | None = None
+    reform_poverty_by_race: OutputCollection[Poverty] | None = None
+
+
+class InequalityResult(BaseModel):
+    """Standardised inequality result returned by a country strategy."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    baseline_inequality: Inequality
+    reform_inequality: Inequality
+
+
+@runtime_checkable
+class AnalysisStrategy(Protocol):
+    """Country-specific strategy for economic impact analysis.
+
+    Each property/method corresponds to a standardised extension point
+    in the shared analysis pipeline.
+    """
+
+    @property
+    def income_variable(self) -> str:
+        """Primary income variable for decile / intra-decile analysis."""
+        ...
+
+    @property
+    def budget_variable_names(self) -> list[str]:
+        """Variable names for budget summary.
+
+        Entities are looked up from the tax-benefit system at runtime.
+        """
+        ...
+
+    @property
+    def programs(self) -> dict[str, ProgramDefinition]:
+        """Program definitions: name -> ProgramDefinition."""
+        ...
+
+    def compute_poverty(
+        self,
+        baseline: Simulation,
+        reform: Simulation,
+    ) -> PovertyResult:
+        """Compute all poverty metrics (overall + demographic breakdowns)."""
+        ...
+
+    def compute_inequality(
+        self,
+        baseline: Simulation,
+        reform: Simulation,
+    ) -> InequalityResult:
+        """Compute inequality metrics."""
+        ...