def _add_plot(fig, ax, plot_data, color_data, pkeys, lw=1, cmap='_auto',
alpha=1.0):
colors = color_data.color
if cmap == '_auto':
cmap = None
if color_data.is_categorical:
colors = np.take(COLOR_CYCLE, colors, mode='wrap')
if plot_data.scatter:
data, = plot_data.trajs
artist = ax.scatter(*data.T, marker='o', c=colors, edgecolor='none',
s=lw*20, cmap=cmap, alpha=alpha, picker=5)
else:
# trajectory plot
if hasattr(colors, 'dtype') and np.issubdtype(colors.dtype, np.number):
# delete lines with NaN colors
mask = np.isfinite(colors)
if mask.all():
trajs = plot_data.trajs
else:
trajs = [t for i,t in enumerate(plot_data.trajs) if mask[i]]
colors = colors[mask]
artist = LineCollection(trajs, linewidths=lw, cmap=cmap)
artist.set_array(colors)
else:
artist = LineCollection(plot_data.trajs, linewidths=lw, cmap=cmap)
artist.set_color(colors)
artist.set_alpha(alpha)
artist.set_picker(True)
artist.set_pickradius(5)
ax.add_collection(artist, autolim=True)
ax.autoscale_view()
# Force ymin -> 0
ax.set_ylim((0, ax.get_ylim()[1]))
def on_pick(event):
if event.artist is not artist:
return
label = pkeys[event.ind[0]]
ax.set_title(label)
fig.canvas.draw_idle()
# XXX: hack, make this more official
if hasattr(fig, '_superman_cb'):
fig.canvas.mpl_disconnect(fig._superman_cb[0])
cb_id = fig.canvas.mpl_connect('pick_event', on_pick)
fig._superman_cb = (cb_id, on_pick)
return artist
