Coverage for dynasor/trajectory/trajectory.py: 100%

1__all__ = ['Trajectory', 'WindowIterator'] 

2 

3from collections import deque 

4from itertools import islice, chain 

5from os.path import isfile 

6from typing import Callable, Optional, Union 

7 

8import numpy as np 

9from numpy.typing import NDArray 

10 

11from dynasor.trajectory.atomic_indices import parse_gromacs_index_file 

12from dynasor.trajectory.ase_trajectory_reader import ASETrajectoryReader 

13from dynasor.trajectory.extxyz_trajectory_reader import ExtxyzTrajectoryReader 

14from dynasor.trajectory.lammps_trajectory_reader import LammpsTrajectoryReader 

15from dynasor.trajectory.mdanalysis_trajectory_reader import MDAnalysisTrajectoryReader 

16from dynasor.trajectory.trajectory_frame import TrajectoryFrame 

17from dynasor.logging_tools import logger 

18 

19 

20class Trajectory: 

21 """Instances of this class hold trajectories in a format suitable for 

22 the computation of correlation functions. They behave as 

23 iterators, where each step returns the next frame as a 

24 :class:`TrajectoryFrame` object. The latter hold information 

25 regarding atomic positions, types, and velocities. 

26 

27 Parameters 

28 ---------- 

29 filename 

30 Name of input file. 

31 trajectory_format 

32 Type of trajectory. Possible values are: 

33 ``'lammps_internal'``, ``'extxyz'``, ``'ase'`` or one of the formats supported by 

34 `MDAnalysis <https://www.mdanalysis.org/>`_ (except for ``'lammpsdump'``, 

35 which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity) 

36 atomic_indices 

37 Specify which indices belong to which atom type. Can be 

38 (1) a dictionary where the keys specify the species and the values 

39 are a list of atomic indices, 

40 (2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or 

41 (3) the path to a gromacs index file. 

42 length_unit 

43 Length unit of trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). Necessary for 

44 correct conversion to internal dynasor units if the trajectory file does not contain unit 

45 information. 

46 If no length unit is specified and the reader cannot read units from the trajectory, 

47 Angstrom is assumed. 

48 time_unit 

49 Time unit of trajectory (``'fs'``, ``'ps'``, ``'ns'``). Necessary for correct conversion to 

50 internal dynasor units if the trajectory file does not contain unit information. 

51 If no time unit is specified and the reader cannot read units from the trajectory, fs is 

52 assumed. 

53 frame_start 

54 First frame to read; must be larger or equal ``0``. 

55 frame_stop 

56 Last frame to read. By default (``None``) the entire trajectory is read. 

57 frame_step 

58 Read every :attr:`frame_step`-th step of the input trajectory. 

59 By default (``1``) every frame is read. Must be larger than ``0``. 

60 

61 """ 

62 def __init__( 

63 self, 

64 filename: str, 

65 trajectory_format: str, 

66 atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None, 

67 length_unit: Optional[str] = None, 

68 time_unit: Optional[str] = None, 

69 frame_start: Optional[int] = 0, 

70 frame_stop: Optional[int] = None, 

71 frame_step: Optional[int] = 1 

72 ): 

73 

74 if frame_start < 0: 

75 raise ValueError('frame_start should be positive') 

76 if frame_step < 0: 

77 raise ValueError('frame_step should be positive') 

78 

79 self._frame_start = frame_start 

80 self._frame_step = frame_step 

81 self._frame_stop = frame_stop 

82 

83 # setup trajectory reader 

84 if not isfile(filename): 

85 raise IOError(f'File {filename} does not exist') 

86 self._filename = filename 

87 

88 if trajectory_format == 'lammps_internal': 

89 reader = LammpsTrajectoryReader 

90 elif trajectory_format == 'extxyz': 

91 reader = ExtxyzTrajectoryReader 

92 elif trajectory_format == 'lammps_mdanalysis': 

93 reader = MDAnalysisTrajectoryReader 

94 trajectory_format = 'lammpsdump' 

95 elif trajectory_format == 'ase': 

96 reader = ASETrajectoryReader 

97 elif trajectory_format == 'lammps': 

98 raise IOError('Ambiguous trajectory format, ' 

99 'did you mean lammps_internal or lammps_mdanalysis?') 

100 else: 

101 reader = MDAnalysisTrajectoryReader 

102 

103 logger.debug(f'Using trajectory reader: {reader.__name__}') 

104 if reader == MDAnalysisTrajectoryReader: 

105 self._reader_obj = reader(self._filename, trajectory_format, 

106 length_unit=length_unit, time_unit=time_unit) 

107 else: 

108 self._reader_obj = reader(self._filename, length_unit=length_unit, time_unit=time_unit) 

109 

110 # Get two frames to set cell etc. 

111 frame0 = next(self._reader_obj) 

112 frame1 = next(self._reader_obj) 

113 self._cell = frame0.cell 

114 self._n_atoms = frame0.n_atoms 

115 self._has_velocities = frame0.velocities is not None 

116 

117 # Make sure cell is not changed during consecutive frames 

118 if not np.allclose(frame0.cell, frame1.cell): 

119 raise ValueError('The cell changes between the first and second frame. ' 

120 'The concept of q-points becomes muddy if the simulation cell is ' 

121 'changing, such as during NPT MD simulations, so trajectories where ' 

122 'the cell changes are not supported by dynasor.') 

123 

124 # setup iterator slice (reuse frame0 and frame1 via chain) 

125 self.number_of_frames_read = 0 

126 self.current_frame_index = 0 

127 self._reader_iter = islice(chain([frame0, frame1], self._reader_obj), 

128 self._frame_start, self._frame_stop, self._frame_step) 

129 

130 # setup atomic indices 

131 if atomic_indices is None: # Default behaviour 

132 atomic_indices = {'X': np.arange(0, self.n_atoms)} 

133 elif isinstance(atomic_indices, str): # Str input 

134 if atomic_indices == 'read_from_trajectory': 

135 if frame0.atom_types is None: 

136 raise ValueError('Could not read atomic indices from the trajectory.') 

137 else: 

138 uniques = np.unique(frame0.atom_types) 

139 atomic_indices = {str(uniques[i]): 

140 (frame0.atom_types == uniques[i]).nonzero()[0] 

141 for i in range(len(uniques))} 

142 else: 

143 atomic_indices = parse_gromacs_index_file(atomic_indices) 

144 elif isinstance(atomic_indices, dict): # dict input 

145 pass 

146 else: 

147 raise ValueError('Could not understand atomic_indices.') 

148 self._atomic_indices = atomic_indices 

149 

150 # sanity checks for atomic_indices 

151 for key, indices in self._atomic_indices.items(): 

152 if np.max(indices) >= self.n_atoms: 

153 raise ValueError('Maximum index in atomic_indices exceeds number of atoms.') 

154 if np.min(indices) < 0: 

155 raise ValueError('Minimum index in atomic_indices is negative.') 

156 if '_' in key: 

157 # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb 

158 raise ValueError('The char "_" is not allowed in atomic_indices.') 

159 

160 # log info on trajectory and atom types etc 

161 logger.info(f'Trajectory file: {self.filename}') 

162 logger.info(f'Total number of particles: {self.n_atoms}') 

163 logger.info(f'Number of atom types: {len(self.atom_types)}') 

164 for atom_type, indices in self._atomic_indices.items(): 

165 logger.info(f'Number of atoms of type {atom_type}: {len(indices)}') 

166 logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}') 

167 

168 def __iter__(self): 

169 return self 

170 

171 def __next__(self): 

172 frame = next(self._reader_iter) 

173 logger.debug(f'Read frame #{frame.frame_index}') 

174 new_frame = TrajectoryFrame( 

175 self.atomic_indices, frame.frame_index, frame.positions, frame.velocities) 

176 self.number_of_frames_read += 1 

177 self.current_frame_index = frame.frame_index 

178 return new_frame 

179 

180 def __str__(self) -> str: 

181 s = ['Trajectory'] 

182 s += ['{:12} : {}'.format('filename', self.filename)] 

183 s += ['{:12} : {}'.format('natoms', self.n_atoms)] 

184 s += ['{:12} : {}'.format('frame_start', self._frame_start)] 

185 s += ['{:12} : {}'.format('frame_stop', self._frame_stop)] 

186 s += ['{:12} : {}'.format('frame_step', self.frame_step)] 

187 s += ['{:12} : {}'.format('frame_index', self.current_frame_index)] 

188 s += ['{:12} : [{}\n {}\n {}]' 

189 .format('cell', self.cell[0], self.cell[1], self.cell[2])] 

190 return '\n'.join(s) 

191 

192 def __repr__(self) -> str: 

193 return str(self) 

194 

195 def _repr_html_(self) -> str: 

196 s = [f'<h3>{self.__class__.__name__}</h3>'] 

197 s += ['<table border="1" class="dataframe">'] 

198 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>'] 

199 s += ['<tbody>'] 

200 s += [f'<tr"><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>'] 

201 s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>'] 

202 s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>'] 

203 s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>'] 

204 s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>'] 

205 s += ['</tbody>'] 

206 s += ['</table>'] 

207 return '\n'.join(s) 

208 

209 @property 

210 def cell(self) -> NDArray[float]: 

211 """ Simulation cell """ 

212 return self._cell 

213 

214 @property 

215 def n_atoms(self) -> int: 

216 """ Number of atoms """ 

217 return self._n_atoms 

218 

219 @property 

220 def has_velocities(self) -> bool: 

221 """ Whether the trajectory provides velocities """ 

222 return self._has_velocities 

223 

224 @property 

225 def filename(self) -> str: 

226 """ The trajectory filename """ 

227 return self._filename 

228 

229 @property 

230 def atomic_indices(self) -> dict[str, list[int]]: 

231 """ Return copy of index arrays """ 

232 atomic_indices = dict() 

233 for name, inds in self._atomic_indices.items(): 

234 atomic_indices[name] = inds.copy() 

235 return atomic_indices 

236 

237 @property 

238 def atom_types(self) -> list[str]: 

239 return sorted(self._atomic_indices.keys()) 

240 

241 @property 

242 def frame_step(self) -> int: 

243 """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot. """ 

244 return self._frame_step 

245 

246 

247def consume(iterator, n): 

248 """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """ 

249 # From python.org 

250 if n is None: 

251 deque(iterator, maxlen=0) 

252 else: 

253 next(islice(iterator, n, n), None) 

254 

255 

256class WindowIterator: 

257 """Sliding window iterator. 

258 

259 Returns consecutive windows (a window is represented as a list 

260 of objects), created from an input iterator. 

261 

262 Parameters 

263 ---------- 

264 itraj 

265 Trajectory object. 

266 width 

267 Length of window (``window_size`` + 1). 

268 window_step 

269 Distance between the start of two consecutive window frames. 

270 element_processor 

271 Optional function applied to each frame before it is stored in the window. 

272 Useful for pre-computing per-frame quantities (e.g., reciprocal-space densities) 

273 so that the work is not repeated for frames shared between consecutive windows. 

274 """ 

275 def __init__(self, 

276 itraj: Trajectory, 

277 width: int, 

278 window_step: Optional[int] = 1, 

279 element_processor: Optional[Callable] = None): 

280 

281 self._raw_it = itraj 

282 if element_processor: 

283 self._it = map(element_processor, self._raw_it) 

284 else: 

285 self._it = self._raw_it 

286 assert window_step >= 1 

287 assert width >= 1 

288 self.width = width 

289 self.window_step = window_step 

290 self._window = None 

291 

292 def __iter__(self): 

293 return self 

294 

295 def __next__(self): 

296 """ Returns next element in sequence. """ 

297 if self._window is None: 

298 self._window = deque(islice(self._it, self.width), self.width) 

299 else: 

300 if self.window_step >= self.width: 

301 self._window.clear() 

302 consume(self._raw_it, self.window_step - self.width) 

303 else: 

304 for _ in range(min((self.window_step, len(self._window)))): 

305 self._window.popleft() 

306 for f in islice(self._it, min((self.window_step, self.width))): 

307 self._window.append(f) 

308 

309 if len(self._window) == 0: 

310 raise StopIteration 

311 

312 return list(self._window)