Add an opt-in anchored Anderson accelerator for the TPI outer loop#1164
Open
marcelolafleur wants to merge 1 commit into
Open
Add an opt-in anchored Anderson accelerator for the TPI outer loop#1164marcelolafleur wants to merge 1 commit into
marcelolafleur wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
jdebacker
reviewed
Jul 3, 2026
| per-element scale is kept.""" | ||
|
|
||
|
|
||
| class AndersonAccelerator(_Accelerator): |
Member
There was a problem hiding this comment.
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
requested changes
Jul 3, 2026
jdebacker
left a comment
Member
There was a problem hiding this comment.
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!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
nusometimes 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 turnnuway 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_methodselects howrun_TPIupdates{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 (newogcore/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, andtpi_trust_radiuscontrol 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 withanderson— 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):nu)run_og_zaf.py(with EAPD-DRB/OG-ZAF#134's regenerated calibration)run_og_idn.pyrun_ogcore_example.pyZero failures were attributable to the accelerator across the campaign — every failure we hit reproduced under the default
picardand 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_preTPfields, 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.pyin this repo is blocked by a pre-existingget_cireshape error forI>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.jsonwithpicardas the default.cc: @rickecon @jdebacker