Skip to content

Add explicit preTP boundary objects to fix aggregate indexing #1117

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
SeaCelo wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

Add explicit preTP boundary objects to fix aggregate indexing #1117

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
33c2a68
Add explicit preTP boundary objects to fix aggregate indexing
SeaCelo Apr 28, 2026
c25e37a
Apply ruff format to test_pretp_boundary
SeaCelo Apr 28, 2026
52a257e
Limit aggregates boundary overrides to mortality and immigration
SeaCelo Apr 28, 2026
798835d
Merge branch 'PSLmodels:master' into feature/pr1073-pretp-boundary
SeaCelo May 7, 2026
6bbbd58
Declare preTP fields in default_parameters.json so downstream callers…
SeaCelo May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions ogcore/aggregates.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,12 +129,14 @@ 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, :])
pop_growth_rate_preTP = getattr(p, "g_n_preTP", p.g_n[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])
pop_growth_rate = p.g_n[0]
imm_extended = np.append(imm_rates_preTP[1:], [0.0])
pop_growth_rate = pop_growth_rate_preTP
else:
part1 = b * np.transpose(p.omega_SS * p.lambdas)
omega_extended = np.append(p.omega_SS[1:], [0.0])
Expand Down Expand Up @@ -186,11 +188,13 @@ 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, :])
pop_growth_rate_preTP = getattr(p, "g_n_preTP", p.g_n[0])
if method == "SS":
if preTP:
omega = p.omega_S_preTP
pop_growth_rate = p.g_n[0]
rho = p.rho[0, :]
pop_growth_rate = pop_growth_rate_preTP
rho = rho_preTP
else:
omega = p.omega_SS
pop_growth_rate = p.g_n_ss
Expand All @@ -205,21 +209,20 @@ 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)
growth = np.hstack((pop_growth_rate_preTP, p.g_n[1 : p.T]))

@jdebacker jdebacker Apr 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@SeaCelo Isn't the appending of $g_{n,preTP}$ inconsistent with Equation 173, which would have ${ g_{n,1},...g_{n,T} }$ in the denominator?

Maybe that equation is incorrect, but I can't see it.

@SeaCelo SeaCelo Apr 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @jdebacker . The change you noted is mostly to have a naming symmetry, not an actual change in the indexing. This is how I was thinking about it in this PR:

The first row of the TPI bequest calculation uses the pre-time-path population distribution. For the row to be internally consistent, the mortality and growth rate in it need to describe the same transition — pre-period into the first model period — so all three inputs refer to the same year boundary. Here is what each input is, before and after this PR:

First-row input Master uses This PR uses
population distribution p.omega_S_preTP (preTP) p.omega_S_preTP (unchanged)
mortality p.rho[0, :] (period 0→1) p.rho_preTP (preTP→0)
population growth rate p.g_n[0] (preTP→0) p.g_n_preTP (preTP→0)

In that reading, g_n_preTP is not an additional term before the sequence in Eq. 173. It is the first term in that sequence, $\tilde{g}_{n,1}$ in the equation's notation: the growth rate from the pre-period into the first TPI period.

In the default flow, g_n_preTP and g_n[0] come from the same source (g_n_path[0] in demographics.py). So np.hstack((g_n_preTP, g_n[1:p.T])) evaluates to the same vector as master's p.g_n[:p.T]. No numerical change to the equation.

The numerical fix in this PR is on rho_preTP (and imm_rates_preTP in get_B) — those are the boundary fields master had pointing at the wrong year. Master's g_n[0] was already correct. The g_n_preTP rename is naming symmetry: all three boundary inputs read as named fields rather than mixing *_preTP with [0] indexing.

@SeaCelo SeaCelo Apr 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One source of why I was confused at first was the math-to-code mapping

I'm reading the theory equation as using model-period notation, where the first TPI bequest period has denominator $\tilde{g}_{n,1}$. In Python, that same transition is zero-indexed and historically lives at p.g_n[0], because it is the growth from the pre-period population into omega[0]. The next transition, first-TPI-period to second-TPI-period, is p.g_n[1].

Theory notation Code object Transition
$\tilde{g}_{n,1}$ g_n_preTP / historically p.g_n[0] preTP → first TPI period
$\tilde{g}_{n,2}$ p.g_n[1] first TPI period → second TPI period
... ... ...
$\tilde{g}_{n,T}$ p.g_n[T-1] (T−1)th TPI period → Tth TPI period

That's why np.hstack((g_n_preTP, g_n[1:p.T])) is the code equivalent of $\lbrace \tilde{g}_{n,1}, \ldots, \tilde{g}_{n,T} \rbrace$.

@jdebacker jdebacker Apr 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@SeaCelo I do think there is a misunderstanding. Let me see if I can clarify.

In the master branch, the indexing should be g_n[0:p.T] = ${ g_{n,1},...g_{n,T} }$. But also in master, since demographics.py doesn't return a g_n_preTP we are assuming that g_n[0]=g_n_preTP.

PR #1073 tries to avoid making that assumption be explicitly computing g_n_preTP. So g_n_preTP is the growth rate in the population from the year before the model start year and g[0]=$g_{n,1}$ is the growth rate in the first period of the model.

@SeaCelo SeaCelo Apr 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @jdebacker. You're right that equation 173's denominator is {g_{n,1}, ..., g_{n,T}} — no preTP term. The hstack I added was a cosmetic rename of p.g_n[0] and didn't change any numbers, but it does make the PR look like it's taking a position on the broader g_n contract, which isn't the goal here.

I'll narrow the change. The PR's only boundary overrides will be on mortality and immigration — the two inputs master was getting from the wrong year. Growth stays exactly as master wrote it. Concretely:

  • Revert the TPI hstack to p.g_n[:p.T]
  • Revert both SS/preTP branches (get_B and get_BQ) to use p.g_n[0] directly
  • Keep g_n_preTP as a returned field from get_pop_objs (symmetric with rho_preTP and imm_rates_preTP), but aggregates.py won't read it

Will push the narrower diff shortly.


if j is not None:
BQ_presum = (b_splus1 * p.lambdas[j]) * (pop * rho)
BQ = BQ_presum.sum(1)
BQ *= (1.0 + r) / (1.0 + p.g_n[: p.T])
BQ *= (1.0 + r) / (1.0 + growth)
else:
BQ_presum = (b_splus1 * np.squeeze(p.lambdas)) * np.tile(
np.reshape(pop * rho, (p.T, p.S, 1)), (1, 1, p.J)
)
BQ = BQ_presum.sum(1)
BQ *= np.tile(
np.reshape((1.0 + r) / (1.0 + p.g_n[: p.T]), (p.T, 1)),
np.reshape((1.0 + r) / (1.0 + growth), (p.T, 1)),
(1, p.J),
)
if p.use_zeta:
Expand Down
77 changes: 77 additions & 0 deletions ogcore/demographics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
17 changes: 17 additions & 0 deletions ogcore/parameters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Loading