diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py index 63b2a94a6..8667edb26 100644 --- a/ogcore/aggregates.py +++ b/ogcore/aggregates.py @@ -129,11 +129,12 @@ def get_B(b, p, method, preTP): B (array_like): aggregate supply of savings """ + imm_rates_preTP = getattr(p, "imm_rates_preTP", p.imm_rates[0, :]) if method == "SS": if preTP: part1 = b * np.transpose(p.omega_S_preTP * p.lambdas) omega_extended = np.append(p.omega_S_preTP[1:], [0.0]) - imm_extended = np.append(p.imm_rates[0, 1:], [0.0]) + imm_extended = np.append(imm_rates_preTP[1:], [0.0]) pop_growth_rate = p.g_n[0] else: part1 = b * np.transpose(p.omega_SS * p.lambdas) @@ -186,11 +187,12 @@ def get_BQ(r, b_splus1, j, p, method, preTP): income group, depending on `use_zeta` value. """ + rho_preTP = getattr(p, "rho_preTP", p.rho[0, :]) if method == "SS": if preTP: omega = p.omega_S_preTP pop_growth_rate = p.g_n[0] - rho = p.rho[0, :] + rho = rho_preTP else: omega = p.omega_SS pop_growth_rate = p.g_n_ss @@ -205,9 +207,7 @@ def get_BQ(r, b_splus1, j, p, method, preTP): pop = np.append( p.omega_S_preTP.reshape(1, p.S), p.omega[: p.T - 1, :], axis=0 ) - rho = np.append( - p.rho[0, :].reshape(1, p.S), p.rho[: p.T - 1, :], axis=0 - ) + rho = np.append(rho_preTP.reshape(1, p.S), p.rho[: p.T - 1, :], axis=0) if j is not None: BQ_presum = (b_splus1 * p.lambdas[j]) * (pop * rho) diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index c36944550..c6f33da4b 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -4874,6 +4874,24 @@ } } }, + "g_n_preTP": { + "title": "Population growth rate from year before model start to start year.", + "description": "Population growth rate from year before model start to start year. Default matches g_n[0] for backward compatibility; downstream calibration pipelines may supply real prior-year data via update_specifications.", + "section_1": "Demographic Parameters", + "notes": "", + "type": "float", + "value": [ + { + "value": 0.0074789353843102225 + } + ], + "validators": { + "range": { + "min": -1.0, + "max": 1.0 + } + } + }, "imm_rates": { "title": "Immigration rates over the time path.", "description": "Immigration rates over the time path.", @@ -4897,6 +4915,25 @@ } } }, + "imm_rates_preTP": { + "title": "Immigration rates by age in the period before model start.", + "description": "Immigration rates by age in the period before model start. Default matches imm_rates[0, :] for backward compatibility; downstream calibration pipelines may supply real prior-year data via update_specifications.", + "section_1": "Demographic Parameters", + "notes": "", + "type": "float", + "number_dims": 1, + "value": [ + { + "value": [0.0062112432351020715, 0.006695391269304812, 0.0072416059873147915, 0.007590577071078806, 0.007646147746699469, 0.00751951109796381, 0.007066814891074261, 0.006564297706210947, 0.006073197537540996, 0.005507000333471546, 0.004947327413337241, 0.004549582594943434, 0.004234250012419245, 0.003974039945508412, 0.003677521830526146, 0.0033802128693471364, 0.0030051433785756396, 0.002719163436517305, 0.002416866691604522, 0.0021354853318490724, 0.001868509517440961, 0.0016393085543080167, 0.001457134576126563, 0.0013421224668364702, 0.0012716481058643412, 0.0009881382611250649, 0.0008422678904918091, 0.0007978012112857263, 0.000663231804968117, 0.000487948388872445, 0.0004329790187167157, 0.0005600407822959095, 0.00045536991638676395, 0.00041371721643232613, 0.0004550292264560817, 0.00045210410909610134, 0.0005586451762040458, 0.0006146312763669243, 0.0005666841212952642, 0.0005114080149799634, 0.00046896874841981473, 0.00048654173068420997, 0.00018708464449557118, -2.3352594037780405e-05, 0.0012131573735734234, 0.0010532417702141814, 8.290111290308402e-07, -0.0001804564659585655, 0.0006446968377394005, 0.0010883872003127893, 0.0006181376601826033, 0.0004414738263082719, 0.0006101868033593997, 0.0006586334230713459, 0.0005632074039269039, 0.0005342957172748685, 0.0006001149929916475, 0.0006290507627871093, 0.0008430988793991255, 0.000920493619977935, 0.0009689086963466627, 0.0012637402224869932, 0.0016061695551225035, 0.0017284030417787173, 0.0015042954657816972, 0.001034592968885015, 0.0008014496492442746, 0.000928625039946825, 0.0016099333592955103, 0.0022322710222315374, 0.0035345112500425146, 0.005324802746303509, 0.007846226075289809, 0.01082203864395205, 0.013081043435610118, 0.01645197200918262, 0.018801001661434714, 0.021272590494471732, 0.021795511437529547, 0.019240097157757045] + } + ], + "validators": { + "range": { + "min": -1.0, + "max": 1.0 + } + } + }, "rho": { "title": "Age-specific mortality rates.", "description": "Age-specific mortality rates.", @@ -4918,6 +4955,25 @@ } } }, + "rho_preTP": { + "title": "Age-specific mortality rates in the period before model start.", + "description": "Age-specific mortality rates in the period before model start. Default matches rho[0, :] for backward compatibility; downstream calibration pipelines may supply real prior-year data via update_specifications.", + "section_1": "Demographic Parameters", + "notes": "", + "type": "float", + "number_dims": 1, + "value": [ + { + "value": [0.0007328533962063233, 0.0008171197422418564, 0.0008788729669386708, 0.0009116525208591186, 0.0009229802325645409, 0.0009273406814831464, 0.0009357102393970917, 0.0009500953703120985, 0.0009759735565497696, 0.0010103554581823992, 0.0010497459745855142, 0.0010896505583020133, 0.0011290564585360041, 0.0011654456037331062, 0.0012043239258441307, 0.0012511773865916398, 0.0013110045346917953, 0.0013828255727224015, 0.0014696353786847194, 0.0015729202258729647, 0.0016901599123994027, 0.0018253151315732463, 0.00198785868083029, 0.002182248418529209, 0.0024054583493769988, 0.002649495698085591, 0.002912359902066397, 0.003196521891964843, 0.0035014404113965503, 0.0038265853053018883, 0.004175977199939118, 0.004546000079867052, 0.004929386691983617, 0.005323440679881664, 0.005732162610849301, 0.00618294920155571, 0.0066650630096521946, 0.007142892752724284, 0.007605334723393753, 0.008079293606839011, 0.008607183802693097, 0.009219540427197637, 0.009916977317307896, 0.01071379466147504, 0.011618680680630211, 0.012651175630041145, 0.013806200722411788, 0.015069231422104257, 0.01643990870680101, 0.017946285365270986, 0.019682839990789902, 0.02164309127204489, 0.02374663347580741, 0.025983719293213747, 0.028428273440626484, 0.03125859803249409, 0.03451112898875497, 0.038087062505788305, 0.04199247322144961, 0.04632683864241138, 0.05133101803558171, 0.05704185782117355, 0.06332015480431863, 0.07017696315508493, 0.0777688358372085, 0.0862903314228447, 0.09590618992081834, 0.1067242554403307, 0.11879857491346835, 0.13212525643676043, 0.14668282278042066, 0.16243247757898882, 0.17932696648992474, 0.19731341445765527, 0.21633944779689218, 0.23535763270336796, 0.2540598444887924, 0.2721191714344474, 0.28919517981339793, 1.0] + } + ], + "validators": { + "range": { + "min": 0.0, + "max": 1.0 + } + } + }, "etr_params": { "title": "Effective tax rate function parameters.", "description": "Effective tax rate function parameters.", diff --git a/ogcore/demographics.py b/ogcore/demographics.py index 13990f692..3b6ca19d2 100644 --- a/ogcore/demographics.py +++ b/ogcore/demographics.py @@ -719,6 +719,7 @@ def get_pop_objs( infer_pop=False, pop_dist=None, pre_pop_dist=None, + pre_mort_rates=None, country_id=UN_COUNTRY_CODE, initial_data_year=START_YEAR - 1, final_data_year=START_YEAR + 2, @@ -752,6 +753,12 @@ def get_pop_objs( pre_pop_dist (array_like): user provided population distribution for the year before the initial year for calibration, length E+S + pre_mort_rates (array_like): user provided mortality rates for + the year before the initial year for calibration, length E+S. + When supplied, used to populate ``rho_preTP`` and to invert + ``imm_rates_preTP`` against ``pre_pop_dist`` -> ``pop_2D[0]``. + When None, fetched from UN data if ``mort_rates`` was also + None; otherwise falls back to the first row of ``mort_rates``. country_id (str): country id for UN data initial_data_year (int): initial year of data to use (not relevant if have user provided data) @@ -778,6 +785,9 @@ def get_pop_objs( path, length T + S """ + user_fert_rates = fert_rates is not None + user_mort_rates = mort_rates is not None + user_imm_rates = imm_rates is not None # TODO: this function does not generalize with T. # It assumes one model period is equal to one calendar year in the # time dimesion (it does adjust for S, however) @@ -1160,6 +1170,70 @@ def get_pop_objs( ), axis=0, ) + rho_preTP = mort_rates_S[0, :] + g_n_preTP = g_n_path[0] + imm_rates_preTP = imm_rates_mat[0, :] + + # Compute a genuine prior-year boundary transition when we have the + # data to do so. The path arrays (rho, imm_rates, g_n, omega) are not + # touched -- only the *_preTP fields, which feed aggregates.py's + # boundary calculation. ``imm_rates_preTP`` is derived by inverting the + # level-space recurrence pre_pop -> pop_2D[0] via ``get_imm_rates``, + # so the boundary identity holds at machine precision for whatever + # mortality is supplied. + if pre_mort_rates is not None: + pre_mort_rates_full = np.asarray(pre_mort_rates) + assert pre_mort_rates_full.shape == (E + S,) + prior_year = initial_data_year - 1 + # Fertility / infmort only enter ``imm_rates_preTP[0]`` (newborn + # age), which aggregates.py never reads. Pass zeros to keep + # ``get_imm_rates`` happy without introducing extra kwargs. + imm_rates_preTP_full = get_imm_rates( + E + S, + min_age, + max_age, + fert_rates=np.zeros((1, E + S)), + mort_rates=pre_mort_rates_full.reshape(1, E + S), + infmort_rates=np.zeros(1), + pop_dist=np.vstack((pre_pop_EpS, pop_2D[0, :])), + country_id=country_id, + start_year=prior_year, + end_year=prior_year, + ) + rho_preTP = pre_mort_rates_full[E:] + imm_rates_preTP = imm_rates_preTP_full[0, E:] + elif not (user_fert_rates or user_mort_rates or user_imm_rates): + prior_year = initial_data_year - 1 + fert_rates_preTP = get_fert( + E + S, + min_age, + max_age, + country_id, + start_year=prior_year, + end_year=prior_year, + ) + mort_rates_preTP_full, infmort_rates_preTP = get_mort( + E + S, + min_age, + max_age, + country_id, + start_year=prior_year, + end_year=prior_year, + ) + imm_rates_preTP_full = get_imm_rates( + E + S, + min_age, + max_age, + fert_rates=fert_rates_preTP, + mort_rates=mort_rates_preTP_full, + infmort_rates=infmort_rates_preTP, + pop_dist=np.vstack((pre_pop_EpS, pop_2D[0, :])), + country_id=country_id, + start_year=prior_year, + end_year=prior_year, + ) + rho_preTP = mort_rates_preTP_full[0, E:] + imm_rates_preTP = imm_rates_preTP_full[0, E:] if GraphDiag: # Check whether original SS population distribution is close to @@ -1315,6 +1389,9 @@ def get_pop_objs( "g_n": g_n_path, "imm_rates": imm_rates_mat, "omega_S_preTP": omega_S_preTP, + "rho_preTP": rho_preTP, + "g_n_preTP": g_n_preTP, + "imm_rates_preTP": imm_rates_preTP, } return pop_dict diff --git a/ogcore/parameters.py b/ogcore/parameters.py index 57df782ca..fbdb43fc0 100644 --- a/ogcore/parameters.py +++ b/ogcore/parameters.py @@ -332,6 +332,23 @@ def compute_default_params(self): ) self.omega_S_preTP = self.omega_SS + # Preserve legacy semantics by default, but expose explicit + # boundary-period demographic objects when they are not provided. + if (not hasattr(self, "omega_S_preTP")) or ( + np.asarray(self.omega_S_preTP).shape[-1] != self.S + ): + self.omega_S_preTP = self.omega[0, :] + if not hasattr(self, "g_n_preTP"): + self.g_n_preTP = self.g_n[0] + if (not hasattr(self, "rho_preTP")) or ( + np.asarray(self.rho_preTP).shape[-1] != self.S + ): + self.rho_preTP = self.rho[0, :] + if (not hasattr(self, "imm_rates_preTP")) or ( + np.asarray(self.imm_rates_preTP).shape[-1] != self.S + ): + self.imm_rates_preTP = self.imm_rates[0, :] + # Create time series of stationarized UBI transfers self.ubi_nom_array = self.get_ubi_nom_objs() diff --git a/tests/test_pretp_boundary.py b/tests/test_pretp_boundary.py new file mode 100644 index 000000000..adc8d1aa0 --- /dev/null +++ b/tests/test_pretp_boundary.py @@ -0,0 +1,506 @@ +import os +import numpy as np +from types import SimpleNamespace + +from ogcore import aggregates as aggr +from ogcore import demographics +from ogcore.parameters import Specifications + + +def make_mock_params(): + """ + Create a small parameter object that isolates the pre-time-path + boundary from the later TPI path. + """ + p = SimpleNamespace() + p.S = 3 + p.J = 1 + p.T = 4 + p.use_zeta = False + p.lambdas = np.array([1.0]).reshape(1, 1) + p.omega_S_preTP = np.array([0.2, 0.3, 0.5]) + p.omega_SS = np.array([0.25, 0.35, 0.4]) + p.omega = np.array( + [ + [0.3, 0.3, 0.4], + [0.25, 0.35, 0.4], + [0.2, 0.4, 0.4], + [0.2, 0.4, 0.4], + ] + ) + # Explicit pre-period objects intentionally differ from the path + # objects so the boundary bug is visible in a small unit test. + p.imm_rates_preTP = np.array([0.11, 0.12, 0.13]) + p.rho_preTP = np.array([0.41, 0.42, 0.43]) + p.g_n_preTP = 0.21 + p.imm_rates = np.array( + [ + [0.01, 0.02, 0.03], + [0.04, 0.05, 0.06], + [0.07, 0.08, 0.09], + [0.1, 0.11, 0.12], + ] + ) + p.rho = np.array( + [ + [0.11, 0.12, 0.13], + [0.21, 0.22, 0.23], + [0.31, 0.32, 0.33], + [0.41, 0.42, 0.43], + ] + ) + p.g_n = np.array([0.01, 0.06, 0.07, 0.08]) + p.g_n_ss = 0.09 + p.g_y = 0.0 + p.delta = 0.1 + p.delta_g = 0.05 + p.infra_investment_leakage_rate = 0.0 + p.alpha_RM_1 = 0.05 + p.alpha_RM_T = 0.05 + p.g_RM = np.array([0.02, 0.02, 0.02, 0.02]) + p.tG1 = 2 + p.tG2 = 4 + return p + + +def step_population(pop_t, fert_t, mort_t, infmort_t, imm_t): + """ + One-period demographic transition used to build a consistent toy path. + """ + pop_tp1 = np.zeros_like(pop_t) + newborns = np.dot(fert_t, pop_t) + pop_tp1[0] = (1.0 - infmort_t) * newborns + imm_t[0] * pop_t[0] + pop_tp1[1:] = pop_t[:-1] * (1.0 - mort_t[:-1]) + pop_t[1:] * imm_t[1:] + return pop_tp1 + + +def test_pretp_ss_aggregates_use_explicit_boundary_objects(): + """ + The preTP SS accounting objects should use the explicit pre-period + demographic rates, not the first row of the TPI path. + """ + p = make_mock_params() + b = np.array([[1.0], [2.0], [3.0]]) + r = 0.04 + + omega_shift = np.append(p.omega_S_preTP[1:], [0.0]) + imm_shift = np.append(p.imm_rates_preTP[1:], [0.0]) + expected_B = ( + b[:, 0] * (p.omega_S_preTP + (omega_shift * imm_shift)) + ).sum() / (1.0 + p.g_n[0]) + expected_BQ = ( + (p.omega_S_preTP * p.rho_preTP * b[:, 0]).sum() + * (1.0 + r) + / (1.0 + p.g_n[0]) + ) + + assert np.allclose(aggr.get_B(b, p, "SS", True), expected_B) + assert np.allclose(aggr.get_BQ(r, b, None, p, "SS", True), expected_BQ) + + +def test_pretp_tpi_bequests_change_only_boundary_period(): + """ + A boundary-only fix should change only the prepended preTP element of + the TPI bequest path. Later periods should keep the existing path + indexing contract. + """ + p = make_mock_params() + r = np.array([0.01, 0.02, 0.03, 0.04]) + b_splus1 = np.array( + [ + [[1.0], [2.0], [3.0]], + [[1.5], [2.5], [3.5]], + [[2.0], [3.0], [4.0]], + [[2.5], [3.5], [4.5]], + ] + ) + + expected = np.zeros(p.T) + expected[0] = ( + (p.omega_S_preTP * p.rho_preTP * b_splus1[0, :, 0]).sum() + * (1.0 + r[0]) + / (1.0 + p.g_n[0]) + ) + for t in range(1, p.T): + expected[t] = ( + (p.omega[t - 1, :] * p.rho[t - 1, :] * b_splus1[t, :, 0]).sum() + * (1.0 + r[t]) + / (1.0 + p.g_n[t]) + ) + + assert np.allclose( + np.asarray(aggr.get_BQ(r, b_splus1, None, p, "TPI", False)).reshape( + -1 + ), + expected, + ) + + +def test_pretp_fields_are_optional_for_legacy_callers(): + """ + Existing callers that do not yet provide explicit preTP demographic + objects should keep the legacy behavior. + """ + p = make_mock_params() + del p.imm_rates_preTP + del p.rho_preTP + del p.g_n_preTP + + b = np.array([[1.0], [2.0], [3.0]]) + r = 0.04 + + omega_shift = np.append(p.omega_S_preTP[1:], [0.0]) + imm_shift = np.append(p.imm_rates[0, 1:], [0.0]) + expected_B = ( + b[:, 0] * (p.omega_S_preTP + (omega_shift * imm_shift)) + ).sum() + expected_B /= 1.0 + p.g_n[0] + expected_BQ = ( + (p.omega_S_preTP * p.rho[0, :] * b[:, 0]).sum() + * (1.0 + r) + / (1.0 + p.g_n[0]) + ) + + assert np.allclose(aggr.get_B(b, p, "SS", True), expected_B) + assert np.allclose(aggr.get_BQ(r, b, None, p, "SS", True), expected_BQ) + + +def test_get_pop_objs_pins_growth_timing_contract(): + """ + The existing downstream code expects the master storage convention: + `g_n[0]` is preTP->0 growth and `g_n[1]` is 0->1 growth. + """ + E = 1 + S = 4 + T = 8 + fert_rates = np.array( + [ + [0.20, 0.0, 0.0, 0.0, 0.0], + [0.25, 0.0, 0.0, 0.0, 0.0], + [0.25, 0.0, 0.0, 0.0, 0.0], + ] + ) + mort_rates = np.array( + [ + [0.05, 0.10, 0.15, 0.20, 1.0], + [0.06, 0.11, 0.16, 0.21, 1.0], + [0.06, 0.11, 0.16, 0.21, 1.0], + ] + ) + infmort_rates = np.array([0.02, 0.03, 0.03]) + imm_rates = np.array( + [ + [0.30, 0.31, 0.32, 0.33, 0.34], + [0.01, 0.02, 0.03, 0.04, 0.05], + [0.01, 0.02, 0.03, 0.04, 0.05], + ] + ) + pre_pop = np.array([50.0, 40.0, 30.0, 20.0, 10.0]) + pop_dist = np.zeros((4, E + S)) + pop_dist[0] = np.array([24.8, 59.9, 45.6, 32.1, 19.4]) + pop_dist[1] = step_population( + pop_dist[0], + fert_rates[0], + mort_rates[0], + infmort_rates[0], + imm_rates[0], + ) + pop_dist[2] = step_population( + pop_dist[1], + fert_rates[1], + mort_rates[1], + infmort_rates[1], + imm_rates[1], + ) + pop_dist[3] = step_population( + pop_dist[2], + fert_rates[2], + mort_rates[2], + infmort_rates[2], + imm_rates[2], + ) + + pop_dict = demographics.get_pop_objs( + E=E, + S=S, + T=T, + min_age=0, + max_age=4, + fert_rates=fert_rates, + mort_rates=mort_rates, + infmort_rates=infmort_rates, + imm_rates=imm_rates, + infer_pop=False, + pop_dist=pop_dist, + pre_pop_dist=pre_pop, + initial_data_year=2020, + final_data_year=2022, + GraphDiag=False, + ) + + pre_growth = (pop_dist[0, -S:].sum() - pre_pop[-S:].sum()) / pre_pop[ + -S: + ].sum() + zero_to_one_growth = ( + pop_dist[1, -S:].sum() - pop_dist[0, -S:].sum() + ) / pop_dist[0, -S:].sum() + + assert np.allclose( + pop_dict["omega_S_preTP"], pre_pop[-S:] / pre_pop[-S:].sum() + ) + assert np.allclose(pop_dict["g_n_preTP"], pre_growth) + assert np.allclose(pop_dict["g_n"][0], pre_growth) + assert np.allclose(pop_dict["g_n"][1], zero_to_one_growth) + # Legacy: with user-supplied rates and no pre_mort_rates kwarg, the + # boundary objects alias the first in-window row. + assert np.allclose(pop_dict["rho_preTP"], mort_rates[0, E:]) + assert np.allclose(pop_dict["imm_rates_preTP"], imm_rates[0, E:]) + + +def test_get_pop_objs_fetches_distinct_prior_year_boundary_rates(monkeypatch): + """ + If OG-Core is sourcing demographic rates itself, it should populate + the preTP boundary objects from a genuine prior-year row rather than + aliasing the first in-window path row. + """ + E = 1 + S = 4 + T = 8 + initial_data_year = 2020 + final_data_year = 2022 + pop_dist = np.array( + [ + [24.8, 59.9, 45.6, 32.1, 19.4], + [12.3008, 42.129, 68.502, 49.353, 32.276], + [3.105952, 12.405332, 39.54987, 59.5158, 40.60267], + [0.78425288, 3.16770152, 12.22724158, 35.6025228, 49.0476155], + ] + ) + pre_pop = np.array([50.0, 40.0, 30.0, 20.0, 10.0]) + main_fert = np.array( + [ + [0.20, 0.0, 0.0, 0.0, 0.0], + [0.25, 0.0, 0.0, 0.0, 0.0], + [0.25, 0.0, 0.0, 0.0, 0.0], + ] + ) + prior_fert = np.array([[0.15, 0.0, 0.0, 0.0, 0.0]]) + main_mort = np.array( + [ + [0.05, 0.10, 0.15, 0.20, 1.0], + [0.06, 0.11, 0.16, 0.21, 1.0], + [0.06, 0.11, 0.16, 0.21, 1.0], + ] + ) + prior_mort = np.array([[0.01, 0.02, 0.03, 0.04, 1.0]]) + main_infmort = np.array([0.02, 0.03, 0.03]) + prior_infmort = np.array([0.005]) + main_imm = np.array( + [ + [0.30, 0.31, 0.32, 0.33, 0.34], + [0.01, 0.02, 0.03, 0.04, 0.05], + [0.01, 0.02, 0.03, 0.04, 0.05], + ] + ) + prior_imm = np.array([[0.20, 0.21, 0.22, 0.23, 0.24]]) + + def _resolve_start_year(args, kwargs, positional_index): + if "start_year" in kwargs: + return kwargs["start_year"] + if len(args) > positional_index: + return args[positional_index] + raise TypeError("start_year not supplied") + + def fake_get_fert(*args, **kwargs): + # get_fert(totpers, min_age, max_age, country_id, start_year, ...) + start_year = _resolve_start_year(args, kwargs, 4) + if start_year == initial_data_year - 1: + return prior_fert + return main_fert + + def fake_get_mort(*args, **kwargs): + # get_mort(totpers, min_age, max_age, country_id, start_year, ...) + start_year = _resolve_start_year(args, kwargs, 4) + if start_year == initial_data_year - 1: + return prior_mort, prior_infmort + return main_mort, main_infmort + + def fake_get_imm_rates(*args, **kwargs): + # get_imm_rates(totpers, min_age, max_age, fert_rates, mort_rates, + # infmort_rates, pop_dist, country_id, start_year, ...) + start_year = _resolve_start_year(args, kwargs, 8) + if start_year == initial_data_year - 1: + return prior_imm + return main_imm + + monkeypatch.setattr(demographics, "get_fert", fake_get_fert) + monkeypatch.setattr(demographics, "get_mort", fake_get_mort) + monkeypatch.setattr(demographics, "get_imm_rates", fake_get_imm_rates) + + pop_dict = demographics.get_pop_objs( + E=E, + S=S, + T=T, + min_age=0, + max_age=4, + fert_rates=None, + mort_rates=None, + infmort_rates=None, + imm_rates=None, + infer_pop=False, + pop_dist=pop_dist, + pre_pop_dist=pre_pop, + initial_data_year=initial_data_year, + final_data_year=final_data_year, + GraphDiag=False, + ) + + assert np.allclose(pop_dict["rho_preTP"], prior_mort[0, E:]) + assert np.allclose(pop_dict["imm_rates_preTP"], prior_imm[0, E:]) + assert np.allclose(pop_dict["rho"][0], main_mort[0, E:]) + assert np.allclose(pop_dict["imm_rates"][0], main_imm[0, E:]) + + +def test_pre_mort_rates_kwarg_collapses_boundary_identity(): + """ + Supplying ``pre_mort_rates`` to ``get_pop_objs`` should populate the + boundary objects from genuine prior-year data so that the preTP -> 0 + demographic identity holds at machine precision (s in [1, S-1]). + + The synthetic universe has rates that vary year over year, so the + legacy fallback (preTP rates = period-0 rates) leaves a residual + proportional to the year-over-year rate change. + """ + E = 0 + S = 5 + T = 8 + initial_data_year = 2020 + final_data_year = 2022 + + mort_yrm1 = np.array([0.05, 0.10, 0.15, 0.30, 1.00]) + imm_yrm1 = np.array([0.02, 0.03, 0.04, 0.05, 0.06]) + fert_yrm1 = np.array([0.20, 0.30, 0.10, 0.00, 0.00]) + infmort_yrm1 = 0.04 + + main_mort = np.array( + [ + [0.07, 0.13, 0.18, 0.34, 1.00], # year 2020 + [0.08, 0.14, 0.19, 0.35, 1.00], # year 2021 + [0.08, 0.14, 0.19, 0.35, 1.00], # year 2022 + ] + ) + main_fert = np.array( + [ + [0.18, 0.28, 0.09, 0.00, 0.00], + [0.17, 0.27, 0.08, 0.00, 0.00], + [0.17, 0.27, 0.08, 0.00, 0.00], + ] + ) + main_infmort = np.array([0.03, 0.03, 0.03]) + main_imm = np.array( + [ + [0.05, 0.06, 0.07, 0.08, 0.09], + [0.05, 0.06, 0.07, 0.08, 0.09], + [0.05, 0.06, 0.07, 0.08, 0.09], + ] + ) + + pre_pop = np.array([100.0, 80.0, 60.0, 40.0, 20.0]) + pop0 = step_population( + pre_pop, fert_yrm1, mort_yrm1, infmort_yrm1, imm_yrm1 + ) + pop1 = step_population( + pop0, main_fert[0], main_mort[0], main_infmort[0], main_imm[0] + ) + pop2 = step_population( + pop1, main_fert[1], main_mort[1], main_infmort[1], main_imm[1] + ) + pop3 = step_population( + pop2, main_fert[2], main_mort[2], main_infmort[2], main_imm[2] + ) + pop_dist = np.vstack((pop0, pop1, pop2, pop3)) + + pop_dict = demographics.get_pop_objs( + E=E, + S=S, + T=T, + min_age=0, + max_age=4, + fert_rates=main_fert, + mort_rates=main_mort, + infmort_rates=main_infmort, + imm_rates=main_imm, + infer_pop=False, + pop_dist=pop_dist, + pre_pop_dist=pre_pop, + pre_mort_rates=mort_yrm1, + initial_data_year=initial_data_year, + final_data_year=final_data_year, + GraphDiag=False, + ) + + # Boundary identity, S-space, s in [1, S-1]: + # omega[0, s] * (1 + g_n_preTP) + # == (1 - rho_preTP[s-1]) * omega_S_preTP[s-1] + # + imm_rates_preTP[s] * omega_S_preTP[s] + omega_0 = pop_dict["omega"][0, :] + omega_S_preTP = pop_dict["omega_S_preTP"] + g_n_preTP = pop_dict["g_n_preTP"] + rho_preTP = pop_dict["rho_preTP"] + imm_preTP = pop_dict["imm_rates_preTP"] + + lhs = omega_0[1:] * (1.0 + g_n_preTP) + rhs = (1.0 - rho_preTP[:-1]) * omega_S_preTP[:-1] + imm_preTP[ + 1: + ] * omega_S_preTP[1:] + assert np.allclose(lhs, rhs, atol=1e-13) + assert np.allclose(rho_preTP, mort_yrm1[E:]) + + +def test_get_pop_objs_dict_passes_through_update_specifications(): + """ + The dict returned by ``get_pop_objs`` must be accepted by + ``Specifications.update_specifications`` without paramtools rejecting any + of the new ``*_preTP`` fields. This is the calibration pipeline that + downstream country repos (OG-USA, OG-ETH, OG-ZAF, etc.) rely on: + + d = Calibration(p).get_dict() + p.update_specifications(d) + + Regression guard for the integration path. paramtools requires every + name in ``d`` to be declared in ``default_parameters.json``; if any of + the new ``rho_preTP``, ``imm_rates_preTP``, ``g_n_preTP`` entries is + missing, this test fails immediately rather than after the downstream + runner does. + """ + data_dir = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "test_io_data" + ) + + def ld(name): + return np.loadtxt(os.path.join(data_dir, name), delimiter=",") + + pop_dict = demographics.get_pop_objs( + E=20, + S=80, + T=320, + min_age=0, + max_age=99, + fert_rates=ld("fert_rates.csv"), + mort_rates=ld("mort_rates.csv"), + infmort_rates=ld("infmort_rates.csv"), + imm_rates=ld("immigration_rates.csv"), + infer_pop=True, + pop_dist=ld("population_distribution.csv")[0, :].reshape(1, 100), + pre_pop_dist=ld("pre_period_population_distribution.csv"), + initial_data_year=2023, + final_data_year=2024, + GraphDiag=False, + ) + + p = Specifications() + p.update_specifications(pop_dict) + + assert np.allclose(p.g_n_preTP, pop_dict["g_n_preTP"]) + assert np.allclose(p.rho_preTP, pop_dict["rho_preTP"]) + assert np.allclose(p.imm_rates_preTP, pop_dict["imm_rates_preTP"])