Flow visualisation in python using curved (path-following) vectors

Question

I would like to plot a vector field with curved arrows in python, as can be done in vfplot (see below) or IDL.

You can get close in matplotlib, but using quiver() limits you to straight vectors (see below left) whereas streamplot() doesn't seem to permit meaningful control over arrow length or arrowhead position (see below right), even when changing integration_direction, density, and maxlength.

So, is there a python library that can do this? Or is there a way of getting matplotlib to do it?

Peru

Could you provide some sample data or function to work with? – Romalda

@ThomasKühn I am interested in this question too. Here you can find example data from matplotlib documentation for streamplot. matplotlib.org/gallery/images_contours_and_fields/… – Kong

Have a look at #44526603 & #37513002 – Lomasi

Should anyone be interested, my poorly-written bare-bones matplotlib bolt-on is here: github.com/kieranmrhunt/curved-quivers – Peru

https://mcmap.net/q/749669/-vertical-wind-visualization-in-vector-form/14105784 – Builtup

Answer 1

If you look at the streamplot.py that is included in matplotlib, on lines 196 - 202 (ish, idk if this has changed between versions - I'm on matplotlib 2.1.2) we see the following:

 ... (to line 195)
    # Add arrows half way along each trajectory.
    s = np.cumsum(np.sqrt(np.diff(tx) ** 2 + np.diff(ty) ** 2))
    n = np.searchsorted(s, s[-1] / 2.)
    arrow_tail = (tx[n], ty[n])
    arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2]))
 ... (after line 196)

changing that part to this will do the trick (changing assignment of n):

 ... (to line 195)
    # Add arrows half way along each trajectory.
    s = np.cumsum(np.sqrt(np.diff(tx) ** 2 + np.diff(ty) ** 2))
    n = np.searchsorted(s, s[-1]) ### THIS IS THE EDITED LINE! ###
    arrow_tail = (tx[n], ty[n])
    arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2]))
 ... (after line 196)

If you modify this to put the arrow at the end, then you could generate the arrows more to your liking.

Additionally, from the docs at the top of the function, we see the following:

*linewidth* : numeric or 2d array
        vary linewidth when given a 2d array with the same shape as velocities.

The linewidth can be a numpy.ndarray, and if you can pre-calculate the desired width of your arrows, you'll be able to modify the pencil width while drawing the arrows. It looks like this part has already been done for you.

So, in combination with shortening the arrows maxlength, increasing the density, and adding start_points, as well as tweaking the function to put the arrow at the end instead of the middle, you could get your desired graph.

With these modifications, and the following code, I was able to get a result much closer to what you wanted:

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

w = 3
Y, X = np.mgrid[-w:w:100j, -w:w:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

fig = plt.figure(figsize=(14, 18))
gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2])

grains = 10
tmp = tuple([x]*grains for x in np.linspace(-2, 2, grains))
xs = []
for x in tmp:
    xs += x
ys = tuple(np.linspace(-2, 2, grains))*grains


seed_points = np.array([list(xs), list(ys)])
# Varying color along a streamline
ax1 = fig.add_subplot(gs[0, 1])

strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=np.array(5*np.random.random_sample((100, 100))**2 + 1), cmap='winter', density=10,
                      minlength=0.001, maxlength = 0.07, arrowstyle='fancy',
                      integration_direction='forward', start_points = seed_points.T)
fig.colorbar(strm.lines)
ax1.set_title('Varying Color')

plt.tight_layout()
plt.show()

tl;dr: go copy the source code, and change it to put the arrows at the end of each path, instead of in the middle. Then use your streamplot instead of the matplotlib streamplot.

Edit: I got the linewidths to vary

Remind

Answer 2

Starting with David Culbreth's modification, I rewrote chunks of the streamplot function to achieve the desired behaviour. Slightly too numerous to specify them all here, but it includes a length-normalising method and disables the trajectory-overlap checking. I've appended two comparisons of the new curved quiver function with the original streamplot and quiver.

Peru

This is fantastic! would you be willing to submit your curved quivers code to be added to the matplotlib codebase? This seems like it should be a core graphing functionality. Here's the link to their contrib page – Remind

Hi @Kieran Hunt. I tried your script at github, and noticed that if you zoom in enough, the arrow is not at the very end of the curve, but at ~90% of the curve length. Any idea why and how to fix? – Hydroquinone

Answer 3

Here's a way to obtain the desired output in vanilla pyplot (i.e., without modifying the streamplot function or anything that fancy). For reminder, the goal is to visualize a vector field with curved arrows whose length is proportional to the norm of the vector.

The trick is to:

1. make streamplot with no arrows that is traced backward from a given point (see)
2. plot a quiver from that point. Make the quiver small enough so that only the arrow is visible
3. repeat 1. and 2. in a loop for every seed and scale the length of the streamplot to be proportional to the norm of the vector.

import matplotlib.pyplot as plt
import numpy as np
w = 3
Y, X = np.mgrid[-w:w:8j, -w:w:8j]

U = -Y
V = X
norm = np.sqrt(U**2 + V**2)
norm_flat = norm.flatten()

start_points = np.array([X.flatten(),Y.flatten()]).T

plt.clf()
scale = .2/np.max(norm)

plt.subplot(121)
plt.title('scaling only the length')
for i in range(start_points.shape[0]):
    plt.streamplot(X,Y,U,V, color='k', start_points=np.array([start_points[i,:]]),minlength=.95*norm_flat[i]*scale, maxlength=1.0*norm_flat[i]*scale,
                integration_direction='backward', density=10, arrowsize=0.0)
plt.quiver(X,Y,U/norm, V/norm,scale=30)
plt.axis('square')



plt.subplot(122)
plt.title('scaling length, arrowhead and linewidth')
for i in range(start_points.shape[0]):
    plt.streamplot(X,Y,U,V, color='k', start_points=np.array([start_points[i,:]]),minlength=.95*norm_flat[i]*scale, maxlength=1.0*norm_flat[i]*scale,
                integration_direction='backward', density=10, arrowsize=0.0, linewidth=.5*norm_flat[i])
plt.quiver(X,Y,U/np.max(norm), V/np.max(norm),scale=30)

plt.axis('square')

Here's the result:

Ova

Answer 4

Just looking at the documentation on streamplot(), found here -- what if you used something like streamplot( ... ,minlength = n/2, maxlength = n) where n is the desired length -- you will need to play with those numbers a bit to get your desired graph

you can control for the points using start_points, as shown in the example provided by @JohnKoch

Here's an example of how I controlled the length with streamplot() -- it's pretty much a straight copy/paste/crop from the example from above.

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

w = 3
Y, X = np.mgrid[-w:w:100j, -w:w:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

fig = plt.figure(figsize=(14, 18))
gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2])

grains = 10
tmp = tuple([x]*grains for x in np.linspace(-2, 2, grains))
xs = []
for x in tmp:
    xs += x
ys = tuple(np.linspace(-2, 2, grains))*grains


seed_points = np.array([list(xs), list(ys)])
arrowStyle = pat.ArrowStyle.Fancy()
# Varying color along a streamline
ax1 = fig.add_subplot(gs[0, 1])
strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=1.5, cmap='winter', density=10,
                      minlength=0.001, maxlength = 0.1, arrowstyle='->',
                      integration_direction='forward', start_points = seed_points.T)
fig.colorbar(strm.lines)
ax1.set_title('Varying Color')

plt.tight_layout()
plt.show()

Edit: made it prettier, though still not quite what we were looking for.

Remind

I had tried that, unfortunately (1) you can't force the arrow heads to be at the ends of the streamlines, and (2) the 'length' doesn't correspond to anything meaningful, like vector magnitude - it's just controlled by local streamline density as far as I can tell. – Peru

I wasn't able to accomplish what you are looking for... this looks like a good opportunity for an implementation of a new plot type... if you get this working, you should submit it to be added as part of matplotlib – Remind