Merge pull request #377 from PolicyEngine/codex/cliff-impact

[codex] Add opt-in cliff impact outputs

Author: anth-volk

changelog.d/cliff-impact.added.md

@@ -0,0 +1 @@
+Added opt-in macro cliff impact outputs for US and UK reform analyses.

src/policyengine/outputs/__init__.py

@@ -4,6 +4,12 @@
     ChangeAggregate,
     ChangeAggregateType,
 )
+from policyengine.outputs.cliff_impact import (
+    CliffImpact,
+    CliffImpactInSimulation,
+    calculate_cliff_impact,
+    configure_cliff_impact_variables,
+)
 from policyengine.outputs.congressional_district_impact import (
     CongressionalDistrictImpact,
     compute_us_congressional_district_impacts,
@@ -76,6 +82,10 @@
     "AggregateType",
     "ChangeAggregate",
     "ChangeAggregateType",
+    "CliffImpact",
+    "CliffImpactInSimulation",
+    "calculate_cliff_impact",
+    "configure_cliff_impact_variables",
     "DecileImpact",
     "calculate_decile_impacts",
     "ProgramStatistics",

src/policyengine/outputs/cliff_impact.py

@@ -0,0 +1,89 @@
+"""Legacy-compatible tax-benefit cliff macro output."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from policyengine.core import Output, Simulation
+from policyengine.outputs.aggregate import (
+    get_aggregate_variable,
+    get_output_entity_data,
+    require_output_column,
+)
+from policyengine.outputs.extra_variables import add_extra_variables
+
+CLIFF_IMPACT_VARIABLES = ("cliff_gap", "is_on_cliff", "is_adult")
+
+
+class CliffImpactInSimulation(BaseModel):
+    cliff_gap: float
+    cliff_share: float
+
+
+class CliffImpact(Output):
+    baseline: CliffImpactInSimulation
+    reform: CliffImpactInSimulation
+
+
+def _cliff_variables_by_entity(
+    simulation: Simulation,
+) -> dict[str, list[str]]:
+    variables_by_entity: dict[str, list[str]] = {}
+    for variable_name in CLIFF_IMPACT_VARIABLES:
+        variable = get_aggregate_variable(
+            simulation,
+            variable_name,
+            "CliffImpact.extra_variables",
+        )
+        variables_by_entity.setdefault(variable.entity, []).append(variable_name)
+    return variables_by_entity
+
+
+def configure_cliff_impact_variables(*simulations: Simulation) -> None:
+    """Materialize cliff columns only for analyses that request them."""
+    for simulation in simulations:
+        add_extra_variables(
+            simulation,
+            _cliff_variables_by_entity(simulation),
+        )
+
+
+def _sum_output_variable(
+    simulation: Simulation,
+    variable_name: str,
+) -> float:
+    context = f"CliffImpact.{variable_name}"
+    variable = get_aggregate_variable(simulation, variable_name, context)
+    data = get_output_entity_data(simulation, variable.entity, context)
+    require_output_column(
+        data,
+        variable_name,
+        variable.entity,
+        simulation,
+        context,
+    )
+    return float(data[variable_name].sum())
+
+
+def _calculate_cliff_impact_in_simulation(
+    simulation: Simulation,
+) -> CliffImpactInSimulation:
+    cliff_gap = _sum_output_variable(simulation, "cliff_gap")
+    people_on_cliffs = _sum_output_variable(simulation, "is_on_cliff")
+    adults = _sum_output_variable(simulation, "is_adult")
+
+    return CliffImpactInSimulation(
+        cliff_gap=cliff_gap,
+        cliff_share=float(people_on_cliffs / adults),
+    )
+
+
+def calculate_cliff_impact(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> CliffImpact:
+    """Calculate legacy macro cliff output from materialized simulations."""
+    return CliffImpact(
+        baseline=_calculate_cliff_impact_in_simulation(baseline_simulation),
+        reform=_calculate_cliff_impact_in_simulation(reform_simulation),
+    )

src/policyengine/outputs/extra_variables.py

@@ -0,0 +1,22 @@
+"""Helpers for conditionally materialized output variables."""
+
+from __future__ import annotations
+
+from policyengine.core import Simulation
+
+
+def add_extra_variables(
+    simulation: Simulation,
+    variables_by_entity: dict[str, list[str]],
+) -> None:
+    """Append extra output variables without dropping caller-supplied extras."""
+    extra_variables = {
+        entity: list(variables)
+        for entity, variables in (simulation.extra_variables or {}).items()
+    }
+    for entity, variables in variables_by_entity.items():
+        existing = extra_variables.setdefault(entity, [])
+        for variable in variables:
+            if variable not in existing:
+                existing.append(variable)
+    simulation.extra_variables = extra_variables

src/policyengine/outputs/labor_supply_response.py

@@ -14,6 +14,7 @@
     get_output_entity_data,
     require_output_column,
 )
+from policyengine.outputs.extra_variables import add_extra_variables
 
 CountryCode = Literal["us", "uk"]
 DecileValues = dict[int, float]
@@ -162,22 +163,6 @@ def _active_lsr_variables(country_code: CountryCode) -> dict[str, list[str]]:
     )
 
 
-def _add_extra_variables(
-    simulation: Simulation,
-    variables_by_entity: dict[str, list[str]],
-) -> None:
-    extra_variables = {
-        entity: list(variables)
-        for entity, variables in (simulation.extra_variables or {}).items()
-    }
-    for entity, variables in variables_by_entity.items():
-        existing = extra_variables.setdefault(entity, [])
-        for variable in variables:
-            if variable not in existing:
-                existing.append(variable)
-    simulation.extra_variables = extra_variables
-
-
 def configure_labor_supply_response_variables(
     baseline_simulation: Simulation,
     reform_simulation: Simulation,
@@ -193,8 +178,8 @@ def configure_labor_supply_response_variables(
         return False
 
     active_variables = _active_lsr_variables(country_code)
-    _add_extra_variables(baseline_simulation, active_variables)
-    _add_extra_variables(reform_simulation, active_variables)
+    add_extra_variables(baseline_simulation, active_variables)
+    add_extra_variables(reform_simulation, active_variables)
     return True
 
 

src/policyengine/tax_benefit_models/uk/analysis.py

@@ -11,9 +11,12 @@
 
 from policyengine.core import OutputCollection, Simulation
 from policyengine.outputs import (
+    CliffImpact,
     LaborSupplyResponse,
     ProgramStatistics,
+    calculate_cliff_impact,
     calculate_labor_supply_response,
+    configure_cliff_impact_variables,
     configure_labor_supply_response_variables,
 )
 from policyengine.outputs.decile_impact import (
@@ -71,6 +74,7 @@ class PolicyReformAnalysis(BaseModel):
     baseline_inequality: Inequality
     reform_inequality: Inequality
     labor_supply_response: LaborSupplyResponse
+    cliff_impact: CliffImpact | None = None
 
 
 def _format_missing_program_variables(missing_variables: set[str]) -> str | None:
@@ -141,13 +145,16 @@ def _validate_program_statistics_config(
 def economic_impact_analysis(
     baseline_simulation: Simulation,
     reform_simulation: Simulation,
+    include_cliff_impacts: bool = False,
 ) -> PolicyReformAnalysis:
     """Perform comprehensive analysis of a UK policy reform."""
     configure_labor_supply_response_variables(
         baseline_simulation,
         reform_simulation,
         country_code="uk",
     )
+    if include_cliff_impacts:
+        configure_cliff_impact_variables(baseline_simulation, reform_simulation)
     _validate_program_statistics_config(baseline_simulation, reform_simulation)
 
     baseline_simulation.ensure()
@@ -224,6 +231,11 @@ def economic_impact_analysis(
         reform_simulation,
         country_code="uk",
     )
+    cliff_impact = (
+        calculate_cliff_impact(baseline_simulation, reform_simulation)
+        if include_cliff_impacts
+        else None
+    )
 
     return PolicyReformAnalysis(
         decile_impacts=decile_impacts,
@@ -235,4 +247,5 @@ def economic_impact_analysis(
         baseline_inequality=baseline_inequality,
         reform_inequality=reform_inequality,
         labor_supply_response=labor_supply_response,
+        cliff_impact=cliff_impact,
     )

src/policyengine/tax_benefit_models/us/analysis.py

@@ -13,9 +13,12 @@
 
 from policyengine.core import OutputCollection, Simulation
 from policyengine.outputs import (
+    CliffImpact,
     LaborSupplyResponse,
     ProgramStatistics,
+    calculate_cliff_impact,
     calculate_labor_supply_response,
+    configure_cliff_impact_variables,
     configure_labor_supply_response_variables,
 )
 from policyengine.outputs.decile_impact import (
@@ -63,6 +66,7 @@ class PolicyReformAnalysis(BaseModel):
     baseline_inequality: Inequality
     reform_inequality: Inequality
     labor_supply_response: LaborSupplyResponse
+    cliff_impact: CliffImpact | None = None
 
 
 def _format_missing_program_variables(missing_variables: set[str]) -> str | None:
@@ -134,6 +138,7 @@ def economic_impact_analysis(
     baseline_simulation: Simulation,
     reform_simulation: Simulation,
     inequality_preset: Union[USInequalityPreset, str] = USInequalityPreset.STANDARD,
+    include_cliff_impacts: bool = False,
 ) -> PolicyReformAnalysis:
     """Perform comprehensive analysis of a US policy reform.
 
@@ -151,6 +156,8 @@ def economic_impact_analysis(
         reform_simulation,
         country_code="us",
     )
+    if include_cliff_impacts:
+        configure_cliff_impact_variables(baseline_simulation, reform_simulation)
     _validate_program_statistics_config(baseline_simulation, reform_simulation)
 
     baseline_simulation.ensure()
@@ -218,6 +225,11 @@ def economic_impact_analysis(
         reform_simulation,
         country_code="us",
     )
+    cliff_impact = (
+        calculate_cliff_impact(baseline_simulation, reform_simulation)
+        if include_cliff_impacts
+        else None
+    )
 
     return PolicyReformAnalysis(
         decile_impacts=decile_impacts,
@@ -227,4 +239,5 @@ def economic_impact_analysis(
         baseline_inequality=baseline_inequality,
         reform_inequality=reform_inequality,
         labor_supply_response=labor_supply_response,
+        cliff_impact=cliff_impact,
     )