From 82c58af64bec6e9d5070e8008e3c9eee0ae73cd3 Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Tue, 28 Jul 2026 12:29:46 -0400 Subject: [PATCH 1/7] Add initial_wealth_ratio to make initial household wealth calibratable --- CHANGELOG.md | 16 ++++++++++++++++ ogcore/TPI.py | 6 +++++- ogcore/default_parameters.json | 18 ++++++++++++++++++ tests/test_TPI.py | 24 ++++++++++++++++++++++++ 4 files changed, 63 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63f4fa4eb..5fd61372d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,22 @@ 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). +## [Unreleased] + +### Added + +- New parameter `initial_wealth_ratio` (default 1.0): the ratio of aggregate + household wealth in the initial period to its steady-state level. The + transition path imposes the steady-state wealth profile on the initial + population rescaled so aggregate initial wealth equals the steady-state + aggregate; when the initial age distribution is far from the stationary one + this hands every initial household a large uniform wealth windfall (younger + population) or confiscation (older population), producing artificial + consumption/investment swings in the first years of any baseline transition. + The new parameter makes initial wealth calibratable to data, like + `initial_debt_ratio` and `initial_Kg_ratio`; the default reproduces the + previous behavior exactly. + ## [0.18.1] - 2026-07-22 12:00:00 ### Bug Fixes diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 5ecf7b936..8bbc74f9e 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -162,7 +162,11 @@ def get_initial_SS_values(p): ss_baseline_vars = utils.safe_read_pickle(baseline_ss) factor = ss_baseline_vars["factor"] B0 = aggr.get_B(ss_baseline_vars["b_sp1"], p, "SS", True) - initial_b = ss_baseline_vars["b_sp1"] * (ss_baseline_vars["B"] / B0) + initial_b = ( + ss_baseline_vars["b_sp1"] + * (ss_baseline_vars["B"] / B0) + * p.initial_wealth_ratio + ) initial_n = ss_baseline_vars["n"] # The DB/NDC/PS pension formulas need the labor supplied before the # time path begins by cohorts alive at t=0. Use the model's initial diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index bc1a58eb4..d628fec67 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -1103,6 +1103,24 @@ } } }, + "initial_wealth_ratio": { + "title": "Ratio of aggregate household wealth in the initial period to its steady-state level", + "description": "Scales the initial distribution of household wealth used to start the transition path. The transition imposes the steady-state wealth profile on the initial population, rescaled so that aggregate initial wealth equals initial_wealth_ratio times the steady-state aggregate. The default of 1.0 reproduces the long-standing behavior.", + "section_1": "Household Parameters", + "notes": "Calibrate from observed household wealth or capital stock data in the model start year, as approximately the ratio of observed wealth-to-GDP to the model's steady-state wealth-to-GDP. With the default of 1.0, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", + "type": "float", + "value": [ + { + "value": 1.0 + } + ], + "validators": { + "range": { + "min": 0.05, + "max": 5.0 + } + } + }, "r_gov_scale": { "title": "Scale parameter to determine government interest rate", "description": "Parameter to scale the market interest rate to find interest rate on government debt.", diff --git a/tests/test_TPI.py b/tests/test_TPI.py index ed424a347..688d8ebed 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -258,6 +258,30 @@ def test_get_initial_SS_values(baseline, param_updates, filename, tmpdir): ) +def test_get_initial_SS_values_initial_wealth_ratio(tmpdir): + """initial_wealth_ratio scales the initial wealth distribution uniformly; + the default of 1.0 reproduces the long-standing behavior.""" + ss_vars = utils.safe_read_pickle( + os.path.join(CUR_PATH, "test_io_data", "OUTPUT", "SS", "SS_vars.pkl") + ) + ss_vars_new = {SS_VAR_NAME_MAPPING[k]: v for k, v in ss_vars.items()} + baseline_dir = os.path.join(tmpdir, "baseline") + utils.mkdirs(os.path.join(baseline_dir, "SS")) + with open(os.path.join(baseline_dir, "SS", "SS_vars.pkl"), "wb") as f: + pickle.dump(ss_vars_new, f) + + initial_b = {} + for ratio in [1.0, 0.7]: + p = Specifications(baseline=True, num_workers=NUM_WORKERS) + p.update_specifications({"initial_wealth_ratio": ratio}) + p.baseline_dir = baseline_dir + p.output_base = baseline_dir + initial_values, _, _, _ = TPI.get_initial_SS_values(p) + initial_b[ratio] = initial_values[4] + + assert np.allclose(initial_b[0.7], 0.7 * initial_b[1.0]) + + def test_firstdoughnutring(): # Test TPI.firstdoughnutring function. Provide inputs to function and # ensure that output returned matches what it has been before. From c0e229506f314305e5f87c33b1058a8b94dbea39 Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Tue, 28 Jul 2026 14:53:27 -0400 Subject: [PATCH 2/7] Anchor initial wealth statically to steady-state GDP; reforms clone the baseline's initial wealth --- CHANGELOG.md | 31 +++++++++++------- ogcore/TPI.py | 60 +++++++++++++++++++++++++++++++--- ogcore/default_parameters.json | 12 +++---- tests/test_TPI.py | 47 ++++++++++++++------------ 4 files changed, 107 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fd61372d..45d5c7dc4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,17 +9,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- New parameter `initial_wealth_ratio` (default 1.0): the ratio of aggregate - household wealth in the initial period to its steady-state level. The - transition path imposes the steady-state wealth profile on the initial - population rescaled so aggregate initial wealth equals the steady-state - aggregate; when the initial age distribution is far from the stationary one - this hands every initial household a large uniform wealth windfall (younger - population) or confiscation (older population), producing artificial - consumption/investment swings in the first years of any baseline transition. - The new parameter makes initial wealth calibratable to data, like - `initial_debt_ratio` and `initial_Kg_ratio`; the default reproduces the - previous behavior exactly. +- New parameter `initial_wealth_ratio` (default 0.0 = disabled): household + wealth to GDP ratio in the initial period of the transition path, anchoring + B(0) = initial_wealth_ratio x steady-state Y. Initial wealth is a + predetermined state, so the anchor is STATIC within the solve, and + steady-state GDP is the anchor base because the steady-state solve has + already pinned it down exactly. Reform runs ignore the parameter and clone + the baseline's initial wealth (read from the baseline's saved transition), + so baseline and reform always share the same initial condition. Anchoring + to initial-period GDP instead was tried and rejected twice: Y(0) is + endogenous, and rescaling the households' initial wealth between + outer-loop iterations -- even damped -- drives the initial cohorts' + root-finding into infeasible negative-consumption roots that satisfy the + extended FOCs and pass the constraint checker. The transition path otherwise imposes the + steady-state wealth profile rescaled so aggregate initial wealth equals the + steady-state aggregate; when the initial age distribution is far from the + stationary one this hands every initial household a large uniform wealth + windfall (younger population) or confiscation (older population), producing + artificial consumption/investment swings in the first years of any baseline + transition. The new parameter makes initial wealth calibratable to data; + the default reproduces the previous behavior exactly. ## [0.18.1] - 2026-07-22 12:00:00 diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 8bbc74f9e..120b7f64c 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -162,11 +162,7 @@ def get_initial_SS_values(p): ss_baseline_vars = utils.safe_read_pickle(baseline_ss) factor = ss_baseline_vars["factor"] B0 = aggr.get_B(ss_baseline_vars["b_sp1"], p, "SS", True) - initial_b = ( - ss_baseline_vars["b_sp1"] - * (ss_baseline_vars["B"] / B0) - * p.initial_wealth_ratio - ) + initial_b = ss_baseline_vars["b_sp1"] * (ss_baseline_vars["B"] / B0) initial_n = ss_baseline_vars["n"] # The DB/NDC/PS pension formulas need the labor supplied before the # time path begins by cohorts alive at t=0. Use the model's initial @@ -237,6 +233,35 @@ def get_initial_SS_values(p): return initial_values, ss_vars, theta, baseline_values +def scale_initial_wealth( + initial_b_shape, B0_shape, target_B0, factor, initial_n, p +): + """ + Rescale the initial wealth distribution to a target aggregate. + + Args: + initial_b_shape (Numpy array): SxJ unscaled initial wealth profile + B0_shape (scalar): aggregate of initial_b_shape over the initial + population + target_B0 (scalar): target aggregate initial wealth + factor (scalar): income scaling factor + initial_n (Numpy array): initial labor supply + p (OG-Core Specifications object): model parameters + + Returns: + (tuple): rescaled initial period values, + (B0, b_sinit, b_splus1init, factor, initial_b, initial_n) + + """ + scale = target_B0 / B0_shape + initial_b = initial_b_shape * scale + b_sinit = np.array( + list(np.zeros(p.J).reshape(1, p.J)) + list(initial_b[:-1]) + ) + b_splus1init = initial_b + return (target_B0, b_sinit, b_splus1init, factor, initial_b, initial_n) + + def firstdoughnutring( guesses, r, @@ -762,6 +787,30 @@ def run_TPI(p, client=None): Kg0_baseline, ) = baseline_values + # Anchor initial household wealth when initial_wealth_ratio is set (> 0). + # Initial wealth is a predetermined state, so the anchor is STATIC within + # the solve (rescaling it between outer-loop iterations -- even damped -- + # drives the initial cohorts' root-finding into infeasible negative- + # consumption roots that satisfy the extended FOCs). A baseline run sets + # aggregate initial wealth to initial_wealth_ratio times steady-state + # GDP, which the steady-state solve has already pinned down exactly; a + # reform run clones the baseline's initial wealth outright (the initial + # state is history -- policy cannot change what households start with). + anchor_initial_wealth = p.initial_wealth_ratio > 0 + if anchor_initial_wealth: + if p.baseline: + target_B0 = p.initial_wealth_ratio * ss_vars["Y"] + else: + baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") + tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) + target_B0 = tpi_baseline_vars["B"][0] + initial_values = scale_initial_wealth( + initial_b, B0, target_B0, factor, initial_n, p + ) + B0, b_sinit, b_splus1init, factor, initial_b, initial_n = ( + initial_values + ) + # Create time path of UBI household benefits and aggregate UBI outlays ubi = p.ubi_nom_array / factor UBI = aggr.get_L(ubi[: p.T], p, "TPI") @@ -1183,6 +1232,7 @@ def run_TPI(p, client=None): ) # Update aggregate variables L[: p.T] = aggr.get_L(n_mat[: p.T], p, "TPI") + B[0] = B0 B[1 : p.T] = aggr.get_B(bmat_splus1[: p.T], p, "TPI", False)[: p.T - 1] w_open = firm.get_w_from_r(p.world_int_rate[: p.T], p, "TPI") diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index d628fec67..bc22ac82d 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -1104,20 +1104,20 @@ } }, "initial_wealth_ratio": { - "title": "Ratio of aggregate household wealth in the initial period to its steady-state level", - "description": "Scales the initial distribution of household wealth used to start the transition path. The transition imposes the steady-state wealth profile on the initial population, rescaled so that aggregate initial wealth equals initial_wealth_ratio times the steady-state aggregate. The default of 1.0 reproduces the long-standing behavior.", + "title": "Aggregate household wealth in the initial period, relative to steady-state GDP", + "description": "Anchors aggregate household wealth in the initial period of the transition path: B(0) = initial_wealth_ratio x steady-state Y, with the age profile keeping the steady-state shape. Steady-state GDP is the anchor base because it is pinned down exactly before the transition solves, making the anchor static (initial wealth is a predetermined state). Reform runs ignore the parameter and clone the baseline run's initial wealth, so baseline and reform always share the same initial condition. The default of 0.0 disables the anchor and reproduces the long-standing behavior, in which aggregate initial wealth is set equal to its steady-state level regardless of the initial population.", "section_1": "Household Parameters", - "notes": "Calibrate from observed household wealth or capital stock data in the model start year, as approximately the ratio of observed wealth-to-GDP to the model's steady-state wealth-to-GDP. With the default of 1.0, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", + "notes": "Calibrate so the solved initial-period wealth-to-GDP ratio matches observed household wealth (capital stock plus domestically held government debt) relative to GDP in the start year: set to the data ratio times the model's Y(0)/Y_ss (one solve iteration pins it; report the delivered B(0)/Y(0)). With the anchor disabled, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", "type": "float", "value": [ { - "value": 1.0 + "value": 0.0 } ], "validators": { "range": { - "min": 0.05, - "max": 5.0 + "min": 0.0, + "max": 20.0 } } }, diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 688d8ebed..9c6df1050 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -258,28 +258,33 @@ def test_get_initial_SS_values(baseline, param_updates, filename, tmpdir): ) -def test_get_initial_SS_values_initial_wealth_ratio(tmpdir): - """initial_wealth_ratio scales the initial wealth distribution uniformly; - the default of 1.0 reproduces the long-standing behavior.""" - ss_vars = utils.safe_read_pickle( - os.path.join(CUR_PATH, "test_io_data", "OUTPUT", "SS", "SS_vars.pkl") +def test_scale_initial_wealth(): + """scale_initial_wealth rescales the wealth profile uniformly to a target + aggregate, keeping the profile's shape and rebuilding the beginning- and + end-of-period views consistently.""" + p = Specifications(baseline=True, num_workers=NUM_WORKERS) + rng = np.random.default_rng(5) + initial_b_shape = rng.uniform(0.1, 2.0, (p.S, p.J)) + B0_shape = 3.0 + initial_n = rng.uniform(0.2, 0.5, (p.S, p.J)) + target_B0 = 4.5 + (B0, b_sinit, b_splus1init, factor, initial_b, n_out) = ( + TPI.scale_initial_wealth( + initial_b_shape, B0_shape, target_B0, 1000.0, initial_n, p + ) ) - ss_vars_new = {SS_VAR_NAME_MAPPING[k]: v for k, v in ss_vars.items()} - baseline_dir = os.path.join(tmpdir, "baseline") - utils.mkdirs(os.path.join(baseline_dir, "SS")) - with open(os.path.join(baseline_dir, "SS", "SS_vars.pkl"), "wb") as f: - pickle.dump(ss_vars_new, f) - - initial_b = {} - for ratio in [1.0, 0.7]: - p = Specifications(baseline=True, num_workers=NUM_WORKERS) - p.update_specifications({"initial_wealth_ratio": ratio}) - p.baseline_dir = baseline_dir - p.output_base = baseline_dir - initial_values, _, _, _ = TPI.get_initial_SS_values(p) - initial_b[ratio] = initial_values[4] - - assert np.allclose(initial_b[0.7], 0.7 * initial_b[1.0]) + assert B0 == target_B0 + assert np.allclose(initial_b, initial_b_shape * (target_B0 / B0_shape)) + assert np.allclose(b_splus1init, initial_b) + assert np.allclose(b_sinit[0, :], np.zeros(p.J)) + assert np.allclose(b_sinit[1:, :], initial_b[:-1, :]) + assert np.allclose(n_out, initial_n) + + +def test_initial_wealth_ratio_default_is_off(): + """The default of 0.0 disables the anchor (legacy behavior).""" + p = Specifications(baseline=True, num_workers=NUM_WORKERS) + assert p.initial_wealth_ratio == 0.0 def test_firstdoughnutring(): From 3258f7b3520a05e6b7390b1120d812121733f830 Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Tue, 28 Jul 2026 16:27:48 -0400 Subject: [PATCH 3/7] Apply ruff format to benchmark README code blocks --- tests/BENCHMARK_README.md | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/BENCHMARK_README.md b/tests/BENCHMARK_README.md index 892f684fb..727b8b562 100644 --- a/tests/BENCHMARK_README.md +++ b/tests/BENCHMARK_README.md @@ -87,17 +87,17 @@ Each benchmark produces a `BenchmarkResult` with the following metrics: ```python @dataclass class BenchmarkResult: - test_name: str # Name of the test - platform: str # Operating system - scheduler: str # Dask scheduler used - num_workers: int # Number of workers - compute_time: float # Execution time in seconds - peak_memory_mb: float # Peak memory usage in MB - avg_memory_mb: float # Average memory usage in MB - data_size_mb: float # Input data size in MB - num_tasks: int # Number of parallel tasks - success: bool # Whether test succeeded - error_message: str # Error details if failed + test_name: str # Name of the test + platform: str # Operating system + scheduler: str # Dask scheduler used + num_workers: int # Number of workers + compute_time: float # Execution time in seconds + peak_memory_mb: float # Peak memory usage in MB + avg_memory_mb: float # Average memory usage in MB + data_size_mb: float # Input data size in MB + num_tasks: int # Number of parallel tasks + success: bool # Whether test succeeded + error_message: str # Error details if failed ``` ### Key Metrics to Monitor @@ -228,7 +228,7 @@ cluster = LocalCluster( n_workers=num_workers, threads_per_worker=2, processes=False, # Use threads, not processes - memory_limit='4GB', + memory_limit="4GB", ) client = Client(cluster) ``` From 6fe54d95e1f49eacec3383f13727a3486108e2ef Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Sun, 9 Aug 2026 10:19:09 -0400 Subject: [PATCH 4/7] Address review: comment no-op B[0] guard, document same-ratio rule, add solve-level anchor test --- ogcore/TPI.py | 4 +++ ogcore/default_parameters.json | 2 +- tests/test_TPI.py | 51 ++++++++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 120b7f64c..85bf83924 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -1232,6 +1232,10 @@ def run_TPI(p, client=None): ) # Update aggregate variables L[: p.T] = aggr.get_L(n_mat[: p.T], p, "TPI") + # B[0] is predetermined (set before the loop, anchored when + # initial_wealth_ratio > 0) and nothing in the loop writes to it, so + # this re-assert is a no-op today. It guards B[0] in case the update + # below is ever refactored to write the full B[: p.T] slice. B[0] = B0 B[1 : p.T] = aggr.get_B(bmat_splus1[: p.T], p, "TPI", False)[: p.T - 1] w_open = firm.get_w_from_r(p.world_int_rate[: p.T], p, "TPI") diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index 60cdb4868..89e3466ef 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -1107,7 +1107,7 @@ "title": "Aggregate household wealth in the initial period, relative to steady-state GDP", "description": "Anchors aggregate household wealth in the initial period of the transition path: B(0) = initial_wealth_ratio x steady-state Y, with the age profile keeping the steady-state shape. Steady-state GDP is the anchor base because it is pinned down exactly before the transition solves, making the anchor static (initial wealth is a predetermined state). Reform runs ignore the parameter and clone the baseline run's initial wealth, so baseline and reform always share the same initial condition. The default of 0.0 disables the anchor and reproduces the long-standing behavior, in which aggregate initial wealth is set equal to its steady-state level regardless of the initial population.", "section_1": "Household Parameters", - "notes": "Calibrate so the solved initial-period wealth-to-GDP ratio matches observed household wealth (capital stock plus domestically held government debt) relative to GDP in the start year: set to the data ratio times the model's Y(0)/Y_ss (one solve iteration pins it; report the delivered B(0)/Y(0)). With the anchor disabled, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", + "notes": "Calibrate so the solved initial-period wealth-to-GDP ratio matches observed household wealth (capital stock plus domestically held government debt) relative to GDP in the start year: set to the data ratio times the model's Y(0)/Y_ss (one solve iteration pins it; report the delivered B(0)/Y(0)). Set the same value in the baseline and reform specifications: a baseline that anchors while the reform keeps the default of 0.0 would give the two runs different initial conditions. With the anchor disabled, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", "type": "float", "value": [ { diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 9c6df1050..f3b11345a 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -287,6 +287,57 @@ def test_initial_wealth_ratio_default_is_off(): assert p.initial_wealth_ratio == 0.0 +@pytest.mark.local +def test_run_TPI_initial_wealth_anchor(tmpdir, dask_client): + """ + A baseline solve delivers aggregate initial wealth equal to + initial_wealth_ratio * steady-state Y, and a reform solve clones the + baseline's initial wealth regardless of its own ratio value. + """ + baseline_dir = os.path.join(tmpdir, "baseline") + p = Specifications( + baseline=True, + baseline_dir=baseline_dir, + output_base=baseline_dir, + num_workers=NUM_WORKERS, + ) + SS.ENFORCE_SOLUTION_CHECKS = True + ss_outputs = SS.run_SS(p, client=dask_client) + utils.mkdirs(os.path.join(baseline_dir, "SS")) + with open(os.path.join(baseline_dir, "SS", "SS_vars.pkl"), "wb") as f: + pickle.dump(ss_outputs, f) + # Anchor mildly below the steady-state wealth-to-GDP ratio so the + # target is feasible for any test calibration while still moving B[0] + # away from its legacy value. + p.initial_wealth_ratio = 0.95 * ss_outputs["B"] / ss_outputs["Y"] + tpi_baseline = TPI.run_TPI(p, client=dask_client) + assert np.allclose( + tpi_baseline["B"][0], + p.initial_wealth_ratio * ss_outputs["Y"], + rtol=1e-10, + ) + utils.mkdirs(os.path.join(baseline_dir, "TPI")) + with open(os.path.join(baseline_dir, "TPI", "TPI_vars.pkl"), "wb") as f: + pickle.dump(tpi_baseline, f) + + # Reform: a deliberately different ratio value, which the clone-baseline + # design must ignore in favor of the baseline's B[0]. + reform_dir = os.path.join(tmpdir, "reform") + p2 = Specifications( + baseline=False, + baseline_dir=baseline_dir, + output_base=reform_dir, + num_workers=NUM_WORKERS, + ) + p2.initial_wealth_ratio = 9.99 + ss_reform = SS.run_SS(p2, client=dask_client) + utils.mkdirs(os.path.join(reform_dir, "SS")) + with open(os.path.join(reform_dir, "SS", "SS_vars.pkl"), "wb") as f: + pickle.dump(ss_reform, f) + tpi_reform = TPI.run_TPI(p2, client=dask_client) + assert np.allclose(tpi_reform["B"][0], tpi_baseline["B"][0], rtol=1e-12) + + def test_firstdoughnutring(): # Test TPI.firstdoughnutring function. Provide inputs to function and # ensure that output returned matches what it has been before. From 82cb8e4d7bdfb6f62914861dff607538d9eccc3e Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Mon, 10 Aug 2026 07:57:12 -0400 Subject: [PATCH 5/7] Slim scale_initial_wealth to computed outputs only; rebuild initial_values at call site --- ogcore/TPI.py | 28 ++++++++++++++++------------ tests/test_TPI.py | 9 ++------- 2 files changed, 18 insertions(+), 19 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 432da7ff2..e36187f99 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -233,9 +233,7 @@ def get_initial_SS_values(p): return initial_values, ss_vars, theta, baseline_values -def scale_initial_wealth( - initial_b_shape, B0_shape, target_B0, factor, initial_n, p -): +def scale_initial_wealth(initial_b_shape, B0_shape, target_B0, p): """ Rescale the initial wealth distribution to a target aggregate. @@ -244,13 +242,11 @@ def scale_initial_wealth( B0_shape (scalar): aggregate of initial_b_shape over the initial population target_B0 (scalar): target aggregate initial wealth - factor (scalar): income scaling factor - initial_n (Numpy array): initial labor supply p (OG-Core Specifications object): model parameters Returns: - (tuple): rescaled initial period values, - (B0, b_sinit, b_splus1init, factor, initial_b, initial_n) + (tuple): rescaled initial period wealth values, + (b_sinit, b_splus1init, initial_b) """ scale = target_B0 / B0_shape @@ -259,7 +255,7 @@ def scale_initial_wealth( list(np.zeros(p.J).reshape(1, p.J)) + list(initial_b[:-1]) ) b_splus1init = initial_b - return (target_B0, b_sinit, b_splus1init, factor, initial_b, initial_n) + return (b_sinit, b_splus1init, initial_b) def firstdoughnutring( @@ -804,11 +800,19 @@ def run_TPI(p, client=None): baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) target_B0 = tpi_baseline_vars["B"][0] - initial_values = scale_initial_wealth( - initial_b, B0, target_B0, factor, initial_n, p + b_sinit, b_splus1init, initial_b = scale_initial_wealth( + initial_b, B0, target_B0, p ) - B0, b_sinit, b_splus1init, factor, initial_b, initial_n = ( - initial_values + B0 = target_B0 + # Rebuild the initial_values tuple consumed by inner_loop with the + # anchored wealth objects; factor and initial_n are unaffected. + initial_values = ( + B0, + b_sinit, + b_splus1init, + factor, + initial_b, + initial_n, ) # Create time path of UBI household benefits and aggregate UBI outlays diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 44a66b928..77eaea811 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -266,19 +266,14 @@ def test_scale_initial_wealth(): rng = np.random.default_rng(5) initial_b_shape = rng.uniform(0.1, 2.0, (p.S, p.J)) B0_shape = 3.0 - initial_n = rng.uniform(0.2, 0.5, (p.S, p.J)) target_B0 = 4.5 - (B0, b_sinit, b_splus1init, factor, initial_b, n_out) = ( - TPI.scale_initial_wealth( - initial_b_shape, B0_shape, target_B0, 1000.0, initial_n, p - ) + b_sinit, b_splus1init, initial_b = TPI.scale_initial_wealth( + initial_b_shape, B0_shape, target_B0, p ) - assert B0 == target_B0 assert np.allclose(initial_b, initial_b_shape * (target_B0 / B0_shape)) assert np.allclose(b_splus1init, initial_b) assert np.allclose(b_sinit[0, :], np.zeros(p.J)) assert np.allclose(b_sinit[1:, :], initial_b[:-1, :]) - assert np.allclose(n_out, initial_n) def test_initial_wealth_ratio_default_is_off(): From b6dd4f66dfeda937ad1e40173cac740f8fbbbb7b Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Mon, 10 Aug 2026 08:42:10 -0400 Subject: [PATCH 6/7] Ensure reforms inherit baseline initial wealth --- ogcore/TPI.py | 25 ++++++++++++++--------- ogcore/default_parameters.json | 4 ++-- tests/test_TPI.py | 37 +++++++++++++++++++++++++++++----- 3 files changed, 50 insertions(+), 16 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index e36187f99..7aacd71ed 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -233,6 +233,18 @@ def get_initial_SS_values(p): return initial_values, ss_vars, theta, baseline_values +def get_initial_wealth_target(p, ss_Y): + """Return the target aggregate initial wealth, or None when disabled.""" + if p.baseline: + if p.initial_wealth_ratio <= 0: + return None + return p.initial_wealth_ratio * ss_Y + + baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") + tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) + return tpi_baseline_vars["B"][0] + + def scale_initial_wealth(initial_b_shape, B0_shape, target_B0, p): """ Rescale the initial wealth distribution to a target aggregate. @@ -783,7 +795,8 @@ def run_TPI(p, client=None): Kg0_baseline, ) = baseline_values - # Anchor initial household wealth when initial_wealth_ratio is set (> 0). + # Anchor baseline initial household wealth when initial_wealth_ratio is + # set (> 0), and always clone that initial wealth in reform runs. # Initial wealth is a predetermined state, so the anchor is STATIC within # the solve (rescaling it between outer-loop iterations -- even damped -- # drives the initial cohorts' root-finding into infeasible negative- @@ -792,14 +805,8 @@ def run_TPI(p, client=None): # GDP, which the steady-state solve has already pinned down exactly; a # reform run clones the baseline's initial wealth outright (the initial # state is history -- policy cannot change what households start with). - anchor_initial_wealth = p.initial_wealth_ratio > 0 - if anchor_initial_wealth: - if p.baseline: - target_B0 = p.initial_wealth_ratio * ss_vars["Y"] - else: - baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") - tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) - target_B0 = tpi_baseline_vars["B"][0] + target_B0 = get_initial_wealth_target(p, ss_vars["Y"]) + if target_B0 is not None: b_sinit, b_splus1init, initial_b = scale_initial_wealth( initial_b, B0, target_B0, p ) diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index d5d0800ff..5c88043c8 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -1105,9 +1105,9 @@ }, "initial_wealth_ratio": { "title": "Aggregate household wealth in the initial period, relative to steady-state GDP", - "description": "Anchors aggregate household wealth in the initial period of the transition path: B(0) = initial_wealth_ratio x steady-state Y, with the age profile keeping the steady-state shape. Steady-state GDP is the anchor base because it is pinned down exactly before the transition solves, making the anchor static (initial wealth is a predetermined state). Reform runs ignore the parameter and clone the baseline run's initial wealth, so baseline and reform always share the same initial condition. The default of 0.0 disables the anchor and reproduces the long-standing behavior, in which aggregate initial wealth is set equal to its steady-state level regardless of the initial population.", + "description": "Anchors aggregate household wealth in the initial period of the transition path: B(0) = initial_wealth_ratio x steady-state Y, with the age profile keeping the steady-state shape. Steady-state GDP is the anchor base because it is pinned down exactly before the transition solves, making the anchor static (initial wealth is a predetermined state). Reform runs ignore the parameter and clone the baseline run's initial wealth, so baseline and reform always share the same initial condition. For baseline runs, the default of 0.0 disables the anchor and reproduces the long-standing behavior, in which aggregate initial wealth is set equal to its steady-state level regardless of the initial population.", "section_1": "Household Parameters", - "notes": "Calibrate so the solved initial-period wealth-to-GDP ratio matches observed household wealth (capital stock plus domestically held government debt) relative to GDP in the start year: set to the data ratio times the model's Y(0)/Y_ss (one solve iteration pins it; report the delivered B(0)/Y(0)). Set the same value in the baseline and reform specifications: a baseline that anchors while the reform keeps the default of 0.0 would give the two runs different initial conditions. With the anchor disabled, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", + "notes": "Calibrate so the solved initial-period wealth-to-GDP ratio matches observed household wealth (capital stock plus domestically held government debt) relative to GDP in the start year: set to the data ratio times the model's Y(0)/Y_ss (one solve iteration pins it; report the delivered B(0)/Y(0)). Set this value on the baseline specification; reform runs clone the baseline's initial wealth even when their own value remains at the default 0.0. With the baseline anchor disabled, an initial age distribution far from the stationary one implies a large uniform wealth windfall (younger population) or confiscation (older population) for all initial households.", "type": "float", "value": [ { diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 77eaea811..d3cf622c7 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -277,17 +277,44 @@ def test_scale_initial_wealth(): def test_initial_wealth_ratio_default_is_off(): - """The default of 0.0 disables the anchor (legacy behavior).""" + """The default of 0.0 leaves the baseline anchor disabled.""" p = Specifications(baseline=True, num_workers=NUM_WORKERS) assert p.initial_wealth_ratio == 0.0 +def test_get_initial_wealth_target(tmpdir): + """Baseline anchoring is opt-in, while reform cloning is unconditional.""" + baseline_dir = os.path.join(tmpdir, "baseline") + p = Specifications( + baseline=True, + baseline_dir=baseline_dir, + num_workers=NUM_WORKERS, + ) + assert TPI.get_initial_wealth_target(p, 2.5) is None + p.initial_wealth_ratio = 2.0 + assert np.allclose(TPI.get_initial_wealth_target(p, 2.5), 5.0) + + utils.mkdirs(os.path.join(baseline_dir, "TPI")) + with open(os.path.join(baseline_dir, "TPI", "TPI_vars.pkl"), "wb") as f: + pickle.dump({"B": np.array([7.0])}, f) + p2 = Specifications( + baseline=False, + baseline_dir=baseline_dir, + output_base=os.path.join(tmpdir, "output"), + num_workers=NUM_WORKERS, + ) + assert p2.initial_wealth_ratio == 0.0 + assert np.allclose(TPI.get_initial_wealth_target(p2, 2.5), 7.0) + p2.initial_wealth_ratio = 9.99 + assert np.allclose(TPI.get_initial_wealth_target(p2, 2.5), 7.0) + + @pytest.mark.local def test_run_TPI_initial_wealth_anchor(tmpdir, dask_client): """ A baseline solve delivers aggregate initial wealth equal to initial_wealth_ratio * steady-state Y, and a reform solve clones the - baseline's initial wealth regardless of its own ratio value. + baseline's initial wealth when its own ratio remains at the default 0.0. """ baseline_dir = os.path.join(tmpdir, "baseline") p = Specifications( @@ -315,8 +342,8 @@ def test_run_TPI_initial_wealth_anchor(tmpdir, dask_client): with open(os.path.join(baseline_dir, "TPI", "TPI_vars.pkl"), "wb") as f: pickle.dump(tpi_baseline, f) - # Reform: a deliberately different ratio value, which the clone-baseline - # design must ignore in favor of the baseline's B[0]. + # Reform: leave the ratio at its default 0.0. The clone-baseline design + # must still use the baseline's B[0]. reform_dir = os.path.join(tmpdir, "reform") p2 = Specifications( baseline=False, @@ -324,7 +351,7 @@ def test_run_TPI_initial_wealth_anchor(tmpdir, dask_client): output_base=reform_dir, num_workers=NUM_WORKERS, ) - p2.initial_wealth_ratio = 9.99 + assert p2.initial_wealth_ratio == 0.0 ss_reform = SS.run_SS(p2, client=dask_client) utils.mkdirs(os.path.join(reform_dir, "SS")) with open(os.path.join(reform_dir, "SS", "SS_vars.pkl"), "wb") as f: From b440d4840364cef32f5919b9449e2979af248e62 Mon Sep 17 00:00:00 2001 From: marcelolafleur Date: Mon, 10 Aug 2026 09:22:50 -0400 Subject: [PATCH 7/7] Simplify reform wealth inheritance --- ogcore/TPI.py | 21 ++++++++------------- tests/test_TPI.py | 27 --------------------------- 2 files changed, 8 insertions(+), 40 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 7aacd71ed..9d2041661 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -233,18 +233,6 @@ def get_initial_SS_values(p): return initial_values, ss_vars, theta, baseline_values -def get_initial_wealth_target(p, ss_Y): - """Return the target aggregate initial wealth, or None when disabled.""" - if p.baseline: - if p.initial_wealth_ratio <= 0: - return None - return p.initial_wealth_ratio * ss_Y - - baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") - tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) - return tpi_baseline_vars["B"][0] - - def scale_initial_wealth(initial_b_shape, B0_shape, target_B0, p): """ Rescale the initial wealth distribution to a target aggregate. @@ -805,7 +793,14 @@ def run_TPI(p, client=None): # GDP, which the steady-state solve has already pinned down exactly; a # reform run clones the baseline's initial wealth outright (the initial # state is history -- policy cannot change what households start with). - target_B0 = get_initial_wealth_target(p, ss_vars["Y"]) + target_B0 = None + if p.baseline: + if p.initial_wealth_ratio > 0: + target_B0 = p.initial_wealth_ratio * ss_vars["Y"] + else: + baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl") + tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi) + target_B0 = tpi_baseline_vars["B"][0] if target_B0 is not None: b_sinit, b_splus1init, initial_b = scale_initial_wealth( initial_b, B0, target_B0, p diff --git a/tests/test_TPI.py b/tests/test_TPI.py index d3cf622c7..29459c64d 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -282,33 +282,6 @@ def test_initial_wealth_ratio_default_is_off(): assert p.initial_wealth_ratio == 0.0 -def test_get_initial_wealth_target(tmpdir): - """Baseline anchoring is opt-in, while reform cloning is unconditional.""" - baseline_dir = os.path.join(tmpdir, "baseline") - p = Specifications( - baseline=True, - baseline_dir=baseline_dir, - num_workers=NUM_WORKERS, - ) - assert TPI.get_initial_wealth_target(p, 2.5) is None - p.initial_wealth_ratio = 2.0 - assert np.allclose(TPI.get_initial_wealth_target(p, 2.5), 5.0) - - utils.mkdirs(os.path.join(baseline_dir, "TPI")) - with open(os.path.join(baseline_dir, "TPI", "TPI_vars.pkl"), "wb") as f: - pickle.dump({"B": np.array([7.0])}, f) - p2 = Specifications( - baseline=False, - baseline_dir=baseline_dir, - output_base=os.path.join(tmpdir, "output"), - num_workers=NUM_WORKERS, - ) - assert p2.initial_wealth_ratio == 0.0 - assert np.allclose(TPI.get_initial_wealth_target(p2, 2.5), 7.0) - p2.initial_wealth_ratio = 9.99 - assert np.allclose(TPI.get_initial_wealth_target(p2, 2.5), 7.0) - - @pytest.mark.local def test_run_TPI_initial_wealth_anchor(tmpdir, dask_client): """