Coverage for dynasor/trajectory/lammps_trajectory_reader.py: 99%

1import numpy as np 

2import re 

3 

4from collections import deque 

5from typing import Optional 

6from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader 

7from dynasor.trajectory.trajectory_frame import ReaderFrame 

8from itertools import count 

9from numpy import array, arange, zeros 

10 

11 

12class LammpsTrajectoryReader(AbstractTrajectoryReader): 

13 """Read LAMMPS trajectory file 

14 

15 This is a naive (and comparatively slow) implementation, 

16 written entirely in python. 

17 

18 Parameters 

19 ---------- 

20 filename 

21 Name of input file. 

22 length_unit 

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

24 time_unit 

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

26 force_unit 

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

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

29 ``'Hartree/Bohr'``). 

30 """ 

31 

32 def __init__( 

33 self, 

34 filename: str, 

35 length_unit: Optional[str] = None, 

36 time_unit: Optional[str] = None, 

37 force_unit: Optional[str] = None 

38 ): 

39 

40 if filename.endswith('.gz'): 

41 import gzip 

42 self._fh = gzip.open(filename, 'rt') 

43 elif filename.endswith('.bz2'): 

44 import bz2 

45 self._fh = bz2.open(filename, 'rt') 

46 else: 

47 self._fh = open(filename, 'r') 

48 

49 self._open = True 

50 regexp = r'^ITEM: (TIMESTEP|NUMBER OF ATOMS|BOX BOUNDS|ATOMS) ?(.*)$' 

51 self._item_re = re.compile(regexp) 

52 

53 self._first_called = False 

54 self._frame_index = count(0) 

55 

56 # set up units 

57 self.set_unit_scaling_factors(length_unit, time_unit, force_unit) 

58 

59 # ITEM: TIMESTEP 

60 # 81000 

61 # ITEM: NUMBER OF ATOMS 

62 # 1536 

63 # ITEM: BOX BOUNDS pp pp pp 

64 # 1.54223 26.5378 

65 # 1.54223 26.5378 

66 # 1.54223 26.5378 

67 # ITEM: ATOMS id type x y z vx vy vz 

68 # 247 1 3.69544 2.56202 3.27701 0.00433856 -0.00099307 -0.00486166 

69 # 249 2 3.73324 3.05962 4.14359 0.00346029 0.00332502 -0.00731005 

70 # 463 1 3.5465 4.12841 5.34888 0.000523332 0.00145597 -0.00418675 

71 

72 def _read_frame_header(self): 

73 while True: 

74 L = self._fh.readline() 

75 m = self._item_re.match(L) 

76 if not m: 

77 if L == '': 

78 self._fh.close() 

79 self._open = False 

80 raise StopIteration 

81 if L.strip() == '': 

82 continue 

83 raise IOError('TRJ_reader: Failed to parse TRJ frame header') 

84 if m.group(1) == 'TIMESTEP': 

85 step = int(self._fh.readline()) 

86 elif m.group(1) == 'NUMBER OF ATOMS': 

87 n_atoms = int(self._fh.readline()) 

88 elif m.group(1) == 'BOX BOUNDS': 

89 bbounds = [deque(map(float, self._fh.readline().split())) 

90 for _ in range(3)] 

91 x = array(bbounds) 

92 cell = np.diag(x[:, 1] - x[:, 0]) 

93 if x.shape == (3, 3): 

94 cell[1, 0] = x[0, 2] 

95 cell[2, 0] = x[1, 2] 

96 cell[2, 1] = x[2, 2] 

97 elif x.shape != (3, 2): 

98 raise IOError('TRJ_reader: Malformed cell bounds') 

99 elif m.group(1) == 'ATOMS': 99 ↛ 73line 99 didn't jump to line 73 because the condition on line 99 was always true

100 cols = tuple(m.group(2).split()) 

101 # At this point, there should be only atomic data left 

102 return (step, n_atoms, cell, cols) 

103 

104 def _get_first(self): 

105 # Read first frame, update state of self, create indices etc 

106 step, N, cell, cols = self._read_frame_header() 

107 self._n_atoms = N 

108 self._step = step 

109 self._cols = cols 

110 self._cell = cell 

111 

112 def _all_in_cols(keys): 

113 for k in keys: 

114 if k not in cols: 

115 return False 

116 return True 

117 

118 self._x_map = None 

119 if _all_in_cols(('id', 'xu', 'yu', 'zu')): 

120 self._x_I = array(deque(map(cols.index, ('xu', 'yu', 'zu')))) 

121 elif _all_in_cols(('id', 'x', 'y', 'z')): 

122 self._x_I = array(deque(map(cols.index, ('x', 'y', 'z')))) 

123 elif _all_in_cols(('id', 'xs', 'ys', 'zs')): 

124 self._x_I = array(deque(map(cols.index, ('xs', 'ys', 'zs')))) 

125 _x_factor = self._cell.diagonal() 

126 # xs.shape == (n, 3) 

127 self._x_map = lambda xs: xs * _x_factor 

128 else: 

129 raise RuntimeError('TRJ file must contain at least atom-id, x, y,' 

130 ' and z coordinates to be useful.') 

131 self._id_I = cols.index('id') 

132 

133 if _all_in_cols(('vx', 'vy', 'vz')): 

134 self._v_I = array(deque(map(cols.index, ('vx', 'vy', 'vz')))) 

135 else: 

136 self._v_I = None 

137 

138 if _all_in_cols(('fx', 'fy', 'fz')): 

139 self._f_I = array(deque(map(cols.index, ('fx', 'fy', 'fz')))) 

140 else: 

141 self._f_I = None 

142 

143 if 'type' in cols: 

144 self._type_I = cols.index('type') 

145 else: 

146 self._type_I = None 

147 

148 data = array([list(map(float, self._fh.readline().split())) for _ in range(N)]) 

149 # data.shape == (N, Ncols) 

150 II = np.asarray(data[:, self._id_I], dtype=np.int_) 

151 # Unless dump is done for group 'all' ... 

152 II[np.argsort(II)] = arange(len(II)) 

153 self._x = zeros((N, 3)) 

154 if self._x_map is None: 

155 self._x[II] = data[:, self._x_I] 

156 else: 

157 self._x[II] = self._x_map(data[:, self._x_I]) 

158 if self._v_I is not None: 

159 self._v = zeros((N, 3)) 

160 self._v[II] = data[:, self._v_I] 

161 if self._f_I is not None: 

162 self._f = zeros((N, 3)) 

163 self._f[II] = data[:, self._f_I] 

164 

165 self._first_called = True 

166 

167 def _get_next(self): 

168 # get next frame, update state of self 

169 step, N, cell, cols = self._read_frame_header() 

170 assert self._n_atoms == N 

171 assert self._cols == cols 

172 self._step = step 

173 self._cell = cell 

174 

175 data = array([deque(map(float, self._fh.readline().split())) 

176 for _ in range(N)]) 

177 II = np.asarray(data[:, self._id_I], dtype=np.int_) - 1 

178 if self._x_map is None: 

179 self._x[II] = data[:, self._x_I] 

180 else: 

181 self._x[II] = self._x_map(data[:, self._x_I]) 

182 if self._v_I is not None: 

183 self._v[II] = data[:, self._v_I] 

184 if self._f_I is not None: 

185 self._f[II] = data[:, self._f_I] 

186 

187 def __iter__(self): 

188 return self 

189 

190 def close(self): 

191 self._open = False 

192 if not self._fh.closed: 

193 self._fh.close() 

194 

195 def __next__(self): 

196 if not self._open: 

197 raise StopIteration 

198 

199 if self._first_called: 

200 self._get_next() 

201 else: 

202 self._get_first() 

203 

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

205 n_atoms=int(self._n_atoms), 

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

207 positions=self.x_factor * self._x, 

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

209 forces=None if self._f_I is None else self.f_factor * self._f 

210 )