feat: enable sim API cliff impacts

Author: anth-volk

projects/policyengine-api-simulation/pyproject.toml

@@ -16,7 +16,7 @@ dependencies = [
     "pydantic-settings (>=2.7.1,<3.0.0)",
     "opentelemetry-instrumentation-fastapi (>=0.51b0,<0.52)",
     "policyengine-fastapi",
-    "policyengine==4.10.0",
+    "policyengine==4.12.0",
     "policyengine-core==3.26.1",
     "policyengine-uk==2.88.20",
     "policyengine-us==1.700.0",

projects/policyengine-api-simulation/src/modal/gateway/models.py

@@ -184,6 +184,11 @@ def validate_end_year(self) -> "BudgetWindowBatchRequest":
             raise ValueError(
                 f"budget-window end_year must be {self.MAX_END_YEAR} or earlier"
             )
+        if self.include_cliffs is True:
+            raise ValueError(
+                "budget-window cliff impacts are not supported; use the single-year "
+                "simulation endpoint with include_cliffs=true"
+            )
         return self
 
 

projects/policyengine-api-simulation/src/policyengine_api_simulation/simulation_macro_output.py

@@ -110,6 +110,16 @@ class LaborSupplyResponseOutput(MacroRootModel[dict[str, Any]]):
     pass
 
 
+class CliffImpactInSimulation(MacroOutputModel):
+    cliff_gap: float
+    cliff_share: float
+
+
+class CliffImpactOutput(MacroOutputModel):
+    baseline: CliffImpactInSimulation
+    reform: CliffImpactInSimulation
+
+
 class GeographicImpactOutput(MacroRootModel[list[dict[str, Any]]]):
     pass
 
@@ -131,4 +141,4 @@ class SingleYearMacroOutput(MacroOutputModel):
     constituency_impact: GeographicImpactOutput | None
     local_authority_impact: GeographicImpactOutput | None
     congressional_district_impact: GeographicImpactOutput | None
-    cliff_impact: None = None
+    cliff_impact: CliffImpactOutput | None = None

projects/policyengine-api-simulation/src/policyengine_api_simulation/simulation_output_builder.py

@@ -14,6 +14,8 @@
     AgePovertyOutput,
     BaselineReformValue,
     BudgetaryImpact,
+    CliffImpactInSimulation,
+    CliffImpactOutput,
     DecileOutput,
     DetailedBudgetOutput,
     DetailedBudgetProgramOutput,
@@ -252,10 +254,15 @@ def __post_init__(self) -> None:
     def analysis(self) -> Any:
         if self._analysis is None:
             self._analysis = self.country_module.economic_impact_analysis(
-                self.baseline, self.reform
+                self.baseline,
+                self.reform,
+                include_cliff_impacts=self._include_cliff_impacts(),
             )
         return self._analysis
 
+    def _include_cliff_impacts(self) -> bool:
+        return self.simulation_params.get("include_cliffs") is True
+
     def build(self) -> SingleYearMacroOutput:
         poverty_outputs = self._build_poverty_outputs()
         wealth_decile = getattr(self.analysis, "wealth_decile_impacts", None)
@@ -280,7 +287,7 @@ def build(self) -> SingleYearMacroOutput:
             congressional_district_impact=(self._build_congressional_district_impact()),
             constituency_impact=self._build_uk_constituency_impact(),
             local_authority_impact=self._build_uk_local_authority_impact(),
-            cliff_impact=None,
+            cliff_impact=self._build_cliff_impact(),
         )
 
     def serialize(self) -> dict[str, Any]:
@@ -448,6 +455,18 @@ def _build_labor_supply_response(self) -> LaborSupplyResponseOutput | None:
         output = _output_model_dump(labor_supply_response)
         return LaborSupplyResponseOutput(output) if isinstance(output, dict) else None
 
+    def _build_cliff_impact(self) -> CliffImpactOutput | None:
+        cliff_impact = getattr(self.analysis, "cliff_impact", None)
+        if isinstance(cliff_impact, CliffImpactOutput):
+            return cliff_impact
+        output = _output_model_dump(cliff_impact)
+        if not isinstance(output, Mapping):
+            return None
+        return CliffImpactOutput(
+            baseline=CliffImpactInSimulation(**output["baseline"]),
+            reform=CliffImpactInSimulation(**output["reform"]),
+        )
+
     def _build_geographic_impact_output(
         self, value: Any
     ) -> GeographicImpactOutput | None:

projects/policyengine-api-simulation/tests/gateway/test_endpoints.py

@@ -229,6 +229,27 @@ def test__given_submission__then_returns_job_id_and_poll_url(
         assert data["poll_url"] == "/jobs/mock-job-id-123"
         assert data["status"] == "submitted"
 
+    def test__given_submission_with_include_cliffs__then_forwards_worker_flag(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-py4-10-0",
+        }
+
+        response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+                "include_cliffs": True,
+            },
+        )
+
+        assert response.status_code == 200
+        assert mock_modal["func"].last_payload["include_cliffs"] is True
+
     def test__given_submission_with_telemetry__then_preserves_run_id(
         self, mock_modal, client: TestClient
     ):
@@ -714,6 +735,31 @@ def test__given_budget_window_submission__then_returns_parent_batch_job_id(
             "policyengine_bundle": expected_bundle("us", "1.500.0"),
         }
 
+    def test__given_budget_window_include_cliffs__then_returns_422(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-py4-10-0",
+        }
+
+        response = client.post(
+            "/simulate/economy/budget-window",
+            json={
+                "country": "us",
+                "region": "us",
+                "scope": "macro",
+                "reform": {},
+                "start_year": "2026",
+                "window_size": 3,
+                "include_cliffs": True,
+            },
+        )
+
+        assert response.status_code == 422
+        assert "cliff impacts are not supported" in response.text
+        assert mock_modal["func"].last_payload is None
+
     def test__given_budget_window_submission__then_initial_poll_returns_seed_state(
         self, mock_modal, client: TestClient
     ):

projects/policyengine-api-simulation/tests/gateway/test_models.py

@@ -429,6 +429,16 @@ def test_budget_window_batch_request_rejects_non_general_target(self):
                 target="cliff",
             )
 
+    def test_budget_window_batch_request_rejects_include_cliffs(self):
+        with pytest.raises(ValidationError, match="cliff impacts are not supported"):
+            BudgetWindowBatchRequest(
+                country="us",
+                region="us",
+                start_year="2026",
+                window_size=3,
+                include_cliffs=True,
+            )
+
     def test_budget_window_batch_request_rejects_max_parallel_above_active_limit(self):
         with pytest.raises(ValidationError):
             BudgetWindowBatchRequest(

projects/policyengine-api-simulation/tests/test_simulation_output_builder.py

@@ -19,6 +19,7 @@
     REFORM_POVERTY_BY_AGE,
     REFORM_POVERTY_BY_GENDER,
     REFORM_POVERTY_BY_RACE,
+    FakeModelOutput,
     fake_analysis,
 )
 from policyengine_api_simulation.simulation_runtime import RegionResolution
@@ -85,20 +86,32 @@ def _simulation_output_builder(
     baseline,
     reform,
     analysis=None,
+    include_cliffs: bool | None = None,
 ) -> SimulationOutputBuilder:
     analysis = analysis or fake_analysis()
+
+    def economic_impact_analysis(
+        baseline_simulation,
+        reform_simulation,
+        *,
+        include_cliff_impacts=False,
+    ):
+        return analysis
+
+    simulation_params = {
+        "country": country,
+        "data_version": "1.115.5" if country == "us" else "1.55.10",
+    }
+    if include_cliffs is not None:
+        simulation_params["include_cliffs"] = include_cliffs
+
     country_module = SimpleNamespace(
         model=SimpleNamespace(version="1.700.0" if country == "us" else "2.88.20"),
-        economic_impact_analysis=lambda baseline_simulation, reform_simulation: (
-            analysis
-        ),
+        economic_impact_analysis=economic_impact_analysis,
     )
     return SimulationOutputBuilder(
         country=country,
-        simulation_params={
-            "country": country,
-            "data_version": "1.115.5" if country == "us" else "1.55.10",
-        },
+        simulation_params=simulation_params,
         country_module=country_module,
         dataset=SimpleNamespace(metadata={}),
         baseline=baseline,
@@ -218,11 +231,19 @@ def test_builder_calls_policyengine_economic_impact_analysis():
     baseline, reform = _macro_baseline_reform()
     analysis = fake_analysis()
     calls = []
+
+    def economic_impact_analysis(
+        baseline_simulation,
+        reform_simulation,
+        *,
+        include_cliff_impacts=False,
+    ):
+        calls.append((baseline_simulation, reform_simulation, include_cliff_impacts))
+        return analysis
+
     country_module = SimpleNamespace(
         model=SimpleNamespace(version="1.700.0"),
-        economic_impact_analysis=lambda baseline_simulation, reform_simulation: (
-            calls.append((baseline_simulation, reform_simulation)) or analysis
-        ),
+        economic_impact_analysis=economic_impact_analysis,
     )
     builder = SimulationOutputBuilder(
         country="us",
@@ -235,7 +256,68 @@ def test_builder_calls_policyengine_economic_impact_analysis():
 
     assert builder.analysis is analysis
     assert builder.analysis is analysis
-    assert calls == [(baseline, reform)]
+    assert calls == [(baseline, reform, False)]
+
+
+def test_builder_passes_include_cliffs_to_policyengine_economic_impact_analysis():
+    baseline, reform = _macro_baseline_reform()
+    analysis = fake_analysis()
+    calls = []
+
+    def economic_impact_analysis(
+        baseline_simulation,
+        reform_simulation,
+        *,
+        include_cliff_impacts=False,
+    ):
+        calls.append(include_cliff_impacts)
+        return analysis
+
+    country_module = SimpleNamespace(
+        model=SimpleNamespace(version="1.700.0"),
+        economic_impact_analysis=economic_impact_analysis,
+    )
+    builder = SimulationOutputBuilder(
+        country="us",
+        simulation_params={
+            "country": "us",
+            "data_version": "1.115.5",
+            "include_cliffs": True,
+        },
+        country_module=country_module,
+        dataset=SimpleNamespace(metadata={}),
+        baseline=baseline,
+        reform=reform,
+    )
+
+    assert builder.analysis is analysis
+    assert calls == [True]
+
+
+def test_builder_serializes_cliff_impact_when_requested(monkeypatch):
+    baseline, reform = _macro_baseline_reform()
+    _stub_policyengine_output_calls(monkeypatch, baseline, reform)
+    analysis = fake_analysis()
+    analysis.cliff_impact = FakeModelOutput(
+        {
+            "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25},
+            "reform": {"cliff_gap": 20.0, "cliff_share": 0.5},
+        }
+    )
+
+    output = _simulation_output_builder(
+        "us",
+        baseline,
+        reform,
+        analysis=analysis,
+        include_cliffs=True,
+    ).build()
+
+    assert output.cliff_impact is not None
+    assert output.model_dump(mode="json")["cliff_impact"] == {
+        "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25},
+        "reform": {"cliff_gap": 20.0, "cliff_share": 0.5},
+    }
 
 
 def test_normalise_policy_converts_legacy_period_range_keys():
@@ -438,11 +520,18 @@ def test_resolve_region_scopes_uk_country_from_national_dataset():
 
 def test_builder_data_version_prefers_resolved_revision_then_dataset_metadata():
     baseline, reform = _macro_baseline_reform()
+
+    def economic_impact_analysis(
+        baseline_simulation,
+        reform_simulation,
+        *,
+        include_cliff_impacts=False,
+    ):
+        return fake_analysis()
+
     country_module = SimpleNamespace(
         model=SimpleNamespace(version="1.700.0"),
-        economic_impact_analysis=lambda baseline_simulation, reform_simulation: (
-            fake_analysis()
-        ),
+        economic_impact_analysis=economic_impact_analysis,
     )
 
     resolved_builder = SimulationOutputBuilder(

projects/policyengine-api-simulation/tests/test_standalone_simulation_contract.py

@@ -62,3 +62,25 @@ def fake_run_simulation_impl(params):
 
     assert response.status_code == 200
     assert response.json() == CURRENT_SINGLE_YEAR_MACRO_RESULT
+
+
+def test_standalone_simulation_route_forwards_include_cliffs(monkeypatch):
+    def fake_run_simulation_impl(params):
+        assert params == {
+            "country": "us",
+            "reform": {},
+            "include_cliffs": True,
+        }
+        return CURRENT_SINGLE_YEAR_MACRO_RESULT
+
+    monkeypatch.setattr(
+        "policyengine_api_simulation.simulation.run_simulation_impl",
+        fake_run_simulation_impl,
+    )
+
+    response = TestClient(app).post(
+        "/simulate/economy/comparison",
+        json={"country": "us", "reform": {}, "include_cliffs": True},
+    )
+
+    assert response.status_code == 200