from dataclasses import dataclass
from typing import Optional, Sequence
import numpy as np
from numpy.typing import NDArray
per_atom_properties = ('positions', 'velocities', 'forces')
"""Names of the per-atom properties that a trajectory frame can carry, in the order in
which a frame reports them."""
@dataclass
class ReaderFrame:
"""Trivial data struct holding MD-data for one time frame.
Parameters
----------
frame_index
Trajectory index of the snapshot (frame).
cell
Simulation cell as 3 row vectors (Å).
n_atoms
Number of atoms.
positions
Particle positions as an array with shape ``(n_atoms, 3)`` (Å).
velocities
Particle velocities as an array with shape ``(n_atoms, 3)`` (Å/fs);
may not be available, depending on reader and trajectory file format.
forces
Forces on the particles as an array with shape ``(n_atoms, 3)`` (eV/Å);
may not be available, depending on reader and trajectory file format.
atom_types
Array with the type of each atom;
may not be available, depending on reader and trajectory file format.
"""
frame_index: int
cell: NDArray[float]
n_atoms: int
positions: NDArray[float]
velocities: Optional[NDArray[float]] = None
forces: Optional[NDArray[float]] = None
atom_types: Optional[NDArray[str]] = None
[docs]
class TrajectoryFrame:
"""
Class holding positions and optionally velocities and forces, split by atom type,
for one snapshot (frame) in a trajectory.
Each quantity is provided as a dictionary keyed by atom type, such that
`positions_by_type['Cs']` is a numpy array with shape `(n_atoms_Cs, 3)` and
`positions_by_type['Pb']` is a numpy array with shape `(n_atoms_Pb, 3)`.
The `get_*_as_array` methods return a quantity for all atoms in a single array
ordered by atom index.
A quantity that the frame does not carry is reported as `None`.
Parameters
----------
atomic_indices
Dictionary specifying which indices (values) belong to which atom type (keys).
frame_index
Trajectory index of the snapshot (frame).
positions
Positions as an array with shape `(n_atoms, 3)`.
velocities
Velocities as an array with shape `(n_atoms, 3)`; defaults to `None`.
forces
Forces as an array with shape `(n_atoms, 3)`; defaults to `None`.
properties
Names of the per-atom properties to keep, a subset of
:data:`per_atom_properties`; defaults to `None`, which keeps every quantity that
is provided.
"""
def __init__(self,
atomic_indices: dict[str, list[int]],
frame_index: int,
positions: Optional[NDArray[float]] = None,
velocities: Optional[NDArray[float]] = None,
forces: Optional[NDArray[float]] = None,
properties: Optional[Sequence[str]] = None):
self._frame_index = frame_index
# Indexing along the atom axis only puts no constraint on the shape of a per-atom
# quantity, and already returns a copy rather than a view of the reader array.
arrays = dict(positions=positions, velocities=velocities, forces=forces)
self._arrays_by_type = dict()
for name in per_atom_properties:
array = arrays[name]
if array is None:
continue
if properties is not None and name not in properties:
continue
self._arrays_by_type[name] = {
atom_type: np.asarray(array)[indices]
for atom_type, indices in atomic_indices.items()}
@property
def positions_by_type(self) -> Optional[dict[str, NDArray[float]]]:
""" Positions split by atom type; ``None`` if the frame carries no positions. """
return self._arrays_by_type.get('positions')
@property
def velocities_by_type(self) -> Optional[dict[str, NDArray[float]]]:
""" Velocities split by atom type; ``None`` if the frame carries no velocities. """
return self._arrays_by_type.get('velocities')
@property
def forces_by_type(self) -> Optional[dict[str, NDArray[float]]]:
""" Forces split by atom type; ``None`` if the frame carries no forces. """
return self._arrays_by_type.get('forces')
def _get_as_array(self,
name: str,
atomic_indices: dict[str, list[int]]) -> NDArray[float]:
"""
Return the per-atom quantity :attr:`name` reassembled into a single array whose
first dimension runs over all atoms.
Parameters
----------
name
Name of the quantity, one of ``'positions'``, ``'velocities'``, ``'forces'``.
atomic_indices
Dictionary specifying which indices (values) belong to which atom type (keys).
"""
arrays_by_type = self._arrays_by_type.get(name)
if arrays_by_type is None:
raise ValueError(f'This frame provides no {name}.')
# check that atomic_indices is complete
n_atoms = np.max([np.max(indices) for indices in atomic_indices.values()]) + 1
all_inds = [i for indices in atomic_indices.values() for i in indices]
if len(all_inds) != n_atoms or len(set(all_inds)) != n_atoms:
raise ValueError('atomic_indices is incomplete')
# Collect the quantity into a single array. Only the shape is taken from the
# per-type arrays, so that the returned array is float64 whatever the precision the
# reader provides, which some of them read as float32.
reference = next(iter(arrays_by_type.values()))
array = np.empty((n_atoms, ) + reference.shape[1:])
for atom_type, indices in atomic_indices.items():
array[indices] = arrays_by_type[atom_type]
return array
[docs]
def get_positions_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
"""
Return the full positions array with shape ``(n_atoms, 3)``.
Parameters
----------
atomic_indices
Dictionary specifying which indices (values) belong to which atom type (keys).
"""
return self._get_as_array('positions', atomic_indices)
[docs]
def get_velocities_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
"""
Return the full velocities array with shape ``(n_atoms, 3)``.
Parameters
----------
atomic_indices
Dictionary specifying which indices (values) belong to which atom type (keys).
"""
return self._get_as_array('velocities', atomic_indices)
[docs]
def get_forces_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
"""
Return the full forces array with shape ``(n_atoms, 3)``.
Parameters
----------
atomic_indices
Dictionary specifying which indices (values) belong to which atom type (keys).
"""
return self._get_as_array('forces', atomic_indices)
@property
def frame_index(self) -> int:
""" Index of the frame. """
return self._frame_index
def __str__(self) -> str:
s = [f'Frame index {self.frame_index}']
for name, arrays_by_type in self._arrays_by_type.items():
for key, val in arrays_by_type.items():
s.append(f' {name:10} : {key} shape : {val.shape}')
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/Shape</th></tr></thead>']
s += ['<tbody>']
s += [f'<tr><td style="text-align: left;">Index</td><td>{self.frame_index}</td></tr>']
for name, arrays_by_type in self._arrays_by_type.items():
for key, val in arrays_by_type.items():
s += [f'<tr><td style="text-align: left;">{name.capitalize()} {key}</td>'
f'<td>{val.shape}</td></tr>']
s += ['</tbody>']
s += ['</table>']
return '\n'.join(s)
