Coverage for /builds/kinetik161/ase/ase/io/gromacs.py: 86.73%

1""" 

2read and write gromacs geometry files 

3""" 

4 

5import numpy as np 

6 

7from ase import units 

8from ase.atoms import Atoms 

9from ase.data import atomic_numbers 

10from ase.utils import reader, writer 

11 

12 

13@reader 

14def read_gromacs(fd): 

15 """ From: 

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

17 C format 

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

19 python: starting from 0, including first excluding last 

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

21 

22 Import gromacs geometry type files (.gro). 

23 Reads atom positions, 

24 velocities(if present) and 

25 simulation cell (if present) 

26 """ 

27 

28 atoms = Atoms() 

29 lines = fd.readlines() 

30 positions = [] 

31 gromacs_velocities = [] 

32 symbols = [] 

33 tags = [] 

34 gromacs_residuenumbers = [] 

35 gromacs_residuenames = [] 

36 gromacs_atomtypes = [] 

37 sym2tag = {} 

38 tag = 0 

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

40 # print(line[0:5]+':'+line[5:11]+':'+line[11:15]+':'+line[15:20]) 

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

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

43 # they may not be space between the fields 

44 # inp = line.split() 

45 

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

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

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

49 positions.append(floatvect) 

50 

51 # read velocities 

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

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

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

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

56 

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

58 if len(vxyz) > 0: 

59 try: 

60 velocities[iv] = float(vxyz) 

61 except ValueError: 

62 raise ValueError("can not convert velocity to float") 

63 else: 

64 velocities = None 

65 

66 if velocities is not None: 

67 # velocities from nm/ps to ase units 

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

69 gromacs_velocities.append(velocities) 

70 

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

72 gromacs_residuenames.append(line[5:11].strip()) 

73 

74 symbol_read = line[11:16].strip()[0:2] 

75 if symbol_read not in sym2tag.keys(): 

76 sym2tag[symbol_read] = tag 

77 tag += 1 

78 

79 tags.append(sym2tag[symbol_read]) 

80 if symbol_read in atomic_numbers: 

81 symbols.append(symbol_read) 

82 elif symbol_read[0] in atomic_numbers: 

83 symbols.append(symbol_read[0]) 

84 elif symbol_read[-1] in atomic_numbers: 

85 symbols.append(symbol_read[-1]) 

86 else: 

87 # not an atomic symbol 

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

89 # the dummy symbol X 

90 symbols.append("X") 

91 

92 gromacs_atomtypes.append(line[11:16].strip()) 

93 

94 line = lines[-1] 

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

96 

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

98 atoms.set_velocities(gromacs_velocities) 

99 elif len(gromacs_velocities) != 0: 

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

101 

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

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

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

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

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

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

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

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

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

111 

112 # determine PBC and unit cell 

113 atoms.pbc = False 

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

115 try: 

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

117 except ValueError: 

118 return atoms 

119 

120 if len(grocell) < 3: 

121 return atoms 

122 

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

124 

125 if len(grocell) >= 9: 

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

127 

128 atoms.cell = cell * 10. 

129 atoms.pbc = True 

130 return atoms 

131 

132 

133@writer 

134def write_gromacs(fileobj, atoms): 

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

136 

137 Writes: 

138 

139 * atom positions, 

140 * velocities (if present, otherwise 0) 

141 * simulation cell (if present) 

142 """ 

143 

144 natoms = len(atoms) 

145 try: 

146 gromacs_residuenames = atoms.get_array('residuenames') 

147 except KeyError: 

148 gromacs_residuenames = [] 

149 for idum in range(natoms): 

150 gromacs_residuenames.append('1DUM') 

151 try: 

152 gromacs_atomtypes = atoms.get_array('atomtypes') 

153 except KeyError: 

154 gromacs_atomtypes = atoms.get_chemical_symbols() 

155 

156 try: 

157 residuenumbers = atoms.get_array('residuenumbers') 

158 except KeyError: 

159 residuenumbers = np.ones(natoms, int) 

160 

161 pos = atoms.get_positions() 

162 pos = pos / 10.0 

163 

164 vel = atoms.get_velocities() 

165 if vel is None: 

166 vel = pos * 0.0 

167 else: 

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

169 

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

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

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

173 count = 1 

174 

175 # gromacs line see 

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

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

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

179 for (resnb, resname, atomtype, xyz, 

180 vxyz) in zip(residuenumbers, gromacs_residuenames, 

181 gromacs_atomtypes, pos, vel): 

182 

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

184 # C FORMATTING GIVEN IN THE GROMACS DOCUMENTATION: 

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

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

187 "{:>8.4f}{:>8.4f}{:>8.4f}\n" 

188 .format(resnb, resname, atomtype, count, 

189 xyz[0], xyz[1], xyz[2], vxyz[0], vxyz[1], vxyz[2])) 

190 

191 fileobj.write(line) 

192 count += 1 

193 

194 # write box geometry 

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

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

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

198 # 

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

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

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

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

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

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

205 else: 

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

207 fileobj.write("\n")