def egm(par, sol, t, m, c, inv_v):
""" apply egm step """
# loop over end-of-period assets
for i_a in prange(par.Na): # parallel
a = par.grid_a[t, i_a]
still_working_next_period = t + 1 <= par.TR - 1
Nshocks = par.Nshocks if still_working_next_period else 1
# loop over shocks
avg_marg_u_plus = 0
avg_v_plus = 0
for i_shock in range(Nshocks):
# a. prep
if still_working_next_period:
fac = par.G * par.L[t] * par.psi_vec[i_shock]
w = par.w[i_shock]
xi = par.xi_vec[i_shock]
else:
fac = par.G * par.L[t]
w = 1
xi = 1
inv_fac = 1.0 / fac
# b. future m and c
m_plus = inv_fac * par.R * a + xi
c_plus = linear_interp.interp_1d(sol.m[t + 1, :], sol.c[t + 1, :],
m_plus)
inv_v_plus = linear_interp.interp_1d(sol.m[t + 1, :],
sol.inv_v[t + 1, :], m_plus)
v_plus = 1.0 / inv_v_plus
# c. average future marginal utility
marg_u_plus = marg_utility(fac * c_plus, par)
avg_marg_u_plus += w * marg_u_plus
avg_v_plus += w * (fac**(1 - par.rho)) * v_plus
# d. current c
c[i_a] = inv_marg_utility(par.beta * par.R * avg_marg_u_plus, par)
# e. current m
m[i_a] = a + c[i_a]
# f. current v
if c[i_a] > 0:
inv_v[i_a] = 1.0 / (utility(c[i_a], par) + par.beta * avg_v_plus)
else:
inv_v[i_a] = 0
def solve_backwards(par, r, w, Va_p, Va, a, c, m, V_p, V, Vbar):
""" perform time iteration step with Va_p from previous iteration """
# a. post-decision
marg_u_plus = (par.beta * par.e_trans) @ Va_p
Vbar[:, :] = (par.beta * par.e_trans) @ V_p
# b. egm loop
for i_e in prange(par.Ne):
# i. egm
c_endo = marg_u_plus[i_e]**(-1 / par.sigma)
m_endo = c_endo + par.a_grid
# ii. interpolation
linear_interp.interp_1d_vec(m_endo, par.a_grid, m[i_e], a[i_e])
a[i_e, 0] = np.fmax(a[i_e, 0], 0)
c[i_e] = m[i_e] - a[i_e]
# iii. envelope condition
Va[i_e] = (1 + r) * c[i_e]**(-par.sigma)
# iv. value function
for i_a in range(par.Na):
a[i_e, i_a] = np.fmax(a[i_e, i_a], 0)
Vbar_now = linear_interp.interp_1d(par.a_grid, Vbar[i_e], a[i_e,
i_a])
V[i_e,
i_a] = c[i_e, i_a]**(1 - par.sigma) / (1 - par.sigma) + Vbar_now
def obj_keep(c,n,m,inv_w,grid_a,d_ubar,alpha,rho):
""" evaluate bellman equation """
# a. end-of-period assets
a = m-c
# b. continuation value
w = -1.0/linear_interp.interp_1d(grid_a,inv_w,a)
# c. total value
value_of_choice = utility.func_nopar(c,n,d_ubar,alpha,rho) + w
return -value_of_choice # we are minimizing
def simulate_forwards(par,sol,r,w,i_e,y,a):
""" simulate forward """
for i in prange(par.simN):
# i. update income
i_e_lag = i_e[i]
pi_e_val = np.random.uniform(0,1)
i_e[i] = choice(pi_e_val,par.e_trans_cumsum[i_e_lag,:])
y[i] = w*par.e_grid[i_e[i]]
# ii. update assets
a[i] = linear_interp.interp_1d(par.a_grid,sol.a[i_e[i]],a[i])
def obj_bellman(c, m, interp_w, par):
""" evaluate bellman equation """
# a. end-of-period assets
a = m - c
# b. continuation value
w = linear_interp.interp_1d(par.grid_a, interp_w, a)
# c. total value
value_of_choice = utility.func(c, par) + w
return -value_of_choice # we are minimizing
def obj_keep(c,n,m,inv_w,grid_b,d_ubar,sigma,P_n_p): # P_n_p
""" evaluate bellman equation """
# a. end-of-period assets
b = m-P_n_p*c
# b. continuation value
w = -1.0/linear_interp.interp_1d(grid_b,inv_w,b)
# c. total value
value_of_choice = utility(c,n,d_ubar,sigma) + w
return -value_of_choice # we are minimizing
def simulate_time_loop(par, sol, sim):
""" simulate model with parallization over households """
# unpack (helps numba)
m = sim.m
p = sim.p
y = sim.y
c = sim.c
a = sim.a
sol_c = sol.c
sol_m = sol.m
# loop over first households and then time
for i in prange(par.simN):
for t in range(par.simT):
# a. solution
if par.simlifecycle == 0:
grid_m = sol_m[0, :]
grid_c = sol_c[0, :]
else:
grid_m = sol_m[t, :]
grid_c = sol_c[t, :]
# b. consumption
c[i, t] = linear_interp.interp_1d(grid_m, grid_c, m[i, t])
a[i, t] = m[i, t] - c[i, t]
# c. next-period
if t < par.simT - 1:
if t + 1 > par.TR - 1:
m[i, t + 1] = par.R * a[i, t] / (par.G * par.L[t]) + 1
p[i, t + 1] = np.log(par.G) + np.log(par.L[t]) + p[i, t]
y[i, t + 1] = p[i, t + 1]
else:
m[i, t + 1] = par.R * a[i, t] / (
par.G * par.L[t] * sim.psi[i, t + 1]) + sim.xi[i,
t + 1]
p[i, t + 1] = np.log(par.G) + np.log(
par.L[t]) + p[i, t] + np.log(sim.psi[i, t + 1])
if sim.xi[i, t + 1] > 0:
y[i, t + 1] = p[i, t + 1] + np.log(sim.xi[i, t + 1])
else:
y[i, t + 1] = -np.inf
def solve_bellman(t,sol,par):
"""solve the bellman equation using the endogenous grid method"""
# unpack (helps numba optimize)
c = sol.c[t]
for ip in prange(par.Np): # in parallel
# a. temporary container (local to each thread)
m_temp = np.zeros(par.Na+1) # m_temp[0] = 0
c_temp = np.zeros(par.Na+1) # c_temp[0] = 0
# b. invert Euler equation
for ia in range(par.Na):
c_temp[ia+1] = utility.inv_marg_func(sol.q[ip,ia],par)
m_temp[ia+1] = par.grid_a[ia] + c_temp[ia+1]
# b. re-interpolate consumption to common grid
if par.do_simple_w: # use an explicit loop
for im in range(par.Nm):
c[ip,im] = linear_interp.interp_1d(m_temp,c_temp,par.grid_m[im])
else: # use a vectorized call (assumming par.grid_m is monotone)
linear_interp.interp_1d_vec_mon_noprep(m_temp,c_temp,par.grid_m,c[ip,:])
def value_of_choice(self, c, t, m):
""" value of choice of c used in vfi """
par = self.par
sol = self.sol
# a. end-of-period assets
a = m - c
# b. next-period cash-on-hand
still_working_next_period = t + 1 <= par.TR - 1
if still_working_next_period:
fac = par.G * par.L[t] * par.psi_vec
w = par.w
xi = par.xi_vec
else:
fac = par.G * par.L[t]
w = 1
xi = 1
m_plus = (par.R / fac) * a + xi
# c. continuation value
if still_working_next_period:
inv_v_plus = np.zeros(m_plus.size)
linear_interp.interp_1d_vec(sol.m[t + 1, :], sol.inv_v[t + 1, :],
m_plus, inv_v_plus)
else:
inv_v_plus = linear_interp.interp_1d(sol.m[t + 1, :],
sol.inv_v[t + 1, :], m_plus)
v_plus = 1 / inv_v_plus
# d. value-of-choice
total = utility(
c, par) + par.beta * np.sum(w * fac**(1 - par.rho) * v_plus)
return -total
def euler(sim, sol, par):
euler = sim.euler
# a. grids
min_m = 0.50
min_n = 0.01
m_max = 5.00
n_max = 5.00
n_grid = np.linspace(min_n, n_max, par.eulerK)
m_grid = np.linspace(min_m, m_max, par.eulerK)
# b. loop over time
for t in range(par.T - 1):
for i_n in range(par.eulerK):
for i_m in range(par.eulerK):
# i. states
n = n_grid[i_n]
m = m_grid[i_m]
m_retire = m_grid[i_m] + n_grid[i_n]
# ii. discrete choice
inv_v_retire = linear_interp.interp_1d(sol.m_ret[t],
sol.inv_v_ret[t],
m_retire)
inv_v = linear_interp.interp_2d(par.grid_n, par.grid_m,
sol.inv_v[t], n, m)
if inv_v_retire > inv_v: continue
# iii. continuous choice
c = np.fmin(
linear_interp.interp_2d(par.grid_n, par.grid_m, sol.c[t],
n, m), m)
d = np.fmax(
linear_interp.interp_2d(par.grid_n, par.grid_m, sol.d[t],
n, m), 0)
a = m - c - d
b = n + d + pens.func(d, par)
if a < 0.001: continue
# iv. shocks
RHS = 0
for i_eta in range(par.Neta):
# o. state variables
n_plus = par.Rb * b
m_plus = par.Ra * a + par.eta[i_eta]
m_retire_plus = m_plus + n_plus
# oo. discrete choice
inv_v_retire = linear_interp.interp_1d(
sol.m_ret[t + 1], sol.inv_v_ret[t + 1], m_retire_plus)
inv_v = linear_interp.interp_2d(par.grid_n, par.grid_m,
sol.inv_v[t + 1], n_plus,
m_plus)
# ooo. continous choice
if inv_v_retire > inv_v:
c_plus = np.fmin(
linear_interp.interp_1d(sol.m_ret[t + 1],
sol.c_ret[t + 1],
m_retire_plus),
m_retire_plus)
else:
c_plus = np.fmin(
linear_interp.interp_2d(par.grid_n, par.grid_m,
sol.c[t + 1], n_plus,
m_plus), m_plus)
# oooo. accumulate
RHS += par.w_eta[
i_eta] * par.beta * par.Ra * utility.marg_func(
c_plus, par)
# v. euler error
euler_raw = c - utility.inv_marg_func(RHS, par)
euler[t, i_m, i_n] = np.log10(np.abs(euler_raw / c) + 1e-16)
def collect_policy(Np,Nn,Nm,Nb,grid_p,grid_n,grid_m,grid_x,grid_b,inv_v_keep,inv_v_adj,c_keep,c_adj,d_adj,sigma_eps,tau,delta,n_max,R,P_d_p,P_n_p): # P_d_p, P_n_p
"""solve bellman equation for keepers using nvfi"""
# unpack output
collected_shape = (Np,Nn,Nb)
c_combined = np.zeros(collected_shape)
d_combined = np.zeros(collected_shape)
n_plus = np.zeros(collected_shape)
a_plus = np.zeros(collected_shape)
#inv_v = np.zeros(collected_shape)
# loop over outer states
for i_p in prange(Np):
# permanent income shock
p = grid_p[i_p]
for i_n in range(Nn):
# outer states
n = grid_n[i_n]
# loop over m state
for i_b in range(Nb): # beginning of period assets
# i. cash on hand
m = R*grid_b[i_b] + p - grid_b[0] # cash on hand non-adjuster
x = m + P_d_p*(1-tau)*n # adjuster cash on hand
# ii. value functions
inv_v_k = linear_interp.interp_1d(grid_m,inv_v_keep[i_p,i_n,:],m)
inv_v_a = linear_interp.interp_1d(grid_x,inv_v_adj[i_p,:],x)
if inv_v_k == 0:
v_keep =-9999.0
else:
v_keep = -1.0/inv_v_k
if inv_v_a == 0:
v_adj=-9999.0
else:
v_adj = -1.0/inv_v_a
v_max = np.fmax(v_keep,v_adj)
# iii. choice probability adjusting
numerator = np.exp((v_adj-v_max)/sigma_eps)
denominator = np.exp((v_adj-v_max)/sigma_eps) + np.exp((v_keep-v_max)/sigma_eps)
P_adj = numerator/denominator
# iv. combined policy functions
# o. adjuster
d_adj_temp = linear_interp.interp_1d(grid_x,d_adj[i_p,:],x)
c_adj_temp = linear_interp.interp_1d(grid_x,c_adj[i_p,:],x)
tot_adj = P_d_p*d_adj_temp + P_n_p*c_adj_temp
if tot_adj > x:
d_adj_temp *= x/tot_adj
c_adj_temp *= x/tot_adj
a_adj_temp = 0.0
n_plus_adj_temp = (1-delta)*d_adj_temp
n_plus_adj_temp = np.fmin(n_plus_adj_temp,n_max)
else:
a_adj_temp = x - tot_adj
n_plus_adj_temp = (1-delta)*d_adj_temp
n_plus_adj_temp = np.fmin(n_plus_adj_temp,n_max)
# oo. keeper
d_keep_temp = n
c_keep_temp = linear_interp.interp_1d(grid_m,c_keep[i_p,i_n,:],m)
if c_keep_temp > m:
c_keep_temp = m/P_n_p # P_n_p*c, have m. Use all m, P_n_p*c=m, c=m/P_n_p
a_keep_temp = 0.0
n_plus_keep_temp = (1-delta)*n
n_plus_keep_temp= np.fmin(n_plus_keep_temp,n_max)
else:
a_keep_temp = m - P_n_p*c_keep_temp
n_plus_keep_temp = (1-delta)*n
n_plus_keep_temp = np.fmin(n_plus_keep_temp,n_max)
# ooo. combined policy function
# i. states
n_plus_adj_temp = (1-delta)*d_adj_temp
n_plus_keep_temp = (1-delta)*d_keep_temp
# ii. policy functions
d_combined[i_p,i_n,i_b] = P_adj*d_adj_temp + (1-P_adj)*d_keep_temp
c_combined[i_p,i_n,i_b] = P_adj*c_adj_temp + (1-P_adj)*c_keep_temp
a_plus[i_p,i_n,i_b] = P_adj*a_adj_temp + (1-P_adj)*a_keep_temp
n_plus[i_p,i_n,i_b] = P_adj*n_plus_adj_temp + (1-P_adj)*n_plus_keep_temp
# oooo. value function
#inv_v[i_p,i_n,i_b] = sigma_eps*np.log(np.exp(inv_v_k/sigma_eps) + np.exp(inv_v_a/sigma_eps))
return d_combined, c_combined, a_plus, n_plus