Structured grid example¶
An example of how to generate a structured grid dataset using numpy arrays. Also shown is a way to visualize this data with the mayavi2 application.
The script can be run like so:
$ mayavi2 -x structured_grid.py
Alternatively, it can be run as:
$ python structured_grid.py
Python source code:
# Authors: Eric Jones <eric at enthought dot com> # Prabhu Ramachandran <prabhu at aero dot iitb dot ac dot in> # Copyright (c) 2007, Enthought, Inc. # License: BSD style. import numpy as np from numpy import cos, sin, pi from tvtk.api import tvtk from mayavi.scripts import mayavi2 def generate_annulus(r=None, theta=None, z=None): """ Generate points for structured grid for a cylindrical annular volume. This method is useful for generating a structured cylindrical mesh for VTK (and perhaps other tools). Parameters ---------- r : array : The radial values of the grid points. It defaults to linspace(1.0, 2.0, 11). theta : array : The angular values of the x axis for the grid points. It defaults to linspace(0,2*pi,11). z: array : The values along the z axis of the grid points. It defaults to linspace(0,0,1.0, 11). Return ------ points : array Nx3 array of points that make up the volume of the annulus. They are organized in planes starting with the first value of z and with the inside "ring" of the plane as the first set of points. The default point array will be 1331x3. """ # Default values for the annular grid. if r is None: r = np.linspace(1.0, 2.0, 11) if theta is None: theta = np.linspace(0, 2*pi, 11) if z is None: z = np.linspace(0.0, 1.0, 11) # Find the x values and y values for each plane. x_plane = (cos(theta)*r[:,None]).ravel() y_plane = (sin(theta)*r[:,None]).ravel() # Allocate an array for all the points. We'll have len(x_plane) # points on each plane, and we have a plane for each z value, so # we need len(x_plane)*len(z) points. points = np.empty([len(x_plane)*len(z),3]) # Loop through the points for each plane and fill them with the # correct x,y,z values. start = 0 for z_plane in z: end = start + len(x_plane) # slice out a plane of the output points and fill it # with the x,y, and z values for this plane. The x,y # values are the same for every plane. The z value # is set to the current z plane_points = points[start:end] plane_points[:,0] = x_plane plane_points[:,1] = y_plane plane_points[:,2] = z_plane start = end return points # Make the data. dims = (51, 25, 25) # Note here that the 'x' axis corresponds to 'theta' theta = np.linspace(0, 2*np.pi, dims) # 'y' corresponds to varying 'r' r = np.linspace(1, 10, dims) z = np.linspace(0, 5, dims) pts = generate_annulus(r, theta, z) # Uncomment the following if you want to add some noise to the data. #pts += np.random.randn(dims*dims*dims, 3)*0.04 sgrid = tvtk.StructuredGrid(dimensions=dims) sgrid.points = pts s = np.sqrt(pts[:,0]**2 + pts[:,1]**2 + pts[:,2]**2) sgrid.point_data.scalars = np.ravel(s.copy()) sgrid.point_data.scalars.name = 'scalars' # Uncomment the next two lines to save the dataset to a VTK XML file. #w = tvtk.XMLStructuredGridWriter(input=sgrid, file_name='sgrid.vts') #w.write() # View the data. @mayavi2.standalone def view(): from mayavi.sources.vtk_data_source import VTKDataSource from mayavi.modules.api import Outline, GridPlane mayavi.new_scene() src = VTKDataSource(data=sgrid) mayavi.add_source(src) mayavi.add_module(Outline()) g = GridPlane() g.grid_plane.axis = 'x' mayavi.add_module(g) g = GridPlane() g.grid_plane.axis = 'y' mayavi.add_module(g) g = GridPlane() g.grid_plane.axis = 'z' mayavi.add_module(g) if __name__ == '__main__': view()