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): __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))