Coverage for /builds/debichem-team/python-ase/ase/io/gromacs.py: 90.18%

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:10]+':'+line[10: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 as exc: 

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

63 from exc 

64 else: 

65 velocities = None 

66 

67 if velocities is not None: 

68 # velocities from nm/ps to ase units 

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

70 gromacs_velocities.append(velocities) 

71 

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

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

74 

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

76 if symbol_read not in sym2tag: 

77 sym2tag[symbol_read] = tag 

78 tag += 1 

79 

80 tags.append(sym2tag[symbol_read]) 

81 if symbol_read in atomic_numbers: 

82 symbols.append(symbol_read) 

83 elif symbol_read[0] in atomic_numbers: 

84 symbols.append(symbol_read[0]) 

85 elif symbol_read[-1] in atomic_numbers: 

86 symbols.append(symbol_read[-1]) 

87 else: 

88 # not an atomic symbol 

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

90 # the dummy symbol X 

91 symbols.append("X") 

92 

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

94 

95 line = lines[-1] 

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

97 

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

99 atoms.set_velocities(gromacs_velocities) 

100 elif len(gromacs_velocities) != 0: 

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

102 

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

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

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

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

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

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

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

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

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

112 

113 # determine PBC and unit cell 

114 atoms.pbc = False 

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

116 try: 

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

118 except ValueError: 

119 return atoms 

120 

121 if len(grocell) < 3: 

122 return atoms 

123 

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

125 

126 if len(grocell) >= 9: 

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

128 

129 atoms.cell = cell * 10. 

130 atoms.pbc = True 

131 return atoms 

132 

133 

134@writer 

135def write_gromacs(fileobj, atoms): 

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

137 

138 Writes: 

139 

140 * atom positions, 

141 * velocities (if present, otherwise 0) 

142 * simulation cell (if present) 

143 """ 

144 

145 natoms = len(atoms) 

146 try: 

147 gromacs_residuenames = atoms.get_array('residuenames') 

148 except KeyError: 

149 gromacs_residuenames = [] 

150 for _ in range(natoms): 

151 gromacs_residuenames.append('1DUM') 

152 try: 

153 gromacs_atomtypes = atoms.get_array('atomtypes') 

154 except KeyError: 

155 gromacs_atomtypes = atoms.get_chemical_symbols() 

156 

157 try: 

158 residuenumbers = atoms.get_array('residuenumbers') 

159 except KeyError: 

160 residuenumbers = np.ones(natoms, int) 

161 

162 pos = atoms.get_positions() 

163 pos = pos / 10.0 

164 

165 vel = atoms.get_velocities() 

166 if vel is None: 

167 vel = pos * 0.0 

168 else: 

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

170 

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

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

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

174 # gromacs line see 

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

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

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

178 for count, (resnb, resname, atomtype, 

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

180 gromacs_residuenames, 

181 gromacs_atomtypes, pos, vel), 

182 start=1): 

183 

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

185 # C FORMATTING GIVEN IN THE GROMACS DOCUMENTATION: 

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

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

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

189 count, xyz[0], xyz[1], 

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

191 vxyz[2])) 

192 

193 fileobj.write(line) 

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(['{:10.5f}'.format(x) for x in grocell])) 

205 fileobj.write('\n') 

206 else: 

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

208 fileobj.write("\n")