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 72.94%. Comparing base (c24e756) to head (64f6f9e).
⚠️ Report is 4 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   72.85%   72.94%   +0.09%     
==========================================
  Files          21       21              
  Lines        5201     5179      -22     
==========================================
- Hits         3789     3778      -11     
+ Misses       1412     1401      -11     

Flag	Coverage Δ
unittests	72.94% <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.