def __init__(self, **kw_args):
super(PPCoordSource, self).__init__(**kw_args)
fn=kw_args.pop('file_name', None)
dir=os.path.dirname(self.file_name)
phparams=os.path.join(dir, "phasingparams.py")
engine=kw_args.pop('engine', None)
self.engine.add_trait('dataarrays', Dict)
self.engine.add_trait('coords', Dict)
if self.engine.coords.has_key(phparams):
self.coords=self.engine.coords[phparams]
else:
self.engine.coords[phparams]=cs.CoordSystem()
self.coords=self.engine.coords[phparams]
self.coords.exec_param_file(phparams)
self.coords.on_trait_change(self._coords_changed, 'T')
if fn is not None:
self.file_name= fn
self._open(self.file_name)
dims=self.dataarray.shape
print "init coords update start"
self.coords.UpdateCoordSystem(dims)
print "init coords update end"
self.sg=tvtk.StructuredGrid()
self.im=tvtk.ImageData()
self.component="Real"
self._component_changed('Real') #can change this for a different default
def create_grid(grid, data, period=1):
s = np.shape(data)
nx = grid["nx"]
ny = grid["ny"]
nz = s[2]
print("data: %d,%d,%d grid: %d,%d\n" % (s[0], s[1], s[2], nx, ny))
dims = (nx, nz, ny)
sgrid = tvtk.StructuredGrid(dimensions=dims)
pts = aligned_points(grid, nz, period)
print(np.shape(pts))
sgrid.points = pts
scalar = np.zeros([nx * ny * nz])
start = 0
for y in range(ny):
end = start + nx * nz
scalar[start:end] = (data[:, y, :]).transpose().ravel()
print(y, " = ", np.max(scalar[start:end]))
start = end
sgrid.point_data.scalars = np.ravel(scalar.copy())
sgrid.point_data.scalars.name = "data"
return sgrid
def _mk_structured_grid(self):
""" Creates a StructuredGrid VTK data set using the factory's
attributes.
"""
# FIXME: We need to figure out the dimensions of the data
# here, if any.
sg = tvtk.StructuredGrid(dimensions=self.scalar_data.shape)
sg.points = c_[self.position_x.ravel(),
self.position_y.ravel(),
self.position_z.ravel(), ]
self._vtk_source = sg
self._mayavi_source = VTKDataSource(data=self._vtk_source)
def structured_grid():
# Make the data.
dims = (3, 4, 3)
r = linspace(5, 15, dims[0])
theta = linspace(0, 0.5 * pi, dims[1])
z = linspace(0, 10, dims[2])
pts = generate_annulus(r, theta, z)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
sgrid.points = pts
s = random.random((dims[0] * dims[1] * dims[2]))
sgrid.point_data.scalars = ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
return sgrid
def setUp(self):
datasets = [
tvtk.ImageData(),
tvtk.StructuredPoints(),
tvtk.RectilinearGrid(),
tvtk.StructuredGrid(),
tvtk.PolyData(),
tvtk.UnstructuredGrid(),
]
exts = ['.vti', '.vti', '.vtr', '.vts', '.vtp', '.vtu']
self.datasets = datasets
self.exts = exts
def generateStructuredGrid():
"""Generates Structured Grid"""
dims = (32, 32, 12)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
r = linspace(1, 10, dims[0])
theta = linspace(0, 2 * numpy.pi, dims[1])
z = linspace(0, 5, dims[2])
pts = generate_annulus(r, theta, z)
sgrid.points = pts
s = sqrt(pts[:, 0]**2 + pts[:, 1]**2 + pts[:, 2]**2)
sgrid.point_data.scalars = numpy.ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
return sgrid
def isosurface(self,points,scalars,contours,dimensions,name='val'):
self.ensure_init()
mlab = self.mlab
tvtk = self.tvtk
sg=tvtk.StructuredGrid(dimensions=dimensions,points=points)
sg.point_data.scalars = scalars
sg.point_data.scalars.name = name
d = mlab.pipeline.add_dataset(sg)
iso = mlab.pipeline.iso_surface(d)
if isinstance(contours,int):
iso.contour.number_of_contours = contours
elif isinstance(contours,list):
iso.contour.auto_contours = False
iso.contour.contours = contours
else:
self.error('isosurface contours must be an int or list\n a '+str(type(contours))+' was provided instead')
#end if
return
def test_tvtk_dataset_name(self):
"Can tvtk datasets can be converted to names correctly."
datasets = [
tvtk.ImageData(),
tvtk.StructuredPoints(),
tvtk.RectilinearGrid(),
tvtk.StructuredGrid(),
tvtk.PolyData(),
tvtk.UnstructuredGrid(),
tvtk.Property(), # Not a dataset!
'foo', # Not a TVTK object.
]
expect = [
'image_data', 'image_data', 'rectilinear_grid', 'structured_grid',
'poly_data', 'unstructured_grid', 'none', 'none'
]
result = [pipeline_info.get_tvtk_dataset_name(d) for d in datasets]
self.assertEqual(result, expect)
# -*- coding: utf-8 -*-
from enthought.tvtk.api import tvtk
import numpy as np
def make_points_array(x, y, z):
return np.c_[x.ravel(), y.ravel(), z.ravel()]
z, y, x = np.mgrid[:3.0, :5.0, :4.0]
x *= (4 - z) / 3
y *= (4 - z) / 3
s1 = tvtk.StructuredGrid()
s1.points = make_points_array(x, y, z)
s1.dimensions = x.shape[::-1]
s1.point_data.scalars = np.arange(0, s1.number_of_points)
s1.point_data.scalars.name = 'scalars'
r, theta, z2 = np.mgrid[2:3:3j, -np.pi / 2:np.pi / 2:6j, 0:4:7j]
x2 = np.cos(theta) * r
y2 = np.sin(theta) * r
s2 = tvtk.StructuredGrid(dimensions=x2.shape[::-1])
s2.points = make_points_array(x2, y2, z2)
s2.point_data.scalars = np.arange(0, s2.number_of_points)
s2.point_data.scalars.name = 'scalars'
# Some vectors vectors = empty(z.shape + (3, ), dtype=float) vectors[..., 0] = (4 - y * 2) vectors[..., 1] = (x * 3 - 12) vectors[..., 2] = sin(z * pi) # We reorder the points, scalars and vectors so this is as per VTK's # requirement of x first, y next and z last. pts = pts.transpose(2, 1, 0, 3).copy() pts.shape = pts.size / 3, 3 scalars = scalars.T.copy() vectors = vectors.transpose(2, 1, 0, 3).copy() vectors.shape = vectors.size / 3, 3 # Create the dataset. sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts) sg.point_data.scalars = scalars.ravel() sg.point_data.scalars.name = 'temperature' sg.point_data.vectors = vectors sg.point_data.vectors.name = 'velocity' # Thats it! # Now visualize the data. d = mlab.pipeline.add_dataset(sg) gx = mlab.pipeline.grid_plane(d) gy = mlab.pipeline.grid_plane(d) gy.grid_plane.axis = 'y' gz = mlab.pipeline.grid_plane(d) gz.grid_plane.axis = 'z' iso = mlab.pipeline.iso_surface(d)
def write_vtk(data,
filename,
scalars=None,
vectors={'V': ('u', 'v', 'w')},
tensors={},
coords=('x', 'y', 'z'),
dims=None,
**kwargs):
''' write data in to vtk file
Parameters
----------
data : dict
mapping of variable name to their numpy array
filename : str
the file to write to (can be any recognized vtk extension)
if extension is missing .vts extension is appended
scalars : list
list of arrays to write as scalars (defaults to data.keys())
vectors : dict
mapping of vector name to vector component names to take from data
tensors : dict
mapping of tensor name to tensor component names to take from data
coords : list
the name of coordinate data arrays (default=('x','y','z'))
dims : 3 tuple
the size along the dimensions for (None means x.shape)
**kwargs : extra arguments for the file writer
example file_type=binary/ascii
'''
x = data[coords[0]]
y = data.get(coords[1], zeros_like(x))
z = data.get(coords[2], zeros_like(x))
if dims is None:
dims = array([1, 1, 1])
dims[:x.ndim] = x.shape
else:
dims = array(dims)
sg = tvtk.StructuredGrid(points=c_[x.flat, y.flat, z.flat],
dimensions=array(dims))
pd = tvtk.PointData()
if scalars is None:
scalars = [i for i in data.keys() if i not in coords]
for v in scalars:
pd.scalars = ravel(data[v])
pd.scalars.name = v
sg.point_data.add_array(pd.scalars)
for vec, vec_vars in vectors.items():
u, v, w = [data[i] for i in vec_vars]
pd.vectors = c_[ravel(u), ravel(v), ravel(w)]
pd.vectors.name = vec
sg.point_data.add_array(pd.vectors)
for ten, ten_vars in tensors.items():
vars = [data[i] for i in ten_vars]
tensors = c_[[ravel(i) for i in vars]].T
pd.tensors = tensors
pd.tensors.name = ten
sg.point_data.add_array(pd.tensors)
write_data(sg, filename, **kwargs)
def test_structured(self):
import numpy as np
from numpy import cos, sin, pi
from enthought.tvtk.api import tvtk
from enthought.mayavi import mlab
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 unstructured
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[0])
# 'y' corresponds to varying 'r'
r = np.linspace(1, 10, dims[1])
z = np.linspace(0, 5, dims[2])
pts = generate_annulus(r, theta, z)
# Uncomment the following if you want to add some noise to the data.
#pts += np.random.randn(dims[0]*dims[1]*dims[2], 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'
contour = mlab.pipeline.contour(sgrid)
mlab.pipeline.surface(contour)
return
"""
from enthought.mayavi import mlab
import numpy as np
from enthought.tvtk.api import tvtk
dims = (4, 4, 4)
x, y, z = np.mgrid[0.:dims[0], 0:dims[1], 0:dims[2]]
x = np.reshape(x.T, (-1,))
y = np.reshape(y.T, (-1,))
z = np.reshape(z.T, (-1,))
y += 0.3*np.sin(x)
z += 0.4*np.cos(x)
x += 0.05*y**3
sgrid = tvtk.StructuredGrid(dimensions=(dims[0], dims[1], dims[2]))
sgrid.points = np.c_[x, y, z]
s = np.random.random((dims[0]*dims[1]*dims[2]))
sgrid.point_data.scalars = np.ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
mlab.figure(1, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
mlab.clf()
mlab.pipeline.surface(sgrid, opacity=0.4)
mlab.pipeline.surface(mlab.pipeline.extract_edges(sgrid), color=(0, 0, 0))
mlab.pipeline.glyph(sgrid, mode='cube', scale_factor=0.2, scale_mode='none')
mlab.savefig('volume_grid.jpg')
mlab.show()
import tables
import numpy
from enthought.mayavi import mlab
from enthought.tvtk.api import tvtk
# open file for reading data
fh = tables.openFile("gyro3D20000.h5")
# read grid
cartMesh = fh.root.grid.cartMesh.read()
dims = cartMesh[:,:,:,0].shape
xyz_pts = cartMesh.transpose(2, 1, 0, 3).copy()
xyz_pts.shape = xyz_pts.size/3, 3
# delete extra memory
del cartMesh
# create structured grid object
sg = tvtk.StructuredGrid(dimensions=dims, points=xyz_pts)
# now read in electron density
density_electron = fh.root.density_electron.read()
density_electron = density_electron.T.copy()
sg.point_data.scalars = density_electron.ravel()
sg.point_data.scalars.name = 'Electron density'
# delete extra memory
del density_electron
# viz it
surf = mlab.pipeline.surface(sg, opacity=0.1)
cut_plane = mlab.pipeline.scalar_cut_plane(surf)
mlab.show()
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[0])
# 'y' corresponds to varying 'r'
r = np.linspace(1, 10, dims[1])
z = np.linspace(0, 5, dims[2])
pts = generate_annulus(r, theta, z)
# Uncomment the following if you want to add some noise to the data.
#pts += np.random.randn(dims[0]*dims[1]*dims[2], 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 enthought.mayavi.sources.vtk_data_source import VTKDataSource
from enthought.mayavi.modules.api import Outline, GridPlane