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 set_log_angle_ticks(ax):
    yticks = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
    xticks = np.log(np.array([5, 10, 20, 30, 60, 90, 180])*np.pi/180.).tolist()
    fpl.adjust_spines(ax, ['left', 'bottom'], yticks=yticks, xticks=xticks, smart_bounds=True)
    
    xticklabels = [str(i) for i in [5, 10, 20, 30, 60, 90, 180]]
    ax.set_xticklabels(xticklabels)
    
    ax.set_xlabel('Retinal size, deg')
    ax.set_ylabel('Speed, m/s')
def heatmap_of_straight_landing_trajecs(dataset_landing, dataset_flyby):

    speeds = []
    angles_subtended = []
    for k, trajec in dataset_landing.trajecs.items():
        if trajec.classification == 'no_saccade_after_deceleration':
            speeds.extend(trajec.speed[0:trajec.frame_of_landing].tolist())
            angles_subtended.extend(trajec.angle_subtended_by_post[0:trajec.frame_of_landing].tolist())
        
    fig = plt.figure()
    ax = fig.add_subplot(111)
    
    fpl.histogram2d(ax, np.log(np.array(angles_subtended)), np.array(speeds), bins=50, normed=False, histrange=None, weights=None, logcolorscale=True, colormap='jet', interpolation='bicubic')
    
    fig_saccades_flyby = plt.figure()
    ax_saccades_flyby = fig_saccades_flyby.add_subplot(111)
    
    
    # now look at speed and retinal size for flyby trajectories at last saccade if headed towards post:
    for k, trajec in dataset_flyby.trajecs.items():
        fd, fs = get_frame_of_decel_and_last_saccade(trajec)
        for sac_range in trajec.sac_ranges:
            if fs in sac_range:
                if trajec.angle_subtended_by_post[sac_range[0]] > 0.*np.pi/180.:
                    
                #if np.abs(trajec.angle_to_post[sac_range[0]]) < 180.*np.pi/180.:
                    ax.plot(np.log(trajec.angle_subtended_by_post[sac_range[0]]), trajec.speed[sac_range[0]], '.', markersize=2, color='white', markeredgecolor='black', linewidth=0.5)

                    angle = sac.get_angle_of_saccade(trajec, sac_range)
                    ax_saccades_flyby.plot(-1*trajec.angle_to_post[sac_range[0]], angle, '.', color='red', markersize=2)

    set_log_angle_ticks(ax)
    ax.set_aspect('auto')
        
    fig.savefig('landing_deceleration_heatmap.pdf')
        
    fpl.adjust_spines(ax_saccades_flyby, ['left', 'bottom'], yticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], xticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], smart_bounds=True)
    ax_saccades_flyby.set_xlim([-np.pi, np.pi])
    ax_saccades_flyby.set_ylim([-np.pi, np.pi])
    
    deg_ticks = ['-180', '-90', '0', '90', '180']
    ax_saccades_flyby.set_xticklabels(deg_ticks)
    ax_saccades_flyby.set_yticklabels(deg_ticks)
    
    ax_saccades_flyby.set_xlabel('Angle to post, deg')
    ax_saccades_flyby.set_ylabel('Turn angle, deg')
    
    fig_saccades_flyby.savefig('saccades_flyby.pdf', format='pdf')
def plot_deceleration_for_flybys(dataset):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    
    fig_saccade_flyby = plt.figure()
    ax_saccade_flyby = fig_saccade_flyby.add_subplot(111)
    
    for k, trajec in dataset.trajecs.items():
        fd, fs = get_frame_of_decel_and_last_saccade(trajec)
        if np.abs(trajec.angle_to_post[fd]) < 180.*np.pi/180.:
            
            if 0:#fs is not None:
                
                if trajec.angle_subtended_by_post[fs] > 20*np.pi/180.:
                
                    if fs > fd:
                        ax.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='green')
                        #ax.plot(np.log(trajec.angle_subtended_by_post[fs]), trajec.speed[fs], '.', color='green')
                        #ax.plot([np.log(trajec.angle_subtended_by_post[fd]), np.log(trajec.angle_subtended_by_post[fs])], [trajec.speed[fd], trajec.speed[fs]], '-', color='black')
                        sac_range = get_saccade_range_from_first_frame(trajec, fs)
                        angle = sac.get_angle_of_saccade(trajec, sac_range)
                        ax_saccade_flyby.plot(trajec.angle_to_post[fs], angle, '.', color='green')
                    else:
                        print trajec.key # saccade before deceleration
                        ax.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='red')
                        sac_range = get_saccade_range_from_first_frame(trajec, fs)
                        angle = sac.get_angle_of_saccade(trajec, sac_range)
                        ax_saccade_flyby.plot(trajec.angle_to_post[fs], angle, '.', color='red')
                    
            if fs is not None:
                if np.abs(fs-fd) > 10 and np.abs(trajec.angle_to_post[fd]) < 60.*np.pi/180.:
                    ax.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
                    
            
            
        
        # try looking for trajectories that slowed down after saccading?
    
    set_log_angle_ticks(ax)
    fig.savefig('deceleration_for_flybys.pdf', format='pdf') 
    
    fpl.adjust_spines(ax_saccade_flyby, ['left', 'bottom'], yticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], xticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], smart_bounds=True)
    
    ax_saccade_flyby.set_xlim([-np.pi, np.pi])
    ax_saccade_flyby.set_ylim([-np.pi, np.pi])
    
    fig_saccade_flyby.savefig('last_saccade_of_flyby.pdf', format='pdf')
def heatmap(ax, dataset, axis='xy'):  

    # collect data
    xpos = np.array([])
    ypos = np.array([])
    zpos = np.array([])
    
    for key, trajec in dataset.trajecs.items():
        xpos = np.hstack( (xpos, trajec.positions[:,0]) )
        ypos = np.hstack( (ypos, trajec.positions[:,1]) )
        zpos = np.hstack( (zpos, trajec.positions[:,2]) )
    
    if axis == 'xy':
        fpl.histogram2d(ax, xpos, ypos, bins=100, logcolorscale=True)
    elif axis == 'xz':
        fpl.histogram2d(ax, xpos, zpos, bins=100, logcolorscale=True)
    elif axis == 'yz':
        fpl.histogram2d(ax, ypos, zpos, bins=100, logcolorscale=True)
    
    if axis == 'xy':
        post = patches.Circle( (0, 0), radius=0.01, facecolor='black', edgecolor='none', alpha=1, linewidth=0)
    elif axis == 'yz':
        post = patches.Rectangle( (-.01,0), width=0.02, facecolor='black', height=.16, edgecolor='none', alpha=1, linewidth=0)
    artists = [post]
    
    if artists is not None:
        for artist in artists:
            ax.add_artist(artist)
    
    if 'x' in axis:
        xticks = [-0.15, 0, 0.15]
    else:
        xticks = None
        
    if 'z' in axis:
        yticks = [0, .15, .30]
    else:
        yticks = None
        
    if axis == 'xy':
        ax.set_xlim(-.15, .15)
        ax.set_ylim(.2,-.8)
        
    fpl.adjust_spines(ax, ['left', 'bottom'], xticks=xticks, yticks=yticks)
    
    ax.set_aspect('equal')
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 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 print_frame_diff_btwn_decel_and_saccade(dataset):
    
    neg = 0
    pos = 0
    none = 0
    
    fig_neg = plt.figure()
    ax_neg = fig_neg.add_subplot(111)
    
    fig_pos = plt.figure()
    ax_pos = fig_pos.add_subplot(111)
    
    fig_none = plt.figure()
    ax_none = fig_none.add_subplot(111)
    
    fig_saccades = plt.figure()
    ax_saccades = fig_saccades.add_subplot(111)
    
    fig_deceleration = plt.figure()
    ax_deceleration = fig_deceleration.add_subplot(111)
        
    for k, trajec in dataset.trajecs.items():
        fd, fs = get_frame_of_decel_and_last_saccade(trajec)
        pos_frame_minus_landing = []
        
        if trajec.classification == 'straight':
            none += 1
            ax_none.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
            ax_deceleration.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
        else:
            try:
                difference = fd - fs
                if difference < 0:
                    neg += 1
                    ax_neg.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
                    ax_neg.plot(np.log(trajec.angle_subtended_by_post[fs]), trajec.speed[fs], '.', color='green')
                    ax_neg.plot([np.log(trajec.angle_subtended_by_post[fd]), np.log(trajec.angle_subtended_by_post[fs])], [trajec.speed[fd], trajec.speed[fs]], '-', color='black')
                    ax_deceleration.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='red')
                    
                    sac_range = get_saccade_range_from_first_frame(trajec, fs)
                    angle = sac.get_angle_of_saccade(trajec, sac_range)
                    print trajec.angle_to_post, angle
                    ax_saccades.plot(trajec.angle_to_post[fs], angle, '.')
                    print trajec.angle_to_post, angle
                    
                elif difference > 0:
                    pos += 1
                    pos_frame_minus_landing.append(trajec.frame_of_landing - fs)
                    ax_pos.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
                    ax_pos.plot(np.log(trajec.angle_subtended_by_post[fs]), trajec.speed[fs], '.', color='green')
                    #ax_pos.plot([np.log(trajec.angle_subtended_by_post[fd]), np.log(trajec.angle_subtended_by_post[fs])], [trajec.speed[fd], trajec.speed[fs]], '-', color='black')
                    ax_pos.plot(np.log(trajec.angle_subtended_by_post[0:trajec.frame_of_landing]), trajec.speed[0:trajec.frame_of_landing], '-', color='black', linewidth=0.5)
                    ax_deceleration.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='green')
            except:
                none += 1
                ax_none.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
                ax_deceleration.plot(np.log(trajec.angle_subtended_by_post[fd]), trajec.speed[fd], '.', color='purple')
            
            
    yticks = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
    xticks = np.log(np.array([10, 20, 30, 60, 90, 180])*np.pi/180.).tolist()
    
    set_log_angle_ticks(ax_neg)
    set_log_angle_ticks(ax_pos)
    set_log_angle_ticks(ax_none)
    set_log_angle_ticks(ax_deceleration)
            
    fpl.adjust_spines(ax_saccades, ['left', 'bottom'], yticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], xticks=[-np.pi, -np.pi/2., 0, np.pi/2., np.pi], smart_bounds=True)
            
    fig_neg.savefig('decel_and_saccade_neg.pdf', format='pdf')
    fig_pos.savefig('decel_and_saccade_pos.pdf', format='pdf')    
    fig_none.savefig('decel_and_saccade_none.pdf', format='pdf')   
    fig_deceleration.savefig('decel_for_all.pdf', format='pdf')    
    fig_saccades.savefig('saccade_for_saccade_after_deceleration.pdf', format='pdf')
    
    
    print 'neg: ', neg, 'pos: ', pos, 'none: ', none
    print pos_frame_minus_landing