Coverage for gpaw/new/pot_calc.py: 86%

1""" 

2== ========== 

3R 

4r 

5G 

6g 

7h 

8x r or h 

9== ========== 

10 

11""" 

12 

13from __future__ import annotations 

14 

15import functools 

16import operator 

17from collections import defaultdict 

18from typing import DefaultDict 

19 

20import numpy as np 

21 

22from gpaw.core.arrays import DistributedArrays 

23from gpaw.core.atom_arrays import AtomArrays 

24from gpaw.core.uniform_grid import UGArray 

25from gpaw.mpi import MPIComm, serial_comm 

26from gpaw.new import trace, zips 

27from gpaw.new.energies import DFTEnergies 

28from gpaw.new.environment import Environment 

29from gpaw.new.logger import indent 

30from gpaw.new.potential import Potential 

31from gpaw.new.xc import Functional 

32from gpaw.setup import Setup 

33from gpaw.spinorbit import soc as soc_terms 

34from gpaw.typing import Array1D, Array2D, Array3D 

35from gpaw.utilities import pack_density, pack_hermitian, unpack_hermitian 

36 

37 

38class PotentialCalculator: 

39 def __init__(self, 

40 xc: Functional, 

41 poisson_solver, 

42 setups: list[Setup], 

43 *, 

44 relpos_ac: Array2D, 

45 environment: Environment, 

46 extensions: list | None = None, 

47 soc: bool = False): 

48 self.poisson_solver = poisson_solver 

49 self.xc = xc 

50 self.setups = setups 

51 self.relpos_ac = relpos_ac 

52 self.soc = soc 

53 self.environment = environment or Environment(len(relpos_ac)) 

54 self.extensions: list = extensions or [] 

55 

56 def __str__(self): 

57 return (f'{self.poisson_solver}\n' 

58 f'xc functional:\n{indent(self.xc)}\n') 

59 

60 def calculate_pseudo_potential(self, 

61 density, 

62 ibzwfs, 

63 vHt_x: DistributedArrays | None 

64 ) -> tuple[dict[str, float], 

65 UGArray, 

66 UGArray, 

67 DistributedArrays, 

68 AtomArrays, 

69 float]: 

70 raise NotImplementedError 

71 

72 def move(self, relpos_ac, atomdist): 

73 for ext in self.extensions: 

74 ext.move_atoms(relpos_ac) 

75 

76 @property 

77 def extensions_force_av(self): 

78 if not self.extensions: 

79 return np.zeros((len(self.setups), 3)) 

80 return functools.reduce(operator.add, [ext.force_contribution() 

81 for ext in self.extensions]) 

82 

83 @property 

84 def extensions_stress_contribution(self): 

85 if not self.extensions: 

86 return np.zeros((3, 3)) 

87 return functools.reduce(operator.add, [ext.stress_contribution() 

88 for ext in self.extensions]) 

89 

90 def calculate_charges(self, vHt_x): 

91 raise NotImplementedError 

92 

93 def restrict(self, a_r, a_R=None): 

94 raise NotImplementedError 

95 

96 def calculate_without_orbitals(self, 

97 density, 

98 ibzwfs=None, 

99 vHt_x: DistributedArrays | None = None, 

100 kpt_band_comm: MPIComm | None = None 

101 ) -> tuple[Potential, 

102 DFTEnergies, 

103 AtomArrays]: 

104 xc = self.xc 

105 if xc.exx_fraction != 0.0: 

106 from gpaw.new.xc import create_functional 

107 self.xc = create_functional('PBE', xc.grid) 

108 potential, energies, V_al = self.calculate( 

109 density, ibzwfs, vHt_x, kpt_band_comm) 

110 if xc.exx_fraction != 0.0: 

111 self.xc = xc 

112 return potential, energies, V_al 

113 

114 @trace 

115 def calculate(self, 

116 density, 

117 ibzwfs=None, 

118 vHt_x: DistributedArrays | None = None, 

119 kpt_band_comm: MPIComm | None = None 

120 ) -> tuple[Potential, DFTEnergies, AtomArrays]: 

121 energies, vt_sR, dedtaut_sr, vHt_x, V_aL, e_stress = ( 

122 self.calculate_pseudo_potential(density, ibzwfs, vHt_x)) 

123 e_kinetic = 0.0 

124 for spin, (vt_R, nt_R) in enumerate(zips(vt_sR, density.nt_sR)): 

125 e_kinetic -= vt_R.integrate(nt_R) 

126 if spin < density.ndensities: 

127 e_kinetic += vt_R.integrate(density.nct_R) 

128 

129 if dedtaut_sr is not None: 

130 dedtaut_sR = self.restrict(dedtaut_sr) 

131 for dedtaut_R, taut_R in zips(dedtaut_sR, 

132 density.taut_sR): 

133 e_kinetic -= dedtaut_R.integrate(taut_R) 

134 e_kinetic += dedtaut_R.integrate(density.tauct_R) 

135 else: 

136 dedtaut_sR = None 

137 

138 energies['kinetic_correction'] = e_kinetic 

139 

140 if kpt_band_comm is None: 

141 if ibzwfs is None: 

142 kpt_band_comm = serial_comm 

143 else: 

144 kpt_band_comm = ibzwfs.kpt_band_comm 

145 dH_asii, corrections = calculate_non_local_potential( 

146 self.setups, 

147 density, 

148 self.xc, 

149 V_aL, 

150 self.soc, 

151 self.extensions, 

152 kpt_band_comm) 

153 

154 for ext in self.extensions: 

155 dct = ext.get_energy_contributions() 

156 for name, e in dct.items(): 

157 assert name not in energies 

158 energies[name] = e 

159 

160 energies['spinorbit'] = 0.0 

161 for key, e in corrections.items(): 

162 if 0: 

163 print(f'{key:10} {energies[key]:15.9f} {e:15.9f}') 

164 energies[key] += e 

165 

166 return (Potential(vt_sR, dH_asii, dedtaut_sR, vHt_x, e_stress), 

167 DFTEnergies(**energies), 

168 V_aL) 

169 

170 

171@trace 

172def calculate_non_local_potential(setups, 

173 density, 

174 xc, 

175 V_aL, 

176 soc: bool, 

177 extensions, 

178 kpt_band_comm: MPIComm 

179 ) -> tuple[AtomArrays, 

180 dict[str, float]]: 

181 dtype = float if density.ncomponents < 4 else complex 

182 D_asii = density.D_asii.to_xp(np) 

183 dH_asii = D_asii.layout.new(dtype=dtype).empty(density.ncomponents) 

184 V_aL = V_aL.to_xp(np) 

185 energy_corrections: DefaultDict[str, float] = defaultdict(float) 

186 rank = 0 

187 for a, D_sii in D_asii.items(): 

188 if rank % kpt_band_comm.size == kpt_band_comm.rank: 

189 V_L = V_aL[a] 

190 setup = setups[a] 

191 dH_sii, corrections = calculate_non_local_potential1( 

192 setup, xc, D_sii, V_L, soc, extensions, a) 

193 dH_asii[a][:] = dH_sii 

194 for key, e in corrections.items(): 

195 energy_corrections[key] += e 

196 else: 

197 dH_asii[a][:] = 0.0 

198 rank += 1 

199 

200 kpt_band_comm.sum(dH_asii.data) 

201 

202 # Sum over domain: 

203 names = ['kinetic_correction', 'coulomb', 'zero', 'xc', 'external', 

204 'spinorbit'] 

205 energies = np.array([energy_corrections[name] for name in names]) 

206 density.D_asii.layout.atomdist.comm.sum(energies) 

207 kpt_band_comm.sum(energies) 

208 

209 return (dH_asii.to_xp(density.D_asii.layout.xp), 

210 dict(zips(names, energies))) 

211 

212 

213def calculate_non_local_potential1(setup: Setup, 

214 xc: Functional, 

215 D_sii: Array3D, 

216 V_L: Array1D, 

217 soc: bool, 

218 extensions, 

219 atom_index) -> tuple[Array3D, 

220 dict[str, float]]: 

221 ncomponents = len(D_sii) 

222 ndensities = 2 if ncomponents == 2 else 1 

223 D_sp = np.array([pack_density(D_ii.real) for D_ii in D_sii]) 

224 

225 D_p = D_sp[:ndensities].sum(0) 

226 

227 dH_p = (setup.K_p + setup.M_p + 

228 setup.MB_p + 2.0 * setup.M_pp @ D_p + 

229 setup.Delta_pL @ V_L) 

230 e_kinetic = setup.K_p @ D_p + setup.Kc 

231 e_zero = setup.MB + setup.MB_p @ D_p 

232 e_coulomb = setup.M + D_p @ (setup.M_p + setup.M_pp @ D_p) 

233 

234 dH_sp = np.zeros_like(D_sp, dtype=float if ncomponents < 4 else complex) 

235 

236 e_soc = 0. 

237 if soc: 

238 dHsoc_sii = soc_terms(setup, xc.xc, D_sp) 

239 e_soc += (D_sii[1:4] * dHsoc_sii).sum().real 

240 dH_sp[1:4] = pack_hermitian(dHsoc_sii) 

241 

242 dH_sp[:ndensities] = dH_p 

243 e_xc = xc.calculate_paw_correction(setup, D_sp, dH_sp) 

244 

245 # e_external = ext_pot.add_paw_correction(setup.Delta_pL[:, 0], dH_sp) 

246 

247 dH_sii = unpack_hermitian(dH_sp) 

248 

249 if setup.hubbard_u is not None: 

250 eU, dHU_sii = setup.hubbard_u.calculate(setup, D_sii) 

251 e_xc += eU 

252 dH_sii += dHU_sii 

253 

254 for extension in extensions: 

255 e_xc += extension.update_non_local_hamiltonian( 

256 D_sii, setup, atom_index, dH_sii) 

257 

258 e_kinetic -= (D_sii * dH_sii).sum().real 

259 

260 return dH_sii, {'kinetic_correction': e_kinetic, 

261 'coulomb': e_coulomb, 

262 'zero': e_zero, 

263 'xc': e_xc, 

264 'external': 0.0, # e_external, 

265 'spinorbit': e_soc}