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.01%. Comparing base (83227e3) to head (575bb01).

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.92%   73.01%   +0.09%     
==========================================
  Files          21       21              
  Lines        5122     5100      -22     
==========================================
- Hits         3735     3724      -11     
+ Misses       1387     1376      -11     

Flag	Coverage Δ
unittests	73.01% <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)