Using python and mayavi to create a 3D streamplot

Question

It is currently quite easy to produce a 2D streamplot with python and matplotlib because streamplot was recently incorporated to matplotlib by Tom Flannaghan and Tony Yu.

Although it is possible to produce certain types of 3D plots with matplotlib, 3D streamplots are not currently supported. However, the python plotting program mayavi (which provides a python interface to vtk-based plotting) is capable of 3D streamplots using its flow() function.

I've created a simple python module to streamplot data in 2D and 3D, with no Z-slope (all dZ = 0) in the 3D data set, to demonstrate the plotting challenge I'm facing with mayavi versus matplotlib on effectively matching data in the xy plane. The commented code and resulting plots are shown below:

import numpy, matplotlib, mayavi, matplotlib.pyplot, mayavi.mlab
#for now, let's produce artificial streamline data sets for 2D & 3D cases where x and y change by +1 at each point, while Z changes by 0 at each point:

#2D data first:
x = numpy.ones((10,10))
y = numpy.ones((10,10))
#and a corresponding meshgrid:
Y,X = numpy.mgrid[-10:10:10j,-10:10:10j]

#now 3D data with Z = 0 (i.e., I want to be able to produce a matching streamplot plane in mayavi, which is normally used for 3D):
xx = numpy.ones((10,10,10))
yy = numpy.ones((10,10,10))
zz = numpy.zeros((10,10,10))
#and a corresponding meshgrid:
ZZ,YY,XX = numpy.mgrid[-10:10:10j,-10:10:10j,-10:10:10j]

#plot the 2D streamplot data with matplotlib:
fig = matplotlib.pyplot.figure()
ax = fig.add_subplot(111,aspect='equal')
speed = numpy.sqrt(x*x + y*y)
ax.streamplot(X, Y, x, y, color=x, linewidth=2, cmap=matplotlib.pyplot.cm.autumn,arrowsize=3)
fig.savefig('test_streamplot_2D.png',dpi=300)

#there's no streamplot 3D available in matplotlib, so try to see how mayavi behaves with a similar 3D data set:
fig = mayavi.mlab.figure(bgcolor=(1.0,1.0,1.0),size=(800,800),fgcolor=(0, 0, 0))
st = mayavi.mlab.flow(XX,YY,ZZ,xx,yy,zz,line_width=4,seedtype='sphere',integration_direction='forward') #sphere is the default seed type
mayavi.mlab.axes(extent = [-10.0,10.0,-10.0,10.0,-1.0,1.0]) #set plot bounds
fig.scene.z_plus_view() #adjust the view for a perspective along z (xy plane flat)
mayavi.mlab.savefig('test_streamplot_3D_attempt_1.png')

#now start to modify the mayavi code to see if I can 'seed in' more streamlines (default only produces a single short streamline, albeit of the correct slope)
#attempt 2 uses a line seed / widget over the specified bounds (points 1 and 2):
fig = mayavi.mlab.figure(bgcolor=(1.0,1.0,1.0),size=(800,800),fgcolor=(0, 0, 0))
st = mayavi.mlab.flow(XX,YY,ZZ,xx,yy,zz,line_width=4,seedtype='line',integration_direction='forward') #line instead of sphere
st.seed.widget.point1 = [0,-10,0]
st.seed.widget.point2 = [0,10,0] #so seed line should go up along y axis
st.seed.widget.resolution = 25 #seems to be the number of seeds points along the seed line
mayavi.mlab.axes(extent = [-10.0,10.0,-10.0,10.0,-1.0,1.0]) #set plot bounds
fig.scene.z_plus_view() #adjust the view for a perspective along z (xy plane flat)
mayavi.mlab.savefig('test_streamplot_3D_attempt_2.png')

#attempt 3 will try to seed a diagonal line across the plot to produce streamlines that cover the full plot:
#would need to use 'both' for integration_direction if I could get the diagonal seed line to work
fig = mayavi.mlab.figure(bgcolor=(1.0,1.0,1.0),size=(800,800),fgcolor=(0, 0, 0)) 
st = mayavi.mlab.flow(XX,YY,ZZ,xx,yy,zz,line_width=4,seedtype='line',integration_direction='forward') 
st.seed.widget.point1 = [-10,10,0] #start seed line at top left corner of plot
st.seed.widget.point2 = [10,-10,0] #end seed line at bottom right corner of plot 
#this fails to produce a diagonal seed line though
st.seed.widget.resolution = 25 #seems to be the number of seeds points along the seed line
mayavi.mlab.axes(extent = [-10.0,10.0,-10.0,10.0,-1.0,1.0]) #set plot bounds
fig.scene.z_plus_view() #adjust the view for a perspective along z (xy plane flat)
mayavi.mlab.savefig('test_streamplot_3D_attempt_3.png')

2D matplotlib result (note the streamlines of slope 1 match sensibly with the dx and dy arrays all filled with values of unity): 3D mayavi result (attempt 1; note that the slope of the single streamline present here is correct, but the length and number of streamlines is obviously quite different from the 2D matplotlib example): 3D mayavi result (attempt 2; note that my use of a line seed with a sufficiently high resolution has produced many more streamlines of appropriate slope, but also note that the start & end x coordinates of the black seed line specified in the code don't match with the plot) Finally, attempt #3 (confusingly) produces the exact same plot as #2 despite the specification of different seed / widget points. So, the question is: how can I get greater control over the position of my seed line to make it diagonal? More broadly, I'd like to be able to feed in an arbitrary array of seed points for more general stream plot 3D (non-planar) problems, but solving the former specific case with the line should get me started.

Some other helpful resources I found while working on this problem, which did not quite solve my issues:

1. matplotlib streamplot co-author Tom Flannaghan's blog post on the topic
2. An example where a custom array of seeds is used with a subclass of mayavi Streamline (which, in turn, is used by flow()), but unfortunately with insufficient implementation detail to reproduce.
3. Examples, including commented source code, for using mayavi to produce 2D streamplots of magnetic field lines.
4. Similar examples with source in 3D, but still not sufficient to solve my plotting issues:
5. The flow() plot resulting from the example code provided in the mayavi standard documentation (with spherical seed widget visible behind the streamlines).

Answer 1

It seems the problem is that the default value for the streamline seed widget's clamp_to_bounds property is set to True. You have to set this to False to actually be able to move the widget.

st.seed.widget.clamp_to_bounds = False

After adding this to the code, the final result looks like this:

You may be familiar with this method of exploring Mayavi from before, but with the possibility of explaining too much, I'll mention this anyways:

The way I found this option, and how I usually find such obscure properties in Mayavi, is by launching the scripts for IPython with pylab enabled. On Ubuntu, the terminal command looks like this:

ipython --pylab=qt

After launching IPython, I run the script using the %run magic command:

%run streamlines.py

After the plot has been made, I now have access to the Mayavi window and may see the internals of the Mayavi pipeline by clicking the Mayavi icon:

In the pipeline you'll find the Streamline object, and if you click it and then select Seed, you'll find the "Clamp to bounds" checkbox. (Sorry about the bad mix of colors in the screenshots - Mayavi has for some reason started to use a dark theme lately...)

To see what this actually does, you may hit the record button to get open a window that shows you the code equivalent of you changing the settings:

There you see the clamp_to_bounds property, and will be able to add it to your script.

Good luck with the streamlines!