Coverage for /builds/debichem-team/python-ase/ase/calculators/psi4.py: 15.93%

1""" 

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

3 

4""" 

5import json 

6import multiprocessing 

7import os 

8import warnings 

9from io import StringIO 

10 

11import numpy as np 

12 

13from ase import io 

14from ase.calculators.calculator import ( 

15 Calculator, 

16 CalculatorSetupError, 

17 InputError, 

18 ReadError, 

19 all_changes, 

20) 

21from ase.config import cfg 

22from ase.units import Bohr, Hartree 

23 

24 

25class Psi4(Calculator): 

26 """ 

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

28 psi4. 

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

30 and quantum chemistry method implemented in psi4 can be input 

31 (i.e. ccsd(t)) 

32 

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

34 calc.psi4 

35 """ 

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

37 discard_results_on_any_change = True 

38 

39 default_parameters = { 

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

41 "method": "hf", 

42 'symmetry': 'c1'} 

43 

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

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

46 **kwargs): 

47 Calculator.__init__(self, restart=restart, 

48 ignore_bad_restart=ignore_bad_restart, label=label, 

49 atoms=atoms, command=command, **kwargs) 

50 import psi4 

51 self.psi4 = psi4 

52 # perform initial setup of psi4 python API 

53 self.set_psi4(atoms=atoms) 

54 

55 def set_psi4(self, atoms=None): 

56 """ 

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

58 defines 

59 """ 

60 

61 # Set the scrath directory for electronic structure files. 

62 # The default is /tmp 

63 if 'PSI_SCRATCH' in cfg: 

64 pass 

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

66 # XXX We should probably not be setting envvars except 

67 # if we are creating new processes. 

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

69 

70 # Input spin settings 

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

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

73 self.parameters['reference']}) 

74 # Memory 

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

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

77 

78 # Threads 

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

80 if nthreads == 'max': 

81 nthreads = multiprocessing.cpu_count() 

82 self.psi4.set_num_threads(nthreads) 

83 

84 # deal with some ASE specific inputs 

85 if 'kpts' in self.parameters: 

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

87 ' require k-points. Please remove this ' 

88 'argument.') 

89 

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

91 # svwn is equivalent to LDA 

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

93 

94 if 'nbands' in self.parameters: 

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

96 

97 if 'smearing' in self.parameters: 

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

99 ' psi4 currently, thus a smearing argument ' 

100 'cannot be utilized. please remove this ' 

101 'argument') 

102 

103 if 'xc' in self.parameters: 

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

105 ' use the `method` argument instead') 

106 

107 if atoms is None: 

108 if self.atoms is None: 

109 return 

110 else: 

111 atoms = self.atoms 

112 if self.atoms is None: 

113 self.atoms = atoms 

114 # Generate the atomic input 

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

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

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

118 geom.append('units angstrom') 

119 

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

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

122 if mult is None: 

123 mult = 1 

124 if charge is not None: 

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

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

127 if charge is None: 

128 charge = 0 

129 

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

131 geom.append('no_reorient') 

132 

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

134 os.mkdir(self.directory) 

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

136 

137 def read(self, label): 

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

139 """ 

140 filename = label + '.dat' 

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

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

143 

144 with open(filename) as fd: 

145 txt = fd.read() 

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

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

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

149 ' lines inserted by this calculator. This likely' 

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

151 'ASE calculator or has since been modified and ' 

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

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

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

155 saved_dict = json.loads(info) 

156 # use io read to recode atoms 

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

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

159 self.parameters = saved_dict['parameters'] 

160 self.results = saved_dict['results'] 

161 # convert forces to numpy array 

162 if 'forces' in self.results: 

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

164 

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

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

167 

168 Calculator.calculate(self, atoms=atoms) 

169 if self.atoms is None: 

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

171 'perform a calculation') 

172 atoms = self.atoms 

173 

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

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

176 self.parameters['multiplicity'] = None 

177 # this inputs all the settings into psi4 

178 self.set_psi4(atoms=atoms) 

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

180 False) 

181 

182 # Set up the method 

183 method = self.parameters['method'] 

184 basis = self.parameters['basis'] 

185 

186 # Do the calculations 

187 if 'forces' in properties: 

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

189 return_wfn=True) 

190 # energy comes for free 

191 energy = wf.energy() 

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

193 # convert to eV/A 

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

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

196 elif 'energy' in properties: 

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

198 molecule=self.molecule) 

199 # convert to eV 

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

201 

202 # dump the calculator info to the psi4 file 

203 save_atoms = self.atoms.copy() 

204 # use io.write to encode atoms 

205 with StringIO() as fd: 

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

207 json_atoms = fd.getvalue() 

208 # convert forces to list for json storage 

209 save_results = self.results.copy() 

210 if 'forces' in save_results: 

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

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

213 'results': save_results, 

214 'atoms': json_atoms} 

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

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

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