Coverage for dynasor/trajectory/mdanalysis_trajectory_reader.py: 97%
108 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
1from itertools import count
2from typing import Optional
3from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader
4from dynasor.trajectory.trajectory_frame import ReaderFrame
5from dynasor.logging_tools import logger
6import MDAnalysis as mda
7from MDAnalysis.coordinates.core import get_reader_for
8from MDAnalysis.coordinates.XYZ import XYZReader
9import mmap
10import numpy as np
11import warnings
14# MDAnalysis declares the native force unit of a trajectory with its own spelling, which is
15# translated here onto the names used by
16# :attr:`AbstractTrajectoryReader.forceunits_to_eV_per_Angstrom_table`. The force units of
17# MDAnalysis that have no counterpart there, such as Newton, are left out, so that a
18# trajectory reporting one of them has its forces passed on unconverted.
19mdanalysis_force_unit_names = {
20 'kJ/(mol*Angstrom)': 'kJ/mol/Angstrom',
21 'kJ/(mol*A)': 'kJ/mol/Angstrom',
22 # the Angstrom here is U+212B ANGSTROM SIGN, which is what MDAnalysis uses, and
23 # not the visually identical U+00C5 LATIN CAPITAL LETTER A WITH RING ABOVE
24 'kJ/(mol*Å)': 'kJ/mol/Angstrom',
25 'kJ/(mol*nm)': 'kJ/mol/nm',
26 'kcal/(mol*Angstrom)': 'kcal/mol/Angstrom',
27}
30class MDAnalysisTrajectoryReader(AbstractTrajectoryReader):
31 """ Read a trajectory using the MDAnalysis Python library.
33 Parameters
34 ----------
35 filename
36 Name of input file.
37 trajectory_format
38 Name of the format to read the file with, for example ``'nc'`` or ``'NCDF'`` for an
39 AMBER netCDF trajectory. It must be a name MDAnalysis knows, and case does not
40 matter. If ``None``, MDAnalysis picks a reader based on the file extension.
41 length_unit
42 Unit of length for the input trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``).
43 time_unit
44 Unit of time for the input trajectory (``'fs'``, ``'ps'``, ``'ns'``).
45 force_unit
46 Unit of force for the input trajectory (``'eV/Angstrom'``, ``'eV/nm'``,
47 ``'kJ/mol/Angstrom'``, ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``,
48 ``'Hartree/Bohr'``). Defaults to the unit that MDAnalysis reports for the
49 trajectory.
51 Raises
52 ------
53 ValueError
54 If MDAnalysis has no reader for :attr:`trajectory_format`,
55 or if the trajectory carries no cell, or a cell that spans no volume.
56 """
58 def __init__(self,
59 filename: str,
60 trajectory_format: Optional[str],
61 length_unit: Optional[str] = None,
62 time_unit: Optional[str] = None,
63 force_unit: Optional[str] = None):
65 self._open = True
66 self._first_called = False
68 # An unknown format name only makes mda.Universe warn and return a topology
69 # without coordinates, which then fails with an unrelated message.
70 if trajectory_format is not None:
71 try:
72 get_reader_for(filename, format=trajectory_format)
73 except ValueError as error:
74 raise ValueError(
75 f'MDAnalysis has no reader for the trajectory format'
76 f' "{trajectory_format}". dynasor reads the formats "lammps_internal",'
77 ' "extxyz", "ase", and "lammps_mdanalysis" itself and passes any other'
78 ' name on to MDAnalysis. Use "NCDF" for an AMBER netCDF trajectory.'
79 f' The message from MDAnalysis is:\n{error}') from error
81 with warnings.catch_warnings():
82 warnings.filterwarnings('ignore', category=UserWarning,
83 message='Guessed all Masses to 1.0')
84 warnings.filterwarnings('ignore', category=UserWarning,
85 message='Reader has no dt information, set to 1.0 ps')
86 try:
87 u = mda.Universe(filename, format=trajectory_format, convert_units=False)
88 except ValueError as error:
89 if trajectory_format is not None: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true
90 raise
91 # Without a format name MDAnalysis goes by the file extension, and it
92 # knows none of the extensions that the readers of dynasor handle.
93 raise ValueError(
94 f'MDAnalysis found no reader for {filename} by its file extension.'
95 f' Name the format with trajectory_format. dynasor reads the formats'
96 ' "lammps_internal", "extxyz", "ase", and "lammps_mdanalysis" itself'
97 ' and passes any other name on to MDAnalysis.'
98 f' The message from MDAnalysis is:\n{error}') from error
99 self._frame_index = count(0)
100 self._trajectory = u.trajectory
102 # dynasor needs a cell that spans a non-zero volume, and several formats read here
103 # provide no cell at all. MDAnalysis maps an invalid cell, such as one with a
104 # non-positive length or an angle outside the range it accepts, onto an all-zero
105 # matrix, which the unit conversion below then turns into None. That fails later and
106 # far from its cause, so such a cell is rejected here as well.
107 cell = self._trajectory.ts.triclinic_dimensions
108 if cell is None or not np.all(np.isfinite(cell)) or np.linalg.det(cell) == 0:
109 if trajectory_format is None:
110 read_with = 'the format MDAnalysis picked from the file extension'
111 else:
112 read_with = f'trajectory format "{trajectory_format}"'
113 message = (f'No valid cell was found in {filename} when read with {read_with}.'
114 ' dynasor requires a cell that spans a non-zero volume.')
115 # The hint applies whenever the xyz reader of MDAnalysis was used, which
116 # happens both for an explicit format name and for a file picked by extension.
117 if isinstance(self._trajectory, XYZReader):
118 message += (' An extended xyz file must be read with'
119 ' trajectory_format="extxyz", since the xyz reader of MDAnalysis'
120 ' ignores the Lattice field.')
121 raise ValueError(message)
123 # netCDF is the only format read here that MDAnalysis maps into memory rather
124 # than streams, and the map is never shrunk as the file is read.
125 self._mmap = None
126 if hasattr(mmap, 'MADV_DONTNEED'): 126 ↛ 130line 126 didn't jump to line 130 because the condition on line 126 was always true
127 self._mmap = getattr(getattr(self._trajectory, 'trjfile', None), '_mm', None)
129 # Set atomic_indices dict, if possible
130 try:
131 self._atom_types = u.atoms.types
132 except mda.exceptions.NoDataError:
133 self._atom_types = None
135 trajectory_length_unit = self._trajectory.units['length']
136 trajectory_time_unit = self._trajectory.units['time']
137 trajectory_force_unit = self._trajectory.units.get('force')
139 if length_unit is not None and length_unit not in self.lengthunits_to_Angstrom_table:
140 raise ValueError(f'Specified length unit {length_unit} is not an available option.')
141 if time_unit is not None and time_unit not in self.timeunits_to_fs_table: 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true
142 raise ValueError(f'Specified time unit {time_unit} is not an available option.')
143 if (force_unit is not None
144 and force_unit not in self.forceunits_to_eV_per_Angstrom_table):
145 raise ValueError(f'Specified force unit {force_unit} is not an available option.')
147 if length_unit is None:
148 length_unit = trajectory_length_unit or 'Angstrom'
149 logger.info(f'Using length unit {length_unit} for the trajectory.')
150 if time_unit is None:
151 time_unit = trajectory_time_unit or 'fs'
152 logger.info(f'Using time unit {time_unit} for the trajectory.')
154 # Forces are optional, so the unit is resolved without logging here and reported by
155 # :class:`Trajectory <dynasor.Trajectory>` only for a trajectory that provides them.
156 # Several formats read here report no force unit at all, the LAMMPS dump reader of
157 # MDAnalysis among them, and those get the same eV/Angstrom default as the other
158 # readers. The unit stays None only when MDAnalysis reports one that dynasor cannot
159 # convert, such as Newton, in which case the forces are passed on unconverted.
160 if force_unit is None:
161 if trajectory_force_unit is None:
162 force_unit = 'eV/Angstrom'
163 else:
164 force_unit = mdanalysis_force_unit_names.get(trajectory_force_unit)
165 self.force_unit = force_unit
167 length_scaling = mda.units.get_conversion_factor('length', length_unit, 'Angstrom')
168 time_scaling = mda.units.get_conversion_factor('time', time_unit, 'fs')
169 force_scaling = 1.0 if force_unit is None \
170 else self.forceunits_to_eV_per_Angstrom_table[force_unit]
172 def convert_units(ts):
173 ts.positions *= length_scaling
174 ts.triclinic_dimensions *= length_scaling
175 if ts.has_velocities:
176 ts.velocities *= length_scaling / time_scaling
177 if ts.has_forces:
178 ts.forces *= force_scaling
179 return ts
180 self._trajectory.add_transformations(convert_units)
182 def _get_next(self):
183 with warnings.catch_warnings():
184 warnings.filterwarnings('ignore', category=UserWarning,
185 message='Reader has no dt information, set to 1.0 ps')
186 if self._first_called:
187 self._trajectory.next()
188 else:
189 self._first_called = True
190 self._positions = self._trajectory.ts.positions
191 self._cell = self._trajectory.ts.triclinic_dimensions
192 self._n_atoms = self._trajectory.ts.n_atoms
193 if self._trajectory.ts.has_velocities:
194 self._velocities = self._trajectory.ts.velocities
195 else:
196 self._velocities = None
197 if self._trajectory.ts.has_forces:
198 self._forces = self._trajectory.ts.forces
199 else:
200 self._forces = None
202 if self._mmap is not None:
203 # This frame has been copied out of the map, so drop the pages it was read
204 # from. Without this the resident set grows by every page touched and ends
205 # up holding the entire trajectory, which is routinely tens of GB.
206 self._mmap.madvise(mmap.MADV_DONTNEED)
208 def __iter__(self):
209 """ Iterates through the trajectory file, frame by frame. """
210 return self
212 def __next__(self):
213 """ Gets next trajectory frame. """
214 if not self._open:
215 raise StopIteration
217 self._get_next()
219 return ReaderFrame(frame_index=next(self._frame_index),
220 cell=self._cell,
221 n_atoms=self._n_atoms,
222 positions=self._positions.copy(),
223 velocities=None if self._velocities is None
224 else self._velocities.copy(),
225 forces=None if self._forces is None else self._forces.copy(),
226 atom_types=self._atom_types
227 )
229 def close(self):
230 """ Closes down, release resources etc. """
231 if self._open:
232 self._trajectory.close()
233 self._open = False