Coverage for /builds/ase/ase/ase/io/gromacs.py: 90.18%

1# fmt: off 

2 

3""" 

4read and write gromacs geometry files 

5""" 

6 

7import numpy as np 

8 

9from ase import units 

10from ase.atoms import Atoms 

11from ase.data import atomic_numbers 

12from ase.utils import reader, writer 

13 

14 

15@reader 

16def read_gromacs(fd): 

17 """ From: 

18 http://manual.gromacs.org/current/online/gro.html 

19 C format 

20 "%5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f" 

21 python: starting from 0, including first excluding last 

22 0:4 5:10 10:15 15:20 20:28 28:36 36:44 44:52 52:60 60:68 

23 

24 Import gromacs geometry type files (.gro). 

25 Reads atom positions, 

26 velocities(if present) and 

27 simulation cell (if present) 

28 """ 

29 

30 atoms = Atoms() 

31 lines = fd.readlines() 

32 positions = [] 

33 gromacs_velocities = [] 

34 symbols = [] 

35 tags = [] 

36 gromacs_residuenumbers = [] 

37 gromacs_residuenames = [] 

38 gromacs_atomtypes = [] 

39 sym2tag = {} 

40 tag = 0 

41 for line in (lines[2:-1]): 

42 # print(line[0:5]+':'+line[5:10]+':'+line[10:15]+':'+line[15:20]) 

43 # it is not a good idea to use the split method with gromacs input 

44 # since the fields are defined by a fixed column number. Therefore, 

45 # they may not be space between the fields 

46 # inp = line.split() 

47 

48 floatvect = float(line[20:28]) * 10.0, \ 

49 float(line[28:36]) * 10.0, \ 

50 float(line[36:44]) * 10.0 

51 positions.append(floatvect) 

52 

53 # read velocities 

54 velocities = np.array([0.0, 0.0, 0.0]) 

55 vx = line[44:52].strip() 

56 vy = line[52:60].strip() 

57 vz = line[60:68].strip() 

58 

59 for iv, vxyz in enumerate([vx, vy, vz]): 

60 if len(vxyz) > 0: 

61 try: 

62 velocities[iv] = float(vxyz) 

63 except ValueError as exc: 

64 raise ValueError("can not convert velocity to float") \ 

65 from exc 

66 else: 

67 velocities = None 

68 

69 if velocities is not None: 

70 # velocities from nm/ps to ase units 

71 velocities *= units.nm / (1000.0 * units.fs) 

72 gromacs_velocities.append(velocities) 

73 

74 gromacs_residuenumbers.append(int(line[0:5])) 

75 gromacs_residuenames.append(line[5:10].strip()) 

76 

77 symbol_read = line[10:15].strip()[0:2] 

78 if symbol_read not in sym2tag: 

79 sym2tag[symbol_read] = tag 

80 tag += 1 

81 

82 tags.append(sym2tag[symbol_read]) 

83 if symbol_read in atomic_numbers: 

84 symbols.append(symbol_read) 

85 elif symbol_read[0] in atomic_numbers: 

86 symbols.append(symbol_read[0]) 

87 elif symbol_read[-1] in atomic_numbers: 

88 symbols.append(symbol_read[-1]) 

89 else: 

90 # not an atomic symbol 

91 # if we can not determine the symbol, we use 

92 # the dummy symbol X 

93 symbols.append("X") 

94 

95 gromacs_atomtypes.append(line[10:15].strip()) 

96 

97 line = lines[-1] 

98 atoms = Atoms(symbols, positions, tags=tags) 

99 

100 if len(gromacs_velocities) == len(atoms): 

101 atoms.set_velocities(gromacs_velocities) 

102 elif len(gromacs_velocities) != 0: 

103 raise ValueError("Some atoms velocities were not specified!") 

104 

105 if not atoms.has('residuenumbers'): 

106 atoms.new_array('residuenumbers', gromacs_residuenumbers, int) 

107 atoms.set_array('residuenumbers', gromacs_residuenumbers, int) 

108 if not atoms.has('residuenames'): 

109 atoms.new_array('residuenames', gromacs_residuenames, str) 

110 atoms.set_array('residuenames', gromacs_residuenames, str) 

111 if not atoms.has('atomtypes'): 

112 atoms.new_array('atomtypes', gromacs_atomtypes, str) 

113 atoms.set_array('atomtypes', gromacs_atomtypes, str) 

114 

115 # determine PBC and unit cell 

116 atoms.pbc = False 

117 inp = lines[-1].split() 

118 try: 

119 grocell = list(map(float, inp)) 

120 except ValueError: 

121 return atoms 

122 

123 if len(grocell) < 3: 

124 return atoms 

125 

126 cell = np.diag(grocell[:3]) 

127 

128 if len(grocell) >= 9: 

129 cell.flat[[1, 2, 3, 5, 6, 7]] = grocell[3:9] 

130 

131 atoms.cell = cell * 10. 

132 atoms.pbc = True 

133 return atoms 

134 

135 

136@writer 

137def write_gromacs(fileobj, atoms): 

138 """Write gromacs geometry files (.gro). 

139 

140 Writes: 

141 

142 * atom positions, 

143 * velocities (if present, otherwise 0) 

144 * simulation cell (if present) 

145 """ 

146 

147 natoms = len(atoms) 

148 try: 

149 gromacs_residuenames = atoms.get_array('residuenames') 

150 except KeyError: 

151 gromacs_residuenames = [] 

152 for _ in range(natoms): 

153 gromacs_residuenames.append('1DUM') 

154 try: 

155 gromacs_atomtypes = atoms.get_array('atomtypes') 

156 except KeyError: 

157 gromacs_atomtypes = atoms.get_chemical_symbols() 

158 

159 try: 

160 residuenumbers = atoms.get_array('residuenumbers') 

161 except KeyError: 

162 residuenumbers = np.ones(natoms, int) 

163 

164 pos = atoms.get_positions() 

165 pos = pos / 10.0 

166 

167 vel = atoms.get_velocities() 

168 if vel is None: 

169 vel = pos * 0.0 

170 else: 

171 vel *= 1000.0 * units.fs / units.nm 

172 

173 # No "#" in the first line to prevent read error in VMD 

174 fileobj.write('A Gromacs structure file written by ASE \n') 

175 fileobj.write('%5d\n' % len(atoms)) 

176 # gromacs line see 

177 # manual.gromacs.org/documentation/current/user-guide/file-formats.html#gro 

178 # (EDH: link seems broken as of 2020-02-21) 

179 # 1WATER OW1 1 0.126 1.624 1.679 0.1227 -0.0580 0.0434 

180 for count, (resnb, resname, atomtype, 

181 xyz, vxyz) in enumerate(zip(residuenumbers, 

182 gromacs_residuenames, 

183 gromacs_atomtypes, pos, vel), 

184 start=1): 

185 

186 # THIS SHOULD BE THE CORRECT, PYTHON FORMATTING, EQUIVALENT TO THE 

187 # C FORMATTING GIVEN IN THE GROMACS DOCUMENTATION: 

188 # >>> %5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f <<< 

189 line = ("{:>5d}{:<5s}{:>5s}{:>5d}{:>8.3f}{:>8.3f}{:>8.3f}" 

190 "{:>8.4f}{:>8.4f}{:>8.4f}\n".format(resnb, resname, atomtype, 

191 count, xyz[0], xyz[1], 

192 xyz[2], vxyz[0], vxyz[1], 

193 vxyz[2])) 

194 

195 fileobj.write(line) 

196 # write box geometry 

197 if atoms.get_pbc().any(): 

198 # gromacs manual (manual.gromacs.org/online/gro.html) says: 

199 # v1(x) v2(y) v3(z) v1(y) v1(z) v2(x) v2(z) v3(x) v3(y) 

200 # 

201 # cell[i,j] is the jth Cartesian coordinate of the ith unit vector 

202 # cell[0,0] cell[1,1] cell[2,2] v1(x) v2(y) v3(z) fv0[0 1 2] 

203 # cell[0,1] cell[0,2] cell[1,0] v1(y) v1(z) v2(x) fv1[0 1 2] 

204 # cell[1,2] cell[2,0] cell[2,1] v2(z) v3(x) v3(y) fv2[0 1 2] 

205 grocell = atoms.cell.flat[[0, 4, 8, 1, 2, 3, 5, 6, 7]] * 0.1 

206 fileobj.write(''.join(['{:10.5f}'.format(x) for x in grocell])) 

207 fileobj.write('\n') 

208 else: 

209 # When we do not have a cell, the cell is specified as an empty line 

210 fileobj.write("\n")