def save_drift_vtk(result, outname, **kwds):
'''
Save a drift result into a VTK file.
'''
outname = osp.splitext(outname)[0]
from tvtk.api import tvtk, write_data
arrs = result.array_data_by_name()
for thing in ['potential', 'gradient', 'velocity']:
points = arrs[thing + '_points']
values = arrs[thing]
npoints = len(points)
ug = tvtk.UnstructuredGrid()
point_type = tvtk.Vertex().cell_type
cell_types = numpy.array([point_type] * npoints)
cell_array = tvtk.CellArray()
cells = numpy.array([npoints] + range(npoints))
cell_array.set_cells(point_type, cells)
ug.set_cells(1, cell_array)
ug.points = points
ug.point_data.scalars = values
ug.point_data.scalars.name = thing
fname = '%s-%s.vtk' % (outname, thing)
print 'writing %s' % fname
write_data(ug, fname)
def grid2poly(grid):
""" Read grid and convert to polydata """
points = grid.nodes
face_coords_mask = np.dstack([grid.faces.mask, grid.faces.mask])
face_coords = np.ma.masked_array(
grid.nodes[grid.faces],
face_coords_mask
)
# select faces that we need to duplicate
# TODO: duplicate these (with nodes moved to other end of the world) (mod something...)
x_max, y_max = face_coords.max(axis=(1)).T
x_min, y_min = face_coords.min(axis=(1)).T
mask = (x_min < -150) & (x_max > 150)
# Select faces
faces = grid.faces[~mask]
n_cells = faces.shape[0]
cell_array = tvtk.CellArray()
counts = (~faces.mask).sum(axis=1)
assert faces.min() >= 0, 'expected 0 based faces'
cell_idx = np.c_[counts, faces.filled(-999)].ravel()
cell_idx = cell_idx[cell_idx != -999]
cell_array.set_cells(n_cells, cell_idx)
# fill in the properties
polydata = tvtk.PolyData()
# fill in z dimension
polydata.points = np.c_[points, np.zeros_like(points[:, 0])]
polydata.polys = cell_array
return polydata, mask
def save_current_vtk(result, outfile, **kwds):
'''
Save current assuming no structure between starting points
'''
from tvtk.api import tvtk, write_data
cres = result
dres = cres.parent_by_type('drift') # drift
values = numpy.hstack([a.data for a in cres.arrays if a.type == 'pscalar'])
points4 = numpy.vstack([a.data for a in dres.arrays if a.type == 'path'])
assert len(values) == len(points4)
points = points4[:, :3]
npoints = len(points)
print 'shapes: %s %s' % (str(values.shape), str(points.shape))
ug = tvtk.UnstructuredGrid()
point_type = tvtk.Vertex().cell_type
cell_types = numpy.array([point_type] * npoints)
cell_array = tvtk.CellArray()
cells = numpy.array([npoints] + range(npoints))
cell_array.set_cells(point_type, cells)
ug.set_cells(1, cell_array)
ug.points = points
ug.point_data.scalars = values
ug.point_data.scalars.name = 'current'
write_data(ug, outfile)
def write(self):
print self._indent, "="*50
points=self.generatePoints()
[cell_types1, offset1, cellcount1, cells1, cellData1] = self.generate3dCells(points)
[cell_types2, offset2, cellcount2, cells2, cellData2] = self.generate2dCells(points)
print self._indent, "="*50
print self._indent, "= Points and cells Created"
cell_types=np.append(cell_types1, cell_types2)
offset=np.append(offset1, offset2)
cellcount= cellcount1 + cellcount2
cells = np.append(cells1,cells2)
cells = cells.flatten()
cellData=np.append(cellData1, cellData2)
cellData=cellData.flatten()
cell_array = tvtk.CellArray()
cell_array.set_cells(cellcount, cells)
points=np.array(np.dstack([self.x*self.scale[0],self.y*self.scale[1],self.z*self.scale[2]]),'f')
points=points.reshape(self.x.shape[0],3)
print self._indent, "= Writing a VTK file"
ug=self.writeUnstructuredGrid(points, cell_types, offset, cell_array, cellData)
print self._indent, "= The file has been successfully written!"
print self._indent, "-"*50
return ug
def make_grid(self):
"""return an unstructured grid, based on contours (xc, yc) with possible
scalar and vector values"""
# Get the contours
xc = self.ds.variables['FlowElemContour_x']
yc = self.ds.variables['FlowElemContour_y']
ncells = xc.shape[0]
# We're using an unstructured grid
grid = tvtk.PolyData()
ncells = xc.shape[0]
X = xc[:]
Y = yc[:]
Z = np.zeros_like(X)
pts = np.c_[X.ravel(), Y.ravel(), Z.ravel()]
# quads all the way down
# quads with different resolutions actually don't work so well
cell_array = tvtk.CellArray()
# For unstructured grids you need to count the number of edges per cell
cell_idx = np.array([[4] + [i * 4 + j for j in range(4)]
for i in range(ncells)]).ravel()
cell_array.set_cells(ncells, cell_idx)
# fill in the properties
grid.points = pts
grid.polys = cell_array
return grid
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 VTKCellDataSet(self):
"""Returns a TVTK `DataSet` representing the cells of this mesh
"""
cvi = self._orderedCellVertexIDs.swapaxes(0, 1)
from fipy.tools import numerix
if isinstance(cvi, numerix.ma.masked_array):
counts = cvi.count(axis=1)[:, None]
cells = numerix.ma.concatenate((counts, cvi), axis=1).compressed()
else:
counts = numerix.array([cvi.shape[1]]*cvi.shape[0])[:, None]
cells = numerix.concatenate((counts, cvi), axis=1).flatten()
try:
from tvtk.api import tvtk
except ImportError as e:
from enthought.tvtk.api import tvtk
num = counts.shape[0]
cps_type = self._VTKCellType
cell_types = numerix.array([cps_type]*num)
cell_array = tvtk.CellArray()
cell_array.set_cells(num, cells)
points = self.vertexCoords
points = self._toVTK3D(points)
ug = tvtk.UnstructuredGrid(points=points)
offset = numerix.cumsum(counts[:, 0]+1)
if len(offset) > 0:
offset -= offset[0]
ug.set_cells(cell_types, offset, cell_array)
return ug
def VTKFaceDataSet(self):
"""Returns a TVTK `DataSet` representing the face centers of this mesh
"""
try:
from tvtk.api import tvtk
except ImportError as e:
from enthought.tvtk.api import tvtk
points = self.faceCenters
points = self._toVTK3D(numerix.array(points))
ug = tvtk.UnstructuredGrid(points=points)
num = len(points)
counts = numerix.array([1] * num)[..., numerix.newaxis]
cells = numerix.arange(self.numberOfFaces)[..., numerix.newaxis]
cells = numerix.concatenate((counts, cells), axis=1)
cell_types = numerix.array([tvtk.Vertex().cell_type]*num)
cell_array = tvtk.CellArray()
cell_array.set_cells(num, cells)
counts = numerix.array([1] * num)
offset = numerix.cumsum(counts+1)
if len(offset) > 0:
offset -= offset[0]
ug.set_cells(cell_types, offset, cell_array)
return ug
def initMixedTvtk2D(self):
"""
Initialise the actual 2 dimensional tvtk object
This is for a mixed data type
"""
poly_type = tvtk.Polygon().cell_type
#poly_type = tvtk.Polyhedron().cell_type
self.ug = tvtk.UnstructuredGrid(points=self.points)
cells = np.array(self.cells[self.cells.mask==False])
cell_array = tvtk.CellArray()
cell_array.set_cells(self.Nc,cells)
offsets = np.cumsum(self.nfaces)
offsets = offsets - offsets[0]
poly_types = [poly_type for ii in range(self.Nc)]
pdb.set_trace()
self.ug.set_cells(np.array(poly_types),offsets, cell_array)
self.ug.cell_data.scalars = self.data
self.ug.cell_data.scalars.name = 'suntans_scalar'
def typeUnstructuredGrid_Plane(xyz, ikle):
cellUG = tvtk.CellArray()
cellUG.set_cells(len(ikle), np.insert(ikle, 0, 3, axis=1).ravel())
typeUG = tvtk.UnstructuredGrid(points=xyz)
typeUG.set_cells(5 * np.ones(len(ikle)), 7 * np.arange(len(ikle)), cellUG)
return typeUG
def _get_vtk_structure(self):
ug = tvtk.UnstructuredGrid()
cell_array, cell_offsets, cell_types = self.vtk_cell_data
n_cells = cell_types.shape[0]
ug.points = self.vtk_X
vtk_cell_array = tvtk.CellArray()
vtk_cell_array.set_cells(n_cells, cell_array)
ug.set_cells(cell_types, cell_offsets, vtk_cell_array)
return ug
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 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 test_cell_array(self):
""" Test if cell array insertion updates number of cells.
Fixes GH Issue 178.
"""
cell_array = tvtk.CellArray()
line1 = tvtk.Line()
self.assertEqual(cell_array.number_of_cells, 0)
cell_array.insert_next_cell(line1)
self.assertEqual(cell_array.number_of_cells, 1)
line2 = tvtk.Line()
cell_array.insert_next_cell(line2)
self.assertEqual(cell_array.number_of_cells, 2)
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 _configure(self, points):
self.vtk_points = tvtk.Points()
self.lines = tvtk.CellArray()
for point in points:
self.vtk_points.insert_next_point(self._to_float_tuple(point))
for i in xrange(len(points)-1):
line = tvtk.Line()
line.point_ids.set_id(0, i)
line.point_ids.set_id(1, i+1)
self.lines.insert_next_cell(line)
self.poly_data = tvtk.PolyData(points=self.vtk_points, lines=self.lines)
self._set_actor()
def _build_triangle(self, cell, points):
faces = tvtk.CellArray()
polygon = tvtk.Polygon()
polygon.point_ids.number_of_ids = 3
polygon.point_ids.set_id(0, 0)
polygon.point_ids.set_id(1, 1)
polygon.point_ids.set_id(2, 2)
faces.insert_next_cell(polygon)
cell_points = self._get_cell_points(cell, points)
poly = Polyhedron(cell_points, faces)
name = 'Triang-{}-{}-{}'.format(*cell)
return poly, name
def _build_tetrahedron(self, cell, points):
faces = tvtk.CellArray()
for i0,i1,i2 in [(0,1,2), (0,3,1), (0,2,3), (1,3,2)]:
polygon = tvtk.Polygon()
polygon.point_ids.number_of_ids = 3
polygon.point_ids.set_id(0, i0)
polygon.point_ids.set_id(1, i1)
polygon.point_ids.set_id(2, i2)
faces.insert_next_cell(polygon)
cell_points = self._get_cell_points(cell, points)
poly = Polyhedron(cell_points, faces)
name = 'Tetra-{}-{}-{}-{}'.format(*cell)
return poly, name
def _build_quadrilateral(self, cell, points):
faces = tvtk.CellArray()
polygon = tvtk.Polygon()
polygon.point_ids.number_of_ids = 4
polygon.point_ids.set_id(0, 0)
polygon.point_ids.set_id(1, 1)
polygon.point_ids.set_id(2, 2)
polygon.point_ids.set_id(3, 3)
faces.insert_next_cell(polygon)
cell_points = self._get_cell_points(cell, points)
poly = Polyhedron(cell_points, faces)
name = 'Quad-{}-{}-{}-{}'.format(*cell)
return poly, name
def setup(self):
self.new_dir()
# make a loop over the XDomainModel
#sim = self.sim
ts = self.tstep
fe_domain = ts.fe_domain
vtk_point_list = []
vtk_cell_list = []
vtk_cell_offset_list = []
vtk_cell_types_list = []
point_offset = 0
cell_offset = 0
for domain in fe_domain:
xdomain = domain.xdomain
print('VAR')
print(self.var)
print('DOMAIN')
print(domain.tmodel.var_dict)
var_function = domain.tmodel.var_dict.get(self.var, None)
if var_function == None or xdomain.hidden:
continue
fets = xdomain.fets
DELTA_x_ab = fets.vtk_expand_operator
n_c = fets.n_nodal_dofs
vtk_points = np.einsum('Ia,ab->Ib', xdomain.x_Eia.reshape(-1, n_c),
DELTA_x_ab)
vtk_point_list.append(vtk_points)
cells, cell_offsets, cell_types = xdomain.get_vtk_cell_data(
'nodes', point_offset, cell_offset)
point_offset += vtk_points.shape[0]
cell_offset += cells.shape[0]
vtk_cell_list.append(cells)
vtk_cell_offset_list.append(cell_offsets)
vtk_cell_types_list.append(cell_types)
if len(vtk_cell_types_list) == 0:
raise ValueError('Empty output for field variable %s in model %s' %
(self.var, domain.tmodel))
vtk_cell_types = np.hstack(vtk_cell_types_list)
vtk_cell_offsets = np.hstack(vtk_cell_offset_list)
vtk_cells = np.hstack(vtk_cell_list)
n_cells = vtk_cell_types.shape[0]
vtk_cell_array = tvtk.CellArray()
vtk_cell_array.set_cells(n_cells, vtk_cells)
self.ug = tvtk.UnstructuredGrid(points=np.vstack(vtk_point_list))
self.ug.set_cells(vtk_cell_types, vtk_cell_offsets, vtk_cell_array)
def _build_voxel(self, cell, points):
faces = tvtk.CellArray()
for i0,i1,i2,i3 in [(0,1,3,2), (1,3,7,5), (5,7,6,4), (4,0,2,6),
(6,2,3,7), (0,1,5,4)]:
polygon = tvtk.Polygon()
polygon.point_ids.number_of_ids = 4
polygon.point_ids.set_id(0, i0)
polygon.point_ids.set_id(1, i1)
polygon.point_ids.set_id(2, i2)
polygon.point_ids.set_id(3, i3)
faces.insert_next_cell(polygon)
cell_points = self._get_cell_points(cell, points)
poly = Polyhedron(cell_points, faces)
name = 'Voxel-{}-{}-{}-{}-{}-{}-{}-{}'.format(*cell)
return poly, name
def get_trunkcone(b, a):
phi_base, theta_base = sph.to(a, b)[1:]
quads = tvtk.CellArray() #vtk.vtkCellArray()
points = tvtk.Points() #vtk.vtkPoints()
Nface = 3
for i in range(Nface + 1):
# rotate
phi, theta = convdir((i % Nface) * 2 * pi / Nface, pi * 0.5, phi_base,
theta_base)
# generate new points
p = tuple(sph.xyz(a[3] * 0.5 * radius_factor, phi, theta, a[:3]))
q = tuple(sph.xyz(b[3] * 0.5 * radius_factor, phi, theta, b[:3]))
# insert points
points.append(p)
points.append(q)
if i >= 1:
# create a face
quad = tvtk.Quad()
n = points.number_of_points - 1
quad.point_ids.set_id(0, n - 3) # p
quad.point_ids.set_id(1, n - 2) # q
quad.point_ids.set_id(2, n) # q
quad.point_ids.set_id(3, n - 1) # p
# insert the new face
quads.insert_next_cell(quad)
# create the actor
polydata = tvtk.PolyData(points=points, polys=quads)
if new_tvtk:
mapper = tvtk.PolyDataMapper()
configure_input(mapper, polydata)
else:
mapper = tvtk.PolyDataMapper(input=polydata)
actor = tvtk.Actor(mapper=mapper)
fig.scene.add_actor(actor)
return actor
def main():
print ''
print '1. Format of surface families should be [[a1, b1, c1], [a2, b2, c2],...,[an, bn, cn]], n >= 1.'
print '2. Distance between the center of the polyhedron and the surface should be positive.'
print '3. Color (r, g, b) is represented by the number from 0 to 1.'
print ''
hkl = surface_index() # surface family
print ''
surf_dist = surface_distance(hkl) # distance
print ''
surf_color = surface_color(hkl) # surface color
print 'surface_color', surf_color
print ''
planes, number_planes = crystal_plane(hkl)
ps_setd_mingled(number_planes, planes, surf_dist)
diameter = 20.0
volume = 4.0 / 3.0 * np.pi * (diameter * 0.5)**3
qfp = enclose_polyhedron(number_planes, planes, volume)
p1 = tvtk.PolyData()
n_planes = len(qfp[2])
for i in np.arange(len(number_planes)):
for surf_index in np.arange(sum(number_planes[0:i]),
sum(number_planes[0:i + 1])):
print 'surf_index', surf_index
p1.points = qfp[2][
surf_index] # the coordinates of the intersection of each plane.
faces = crystal_surface(qfp[2][surf_index])
cells = tvtk.CellArray(
) # create a new CellArray object to assign the polys property.
cells.set_cells(
1, faces
) # the first parameter is the number of faces (here is 1),
p1.polys = cells # and the second parameter is an array describing the composition of each face.
p1.point_data.scalars = np.linspace(0.0, 1.0, len(p1.points))
mlab.figure(number_planes, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
mlab.pipeline.surface(p1,
representation='surface',
opacity=1.0,
color=surf_color[i])
axe = tvtk.AxesActor(total_length=(3, 3, 3))
mlab.show()
def vtkCone(p, q):
from math import pi
phi_base, theta_base = misc.Spherical.to(q, p)[1:]
quads = tvtk.CellArray() #vtk.vtkCellArray()
points = tvtk.Points() #vtk.vtkPoints()
for i in range(11):
# rotate
phi, theta = ConvertDirection((i % 10) * 2 * pi / 10, pi * .5,
phi_base, theta_base, True)
# generate new points
_p = tuple(
misc.Spherical.xyz(p[3] * .5 * cone_factor, phi, theta, p[0:3]))
_q = tuple(
misc.Spherical.xyz(q[3] * .5 * cone_factor, phi, theta, q[0:3]))
# insert points
points.append(_p)
points.append(_q)
if i >= 1:
# create a face
quad = tvtk.Quad()
n = points.number_of_points - 1
quad.point_ids.set_id(0, n - 3) # p
quad.point_ids.set_id(1, n - 2) # q
quad.point_ids.set_id(2, n) # q
quad.point_ids.set_id(3, n - 1) # p
# insert the new face
quads.insert_next_cell(quad)
# create the actor
polydata = tvtk.PolyData(points=points, polys=quads)
mapper = tvtk.PolyDataMapper(input=polydata)
actor = tvtk.Actor(mapper=mapper)
return actor
def save_raster_vtk(ses, outname, res_id):
'''
Save a drift result into a VTK file.
'''
print 'Saving raster results...'
#res_id = input('Enter the raster result ID: ')
result = get_result(ses, id=res_id)
if result is None:
print 'No matching results for ID = {}'.format(res_id)
return
arrs = result.array_data_by_name()
points = arrs['points'].T
from tvtk.api import tvtk, write_data
for thing in ['pfield','efield']:
values = arrs[thing]
npoints = len(points)
ug = tvtk.UnstructuredGrid()
point_type = tvtk.Vertex().cell_type
cell_types = numpy.array([point_type]*npoints)
cell_array = tvtk.CellArray()
cells = numpy.array([npoints]+range(npoints))
cell_array.set_cells(point_type, cells)
ug.set_cells(1, cell_array)
if thing == 'pfield':
ug.points = points
ug.point_data.scalars = values.reshape(npoints)
ug.point_data.scalars.name = thing
else:
ug.points = points
field = numpy.asarray([[x,y,z] for x,y,z in zip(values[0,:,:,:].reshape(npoints),values[1,:,:,:].reshape(npoints),values[2,:,:,:].reshape(npoints))])
ug.point_data.vectors = field
ug.point_data.vectors.name = thing
fname = '%s-%s.vtk' % (outname, thing)
write_data(ug, fname)
def main():
xo = yo = zo = 0.0
while True:
a, b, c = state.uniform(size=3)
xo = bounds[0] + (bounds[1] - bounds[0]) * a
yo = bounds[2] + (bounds[3] - bounds[2]) * b
zo = bounds[4] + (bounds[5] - bounds[4]) * c
r = pow(250 - xo, 2) + pow(250 - yo, 2)
print "coordinate: {},{},{}, {}".format(xo, yo, zo, r)
# is_inside_surface ==1 when point is INSIDE enclosed surface
if r < rad_om2 and not enclosed.is_inside_surface(xo, yo, zo):
break
idlist = tvtk.IdList()
points = tvtk.Points()
for pid in range(100001):
while True:
xr, yr, zr = eps * (state.uniform(size=3) - 0.5)
r = pow(250 - (xo + xr), 2) + pow(250 - (yo + yr), 2)
if (r < rad_om2 and not
enclosed.is_inside_surface(xo+xr, yo+yr, zo+zr)):
break
(xo, yo, zo) = (xo + xr, yo + yr, zo + zr)
idlist.insert_next_id(pid)
points.insert_next_point((xo, yo, zo))
points.update_traits()
# write out path trace to vtk every 5000 timesteps
if pid % 5000 == 0:
_array = tvtk.CellArray()
rand_walk = tvtk.PolyData(points=points, lines=_array)
rand_walk.insert_next_cell(4, idlist) # VTK_POLY_LINE == 4
writer = tvtk.PolyDataWriter(file_name='run_{}.vtk'.format(pid),
input=rand_walk)
writer.update()
def to_polydata(self):
"""convert grid to polydata"""
grid = self.ugrid
faces = grid['faces']
points = grid['points']
n_cells = faces.shape[0]
points.shape, n_cells
cell_array = tvtk.CellArray()
counts = (~faces.mask).sum(axis=1)
faces_zero_based = faces - 1
cell_idx = np.c_[counts, faces_zero_based.filled(-999)].ravel()
cell_idx = cell_idx[cell_idx != -999]
cell_array.set_cells(n_cells, cell_idx)
# fill in the properties
polydata = tvtk.PolyData()
polydata.points = points
polydata.polys = cell_array
return polydata
def vtk_cell_array(self):
cell_array = tvtk.CellArray()
cell_array.set_cells(self.get_cell_count(),
self.vtk_connectivity())
return cell_array
def _close_helix(c0, c1, c2):
"""
Close the start and end of the given open helix coordinates, return the
result as a <tvtk.PolyData> instance.
"""
# We do not rely on the triangulation of mlab.mesh(), as (1) it seems to
# fail for small numbers of points, (2) we have to close the two ends of
# the tube manually anyway
assert c0.shape == c1.shape == c2.shape
n_circles = c0.shape[0]
points_per_circle = c0.shape[1] - 1
cells_per_circle = points_per_circle * 2
open_num = n_circles * points_per_circle
# Make the coordinates one array, remove overlapping seam points, flatten
open_points = np.empty((n_circles, points_per_circle, 3), dtype=np.float)
open_points[..., 0] = c0[:, :-1]
open_points[..., 1] = c1[:, :-1]
open_points[..., 2] = c2[:, :-1]
open_points = open_points.reshape(-1, 3)
# Create triangles between circles of the tube
open_cells = np.empty((cells_per_circle * (n_circles - 1), 3),
dtype=np.int)
for i_circle in range(n_circles - 1):
current_circle_ids = np.arange((i_circle + 0) * points_per_circle,
(i_circle + 1) * points_per_circle)
next_circle_ids = np.arange((i_circle + 1) * points_per_circle,
(i_circle + 2) * points_per_circle)
# Connect the neighboring circles' points, i.e. create the cells
# (a, b, c) and (b, d, c) according to the following scheme:
#
# current a-c ...
# |/|
# next b-d ...
#
current_cells = np.empty((cells_per_circle, 3), dtype=np.int)
for i_vertex in range(points_per_circle):
a = current_circle_ids[i_vertex]
c = current_circle_ids[(i_vertex + 1) % points_per_circle]
b = next_circle_ids[i_vertex]
d = next_circle_ids[(i_vertex + 1) % points_per_circle]
current_cells[2 * i_vertex] = a, b, c
current_cells[2 * i_vertex + 1] = b, d, c
open_cells[cells_per_circle * i_circle:cells_per_circle *
(i_circle + 1)] = current_cells
# Calculate the centers for the first and last circle (we can simply use
# the mean of the points, as they are evenly spaced)
first_circle = open_points[:points_per_circle]
first_center = np.mean(first_circle, axis=0)
last_circle = open_points[-points_per_circle:]
last_center = np.mean(last_circle, axis=0)
# Get the ids for the circle points
first_ids = np.arange(0, points_per_circle)
last_ids = np.arange(open_num - points_per_circle, open_num)
# Create a new point array that additionally holds the center points
closed_points = np.empty((open_num + 2, 3), dtype=np.float)
closed_points[:-2] = open_points
closed_points[-2] = first_center
closed_points[-1] = last_center
first_center_id = open_num
last_center_id = open_num + 1
# Create the cells that close the first and last circle
first_cells = np.empty((points_per_circle, 3), dtype=np.int)
for i in range(points_per_circle):
first_cells[i] = first_center_id, first_ids[i], first_ids[
(i + 1) % points_per_circle]
last_cells = np.empty((points_per_circle, 3), dtype=np.int)
for i in range(points_per_circle):
last_cells[i] = last_center_id, last_ids[i], last_ids[
(i + 1) % points_per_circle]
# Stack the cells and convert them to tvtk.CellArray, then create and
# return new tvtk.PolyData instance
closed_cells = tvtk.CellArray()
closed_cells.from_array(np.vstack((open_cells, first_cells, last_cells)))
closed_poly_data = tvtk.PolyData(points=closed_points, polys=closed_cells)
return closed_poly_data
def plot(self):
'''
Plot a 3D visualisation of the Voronoi grid using mayavi.
This method requires mayavi to be installed and also needs the vertices
information to be available (see the class constructor).
Note that in order for this method to work in an interactive IPython session,
a series of environment variables and proper switches need to be used
depending on your system configuration. For instance, on a Linux machine
with PyQt4 and a recent IPython version, the following bash startup
command for IPython can be used:
``ETS_TOOLKIT=qt4 QT_API=pyqt ipython --gui=qt4``
This sets both the mayavi and the IPython GUI toolkit to qt4, and the ``QT_API``
variable is used to specify that we want the ``pyqt`` API (as opposed to the
``pyside`` alternative API - PySide is an alternative implementation of PyQt).
It should be possible to get this method working on different configurations,
but the details will be highly system-dependent.
'''
if not self._with_vertices:
raise ValueError(
'the class must be constructed with \'with_vertices=True\' in order to support plotting'
)
import numpy as np
try:
from tvtk.api import tvtk
from mayavi.api import Engine
from mayavi import mlab
from mayavi.sources.vtk_data_source import VTKDataSource
from mayavi.modules.surface import Surface
from mayavi.modules.scalar_cut_plane import ScalarCutPlane
except ImportError:
raise ImportError(
'the plot method requires Mayavi, please make sure it is correctly installed'
)
# Shortcut.
vertices = self._neighbours_table['vertices']
# This is a list of all the vertices composing all voronoi cells.
# points = [[x1,y1,z1],[x2,y2,z2],...]
points = []
# Array to describe each voronoi cell in terms of the points list above. E.g.,
# cells = [4,0,1,2,3,5,4,5,6,7,8]
# This describes two cells, the first with 4 vertices whose indices in the points array
# are 0,1,2,3, the second with 5 vertices whose indices are 4,5,6,7,8.
cells = []
cur_cell_idx = 0
# Indices in the cells array where each new cell starts. In the example above,
# offset = [0,5]
offset = []
cur_offset = 0
# Array of cell types. Cells will all be of the same type.
cell_types = []
# Build the above quantities.
for v in vertices:
# Drop the empty vertices coordinates, signalled by NaN.
arr = v[~np.isnan(v)]
assert (len(arr) % 3 == 0)
tmp = np.split(arr, len(arr) / 3)
# Append the vertices.
points = points + tmp
# Append the cell description.
cells = cells + \
[len(tmp)] + range(cur_cell_idx, cur_cell_idx + len(tmp))
cur_cell_idx += len(tmp)
# Append the offset info.
offset.append(cur_offset)
cur_offset += len(tmp) + 1
# Append the cell type.
cell_types.append(tvtk.ConvexPointSet().cell_type)
# Cache the sites' positions.
sites_arr = self._neighbours_table['coordinates']
# Setup the Mayavi engine and figure.
e = Engine()
e.start()
fig = mlab.figure(engine=e)
# Plot the sites.
mlab.points3d(sites_arr[:, 0],
sites_arr[:, 1],
sites_arr[:, 2],
figure=fig)
# Plot the cells with coloured surfaces.
# This is just an array of scalars to assign a "temperature" to each cell vertex, which will be
# used for coloring purposes.
temperature = np.arange(0, len(points) * 10, 10, 'd')
# Initialise the array of cells.
cell_array = tvtk.CellArray()
cell_array.set_cells(len(vertices), np.array(cells))
# Initialise the unstructured grid object.
ug = tvtk.UnstructuredGrid(points=np.array(points))
ug.set_cells(np.array(cell_types), np.array(offset), cell_array)
ug.point_data.scalars = temperature
ug.point_data.scalars.name = 'temperature'
# Create a data source from the unstructured grid object.
src = VTKDataSource(data=ug)
# Add the source to the engine.
e.add_source(src)
# Create a surface object with opacity 0.5
surf = Surface()
surf.actor.property.opacity = 0.5
# Add the surface object to the engine.
e.add_module(surf)
# Add a cut plane as well.
e.add_module(ScalarCutPlane())
# Create another representation of the grid, this time using only white wireframe
# to highlight to shape of the cells.
# Rebuild the ug.
ug = tvtk.UnstructuredGrid(points=np.array(points))
ug.set_cells(np.array(cell_types), np.array(offset), cell_array)
src = VTKDataSource(data=ug)
e.add_source(src)
surf = Surface()
surf.actor.property.representation = 'wireframe'
e.add_module(surf)
cp = ScalarCutPlane()
e.add_module(cp)
