def step4(ap_endo, c_endo, z_grid, b_grid, a_grid, ra, rb, chi0, chi1, chi2):
# b(z, b', a)
zzz = z_grid[:, np.newaxis, np.newaxis]
bbb = b_grid[np.newaxis, :, np.newaxis]
aaa = a_grid[np.newaxis, np.newaxis, :]
b_endo = (c_endo + ap_endo + bbb - (1 + ra) * aaa + Psi_fun(ap_endo, aaa, ra, chi0, chi1, chi2) -
zzz) / (1 + rb)
# b'(z, b, a), a'(z, b, a)
# assert np.min(np.diff(b_endo, axis=1)) > 0, 'b(bp) is not increasing'
# assert np.min(np.diff(ap_endo, axis=1)) > 0, 'ap(bp) is not increasing'
i, pi = utils.interpolate_coord(b_endo.swapaxes(1, 2), b_grid)
ap = utils.apply_coord(i, pi, ap_endo.swapaxes(1, 2)).swapaxes(1, 2)
bp = utils.apply_coord(i, pi, b_grid).swapaxes(1, 2)
return bp, ap
def time_iteration(sigma, theta, alpha, delta, r, lump, beta, Pi_e, d_grid,
z_grid, b_grid, k_grid, zzz, lll, bbb, ddd, Ne, Nb, Nd, Nk,
rhs, Vb, Vd, P_n_p, P_d_p):
# output containers
sol_shape = (Ne, Nb, Nd)
Wb = np.zeros(sol_shape)
Wd = np.zeros(sol_shape)
d_unc = np.zeros(sol_shape)
b_unc = np.zeros(sol_shape)
d_con = np.zeros(sol_shape)
''' computes one backward EGM step '''
# a. pre-compute post-decision value functions and LHS offirst-order optimality equation
# i. post-decision functions
Wb[:, :, :] = (Vb.T @ (beta * Pi_e.T)).T
Wd[:, :, :] = (Vd.T @ (beta * Pi_e.T)).T
# ii. LHS of FOC optimality equation in unconstrained and constrained case
lhs_unc = Wd / Wb
lhs_con = lhs_unc[:, 0, :]
lhs_con = lhs_con[:,
np.newaxis, :] / (1 + k_grid[np.newaxis, :, np.newaxis])
# b. get unconstrained solution to d'(z,b,d), a'(z,b,d), c(z,b,d)
# i. get d'(z,b',d), b(z,b',d), c(z,b,d) with EGM
dp_endo_unc, _, b_endo_unc, _ = solve_unconstrained(
sigma, theta, alpha, delta, r, d_grid, b_grid, z_grid, lump, Ne, Nb,
Nd, lhs_unc, rhs, Wb, P_d_p, P_n_p)
# ii. use d'(z,b',d), b(z,b',d) to interpolate to d'(z,b,d), b'(z,b,d)
i, pi = utils.interpolate_coord(b_endo_unc.swapaxes(
1, 2), b_grid) # gives interpolation weights from a' grid to a grid
d_unc[:, :, :] = utils.apply_coord(i, pi, dp_endo_unc.swapaxes(
1, 2)).swapaxes(1, 2) # apply weights to d'(z,a',d)->d'(z,a,d)
b_unc[:, :, :] = utils.apply_coord(i, pi, b_grid).swapaxes(
1, 2) # apply weights to a'(z,a',d)->a'(z,a,d)
# c. get constrained solution to d'(z,0,d), b'(z,0,d), c(z,0,d)
# i. get d'(z,k,d), b(z,k,d), c(z,k,d) with EGM
dp_endo_con, _, b_endo_con, _ = solve_constrained(sigma, theta, alpha,
delta, r, d_grid, b_grid,
k_grid, z_grid, lump, Ne,
Nk, Nd, lhs_con, rhs, Wb,
P_n_p, P_d_p)
# ii. get d'(z,b,d) by interpolating using b'(z,k,d)
d_con[:] = utils.interpolate_y(b_endo_con[:, ::-1, :].swapaxes(1, 2),
b_grid, dp_endo_con[:, ::-1, :].swapaxes(
1, 2)).swapaxes(1, 2)
# d. collect policy functions d',b',c by combining unconstrained and constrained solutions
d, b, c = collect_policy(delta, alpha, r, Ne, Nb, Nd, b_grid, zzz, lll,
ddd, bbb, d_unc, b_unc, d_con, P_n_p, P_d_p)
c_d = d - (1 - delta) * ddd # durable investment policy fuctions
# e. update Vb, Vd
# i. compute marginal utilities
c[c < 0] = 1e-8 # for numerical stability while converging
uc = theta * c**(theta - 1) * (ddd)**(1 - theta) * (
(c**theta * (ddd)**(1 - theta)))**(-sigma)
ud = (1 - theta) * c**theta * (ddd)**(-theta) * (
(c**theta * (ddd)**(1 - theta)))**(-sigma)
# ii. compute Vb, Vd using envelope conditions
Vb = (1 / P_n_p) * (1 + r) * uc
Vd = ud + (1 / P_n_p) * uc * (P_d_p * (1 - delta) - 0.5 * alpha *
((1 - delta)**2 - (d / ddd)**2))
return Vb, Vd, d, b, c, c_d