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:
- matplotlib streamplot co-author Tom Flannaghan's blog post on the topic
- 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.
- Examples, including commented source code, for using mayavi to produce 2D streamplots of magnetic field lines.
- Similar examples with source in 3D, but still not sufficient to solve my plotting issues:
- The flow() plot resulting from the example code provided in the mayavi standard documentation (with spherical seed widget visible behind the streamlines).