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 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 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 _dump_arrays(self, filename):
from tvtk.api import tvtk
n = self.numPoints
cells = np.arange(n)
cells.shape = (n, 1)
cell_type = tvtk.Vertex().cell_type
ug = tvtk.UnstructuredGrid(points=self.points.transpose())
ug.set_cells(cell_type, cells)
from mayavi.core.dataset_manager import DatasetManager
dsm = DatasetManager(dataset=ug)
for name, field in self.data:
dsm.add_array(field.transpose(), name)
dsm.activate(name)
from tvtk.api import write_data
write_data(ug, filename)
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)
class VtkGridMixIn(object):
_EDGE_COUNT_TO_TYPE = {
1: tvtk.Vertex().cell_type,
2: tvtk.Line().cell_type,
3: tvtk.Triangle().cell_type,
4: tvtk.Quad().cell_type,
}
def to_vtk(self):
points = self.vtk_points()
cell_types = self.vtk_cell_types()
cell_array = self.vtk_cell_array()
offsets = self.vtk_offsets()
vtk_grid = tvtk.UnstructuredGrid(points=points)
vtk_grid.set_cells(cell_types, offsets, cell_array)
return vtk_grid
def vtk_points(self):
pad = np.zeros((3 - self._coords.shape[0], self._coords.shape[1]))
return np.vstack([self._coords, pad]).T
def vtk_cell_array(self):
cell_array = tvtk.CellArray()
cell_array.set_cells(self.get_cell_count(),
self.vtk_connectivity())
return cell_array
def vtk_cell_types(self):
cell_types = np.empty(self.get_cell_count(), dtype=int)
for (id, n_nodes) in enumerate(self.nodes_per_cell()):
try:
cell_types[id] = self._EDGE_COUNT_TO_TYPE[n_nodes]
except KeyError:
cell_types[id] = tvtk.Polygon().cell_type
return cell_types
def vtk_connectivity(self):
cells = np.empty(self.get_vertex_count() + self.get_cell_count(),
dtype=int)
cell_nodes = self.get_connectivity()
offset = 0
for n_nodes in self.nodes_per_cell():
cells[offset] = n_nodes
offset += n_nodes + 1
offset = 1
for cell in self.vtk_offsets():
n_nodes = cells[offset - 1]
cells[offset:offset + n_nodes] = cell_nodes[cell:cell +
n_nodes]
offset += n_nodes + 1
return cells
def vtk_offsets(self):
offsets = np.empty(self.get_cell_count(), dtype=int)
(offsets[0], offsets[1:]) = (0, self._offset[:-1])
return offsets
def vtk_write(self, file_name):
writer = tvtk.XMLUnstructuredGridWriter()
writer.set_input(self.to_vtk())
writer.file_name = file_name
writer.write()
def save_step_vtk(ses, outname, res_id):
'''
Save a step result into a VTK file.
'''
print 'Saving stepping results...'
#res_id = input('Enter the step result ID: ')
step_res = get_result(ses, id=res_id)
if step_res is None:
print 'No matching results for ID = {}'.format(res_id)
return
# we need the raster data
vel_res = get_result(ses, id=step_res.parent_id)
if vel_res is None:
print 'No matching results for ID = {}'.format(step_res.parent_id)
return
rast_res = get_result(ses, id=vel_res.parent_id)
if rast_res is None:
print 'No matching results for ID = {}'.format(vel_res.parent_id)
return
# get the field and linspaces
field, linspaces = None, None
for arr in rast_res.data:
if arr.typename == 'scalar':
field = arr.data
if arr.typename == 'linspace':
linspaces = arr.data
assert(field is not None and linspaces is not None)
field = Scalar(field, linspaces)
# exporting initial position
vtxs = [x for x in step_res.data if 'points' in x.typename]
points = list()
pot = list()
for vtx in vtxs:
for pt in vtx.data:
points.append(pt)
pot.append(field(pt))
points = numpy.asarray(points)
pot = numpy.asarray(pot)
from tvtk.api import tvtk, write_data
if len(vtxs):
ug = tvtk.UnstructuredGrid()
point_type = tvtk.Vertex().cell_type
npoints = len(points)
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 = pot
ug.point_data.scalars.name = 'potential'
fname = '%s-%s.vtk' % (outname, 'vtxs')
write_data(ug, fname)
# exporting paths
paths = [x for x in step_res.data if 'tuples' in x.typename]
points = list()
pot = list()
vel = list()
for path in paths:
for pt in path.data:
xyz = pt[0:3]
uvw = pt[3:6]
points.append(xyz)
pot.append(field(xyz))
mag = numpy.sqrt(uvw.dot(uvw))
vel.append(mag)
points = numpy.asarray(points)
pot = numpy.asarray(pot)
vel = numpy.asarray(vel)
if len(paths):
for flv,data in {'potential':pot, 'velocity':vel}.iteritems():
ug = tvtk.UnstructuredGrid()
point_type = tvtk.Vertex().cell_type
npoints = len(points)
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 = data
ug.point_data.scalars.name = flv
fname = '%s-%s-%s.vtk' % (outname, 'paths', flv)
write_data(ug, fname)
dbfile = sys.argv[1] step_res_id = sys.argv[2] outname = sys.argv[3] ses = larf.store.session(dbfile) sres = larf.store.result_typed(ses, 'stepping', step_res_id) sarrs = sres.array_data_by_name() points = sarrs['potential_points'] potarr = sarrs['potential'] 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 = potarr ug.point_data.scalars.name = 'potential' ## fixme: not currently saving 4-points # ug.point_data.add_array(times) # ug.point_data.get_array(1).name = 'time'
"""Returns a TVTK `DataSet` representing the face centers of this mesh
"""
try:
from tvtk.api import tvtk
except ImportError, 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 _toVTK3D(self, arr, rank=1):
if arr.dtype.name is 'bool':
# VTK can't do bool, and the exception isn't properly
# thrown back to the user
