"""This creates matlab like one liners that make it easy to visualize
data from the Python interpreter.
Right now this provides the following useful functions.
1. surf(x, y, f) -- samples f along x, y and plots a surface.
2. view(arr) -- Views array as structured points.
3. viewi(arr) -- Views array as an image.
4. sampler(xa, ya, func) -- Samples func along array of ordered
points (xa, ya).
Author:
Prabhu Ramachandran <prabhu_r@users.sf.net>
License:
BSD -- http://www.opensource.org/licenses/bsd-license.html
"""
import mayavi
import vtk
import Numeric
try:
from vtk.util import vtkConstants
except ImportError:
class vtkConstants:
pass
vtkConstants.VTK_CHAR=2
vtkConstants.VTK_UNSIGNED_CHAR = 3
vtkConstants.VTK_SHORT = 4
vtkConstants.VTK_UNSIGNED_SHORT = 5
vtkConstants.VTK_INT = 6
vtkConstants.VTK_UNSIGNED_INT = 7
vtkConstants.VTK_LONG = 8
vtkConstants.VTK_UNSIGNED_LONG = 9
vtkConstants.VTK_FLOAT =10
vtkConstants.VTK_DOUBLE =11
def array2vtk(z):
"""Converts a Numeric Array to a VTK array object directly. The
resulting array copies the data in the passed Numeric array. The
array can therefore be deleted safely. This works for real arrays.
XXX what should be done for complex arrays?
"""
arr_vtk = {'c':vtkConstants.VTK_UNSIGNED_CHAR,
'b':vtkConstants.VTK_UNSIGNED_CHAR,
'1':vtkConstants.VTK_CHAR,
's':vtkConstants.VTK_SHORT,
'i':vtkConstants.VTK_INT,
'l':vtkConstants.VTK_LONG,
'f':vtkConstants.VTK_FLOAT,
'd':vtkConstants.VTK_DOUBLE,
'F':vtkConstants.VTK_FLOAT,
'D':vtkConstants.VTK_DOUBLE }
# A dummy array used to create others.
f = vtk.vtkFloatArray()
# First create an array of the right type by using the typecode.
tmp = f.CreateDataArray(arr_vtk[z.typecode()])
tmp.SetReferenceCount(2) # Prevents memory leak.
zf = Numeric.ravel(z)
tmp.SetNumberOfTuples(len(zf))
tmp.SetNumberOfComponents(1)
tmp.SetVoidArray(zf, len(zf), 1)
# Now create a new array that is a DeepCopy of tmp. This is
# required because tmp does not copy the data from the NumPy array
# and will point to garbage if the NumPy array is deleted.
arr = f.CreateDataArray(arr_vtk[z.typecode()])
arr.SetReferenceCount(2) # Prevents memory leak.
arr.DeepCopy(tmp)
return arr
def _create_structured_points_pyvtk(x, y, z):
"""Creates a vtkStructuredPoints object given input data in the
form of arrays. This uses pyvtk to do the job and generates a
temporary file in the process.
Input Arguments:
x -- Array of x-coordinates. These should be regularly spaced.
y -- Array of y-coordinates. These should be regularly spaced.
z -- Array of z values for the x, y values given. The values
should be computed such that the z values are computed as x
varies fastest and y next.
"""
import pyvtk
import tempfile, os
nx = len(x)
ny = len(y)
nz = len(z)
assert nx*ny == nz, "len(x)*len(y) != len(z). "\
"You passed nx=%d, ny=%d, nz=%d"%(nx, ny, nz)
xmin, ymin = x[0], y[0]
dx, dy= (x[1] - x[0]), (y[1] - y[0])
# create a vtk data file
sp = pyvtk.StructuredPoints ((nx, ny, 1), (xmin, ymin, 0), (dx, dy, 1))
pd = pyvtk.PointData(pyvtk.Scalars(z, name='Scalars',
lookup_table="default"))
d = pyvtk.VtkData(sp, pd, "Surf data")
file_name = tempfile.mktemp(suffix='.vtk')
d.tofile(file_name, format='ascii')
# read the created file - yes this is circuitous but works for now.
reader = vtk.vtkStructuredPointsReader()
reader.SetFileName(file_name)
reader.Update()
# cleanup.
os.remove(file_name)
return reader.GetOutput()
def _create_structured_points_direct(x, y, z=None):
"""Creates a vtkStructuredPoints object given input data in the
form of Numeric arrays.
Input Arguments:
x -- Array of x-coordinates. These should be regularly spaced.
y -- Array of y-coordinates. These should be regularly spaced.
z -- Array of z values for the x, y values given. The values
should be computed such that the z values are computed as x
varies fastest and y next. If z is None then no scalars are
associated with the structured points. Only the structured
points data set is created.
"""
nx = len(x)
ny = len(y)
if z:
nz = Numeric.size(z)
assert nx*ny == nz, "len(x)*len(y) != len(z)"\
"You passed nx=%d, ny=%d, nz=%d"%(nx, ny, nz)
xmin, ymin = x[0], y[0]
dx, dy= (x[1] - x[0]), (y[1] - y[0])
sp = vtk.vtkStructuredPoints()
sp.SetDimensions(nx, ny, 1)
sp.SetOrigin(xmin, ymin, 0)
sp.SetSpacing(dx, dy, 1)
if z:
sc = array2vtk(z)
sp.GetPointData().SetScalars(sc)
return sp
def sampler(xa, ya, func,f_args=(),f_keyw=None):
"""Samples a function at an array of ordered points (with equal
spacing) and returns an array of scalars as per VTK's requirements
for a structured points data set, i.e. x varying fastest and y
varying next.
Input Arguments:
xa -- Array of x points.
ya -- Array if y points.
func -- function of x, and y to sample.
f_args -- a tuple of additional positional arguments for func()
(default is empty)
f_keyw -- a dict of additional keyword arguments for func()
(default is empty)
"""
if f_keyw is None:
f_keyw = {}
ret = func(xa[:,Numeric.NewAxis] +
Numeric.zeros(len(ya), ya.typecode()),
Numeric.transpose(ya[:,Numeric.NewAxis] +
Numeric.zeros(len(xa), xa.typecode()) ),
*f_args, **f_keyw
)
return Numeric.transpose(ret)
def _check_sanity(x, y, z):
"""Checks the given arrays to see if they are suitable for
surf."""
msg = "Only ravelled or 2D arrays can be viewed! "\
"This array has shape %s" % str(z.shape)
assert len(z.shape) <= 2, msg
if len( z.shape ) == 2:
msg = "len(x)*len(y) != len(z.flat). You passed "\
"nx=%d, ny=%d, shape of z=%s"%(len(x), len(y), z.shape)
assert z.shape[0]*z.shape[1] == len(x)*len(y), msg
msg = "length of y(%d) and x(%d) must match shape of z "\
"%s. (Maybe you need to swap x and y?)"%(len(y), len(x),
str(z.shape))
assert z.shape == (len(y), len(x)), msg
def squeeze(a):
"Returns a with any ones from the shape of a removed"
a = Numeric.asarray(a)
b = Numeric.asarray(a.shape)
val = Numeric.reshape(a,
tuple(Numeric.compress(Numeric.not_equal(b, 1), b)))
return val
def surf(x, y, z, warp=1, scale=[1.0, 1.0, 1.0], viewer=None,
f_args=(), f_keyw=None):
"""Creates a surface given regularly spaced values of x, y and the
corresponding z as arrays. Also works if z is a function.
Currently works only for regular data - can be enhanced later.
Input Arguments:
x -- Array of x points (regularly spaced)
y -- Array if y points (regularly spaced)
z -- A 2D array for the x and y points with x varying fastest
and y next. Also will work if z is a callable which supports
x and y arrays as the arguments.
warp -- If true, warp the data to show a 3D surface
(default = 1).
scale -- Scale the x, y and z axis as per passed values.
Defaults to [1.0, 1.0, 1.0].
viewer -- An optional viewer (defaults to None). If provided
it will use this viewer instead of creating a new MayaVi
window each time.
f_args -- a tuple of additional positional arguments for func()
(default is empty)
f_keyw -- a dict of additional keyword arguments for func()
(default is empty)
"""
if f_keyw is None:
f_keyw = {}
if callable(z):
zval = Numeric.ravel(sampler(x, y, z, f_args=f_args, f_keyw=f_keyw))
x, y = squeeze(x), squeeze(y)
else:
x, y = squeeze(x), squeeze(y)
_check_sanity(x, y, z)
zval = Numeric.ravel(z)
assert len(zval) > 0, "z is empty - nothing to plot!"
xs = x*scale[0]
ys = y*scale[1]
data = _create_structured_points_direct(xs, ys, zval)
# do the mayavi stuff.
if not viewer:
v = mayavi.mayavi()
else:
v = viewer
v.open_vtk_data(data)
if warp:
f = v.load_filter('WarpScalar', 0)
f.fil.SetScaleFactor(scale[2])
n = v.load_filter('PolyDataNormals', 0)
n.fil.SetFeatureAngle(45)
m = v.load_module('SurfaceMap', 0)
if not viewer:
a = v.load_module('Axes', 0)
a.axes.SetCornerOffset(0.0)
if (min(scale) != max(scale)) or (scale[0] != 1.0):
a.axes.UseRangesOn()
a.axes.SetRanges(x[0], x[-1], y[0], y[-1], min(zval), max(zval))
o = v.load_module('Outline', 0)
v.Render()
return v
def view(arr, smooth=0, warp=0, scale=[1.0, 1.0, 1.0]):
"""Allows one to view a 2D Numeric array. Note that the view will
be set to the way we normally think of matrices with with 0, 0 at
the top left of the screen.
Input Arguments:
arr -- Array to be viewed.
smooth -- If true, view the array as VTK point data i.e. the
matrix entries will be treated as scalar data at the integral
values along the axes and between these points the colours will
be interpolated providing a smooth appearance for the data. If
set to false the data is viewed as cells with no interpolation
and each cell of the matrix is viewed in a separate color with
no interpolation. Note that if smooth is set to true you will
be able to do more fancy things with the visualization (like
contouring the data, warping it to get a 3d surface etc.) than
when smooth is off. If warp is set to true then this option is
set to true. (default=false)
warp -- If true, warp the data to show a 3D surface (default =
0). If set the smooth option has no effect. Note that this is
an expensive operation so use it sparingly for large arrays.
scale -- Scale the x, y and z axis as per passed values.
Defaults to [1.0, 1.0, 1.0].
"""
assert len(arr.shape) == 2, "Only 2D arrays can be viewed!"
ny, nx = arr.shape
if warp:
smooth=1
if not smooth:
nx += 1
ny += 1
dx, dy, junk = Numeric.array(scale)*1.0
xa = Numeric.arange(0, nx*scale[0] - 0.1*dx, dx, 'f')
ya = Numeric.arange(0, ny*scale[1] - 0.1*dy, dy, 'f')
if smooth:
data = _create_structured_points_direct(xa, ya, arr)
else:
data = _create_structured_points_direct(xa, ya)
sc = array2vtk(arr)
data.GetCellData().SetScalars(sc)
# do the mayavi stuff.
v = mayavi.mayavi()
v.open_vtk_data(data)
if warp:
f = v.load_filter('WarpScalar', 0)
f.fil.SetScaleFactor(scale[2])
n = v.load_filter('PolyDataNormals', 0)
n.fil.SetFeatureAngle(45)
m = v.load_module('SurfaceMap', 0)
a = v.load_module('Axes', 0)
a.axes.SetCornerOffset(0.0)
a.axes.YAxisVisibilityOff()
a.axes.UseRangesOn()
arr_flat = Numeric.ravel(arr)
a.axes.SetRanges(0, nx, 0, ny, min(arr_flat), max(arr_flat))
o = v.load_module('Outline', 0)
v.renwin.update_view(0, 0, -1, 0, -1, 0)
v.Render()
return v
def viewi(arr, smooth=0, lut='br', scale=[1.0, 1.0, 1.0]):
"""Allows one to view a 2D Numeric array as an image. This works
best for very large arrays (like 1024x1024 arrays). The
implementation of this function is a bit of a hack and many of
MayaVi's features cannot be used. For instance you cannot change
the lookup table color and expect the color of the image to
change.
Note that the view will be set to the way we normally think of
matrices with with 0, 0 at the top left of the screen.
Input Arguments:
arr -- Array to be viewed.
smooth -- If true, perform interpolation on the image. If
false do not perform interpolation. (default=0)
lut -- Specifies the lookuptable to map the data with. This
should be on of ['br', 'rb', 'bw', 'wb'] 'br' refers to
blue-to-red, 'rb' to red-to-blue, 'bw' to black-to-white and
'wb' to white-black. (default='br')
scale -- Scale the x, y and z axis as per passed values.
Defaults to [1.0, 1.0, 1.0].
"""
valid_lut = {'br': 1, 'rb':2, 'bw':3, 'wb':4}
assert len(arr.shape) == 2, "Only 2D arrays can be viewed!"
assert lut in valid_lut.keys(), \
"lut must be one of %s!"%(valid_lut.keys())
ny, nx = arr.shape
dx, dy, junk = Numeric.array(scale)*1.0
xa = Numeric.arange(0, nx*scale[0] - 0.1*dx, dx, 'f')
ya = Numeric.arange(0, ny*scale[1] - 0.1*dy, dy, 'f')
arr_flat = Numeric.ravel(arr)
min_val = min(arr_flat)
max_val = max(arr_flat)
sp = _create_structured_points_direct(xa, ya)
v = mayavi.mayavi()
v.open_vtk_data(sp)
mm = v.get_current_dvm().get_current_module_mgr()
luth = mm.get_scalar_lut_handler()
luth.lut_var.set(valid_lut[lut])
luth.change_lut()
l = luth.get_lut()
l.SetRange(min_val, max_val)
za = array2vtk(arr_flat)
a = l.MapScalars(za, 0, 0)
sp.GetPointData().SetScalars(a)
sp.SetScalarTypeToUnsignedChar()
sp.SetNumberOfScalarComponents(4)
luth.legend_on.set(1)
luth.legend_on_off()
luth.set_data_range([min_val, max_val])
luth.range_on_var.set(1)
luth.set_range_on()
ia = vtk.vtkImageActor()
if smooth:
ia.InterpolateOn()
else:
ia.InterpolateOff()
ia.SetInput(sp)
v.renwin.add_actors(ia)
a = v.load_module('Axes', 0)
a.axes.SetCornerOffset(0.0)
a.axes.YAxisVisibilityOff()
a.axes.UseRangesOn()
a.axes.SetRanges(0, nx, 0, ny, min_val, max_val)
o = v.load_module('Outline', 0)
v.renwin.update_view(0, 0, -1, 0, -1, 0)
t = v.renwin.tkwidget
t.UpdateRenderer(0,0)
t.Pan(0,-25)
v.Render()
return v
def main():
""" A simple example. Note that the Tkinter lines are there only
because this code will be run standalone. On the interpreter,
simply invoking surf and view would do the job."""
import Tkinter
r = Tkinter.Tk()
r.withdraw()
def f(x, y):
return Numeric.sin(x*y)/(x*y)
x = Numeric.arange(-7., 7.05, 0.1)
y = Numeric.arange(-5., 5.05, 0.05)
v = surf(x, y, f)
import RandomArray
z = RandomArray.random((50, 25))
v1 = view(z)
v2 = view(z, warp=1)
z_large = RandomArray.random((1024, 512))
v3 = viewi(z_large)
# A hack for stopping Python when all windows are closed.
v.master = r
v1.master = r
v2.master = r
#v3.master = r
r.mainloop()
if __name__ == "__main__":
main()
|