diff --git a/CHANGELOG.md b/CHANGELOG.md index 3eaa96816..c1c0ffca2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 47ec882c1..3c3797532 100644 --- a/ogcore/TPI.py +++ b/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, diff --git a/ogcore/__init__.py b/ogcore/__init__.py index cac2d795f..da108187c 100644 --- 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.15.11" +__version__ = "0.15.12" diff --git a/ogcore/household.py b/ogcore/household.py index 7bd19e8f0..d7ec57eb4 100644 --- a/ogcore/household.py +++ b/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: diff --git a/ogcore/txfunc.py b/ogcore/txfunc.py index d2c112942..38a987152 100644 --- a/ogcore/txfunc.py +++ b/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]), diff --git a/pyproject.toml b/pyproject.toml index 2c209ef32..e601a6c3f 100644 --- a/pyproject.toml +++ b/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"}, ] diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 6ca928eae..499946c74 100644 --- a/tests/test_TPI.py +++ b/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. diff --git a/uv.lock b/uv.lock index c35a1543e..1f9d84210 100644 --- a/uv.lock +++ b/uv.lock @@ -1734,7 +1734,7 @@ wheels = [ [[package]] name = "ogcore" -version = "0.15.11" +version = "0.15.12" source = { editable = "." } dependencies = [ { name = "dask" },