orientation_trajectory_plot.py

#########################################################################################################
## OVERVIEW
#########################################################################################################

# Plot trajectory with oriented, colored wedges. See file end for example.

#########################################################################################################
## IMPORTS
#########################################################################################################

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.collections import PatchCollection

#########################################################################################################
## FUNCTIONS
#########################################################################################################

def get_wedges_for_heading_plot(x, y, color, orientation, size_radius=0.1, size_angle=20, colormap='jet', colornorm=None, size_radius_range=(0.01,.1), size_radius_norm=None, edgecolor='none', alpha=1, flip=True, deg=True, nskip=0, center_offset_fraction=0.75):
    '''
    Returns a Patch Collection of Wedges, with arbitrary color and orientation
    
    Outputs:
    Patch Collection
    
    Inputs:
    x, y        - x and y positions (np.array or list, each of length N)
    color       - values to color wedges by (np.array or list, length N), OR color string. 
       colormap - specifies colormap to use (string, eg. 'jet')
       norm     - specifies range you'd like to normalize to, 
                  if none, scales to min/max of color array (2-tuple, eg. (0,1) )
    orientation - angles are in degrees, use deg=False to convert radians to degrees
    size_radius - radius of wedge, in same units as x, y. Can be list or np.array, length N, for changing sizes
       size_radius_norm - specifies range you'd like to normalize size_radius to, if size_radius is a list/array
                  should be tuple, eg. (0.01, .1)
    size_angle  - angular extent of wedge, degrees. Can be list or np.array, length N, for changing sizes
    edgecolor   - color for lineedges, string or np.array of length N
    alpha       - transparency (single value, between 0 and 1)
    flip        - flip orientations by 180 degrees, default = True
    nskip       - allows you to skip between points to make the points clearer, nskip=1 skips every other point
    center_offset_fraction  - (float in range (0,1) ) - 0 means (x,y) is at the tip, 1 means (x,y) is at the edge
    '''
    cmap = plt.get_cmap(colormap)
    
    # norms
    if colornorm is None and type(color) is not str:
        colornorm = plt.Normalize(np.min(color), np.max(color))
    elif type(color) is not str:
        colornorm = plt.Normalize(colornorm[0], colornorm[1])
    if size_radius_norm is None:
        size_radius_norm = plt.Normalize(np.min(size_radius), np.max(size_radius), clip=True)
    else:
        size_radius_norm = plt.Normalize(size_radius_norm[0], size_radius_norm[1], clip=True)
        
    indices_to_plot = np.arange(0, len(x), nskip+1)
        
    # fix orientations
    if type(orientation) is list:
        orientation = np.array(orientation)
    if deg is False:
        orientation = orientation*180./np.pi
    if flip:
        orientation += 180
    
    flycons = []
    n = 0
    for i in indices_to_plot:
        # wedge parameters
        if type(size_radius) is list or type(size_radius) is np.array or type(size_radius) is np.ndarray: 
            r = size_radius_norm(size_radius[i])*(size_radius_range[1]-size_radius_range[0]) + size_radius_range[0] 
        else: r = size_radius
        
        if type(size_angle) is list or type(size_angle) is np.array or type(size_angle) is np.ndarray: 
            angle_swept = size_radius[i]
        else: angle_swept = size_radius
        theta1 = orientation[i] - size_angle/2.
        theta2 = orientation[i] + size_angle/2.
        
        center = [x[i], y[i]]
        center[0] -= np.cos(orientation[i]*np.pi/180.)*r*center_offset_fraction
        center[1] -= np.sin(orientation[i]*np.pi/180.)*r*center_offset_fraction
        
        wedge = patches.Wedge(center, r, theta1, theta2)
        flycons.append(wedge)
        
    # add collection and color it
    pc = PatchCollection(flycons, cmap=cmap, norm=colornorm)
    
    # set properties for collection
    pc.set_edgecolors(edgecolor)
    if type(color) is list or type(color) is np.array or type(color) is np.ndarray:
        if type(color) is list:
            color = np.asarray(color)
        pc.set_array(color[indices_to_plot])
    else:
        pc.set_facecolors(color)
    pc.set_alpha(alpha)
    
    return pc

def colorline_with_heading(ax, x, y, color, orientation, size_radius=0.1, size_angle=20, colormap='jet', colornorm=None, size_radius_range=(0.01,.1), size_radius_norm=None, edgecolor='none', alpha=1, flip=True, deg=True, nskip=0, use_center='center', show_centers=True, center_offset_fraction=0.75, center_point_size=2):
    '''
    Plots a trajectory with colored wedge shapes to indicate orientation. 
    See function get_wedges_for_heading_plot for details
    
    Additional options:
    
    show_centers      - (bool) - show a black dot where the actual point is - shows where the center of the wedge is 
    center_point_size - markersize for center, if show_centers
    '''
        
    pc = get_wedges_for_heading_plot(x, y, color, orientation, size_radius=size_radius, size_angle=size_angle, colormap=colormap, colornorm=colornorm, size_radius_range=size_radius_range, size_radius_norm=size_radius_norm, edgecolor=edgecolor, alpha=alpha, flip=flip, deg=deg, nskip=nskip, center_offset_fraction=center_offset_fraction)
        
    ax.add_collection(pc)
    
    if show_centers:
        indices_to_plot = np.arange(0, len(x), nskip+1)
        ax.plot(x[indices_to_plot],y[indices_to_plot],'.', color='black', markersize=center_point_size)

def plot_trajec_with_orientation(xpos, ypos, theta, color, ax=None, size_radius=5, nskip = 190, colormap='bone_r'):
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111)
    
    colorline_with_heading(ax, xpos, ypos, color, theta, 
                                    nskip=nskip, size_radius=size_radius, deg=False, colormap=colormap, center_point_size=0.0001,
                                    colornorm=[0.05*color.max(),color.max()], show_centers=False)

    ax.set_aspect('equal')
    xrang = xpos.max() - xpos.min()
    xrang = np.max([xrang, 0.1])
    yrang = ypos.max() - ypos.min()
    yrang = np.max([yrang, 0.1])
    ax.set_xlim(xpos.min()-0.1*xrang, xpos.max()+0.1*xrang)
    ax.set_ylim(ypos.min()-0.1*yrang, ypos.max()+0.1*yrang)

    
#########################################################################################################
## EXAMPLE
#########################################################################################################

xpos = np.linspace(0, 10, 1000)
ypos = np.sin(x)
theta = np.tan(y)
color = x

plot_trajec_with_orientation(xpos, ypos, theta, color, size_radius=1, nskip=50)