def merge(gather_dir=os.getcwd(),
combine_type='run',
base_name='',
out_dir='',
**params):
"""
Merges the epoched*.pkl objects in gather_dir
Arguments:
gather_dir: path to folder containing epoched files
combine_type: 'run' to combine averages of each run
'trial' to combine every trial
Outputs:
merged.pkl, merged.mat containing the merged Epoched object
A plot is also generated
"""
files = [
f for f in os.listdir(gather_dir)
if ('mat' or 'pkl' in f) and 'epoched' in f
]
merged = None
for f in files:
run = read_epoch(os.path.join(gather_dir, f))
if merged is None:
merged = Epoched(run.n_categs, run.n_samples, 0)
merged.names = run.names
merged.num_trials = [0 for i in range(len(run.num_trials))]
merged.num_rejected = [0 for i in range(len(run.num_rejected))]
if combine_type == 'run':
with warnings.catch_warnings():
warnings.simplefilter("ignore")
avg = np.nanmean(run.matrix, axis=2, keepdims=True)
merged.matrix = np.concatenate((merged.matrix, avg), axis=2)
elif combine_type == 'trial':
merged.matrix = np.concatenate((merged.matrix, run.matrix), axis=2)
if len(run.num_trials) == len(merged.num_trials):
merged.num_trials = [
x + y for x, y in zip(run.num_trials, merged.num_trials)
]
if len(run.num_rejected) == len(merged.num_rejected):
merged.num_rejected = [
x + y for x, y in zip(run.num_rejected, merged.num_rejected)
]
spio.savemat(
make_path('merged', '.mat', out_dir=out_dir, base_name=base_name),
{'merged': merged})
save_pkl(make_path('merged', '.pkl', out_dir=out_dir, base_name=base_name),
merged)
plot_conds(merged, out_dir=out_dir, base_name=base_name, **params)
def train(map_name, num_timesteps, batch_steps, seed, network, ar, lr,
lrschedule, screen_size, minimap_size, step_mul, num_cpu, optimizer,
ent_coef, vl_coef, max_grad_norm):
maps.get(map_name) # Assert the map exists.
log_path = './experiments/%s/' % (time.strftime("%m%d_%H%M_") + map_name)
make_path(log_path)
make_path("%sreplay" % log_path)
def make_env(rank):
def _thunk():
agent_interface = features.parse_agent_interface_format(
feature_screen=64, feature_minimap=64)
env = sc2_env.SC2Env(
map_name=map_name,
step_mul=step_mul,
agent_interface_format=agent_interface,
# screen_size_px=(screen_size, screen_size),
# minimap_size_px=(minimap_size, minimap_size),
visualize=False)
return env
return _thunk
set_global_seeds(seed)
log_file = open("%sconfig.log" % log_path, "a+")
log_file.write("Map Name: %s\n" % map_name)
log_file.write("Optimizer: %s\n" % optimizer)
log_file.write("Network: %s\n" % network)
log_file.write("Learning Rate: %f\n" % lr)
log_file.write("Entropy Coefficient: %f\n" % ent_coef)
log_file.write("Value Function Coefficient: %f\n" % vl_coef)
log_file.write("Maximum Gradient Norm: %f\n" % max_grad_norm)
log_file.write("Screen Size: %d\n" % screen_size)
log_file.write("Minimap Size: %d\n" % minimap_size)
log_file.write("Batch Steps: %d\n" % batch_steps)
log_file.close()
learn(network,
log_path,
make_env,
total_timesteps=num_timesteps,
nsteps=batch_steps,
ent_coef=ent_coef,
max_grad_norm=max_grad_norm,
optimizer=optimizer,
vl_coef=vl_coef,
ar=ar,
lr=lr,
num_cpu=num_cpu)
def plot_line(pupil_data, time_list, time_with_ecog, start_index, channel_name,
merged, out_dir, base_name, pupil_data_smooth_matrix):
for count, i in enumerate(conds_to_plot):
to_plot = pupil_data_smooth_matrix[count]
fig1, ax1 = plots.subplots()
ax1.set_xlabel('Time (ms)')
x = time_list
y = to_plot[0:len(time_list)]
ax1.plot(x, y, color = plot_colors_line[count], lw = 0.8, ls=line_style)
ax1.plot(np.nan, color=plot_colors_line[count+3], lw=0.8, ls= '-')
ax1.set_ylabel('Baseline Corrected Pupil Diameter')
ax1.set_xlim([0,0.8])
a = time_with_ecog[0, start_index:]
b = time_with_ecog[count+1, start_index:]
ax2 = ax1.twinx()
ax2.plot(a, b, color = plot_colors_line[count+3], lw= 0.8, ls='-')
ax2.set_ylabel('HFB Power')
plots.title('Subject ' + str(subject) + ', Channel ' + str(channel_name) +
', Condition '+ str(merged.names[i]), fontsize = 10)
ax1.legend(labels = ['Pupil', 'ECoG'], loc='upper right', fontsize=7)
plots.gcf().subplots_adjust(left=0.15)
plots.savefig(make_path('pupil_ecog_correlation_'+str(merged.names[i])+str(subject), '.png',
out_dir=out_dir, base_name=str(channel_name)),dpi=600,bbox_inches='tight')
if all_channels:
if show:
plots.show()
else:
plots.close()
return fig1
def plot_scatter(pupil_data, desired_indecies, ecog_matrix, merged,
channel_name, out_dir, base_name, pupil_data_smooth_matrix):
labels = []
n = len(ecog_matrix[1, :])
fig, ax = plots.subplots()
for count, i in enumerate(conds_to_plot):
#use this if you want to plot not-smoothed data
#x = pupil_data[count, 0:len(desired_indecies)]
to_plot = pupil_data_smooth_matrix[count]
x = to_plot[0:len(desired_indecies)]
y = ecog_matrix[count + 1, :]
slope, intercept, r_value, p_value, std_err = stats.linregress(x, y)
r_value = np.round(r_value, 3)
p_value = np.round(p_value, 4)
ax.plot(x,
y,
marker='o',
color=plot_colors_scatter[count],
linestyle='None',
label=merged.names[i],
ms=1)
labels = np.append(labels,
merged.names[i] + ', r-value = ' + str(r_value))
#+', p-value= ' + str(p_value))
ax.legend(labels=labels, loc='upper left', fontsize=7)
for count, i in enumerate(conds_to_plot):
#x = pupil_data[count, 0:len(desired_indecies)]
to_plot = pupil_data_smooth_matrix[count]
x = to_plot[0:len(desired_indecies)]
y = ecog_matrix[count + 1, :]
slope, intercept, r_value, p_value, std_err = stats.linregress(x, y)
ax.plot(x,
intercept + slope * x,
'r',
color=plot_colors_scatter[count],
lw=-0.8)
ax.annotate('n= ' + str(n),
xy=(1, 0),
xycoords='axes fraction',
fontsize=10,
horizontalalignment='right',
verticalalignment='bottom')
ax.set_xlabel('Baseline Corrected Pupil Diameter (mm)')
ax.set_ylabel('HFB Power')
plots.title('Pupil/ECoG Correlation: ' + 'Subject ' + str(subject) +
', Channel ' + str(channel_name),
fontsize=12)
plots.savefig(make_path('pupil_ecog_correlation_scatter',
'.png',
out_dir=out_dir,
base_name=str(channel_name)),
dpi=600)
plots.close()
return fig
def save_params(params):
with open(make_path('parameters', '.py', **params), 'w') as f:
f.write('params = {}\n')
for k, v in params.items():
if type(v) == str:
f.write("params['" + str(k) + "']" + " = " + "'" + str(v) +
"'" + '\n')
else:
f.write("params['" + str(k) + "']" + " = " + str(v) + '\n')
def qualitycheck(raw, out_dir='', base_name='', **params):
'''
Does a quality check of the raw data. Using paramters in **params.
'''
print('\nChecking the quality...\n')
plot_raw(raw, out_dir=out_dir, base_name=base_name)
check_outliers(raw, out_dir=out_dir, base_name=base_name, **params)
calculate_stats(raw, out_dir=out_dir, base_name=base_name)
raw.to_csv(make_path('outliers_removed', '.csv', out_dir=out_dir,
base_name=base_name), index=False)
def plot_raw(raw, name='raw_plot', title='Raw Data Plot', out_dir='', base_name='', **params):
'''
Creates a scatter plot from the raw data and saves it to a file.
Takes in the dataframe and requires params['out_dir'], and params['base_name']
'''
s = [4 for _ in range(len(raw))]
# ax = raw.plot(x='Time', y='Pupil', title=title)
ax = raw.plot(x='Time', y='Pupil', kind='scatter', title=title, s=s)
ax.set(xlabel="Time", ylabel="Pupil Diameter (mm)")
plt.savefig(make_path(name, '.png', out_dir=out_dir,
base_name=base_name, **params))
plt.cla()
plt.clf()
plt.close()
def calculate_stats(raw, fname='descriptive_stats', plot=True, out_dir='', base_name='', **params):
'''
Calculates descriptive statistics on the raw data. Save the results into a
file called descriptive_stats. Another plot with these stats is saved in
a file called raw_plot_with_stats_outliers_removed.
'''
stats = raw.Pupil.describe()
if plot:
ax = raw.plot(x='Time', y='Pupil',
title='Raw Data (Outliers Removed)')
plt.axhline(stats['mean'], color='r', label='mean')
plt.axhline(stats['25%'], color='g', label='25%')
plt.axhline(stats['50%'], color='g', label='median')
plt.axhline(stats['75%'], color='g', label='75%')
ax.set(xlabel="Time (ms)", ylabel="Pupil Diameter (mm)")
plt.legend()
plt.savefig(make_path('raw_plot_with_stats_outliers_removed', '.png',
out_dir=out_dir, base_name=base_name))
plt.cla()
plt.clf()
plt.close()
stats.to_csv(
make_path(fname, '.csv', out_dir=out_dir, base_name=base_name))
def check_outliers(raw, impossible_upper=5, impossible_lower=1.5, out_dir='', base_name='', **params):
'''
Counts and removes outliers in the raw data. Saves a report of the number
of unusable samples based on the largest and smallest values specified in
params['impossible_upper'] and params['impossible_lower'].
'''
nsamples = len(raw)
raw.loc[(raw.Pupil < impossible_lower) | (
raw.Pupil > impossible_upper), 'Pupil'] = np.nan
count = raw.Pupil.count()
string = '{} out of {} samples were within acceptable range({} - {})\nPercentage of usable data: {}%'.format(
count, nsamples, impossible_lower, impossible_upper, count / nsamples * 100)
print(string)
with open(make_path('outlier_report', '.txt', base_name=base_name, out_dir=out_dir, **params), 'w') as f:
f.write(string)
def plot_conds(epoched,
conds_to_plot='all',
plot_colors=def_cols,
plot_style=def_style,
plot_error=True,
sample_rate=250,
back_time=60,
plot_title='Pupil Diameter',
plot_fname='pupil_diameter_plot',
out_dir='',
base_name='',
y_label='',
plot_nums=True,
**params):
'''
Averages across trials in each condition and saves a plot
Arguments:
epoched: An Epoched object, like the one one outputed by epoch()
conds_to_plot: A list of indices of conditions to plot or 'all'
To plot all conditions
plot_colors: list of colors in hex format. E.g. '#00FF00'
plot_style: list of style specs like ['-', ':']
plot_error: True or False
sample_rate: Sample rate of the eye tracker (hz)
back_time: Time to plot before event (ms)
plot_title: Title of the plot
plot_fname: File name of the saved image
out_dir: Outputs will be saved to this directory
base_name: string is appended to output files
'''
print('\nPlotting...\n')
# Calulates the mean and errors (ignoring nans), supress warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore")
errors = stat.sem(epoched.matrix, axis=2, ddof=1, nan_policy='omit')
flattened = np.nanmean(epoched.matrix, axis=2)
if conds_to_plot == 'all':
conds_to_plot = [i for i in range(epoched.n_categs)]
assert len(plot_colors) >= len(
conds_to_plot), 'Require more colors than plotted conditions'
assert len(plot_style) > len(
conds_to_plot), 'Require more styles than plotted conditions'
x = [
1000 * i / sample_rate - back_time
for i in range(epoched.total_samples)
]
fig, ax = plt.subplots()
fig.set_size_inches(8, 5)
# i is the index of the condition, count is the order of plotting
for count, i in enumerate(conds_to_plot):
y = flattened[i, :]
ax.plot(x,
y,
label=epoched.names[i],
color=plot_colors[count],
ls=plot_style[count])
# Plot error bars
if plot_error:
err = errors[i, :]
ax.fill_between(x,
y - err,
y + err,
alpha=0.25,
color=plot_colors[count])
num_plotted = [
x - y for x, y in zip(epoched.num_trials, epoched.num_rejected)
]
if plot_nums:
plot_title += '\n # trials plotted: ' + str(num_plotted)
# Formatting...
ax.set_xlabel('Time (ms)')
ax.set_ylabel(y_label)
plt.axvline(x=0, lw=0.5, color='0')
plt.xlim((x[0], x[-1]))
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.9, box.height])
ax.legend(bbox_to_anchor=(1, 1.04), frameon=False)
plt.title(plot_title)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
plt.savefig(make_path(plot_fname,
'.png',
out_dir=out_dir,
base_name=base_name),
bbox_inches='tight',
dpi=600)
def epoch(pupil_data,
events_path,
sample_rate=250,
epoch_time=200,
back_time=60,
out_dir='',
base_name='',
baseline_type='no',
**params):
'''
Epochs the pupil_data according to behavioural events specified
by the events file in events_path.
Arguments:
pupil_data: A pupil_data object like the one loaded by read_pupil
events_path <str>: A path to the events file (see tutorial)
sample_rate <int>: The sample rate of your pupil recording
epoch_time <int>: Time in ms to epoch
back_time <int>: Time before the event to plot
Returns:
epoched <Epoched>: an Epoched object defined above
'''
print('\nEpoching...\n')
# Find all the events
pupil_events = pupil_data[pupil_data['Content'] != '-']
num_events = len(pupil_events)
print('There are {} events in the pupil data, does '
'this look correct? If not, the first two events '
'of pupil may not have been recorded...'.format(num_events))
# Reads behavioral events
# TODO: from python
categories = read_events(events_path, pupil_events.Time.iat[0])
samples_per_epoch = get_nsamples(sample_rate, epoch_time, back_time)
# Initialize output variables
epoched = Epoched(len(categories), samples_per_epoch, num_events)
# Extra: Report median miss time.
# Get the baseline dictionary mapping from time to baseline values.
bl_events = None
if type(baseline_type) == tuple:
bl_events = get_baseline_events(pupil_events,
categories[baseline_type[1]])
for c, category in enumerate(categories):
epoched.names.append(category.name)
# +1 to exclude the content sample, which is out of sync.
onsets = list(
map(lambda x: get_nearest_ind(pupil_events, x) + 1,
category.start))
rejected_inds = []
rejected = 0
for t, onset in enumerate(onsets):
# Note the -1 to avoid the content sample
pre = pupil_data.Pupil.iloc[onset - samples_per_epoch[0] -
1:onset - 1].values
post = pupil_data.Pupil.iloc[onset:onset +
samples_per_epoch[1]].values
baseline = get_baseline(pupil_data, onset, baseline_type,
sample_rate, bl_events)
trial = np.concatenate((pre, post)) - baseline
if np.isnan(np.sum(trial)): # Checks if there is a nan
rejected_inds.append(t)
rejected += 1
elif len(trial) == sum(samples_per_epoch):
epoched.matrix[c, :, t - rejected] = trial
epoched.num_trials.append(len(onsets))
epoched.num_rejected.append(rejected)
epoched.rejected[category.name] = rejected_inds
# Save .mat and .pkl of epoched data
spio.savemat(
make_path('epoched', '.mat', out_dir=out_dir, base_name=base_name),
{'epoched': epoched})
save_pkl(
make_path('epoched', '.pkl', out_dir=out_dir, base_name=base_name),
epoched)
return epoched