<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Test Parallel Processing
This script is used to test whether the parallel processing works in the docker image. It includes two parts
Step1:  Part 1
Step2:  Part 2
Step3:  Test the docker image
Step4:  Test the docker image
Step5:  Check convergence (visual inspection)
plot multiple chain and visually inspect how the convergence goes.
Step7:  Test the docker image
<ASSISTANT_TASK:>
Python Code:
# check which python is in use.
import sys
print('Notebook is running:', sys.executable) # /opt/conda/bin/python
# or uncomment the code below
from platform import python_version # 3.7.6
print('The current Python version is', python_version())
import hddm, IPython
print('The current HDDM version is', hddm.__version__) # 0.8.0
import kabuki
print('The current kabuki version is', kabuki.__version__) # 0.6.3
import matplotlib
print('The current matplotlib version is', matplotlib.__version__) # 0.3.1.3
# Warning:`IPython.parallel` package has been deprecated since IPython 4.0. 
#print('The current IPython version is', IPython.__version__) 
import os
cur_dir = os.getcwd()
print('The current working directory is', cur_dir)
Explanation: Test Parallel Processing
This script is used to test whether the parallel processing works in the docker image. It includes two parts:
* Part 1: test parallel processing using easy calculations
* Part 2: test parallel processing using HDDM's modelling
Please make sure that you started multiple engines at the IPython Cluster
Note: Example data are from Hu et al. (2020)'s Experiment 2. This data can be downloaded from github too.
End of explanation
import ipyparallel as ipp
rc = ipp.Client()
rc.ids
def foo():
    import time
    time.sleep(5)
    return 'foo'
def bar():
    import time
    time.sleep(10) 
    return 'bar'  
res1 = rc[0].apply(foo)
res2 = rc[1].apply(bar)
results = [res1, res2]
while not all(map(lambda ar: ar.ready(), results)):
    pass
print(res1.get(), res2.get())
dv = rc[:]
res = dv.map(lambda x, y, z: x + y + z, range(10), range(10), range(10))
res.get()
@dv.remote(block = True)
def f1(n):
    import numpy as np
    return np.random.rand(n)
f1(4)
Explanation: Part 1: testing ipyparallel
Test computing using examples from this tutorial.
 There should no error for the code below
End of explanation
%matplotlib inline
# Preparation
import os, time, csv, datetime
import kabuki, hddm
from kabuki.analyze import gelman_rubin
from datetime import date
import random
import pandas as pd
import matplotlib.pyplot as plt
# check the data
# load data
df = hddm.load_csv('df_example.csv')
df.head(10)  # uncomment if you want to check the head of the data
# Uncomment the codes below if you are interested in checking the distribution of raw data
# plot and save the distribution of raw data:
fig = plt.figure()
ax = fig.add_subplot(111, xlabel='RT', ylabel='count', title='RT distributions')
for i, subj_data in df.groupby('subj_idx'):
    subj_data.rt.hist(bins=20, histtype='step', ax=ax)
# define a function to run model in parallel
def run_model(id):
    print('running model%i'%id);
    
    import hddm
    import random
    #import os
    
    #cur_dir = os.getcwd()
    exp_name = 'example'
    print('running models %i'%id, 'for for exp', exp_name)
    
    # USE the absolute directory in docker.
    dbname = '/home/jovyan/example/df_' + exp_name + '_chain_vaz_test_%i.db'%id # define the database name, which uses pickle format
    mname  = '/home/jovyan/example/df_' + exp_name + '_chain_vaz_test_%i'%id    # define the name for the model
    fname  = '/home/jovyan/example/df_' + exp_name + '.csv'
    df = hddm.load_csv(fname)
    df_subj = df['subj_idx'].unique()
    # random select without repetition
    random.seed(10)
    df_test_list = []
    for i in range(10):
        pos = random.randint(0, (len(df_subj)-1))
        df_test_list.append(df_subj[pos])  
    
    df_test = df[df['subj_idx'].isin(df_test_list)]
    
    m = hddm.HDDMStimCoding(df_test, 
                            include='z', 
                            stim_col='stim', 
                            depends_on={'v':['match','val', 'id']},
                            split_param='v', 
                            drift_criterion=False,
                            p_outlier=0.05)
    m.find_starting_values()
    m.sample(500, burn=100, dbname=dbname, db='pickle')
    m.save(mname)
    
    return m
Explanation: Part 2: Test parallel processing using HDDM's func
End of explanation
start_time = time.time()  # the start time of the processing
m = run_model(0)
print("\nRunning 1 chain used: %f seconds." % (time.time() - start_time))
# m = hddm.load("df_example_chain_vaz_test_0")
print("M_df1a_v_corr DIC: %f" % m.dic) 
m.plot_posterior_predictive()
Explanation: Test the docker image: Single chain
End of explanation
from ipyparallel import Client
v = Client()[:]
start_time = time.time()  # the start time of the processing
jobs = v.map(run_model, range(4)) # 4 is the number of CPUs
models = jobs.get()
print("\nRunning 4 chains used: %f seconds." % (time.time() - start_time))
print("M_df1a_v_corr DIC: %f" % models[0].dic) 
gelman_rubin(models)
# Create a new model concatenating all individual models.
import kabuki
combined_model = kabuki.utils.concat_models(models)
Explanation: Test the docker image: Multiple chains in parallel
End of explanation
def plot_grp_trace(models, param):
    import seaborn as sns
    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    
    df_grp_traces = []
    for i in range(4):
        df_m = models[i]
        df_grp_trace = df_m.get_group_traces()
        df_grp_trace['chain'] = i
        print('chain', i, df_grp_trace.shape)
        df_grp_traces.append(df_grp_trace)
    df_grp_traces = pd.concat(df_grp_traces)
    
    # select the columns included that parameter
    filter_col = [col for col in df_grp_traces if col.startswith(param)]
    filter_col.append('chain') # add the chain column to the data
    
    #df_tmp = df_grp_traces[df_grp_traces.columns.intersection(filter_col)]
    df_tmp = df_grp_traces.loc[:,filter_col]
    df_tmp["sample_id"] = df_tmp.index
    
    df_melt = pd.melt(df_tmp,
                      id_vars=['sample_id', 'chain'],
                      value_vars=filter_col[:-1],
                      var_name='cond',
                      value_name='value')
    
    sns.set_palette("Set2")
    for indx, df_tmp in df_melt.groupby('cond'):
        fig, axs = plt.subplots(ncols=2,figsize=(20, 9))
        print(indx)
        sns.lineplot(x='sample_id', y='value', hue="chain", data=df_tmp, ax=axs[0], palette='Set2')
        sns.distplot(df_tmp.loc[df_tmp['chain'] == 0, 'value'], hist = False, kde = True,
                         kde_kws = {'shade': False, 'linewidth': 3},ax=axs[1], label = '0')
        sns.distplot(df_tmp.loc[df_tmp['chain'] == 1, 'value'], hist = False, kde = True,
                         kde_kws = {'shade': False, 'linewidth': 3},ax=axs[1], label = '1')
        sns.distplot(df_tmp.loc[df_tmp['chain'] == 2, 'value'], hist = False, kde = True,
                         kde_kws = {'shade': False, 'linewidth': 3},ax=axs[1], label = '2')
        sns.distplot(df_tmp.loc[df_tmp['chain'] == 3, 'value'], hist = False, kde = True,
                         kde_kws = {'shade': False, 'linewidth': 3},ax=axs[1], label = '3')
        #axs[0].set_ylim([0.5, 3])
        axs[0].set_title(indx, fontsize=20)
        #axs[1].set_xlim([0.5, 3])
        axs[1].set_title(indx, fontsize=20)
        plt.show()
plot_grp_trace(models=df1a_models, param='z')
Explanation: Check convergence (visual inspection)
plot multiple chain and visually inspect how the convergence goes.
End of explanation
import sys
import time
from IPython.display import clear_output
def wait_watching_stdout(ar, dt=30):
    
    ar: vmap output of the models being run 
    dt: number of seconds between checking output, you can make is shorter or longer.
    
    while not ar.ready():
        stdouts = ar.stdout
        if not any(stdouts):
            continue
        # clear_output doesn't do much in terminal environments
        clear_output()
        print('-' * 30)
        print("%.3fs elapsed" % ar.elapsed)
        print("")
        for out in ar.stdout: print(out);
        sys.stdout.flush()
        time.sleep(dt)
from ipyparallel import Client
start_time = time.time()  # the start time of the processing
v = Client()[:]
jobs = v.map(run_model, range(4)) # 4 is the number of CPUs
wait_watching_stdout(jobs)
models = jobs.get()
m1_time = time.time() 
print("\nRunning 4 chains used: %f seconds." % (m1_time - start_time))
Explanation: Test the docker image: Multiple chains with progress bar
Sometimes you want to know the progress of the multiple chains. To print out the progress, you need to define another function (See this post in hddm discussion group).
End of explanation
<END_TASK>