Source code for ase.visualize.vtk.module

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