import numpy as np
from vtk import vtkProp3D, vtkPolyDataMapper, vtkActor, vtkLODActor, \
vtkPointSet, vtkGlyph3D, vtkRenderer
from ase.visualize.vtk.sources import vtkCustomGlyphSource, \
vtkClampedGlyphSource
# -------------------------------------------------------------------
class vtkModule:
"""Modules represent a unified collection of VTK objects needed for
introducing basic visualization concepts such as surfaces or shapes.
A common trait of all modules is the need for an actor representation and
corresponding generic properties such as lighting, color and transparency.
"""
def __init__(self, vtk_act, vtk_property=None):
"""Construct basic VTK-representation of a module.
vtk_act: Instance of vtkActor or subclass thereof
A vtkActor represents an entity in a rendering scene.
vtk_property = None: Instance of vtkProperty or subclass thereof
A vtkProperty represents e.g. lighting and other surface
properties of a geometric object, in this case the actor.
"""
self.vtk_act = None
self.set_actor(vtk_act)
if vtk_property is not None:
self.set_property(vtk_property)
def set_actor(self, vtk_act):
"""Set the actor representing this module in a rendering scene."""
assert isinstance(vtk_act, vtkActor)
self.vtk_act = vtk_act
def set_property(self, vtk_property):
"""Set the property of the actor representing this module."""
self.vtk_act.SetProperty(vtk_property)
def get_actor(self):
"""Return the actor representing this module in a rendering scene."""
return self.vtk_act
# -------------------------------------------------------------------
class vtkLODModule(vtkModule):
_vtk_actor_class = vtkLODActor
def get_lod(self):
return 100
def set_actor(self, vtk_act):
vtkModule.set_actor(self, vtk_act)
if isinstance(vtk_act, vtkLODActor):
vtk_act.SetNumberOfCloudPoints(self.get_lod())
class vtkPolyDataModule(vtkModule):
_vtk_actor_class = vtkActor
__doc__ = vtkModule.__doc__ + """
Poly data modules are based on polygonal data sets, which can be mapped
into graphics primitives suitable for rendering within the VTK framework.
"""
def __init__(self, vtk_polydata, vtk_property=None):
"""Construct VTK-representation of a module containing polygonals.
vtk_polydata: Instance of vtkPolyData, subclass thereof or
vtkPolyDataAlgorithm, which produces vtkPolyData as output.
A vtkPolyData represents a polygonal data set consisting of
point and cell attributes, which can be mapped to graphics
primitives for subsequent rendering.
vtk_property = None: Instance of vtkProperty or subclass thereof
A vtkProperty represents e.g. lighting and other surface
properties of a geometric object, in this case the polydata.
"""
vtkModule.__init__(self, self._vtk_actor_class(), vtk_property)
self.vtk_dmap = vtkPolyDataMapper()
self.vtk_dmap.SetInputConnection(vtk_polydata.GetOutputPort())
self.vtk_act.SetMapper(self.vtk_dmap)
class vtkGlyphModule(vtkPolyDataModule):
[docs] __doc__ = vtkPolyDataModule.__doc__ + """
Glyph modules construct these polygonal data sets by replicating a glyph
source across a specific set of points, using available scalar or vector
point data to scale and orientate the glyph source if desired.
Example:
>>> atoms = molecule('C60')
>>> cell = vtkUnitCellModule(atoms)
>>> apos = vtkAtomicPositions(atoms.get_positions(), cell)
>>> vtk_ugd = apos.get_unstructured_grid()
>>> glyph_source = vtkAtomSource('C')
>>> glyph_module = vtkGlyphModule(vtk_ugd, glyph_source)
"""
def __init__(self, vtk_pointset, glyph_source,
scalemode=None, colormode=None):
"""Construct VTK-representation of a module containing glyphs.
These glyphs share a common source, defining their geometrical
shape, which is cloned and oriented according to the input data.
vtk_pointset: Instance of vtkPointSet or subclass thereof
A vtkPointSet defines a set of positions, which may then be
assigned specific scalar of vector data across the entire set.
glyph_source: Instance of ~vtk.vtkCustomGlyphSource or subclass thereof
Provides the basic shape to distribute over the point set.
"""
assert isinstance(vtk_pointset, vtkPointSet)
assert isinstance(glyph_source, vtkCustomGlyphSource)
# Create VTK Glyph3D based on unstructured grid
self.vtk_g3d = vtkGlyph3D()
self.vtk_g3d.SetInput(vtk_pointset)
self.vtk_g3d.SetSource(glyph_source.get_output())
self.vtk_g3d.SetScaleFactor(glyph_source.get_scale())
# Clamping normalizes the glyph scaling to within the range [0,1].
# Setting a scale factor will then scale these by the given factor.
if isinstance(glyph_source, vtkClampedGlyphSource):
self.vtk_g3d.SetClamping(True)
self.vtk_g3d.SetRange(glyph_source.get_range())
if scalemode is 'off':
self.vtk_g3d.SetScaleModeToDataScalingOff()
elif scalemode is 'scalar':
self.vtk_g3d.SetScaleModeToScaleByScalar()
elif scalemode is 'vector':
self.vtk_g3d.SetScaleModeToScaleByVector()
elif scalemode is not None:
raise ValueError('Unrecognized scale mode \'%s\'.' % scalemode)
if colormode is 'scale':
self.vtk_g3d.SetColorModeToColorByScale()
elif colormode is 'scalar':
self.vtk_g3d.SetColorModeToColorByScalar()
elif colormode is 'vector':
self.vtk_g3d.SetColorModeToColorByVector()
elif colormode is not None:
raise ValueError('Unrecognized scale mode \'%s\'.' % scalemode)
vtkPolyDataModule.__init__(self, self.vtk_g3d, glyph_source.get_property())
# -------------------------------------------------------------------
class vtkLabelModule(vtkModule):
def __init__(self, vtk_pointset, vtk_property=None):
vtk_Module.__init__(self, vtkActor(), vtk_property)
assert isinstance(vtk_pointset, vtkPointSet)
self.vtk_dmap = vtkLabeledDataMapper()
#self.vtk_dmap.SetLabelModeToLabelIds() #TODO XXX strings!!!
self.vtk_dmap.GetLabelTextProperty().SetFontSize(12)
self.vtk_dmap.GetLabelTextProperty().SetJustificationToRight()
self.vtk_dmap.SetInputConnection(vtk_pointset.GetOutputPort())
self.vtk_act.SetMapper(self.vtk_dmap)
#vtk_g3d.GetPointIdsName...
# -------------------------------------------------------------------
class vtkModuleAnchor:
"""
TODO XXX
"""
def __init__(self):
"""Construct an anchoring point for various VTK modules.
"""
self.modules = {}
def get_module(self, name):
if name not in self.modules.keys():
raise RuntimeError('Module \'%s\' does not exists.' % name)
return self.modules[name]
def add_module(self, name, module):
if not isinstance(module, vtkModule):
raise ValueError('Module must be instance of vtkModule.')
if name in self.modules.keys():
raise RuntimeError('Module \'%s\' already exists.' % name)
self.modules[name] = module
def get_actor(self, name):
return self.get_module(name).get_actor()
def add_actors_to_renderer(self, vtk_renderer, name=None):
assert isinstance(vtk_renderer, vtkRenderer)
if name is None:
for module in self.modules.values():
vtk_renderer.AddActor(module.get_actor())
else:
vtk_renderer.AddActor(self.get_actor(name))