__all__ = ['Trajectory', 'WindowIterator']
from collections import deque
from itertools import islice, chain
from os.path import isfile
from typing import Callable, Optional, Sequence, Union
import numpy as np
from numpy.typing import NDArray
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 per_atom_properties, TrajectoryFrame
from dynasor.logging_tools import logger
[docs]
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 <dynasor.trajectory.trajectory_frame.TrajectoryFrame>`
object. The latter hold information regarding atomic positions, types, velocities,
and forces.
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 <https://www.mdanalysis.org/>`_ (except for ``'lammpsdump'``,
which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity).
A format handled by MDAnalysis must be a name MDAnalysis knows, for example
``'nc'`` or ``'NCDF'`` for an AMBER netCDF trajectory. Case does not matter.
If ``None``, MDAnalysis picks a reader based on the file extension. Name the
format for an extended xyz file, since the xyz reader of MDAnalysis ignores the
``Lattice`` field that dynasor needs.
atomic_indices
Specify which indices belong to which atom type. Can be
(1) a dictionary where the keys specify the species and the values
are a list of atomic indices,
(2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or
(3) the path to a gromacs 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.
If no length unit is specified and the reader cannot read units from the trajectory,
Angstrom is assumed.
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.
If no time unit is specified and the reader cannot read units from the trajectory, fs is
assumed.
frame_start
First frame to read; must be larger or equal ``0``.
frame_stop
Exclusive upper bound on the frame indices to read. For example, ``frame_stop=3`` reads
frames 0, 1, and 2 when using the default ``frame_start`` and ``frame_step``. 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``.
force_unit
Force unit of trajectory (``'eV/Angstrom'``, ``'eV/nm'``, ``'kJ/mol/Angstrom'``,
``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, ``'Hartree/Bohr'``).
Necessary for correct conversion to internal dynasor units if the trajectory file does
not contain unit information. If no force unit is specified and the reader cannot read
units from the trajectory, eV/Angstrom is assumed. A unit dynasor cannot convert leaves
the forces unconverted, with a warning.
properties
Names of the per-atom properties to read, any of ``'positions'``, ``'velocities'``,
and ``'forces'``, given as a sequence or, for a single one, as a string. A property
that is not read is never split by atom type and hence costs no memory, which is
worth doing for a large system, since a window of a dynamic structure factor
calculation holds ``window_size + 1`` frames at a time. Note what the analysis
functions require: the structure factors need positions, currents need velocities in
addition, and the spectral energy density needs velocities alone.
The first frame of the trajectory decides what is read, so it provides the default
and it is what a request is checked against.
"""
def __init__(
self,
filename: str,
trajectory_format: Optional[str] = None,
atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None,
length_unit: Optional[str] = None,
time_unit: Optional[str] = None,
frame_start: Optional[int] = 0,
frame_stop: Optional[int] = None,
frame_step: Optional[int] = 1,
force_unit: Optional[str] = None,
properties: Optional[Union[str, Sequence[str]]] = None
):
if frame_start < 0:
raise ValueError('frame_start should be positive')
if frame_stop is not None and frame_stop < 0:
raise ValueError('frame_stop should be non-negative')
if frame_stop is not None and frame_stop <= frame_start:
raise ValueError('frame_stop should be larger than frame_start')
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,
force_unit=force_unit)
else:
self._reader_obj = reader(self._filename, length_unit=length_unit,
time_unit=time_unit, force_unit=force_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
available = tuple(name for name in per_atom_properties
if getattr(frame0, name) is not None)
self._properties = self._resolve_properties(properties, available)
# 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 = {str(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
if len(self._atomic_indices) == 0:
raise ValueError('atomic_indices does not contain any atom types.')
for key, indices in self._atomic_indices.items():
if len(indices) == 0:
raise ValueError(f'No indices in atomic_indices for atom type {key}.')
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'Per-atom properties read: {", ".join(self.properties)}')
logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}')
if self.has_forces:
if self._reader_obj.force_unit is None:
logger.warning('The trajectory provides forces in a unit that dynasor cannot'
' convert, so they are passed on unconverted. Specify force_unit'
' to have them converted to eV/Angstrom.')
else:
logger.info(f'Reading forces in {self._reader_obj.force_unit}, converted to'
' eV/Angstrom.')
@staticmethod
def _resolve_properties(properties, available) -> tuple[str, ...]:
"""
Return the per-atom properties to read, in the order of
:data:`per_atom_properties`.
Parameters
----------
properties
The requested properties, or ``None`` for the properties that the first frame
provides.
available
The properties that the first frame of the trajectory provides, which is both
the default and what a request is checked against.
"""
if properties is None:
return available
if isinstance(properties, str):
properties = (properties, )
properties = tuple(properties)
if len(properties) == 0:
raise ValueError('properties must name at least one per-atom property; the '
f'available options are {per_atom_properties}.')
for name in properties:
if name not in per_atom_properties:
raise ValueError(f'{name} is not a per-atom property; the options are '
f'{per_atom_properties}.')
if name not in available:
raise ValueError(f'The trajectory does not provide {name}.')
return tuple(name for name in per_atom_properties if name in properties)
def __iter__(self):
return self
def __next__(self):
frame = next(self._reader_iter)
logger.debug(f'Read frame #{frame.frame_index}')
# TrajectoryFrame only reads the indices, so avoid the per-frame copy made by the property
new_frame = TrajectoryFrame(
self._atomic_indices, frame.frame_index, frame.positions, frame.velocities,
frame.forces, properties=self._properties)
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'<h3>{self.__class__.__name__}</h3>']
s += ['<table border="1" class="dataframe">']
s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>']
s += ['<tbody>']
s += [f'<tr><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>']
s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>']
s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>']
s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>']
s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>']
s += ['</tbody>']
s += ['</table>']
return '\n'.join(s)
@property
def cell(self) -> NDArray[float]:
""" Simulation cell """
return self._cell
@property
def n_atoms(self) -> int:
""" Number of atoms """
return self._n_atoms
@property
def has_velocities(self) -> bool:
""" Whether the trajectory provides velocities """
return 'velocities' in self.properties
@property
def has_positions(self) -> bool:
""" Whether the trajectory provides positions """
return 'positions' in self.properties
@property
def has_forces(self) -> bool:
""" Whether the trajectory provides forces """
return 'forces' in self.properties
@property
def properties(self) -> tuple[str, ...]:
""" Per-atom properties that are read from the trajectory """
return self._properties
@property
def filename(self) -> str:
""" The trajectory filename """
return self._filename
@property
def atomic_indices(self) -> dict[str, list[int]]:
""" 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) -> int:
""" 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 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
Optional function applied to each frame before it is stored in the window.
Useful for pre-computing per-frame quantities (e.g., reciprocal-space densities)
so that the work is not repeated for frames shared between consecutive windows.
"""
def __init__(self,
itraj: Trajectory,
width: int,
window_step: Optional[int] = 1,
element_processor: Optional[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)
