Coverage for dynasor/trajectory/extxyz_trajectory_reader.py: 90%

1import concurrent.futures 

2import numpy as np 

3import io 

4from itertools import count 

5from typing import Optional 

6from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

7from dynasor.trajectory.abstract_trajectory_reader import get_forces_from_atoms 

8from dynasor.trajectory.trajectory_frame import ReaderFrame 

9from ase.io import read 

10from ase import Atoms 

11 

12 

13def frame_text_to_atoms(frame_text): 

14 return read(io.StringIO(frame_text), format='extxyz') 

15 

16 

17def iter_frame_texts(f): 

18 while True: 

19 natoms_line = f.readline() 

20 if not natoms_line: 

21 break 

22 

23 comment_line = f.readline() 

24 if not comment_line: 24 ↛ 25line 24 didn't jump to line 25 because the condition on line 24 was never true

25 raise ValueError('Unexpected EOF while reading extxyz comment line') 

26 

27 try: 

28 natoms = int(natoms_line.strip()) 

29 except ValueError: 

30 raise ValueError(f'Invalid extxyz atom count line: {natoms_line!r}') 

31 

32 atom_lines = [] 

33 for _ in range(natoms): 

34 line = f.readline() 

35 if not line: 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true

36 raise ValueError('Unexpected EOF while reading extxyz atom lines') 

37 atom_lines.append(line) 

38 

39 yield natoms_line + comment_line + ''.join(atom_lines) 

40 

41 

42def iread(f, max_workers: Optional[int] = None) -> Atoms: 

43 frame_iterator = iter_frame_texts(f) 

44 with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as ex: 

45 buff = [] 

46 n_submit = ex._max_workers 

47 

48 for _ in range(n_submit): 

49 try: 

50 frame_text = next(frame_iterator) 

51 buff.append(ex.submit(frame_text_to_atoms, frame_text)) 

52 except StopIteration: 

53 break 

54 

55 while buff: 

56 res = buff.pop(0) 

57 

58 try: 

59 frame_text = next(frame_iterator) 

60 buff.append(ex.submit(frame_text_to_atoms, frame_text)) 

61 except StopIteration: 

62 pass 

63 

64 yield res.result() 

65 

66 

67class ExtxyzTrajectoryReader(AbstractTrajectoryReader): 

68 """Read extend xyz trajectory file, typically produced by GPUMD 

69 

70 This is a naive (and comparatively slow) parallel implementation which 

71 relies on the ASE xyz reader. 

72 

73 Parameters 

74 ---------- 

75 filename 

76 Name of input file. 

77 length_unit 

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

79 time_unit 

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

81 max_workers 

82 Number of working processes; defaults to ``None``, which means that the number of 

83 processors on the machine is used. 

84 force_unit 

85 Unit of force for the input trajectory (``'eV/Angstrom'``, ``'eV/nm'``, 

86 ``'kJ/mol/Angstrom'``, ``'kJ/mol/nm'``, ``'kcal/mol/Angstrom'``, ``'kcal/mol/nm'``, 

87 ``'Hartree/Bohr'``). 

88 """ 

89 

90 def __init__(self, 

91 filename: str, 

92 length_unit: Optional[str] = None, 

93 time_unit: Optional[str] = None, 

94 max_workers: Optional[int] = None, 

95 force_unit: Optional[str] = None): 

96 

97 # setup generator object 

98 self._fobj = open(filename, 'r') 

99 self._generator_xyz = iread(self._fobj, max_workers=max_workers) 

100 self._open = True 

101 self._frame_index = count(0) 

102 

103 # set up units 

104 self.set_unit_scaling_factors(length_unit, time_unit, force_unit) 

105 

106 def _get_next(self): 

107 try: 

108 atoms = next(self._generator_xyz) 

109 except StopIteration: 

110 self._fobj.close() 

111 self._open = False 

112 raise 

113 except Exception: 

114 self._fobj.close() 

115 self._open = False 

116 raise 

117 

118 self._atom_types = np.array(list(atoms.symbols)) 

119 self._n_atoms = len(atoms) 

120 self._cell = atoms.cell[:] 

121 self._x = atoms.positions 

122 if 'vel' in atoms.arrays: 

123 self._v = atoms.arrays['vel'] 

124 else: 

125 self._v = None 

126 self._f = get_forces_from_atoms(atoms) 

127 

128 def __iter__(self): 

129 return self 

130 

131 def close(self): 

132 if not self._fobj.closed: 

133 self._fobj.close() 

134 self._open = False 

135 

136 def __next__(self): 

137 if not self._open: 

138 raise StopIteration 

139 

140 self._get_next() 

141 

142 return ReaderFrame(frame_index=next(self._frame_index), 

143 n_atoms=int(self._n_atoms), 

144 cell=self.x_factor * self._cell.copy('F'), 

145 positions=self.x_factor * self._x, 

146 velocities=None if self._v is None else self.v_factor * self._v, 

147 forces=None if self._f is None else self.f_factor * self._f, 

148 atom_types=self._atom_types 

149 )