Merge remote-tracking branch 'upstream/master' into demog_J

jdebacker · jdebacker · commit a0c9ade221e9 · 2026-07-15T08:54:08.000-04:00
diff --git a/.github/workflows/check_catalog.yml b/.github/workflows/check_catalog.yml
@@ -18,7 +18,7 @@ jobs:
   catalog-in-sync:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
       - name: Emit catalog from each installer
         run: |
           bash scripts/install.sh --list-json > /tmp/catalog_sh.json
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.16.4] - 2026-07-02 12:00:00
+
+### Added
+
+- Adds an opt-in accelerated update rule for the TPI outer loop, selected by the new `TPI_outer_method` parameter (default `"picard"`, which leaves the model's historical damped functional iteration and all model solutions unchanged). Setting `TPI_outer_method="anderson"` uses limited-memory Anderson acceleration on the outer-loop residual history, guarded by a trust region anchored to the always-feasible damped point (controlled by `TPI_anderson_m`, `TPI_anderson_beta`, and `TPI_trust_radius`). On a stiff multi-industry reform that limit-cycles under constant dampening, this converged to the same equilibrium in 53 outer iterations vs 126 for constant `nu=0.1` (about −37% wall-clock vs the best damped-`nu` schedule). The new solver lives in `ogcore/solvers.py`. See PR [#1164](https://github.com/PSLmodels/OG-Core/pull/1164).
+
 ## [0.16.3] - 2026-06-25 15:00:00
 
 ### Added
diff --git a/docs/book/content/api/public_api.rst b/docs/book/content/api/public_api.rst
@@ -24,6 +24,7 @@ There is also a link to the source code for each documented member.
    parameter_tables
    parameters
    pensions
+   solvers
    tax
    txfunc
    utils
diff --git a/docs/book/content/api/solvers.rst b/docs/book/content/api/solvers.rst
@@ -0,0 +1,17 @@
+.. _solvers:
+
+Solvers
+=================================================
+
+**solvers.py classes and functions**
+
+ogcore.solvers
+------------------------------------------
+
+.. currentmodule:: ogcore.solvers
+
+.. autoclass:: AndersonAccelerator
+  :members: update, reset
+
+.. automodule:: ogcore.solvers
+  :members: make_outer_updater, pack_outer_vars, unpack_outer_vars
diff --git a/examples/run_ogcore_example.py b/examples/run_ogcore_example.py
@@ -85,6 +85,10 @@ def main():
     """
     # update the effective corporate income tax rate
     og_spec.update({"cit_rate": [[0.35]]})
+    # Optional: accelerate the TPI outer loop with anchored Anderson
+    # acceleration, which solves to the same equilibrium in roughly half
+    # the outer-loop iterations (see the TPI_outer_method parameter).
+    # og_spec.update({"TPI_outer_method": "anderson"})
     p2 = Specifications(
         baseline=False,
         num_workers=num_workers,
diff --git a/ogcore/SS.py b/ogcore/SS.py
@@ -912,7 +912,7 @@ def SS_solver(
         (p.S, 1),
     )
     capital_noncompliance_rate_2D = np.tile(
-        np.reshape(p.labor_income_tax_noncompliance_rate[-1, :], (1, p.J)),
+        np.reshape(p.capital_income_tax_noncompliance_rate[-1, :], (1, p.J)),
         (p.S, 1),
     )
     income_tax_filer_2D = np.tile(
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
@@ -14,7 +14,7 @@
 import pickle
 import scipy.optimize as opt
 from scipy.sparse import csr_matrix
-from ogcore import tax, utils, household, firm, fiscal, pensions
+from ogcore import tax, utils, household, firm, fiscal, pensions, solvers
 from ogcore import aggregates as aggr
 from ogcore.constants import SHOW_RUNTIME
 import os
@@ -924,6 +924,27 @@ 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)
+    # 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
+    # ogcore.solvers.
+    outer_updater = solvers.make_outer_updater(
+        getattr(p, "TPI_outer_method", "picard"), p
+    )
+    # Optional trust region for the accelerated step ("anchored" Anderson):
+    # clamp each step to within trust_radius x the damped-step length of the
+    # always-feasible damped point, growing the radius after an improving step
+    # and shrinking it (with a reset) after a worsening one. This keeps the
+    # accelerated iterate near a feasible point -- the standard cure for the
+    # overshoot-into-infeasible-region divergence of unguarded accelerators.
+    # Defaults to 1.0 (anchored) when accelerating; a non-positive radius
+    # disables the trust region (unguarded accelerator, for testing only).
+    trust_radius = getattr(p, "TPI_trust_radius", 1.0)
+    if trust_radius is not None and trust_radius <= 0:
+        trust_radius = None
+    trust_radius_min = getattr(p, "TPI_trust_radius_min", 0.1)
+    trust_radius_max = getattr(p, "TPI_trust_radius_max", 10.0)
+    prev_accel_dist = np.inf
 
     # Before scattering, temporarily remove unpicklable schema objects
     schema_backup = {}
@@ -1366,17 +1387,80 @@ def run_TPI(p, client=None):
         RM = np.concatenate([RM, np.ones(p.S) * RM[-1]])
 
         # update vars for next iteration
-        w[: p.T] = utils.convex_combo(wnew[: p.T], w[: p.T], p.nu)
-        r[: p.T] = utils.convex_combo(rnew[: p.T], r[: p.T], p.nu)
+        if outer_updater is None:
+            # "picard": the historical damped functional-iteration step,
+            # unchanged, so the default behavior (and golden outputs) is
+            # preserved exactly.
+            w[: p.T] = utils.convex_combo(wnew[: p.T], w[: p.T], p.nu)
+            r[: p.T] = utils.convex_combo(rnew[: p.T], r[: p.T], p.nu)
+            r_p[: p.T] = utils.convex_combo(r_p_new[: p.T], r_p[: p.T], p.nu)
+            p_m[: p.T, :] = utils.convex_combo(
+                new_p_m[: p.T, :], p_m[: p.T, :], p.nu
+            )
+            BQ[: p.T] = utils.convex_combo(BQnew[: p.T], BQ[: p.T], p.nu)
+            if not p.baseline_spending:
+                TR[: p.T] = utils.convex_combo(TR_new[: p.T], TR[: p.T], p.nu)
+        else:
+            # Accelerated outer step on the packed macro/price vector
+            # {r_p, r, w, p_m, BQ[, TR]}; the update rule (Anderson) uses the
+            # recent residual history. Auxiliaries stay damped below.
+            blocks = [
+                (r_p, r_p_new),
+                (r, rnew),
+                (w, wnew),
+                (p_m, new_p_m),
+                (BQ, BQnew),
+            ]
+            if not p.baseline_spending:
+                blocks.append((TR, TR_new))
+            x, gx = solvers.pack_outer_vars(blocks, p.T)
+            # True fixed-point residual (implied vs the PRE-update guess). The
+            # post-update TPIdist below is spuriously ~0 for steps that set
+            # x_next ~= gx (e.g. Anderson's undamped first step), so it is
+            # overridden with this to avoid declaring false convergence.
+            accel_dist = float(np.max(utils.pct_diff_func(gx, x)))
+            # Anchored/trust-region control: grow the radius after an improving
+            # accelerated step and shrink it (resetting the memory) after a
+            # worsening one, using the residual trend as the accept/reject
+            # signal. None => unclamped.
+            if trust_radius is not None and TPIiter > 0:
+                if accel_dist <= prev_accel_dist:
+                    trust_radius = min(trust_radius_max, trust_radius * 1.5)
+                else:
+                    trust_radius = max(trust_radius_min, trust_radius * 0.5)
+                    outer_updater.reset()
+                    logger.info(
+                        f"accel step worsened; trust radius -> {trust_radius}"
+                    )
+            prev_accel_dist = accel_dist
+            # Accelerate the DAMPED map (1-nu)x + nu*G(x), which is contractive
+            # at the calibrated nu, rather than the raw map G -- G is
+            # non-contractive on the stiff case (why Picard needs a low nu), so
+            # accelerating it directly overshoots and diverges. Same fixed
+            # point; convergence is still measured on the raw residual above.
+            gx_damped = (1.0 - p.nu) * x + p.nu * gx
+            x_next = outer_updater.update(x, gx_damped)
+            if not np.all(np.isfinite(x_next)):
+                # accelerated step blew up -> damped Picard fallback + reset
+                x_next = p.nu * gx + (1.0 - p.nu) * x
+                outer_updater.reset()
+                logger.info(
+                    "accelerated step non-finite; Picard fallback + reset"
+                )
+            elif trust_radius is not None:
+                # Clamp the accelerated step to the trust region around the
+                # always-feasible damped point gx_damped, sized relative to the
+                # damped step length so it tightens as the solve converges.
+                dev = x_next - gx_damped
+                dev_norm = float(np.linalg.norm(dev))
+                max_dev = trust_radius * float(np.linalg.norm(gx_damped - x))
+                if dev_norm > max_dev > 0.0:
+                    x_next = gx_damped + dev * (max_dev / dev_norm)
+            solvers.unpack_outer_vars(x_next, blocks, p.T)
+        # Auxiliaries (unchanged): government rate, debt, and the household
+        # policy warm-starts stay on the damped update.
         r_gov[: p.T] = utils.convex_combo(r_gov_new[: p.T], r_gov[: p.T], p.nu)
-        r_p[: p.T] = utils.convex_combo(r_p_new[: p.T], r_p[: p.T], p.nu)
-        p_m[: p.T, :] = utils.convex_combo(
-            new_p_m[: p.T, :], p_m[: p.T, :], p.nu
-        )
-        BQ[: p.T] = utils.convex_combo(BQnew[: p.T], BQ[: p.T], p.nu)
         D[: p.T] = Dnew[: p.T]
-        if not p.baseline_spending:
-            TR[: p.T] = utils.convex_combo(TR_new[: p.T], TR[: p.T], p.nu)
         guesses_b = utils.convex_combo(b_mat, guesses_b, p.nu)
         guesses_n = utils.convex_combo(n_mat, guesses_n, p.nu)
         logger.info(
@@ -1414,15 +1498,23 @@ def run_TPI(p, client=None):
             + list(utils.pct_diff_func(BQnew[: p.T], BQ[: p.T]).flatten())
             + list(utils.pct_diff_func(TR_new[: p.T], TR[: p.T]))
         ).max()
+        if outer_updater is not None:
+            # accelerated methods: use the true residual accel_dist, computed
+            # above from the pre-update guess, because the post-update
+            # distance can be spuriously ~0 for accelerated steps.
+            TPIdist = accel_dist
 
         TPIdist_vec[TPIiter] = TPIdist
-        # After T=10, if cycling occurs, drop the value of nu
-        # wait til after T=10 or so, because sometimes there is a jump up
-        # in the first couple iterations
-        # if TPIiter > 10:
-        #     if TPIdist_vec[TPIiter] - TPIdist_vec[TPIiter - 1] > 0:
-        #         nu /= 2
-        #         print 'New Value of nu:', nu
+        # Accelerated safety net: if a step raised the distance sharply or went
+        # non-finite, reset the accelerator so the next step is a fresh damped
+        # restart (prevents runaway on the stiff case). Never fires for picard.
+        if outer_updater is not None and TPIiter > 0:
+            if (
+                not np.isfinite(TPIdist)
+                or TPIdist > 10.0 * TPIdist_vec[TPIiter - 1]
+            ):
+                outer_updater.reset()
+                logger.info("accelerated step diverged; reset accelerator")
         TPIiter += 1
         logger.info(f"Iteration: {TPIiter}")
         logger.info(f"Distance: {TPIdist}")
diff --git a/ogcore/__init__.py b/ogcore/__init__.py
@@ -21,4 +21,4 @@
 from ogcore.txfunc import *  # noqa: F403
 from ogcore.utils import *  # noqa: F403
 
-__version__ = "0.16.3"
+__version__ = "0.16.4"
diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json
@@ -956,9 +956,9 @@
     },
     "alpha_FA": {
         "title": "Foreign aid payments to domestic government as a share of GDP",
-        "description": "Foreign aid payments to domestic government as a share of GDP.",
+        "description": "Foreign aid payments to domestic government as a share of GDP. Set value for base year, click '+' to add value for next year.  All future years not specified are set to last value entered.",
         "short_description": "Foreign aid as a share of GDP",
-        "param_notation": "$\\alpha_{FA}$",
+        "param_notation": "$\\alpha_{FA,t}$",
         "section_1": "Fiscal Policy Parameters",
         "section_2": "Fiscal Policy Parameters",
         "notes": "",
@@ -4590,6 +4590,84 @@
             }
         }
     },
+    "TPI_outer_method": {
+        "title": "Update rule for the TPI outer (price/aggregate) loop",
+        "description": "Update rule for the transition-path outer loop. 'picard' (default) is the model's historical damped functional iteration x <- (1-nu) x + nu G(x) (see nu), which leaves model solutions unchanged. 'anderson' uses limited-memory Anderson acceleration on the residual history to take larger, better-directed steps, guarded by a trust region anchored to the damped point (TPI_trust_radius).",
+        "short_description": "TPI outer-loop update rule",
+        "section_1": "Model Solution Parameters",
+        "notes": "Opt-in solver acceleration. The default ('picard') reproduces the constant-nu behavior exactly.",
+        "type": "str",
+        "value": [
+            {
+                "value": "picard"
+            }
+        ],
+        "validators": {
+            "choice": {
+                "choices": [
+                    "picard",
+                    "anderson"
+                ]
+            }
+        }
+    },
+    "TPI_anderson_m": {
+        "title": "Anderson acceleration memory depth",
+        "description": "Number of past iterate/residual differences retained by the Anderson accelerator when TPI_outer_method='anderson'. Ignored otherwise.",
+        "short_description": "Anderson memory depth (m)",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "int",
+        "value": [
+            {
+                "value": 5
+            }
+        ],
+        "validators": {
+            "range": {
+                "min": 1,
+                "max": 50
+            }
+        }
+    },
+    "TPI_anderson_beta": {
+        "title": "Anderson acceleration mixing parameter",
+        "description": "Mixing weight for the Anderson step when TPI_outer_method='anderson'. beta=1 is undamped; beta<1 adds damping for robustness far from the solution. Ignored otherwise.",
+        "short_description": "Anderson mixing (beta)",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "float",
+        "value": [
+            {
+                "value": 1.0
+            }
+        ],
+        "validators": {
+            "range": {
+                "min": 0.1,
+                "max": 1.0
+            }
+        }
+    },
+    "TPI_trust_radius": {
+        "title": "Trust radius for the anchored Anderson step",
+        "description": "Initial trust radius for the accelerated TPI step, as a multiple of the damped functional-iteration step length around the always-feasible damped point. Grown after an improving iteration and shrunk (with a reset) after a worsening one. A non-positive value disables the trust region (unguarded accelerator; not recommended). Ignored when TPI_outer_method='picard'.",
+        "short_description": "Anchored-Anderson trust radius",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "float",
+        "value": [
+            {
+                "value": 1.0
+            }
+        ],
+        "validators": {
+            "range": {
+                "min": 0.0,
+                "max": 100.0
+            }
+        }
+    },
     "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.",
diff --git a/ogcore/parameters.py b/ogcore/parameters.py
@@ -181,6 +181,7 @@ def compute_default_params(self):
             "alpha_G",
             "alpha_T",
             "alpha_I",
+            "alpha_FA",
             "alpha_bs_G",
             "alpha_bs_T",
             "alpha_bs_I",
diff --git a/ogcore/solvers.py b/ogcore/solvers.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/tests/test_TPI.py b/tests/test_TPI.py
diff --git a/tests/test_parameters.py b/tests/test_parameters.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -912,7 +912,7 @@ def SS_solver(`
`912`	`912`	`(p.S, 1),`
`913`	`913`	`)`
`914`	`914`	`capital_noncompliance_rate_2D = np.tile(`
`915`		`- np.reshape(p.labor_income_tax_noncompliance_rate[-1, :], (1, p.J)),`
	`915`	`+ np.reshape(p.capital_income_tax_noncompliance_rate[-1, :], (1, p.J)),`
`916`	`916`	`(p.S, 1),`
`917`	`917`	`)`
`918`	`918`	`income_tax_filer_2D = np.tile(`