Merge pull request #1128 from SeaCelo/perf/cache-e-long

3.3x speedup: cache e_long, skip redundant tax-parameter conversions

Author: jdebacker

CHANGELOG.md

@@ -5,6 +5,15 @@ 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.15.12] - 2026-05-14 12:00:00
+
+### Added
+
+- Caches the `e_long` array in `household.FOC_savings` and `household.FOC_labor` rather than rebuilding it on every call. `e_long` is a pure function of `p.e`, `p.S`, and `p.J`, none of which change during a solve, so a `_get_e_long` helper now builds it once per worker and reuses it. Profiling identified this rebuild as the single most expensive operation in a TPI run. The change is a pure cache — model output is bit-for-bit identical to master — and gives roughly a 3x speedup on a single-reform TPI run. See PR [#1128](https://github.com/PSLmodels/OG-Core/pull/1128).
+- Builds the per-period tax-parameter slices in `TPI.inner_loop` as numpy arrays via a new `_params_to_array` helper, and switches `txfunc.get_tax_rates` to `np.asarray`, so the repeated per-call list-to-array conversion is skipped on the hot TPI path. `mono` and `mono2D` tax functions store callables rather than numbers, so their nested-list form is passed through unchanged. Profiling identified this conversion as the next hot spot after the `e_long` rebuild. Model output is bit-for-bit identical to master, and the change gives roughly a further 10% speedup on a single-reform TPI run (about 3.3x cumulative versus master). See PR [#1128](https://github.com/PSLmodels/OG-Core/pull/1128).
+- Changes the minimum of the allowable range for `tau_c` in `default_parameters.py` to allow for government consumption subsidies.
+- Fixes a bug in the `parameter_plots.plot_fert_rates` function. See PR [#1127](https://github.com/PSLmodels/OG-Core/pull/1127).
+
 ## [0.15.11] - 2026-05-08 12:00:00
 
 ### Added

ogcore/TPI.py

@@ -363,6 +363,19 @@ def twist_doughnut(
     return list(error1.flatten()) + list(error2.flatten())
 
 
+def _params_to_array(nested, tax_func_type):
+    """Convert a nested tax-parameter slice to a numpy array.
+
+    For numeric tax functions this lets ``get_tax_rates`` skip a costly
+    per-call list-to-array conversion. ``mono`` and ``mono2D`` store
+    callables rather than numbers, so their nested-list form is returned
+    unchanged.
+    """
+    if tax_func_type in ("mono", "mono2D"):
+        return nested
+    return np.array(nested)
+
+
 def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
     """
     Given path of economic aggregates and factor prices, solves
@@ -456,27 +469,25 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
         ubi_to_use = np.diag(ubi[: p.S, :, j], p.S - (s + 2))
 
         num_params = len(p.etr_params[0][0])
+        # Convert the per-age tax-parameter slices to arrays here, once,
+        # rather than letting get_tax_rates re-convert them from nested
+        # lists on each of its many calls. mono/mono2D store callables,
+        # not numbers, so _params_to_array leaves them as nested lists.
         temp_etr = [
             [p.etr_params[t][p.S - s - 2 + t][i] for i in range(num_params)]
             for t in range(s + 2)
         ]
-        etr_params_to_use = [
-            [temp_etr[i][j] for j in range(num_params)] for i in range(s + 2)
-        ]
+        etr_params_to_use = _params_to_array(temp_etr, p.tax_func_type)
         temp_mtrx = [
             [p.mtrx_params[t][p.S - s - 2 + t][i] for i in range(num_params)]
             for t in range(s + 2)
         ]
-        mtrx_params_to_use = [
-            [temp_mtrx[i][j] for j in range(num_params)] for i in range(s + 2)
-        ]
+        mtrx_params_to_use = _params_to_array(temp_mtrx, p.tax_func_type)
         temp_mtry = [
             [p.mtry_params[t][p.S - s - 2 + t][i] for i in range(num_params)]
             for t in range(s + 2)
         ]
-        mtry_params_to_use = [
-            [temp_mtry[i][j] for j in range(num_params)] for i in range(s + 2)
-        ]
+        mtry_params_to_use = _params_to_array(temp_mtry, p.tax_func_type)
 
         solutions = opt.root(
             twist_doughnut,
@@ -521,18 +532,27 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
 
         # initialize array of diagonal elements
         num_params = len(p.etr_params[t][0])
-        etr_params_to_use = [
-            [p.etr_params[t + s][s][i] for i in range(num_params)]
-            for s in range(p.S)
-        ]
-        mtrx_params_to_use = [
-            [p.mtrx_params[t + s][s][i] for i in range(num_params)]
-            for s in range(p.S)
-        ]
-        mtry_params_to_use = [
-            [p.mtry_params[t + s][s][i] for i in range(num_params)]
-            for s in range(p.S)
-        ]
+        etr_params_to_use = _params_to_array(
+            [
+                [p.etr_params[t + s][s][i] for i in range(num_params)]
+                for s in range(p.S)
+            ],
+            p.tax_func_type,
+        )
+        mtrx_params_to_use = _params_to_array(
+            [
+                [p.mtrx_params[t + s][s][i] for i in range(num_params)]
+                for s in range(p.S)
+            ],
+            p.tax_func_type,
+        )
+        mtry_params_to_use = _params_to_array(
+            [
+                [p.mtry_params[t + s][s][i] for i in range(num_params)]
+                for s in range(p.S)
+            ],
+            p.tax_func_type,
+        )
 
         solutions = opt.root(
             twist_doughnut,

ogcore/__init__.py

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

ogcore/household.py

@@ -18,6 +18,25 @@
 """
 
 
+def _get_e_long(p):
+    """Return ``p.e`` extended for the TPI transition window.
+
+    ``e_long`` is a pure function of ``p.e`` / ``p.S`` / ``p.J`` -- none of
+    which change during a solve -- so it is built once and cached on the
+    parameters object. ``FOC_savings`` and ``FOC_labor`` previously rebuilt
+    this array on every call, which profiling identified as the single most
+    expensive operation in a TPI run.
+    """
+    e_long = getattr(p, "_e_long_cache", None)
+    if e_long is None:
+        e_long = np.concatenate(
+            (p.e, np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1))),
+            axis=0,
+        )
+        p._e_long_cache = e_long
+    return e_long
+
+
 def marg_ut_cons(c, sigma):
     r"""
     Compute the marginal utility of consumption.
@@ -493,13 +512,7 @@ def FOC_savings(
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, j]
             wealth_tax_filer = p.wealth_tax_filer[t : t + length, j]
-            e_long = np.concatenate(
-                (
-                    p.e,
-                    np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)),
-                ),
-                axis=0,
-            )
+            e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0))
     else:
         chi_b = p.chi_b
@@ -521,13 +534,7 @@ def FOC_savings(
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, :]
             wealth_tax_filer = p.wealth_tax_filer[t : t + length, :]
-            e_long = np.concatenate(
-                (
-                    p.e,
-                    np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)),
-                ),
-                axis=0,
-            )
+            e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, :], max(p.S - length, 0))
     e = np.squeeze(e)
     if method == "SS":
@@ -707,13 +714,7 @@ def FOC_labor(
                 t : t + length, j
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, j]
-            e_long = np.concatenate(
-                (
-                    p.e,
-                    np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)),
-                ),
-                axis=0,
-            )
+            e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0))
     else:
         if method == "SS":
@@ -729,13 +730,7 @@ def FOC_labor(
                 t : t + length, :
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, :]
-            e_long = np.concatenate(
-                (
-                    p.e,
-                    np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)),
-                ),
-                axis=0,
-            )
+            e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0))
     if method == "TPI":
         if b.ndim == 2:

ogcore/txfunc.py

@@ -83,8 +83,9 @@ def get_tax_rates(
     income = X + Y
     if tax_func_type != "mono":
         # easier to use arrays for calculations below, except when
-        # can't (bc lists of functions)
-        params = np.array(params)
+        # can't (bc lists of functions). asarray avoids a copy when the
+        # caller already passes an array (the hot TPI path does).
+        params = np.asarray(params)
     if tax_func_type == "GS":
         phi0, phi1, phi2 = (
             np.squeeze(params[..., 0]),

pyproject.toml

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

tests/test_TPI.py

@@ -4,6 +4,7 @@
     - test_get_initial_SS_values(), 3 parameterizations
     - test_firstdoughnutring(), 1 parameterization
     - test_twist_doughnut(), 2 parameterizations
+    - test_params_to_array(), 7 parameterizations
     - test_inner_loop(), 1 parameterization
     - test_run_TPI_full_run(), 11 parameterizations, local only
     - test_run_TPI(), 2 parameterizations, local only
@@ -364,6 +365,26 @@ def test_twist_doughnut(file_inputs, file_outputs):
     assert np.allclose(np.array(test_list), np.array(expected_list), atol=1e-5)
 
 
+@pytest.mark.parametrize(
+    "tax_func_type",
+    ["DEP", "DEP_totalinc", "GS", "HSV", "linear", "mono", "mono2D"],
+    ids=["DEP", "DEP_totalinc", "GS", "HSV", "linear", "mono", "mono2D"],
+)
+def test_params_to_array(tax_func_type):
+    # Test TPI._params_to_array helper.  For numeric tax functions the
+    # nested list is converted to a numpy array; for mono/mono2D (which
+    # store callables, not numbers) the nested list is returned unchanged.
+    nested = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+    result = TPI._params_to_array(nested, tax_func_type)
+    if tax_func_type in ("mono", "mono2D"):
+        assert result is nested
+    else:
+        expected = np.array(nested)
+        assert isinstance(result, np.ndarray)
+        assert np.array_equal(result, expected)
+        assert result.dtype == expected.dtype
+
+
 def test_inner_loop():
     # Test TPI.inner_loop function.  Provide inputs to function and
     # ensure that output returned matches what it has been before.