-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathstep2b_data_quality-ERFs.py
122 lines (101 loc) · 4.09 KB
/
step2b_data_quality-ERFs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#!/usr/bin/env python3
"""
@ Lina Teichmann
INPUTS:
call from command line with following inputs:
-bids_dir
OUTPUTS:
Plots the ERFs of the repeat trials.
NOTES:
If it doesn't exist, the script makes a figures folder in the BIDS derivatives folder
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import mne, os
import pandas as pd
import seaborn as sns
#*****************************#
### PARAMETERS ###
#*****************************#
n_participants = 4
n_sessions = 12
n_images = 200
channel_picks = ['O','T','P']
title_names = ['Occipital','Temporal','Parietal']
colors = ['mediumseagreen','steelblue','goldenrod','indianred','grey']
plt.rcParams['font.size'] = '16'
plt.rcParams['font.family'] = 'Helvetica'
#*****************************#
### HELPER FUNCTIONS ###
#*****************************#
def load_epochs(preproc_dir,all_epochs = []):
for p in range(1,n_participants+1):
epochs = mne.read_epochs(f'{preproc_dir}/preprocessed_P{str(p)}-epo.fif', preload=False)
all_epochs.append(epochs)
return all_epochs
# helper function
def plot_erfs(epochs,n_sessions,name,color,ax,ax2,lab):
ctf_layout = mne.find_layout(epochs.info)
picks_epochs = [epochs.ch_names[i] for i in np.where([s[2]==name for s in epochs.ch_names])[0]]
picks = np.where([i[2]==name for i in ctf_layout.names])[0]
# get evoked data
for s in range(n_sessions):
evoked = epochs[(epochs.metadata['trial_type']=='test') & (epochs.metadata['session_nr']==s+1)].average()
evoked.pick_channels(ch_names=picks_epochs)
ax.plot(epochs.times*1000,np.mean(evoked.data.T,axis=1),color=color,lw=0.5,alpha=0.4)
evoked = epochs[(epochs.metadata['trial_type']=='test')].average()
evoked.pick_channels(ch_names=picks_epochs)
# plot ERFs for selected sensor group
ax.plot(epochs.times*1000,np.mean(evoked.data.T,axis=1),color=color,lw=1,label=lab)
ax.set_xlim([epochs.times[0]*1000,epochs.times[len(epochs.times)-1]*1000])
ax.set_ylim([-0.6,0.6])
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
# plot sensor locations
ax2.plot(ctf_layout.pos[:,0],ctf_layout.pos[:,1],color='gainsboro',marker='.',linestyle='',markersize=5)
ax2.plot(ctf_layout.pos[picks,0],ctf_layout.pos[picks,1],color='grey',marker='.',linestyle='',markersize=5)
ax2.set_aspect('equal')
plt.axis('off')
# Make the ERF plot
def make_figure(all_epochs,fig_dir):
fig = plt.figure(num=1,tight_layout=True,figsize = (11,6))
gs = GridSpec(3, 5, figure=fig)
for i,ch in enumerate(channel_picks):
for p in range(n_participants):
ax = fig.add_subplot(gs[i, p])
if i == 0:
ax.set_title('M' + str(p+1))
if i == 2:
ax.set_xlabel('time (ms)')
else:
plt.setp(ax.get_xticklabels(), visible=False)
if p == 0:
ax.set_ylabel('fT')
else:
plt.setp(ax.get_yticklabels(), visible=False)
ax2=fig.add_subplot(gs[i, -1])
plot_erfs(all_epochs[p],12,ch,colors[p],ax,ax2,'Sub' + str(p+1))
ax2.set_title(title_names[i])
plt.savefig(f'{fig_dir}/data_quality-ERFs.pdf',dpi=1000)
#*****************************#
### COMMAND LINE INPUTS ###
#*****************************#
if __name__=='__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"-bids_dir",
required=True,
help='path to bids root',
)
args = parser.parse_args()
bids_dir = args.bids_dir
preproc_dir = f'{bids_dir}/derivatives/preprocessed/'
sourcedata_dir = f'{bids_dir}/sourcedata/'
fig_dir = f'{bids_dir}/derivatives/figures/'
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
####### Run ########
all_epochs = load_epochs(preproc_dir)
make_figure(all_epochs,fig_dir)