Coverage for /builds/kinetik161/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 (Calculator, CalculatorSetupError, 

15 InputError, ReadError, all_changes) 

16from ase.config import cfg 

17from ase.units import Bohr, Hartree 

18 

19 

20class Psi4(Calculator): 

21 """ 

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

23 psi4. 

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

25 and quantum chemistry method implemented in psi4 can be input 

26 (i.e. ccsd(t)) 

27 

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

29 calc.psi4 

30 """ 

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

32 discard_results_on_any_change = True 

33 

34 default_parameters = { 

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

36 "method": "hf", 

37 'symmetry': 'c1'} 

38 

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

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

41 **kwargs): 

42 Calculator.__init__(self, restart=restart, 

43 ignore_bad_restart=ignore_bad_restart, label=label, 

44 atoms=atoms, command=command, **kwargs) 

45 import psi4 

46 self.psi4 = psi4 

47 # perform initial setup of psi4 python API 

48 self.set_psi4(atoms=atoms) 

49 

50 def set_psi4(self, atoms=None): 

51 """ 

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

53 defines 

54 """ 

55 

56 # Set the scrath directory for electronic structure files. 

57 # The default is /tmp 

58 if 'PSI_SCRATCH' in cfg: 

59 pass 

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

61 # XXX We should probably not be setting envvars except 

62 # if we are creating new processes. 

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

64 

65 # Input spin settings 

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

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

68 self.parameters['reference']}) 

69 # Memory 

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

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

72 

73 # Threads 

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

75 if nthreads == 'max': 

76 nthreads = multiprocessing.cpu_count() 

77 self.psi4.set_num_threads(nthreads) 

78 

79 # deal with some ASE specific inputs 

80 if 'kpts' in self.parameters: 

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

82 ' require k-points. Please remove this ' 

83 'argument.') 

84 

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

86 # svwn is equivalent to LDA 

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

88 

89 if 'nbands' in self.parameters: 

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

91 

92 if 'smearing' in self.parameters: 

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

94 ' psi4 currently, thus a smearing argument ' 

95 'cannot be utilized. please remove this ' 

96 'argument') 

97 

98 if 'xc' in self.parameters: 

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

100 ' use the `method` argument instead') 

101 

102 if atoms is None: 

103 if self.atoms is None: 

104 return None 

105 else: 

106 atoms = self.atoms 

107 if self.atoms is None: 

108 self.atoms = atoms 

109 # Generate the atomic input 

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

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

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

113 geom.append('units angstrom') 

114 

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

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

117 if mult is None: 

118 mult = 1 

119 if charge is not None: 

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

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

122 if charge is None: 

123 charge = 0 

124 

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

126 geom.append('no_reorient') 

127 

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

129 os.mkdir(self.directory) 

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

131 

132 def read(self, label): 

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

134 """ 

135 filename = label + '.dat' 

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

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

138 

139 with open(filename) as fd: 

140 txt = fd.read() 

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

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

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

144 ' lines inserted by this calculator. This likely' 

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

146 'ASE calculator or has since been modified and ' 

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

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

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

150 saved_dict = json.loads(info) 

151 # use io read to recode atoms 

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

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

154 self.parameters = saved_dict['parameters'] 

155 self.results = saved_dict['results'] 

156 # convert forces to numpy array 

157 if 'forces' in self.results: 

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

159 

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

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

162 

163 Calculator.calculate(self, atoms=atoms) 

164 if self.atoms is None: 

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

166 'perform a calculation') 

167 atoms = self.atoms 

168 

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

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

171 self.parameters['multiplicity'] = None 

172 # this inputs all the settings into psi4 

173 self.set_psi4(atoms=atoms) 

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

175 False) 

176 

177 # Set up the method 

178 method = self.parameters['method'] 

179 basis = self.parameters['basis'] 

180 

181 # Do the calculations 

182 if 'forces' in properties: 

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

184 return_wfn=True) 

185 # energy comes for free 

186 energy = wf.energy() 

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

188 # convert to eV/A 

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

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

191 elif 'energy' in properties: 

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

193 molecule=self.molecule) 

194 # convert to eV 

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

196 

197 # dump the calculator info to the psi4 file 

198 save_atoms = self.atoms.copy() 

199 # use io.write to encode atoms 

200 with StringIO() as fd: 

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

202 json_atoms = fd.getvalue() 

203 # convert forces to list for json storage 

204 save_results = self.results.copy() 

205 if 'forces' in save_results: 

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

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

208 'results': save_results, 

209 'atoms': json_atoms} 

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

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

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