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 generate3dCells(self, points):
boundaryz = np.max(self.z)
boundaryx = np.max(self.x)
boundaryy = np.max(self.y)
x = self.x.reshape(self.x.shape[0], 1)
y = self.y.reshape(self.y.shape[0], 1)
z = self.z.reshape(self.z.shape[0], 1)
a = np.max(self.z) + 1
b = (1 + np.max(self.z)) * (np.max(self.y) + 1)
count = self.x.shape[0]
k = np.arange(count)
i2 = np.asarray([
i for i in k if boundaryz not in z[i] and boundaryy not in y[i]
and boundaryx not in x[i]
])
i1 = np.ones(np.size(i2)) * 8
i3 = np.asarray([i + b for i in i2])
i4 = np.asarray([i + a + b for i in i2])
i5 = np.asarray([i + a for i in i2])
i6 = np.asarray([i + 1 for i in i2])
i7 = np.asarray([i + b + 1 for i in i2])
i8 = np.asarray([i + a + b + 1 for i in i2])
i9 = np.asarray([i + a + 1 for i in i2])
cell_data = self.image
cell_data = cell_data.flatten()
index = np.where(cell_data > 0)
i1 = np.delete(i1, index)
i2 = np.delete(i2, index)
i3 = np.delete(i3, index)
i4 = np.delete(i4, index)
i5 = np.delete(i5, index)
i6 = np.delete(i6, index)
i7 = np.delete(i7, index)
i8 = np.delete(i8, index)
i9 = np.delete(i9, index)
cell_data = np.delete(cell_data, index)
cellcount = i1.shape[0]
cells = np.array(np.dstack([i1, i2, i3, i4, i5, i6, i7, i8, i9]))
cells = cells.reshape(i1.shape[0], 9)
cells = cells.flatten()
hex_type = tvtk.Hexahedron().cell_type # VTK_HEXAHEDRON == 12
cell_types = np.ones(cellcount) * hex_type
offset = np.arange(0, cells.shape[0], 9)
return [cell_types, offset, cellcount, cells, cell_data]
def generate_unstrgrid_mesh(self, filter=0):
points = global_variable.NODE_COORDINATES
self.index = where(self.density >= (1.0 - filter))[0].tolist()
cells = (global_variable.ELEMENT_ATTRIBUTES[self.index, :] - 1)
if global_variable.GRID_TYPE == 'Polygon':
cell_tpye = tvtk.Polygon().cell_type
if global_variable.GRID_TYPE == 'Hexahedron':
cell_tpye = tvtk.Hexahedron().cell_type
grid = tvtk.UnstructuredGrid(points=points)
grid.set_cells(cell_tpye, cells)
return grid
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 write_vtk_mesh(mesh, fileName):
node = mesh.entity('node')
GD = mesh.geo_dimension()
if GD == 2:
node = np.concatenate(
(node, np.zeros((node.shape[0], 1), dtype=mesh.ftype)), axis=1)
ug = tvtk.UnstructuredGrid(points=node)
if mesh.meshtype is 'hex':
cell_type = tvtk.Hexahedron().cell_type
cell = mesh.ds.cell
elif mesh.meshtype is 'tri':
cell_type = tvtk.Triangle().cell_type
cell = mesh.ds.cell
for key, value in mesh.cellData.items():
i = ug.cell_data.add_array(value)
ug.cell_data.get_array(i).name = key
for key, value in mesh.pointData.items():
i = ug.point_data.add_array(value)
ug.point_data.get_array(i).name = key
elif mesh.meshtype is 'polyhedron':
cell_type = tvtk.Polygon().cell_type
NF, faces = mesh.to_vtk()
cell = tvtk.CellArray()
cell.set_cells(NF, faces)
elif mesh.meshtype is 'polygon':
cell_type = tvtk.Polygon().cell_type
NC, cells = mesh.to_vtk()
cell = tvtk.CellArray()
cell.set_cells(NC, cells)
elif mesh.meshtype is 'tet':
cell_type = tvtk.Tetra().cell_type
cell = mesh.ds.cell
for key, value in mesh.cellData.items():
i = ug.cell_data.add_array(value)
ug.cell_data.get_array(i).name = key
for key, value in mesh.pointData.items():
i = ug.point_data.add_array(value)
ug.point_data.get_array(i).name = key
ug.set_cells(cell_type, cell)
write_data(ug, fileName)
def write_vtk_mesh(mesh, fileName):
point = mesh.point
if point.shape[1] == 2:
point = np.concatenate(
(point, np.zeros((point.shape[0], 1), dtype=np.float)), axis=1)
ug = tvtk.UnstructuredGrid(points=point)
if mesh.meshtype is 'hex':
cell_type = tvtk.Hexahedron().cell_type
cell = mesh.ds.cell
elif mesh.meshtype is 'tri':
cell_type = tvtk.Triangle().cell_type
cell = mesh.ds.cell
for key, value in mesh.cellData.items():
i = ug.cell_data.add_array(value)
ug.cell_data.get_array(i).name = key
for key, value in mesh.pointData.items():
i = ug.point_data.add_array(value)
ug.point_data.get_array(i).name = key
elif mesh.meshtype is 'polyhedron':
cell_type = tvtk.Polygon().cell_type
NF, faces = mesh.to_vtk()
cell = tvtk.CellArray()
cell.set_cells(NF, faces)
ug.cell_data.scalars = mesh.cellData['flag']
elif mesh.meshtype is 'polygon':
cell_type = tvtk.Polygon().cell_type
NC, cells = mesh.to_vtk()
cell = tvtk.CellArray()
cell.set_cells(NC, cells)
elif mesh.meshtype is 'tet':
cell_type = tvtk.Tetra().cell_type
cell = mesh.ds.cell
for key, value in mesh.cellData.items():
i = ug.cell_data.add_array(value)
ug.cell_data.get_array(i).name = key
for key, value in mesh.pointData.items():
i = ug.point_data.add_array(value)
ug.point_data.get_array(i).name = key
ug.set_cells(cell_type, cell)
write_data(ug, fileName)
def _readData(self):
from tvtk.api import tvtk
import h5py
import numpy
filename = "output/%s_%s.h5" % (cell, resolution)
h5 = h5py.File(filename, 'r', driver="sec2")
elastThick = 40.0e+3
eqslip = 4.0
cells = h5['topology/cells'][:]
(ncells, ncorners) = cells.shape
vertices = h5['geometry/vertices'][:] / elastThick
(nvertices, spaceDim) = vertices.shape
disp = h5['vertex_fields/displacement'][tindex, :, :] / eqslip
h5.close()
if cell == "tet4":
assert (spaceDim == 3)
assert (ncorners == 4)
cellType = tvtk.Tetra().cell_type
elif cell == "hex8":
assert (spaceDim == 3)
assert (ncorners == 8)
cellType = tvtk.Hexahedron().cell_type
else:
raise ValueError("Unknown cell '%s'." % cell)
if (resolution != "6.7km" and resolution != "20km"):
raise ValueError("Unavailable resolution '%s'." % resolution)
data = tvtk.UnstructuredGrid()
data.points = vertices
data.set_cells(cellType, cells)
data.point_data.vectors = disp
data.point_data.vectors.name = "Displacement [m]"
return data
PropID = np.zeros(mc.nel) # Property ID for nod in mc.nodes: if len(nod.loc)==2: xyz[nod.localID]= np.array([nod.loc[0]+nod.Dof(dc)[0],nod.loc[1]+nod.Dof(dc)[1],0.]) uvw[nod.localID]= np.array([nod.Dof(dc)[0],nod.Dof(dc)[1],0.]) elif len(nod.loc)==3: xyz[nod.localID]= np.array([nod.loc[0]+nod.Dof(dc)[0],nod.loc[1]+nod.Dof(dc)[1],nod.loc[2]+nod.Dof(dc)[2]]) uvw[nod.localID]= np.array([nod.Dof(dc)[0],nod.Dof(dc)[1],nod.Dof(dc)[2]]) else: raise "node.loc has wrong length for plotting" ug = tvtk.UnstructuredGrid(points=xyz) line_type = tvtk.Line().cell_type quad_type = tvtk.Quad().cell_type hexa_type = tvtk.Hexahedron().cell_type count = 0; for el in mc.elements: if(el.type=='CBAR' or el.type=='CBAR1' or el.type=='CBARX' or el.type=='CBEAMX'): el_type = line_type Jg[count]= 1. elif(el.type=='CQUAD'): el_type = quad_type Jg[count]=np.mean([eh.getJg(el,ip,dc,2) for ip in range(4)]) elif(el.type=='CHEXA'): el_type = hexa_type Jg[count]=np.mean([eh.getJg(el,ip,dc,3) for ip in range(8)])
sgrid.point_data.scalars = np.ravel(tfield.transpose())
sgrid.point_data.scalars.name = filter_results[0].full_name
writer = tvtk.XMLStructuredGridWriter(input=sgrid,
file_name=os.path.join(
run_dir, "vtk", "test.vts"))
writer.write()
# cell-based mesh (use unstructured mesh, polydata is for surfaces and not
# volumes)
sdims, spts, vertices = smesh_cells(gchemgrid.e_lon_4x5, gchemgrid.e_lat_4x5,
gchemgrid.e_km_geos5_r * 1e3, lev_scale,
lev_offset, True)
sgrid = tvtk.UnstructuredGrid(points=spts)
tet_type = tvtk.Hexahedron().cell_type # cells = hexahedrons
sgrid.set_cells(tet_type, vertices)
ijavg = ctm_f.filter(category=field_category)
for tracer in ijavg:
print "export diagnostic %s to VTK" % tracer.full_name
vals_array = tracer.values * units_conv[mr_units]
#ar = sgrid.cell_data.scalars = np.ravel(vals_array.transpose())
#sgrid.cell_data.scalars.name = tracer.full_name
ar = sgrid.cell_data.add_array(np.ravel(vals_array.transpose()))
sgrid.cell_data.get_array(ar).name = tracer.full_name
outfile = os.path.join(run_dir, "vtk", "test_ijavg_cell.vtu")
if os.path.exists(outfile):
os.remove(outfile)
writer = tvtk.XMLUnstructuredGridWriter(input=sgrid, file_name=outfile)
def exa2vtk(field, downsample):
"""A function for converting exadata to vtk data
Parameters
----------
field : str
a dataset in nekdata format
downsample : bool
flag T/F
"""
#
if (downsample):
ixs = field.lr1[0] - 1
iys = field.lr1[1] - 1
izs = max(field.lr1[2] - 1, 1)
else:
ixs = 1
iys = 1
izs = 1
#
iix = range(0, field.lr1[0], ixs)
nix = len(iix)
iiy = range(0, field.lr1[1], iys)
niy = len(iiy)
iiz = range(0, field.lr1[2], izs)
niz = len(iiz)
#
nppel = nix * niy * niz
nepel = (nix - 1) * (niy - 1) * max((niz - 1), 1)
nel = field.nel * nepel
#
if (field.ndim == 3):
nvert = 8
cellType = tvtk.Hexahedron().cell_type
else:
nvert = 4
cellType = tvtk.Quad().cell_type
#
ct = np.array(nel * [cellType])
of = np.arange(0, nvert * nel, nvert)
ce = np.zeros(nel * (nvert + 1))
ce[range(0, nel * (nvert + 1), nvert + 1)] = nvert
if (field.var[0] != 0):
r = np.zeros((nvert * nel, 3))
if (field.var[1] != 0):
v = np.zeros((nvert * nel, 3))
if (field.var[2] == 1):
p = np.zeros(nvert * nel)
if (field.var[3] == 1):
T = np.zeros(nvert * nel)
if (field.var[4] != 0):
S = np.zeros((nvert * nel, field.var[4]))
#
ice = -(nvert + 1)
for iel in range(field.nel):
for iz in range(niz):
for iy in range(niy):
for ix in range(nix):
if (field.var[0] == 3):
r[iel * nppel + ix + iy * nix +
iz * nix * niy, :] = field.elem[iel].pos[:, iiz[iz],
iiy[iy],
iix[ix]]
if (field.var[1] == 3):
v[iel * nppel + ix + iy * nix +
iz * nix * niy, :] = field.elem[iel].vel[:, iiz[iz],
iiy[iy],
iix[ix]]
if (field.var[2] == 1):
p[iel * nppel + ix + iy * nix +
iz * nix * niy] = field.elem[iel].pres[:, iiz[iz],
iiy[iy],
iix[ix]]
if (field.var[3] == 1):
T[iel * nppel + ix + iy * nix +
iz * nix * niy] = field.elem[iel].temp[:, iiz[iz],
iiy[iy],
iix[ix]]
if (field.var[4] != 0):
S[iel * nppel + ix + iy * nix +
iz * nix * niy, :] = field.elem[iel].scal[:, iiz[iz],
iiy[iy],
iix[ix]]
if (field.var[0] == 3):
for iz in max(range(niz - 1), [0]):
for iy in range(niy - 1):
for ix in range(nix - 1):
ice = ice + nvert + 1
for face in range(field.ndim - 1):
ce[ice + face * 4 +
1] = iel * nppel + ix + iy * nix + (
iz + face) * nix * niy
ce[ice + face * 4 +
2] = iel * nppel + ix + 1 + iy * nix + (
iz + face) * nix * niy
ce[ice + face * 4 + 3] = iel * nppel + ix + 1 + (
iy + 1) * nix + (iz + face) * nix * niy
ce[ice + face * 4 + 4] = iel * nppel + ix + (
iy + 1) * nix + (iz + face) * nix * niy
# create the array of cells
ca = tvtk.CellArray()
ca.set_cells(nel, ce)
# create the unstructured dataset
dataset = tvtk.UnstructuredGrid(points=r)
# set the cell types
dataset.set_cells(ct, of, ca)
# set the data
dataset.point_data.vectors = v
dataset.point_data.vectors.name = 'vel'
if (field.var[2] == 1):
dataset.point_data.scalars = p
dataset.point_data.scalars.name = 'pres'
if (field.var[3] == 1):
dataset.point_data.add_array(T)
dataset.point_data.get_array(2).name = 'temp'
if (field.var[4] != 0):
for ii in range(field.var[4]):
dataset.point_data.add_array(S[:, ii])
dataset.point_data.get_array(ii + 3).name = 'scal_%d' % (ii + 1)
#
dataset.point_data.update()
#
return dataset
def exa2vtk(field, downsample=False):
"""A function for converting exadata to `Traited VTK`_ dataset. The
returned dataset can be manipulated with libraries which accept a VTK
object, for example Mayavi_.
.. _Traited VTK: https://docs.enthought.com/mayavi/tvtk/README.html
Example
-------
This also requires you to have a GUI toolkit installed: either PyQt4,
PySide, PySide2, PyQt5 or wxPython.
.. code-block:: python
import pymech as pm
from pymech.vtksuite import exa2vtk
from mayavi import mlab
field = pm.readnek("tests/nek/channel3D_0.f00001")
dataset = exa2vtk(field)
mlab.pipeline.add_dataset(dataset)
Instead of MayaVi_ you could use also use something high-level like PyVista_
to wrap the underlying VTK object and later visualize them.
.. code-block:: python
import pyvista as pv
dataset = pv.wrap(dataset._vtk_obj)
dataset.plot()
.. _MayaVi: https://docs.enthought.com/mayavi/mayavi/mlab.html
.. _PyVista: https://docs.pyvista.org/getting-started/simple.html
Parameters
----------
field : exadata
a dataset in nekdata format
downsample : bool
flag T/F
Returns
-------
dataset : tvtk.tvtk_classes.unstructured_grid.UnstructuredGrid
a VTK dataset
"""
#
if downsample:
ixs = field.lr1[0] - 1
iys = field.lr1[1] - 1
izs = max(field.lr1[2] - 1, 1)
else:
ixs = 1
iys = 1
izs = 1
#
iix = range(0, field.lr1[0], ixs)
nix = len(iix)
iiy = range(0, field.lr1[1], iys)
niy = len(iiy)
iiz = range(0, field.lr1[2], izs)
niz = len(iiz)
#
nppel = nix * niy * niz
nepel = (nix - 1) * (niy - 1) * max((niz - 1), 1)
nel = field.nel * nepel
#
if field.ndim == 3:
nvert = 8
cellType = tvtk.Hexahedron().cell_type
else:
nvert = 4
cellType = tvtk.Quad().cell_type
#
ct = np.array(nel * [cellType])
of = np.arange(0, nvert * nel, nvert)
ce = np.zeros(nel * (nvert + 1))
ce[np.arange(nel) * (nvert + 1)] = nvert
if field.var[0] != 0:
r = np.zeros((nvert * nel, 3))
if field.var[1] != 0:
v = np.zeros((nvert * nel, 3))
if field.var[2] == 1:
p = np.zeros(nvert * nel)
if field.var[3] == 1:
T = np.zeros(nvert * nel)
if field.var[4] != 0:
S = np.zeros((nvert * nel, field.var[4]))
#
ice = -(nvert + 1)
for iel in range(field.nel):
for (iz, ez), (iy, ey), (ix, ex) in product(enumerate(iiz),
enumerate(iiy),
enumerate(iix)):
iarray = iel * nppel + ix + iy * nix + iz * (nix * niy)
# Downsample copy into a column vector
if field.var[0] == 3:
r[iarray, :] = field.elem[iel].pos[:, ez, ey, ex]
if field.var[1] == 3:
v[iarray, :] = field.elem[iel].vel[:, ez, ey, ex]
if field.var[2] == 1:
p[iarray] = field.elem[iel].pres[:, ez, ey, ex]
if field.var[3] == 1:
T[iarray] = field.elem[iel].temp[:, ez, ey, ex]
if field.var[4] != 0:
S[iarray, :] = field.elem[iel].scal[:, ez, ey, ex]
if field.var[0] == 3:
for iz, iy, ix in product(range(max(niz - 1, 1)), range(niy - 1),
range(nix - 1)):
ice = ice + nvert + 1
for face in range(field.ndim - 1):
cell_id = iel * nppel + ix + iy * nix + (iz +
face) * nix * niy
ce[ice + face * 4 + 1] = cell_id
ce[ice + face * 4 + 2] = cell_id + 1
ce[ice + face * 4 + 3] = cell_id + nix + 1
ce[ice + face * 4 + 4] = cell_id + nix
# create the array of cells
ca = tvtk.CellArray()
ca.set_cells(nel, ce)
# create the unstructured dataset
dataset = tvtk.UnstructuredGrid(points=r)
# set the cell types
dataset.set_cells(ct, of, ca)
# set the data
idata = 0
if field.var[1] != 0:
dataset.point_data.vectors = v
dataset.point_data.vectors.name = "vel"
idata += 1
if field.var[2] == 1:
dataset.point_data.scalars = p
dataset.point_data.scalars.name = "pres"
idata += 1
if field.var[3] == 1:
dataset.point_data.add_array(T)
dataset.point_data.get_array(idata).name = "temp"
idata += 1
if field.var[4] != 0:
for ii in range(field.var[4]):
dataset.point_data.add_array(S[:, ii])
dataset.point_data.get_array(ii +
idata).name = "scal_%d" % (ii + 1)
#
dataset.point_data.update()
#
return dataset
7,
8,
9,
10,
11, # hex
3,
3,
1,
5
])
# The offsets for the cells, i.e. the indices where the cells
# start.
# offset = np.array([0, 5, 14])
offset = np.array([0, 5, 14])
tetra_type = tvtk.Tetra().cell_type # VTK_TETRA == 10
hex_type = tvtk.Hexahedron().cell_type # VTK_HEXAHEDRON == 12
cell_types = np.array([tetra_type, hex_type, 5])
# Create the array of cells unambiguously.
cell_array = tvtk.CellArray()
cell_array.set_cells(3, cells)
print('cell_array')
print(cell_array)
# Now create the UG.
ug = tvtk.UnstructuredGrid(points=surf_points)
# Now just set the cell types and reuse the ug locations and cells.
ug.set_cells(cell_types, offset, cell_array)
mlab.pipeline.surface(ug)
# mlab.show()
mlab.show()
