<jupyter_start><jupyter_text>State $$x = [w,n,m,s,e,o]$$   $w$: wealth level    size: 20   $n$: 401k level      size: 10   $m$: mortgage level  size: 10   $s$: economic state  size: 8   $e$: employment state size: 2   $o$: housing state:  size: 2    Action$c$: consumption amount size: 20   $b$: bond investment size: 20   $k$: stock investment derived from budget constrain once $c$ and $b$ are determined.    $h$: housing consumption size, related to housing status and consumption level   If $O = 1$, the agent owns a house:    $A = [c, b, k, h=H, action = 1]$ sold the house    $A = [c, b, k, h=H, action = 0]$ keep the house   If $O = 0$, the agent do not own a house:   $A = [c, b, k, h= \frac{c}{\alpha} \frac{1-\alpha}{pr}, action = 0]$ keep renting the house   $A = [c, b, k, h= \frac{c}{\alpha} \frac{1-\alpha}{pr}, action = 1]$ buy a housing with H unit      Housing20% down payment of mortgage, fix mortgage rate, single housing unit available, from age between 20 and 50, agents could choose to buy a house, and could choose to sell the house at any moment.  $H = 750$<jupyter_code>nX = Xs.shape[0]
nA = As.shape[0]
Xs.shape, As.shape<jupyter_output><empty_output><jupyter_text>Earning function part<jupyter_code>#Define the earning function, which applies for both employment status and 8 econ states
@partial(jit, static_argnums=(0,))
def y(t, x):
    '''
        x = [w,n,m,s,e,o]
        x = [0,1,2,3,4,5]
    '''
    if t <= T_R:
        return detEarning[t] * (1+gGDP[jnp.array(x[3], dtype = jnp.int8)]) * x[4] + (1-x[4]) * welfare
    else:
        return detEarning[-1]
    
#Earning after tax and fixed by transaction in and out from 401k account 
@partial(jit, static_argnums=(0,))
def yAT(t,x):
    yt = y(t, x)
    if t <= T_R:
        # yi portion of the income will be put into the 401k if employed
        return (1-tau_L)*(yt * (1-yi))*x[4] + (1-x[4])*yt
    else:
        # t > T_R, n/discounting amount will be withdraw from the 401k 
        return (1-tau_R)*yt + x[1]*Dn[t]
    
#Define the evolution of the amount in 401k account 
@partial(jit, static_argnums=(0,))
def gn(t, x, r = r_bar):
    if t <= T_R:
        # if the person is employed, then yi portion of his income goes into 401k 
        n_cur = x[1] + y(t, x) * yi * x[4]
    else:
        # t > T_R, n*Dn amount will be withdraw from the 401k 
        n_cur = x[1] - x[1]*Dn[t]
        # the 401 grow with the rate r 
    return (1+r)*n_cur
#Define the utility function
@jit
def u(c):
    return (jnp.power(c, 1-gamma) - 1)/(1 - gamma)

#Define the bequeath function, which is a function of bequeath wealth
@jit
def uB(tb):
    return B*u(tb)

#Reward function depends on the housing and non-housing consumption
@jit
def R(x,a):
    '''
    Input:
        x = [w,n,m,s,e,o]
        x = [0,1,2,3,4,5]
        a = [c,b,k,h,action]
        a = [0,1,2,3,4]
    '''
    c = a[:,0]
    h = a[:,3]
    C = jnp.power(c, alpha) * jnp.power(h, 1-alpha)
    return u(C)
# pc*qc / (ph*qh) = alpha/(1-alpha)

@partial(jit, static_argnums=(0,))
def feasibleActions(t, x):
    # owner
    sell = As[:,2]
    budget1 = yAT(t,x) + x[0] - (1-sell)*(x[2]*(1+rh)-x[2]*Dm[t]) + sell*(H*pt - x[2] - c_s) + (1-sell)*(((t<=T_R)*tau_L + (t>T_R)*tau_R)*x[2]*rh)
    # last term is the tax deduction of the interest portion of mortgage payment    
    h = jnp.ones(nA)*H*(1+kappa)*(1-sell) + sell*jnp.clip(budget1*As[:,0]*(1-alpha)/pr, a_max = 500)
    c = budget1*As[:,0]*(1-sell) + sell*(budget1*As[:,0] - h*pr)
    budget2 = budget1*(1-As[:,0])
    k = budget2*As[:,1]*(1-Kc)
    b = budget2*(1-As[:,1])
    owner_action = jnp.column_stack((c,b,k,h,sell))   
    # renter
    buy = As[:,2]
    budget1 = yAT(t,x) + x[0] - buy*(H*pt*0.2 + c_h)
    h = jnp.clip(budget1*As[:,0]*(1-alpha)/pr, a_max = 500)*(1-buy) + buy*jnp.ones(nA)*H*(1+kappa)
    c = (budget1*As[:,0] - h*pr)*(1-buy) + buy*budget1*As[:,0]
    budget2 = budget1*(1-As[:,0])
    k = budget2*As[:,1]*(1-Kc)
    b = budget2*(1-As[:,1])
    renter_action = jnp.column_stack((c,b,k,h,buy))
    
    actions = x[5]*owner_action + (1-x[5])*renter_action
    return actions

@partial(jit, static_argnums=(0,))
def transition(t,a,x):
    '''
        Input:
            x = [w,n,m,s,e,o]
            x = [0,1,2,3,4,5]
            a = [c,b,k,h,action]
            a = [0,1,2,3,4]
        Output:
            w_next
            n_next
            m_next
            s_next
            e_next
            o_next
            
            prob_next
    '''
    nA = a.shape[0]
    s = jnp.array(x[3], dtype = jnp.int8)
    e = jnp.array(x[4], dtype = jnp.int8)
    # actions taken
    b = a[:,1]
    k = a[:,2]
    action = a[:,4]
    w_next = ((1+r_b[s])*b + jnp.outer(k,(1+r_k)).T).T.flatten().repeat(nE)
    n_next = gn(t, x)*jnp.ones(w_next.size)
    s_next = jnp.tile(jnp.arange(nS),nA).repeat(nE)
    e_next = jnp.column_stack((e.repeat(nA*nS),(1-e).repeat(nA*nS))).flatten()
    # job status changing probability and econ state transition probability
    pe = Pe[s, e]
    ps = jnp.tile(Ps[s], nA)
    prob_next = jnp.column_stack(((1-pe)*ps,pe*ps)).flatten()
    
    # owner
    m_next_own = ((1-action)*x[2]*Dm[t]).repeat(nS*nE)
    o_next_own = (x[5] - action).repeat(nS*nE)
    # renter
    m_next_rent = (action*H*pt*0.8*Dm[t]).repeat(nS*nE)
    o_next_rent = action.repeat(nS*nE)
    
    m_next = x[5] * m_next_own + (1-x[5]) * m_next_rent
    o_next = x[5] * o_next_own + (1-x[5]) * o_next_rent   
    return jnp.column_stack((w_next,n_next,m_next,s_next,e_next,o_next,prob_next))

# used to calculate dot product
@jit
def dotProduct(p_next, uBTB):
    return (p_next*uBTB).reshape((p_next.shape[0]//(nS*nE), (nS*nE))).sum(axis = 1)

# define approximation of fit
@jit
def fit(v, xp):
    return map_coordinates(v,jnp.vstack((xp[:,0]/scaleW,
                                                      xp[:,1]/scaleN,
                                                      xp[:,2]/scaleM,
                                                      xp[:,3],
                                                      xp[:,4],
                                                      xp[:,5])),
                                                     order = 1, mode = 'nearest')

@partial(jit, static_argnums=(0,))
def V(t,V_next,x):
    '''
    x = [w,n,m,s,e,o]
    x = [0,1,2,3,4,5]
    xp:
        w_next    0
        n_next    1
        m_next    2
        s_next    3
        e_next    4
        o_next    5
        prob_next 6
    '''
    actions = feasibleActions(t,x)
    xp = transition(t,actions,x)
    # bequeath utility
    TB = xp[:,0]+x[1]*(1+r_bar)+xp[:,5]*(H*pt-x[2]*(1+rh))
    bequeathU = uB(TB)
    if t == T_max-1:
        Q = R(x,actions) + beta * dotProduct(xp[:,6], bequeathU)
    else:
        Q = R(x,actions) + beta * dotProduct(xp[:,6], Pa[t]*fit(V_next, xp) + (1-Pa[t])*bequeathU)
    Q = jnp.nan_to_num(Q, nan = -100)
    v = Q.max()
    cbkha = actions[Q.argmax()]
    return v, cbkha
%%time
for t in tqdm(range(T_max-1,T_min-1, -1)):
    if t == T_max-1:
        v,cbkha = vmap(partial(V,t,Vgrid[:,:,:,:,:,:,t]))(Xs)
    else:
        v,cbkha = vmap(partial(V,t,Vgrid[:,:,:,:,:,:,t+1]))(Xs)
    Vgrid[:,:,:,:,:,:,t] = v.reshape(dim)
    cgrid[:,:,:,:,:,:,t] = cbkha[:,0].reshape(dim)
    bgrid[:,:,:,:,:,:,t] = cbkha[:,1].reshape(dim)
    kgrid[:,:,:,:,:,:,t] = cbkha[:,2].reshape(dim)
    hgrid[:,:,:,:,:,:,t] = cbkha[:,3].reshape(dim)
    agrid[:,:,:,:,:,:,t] = cbkha[:,4].reshape(dim)
'''
    x = [w,n,m,s,e,o]
'''
wealthLevel = 0
retirement = 5
mortgage = 5
econState = 20
employ = 0
house = 1

# plt.figure(figsize = [12,6])
# plt.plot(cgrid[wealthLevel,retirement,mortgage,econState,employ,house,:], label = "consumption")
# plt.plot(bgrid[wealthLevel,retirement,mortgage,econState,employ,house,:], label = "bond")
# plt.plot(kgrid[wealthLevel,retirement,mortgage,econState,employ,house,:], label = "stock")
# plt.plot(hgrid[wealthLevel,retirement,mortgage,econState,employ,house,:], label = "housing")
plt.plot(agrid[wealthLevel,retirement,mortgage,econState,1,house,:], label = "action_employed")
plt.plot(agrid[wealthLevel,retirement,mortgage,econState,0,house,:], label = "action_unemployed")
legend()
np.save("Value",Vgrid)
np.save("cgrid",cgrid)
np.save("bgrid",bgrid)
np.save("kgrid",kgrid)
np.save("hgrid",hgrid)
np.save("agrid",agrid)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19036752342061014, 'fraction_numerical': 0.0313475858755705, 'mean_word_length': 2.697602131438721, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '18308', 'license': 'MIT', 'n_tokens_mistral': 3569, 'n_tokens_neox': 3482, 'n_words': 717, 'path': '20210516/.ipynb_checkpoints/housing-time-checkpoint.ipynb', 'repo_name': 'dongxulee/lifeCycle', 'size': '8326'}