def compute_q(sol, par, t, t_plus):
""" compute the post-decision function q """
# unpack (helps numba optimize)
q = sol.q
# loop over delta
for idelta in prange(par.Ndelta):
# clean-up
q[t, idelta, :] = 0
# temp
m_plus = np.empty(par.Na)
c_plus = np.empty(par.Na)
# loop over shock
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]
weight = psi_w * xi_w
# ii. next-period extra income component
delta_plus = par.grid_delta[idelta] / (psi * par.G)
if t == 0:
delta_plus *= par.zeta
# ii. next-period cash-on-hand
for ia in range(par.Na):
m_plus[ia] = par.R * par.grid_a[ia] / (psi * par.G) + xi
# iii. next-period consumption
if par.Ndelta > 1:
prep = linear_interp.interp_2d_prep(par.grid_delta, delta_plus,
par.Na)
else:
prep = linear_interp.interp_1d_prep(par.Na)
if par.Ndelta > 1:
linear_interp.interp_2d_only_last_vec_mon(
prep, par.grid_delta, par.grid_m, sol.c[t_plus],
delta_plus, m_plus, c_plus)
else:
linear_interp.interp_1d_vec_mon(prep, par.grid_m,
sol.c[t_plus,
idelta], m_plus, c_plus)
# iv. accumulate all
for ia in range(par.Na):
q[t, idelta,
ia] += weight * par.R * par.beta * utility.marg_func(
par.G * psi * c_plus[ia], par)
def shocks_GH(t, inc_no_shock, inc, w, c, m, v, par, d_plus):
""" compute v_plus_raw and avg_marg_u_plus using GaussHermite integration if necessary """
# a. initialize
c_plus_interp = np.zeros((4, inc_no_shock.size))
v_plus_interp = np.zeros((4, inc_no_shock.size))
v_plus_raw = np.zeros(inc_no_shock.size)
avg_marg_u_plus = np.zeros(inc_no_shock.size)
prep = linear_interp.interp_1d_prep(
len(inc_no_shock
)) # save the position of numbers to speed up interpolation
# b. loop over GH-nodes
for i in range(len(w)):
m_plus = inc_no_shock + inc[i]
# 1. interpolate
for d in d_plus:
linear_interp.interp_1d_vec_mon(prep, m[:], c[d, :], m_plus,
c_plus_interp[d, :])
linear_interp.interp_1d_vec_mon_rep(prep, m[:], v[d, :], m_plus,
v_plus_interp[d, :])
# 2. logsum and v_plus_raw
if len(d_plus) == 1: # no taste shocks
v_plus_raw += w[i] * v_plus_interp[d_plus[0], :]
avg_marg_u_plus += w[i] * utility.marg_func(
c_plus_interp[d_plus[0], :], par)
elif len(d_plus) == 2: # taste shocks
logsum, prob = funs.logsum2(v_plus_interp[d_plus, :], par)
v_plus_raw += w[i] * logsum[0, :]
marg_u_plus = prob[d_plus[0], :] * utility.marg_func(
c_plus_interp[d_plus[0], :],
par) + (1 - prob[d_plus[0], :]) * utility.marg_func(
c_plus_interp[d_plus[1], :], par)
avg_marg_u_plus += w[i] * marg_u_plus
elif len(d_plus) == 4: # both are working
logsum, prob = funs.logsum4(v_plus_interp[d_plus, :], par)
v_plus_raw += w[i] * logsum[0, :]
marg_u_plus = (
prob[0, :] * utility.marg_func(c_plus_interp[0, :], par) +
prob[1, :] * utility.marg_func(c_plus_interp[1, :], par) +
prob[2, :] * utility.marg_func(c_plus_interp[2, :], par) +
prob[3, :] * utility.marg_func(c_plus_interp[3, :], par))
avg_marg_u_plus += w[i] * marg_u_plus
# c. return results
return v_plus_raw, avg_marg_u_plus
def ConsValue_c(t, ad, st_h, st_w, ra_h, ra_w, d_h, d_w, m, sol, par, idx):
""" interpolate consumption and value for couple model """
# a. unpack solution
D = transitions.d_plus_c(t - 1, ad, d_h, d_w,
par) # t-1 so we get choice set today
ra_look_h = transitions.ra_look_up(t, st_h, ra_h, d_h, par)
ra_look_w = transitions.ra_look_up(t + ad, st_w, ra_w, d_w, par)
ad_idx = ad + par.ad_min
c_sol = sol.c[t, ad_idx, st_h, st_w, ra_look_h, ra_look_w]
m_sol = sol.m[:]
v_sol = sol.v[t, ad_idx, st_h, st_w, ra_look_h, ra_look_w]
# a. initialize interpolation
prep = linear_interp.interp_1d_prep(idx.size)
c_interp = np.zeros(
(idx.size,
4)) # note c_interp and v_interp are transposed of each other
v_interp = np.zeros((4, idx.size))
# b. sort m (so interp is faster)
idx_unsort = np.argsort(np.argsort(m[idx, t])) # index to unsort
m_sort = np.sort(m[idx, t])
# c. interpolate and sort back
if d_h == 0 and d_w == 0:
for d in D:
linear_interp.interp_1d_vec_mon(prep, m_sol[:], c_sol[d], m_sort,
c_interp[:, d])
# sort back
c_interp = c_interp[idx_unsort]
else:
for d in D:
linear_interp.interp_1d_vec_mon(prep, m_sol[:], c_sol[d], m_sort,
c_interp[:, d])
linear_interp.interp_1d_vec_mon(prep, m_sol[:], v_sol[d], m_sort,
v_interp[d])
# sort back
c_interp = c_interp[idx_unsort]
v_interp = v_interp[:, idx_unsort]
# d. return
return c_interp, v_interp
def ConsValue(t, ma, st, ra, ds, m, sol, par, idx, ad=0, ad_min=0):
""" interpolate consumption and value (ad and ad_min are only so the couple model can look up in this function) """
# a. unpack solution
ad_idx = ad + ad_min
D = transitions.d_plus(t + ad - 1, ds,
par) # t-1 so we get choice set today
ra_look = transitions.ra_look_up(t + ad, st, ra, ds, par)
c_sol = sol.c[t + ad_idx, ma, st, ra_look]
m_sol = sol.m[:]
v_sol = sol.v[t + ad_idx, ma, st, ra_look]
# a. initialize interpolation
prep = linear_interp.interp_1d_prep(idx.size)
c_interp = np.zeros(
(idx.size,
D.size)) # note c_interp and v_interp are transposed of each other
v_interp = np.zeros((D.size, idx.size))
# b. sort m (so interp is faster)
idx_unsort = np.argsort(np.argsort(m[idx, t])) # index to unsort
m_sort = np.sort(m[idx, t])
# c. interpolate and sort back
if ds == 1:
for d in D:
linear_interp.interp_1d_vec_mon(prep, m_sol[:], c_sol[d], m_sort,
c_interp[:, d])
linear_interp.interp_1d_vec_mon(prep, m_sol[:], v_sol[d], m_sort,
v_interp[d])
# sort back
c_interp = c_interp[idx_unsort]
v_interp = v_interp[:, idx_unsort]
elif ds == 0:
for d in D:
linear_interp.interp_1d_vec_mon(prep, m_sol[:], c_sol[d], m_sort,
c_interp[:, d])
# sort back
c_interp = c_interp[idx_unsort]
# d. return
return c_interp, v_interp
def compute(t, sol, par, G2EGM=True):
# unpack
w = sol.w[t]
wa = sol.wa[t]
if G2EGM:
wb = sol.wb[t]
# loop over outermost post-decision state
for i_b in range(par.Nb_pd):
# a. initialize
w[i_b, :] = 0
wa[i_b, :] = 0
if G2EGM:
wb[i_b, :] = 0
# b. working memoery
inv_v_plus = np.zeros(par.Na_pd)
inv_vm_plus = np.zeros(par.Na_pd)
if G2EGM:
inv_vn_plus = np.zeros(par.Na_pd)
inv_v_ret_plus = np.zeros(par.Na_pd)
inv_vm_ret_plus = np.zeros(par.Na_pd)
if G2EGM:
inv_vn_ret_plus = np.zeros(par.Na_pd)
# c. loop over shocks
for i_eta in range(par.Neta):
# i. next period states
m_plus = par.Ra * par.grid_a_pd + par.eta[i_eta]
n_plus = par.Rb * par.grid_b_pd[i_b]
m_plus_ret = m_plus + n_plus
# ii. prepare interpolation in p direction
prep = linear_interp.interp_2d_prep(par.grid_n, n_plus, par.Na_pd)
prep_ret = linear_interp.interp_1d_prep(par.Na_pd)
# iii. interpolations
# work
linear_interp.interp_2d_only_last_vec_mon(prep, par.grid_n,
par.grid_m,
sol.inv_v[t + 1], n_plus,
m_plus, inv_v_plus)
linear_interp.interp_2d_only_last_vec_mon_rep(
prep, par.grid_n, par.grid_m, sol.inv_vm[t + 1], n_plus,
m_plus, inv_vm_plus)
if G2EGM:
linear_interp.interp_2d_only_last_vec_mon_rep(
prep, par.grid_n, par.grid_m, sol.inv_vn[t + 1], n_plus,
m_plus, inv_vn_plus)
# retire
linear_interp.interp_1d_vec_mon(prep_ret, sol.m_ret[t + 1],
sol.inv_v_ret[t + 1], m_plus_ret,
inv_v_ret_plus)
linear_interp.interp_1d_vec_mon_rep(prep_ret, sol.m_ret[t + 1],
sol.inv_vm_ret[t + 1],
m_plus_ret, inv_vm_ret_plus)
if G2EGM:
linear_interp.interp_1d_vec_mon_rep(prep_ret, sol.m_ret[t + 1],
sol.inv_vn_ret[t + 1],
m_plus_ret,
inv_vn_ret_plus)
# iv. accumulate
for i_a in range(par.Na_pd):
if inv_v_ret_plus[i_a] > inv_v_plus[i_a]:
w_now = -1.0 / inv_v_ret_plus[i_a]
wa_now = 1.0 / inv_vm_ret_plus[i_a]
if G2EGM:
wb_now = 1.0 / inv_vn_ret_plus[i_a]
else:
w_now = -1.0 / inv_v_plus[i_a]
wa_now = 1.0 / inv_vm_plus[i_a]
if G2EGM:
wb_now = 1.0 / inv_vn_plus[i_a]
w[i_b, i_a] += par.w_eta[i_eta] * par.beta * w_now
wa[i_b, i_a] += par.w_eta[i_eta] * par.Ra * par.beta * wa_now
if G2EGM:
wb[i_b,
i_a] += par.w_eta[i_eta] * par.Rb * par.beta * wb_now