Coverage for dynasor/trajectory/trajectory.py: 100%
202 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19: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.trajectory_frame import per_atom_properties, TrajectoryFrame
17from dynasor.logging_tools import logger
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 <dynasor.trajectory.trajectory_frame.TrajectoryFrame>`
25 object. The latter hold information regarding atomic positions, types, velocities,
26 and forces.
28 Parameters
29 ----------
30 filename
31 Name of input file.
32 trajectory_format
33 Type of trajectory. Possible values are:
34 ``'lammps_internal'``, ``'extxyz'``, ``'ase'`` or one of the formats supported by
35 `MDAnalysis <https://www.mdanalysis.org/>`_ (except for ``'lammpsdump'``,
36 which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity).
37 A format handled by MDAnalysis must be a name MDAnalysis knows, for example
38 ``'nc'`` or ``'NCDF'`` for an AMBER netCDF trajectory. Case does not matter.
39 If ``None``, MDAnalysis picks a reader based on the file extension. Name the
40 format for an extended xyz file, since the xyz reader of MDAnalysis ignores the
41 ``Lattice`` field that dynasor needs.
42 atomic_indices
43 Specify which indices belong to which atom type. Can be
44 (1) a dictionary where the keys specify the species and the values
45 are a list of atomic indices,
46 (2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or
47 (3) the path to a gromacs index file.
48 length_unit
49 Length unit of trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). Necessary for
50 correct conversion to internal dynasor units if the trajectory file does not contain unit
51 information.
52 If no length unit is specified and the reader cannot read units from the trajectory,
53 Angstrom is assumed.
54 time_unit
55 Time unit of trajectory (``'fs'``, ``'ps'``, ``'ns'``). Necessary for correct conversion to
56 internal dynasor units if the trajectory file does not contain unit information.
57 If no time unit is specified and the reader cannot read units from the trajectory, fs is
58 assumed.
59 frame_start
60 First frame to read; must be larger or equal ``0``.
61 frame_stop
62 Exclusive upper bound on the frame indices to read. For example, ``frame_stop=3`` reads
63 frames 0, 1, and 2 when using the default ``frame_start`` and ``frame_step``. By default
64 (``None``), the entire trajectory is read.
65 frame_step
66 Read every :attr:`frame_step`-th step of the input trajectory.
67 By default (``1``) every frame is read. Must be larger than ``0``.
68 force_unit
69 Force unit of trajectory (``'eV/Angstrom'``, ``'eV/nm'``, ``'kJ/mol/Angstrom'``,
70 ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, ``'Hartree/Bohr'``).
71 Necessary for correct conversion to internal dynasor units if the trajectory file does
72 not contain unit information. If no force unit is specified and the reader cannot read
73 units from the trajectory, eV/Angstrom is assumed. A unit dynasor cannot convert leaves
74 the forces unconverted, with a warning.
75 properties
76 Names of the per-atom properties to read, any of ``'positions'``, ``'velocities'``,
77 and ``'forces'``, given as a sequence or, for a single one, as a string. A property
78 that is not read is never split by atom type and hence costs no memory, which is
79 worth doing for a large system, since a window of a dynamic structure factor
80 calculation holds ``window_size + 1`` frames at a time. Note what the analysis
81 functions require: the structure factors need positions, currents need velocities in
82 addition, and the spectral energy density needs velocities alone.
83 The first frame of the trajectory decides what is read, so it provides the default
84 and it is what a request is checked against.
86 """
87 def __init__(
88 self,
89 filename: str,
90 trajectory_format: Optional[str] = None,
91 atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None,
92 length_unit: Optional[str] = None,
93 time_unit: Optional[str] = None,
94 frame_start: Optional[int] = 0,
95 frame_stop: Optional[int] = None,
96 frame_step: Optional[int] = 1,
97 force_unit: Optional[str] = None,
98 properties: Optional[Union[str, Sequence[str]]] = None
99 ):
101 if frame_start < 0:
102 raise ValueError('frame_start should be positive')
103 if frame_stop is not None and frame_stop < 0:
104 raise ValueError('frame_stop should be non-negative')
105 if frame_stop is not None and frame_stop <= frame_start:
106 raise ValueError('frame_stop should be larger than frame_start')
107 if frame_step <= 0:
108 raise ValueError('frame_step should be positive')
110 self._frame_start = frame_start
111 self._frame_step = frame_step
112 self._frame_stop = frame_stop
114 # setup trajectory reader
115 if not isfile(filename):
116 raise IOError(f'File {filename} does not exist')
117 self._filename = filename
119 if trajectory_format == 'lammps_internal':
120 reader = LammpsTrajectoryReader
121 elif trajectory_format == 'extxyz':
122 reader = ExtxyzTrajectoryReader
123 elif trajectory_format == 'lammps_mdanalysis':
124 reader = MDAnalysisTrajectoryReader
125 trajectory_format = 'lammpsdump'
126 elif trajectory_format == 'ase':
127 reader = ASETrajectoryReader
128 elif trajectory_format == 'lammps':
129 raise IOError('Ambiguous trajectory format, '
130 'did you mean lammps_internal or lammps_mdanalysis?')
131 else:
132 reader = MDAnalysisTrajectoryReader
134 logger.debug(f'Using trajectory reader: {reader.__name__}')
135 if reader == MDAnalysisTrajectoryReader:
136 self._reader_obj = reader(self._filename, trajectory_format,
137 length_unit=length_unit, time_unit=time_unit,
138 force_unit=force_unit)
139 else:
140 self._reader_obj = reader(self._filename, length_unit=length_unit,
141 time_unit=time_unit, force_unit=force_unit)
143 # Get two frames to set cell etc.
144 frame0 = next(self._reader_obj)
145 frame1 = next(self._reader_obj)
146 self._cell = frame0.cell
147 self._n_atoms = frame0.n_atoms
148 available = tuple(name for name in per_atom_properties
149 if getattr(frame0, name) is not None)
150 self._properties = self._resolve_properties(properties, available)
152 # Make sure cell is not changed during consecutive frames
153 if not np.allclose(frame0.cell, frame1.cell):
154 raise ValueError('The cell changes between the first and second frame. '
155 'The concept of q-points becomes muddy if the simulation cell is '
156 'changing, such as during NPT MD simulations, so trajectories where '
157 'the cell changes are not supported by dynasor.')
159 # setup iterator slice (reuse frame0 and frame1 via chain)
160 self.number_of_frames_read = 0
161 self.current_frame_index = 0
162 self._reader_iter = islice(chain([frame0, frame1], self._reader_obj),
163 self._frame_start, self._frame_stop, self._frame_step)
165 # setup atomic indices
166 if atomic_indices is None: # Default behaviour
167 atomic_indices = {'X': np.arange(0, self.n_atoms)}
168 elif isinstance(atomic_indices, str): # Str input
169 if atomic_indices == 'read_from_trajectory':
170 if frame0.atom_types is None:
171 raise ValueError('Could not read atomic indices from the trajectory.')
172 else:
173 uniques = np.unique(frame0.atom_types)
174 atomic_indices = {str(uniques[i]):
175 (frame0.atom_types == uniques[i]).nonzero()[0]
176 for i in range(len(uniques))}
177 else:
178 atomic_indices = parse_gromacs_index_file(atomic_indices)
179 elif isinstance(atomic_indices, dict): # dict input
180 pass
181 else:
182 raise ValueError('Could not understand atomic_indices.')
183 self._atomic_indices = atomic_indices
185 # sanity checks for atomic_indices
186 if len(self._atomic_indices) == 0:
187 raise ValueError('atomic_indices does not contain any atom types.')
188 for key, indices in self._atomic_indices.items():
189 if len(indices) == 0:
190 raise ValueError(f'No indices in atomic_indices for atom type {key}.')
191 if np.max(indices) >= self.n_atoms:
192 raise ValueError('Maximum index in atomic_indices exceeds number of atoms.')
193 if np.min(indices) < 0:
194 raise ValueError('Minimum index in atomic_indices is negative.')
195 if '_' in key:
196 # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb
197 raise ValueError('The char "_" is not allowed in atomic_indices.')
199 # log info on trajectory and atom types etc
200 logger.info(f'Trajectory file: {self.filename}')
201 logger.info(f'Total number of particles: {self.n_atoms}')
202 logger.info(f'Number of atom types: {len(self.atom_types)}')
203 for atom_type, indices in self._atomic_indices.items():
204 logger.info(f'Number of atoms of type {atom_type}: {len(indices)}')
205 logger.info(f'Per-atom properties read: {", ".join(self.properties)}')
206 logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}')
207 if self.has_forces:
208 if self._reader_obj.force_unit is None:
209 logger.warning('The trajectory provides forces in a unit that dynasor cannot'
210 ' convert, so they are passed on unconverted. Specify force_unit'
211 ' to have them converted to eV/Angstrom.')
212 else:
213 logger.info(f'Reading forces in {self._reader_obj.force_unit}, converted to'
214 ' eV/Angstrom.')
216 @staticmethod
217 def _resolve_properties(properties, available) -> tuple[str, ...]:
218 """
219 Return the per-atom properties to read, in the order of
220 :data:`per_atom_properties`.
222 Parameters
223 ----------
224 properties
225 The requested properties, or ``None`` for the properties that the first frame
226 provides.
227 available
228 The properties that the first frame of the trajectory provides, which is both
229 the default and what a request is checked against.
230 """
231 if properties is None:
232 return available
234 if isinstance(properties, str):
235 properties = (properties, )
236 properties = tuple(properties)
238 if len(properties) == 0:
239 raise ValueError('properties must name at least one per-atom property; the '
240 f'available options are {per_atom_properties}.')
241 for name in properties:
242 if name not in per_atom_properties:
243 raise ValueError(f'{name} is not a per-atom property; the options are '
244 f'{per_atom_properties}.')
245 if name not in available:
246 raise ValueError(f'The trajectory does not provide {name}.')
247 return tuple(name for name in per_atom_properties if name in properties)
249 def __iter__(self):
250 return self
252 def __next__(self):
253 frame = next(self._reader_iter)
254 logger.debug(f'Read frame #{frame.frame_index}')
255 # TrajectoryFrame only reads the indices, so avoid the per-frame copy made by the property
256 new_frame = TrajectoryFrame(
257 self._atomic_indices, frame.frame_index, frame.positions, frame.velocities,
258 frame.forces, properties=self._properties)
259 self.number_of_frames_read += 1
260 self.current_frame_index = frame.frame_index
261 return new_frame
263 def __str__(self) -> str:
264 s = ['Trajectory']
265 s += ['{:12} : {}'.format('filename', self.filename)]
266 s += ['{:12} : {}'.format('natoms', self.n_atoms)]
267 s += ['{:12} : {}'.format('frame_start', self._frame_start)]
268 s += ['{:12} : {}'.format('frame_stop', self._frame_stop)]
269 s += ['{:12} : {}'.format('frame_step', self.frame_step)]
270 s += ['{:12} : {}'.format('frame_index', self.current_frame_index)]
271 s += ['{:12} : [{}\n {}\n {}]'
272 .format('cell', self.cell[0], self.cell[1], self.cell[2])]
273 return '\n'.join(s)
275 def __repr__(self) -> str:
276 return str(self)
278 def _repr_html_(self) -> str:
279 s = [f'<h3>{self.__class__.__name__}</h3>']
280 s += ['<table border="1" class="dataframe">']
281 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>']
282 s += ['<tbody>']
283 s += [f'<tr><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>']
284 s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>']
285 s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>']
286 s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>']
287 s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>']
288 s += ['</tbody>']
289 s += ['</table>']
290 return '\n'.join(s)
292 @property
293 def cell(self) -> NDArray[float]:
294 """ Simulation cell """
295 return self._cell
297 @property
298 def n_atoms(self) -> int:
299 """ Number of atoms """
300 return self._n_atoms
302 @property
303 def has_velocities(self) -> bool:
304 """ Whether the trajectory provides velocities """
305 return 'velocities' in self.properties
307 @property
308 def has_positions(self) -> bool:
309 """ Whether the trajectory provides positions """
310 return 'positions' in self.properties
312 @property
313 def has_forces(self) -> bool:
314 """ Whether the trajectory provides forces """
315 return 'forces' in self.properties
317 @property
318 def properties(self) -> tuple[str, ...]:
319 """ Per-atom properties that are read from the trajectory """
320 return self._properties
322 @property
323 def filename(self) -> str:
324 """ The trajectory filename """
325 return self._filename
327 @property
328 def atomic_indices(self) -> dict[str, list[int]]:
329 """ Return copy of index arrays """
330 atomic_indices = dict()
331 for name, inds in self._atomic_indices.items():
332 atomic_indices[name] = inds.copy()
333 return atomic_indices
335 @property
336 def atom_types(self) -> list[str]:
337 return sorted(self._atomic_indices.keys())
339 @property
340 def frame_step(self) -> int:
341 """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot. """
342 return self._frame_step
345def consume(iterator, n):
346 """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """
347 # From python.org
348 if n is None:
349 deque(iterator, maxlen=0)
350 else:
351 next(islice(iterator, n, n), None)
354class WindowIterator:
355 """Sliding window iterator.
357 Returns consecutive windows (a window is represented as a list
358 of objects), created from an input iterator.
360 Parameters
361 ----------
362 itraj
363 Trajectory object.
364 width
365 Length of window (``window_size`` + 1).
366 window_step
367 Distance between the start of two consecutive window frames.
368 element_processor
369 Optional function applied to each frame before it is stored in the window.
370 Useful for pre-computing per-frame quantities (e.g., reciprocal-space densities)
371 so that the work is not repeated for frames shared between consecutive windows.
372 """
373 def __init__(self,
374 itraj: Trajectory,
375 width: int,
376 window_step: Optional[int] = 1,
377 element_processor: Optional[Callable] = None):
379 self._raw_it = itraj
380 if element_processor:
381 self._it = map(element_processor, self._raw_it)
382 else:
383 self._it = self._raw_it
384 assert window_step >= 1
385 assert width >= 1
386 self.width = width
387 self.window_step = window_step
388 self._window = None
390 def __iter__(self):
391 return self
393 def __next__(self):
394 """ Returns next element in sequence. """
395 if self._window is None:
396 self._window = deque(islice(self._it, self.width), self.width)
397 else:
398 if self.window_step >= self.width:
399 self._window.clear()
400 consume(self._raw_it, self.window_step - self.width)
401 else:
402 for _ in range(min((self.window_step, len(self._window)))):
403 self._window.popleft()
404 for f in islice(self._it, min((self.window_step, self.width))):
405 self._window.append(f)
407 if len(self._window) == 0:
408 raise StopIteration
410 return list(self._window)