Extract shared program-statistics validation and build helpers

changelog.d/263.changed.md

@@ -0,0 +1 @@
+Extracted the duplicated US and UK program-statistics validation and build logic into shared `validate_program_statistics_config` and `build_program_statistics` helpers in `policyengine.outputs.program_statistics`.

src/policyengine/outputs/__init__.py

@@ -62,7 +62,11 @@
     calculate_us_poverty_by_race,
     calculate_us_poverty_rates,
 )
-from policyengine.outputs.program_statistics import ProgramStatistics
+from policyengine.outputs.program_statistics import (
+    ProgramStatistics,
+    build_program_statistics,
+    validate_program_statistics_config,
+)
 from policyengine.outputs.uk_geography_assets import (
     CONSTITUENCY_ASSET_SPEC,
     LOCAL_AUTHORITY_ASSET_SPEC,
@@ -89,6 +93,8 @@
     "DecileImpact",
     "calculate_decile_impacts",
     "ProgramStatistics",
+    "build_program_statistics",
+    "validate_program_statistics_config",
     "IntraDecileImpact",
     "compute_intra_decile_impacts",
     "HoursResponse",

src/policyengine/outputs/program_statistics.py

@@ -2,14 +2,16 @@
 
 from typing import Optional
 
+import pandas as pd
 from pydantic import ConfigDict
 
-from policyengine.core import Output, Simulation
+from policyengine.core import Output, OutputCollection, Simulation
 from policyengine.outputs.aggregate import Aggregate, AggregateType
 from policyengine.outputs.change_aggregate import (
     ChangeAggregate,
     ChangeAggregateType,
 )
+from policyengine.utils.errors import format_conditional_error_detail
 
 
 class ProgramStatistics(Output):
@@ -108,3 +110,128 @@ def run(self):
         self.reform_count = float(reform_count.result)
         self.winners = float(winners.result)
         self.losers = float(losers.result)
+
+
+def _format_missing_program_variables(missing_variables: set[str]) -> Optional[str]:
+    """Format the optional missing-variable detail for program statistics."""
+    return format_conditional_error_detail(
+        "Missing model variables",
+        missing_variables,
+    )
+
+
+def _program_statistics_config_error_message(
+    country_label: str,
+    missing_variables: set[str],
+    missing_outputs: set[tuple[str, str]],
+) -> str:
+    lines = [f"{country_label} program statistics config is invalid:"]
+
+    missing_variables_message = _format_missing_program_variables(missing_variables)
+    if missing_variables_message is not None:
+        lines.append(missing_variables_message)
+
+    if missing_outputs:
+        formatted = ", ".join(
+            f"{program_name} on {entity}"
+            for program_name, entity in sorted(missing_outputs)
+        )
+        lines.append("Variables not materialized in simulation outputs: " + formatted)
+        lines.append(
+            "Add them to the model version's entity_variables or pass them "
+            "via Simulation.extra_variables before running the simulation."
+        )
+
+    return "\n".join(lines)
+
+
+def validate_program_statistics_config(
+    programs: dict[str, dict],
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+    country_label: str,
+) -> None:
+    """Validate program-statistics variables before running simulations.
+
+    ``programs`` maps each program-statistics variable name to its metadata.
+    ``country_label`` (for example ``"US"`` or ``"UK"``) only shapes the error
+    message; the validation logic itself is country-agnostic. Raises
+    ``ValueError`` if any program variable is missing from the model or is not
+    materialized in the simulation outputs.
+    """
+    missing_variables: set[str] = set()
+    missing_outputs: set[tuple[str, str]] = set()
+
+    simulations = (baseline_simulation, reform_simulation)
+    for program_name in programs:
+        for simulation in simulations:
+            model_version = simulation.tax_benefit_model_version
+            try:
+                variable = model_version.get_variable(program_name)
+            except ValueError:
+                missing_variables.add(program_name)
+                continue
+
+            resolved_variables = model_version.resolve_entity_variables(simulation)
+            if program_name not in resolved_variables.get(variable.entity, []):
+                missing_outputs.add((program_name, variable.entity))
+
+    if not missing_variables and not missing_outputs:
+        return
+
+    raise ValueError(
+        _program_statistics_config_error_message(
+            country_label,
+            missing_variables,
+            missing_outputs,
+        ),
+    )
+
+
+def build_program_statistics(
+    programs: dict[str, dict],
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> OutputCollection[ProgramStatistics]:
+    """Run program statistics for each configured program.
+
+    ``programs`` maps each program-statistics variable name to its metadata
+    (currently just ``is_tax``). Each program's entity is derived from the
+    model's variable metadata, so the set of programs cannot silently drift when
+    a country package moves a variable between entities. Returns the collection
+    of ``ProgramStatistics`` with an assembled dataframe of per-program totals,
+    counts, winners, and losers.
+    """
+    model_version = baseline_simulation.tax_benefit_model_version
+    program_statistics = []
+    for program_name, program_info in programs.items():
+        stats = ProgramStatistics(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            program_name=program_name,
+            entity=model_version.get_variable(program_name).entity,
+            is_tax=program_info["is_tax"],
+        )
+        stats.run()
+        program_statistics.append(stats)
+
+    program_df = pd.DataFrame(
+        [
+            {
+                "baseline_simulation_id": p.baseline_simulation.id,
+                "reform_simulation_id": p.reform_simulation.id,
+                "program_name": p.program_name,
+                "entity": p.entity,
+                "is_tax": p.is_tax,
+                "baseline_total": p.baseline_total,
+                "reform_total": p.reform_total,
+                "change": p.change,
+                "baseline_count": p.baseline_count,
+                "reform_count": p.reform_count,
+                "winners": p.winners,
+                "losers": p.losers,
+            }
+            for p in program_statistics
+        ]
+    )
+    return OutputCollection(outputs=program_statistics, dataframe=program_df)

src/policyengine/tax_benefit_models/uk/analysis.py

@@ -6,18 +6,19 @@
 
 from __future__ import annotations
 
-import pandas as pd
 from pydantic import BaseModel
 
 from policyengine.core import OutputCollection, Simulation
 from policyengine.outputs import (
     CliffImpact,
     LaborSupplyResponse,
     ProgramStatistics,
+    build_program_statistics,
     calculate_cliff_impact,
     calculate_labor_supply_response,
     configure_cliff_impact_variables,
     configure_labor_supply_response_variables,
+    validate_program_statistics_config,
 )
 from policyengine.outputs.decile_impact import (
     DecileImpact,
@@ -35,7 +36,6 @@
     Poverty,
     calculate_uk_poverty_rates,
 )
-from policyengine.utils.errors import format_conditional_error_detail
 
 # Map of UK program-statistics variable name -> program metadata. The
 # entity for each program is derived from the variable's own metadata at
@@ -77,68 +77,16 @@ class PolicyReformAnalysis(BaseModel):
     cliff_impact: CliffImpact | None = None
 
 
-def _format_missing_program_variables(missing_variables: set[str]) -> str | None:
-    """Format the optional missing-variable detail for program statistics."""
-    return format_conditional_error_detail(
-        "Missing model variables",
-        missing_variables,
-    )
-
-
-def _uk_program_statistics_config_error_message(
-    missing_variables: set[str],
-    missing_outputs: set[tuple[str, str]],
-) -> str:
-    lines = ["UK program statistics config is invalid:"]
-
-    missing_variables_message = _format_missing_program_variables(missing_variables)
-    if missing_variables_message is not None:
-        lines.append(missing_variables_message)
-
-    if missing_outputs:
-        formatted = ", ".join(
-            f"{program_name} on {entity}"
-            for program_name, entity in sorted(missing_outputs)
-        )
-        lines.append("Variables not materialized in simulation outputs: " + formatted)
-        lines.append(
-            "Add them to the model version's entity_variables or pass them "
-            "via Simulation.extra_variables before running the simulation."
-        )
-
-    return "\n".join(lines)
-
-
 def _validate_program_statistics_config(
     baseline_simulation: Simulation,
     reform_simulation: Simulation,
 ) -> None:
     """Validate UK program-stat variables before running simulations."""
-    missing_variables: set[str] = set()
-    missing_outputs: set[tuple[str, str]] = set()
-
-    simulations = (baseline_simulation, reform_simulation)
-    for program_name in UK_PROGRAMS:
-        for simulation in simulations:
-            model_version = simulation.tax_benefit_model_version
-            try:
-                variable = model_version.get_variable(program_name)
-            except ValueError:
-                missing_variables.add(program_name)
-                continue
-
-            resolved_variables = model_version.resolve_entity_variables(simulation)
-            if program_name not in resolved_variables.get(variable.entity, []):
-                missing_outputs.add((program_name, variable.entity))
-
-    if not missing_variables and not missing_outputs:
-        return
-
-    raise ValueError(
-        _uk_program_statistics_config_error_message(
-            missing_variables,
-            missing_outputs,
-        ),
+    validate_program_statistics_config(
+        UK_PROGRAMS,
+        baseline_simulation,
+        reform_simulation,
+        "UK",
     )
 
 
@@ -186,40 +134,10 @@ def economic_impact_analysis(
         entity="household",
     )
 
-    model_version = baseline_simulation.tax_benefit_model_version
-    program_statistics = []
-    for program_name, program_info in UK_PROGRAMS.items():
-        stats = ProgramStatistics(
-            baseline_simulation=baseline_simulation,
-            reform_simulation=reform_simulation,
-            program_name=program_name,
-            entity=model_version.get_variable(program_name).entity,
-            is_tax=program_info["is_tax"],
-        )
-        stats.run()
-        program_statistics.append(stats)
-
-    program_df = pd.DataFrame(
-        [
-            {
-                "baseline_simulation_id": p.baseline_simulation.id,
-                "reform_simulation_id": p.reform_simulation.id,
-                "program_name": p.program_name,
-                "entity": p.entity,
-                "is_tax": p.is_tax,
-                "baseline_total": p.baseline_total,
-                "reform_total": p.reform_total,
-                "change": p.change,
-                "baseline_count": p.baseline_count,
-                "reform_count": p.reform_count,
-                "winners": p.winners,
-                "losers": p.losers,
-            }
-            for p in program_statistics
-        ]
-    )
-    program_collection = OutputCollection(
-        outputs=program_statistics, dataframe=program_df
+    program_collection = build_program_statistics(
+        UK_PROGRAMS,
+        baseline_simulation,
+        reform_simulation,
     )
 
     baseline_poverty = calculate_uk_poverty_rates(baseline_simulation)