1import os 

2import numpy as np 

3from copy import deepcopy 

4 

5from ase.calculators.calculator import KPoints, kpts2kpts 

6 

7_special_kws = ['center', 'autosym', 'autoz', 'theory', 'basis', 'xc', 'task', 

8 'set', 'symmetry', 'label', 'geompar', 'basispar', 'kpts', 

9 'bandpath', 'restart_kw'] 

10 

11_system_type = {1: 'polymer', 2: 'surface', 3: 'crystal'} 

12 

13 

14def _get_geom(atoms, **params): 

15 geom_header = ['geometry units angstrom'] 

16 for geomkw in ['center', 'autosym', 'autoz']: 

17 geom_header.append(geomkw if params.get(geomkw) else 'no' + geomkw) 

18 if 'geompar' in params: 

19 geom_header.append(params['geompar']) 

20 geom = [' '.join(geom_header)] 

21 

22 outpos = atoms.get_positions() 

23 pbc = atoms.pbc 

24 if np.any(pbc): 

25 scpos = atoms.get_scaled_positions() 

26 for i, pbci in enumerate(pbc): 

27 if pbci: 

28 outpos[:, i] = scpos[:, i] 

29 npbc = pbc.sum() 

30 cellpars = atoms.cell.cellpar() 

31 geom.append(' system {} units angstrom'.format(_system_type[npbc])) 

32 if npbc == 3: 

33 geom.append(' lattice_vectors') 

34 for row in atoms.cell: 

35 geom.append(' {:20.16e} {:20.16e} {:20.16e}'.format(*row)) 

36 else: 

37 if pbc[0]: 

38 geom.append(' lat_a {:20.16e}'.format(cellpars[0])) 

39 if pbc[1]: 

40 geom.append(' lat_b {:20.16e}'.format(cellpars[1])) 

41 if pbc[2]: 

42 geom.append(' lat_c {:20.16e}'.format(cellpars[2])) 

43 if pbc[1] and pbc[2]: 

44 geom.append(' alpha {:20.16e}'.format(cellpars[3])) 

45 if pbc[0] and pbc[2]: 

46 geom.append(' beta {:20.16e}'.format(cellpars[4])) 

47 if pbc[1] and pbc[0]: 

48 geom.append(' gamma {:20.16e}'.format(cellpars[5])) 

49 geom.append(' end') 

50 

51 for i, atom in enumerate(atoms): 

52 geom.append(' {:<2} {:20.16e} {:20.16e} {:20.16e}' 

53 ''.format(atom.symbol, *outpos[i])) 

54 symm = params.get('symmetry') 

55 if symm is not None: 

56 geom.append(' symmetry {}'.format(symm)) 

57 geom.append('end') 

58 return geom 

59 

60 

61def _get_basis(theory, **params): 

62 if 'basis' not in params: 

63 if theory in ['pspw', 'band', 'paw']: 

64 return [] 

65 basis_in = params.get('basis', '3-21G') 

66 if 'basispar' in params: 

67 header = 'basis {} noprint'.format(params['basispar']) 

68 else: 

69 header = 'basis noprint' 

70 basis_out = [header] 

71 if isinstance(basis_in, str): 

72 basis_out.append(' * library {}'.format(basis_in)) 

73 else: 

74 for symbol, ibasis in basis_in.items(): 

75 basis_out.append('{:>4} library {}'.format(symbol, ibasis)) 

76 basis_out.append('end') 

77 return basis_out 

78 

79 

80_special_keypairs = [('nwpw', 'simulation_cell'), 

81 ('nwpw', 'carr-parinello'), 

82 ('nwpw', 'brillouin_zone'), 

83 ('tddft', 'grad'), 

84 ] 

85 

86 

87def _format_brillouin_zone(array, name=None): 

88 out = [' brillouin_zone'] 

89 if name is not None: 

90 out += [' zone_name {}'.format(name)] 

91 template = ' kvector' + ' {:20.16e}' * array.shape[1] 

92 for row in array: 

93 out.append(template.format(*row)) 

94 out.append(' end') 

95 return out 

96 

97 

98def _get_bandpath(bp): 

99 if bp is None: 

100 return [] 

101 out = ['nwpw'] 

102 out += _format_brillouin_zone(bp.kpts, name=bp.path) 

103 out += [' zone_structure_name {}'.format(bp.path), 

104 'end', 

105 'task band structure'] 

106 return out 

107 

108 

109def _format_line(key, val): 

110 if val is None: 

111 return key 

112 if isinstance(val, bool): 

113 return '{} .{}.'.format(key, str(val).lower()) 

114 else: 

115 return ' '.join([key, str(val)]) 

116 

117 

118def _format_block(key, val, nindent=0): 

119 prefix = ' ' * nindent 

120 prefix2 = ' ' * (nindent + 1) 

121 if val is None: 

122 return [prefix + key] 

123 

124 if not isinstance(val, dict): 

125 return [prefix + _format_line(key, val)] 

126 

127 out = [prefix + key] 

128 for subkey, subval in val.items(): 

129 if (key, subkey) in _special_keypairs: 

130 if (key, subkey) == ('nwpw', 'brillouin_zone'): 

131 out += _format_brillouin_zone(subval) 

132 else: 

133 out += _format_block(subkey, subval, nindent + 1) 

134 else: 

135 if isinstance(subval, dict): 

136 subval = ' '.join([_format_line(a, b) 

137 for a, b in subval.items()]) 

138 out.append(prefix2 + ' '.join([_format_line(subkey, subval)])) 

139 out.append(prefix + 'end') 

140 return out 

141 

142 

143def _get_other(**params): 

144 out = [] 

145 for kw, block in params.items(): 

146 if kw in _special_kws: 

147 continue 

148 out += _format_block(kw, block) 

149 return out 

150 

151 

152def _get_set(**params): 

153 return ['set ' + _format_line(key, val) for key, val in params.items()] 

154 

155 

156_gto_theories = ['tce', 'ccsd', 'mp2', 'tddft', 'scf', 'dft'] 

157_pw_theories = ['band', 'pspw', 'paw'] 

158_all_theories = _gto_theories + _pw_theories 

159 

160 

161def _get_theory(**params): 

162 # Default: user-provided theory 

163 theory = params.get('theory') 

164 if theory is not None: 

165 return theory 

166 

167 # Check if the user passed a theory to xc 

168 xc = params.get('xc') 

169 if xc in _all_theories: 

170 return xc 

171 

172 # Check for input blocks that correspond to a particular level of 

173 # theory. Correlated theories (e.g. CCSD) are checked first. 

174 for kw in _gto_theories: 

175 if kw in params: 

176 return kw 

177 

178 # If the user passed an 'nwpw' block, then they want a plane-wave 

179 # calculation, but what kind? If they request k-points, then 

180 # they want 'band', otherwise assume 'pspw' (if the user wants 

181 # to use 'paw', they will have to ask for it specifically). 

182 nwpw = params.get('nwpw') 

183 if nwpw is not None: 

184 if 'monkhorst-pack' in nwpw or 'brillouin_zone' in nwpw: 

185 return 'band' 

186 return 'pspw' 

187 

188 # When all else fails, default to dft. 

189 return 'dft' 

190 

191 

192_xc_conv = dict(lda='slater pw91lda', 

193 pbe='xpbe96 cpbe96', 

194 revpbe='revpbe cpbe96', 

195 rpbe='rpbe cpbe96', 

196 pw91='xperdew91 perdew91', 

197 ) 

198 

199 

200def _update_mult(magmom_tot, **params): 

201 theory = params['theory'] 

202 if magmom_tot == 0: 

203 magmom_mult = 1 

204 else: 

205 magmom_mult = np.sign(magmom_tot) * (abs(magmom_tot) + 1) 

206 if 'scf' in params: 

207 for kw in ['nopen', 'singlet', 'doublet', 'triplet', 'quartet', 

208 'quintet', 'sextet', 'septet', 'octet']: 

209 if kw in params['scf']: 

210 break 

211 else: 

212 params['scf']['nopen'] = magmom_tot 

213 elif theory in ['scf', 'mp2', 'ccsd', 'tce']: 

214 params['scf'] = dict(nopen=magmom_tot) 

215 

216 if 'dft' in params: 

217 if 'mult' not in params['dft']: 

218 params['dft']['mult'] = magmom_mult 

219 elif theory in ['dft', 'tddft']: 

220 params['dft'] = dict(mult=magmom_mult) 

221 

222 if 'nwpw' in params: 

223 if 'mult' not in params['nwpw']: 

224 params['nwpw']['mult'] = magmom_mult 

225 elif theory in ['pspw', 'band', 'paw']: 

226 params['nwpw'] = dict(mult=magmom_mult) 

227 

228 return params 

229 

230 

231def _get_kpts(atoms, **params): 

232 """Converts top-level 'kpts' argument to native keywords""" 

233 kpts = params.get('kpts') 

234 if kpts is None: 

235 return params 

236 

237 nwpw = params.get('nwpw', dict()) 

238 

239 if 'monkhorst-pack' in nwpw or 'brillouin_zone' in nwpw: 

240 raise ValueError("Redundant k-points specified!") 

241 

242 if isinstance(kpts, KPoints): 

243 nwpw['brillouin_zone'] = kpts.kpts 

244 elif isinstance(kpts, dict): 

245 if kpts.get('gamma', False) or 'size' not in kpts: 

246 nwpw['brillouin_zone'] = kpts2kpts(kpts, atoms).kpts 

247 else: 

248 nwpw['monkhorst-pack'] = ' '.join(map(str, kpts['size'])) 

249 elif isinstance(kpts, np.ndarray): 

250 nwpw['brillouin_zone'] = kpts 

251 else: 

252 nwpw['monkhorst-pack'] = ' '.join(map(str, kpts)) 

253 

254 params['nwpw'] = nwpw 

255 return params 

256 

257 

258def write_nwchem_in(fd, atoms, properties=None, echo=False, **params): 

259 """Writes NWChem input file. 

260 

261 Parameters 

262 ---------- 

263 fd 

264 file descriptor 

265 atoms 

266 atomic configuration 

267 properties 

268 list of properties to compute; by default only the 

269 calculation of the energy is requested 

270 echo 

271 if True include the `echo` keyword at the top of the file, 

272 which causes the content of the input file to be included 

273 in the output file 

274 params 

275 dict of instructions blocks to be included 

276 """ 

277 params = deepcopy(params) 

278 

279 if properties is None: 

280 properties = ['energy'] 

281 

282 if 'stress' in properties: 

283 if 'set' not in params: 

284 params['set'] = dict() 

285 params['set']['includestress'] = True 

286 

287 task = params.get('task') 

288 if task is None: 

289 if 'stress' in properties or 'forces' in properties: 

290 task = 'gradient' 

291 else: 

292 task = 'energy' 

293 

294 params = _get_kpts(atoms, **params) 

295 

296 theory = _get_theory(**params) 

297 params['theory'] = theory 

298 xc = params.get('xc') 

299 if 'xc' in params: 

300 xc = _xc_conv.get(params['xc'].lower(), params['xc']) 

301 if theory in ['dft', 'tddft']: 

302 if 'dft' not in params: 

303 params['dft'] = dict() 

304 params['dft']['xc'] = xc 

305 elif theory in ['pspw', 'band', 'paw']: 

306 if 'nwpw' not in params: 

307 params['nwpw'] = dict() 

308 params['nwpw']['xc'] = xc 

309 

310 magmom_tot = int(atoms.get_initial_magnetic_moments().sum()) 

311 params = _update_mult(magmom_tot, **params) 

312 

313 label = params.get('label', 'nwchem') 

314 perm = os.path.abspath(params.pop('perm', label)) 

315 scratch = os.path.abspath(params.pop('scratch', label)) 

316 restart_kw = params.get('restart_kw', 'start') 

317 if restart_kw not in ('start', 'restart'): 

318 raise ValueError("Unrecognised restart keyword: {}!" 

319 .format(restart_kw)) 

320 short_label = label.rsplit('/', 1)[-1] 

321 if echo: 

322 out = ['echo'] 

323 else: 

324 out = [] 

325 out.extend(['title "{}"'.format(short_label), 

326 'permanent_dir {}'.format(perm), 

327 'scratch_dir {}'.format(scratch), 

328 '{} {}'.format(restart_kw, short_label), 

329 '\n'.join(_get_geom(atoms, **params)), 

330 '\n'.join(_get_basis(**params)), 

331 '\n'.join(_get_other(**params)), 

332 '\n'.join(_get_set(**params.get('set', dict()))), 

333 'task {} {}'.format(theory, task), 

334 '\n'.join(_get_bandpath(params.get('bandpath', None)))]) 

335 

336 fd.write('\n\n'.join(out))