diff --git a/changelog.d/263.changed.md b/changelog.d/263.changed.md new file mode 100644 index 00000000..e08f81df --- /dev/null +++ b/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`. diff --git a/src/policyengine/outputs/__init__.py b/src/policyengine/outputs/__init__.py index b6ebd37b..ca85f2e9 100644 --- a/src/policyengine/outputs/__init__.py +++ b/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", diff --git a/src/policyengine/outputs/program_statistics.py b/src/policyengine/outputs/program_statistics.py index ccb4f1e1..b830b5a1 100644 --- a/src/policyengine/outputs/program_statistics.py +++ b/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) diff --git a/src/policyengine/tax_benefit_models/uk/analysis.py b/src/policyengine/tax_benefit_models/uk/analysis.py index 0d962026..266736fc 100644 --- a/src/policyengine/tax_benefit_models/uk/analysis.py +++ b/src/policyengine/tax_benefit_models/uk/analysis.py @@ -6,7 +6,6 @@ from __future__ import annotations -import pandas as pd from pydantic import BaseModel from policyengine.core import OutputCollection, Simulation @@ -14,10 +13,12 @@ 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) diff --git a/src/policyengine/tax_benefit_models/us/analysis.py b/src/policyengine/tax_benefit_models/us/analysis.py index 1d4bb5d1..bd271214 100644 --- a/src/policyengine/tax_benefit_models/us/analysis.py +++ b/src/policyengine/tax_benefit_models/us/analysis.py @@ -8,7 +8,6 @@ from typing import Union -import pandas as pd from pydantic import BaseModel from policyengine.core import OutputCollection, Simulation @@ -16,10 +15,12 @@ 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, @@ -34,7 +35,6 @@ Poverty, calculate_us_poverty_rates, ) -from policyengine.utils.errors import format_conditional_error_detail # Map of US program-statistics variable name -> program metadata. The # entity for each program is derived from the variable's own metadata @@ -69,68 +69,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 _program_statistics_config_error_message( - missing_variables: set[str], - missing_outputs: set[tuple[str, str]], -) -> str: - lines = ["US 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 US 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 US_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( - missing_variables, - missing_outputs, - ), + validate_program_statistics_config( + US_PROGRAMS, + baseline_simulation, + reform_simulation, + "US", ) @@ -176,40 +124,10 @@ def economic_impact_analysis( income_variable="household_net_income", ) - model_version = baseline_simulation.tax_benefit_model_version - program_statistics = [] - for program_name, program_info in US_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( + US_PROGRAMS, + baseline_simulation, + reform_simulation, ) baseline_poverty = calculate_us_poverty_rates(baseline_simulation) diff --git a/tests/test_cliff_impact_analysis.py b/tests/test_cliff_impact_analysis.py index ab2a0b3d..bd33ba51 100644 --- a/tests/test_cliff_impact_analysis.py +++ b/tests/test_cliff_impact_analysis.py @@ -10,6 +10,7 @@ LaborSupplyResponse, ProgramStatistics, ) +from policyengine.outputs import program_statistics as program_statistics_module from policyengine.outputs.inequality import Inequality from policyengine.tax_benefit_models.uk import analysis as uk_analysis from policyengine.tax_benefit_models.us import analysis as us_analysis @@ -70,7 +71,9 @@ def fake_program_statistics(**kwargs): "_validate_program_statistics_config", lambda baseline_simulation, reform_simulation: None, ) - monkeypatch.setattr(analysis_module, "ProgramStatistics", fake_program_statistics) + monkeypatch.setattr( + program_statistics_module, "ProgramStatistics", fake_program_statistics + ) monkeypatch.setattr( analysis_module, "configure_labor_supply_response_variables", diff --git a/tests/test_program_statistics.py b/tests/test_program_statistics.py new file mode 100644 index 00000000..bdf415ae --- /dev/null +++ b/tests/test_program_statistics.py @@ -0,0 +1,103 @@ +"""Unit tests for the shared program-statistics helpers. + +These exercise :func:`validate_program_statistics_config` directly with +lightweight fakes so the country-agnostic validation logic (and, in particular, +the ``country_label`` threaded into the error message) is covered without +building a full country simulation. +""" + +import pytest + +from policyengine.outputs.program_statistics import ( + validate_program_statistics_config, +) + + +class _FakeVariable: + def __init__(self, entity: str): + self.entity = entity + + +class _FakeModelVersion: + def __init__( + self, + variable_entities: dict[str, str], + resolved_entity_variables: dict[str, list[str]], + ): + self._variable_entities = variable_entities + self._resolved_entity_variables = resolved_entity_variables + + def get_variable(self, name: str) -> _FakeVariable: + if name not in self._variable_entities: + raise ValueError(name) + return _FakeVariable(self._variable_entities[name]) + + def resolve_entity_variables(self, simulation) -> dict[str, list[str]]: + return self._resolved_entity_variables + + +class _FakeSimulation: + def __init__(self, model_version: _FakeModelVersion): + self.tax_benefit_model_version = model_version + + +def _simulation( + variable_entities: dict[str, str], + resolved_entity_variables: dict[str, list[str]], +) -> _FakeSimulation: + return _FakeSimulation( + _FakeModelVersion(variable_entities, resolved_entity_variables) + ) + + +def test_validate_config_passes_when_all_present(): + programs = {"income_tax": {"is_tax": True}} + simulation = _simulation( + {"income_tax": "tax_unit"}, + {"tax_unit": ["income_tax"]}, + ) + + # Should not raise for either simulation. + validate_program_statistics_config( + programs, + simulation, + simulation, + "Ruritania", + ) + + +def test_validate_config_missing_variable_reports_label(): + programs = {"income_tax": {"is_tax": True}} + simulation = _simulation({}, {}) + + with pytest.raises(ValueError) as exc_info: + validate_program_statistics_config( + programs, + simulation, + simulation, + "Ruritania", + ) + + message = str(exc_info.value) + assert message.startswith("Ruritania program statistics config is invalid") + assert "income_tax" in message + + +def test_validate_config_unmaterialized_variable_reports_label(): + programs = {"income_tax": {"is_tax": True}} + simulation = _simulation( + {"income_tax": "tax_unit"}, + {"tax_unit": []}, + ) + + with pytest.raises(ValueError) as exc_info: + validate_program_statistics_config( + programs, + simulation, + simulation, + "Ruritania", + ) + + message = str(exc_info.value) + assert message.startswith("Ruritania program statistics config is invalid") + assert "income_tax on tax_unit" in message diff --git a/tests/test_uk_analysis.py b/tests/test_uk_analysis.py index ca943c71..85c22c80 100644 --- a/tests/test_uk_analysis.py +++ b/tests/test_uk_analysis.py @@ -4,6 +4,7 @@ from policyengine.core import OutputCollection from policyengine.outputs import LaborSupplyResponse, ProgramStatistics +from policyengine.outputs import program_statistics as program_statistics_module from policyengine.outputs.inequality import Inequality from policyengine.tax_benefit_models.uk import analysis as uk_analysis @@ -90,7 +91,9 @@ def fake_inequality(simulation): "_validate_program_statistics_config", lambda baseline_simulation, reform_simulation: None, ) - monkeypatch.setattr(uk_analysis, "ProgramStatistics", fake_program_statistics) + monkeypatch.setattr( + program_statistics_module, "ProgramStatistics", fake_program_statistics + ) monkeypatch.setattr(uk_analysis, "calculate_uk_poverty_rates", fake_poverty_rates) monkeypatch.setattr(uk_analysis, "calculate_uk_inequality", fake_inequality) monkeypatch.setattr(