def v_iter_single(setup, t, EV, female, ushift, force_f32=False):
agrid_s = setup.agrid_s
sgrid_s = setup.sgrid_s
dtype = setup.dtype
using_pre_u = setup.usinglef_precomputed_u if female else setup.usinglem_precomputed_u
using_pre_x = setup.usinglef_precomputed_x if female else setup.usinglem_precomputed_x
zvals = setup.exogrid.zf_t[t] if female else setup.exogrid.zm_t[t]
ztrend = setup.pars['f_wage_trend'][t] if female else setup.pars[
'm_wage_trend'][t]
#sigma = setup.pars['crra_power']
beta = setup.pars['beta_t'][t]
R = setup.pars['R_t'][t]
dtype_here = np.float32 if force_f32 else dtype
ls = np.array([setup.ls_levels[-1]], dtype=dtype) if female else np.array(
[setup.mlevel], dtype=dtype)
money_t = (R * agrid_s, np.exp(zvals + ztrend), np.zeros_like(zvals))
if EV is None:
EV = np.zeros((agrid_s.size, zvals.size), dtype=dtype_here)
else:
EV = EV.astype(dtype_here, copy=False)
assert EV.dtype == dtype_here
V_0, c_opt, x_opt, s_opt, i_opt, _, _ = \
v_optimize_couple(money_t,sgrid_s,(setup.vsgrid_s,EV[:,:,None,None]),setup.mgrid,
using_pre_u[:,None,None],
using_pre_x[:,None,None],
ls,beta,ushift,dtype=dtype)
V_0, c_opt, x_opt, s_opt, i_opt = \
(x.squeeze(axis=2) for x in [V_0, c_opt, x_opt, s_opt, i_opt])
# wf = money_t[1].reshape((1,money_t[1].size))*0.1
# wm = money_t[2].reshape((1,money_t[2].size))
# asset_income = money_t[0].reshape((money_t[0].size,1))
# money = wf + wm + asset_income
# aaa=money-c_opt-x_opt-s_opt
EVexp = setup.vsgrid_s.apply_preserve_shape(EV)
V_ret = setup.u_single_pub(
c_opt, x_opt,
ls) + ushift + beta * np.take_along_axis(EVexp, i_opt, 0)
assert V_ret.dtype == dtype
def r(x):
return x
return r(V_ret), r(c_opt), r(x_opt), r(s_opt)
def v_iter_single(setup, t, EV, female, ushift):
agrid_s = setup.agrid_s if nogpu else setup.cupy.agrid_s
sgrid_s = setup.sgrid_s if nogpu else setup.cupy.sgrid_s
dtype = setup.dtype
zvals = co(setup.exogrid.zf_t[t]) if female else co(setup.exogrid.zm_t[t])
ztrend = setup.pars['f_wage_trend'][t] if female else setup.pars[
'm_wage_trend'][t]
beta = setup.pars['beta_t'][t]
R = setup.pars['R_t'][t]
money_t = (R * agrid_s, np.exp(zvals + ztrend), np.zeros_like(zvals))
if EV is None:
EV = np.zeros((agrid_s.size, zvals.size), dtype=setup.dtype)
sname = 'Female, single' if female else 'Male, single'
taxfun = setup.taxes[sname] if t < setup.pars['Tret'] else lambda x: 0.0 * x
uu = setup.u_precomputed[sname][
'u'] if nogpu else setup.cupy.u_precomputed[sname]['u']
ux = setup.u_precomputed[sname][
'x'] if nogpu else setup.cupy.u_precomputed[sname]['x']
mgrid = setup.mgrid_s if nogpu else setup.cupy.mgrid_s
ls = setup.ls_levels[sname]
vsgrid = setup.vsgrid_s if nogpu else setup.cupy.vsgrid_s
V, c, x, s, i_opt, ils, V_all_l = \
v_optimize_couple(money_t,sgrid_s,(vsgrid,EV[...,None,None]),mgrid,uu,ux,
ls,beta,ushift,dtype=dtype,taxfun=taxfun)
assert np.all(x < 1e-5) # no x and no labor supply
V_ret, c_opt, s_opt, i_opt = (q.squeeze(axis=2) for q in (V, c, s, i_opt))
c_refine = R * agrid_s[:, None] + np.exp(zvals + ztrend)[None, :] - s_opt
vs = setup.vsgrid_s if nogpu else setup.cupy.vsgrid_s
EVexp = vs.apply_preserve_shape(EV)
V_refine = setup.u(c_refine) + ushift + beta * np.take_along_axis(
EVexp, i_opt, 0)
def r(x):
return x.astype(dtype, copy=False)
return r(V_refine), r(c_opt), r(s_opt)
def v_iter_single_mom(setup,t,EV,ushift):
agrid_s = setup.agrid_s if nogpu else setup.cupy.agrid_s
sgrid_s = setup.sgrid_s if nogpu else setup.cupy.sgrid_s
ls = setup.ls_levels['Female and child']
taxfun = setup.taxes['Female and child']
dtype = setup.dtype
zvals = co(setup.exogrid.zf_t[t])
ztrend = setup.pars['f_wage_trend'][t]
R = setup.pars['R_t'][t]
beta = setup.pars['beta_t'][t]
wm0 = np.zeros_like(zvals + ztrend)
money_t = (R*agrid_s,np.exp(zvals + ztrend),wm0)
if EV is None:
EV = np.zeros((agrid_s.size,zvals.size,ls.size),dtype=setup.dtype)
vs = setup.vsgrid_s if nogpu else setup.cupy.vsgrid_s
# I add virtual theta axis
EV_t = (vs, EV[...,None,:])
uu = setup.u_precomputed['Female and child']['u'] if nogpu else setup.cupy.u_precomputed['Female and child']['u']
ux = setup.u_precomputed['Female and child']['x'] if nogpu else setup.cupy.u_precomputed['Female and child']['x']
V, c, x, s, i_opt, ils, V_all_l = \
v_optimize_couple(money_t,sgrid_s,EV_t,setup.mgrid_s,uu,ux,
ls,beta,ushift,taxfun=taxfun,dtype=dtype)
# remove the virtual axis
V_old, c, x, s, i_opt, ils, V_all_l = (q.squeeze(axis=2) for q in (V, c, x, s, i_opt, ils, V_all_l))
u = setup.u_single_k(c,x,ils,ushift)
EVexp = vs.apply_preserve_shape(EV)
V_refine = u + beta*np.take_along_axis(np.take_along_axis(EVexp,i_opt[...,None],0),ils[...,None],2).squeeze(axis=2)
#print('After retirement (sm) max diff is {}'.format(np.max(np.abs(V_old-V_refine))))
def r(x): return x.astype(dtype,copy=False)
return r(V_refine), r(c), r(x), r(s), r(ils)
def v_iter_couple(setup,t,EV_tuple,ushift,haschild,nbatch=nbatch_def,verbose=False):
if verbose: start = default_timer()
agrid = setup.agrid_c
sgrid = setup.sgrid_c
dtype = setup.dtype
key = 'Couple and child' if haschild else 'Couple, no children'
taxfun = setup.taxes[key] if t < setup.pars['Tret'] else lambda x : 0.0*x
ls = setup.ls_levels[key]
uu, ux = setup.u_precomputed[key]['u'],setup.u_precomputed[key]['x']
upart, ucouple = (setup.u_part_k, setup.u_couple_k) if haschild else (setup.u_part_nk, setup.u_couple_nk)
nls = len(ls)
# type conversion is here
zf = setup.exogrid.all_t[t][:,0]
zm = setup.exogrid.all_t[t][:,1]
zftrend = setup.pars['f_wage_trend'][t]
zmtrend = setup.pars['m_wage_trend'][t]
psi = setup.exogrid.all_t[t][:,2]
beta = setup.pars['beta_t'][t]
sigma = setup.pars['crra_power']
R = setup.pars['R_t'][t]
wf = np.exp(zf + zftrend)
wm = np.exp(zm + zmtrend)
#labor_income = np.exp(zf) + np.exp(zm)
#money = R*agrid[:,None] + wf[None,:]
nexo = setup.pars['nexo_t'][t]
shp = (setup.na,nexo,setup.ntheta)
if EV_tuple is None:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = np.zeros((4,) + shp + (nls,),dtype=setup.dtype)
else:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = EV_tuple
# type conversion
sgrid,sigma,beta = (dtype(x) for x in (sgrid,sigma,beta))
V_couple, c_opt, s_opt, x_opt = np.empty((4,)+shp,dtype)
i_opt, il_opt = np.empty(shp,np.int16), np.empty(shp,np.int16)
V_all_l = np.empty(shp+(nls,),dtype=dtype)
theta_val = dtype(setup.thetagrid)
# the original problem is max{umult*u(c) + beta*EV}
# we need to rescale the problem to max{u(c) + beta*EV_resc}
istart = 0
ifinish = nbatch if nbatch < nexo else nexo
# this natually splits everything onto slices
for ibatch in range(int(np.ceil(nexo/nbatch))):
#money_i = money[:,istart:ifinish]
assert ifinish > istart
money_t = (R*agrid, wf[istart:ifinish], wm[istart:ifinish])
EV_t = (setup.vsgrid_c,EVr_by_l[:,istart:ifinish,:,:])
V_pure_i, c_opt_i, x_opt_i, s_opt_i, i_opt_i, il_opt_i, V_all_l_i = \
v_optimize_couple(money_t,sgrid,EV_t,setup.mgrid_c,uu,ux,
ls,beta,ushift,taxfun=taxfun,dtype=dtype)
V_ret_i = V_pure_i + psi[None,istart:ifinish,None]
V_couple[:,istart:ifinish,:] = V_ret_i # this estimate of V can be improved
c_opt[:,istart:ifinish,:] = c_opt_i
s_opt[:,istart:ifinish,:] = s_opt_i
i_opt[:,istart:ifinish,:] = i_opt_i
x_opt[:,istart:ifinish,:] = x_opt_i
il_opt[:,istart:ifinish,:] = il_opt_i
V_all_l[:,istart:ifinish,:,:] = V_all_l_i # we need this for l choice so it is ok
istart = ifinish
ifinish = ifinish+nbatch if ifinish+nbatch < nexo else nexo
if verbose: print('Batch {} done at {} sec'.format(ibatch,default_timer()-start))
assert np.all(c_opt > 0)
psi_r = psi[None,:,None].astype(setup.dtype,copy=False)
# finally obtain value functions of partners
uf, um = upart(c_opt,x_opt,il_opt,theta_val[None,None,:],ushift,psi_r)
uc = ucouple(c_opt,x_opt,il_opt,theta_val[None,None,:],ushift,psi_r)
EVf_all, EVm_all, EVc_all = (setup.vsgrid_c.apply_preserve_shape(x) for x in (EV_fem_by_l, EV_mal_by_l,EVc_by_l))
V_fem = uf + beta*np.take_along_axis(np.take_along_axis(EVf_all,i_opt[...,None],0),il_opt[...,None],3).squeeze(axis=3)
V_mal = um + beta*np.take_along_axis(np.take_along_axis(EVm_all,i_opt[...,None],0),il_opt[...,None],3).squeeze(axis=3)
V_all = uc + beta*np.take_along_axis(np.take_along_axis(EVc_all,i_opt[...,None],0),il_opt[...,None],3).squeeze(axis=3)
def r(x): return x
'''
if V_all.shape[-1] > 1:
V_weighted = V_fem*setup.thetagrid[None,None,:] + V_mal*(1-setup.thetagrid[None,None,:])
vd_max = np.max(np.abs(V_weighted - V_all))
vd_mean = np.mean(np.abs(V_weighted - V_all))
print('max diff is {}, mean diff is {}'.format(vd_max,vd_mean))
'''
assert V_all.dtype==EVc_all.dtype==V_couple.dtype
'''
if uf.shape[-1] > 1:
assert np.all(np.diff(uf,axis=-1)>0), 'Monotonicity is violated for females in u'
assert np.all(np.diff(um,axis=-1)<0), 'Monotonicity is violated for males in u'
if V_fem.shape[-1] > 1:
assert np.all(np.diff(V_fem,axis=-1)>0), 'Monotonicity is violated for females in V'
assert np.all(np.diff(V_mal,axis=-1)<0), 'Monotonicity is violated for males in V'
'''
return r(V_all), r(V_fem), r(V_mal), r(c_opt), r(x_opt), r(s_opt), il_opt, r(V_all_l)
def v_iter_couple(model,
t,
EV_tuple,
ushift,
haschild,
nbatch=500,
verbose=False):
setup = model.setup
if verbose: start = default_timer()
if gpu:
np = cp
else:
np = onp
nogpu = not gpu
agrid = setup.agrid_c if nogpu else setup.cupy.agrid_c
sgrid = setup.sgrid_c if nogpu else setup.cupy.sgrid_c
dtype = setup.dtype
key = 'Couple and child' if haschild else 'Couple, no children'
taxfun = setup.taxes[key] if t < setup.pars['Tret'] else lambda x: 0.0 * x
ls = setup.ls_levels[key]
if nogpu:
uu, ux, mgrid = setup.u_precomputed[key]['u'], setup.u_precomputed[
key]['x'], setup.mgrid_c
else:
uu, ux, mgrid = setup.cupy.u_precomputed[key][
'u'], setup.cupy.u_precomputed[key]['x'], setup.cupy.mgrid_c
upart, ucouple = (setup.u_part_k,
setup.u_couple_k) if haschild else (setup.u_part_nk,
setup.u_couple_nk)
nls = len(ls)
# type conversion is here
all_t = np.array(setup.exogrid.all_t[t], copy=False)
zf = all_t[:, 0]
zm = all_t[:, 1]
zftrend = setup.pars['f_wage_trend'][t]
zmtrend = setup.pars['m_wage_trend'][t]
psi = all_t[:, 2]
beta = setup.pars['beta_t'][t]
sigma = setup.pars['crra_power']
R = setup.pars['R_t'][t]
wf = np.exp(zf + zftrend)
wm = np.exp(zm + zmtrend)
nexo = setup.pars['nexo_t'][t]
shp = (setup.na, nexo, setup.ntheta)
if EV_tuple is None:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = np.zeros(
(4, ) + shp + (nls, ), dtype=setup.dtype)
else:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = EV_tuple
# type conversion
sigma, beta = (dtype(x) for x in (sigma, beta))
V_couple, c_opt, s_opt, x_opt = np.empty((4, ) + shp, dtype)
i_opt, il_opt = np.empty(shp, np.int16), np.empty(shp, np.int16)
sgrid = sgrid.astype(dtype, copy=False)
theta_val = dtype(setup.thetagrid) if nogpu else setup.cupy.thetagrid
# the original problem is max{umult*u(c) + beta*EV}
# we need to rescale the problem to max{u(c) + beta*EV_resc}
istart = 0
ifinish = nbatch if nbatch < nexo else nexo
# this natually splits everything onto slices
vs = setup.vsgrid_c if nogpu else setup.cupy.vsgrid_c
for ibatch in range(int(np.ceil(nexo / nbatch))):
#money_i = money[:,istart:ifinish]
assert ifinish > istart
money_t = (R * agrid, wf[istart:ifinish], wm[istart:ifinish])
EV_t = (vs, EVr_by_l[:, istart:ifinish, :, :])
V_pure_i, c_opt_i, x_opt_i, s_opt_i, i_opt_i, il_opt_i, V_all_l_i = \
v_optimize_couple(money_t,sgrid,EV_t,mgrid,uu,ux,
ls,beta,ushift,taxfun=taxfun,dtype=dtype)
V_ret_i = V_pure_i + psi[None, istart:ifinish, None]
V_couple[:, istart:
ifinish, :] = V_ret_i # this estimate of V can be improved
c_opt[:, istart:ifinish, :] = c_opt_i
s_opt[:, istart:ifinish, :] = s_opt_i
i_opt[:, istart:ifinish, :] = i_opt_i
x_opt[:, istart:ifinish, :] = x_opt_i
il_opt[:, istart:ifinish, :] = il_opt_i
#V_all_l[:,istart:ifinish,:,:] = V_all_l_i # we need this for l choice so it is ok
istart = ifinish
ifinish = ifinish + nbatch if ifinish + nbatch < nexo else nexo
if verbose:
print('Batch {} done at {} sec'.format(ibatch,
default_timer() - start))
assert np.all(c_opt > 0)
psi_r = psi[None, :, None].astype(setup.dtype, copy=False)
# finally obtain value functions of partners
uf, um = upart(c_opt, x_opt, il_opt, theta_val[None, None, :], ushift,
psi_r)
uc = ucouple(c_opt, x_opt, il_opt, theta_val[None, None, :], ushift, psi_r)
EVf_all, EVm_all, EVc_all = (vs.apply_preserve_shape(x)
for x in (EV_fem_by_l, EV_mal_by_l, EVc_by_l))
V_fem = uf + beta * np.take_along_axis(
np.take_along_axis(EVf_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
V_mal = um + beta * np.take_along_axis(
np.take_along_axis(EVm_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
V_all = uc + beta * np.take_along_axis(
np.take_along_axis(EVc_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
def r(x):
return x
assert V_all.dtype == EVc_all.dtype == V_couple.dtype
return r(V_all), r(V_fem), r(V_mal), r(c_opt), r(x_opt), r(s_opt), il_opt
def v_iter_couple(setup,
t,
EV_tuple,
ushift,
nbatch=nbatch_def,
verbose=False,
force_f32=False):
if verbose: start = default_timer()
agrid = setup.agrid_c
sgrid = setup.sgrid_c
dtype = setup.dtype
ls = setup.ls_levels
nls = len(ls)
# type conversion is here
zf = setup.exogrid.all_t[t][:, 0]
zm = setup.exogrid.all_t[t][:, 1]
zftrend = setup.pars['f_wage_trend'][t]
zmtrend = setup.pars['m_wage_trend'][t]
psi = setup.exogrid.all_t[t][:, 2]
beta = setup.pars['beta_t'][t]
sigma = setup.pars['crra_power']
R = setup.pars['R_t'][t]
nexo = setup.pars['nexo_t'][t]
shp = (setup.na, nexo, setup.ntheta)
wf = np.exp(zf + zftrend)
wm = np.exp(zm + zmtrend)
dtype_here = np.float32 if force_f32 else dtype
if EV_tuple is None:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = np.zeros(
((4, ) + shp + (nls, )), dtype=dtype)
else:
EVr_by_l, EVc_by_l, EV_fem_by_l, EV_mal_by_l = EV_tuple
# type conversion
sgrid, sigma, beta = (dtype(x) for x in (sgrid, sigma, beta))
V_couple, c_opt, s_opt, x_opt = np.empty((4, ) + shp, dtype)
i_opt, il_opt = np.empty(shp, np.int16), np.empty(shp, np.int16)
V_all_l = np.empty(shp + (nls, ), dtype=dtype)
theta_val = dtype(setup.thetagrid)
# the original problem is max{umult*u(c) + beta*EV}
# we need to rescale the problem to max{u(c) + beta*EV_resc}
istart = 0
ifinish = nbatch if nbatch < nexo else nexo
#Time husband contribute to build Q
mt = 1.0 - setup.mlevel
# this natually splits everything onto slices
for ibatch in range(int(np.ceil(nexo / nbatch))):
#money_i = money[:,istart:ifinish]
assert ifinish > istart
money_t = (R * agrid, wf[istart:ifinish], wm[istart:ifinish])
EV_t = (setup.vsgrid_c, EVr_by_l[:, istart:ifinish, :, :])
V_pure_i, c_opt_i, x_opt_i, s_opt_i, i_opt_i, il_opt_i, V_all_l_i = \
v_optimize_couple(money_t,sgrid,EV_t,setup.mgrid,
setup.ucouple_precomputed_u,setup.ucouple_precomputed_x,
ls,beta,ushift,dtype=dtype_here,mt=mt)
V_ret_i = V_pure_i + psi[None, istart:ifinish, None]
# if dtype_here != dtype type conversion happens here
V_couple[:, istart:
ifinish, :] = V_ret_i # this estimate of V can be improved
c_opt[:, istart:ifinish, :] = c_opt_i
s_opt[:, istart:ifinish, :] = s_opt_i
i_opt[:, istart:ifinish, :] = i_opt_i
x_opt[:, istart:ifinish, :] = x_opt_i
il_opt[:, istart:ifinish, :] = il_opt_i
V_all_l[:, istart:
ifinish, :, :] = V_all_l_i # we need this for l choice so it is ok
istart = ifinish
ifinish = ifinish + nbatch if ifinish + nbatch < nexo else nexo
if verbose:
print('Batch {} done at {} sec'.format(ibatch,
default_timer() - start))
assert np.all(c_opt > 0)
psi_r = psi[None, :, None].astype(setup.dtype, copy=False)
# finally obtain value functions of partners
uf, um = setup.u_part(c_opt, x_opt, il_opt, theta_val[None, None, :],
ushift, psi_r)
uc = setup.u_couple(c_opt, x_opt, il_opt, theta_val[None, None, :], ushift,
psi_r)
EVf_all, EVm_all, EVc_all = (setup.vsgrid_c.apply_preserve_shape(x)
for x in (EV_fem_by_l, EV_mal_by_l, EVc_by_l))
V_fem = uf + beta * np.take_along_axis(
np.take_along_axis(EVf_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
V_mal = um + beta * np.take_along_axis(
np.take_along_axis(EVm_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
V_all = uc + beta * np.take_along_axis(
np.take_along_axis(EVc_all, i_opt[..., None], 0), il_opt[..., None],
3).squeeze(axis=3)
#def r(x): return x.astype(dtype)
def r(x):
return x
assert V_all.dtype == dtype
assert V_fem.dtype == dtype
assert V_mal.dtype == dtype
assert c_opt.dtype == dtype
assert x_opt.dtype == dtype
assert s_opt.dtype == dtype
try:
assert np.allclose(V_all, V_couple, atol=1e-4, rtol=1e-3)
except:
#print('max difference in V is {}'.format(np.max(np.abs(V_all-V_couple))))
pass
return r(V_all), r(V_fem), r(V_mal), r(c_opt), r(x_opt), r(
s_opt), il_opt, r(V_all_l)
