#!/home/a.ghaderi/.conda/envs/envjm/bin/python
# Model 2
import pystan 
import pandas as pd
import numpy as np
import sys
sys.path.append('../../')
import utils

parts = 1 
data = utils.get_data()          #loading dateset
data = data[data['participant']==parts]

mis = np.where((data['n200lat']<.101)|(data['n200lat']>.248))[0] # missing data for n200lat
obs = np.where((data['n200lat']>.101)&(data['n200lat']<.248))[0] # observation and missing data for n200lat

N_mis = mis.shape[0]  # number of missing data
N_obs = obs.shape[0]   # number of observed data

modelfile = '../../stans/res_nonhier.stan' #reading the model span
f = open(modelfile, 'r')
model_wiener = f.read()
sm = pystan.StanModel(model_code=model_wiener)# Compile the model stan

ncohers = 2  #Number of coherence conditions
nspats = 2   #Number of spatial conditions
nconds = 4   #Number of conditions
y = data['y'].to_numpy()
cond_coher = data['cond_coher'].to_numpy()
cond_spat = data['cond_spat'].to_numpy()
conds = data['conds'].to_numpy()
n200lat = data['n200lat'].to_numpy()

#set inistial data for molde span
data_winner =  {'N_obs':N_obs,      #Number of trial-level observations
                'N_mis':N_mis,      #Number of trial-level mising data
                'ncohers':ncohers,  #Number of coherence conditions
                'nspats':nspats,    #Number of spatial conditions
                'nconds':nconds,    #Number of conditions
                'y':np.concatenate([y[obs],y[mis]]),     #acc*rt in seconds for obervation and missing data
                'cond_coher':np.concatenate([cond_coher[obs],cond_coher[mis]]),     #Coherence index for each trial 
                'cond_spat':np.concatenate([cond_spat[obs],cond_spat[mis]]),        #sptial index for each trial 
                'conds':np.concatenate([conds[obs],conds[mis]]),        #sptial index for each trial
                'n200lat_obs':n200lat[obs]};    #n200 latency for each trial observation

# setting MCMC arguments
niter = 10000
nwarmup = 4000
nchains = 1
thin = 1

initials = [] # initial sampling
for c in range(0, nchains):
    chaininit = {
        'delta': np.random.uniform(1, 3, size=ncohers),
        'alpha': np.random.uniform(.5, 1.),
        'eta': np.random.uniform(.01, .2),
        'res': np.random.uniform(.01, .02, size=nspats),     
        'n200sub': np.random.uniform(.11, .2, size=nconds),
        'lambda': np.random.uniform(.01, .02),
        'n200lat_mis': np.random.uniform(.11, .2, size = N_mis)
    }
    initials.append(chaininit)    
      
# Train the model and generate samples
fit = sm.sampling(data=data_winner, iter=niter, chains=nchains, warmup=nwarmup, thin=thin, init=initials)

utils.to_pickle(stan_model=sm, stan_fit=fit, save_path='../../save/nonhier/'+str(parts)+'_res_nonhier.pkl')


[((122, 147), 'sys.path.append', 'sys.path.append', (['"""../../"""'], {}), "('../../')\n", (137, 147), False, 'import sys\n'), ((180, 196), 'utils.get_data', 'utils.get_data', ([], {}), '()\n', (194, 196), False, 'import utils\n'), ((683, 724), 'pystan.StanModel', 'pystan.StanModel', ([], {'model_code': 'model_wiener'}), '(model_code=model_wiener)\n', (699, 724), False, 'import pystan\n'), ((270, 333), 'numpy.where', 'np.where', (["((data['n200lat'] < 0.101) | (data['n200lat'] > 0.248))"], {}), "((data['n200lat'] < 0.101) | (data['n200lat'] > 0.248))\n", (278, 333), True, 'import numpy as np\n'), ((362, 425), 'numpy.where', 'np.where', (["((data['n200lat'] > 0.101) & (data['n200lat'] < 0.248))"], {}), "((data['n200lat'] > 0.101) & (data['n200lat'] < 0.248))\n", (370, 425), True, 'import numpy as np\n'), ((1442, 1474), 'numpy.concatenate', 'np.concatenate', (['[y[obs], y[mis]]'], {}), '([y[obs], y[mis]])\n', (1456, 1474), True, 'import numpy as np\n'), ((1559, 1609), 'numpy.concatenate', 'np.concatenate', (['[cond_coher[obs], cond_coher[mis]]'], {}), '([cond_coher[obs], cond_coher[mis]])\n', (1573, 1609), True, 'import numpy as np\n'), ((1675, 1723), 'numpy.concatenate', 'np.concatenate', (['[cond_spat[obs], cond_spat[mis]]'], {}), '([cond_spat[obs], cond_spat[mis]])\n', (1689, 1723), True, 'import numpy as np\n'), ((1785, 1825), 'numpy.concatenate', 'np.concatenate', (['[conds[obs], conds[mis]]'], {}), '([conds[obs], conds[mis]])\n', (1799, 1825), True, 'import numpy as np\n'), ((2124, 2161), 'numpy.random.uniform', 'np.random.uniform', (['(1)', '(3)'], {'size': 'ncohers'}), '(1, 3, size=ncohers)\n', (2141, 2161), True, 'import numpy as np\n'), ((2180, 2207), 'numpy.random.uniform', 'np.random.uniform', (['(0.5)', '(1.0)'], {}), '(0.5, 1.0)\n', (2197, 2207), True, 'import numpy as np\n'), ((2222, 2250), 'numpy.random.uniform', 'np.random.uniform', (['(0.01)', '(0.2)'], {}), '(0.01, 0.2)\n', (2239, 2250), True, 'import numpy as np\n'), ((2265, 2307), 'numpy.random.uniform', 'np.random.uniform', (['(0.01)', '(0.02)'], {'size': 'nspats'}), '(0.01, 0.02, size=nspats)\n', (2282, 2307), True, 'import numpy as np\n'), ((2331, 2372), 'numpy.random.uniform', 'np.random.uniform', (['(0.11)', '(0.2)'], {'size': 'nconds'}), '(0.11, 0.2, size=nconds)\n', (2348, 2372), True, 'import numpy as np\n'), ((2390, 2419), 'numpy.random.uniform', 'np.random.uniform', (['(0.01)', '(0.02)'], {}), '(0.01, 0.02)\n', (2407, 2419), True, 'import numpy as np\n'), ((2442, 2482), 'numpy.random.uniform', 'np.random.uniform', (['(0.11)', '(0.2)'], {'size': 'N_mis'}), '(0.11, 0.2, size=N_mis)\n', (2459, 2482), True, 'import numpy as np\n')]