Number of arrowheads on matplotlib streamplot

Question

Is there anyway to increase the number of arrowheads on a matplotlib streamplot? Right now it appears as if three is only one arrowhead per streamline, which is a problem if I want to change to x/y axes limits to zoom in on the data.

Answer 1

Building on @Richard_wth's answer, I wrote a function to provide control on the location of the arrows on a streamplot. One can choose n arrows per streamline, or choose to have the arrows equally spaced on a streamline.

First, you do a normal streamplot, until you are happy with the location and number of streamlines. You keep the returned argument sp. For instance:

sp = ax.streamplot(x,y,u,v,arrowstyle='-',density=10)

What's important here is to have arrowstyle='-' so that arrows are not displayed.

Then, you can call the function streamQuiver (provided below) to control the arrows on the each streamline. If you want 3 arrows per streamline:

streamQuiver(ax, sp, n=3, ...)

If you want a streamline every 1.5 curvilinear length:

streamQuiver(ax, sp, spacing=1.5, ...)

where ... are options that would be passed to quiver. The function streamQuiver is probably not fully bulletproof and may need some additional handling for particular cases. It relies on 4 subfunctions:

• curve_coord to get the curvilinear length along a path
• curve extract to extract equidistant point along a path
• seg_to_lines to convert the segments from streamplot into continuous lines. There might be a better way to do that!
• lines_to_arrows: this is the main function that extract arrows on each lines

Here's an example where the arrows are at equidistant points on each streamlines.

import numpy as np
import matplotlib.pyplot as plt

def streamQuiver(ax,sp,*args,spacing=None,n=5,**kwargs):
    """ Plot arrows from streamplot data  
    The number of arrows per streamline is controlled either by `spacing` or by `n`.
    See `lines_to_arrows`.
    """
    def curve_coord(line=None):
        """ return curvilinear coordinate """
        x=line[:,0]
        y=line[:,1]
        s     = np.zeros(x.shape)
        s[1:] = np.sqrt((x[1:]-x[0:-1])**2+ (y[1:]-y[0:-1])**2)
        s     = np.cumsum(s)                                  
        return s

    def curve_extract(line,spacing,offset=None):
        """ Extract points at equidistant space along a curve"""
        x=line[:,0]
        y=line[:,1]
        if offset is None:
            offset=spacing/2
        # Computing curvilinear length
        s = curve_coord(line)
        offset=np.mod(offset,s[-1]) # making sure we always get one point
        # New (equidistant) curvilinear coordinate
        sExtract=np.arange(offset,s[-1],spacing)
        # Interpolating based on new curvilinear coordinate
        xx=np.interp(sExtract,s,x);
        yy=np.interp(sExtract,s,y);
        return np.array([xx,yy]).T

    def seg_to_lines(seg):
        """ Convert a list of segments to a list of lines """ 
        def extract_continuous(i):
            x=[]
            y=[]
            # Special case, we have only 1 segment remaining:
            if i==len(seg)-1:
                x.append(seg[i][0,0])
                y.append(seg[i][0,1])
                x.append(seg[i][1,0])
                y.append(seg[i][1,1])
                return i,x,y
            # Looping on continuous segment
            while i<len(seg)-1:
                # Adding our start point
                x.append(seg[i][0,0])
                y.append(seg[i][0,1])
                # Checking whether next segment continues our line
                Continuous= all(seg[i][1,:]==seg[i+1][0,:])
                if not Continuous:
                    # We add our end point then
                    x.append(seg[i][1,0])
                    y.append(seg[i][1,1])
                    break
                elif i==len(seg)-2:
                    # we add the last segment
                    x.append(seg[i+1][0,0])
                    y.append(seg[i+1][0,1])
                    x.append(seg[i+1][1,0])
                    y.append(seg[i+1][1,1])
                i=i+1
            return i,x,y
        lines=[]
        i=0
        while i<len(seg):
            iEnd,x,y=extract_continuous(i)
            lines.append(np.array( [x,y] ).T)
            i=iEnd+1
        return lines

    def lines_to_arrows(lines,n=5,spacing=None,normalize=True):
        """ Extract "streamlines" arrows from a set of lines 
        Either: `n` arrows per line
            or an arrow every `spacing` distance
        If `normalize` is true, the arrows have a unit length
        """
        if spacing is None:
            # if n is provided we estimate the spacing based on each curve lenght)
            spacing = [ curve_coord(l)[-1]/n for l in lines]
        try:
            len(spacing)
        except:
            spacing=[spacing]*len(lines)

        lines_s=[curve_extract(l,spacing=sp,offset=sp/2)         for l,sp in zip(lines,spacing)]
        lines_e=[curve_extract(l,spacing=sp,offset=sp/2+0.01*sp) for l,sp in zip(lines,spacing)]
        arrow_x  = [l[i,0] for l in lines_s for i in range(len(l))]
        arrow_y  = [l[i,1] for l in lines_s for i in range(len(l))]
        arrow_dx = [le[i,0]-ls[i,0] for ls,le in zip(lines_s,lines_e) for i in range(len(ls))]
        arrow_dy = [le[i,1]-ls[i,1] for ls,le in zip(lines_s,lines_e) for i in range(len(ls))]

        if normalize:
            dn = [ np.sqrt(ddx**2 + ddy**2) for ddx,ddy in zip(arrow_dx,arrow_dy)]
            arrow_dx = [ddx/ddn for ddx,ddn in zip(arrow_dx,dn)] 
            arrow_dy = [ddy/ddn for ddy,ddn in zip(arrow_dy,dn)] 
        return  arrow_x,arrow_y,arrow_dx,arrow_dy 

    # --- Main body of streamQuiver
    # Extracting lines
    seg   = sp.lines.get_segments() # list of (2, 2) numpy arrays
    lines = seg_to_lines(seg)       # list of (N,2) numpy arrays
    # Convert lines to arrows
    ar_x, ar_y, ar_dx, ar_dy = lines_to_arrows(lines,spacing=spacing,n=n,normalize=True)
    # Plot arrows
    qv=ax.quiver(ar_x, ar_y, ar_dx, ar_dy, *args, angles='xy', **kwargs)
    return qv

# --- Example
x = np.linspace(-1,1,100)
y = np.linspace(-1,1,100)
X,Y=np.meshgrid(x,y)
u = -np.sin(np.arctan2(Y,X))
v =  np.cos(np.arctan2(Y,X))

xseed=np.linspace(0.1,1,4)

fig=plt.figure()
ax=fig.add_subplot(111)
sp = ax.streamplot(x,y,u,v,color='k',arrowstyle='-',start_points=np.array([xseed,xseed*0]).T,density=30)
qv = streamQuiver(ax,sp,spacing=0.5, scale=60)
plt.show()

Answer 2

I have found a way to customize the number of arrowheads on streamline plot.

The idea is to plot streamline and arrows separately:

• plt.streamplot returns a stream_container with two attributes: lines and arrows. The lines contain line segments that can be used to reconstruct streamline without arrows.
• plt.quiver can be used to plot gradient fields. With the proper scaling, the length of the arrows is neglectable, leaving only arrowheads.

Thus, we only need to define the positions of arrows using the line segments and pass them to plt.quiver.

Here is a toy example:

import matplotlib.pyplot as plt
from matplotlib import collections as mc
import numpy as np

# get line segments
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
sp = ax.streamplot(x, y, u, v, start_points=start_points, density=10)
seg = sps.lines.get_segments()  # seg is a list of (2, 2) numpy arrays
lc = mc.LineCollection(seg, ...)

# define arrows
# here I define one arrow every 50 segments
# you could also select segs based on some criterion, e.g. intersect with certain lines
period = 50
arrow_x = np.array([seg[i][0, 0] for i in range(0, len(seg), period)])
arrow_y = np.array([seg[i][0, 1] for i in range(0, len(seg), period)])
arrow_dx = np.array([seg[i][1, 0] - seg[i][0, 0] for i in range(0, len(seg), period)])
arrow_dy = np.array([seg[i][1, 1] - seg[i][0, 1] for i in range(0, len(seg), period)])

# plot the final streamline
fig = plt.figure(figsize=(12.8, 10.8))
ax = fig.add_subplot(1, 1, 1)
ax.add_collection(lc)
ax.autoscale()
ax.quiver(
    arrow_x, arrow_y, arrow_dx, arrow_dy, angles='xy',  # arrow position
    scale=0.2, scale_units='inches', units='y', minshaft=0,  # arrow scaling
    headwidth=6, headlength=10, headaxislength=9)  # arrow style
fig.show()

There is more than one way to scale the arrows so that they appear to have zero length.

Answer 3

I'm not sure about just increasing the number of arrowheads - but you can increase the density of streamlines with the density parameter in the streamplot function, here's the documentation:

*density* : float or 2-tuple
    Controls the closeness of streamlines. When `density = 1`, the domain
    is divided into a 30x30 grid---*density* linearly scales this grid.
    Each cell in the grid can have, at most, one traversing streamline.
    For different densities in each direction, use [density_x, density_y].

Here is an example:

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(0,20,1)
y = np.arange(0,20,1)

u=np.random.random((x.shape[0], y.shape[0]))
v=np.random.random((x.shape[0], y.shape[0]))


fig, ax = plt.subplots(2,2)

ax[0,0].streamplot(x,y,u,v,density=1)
ax[0,0].set_title('Original')

ax[0,1].streamplot(x,y,u,v,density=4)
ax[0,1].set_xlim(5,10)
ax[0,1].set_ylim(5,10)
ax[0,1].set_title('Zoomed, higher density')

ax[1,1].streamplot(x,y,u,v,density=1)
ax[1,1].set_xlim(5,10)
ax[1,1].set_ylim(5,10)
ax[1,1].set_title('Zoomed, same density')

ax[1,0].streamplot(x,y,u,v,density=4)
ax[1,0].set_title('Original, higher density')

fig.show()