Coverage for /builds/ase/ase/ase/calculators/psi4.py: 15.93%

1# fmt: off 

2 

3""" 

4authors: Ben Comer (Georgia Tech), Xiangyun (Ray) Lei (Georgia Tech) 

5 

6""" 

7import json 

8import multiprocessing 

9import os 

10import warnings 

11from io import StringIO 

12 

13import numpy as np 

14 

15from ase import io 

16from ase.calculators.calculator import ( 

17 Calculator, 

18 CalculatorSetupError, 

19 InputError, 

20 ReadError, 

21 all_changes, 

22) 

23from ase.config import cfg 

24from ase.units import Bohr, Hartree 

25 

26 

27class Psi4(Calculator): 

28 """ 

29 An ase calculator for the popular open source Q-chem code 

30 psi4. 

31 method is the generic input for whatever method you wish to use, thus 

32 and quantum chemistry method implemented in psi4 can be input 

33 (i.e. ccsd(t)) 

34 

35 also note that you can always use the in-built psi4 module through: 

36 calc.psi4 

37 """ 

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

39 discard_results_on_any_change = True 

40 

41 default_parameters = { 

42 "basis": "aug-cc-pvtz", 

43 "method": "hf", 

44 'symmetry': 'c1'} 

45 

46 def __init__(self, restart=None, ignore_bad_restart=False, 

47 label='psi4-calc', atoms=None, command=None, 

48 **kwargs): 

49 Calculator.__init__(self, restart=restart, 

50 ignore_bad_restart=ignore_bad_restart, label=label, 

51 atoms=atoms, command=command, **kwargs) 

52 import psi4 

53 self.psi4 = psi4 

54 # perform initial setup of psi4 python API 

55 self.set_psi4(atoms=atoms) 

56 

57 def set_psi4(self, atoms=None): 

58 """ 

59 This function sets the imported psi4 module to the settings the user 

60 defines 

61 """ 

62 

63 # Set the scrath directory for electronic structure files. 

64 # The default is /tmp 

65 if 'PSI_SCRATCH' in cfg: 

66 pass 

67 elif self.parameters.get('PSI_SCRATCH'): 

68 # XXX We should probably not be setting envvars except 

69 # if we are creating new processes. 

70 os.environ['PSI_SCRATCH'] = self.parameters['PSI_SCRATCH'] 

71 

72 # Input spin settings 

73 if self.parameters.get('reference') is not None: 

74 self.psi4.set_options({'reference': 

75 self.parameters['reference']}) 

76 # Memory 

77 if self.parameters.get('memory') is not None: 

78 self.psi4.set_memory(self.parameters['memory']) 

79 

80 # Threads 

81 nthreads = self.parameters.get('num_threads', 1) 

82 if nthreads == 'max': 

83 nthreads = multiprocessing.cpu_count() 

84 self.psi4.set_num_threads(nthreads) 

85 

86 # deal with some ASE specific inputs 

87 if 'kpts' in self.parameters: 

88 raise InputError('psi4 is a non-periodic code, and thus does not' 

89 ' require k-points. Please remove this ' 

90 'argument.') 

91 

92 if self.parameters['method'] == 'LDA': 

93 # svwn is equivalent to LDA 

94 self.parameters['method'] = 'svwn' 

95 

96 if 'nbands' in self.parameters: 

97 raise InputError('psi4 does not support the keyword "nbands"') 

98 

99 if 'smearing' in self.parameters: 

100 raise InputError('Finite temperature DFT is not implemented in' 

101 ' psi4 currently, thus a smearing argument ' 

102 'cannot be utilized. please remove this ' 

103 'argument') 

104 

105 if 'xc' in self.parameters: 

106 raise InputError('psi4 does not accept the `xc` argument please' 

107 ' use the `method` argument instead') 

108 

109 if atoms is None: 

110 if self.atoms is None: 

111 return 

112 else: 

113 atoms = self.atoms 

114 if self.atoms is None: 

115 self.atoms = atoms 

116 # Generate the atomic input 

117 geomline = '{}\t{:.15f}\t{:.15f}\t{:.15f}' 

118 geom = [geomline.format(atom.symbol, *atom.position) for atom in atoms] 

119 geom.append('symmetry {}'.format(self.parameters['symmetry'])) 

120 geom.append('units angstrom') 

121 

122 charge = self.parameters.get('charge') 

123 mult = self.parameters.get('multiplicity') 

124 if mult is None: 

125 mult = 1 

126 if charge is not None: 

127 warnings.warn('A charge was provided without a spin ' 

128 'multiplicity. A multiplicity of 1 is assumed') 

129 if charge is None: 

130 charge = 0 

131 

132 geom.append(f'{charge} {mult}') 

133 geom.append('no_reorient') 

134 

135 if not os.path.isdir(self.directory): 

136 os.mkdir(self.directory) 

137 self.molecule = self.psi4.geometry('\n'.join(geom)) 

138 

139 def read(self, label): 

140 """Read psi4 outputs made from this ASE calculator 

141 """ 

142 filename = label + '.dat' 

143 if not os.path.isfile(filename): 

144 raise ReadError('Could not find the psi4 output file: ' + filename) 

145 

146 with open(filename) as fd: 

147 txt = fd.read() 

148 if '!ASE Information\n' not in txt: 

149 raise Exception('The output file {} could not be read because ' 

150 'the file does not contain the "!ASE Information"' 

151 ' lines inserted by this calculator. This likely' 

152 ' means the output file was not made using this ' 

153 'ASE calculator or has since been modified and ' 

154 'thus cannot be read.'.format(filename)) 

155 info = txt.split('!ASE Information\n')[1] 

156 info = info.split('!')[0] 

157 saved_dict = json.loads(info) 

158 # use io read to recode atoms 

159 with StringIO(str(saved_dict['atoms'])) as g: 

160 self.atoms = io.read(g, format='json') 

161 self.parameters = saved_dict['parameters'] 

162 self.results = saved_dict['results'] 

163 # convert forces to numpy array 

164 if 'forces' in self.results: 

165 self.results['forces'] = np.array(self.results['forces']) 

166 

167 def calculate(self, atoms=None, properties=['energy'], 

168 system_changes=all_changes, symmetry='c1'): 

169 

170 Calculator.calculate(self, atoms=atoms) 

171 if self.atoms is None: 

172 raise CalculatorSetupError('An Atoms object must be provided to ' 

173 'perform a calculation') 

174 atoms = self.atoms 

175 

176 if atoms.get_initial_magnetic_moments().any(): 

177 self.parameters['reference'] = 'uhf' 

178 self.parameters['multiplicity'] = None 

179 # this inputs all the settings into psi4 

180 self.set_psi4(atoms=atoms) 

181 self.psi4.core.set_output_file(self.label + '.dat', 

182 False) 

183 

184 # Set up the method 

185 method = self.parameters['method'] 

186 basis = self.parameters['basis'] 

187 

188 # Do the calculations 

189 if 'forces' in properties: 

190 grad, wf = self.psi4.driver.gradient(f'{method}/{basis}', 

191 return_wfn=True) 

192 # energy comes for free 

193 energy = wf.energy() 

194 self.results['energy'] = energy * Hartree 

195 # convert to eV/A 

196 # also note that the gradient is -1 * forces 

197 self.results['forces'] = -1 * np.array(grad) * Hartree / Bohr 

198 elif 'energy' in properties: 

199 energy = self.psi4.energy(f'{method}/{basis}', 

200 molecule=self.molecule) 

201 # convert to eV 

202 self.results['energy'] = energy * Hartree 

203 

204 # dump the calculator info to the psi4 file 

205 save_atoms = self.atoms.copy() 

206 # use io.write to encode atoms 

207 with StringIO() as fd: 

208 io.write(fd, save_atoms, format='json') 

209 json_atoms = fd.getvalue() 

210 # convert forces to list for json storage 

211 save_results = self.results.copy() 

212 if 'forces' in save_results: 

213 save_results['forces'] = save_results['forces'].tolist() 

214 save_dict = {'parameters': self.parameters, 

215 'results': save_results, 

216 'atoms': json_atoms} 

217 self.psi4.core.print_out('!ASE Information\n') 

218 self.psi4.core.print_out(json.dumps(save_dict)) 

219 self.psi4.core.print_out('!')