Coverage for dynasor/trajectory/trajectory.py: 100%
172 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 08:27 +0000
1__all__ = ['Trajectory', 'WindowIterator']
3from collections import deque
4from itertools import islice, chain
5from os.path import isfile
6from typing import Callable, Optional, 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 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` object. The latter hold information
25 regarding atomic positions, types, and velocities.
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 Exclusive upper bound on the frame indices to read. For example, ``frame_stop=3`` reads
57 frames 0, 1, and 2 when using the default ``frame_start`` and ``frame_step``. By default
58 (``None``), the entire trajectory is read.
59 frame_step
60 Read every :attr:`frame_step`-th step of the input trajectory.
61 By default (``1``) every frame is read. Must be larger than ``0``.
63 """
64 def __init__(
65 self,
66 filename: str,
67 trajectory_format: str,
68 atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None,
69 length_unit: Optional[str] = None,
70 time_unit: Optional[str] = None,
71 frame_start: Optional[int] = 0,
72 frame_stop: Optional[int] = None,
73 frame_step: Optional[int] = 1
74 ):
76 if frame_start < 0:
77 raise ValueError('frame_start should be positive')
78 if frame_stop is not None and frame_stop < 0:
79 raise ValueError('frame_stop should be non-negative')
80 if frame_stop is not None and frame_stop <= frame_start:
81 raise ValueError('frame_stop should be larger than frame_start')
82 if frame_step <= 0:
83 raise ValueError('frame_step should be positive')
85 self._frame_start = frame_start
86 self._frame_step = frame_step
87 self._frame_stop = frame_stop
89 # setup trajectory reader
90 if not isfile(filename):
91 raise IOError(f'File {filename} does not exist')
92 self._filename = filename
94 if trajectory_format == 'lammps_internal':
95 reader = LammpsTrajectoryReader
96 elif trajectory_format == 'extxyz':
97 reader = ExtxyzTrajectoryReader
98 elif trajectory_format == 'lammps_mdanalysis':
99 reader = MDAnalysisTrajectoryReader
100 trajectory_format = 'lammpsdump'
101 elif trajectory_format == 'ase':
102 reader = ASETrajectoryReader
103 elif trajectory_format == 'lammps':
104 raise IOError('Ambiguous trajectory format, '
105 'did you mean lammps_internal or lammps_mdanalysis?')
106 else:
107 reader = MDAnalysisTrajectoryReader
109 logger.debug(f'Using trajectory reader: {reader.__name__}')
110 if reader == MDAnalysisTrajectoryReader:
111 self._reader_obj = reader(self._filename, trajectory_format,
112 length_unit=length_unit, time_unit=time_unit)
113 else:
114 self._reader_obj = reader(self._filename, length_unit=length_unit, time_unit=time_unit)
116 # Get two frames to set cell etc.
117 frame0 = next(self._reader_obj)
118 frame1 = next(self._reader_obj)
119 self._cell = frame0.cell
120 self._n_atoms = frame0.n_atoms
121 self._has_velocities = frame0.velocities is not None
123 # Make sure cell is not changed during consecutive frames
124 if not np.allclose(frame0.cell, frame1.cell):
125 raise ValueError('The cell changes between the first and second frame. '
126 'The concept of q-points becomes muddy if the simulation cell is '
127 'changing, such as during NPT MD simulations, so trajectories where '
128 'the cell changes are not supported by dynasor.')
130 # setup iterator slice (reuse frame0 and frame1 via chain)
131 self.number_of_frames_read = 0
132 self.current_frame_index = 0
133 self._reader_iter = islice(chain([frame0, frame1], self._reader_obj),
134 self._frame_start, self._frame_stop, self._frame_step)
136 # setup atomic indices
137 if atomic_indices is None: # Default behaviour
138 atomic_indices = {'X': np.arange(0, self.n_atoms)}
139 elif isinstance(atomic_indices, str): # Str input
140 if atomic_indices == 'read_from_trajectory':
141 if frame0.atom_types is None:
142 raise ValueError('Could not read atomic indices from the trajectory.')
143 else:
144 uniques = np.unique(frame0.atom_types)
145 atomic_indices = {str(uniques[i]):
146 (frame0.atom_types == uniques[i]).nonzero()[0]
147 for i in range(len(uniques))}
148 else:
149 atomic_indices = parse_gromacs_index_file(atomic_indices)
150 elif isinstance(atomic_indices, dict): # dict input
151 pass
152 else:
153 raise ValueError('Could not understand atomic_indices.')
154 self._atomic_indices = atomic_indices
156 # sanity checks for atomic_indices
157 if len(self._atomic_indices) == 0:
158 raise ValueError('atomic_indices does not contain any atom types.')
159 for key, indices in self._atomic_indices.items():
160 if len(indices) == 0:
161 raise ValueError(f'No indices in atomic_indices for atom type {key}.')
162 if np.max(indices) >= self.n_atoms:
163 raise ValueError('Maximum index in atomic_indices exceeds number of atoms.')
164 if np.min(indices) < 0:
165 raise ValueError('Minimum index in atomic_indices is negative.')
166 if '_' in key:
167 # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb
168 raise ValueError('The char "_" is not allowed in atomic_indices.')
170 # log info on trajectory and atom types etc
171 logger.info(f'Trajectory file: {self.filename}')
172 logger.info(f'Total number of particles: {self.n_atoms}')
173 logger.info(f'Number of atom types: {len(self.atom_types)}')
174 for atom_type, indices in self._atomic_indices.items():
175 logger.info(f'Number of atoms of type {atom_type}: {len(indices)}')
176 logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}')
178 def __iter__(self):
179 return self
181 def __next__(self):
182 frame = next(self._reader_iter)
183 logger.debug(f'Read frame #{frame.frame_index}')
184 # TrajectoryFrame only reads the indices, so avoid the per-frame copy made by the property
185 new_frame = TrajectoryFrame(
186 self._atomic_indices, frame.frame_index, frame.positions, frame.velocities)
187 self.number_of_frames_read += 1
188 self.current_frame_index = frame.frame_index
189 return new_frame
191 def __str__(self) -> str:
192 s = ['Trajectory']
193 s += ['{:12} : {}'.format('filename', self.filename)]
194 s += ['{:12} : {}'.format('natoms', self.n_atoms)]
195 s += ['{:12} : {}'.format('frame_start', self._frame_start)]
196 s += ['{:12} : {}'.format('frame_stop', self._frame_stop)]
197 s += ['{:12} : {}'.format('frame_step', self.frame_step)]
198 s += ['{:12} : {}'.format('frame_index', self.current_frame_index)]
199 s += ['{:12} : [{}\n {}\n {}]'
200 .format('cell', self.cell[0], self.cell[1], self.cell[2])]
201 return '\n'.join(s)
203 def __repr__(self) -> str:
204 return str(self)
206 def _repr_html_(self) -> str:
207 s = [f'<h3>{self.__class__.__name__}</h3>']
208 s += ['<table border="1" class="dataframe">']
209 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>']
210 s += ['<tbody>']
211 s += [f'<tr><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>']
212 s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>']
213 s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>']
214 s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>']
215 s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>']
216 s += ['</tbody>']
217 s += ['</table>']
218 return '\n'.join(s)
220 @property
221 def cell(self) -> NDArray[float]:
222 """ Simulation cell """
223 return self._cell
225 @property
226 def n_atoms(self) -> int:
227 """ Number of atoms """
228 return self._n_atoms
230 @property
231 def has_velocities(self) -> bool:
232 """ Whether the trajectory provides velocities """
233 return self._has_velocities
235 @property
236 def filename(self) -> str:
237 """ The trajectory filename """
238 return self._filename
240 @property
241 def atomic_indices(self) -> dict[str, list[int]]:
242 """ Return copy of index arrays """
243 atomic_indices = dict()
244 for name, inds in self._atomic_indices.items():
245 atomic_indices[name] = inds.copy()
246 return atomic_indices
248 @property
249 def atom_types(self) -> list[str]:
250 return sorted(self._atomic_indices.keys())
252 @property
253 def frame_step(self) -> int:
254 """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot. """
255 return self._frame_step
258def consume(iterator, n):
259 """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """
260 # From python.org
261 if n is None:
262 deque(iterator, maxlen=0)
263 else:
264 next(islice(iterator, n, n), None)
267class WindowIterator:
268 """Sliding window iterator.
270 Returns consecutive windows (a window is represented as a list
271 of objects), created from an input iterator.
273 Parameters
274 ----------
275 itraj
276 Trajectory object.
277 width
278 Length of window (``window_size`` + 1).
279 window_step
280 Distance between the start of two consecutive window frames.
281 element_processor
282 Optional function applied to each frame before it is stored in the window.
283 Useful for pre-computing per-frame quantities (e.g., reciprocal-space densities)
284 so that the work is not repeated for frames shared between consecutive windows.
285 """
286 def __init__(self,
287 itraj: Trajectory,
288 width: int,
289 window_step: Optional[int] = 1,
290 element_processor: Optional[Callable] = None):
292 self._raw_it = itraj
293 if element_processor:
294 self._it = map(element_processor, self._raw_it)
295 else:
296 self._it = self._raw_it
297 assert window_step >= 1
298 assert width >= 1
299 self.width = width
300 self.window_step = window_step
301 self._window = None
303 def __iter__(self):
304 return self
306 def __next__(self):
307 """ Returns next element in sequence. """
308 if self._window is None:
309 self._window = deque(islice(self._it, self.width), self.width)
310 else:
311 if self.window_step >= self.width:
312 self._window.clear()
313 consume(self._raw_it, self.window_step - self.width)
314 else:
315 for _ in range(min((self.window_step, len(self._window)))):
316 self._window.popleft()
317 for f in islice(self._it, min((self.window_step, self.width))):
318 self._window.append(f)
320 if len(self._window) == 0:
321 raise StopIteration
323 return list(self._window)