feat: Fix economic_impact_analysis with correct programs, budget summary, and poverty demographics

Fixes #258

- Add BudgetSummaryItem output and compute_budget_summary() for per-variable budget analysis
- Add CountryConfig strategy pattern (US_CONFIG, UK_CONFIG) replacing country conditionals
- Add compute_decile_impacts() accepting already-run simulations
- Add compute_program_statistics() shared function for both US and UK
- Add PolicyReformAnalysis unified result container
- Fix US program names and entities to match API (employee_payroll_tax, ssi on spm_unit)
- Add intra-decile impacts, poverty by demographics (age/gender/race) to analysis
- Fix StopIteration crashes in Aggregate/ChangeAggregate with clear ValueError messages
- Add model_rebuild() for ProgramStatistics/ProgrammeStatistics and BudgetSummaryItem
- Update outputs __init__.py exports

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: anth-volk

src/policyengine/outputs/__init__.py

@@ -1,5 +1,9 @@
 from policyengine.core import Output, OutputCollection
 from policyengine.outputs.aggregate import Aggregate, AggregateType
+from policyengine.outputs.budget_summary import (
+    BudgetSummaryItem,
+    compute_budget_summary,
+)
 from policyengine.outputs.change_aggregate import (
     ChangeAggregate,
     ChangeAggregateType,
@@ -12,9 +16,11 @@
     ConstituencyImpact,
     compute_uk_constituency_impacts,
 )
+from policyengine.outputs.country_config import UK_CONFIG, US_CONFIG, CountryConfig
 from policyengine.outputs.decile_impact import (
     DecileImpact,
     calculate_decile_impacts,
+    compute_decile_impacts,
 )
 from policyengine.outputs.inequality import (
     UK_INEQUALITY_INCOME_VARIABLE,
@@ -31,6 +37,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 +55,7 @@
     calculate_us_poverty_by_race,
     calculate_us_poverty_rates,
 )
+from policyengine.outputs.program_statistics import compute_program_statistics
 
 __all__ = [
     "Output",
@@ -86,4 +94,12 @@
     "compute_uk_constituency_impacts",
     "LocalAuthorityImpact",
     "compute_uk_local_authority_impacts",
+    "BudgetSummaryItem",
+    "compute_budget_summary",
+    "compute_decile_impacts",
+    "compute_program_statistics",
+    "PolicyReformAnalysis",
+    "CountryConfig",
+    "US_CONFIG",
+    "UK_CONFIG",
 ]

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/budget_summary.py

@@ -0,0 +1,94 @@
+"""Budget summary output — totals for key budget variables under baseline and reform."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+from pydantic import ConfigDict
+
+from policyengine.core import Output, OutputCollection
+from policyengine.outputs.aggregate import Aggregate, AggregateType
+
+if TYPE_CHECKING:
+    from policyengine.core.simulation import Simulation
+
+
+class BudgetSummaryItem(Output):
+    """One row of the budget summary — totals for a single variable."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    baseline_simulation: Simulation
+    reform_simulation: Simulation
+    variable_name: str
+    entity: str
+
+    # Results populated by run()
+    baseline_total: float | None = None
+    reform_total: float | None = None
+    change: float | None = None
+
+    def run(self):
+        baseline_agg = Aggregate(
+            simulation=self.baseline_simulation,
+            variable=self.variable_name,
+            aggregate_type=AggregateType.SUM,
+            entity=self.entity,
+        )
+        baseline_agg.run()
+
+        reform_agg = Aggregate(
+            simulation=self.reform_simulation,
+            variable=self.variable_name,
+            aggregate_type=AggregateType.SUM,
+            entity=self.entity,
+        )
+        reform_agg.run()
+
+        self.baseline_total = float(baseline_agg.result)
+        self.reform_total = float(reform_agg.result)
+        self.change = self.reform_total - self.baseline_total
+
+
+def compute_budget_summary(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+    variables: dict[str, str],
+) -> OutputCollection[BudgetSummaryItem]:
+    """Compute budget totals for each variable under baseline and reform.
+
+    Args:
+        baseline_simulation: Already-run baseline simulation.
+        reform_simulation: Already-run reform simulation.
+        variables: Mapping of variable name to entity,
+            e.g. ``{"household_tax": "household"}``.
+
+    Returns:
+        OutputCollection of BudgetSummaryItem objects with a DataFrame.
+    """
+    results: list[BudgetSummaryItem] = []
+    for var_name, entity in variables.items():
+        item = BudgetSummaryItem(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            variable_name=var_name,
+            entity=entity,
+        )
+        item.run()
+        results.append(item)
+
+    df = pd.DataFrame(
+        [
+            {
+                "variable_name": r.variable_name,
+                "entity": r.entity,
+                "baseline_total": r.baseline_total,
+                "reform_total": r.reform_total,
+                "change": r.change,
+            }
+            for r in results
+        ]
+    )
+
+    return OutputCollection(outputs=results, dataframe=df)

src/policyengine/outputs/change_aggregate.py

@@ -59,10 +59,15 @@ def run(self):
 
         # Get variable object
         var_obj = next(
-            v
-            for v in self.baseline_simulation.tax_benefit_model_version.variables
-            if v.name == self.variable
+            (
+                v
+                for v in self.baseline_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
@@ -123,10 +128,17 @@ def run(self):
         # Apply filter_variable filters
         if self.filter_variable is not None:
             filter_var_obj = next(
-                v
-                for v in self.baseline_simulation.tax_benefit_model_version.variables
-                if v.name == self.filter_variable
+                (
+                    v
+                    for v in self.baseline_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 = (

src/policyengine/outputs/country_config.py

@@ -0,0 +1,84 @@
+"""Country configuration strategy — holds all country-specific parameters."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class CountryConfig:
+    """All country-specific parameters needed by compute functions.
+
+    Individual compute functions read the fields they need from this
+    config rather than accepting a ``country_id`` string and branching.
+    """
+
+    country_id: str
+    income_variable: str
+    programs: dict[str, dict] = field(default_factory=dict)
+    budget_variables: dict[str, str] = field(default_factory=dict)
+    poverty_variables: dict[str, str] = field(default_factory=dict)
+    poverty_entity: str = "person"
+    poverty_breakdowns: tuple[str, ...] = ()
+    inequality_income_variable: str | None = None
+    inequality_entity: str = "household"
+
+
+US_CONFIG = CountryConfig(
+    country_id="us",
+    income_variable="household_net_income",
+    programs={
+        "income_tax": {"entity": "tax_unit", "is_tax": True},
+        "employee_payroll_tax": {"entity": "person", "is_tax": True},
+        "snap": {"entity": "spm_unit", "is_tax": False},
+        "tanf": {"entity": "spm_unit", "is_tax": False},
+        "ssi": {"entity": "spm_unit", "is_tax": False},
+        "social_security": {"entity": "person", "is_tax": False},
+    },
+    budget_variables={
+        "household_tax": "household",
+        "household_benefits": "household",
+        "household_net_income": "household",
+        "household_state_income_tax": "tax_unit",
+    },
+    poverty_variables={
+        "spm": "spm_unit_is_in_spm_poverty",
+        "spm_deep": "spm_unit_is_in_deep_spm_poverty",
+    },
+    poverty_entity="person",
+    poverty_breakdowns=("age", "gender", "race"),
+    inequality_income_variable="household_net_income",
+    inequality_entity="household",
+)
+
+UK_CONFIG = CountryConfig(
+    country_id="uk",
+    income_variable="equiv_hbai_household_net_income",
+    programs={
+        "income_tax": {"entity": "person", "is_tax": True},
+        "national_insurance": {"entity": "person", "is_tax": True},
+        "vat": {"entity": "household", "is_tax": True},
+        "council_tax": {"entity": "household", "is_tax": True},
+        "universal_credit": {"entity": "person", "is_tax": False},
+        "child_benefit": {"entity": "person", "is_tax": False},
+        "pension_credit": {"entity": "person", "is_tax": False},
+        "income_support": {"entity": "person", "is_tax": False},
+        "working_tax_credit": {"entity": "person", "is_tax": False},
+        "child_tax_credit": {"entity": "person", "is_tax": False},
+    },
+    budget_variables={
+        "household_tax": "household",
+        "household_benefits": "household",
+        "household_net_income": "household",
+    },
+    poverty_variables={
+        "absolute_bhc": "in_poverty_bhc",
+        "absolute_ahc": "in_poverty_ahc",
+        "relative_bhc": "in_relative_poverty_bhc",
+        "relative_ahc": "in_relative_poverty_ahc",
+    },
+    poverty_entity="person",
+    poverty_breakdowns=("age", "gender"),
+    inequality_income_variable="equiv_hbai_household_net_income",
+    inequality_entity="household",
+)

src/policyengine/outputs/decile_impact.py

@@ -96,6 +96,63 @@ def run(self):
         self.count_no_change = float((absolute_change[mask] == 0).sum())
 
 
+def compute_decile_impacts(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+    income_variable: str = "equiv_hbai_household_net_income",
+    entity: str | None = None,
+    quantiles: int = 10,
+) -> OutputCollection[DecileImpact]:
+    """Calculate decile-by-decile impact using already-run simulations.
+
+    Unlike ``calculate_decile_impacts`` this does **not** create new
+    Simulation objects — it works directly with the provided ones.
+
+    Args:
+        baseline_simulation: Already-run baseline simulation.
+        reform_simulation: Already-run reform simulation.
+        income_variable: Variable to measure income changes.
+        entity: Entity to aggregate on (default: variable's entity).
+        quantiles: Number of quantiles (default 10 for deciles).
+
+    Returns:
+        OutputCollection of DecileImpact objects with a DataFrame.
+    """
+    results = []
+    for decile in range(1, quantiles + 1):
+        impact = DecileImpact(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            income_variable=income_variable,
+            entity=entity,
+            decile=decile,
+            quantiles=quantiles,
+        )
+        impact.run()
+        results.append(impact)
+
+    df = pd.DataFrame(
+        [
+            {
+                "baseline_simulation_id": r.baseline_simulation.id,
+                "reform_simulation_id": r.reform_simulation.id,
+                "income_variable": r.income_variable,
+                "decile": r.decile,
+                "baseline_mean": r.baseline_mean,
+                "reform_mean": r.reform_mean,
+                "absolute_change": r.absolute_change,
+                "relative_change": r.relative_change,
+                "count_better_off": r.count_better_off,
+                "count_worse_off": r.count_worse_off,
+                "count_no_change": r.count_no_change,
+            }
+            for r in results
+        ]
+    )
+
+    return OutputCollection(outputs=results, dataframe=df)
+
+
 def calculate_decile_impacts(
     dataset: Dataset,
     tax_benefit_model_version: TaxBenefitModelVersion,