Skip to content

Add an opt-in anchored Anderson accelerator for the TPI outer loop#1164

Open
marcelolafleur wants to merge 1 commit into
PSLmodels:masterfrom
marcelolafleur:anchored-anderson-tpi
Open

Add an opt-in anchored Anderson accelerator for the TPI outer loop#1164
marcelolafleur wants to merge 1 commit into
PSLmodels:masterfrom
marcelolafleur:anchored-anderson-tpi

Conversation

@marcelolafleur

@marcelolafleur marcelolafleur commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

This is another PR to both improve the solver robustness and accelerate it. This has been especially important for the multi-sector calibrations being created as they take longer and sometimes are in knife-edge solutions that a fixed value of nu sometimes can't solve:

Some reforms make the TPI outer loop limit-cycle: the damped fixed-point iteration oscillates and never reaches mindist_TPI, and the usual fix is to turn nu way down and accept a much slower solve. This PR adds an opt-in accelerated update rule that solves those cases — and easy ones — in 2–3× fewer iterations, while leaving the default solve untouched.

A pluggable outer-loop update rule. A new parameter TPI_outer_method selects how run_TPI updates {r_p, r, w, p_m, BQ, TR} between iterations. The default, "picard", is the existing damped step, and that code path is unchanged — model output under the default is identical to master. Setting "anderson" applies limited-memory Anderson acceleration (new ogcore/solvers.py) to the damped map: it uses the recent residual history to take larger, better-directed steps toward the same fixed point. Convergence for accelerated runs is measured on the true residual, so a bold step can't masquerade as convergence.

Anchored, because plain Anderson diverges here. Unguarded Anderson overshoots on stiff reforms — extrapolated steps land in regions where the household problem breaks down (we measured divergence to ~1e7 on a hard multi-industry case). So each accelerated step is clamped to a trust region around the always-feasible damped iterate: the radius grows while steps improve the residual and shrinks (resetting the accelerator's memory) when they don't, with a damped-Picard fallback if a step goes non-finite. tpi_anderson_m, tpi_anderson_beta, and tpi_trust_radius control this; the trust region is on by default, so "anderson" is safe out of the box. Downstream repos opt in with one parameter on the run they want accelerated.

Tested. Each case runs a country's standard example unmodified, twice — reform solved with picard (reference) and with anderson — from the same baseline, and checks the two reform equilibria agree (max % difference of the Y/C/K/L/r/w paths, judged against each model's own resource-constraint tolerance; late-path agreement of ~1e-7–1e-9 confirms the same steady state):

model / case picard anderson same equilibrium
OG-PHL multi-industry (M=8, capital-intensity reform — limit-cycles at the calibrated nu) 126 iters 53 iters (−58%, −61% wall) max path diff 5.3e-5
OG-ZAF run_og_zaf.py (with EAPD-DRB/OG-ZAF#134's regenerated calibration) 22 11 (−50%, −52% wall) late-path ~9e-7
OG-IDN run_og_idn.py 23 8 max path diff 2.7e-6
OG-Core run_ogcore_example.py 31 16 late-path ~1e-9

Zero failures were attributable to the accelerator across the campaign — every failure we hit reproduced under the default picard and traced to a pre-existing issue.

One pre-existing issue surfaced by the campaign is worth flagging: country parameter JSONs frozen before #1073 lack the new g_n_preTP/rho_preTP/imm_rates_preTP fields, so on ogcore ≥0.16.3 they silently fall back to this repo's US defaults and the transition resource constraint misses by ~10–35% under the default solver. I've been regenerating the country baselines to address this — merged in OG-PHL (EAPD-DRB/OG-PHL#67), OG-IDN (EAPD-DRB/OG-IDN#52), and OG-ETH (EAPD-DRB/OG-ETH#62, EAPD-DRB/OG-ETH#63); still open in OG-ZAF (EAPD-DRB/OG-ZAF#134 — the ZAF row above runs on that PR's calibration). The PHL multi-industry example got the same fix on EAPD-DRB/OG-PHL#63. Separately, multi_industry_example.py in this repo is blocked by a pre-existing get_ci reshape error for I>1, and OG-USA wasn't run (no local environment).

Unit tests cover the accelerator math, the dispatch, and the picard default; the full non-local suite passes. New parameters are documented in default_parameters.json with picard as the default.

cc: @rickecon @jdebacker

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 52.54237% with 56 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.00%. Comparing base (903cfc2) to head (7c7b036).

Files with missing lines Patch % Lines
ogcore/TPI.py 1.92% 51 Missing ⚠️
ogcore/solvers.py 92.30% 5 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##           master    #1164      +/-   ##
==========================================
- Coverage   73.37%   73.00%   -0.37%     
==========================================
  Files          21       22       +1     
  Lines        5374     5483     +109     
==========================================
+ Hits         3943     4003      +60     
- Misses       1431     1480      +49     
Flag Coverage Δ
unittests 73.00% <52.54%> (-0.37%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
ogcore/__init__.py 100.00% <100.00%> (ø)
ogcore/solvers.py 92.30% <92.30%> (ø)
ogcore/TPI.py 35.84% <1.92%> (-3.16%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Comment thread ogcore/solvers.py
per-element scale is kept."""


class AndersonAccelerator(_Accelerator):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's add Google style docstrings with Args and Returns. We'll also want to make sure the docs build for this module in /docs/book/content/api.

@jdebacker jdebacker left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a very nice addition. I requested some additions for documentation. Once that's added and I can run on my machine, I think this is ready to merge. Thanks @marcelolafleur!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants