def plot_trajectories_from_dataset(dataset, keys, interpolated_data='hide'):
'''
interpolated_data - if trajectories contain interpolated data, use 'hide' to hide that data, 'dotted' to show it as a dotted line, or 'show' to show it in its entirety
'''
import fly_plot_lib.plot as fpl # download at: https://github.com/florisvb/FlyPlotLib
from fly_plot_lib.colormaps import viridis as viridis
path = dataset.config.path
bgimg_filename = get_filename(path, 'bgimg_N1.png')
bgimg = plt.imread(bgimg_filename)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_aspect('equal')
ax.imshow(bgimg, cmap='gray')
for key in keys:
trajec = dataset.trajec(key)
if 'interpolated' in trajec.__dict__.keys():
interpolated_indices = np.where(trajec.interpolated==1)[0]
if interpolated_data == 'dotted':
r = np.arange(0, len(interpolated_indices), 5)
interpolated_indices = interpolated_indices[r]
if interpolated_data != 'show':
trajec.position_x[interpolated_indices] = np.nan
trajec.position_y[interpolated_indices] = np.nan
l = -1
fpl.colorline(ax, trajec.position_x[0:l], trajec.position_y[0:l], trajec.time_epoch[0:l]-trajec.time_epoch[0:l][0], linewidth=2, colormap='none', norm=None, zorder=1, alpha=1, linestyle='solid', cmap=viridis)
fpl.adjust_spines(ax, [])
ax.set_frame_on(False)
def plot_digits_distribution(n):
fig = plt.figure()
ax = fig.add_subplot(111)
bins = np.arange(-.5, 10.5, 1)
npages = [30, 60, 400, 1500]
colors = ['black', 'blue', 'green', 'red']
digits = []
for npage in npages:
d = get_digits(n, npage)
digits.append(np.array(d))
xticks = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
fpl.histogram(ax,
digits,
bins=bins,
colors=colors,
normed=True,
show_smoothed=False)
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_ylim(0, 0.15)
ax.set_xlim(-1, 10)
ax.set_xlabel('digit')
ax.set_ylabel('probability')
def plot_numerical_grammian(eps=1e-5):
fig = plt.figure(figsize=(3,3))
ax = fig.add_axes([0.2,0.2,0.7,0.7])
cn, distance = numerical_grammian(nsamples=2000, control='exp', eps=eps)
ax.plot(distance, np.log(cn), 'red')
cn, distance = numerical_grammian(nsamples=2000, control='constantaccel', eps=eps, accel=-0.1)
print distance.shape, len(cn)
ax.plot(distance, np.log(cn), 'blue')
cn, distance = numerical_grammian(nsamples=2000, control='constantaccel', eps=eps, accel=-0.01)
print distance.shape, len(cn)
ax.plot(distance, np.log(cn), 'lightblue')
cn, distance = numerical_grammian(nsamples=2000, control='none', eps=eps)
ax.plot(distance, np.log(cn), 'black')
yticks = np.log(np.array([1e0,1e4,1e8,1e12,1e16]))
fpl.adjust_spines(ax, ['left', 'bottom'], yticks=yticks)
ax.set_yticklabels(['$10^{0}$','$10^{4}$', '$10^{8}$', '$10^{12}$', '$10^{16}$'])
ax.set_xlabel('distance to target, m')
ax.set_ylabel('condition number')
flytext.set_fontsize(fig, 8)
fname = 'condition_numbers_' + str(np.log10(eps)) + '.pdf'
fig.savefig(fname, format='pdf')
def example_colored_cartesian_spagetti(dataset, axis='xy', xlim=(-0.2, .2), ylim=(-0.75, .25), zlim=(-.15, -.15), keys=None, keys_to_highlight=[], show_saccades=False, colormap='jet', color_attribute='speed', norm=(0,0.5), artists=None):
fig = plt.figure()
ax = fig.add_subplot(111)
if axis=='xy': # xy plane
ax.set_ylim(ylim[0], ylim[1])
ax.set_xlim(xlim[0], xlim[1])
ax.set_autoscale_on(True)
ax.set_aspect('equal')
axes=[0,1]
cartesian_spagetti(ax, dataset, keys=keys, nkeys=300, start_key=0, axes=axes, show_saccades=show_saccades, keys_to_highlight=[], colormap=colormap, color_attribute=color_attribute, norm=norm, show_start=False)
post = patches.Circle( (0, 0), radius=0.01, facecolor='black', edgecolor='none', alpha=1)
artists = [post]
if artists is not None:
for artist in artists:
ax.add_artist(artist)
#prep_cartesian_spagetti_for_saving(ax)
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=[-.2, 0, .2], yticks=[-.75, -.5, -.25, 0, .25])
ax.set_xlabel('x axis, m')
ax.set_ylabel('y axis, m')
ax.set_title('xy plot, color=speed from 0-0.5 m/s')
fig.set_size_inches(8,8)
fig.savefig('example_colored_xy_spagetti_plot.pdf', format='pdf')
return ax
def plot_speed_vs_eccentricity(dataset, config):
orientations, airheadings, groundheadings, eccentricities, speeds, airspeeds = get_orientation_data(dataset, config, visual_stimulus='none')
fig = plt.figure()
ax = fig.add_subplot(111)
xticks = [0,1]
yticks = [0,1]
ax.set_xlim(xticks[0], xticks[-1])
ax.set_ylim(yticks[0], yticks[-1])
fpl.scatter(ax, speeds, eccentricities, color='black', colornorm=[0,0.8], radius=.003, xlim=[xticks[0], xticks[-1]], ylim=[yticks[0], yticks[-1]])
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=yticks)
ax.set_xlabel('ground speed')
ax.set_ylabel('eccentricity')
savename = 'speed_vs_eccentricity.pdf'
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def plot_orientation_vs_groundheading(dataset, config, odor_stimulus='on', odor=True):
orientations, airheadings, groundheadings, eccentricities, speeds, airspeeds = get_orientation_data(dataset, config, visual_stimulus='none', odor_stimulus=odor_stimulus, odor=odor)
fig = plt.figure()
ax = fig.add_subplot(111)
xticks = [-np.pi, -np.pi/2., 0, np.pi/2., np.pi]
ax.set_xlim(xticks[0], xticks[-1])
ax.set_ylim(xticks[0], xticks[-1])
fpl.scatter(ax, groundheadings, orientations, color=eccentricities, colornorm=[0,0.8], radius=.01, xlim=[xticks[0], xticks[-1]], ylim=[xticks[0], xticks[-1]])
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=xticks)
ax.set_xlabel('groundspeed heading')
xticklabels = ['-180', '-90', 'upwind', '90', '180']
ax.set_xticklabels(xticklabels)
ax.set_yticklabels(xticklabels)
ax.set_ylabel('body orientation')
savename = 'orientation_vs_groundheading.pdf'
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def make_histograms(config, data, colors, savename='orientation_histogram.pdf'):
print data
fig = plt.figure(figsize=(5,4))
ax = fig.add_subplot(111)
bins = np.linspace(-np.pi,np.pi,50)
fpl.histogram(ax, data.values(), bins=bins, bin_width_ratio=1, colors=colors.values(), edgecolor='none', bar_alpha=1, curve_fill_alpha=0.2, curve_line_alpha=1, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, smoothing_bins_to_exclude=[])
xticks = [-np.pi, -np.pi/2., 0, np.pi/2., np.pi]
ax.set_xlim(xticks[0], xticks[-1])
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_xlabel('heading')
xticklabels = ['-180', '-90', 'upwind', '90', '180']
ax.set_xticklabels(xticklabels)
ax.set_ylabel('occurences, normalized')
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def make_histograms(config, data, colors, savename='histogram.pdf'):
print data
fig = plt.figure(figsize=(5,4))
ax = fig.add_subplot(111)
bins = np.linspace(-.5,.5,150)
fpl.histogram(ax, data.values(), bins=bins, bin_width_ratio=1, colors=colors.values(), edgecolor='none', bar_alpha=1, curve_fill_alpha=0.2, curve_line_alpha=1, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, smoothing_bins_to_exclude=[74,75,76])
xticks = [-.5, 0, .5]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
#ax.set_xticklabels(xticklabels)
ax.set_xlabel('altitudes')
ax.set_ylabel('occurences, normalized')
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def plot_trajectories_from_dataset(dataset, keys, interpolated_data='hide'):
'''
interpolated_data - if trajectories contain interpolated data, use 'hide' to hide that data, 'dotted' to show it as a dotted line, or 'show' to show it in its entirety
'''
import matplotlib.pyplot as plt
# download at: https://github.com/florisvb/FlyPlotLib
import fly_plot_lib.plot as fpl
from fly_plot_lib.colormaps import viridis as viridis
path = dataset.config.path
bgimg_filename = get_filename(path, 'bgimg_N1.png')
bgimg = plt.imread(bgimg_filename)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_aspect('equal')
ax.imshow(bgimg, cmap='gray')
for key in keys:
trajec = dataset.trajec(key)
if 'interpolated' in trajec.__dict__.keys():
interpolated_indices = np.where(trajec.interpolated==1)[0]
if interpolated_data == 'dotted':
r = np.arange(0, len(interpolated_indices), 5)
interpolated_indices = interpolated_indices[r]
if interpolated_data != 'show':
trajec.position_x[interpolated_indices] = np.nan
trajec.position_y[interpolated_indices] = np.nan
l = -1
fpl.colorline(ax, trajec.position_x[0:l], trajec.position_y[0:l], trajec.time_epoch[0:l]-trajec.time_epoch[0:l][0], linewidth=2, colormap='none', norm=None, zorder=1, alpha=1, linestyle='solid', cmap=viridis)
fpl.adjust_spines(ax, [])
ax.set_frame_on(False)
def plot_residency_time_histogram(config, dataset, save_figure_path=''):
odor_stimulus = 'on'
threshold_distance_min = 0.05
# no odor
keys_odor_off = opa.get_keys_with_odor_before_post(config, dataset, threshold_odor=10, threshold_distance=-1, odor_stimulus='none', upwind_only=True, threshold_distance_min=0.1, odor=True, post_behavior='landing')
residency_time_odor_off = []
for key in keys_odor_off:
trajec = dataset.trajecs[key]
if trajec.residency_time is not None:
residency_time_odor_off.append(trajec.residency_time)
# odor on, odor experienced
keys_odor_on_true = opa.get_keys_with_odor_before_post(config, dataset, threshold_odor=10, threshold_distance=-1, odor_stimulus='on', upwind_only=True, threshold_distance_min=0.1, odor=True, post_behavior='landing')
residency_time_odor_on_true = []
for key in keys_odor_on_true:
trajec = dataset.trajecs[key]
if trajec.residency_time is not None:
residency_time_odor_on_true.append(trajec.residency_time)
# odor on, odor not experienced
keys_odor_on_false = opa.get_keys_with_odor_before_post(config, dataset, threshold_odor=10, threshold_distance=-1, odor_stimulus='on', upwind_only=False, threshold_distance_min=0.1, odor=False, post_behavior='landing')
residency_time_odor_on_false = []
for key in keys_odor_on_false:
trajec = dataset.trajecs[key]
if trajec.residency_time is not None:
residency_time_odor_on_false.append(trajec.residency_time)
data = [np.array(residency_time_odor_off), np.array(residency_time_odor_on_true), np.array(residency_time_odor_on_false)]
colors = ['black', 'red', 'blue']
nbins = 30 # note: if show_smoothed=True with default butter filter, nbins needs to be > ~15
fig = plt.figure()
ax = fig.add_subplot(111)
fpl.histogram(ax, data, bins=nbins, bin_width_ratio=0.8, colors=colors, edgecolor='none', bar_alpha=1, curve_fill_alpha=0.4, curve_line_alpha=0, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, bootstrap_std=False, exponential_histogram=False)
#xticks = [0,.02,.04,.06,.08,.15]
fpl.adjust_spines(ax, ['left', 'bottom'])
#ax.set_xlim(xticks[0], xticks[-1])
ax.set_xlabel('residency time, frames')
ax.set_ylabel('Occurences, normalized')
ax.set_title('Residency time')
if save_figure_path == '':
save_figure_path = os.path.join(config.path, config.figure_path, 'activity/')
fig.set_size_inches(8,8)
figname = save_figure_path + 'residency_time_on_post_histogram' + '.pdf'
fig.savefig(figname, format='pdf')
def fit_1d_gaussian(odor_dataset, t, axis=0, keys=None, plot=True, lims=[-1,1], ignore_parameter_names=[]):
if keys is None:
keys = odor_dataset.odor_traces.keys()
ordinate = []
odor = []
odor_std = []
for key in keys:
odor_trace = odor_dataset.odor_traces[key]
ordinate.append( odor_trace.position[axis] )
index_at_t = np.argmin( np.abs(odor_trace.timestamps - t) )
odor.append( odor_trace.voltage[index_at_t] )
odor_std.append( odor_trace.voltage_std[index_at_t] )
ordinate = np.array(ordinate)
odor = np.array(odor)
odor_std = np.array(odor_std)
print ordinate
print odor
# now fit gaussian to data
gm = data_fit.models.GaussianModel1D()
inputs = [ordinate]
#return odor, ordinate
gm.fit_with_guess(odor, inputs, ignore_parameter_names=ignore_parameter_names)
if plot:
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(ordinate, odor, 'ob')
for i, pt in enumerate(ordinate):
ax.vlines(pt, odor[i]-odor_std[i], odor[i]+odor_std[i], linewidth=2)
x = np.arange(lims[0], lims[1], 0.001)
vals = gm.get_val(x)
ax.plot(x, vals)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('x position, m')
ax.set_ylabel('odor value, ethanol')
ax.set_title('mean and std dev of measured odor values and gaussian fit')
return gm, ordinate, odor, odor_std
def plot_distance_histogram(config, dataset, save_figure_path=''):
# no odor
keys_no_odor = opa.get_keys_with_odor_before_post(config, dataset, threshold_odor=10, odor_stimulus='none', threshold_distance_min=0.1, odor=True)
data_no_odor = []
for key in keys_no_odor:
trajec = dataset.trajecs[key]
if np.max(trajec.distance_to_post) > 0.1:
for f, d in enumerate(trajec.distance_to_post):
if trajec.positions[f,0] > 0:
data_no_odor.append(d)
# odor on, odor experienced
keys_odor_on_true = opa.get_keys_with_odor_before_post(config, dataset, threshold_odor=10, odor_stimulus='on', threshold_distance_min=0.1, odor=True)
data_odor_on_true = []
for key in keys_odor_on_true:
trajec = dataset.trajecs[key]
if np.max(trajec.distance_to_post) > 0.1:
for f, d in enumerate(trajec.distance_to_post):
if trajec.positions[f,0] > 0:
data_odor_on_true.append(d)
print len(data_no_odor)
print len(data_odor_on_true)
data = [np.array(data_no_odor), np.array(data_odor_on_true)]
colors = ['black', 'red']
nbins = 30
bins = np.linspace(0,0.1,nbins)
fig = plt.figure()
ax = fig.add_subplot(111)
fpl.histogram(ax, data, bins=bins, bin_width_ratio=0.8, colors=colors, edgecolor='none', bar_alpha=1, curve_fill_alpha=0.4, curve_line_alpha=0, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, bootstrap_std=True, exponential_histogram=False, n_bootstrap_samples=10000, smoothing_range=[0.005,0.1])
xticks = [0,.02,.04,.06,.08,.1]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_xlim(xticks[0], xticks[-1])
ax.set_xlabel('Distance to post, m')
ax.set_ylabel('Occurences, normalized')
title_text = 'Distance to post. red (upwind, odor) N = ' + str(len(keys_odor_on_true)) + '; black (upwind, no odor) N = ' + str(len(keys_no_odor))
ax.set_title(title_text)
if save_figure_path == '':
save_figure_path = os.path.join(config.path, config.figure_path, 'activity/')
fig.set_size_inches(8,8)
figname = save_figure_path + 'distance_to_post_histogram' + '.pdf'
fig.savefig(figname, format='pdf')
def plot_odor_trace_from_file(filename, ax=None):
f = open(filename)
odor_trace = pickle.load(f)
if ax is None:
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(odor_trace.timestamps - odor_trace.timestamps[0], odor_trace.signal, 'red')
ax.plot(odor_trace.timestamps - odor_trace.timestamps[0], odor_trace.voltage, 'black')
ax.set_ylim(odor_trace.voltage[0]-10, np.max(odor_trace.voltage)+10)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('time, sec')
ax.set_ylabel('control signal and odor response')
ax.set_title('raw odor traces + signal, and mean (red)')
def plot_activity_histogram(dataset, save_figure_path=''):
keys = get_keys(dataset)
print 'number of keys: ', len(keys)
if len(keys) < 1:
print 'No data'
return
odor_local_time = []
no_odor_local_time = []
for i, key in enumerate(keys):
trajec = dataset.trajecs[key]
if trajec.odor is not False:
odor_local_time.append( trajec.timestamp_local_float )
else:
no_odor_local_time.append( trajec.timestamp_local_float )
odor_local_time = np.array(odor_local_time)
no_odor_local_time = np.array(no_odor_local_time)
fig = plt.figure()
ax = fig.add_subplot(111)
nbins = 24 # note: if show_smoothed=True with default butter filter, nbins needs to be > ~15
bins = np.linspace(0,24,nbins)
data = []
colors = []
if len(odor_local_time) > 0:
data.append(odor_local_time)
colors.append('red')
if len(no_odor_local_time) > 0:
data.append(no_odor_local_time)
colors.append('blue')
fpl.histogram(ax, data, bins=bins, bin_width_ratio=0.8, colors=colors, edgecolor='none', bar_alpha=1, curve_fill_alpha=0.4, curve_line_alpha=0, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=False, normed_occurences=False, bootstrap_std=False, exponential_histogram=False)
xticks = np.linspace(bins[0], bins[-1], 5, endpoint=True)
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_xlabel('Time of day, hours')
ax.set_ylabel('Occurences, normalized')
ax.set_title('Activity, as measured by number of trajectories: odor=red, no odor=blue')
fig.set_size_inches(8,8)
figname = save_figure_path + 'activity_histogram' + '.pdf'
fig.savefig(figname, format='pdf')
def show_start_stop(dataset):
fig = plt.figure()
ax = fig.add_subplot(111)
artists = []
xpos = []
ypos = []
for key, trajec in dataset.trajecs.items():
if 1:
x = trajec.positions[0,0]
y = trajec.positions[0,1]
start = patches.Circle( (x, y), radius=0.003, facecolor='green', edgecolor='none', alpha=1, linewidth=0)
x = trajec.positions[-1,0]
y = trajec.positions[-1,1]
stop = patches.Circle( (x, y), radius=0.003, facecolor='red', edgecolor='none', alpha=1, linewidth=0)
#artists.append(start)
artists.append(stop)
if 0:
xpos.append(trajec.positions[-1,0])
ypos.append(trajec.positions[-1,1])
if 1:
for artist in artists:
ax.add_artist(artist)
#fpl.histogram2d(ax, np.array(xpos), np.array(ypos), bins=100, logcolorscale=True, xextent=[-.2,.2], yextent=[-.75,.25])
ax.set_aspect('equal')
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=[-.2, 0, .2], yticks=[-.75, -.5, -.25, 0, .25])
ax.set_xlabel('x axis, m')
ax.set_ylabel('y axis, m')
ax.set_title('xy plot, color=speed from 0-0.5 m/s')
fig.set_size_inches(8,8)
fig.savefig('start_stop.pdf', format='pdf')
def plot_all_data_files(path, colors=None):
all_data = open_all_data_files(path)
fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111)
ax.set_rasterization_zorder(0)
if colors is None:
colors = ['red' for i in range(len(all_data))]
for i, data in enumerate(all_data):
all_positions = []
all_pid_vals = []
pos = np.vstack(data['flydra_position'])
vals = np.array(data['pid'])
indices_where_no_plume = np.where(pos[:,1]<-.04)[0]
print np.mean(vals[indices_where_no_plume])
vals = vals - np.mean(vals[indices_where_no_plume])
print np.mean(vals[indices_where_no_plume])
#ax.plot(pos[:,1], vals, color='gray')
all_positions.extend(pos[:,1].tolist())
all_pid_vals.extend(vals)
bins, mean_binned_values = get_binned_data(all_positions, all_pid_vals)
print i
try:
ax.plot(bins, mean_binned_values, color=colors[i], zorder=10)
except:
ax.plot(bins, mean_binned_values, color='pink', zorder=10)
ax.set_xlim(-.05, .1)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('crosswind position, m')
ax.set_ylabel('pid value')
filename = 'odor_plume_plot.pdf'
filename_with_path = os.path.join(path, filename)
fig.savefig(filename_with_path, format='pdf')
def plot_all_data_files(path, colors=None):
all_data = open_all_data_files(path)
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)
ax.set_rasterization_zorder(0)
if colors is None:
colors = ['red' for i in range(len(all_data))]
for i, data in enumerate(all_data):
all_positions = []
all_pid_vals = []
pos = np.vstack(data['flydra_position'])
vals = np.array(data['pid'])
indices_where_no_plume = np.where(pos[:, 1] < -.04)[0]
print np.mean(vals[indices_where_no_plume])
vals = vals - np.mean(vals[indices_where_no_plume])
print np.mean(vals[indices_where_no_plume])
#ax.plot(pos[:,1], vals, color='gray')
all_positions.extend(pos[:, 1].tolist())
all_pid_vals.extend(vals)
bins, mean_binned_values = get_binned_data(all_positions, all_pid_vals)
print i
try:
ax.plot(bins, mean_binned_values, color=colors[i], zorder=10)
except:
ax.plot(bins, mean_binned_values, color='pink', zorder=10)
ax.set_xlim(-.05, .1)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('crosswind position, m')
ax.set_ylabel('pid value')
filename = 'odor_plume_plot.pdf'
filename_with_path = os.path.join(path, filename)
fig.savefig(filename_with_path, format='pdf')
def plot_mean_odor_trace(odor_dataset, ignore_traces=[1,2,3,4], ax=None):
if ax is None:
fig = plt.figure()
ax = fig.add_subplot(111)
ignore_traces_strings = ['_' + str(i) for i in ignore_traces]
for key, odor_trace in odor_dataset.odor_traces.items():
if key[-2:] in ignore_traces_strings:
continue
color = 'gray'
ax.plot(odor_dataset.mean_odor_trace.timestamps, odor_trace.signal, color)
ax.plot(odor_dataset.mean_odor_trace.timestamps, odor_trace.voltage, color)
ax.plot(odor_dataset.mean_odor_trace.timestamps, odor_dataset.mean_odor_trace.signal, 'red')
ax.plot(odor_dataset.mean_odor_trace.timestamps, odor_dataset.mean_odor_trace.voltage, 'red')
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('time, sec')
ax.set_ylabel('control signal and odor response')
ax.set_title('raw odor traces + signal, and mean (red)')
def set_spines_and_labels(ax_xy, ax_xz, ax_yz):
yticks = [-.15, 0, .15]
xticks = [-.2, 0, 1]
zticks = [-.15, 0, .15]
ax_xy.set_xlim(xticks[0], xticks[-1])
ax_xy.set_ylim(yticks[0], yticks[-1])
ax_xz.set_xlim(xticks[0], xticks[-1])
ax_xz.set_ylim(zticks[0], zticks[-1])
ax_yz.set_xlim(yticks[0], yticks[-1])
ax_yz.set_ylim(zticks[0], zticks[-1])
fpl.adjust_spines(ax_xy, ['left'], xticks=xticks, yticks=yticks)
fpl.adjust_spines(ax_xz, ['left', 'bottom'], xticks=xticks, yticks=zticks)
fpl.adjust_spines(ax_yz, ['right', 'bottom'], xticks=yticks, yticks=zticks)
ax_xy.set_xlabel('')
ax_xy.set_ylabel('y axis')
ax_xz.set_ylabel('z axis')
ax_xz.set_xlabel('x axis, upwind negative')
ax_yz.set_xlabel('y axis')
ax_yz.yaxis.set_label_position('right')
ax_yz.set_ylabel('z axis')
ax_xy.set_aspect('equal')
ax_xz.set_aspect('equal')
ax_yz.set_aspect('equal')
def plot_slipangles(dataset, config, visual_stimulus='none', odor_stimulus='on', odor=True):
orientations, airheadings, groundheadings, eccentricities, speeds, airspeeds = get_orientation_data(dataset, config, visual_stimulus='none', odor_stimulus=odor_stimulus, odor=odor, eccentricity_threshold=0.8)
airslip = airheadings - orientations
groundslip = groundheadings - orientations
data = {'airslip': airslip, 'groundslip': groundslip}
#data = {'orientation': orientations-airheadings, 'airheadings': orientations-groundheadings, 'groundheadings': groundheadings}
color = {'airslip': 'green', 'groundslip': 'red'}
savename = 'slipangle_histogram.pdf'
fig = plt.figure(figsize=(5,4))
ax = fig.add_subplot(111)
bins = np.linspace(-np.pi,np.pi,50)
fpl.histogram(ax, data.values(), bins=bins, bin_width_ratio=1, colors=color.values(), edgecolor='none', bar_alpha=1, curve_fill_alpha=0.2, curve_line_alpha=1, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, smoothing_bins_to_exclude=[])
xticks = [-np.pi, -np.pi/2., 0, np.pi/2., np.pi]
ax.set_xlim(xticks[0], xticks[-1])
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_xlabel('slip angle')
xticklabels = ['-180', '-90', '0', '90', '180']
ax.set_xticklabels(xticklabels)
ax.set_ylabel('occurences, normalized')
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def make_histograms(config, data, colors, axis, savename='histogram.pdf'):
for key, item in data.items():
print key, len(item)
fig = plt.figure(figsize=(5,4))
ax = fig.add_subplot(111)
if axis=='z':
bins = np.linspace(-0.5,0.5,50,endpoint=True)
bins_to_exclude = [24]
elif axis=='xy':
bins = np.linspace(0.05,1,50,endpoint=True)
bins_to_exclude = []
fpl.histogram(ax, data.values(), bins=bins, bin_width_ratio=1, colors=colors.values(), edgecolor='none', bar_alpha=1, curve_fill_alpha=0.2, curve_line_alpha=1, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, smoothing_bins_to_exclude=bins_to_exclude)
if axis=='z':
xticks = [-0.5, -.25, 0, .25, .5]
elif axis == 'xy':
xticks = [0,.2,.4,.6,.8,1.]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
#ax.set_xticklabels(xticklabels)
ax.set_xlabel('speed in ' + axis)
ax.set_ylabel('occurences, normalized')
ax.set_xlim(xticks[0], xticks[-1])
path = config.path
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path=os.path.join(figure_path, 'odor_traces/')
figure_path = os.path.join(path, config.figure_path)
save_figure_path = os.path.join(figure_path, 'odor_traces/')
fig_name_with_path = os.path.join(save_figure_path, savename)
print 'SAVING TO: ', fig_name_with_path
fig.savefig(fig_name_with_path, format='pdf')
def prep_cartesian_spagetti_for_saving(ax):
fig.set_size_inches(fig_width,fig_height)
rect = patches.Rectangle( [-.25, -.15], .5, .3, facecolor='none', edgecolor='gray', clip_on=False, linewidth=0.2)
ax.add_artist(rect)
offset = 0.00
dxy = 0.05
#xarrow = patches.FancyArrowPatch(posA=(-.25+offset, -.15+offset), posB=(-.25+offset+dxy, -.15+offset), arrowstyle='simple')
#patches.Arrow( -.25+offset, -.15+offset, dxy, 0, color='black', width=0.002)
xarrow = patches.FancyArrowPatch((-.25+offset, -.15+offset), (-.25+offset+dxy, -.15+offset), arrowstyle="-|>", mutation_scale=10, color='gray', shrinkA=0, clip_on=False)
ax.add_patch(xarrow)
yarrow = patches.FancyArrowPatch((-.25+offset, -.15+offset), (-.25+offset, -.15+offset+dxy), arrowstyle="-|>", mutation_scale=10, color='gray', shrinkA=0, clip_on=False)
ax.add_artist(yarrow)
text_offset = -.011
ax.text(-.25+offset+dxy+text_offset, -.15+offset+.005, 'x', verticalalignment='bottom', horizontalalignment='left', color='gray', weight='bold')
ax.text(-.25+offset+.005, -.15+offset+dxy+text_offset, 'y', verticalalignment='bottom', horizontalalignment='left', color='gray', weight='bold')
scale_bar_offset = 0.01
ax.hlines(-0.15+scale_bar_offset, 0.25-scale_bar_offset-.1, 0.25-scale_bar_offset, linewidth=1, color='gray')
ax.text(0.25-scale_bar_offset-.1/2., -0.15+scale_bar_offset+.002, '10cm', horizontalalignment='center', verticalalignment='bottom', color='gray')
ax.set_aspect('equal')
scaling = .5/.75
margin = 0.04
aspect_ratio = 3/5. # height/width
fig_width = 7.204*scaling
plt_width = fig_width - 2*margin*(1-aspect_ratio)
fig_height = plt_width*aspect_ratio + 2*margin
fig = ax.figure
fig.set_size_inches(fig_width,fig_height)
fig.subplots_adjust(bottom=margin, top=1-margin, right=1, left=0)
ax.set_axis_off()
fpl.adjust_spines(ax, ['left', 'bottom'])
def plot_digits_distribution(n):
fig = plt.figure()
ax = fig.add_subplot(111)
bins = np.arange(-.5,10.5,1)
npages = [30,60,400,1500]
colors = ['black', 'blue', 'green', 'red']
digits = []
for npage in npages:
d = get_digits(n, npage)
digits.append(np.array(d))
xticks = [0,1,2,3,4,5,6,7,8,9]
fpl.histogram(ax, digits, bins=bins, colors=colors, normed=True, show_smoothed=False)
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_ylim(0,0.15)
ax.set_xlim(-1,10)
ax.set_xlabel('digit')
ax.set_ylabel('probability')
def get_axes(fig=None):
figure_padding = 0.25
subplot_padding = 0.08
if fig is None:
fig = plt.figure(figsize=(8,4.5))
x = 1.2
y = .3
z = .3
aspect_ratio = (y+z+subplot_padding)/(x+y+subplot_padding)
gs1 = gridspec.GridSpec(2, 2, width_ratios=[x,y])
gs1.update(left=figure_padding*aspect_ratio, right=1-figure_padding*aspect_ratio, wspace=subplot_padding, hspace=subplot_padding, top=1-figure_padding+subplot_padding, bottom=figure_padding-subplot_padding)
ax_xy = plt.subplot(gs1[0, 0])
ax_xz = plt.subplot(gs1[1, 0])
ax_yz = plt.subplot(gs1[1, 1])
if 1:
yticks = [-.15, 0, .15]
xticks = [-.2, 0, 1]
zticks = [-.15, 0, .15]
ax_xy.set_ylabel('y axis')
ax_xz.set_ylabel('z axis')
ax_xz.set_xlabel('x axis')
ax_yz.set_xlabel('x axis')
ax_yz.yaxis.set_label_position('right')
ax_yz.set_ylabel('z axis')
fpl.adjust_spines(ax_xy, ['left'], yticks=yticks)
fpl.adjust_spines(ax_xz, ['left', 'bottom'], xticks=xticks, yticks=zticks)
fpl.adjust_spines(ax_yz, ['right', 'bottom'], xticks=yticks, yticks=zticks)
ax_xy.set_aspect('equal')
ax_xz.set_aspect('equal')
ax_yz.set_aspect('equal')
return [ax_xy, ax_xz, ax_yz]
mean_accels_in_odor.extend((np.diff(trajec.speed[frames_in_odor])*trajec.fps).tolist())
mean_accels_not_in_odor.extend((np.diff(trajec.speed[frames_not_in_odor])*trajec.fps).tolist())
data_in_odor.append(np.array(mean_accels_in_odor))
data_not_in_odor.append(np.array(mean_accels_not_in_odor))
print odor_stimulus, ': ', np.mean(mean_accels), ' +/- ', np.std(mean_accels)
# in odor
fig = plt.figure(figsize=(4,2))
ax = fig.add_subplot(111)
nbins = 75
bins = np.linspace(-.05,.05,nbins)*100
fpl.histogram(ax, data_in_odor, bins=bins, bin_width_ratio=0.8, colors=['black', 'red'], edgecolor='none', normed=True, show_smoothed=True, bar_alpha=1, curve_line_alpha=0)
ax.set_xlim(bins[0], bins[-1])
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('Acceleration, m/s2')
ax.set_ylabel('Occurences, normalized')
save_figure_path = os.path.join(config.path, config.figure_path, 'activity/')
figname = save_figure_path + 'accceleration_histogram_in_odor' + '.pdf'
fig.savefig(figname, format='pdf')
# not in odor
fig = plt.figure(figsize=(4,2))
ax = fig.add_subplot(111)
nbins = 75
bins = np.linspace(-.05,.05,nbins)*100
fpl.histogram(ax, data_not_in_odor, bins=bins, bin_width_ratio=0.8, colors=['black', 'blue'], edgecolor='none', normed=True, show_smoothed=True, bar_alpha=1, curve_line_alpha=0)
ax.set_xlim(bins[0], bins[-1])
fpl.adjust_spines(ax, ['left', 'bottom'])
dhuman = np.linspace(1e-3,10,100)
theta = 0.003*np.pi/180.
l = 65e-3
human_error = stereoscopic_error(dhuman,theta,l)
ax1.plot(np.log(dhuman),human_error,'blue')
dfly = np.linspace(1e-3,.02,100)
theta = 5*np.pi/180.
l = 0.3e-3
fly_error = stereoscopic_error(dfly,theta,l)
ax1.plot(np.log(dfly),fly_error,'red')
human_xticks = np.log(np.array([1e-3,1e-2,1e-1,1,10]))
xticklabels = ['$10^{-3}$', '$10^{-2}$', '$10^{-1}$', '$1$', '$10$']
ax1.set_ylim(-.001,.12)
ax1.set_xlim(human_xticks[0], human_xticks[-1])
human_yticks = [0,0.03,0.06,0.09,0.12]
fpl.adjust_spines(ax1, ['left', 'bottom'], xticks=human_xticks, yticks=human_yticks)
ax1.set_xticklabels(xticklabels)
ax1.set_ylabel('Min. error in stereo depth estimate, m')
ax1.set_xlabel('Distance to object, m (log scale)')
fpltext.set_fontsize(fig, 8)
fig.savefig('stereo_error.pdf', format='pdf')
def play_movie_3_axes(x, y, z, color, orientation=None, save=False, save_movie_path='', nskip=0, images=[None, None, None], images_extents=[None, None, None], artists=[[], [], []], lims=[None, None, None], colornorm=None, colormap='jet', ghost_tail=20):
'''
Show an animation of an x,y,z trajectory with color and a tail. The axes are referred to in the following order: (xy, xz, yz)
x, y, z -- lists or arrays of length N
color -- list or array of length N, or color name string for constant color
orientation -- list of length N, or None. If None, movie uses circles, else uses oriented wedges from fly_plot_lib.get_wedges...
artists -- list of length 3, corresponding to x/y/z axes, each component should be a list of matplotlib artists
lims -- list of limits for the three axes (xy, xz, yz), if None, automatically generate limits based on trajectory
ghost_tail -- number of frames of 'tail' to show
TODO: test and animate images, add image function instead of requiring premade list
'''
# prep plot
fig = plt.figure()
ax_xy = fig.add_subplot(221)
ax_xz = fig.add_subplot(223)
ax_yz = fig.add_subplot(224)
axes = [ax_xy, ax_xz, ax_yz]
anim_params = {'frame': -1*(1+nskip)}
if colornorm is None:
colornorm = [np.min(color), np.max(color)]
norm = matplotlib.colors.Normalize(colornorm[0], colornorm[1])
color_mappable = matplotlib.cm.ScalarMappable(norm, plt.get_cmap('jet'))
if orientation is None:
radius = 0.01
alpha = 1
radiusnorm = None
maxradius = 1
minradius = 0
flies_xy = fpl.get_circles_for_scatter(x, y, color=color, colormap=colormap, radius=radius, colornorm=colornorm, alpha=alpha, radiusnorm=radiusnorm, maxradius=maxradius, minradius=minradius)
flies_xz = fpl.get_circles_for_scatter(x, z, color=color, colormap=colormap, radius=radius, colornorm=colornorm, alpha=alpha, radiusnorm=radiusnorm, maxradius=maxradius, minradius=minradius)
flies_yz = fpl.get_circles_for_scatter(y, z, color=color, colormap=colormap, radius=radius, colornorm=colornorm, alpha=alpha, radiusnorm=radiusnorm, maxradius=maxradius, minradius=minradius)
def get_images(f):
ims = [None for i in range(3)]
origin = 'lower'
zorder_image = 0
alpha_image = 0.5
for i, image in enumerate(images):
if image is not None:
if image_extents[i] is None:
print 'Must include extent for image'
raise(ValueError)
if type(image) is list:
im = image[f]
else:
im = image
ims[i] = im
return ims
# initialize images:
ims = get_images(0)
ax_ims = [None for i in range(3)]
for i, ax in enumerate(axes):
if ims[i] is not None:
im = ax.imshow(ims[i], extent=image_extens[i], origin='lower', cmap=plt.get_cmap(colormap), norm=norm, alpha=0.5, zorder=0)
ax_ims[i] = im
# add artists
for i, ax in enumerate(axes):
for artist in artists[i]:
ax.add_artist(artist)
ax_xy.add_collection(flies_xy)
ax_xz.add_collection(flies_xz)
ax_yz.add_collection(flies_yz)
def init_plot():
flies_xy.set_color('none')
flies_xz.set_color('none')
flies_yz.set_color('none')
flies_xy.set_edgecolors('none')
flies_xz.set_edgecolors('none')
flies_yz.set_edgecolors('none')
return flies_xy, flies_xz, flies_yz
def updatefig(*args):
#x = anim_params['x']
#y = anim_params['y']
#z = anim_params['z']
#color = anim_params['color']
print anim_params['frame']
anim_params['frame'] += 1 + nskip
frame_end = anim_params['frame'] + ghost_tail
if frame_end > len(x):
frame_end = len(x)
if anim_params['frame'] >= len(x)-1:
anim_params['frame'] = 0
frame_end = anim_params['frame'] + ghost_tail
frames = np.arange(anim_params['frame'], frame_end, 1).tolist()
colors = ['none' for i in range(len(x))]
for f in frames:
colors[f] = color_mappable.to_rgba(color[f])
flies_xy.set_facecolors(colors)
flies_xz.set_facecolors(colors)
flies_yz.set_facecolors(colors)
if save:
print anim_params['frame']
frame_prefix = '_tmp'
frame_prefix = os.path.join(save_movie_path, frame_prefix)
strnum = str(anim_params['frame'])
num_frame_digits = np.ceil(len(x) / 10.) + 1
while len(strnum) < num_frame_digits:
strnum = '0' + strnum
frame_name = frame_prefix + '_' + strnum + '_' + '.png'
fig.savefig(frame_name, format='png')
return flies_xy, flies_xz, flies_yz
# generate and set limits:
maxs = [np.max(x), np.max(y), np.max(z)]
mins = [np.min(x), np.min(y), np.min(z)]
for i, lim in enumerate(lims):
if lim is None:
lims[i] = [mins[i], maxs[i]]
ax_xy.set_xlim(lims[0][0], lims[0][1])
ax_xy.set_ylim(lims[1][0], lims[1][1])
ax_xz.set_xlim(lims[0][0], lims[0][1])
ax_xz.set_ylim(lims[2][0], lims[2][1])
ax_yz.set_xlim(lims[1][0], lims[1][1])
ax_yz.set_ylim(lims[2][0], lims[2][1])
for ax in axes:
ax.set_aspect('equal')
#
fpl.adjust_spines(ax_xy, ['left'], yticks=[-.15, 0, .15])
fpl.adjust_spines(ax_xz, ['left', 'bottom'], xticks=[-.2, 0, 1], yticks=[-.15, 0, .15])
fpl.adjust_spines(ax_yz, ['right', 'bottom'], xticks=[-.15, 0, .15], yticks=[-.15, 0, .15])
# top left plot
ax_xy.set_ylabel('y position, m')
# bottom left plot
ax_xz.set_xlabel('x position, m')
ax_xz.set_ylabel('z position, m')
# bottom right plot
ax_yz.set_xlabel('y position, m')
#ax_yz.set_ylabel('z position, m') # I can't figure out how to make this appear on the right side...
ani = animation.FuncAnimation(fig, updatefig, init_func=init_plot, fargs=anim_params, interval=50, blit=True)
plt.show()
def plot_colored_cartesian_spagetti(config, dataset, axis='xy', xlim=(-0.2, .2), ylim=(-0.75, .25), zlim=(0, 0.3), keys=None, keys_to_highlight=[], show_saccades=False, colormap='jet', color_attribute='speed', norm=(0,0.5), artists=None, save_figure_path='', figname=None, show_start=False):
if keys is None:
keys = get_keys(dataset)
print 'plotting spagetti, axis: ', axis
print 'number of keys: ', len(keys)
if len(keys) < 1:
print 'No data'
return
fig = plt.figure()
ax = fig.add_subplot(111)
height = config.post_center[2]-config.ticks['z'][0]
print 'ARTISTS STARTING'
print artists
if axis=='xy': # xy plane
ax.set_ylim(ylim[0], ylim[1])
ax.set_xlim(xlim[0], xlim[1])
ax.set_autoscale_on(True)
ax.set_aspect('equal')
axes=[0,1]
fap.cartesian_spagetti(ax, dataset, keys=keys, nkeys=10, start_key=0, axes=axes, show_saccades=show_saccades, keys_to_highlight=[], colormap=colormap, color_attribute=color_attribute, norm=norm, show_start=show_start)
post = patches.Circle(config.post_center[0:2], config.post_radius, color='black')
if axis=='yz': # yz plane
ax.set_ylim(zlim[0], zlim[1])
ax.set_xlim(ylim[0], ylim[1])
ax.set_autoscale_on(True)
ax.set_aspect('equal')
axes=[1,2]
fap.cartesian_spagetti(ax, dataset, keys=keys, nkeys=10, start_key=0, axes=axes, show_saccades=show_saccades, keys_to_highlight=[], colormap=colormap, color_attribute=color_attribute, norm=norm, show_start=show_start)
post = patches.Rectangle([-1*config.post_radius, config.ticks['z'][0]], config.post_radius*2, height, color='black')
if axis=='xz': # xz plane
ax.set_ylim(zlim[0], zlim[1])
ax.set_xlim(xlim[0], xlim[1])
ax.set_autoscale_on(True)
ax.set_aspect('equal')
axes=[0,2]
fap.cartesian_spagetti(ax, dataset, keys=keys, nkeys=10, start_key=0, axes=axes, show_saccades=show_saccades, keys_to_highlight=[], colormap=colormap, color_attribute=color_attribute, norm=norm, show_start=show_start)
post = patches.Rectangle([-1*config.post_radius, config.ticks['z'][0]], config.post_radius*2, height, color='black')
if artists is None:
artists = []
artists.append(post)
if artists is not None:
for artist in artists:
ax.add_artist(artist)
#prep_cartesian_spagetti_for_saving(ax)
xticks = config.ticks['x']
yticks = config.ticks['y']
zticks = config.ticks['z']
if axis=='xy':
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=yticks)
ax.set_xlabel('x axis, m')
ax.set_ylabel('y axis, m')
ax.set_title('xy plot, color=speed from 0-0.5 m/s')
if axis=='yz':
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=yticks, yticks=zticks)
ax.set_xlabel('y axis, m')
ax.set_ylabel('z axis, m')
ax.set_title('yz plot, color=speed from 0-0.5 m/s')
if axis=='xz':
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=zticks)
ax.set_xlabel('x axis, m')
ax.set_ylabel('z axis, m')
ax.set_title('xz plot, color=speed from 0-0.5 m/s')
fig.set_size_inches(8,8)
if figname is None:
figname = save_figure_path + 'spagetti_' + axis + '.pdf'
else:
figname = os.path.join(save_figure_path, figname)
fig.savefig(figname, format='pdf')
return ax
def play_movie(x,
y,
images=None,
extent=None,
aspect='equal',
color='blue',
edgecolor='none',
orientation=None,
save=False,
save_movie_path='',
nskip=0,
artists=[],
xlim=None,
ylim=None,
colornorm=None,
colormap='jet',
ghost_tail=20,
ax=None,
wedge_radius=0.01,
circle_radius=0.005,
deg=False,
flip=False,
imagecolormap='jet',
mono=True,
flipimgx=False,
flipimgy=False,
strobe=False,
sync_frames=None,
figsize=(5, 3),
dpi=72):
'''
Show an animation of N x,y trajectories with color, orientation, and a tail. And optionally background images.
x -- list or np.array for x position of trajectory, OR, if multiple trajectories, list of lists/np.arrays
y -- list or np.array for y position of trajectory, OR, if multiple trajectories, list of lists/np.arrays
images -- - list of images to set as the background for the animation
- or a single static image
- or path to a directory with a sequence of pyplot.imread readable images (jpeg, png, etc), numbered in order
- or a type with a __call__ attribute (eg. a function or class method).
The function __call__ attribute should take a single input, the frame number, and return the image to be used.
extent, aspect -- see matplotlib.pyplot.imshow for details - note "origin" from imshow does not work as expected, use flipimgx and flipimgy instead
flipimgx -- flip the image in along the "x" axis (eg. reverse image columns), default: False.
flipimgy -- flip the image in along the "y" axis (eg. reverse image rows), default: False.
imagecolormap -- colormap for images, eg. 'jet' or 'gray'
mono -- if the image is, or should be, mono (grayscale) set this to True (default: True). Required for colormaps to work properly
color -- list, OR list of lists/np.arrays, OR string
colornorm -- [min, max] to use for normalizing color
If None, uses first trajectory to set colornorm
orientation -- list, OR list of lists/np.arrays, OR None.
If None, movie uses circles, else uses oriented wedges from fly_plot_lib.get_wedges...
artists -- optional list of matplotlib artists
xlim -- xlim for plot, if None, automatically generate limits based on trajectory
ylim -- ylim for plot, if None, automatically generate limits based on trajectory
ghost_tail -- number of frames of 'tail' to show
ax -- matplotlib axis, optional
wedge_radius -- size_radius used in fpl.get_wedges
circle_radius -- radius used in fpl.get_circles
currently changing size is NOT supported for animations.
deg -- (bool) is orientation given in degrees? True means yes. Passed to fpl.get_wedges...
flip -- (bool) flip orientation? True means yes. Passed to fpl.get_wedges...
'''
x = copy.copy(x)
y = copy.copy(y)
orientation = copy.copy(orientation)
if orientation is None:
orientation = [None for i in range(len(x))]
plt.close('all')
# prep plot
if ax is None:
fig = plt.figure(figsize=figsize, dpi=dpi)
ax = fig.add_subplot(111)
anim_params = {
'frame': -1 * (1 + nskip) + 1,
'movie_finished': False,
'strobe': strobe,
'strobeimg': None,
'frames': []
}
# fix format for single trajectories
if type(x) is list:
if type(x[0]) is list or type(x[0]) is np.ndarray:
print
print type(x[0])
print str(len(x)) + ' flies!'
if sync_frames is not None:
largest_sync_frame = np.max(sync_frames)
first_frames = []
for i, xi in enumerate(x):
padding = [np.nan] * (largest_sync_frame - sync_frames[i])
x[i] = np.hstack((padding, x[i]))
y[i] = np.hstack((padding, y[i]))
color[i] = np.hstack((padding, color[i]))
if orientation[0] is not None:
orientation[i] = np.hstack((padding, orientation[i]))
first_frames.append(len(padding))
else:
first_frames = [0] * len(x)
final_frames = []
longest_trajectory = 0
for i, xi in enumerate(x):
if len(xi) > longest_trajectory:
longest_trajectory = len(xi)
for i, xi in enumerate(x):
final_frames.append(len(xi))
for xi in x:
print len(xi)
else: # type(x[0]) is float or type(x[0]) is int or type(x[0]) is long or type(x[0]):
print
print 'Just one fly!'
x = [x]
y = [y]
color = [color]
edgecolor = [edgecolor]
orientation = [orientation]
first_frames = [0]
final_frames = [len(x[0])]
else:
print
print 'Not sure what format x is!'
anim_params.setdefault('final_frames', final_frames)
anim_params.setdefault('first_frames', first_frames)
print 'length color: ', len(color), 'n flies: ', len(x)
if len(color) != len(x):
only_color = color[0]
color = [only_color for i in range(len(x))]
if type(edgecolor) is not list:
edgecolor = [edgecolor]
if len(edgecolor) != len(x):
only_edgecolor = edgecolor[0]
edgecolor = [only_edgecolor for i in range(len(x))]
if type(color[0]) is not str:
if colornorm is None:
colornorm = [np.min(color), np.max(color)]
norm = matplotlib.colors.Normalize(colornorm[0], colornorm[1])
color_mappable = matplotlib.cm.ScalarMappable(norm,
plt.get_cmap(colormap))
else:
colornorm = None
colormap = None
norm = None
color_mappable = None
flies = []
for i, xpos in enumerate(x):
if orientation[i] is None: # use circles from scatter
radius = circle_radius
alpha = 1
radiusnorm = None
maxradius = 1
minradius = 0
fly = fpl.get_circles_for_scatter(x[i],
y[i],
color=color[i],
colormap=colormap,
radius=radius,
colornorm=colornorm,
alpha=alpha,
radiusnorm=radiusnorm,
maxradius=maxradius,
minradius=minradius,
edgecolor=edgecolor)
elif orientation[i] is not None: # use wedges from get_wedges
fly = fpl.get_wedges_for_heading_plot(x[i],
y[i],
color[i],
orientation[i],
size_radius=wedge_radius,
size_angle=20,
colormap=colormap,
colornorm=colornorm,
alpha=1,
flip=flip,
deg=deg,
nskip=0,
center_offset_fraction=0.75,
edgecolor=edgecolor)
flies.append(fly)
# add artists
for artist in artists:
ax.add_artist(artist)
# add images
if images is not None:
frame = 0
if hasattr(images, '__call__'):
imgdata = images(frame)
if len(imgdata.shape) > 2:
imgdata = imgdata[:, :, 0]
if flipimgx:
if flipimgy:
imgdata = imgdata[::-1, ::-1]
else:
imgdata = imgdata[::-1, :]
elif flipimgy:
imgdata = imgdata[:, ::-1]
else:
imgdata = imgdata[:, :]
else:
imgdata = get_image_data(images,
frame,
mono=mono,
flipimgx=flipimgx,
flipimgy=flipimgy)
img = ax.imshow(imgdata,
extent=extent,
cmap=plt.get_cmap(imagecolormap),
zorder=-10)
for fly in flies:
ax.add_collection(fly)
def init_plot():
for fly in flies:
fly.set_color('none')
if images is None:
return flies
else:
animation_objects = [fly for fly in flies]
animation_objects.insert(0, img)
return animation_objects
def updatefig(*args):
anim_params['frame'] += 1 + nskip
if anim_params['frame'] >= len(x[0]) - 1:
anim_params['frame'] = 1
anim_params['movie_finished'] = True
anim_params['frames'].append(anim_params['frame'])
print anim_params['frame']
for i, fly in enumerate(flies):
frame_start = anim_params['frame'] - ghost_tail
frame_end = anim_params['frame']
if frame_start < anim_params['first_frames'][i]:
frame_start = anim_params['first_frames'][i]
if frame_end > anim_params['final_frames'][i]:
frame_end = anim_params['final_frames'][i]
if frame_end < frame_start + nskip:
continue
frames = np.arange(frame_start, frame_end, nskip + 1).tolist()
colors = ['none' for f in range(len(x[i]))]
edgecolors = ['none' for f in range(len(x[i]))]
for f in frames:
try:
colors[f] = color_mappable.to_rgba(color[i][f])
except:
colors[f] = color[i]
edgecolors[f] = edgecolor[i]
fly.set_edgecolors(edgecolors)
fly.set_facecolors(colors)
if images is not None:
if hasattr(images, '__call__'):
imgdata = images(frames[-1])
if len(imgdata.shape) > 2:
imgdata = imgdata[:, :, 0]
if flipimgx:
if flipimgy:
imgdata = imgdata[::-1, ::-1]
else:
imgdata = imgdata[::-1, :]
elif flipimgy:
imgdata = imgdata[:, ::-1]
else:
imgdata = imgdata[:, :]
else:
imgdata = get_image_data(images,
frames[-1],
mono=mono,
flipimgx=flipimgx,
flipimgy=flipimgy)
if anim_params['strobe']:
if anim_params['strobeimg'] is None:
anim_params['strobeimg'] = imgdata
else:
anim_params['strobeimg'] = nim.darken(
[anim_params['strobeimg'], imgdata])
img.set_array(anim_params['strobeimg'])
else:
img.set_array(imgdata)
if save and not anim_params['movie_finished']:
print 'saving frame: ', str(
anim_params['frame']
), ' -- if the animation you see is strange, do not worry, look at the pngs'
frame_prefix = '_tmp'
frame_prefix = os.path.join(save_movie_path, frame_prefix)
strnum = str(anim_params['frame'])
num_frame_digits = np.ceil(np.log10(len(x[0]))) + 1
while len(strnum) < num_frame_digits:
strnum = '0' + strnum
frame_name = frame_prefix + '_' + strnum + '_' + '.png'
flytext.set_fontsize(fig, 8)
fig.savefig(frame_name, format='png')
if save and anim_params['movie_finished']:
print
print 'Movie finished saving! Close the plot screen now.'
print 'PNGs are at: ', save_movie_path
print 'To turn the PNGs into a movie, you can run this command from inside the directory with the tmp files: '
print 'mencoder \'mf://*.png\' -mf type=png:fps=30 -ovc lavc -lavcopts vcodec=mpeg4 -oac copy -o animation.avi'
if images is None:
return flies
else:
animation_objects = [fly for fly in flies]
animation_objects.insert(0, img)
return animation_objects
# generate and set limits:
if xlim is None:
xlim = [np.min(x[0]), np.max(x[0])]
if ylim is None:
ylim = [np.min(y[0]), np.max(y[0])]
ax.set_xlim(xlim[0], xlim[1])
ax.set_ylim(ylim[0], ylim[1])
ax.set_aspect(aspect)
#
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xlim, yticks=ylim)
ax.set_xlabel('x position, m')
ax.set_ylabel('y position, m')
ani = animation.FuncAnimation(fig,
updatefig,
init_func=init_plot,
fargs=anim_params,
interval=50,
blit=True)
plt.show()
def plot_odor_heading_book(pp, threshold_odor, path, config, dataset, keys=None):
fig = plt.figure(figsize=(4,4))
ax = fig.add_subplot(111)
saccade_angles_after_odor = []
heading_at_saccade_initiation = []
heading_after_saccade = []
for key in keys:
trajec = dataset.trajecs[key]
frames_in_odor = np.where(trajec.odor > threshold_odor)[0]
odor_blocks = hf.find_continuous_blocks(frames_in_odor, 5, return_longest_only=False)
for block in odor_blocks:
middle_of_block = int(np.mean(block))
# find next saccade
first_sac = None
second_sac = None
for sac in trajec.saccades:
if sac[0] > middle_of_block:
if first_sac is None:
first_sac = sac
elif second_sac is None:
if trajec.odor[sac[0]] < threshold_odor:
second_sac = sac
break
if first_sac is not None:
next_sac = first_sac
angle_of_saccade = tac.get_angle_of_saccade(trajec, next_sac)
heading_prior_to_saccade = trajec.heading_smooth[next_sac[0]]
# flip heading
if heading_prior_to_saccade < 0:
heading_prior_to_saccade += np.pi
else:
heading_prior_to_saccade -= np.pi
# flip saccade angle
if angle_of_saccade < 0:
angle_of_saccade += np.pi
else:
angle_of_saccade -= np.pi
saccade_angles_after_odor.append(angle_of_saccade)
heading_at_saccade_initiation.append(heading_prior_to_saccade)
heading_after_saccade.append(heading_prior_to_saccade + angle_of_saccade)
saccade_angles_after_odor = np.array(saccade_angles_after_odor)
heading_at_saccade_initiation = np.array(heading_at_saccade_initiation)
heading_after_saccade = np.array(heading_after_saccade)
ax.plot(heading_at_saccade_initiation*180./np.pi, saccade_angles_after_odor*180./np.pi, '.')
#ax.plot(heading_at_saccade_initiation*180./np.pi, heading_after_saccade*180./np.pi, '.')
xticks = [-180, -90, 0, 90, 180]
yticks = [-180, -90, 0, 90, 180]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=yticks)
ax.set_xlabel('Heading before saccade')
ax.set_ylabel('Angle of saccade')
title_text = 'Odor: ' + trajec.odor_stimulus.title()
ax.set_title(title_text)
ax.text(0,-180, 'Upwind', horizontalalignment='center', verticalalignment='top')
ax.text(90,-180, 'Starboard', horizontalalignment='center', verticalalignment='top')
ax.text(-90,-180, 'Port', horizontalalignment='center', verticalalignment='top')
ax.text(-180,90, 'Starboard', horizontalalignment='left', verticalalignment='center', rotation='vertical')
ax.text(-180,-90, 'Port', horizontalalignment='left', verticalalignment='center', rotation='vertical')
pp.savefig()
plt.close('all')
# angle of saccade histogram
fig = plt.figure()
ax = fig.add_subplot(111)
fpl.histogram_stack(ax, [saccade_angles_after_odor*180./np.pi], bins=20, bin_width_ratio=0.9, colors=['red'], edgecolor='none', normed=True)
ax.set_xlabel('Angle of Saccade')
ax.set_ylabel('Occurences, normalized')
xticks = [-180, -90, 0, 90, 180]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_title(title_text)
pp.savefig()
plt.close('all')
def analyze_statistical_exploration():
ndatapoints = [5, 10, 20, 50, 100, 200, 500, 1000, 10000]
medians = []
stdevs = []
for n in ndatapoints:
stats = explore_statistical_measures_across_repetitions(iterations=20,
ndatapoints=n)
medians.append(np.mean(stats, axis=0))
stdevs.append(np.std(stats, axis=0))
medians = np.array(medians)
stdevs = np.array(stdevs)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.semilogx()
ax.plot(ndatapoints, medians[:, 0], color='black', linewidth=4)
ax.vlines(ndatapoints,
medians[:, 0] - stdevs[:, 0],
medians[:, 0] + stdevs[:, 0],
color='black',
linewidth=2)
ax.plot(np.array(ndatapoints) + 0.1 * np.array(ndatapoints),
medians[:, 1],
color='red',
linewidth=4)
ax.vlines(np.array(ndatapoints) + 0.1 * np.array(ndatapoints),
medians[:, 1] - stdevs[:, 1],
medians[:, 1] + stdevs[:, 1],
color='red',
linewidth=2)
ax.plot(np.array(ndatapoints) - 0.1 * np.array(ndatapoints),
medians[:, 2],
color='green',
linewidth=4)
ax.vlines(np.array(ndatapoints) - 0.1 * np.array(ndatapoints),
medians[:, 2] - stdevs[:, 2],
medians[:, 2] + stdevs[:, 2],
color='green',
linewidth=2)
ax.set_xlim(0, 1000)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('number of data points')
ax.set_ylabel('statistical measure')
fig.savefig('stats.pdf', format='pdf')
fig = plt.figure(figsize=(10, 8))
bins = np.linspace(0, 30, 30)
for i, n in enumerate(ndatapoints):
ax = fig.add_subplot(1, len(ndatapoints), i + 1)
data1, data2 = get_two_datasets(n)
fpl.histogram(ax, [data1, data2],
bins=bins,
colors=['blue', 'orange'],
show_smoothed=False,
normed=True,
bin_width_ratio=0.9)
ax.set_xlim(bins[0], bins[-1])
fpl.adjust_spines(ax, ['bottom'], xticks=[0, 30])
fig.savefig('distributions.pdf', format='pdf')
def plot_deceleration(config, dataset, save=True):
keys = fad.get_keys_with_attr(dataset, "post_behavior", "landing")
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111)
speeds = []
angles = []
for key in keys:
trajec = dataset.trajecs[key]
# get landing frame
for i, behavior in enumerate(trajec.post_behavior):
if behavior == "landing":
landing_frame = trajec.post_behavior_frames[i]
break
if landing_frame > 60:
frame0 = landing_frame - 60
frame1 = np.min([landing_frame + 5, trajec.length])
frames = np.arange(frame0, frame1).tolist()
angle_subtended_by_post = 2 * np.sin(config.post_radius / (trajec.distance_to_post + config.post_radius))
# find frame of deceleration
# step backwards from frame of landing until acceleration is positive
accel = -1
f = landing_frame
while accel < 0:
f -= 1
accel = trajec.speed_xy[f] - trajec.speed_xy[f - 1]
speeds.append(trajec.speed[f])
angles.append(np.log(angle_subtended_by_post[f]))
ax.plot(np.log(angle_subtended_by_post[frames]), trajec.speed[frames], color="black", linewidth=0.5)
ax.plot(np.log(angle_subtended_by_post[f]), trajec.speed[f], ".", color="purple", markersize=8, zorder=10)
speeds = np.array(speeds)
angles = np.array(angles)
lm = data_fit.models.LinearModel()
lm.fit(speeds, inputs=angles)
xvals = np.linspace(np.log(5 * np.pi / 180.0), np.log(np.pi / 2.0), 10)
yvals = lm.get_val(xvals)
ax.plot(xvals, yvals, color="purple", linewidth=2)
angle_ticks = [5, 10, 30, 60, 90, 180]
xticks = np.log(np.array(angle_ticks) * np.pi / 180.0)
yticks = [0, 0.2, 0.4, 0.6, 0.8]
fpl.adjust_spines(ax, ["left", "bottom"], xticks=xticks, yticks=yticks)
ax.set_xticklabels(angle_ticks)
ax.set_xlabel("Retinal size of post")
ax.set_ylabel("Ground speed")
if save:
figure_path = os.path.join(config.path, config.figure_path)
save_figure_path = os.path.join(figure_path, "activity/")
figname = save_figure_path + "deceleration_for_landings" + ".pdf"
plt.savefig(figname, format="pdf")
def fit_1d_gaussian_time_varying(odor_dataset, tmin=15.2, tmax=18, tres=0.1, num=None, colormap='jet', tres_for_plot=0.5, axis=0, keys=None, lims=None, ignore_parameter_names=[], plot=True):
if lims is None:
if axis==0:
lims = [-.3,1]
elif axis==1:
lims = [-.1,.1]
if keys is None:
sets, keys_to_position = od.find_sets(odor_dataset,axis)
lengths_of_sets = np.asarray([len(s) for s in sets.values()])
set_to_use = np.argmax(lengths_of_sets)
if num is not None:
set_to_use = num
keys = sets[sets.keys()[set_to_use]]
timestamps = np.arange(tmin, tmax, tres)
ordinate = []
odor = []
odor_std = []
for t in timestamps:
odor_data_at_time_t = []
odor_std_data_at_time_t = []
for key in keys:
odor_trace = odor_dataset.odor_traces[key]
odor_at_time_t = np.interp(t, odor_trace.timestamps, odor_trace.voltage)
odor_data_at_time_t.append(odor_at_time_t)
odor_std_at_time_t = np.interp(t, odor_trace.timestamps, odor_trace.voltage_std)
odor_std_data_at_time_t.append(odor_std_at_time_t)
if t == timestamps[0]:
ordinate.append( odor_trace.position[axis] )
odor.append(np.array(odor_data_at_time_t))
odor_std.append(np.array(odor_std_data_at_time_t))
ordinate = np.array(ordinate)
# now fit gaussian to data
gm = data_fit.models.GaussianModel1D_TimeVarying()
inputs = [ordinate]
gm.fit(timestamps, odor, inputs)
if plot:
fig = plt.figure()
ax = fig.add_subplot(111)
norm = matplotlib.colors.Normalize(tmin, tmax)
cmap = matplotlib.cm.ScalarMappable(norm, colormap)
t_arr = np.arange(tmin, tmax, tres_for_plot)
for t in t_arr:
# find index of timestamps array that corresponds to this t
index = np.argmin(np.abs(timestamps-t))
color = cmap.to_rgba(t)
for i, pt in enumerate(ordinate):
ax.vlines(pt, odor[index][i]-odor_std[index][i], odor[index][i]+odor_std[index][i], color='black', linewidth=1)
ax.plot(ordinate, odor[index], 'o', color=color, markersize=8)
x = np.arange(lims[0], lims[1], 0.001)
vals = gm.get_val([t, [x]])
ax.plot(x, vals, color=color, linewidth=2)
fpl.adjust_spines(ax, ['left', 'bottom'])
ax.set_xlabel('x position, m')
ax.set_ylabel('odor value, ethanol')
ax.set_title('mean and std dev of measured odor values and time varying gaussian fit\ncolor=time')
return gm
data = [np.array(time_to_saccade_odor_off), np.array(time_to_saccade_odor_on_true)]
print [len(d) for d in data]
colors = ['black', 'red']
nbins = 25 # note: if show_smoothed=True with default butter filter, nbins needs to be > ~15
bins = np.linspace(0,1,nbins)
fig = plt.figure()
ax = fig.add_subplot(111)
fpl.histogram(ax, data, bins=bins, bin_width_ratio=0.8, colors=colors, edgecolor='none', bar_alpha=1, curve_fill_alpha=0.4, curve_line_alpha=0, curve_butter_filter=[3,0.3], return_vals=False, show_smoothed=True, normed=True, normed_occurences=False, bootstrap_std=True, exponential_histogram=False)
xticks = [0,0.2,0.4,0.5, 1.]
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks)
ax.set_xlim(xticks[0], xticks[-1])
ax.set_xlabel('Time to saccade after leaving odor, secs')
ax.set_ylabel('Occurences, normalized')
ax.set_title('Time to saccade')
if save_figure_path == '':
save_figure_path = os.path.join(config.path, config.figure_path, 'activity/')
fig.set_size_inches(8,8)
figname = save_figure_path + 'time_to_saccade_histogram' + '.pdf'
fig.savefig(figname, format='pdf')
def main(config, dataset, save_figure_path=''):
def play_movie(x, y, images=None, extent=None, aspect='equal', color='blue', edgecolor='none', orientation=None, save=False, save_movie_path='', nskip=0, artists=[[]], xlim=None, ylim=None, colornorm=None, colormap='jet', ghost_tail=20, ax=None, wedge_radius=0.01, circle_radius=0.005, deg=False, flip=False, imagecolormap='jet', mono=True, flipimgx=False, flipimgy=False):
'''
Show an animation of N x,y trajectories with color, orientation, and a tail. And optionally background images.
x -- list or np.array for x position of trajectory, OR, if multiple trajectories, list of lists/np.arrays
y -- list or np.array for y position of trajectory, OR, if multiple trajectories, list of lists/np.arrays
images -- - list of images to set as the background for the animation
- or a single static image
- or path to a directory with a sequence of pyplot.imread readable images (jpeg, png, etc), numbered in order
- or a type with a __call__ attribute (eg. a function or class method).
The function __call__ attribute should take a single input, the frame number, and return the image to be used.
extent, aspect -- see matplotlib.pyplot.imshow for details - note "origin" from imshow does not work as expected, use flipimgx and flipimgy instead
flipimgx -- flip the image in along the "x" axis (eg. reverse image columns), default: False.
flipimgy -- flip the image in along the "y" axis (eg. reverse image rows), default: False.
imagecolormap -- colormap for images, eg. 'jet' or 'gray'
mono -- if the image is, or should be, mono (grayscale) set this to True (default: True). Required for colormaps to work properly
color -- list, OR list of lists/np.arrays, OR string
colornorm -- [min, max] to use for normalizing color
If None, uses first trajectory to set colornorm
orientation -- list, OR list of lists/np.arrays, OR None.
If None, movie uses circles, else uses oriented wedges from fly_plot_lib.get_wedges...
artists -- optional list of matplotlib artists
xlim -- xlim for plot, if None, automatically generate limits based on trajectory
ylim -- ylim for plot, if None, automatically generate limits based on trajectory
ghost_tail -- number of frames of 'tail' to show
ax -- matplotlib axis, optional
wedge_radius -- size_radius used in fpl.get_wedges
circle_radius -- radius used in fpl.get_circles
currently changing size is NOT supported for animations.
deg -- (bool) is orientation given in degrees? True means yes. Passed to fpl.get_wedges...
flip -- (bool) flip orientation? True means yes. Passed to fpl.get_wedges...
'''
plt.close('all')
# prep plot
if ax is None:
fig = plt.figure()
ax = fig.add_subplot(111)
anim_params = {'frame': -1*(1+nskip)+1, 'movie_finished': False}
# fix format for single trajectories
if type(x) is list:
if type(x[0]) is list or type(x[0]) is np.ndarray:
print
print type(x[0])
print str(len(x)) + ' flies!'
else:# type(x[0]) is float or type(x[0]) is int or type(x[0]) is long or type(x[0]):
print
print 'Just one fly!'
x = [x]
y = [y]
color = [color]
edgecolor = [edgecolor]
orientation = [orientation]
else:
print
print str(len(x)) + ' flies!'
if len(color) != len(x):
only_color = color[0]
color = [only_color for i in range(len(x))]
if type(edgecolor) is not list:
edgecolor = [edgecolor]
if len(edgecolor) != len(x):
only_edgecolor = edgecolor[0]
edgecolor = [only_edgecolor for i in range(len(x))]
if type(color[0]) is not str:
if colornorm is None:
colornorm = [np.min(color), np.max(color)]
norm = matplotlib.colors.Normalize(colornorm[0], colornorm[1])
color_mappable = matplotlib.cm.ScalarMappable(norm, plt.get_cmap('jet'))
else:
colornorm = None
colormap= None
norm = None
color_mappable = None
flies = []
for i, xpos in enumerate(x):
if orientation[i] is None: # use circles from scatter
radius = circle_radius
alpha = 1
radiusnorm = None
maxradius = 1
minradius = 0
fly = fpl.get_circles_for_scatter(x[i], y[i], color=color[i], colormap=colormap, radius=radius, colornorm=colornorm, alpha=alpha, radiusnorm=radiusnorm, maxradius=maxradius, minradius=minradius, edgecolor=edgecolor)
elif orientation[i] is not None: # use wedges from get_wedges
fly = fpl.get_wedges_for_heading_plot(x[i], y[i], color[i], orientation[i], size_radius=wedge_radius, size_angle=20, colormap=colormap, colornorm=colornorm, alpha=1, flip=flip, deg=deg, nskip=0, center_offset_fraction=0.75, edgecolor=edgecolor)
flies.append(fly)
# add artists
for artist in artists[i]:
ax.add_artist(artist)
# add images
if images is not None:
frame = 0
if hasattr(images, '__call__'):
imgdata = images(frame)
if len(imgdata.shape) > 2:
imgdata = imgdata[:,:,0]
if flipimgx:
if flipimgy:
imgdata = imgdata[::-1,::-1]
else:
imgdata = imgdata[::-1,:]
elif flipimgy:
imgdata = imgdata[:,::-1]
else:
imgdata = imgdata[:,:]
else:
imgdata = get_image_data(images, frame, mono=mono, flipimgx=flipimgx, flipimgy=flipimgy)
img = ax.imshow( imgdata, extent=extent, cmap=plt.get_cmap(imagecolormap), zorder=-10)
for fly in flies:
ax.add_collection(fly)
def init_plot():
for fly in flies:
fly.set_color('none')
if images is None:
return flies
else:
animation_objects = [fly for fly in flies]
animation_objects.insert(0, img)
return animation_objects
def updatefig(*args):
anim_params['frame'] += 1 + nskip
print anim_params['frame']
if anim_params['frame'] >= len(x[0])-1:
anim_params['frame'] = 1
anim_params['movie_finished'] = True
frame_start = anim_params['frame'] - ghost_tail
frame_end = anim_params['frame']
if frame_start < 0:
frame_start = 0
frames = np.arange(frame_start, frame_end, 1).tolist()
for i, fly in enumerate(flies):
if frame_end > len(x[i]):
colors = ['none' for f in range(len(x[i]))]
else:
colors = ['none' for f in range(len(x[i]))]
for f in frames:
try:
colors[f] = color_mappable.to_rgba(color[i][f])
except:
colors[f] = color[i]
if frame_end > len(x[i]):
edgecolors = ['none' for f in range(len(x[i]))]
else:
edgecolors = ['none' for f in range(len(x[i]))]
for f in frames:
edgecolors[f] = edgecolor[i]
fly.set_edgecolors(edgecolors)
fly.set_facecolors(colors)
if images is not None:
if hasattr(images, '__call__'):
imgdata = images(frames[-1])
if len(imgdata.shape) > 2:
imgdata = imgdata[:,:,0]
if flipimgx:
if flipimgy:
imgdata = imgdata[::-1,::-1]
else:
imgdata = imgdata[::-1,:]
elif flipimgy:
imgdata = imgdata[:,::-1]
else:
imgdata = imgdata[:,:]
else:
imgdata = get_image_data(images, frames[-1], mono=mono, flipimgx=flipimgx, flipimgy=flipimgy)
img.set_array(imgdata)
if save and not anim_params['movie_finished']:
print 'saving frame: ', str(anim_params['frame']), ' -- if the animation you see is strange, do not worry, look at the pngs'
frame_prefix = '_tmp'
frame_prefix = os.path.join(save_movie_path, frame_prefix)
strnum = str(anim_params['frame'])
num_frame_digits = np.ceil( np.log10(len(x[0]))) + 1
while len(strnum) < num_frame_digits:
strnum = '0' + strnum
frame_name = frame_prefix + '_' + strnum + '_' + '.png'
fig.savefig(frame_name, format='png')
if save and anim_params['movie_finished']:
print
print 'Movie finished saving! Close the plot screen now.'
print 'PNGs are at: ', save_movie_path
print 'To turn the PNGs into a movie, you can run this command from inside the directory with the tmp files: '
print 'mencoder \'mf://*.png\' -mf type=png:fps=30 -ovc lavc -lavcopts vcodec=mpeg4 -oac copy -o animation.avi'
if images is None:
return flies
else:
animation_objects = [fly for fly in flies]
animation_objects.insert(0, img)
return animation_objects
# generate and set limits:
if xlim is None:
xlim = [np.min(x[0]), np.max(x[0])]
if ylim is None:
ylim = [np.min(y[0]), np.max(y[0])]
ax.set_xlim(xlim[0], xlim[1])
ax.set_ylim(ylim[0], ylim[1])
ax.set_aspect(aspect)
#
fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xlim, yticks=ylim)
ax.set_xlabel('x position, m')
ax.set_ylabel('y position, m')
ani = animation.FuncAnimation(fig, updatefig, init_func=init_plot, fargs=anim_params, interval=50, blit=True)
plt.show()