In [1]:
%matplotlib qt
import numpy as np
from mayavi import mlab

from scipy.integrate import odeint


WARNING:traits.has_traits:DEPRECATED: traits.has_traits.wrapped_class, 'the 'implements' class advisor has been deprecated. Use the 'provides' class decorator.

Lorenz Attractor - 3D line and point plotting demo

Lorenz attractor is a 3D differential equation that we will use to demonstrate mayavi's 3D plotting ability. We will look at some ways to make plotting lots of data more efficient.


In [2]:
# setup parameters for Lorenz equations
sigma=10
beta=8/3.
rho=28

def lorenz(x, t, ):
    dx = np.zeros(3)
    dx[0] = -sigma*x[0] + sigma*x[1]
    dx[1] = rho*x[0] - x[1] - x[0]*x[2]
    dx[2] = -beta*x[2] + x[0]*x[1]
    return dx

In [3]:
# solve for a specific particle
# initial condition
y0 = np.ones(3) + .01

# time steps to compute location
n_time = 20000
t = np.linspace(0,200,n_time)

# solve the ODE 
y = odeint( lorenz, y0, t )

y.shape


Out[3]:
(20000, 3)

Rendering Points and Lines

Mayavi has several ways to render 3D line and point data. The default is to use surfaces, which uses more resources. There are kwargs that can be changed to make it render with 2-D lines and points that make plotting large amounts of data more efficient.

LinePlot


In [4]:
# plot the data as a line
# change the tube radius to see the difference
mlab.figure('Line')
mlab.clf()
mlab.plot3d(y[:,0], y[:,1], y[:,2], tube_radius=.1)
mlab.colorbar()


/Users/ahaefner/anaconda/lib/python2.7/site-packages/traits/has_traits.py:1766: FutureWarning: comparison to `None` will result in an elementwise object comparison in the future.
  setattr( self, name, value )

In [5]:
# plot the data as a line, with color representing the time evolution
mlab.figure('Line')
mlab.clf()
mlab.plot3d(y[:,0], y[:,1], y[:,2], t, tube_radius=None, )
mlab.colorbar()


Out[5]:
<mayavi.core.lut_manager.LUTManager at 0x10f7064d0>

Point Plot


In [6]:
# plot the data as a line, with color representing the time evolution
mlab.figure()

# By default, mayavi will plot points as spheres, so each point will 
# be represented by a surface. 
# Using mode='2dvertex' is needed for plotting large numbers of points.
mlab.figure('Points')
mlab.clf()
mlab.points3d(y[:,0], y[:,1], y[:,2], t, mode='2dvertex')
mlab.colorbar( title='time')
mlab.axes()


Out[6]:
<mayavi.modules.axes.Axes at 0x126e09ef0>

Line + Point Plot


In [7]:
# plot the data as a line, with color representing the time evolution
mlab.figure('Line and Points')
mlab.clf()

# plot the data as a line, with color representing the time evolution
mlab.plot3d(y[:,0], y[:,1], y[:,2], t, tube_radius=None, line_width=1 )
mlab.colorbar()

# By default, mayavi will plot points as spheres, so each point will 
# be represented by a surface. 
# Using mode='2dvertex' is needed for plotting large numbers of points.
mlab.points3d(y[:,0], y[:,1], y[:,2], t, scale_factor=.3, scale_mode='none')
              #mode='2dvertex')
mlab.colorbar( title='time')


Out[7]:
<mayavi.core.lut_manager.LUTManager at 0x12c0bc3b0>

Contour Plot

Let's see how long the particle spends in each location


In [8]:
h3d = np.histogramdd(y, bins=50)

# generate the midpoint coordinates
xg,yg,zg = h3d[1]
xm = xg[1:] - .5*(xg[1]-xg[0])
ym = yg[1:] - .5*(yg[1]-yg[0])
zm = zg[1:] - .5*(zg[1]-zg[0])
xg, yg, zg = np.meshgrid(xm, ym, zm)

mlab.figure('contour')
mlab.clf()
mlab.contour3d( h3d[0], opacity=.5, contours=25 )


Out[8]:
<mayavi.modules.iso_surface.IsoSurface at 0x12dc51470>

Animation

Animation can be accomplished with a mlab.animate decorator. You must define a function that yields to the animate decorator. The yield defines when mayavi will rerender the image.


In [11]:
# plot the data as a line
mlab.figure('Animate')
mlab.clf()
# mlab.plot3d(y[:,0], y[:,1], y[:,2], tube_radius=None)
# mlab.colorbar()

a = mlab.points3d(y0[0], y0[1], y0[2], mode='2dvertex')

In [12]:
# number of points to plot
# n_plot = n_time
n_plot = 1000

@mlab.animate(delay=10, ui=True )
def anim():
    for i in range(n_time):
        # a.mlab_source.set(x=y[i,0],y=y[i,1],z=y[i,2], color=(1,0,0))
        mlab.points3d(y[i,0],y[i,1],y[i,2], mode='2dvertex', reset_zoom=False)
        yield
        
anim()


Out[12]:
<mayavi.tools.animator.Animator at 0x1403e6650>

In [ ]: