__all__ = ['Trajectory', 'WindowIterator'] import numpy as np from collections import deque from itertools import islice, chain from os.path import isfile from typing import Callable, Dict, Union, List from dynasor.trajectory.atomic_indices import parse_gromacs_index_file from dynasor.trajectory.ase_trajectory_reader import ASETrajectoryReader from dynasor.trajectory.extxyz_trajectory_reader import ExtxyzTrajectoryReader from dynasor.trajectory.lammps_trajectory_reader import LammpsTrajectoryReader from dynasor.trajectory.mdanalysis_trajectory_reader import MDAnalysisTrajectoryReader from dynasor.trajectory.trajectory_frame import TrajectoryFrame from dynasor.logging_tools import logger class Trajectory: """Instances of this class hold trajectories in a format suitable for the computation of correlation functions. They behave as iterators, where each step returns the next frame as a :class:`TrajectoryFrame` object. The latter hold information regarding atomic positions, types, and velocities. Parameters ---------- filename Name of input file trajectory_format Type of trajectory. Possible values are: ``'lammps_internal'``, ``'extxyz'``, ``'ase'`` or one of the formats supported by `MDAnalysis `_ (except for ``'lammpsdump'``, which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity) atomic_indices Specify which indices belong to which atom type. Can be (1) a dictionary where the keys specicy species and the values are list of atomic indices, (2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or (3) the path to an index file. length_unit Length unit of trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). Necessary for correct conversion to internal dynasor units if the trajectory file does not contain unit information. time_unit Time unit of trajectory (``'fs'``, ``'ps'``, ``'ns'``). Necessary for correct conversion to internal dynasor units if the trajectory file does not contain unit information. frame_start First frame to read; must be larger or equal ``0``. frame_stop Last frame to read. By default (``None``) the entire trajectory is read. frame_step Read every :attr:`frame_step`-th step of the input trajectory. By default (``1``) every frame is read. Must be larger than ``0``. """ def __init__( self, filename: str, trajectory_format: str, atomic_indices: Union[str, Dict[str, List[int]]] = None, length_unit: str = 'Angstrom', time_unit: str = 'fs', frame_start: int = 0, frame_stop: int = None, frame_step: int = 1 ): if frame_start < 0: raise ValueError('frame_start should be positive') if frame_step < 0: raise ValueError('frame_step should be positive') self._frame_start = frame_start self._frame_step = frame_step self._frame_stop = frame_stop # setup trajectory reader if not isfile(filename): raise IOError(f'File {filename} does not exist') self._filename = filename if trajectory_format == 'lammps_internal': reader = LammpsTrajectoryReader elif trajectory_format == 'extxyz': reader = ExtxyzTrajectoryReader elif trajectory_format == 'lammps_mdanalysis': reader = MDAnalysisTrajectoryReader trajectory_format = 'lammpsdump' elif trajectory_format == 'ase': reader = ASETrajectoryReader elif trajectory_format == 'lammps': raise IOError('Ambiguous trajectory format, ' 'did you mean lammps_internal or lammps_mdanalysis?') else: reader = MDAnalysisTrajectoryReader logger.debug(f'Using trajectory reader: {reader.__name__}') if reader == MDAnalysisTrajectoryReader: self._reader_obj = reader(self._filename, trajectory_format, length_unit=length_unit, time_unit=time_unit) else: self._reader_obj = reader(self._filename, length_unit=length_unit, time_unit=time_unit) # Get two frames to set cell etc. frame0 = next(self._reader_obj) frame1 = next(self._reader_obj) self._cell = frame0.cell self._n_atoms = frame0.n_atoms # Make sure cell is not changed during consecutive frames if not np.allclose(frame0.cell, frame1.cell): raise ValueError('The cell changes between the first and second frame. ' 'The concept of q-points becomes muddy if the simulation cell is ' 'changing, such as during NPT MD simulations, so trajectories where ' 'the cell changes are not supported by dynasor.') # setup iterator slice (reuse frame0 and frame1 via chain) self.number_of_frames_read = 0 self.current_frame_index = 0 self._reader_iter = islice(chain([frame0, frame1], self._reader_obj), self._frame_start, self._frame_stop, self._frame_step) # setup atomic indices if atomic_indices is None: # Default behaviour atomic_indices = {'X': np.arange(0, self.n_atoms)} elif isinstance(atomic_indices, str): # Str input if atomic_indices == 'read_from_trajectory': if frame0.atom_types is None: raise ValueError('Could not read atomic indices from the trajectory.') else: uniques = np.unique(frame0.atom_types) atomic_indices = {uniques[i]: (frame0.atom_types == uniques[i]).nonzero()[0] for i in range(len(uniques))} else: atomic_indices = parse_gromacs_index_file(atomic_indices) elif isinstance(atomic_indices, dict): # Dict input pass else: raise ValueError('Could not understand atomic_indices.') self._atomic_indices = atomic_indices # sanity checks for atomic_indices for key, indices in self._atomic_indices.items(): if np.max(indices) > self.n_atoms: raise ValueError('maximum index in atomic_indices exceeds number of atoms') if np.min(indices) < 0: raise ValueError('minimum index in atomic_indices is negative') if '_' in key: # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb raise ValueError('The char "_" is not allowed in atomic_indices.') # log info on trajectory and atom types etc logger.info(f'Trajectory file: {self.filename}') logger.info(f'Total number of particles: {self.n_atoms}') logger.info(f'Number of atom types: {len(self.atom_types)}') for atom_type, indices in self._atomic_indices.items(): logger.info(f'Number of atoms of type {atom_type}: {len(indices)}') logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}') def __iter__(self): return self def __next__(self): frame = next(self._reader_iter) new_frame = TrajectoryFrame(self.atomic_indices, frame.frame_index, frame.positions, frame.velocities) self.number_of_frames_read += 1 self.current_frame_index = frame.frame_index return new_frame def __str__(self) -> str: s = ['Trajectory'] s += ['{:12} : {}'.format('filename', self.filename)] s += ['{:12} : {}'.format('natoms', self.n_atoms)] s += ['{:12} : {}'.format('frame_start', self._frame_start)] s += ['{:12} : {}'.format('frame_stop', self._frame_stop)] s += ['{:12} : {}'.format('frame_step', self.frame_step)] s += ['{:12} : {}'.format('frame_index', self.current_frame_index)] s += ['{:12} : [{}\n {}\n {}]' .format('cell', self.cell[0], self.cell[1], self.cell[2])] return '\n'.join(s) def __repr__(self) -> str: return str(self) def _repr_html_(self) -> str: s = [f'

{self.__class__.__name__}

'] s += [''] s += [''] s += [''] s += [f''] s += [f''] s += [f''] s += [f''] s += [f''] s += [''] s += ['
FieldValue
File name{self.filename}
Number of atoms{self.n_atoms}
Cell metric{self.cell}
Frame step{self.frame_step}
Atom types{self.atom_types}
'] return '\n'.join(s) @property def cell(self): """ Simulation cell """ return self._cell @property def n_atoms(self): """ Number of atoms """ return self._n_atoms @property def filename(self): """ The trajectory filename """ return self._filename @property def atomic_indices(self): """ Return copy of index arrays """ atomic_indices = dict() for name, inds in self._atomic_indices.items(): atomic_indices[name] = inds.copy() return atomic_indices @property def atom_types(self) -> List[str]: return sorted(self._atomic_indices.keys()) @property def frame_step(self): """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot """ return self._frame_step def consume(iterator, n): """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """ # From the python.org if n is None: deque(iterator, maxlen=0) else: next(islice(iterator, n, n), None) class WindowIterator: """Sliding window iterator. Returns consecutive windows (a window is represented as a list of objects), created from an input iterator. Parameters ---------- itraj Trajectory object width Length of window (``window_size`` + 1) window_step Distance between the start of two consecutive window frames element_processor Enables processing each non-discarded object; useful if ``window_step > width`` and ``map_item`` is expensive (as compared to directly passing ``map(fun, itraj)`` as ``itraj``); if ``window_step < width``, you could as well directly pass ``map(fun, itraj)``. """ def __init__(self, itraj: Trajectory, width: int, window_step: int = 1, element_processor: Callable = None): self._raw_it = itraj if element_processor: self._it = map(element_processor, self._raw_it) else: self._it = self._raw_it assert window_step >= 1 assert width >= 1 self.width = width self.window_step = window_step self._window = None def __iter__(self): return self def __next__(self): """ Returns next element in sequence. """ if self._window is None: self._window = deque(islice(self._it, self.width), self.width) else: if self.window_step >= self.width: self._window.clear() consume(self._raw_it, self.window_step - self.width) else: for _ in range(min((self.window_step, len(self._window)))): self._window.popleft() for f in islice(self._it, min((self.window_step, self.width))): self._window.append(f) if len(self._window) == 0: raise StopIteration return list(self._window)