def _readVtkFile(self):
"""Function to read tensor from a file and store the info in a numpy array.
"""
from enthought.mayavi.sources.vtk_file_reader import VTKFileReader
from enthought.tvtk.api import tvtk
reader = VTKFileReader()
reader.initialize(self.vtkInputFile)
data = reader.outputs[0]
# Get vertex and cell info.
cellVtk = data.get_cells()
numCells = cellVtk.number_of_cells
cellArray = cellVtk.to_array()
self.cells = tvtk.CellArray()
self.cells.set_cells(numCells, cellArray)
self.vertArray = data._get_points().to_array()
self.cellType = data.get_cell_type(0)
(numVerts, self.spaceDim) = self.vertArray.shape
# Get cell fields and extract tensor.
cellData = data._get_cell_data()
numCellDataArrays = cellData._get_number_of_arrays()
tensor = cellData.get_array(self.vtkTensorIndex).to_array()
(self.numTensorPoints, numCols) = tensor.shape
sxx = tensor[:, self.vtkTensorComponentsOrder[0]]
syy = tensor[:, self.vtkTensorComponentsOrder[1]]
szz = tensor[:, self.vtkTensorComponentsOrder[2]]
sxy = tensor[:, self.vtkTensorComponentsOrder[3]]
syz = tensor[:, self.vtkTensorComponentsOrder[4]]
sxz = tensor[:, self.vtkTensorComponentsOrder[5]]
self.tensorSorted = numpy.column_stack((sxx, syy, szz, sxy, syz, sxz))
return
def mixed_type_ug():
"""A slightly more complex example of how to generate an
unstructured grid with different cell types. Returns a created
unstructured grid.
"""
points = array(
[
[0, 0, 0],
[1, 0, 0],
[0, 1, 0],
[0, 0, 1], # tetra
[2, 0, 0],
[3, 0, 0],
[3, 1, 0],
[2, 1, 0],
[2, 0, 1],
[3, 0, 1],
[3, 1, 1],
[2, 1, 1], # Hex
],
'f')
# shift the points so we can show both.
points[:, 1] += 2.0
# The cells
cells = array([
4,
0,
1,
2,
3, # tetra
8,
4,
5,
6,
7,
8,
9,
10,
11 # hex
])
# The offsets for the cells, i.e. the indices where the cells
# start.
offset = array([0, 5])
tetra_type = tvtk.Tetra().cell_type # VTK_TETRA == 10
hex_type = tvtk.Hexahedron().cell_type # VTK_HEXAHEDRON == 12
cell_types = array([tetra_type, hex_type])
# Create the array of cells unambiguously.
cell_array = tvtk.CellArray()
cell_array.set_cells(2, cells)
# Now create the UG.
ug = tvtk.UnstructuredGrid(points=points)
# Now just set the cell types and reuse the ug locations and cells.
ug.set_cells(cell_types, offset, cell_array)
return ug
def unstructured_grid():
points = array(
[
[0, 1.2, 0.6],
[1, 0, 0],
[0, 1, 0],
[1, 1, 1], # tetra
[1, 0, -0.5],
[2, 0, 0],
[2, 1.5, 0],
[0, 1, 0],
[1, 0, 0],
[1.5, -0.2, 1],
[1.6, 1, 1.5],
[1, 1, 1], # Hex
],
'f')
# The cells
cells = array([
4,
0,
1,
2,
3, # tetra
8,
4,
5,
6,
7,
8,
9,
10,
11 # hex
])
# The offsets for the cells, i.e. the indices where the cells
# start.
offset = array([0, 5])
tetra_type = tvtk.Tetra().cell_type # VTK_TETRA == 10
hex_type = tvtk.Hexahedron().cell_type # VTK_HEXAHEDRON == 12
cell_types = array([tetra_type, hex_type])
# Create the array of cells unambiguously.
cell_array = tvtk.CellArray()
cell_array.set_cells(2, cells)
# Now create the UG.
ug = tvtk.UnstructuredGrid(points=points)
# Now just set the cell types and reuse the ug locations and cells.
ug.set_cells(cell_types, offset, cell_array)
scalars = random.random(points.shape[0])
ug.point_data.scalars = scalars
ug.point_data.scalars.name = 'scalars'
return ug
def create_dataset(self):
"""Create a tvtk.UnstructuredGrid dataset from the Mesh instance of the
file source."""
mesh = self.mesh
n_nod, dim = self.n_nod, self.dim
if dim < 3:
nod_zz = nm.zeros((n_nod, 3 - dim), dtype=mesh.coors.dtype)
points = nm.c_[mesh.coors, nod_zz]
else:
points = mesh.coors
dataset = tvtk.UnstructuredGrid(points=points)
cell_types = []
cells = []
offset = [0]
n_cells = [1]
for ig, desc in enumerate(mesh.descs):
conn = mesh.get_conn(desc)
n_cell, n_cv = conn.shape
n_cells.append(n_cell)
aux = nm.empty(n_cell, dtype=nm.int32)
aux.fill(vtk_cell_types[desc])
cell_types.append(aux)
aux = nm.empty((n_cell, n_cv + 1), dtype=nm.int32)
aux[:, 0] = n_cv
aux[:, 1:] = conn
cells.append(aux.ravel())
offset.append(aux.shape[1])
cells = nm.concatenate(cells)
cell_types = nm.concatenate(cell_types)
offset = nm.repeat(offset, n_cells)
offset = nm.cumsum(offset)[:-1]
cell_array = tvtk.CellArray()
cell_array.set_cells(mesh.n_el, cells)
dataset.set_cells(cell_types, offset, cell_array)
return dataset
def _readStress(self, vtkFile):
"""
Function to read stresses from a file and store the info in a numpy array.
"""
from enthought.mayavi.sources.vtk_file_reader import VTKFileReader
from enthought.tvtk.api import tvtk
reader = VTKFileReader()
reader.initialize(vtkFile)
data = reader.outputs[0]
# Get vertex and cell info if it hasn't already been done
if not self.readMesh:
cellVtk = data.get_cells()
self.numVertsPerCell = cellVtk._get_max_cell_size()
self.numCells = cellVtk.number_of_cells
cellArray = cellVtk.to_array()
self.cells = tvtk.CellArray()
self.cells.set_cells(self.numCells, cellArray)
self.vertArray = data._get_points().to_array()
self.cellType = data.get_cell_type(0)
(self.numVerts, self.spaceDim) = self.vertArray.shape
self.readMesh = True
# Get cell fields and extract stresses
cellData = data._get_cell_data()
numCellDataArrays = cellData._get_number_of_arrays()
stress = cellData.get_array(self.stressIndex).to_array()
(self.numStressPoints, numCols) = stress.shape
sxx = stress[:,self.stressComponentsOrder[0]]
syy = stress[:,self.stressComponentsOrder[1]]
szz = stress[:,self.stressComponentsOrder[2]]
sxy = stress[:,self.stressComponentsOrder[3]]
syz = stress[:,self.stressComponentsOrder[4]]
sxz = stress[:,self.stressComponentsOrder[5]]
stressOrdered = numpy.column_stack((sxx, syy, szz, sxy, syz, sxz))
return stressOrdered
def _getVec(self, vtkFile, vecName):
"""
Function to read a vector from a file and store it in a numpy array.
"""
from enthought.mayavi.sources.vtk_file_reader import VTKFileReader
from enthought.tvtk.api import tvtk
reader = VTKFileReader()
reader.initialize(vtkFile)
data = reader.outputs[0]
# Get vertex and cell info if it hasn't already been done
if not self.readMesh:
cellVtk = data.get_cells()
self.numVertsPerCell = cellVtk._get_max_cell_size()
self.numCells = cellVtk.number_of_cells
cellArray = cellVtk.to_array()
self.cells = tvtk.CellArray()
self.cells.set_cells(self.numCells, cellArray)
self.vertArray = data._get_points().to_array()
self.cellType = data.get_cell_type(0)
(self.numVerts, self.spaceDim) = self.vertArray.shape
self.readMesh = True
# Get vertex fields and extract the requested vector.
vertData = data._get_point_data()
numArrays = vertData._get_number_of_arrays()
gotArray = False
for vertDataArray in range(numArrays):
arrayName = vertData.get_array_name(vertDataArray)
if (arrayName == vecName):
vector = vertData.get_array(vertDataArray).to_array()
gotArray = True
if gotArray:
break
else:
raise IOError("Unable to find vector '%s'." % vecName)
return vector
def create_dataset(self):
"""Create a tvtk.UnstructuredGrid dataset from the Mesh instance of the
file source."""
mesh = self.mesh
n_nod, dim = self.n_nod, self.dim
n_el, n_els, n_e_ps = mesh.n_el, mesh.n_els, mesh.n_e_ps
if dim == 2:
nod_zz = nm.zeros((n_nod, 1), dtype=mesh.coors.dtype)
points = nm.c_[mesh.coors, nod_zz]
else:
points = mesh.coors
dataset = tvtk.UnstructuredGrid(points=points)
cell_types = []
cells = []
offset = [0]
for ig, conn in enumerate(mesh.conns):
cell_types += [vtk_cell_types[mesh.descs[ig]]] * n_els[ig]
nn = nm.array([conn.shape[1]] * n_els[ig])
aux = nm.c_[nn[:, None], conn]
cells.extend(aux.ravel())
offset.extend([aux.shape[1]] * n_els[ig])
cells = nm.array(cells)
cell_types = nm.array(cell_types)
offset = nm.cumsum(offset)[:-1]
cell_array = tvtk.CellArray()
cell_array.set_cells(n_el, cells)
dataset.set_cells(cell_types, offset, cell_array)
return dataset
def _diffFiles(self, vtkFile1, vtkFile2, vtkFileOut, dt):
"""
Function to compute field differences between two VTK files, divide the
differences by dt, and output the results to a new VTK file.
"""
from enthought.mayavi.sources.vtk_file_reader import VTKFileReader
from enthought.tvtk.api import tvtk
# Set up input files
reader1 = VTKFileReader()
reader2 = VTKFileReader()
reader1.initialize(vtkFile1)
reader2.initialize(vtkFile2)
data1 = reader1.outputs[0]
data2 = reader2.outputs[0]
# Get vertex and cell info if it hasn't already been done
if not self.readMesh:
cellVtk = data1.get_cells()
self.numVertsPerCell = cellVtk._get_max_cell_size()
self.numCells = cellVtk.number_of_cells
cellArray = cellVtk.to_array()
self.cells = tvtk.CellArray()
self.cells.set_cells(self.numCells, cellArray)
self.vertArray = data1._get_points().to_array()
self.cellType = data1.get_cell_type(0)
(self.numVerts, self.spaceDim) = self.vertArray.shape
self.readMesh = True
# Set up mesh info for VTK file
mesh = tvtk.UnstructuredGrid(points=self.vertArray)
mesh.set_cells(self.cellType, self.cells)
# Get vertex fields and compute differences if the fields exist
vertData1 = data1._get_point_data()
numVertDataArrays = vertData1._get_number_of_arrays()
if numVertDataArrays != 0:
vertData2 = data2._get_point_data()
# This is very kludgy because I haven't yet figured out how to include
# multiple scalar or vector fields, and I also don't know how to put in
# a name for a general array (represented as a field).
numScalarsUsed = 0
numVectorsUsed = 0
for vertDataArray in range(numVertDataArrays):
array1 = vertData1.get_array(vertDataArray).to_array()
(numPoints, numCols) = array1.shape
arrayName = vertData1.get_array_name(vertDataArray) + "/dt"
array2 = vertData2.get_array(vertDataArray).to_array()
arrayOut = (array2 - array1) / dt
# This is wrong if we have a scalar field with 3 components
if (numCols == 3 and numVectorsUsed == 0):
mesh.point_data.vectors = arrayOut
mesh.point_data.vectors.name = arrayName
numVectorsUsed += 1
elif numScalarsUsed == 0:
mesh.point_data.scalars = arrayOut
mesh.point_data.scalars.name = arrayName
numScalarsUsed += 1
# Kludge to add a general array
else:
mesh.point_data.add_array(arrayOut)
# Get cell fields and compute differences if the fields exist
cellData1 = data1._get_cell_data()
numCellDataArrays = cellData1._get_number_of_arrays()
if numCellDataArrays != 0:
cellData2 = data2._get_cell_data()
# This is very kludgy because I haven't yet figured out how to include
# multiple scalar or vector fields, and I also don't know how to put in
# a name for a general array (represented as a field).
numScalarsUsed = 0
numVectorsUsed = 0
for cellDataArray in range(numCellDataArrays):
array1 = cellData1.get_array(cellDataArray).to_array()
(numPoints, numCols) = array1.shape
arrayName = cellData1.get_array_name(cellDataArray) + "/dt"
array2 = cellData2.get_array(cellDataArray).to_array()
arrayOut = (array2 - array1) / dt
# This is wrong if we have a scalar field with 3 components
if (numCols == 3 and numVectorsUsed == 0):
mesh.cell_data.vectors = arrayOut
mesh.cell_data.vectors.name = arrayName
numVectorsUsed += 1
elif numScalarsUsed == 0:
mesh.cell_data.scalars = arrayOut
mesh.cell_data.scalars.name = arrayName
numScalarsUsed += 1
# Kludge to add a general array
else:
mesh.cell_data.add_array(arrayOut)
# Write results to VTK file
#w = tvtk.UnstructuredGridWriter(file_name=vtkFileOut, input=mesh)
w = tvtk.XMLDataSetWriter(file_name=vtkFileOut, input=mesh)
w.write()
return
def mshplot3D(mesh, scale=[1, 1, 1], elm_ids=[], labels=[0, 0, 0]):
''' Plot 3D meshes (and optionally label vertices, elements, and edges)
@param mesh A 3D mesh object
@param elm_ids Numpy array of elements that will be plotted. If left empty,
all elements are plotted.
@param labels (\c bool) List of three elements indicateting what should be
labeled. By default, nothing is labeled
labels[0]=True labels vertices
labels[1]=True labels elements
labels[2]=True labels edges
@param scale (\c int) List indicating the scaling of the vertices -- this
is used to scale the z-direction for a thin mesh. By
default, there is no scaling, i.e. scale = [1, 1, 1]
@see msh.mesh.Mesh3D
@author Matt Ueckermann
'''
if type(elm_ids) is not int:
if len(elm_ids) == 0:
elm_ids = np.arange(len(mesh.elm))
#Figure out how many vertices are in each element type
n_vert_in_type = [len(int_el_pqr(element=element, dim=3)[0]) \
for element in mesh.u_elm_type]
#Now create the TVTK data-structures for the 'cells' and 'offsets'
#cells give [#verts, vert1id, vert2id...,vertnid, #verts ...]
#offsets gives the id's in the cells list where #verts are listed. So above
# it is [0, n+1]
cells = np.array([], dtype=int)
offset = np.array([], dtype=int)
for elm, elm_type in zip(mesh.elm[elm_ids, :], mesh.elm_type[elm_ids]):
n_vt = n_vert_in_type[elm_type]
offset = np.append(offset, len(cells))
cells = np.append(cells, [n_vt] + elm[:n_vt].tolist())
#Also have to create a list of element-types -- to do that we have to
#convert from my numbering system to the TVTK numbering system
if mesh.dim == 3:
type_convert = np.array([tvtk.Tetra().cell_type,\
tvtk.Hexahedron().cell_type, tvtk.Wedge().cell_type])
elif mesh.dim == 2:
type_convert = np.array([tvtk.Triangle().cell_type,\
tvtk.Quad().cell_type])
cell_types = mesh.u_elm_type[mesh.elm_type[elm_ids]]
cell_types = type_convert[cell_types]
#To help visualize the mesh, we color it be the distance from the origin
if mesh.dim == 3:
x, y, z = mesh.vert[:].T
elif mesh.dim == 2:
x, y = mesh.vert[:, :2].T
z = np.zeros_like(x)
dist = np.sqrt(x**2 + y**2 + z**2)
#and we scale the x, y, z, coordinates according the the assigned 'scale'
points = np.column_stack((x * scale[0], y * scale[1], z * scale[2]))
#Now we create the data-structures
cell_array = tvtk.CellArray()
cell_array.set_cells(len(offset), cells)
ug = tvtk.UnstructuredGrid(points=points)
ug.set_cells(cell_types, offset, cell_array)
ug.point_data.scalars = dist.ravel()
ug.point_data.scalars.name = 'Distance from Origin'
#Next we set the new data-structures as a mayavi source
src = sources.vtk_data_source.VTKDataSource(data=ug)
#Extract the edges from the grid
edges = mlab.pipeline.extract_edges(src)
#Use a shorter name for the elements
elm = mesh.elm[elm_ids, :]
if any(labels) and len(elm) > 20:
string = "WARNING:\nAre you sure you want to label more than 20" + \
"elements in 3D? -- Labels are very memory" + \
"(apparently -- who would have guessed?) \nY(es) to proceed:"
answer = raw_input(string)
if answer.lower() not in ['yes', 'y', 'yes']:
labels = [0, 0, 0]
#Next add labels if desired:
maxdist = np.max(dist) / 2.
if labels[0]: #Vertex labels
for i in xrange(len(offset)):
for vnum in elm[i, :]:
if vnum >= 0:
mlab.text3d(points[vnum, 0], points[vnum, 1], \
points[vnum, 2], '%d' % vnum, \
color=(0, 0, 0), scale=0.02*maxdist)
if labels[1]: #element labels
for i in xrange(len(offset)):
if elm_ids[i] < 0:
elm_label = mesh.numelm + elm_ids[i]
else:
elm_label = elm_ids[i]
x = np.mean(points[elm[i, :cells[offset[i]]], 0])
y = np.mean(points[elm[i, :cells[offset[i]]], 1])
z = np.mean(points[elm[i, :cells[offset[i]]], 2])
mlab.text3d(x, y, z, '%d' % elm_label, color=(0.3, 0.3, 0.9), \
scale=0.03*maxdist)
if labels[2]: #edge labels
if len(elm_ids) < len(mesh.elm):
elm2ed = connect_elm2ed(mesh.elm2elm[:], mesh.ed2ed)
ed_ids = np.unique(elm2ed[elm_ids])
ed_ids = ed_ids[ed_ids >= 0]
ed2ed = mesh.ed2ed[ed_ids, :]
else:
ed_ids = np.arange(len(mesh.ed2ed))
ed2ed = mesh.ed2ed[:]
for i in xrange(len(ed2ed)):
n = 8
if ed2ed[i, -1] < 0:
n = 7
x = np.mean(points[ed2ed[i, 4:n], 0])
y = np.mean(points[ed2ed[i, 4:n], 1])
z = np.mean(points[ed2ed[i, 4:n], 2])
mlab.text3d(x, y, z, '%d' % ed_ids[i], color=(0.3, 0.9, 0.3), \
scale=0.03*maxdist)
#And plot them!
mlab.pipeline.surface(edges, opacity=0.4, line_width=2)
mlab.axes()
mlab.title('3D mesh', size=0.5, height=0.95)
#mlab.show()
return cells, offset, cell_types
# -*- coding: utf-8 -*- from enthought.tvtk.api import tvtk import numpy as np p1 = tvtk.PolyData() p1.points = [(1, 1, 0), (1, -1, 0), (-1, -1, 0), (-1, 1, 0), (0, 0, 2)] faces = [4, 0, 1, 2, 3, 3, 4, 0, 1, 3, 4, 1, 2, 3, 4, 2, 3, 3, 4, 3, 0] cells = tvtk.CellArray() cells.set_cells(5, faces) p1.polys = cells p1.point_data.scalars = np.linspace(0.0, 1.0, len(p1.points)) N = 10 a, b = np.mgrid[0:np.pi:N * 1j, 0:np.pi:N * 1j] x = np.sin(a) * np.cos(b) y = np.sin(a) * np.sin(b) z = np.cos(a) from tvtk_structuredgrid import make_points_array points = make_points_array(x, y, z) faces = np.zeros(((N - 1)**2, 4), np.int) t1, t2 = np.mgrid[:(N - 1) * N:N, :N - 1] faces[:, 0] = (t1 + t2).ravel() faces[:, 1] = faces[:, 0] + 1 faces[:, 2] = faces[:, 1] + N faces[:, 3] = faces[:, 0] + N p2 = tvtk.PolyData(points=points, polys=faces) p2.point_data.scalars = np.linspace(0.0, 1.0, len(p2.points))