Coverage for dynasor/trajectory/mdanalysis_trajectory_reader.py: 97%

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 

7import mmap 

8import warnings 

9 

10 

11class MDAnalysisTrajectoryReader(AbstractTrajectoryReader): 

12 """ Read a trajectory using the MDAnalysis Python library. 

13 

14 Parameters 

15 ---------- 

16 filename 

17 Name of input file. 

18 trajectory_format 

19 Type of trajectory. See MDAnalysis for the available formats. 

20 length_unit 

21 Unit of length for the input trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). 

22 time_unit 

23 Unit of time for the input trajectory (``'fs'``, ``'ps'``, ``'ns'``). 

24 """ 

25 

26 def __init__(self, 

27 filename: str, 

28 trajectory_format: str, 

29 length_unit: Optional[str] = None, 

30 time_unit: Optional[str] = None): 

31 

32 self._open = True 

33 self._first_called = False 

34 

35 with warnings.catch_warnings(): 

36 warnings.filterwarnings('ignore', category=UserWarning, 

37 message='Guessed all Masses to 1.0') 

38 warnings.filterwarnings('ignore', category=UserWarning, 

39 message='Reader has no dt information, set to 1.0 ps') 

40 u = mda.Universe(filename, format=trajectory_format, convert_units=False) 

41 self._frame_index = count(0) 

42 self._trajectory = u.trajectory 

43 

44 # netCDF is the only format read here that MDAnalysis maps into memory rather 

45 # than streams, and the map is never shrunk as the file is read. 

46 self._mmap = None 

47 if hasattr(mmap, 'MADV_DONTNEED'): 47 ↛ 51line 47 didn't jump to line 51 because the condition on line 47 was always true

48 self._mmap = getattr(getattr(self._trajectory, 'trjfile', None), '_mm', None) 

49 

50 # Set atomic_indices dict, if possible 

51 try: 

52 self._atom_types = u.atoms.types 

53 except mda.exceptions.NoDataError: 

54 self._atom_types = None 

55 

56 trajectory_length_unit = self._trajectory.units['length'] 

57 trajectory_time_unit = self._trajectory.units['time'] 

58 

59 if length_unit is not None and length_unit not in self.lengthunits_to_Angstrom_table: 

60 raise ValueError(f'Specified length unit {length_unit} is not an available option.') 

61 if time_unit is not None and time_unit not in self.timeunits_to_fs_table: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true

62 raise ValueError(f'Specified time unit {time_unit} is not an available option.') 

63 

64 if length_unit is None: 

65 length_unit = trajectory_length_unit or 'Angstrom' 

66 logger.info(f'Using length unit {length_unit} for the trajectory.') 

67 if time_unit is None: 

68 time_unit = trajectory_time_unit or 'fs' 

69 logger.info(f'Using time unit {time_unit} for the trajectory.') 

70 

71 length_scaling = mda.units.get_conversion_factor('length', length_unit, 'Angstrom') 

72 time_scaling = mda.units.get_conversion_factor('time', time_unit, 'fs') 

73 

74 def convert_units(ts): 

75 ts.positions *= length_scaling 

76 ts.triclinic_dimensions *= length_scaling 

77 if ts.has_velocities: 

78 ts.velocities *= length_scaling / time_scaling 

79 return ts 

80 self._trajectory.add_transformations(convert_units) 

81 

82 def _get_next(self): 

83 with warnings.catch_warnings(): 

84 warnings.filterwarnings('ignore', category=UserWarning, 

85 message='Reader has no dt information, set to 1.0 ps') 

86 if self._first_called: 

87 self._trajectory.next() 

88 else: 

89 self._first_called = True 

90 self._positions = self._trajectory.ts.positions 

91 self._cell = self._trajectory.ts.triclinic_dimensions 

92 self._n_atoms = self._trajectory.ts.n_atoms 

93 if self._trajectory.ts.has_velocities: 

94 self._velocities = self._trajectory.ts.velocities 

95 else: 

96 self._velocities = None 

97 

98 if self._mmap is not None: 

99 # This frame has been copied out of the map, so drop the pages it was read 

100 # from. Without this the resident set grows by every page touched and ends 

101 # up holding the entire trajectory, which is routinely tens of GB. 

102 self._mmap.madvise(mmap.MADV_DONTNEED) 

103 

104 def __iter__(self): 

105 """ Iterates through the trajectory file, frame by frame. """ 

106 return self 

107 

108 def __next__(self): 

109 """ Gets next trajectory frame. """ 

110 if not self._open: 

111 raise StopIteration 

112 

113 self._get_next() 

114 

115 if self._velocities is not None: 

116 frame = ReaderFrame(frame_index=next(self._frame_index), 

117 cell=self._cell, 

118 n_atoms=self._n_atoms, 

119 positions=self._positions.copy(), 

120 velocities=self._velocities.copy(), 

121 atom_types=self._atom_types 

122 ) 

123 else: 

124 frame = ReaderFrame(frame_index=next(self._frame_index), 

125 cell=self._cell, 

126 n_atoms=self._n_atoms, 

127 positions=self._positions.copy(), 

128 atom_types=self._atom_types 

129 ) 

130 

131 return frame 

132 

133 def close(self): 

134 """ Closes down, release resources etc. """ 

135 if self._open: 

136 self._trajectory.close() 

137 self._open = False