/ase/visualize/vtk/atoms.py

https://gitlab.com/vote539/ase · Python · 127 lines · 57 code · 27 blank · 43 comment · 7 complexity · 7875a488bdca606d85328cdcccebe476 MD5 · raw file

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