Merge pull request PSLmodels#1144 from SeaCelo/feature/sparse-foc-jacobian

Add sparse FOC Jacobian to the household solve

Author: jdebacker

CHANGELOG.md

@@ -5,6 +5,11 @@ 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.1] - 2026-06-04 12:00:00
+
+### Added
+
+- Adds a `use_sparse_FOC_jac` `Specifications` parameter (default `True`) that accelerates the time path iteration (TPI) household solve. With it on, `scipy.optimize.root` is given a sparse (banded) finite-difference Jacobian for the stacked household Euler and labor first order conditions: the sparsity pattern is auto-detected once per problem size and the solver then needs far fewer function evaluations per Jacobian build (about 20x fewer on the default S=80 cohort solve), with an automatic fallback to dense finite differences if the Jacobian is not sparse enough to benefit or if a solve fails. The result matches the legacy dense-finite-difference solution to within the model's resource-constraint accuracy floor on every calibration tested (OG-Core standard example, OG-ETH, OG-ZAF, OG-PHL, OG-IDN), giving roughly a 1.9-2.4x TPI speedup at no accuracy cost. Set `use_sparse_FOC_jac=False` to recover bit-identical agreement with v0.16.0 and earlier.
 
 ## [0.16.0] - 2026-06-02 12:00:00
 

docs/book/content/theory/derivations.md

@@ -114,3 +114,16 @@ In the Cobb-Douglas unit elasticity case ($\varepsilon=1$) of the CES production
 ```
 
 Again, even if this simple case, we cannot solve for $r$ as a function of $w$ for the reasons above.
+
+
+(SecAppDerivHHjac)=
+## Sparsity of the household equation Jacobian
+
+  Holding fixed the prices and policies a type-$j$ cohort faces, its $2S$ stationarized necessary conditions {eq}`EqStnrz_eul_n`, {eq}`EqStnrz_eul_b`, and {eq}`EqStnrz_eul_bS` in the $2S$ unknowns $\{n_{j,s},\hat b_{j,s+1}\}_{s=E+1}^{E+S}$ have a banded Jacobian. From the budget constraint {eq}`EqStnrzHHBC`, stationarized consumption at age $s$ depends on only three unknowns,
+
+  ```{math}
+  :label: EqAppDerivHHjac_cons
+    \hat c_{j,s} = \frac{1}{p}\Bigl[(1+r_p)\hat b_{j,s} + \hat w\,e_{j,s}\,n_{j,s} - \widehat{tax}_{j,s} - e^{g_y}\hat b_{j,s+1}\Bigr] + X_{j,s},
+  ```
+
+  where $\widehat{tax}_{j,s}$ depends only on $(\hat b_{j,s}, n_{j,s})$ through labor and capital income (already in the active set, so it adds no further coupling), and $X_{j,s}$ collects terms fixed in the inner solve (bequests $\hat{bq}_{j,s}$, remittances $\hat{rm}_{j,s}$, government transfers $\hat{tr}_{j,s}$, UBI $\hat{ubi}_{j,s}$, the pension benefit $\theta_j$, and the $\hat c_{min,i}$ terms). The labor Euler equation {eq}`EqStnrz_eul_n` at age $s$ therefore depends on $\{\hat b_{j,s},\hat b_{j,s+1},n_{j,s}\}$ alone, and the savings Euler equation {eq}`EqStnrz_eul_b`---which links $\hat c_{j,s}$ to $\hat b_{j,s+1}$ and $\hat c_{j,s+1}$---depends on $\{\hat b_{j,s},\hat b_{j,s+1},\hat b_{j,s+2},n_{j,s},n_{j,s+1}\}$. The marginal tax rates $\tau^{mtrx}_s$ and $\tau^{mtry}_{s+1}$ are functions of own-age income (already in these sets), so they add no further coupling, and the terminal condition {eq}`EqStnrz_eul_bS` is sparser still. Each of the $2S$ equations therefore depends on at most five of the $2S$ unknowns, regardless of $S$, so the Jacobian has at most $10S$ nonzero entries rather than the $(2S)^2 = 4S^2$ of a fully coupled system. This is the per-cohort counterpart to the dense $2JS$ system noted at the start of Chapter {ref}`Chap_Eqm`: cohorts couple only through prices, which are held fixed in the inner solve. A finite-difference Jacobian can then be built from a number of evaluations set by the bandwidth---about seven at $S = 80$---rather than $2S$, by probing together unknowns that affect no common equation (Figure {numref}`FigHHjacSparsity`).

docs/book/content/theory/equilibrium.md

@@ -25,6 +25,15 @@ In all of the specifications of `OG-Core`, we use a two-stage fixed point algori
 
 Our approach is to choose the minimum number of macroeconomic variables in an outer loop in order to be able to solve the household's $2JS$ Euler equations in terms of only the $\bar{n}_{j,s}$ and $\bar{b}_{j,s+1}$ variables directly, holding all other variables constant. The household system of Euler equations has a provable root solution and is orders of magnitude more tractable (less nonlinear) to solve holding these outer loop variables constant.
 
+Moreover, with the outer-loop variables held fixed, each cohort's system of $2S$ Euler equations is not only less nonlinear but structurally sparse: every equation involves at most five of the $2S$ unknowns---a household's own age and its immediate neighbors. The root finder normally probes each unknown separately when building each step ($2S = 160$ evaluations of the system when $S = 80$), but with most equations depending on only a handful of unknowns, those affecting no common equation can be probed together, cutting the count to about seven at $S = 80$---a number set by how many neighbors couple, not by $S$. The parameter `use_sparse_FOC_jac` (default `True`) controls this; set it to `False` to use the legacy dense-finite-difference Jacobian on every call. The structure is derived in Appendix {ref}`SecAppDerivHHjac`.
+
+```{figure} ./images/HH_jac_sparsity.png
+---
+name: FigHHjacSparsity
+---
+Sparsity pattern of the household equation Jacobian, at $S = 12$. Left: the standard finite-difference solve treats every entry of the $2S\times 2S$ matrix as live ($(2S)^2 = 576$ entries). Right: the actual structure---each Euler equation depends only on a household's own age and its immediate neighbors, leaving most entries zero (92 of 576 here; 636 of 25{,}600 at the default $S = 80$).
+```
+
 The steady-state solution method for each of the cases above is associated with a solution method that has a subset of the following outer-loop variables $\Bigl\{\bar{r}_p, \bar{r}, \bar{w}, \{\bar{p}_m\}_{m=1}^{M-1}, \bar{Y}, \overline{TR}, \overline{BQ}, factor\Bigr\}$.
 
 

docs/book/content/theory/images/HH_jac_sparsity.png

@@ -13,6 +13,7 @@
 import numpy as np
 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 aggregates as aggr
 from ogcore.constants import SHOW_RUNTIME
@@ -21,6 +22,19 @@
 import logging
 from distributed import wait
 
+# ``approx_derivative`` and ``group_columns`` live in a private scipy module
+# (scipy.optimize._numdiff). Guard the import so a future scipy change disables
+# the optional sparse-FOC-Jacobian path -- falling back to dense finite
+# differences -- instead of breaking TPI entirely.
+try:
+    from scipy.optimize._numdiff import approx_derivative, group_columns
+
+    _HAVE_SPARSE_FD = True
+except Exception:  # pragma: no cover
+    approx_derivative = None
+    group_columns = None
+    _HAVE_SPARSE_FD = False
+
 logger = logging.getLogger(__name__)
 
 if not SHOW_RUNTIME:
@@ -37,6 +51,91 @@
 ENFORCE_SOLUTION_CHECKS = True
 
 
+_BANDED_JAC_CACHE = {}
+
+
+def _detect_jac_sparsity(fun, x0, args, thr_rel=1e-9):
+    """Auto-detect the Jacobian sparsity of `fun` near x0 and return a
+    (csr_pattern, groups) coloring for sparse finite differences, or None if
+    the Jacobian is too dense to benefit (caller then uses plain dense FD).
+
+    One-time per problem dimension (cached). The pattern is unioned over two
+    nearby points so a coupling that happens to vanish at one point is not
+    missed. Portable across countries/calibrations/pensions: it reads the
+    *actual* structure each run rather than assuming a stencil, so a
+    history-dependent pension (dense columns) is detected automatically and
+    falls back to dense finite differences.
+
+    Args:
+        fun (function): the residual function whose Jacobian is probed
+        x0 (array_like): point at which to probe the sparsity
+        args (tuple): extra arguments passed to ``fun``
+        thr_rel (float): relative threshold below which an entry is treated
+            as a structural zero
+
+    Returns:
+        result (tuple or None): ``(csr_pattern, groups)`` coloring for sparse
+            finite differences, or None if the Jacobian is not sparse enough
+            to benefit
+    """
+    n = len(x0)
+    key = (getattr(fun, "__name__", repr(fun)), n)
+    cached = _BANDED_JAC_CACHE.get(key, "miss")
+    if cached != "miss":
+        return cached
+    result = None
+    try:
+        x = np.asarray(x0, dtype=float)
+        mask = np.zeros((n, n), dtype=bool)
+        for xp in (x, x * 1.05 + 0.01):
+            Jd = approx_derivative(fun, xp, args=args, method="2-point")
+            sc = max(float(np.max(np.abs(Jd))), 1e-300)
+            mask |= np.abs(Jd) > thr_rel * sc
+        P = csr_matrix(mask.astype(int))
+        groups = group_columns(P)
+        if int(groups.max()) + 1 < n // 2:
+            result = (P, groups)
+        # else: pattern not sparse enough to benefit (expected at small
+        # firstdoughnutring dims); silently return None so the caller uses
+        # dense FD without surfacing a noisy notice.
+    except Exception:
+        result = None
+    _BANDED_JAC_CACHE[key] = result
+    return result
+
+
+def _banded_jac_for(fun, x0, args, enabled):
+    """Return a sparse finite-difference `jac` callable for ``opt.root`` when
+    ``enabled`` (``p.use_sparse_FOC_jac``) and the Jacobian is sparse enough to
+    benefit; otherwise None, so the caller uses the default dense solver. Used
+    at the household full-lifetime and triangle root-finds.
+
+    Args:
+        fun (function): the residual function passed to ``opt.root``
+        x0 (array_like): initial guess for the root-find
+        args (tuple): extra arguments passed to ``fun``
+        enabled (bool): value of ``p.use_sparse_FOC_jac``
+
+    Returns:
+        _jac (function or None): a callable returning the dense Jacobian built
+            by sparse finite differences, or None to use dense finite
+            differences
+    """
+    if not enabled or not _HAVE_SPARSE_FD:
+        return None
+    bj = _detect_jac_sparsity(fun, x0, args)
+    if bj is None:
+        return None
+    P, groups = bj
+
+    def _jac(_x, *_a):
+        return approx_derivative(
+            fun, _x, args=_a, method="2-point", sparsity=(P, groups)
+        ).toarray()
+
+    return _jac
+
+
 def get_initial_SS_values(p):
     """
     Get values of variables for the initial period and the steady state
@@ -489,32 +588,47 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
         ]
         mtry_params_to_use = _params_to_array(temp_mtry, p.tax_func_type)
 
-        solutions = opt.root(
-            twist_doughnut,
-            list(b_guesses_to_use) + list(n_guesses_to_use),
-            args=(
-                r_p,
-                w,
-                p_tilde,
-                p_i,
-                bq_to_use,
-                rm_to_use,
-                tr_to_use,
-                theta_to_use,
-                factor,
-                ubi_to_use,
-                j,
-                s,
-                0,
-                etr_params_to_use,
-                mtrx_params_to_use,
-                mtry_params_to_use,
-                initial_b,
-                p,
-            ),
-            method=p.FOC_root_method,
-            tol=MINIMIZER_TOL,
+        _tri_x0 = list(b_guesses_to_use) + list(n_guesses_to_use)
+        _tri_args = (
+            r_p,
+            w,
+            p_tilde,
+            p_i,
+            bq_to_use,
+            rm_to_use,
+            tr_to_use,
+            theta_to_use,
+            factor,
+            ubi_to_use,
+            j,
+            s,
+            0,
+            etr_params_to_use,
+            mtrx_params_to_use,
+            mtry_params_to_use,
+            initial_b,
+            p,
+        )
+        _tri_jac = _banded_jac_for(
+            twist_doughnut, _tri_x0, _tri_args, p.use_sparse_FOC_jac
         )
+        try:
+            solutions = opt.root(
+                twist_doughnut,
+                _tri_x0,
+                args=_tri_args,
+                method=p.FOC_root_method,
+                tol=MINIMIZER_TOL,
+                jac=_tri_jac,
+            )
+        except Exception:
+            solutions = opt.root(
+                twist_doughnut,
+                _tri_x0,
+                args=_tri_args,
+                method=p.FOC_root_method,
+                tol=MINIMIZER_TOL,
+            )
 
         b_vec = solutions.x[: int(len(solutions.x) / 2)]
         b_mat[ind2, p.S - (s + 2) + ind2] = b_vec
@@ -554,32 +668,52 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
             p.tax_func_type,
         )
 
-        solutions = opt.root(
-            twist_doughnut,
-            list(b_guesses_to_use) + list(n_guesses_to_use),
-            args=(
-                r_p,
-                w,
-                p_tilde,
-                p_i,
-                bq_to_use,
-                rm_to_use,
-                tr_to_use,
-                theta_to_use,
-                factor,
-                ubi_to_use,
-                j,
-                None,
-                t,
-                etr_params_to_use,
-                mtrx_params_to_use,
-                mtry_params_to_use,
-                initial_b,
-                p,
-            ),
-            method=p.FOC_root_method,
-            tol=MINIMIZER_TOL,
+        _td_args = (
+            r_p,
+            w,
+            p_tilde,
+            p_i,
+            bq_to_use,
+            rm_to_use,
+            tr_to_use,
+            theta_to_use,
+            factor,
+            ubi_to_use,
+            j,
+            None,
+            t,
+            etr_params_to_use,
+            mtrx_params_to_use,
+            mtry_params_to_use,
+            initial_b,
+            p,
         )
+        _td_x0 = list(b_guesses_to_use) + list(n_guesses_to_use)
+        # When p.use_sparse_FOC_jac is set, supply a banded finite-difference
+        # Jacobian (far fewer evaluations per build); otherwise _bj_jac is None
+        # and opt.root uses its default dense finite differences. The except
+        # clause is a safety net: any solver failure retries with the dense
+        # Jacobian, so the result is unchanged.
+        _bj_jac = _banded_jac_for(
+            twist_doughnut, _td_x0, _td_args, p.use_sparse_FOC_jac
+        )
+        try:
+            solutions = opt.root(
+                twist_doughnut,
+                _td_x0,
+                args=_td_args,
+                method=p.FOC_root_method,
+                tol=MINIMIZER_TOL,
+                jac=_bj_jac,
+            )
+        except Exception:
+            solutions = opt.root(
+                twist_doughnut,
+                _td_x0,
+                args=_td_args,
+                method=p.FOC_root_method,
+                tol=MINIMIZER_TOL,
+            )
         euler_errors[t, :] = solutions.fun
 
         b_vec = solutions.x[: p.S]
@@ -603,6 +737,8 @@ def run_TPI(p, client=None):
             results
 
     """
+    if p.use_sparse_FOC_jac and not _HAVE_SPARSE_FD:
+        print("[TPI] sparse FOC jacobian unavailable; using dense")
     # unpack tuples of parameters
     initial_values, ss_vars, theta, baseline_values = get_initial_SS_values(p)
     B0, b_sinit, b_splus1init, factor, initial_b, initial_n = initial_values

ogcore/TPI.py

@@ -21,4 +21,4 @@
 from ogcore.txfunc import *  # noqa: F403
 from ogcore.utils import *  # noqa: F403
 
-__version__ = "0.16.0"
+__version__ = "0.16.1"

ogcore/__init__.py

@@ -4392,6 +4392,23 @@
            "choice": {"choices": [true, false]}
             }
     },
+    "use_sparse_FOC_jac": {
+        "title": "Flag to use a sparse (banded) finite-difference Jacobian in the household first order condition root finder",
+        "description": "Flag to use a sparse (banded) finite-difference Jacobian in the household first order condition root finder. When True (the default), the sparsity pattern of the stacked Euler/labor first order conditions is auto-detected once per problem size and supplied to scipy.optimize.root, which then needs far fewer function evaluations per Jacobian build than the default dense finite differences. The solver falls back to the dense finite-difference Jacobian automatically if the Jacobian is not sparse enough to benefit or if a solve fails. Set to False to use the legacy dense Jacobian on every call.",
+        "short_description": "Use a banded finite-difference Jacobian in the household solver",
+        "param_notation": "$\\texttt{use_sparse_FOC_jac}$",
+        "section_1": "Model Solution Parameters",
+        "notes": "",
+        "type": "bool",
+        "value": [
+            {
+                "value": true
+            }
+        ],
+        "validators": {
+            "choice": {"choices": [true, false]}
+        }
+    },
     "start_year": {
         "title": "Calendar year in which to start model analysis",
         "description": "Calendar year in which to start model analysis.",

ogcore/default_parameters.json

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ogcore"
-version = "0.16.0"
+version = "0.16.1"
 authors = [
     {name = "Jason DeBacker and Richard W. Evans"},
 ]