/ase/visualize/vtk/atoms.py

  1
  2import numpy as np
  3
  4from ase import Atoms
  5
  6from ase.visualize.vtk.sources import vtkAtomSource, vtkForceSource, \
  7                                      vtkVelocitySource
  8from ase.visualize.vtk.cell import vtkUnitCellModule, vtkAxesModule
  9from ase.visualize.vtk.grid import vtkAtomicPositions
 10from ase.visualize.vtk.module import vtkModuleAnchor, vtkGlyphModule
 11
 12# -------------------------------------------------------------------
 13
 14class vtkAtoms(vtkModuleAnchor, vtkAtomicPositions):
 15    """Provides fundamental representation for ``Atoms``-specific data in VTK.
 16
 17    The ``vtkAtoms`` class plots atoms during simulations, extracting the
 18    relevant information from the list of atoms. It is created using
 19    the list of atoms as an argument to the constructor. Then one or more
 20    visualization modules can be attached using add_module(name, module).
 21
 22    Example:
 23
 24    >>> va = vtkAtoms(atoms)
 25    >>> va.add_forces()
 26    >>> va.add_axes()
 27    >>> XXX va.add_to_renderer(vtk_ren)
 28
 29    """
 30    def __init__(self, atoms, scale=1):
 31        """Construct a fundamental VTK-representation of atoms.
 32
 33        atoms: Atoms object or list of Atoms objects
 34            The atoms to be plotted.
 35
 36        scale = 1: float or int
 37            Relative scaling of all Atoms-specific visualization.
 38
 39        """
 40        assert isinstance(atoms, Atoms)
 41        self.atoms = atoms
 42
 43        self.scale = scale
 44
 45        vtkModuleAnchor.__init__(self)
 46        vtkAtomicPositions.__init__(self, self.atoms.get_positions(),
 47                                    vtkUnitCellModule(self.atoms))
 48
 49        self.force = None
 50        self.velocity = None
 51
 52        symbols = self.atoms.get_chemical_symbols()
 53        for symbol in np.unique(symbols):
 54            # Construct mask for all atoms with this symbol
 55            mask = np.array(symbols) == symbol
 56            if mask.all():
 57                subset = None
 58            else:
 59                subset = np.argwhere(mask).ravel()
 60
 61            # Get relevant VTK unstructured grid
 62            vtk_ugd = self.get_unstructured_grid(subset)
 63
 64            # Create atomic glyph source for this symbol
 65            glyph_source = vtkAtomSource(symbol, self.scale)
 66
 67            # Create glyph module and anchor it
 68            self.add_module(symbol, vtkGlyphModule(vtk_ugd, glyph_source))
 69
 70    def has_forces(self):
 71        return self.force is not None
 72
 73    def has_velocities(self):
 74        return self.velocity is not None
 75
 76    def add_cell(self):
 77        """Add a box outline of the cell using atoms.get_cell(). The existing
 78        ``vtkUnitCellModule`` is added to the module anchor under ``cell``."""
 79        self.add_module('cell', self.cell)
 80
 81    def add_axes(self):
 82        """Add an orientation indicator for the cartesian axes. An appropriate
 83        ``vtkAxesModule`` is added to the module anchor under ``axes``."""
 84        self.add_module('axes', vtkAxesModule(self.cell))
 85
 86    def add_forces(self):
 87        """Add force vectors for the atoms using atoms.get_forces(). A
 88        ``vtkGlyphModule`` is added to the module anchor under ``force``."""
 89        if self.has_forces():
 90            raise RuntimeError('Forces already present.')
 91        elif self.has_velocities():
 92            raise NotImplementedError('Can\'t add forces due to velocities.')
 93
 94        # Add forces to VTK unstructured grid as vector data
 95        vtk_fda = self.add_vector_property(self.atoms.get_forces(), 'force')
 96
 97        # Calculate max norm of the forces
 98        fmax = vtk_fda.GetMaxNorm()
 99
100        # Get relevant VTK unstructured grid
101        vtk_ugd = self.get_unstructured_grid()
102
103        self.force = vtkGlyphModule(vtk_ugd, vtkForceSource(fmax, self.scale),
104                                    scalemode='vector', colormode=None)
105        self.add_module('force', self.force)
106
107    def add_velocities(self):
108        """Add velocity vectors for the atoms using atoms.get_velocities(). A
109        ``vtkGlyphModule`` is added to the module anchor under ``velocity``."""
110        if self.has_velocities():
111            raise RuntimeError('Velocities already present.')
112        elif self.has_forces():
113            raise NotImplementedError('Can\'t add velocities due to forces.')
114
115        # Add velocities to VTK unstructured grid as vector data
116        vtk_vda = self.add_vector_property(self.atoms.get_velocities(), 'velocity')
117
118        # Calculate max norm of the velocities
119        vmax = vtk_vda.GetMaxNorm()
120
121        # Get relevant VTK unstructured grid
122        vtk_ugd = self.get_unstructured_grid()
123
124        self.velocity = vtkGlyphModule(vtk_ugd, vtkVelocitySource(vmax, self.scale),
125                                       scalemode='vector', colormode=None)
126        self.add_module('velocity', self.velocity)
127