Skip to content

Commit

Permalink
Merge pull request #897 from jdebacker/chi_n
Browse files Browse the repository at this point in the history
Merging
  • Loading branch information
rickecon authored Feb 6, 2024
2 parents ba50ecc + 272c969 commit 42cd461
Show file tree
Hide file tree
Showing 7 changed files with 66 additions and 19 deletions.
2 changes: 1 addition & 1 deletion ogcore/SS.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ def euler_equation_solver(guesses, *args):
tr,
ubi,
theta,
p.chi_n,
p.chi_n[-1, :],
p.etr_params[-1],
p.mtrx_params[-1],
None,
Expand Down
4 changes: 2 additions & 2 deletions ogcore/TPI.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ def firstdoughnutring(
np.array([tr]),
np.array([ubi]),
theta[j],
p.chi_n[-1],
p.chi_n[0, -1],
p.etr_params[0][-1],
p.mtrx_params[0][-1],
None,
Expand Down Expand Up @@ -258,7 +258,7 @@ def twist_doughnut(
r_s = r[t : t + length]
p_tilde_s = p_tilde[t : t + length]
n_s = n_guess
chi_n_s = p.chi_n[-length:]
chi_n_s = np.diag(p.chi_n[t : t + p.S, :], max(p.S - length, 0))
rho_s = np.diag(p.rho[t : t + p.S, :], max(p.S - length, 0))

error1 = household.FOC_savings(
Expand Down
42 changes: 30 additions & 12 deletions ogcore/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,30 @@ def compute_default_params(self):
np.ones(self.T + self.S - self.BW) * self.frac_tax_payroll[-1],
)

# Interpolate chi_n and create omega_SS_80 if necessary
if self.S < 80 and self.chi_n.shape[-1] == 80:
self.age_midp_80 = np.linspace(20.5, 99.5, 80)
reshape_chi_n = np.zeros((self.T + self.S, self.S))
for t in range(self.T + self.S):
if self.chi_n.ndim == 1:
self.chi_n_interp = si.interp1d(
self.age_midp_80,
np.squeeze(self.chi_n),
kind="cubic",
)
else:
self.chi_n_interp = si.interp1d(
self.age_midp_80,
np.squeeze(self.chi_n[t, :]),
kind="cubic",
)
self.newstep = 80.0 / self.S
self.age_midp_S = np.linspace(
20 + 0.5 * self.newstep, 100 - 0.5 * self.newstep, self.S
)
reshape_chi_n[t, :] = self.chi_n_interp(self.age_midp_S)
self.chi_n = reshape_chi_n

# Extend parameters that may vary over the time path
tp_param_list = [
"alpha_G",
Expand Down Expand Up @@ -262,6 +286,12 @@ def compute_default_params(self):
param_in, dims=(self.T, self.S, self.J), item="e"
)
setattr(self, "e", param_out)
# Extrapolate chi_n over T + S
param_in = getattr(self, "chi_n")
param_out = extrapolate_arrays(
param_in, dims=(self.T + self.S, self.S), item="chi_n"
)
setattr(self, "chi_n", param_out)

# make sure zeta matrix sums to one (e.g., default off due to rounding)
self.zeta = self.zeta / self.zeta.sum()
Expand Down Expand Up @@ -319,18 +349,6 @@ def compute_default_params(self):
)
self.omega_S_preTP = self.omega_SS

# Interpolate chi_n and create omega_SS_80 if necessary
if self.S < 80 and len(self.chi_n) == 80:
self.age_midp_80 = np.linspace(20.5, 99.5, 80)
self.chi_n_interp = si.interp1d(
self.age_midp_80, np.squeeze(self.chi_n), kind="cubic"
)
self.newstep = 80.0 / self.S
self.age_midp_S = np.linspace(
20 + 0.5 * self.newstep, 100 - 0.5 * self.newstep, self.S
)
self.chi_n = self.chi_n_interp(self.age_midp_S)

# Create time series of stationarized UBI transfers
self.ubi_nom_array = self.get_ubi_nom_objs()

Expand Down
9 changes: 5 additions & 4 deletions ogcore/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -932,20 +932,21 @@ def extrapolate_arrays(param_in, dims=None, item="Parameter Name"):
param_out = np.ones((dims)) * param_in[0]
# case where user enters just one year for all types in 2nd dim
if param_in.shape[0] == dims[1]:
param_in = np.tile(param_in.reshape(1, dims[1]), (dims[0], 1))
param_out = np.tile(param_in.reshape(1, dims[1]), (dims[0], 1))
else:
# case where user enters multiple years, but not full time
# path
# will assume they implied values the same across 2nd dim
# and will fill in all periods
param_in = np.concatenate(
param_out = np.concatenate(
(
param_in,
np.ones((dims[0] - param_in.size)) * param_in[-1],
)
)
param_in = np.tile(param_in.reshape(dims[0], 1), (1, dims[1]))
param_out = np.squeeze(param_in, axis=2)
param_out = np.tile(
param_out.reshape(dims[0], 1), (1, dims[1])
)
elif param_in.ndim == 2:
# case where enter values along 2 dimensions, but those aren't
# complete
Expand Down
8 changes: 8 additions & 0 deletions tests/test_aggregates.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
"S": 40,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
"eta": (np.ones((40, 2)) / (40 * 2)),
Expand Down Expand Up @@ -53,6 +54,7 @@ def test_get_L(n, p, method, expected):
"S": 40,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((40, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -134,6 +136,7 @@ def test_get_I(b_splus1, K_p1, K, p, method, expected):
"S": 40,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((40, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -191,6 +194,7 @@ def test_get_B(b, p, method, PreTP, expected):
"S": 40,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((40, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -273,6 +277,7 @@ def test_get_BQ(r, b_splus1, j, p, method, PreTP, expected):
"S": 40,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((40, 2)),
"M": 3,
"labor_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -323,6 +328,7 @@ def test_get_C(c, p, method, expected):
"S": 20,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((20, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -380,6 +386,7 @@ def test_get_C(c, p, method, expected):
"S": 20,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((20, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down Expand Up @@ -415,6 +422,7 @@ def test_get_C(c, p, method, expected):
"S": 20,
"rho": rho_vec.tolist(),
"J": 2,
"chi_n": np.ones(2),
"e": np.ones((20, 2)),
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
Expand Down
13 changes: 13 additions & 0 deletions tests/test_firm.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
"epsilon": [1.0],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand All @@ -52,6 +53,7 @@
"epsilon": [1.0],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p5.J)) / (3 * p5.J)),
"e": np.ones((3, p5.J)),
Expand All @@ -66,6 +68,7 @@
"epsilon": [1.0],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p5.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p5.J)) / (3 * p5.J)),
Expand All @@ -87,6 +90,7 @@
"T": 3,
"S": 3,
"M": 2,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand All @@ -109,6 +113,7 @@
"T": 3,
"S": 3,
"M": 2,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand All @@ -127,6 +132,7 @@
"T": 3,
"S": 3,
"M": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p5.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p5.J)) / (3 * p5.J)),
Expand All @@ -153,6 +159,7 @@
"T": 3,
"S": 3,
"M": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p5.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p5.J)) / (3 * p5.J)),
Expand Down Expand Up @@ -274,6 +281,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
"delta_annual": 0.5,
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand All @@ -295,6 +303,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
"inv_tax_credit": [[0.07]],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p5.J)) / (3 * p5.J)),
Expand Down Expand Up @@ -370,6 +379,7 @@ def test_get_r(Y, K, p_m, p, method, expected):
"epsilon": [1.2],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand Down Expand Up @@ -460,6 +470,7 @@ def test_get_w(Y, L, p_m, p, method, expected):
"cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand Down Expand Up @@ -569,6 +580,7 @@ def test_get_KLratio(r, p, method, expected):
"cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand Down Expand Up @@ -651,6 +663,7 @@ def test_get_w_from_r(r, p, method, expected):
"cit_rate": [[0.5]],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p4.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p4.J)) / (3 * p4.J)),
Expand Down
7 changes: 7 additions & 0 deletions tests/test_tax.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
"capital_income_tax_noncompliance_rate": [[0.0]],
"J": 1,
"T": 4,
"chi_n": np.ones(4),
"eta": (np.ones((4, 1)) / (4 * 1)),
"e": np.ones((4, 1)),
}
Expand Down Expand Up @@ -81,6 +82,7 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
"lambdas": [1.0],
"J": 1,
"T": 3,
"chi_n": np.ones(3),
"e": np.ones((3, 1)),
"eta": (np.ones((3, 1)) / (3 * 1)),
"h_wealth": [2],
Expand All @@ -98,6 +100,7 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
"lambdas": [1.0],
"J": 1,
"T": 3,
"chi_n": np.ones(3),
"e": np.ones((3, 1)),
"eta": (np.ones((3, 1)) / (3 * 1)),
"h_wealth": [1.2, 1.1, 2.3],
Expand Down Expand Up @@ -132,6 +135,7 @@ def test_ETR_wealth(b, p, expected):
"lambdas": [1.0],
"J": 1,
"T": 3,
"chi_n": np.ones(3),
"e": np.ones((3, 1)),
"eta": (np.ones((3, 1)) / (3 * 1)),
"h_wealth": [3],
Expand All @@ -150,6 +154,7 @@ def test_ETR_wealth(b, p, expected):
"lambdas": [1.0],
"J": 1,
"T": 3,
"chi_n": np.ones(3),
"e": np.ones((3, 1)),
"eta": (np.ones((3, 1)) / (3 * 1)),
"h_wealth": [1.2, 1.1, 2.3],
Expand Down Expand Up @@ -885,6 +890,7 @@ def test_MTR_income(etr_params, mtr_params, params, mtr_capital, expected):
"inv_tax_credit": [[0.02]],
"T": 3,
"S": 3,
"chi_n": np.ones(3),
"e": np.ones((3, p1.J)),
"rho": rho_vec.tolist(),
"eta": (np.ones((3, p1.J)) / (3 * p1.J)),
Expand Down Expand Up @@ -1214,6 +1220,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected):
"T": 3,
"S": 3,
"J": 2,
"chi_n": np.ones(3),
"e": np.ones((3, 2)),
"rho": rho_vec.tolist(),
"lambdas": [0.65, 0.35],
Expand Down

0 comments on commit 42cd461

Please sign in to comment.