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 get_gaussian_model_of_data_files(path):

    xpositions, ypositions, zpositions, pid_vals = get_means_of_data_files(
        path)

    ypositions = compile_data_to_list(ypositions)
    zpositions = compile_data_to_list(zpositions)
    pid_vals = compile_data_to_list(pid_vals)

    gm = data_fit.models.GaussianModel2D()
    gm.fit_with_guess(pid_vals, [ypositions, zpositions])

    #gm.parameters['magnitude'] = 200

    if 1:

        fig = plt.figure()
        ax = fig.add_subplot(111, rasterized=True)

        im, extent = gm.get_array_2d([-.15, .15], [-.15, .15], 0.001)
        ax.imshow(im, extent=extent, origin='lower')
        #ax.plot(x,y,'b.')

        fpl.scatter(ax,
                    ypositions,
                    zpositions,
                    pid_vals,
                    colormap='jet',
                    colornorm=[0, gm.parameters['magnitude']],
                    radius=0.001,
                    use_ellipses=False)

        ax.set_xlim([-.15, .15])
        ax.set_ylim([-.15, .15])
        ax.set_xlabel('crosswind position, meters')
        ax.set_ylabel('altitude position, meters')
        ax.set_aspect('equal')
        #plt.show()

        filename = 'odor_plume_plot_2d.pdf'
        filename_with_path = os.path.join(path, filename)

        fig.savefig(filename_with_path, format='pdf')

    return gm
def get_gaussian_model_of_data_files(path):
    
    xpositions, ypositions, zpositions, pid_vals = get_means_of_data_files(path)
    
    ypositions = compile_data_to_list(ypositions)
    zpositions = compile_data_to_list(zpositions)
    pid_vals = compile_data_to_list(pid_vals)

    gm = data_fit.models.GaussianModel2D()
    gm.fit_with_guess(pid_vals, [ypositions, zpositions])
    
    #gm.parameters['magnitude'] = 200
    
    if 1:
    
        fig = plt.figure()
        ax = fig.add_subplot(111, rasterized=True)
        
        im, extent = gm.get_array_2d([-.15, .15], [-.15,.15], 0.001)
        ax.imshow(im, extent=extent, origin='lower')
        #ax.plot(x,y,'b.')
    
        fpl.scatter(ax, ypositions, zpositions, pid_vals, colormap='jet', colornorm=[0,gm.parameters['magnitude']], radius=0.001, use_ellipses=False)
        
        ax.set_xlim([-.15,.15])
        ax.set_ylim([-.15,.15])
        ax.set_xlabel('crosswind position, meters')
        ax.set_ylabel('altitude position, meters')
        ax.set_aspect('equal')
        #plt.show()
        
        filename = 'odor_plume_plot_2d.pdf'
        filename_with_path = os.path.join(path, filename)
        
        fig.savefig(filename_with_path, format='pdf')
    
    return gm
    def make_plot(ax, births, deaths, keys, axis="yz"):

        if axis == "yz":
            x_births = births[:, [1]]
            y_births = births[:, [2]]
            x_deaths = deaths[:, [1]]
            y_deaths = deaths[:, [2]]

        if axis == "xy":
            x_births = births[:, [0]]
            y_births = births[:, [1]]
            x_deaths = deaths[:, [0]]
            y_deaths = deaths[:, [1]]

        if axis == "xz":
            x_births = births[:, [0]]
            y_births = births[:, [2]]
            x_deaths = deaths[:, [0]]
            y_deaths = deaths[:, [2]]

        fpl.scatter(
            ax, x_births.flatten(), y_births.flatten(), color="green", radius=0.005, alpha=1, use_ellipses=False
        )
        fpl.scatter(ax, x_deaths.flatten(), y_deaths.flatten(), color="red", radius=0.005, alpha=1, use_ellipses=False)
def plot_odor_heading_book(pp, threshold_odor, path, config, dataset, odor_stimulus, keys=None, axis='xy'):

    fig = plt.figure(figsize=(4,4))
    ax = fig.add_subplot(111)
    

    saccade_angles_after_odor = []
    heading_at_saccade_initiation = []
    odor_at_saccade = []
    saccade_number = []
                
    if 1:
        for key in keys:
            trajec = dataset.trajecs[key]
            
            #if trajec.positions[0,0] < 0.2:
            #    continue
            
            frames_in_odor = np.where(trajec.odor > threshold_odor)[0]
            odor_blocks = hf.find_continuous_blocks(frames_in_odor, 5, return_longest_only=False)
            
            b = 0
            for block in odor_blocks:
                if len(block) < 5:
                    continue
                first_sac = None
                
                if axis == 'xy':
                    saccades = trajec.saccades
                elif axis == 'altitude':
                    saccades = trajec.saccades_z
                
                for sac in saccades:
                    if trajec.positions[sac[0],0] < -0.1 or trajec.positions[sac[0],0] > 0.9:
                        continue
                    if np.abs(trajec.positions[sac[0],1]) > 0.05:
                        continue
                    if trajec.positions[sac[0],2] > 0.05 or trajec.positions[sac[0],2] < -0.01:
                        continue
                        
                    if sac[0] > block[0]:
                        if first_sac is None:
                            if trajec.time_fly[sac[0]] - trajec.time_fly[block[-1]] > 0.5:
                                break
                            first_sac = sac
                            break
                        
                if first_sac is not None:
                    next_sac = first_sac
                    if axis == 'xy':
                        angle_of_saccade = tac.get_angle_of_saccade(trajec, next_sac)
                        heading_prior_to_saccade = trajec.heading_smooth[next_sac[0]]
                    elif axis == 'altitude':
                        angle_of_saccade = tac.get_angle_of_saccade_z(trajec, next_sac)
                        heading_prior_to_saccade = trajec.heading_altitude_smooth[next_sac[0]]
                    
                        
                    saccade_angles_after_odor.append(angle_of_saccade)
                    heading_at_saccade_initiation.append(heading_prior_to_saccade)
                    odor_at_saccade.append(trajec.odor[next_sac[0]])
                    b += 1
                    saccade_number.append(b)
        
    saccade_angles_after_odor = np.array(saccade_angles_after_odor)
    heading_at_saccade_initiation = np.array(heading_at_saccade_initiation)
    odor_at_saccade = np.array(odor_at_saccade)
    saccade_number = np.array(saccade_number)
    
    print odor_stimulus, saccade_angles_after_odor.shape
    
    #ax.plot(heading_at_saccade_initiation*180./np.pi, saccade_angles_after_odor*180./np.pi, '.', markersize=3)
    fpl.scatter(ax, heading_at_saccade_initiation*180./np.pi, saccade_angles_after_odor*180./np.pi, color='black', radius=3, colornorm=[0,5])
    
    xpts = np.linspace(-180,180, 100)
    ax.plot(xpts, -1*xpts, color='red', zorder=-10)
    
    #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: ' + odor_stimulus + ' Visual Stim: ' + trajec.visual_stimulus
    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')