Coverage for dynasor/trajectory/trajectory.py: 100%
229 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-02 07:46 +0000
1__all__ = ['Trajectory', 'WindowIterator']
3from collections import deque
4from itertools import islice, chain
5from os.path import isfile
6from typing import Callable, Optional, Sequence, Union
8import numpy as np
9from numpy.typing import NDArray
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.ovito_trajectory_reader import OvitoTrajectoryReader
17from dynasor.trajectory.trajectory_frame import per_atom_properties, TrajectoryFrame
18from dynasor.logging_tools import logger
21class Trajectory:
22 """Instances of this class hold trajectories in a format suitable for
23 the computation of correlation functions. They behave as
24 iterators, where each step returns the next frame as a
25 :class:`TrajectoryFrame <dynasor.trajectory.trajectory_frame.TrajectoryFrame>`
26 object. The latter hold information regarding atomic positions, types, velocities,
27 and forces.
29 Parameters
30 ----------
31 filename
32 Name of input file.
33 trajectory_format
34 Type of trajectory. Possible values are:
35 ``'lammps_internal'``, ``'extxyz'``, ``'ase'``, ``'ovito'`` or one of the formats
36 supported by `MDAnalysis <https://www.mdanalysis.org/>`_ (except for ``'lammpsdump'``,
37 which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity).
38 ``'ovito'`` reads any format `OVITO <https://www.ovito.org/>`_ knows, with the
39 particles sorted by their identifier when the file stores one.
40 A format handled by MDAnalysis must be a name MDAnalysis knows, for example
41 ``'nc'`` or ``'NCDF'`` for an AMBER netCDF trajectory. Case does not matter.
42 If ``None``, MDAnalysis picks a reader based on the file extension. Name the
43 format for an extended xyz file, since the xyz reader of MDAnalysis ignores the
44 ``Lattice`` field that dynasor needs.
45 atomic_indices
46 Specify which indices belong to which atom type. Can be
47 (1) a dictionary where the keys specify the species and the values
48 are a list of atomic indices,
49 (2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or
50 (3) the path to a gromacs index file.
51 length_unit
52 Length unit of trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). Necessary for
53 correct conversion to internal dynasor units if the trajectory file does not contain unit
54 information.
55 If no length unit is specified and the reader cannot read units from the trajectory,
56 Angstrom is assumed.
57 time_unit
58 Time unit of trajectory (``'fs'``, ``'ps'``, ``'ns'``). Necessary for correct conversion to
59 internal dynasor units if the trajectory file does not contain unit information.
60 If no time unit is specified and the reader cannot read units from the trajectory, fs is
61 assumed.
62 frame_start
63 First frame to read; must be larger or equal ``0``.
64 frame_stop
65 Exclusive upper bound on the frame indices to read. For example, ``frame_stop=3`` reads
66 frames 0, 1, and 2 when using the default ``frame_start`` and ``frame_step``. By default
67 (``None``), the entire trajectory is read.
68 frame_step
69 Read every :attr:`frame_step`-th step of the input trajectory.
70 By default (``1``) every frame is read. Must be larger than ``0``.
71 force_unit
72 Force unit of trajectory (``'eV/Angstrom'``, ``'eV/nm'``, ``'kJ/mol/Angstrom'``,
73 ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, ``'Hartree/Bohr'``).
74 Necessary for correct conversion to internal dynasor units if the trajectory file does
75 not contain unit information. If no force unit is specified and the reader cannot read
76 units from the trajectory, eV/Angstrom is assumed. A unit dynasor cannot convert leaves
77 the forces unconverted, with a warning.
78 properties
79 Names of the per-atom properties to read, any of ``'positions'``, ``'velocities'``,
80 and ``'forces'``, given as a sequence or, for a single one, as a string. A property
81 that is not read is never split by atom type and hence costs no memory, which is
82 worth doing for a large system, since a window of a dynamic structure factor
83 calculation holds ``window_size + 1`` frames at a time. Note what the analysis
84 functions require: the structure factors need positions, currents need velocities in
85 addition, and the spectral energy density needs velocities alone.
86 The first frame of the trajectory decides what is read, so it provides the default
87 and it is what a request is checked against.
89 """
90 def __init__(
91 self,
92 filename: str,
93 trajectory_format: Optional[str] = None,
94 atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None,
95 length_unit: Optional[str] = None,
96 time_unit: Optional[str] = None,
97 frame_start: Optional[int] = 0,
98 frame_stop: Optional[int] = None,
99 frame_step: Optional[int] = 1,
100 force_unit: Optional[str] = None,
101 properties: Optional[Union[str, Sequence[str]]] = None
102 ):
104 if frame_start < 0:
105 raise ValueError('frame_start should be positive')
106 if frame_stop is not None and frame_stop < 0:
107 raise ValueError('frame_stop should be non-negative')
108 if frame_stop is not None and frame_stop <= frame_start:
109 raise ValueError('frame_stop should be larger than frame_start')
110 if frame_step <= 0:
111 raise ValueError('frame_step should be positive')
113 self._frame_start = frame_start
114 self._frame_step = frame_step
115 self._frame_stop = frame_stop
117 # setup trajectory reader
118 if not isfile(filename):
119 raise IOError(f'File {filename} does not exist')
120 self._filename = filename
122 if trajectory_format == 'lammps_internal':
123 reader = LammpsTrajectoryReader
124 elif trajectory_format == 'extxyz':
125 reader = ExtxyzTrajectoryReader
126 elif trajectory_format == 'lammps_mdanalysis':
127 reader = MDAnalysisTrajectoryReader
128 trajectory_format = 'lammpsdump'
129 elif trajectory_format == 'ase':
130 reader = ASETrajectoryReader
131 elif trajectory_format == 'ovito':
132 reader = OvitoTrajectoryReader
133 elif trajectory_format == 'lammps':
134 raise IOError('Ambiguous trajectory format, '
135 'did you mean lammps_internal or lammps_mdanalysis?')
136 else:
137 reader = MDAnalysisTrajectoryReader
139 logger.debug(f'Using trajectory reader: {reader.__name__}')
140 if reader == MDAnalysisTrajectoryReader:
141 self._reader_obj = reader(self._filename, trajectory_format,
142 length_unit=length_unit, time_unit=time_unit,
143 force_unit=force_unit)
144 else:
145 self._reader_obj = reader(self._filename, length_unit=length_unit,
146 time_unit=time_unit, force_unit=force_unit)
147 self._n_total_frames = self._reader_obj.number_of_total_frames
149 # Get two frames to set cell etc.
150 frame0 = next(self._reader_obj)
151 frame1 = next(self._reader_obj)
152 self._cell = frame0.cell
153 self._n_atoms = frame0.n_atoms
154 available = tuple(name for name in per_atom_properties
155 if getattr(frame0, name) is not None)
156 self._properties = self._resolve_properties(properties, available)
158 # Make sure cell is not changed during consecutive frames
159 if not np.allclose(frame0.cell, frame1.cell):
160 raise ValueError('The cell changes between the first and second frame. '
161 'The concept of q-points becomes muddy if the simulation cell is '
162 'changing, such as during NPT MD simulations, so trajectories where '
163 'the cell changes are not supported by dynasor.')
165 # setup iterator slice (reuse frame0 and frame1 via chain)
166 self.number_of_frames_read = 0
167 self.current_frame_index = 0
168 self._reader_iter = islice(chain([frame0, frame1], self._reader_obj),
169 self._frame_start, self._frame_stop, self._frame_step)
171 # setup atomic indices
172 if atomic_indices is None: # Default behaviour
173 atomic_indices = {'X': np.arange(0, self.n_atoms)}
174 elif isinstance(atomic_indices, str): # Str input
175 if atomic_indices == 'read_from_trajectory':
176 if frame0.atom_types is None:
177 raise ValueError('Could not read atomic indices from the trajectory.')
178 else:
179 uniques = np.unique(frame0.atom_types)
180 atomic_indices = {str(uniques[i]):
181 (frame0.atom_types == uniques[i]).nonzero()[0]
182 for i in range(len(uniques))}
183 else:
184 atomic_indices = parse_gromacs_index_file(atomic_indices)
185 elif isinstance(atomic_indices, dict): # dict input
186 pass
187 else:
188 raise ValueError('Could not understand atomic_indices.')
189 self._atomic_indices = atomic_indices
191 # sanity checks for atomic_indices
192 if len(self._atomic_indices) == 0:
193 raise ValueError('atomic_indices does not contain any atom types.')
194 for key, indices in self._atomic_indices.items():
195 if len(indices) == 0:
196 raise ValueError(f'No indices in atomic_indices for atom type {key}.')
197 if np.max(indices) >= self.n_atoms:
198 raise ValueError('Maximum index in atomic_indices exceeds number of atoms.')
199 if np.min(indices) < 0:
200 raise ValueError('Minimum index in atomic_indices is negative.')
201 if '_' in key:
202 # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb
203 raise ValueError('The char "_" is not allowed in atomic_indices.')
205 # log info on trajectory and atom types etc
206 logger.info(f'Trajectory file: {self.filename}')
207 logger.info(f'Total number of particles: {self.n_atoms}')
208 logger.info(f'Number of atom types: {len(self.atom_types)}')
209 for atom_type, indices in self._atomic_indices.items():
210 logger.info(f'Number of atoms of type {atom_type}: {len(indices)}')
211 logger.info(f'Per-atom properties read: {", ".join(self.properties)}')
212 logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}')
213 if self.has_forces:
214 if self._reader_obj.force_unit is None:
215 logger.warning('The trajectory provides forces in a unit that dynasor cannot'
216 ' convert, so they are passed on unconverted. Specify force_unit'
217 ' to have them converted to eV/Angstrom.')
218 else:
219 logger.info(f'Reading forces in {self._reader_obj.force_unit}, converted to'
220 ' eV/Angstrom.')
222 @staticmethod
223 def _resolve_properties(properties, available) -> tuple[str, ...]:
224 """
225 Return the per-atom properties to read, in the order of
226 :data:`per_atom_properties`.
228 Parameters
229 ----------
230 properties
231 The requested properties, or ``None`` for the properties that the first frame
232 provides.
233 available
234 The properties that the first frame of the trajectory provides, which is both
235 the default and what a request is checked against.
236 """
237 if properties is None:
238 return available
240 if isinstance(properties, str):
241 properties = (properties, )
242 properties = tuple(properties)
244 if len(properties) == 0:
245 raise ValueError('properties must name at least one per-atom property; the '
246 f'available options are {per_atom_properties}.')
247 for name in properties:
248 if name not in per_atom_properties:
249 raise ValueError(f'{name} is not a per-atom property; the options are '
250 f'{per_atom_properties}.')
251 if name not in available:
252 raise ValueError(f'The trajectory does not provide {name}.')
253 return tuple(name for name in per_atom_properties if name in properties)
255 def __iter__(self):
256 return self
258 def __next__(self):
259 try:
260 frame = next(self._reader_iter)
261 except StopIteration:
262 self.close()
263 raise
264 logger.debug(f'Read frame #{frame.frame_index}')
265 # TrajectoryFrame only reads the indices, so avoid the per-frame copy made by the property
266 new_frame = TrajectoryFrame(
267 self._atomic_indices, frame.frame_index, frame.positions, frame.velocities,
268 frame.forces, properties=self._properties)
269 self.number_of_frames_read += 1
270 self.current_frame_index = frame.frame_index
271 return new_frame
273 def close(self):
274 """Close the underlying trajectory reader. The trajectory yields no further frames."""
275 self._reader_obj.close()
276 self._reader_iter = iter(())
278 def __str__(self) -> str:
279 s = ['Trajectory']
280 s += ['{:12} : {}'.format('filename', self.filename)]
281 s += ['{:12} : {}'.format('natoms', self.n_atoms)]
282 s += ['{:12} : {}'.format('frame_start', self._frame_start)]
283 s += ['{:12} : {}'.format('frame_stop', self._frame_stop)]
284 s += ['{:12} : {}'.format('frame_step', self.frame_step)]
285 s += ['{:12} : {}'.format('frame_index', self.current_frame_index)]
286 s += ['{:12} : [{}\n {}\n {}]'
287 .format('cell', self.cell[0], self.cell[1], self.cell[2])]
288 return '\n'.join(s)
290 def __repr__(self) -> str:
291 return str(self)
293 def _repr_html_(self) -> str:
294 s = [f'<h3>{self.__class__.__name__}</h3>']
295 s += ['<table border="1" class="dataframe">']
296 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>']
297 s += ['<tbody>']
298 s += [f'<tr><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>']
299 s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>']
300 s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>']
301 s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>']
302 s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>']
303 s += ['</tbody>']
304 s += ['</table>']
305 return '\n'.join(s)
307 @property
308 def cell(self) -> NDArray[float]:
309 """ Simulation cell """
310 return self._cell
312 @property
313 def n_atoms(self) -> int:
314 """ Number of atoms """
315 return self._n_atoms
317 @property
318 def n_total_frames(self) -> Optional[int]:
319 """ Total number of frames in the trajectory file, whatever the frame selection, if
320 the reader knows it when the file is opened (the OVITO-based reader scans the file
321 for frames then); ``None`` otherwise. """
322 return self._n_total_frames
324 @property
325 def has_velocities(self) -> bool:
326 """ Whether the trajectory provides velocities """
327 return 'velocities' in self.properties
329 @property
330 def has_positions(self) -> bool:
331 """ Whether the trajectory provides positions """
332 return 'positions' in self.properties
334 @property
335 def has_forces(self) -> bool:
336 """ Whether the trajectory provides forces """
337 return 'forces' in self.properties
339 @property
340 def properties(self) -> tuple[str, ...]:
341 """ Per-atom properties that are read from the trajectory """
342 return self._properties
344 @property
345 def filename(self) -> str:
346 """ The trajectory filename """
347 return self._filename
349 @property
350 def atomic_indices(self) -> dict[str, list[int]]:
351 """ Return copy of index arrays """
352 atomic_indices = dict()
353 for name, inds in self._atomic_indices.items():
354 atomic_indices[name] = inds.copy()
355 return atomic_indices
357 @property
358 def atom_types(self) -> list[str]:
359 return sorted(self._atomic_indices.keys())
361 @property
362 def frame_step(self) -> int:
363 """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot. """
364 return self._frame_step
366 def _log_progress(self, frame_index: int, window_start: Optional[int] = None) -> None:
367 """ Log the progress of the pass at ``INFO`` level, for the frame with index
368 ``frame_index`` that was just processed.
370 Depending on what is known about the number of frames, this looks like one of::
372 Frame 1200 (1201 frames so far)
373 Frame 1200 (1201 of at most 2500 frames)
374 Frame 1200 (1201 of 2500 frames)
375 Frame 1200 (1201 of 2500 frames, 3000 in file)
377 The second form is used when only :attr:`frame_stop` bounds the pass, for a reader
378 that does not know the length of the file.
380 The counts refer to the frames that the frame selection yields, so with
381 ``frame_step=3`` the third form reads ``Frame 1200 (401 of 834 frames)``.
382 The last form is used when :attr:`frame_stop` is smaller than :attr:`n_total_frames`,
383 that is, only part of the file is processed.
385 Parameters
386 ----------
387 frame_index
388 Index of the frame just processed, as the frame reports it in
389 ``frame.frame_index``.
390 window_start
391 If given, logs the range from ``window_start`` to ``frame_index``
392 instead of a single frame, e.g. ``Processing window 40 to 50 (51 of 100
393 frames)``. Used when progress is tracked per sliding window (as in
394 :func:`compute_dynamic_structure_factors
395 <dynasor.compute_dynamic_structure_factors>`) rather than per frame.
396 """
397 n_read = (frame_index - self._frame_start) // self._frame_step + 1
398 bounds = [b for b in (self._frame_stop, self._n_total_frames) if b is not None]
399 if bounds:
400 n_frames = len(range(self._frame_start, min(bounds), self._frame_step))
401 # frame_stop alone may lie beyond a file whose length the reader does not know
402 bound = 'at most ' if self._n_total_frames is None else ''
403 count = f'{n_read} of {bound}{n_frames} frames'
404 if self._n_total_frames is not None and min(bounds) < self._n_total_frames:
405 count += f', {self._n_total_frames} in file'
406 else:
407 count = f'{n_read} frame{"" if n_read == 1 else "s"} so far'
408 if window_start is None:
409 logger.info(f'Frame {frame_index} ({count})')
410 else:
411 logger.info(f'Processing window {window_start} to {frame_index} ({count})')
414def consume(iterator, n):
415 """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """
416 # From python.org
417 if n is None:
418 deque(iterator, maxlen=0)
419 else:
420 next(islice(iterator, n, n), None)
423class WindowIterator:
424 """Sliding window iterator.
426 Returns consecutive windows (a window is represented as a list
427 of objects), created from an input iterator.
429 Parameters
430 ----------
431 itraj
432 Trajectory object.
433 width
434 Length of window (``window_size`` + 1).
435 window_step
436 Distance between the start of two consecutive window frames.
437 element_processor
438 Optional function applied to each frame before it is stored in the window.
439 Useful for pre-computing per-frame quantities (e.g., reciprocal-space densities)
440 so that the work is not repeated for frames shared between consecutive windows.
441 """
442 def __init__(self,
443 itraj: Trajectory,
444 width: int,
445 window_step: Optional[int] = 1,
446 element_processor: Optional[Callable] = None):
448 self._raw_it = itraj
449 if element_processor:
450 self._it = map(element_processor, self._raw_it)
451 else:
452 self._it = self._raw_it
453 assert window_step >= 1
454 assert width >= 1
455 self.width = width
456 self.window_step = window_step
457 self._window = None
459 def __iter__(self):
460 return self
462 def __next__(self):
463 """ Returns next element in sequence. """
464 if self._window is None:
465 self._window = deque(islice(self._it, self.width), self.width)
466 else:
467 if self.window_step >= self.width:
468 self._window.clear()
469 consume(self._raw_it, self.window_step - self.width)
470 else:
471 for _ in range(min((self.window_step, len(self._window)))):
472 self._window.popleft()
473 for f in islice(self._it, min((self.window_step, self.width))):
474 self._window.append(f)
476 if len(self._window) == 0:
477 raise StopIteration
479 return list(self._window)