From 1723aaaf576f0cead5a66e6c23dd1b82e600ba4b Mon Sep 17 00:00:00 2001 From: arihantlodha-cmd Date: Sun, 23 Aug 2026 14:15:43 +0900 Subject: [PATCH] Allow a supplied initial distribution of wealth in the transition The time path solver initializes the t=0 distribution of household wealth from the steady-state distribution, scaled so its aggregate matches the steady-state B. That assumes the economy starts already at its steady-state age-wealth profile. For a country whose model start-year age structure is far from its steady state, such as a rapidly aging population, that profile is a poor initial condition and can leave a small resource-constraint residual in the first periods of the path. This adds an optional hook to supply an observed initial distribution instead. A new boolean use_initial_b_dist gates a new SxJ array initial_b_dist; when the flag is set, get_initial_SS_values takes the shape of the initial wealth distribution from initial_b_dist rather than the steady state. The aggregate is still rescaled to the steady-state B, so only the cross-sectional shape is overridden. The default leaves the flag off and reproduces the previous initial_b exactly, so existing runs are unchanged. Adds test_get_initial_SS_values_custom_dist covering the aggregate-preserving and shape-override behavior. --- ogcore/TPI.py | 16 +++++++++++-- ogcore/default_parameters.json | 35 ++++++++++++++++++++++++++++ ogcore/parameters.py | 5 ++++ tests/test_TPI.py | 42 ++++++++++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 2 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 2032d42c3..21394be39 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -161,8 +161,20 @@ def get_initial_SS_values(p): baseline_ss = os.path.join(p.baseline_dir, "SS", "SS_vars.pkl") 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) + # Shape of the initial (t=0) distribution of wealth across ages and + # ability types. By default this is the steady-state distribution + # (ss_baseline_vars["b_sp1"]), which assumes the economy begins already + # at its steady-state age-wealth profile. A country model whose starting + # age structure is far from its steady state (e.g., a rapidly aging + # population) can instead supply an observed distribution via + # initial_b_dist and use_initial_b_dist; either way the aggregate is + # scaled to the steady-state B below, so only the shape is overridden. + if p.use_initial_b_dist: + b_dist = p.initial_b_dist + else: + b_dist = ss_baseline_vars["b_sp1"] + B0 = aggr.get_B(b_dist, p, "SS", True) + initial_b = b_dist * (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 diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index 5ab5a94ca..1baea1d03 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -1085,6 +1085,41 @@ } } }, + "use_initial_b_dist": { + "title": "Flag to use a supplied initial distribution of wealth", + "description": "If true, the initial (t=0) distribution of wealth across ages and ability types is taken from initial_b_dist rather than from the steady-state distribution. Useful when the model start-year age structure is far from its steady state.", + "short_description": "Whether to use a supplied initial wealth distribution", + "param_notation": "$\\texttt{use_initial_b_dist}$", + "section_1": "Household Parameters", + "section_2": "Endowments", + "notes": "", + "type": "bool", + "value": [ + { + "value": false + } + ] + }, + "initial_b_dist": { + "title": "Initial distribution of wealth over ages and ability types", + "description": "Shape of the initial (t=0) distribution of household wealth across ages (S) and ability types (J), used only when use_initial_b_dist is true. Only the shape matters: the aggregate is rescaled to the steady-state level inside the transition solver, so units are arbitrary.", + "short_description": "Supplied initial wealth distribution by age and type", + "param_notation": "$b_{j,s,0}$", + "section_1": "Household Parameters", + "section_2": "Endowments", + "notes": "", + "type": "float", + "number_dims": 2, + "value": [ + { + "value": [ + [ + 0.0 + ] + ] + } + ] + }, "initial_Kg_ratio": { "title": "Government capital (public capital) to GDP ratio in the initial period", "description": "Government capital (aka infrastructure) to GDP ratio in the initial period.", diff --git a/ogcore/parameters.py b/ogcore/parameters.py index 6e4c13ddf..5a1c58b3f 100644 --- a/ogcore/parameters.py +++ b/ogcore/parameters.py @@ -303,6 +303,11 @@ def compute_default_params(self): param_in, dims=(self.T, self.S, self.J), item="e" ) setattr(self, "e", param_out) + # shape the supplied initial wealth distribution to (S, J); when it is + # not used (use_initial_b_dist is False) this is the harmless default + self.initial_b_dist = extrapolate_array( + self.initial_b_dist, dims=(self.S, self.J), item="initial_b_dist" + ) # Extrapolate chi_n over T + S param_in = getattr(self, "chi_n") param_out = extrapolate_array( diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 2a2b680d1..da31856a3 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -258,6 +258,48 @@ def test_get_initial_SS_values(baseline, param_updates, filename, tmpdir): ) +def test_get_initial_SS_values_custom_dist(tmpdir): + """ + A supplied initial wealth distribution (use_initial_b_dist) overrides the + shape of initial_b but leaves its aggregate scaled to the steady-state B, + and leaving the flag off reproduces the default behavior exactly. + """ + old_baseline_dir = os.path.join(CUR_PATH, "test_io_data", "OUTPUT") + ss_vars = utils.safe_read_pickle( + os.path.join(old_baseline_dir, "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) + + # default (flag off) + p_off = Specifications(baseline=True, num_workers=NUM_WORKERS) + p_off.baseline_dir = p_off.output_base = baseline_dir + iv_off = TPI.get_initial_SS_values(p_off)[0] + B0_off, initial_b_off = iv_off[0], iv_off[4] + + # supplied distribution (flag on) + rng = np.random.default_rng(1) + custom = rng.random((p_off.S, p_off.J)) + 0.05 + p_on = Specifications(baseline=True, num_workers=NUM_WORKERS) + p_on.update_specifications( + {"use_initial_b_dist": True, "initial_b_dist": custom.tolist()} + ) + p_on.baseline_dir = p_on.output_base = baseline_dir + iv_on = TPI.get_initial_SS_values(p_on)[0] + B0_on, initial_b_on = iv_on[0], iv_on[4] + + # aggregate is preserved (both scaled to the same steady-state B) + assert np.isclose(B0_on, B0_off, rtol=1e-10) + # the shape is actually overridden + assert not np.allclose(initial_b_on, initial_b_off) + # and the resulting per-capita distribution is proportional to the input + ratio = initial_b_on / custom + assert np.allclose(ratio, ratio.flat[0], rtol=1e-9) + + def test_firstdoughnutring(): # Test TPI.firstdoughnutring function. Provide inputs to function and # ensure that output returned matches what it has been before.