Add initial_wealth_ratio to calibrate initial household wealth

CHANGELOG.md

@@ -9,6 +9,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- 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.
 - Stall detection for the TPI outer loop (Issue #1177): when the best
   distance has not improved over the last `TPI_stall_window` iterations
   (default 50; 0 disables), `run_TPI` logs a diagnosis distinguishing a

ogcore/TPI.py

@@ -233,6 +233,31 @@ 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, 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
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        (tuple): rescaled initial period wealth values,
+            (b_sinit, b_splus1init, initial_b)
+
+    """
+    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 (b_sinit, b_splus1init, initial_b)
+
+
 def firstdoughnutring(
     guesses,
     r,
@@ -758,6 +783,40 @@ def run_TPI(p, client=None):
         Kg0_baseline,
     ) = baseline_values
 
+    # 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-
+    # 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).
+    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
+        )
+        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
     ubi = p.ubi_nom_array / factor
     UBI = aggr.get_L(ubi[: p.T], p, "TPI")
@@ -1181,6 +1240,11 @@ 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")
 

ogcore/default_parameters.json

@@ -1103,6 +1103,24 @@
             }
         }
     },
+    "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. 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 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": [
+            {
+                "value": 0.0
+            }
+        ],
+        "validators": {
+            "range": {
+                "min": 0.0,
+                "max": 20.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.",

tests/test_TPI.py

@@ -258,6 +258,81 @@ def test_get_initial_SS_values(baseline, param_updates, filename, tmpdir):
         )
 
 
+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
+    target_B0 = 4.5
+    b_sinit, b_splus1init, initial_b = TPI.scale_initial_wealth(
+        initial_b_shape, B0_shape, target_B0, p
+    )
+    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, :])
+
+
+def test_initial_wealth_ratio_default_is_off():
+    """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
+
+
+@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 when its own ratio remains at the default 0.0.
+    """
+    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: 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,
+        baseline_dir=baseline_dir,
+        output_base=reform_dir,
+        num_workers=NUM_WORKERS,
+    )
+    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:
+        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.