Allow a supplied initial distribution of wealth in the transition

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

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.",

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(

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.