Coverage for dynasor/trajectory/trajectory_frame.py: 100%
72 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
1from dataclasses import dataclass
2from typing import Optional, Sequence
3import numpy as np
4from numpy.typing import NDArray
7per_atom_properties = ('positions', 'velocities', 'forces')
8"""Names of the per-atom properties that a trajectory frame can carry, in the order in
9which a frame reports them."""
12@dataclass
13class ReaderFrame:
14 """Trivial data struct holding MD-data for one time frame.
16 Parameters
17 ----------
18 frame_index
19 Trajectory index of the snapshot (frame).
20 cell
21 Simulation cell as 3 row vectors (Å).
22 n_atoms
23 Number of atoms.
24 positions
25 Particle positions as an array with shape ``(n_atoms, 3)`` (Å).
26 velocities
27 Particle velocities as an array with shape ``(n_atoms, 3)`` (Å/fs);
28 may not be available, depending on reader and trajectory file format.
29 forces
30 Forces on the particles as an array with shape ``(n_atoms, 3)`` (eV/Å);
31 may not be available, depending on reader and trajectory file format.
32 atom_types
33 Array with the type of each atom;
34 may not be available, depending on reader and trajectory file format.
35 """
36 frame_index: int
37 cell: NDArray[float]
38 n_atoms: int
39 positions: NDArray[float]
40 velocities: Optional[NDArray[float]] = None
41 forces: Optional[NDArray[float]] = None
42 atom_types: Optional[NDArray[str]] = None
45class TrajectoryFrame:
46 """
47 Class holding positions and optionally velocities and forces, split by atom type,
48 for one snapshot (frame) in a trajectory.
50 Each quantity is provided as a dictionary keyed by atom type, such that
51 `positions_by_type['Cs']` is a numpy array with shape `(n_atoms_Cs, 3)` and
52 `positions_by_type['Pb']` is a numpy array with shape `(n_atoms_Pb, 3)`.
53 The `get_*_as_array` methods return a quantity for all atoms in a single array
54 ordered by atom index.
56 A quantity that the frame does not carry is reported as `None`.
58 Parameters
59 ----------
60 atomic_indices
61 Dictionary specifying which indices (values) belong to which atom type (keys).
62 frame_index
63 Trajectory index of the snapshot (frame).
64 positions
65 Positions as an array with shape `(n_atoms, 3)`.
66 velocities
67 Velocities as an array with shape `(n_atoms, 3)`; defaults to `None`.
68 forces
69 Forces as an array with shape `(n_atoms, 3)`; defaults to `None`.
70 properties
71 Names of the per-atom properties to keep, a subset of
72 :data:`per_atom_properties`; defaults to `None`, which keeps every quantity that
73 is provided.
74 """
76 def __init__(self,
77 atomic_indices: dict[str, list[int]],
78 frame_index: int,
79 positions: Optional[NDArray[float]] = None,
80 velocities: Optional[NDArray[float]] = None,
81 forces: Optional[NDArray[float]] = None,
82 properties: Optional[Sequence[str]] = None):
83 self._frame_index = frame_index
85 # Indexing along the atom axis only puts no constraint on the shape of a per-atom
86 # quantity, and already returns a copy rather than a view of the reader array.
87 arrays = dict(positions=positions, velocities=velocities, forces=forces)
88 self._arrays_by_type = dict()
89 for name in per_atom_properties:
90 array = arrays[name]
91 if array is None:
92 continue
93 if properties is not None and name not in properties:
94 continue
95 self._arrays_by_type[name] = {
96 atom_type: np.asarray(array)[indices]
97 for atom_type, indices in atomic_indices.items()}
99 @property
100 def positions_by_type(self) -> Optional[dict[str, NDArray[float]]]:
101 """ Positions split by atom type; ``None`` if the frame carries no positions. """
102 return self._arrays_by_type.get('positions')
104 @property
105 def velocities_by_type(self) -> Optional[dict[str, NDArray[float]]]:
106 """ Velocities split by atom type; ``None`` if the frame carries no velocities. """
107 return self._arrays_by_type.get('velocities')
109 @property
110 def forces_by_type(self) -> Optional[dict[str, NDArray[float]]]:
111 """ Forces split by atom type; ``None`` if the frame carries no forces. """
112 return self._arrays_by_type.get('forces')
114 def _get_as_array(self,
115 name: str,
116 atomic_indices: dict[str, list[int]]) -> NDArray[float]:
117 """
118 Return the per-atom quantity :attr:`name` reassembled into a single array whose
119 first dimension runs over all atoms.
121 Parameters
122 ----------
123 name
124 Name of the quantity, one of ``'positions'``, ``'velocities'``, ``'forces'``.
125 atomic_indices
126 Dictionary specifying which indices (values) belong to which atom type (keys).
127 """
128 arrays_by_type = self._arrays_by_type.get(name)
129 if arrays_by_type is None:
130 raise ValueError(f'This frame provides no {name}.')
132 # check that atomic_indices is complete
133 n_atoms = np.max([np.max(indices) for indices in atomic_indices.values()]) + 1
134 all_inds = [i for indices in atomic_indices.values() for i in indices]
135 if len(all_inds) != n_atoms or len(set(all_inds)) != n_atoms:
136 raise ValueError('atomic_indices is incomplete')
138 # Collect the quantity into a single array. Only the shape is taken from the
139 # per-type arrays, so that the returned array is float64 whatever the precision the
140 # reader provides, which some of them read as float32.
141 reference = next(iter(arrays_by_type.values()))
142 array = np.empty((n_atoms, ) + reference.shape[1:])
143 for atom_type, indices in atomic_indices.items():
144 array[indices] = arrays_by_type[atom_type]
145 return array
147 def get_positions_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
148 """
149 Return the full positions array with shape ``(n_atoms, 3)``.
151 Parameters
152 ----------
153 atomic_indices
154 Dictionary specifying which indices (values) belong to which atom type (keys).
155 """
156 return self._get_as_array('positions', atomic_indices)
158 def get_velocities_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
159 """
160 Return the full velocities array with shape ``(n_atoms, 3)``.
162 Parameters
163 ----------
164 atomic_indices
165 Dictionary specifying which indices (values) belong to which atom type (keys).
166 """
167 return self._get_as_array('velocities', atomic_indices)
169 def get_forces_as_array(self, atomic_indices: dict[str, list[int]]) -> NDArray[float]:
170 """
171 Return the full forces array with shape ``(n_atoms, 3)``.
173 Parameters
174 ----------
175 atomic_indices
176 Dictionary specifying which indices (values) belong to which atom type (keys).
177 """
178 return self._get_as_array('forces', atomic_indices)
180 @property
181 def frame_index(self) -> int:
182 """ Index of the frame. """
183 return self._frame_index
185 def __str__(self) -> str:
186 s = [f'Frame index {self.frame_index}']
187 for name, arrays_by_type in self._arrays_by_type.items():
188 for key, val in arrays_by_type.items():
189 s.append(f' {name:10} : {key} shape : {val.shape}')
190 return '\n'.join(s)
192 def __repr__(self) -> str:
193 return str(self)
195 def _repr_html_(self) -> str:
196 s = [f'<h3>{self.__class__.__name__}</h3>']
197 s += ['<table border="1" class="dataframe">']
198 s += ['<thead><tr><th style="text-align: left;">Field</th>'
199 '<th>Value/Shape</th></tr></thead>']
200 s += ['<tbody>']
201 s += [f'<tr><td style="text-align: left;">Index</td><td>{self.frame_index}</td></tr>']
202 for name, arrays_by_type in self._arrays_by_type.items():
203 for key, val in arrays_by_type.items():
204 s += [f'<tr><td style="text-align: left;">{name.capitalize()} {key}</td>'
205 f'<td>{val.shape}</td></tr>']
206 s += ['</tbody>']
207 s += ['</table>']
208 return '\n'.join(s)