def solve_acon(out_c,out_d,out_v,w,wb,par):
num = 3
# i. allocate
c = np.zeros((par.Nc_acon,par.Nb_acon))
d = np.zeros((par.Nc_acon,par.Nb_acon))
a = np.zeros((par.Nc_acon,par.Nb_acon))
b = np.zeros((par.Nc_acon,par.Nb_acon))
w_acon = np.zeros((par.Nc_acon,par.Nb_acon))
for i_b in range(par.Nb_acon):
# ii. setup
wb_acon = linear_interp.interp_2d(par.grid_b_pd,par.grid_a_pd,wb,par.b_acon[i_b],0)
c_min = utility.inv_marg_func((par.chi+1)*wb_acon,par)
c_max = utility.inv_marg_func(wb_acon,par)
c[:,i_b] = misc.nonlinspace_jit(c_min,c_max,par.Nc_acon,par.phi_m)
# iii. choices
d[:,i_b] = par.chi/(utility.marg_func(c[:,i_b],par)/(wb_acon)-1)-1
# iv. post-decision states and value function
b[:,i_b] = par.b_acon[i_b]
w_acon[:,i_b] = linear_interp.interp_2d(par.grid_b_pd,par.grid_a_pd,w,par.b_acon[i_b],0)
# v. states and value
m,n,v = inv_mn_and_v(c,d,a,b,w_acon,par)
# vi. upperenvelope and interp to common
upperenvelope.compute(out_c,out_d,out_v,m,n,c,d,v,num,w,par)
def solve_bellman(sol, par, t):
"""solve the bellman equation using the endogenous grid method"""
# unpack (helps numba optimize)
c = sol.c
# loop over delta state
for idelta in prange(par.Ndelta):
# temp
m_temp = np.zeros(par.Na + 1)
c_temp = np.zeros(par.Na + 1)
# a. find consumption
for ia in range(par.Na):
c_temp[ia + 1] = utility.inv_marg_func(sol.q[t, idelta, ia], par)
m_temp[ia + 1] = par.grid_a[ia] + c_temp[ia + 1]
if t == 0:
m_temp[ia + 1] -= par.grid_delta[idelta]
if t == 0: # assuming the borrowing constraint is binding for low enough m
c_temp[0] = par.grid_delta[idelta]
# b. interpolate to common grid
linear_interp.interp_1d_vec_mon_noprep(m_temp, c_temp, par.grid_m,
c[t, idelta, :])
def solve_pure_c(t,sol,par):
w = sol.w[t]
wa = sol.wa[t]
# unpack
inv_v = sol.inv_v_pure_c[t]
c = sol.c_pure_c[t]
for i_b in range(par.Nb_pd):
# temporary containers
temp_c = np.zeros(par.Na_pd)
temp_m = np.zeros(par.Na_pd)
temp_v = np.zeros(par.Na_pd)
# use euler equation
for i_a in range(par.Na_pd):
temp_c[i_a] = utility.inv_marg_func(wa[i_b,i_a],par)
temp_m[i_a] = par.grid_a_pd[i_a] + temp_c[i_a]
# upperenvelope
negm_upperenvelope(par.grid_a_pd,temp_m,temp_c,w[i_b],
par.grid_l,c[i_b,:],temp_v,par)
# negative inverse
for i_m in range(par.Na_pd):
inv_v[i_b,i_m] = -1/temp_v[i_m]
def solve(sol,par,G2EGM=True):
# a. last_period
t = par.T-1
sol.m_ret[t,:] = par.grid_m_ret
sol.c_ret[t,:] = sol.m_ret[t,:]
v = utility.func_ret(sol.c_ret[t,:],par)
sol.inv_v_ret[t,:] = -1.0/v
vm = utility.marg_func(sol.c_ret[t,:],par)
sol.inv_vm_ret[t,:] = 1.0/vm
if G2EGM:
sol.inv_vn_ret[t,:] = sol.inv_vm_ret[t,:]
# b. backwards inducation
for j in range(2,par.T+1):
t = par.T-j
# i. optimal c choice
m_plus = par.Ra*par.grid_a_ret + par.yret
c_plus = np.zeros(m_plus.shape)
linear_interp.interp_1d_vec(sol.m_ret[t+1,:],sol.c_ret[t+1,:],m_plus,c_plus)
vm_plus = utility.marg_func(c_plus,par)
q = par.beta*par.Ra*vm_plus
sol.c_ret[t,par.Nmcon_ret:] = utility.inv_marg_func(q,par)
# ii. constraint
sol.c_ret[t,:par.Nmcon_ret] = nonlinspace_jit(1e-6,sol.c_ret[t,par.Nmcon_ret]*0.999,par.Nmcon_ret,par.phi_m)
# iii. end-of-period assets and value-of-choice
sol.a_ret[t,par.Nmcon_ret:] = par.grid_a_ret
inv_v_plus = np.zeros(m_plus.shape)
linear_interp.interp_1d_vec(sol.m_ret[t+1,:],sol.inv_v_ret[t+1,:],m_plus,inv_v_plus)
v_plus = -1.0/inv_v_plus
v1 = utility.func_ret(sol.c_ret[t,:par.Nmcon_ret],par) + par.beta*v_plus[0]
v2 = utility.func_ret(sol.c_ret[t,par.Nmcon_ret:],par) + par.beta*v_plus
sol.inv_v_ret[t,:par.Nmcon_ret] = -1.0/v1
sol.inv_v_ret[t,par.Nmcon_ret:] = -1.0/v2
# iv. endogenous grid
sol.m_ret[t,:] = sol.a_ret[t,:] + sol.c_ret[t,:]
# v. marginal v
vm = utility.marg_func(sol.c_ret[t,:],par)
sol.inv_vm_ret[t,:] = 1.0/vm
if G2EGM:
sol.inv_vn_ret[t,:] = sol.inv_vm_ret[t,:]
def solve_ucon(out_c, out_d, out_v, w, wa, wb, par):
num = 1
# i. choices
c = utility.inv_marg_func(wa, par)
d = (par.chi * wb) / (wa - wb) - 1
# ii. states and value
a = par.grid_a_pd_nd
b = par.grid_b_pd_nd
m, n, v = inv_mn_and_v(c, d, a, b, w, par)
# iii. upperenvelope and interp to common
upperenvelope.compute(out_c, out_d, out_v, m, n, c, d, v, num, w, par)
def solve_bellman_c(t, ad, st_h, st_w, ra_h, ra_w, D_h, D_w, par, a, sol_c,
sol_m, sol_v, single_sol_v_plus_raw,
single_sol_avg_marg_u_plus):
""" solve the bellman equation for singles"""
# compute post decision
v_raw, q = post_decision.compute_c(t, ad, st_h, st_w, ra_h, ra_w, D_h, D_w,
par, a, sol_c, sol_m, sol_v,
single_sol_v_plus_raw,
single_sol_avg_marg_u_plus)
# unpack solution
ad_min = par.ad_min
ad_idx = ad + ad_min
c = sol_c[t, ad_idx, st_h, st_w, ra_h, ra_w]
m = sol_m[:]
v = sol_v[t, ad_idx, st_h, st_w, ra_h, ra_w]
# loop over the choices
for d_h in D_h:
for d_w in D_w:
# a. indices
d = transitions.d_c(d_h, d_w) # joint index
# b. raw solution
c_raw = utility.inv_marg_func(q[d], par)
m_raw = a + c_raw
# d. upper envelope
envelope_c(
a,
m_raw,
c_raw,
v_raw[d],
m, # input
c[d],
v[d], # output
d_h,
d_w,
st_h,
st_w,
par) # args for utility function
def solve_bellman(t, ma, st, ra, D, sol_c, sol_m, sol_v, sol_v_plus_raw,
sol_avg_marg_u_plus, a, par):
""" solve the bellman equation for singles"""
# compute post decision (and store results, since they are needed in couple model)
v_plus_raw, avg_marg_u_plus = post_decision.compute(
t, ma, st, ra, D, sol_c, sol_m, sol_v, a, par)
sol_v_plus_raw[t + 1, ra] = v_plus_raw
sol_avg_marg_u_plus[t + 1, ra] = avg_marg_u_plus
# unpack
c = sol_c[t, ra]
m = sol_m[:]
v = sol_v[t, ra]
pi_plus = transitions.survival_lookup_single(t + 1, ma, st, par)
# loop over the choices
for d in D:
# a. post decision
q = par.beta * (par.R * pi_plus * avg_marg_u_plus[d] +
(1 - pi_plus) * par.gamma)
# b. raw solution
c_raw = utility.inv_marg_func(q, par)
m_raw = a + c_raw
v_raw = par.beta * (pi_plus * v_plus_raw[d] +
(1 - pi_plus) * par.gamma * a
) # without utility (added in envelope)
# c. upper envelope
envelope(
a,
m_raw,
c_raw,
v_raw,
m, # input
c[d],
v[d], # output
d,
ma,
st,
par) # args for utility function
def solve_dcon(out_c,out_d,out_v,w,wa,par):
num = 2
# i. decisions
c = utility.inv_marg_func(wa,par)
d = par.d_dcon
# ii. states and value
a = par.grid_a_pd_nd
b = par.grid_b_pd_nd
m,n,v = inv_mn_and_v(c,d,a,b,w,par)
# iii. value of deviating a bit from the constraint
valt = np.zeros(v.shape)
deviate_d_con(valt,n,c,a,w,par)
# v. upperenvelope and interp to common
upperenvelope.compute(out_c,out_d,out_v,m,n,c,d,v,num,w,par,valt)
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 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 euler_errors(sim, sol, par):
# unpack
euler_error = sim.euler_error
euler_error_c = sim.euler_error_c
for i in prange(par.simN):
discrete_plus = np.zeros(1)
d_plus = np.zeros(1)
d1_plus = np.zeros(1)
d2_plus = np.zeros(1)
c_plus = np.zeros(1)
a_plus = np.zeros(1)
for t in range(par.T - 1):
constrained = sim.a[t, i] < par.euler_cutoff
if constrained:
euler_error[t, i] = np.nan
euler_error_c[t, i] = np.nan
continue
else:
RHS = 0.0
for ishock in range(par.Nshocks):
# i. shocks
psi = par.psi[ishock]
psi_w = par.psi_w[ishock]
xi = par.xi[ishock]
xi_w = par.xi_w[ishock]
# ii. next-period states
p_plus = trans.p_plus_func(sim.p[t, i], psi, par)
if par.do_2d:
n1_plus = trans.n1_plus_func(sim.d1[t, i], par)
n2_plus = trans.n2_plus_func(sim.d2[t, i], par)
else:
n_plus = trans.n_plus_func(sim.d[t, i], par)
m_plus = trans.m_plus_func(sim.a[t, i], p_plus, xi, par)
# iii. weight
weight = psi_w * xi_w
# iv. next-period choices
if par.do_2d:
optimal_choice_2d(t + 1, p_plus, n1_plus, n2_plus,
m_plus, discrete_plus, d1_plus,
d2_plus, c_plus, a_plus, sol, par)
else:
optimal_choice(t + 1, p_plus, n_plus, m_plus,
discrete_plus, d_plus, c_plus, a_plus,
sol, par)
# v. next-period marginal utility
if par.do_2d:
RHS += weight * par.beta * par.R * utility.marg_func_2d(
c_plus[0], d1_plus[0], d2_plus[0], par)
else:
RHS += weight * par.beta * par.R * utility.marg_func(
c_plus[0], d_plus[0], par)
if par.do_2d:
euler_error[t, i] = sim.c[t, i] - utility.inv_marg_func_2d(
RHS, sim.d1[t, i], sim.d2[t, i], par)
else:
euler_error[t, i] = sim.c[t, i] - utility.inv_marg_func(
RHS, sim.d[t, i], par)
euler_error_c[t, i] = sim.c[t, i]
