Deaton
Deaton

Reputation: 353

How do I specify an arrow-like linestyle in Matplotlib?

I would like to display a set of xy-data in Matplotlib in such a way as to indicate a particular path. Ideally, the linestyle would be modified to use an arrow-like patch. I have created a mock-up, shown below (using Omnigraphsketcher). It seems like I should be able to override one of the common linestyle declarations ('-', '--', ':', etc) to this effect.

Note that I do NOT want to simply connect each datapoint with a single arrow---the actually data points are not uniformly spaced and I need consistent arrow spacing.

enter image description here

Upvotes: 24

Views: 21872

Answers (5)

Suuuehgi
Suuuehgi

Reputation: 4990

If you can live without the fancy going-around-the-edge / fixed-length arrows, here's a poor man's version sub-diving the segments in approx. ds long segments. If ds is not too large, the variance is negligible in my eyes.

from matplotlib import pyplot as plt
import numpy as np

np.random.seed(42)
x, y = 10*np.random.rand(2,8)

# length of line segment
ds=1

# number of line segments per interval
Ns = np.round(np.sqrt( (x[1:]-x[:-1])**2 + (y[1:]-y[:-1])**2 ) / ds).astype(int)

# sub-divide intervals w.r.t. Ns
subdiv = lambda x, Ns=Ns: np.concatenate([ np.linspace(x[ii], x[ii+1], Ns[ii]) for ii, _ in enumerate(x[:-1]) ])

x, y = subdiv(x), subdiv(y)

plt.quiver(x[:-1], y[:-1], x[1:]-x[:-1], y[1:]-y[:-1], scale_units='xy', angles='xy', scale=1, width=.004, headlength=4, headwidth=4)

plot

Upvotes: 0

Prof Huster
Prof Huster

Reputation: 1018

Here is a modified and streamlined version of Duarte's code. I had problems when I ran his code with various data sets and aspect ratios, so I cleaned it up and used FancyArrowPatches for the arrows. Note the example plot has a scale 1,000,000 times different in x than y.

I also changed to drawing the arrow in display coordinates so the different scaling on the x and y axes would not change the arrow lengths.

Along the way I found a bug in matplotlib's FancyArrowPatch that bombs when plotting a purely vertical arrow. I found a work-around that is in my code.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches


def arrowplot(axes, x, y, nArrs=30, mutateSize=10, color='gray', markerStyle='o'): 
    '''arrowplot : plots arrows along a path on a set of axes
        axes   :  the axes the path will be plotted on
        x      :  list of x coordinates of points defining path
        y      :  list of y coordinates of points defining path
        nArrs  :  Number of arrows that will be drawn along the path
        mutateSize :  Size parameter for arrows
        color  :  color of the edge and face of the arrow head
        markerStyle : Symbol
    
        Bugs: If a path is straight vertical, the matplotlab FanceArrowPatch bombs out.
          My kludge is to test for a vertical path, and perturb the second x value
          by 0.1 pixel. The original x & y arrays are not changed
    
        MHuster 2016, based on code by 
    '''
    # recast the data into numpy arrays
    x = np.array(x, dtype='f')
    y = np.array(y, dtype='f')
    nPts = len(x)

    # Plot the points first to set up the display coordinates
    axes.plot(x,y, markerStyle, ms=5, color=color)

    # get inverse coord transform
    inv = axes.transData.inverted()

    # transform x & y into display coordinates
    # Variable with a 'D' at the end are in display coordinates
    xyDisp = np.array(axes.transData.transform(list(zip(x,y))))
    xD = xyDisp[:,0]
    yD = xyDisp[:,1]

    # drD is the distance spanned between pairs of points
    # in display coordinates
    dxD = xD[1:] - xD[:-1]
    dyD = yD[1:] - yD[:-1]
    drD = np.sqrt(dxD**2 + dyD**2)

    # Compensating for matplotlib bug
    dxD[np.where(dxD==0.0)] = 0.1


    # rtotS is the total path length
    rtotD = np.sum(drD)

    # based on nArrs, set the nominal arrow spacing
    arrSpaceD = rtotD / nArrs

    # Loop over the path segments
    iSeg = 0
    while iSeg < nPts - 1:
        # Figure out how many arrows in this segment.
        # Plot at least one.
        nArrSeg = max(1, int(drD[iSeg] / arrSpaceD + 0.5))
        xArr = (dxD[iSeg]) / nArrSeg # x size of each arrow
        segSlope = dyD[iSeg] / dxD[iSeg]
        # Get display coordinates of first arrow in segment
        xBeg = xD[iSeg]
        xEnd = xBeg + xArr
        yBeg = yD[iSeg]
        yEnd = yBeg + segSlope * xArr
        # Now loop over the arrows in this segment
        for iArr in range(nArrSeg):
            # Transform the oints back to data coordinates
            xyData = inv.transform(((xBeg, yBeg),(xEnd,yEnd)))
            # Use a patch to draw the arrow
            # I draw the arrows with an alpha of 0.5
            p = patches.FancyArrowPatch( 
                xyData[0], xyData[1], 
                arrowstyle='simple',
                mutation_scale=mutateSize,
                color=color, alpha=0.5)
            axes.add_patch(p)
            # Increment to the next arrow
            xBeg = xEnd
            xEnd += xArr
            yBeg = yEnd
            yEnd += segSlope * xArr
        # Increment segment number
        iSeg += 1
    
if __name__ == '__main__':
    import numpy as np
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(111)
    # my random data
    xScale = 1e6
    np.random.seed(1)
    x = np.random.random(10) * xScale
    y = np.random.random(10)
    arrowplot(ax, x, y, nArrs=4*(len(x)-1), mutateSize=10, color='red')
    xRng = max(x) - min(x)
    ax.set_xlim(min(x) - 0.05*xRng, max(x) + 0.05*xRng)
    yRng = max(y) - min(y)
    ax.set_ylim(min(y) - 0.05*yRng, max(y) + 0.05*yRng)
    plt.show()

enter image description here

Upvotes: 3

user1769889
user1769889

Reputation: 175

Vectorized version of Yann's answer:

import numpy as np
import matplotlib.pyplot as plt

def distance(data):
    return np.sum((data[1:] - data[:-1]) ** 2, axis=1) ** .5

def draw_path(path):
    HEAD_WIDTH = 2
    HEAD_LEN = 3

    fig = plt.figure()
    axes = fig.add_subplot(111)

    x = path[:,0]
    y = path[:,1]
    axes.plot(x, y)

    theta = np.arctan2(y[1:] - y[:-1], x[1:] - x[:-1])
    dist = distance(path) - HEAD_LEN

    x = x[:-1]
    y = y[:-1]
    ax = x + dist * np.sin(theta)
    ay = y + dist * np.cos(theta)

    for x1, y1, x2, y2 in zip(x,y,ax-x,ay-y):
        axes.arrow(x1, y1, x2, y2, head_width=HEAD_WIDTH, head_length=HEAD_LEN)
    plt.show()

Upvotes: 3

Pedro M Duarte
Pedro M Duarte

Reputation: 28093

Very nice answer by Yann, but by using arrow the resulting arrows can be affected by the axes aspect ratio and limits. I have made a version that uses axes.annotate() instead of axes.arrow(). I include it here for others to use.

In short this is used to plot arrows along your lines in matplotlib. The code is shown below. It can still be improved by adding the possibility of having different arrowheads. Here I only included control for the width and length of the arrowhead.

import numpy as np
import matplotlib.pyplot as plt


def arrowplot(axes, x, y, narrs=30, dspace=0.5, direc='pos', \
                          hl=0.3, hw=6, c='black'): 
    ''' narrs  :  Number of arrows that will be drawn along the curve

        dspace :  Shift the position of the arrows along the curve.
                  Should be between 0. and 1.

        direc  :  can be 'pos' or 'neg' to select direction of the arrows

        hl     :  length of the arrow head 

        hw     :  width of the arrow head        

        c      :  color of the edge and face of the arrow head  
    '''

    # r is the distance spanned between pairs of points
    r = [0]
    for i in range(1,len(x)):
        dx = x[i]-x[i-1] 
        dy = y[i]-y[i-1] 
        r.append(np.sqrt(dx*dx+dy*dy))
    r = np.array(r)

    # rtot is a cumulative sum of r, it's used to save time
    rtot = []
    for i in range(len(r)):
        rtot.append(r[0:i].sum())
    rtot.append(r.sum())

    # based on narrs set the arrow spacing
    aspace = r.sum() / narrs

    if direc is 'neg':
        dspace = -1.*abs(dspace) 
    else:
        dspace = abs(dspace)

    arrowData = [] # will hold tuples of x,y,theta for each arrow
    arrowPos = aspace*(dspace) # current point on walk along data
                                 # could set arrowPos to 0 if you want
                                 # an arrow at the beginning of the curve

    ndrawn = 0
    rcount = 1 
    while arrowPos < r.sum() and ndrawn < narrs:
        x1,x2 = x[rcount-1],x[rcount]
        y1,y2 = y[rcount-1],y[rcount]
        da = arrowPos-rtot[rcount]
        theta = np.arctan2((x2-x1),(y2-y1))
        ax = np.sin(theta)*da+x1
        ay = np.cos(theta)*da+y1
        arrowData.append((ax,ay,theta))
        ndrawn += 1
        arrowPos+=aspace
        while arrowPos > rtot[rcount+1]: 
            rcount+=1
            if arrowPos > rtot[-1]:
                break

    # could be done in above block if you want
    for ax,ay,theta in arrowData:
        # use aspace as a guide for size and length of things
        # scaling factors were chosen by experimenting a bit

        dx0 = np.sin(theta)*hl/2. + ax
        dy0 = np.cos(theta)*hl/2. + ay
        dx1 = -1.*np.sin(theta)*hl/2. + ax
        dy1 = -1.*np.cos(theta)*hl/2. + ay

        if direc is 'neg' :
          ax0 = dx0 
          ay0 = dy0
          ax1 = dx1
          ay1 = dy1 
        else:
          ax0 = dx1 
          ay0 = dy1
          ax1 = dx0
          ay1 = dy0 

        axes.annotate('', xy=(ax0, ay0), xycoords='data',
                xytext=(ax1, ay1), textcoords='data',
                arrowprops=dict( headwidth=hw, frac=1., ec=c, fc=c))


    axes.plot(x,y, color = c)
    axes.set_xlim(x.min()*.9,x.max()*1.1)
    axes.set_ylim(y.min()*.9,y.max()*1.1)


if __name__ == '__main__':
    fig = plt.figure()
    axes = fig.add_subplot(111)

    # my random data
    scale = 10 
    np.random.seed(101)
    x = np.random.random(10)*scale
    y = np.random.random(10)*scale
    arrowplot(axes, x, y ) 

    plt.show()

The resulting figure can be seen here:

enter image description here

Upvotes: 8

Yann
Yann

Reputation: 35503

Here's a starting off point:

  1. Walk along your line at fixed steps (aspace in my example below) .

    A. This involves taking steps along the line segments created by two sets of points (x1,y1) and (x2,y2).

    B. If your step is longer than the line segment, shift to the next set of points.

  2. At that point determine the angle of the line.

  3. Draw an arrow with an inclination corresponding to the angle.

I wrote a little script to demonstrate this:

import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
axes = fig.add_subplot(111)

# my random data
scale = 10 
np.random.seed(101)
x = np.random.random(10)*scale
y = np.random.random(10)*scale

# spacing of arrows
aspace = .1 # good value for scale of 1
aspace *= scale

# r is the distance spanned between pairs of points
r = [0]
for i in range(1,len(x)):
    dx = x[i]-x[i-1]
    dy = y[i]-y[i-1]
    r.append(np.sqrt(dx*dx+dy*dy))
r = np.array(r)

# rtot is a cumulative sum of r, it's used to save time
rtot = []
for i in range(len(r)):
    rtot.append(r[0:i].sum())
rtot.append(r.sum())

arrowData = [] # will hold tuples of x,y,theta for each arrow
arrowPos = 0 # current point on walk along data
rcount = 1 
while arrowPos < r.sum():
    x1,x2 = x[rcount-1],x[rcount]
    y1,y2 = y[rcount-1],y[rcount]
    da = arrowPos-rtot[rcount] 
    theta = np.arctan2((x2-x1),(y2-y1))
    ax = np.sin(theta)*da+x1
    ay = np.cos(theta)*da+y1
    arrowData.append((ax,ay,theta))
    arrowPos+=aspace
    while arrowPos > rtot[rcount+1]: 
        rcount+=1
        if arrowPos > rtot[-1]:
            break

# could be done in above block if you want
for ax,ay,theta in arrowData:
    # use aspace as a guide for size and length of things
    # scaling factors were chosen by experimenting a bit
    axes.arrow(ax,ay,
               np.sin(theta)*aspace/10,np.cos(theta)*aspace/10, 
               head_width=aspace/8)


axes.plot(x,y)
axes.set_xlim(x.min()*.9,x.max()*1.1)
axes.set_ylim(y.min()*.9,y.max()*1.1)

plt.show()

This example results in this figure: enter image description here

There's plenty of room for improvement here, for starters:

  1. One can use FancyArrowPatch to customize the look of the arrows.
  2. One can add a further test when creating the arrows to make sure they don't extend beyond the line. This will be relevant to arrows created at or near a vertex where the line changes direction sharply. This is the case for the right most point above.
  3. One can make a method from this script that will work across a broader range of cases, ie make it more portable.

While looking into this, I discovered the quiver plotting method. It might be able to replace the above work, but it wasn't immediately obvious that this was guaranteed.

Upvotes: 10

Related Questions