Merge pull request #1178 from marcelolafleur/feature/tpi-stall-detection

Stop the TPI outer loop early when it has stalled, with a diagnosis

Author: jdebacker

CHANGELOG.md

@@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Stall detection for the TPI outer loop (Issue #1177): when the best
+  distance has not improved over the last `TPI_stall_window` iterations
+  (default 50; 0 disables), `run_TPI` logs a diagnosis distinguishing a
+  cycling outer loop (suggesting a lower `nu` or
+  `TPI_outer_method="anderson"`) from a diverging economy (usually an
+  inconsistent fiscal block, which solver settings cannot fix). The
+  default `TPI_stall_action="warn"` only logs, leaving model solutions
+  unchanged; `"stop"` also ends the loop early, so a hopeless run fails
+  through the existing non-convergence checks instead of spending the
+  rest of `maxiter`. The window check lives in
+  `ogcore.solvers.diagnose_stall` and works for both the picard and
+  anderson update rules. The diagnosis is re-logged if it changes (e.g.
+  escalates from cycling to diverging), and a run that ends unconverged
+  while stalled carries the diagnosis in the `RuntimeError` message, so
+  it reaches users who only see the traceback.
 ### Bug Fixes
 
 - Fixes Issue [#1186](https://github.com/PSLmodels/OG-Core/issues/1186):
@@ -720,7 +737,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [0.17.0]: https://github.com/PSLmodels/OG-Core/compare/v0.16.4...v0.17.0
 [0.16.4]: https://github.com/PSLmodels/OG-Core/compare/v0.16.3...v0.16.4
 [0.16.3]: https://github.com/PSLmodels/OG-Core/compare/v0.16.2...v0.16.3
-[0.16.3]: https://github.com/PSLmodels/OG-Core/compare/v0.16.2...v0.16.3
 [0.16.2]: https://github.com/PSLmodels/OG-Core/compare/v0.16.1...v0.16.2
 [0.16.1]: https://github.com/PSLmodels/OG-Core/compare/v0.16.0...v0.16.1
 [0.16.0]: https://github.com/PSLmodels/OG-Core/compare/v0.15.13...v0.16.0

ogcore/TPI.py

@@ -931,6 +931,8 @@ def run_TPI(p, client=None):
     TPIdist = 10
     euler_errors = np.zeros((p.T, 2 * p.S, p.J))
     TPIdist_vec = np.zeros(p.maxiter)
+    stall = None
+    stall_reported = None
     # Pluggable outer-loop update rule. Default "picard" -> None -> the native
     # damped functional-iteration path below (unchanged, so golden outputs
     # are preserved); "anderson" accelerates using the residual history. See
@@ -1525,6 +1527,37 @@ def run_TPI(p, client=None):
         TPIiter += 1
         logger.info(f"Iteration: {TPIiter}")
         logger.info(f"Distance: {TPIdist}")
+        # Stall detection: when the best distance has stopped improving
+        # over a window of iterations, further iterations repeat the same
+        # pattern and cannot reach the tolerance -- diagnose the cause and,
+        # if TPI_stall_action="stop", end the loop early (the solution
+        # checks after the loop then fail the run as any non-convergence).
+        # In the default warn mode the diagnosis is logged once per stall,
+        # and again only if it changes (e.g. escalates to diverging).
+        stall = solvers.diagnose_stall(
+            TPIdist_vec, TPIiter, p.TPI_stall_window
+        )
+        if stall != stall_reported:
+            stall_reported = stall
+            if stall == "diverging":
+                logger.error(
+                    "TPI stalled and diverging: the best distance over "
+                    f"the last {p.TPI_stall_window} iterations is far "
+                    "above the earlier best. This usually signals an "
+                    "inconsistent fiscal block (spending, revenue, and "
+                    "debt_ratio_ss), not a solver problem."
+                )
+            elif stall == "oscillating":
+                logger.error(
+                    "TPI stalled: the best distance has not improved "
+                    f"over the last {p.TPI_stall_window} iterations "
+                    f"(current {TPIdist:.2e}, tolerance "
+                    f"{p.mindist_TPI}). The outer loop is cycling; try "
+                    "a lower nu, or TPI_outer_method='anderson' if not "
+                    "already enabled."
+                )
+        if stall is not None and p.TPI_stall_action == "stop":
+            break
 
     # Compute effective and marginal tax rates for all agents
     num_params = len(p.mtrx_params[0][0])
@@ -1788,9 +1821,18 @@ def run_TPI(p, client=None):
     if (
         (TPIiter >= p.maxiter) or (np.absolute(TPIdist) > p.mindist_TPI)
     ) and ENFORCE_SOLUTION_CHECKS:
-        raise RuntimeError(
-            "Transition path equlibrium not found" + " (TPIdist)"
-        )
+        msg = "Transition path equlibrium not found (TPIdist)"
+        if stall == "oscillating":
+            msg += (
+                "; the outer loop stalled cycling -- try a lower nu, or "
+                "TPI_outer_method='anderson'"
+            )
+        elif stall == "diverging":
+            msg += (
+                "; the outer loop stalled diverging -- check the fiscal "
+                "block (spending, revenue, debt_ratio_ss)"
+            )
+        raise RuntimeError(msg)
 
     if (np.any(np.absolute(RC_error) >= p.RC_TPI)) and ENFORCE_SOLUTION_CHECKS:
         raise RuntimeError(

ogcore/default_parameters.json

@@ -4734,6 +4734,46 @@
             }
         }
     },
+    "TPI_stall_window": {
+        "title": "Iteration window for TPI stall detection",
+        "description": "Number of trailing TPI outer-loop iterations over which the best distance must improve on the best from before the window. When it does not, the loop has stalled (cycling or diverging) and a diagnosis is logged; see TPI_stall_action for whether the loop also stops. A value of 0 disables stall detection.",
+        "short_description": "TPI stall-detection window",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "int",
+        "value": [
+            {
+                "value": 50
+            }
+        ],
+        "validators": {
+            "range": {
+                "min": 0,
+                "max": 500
+            }
+        }
+    },
+    "TPI_stall_action": {
+        "title": "Action when TPI stall detection fires",
+        "description": "What to do when stall detection (TPI_stall_window) diagnoses a stalled TPI outer loop. 'warn' (default) logs the diagnosis once and lets the loop continue, leaving model solutions unchanged; 'stop' also ends the loop early, so the run fails through the usual non-convergence checks instead of spending the rest of maxiter.",
+        "short_description": "TPI stall action",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "str",
+        "value": [
+            {
+                "value": "warn"
+            }
+        ],
+        "validators": {
+            "choice": {
+                "choices": [
+                    "warn",
+                    "stop"
+                ]
+            }
+        }
+    },
     "SS_root_method": {
         "title": "Root finding algorithm for outer loop of the SS solution",
         "description": "Root finding algorithm for outer loop of the SS solution.",

ogcore/solvers.py

@@ -191,6 +191,49 @@ def make_outer_updater(method, p):
     raise ValueError(f"unknown TPI_outer_method: {method!r}")
 
 
+def diagnose_stall(dist_vec, iteration, window, tol=0.05, deterioration=2.0):
+    """
+    Check the TPI outer loop's distance history for a stall.
+
+    The loop has stalled when the best distance over the most recent
+    ``window`` iterations is no better (by a relative margin ``tol``)
+    than the best from before that window: the remaining iterations
+    repeat the same pattern and cannot reach the tolerance. A recent
+    best that is far worse than the earlier best (``deterioration``)
+    means the path is drifting away -- a diverging economy, typically
+    an inconsistent fiscal block; a bounce around the earlier best
+    means the outer loop is cycling around the solution, typically
+    ``nu`` too large for the problem.
+
+    Args:
+        dist_vec (Numpy array): per-iteration outer-loop distances,
+            filled for the first ``iteration`` entries
+        iteration (int): number of completed outer-loop iterations
+        window (int): number of trailing iterations that must improve
+            on the earlier best; a non-positive value disables the
+            check
+        tol (float): minimum relative improvement of the recent best
+            over the earlier best that counts as progress
+        deterioration (float): factor by which the recent best must
+            exceed the earlier best to be read as divergence rather
+            than cycling
+
+    Returns:
+        diagnosis (str or None): None while the loop is progressing
+            (or the check is disabled), otherwise "diverging" or
+            "oscillating"
+    """
+    if window <= 0 or iteration < 2 * window:
+        return None
+    recent = dist_vec[iteration - window : iteration]
+    earlier_best = dist_vec[: iteration - window].min()
+    if recent.min() < (1.0 - tol) * earlier_best:
+        return None
+    if recent.min() >= deterioration * earlier_best:
+        return "diverging"
+    return "oscillating"
+
+
 def _selftest():
     """
     Validate the accelerator math on a linear contraction fixed point,

tests/test_TPI.py

@@ -335,6 +335,43 @@ def test_anderson_scaling_and_reset():
     assert u._F == [] and u._X == []
 
 
+def test_stall_defaults_are_warn_only():
+    # Stall detection defaults: a 50-iteration window and warn-only
+    # action, so model solutions are unchanged.
+    p = Specifications()
+    assert p.TPI_stall_window == 50
+    assert p.TPI_stall_action == "warn"
+
+
+def test_diagnose_stall_progressing_and_disabled():
+    # A steadily improving distance history is never a stall, and a
+    # non-positive window (or too-short history) disables the check.
+    improving = 10.0 * 0.9 ** np.arange(100)
+    assert solvers.diagnose_stall(improving, 100, 20, 0.05) is None
+    stuck = np.full(100, 5.0)
+    assert solvers.diagnose_stall(stuck, 100, 0, 0.05) is None
+    assert solvers.diagnose_stall(stuck, 30, 20, 0.05) is None
+
+
+def test_diagnose_stall_oscillating():
+    # A distance that bounces in a band without improving on the earlier
+    # best is diagnosed as the outer loop cycling.
+    rng = np.arange(100)
+    bouncing = 0.3 + 0.25 * (-1.0) ** rng
+    assert solvers.diagnose_stall(bouncing, 100, 20, 0.05) == "oscillating"
+
+
+def test_diagnose_stall_diverging():
+    # A recent best far above the earlier best is diagnosed as a
+    # diverging economy rather than a cycling solver -- whether the
+    # growth is smooth or a widening bounce.
+    growing = 0.1 * 1.05 ** np.arange(100)
+    assert solvers.diagnose_stall(growing, 100, 20, 0.05) == "diverging"
+    rng = np.arange(100)
+    growing_bounce = 0.05 * 1.1**rng * (1.0 + 0.3 * (-1.0) ** rng)
+    assert solvers.diagnose_stall(growing_bounce, 100, 20, 0.05) == "diverging"
+
+
 file_in1 = os.path.join(
     CUR_PATH, "test_io_data", "twist_doughnut_inputs_2.pkl"
 )