Coverage for gpaw/solvation/hamiltonian.py: 98%

1from gpaw.hamiltonian import RealSpaceHamiltonian 

2from gpaw.solvation.poisson import WeightedFDPoissonSolver 

3from gpaw.fd_operators import Gradient 

4from gpaw.io.logger import indent 

5from ase.units import Ha 

6import numpy as np 

7 

8 

9class SolvationRealSpaceHamiltonian(RealSpaceHamiltonian): 

10 """Realspace Hamiltonian with continuum solvent model. 

11 

12 See also Section III of 

13 A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014). 

14 """ 

15 

16 def __init__( 

17 self, 

18 # solvation related arguments: 

19 cavity, dielectric, interactions, 

20 # RealSpaceHamiltonian arguments: 

21 gd, finegd, nspins, collinear, setups, timer, xc, world, 

22 redistributor, vext=None, psolver=None, 

23 stencil=3): 

24 """Constructor of SolvationRealSpaceHamiltonian class. 

25 

26 Additional arguments not present in RealSpaceHamiltonian: 

27 cavity -- A Cavity instance. 

28 dielectric -- A Dielectric instance. 

29 interactions -- A list of Interaction instances. 

30 """ 

31 self.cavity = cavity 

32 self.dielectric = dielectric 

33 self.interactions = interactions 

34 cavity.set_grid_descriptor(finegd) 

35 dielectric.set_grid_descriptor(finegd) 

36 for ia in interactions: 

37 ia.set_grid_descriptor(finegd) 

38 if psolver is None: 

39 psolver = WeightedFDPoissonSolver() 

40 psolver.set_dielectric(self.dielectric) 

41 self.gradient = None 

42 RealSpaceHamiltonian.__init__( 

43 self, 

44 gd, finegd, nspins, collinear, setups, timer, xc, world, 

45 vext=vext, psolver=psolver, 

46 stencil=stencil, redistributor=redistributor) 

47 

48 for ia in interactions: 

49 setattr(self, 'e_' + ia.subscript, None) 

50 self.new_atoms = None 

51 self.vt_ia_g = None 

52 self.e_el_free = None 

53 self.e_el_extrapolated = None 

54 

55 def __str__(self): 

56 s = RealSpaceHamiltonian.__str__(self) + '\n' 

57 s += ' Solvation:\n' 

58 components = [self.cavity, self.dielectric] + self.interactions 

59 s += ''.join([indent(str(c), 2) for c in components]) 

60 return s 

61 

62 def estimate_memory(self, mem): 

63 RealSpaceHamiltonian.estimate_memory(self, mem) 

64 solvation = mem.subnode('Solvation') 

65 for name, obj in [ 

66 ('Cavity', self.cavity), 

67 ('Dielectric', self.dielectric), 

68 ] + [('Interaction: ' + ia.subscript, ia) for ia in self.interactions]: 

69 obj.estimate_memory(solvation.subnode(name)) 

70 

71 def update_atoms(self, atoms, log): 

72 self.new_atoms = atoms.copy() 

73 log('Solvation position-dependent initialization:') 

74 self.cavity.update_atoms(atoms, log) 

75 for ia in self.interactions: 

76 ia.update_atoms(atoms, log) 

77 

78 def initialize(self): 

79 self.gradient = [ 

80 Gradient(self.finegd, i, 1.0, self.poisson.nn) for i in (0, 1, 2) 

81 ] 

82 self.vt_ia_g = self.finegd.zeros() 

83 self.cavity.allocate() 

84 self.dielectric.allocate() 

85 for ia in self.interactions: 

86 ia.allocate() 

87 RealSpaceHamiltonian.initialize(self) 

88 

89 def update(self, density, wfs=None, kin_en_using_band=True): 

90 self.timer.start('Hamiltonian') 

91 if self.vt_sg is None: 

92 self.timer.start('Initialize Hamiltonian') 

93 self.initialize() 

94 self.timer.stop('Initialize Hamiltonian') 

95 

96 cavity_changed = self.cavity.update(self.new_atoms, density) 

97 if cavity_changed: 

98 self.cavity.update_vol_surf() 

99 self.dielectric.update(self.cavity) 

100 

101 # e_coulomb, Ebar, Eext, Exc = 

102 finegd_energies = self.update_pseudo_potential(density) 

103 self.finegd.comm.sum(finegd_energies) 

104 ia_changed = [ 

105 ia.update( 

106 self.new_atoms, 

107 density, 

108 self.cavity if cavity_changed else None) 

109 for ia in self.interactions] 

110 if np.any(ia_changed): 

111 self.vt_ia_g.fill(.0) 

112 for ia in self.interactions: 

113 if ia.depends_on_el_density: 

114 self.vt_ia_g += ia.delta_E_delta_n_g 

115 if self.cavity.depends_on_el_density: 

116 self.vt_ia_g += (ia.delta_E_delta_g_g * 

117 self.cavity.del_g_del_n_g) 

118 if len(self.interactions) > 0: 

119 for vt_g in self.vt_sg: 

120 vt_g += self.vt_ia_g 

121 Eias = np.array([ia.E for ia in self.interactions]) 

122 

123 Ekin1 = self.gd.comm.sum_scalar(self.calculate_kinetic_energy(density)) 

124 W_aL = self.calculate_atomic_hamiltonians(density) 

125 atomic_energies = self.update_corrections(density, W_aL) 

126 self.world.sum(atomic_energies) 

127 

128 energies = atomic_energies 

129 energies[1:] += finegd_energies 

130 energies[0] += Ekin1 

131 

132 if not kin_en_using_band: 

133 assert wfs is not None 

134 with self.timer('New Kinetic Energy'): 

135 energies[0] = \ 

136 self.calculate_kinetic_energy_directly(density, 

137 wfs) 

138 

139 (self.e_kinetic0, self.e_coulomb, self.e_zero, 

140 self.e_external, self.e_xc) = energies 

141 

142 self.finegd.comm.sum(Eias) 

143 

144 self.cavity.communicate_vol_surf(self.world) 

145 for E, ia in zip(Eias, self.interactions): 

146 setattr(self, 'e_' + ia.subscript, E) 

147 

148 self.new_atoms = None 

149 self.timer.stop('Hamiltonian') 

150 

151 def update_pseudo_potential(self, density): 

152 ret = RealSpaceHamiltonian.update_pseudo_potential(self, density) 

153 if not self.cavity.depends_on_el_density: 

154 return ret 

155 del_g_del_n_g = self.cavity.del_g_del_n_g 

156 # XXX optimize numerics 

157 del_eps_del_g_g = self.dielectric.del_eps_del_g_g 

158 Veps = -1. / (8. * np.pi) * del_eps_del_g_g * del_g_del_n_g 

159 Veps *= self.grad_squared(self.vHt_g) 

160 for vt_g in self.vt_sg: 

161 vt_g += Veps 

162 return ret 

163 

164 def calculate_forces(self, dens, F_av): 

165 # XXX reorganize 

166 self.el_force_correction(dens, F_av) 

167 for ia in self.interactions: 

168 if self.cavity.depends_on_atomic_positions: 

169 for a, F_v in enumerate(F_av): 

170 del_g_del_r_vg = self.cavity.get_del_r_vg(a, dens) 

171 for v in (0, 1, 2): 

172 F_v[v] -= self.finegd.integrate( 

173 ia.delta_E_delta_g_g * del_g_del_r_vg[v], 

174 global_integral=False) 

175 if ia.depends_on_atomic_positions: 

176 for a, F_v in enumerate(F_av): 

177 del_E_del_r_vg = ia.get_del_r_vg(a, dens) 

178 for v in (0, 1, 2): 

179 F_v[v] -= self.finegd.integrate( 

180 del_E_del_r_vg[v], 

181 global_integral=False) 

182 return RealSpaceHamiltonian.calculate_forces(self, dens, F_av) 

183 

184 def el_force_correction(self, dens, F_av): 

185 if not self.cavity.depends_on_atomic_positions: 

186 return 

187 del_eps_del_g_g = self.dielectric.del_eps_del_g_g 

188 fixed = 1 / (8 * np.pi) * del_eps_del_g_g * \ 

189 self.grad_squared(self.vHt_g) # XXX grad_vHt_g inexact in bmgs 

190 for a, F_v in enumerate(F_av): 

191 del_g_del_r_vg = self.cavity.get_del_r_vg(a, dens) 

192 for v in (0, 1, 2): 

193 F_v[v] += self.finegd.integrate( 

194 fixed * del_g_del_r_vg[v], 

195 global_integral=False) 

196 

197 def get_energy(self, e_entropy, wfs, kin_en_using_band=True, e_sic=None): 

198 RealSpaceHamiltonian.get_energy(self, e_entropy, wfs, 

199 kin_en_using_band, e_sic) 

200 # The total energy calculated by the parent class includes the 

201 # solvent electrostatic contributions but not the interaction 

202 # energies. We add those here and store the electrostatic energies. 

203 self.e_el_free = self.e_total_free 

204 self.e_el_extrapolated = self.e_total_extrapolated 

205 for ia in self.interactions: 

206 self.e_total_free += getattr(self, 'e_' + ia.subscript) 

207 self.e_total_extrapolated = (self.e_total_free + 

208 wfs.occupations.extrapolate_factor * 

209 e_entropy) 

210 return self.e_total_free 

211 

212 def grad_squared(self, x): 

213 # XXX ugly 

214 gs = np.empty_like(x) 

215 tmp = np.empty_like(x) 

216 self.gradient[0].apply(x, gs) 

217 np.square(gs, gs) 

218 self.gradient[1].apply(x, tmp) 

219 np.square(tmp, tmp) 

220 gs += tmp 

221 self.gradient[2].apply(x, tmp) 

222 np.square(tmp, tmp) 

223 gs += tmp 

224 return gs 

225 

226 def summary(self, wfs, log): 

227 # This is almost duplicate code to gpaw/hamiltonian's 

228 # Hamiltonian.summary, but with the cavity and interactions added. 

229 

230 log('Energy contributions relative to reference atoms:', 

231 f'(reference = {self.setups.Eref * Ha:.6f})\n') 

232 

233 energies = [('Kinetic: ', self.e_kinetic), 

234 ('Potential: ', self.e_coulomb), 

235 ('External: ', self.e_external), 

236 ('XC: ', self.e_xc), 

237 ('Entropy (-ST):', self.e_entropy), 

238 ('Local: ', self.e_zero)] 

239 

240 if len(self.interactions) > 0: 

241 energies += [('Interactions', None)] 

242 for ia in self.interactions: 

243 energies += [(f' {ia.subscript:s}:', 

244 getattr(self, 'e_' + ia.subscript))] 

245 

246 for name, e in energies: 

247 if e is not None: 

248 log('%-14s %+11.6f' % (name, Ha * e)) 

249 else: 

250 log('%-14s' % (name)) 

251 

252 log('--------------------------') 

253 log('Free energy: %+11.6f' % (Ha * self.e_total_free)) 

254 log('Extrapolated: %+11.6f' % (Ha * self.e_total_extrapolated)) 

255 log() 

256 self.xc.summary(log) 

257 

258 try: 

259 workfunctions = self.get_workfunctions(wfs) 

260 except ValueError: 

261 pass 

262 else: 

263 log('Dipole-layer corrected work functions: {:.6f}, {:.6f} eV' 

264 .format(*np.array(workfunctions) * Ha)) 

265 log() 

266 

267 self.cavity.summary(log)