def add_forces(self):
"""Add force vectors for the atoms using atoms.get_forces(). A
``vtkGlyphModule`` is added to the module anchor under ``force``."""
if self.has_forces():
raise RuntimeError('Forces already present.')
elif self.has_velocities():
raise NotImplementedError('Can\'t add forces due to velocities.')
# Add forces to VTK unstructured grid as vector data
vtk_fda = self.add_vector_property(self.atoms.get_forces(), 'force')
# Calculate max norm of the forces
fmax = vtk_fda.GetMaxNorm()
# Get relevant VTK unstructured grid
vtk_ugd = self.get_unstructured_grid()
self.force = vtkGlyphModule(vtk_ugd, vtkForceSource(fmax, self.scale),
scalemode='vector', colormode=None)
self.add_module('force', self.force)
def __init__(self, atoms, scale=1):
"""Construct a fundamental VTK-representation of atoms.
atoms: Atoms object or list of Atoms objects
The atoms to be plotted.
scale = 1: float or int
Relative scaling of all Atoms-specific visualization.
"""
assert isinstance(atoms, Atoms)
self.atoms = atoms
self.scale = scale
vtkModuleAnchor.__init__(self)
vtkAtomicPositions.__init__(self, self.atoms.get_positions(),
vtkUnitCellModule(self.atoms))
self.force = None
self.velocity = None
symbols = self.atoms.get_chemical_symbols()
for symbol in np.unique(symbols):
# Construct mask for all atoms with this symbol
mask = np.array(symbols) == symbol
if mask.all():
subset = None
else:
subset = np.argwhere(mask).ravel()
# Get relevant VTK unstructured grid
vtk_ugd = self.get_unstructured_grid(subset)
# Create atomic glyph source for this symbol
glyph_source = vtkAtomSource(symbol, self.scale)
# Create glyph module and anchor it
self.add_module(symbol, vtkGlyphModule(vtk_ugd, glyph_source))
