使用python和mayavi创建3D streamplot

时间:2013-11-01 03:08:50

标签: python matplotlib plot vtk mayavi

目前使用python和matplotlib制作2D streamplot非常容易,因为streamplot最近由Tom Flannaghan和Tony Yu合并到matplotlib。

虽然可以使用matplotlib生成某些类型的3D绘图,但目前不支持3D streamplots。然而,python绘图程序mayavi(提供基于vtk的绘图的python接口)能够使用其flow()函数进行3D streamplots。

我已经创建了一个简单的python模块来对2D和3D中的数据进行streamplot,在3D数据集中没有Z-slope(所有dZ = 0),以展示我与mayavi对比matplotlib所面临的绘图挑战有效地匹配xy平面中的数据。评论的代码和结果图如下所示:

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结果(注意斜率1的流线与dx和dy数组完全匹配,所有dl和dy数组都填充了单位值): 2D matplotlib streamplot 3D mayavi结果(尝试1; 请注意,此处存在的单个流线的斜率是正确的,但流线的长度和数量显然与2D matplotlib示例完全不同): enter image description here 3D mayavi结果(尝试2; 请注意,我使用具有足够高分辨率的线种子产生了更多适当斜率的流线,但也注意到黑色种子线的开始和结束x坐标代码与情节不匹配enter image description here 最后,尝试#3(令人困惑)产生与#2完全相同的图,尽管规定了不同的种子/小部件点。 所以,问题是:如何更好地控制种子线的位置以使其成为对角线?更广泛地说,我希望能够提供任意数量的种子点以获得更一般的流图3D(非平面)问题,但是解决前一个特定情况的线应该让我开始。

我在处理这个问题时发现了一些其他有用的资源,这些资源并没有完全解决我的问题:

  1. matplotlib streamplot共同作者Tom Flannaghan关于该主题的blog post
  2. 一个example,其中一个自定义种子数组与mayavi Streamline的子类一起使用(反过来,由flow()使用),但遗憾的是实现细节不足以重现。
  3. 使用mayavi生成2D streamplots of magnetic field lines的示例,包括注释的源代码。
  4. Similar examples with source in 3D,但仍不足以解决我的密谋问题:
  5. mayavi标准文档中提供的flow()plot resulting from the example code(在流线后面可以看到球形种子窗口小部件)。

1 个答案:

答案 0 :(得分:5)

似乎问题是流线种子窗口小部件的clamp_to_bounds属性的默认值设置为True。您必须将此设置为False才能实际移动窗口小部件。

st.seed.widget.clamp_to_bounds = False

将此内容添加到代码后,最终结果如下所示:

Final result

你可能已经熟悉了以前探索Mayavi的这种方法,但是有可能解释得太多,我会提到这个:

我找到这个选项的方式,以及我通常如何在Mayavi中找到这些不起眼的属性,是通过启用pylab启动IPython的脚本。在Ubuntu上,终端命令如下所示:

ipython --pylab=qt

启动IPython后,我使用%run magic命令运行脚本:

%run streamlines.py

完成绘图后,我现在可以访问Mayavi窗口,可以通过单击Mayavi图标看到Mayavi管道的内部:

enter image description here

在管道中,您将找到Streamline对象,如果单击它然后选择Seed,您将找到“Clamp to bounds”复选框。 (很抱歉屏幕截图中的颜色组合不一 - Mayavi最近开始使用黑色主题......)

enter image description here

要查看其实际功能,您可以点击记录按钮enter image description here打开一个窗口,向您显示与更改设置等效的代码:

enter image description here

在那里你可以看到clamp_to_bounds属性,并且可以将它添加到你的脚本中。

祝你好运!