Coverage for dynasor/trajectory/ovito_trajectory_reader.py: 97%

1from typing import Optional 

2 

3import numpy as np 

4from numpy.typing import NDArray 

5 

6from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

7from dynasor.trajectory.trajectory_frame import ReaderFrame 

8 

9 

10class OvitoTrajectoryReader(AbstractTrajectoryReader): 

11 """Read trajectory file using OVITO. 

12 

13 Any format OVITO reads can be used. Particles are sorted by their identifier when the 

14 file stores one, type names are taken from the file (numeric LAMMPS types become 

15 ``'1'``, ``'2'``, ...), and velocities and forces are read when the file provides 

16 them. The number of frames is known once the file is opened. 

17 

18 Parameters 

19 ---------- 

20 filename 

21 Name of input file. 

22 length_unit 

23 Unit of length for the input trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). 

24 time_unit 

25 Unit of time for the input trajectory (``'fs'``, ``'ps'``, ``'ns'``). 

26 force_unit 

27 Unit of force for the input trajectory (``'eV/Angstrom'``, ``'eV/nm'``, 

28 ``'kJ/mol/Angstrom'``, ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, 

29 ``'Hartree/Bohr'``). 

30 """ 

31 

32 def __init__( 

33 self, 

34 filename: str, 

35 length_unit: Optional[str] = None, 

36 time_unit: Optional[str] = None, 

37 force_unit: Optional[str] = None 

38 ): 

39 # Imported here so that `import dynasor` does not need ovito and the graphics 

40 # libraries it loads 

41 from ovito.io import import_file 

42 

43 self._pipeline = import_file(filename, sort_particles=True) 

44 self._frame_index = 0 

45 self._open = True 

46 self.set_unit_scaling_factors(length_unit, time_unit, force_unit) 

47 

48 first = self._pipeline.compute(0) 

49 self._n_atoms = int(first.particles.count) 

50 self._number_of_total_frames = self._pipeline.source.num_frames 

51 self._atom_types = self._read_atom_types(first) 

52 self._velocity_property = self._find_property(first, ('Velocity', 'vel')) 

53 self._force_property = self._find_property(first, ('Force', 'forces')) 

54 

55 @staticmethod 

56 def _find_property(data, names: tuple[str, ...]) -> Optional[str]: 

57 """ Return the first of ``names`` that is a particle property of ``data``. """ 

58 return next((name for name in names if name in data.particles), None) 

59 

60 @staticmethod 

61 def _read_atom_types(data) -> Optional[NDArray[str]]: 

62 """ Type name of each particle in ``data``, or ``None`` without types. Read once, 

63 from the first frame, since the type of a particle does not change. """ 

64 if 'Particle Type' not in data.particles: 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true

65 return None 

66 type_property = data.particles['Particle Type'] 

67 # OVITO names numeric LAMMPS types '', so those fall back to the number. A dict 

68 # keyed by type id needs no table sized by the largest id, which a declared type 

69 # absent from the frame or a large id would overrun. 

70 names = {t.id: t.name or str(t.id) for t in type_property.types} 

71 ids, inverse = np.unique(type_property.array, return_inverse=True) 

72 return np.array([names.get(i, str(i)) for i in ids], dtype=str)[inverse] 

73 

74 def __iter__(self): 

75 return self 

76 

77 def close(self): 

78 """Release the OVITO pipeline and stop iteration.""" 

79 self._open = False 

80 self._pipeline = None 

81 

82 def __next__(self): 

83 if not self._open or self._frame_index >= self._number_of_total_frames: 

84 self.close() 

85 raise StopIteration 

86 frame_index = self._frame_index 

87 self._frame_index += 1 

88 

89 data = self._pipeline.compute(frame_index) 

90 if data.particles.count != self._n_atoms: 

91 raise RuntimeError('Number of atoms changed between frames.') 

92 velocities = forces = None 

93 if self._velocity_property is not None: 

94 velocities = self.v_factor * np.asarray(data.particles[self._velocity_property], 

95 dtype=float) 

96 if self._force_property is not None: 

97 forces = self.f_factor * np.asarray(data.particles[self._force_property], 

98 dtype=float) 

99 # OVITO stores the cell vectors as the first three columns of a 3x4 matrix; the 

100 # fourth column, the cell origin, is dropped 

101 cell = np.array(data.cell[:, :3], dtype=float).T 

102 return ReaderFrame(frame_index=frame_index, 

103 n_atoms=self._n_atoms, 

104 cell=self.x_factor * cell, 

105 positions=self.x_factor * np.asarray(data.particles['Position'], 

106 dtype=float), 

107 atom_types=self._atom_types, 

108 velocities=velocities, 

109 forces=forces)