<jupyter_start><jupyter_text>simulation with real shock based on empirical data<jupyter_code>import pandas as pd
df_1999 = pd.read_csv("df_1999_30to60.csv")
df = df_1999[(df_1999["skillLevel"] == "High")&(df_1999["financeExperience"] == "Yes")]
df["ab"] = 30
df["wealth"] = df["liquidWealth"] + df["investmentAmount"]
codes = {'employed':1, 'unemployed': 0, "retired": 0}
df["employmentStatus"] = df["employmentStatus"].map(codes)
codes = {'owner':1, 'renter': 0}
df["ownership"] = df["ownership"].map(codes)
initialStates = df[["ageHead","wealth","ab","year","employmentStatus","ownership","participation"]]
initialStates["year"] = imaginedEconState[0]
initialStates = jnp.array(initialStates.values)

# risk free interest rate depending on current S state 
bondReturn = jnp.array(econRate[:,2])
# stock return depending on current S state
stockReturn = jnp.array(econRate[:,1])

@partial(jit, static_argnums=(0,))
def transition_real(t,age,a,x):
    '''
        Input:
            x = [w,ab,s,e,o,z] single action 
            x = [0,1, 2,3,4,5] 
            a = [c,b,k,h,action] single state
            a = [0,1,2,3,4]
        Output:
            w_next
            ab_next
            s_next
            e_next
            o_next
            z_next
            
            prob_next
    '''
    s = jnp.array(x[2], dtype = jnp.int8)
    e = jnp.array(x[3], dtype = jnp.int8)
    # actions taken
    b = a[1]
    k = a[2]
    action = a[4]
    w_next = ((1+bondReturn[t])*b + (1+stockReturn[t])*k).repeat(nE)
    ab_next = (1-x[4])*(t*(action == 1)).repeat(nE) + x[4]*(x[1]*jnp.ones(nE))
    s_next = econ[t+1].repeat(nE)
    e_next = jnp.array([e,(1-e)])*(t+age-20<T_R) + jnp.array([0,0])*(t+age-20>=T_R)
    z_next = x[5]*jnp.ones(nE) + ((1-x[5]) * (k > 0)).repeat(nE)
    # job status changing probability and econ state transition probability
    pe = Pe[s, e]
    prob_next = jnp.array([1-pe, pe])
    # owner
    o_next_own = (x[4] - action).repeat(nE)
    # renter
    o_next_rent = action.repeat(nE)
    o_next = x[4] * o_next_own + (1-x[4]) * o_next_rent   
    return jnp.column_stack((w_next,ab_next,s_next,e_next,o_next,z_next,prob_next))

'''
    # [w,ab,s,e,o,z]
    # w explicitly 
    # assume ab = 30 the strong assumption we made 
    # s is known 
    # e is known 
    # o is known
    # z is known
'''
from jax import random

def simulation(key, period = yearCount):
    x = initialStates[key.sum()%initialStates.shape[0]][1:]
    age = int(initialStates[key.sum()%initialStates.shape[0]][0])
    path = []
    move = []
    for t in range(0, period):
        key, subkey = random.split(key)
        if t == T_max-1:
            _,a = V_solve(t + age - 20,Vgrid[:,:,:,:,:,:,t + age - 20],x)
        else:
            _,a = V_solve(t + age - 20,Vgrid[:,:,:,:,:,:,t + age - 20],x)
        xp = transition_real(t,age,a,x)           
        p = xp[:,-1]
        x_next = xp[:,:-1]
        path.append(x)
        move.append(a)
        x = x_next[random.choice(a = nE, p=p, key = subkey)]
    path.append(x)
    return jnp.array(path), jnp.array(move)

# total number of agents
num = initialStates.shape[0] * 10
# simulation part 
keys = vmap(random.PRNGKey)(jnp.arange(num))
path = []
move = []
for i in tqdm(range(len(keys))):
    p,m = simulation(keys[i])
    path.append(p)
    move.append(m)
Paths = jnp.array(path)
Moves = jnp.array(move)
# x = [w,ab,s,e,o,z]
# x = [0,1, 2,3,4,5]
ws = Paths[:,:,0].T
ab = Paths[:,:,1].T
ss = Paths[:,:,2].T
es = Paths[:,:,3].T
os = Paths[:,:,4].T
zs = Paths[:,:,5].T
cs = Moves[:,:,0].T
bs = Moves[:,:,1].T
ks = Moves[:,:,2].T
hs = Moves[:,:,3].T
ms = Ms[jnp.append(jnp.array([0]),jnp.arange(yearCount)).reshape(-1,1) - jnp.array(ab, dtype = jnp.int8)]*os
np.save("w1", ws)
np.save("waseozcbkhm1", np.array([ws,ab,ss,es,os,zs,cs,bs,ks,hs,ms]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.20052151238591917, 'fraction_numerical': 0.03129074315514994, 'mean_word_length': 3.142548596112311, 'pattern_counts': {'":': 1, '<': 6, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '214451', 'license': 'MIT', 'n_tokens_mistral': 1634, 'n_tokens_neox': 1586, 'n_words': 288, 'path': '20211101/simulation_richLow.ipynb', 'repo_name': 'dongxulee/lifeCycleRefine', 'size': '3835'}