def plot_trajectory(trajectories, name='trajectory', pos='left'):
print "Plotting trajectory: %s" % name
plt.figure(1, figsize=(4, 4))
ax = plt.subplot(111)
trajx = np.concatenate([t[0] for t in trajectories])
trajy = np.concatenate([t[1] for t in trajectories])
annote = set()
for t in trajectories:
for a in t[2]:
annote.add(a)
cmap = plt.get_cmap('Greys')
t = np.arange(trajx.shape[0]) * 0.001
points = np.array([trajy, trajx]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=cmap,
norm=plt.Normalize(-t[-1] * 0.5, t[-1] * 0.75))
lc.set_array(t)
lc.set_linewidth(3)
lc.set_rasterized(True)
ax.add_collection(lc)
for (a_t, l) in annote:
ix = int(a_t / 0.001) - 1
xy = (trajy[ix], trajx[ix])
xytext = (-30 if xy[0] > 0 else 30, -30 if xy[1] > 0 else 30)
if l == 'Start of next trial' and xy[1] > 0:
l = 'Start of\nnext trial'
xytext = (-50, 5)
if l == 'Release\n(premature)':
xytext = (-10, -35)
plt.annotate(l, xy, xycoords='data', xytext=xytext, ha='center',
va='center', textcoords='offset points',
arrowprops={'arrowstyle': '->',
'connectionstyle': 'arc3, rad=0.2'})
plt.axhline(0.0, color='k', ls=":")
plt.axvline(0.0, color='k', ls=":")
if 'left' in pos:
plt.ylabel("Task state (arbitrary units)")
ax.get_yaxis().tick_left()
else:
plt.yticks(())
ax.get_xaxis().tick_bottom()
plt.xlabel("Relative time in task state (arbitrary units)")
plt.axis((-1.5, 1.5, -2.0, 1.5))
plt.subplots_adjust(bottom=0.12, top=0.97, left=0.17, right=0.97)
save_or_show('plots/' + name + '_traj')
