1""" 

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

3 

4""" 

5from io import StringIO 

6from ase.calculators.calculator import Calculator, all_changes 

7from ase.calculators.calculator import InputError, ReadError 

8from ase.calculators.calculator import CalculatorSetupError 

9import multiprocessing 

10from ase import io 

11import numpy as np 

12import json 

13from ase.units import Bohr, Hartree 

14import warnings 

15import os 

16 

17 

18class Psi4(Calculator): 

19 """ 

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

21 psi4. 

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

23 and quantum chemistry method implemented in psi4 can be input 

24 (i.e. ccsd(t)) 

25 

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

27 calc.psi4 

28 """ 

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

30 discard_results_on_any_change = True 

31 

32 default_parameters = { 

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

34 "method": "hf", 

35 'symmetry': 'c1'} 

36 

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

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

39 **kwargs): 

40 Calculator.__init__(self, restart=restart, 

41 ignore_bad_restart=ignore_bad_restart, label=label, 

42 atoms=atoms, command=command, **kwargs) 

43 import psi4 

44 self.psi4 = psi4 

45 # perform initial setup of psi4 python API 

46 self.set_psi4(atoms=atoms) 

47 

48 def set_psi4(self, atoms=None): 

49 """ 

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

51 defines 

52 """ 

53 

54 # Set the scrath directory for electronic structure files. 

55 # The default is /tmp 

56 if 'PSI_SCRATCH' in os.environ: 

57 pass 

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

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

60 

61 # Input spin settings 

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

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

64 self.parameters['reference']}) 

65 # Memory 

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

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

68 

69 # Threads 

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

71 if nthreads == 'max': 

72 nthreads = multiprocessing.cpu_count() 

73 self.psi4.set_num_threads(nthreads) 

74 

75 # deal with some ASE specific inputs 

76 if 'kpts' in self.parameters: 

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

78 ' require k-points. Please remove this ' 

79 'argument.') 

80 

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

82 # svwn is equivalent to LDA 

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

84 

85 if 'nbands' in self.parameters: 

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

87 

88 if 'smearing' in self.parameters: 

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

90 ' psi4 currently, thus a smearing argument ' 

91 'cannot be utilized. please remove this ' 

92 'argument') 

93 

94 if 'xc' in self.parameters: 

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

96 ' use the `method` argument instead') 

97 

98 if atoms is None: 

99 if self.atoms is None: 

100 return None 

101 else: 

102 atoms = self.atoms 

103 if self.atoms is None: 

104 self.atoms = atoms 

105 # Generate the atomic input 

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

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

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

109 geom.append('units angstrom') 

110 

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

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

113 if mult is None: 

114 mult = 1 

115 if charge is not None: 

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

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

118 if charge is None: 

119 charge = 0 

120 

121 geom.append('{} {}'.format(charge, mult)) 

122 geom.append('no_reorient') 

123 

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

125 os.mkdir(self.directory) 

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

127 

128 def read(self, label): 

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

130 """ 

131 filename = label + '.dat' 

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

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

134 

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

136 txt = fd.read() 

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

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

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

140 ' lines inserted by this calculator. This likely' 

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

142 'ASE calculator or has since been modified and ' 

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

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

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

146 saved_dict = json.loads(info) 

147 # use io read to recode atoms 

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

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

150 self.parameters = saved_dict['parameters'] 

151 self.results = saved_dict['results'] 

152 # convert forces to numpy array 

153 if 'forces' in self.results: 

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

155 

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

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

158 

159 Calculator.calculate(self, atoms=atoms) 

160 if self.atoms is None: 

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

162 'perform a calculation') 

163 atoms = self.atoms 

164 

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

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

167 self.parameters['multiplicity'] = None 

168 # this inputs all the settings into psi4 

169 self.set_psi4(atoms=atoms) 

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

171 False) 

172 

173 # Set up the method 

174 method = self.parameters['method'] 

175 basis = self.parameters['basis'] 

176 

177 # Do the calculations 

178 if 'forces' in properties: 

179 grad, wf = self.psi4.driver.gradient('{}/{}'.format(method, basis), 

180 return_wfn=True) 

181 # energy comes for free 

182 energy = wf.energy() 

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

184 # convert to eV/A 

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

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

187 elif 'energy' in properties: 

188 energy = self.psi4.energy('{}/{}'.format(method, basis), 

189 molecule=self.molecule) 

190 # convert to eV 

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

192 

193 # dump the calculator info to the psi4 file 

194 save_atoms = self.atoms.copy() 

195 # use io.write to encode atoms 

196 with StringIO() as fd: 

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

198 json_atoms = fd.getvalue() 

199 # convert forces to list for json storage 

200 save_results = self.results.copy() 

201 if 'forces' in save_results: 

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

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

204 'results': save_results, 

205 'atoms': json_atoms} 

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

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

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