1"""This module defines an ASE interface to Amber16. 

2 

3Usage: (Tested only with Amber16, http://ambermd.org/) 

4 

5Before usage, input files (infile, topologyfile, incoordfile) 

6 

7""" 

8 

9import subprocess 

10import numpy as np 

11 

12from ase.calculators.calculator import Calculator, FileIOCalculator 

13import ase.units as units 

14from scipy.io import netcdf 

15 

16 

17class Amber(FileIOCalculator): 

18 """Class for doing Amber classical MM calculations. 

19 

20 Example: 

21 

22 mm.in:: 

23 

24 Minimization with Cartesian restraints 

25 &cntrl 

26 imin=1, maxcyc=200, (invoke minimization) 

27 ntpr=5, (print frequency) 

28 &end 

29 """ 

30 

31 implemented_properties = ['energy', 'forces'] 

32 discard_results_on_any_change = True 

33 

34 def __init__(self, restart=None, 

35 ignore_bad_restart_file=FileIOCalculator._deprecated, 

36 label='amber', atoms=None, command=None, 

37 amber_exe='sander -O ', 

38 infile='mm.in', outfile='mm.out', 

39 topologyfile='mm.top', incoordfile='mm.crd', 

40 outcoordfile='mm_dummy.crd', mdcoordfile=None, 

41 **kwargs): 

42 """Construct Amber-calculator object. 

43 

44 Parameters 

45 ========== 

46 label: str 

47 Name used for all files. May contain a directory. 

48 atoms: Atoms object 

49 Optional Atoms object to which the calculator will be 

50 attached. When restarting, atoms will get its positions and 

51 unit-cell updated from file. 

52 label: str 

53 Prefix to use for filenames (label.in, label.txt, ...). 

54 amber_exe: str 

55 Name of the amber executable, one can add options like -O 

56 and other parameters here 

57 infile: str 

58 Input filename for amber, contains instuctions about the run 

59 outfile: str 

60 Logfilename for amber 

61 topologyfile: str 

62 Name of the amber topology file 

63 incoordfile: str 

64 Name of the file containing the input coordinates of atoms 

65 outcoordfile: str 

66 Name of the file containing the output coordinates of atoms 

67 this file is not used in case minisation/dynamics is done by ase. 

68 It is only relevant 

69 if you run MD/optimisation many steps with amber. 

70 

71 """ 

72 

73 self.out = 'mm.log' 

74 

75 self.positions = None 

76 self.atoms = None 

77 

78 self.set(**kwargs) 

79 

80 self.amber_exe = amber_exe 

81 self.infile = infile 

82 self.outfile = outfile 

83 self.topologyfile = topologyfile 

84 self.incoordfile = incoordfile 

85 self.outcoordfile = outcoordfile 

86 self.mdcoordfile = mdcoordfile 

87 if command is not None: 

88 self.command = command 

89 else: 

90 self.command = (self.amber_exe + 

91 ' -i ' + self.infile + 

92 ' -o ' + self.outfile + 

93 ' -p ' + self.topologyfile + 

94 ' -c ' + self.incoordfile + 

95 ' -r ' + self.outcoordfile) 

96 if self.mdcoordfile is not None: 

97 self.command = self.command + ' -x ' + self.mdcoordfile 

98 

99 FileIOCalculator.__init__(self, restart, ignore_bad_restart_file, 

100 label, atoms, **kwargs) 

101 

102 def write_input(self, atoms=None, properties=None, system_changes=None): 

103 """Write updated coordinates to a file.""" 

104 

105 FileIOCalculator.write_input(self, atoms, properties, system_changes) 

106 self.write_coordinates(atoms) 

107 

108 def read_results(self): 

109 """ read energy and forces """ 

110 self.read_energy() 

111 self.read_forces() 

112 

113 def write_coordinates(self, atoms, filename=''): 

114 """ write amber coordinates in netCDF format, 

115 only rectangular unit cells are allowed""" 

116 if filename == '': 

117 filename = self.incoordfile 

118 fout = netcdf.netcdf_file(filename, 'w') 

119 # dimension 

120 fout.Conventions = 'AMBERRESTART' 

121 fout.ConventionVersion = "1.0" 

122 fout.title = 'Ase-generated-amber-restart-file' 

123 fout.application = "AMBER" 

124 fout.program = "ASE" 

125 fout.programVersion = "1.0" 

126 fout.createDimension('cell_spatial', 3) 

127 fout.createDimension('label', 5) 

128 fout.createDimension('cell_angular', 3) 

129 fout.createDimension('time', 1) 

130 time = fout.createVariable('time', 'd', ('time',)) 

131 time.units = 'picosecond' 

132 time[0] = 0 

133 fout.createDimension('spatial', 3) 

134 spatial = fout.createVariable('spatial', 'c', ('spatial',)) 

135 spatial[:] = np.asarray(list('xyz')) 

136 # spatial = 'xyz' 

137 

138 natom = len(atoms) 

139 fout.createDimension('atom', natom) 

140 coordinates = fout.createVariable('coordinates', 'd', 

141 ('atom', 'spatial')) 

142 coordinates.units = 'angstrom' 

143 coordinates[:] = atoms.get_positions()[:] 

144 

145 if atoms.get_velocities() is not None: 

146 velocities = fout.createVariable('velocities', 'd', 

147 ('atom', 'spatial')) 

148 velocities.units = 'angstrom/picosecond' 

149 velocities[:] = atoms.get_velocities()[:] 

150 

151 # title 

152 cell_angular = fout.createVariable('cell_angular', 'c', 

153 ('cell_angular', 'label')) 

154 cell_angular[0] = np.asarray(list('alpha')) 

155 cell_angular[1] = np.asarray(list('beta ')) 

156 cell_angular[2] = np.asarray(list('gamma')) 

157 

158 # title 

159 cell_spatial = fout.createVariable('cell_spatial', 'c', 

160 ('cell_spatial',)) 

161 cell_spatial[0], cell_spatial[1], cell_spatial[2] = 'a', 'b', 'c' 

162 

163 # data 

164 cell_lengths = fout.createVariable('cell_lengths', 'd', 

165 ('cell_spatial',)) 

166 cell_lengths.units = 'angstrom' 

167 cell_lengths[0] = atoms.get_cell()[0, 0] 

168 cell_lengths[1] = atoms.get_cell()[1, 1] 

169 cell_lengths[2] = atoms.get_cell()[2, 2] 

170 

171 cell_angles = fout.createVariable('cell_angles', 'd', 

172 ('cell_angular',)) 

173 box_alpha, box_beta, box_gamma = 90.0, 90.0, 90.0 

174 cell_angles[0] = box_alpha 

175 cell_angles[1] = box_beta 

176 cell_angles[2] = box_gamma 

177 

178 cell_angles.units = 'degree' 

179 fout.close() 

180 

181 def read_coordinates(self, atoms, filename=''): 

182 """Import AMBER16 netCDF restart files. 

183 

184 Reads atom positions and 

185 velocities (if available), 

186 and unit cell (if available) 

187 

188 This may be useful if you have run amber many steps and 

189 want to read new positions and velocities 

190 """ 

191 

192 if filename == '': 

193 filename = self.outcoordfile 

194 

195 from scipy.io import netcdf 

196 import numpy as np 

197 import ase.units as units 

198 

199 fin = netcdf.netcdf_file(filename, 'r') 

200 all_coordinates = fin.variables['coordinates'][:] 

201 get_last_frame = False 

202 if hasattr(all_coordinates, 'ndim'): 

203 if all_coordinates.ndim == 3: 

204 get_last_frame = True 

205 elif hasattr(all_coordinates, 'shape'): 

206 if len(all_coordinates.shape) == 3: 

207 get_last_frame = True 

208 if get_last_frame: 

209 all_coordinates = all_coordinates[-1] 

210 atoms.set_positions(all_coordinates) 

211 if 'velocities' in fin.variables: 

212 all_velocities = fin.variables['velocities'][:] / (1000 * units.fs) 

213 if get_last_frame: 

214 all_velocities = all_velocities[-1] 

215 atoms.set_velocities(all_velocities) 

216 if 'cell_lengths' in fin.variables: 

217 all_abc = fin.variables['cell_lengths'] 

218 if get_last_frame: 

219 all_abc = all_abc[-1] 

220 a, b, c = all_abc 

221 all_angles = fin.variables['cell_angles'] 

222 if get_last_frame: 

223 all_angles = all_angles[-1] 

224 alpha, beta, gamma = all_angles 

225 

226 if (all(angle > 89.99 for angle in [alpha, beta, gamma]) and 

227 all(angle < 90.01 for angle in [alpha, beta, gamma])): 

228 atoms.set_cell( 

229 np.array([[a, 0, 0], 

230 [0, b, 0], 

231 [0, 0, c]])) 

232 atoms.set_pbc(True) 

233 else: 

234 raise NotImplementedError('only rectangular cells are' 

235 ' implemented in ASE-AMBER') 

236 

237 else: 

238 atoms.set_pbc(False) 

239 

240 def read_energy(self, filename='mden'): 

241 """ read total energy from amber file """ 

242 with open(filename, 'r') as fd: 

243 lines = fd.readlines() 

244 self.results['energy'] = \ 

245 float(lines[16].split()[2]) * units.kcal / units.mol 

246 

247 def read_forces(self, filename='mdfrc'): 

248 """ read forces from amber file """ 

249 fd = netcdf.netcdf_file(filename, 'r') 

250 try: 

251 forces = fd.variables['forces'] 

252 self.results['forces'] = forces[-1, :, :] \ 

253 / units.Ang * units.kcal / units.mol 

254 finally: 

255 fd.close() 

256 

257 def set_charges(self, selection, charges, parmed_filename=None): 

258 """ Modify amber topology charges to contain the updated 

259 QM charges, needed in QM/MM. 

260 Using amber's parmed program to change charges. 

261 """ 

262 qm_list = list(selection) 

263 with open(parmed_filename, 'w') as fout: 

264 fout.write('# update the following QM charges \n') 

265 for i, charge in zip(qm_list, charges): 

266 fout.write('change charge @' + str(i + 1) + ' ' + 

267 str(charge) + ' \n') 

268 fout.write('# Output the topology file \n') 

269 fout.write('outparm ' + self.topologyfile + ' \n') 

270 parmed_command = ('parmed -O -i ' + parmed_filename + 

271 ' -p ' + self.topologyfile + 

272 ' > ' + self.topologyfile + '.log 2>&1') 

273 subprocess.check_call(parmed_command, shell=True, cwd=self.directory) 

274 

275 def get_virtual_charges(self, atoms): 

276 with open(self.topologyfile, 'r') as fd: 

277 topology = fd.readlines() 

278 for n, line in enumerate(topology): 

279 if '%FLAG CHARGE' in line: 

280 chargestart = n + 2 

281 lines1 = topology[chargestart:(chargestart 

282 + (len(atoms) - 1) // 5 + 1)] 

283 mm_charges = [] 

284 for line in lines1: 

285 for el in line.split(): 

286 mm_charges.append(float(el) / 18.2223) 

287 charges = np.array(mm_charges) 

288 return charges 

289 

290 def add_virtual_sites(self, positions): 

291 return positions # no virtual sites 

292 

293 def redistribute_forces(self, forces): 

294 return forces 

295 

296 

297def map(atoms, top): 

298 p = np.zeros((2, len(atoms)), dtype="int") 

299 

300 elements = atoms.get_chemical_symbols() 

301 unique_elements = np.unique(atoms.get_chemical_symbols()) 

302 

303 for i in range(len(unique_elements)): 

304 idx = 0 

305 for j in range(len(atoms)): 

306 if elements[j] == unique_elements[i]: 

307 idx += 1 

308 symbol = unique_elements[i] + np.str(idx) 

309 for k in range(len(atoms)): 

310 if top.atoms[k].name == symbol: 

311 p[0, k] = j 

312 p[1, j] = k 

313 break 

314 return p 

315 

316 

317try: 

318 import sander 

319 have_sander = True 

320except ImportError: 

321 have_sander = False 

322 

323 

324class SANDER(Calculator): 

325 """ 

326 Interface to SANDER using Python interface 

327 

328 Requires sander Python bindings from http://ambermd.org/ 

329 """ 

330 implemented_properties = ['energy', 'forces'] 

331 

332 def __init__(self, atoms=None, label=None, top=None, crd=None, 

333 mm_options=None, qm_options=None, permutation=None, **kwargs): 

334 if not have_sander: 

335 raise RuntimeError("sander Python module could not be imported!") 

336 Calculator.__init__(self, label, atoms) 

337 self.permutation = permutation 

338 if qm_options is not None: 

339 sander.setup(top, crd.coordinates, crd.box, mm_options, qm_options) 

340 else: 

341 sander.setup(top, crd.coordinates, crd.box, mm_options) 

342 

343 def calculate(self, atoms, properties, system_changes): 

344 Calculator.calculate(self, atoms, properties, system_changes) 

345 if system_changes: 

346 if 'energy' in self.results: 

347 del self.results['energy'] 

348 if 'forces' in self.results: 

349 del self.results['forces'] 

350 if 'energy' not in self.results: 

351 if self.permutation is None: 

352 crd = np.reshape(atoms.get_positions(), (1, len(atoms), 3)) 

353 else: 

354 crd = np.reshape(atoms.get_positions() 

355 [self.permutation[0, :]], (1, len(atoms), 3)) 

356 sander.set_positions(crd) 

357 e, f = sander.energy_forces() 

358 self.results['energy'] = e.tot * units.kcal / units.mol 

359 if self.permutation is None: 

360 self.results['forces'] = (np.reshape(np.array(f), 

361 (len(atoms), 3)) * 

362 units.kcal / units.mol) 

363 else: 

364 ff = np.reshape(np.array(f), (len(atoms), 3)) * \ 

365 units.kcal / units.mol 

366 self.results['forces'] = ff[self.permutation[1, :]] 

367 if 'forces' not in self.results: 

368 if self.permutation is None: 

369 crd = np.reshape(atoms.get_positions(), (1, len(atoms), 3)) 

370 else: 

371 crd = np.reshape(atoms.get_positions()[self.permutation[0, :]], 

372 (1, len(atoms), 3)) 

373 sander.set_positions(crd) 

374 e, f = sander.energy_forces() 

375 self.results['energy'] = e.tot * units.kcal / units.mol 

376 if self.permutation is None: 

377 self.results['forces'] = (np.reshape(np.array(f), 

378 (len(atoms), 3)) * 

379 units.kcal / units.mol) 

380 else: 

381 ff = np.reshape(np.array(f), (len(atoms), 3)) * \ 

382 units.kcal / units.mol 

383 self.results['forces'] = ff[self.permutation[1, :]]