Restore the resource-constraint check on interior periods

oguk/api.py

@@ -982,6 +982,11 @@ def _build_specs(
     # boundary-condition discontinuity in fiscal.py that causes a large
     # RC error at that single period. All other periods are well within
     # 1e-4. Setting RC_TPI=0.2 allows TPI to complete.
+    #
+    # OG-Core compares RC_TPI against every period with np.any, so this
+    # also waives interior violations up to 0.2 (PSLmodels/OG-Core#1210).
+    # _check_interior_resource_constraint restores the tight check on
+    # t < T-1 after the solve.
     p.RC_TPI = 0.2
     return p
 
@@ -1008,6 +1013,107 @@ def tpi_outer_method(multi_sector: bool) -> str:
     return "picard" if multi_sector else "anderson"
 
 
+# Tolerances for the resource-constraint error, by region of the path.
+#
+# Both ends of a transition path carry artifacts that are not solution
+# failures. At t = T-1 the path is truncated and does not reach a true
+# steady state (PSLmodels/OG-Core#1216: I_d[T-1] is formed from a
+# steady-state-filled K_d[T] against the actual b_sp1[T-1], so the whole
+# gap lands there). At t = 0 the initial conditions are imposed rather
+# than solved. Measured over four production runs (S=80, J=7, T=60, both
+# solvers, baseline and a CIT reform):
+#
+#     t = 0       6.7e-03      initial-condition artifact
+#     t = 1       ~1e-07
+#     t = 2       6.3e-04      largest genuine interior value
+#     t >= 3      <= 3e-06
+#     t = T-1     1.58e-01     truncation artifact
+#
+# The terminal period is exempt. t = 0 is checked against a looser
+# INITIAL_RC_TOL rather than skipped, because it is the impact year and a
+# genuine period-0 calibration inconsistency should not be invisible.
+# INTERIOR_RC_TOL = 5e-3 gives ~8x headroom over the largest measured
+# genuine interior value while still catching both the terminal magnitude
+# (0.16) and the 0.109 magnitude that motivated RC_TPI = 0.2.
+INTERIOR_RC_TOL = 5e-3
+INITIAL_RC_TOL = 1e-2
+
+
+def _check_interior_resource_constraint(
+    tpi_vars: dict,
+    tol: float = INTERIOR_RC_TOL,
+    label: str = "",
+    initial_tol: float = INITIAL_RC_TOL,
+) -> None:
+    """Raise if the resource constraint is violated away from the terminal period.
+
+    ``_build_specs`` sets ``RC_TPI = 0.2`` so that TPI can complete despite a
+    known boundary discontinuity at ``t = T-1``. Because OG-Core applies that
+    tolerance to every period at once (``np.any``), a genuine interior
+    violation of the same magnitude would pass silently. This restores the
+    tight check on the interior of the path.
+
+    Args:
+        tpi_vars (dict): the TPI output dict (as saved to TPI_vars.pkl),
+            containing "resource_constraint_error".
+        tol (float): maximum absolute error permitted at interior periods.
+        label (str): "baseline" or "reform", used in the error message.
+        initial_tol (float): maximum absolute error permitted at t = 0, whose
+            initial conditions are imposed rather than solved.
+
+    Raises:
+        RuntimeError: if any interior period exceeds ``tol``.
+    """
+    rc = tpi_vars.get("resource_constraint_error")
+    if rc is None:
+        return
+    rc = np.absolute(np.asarray(rc, dtype=float))
+    n_periods = rc.shape[0]
+    prefix = f"{label} " if label else ""
+    # Collapse any trailing axes (RC_error is (T, M)) so the result is
+    # indexed by period.
+    by_period = rc.reshape(n_periods, -1).max(axis=1)
+
+    # Non-finite errors must never pass: `nan >= tol` is False, so a
+    # diverged path would otherwise slip through the comparisons below --
+    # and that is exactly the case this check exists to catch.
+    if not np.isfinite(by_period).all():
+        bad = int(np.argmax(~np.isfinite(by_period)))
+        raise RuntimeError(
+            f"{prefix}transition path has a non-finite resource-constraint "
+            f"error at period {bad} of {n_periods}. The solve did not "
+            "produce a usable path."
+        )
+
+    # t = 0 has its initial conditions imposed rather than solved, so it
+    # carries a larger artifact than the interior -- but it is the impact
+    # year, so it is checked against a looser tolerance rather than skipped.
+    # Only when t = 0 is not itself the (exempt) terminal period.
+    if n_periods >= 2 and by_period[0] >= initial_tol:
+        raise RuntimeError(
+            f"{prefix}transition path violates the resource constraint in "
+            f"the initial period: |RC error| = {by_period[0]:.3e} "
+            f"(tolerance {initial_tol:.0e})."
+        )
+
+    # The terminal period is excluded: the path is truncated at T and does
+    # not reach a true steady state (PSLmodels/OG-Core#1216).
+    if n_periods < 3:
+        return
+    interior = by_period[1:-1]
+    worst = int(np.argmax(interior))
+    if interior[worst] >= tol:
+        raise RuntimeError(
+            f"{prefix}transition path violates the resource constraint on "
+            f"the interior of the path: max |RC error| = "
+            f"{interior[worst]:.3e} at period {worst + 1} of {n_periods} "
+            f"(tolerance {tol:.0e}). The terminal period is excluded as a "
+            "truncation artifact and t = 0 is checked separately; an "
+            "interior violation points to an inconsistent calibration "
+            "(spending, revenue, debt_ratio_ss)."
+        )
+
+
 def _ss_dict_to_result(ss: dict) -> SteadyStateResult:
     """Convert OG-Core SS output dict to SteadyStateResult."""
     return SteadyStateResult(
@@ -1156,6 +1262,7 @@ def run_transition_path(
 
         with open(os.path.join(base_dir, "TPI", "TPI_vars.pkl"), "rb") as f:
             tpi_base = pickle.load(f)
+        _check_interior_resource_constraint(tpi_base, label="baseline")
         baseline_tp = _tpi_dict_to_result(tpi_base, start_year)
 
         # Reform
@@ -1187,6 +1294,7 @@ def run_transition_path(
 
             with open(os.path.join(reform_dir, "TPI", "TPI_vars.pkl"), "rb") as f:
                 tpi_reform = pickle.load(f)
+            _check_interior_resource_constraint(tpi_reform, label="reform")
             reform_tp = _tpi_dict_to_result(tpi_reform, start_year)
 
         return baseline_tp, reform_tp

oguk/tests/test_resource_constraint.py

@@ -0,0 +1,142 @@
+"""Tests for the interior resource-constraint check.
+
+``_build_specs`` sets ``RC_TPI = 0.2`` so TPI can complete despite a known
+boundary discontinuity at the terminal period. OG-Core applies that tolerance
+to every period at once, so these tests pin down that an interior violation
+of the same magnitude is still caught.
+"""
+
+import numpy as np
+import pytest
+
+from oguk.api import (
+    INITIAL_RC_TOL,
+    INTERIOR_RC_TOL,
+    _check_interior_resource_constraint,
+)
+
+
+def _measured_profile():
+    """The RC error profile actually observed on a production run.
+
+    Measured over four runs at S=80, J=7, T=60 (baseline and a CIT reform,
+    each under both TPI_outer_method settings). See the comment on
+    INTERIOR_RC_TOL in oguk/api.py.
+    """
+    rc = np.full(60, 3e-06)
+    rc[0] = 6.711e-03  # initial-condition artifact
+    rc[1] = 3.04e-07
+    rc[2] = 6.314e-04  # largest genuine interior value
+    rc[-1] = 1.580e-01  # truncation artifact
+    return rc
+
+
+def test_real_production_profile_passes():
+    """The check must not fire on a run that is actually fine.
+
+    This is the regression test for the first version of this helper, which
+    used a 1e-4 tolerance over rc[:-1] and would have raised on every real
+    transition path because rc[0] is 6.7e-03.
+    """
+    _check_interior_resource_constraint(_rc(_measured_profile()))
+
+
+def test_interior_violation_on_top_of_real_profile_raises():
+    rc = _measured_profile()
+    rc[25] = 0.15
+    with pytest.raises(RuntimeError, match="period 25"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_period_one_is_interior_and_not_exempt():
+    rc = _measured_profile()
+    rc[1] = 0.05
+    with pytest.raises(RuntimeError, match="period 1"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_initial_period_is_checked_not_skipped():
+    """t=0 gets a looser tolerance, but a real violation there still raises."""
+    rc = _measured_profile()
+    rc[0] = 0.05
+    with pytest.raises(RuntimeError, match="initial period"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_measured_initial_value_is_within_its_tolerance():
+    assert _measured_profile()[0] < INITIAL_RC_TOL
+
+
+@pytest.mark.parametrize("bad", [np.nan, -np.nan, np.inf, -np.inf])
+def test_non_finite_error_always_raises(bad):
+    """`nan >= tol` is False, so a diverged path must be caught explicitly."""
+    rc = _measured_profile()
+    rc[10] = bad
+    with pytest.raises(RuntimeError, match="non-finite"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_non_finite_in_exempt_terminal_period_still_raises():
+    rc = _measured_profile()
+    rc[-1] = np.nan
+    with pytest.raises(RuntimeError, match="non-finite"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def _rc(values):
+    return {"resource_constraint_error": np.array(values, dtype=float)}
+
+
+def test_terminal_only_violation_passes():
+    """A large error confined to the last period is a truncation artifact."""
+    rc = np.full(60, 1e-8)
+    rc[-1] = 0.109  # the magnitude that motivated RC_TPI = 0.2
+    _check_interior_resource_constraint(_rc(rc))
+
+
+def test_interior_violation_raises():
+    """An interior error of the same magnitude must not pass silently."""
+    rc = np.full(60, 1e-8)
+    rc[10] = 0.109
+    rc[-1] = 0.109
+    with pytest.raises(RuntimeError, match="period 10"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_sign_is_ignored():
+    """The check is on absolute error."""
+    rc = np.full(60, 1e-8)
+    rc[3] = -0.5
+    with pytest.raises(RuntimeError, match="period 3"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_interior_within_tolerance_passes():
+    rc = np.full(60, INTERIOR_RC_TOL / 10)
+    rc[-1] = 0.2
+    _check_interior_resource_constraint(_rc(rc))
+
+
+def test_trailing_axes_are_collapsed():
+    """resource_constraint_error may carry trailing axes; period is axis 0."""
+    rc = np.full((60, 3), 1e-8)
+    rc[7, 2] = 0.4
+    with pytest.raises(RuntimeError, match="period 7"):
+        _check_interior_resource_constraint(_rc(rc))
+
+
+def test_label_appears_in_message():
+    rc = np.full(10, 1e-8)
+    rc[1] = 1.0
+    with pytest.raises(RuntimeError, match="reform transition path"):
+        _check_interior_resource_constraint(_rc(rc), label="reform")
+
+
+@pytest.mark.parametrize("missing", [{}, {"resource_constraint_error": None}])
+def test_missing_key_is_a_no_op(missing):
+    _check_interior_resource_constraint(missing)
+
+
+def test_single_period_is_a_no_op():
+    """With only a terminal period there is no interior to check."""
+    _check_interior_resource_constraint(_rc([0.5]))