Skip to content

Commit 46f583e

Browse files
jdebackerjdebacker
committed
format
1 parent 4d7c59a commit 46f583e

3 files changed

Lines changed: 76 additions & 48 deletions

File tree

‎ogcore/SS.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -515,8 +515,7 @@ def inner_loop(outer_loop_vars, p, client):
515515
# (S, J) nssmat broadcasts into an (S, S) outer product, scaling the
516516
# income sum by S. p.e[-1, :, :] is already (S, J).
517517
average_income_model = (
518-
(new_r_p * b_s + new_w * p.e[-1, :, :] * nssmat)
519-
* p.omega_SS
518+
(new_r_p * b_s + new_w * p.e[-1, :, :] * nssmat) * p.omega_SS
520519
).sum()
521520
if p.baseline:
522521
new_factor = p.mean_income_data / average_income_model

‎ogcore/aggregates.py‎

Lines changed: 39 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ def get_L(n, p, method):
4040
L_presum = p.e[-1, :, :] * p.omega_SS * n
4141
L = L_presum.sum()
4242
elif method == "TPI":
43-
L_presum = (n * (p.e * p.omega[: p.T, :, :]))
43+
L_presum = n * (p.e * p.omega[: p.T, :, :])
4444
L = L_presum.sum(1).sum(1)
4545
return L
4646

@@ -67,28 +67,28 @@ def get_I(b_splus1, K_p1, K, p, method):
6767
6868
"""
6969
if method == "SS":
70-
omega_extended = np.append(p.omega_SS[1:, :], [0.0], axis=0)
71-
imm_extended = np.append(p.imm_rates[-1, 1:, :], [0.0], axis=0)
72-
part2 = (
73-
(
74-
b_splus1 * omega_extended * imm_extended
75-
).sum()
76-
) / (1 + p.g_n_ss)
70+
omega_extended = np.append(
71+
p.omega_SS[1:, :], np.zeros((1, p.J)), axis=0
72+
)
73+
imm_extended = np.append(
74+
p.imm_rates[-1, 1:, :], np.zeros((1, p.J)), axis=0
75+
)
76+
part2 = ((b_splus1 * omega_extended * imm_extended).sum()) / (
77+
1 + p.g_n_ss
78+
)
7779
aggI = (1 + p.g_n_ss) * np.exp(p.g_y) * (K_p1 - part2) - (
7880
1.0 - p.delta
7981
) * K
8082
elif method == "TPI":
81-
omega_shift = np.append(p.omega[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1)
83+
omega_shift = np.append(
84+
p.omega[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1
85+
)
8286
imm_shift = np.append(
8387
p.imm_rates[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1
8488
)
85-
part2 = (
86-
(
87-
b_splus1 * imm_shift * omega_shift
88-
)
89-
.sum(1)
90-
.sum(1)
91-
) / (1 + np.squeeze(np.hstack((p.g_n[: p.T - 1], p.g_n_ss))))
89+
part2 = ((b_splus1 * imm_shift * omega_shift).sum(1).sum(1)) / (
90+
1 + np.squeeze(np.hstack((p.g_n[: p.T - 1], p.g_n_ss)))
91+
)
9292
aggI = (
9393
1 + np.squeeze(np.hstack((p.g_n[: p.T - 1], p.g_n_ss)))
9494
) * np.exp(p.g_y) * (K_p1 - part2) - (1.0 - p.delta) * K
@@ -123,21 +123,31 @@ def get_B(b, p, method, preTP):
123123
if method == "SS":
124124
if preTP:
125125
part1 = b * p.omega_S_preTP
126-
omega_extended = np.append(p.omega_S_preTP[1:, :], [0.0], axis=0)
127-
imm_extended = np.append(p.imm_rates_preTP[1:, :], [0.0], axis=0)
126+
omega_extended = np.append(
127+
p.omega_S_preTP[1:, :], np.zeros((1, p.J)), axis=0
128+
)
129+
imm_extended = np.append(
130+
p.imm_rates_preTP[1:, :], np.zeros((1, p.J)), axis=0
131+
)
128132
pop_growth_rate = p.g_n_preTP
129133
else:
130134
part1 = b * p.omega_SS
131-
omega_extended = np.append(p.omega_SS[1:, :], [0.0], axis=0)
132-
imm_extended = np.append(p.imm_rates[-1, 1:, :], [0.0], axis=0)
135+
omega_extended = np.append(
136+
p.omega_SS[1:, :], np.zeros((1, p.J)), axis=0
137+
)
138+
imm_extended = np.append(
139+
p.imm_rates[-1, 1:, :], np.zeros((1, p.J)), axis=0
140+
)
133141
pop_growth_rate = p.g_n_ss
134142
part2 = b * omega_extended * imm_extended
135143
B_presum = part1 + part2
136144
B = B_presum.sum()
137145
B /= 1.0 + pop_growth_rate
138146
elif method == "TPI":
139147
part1 = b * p.omega[: p.T, :, :]
140-
omega_shift = np.append(p.omega[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1)
148+
omega_shift = np.append(
149+
p.omega[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1
150+
)
141151
imm_shift = np.append(
142152
p.imm_rates[: p.T, 1:, :], np.zeros((p.T, 1, p.J)), axis=1
143153
)
@@ -183,21 +193,23 @@ def get_BQ(r, b_splus1, j, p, method, preTP):
183193
pop_growth_rate = p.g_n_ss
184194
rho = p.rho[-1, :]
185195
if j is not None:
186-
BQ_presum = omega * rho * b_splus1
196+
BQ_presum = omega[:, j] * rho[:, j] * b_splus1
187197
else:
188198
BQ_presum = omega * rho * b_splus1
189199
BQ = BQ_presum.sum(0)
190200
BQ *= (1.0 + r) / (1.0 + pop_growth_rate)
191201
elif method == "TPI":
192202
pop = np.append(
193-
p.omega_S_preTP.reshape(1, p.S), p.omega[: p.T - 1, :], axis=0
203+
p.omega_S_preTP.reshape(1, p.S, p.J),
204+
p.omega[: p.T - 1, :, :],
205+
axis=0,
194206
)
195207
rho = np.append(
196-
p.rho_preTP.reshape(1, p.S), p.rho[: p.T - 1, :], axis=0
208+
p.rho_preTP.reshape(1, p.S, p.J), p.rho[: p.T - 1, :, :], axis=0
197209
)
198210

199211
if j is not None:
200-
BQ_presum = b_splus1 * pop[:, j] * rho[:, j]
212+
BQ_presum = b_splus1 * pop[:, :, j] * rho[:, :, j]
201213
BQ = BQ_presum.sum(1)
202214
BQ *= (1.0 + r) / (1.0 + np.append(p.g_n_preTP, p.g_n[: p.T - 1]))
203215
else:
@@ -289,13 +301,9 @@ def get_C(c, p, method):
289301
"""
290302

291303
if method == "SS":
292-
aggC = (
293-
(c * p.omega_SS).sum(-1).sum(-1)
294-
)
304+
aggC = (c * p.omega_SS).sum(-1).sum(-1)
295305
elif method == "TPI":
296-
aggC = (
297-
(c * p.omega[: p.T, :, :]).sum(-1).sum(-1)
298-
)
306+
aggC = (c * p.omega[: p.T, :, :]).sum(-1).sum(-1)
299307
return aggC
300308

301309

‎ogcore/demographics.py‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -900,10 +900,10 @@ def expand_pop_obj_J(
900900
Expand aggregate demographic objects to age x income-group objects.
901901
902902
The aggregate population path and rates are left unchanged. If
903-
income_percentiles is None and no income-specific inputs are provided, the
904-
original aggregate objects are returned. Otherwise, income_percentiles gives
905-
the initial population distribution across J for every age and the income
906-
shares of newborns in every period.
903+
income_percentiles is None and no income-specific inputs are
904+
provided, the original aggregate objects are returned. Otherwise,
905+
income_percentiles gives the initial population distribution across
906+
J for every age and the income shares of newborns in every period.
907907
908908
Args:
909909
omega_path_lev (Numpy array): T+S x E+S aggregate population levels.
@@ -926,7 +926,8 @@ def expand_pop_obj_J(
926926
fert_gradient (array_like): log-odds fertility slopes by age.
927927
mort_gradient (array_like): log-odds mortality slopes by age.
928928
infmort_gradient (array_like): log-odds infant mortality slopes.
929-
imm_pctiles (array_like): immigrant income shares by period, age, and J.
929+
imm_pctiles (array_like): immigrant income shares by period,
930+
age, and J.
930931
931932
Returns:
932933
dict: demographic objects with the same keys needed by get_pop_objs.
@@ -938,7 +939,10 @@ def expand_pop_obj_J(
938939
infmort_gradient,
939940
imm_pctiles,
940941
)
941-
assert income_percentiles is not None, "income_percentiles must be provided when using income-specific inputs."
942+
assert income_percentiles is not None, (
943+
"income_percentiles must be provided when using "
944+
+ "income-specific inputs."
945+
)
942946
income_shares, pct_midpoints = _income_shares_and_midpoints(
943947
income_percentiles
944948
)
@@ -951,10 +955,24 @@ def expand_pop_obj_J(
951955
if all_income_inputs_none:
952956
print("In the all are none case.")
953957
return {
954-
"omega_path_S": omega_path_S.reshape(omega_path_S.shape[0], omega_path_S.shape[1], 1) * income_shares.reshape(1, 1, J),
955-
"omega_SS": omega_SS.reshape(omega_SS.shape[0], 1) * income_shares.reshape(1, J),
956-
"mort_rates_S": np.tile(mort_rates_S.reshape(mort_rates_S.shape[0], mort_rates_S.shape[1], 1), (1, 1, J)),
957-
"imm_rates_mat": np.tile(imm_rates_mat.reshape(imm_rates_mat.shape[0], imm_rates_mat.shape[1], 1), (1, 1, J)),
958+
"omega_path_S": omega_path_S.reshape(
959+
omega_path_S.shape[0], omega_path_S.shape[1], 1
960+
)
961+
* income_shares.reshape(1, 1, J),
962+
"omega_SS": omega_SS.reshape(omega_SS.shape[0], 1)
963+
* income_shares.reshape(1, J),
964+
"mort_rates_S": np.tile(
965+
mort_rates_S.reshape(
966+
mort_rates_S.shape[0], mort_rates_S.shape[1], 1
967+
),
968+
(1, 1, J),
969+
),
970+
"imm_rates_mat": np.tile(
971+
imm_rates_mat.reshape(
972+
imm_rates_mat.shape[0], imm_rates_mat.shape[1], 1
973+
),
974+
(1, 1, J),
975+
),
958976
}
959977

960978
fert_slopes = _format_age_gradient(
@@ -1114,14 +1132,17 @@ def get_pop_objs(
11141132
pop_dist (array_like): user provided population distribution,
11151133
dimensions are T0+1 x E+S
11161134
fert_gradient (array_like): user provided fertility rate gradient,
1117-
dimensions are S, represents the log-odds slope in the fertility rate
1135+
dimensions are S, represents the log-odds slope in the
1136+
fertility rate
11181137
per percentile of the lifetime income distribution.
11191138
mort_gradient (array_like): user provided mortality rate gradient,
1120-
dimensions are S, represents the log-odds slope in the mortality rate
1121-
per percentile of the lifetime income distribution.
1122-
infmort_gradient (array_like): user provided infant mortality rate gradient,
1123-
dimensions are S, represents the log-odds slope in the infant mortality rate
1139+
dimensions are S, represents the log-odds slope in the
1140+
mortality rate
11241141
per percentile of the lifetime income distribution.
1142+
infmort_gradient (array_like): user provided infant mortality
1143+
rate gradient, dimensions are S, represents the log-odds
1144+
slope in the infant mortality rate per percentile of the
1145+
lifetime income distribution.
11251146
imm_pctiles (array_like): user provided lifetime income distribution
11261147
for new immigrants, shape is num_per x S x J, where num_per
11271148
is the number of years between initial and final_data_year

0 commit comments

Comments
 (0)