Better handle the pre-time path population distribution

[jdebacker]

The population distribution before the start year of the model is necessary to compute savings and bequests in period 0.

This PR adds explicitly the parameters for mortality, immigration, and population growth rates from period 0 to period 1 in order to make for a consistent transition from this "pre model period" to period 0.

cc @rickecon

[codecov-commenter]

Codecov Report

❌ Patch coverage is 76.47059% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.19%. Comparing base (df883b4) to head (59ae3e5).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
ogcore/parameters.py	0.00%	3 Missing ⚠️
ogcore/demographics.py	90.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1073      +/-   ##
==========================================
+ Coverage   73.09%   73.19%   +0.09%     
==========================================
  Files          21       21              
  Lines        5208     5186      -22     
==========================================
- Hits         3807     3796      -11     
+ Misses       1401     1390      -11     

Flag	Coverage Δ
unittests	73.19% <76.47%> (+0.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
ogcore/aggregates.py	100.00% <100.00%> (ø)
ogcore/demographics.py	54.42% <90.00%> (+0.78%)	⬆️
ogcore/parameters.py	81.59% <0.00%> (-1.53%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jdebacker]

All local tests pass:

tests/test_SS.py ...........................                             [100%]
tests/test_TPI.py ...........................                            [100%]
tests/test_aggregates.py .........................................       [100%]
tests/test_basic.py ....                                                 [100%]
tests/test_demographics.py ................                              [100%]
tests/test_elliptical_u_est.py .......                                   [100%]
tests/test_execute.py .                                                  [100%]
tests/test_firm.py ..................................................... [ 76%]
................                                                         [100%]
tests/test_fiscal.py ......................                              [100%]
tests/test_household.py ................................................ [ 96%]
..                                                                       [100%]
tests/test_output_plots.py ............................................. [ 95%]
..                                                                       [100%]
tests/test_output_tables.py ..............                               [100%]
tests/test_parameter_plots.py ........................................   [100%]
tests/test_parameter_tables.py .......                                   [100%]
tests/test_parameters.py ..............                                  [100%]
tests/test_pensions.py .................................                 [100%]
tests/test_run_example.py .                                          [100%]
tests/test_run_ogcore.py .                                               [100%]
tests/test_tax.py ...................................                    [100%]
tests/test_txfunc.py ..............................                      [100%]
tests/test_user_inputs.py .........                                      [100%]
tests/test_utils.py .................................................... [ 62%]
...............................                                          [100%]

[jdebacker]

When running the OG-USA example (and updating the demographic parameters), I find the following resource constraint issues:

OG-Core built from the master branch:

array([[ 6.69628260e-04],
       [ 1.77852171e-07],
       [ 1.92564580e-07],
       [ 1.95435097e-07],
       [ 1.93660425e-07],
       [ 1.89483083e-07],
...

Here we see small resource constraint errors, except in the first period. The changes in the prepop branch are intended to fix that by using the actual population parameters (immigration rates, mortality rates) from the period before the time path begins.

OG-Core built from the prepop branch:

array([[ 8.58124249e-06],
       [-9.61640895e-04],
       [ 8.94519232e-04],
       [ 1.81341695e-03],
       [ 2.19727225e-03],
       [ 1.99185083e-03],
       [ 1.50901924e-03],
...

Here we see the initial period resource constraint error become smaller, which is what we wanted to correct here. However, the errors after that are larger and remain in the 1e-3 to 1e-4 range for more than 80 periods. Only until the "fix" period (when immigration rates are forced to be constant) do we get down to errors in the 1e-6 range.

I'm unsure what is going on here. All tests pass on this branch, but those all use the assumption that the pre-TP mortality, immigration, population distribution, and population growth rate as the same as in the first period of the time path. When running OG-USA we are using the empirical values which have a difference. Still, it's odd that the errors come from periods after the first period and persist so long.

cc @rickecon

[jdebacker]

@rickecon Can you review these changes and see if you can pinpoint where the new RC errors are coming from? I've gone over these changes manually several time and have had Codex and Claude Code agents trying to diagnose, but not luck.

Here's what my run of OG-USA with updated demographics looks like (snippet of how set the baseline, which is the only thing of interest here):

"""
    ---------------------------------------------------------------------------
    Run baseline policy
    ---------------------------------------------------------------------------
    """
    # Set up baseline parameterization
    p = Specifications(
        baseline=True,
        num_workers=num_workers,
        baseline_dir=base_dir,
        output_base=base_dir,
    )
    # Update parameters for baseline from default json file
    with importlib.resources.open_text(
        "ogusa", "ogusa_default_parameters.json"
    ) as file:
        defaults = json.load(file)
    p.update_specifications(defaults)
    p.tax_func_type = "HSV"
    p.age_specific = True
    c = Calibration(p, estimate_tax_functions=False, estimate_pop=True, client=client)
    d = c.get_dict()
    # # additional parameters to change
    updated_params = {
        "omega_S_preTP": d["omega_S_preTP"],
        "omega_SS": d["omega_SS"],
        "omega": d["omega"],
        "g_n_ss": d["g_n_ss"],
        "rho": d["rho"],
        "g_n": d["g_n"],
        "imm_rates": d["imm_rates"],
        "g_n_preTP": d["g_n_preTP"],
        "imm_rates_preTP": d["imm_rates_preTP"],
        "rho_preTP": d["rho_preTP"],
    }
    p.update_specifications(updated_params)
    # Run model
    start_time = time.time()
    runner(p, time_path=True, client=client)
    print("run time = ", time.time() - start_time)

[SeaCelo]

@jdebacker @rickecon: I dug into the period-1+ drift you flagged. The boundary fix idea is right — master does have a real period-0 inconsistency — but the drift across the rest of the path is coming from how growth rates are indexed in this branch.

In master, g_n[t] is the growth from t-1 to t. In this branch it's the growth from t to t+1. The aggregate formulas in aggregates.py weren't updated for the new convention, so they read the wrong period's growth rate everywhere.

You can confirm this without running the model. The population accounting itself fails the consistency check across t=1 to fixper, in the exact shape you reported.

The reindex is a one-line revert. But this PR also bundles other changes that aren't strictly part of the boundary fix — the timeline extension, the slicing on returned arrays, the new fixper formula, the pre_pop_dist removal. Rather than untangle them, I'd rather propose a smaller PR with just the boundary fix. I'll open it shortly and link back.

[SeaCelo]

Opened #1117 with the smaller version.

[SeaCelo]

@jdebacker — here's what I think is going on and why you're seeing that drift.

This PR changes one line in get_pop_objs that stores the population growth rates. Master writes them starting at position 1 of the array; this PR writes them starting at position 0. That moves every growth rate one slot earlier in the array, and the boundary value — the growth rate from the year before the model starts to the first model year — that used to sit at the front of the array now lives in the new g_n_preTP field instead.

The drift is caused by aggregates.py not being fully updated to match. The steady-state code was updated to read from g_n_preTP (lines 137 and 192). The transition-path code wasn't. It still pulls the growth rate from the same positions in the array as master, but those positions now hold a different year's value, and the new g_n_preTP field is never read on the transition-path side. So every transition-path year picks up a growth rate that's one year off, and the boundary value never reaches the formula at all.

The transition-path updates would be:

1. aggregates.py line 160 (get_B TPI). The growth-rate slice currently starts at position 1, which was the right place under master's storage. Under this PR's storage, the same slice now starts a year too late, so shift it back to position 0:

   • Replace np.hstack((p.g_n[1 : p.T], p.g_n_ss)) with np.hstack((p.g_n[: p.T - 1], p.g_n_ss)).

2. aggregates.py lines 215 and 222 (get_BQ TPI). The growth-rate slice currently starts at position 0 to pick up the boundary value. Under this PR's storage, the boundary value isn't in the array anymore — it's in g_n_preTP. So prepend g_n_preTP and shift the rest of the slice back by one:

   • Replace 1.0 + p.g_n[: p.T] with 1.0 + np.append(p.g_n_preTP, p.g_n[: p.T - 1]).

get_I reads the growth rate the same way at lines 98, 100, and 105 — those would need the same shift to match the new storage layout.

(The test expected values in test_aggregates.py also encode master's indexing, so they'd need updating to use g_n_preTP at the boundary too.)

cc @rickecon