Coverage for gpaw/xc/fxc.py: 84%

573 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-14 00:18 +0000

1from abc import ABC, abstractmethod 

2from time import time 

3from pathlib import Path 

4 

5import ase.io.ulm as ulm 

6import numpy as np 

7from ase.units import Ha 

8from gpaw.response import timer 

9from scipy.special import p_roots, sici 

10 

11from gpaw.blacs import BlacsGrid, Redistributor 

12from gpaw.fd_operators import Gradient 

13from gpaw.kpt_descriptor import KPointDescriptor 

14from gpaw.pw.descriptor import PWDescriptor 

15from gpaw.response.qpd import SingleQPWDescriptor 

16from gpaw.utilities.blas import axpy, gemmdot 

17from gpaw.xc.rpa import RPACorrelation 

18from gpaw.heg import HEG 

19from gpaw.xc.fxc_kernels import ( 

20 get_fHxc_Gr, get_pbe_fxc, get_fspinHxc_Gr_rALDA, get_fspinHxc_Gr_rAPBE) 

21 

22 

23def get_chi0v_spinsum(chi0_sGG, G_G): 

24 nG = chi0_sGG.shape[-1] 

25 chi0v = np.zeros((nG, nG), dtype=complex) 

26 for chi0_GG in chi0_sGG: 

27 chi0v += chi0_GG / G_G / G_G[:, np.newaxis] 

28 chi0v *= 4 * np.pi 

29 return chi0v 

30 

31 

32def get_chi0v_foreach_spin(chi0_sGG, G_G): 

33 ns, nG = chi0_sGG.shape[:2] 

34 

35 chi0v_sGsG = np.zeros((ns, nG, ns, nG), dtype=complex) 

36 for s in range(ns): 

37 chi0v_sGsG[s, :, s, :] = chi0_sGG[s] / G_G / G_G[:, np.newaxis] 

38 chi0v_sGsG *= 4 * np.pi 

39 return chi0v_sGsG.reshape(ns * nG, ns * nG) 

40 

41 

42class FXCCache: 

43 def __init__(self, comm, tag, xc, ecut): 

44 self.comm = comm 

45 self.tag = tag 

46 self.xc = xc 

47 self.ecut = ecut 

48 

49 @property 

50 def prefix(self): 

51 return f'{self.tag}_{self.xc}_{self.ecut}' 

52 

53 def handle(self, iq): 

54 return Handle(self, iq) 

55 

56 

57class Handle: 

58 def __init__(self, cache, iq): 

59 self.cache = cache 

60 self.iq = iq 

61 self.comm = cache.comm 

62 

63 @property 

64 def _path(self): 

65 return Path(f'fhxc_{self.cache.prefix}_{self.iq}.ulm') 

66 

67 def exists(self): 

68 if self.comm.rank == 0: 

69 exists = int(self._path.exists()) 

70 else: 

71 exists = 0 

72 exists = self.comm.sum_scalar(exists) 

73 return bool(exists) 

74 

75 def read(self): 

76 from gpaw.mpi import broadcast 

77 if self.comm.rank == 0: 

78 array = self._read_master() 

79 assert array is not None 

80 else: 

81 array = None 

82 

83 # The shape of the array is only known on rank0, 

84 # so we cannot use the in-place broadcast. Therefore 

85 # we use the standalone function. 

86 array = broadcast(array, root=0, comm=self.comm) 

87 assert array is not None 

88 return array 

89 

90 def write(self, array): 

91 if self.comm.rank == 0: 

92 assert array is not None 

93 self._write_master(array) 

94 self.comm.barrier() 

95 

96 def _read_master(self): 

97 with ulm.open(self._path) as reader: 

98 return reader.array 

99 

100 def _write_master(self, array): 

101 assert array is not None 

102 with ulm.open(self._path, 'w') as writer: 

103 writer.write(array=array) 

104 

105 

106class FXCCorrelation: 

107 def __init__(self, 

108 calc, 

109 xc='RPA', 

110 nlambda=8, 

111 frequencies=None, 

112 weights=None, 

113 density_cut=1.e-6, 

114 unit_cells=None, 

115 tag=None, 

116 avg_scheme=None, 

117 *, 

118 ecut, 

119 **kwargs): 

120 

121 self.ecut = ecut 

122 if isinstance(ecut, (float, int)): 

123 self.ecut_max = ecut 

124 else: 

125 self.ecut_max = max(ecut) 

126 

127 self.rpa = RPACorrelation( 

128 calc, 

129 xc=xc, 

130 nlambda=nlambda, 

131 frequencies=frequencies, 

132 weights=weights, 

133 calculate_q=self.calculate_q_fxc, 

134 ecut=self.ecut, 

135 **kwargs) 

136 

137 self.gs = self.rpa.gs 

138 self.context = self.rpa.context 

139 

140 self.l_l, self.weight_l = p_roots(nlambda) 

141 self.l_l = (self.l_l + 1.0) * 0.5 

142 self.weight_l *= 0.5 

143 self.xc = xc 

144 self.density_cut = density_cut 

145 if unit_cells is None: 

146 unit_cells = self.gs.kd.N_c 

147 self.unit_cells = unit_cells 

148 

149 self.xcflags = XCFlags(self.xc) 

150 self.avg_scheme = self.xcflags.choose_avg_scheme(avg_scheme) 

151 

152 if tag is None: 

153 tag = self.gs.atoms.get_chemical_formula(mode='hill') 

154 if self.avg_scheme is not None: 

155 tag += '_' + self.avg_scheme 

156 

157 self.cache = FXCCache(self.context.comm, tag, self.xc, self.ecut_max) 

158 

159 @property 

160 def blockcomm(self): 

161 # Cannot be aliased as attribute 

162 # because rpa gets blockcomm during calculate 

163 return self.rpa.wblocks.blockcomm 

164 

165 def _calculate_kernel(self): 

166 # Find the first q vector to calculate kernel for 

167 # (density averaging scheme always calculates all q points anyway) 

168 

169 q_empty = None 

170 

171 for iq in reversed(range(len(self.rpa.integral.ibzq_qc))): 

172 handle = self.cache.handle(iq) 

173 

174 if not handle.exists(): 

175 q_empty = iq 

176 

177 if q_empty is None: 

178 self.context.print('%s kernel already calculated\n' % 

179 self.xc) 

180 return 

181 

182 kernelkwargs = dict( 

183 gs=self.gs, 

184 xc=self.xc, 

185 ibzq_qc=self.rpa.integral.ibzq_qc, 

186 ecut=self.ecut_max, 

187 context=self.context) 

188 

189 if self.avg_scheme == 'wavevector': 

190 self.context.print('Calculating %s kernel starting from ' 

191 'q point %s \n' % (self.xc, q_empty)) 

192 kernelkwargs.update(q_empty=q_empty) 

193 kernel = KernelWave(**kernelkwargs) 

194 else: 

195 kernel = KernelDens(**kernelkwargs, 

196 unit_cells=self.unit_cells, 

197 density_cut=self.density_cut) 

198 

199 for iq, fhxc_GG in kernel.calculate_fhxc(): 

200 if self.context.comm.rank == 0: 

201 assert isinstance(fhxc_GG, np.ndarray), str(fhxc_GG) 

202 self.cache.handle(iq).write(fhxc_GG) 

203 

204 @timer('FXC') 

205 def calculate(self, *, nbands=None): 

206 # kernel not required for RPA 

207 if self.xc != 'RPA': 

208 self._calculate_kernel() 

209 

210 data = self.rpa.calculate_all_contributions( 

211 spin=self.gs.nspins > 1, nbands=nbands) 

212 return data.energy_i * Ha # energies in eV 

213 

214 @timer('Chi0(q)') 

215 def calculate_q_fxc(self, chi0_s, m1, m2, gcut): 

216 for s, chi0 in enumerate(chi0_s): 

217 self.rpa.chi0calc.update_chi0(chi0, m1=m1, m2=m2, spins=[s]) 

218 

219 qpd = chi0_s[0].qpd 

220 chi0_swGG = np.array([ 

221 chi0.body.get_distributed_frequencies_array() for chi0 in chi0_s 

222 ]) 

223 wblocks = chi0_s[0].body.get_distributed_frequencies_blocks1d() 

224 if wblocks.blockcomm.size > 1: # why??? 

225 chi0_swGG = np.swapaxes(chi0_swGG, 2, 3) 

226 

227 # XXX Gamma-point code is NOT well tested! 

228 # Changed from qpd.kd.gamma to qpd.optical_limit cf. #1178. 

229 # This if/else was pasted from RPA where bug was also fixed. 

230 # We have not added regression test for fxc and the change 

231 # causes no test failures. 

232 if not chi0.qpd.optical_limit: 

233 energy_w = self.calculate_fxc_energies(qpd, chi0_swGG, gcut) 

234 else: 

235 chi0_swvv = [chi0.chi0_Wvv[wblocks.myslice] for chi0 in chi0_s] 

236 chi0_swxvG = [chi0.chi0_WxvG[wblocks.myslice] for chi0 in chi0_s] 

237 energy_w = self.calculate_optical_limit_fxc_energies( 

238 qpd, chi0_swGG, chi0_swvv, chi0_swxvG, gcut 

239 ) 

240 return wblocks.all_gather(energy_w) 

241 

242 def calculate_optical_limit_fxc_energies( 

243 self, qpd, chi0_swGG, chi0_swvv, chi0_swxvG, gcut): 

244 # For some reason, we "only" average out cartesian directions, instead 

245 # of performing an actual integral over the q-point volume as in rpa... 

246 energy_w = np.zeros(chi0_swGG.shape[1]) 

247 for v in range(3): 

248 for chi0_wGG, chi0_wvv, chi0_wxvG in zip( 

249 chi0_swGG, chi0_swvv, chi0_swxvG): 

250 chi0_wGG[:, 0] = chi0_wxvG[:, 0, v] 

251 chi0_wGG[:, :, 0] = chi0_wxvG[:, 1, v] 

252 chi0_wGG[:, 0, 0] = chi0_wvv[:, v, v] 

253 energy_w += self.calculate_fxc_energies(qpd, chi0_swGG, gcut) / 3 

254 return energy_w 

255 

256 def calculate_energy_contribution(self, chi0v_sGsG, fv_sGsG, nG): 

257 """Calculate contribution to energy from a single frequency point. 

258 

259 The RPA correlation energy is the integral over all frequencies 

260 from 0 to infinity of this expression.""" 

261 

262 e = 0.0 

263 assert len(chi0v_sGsG) % nG == 0 

264 ns = len(chi0v_sGsG) // nG 

265 

266 for l, weight in zip(self.l_l, self.weight_l): 

267 chiv = np.linalg.solve( 

268 np.eye(nG * ns) - l * np.dot(chi0v_sGsG, fv_sGsG), 

269 chi0v_sGsG).real # this is SO slow 

270 

271 chiv = chiv.reshape(ns, nG, ns, nG) 

272 for s1 in range(ns): 

273 for s2 in range(ns): 

274 e -= np.trace(chiv[s1, :, s2, :]) * weight 

275 

276 e += np.trace(chi0v_sGsG.real) 

277 return e 

278 

279 @timer('Energy') 

280 def calculate_fxc_energies(self, qpd, chi0_swGG, gcut): 

281 """Evaluate correlation energy from chi0 and the kernel fhxc""" 

282 ibzq_qc = self.rpa.integral.ibzq_qc 

283 ibzq2_q = [ 

284 np.dot(ibzq_qc[i] - qpd.q_c, 

285 ibzq_qc[i] - qpd.q_c) 

286 for i in range(len(ibzq_qc)) 

287 ] 

288 

289 qi = np.argsort(ibzq2_q)[0] 

290 

291 G_G = gcut.cut(qpd.G2_qG[0]**0.5) # |G+q| 

292 

293 nG = len(G_G) 

294 ns = len(chi0_swGG) 

295 

296 # There are three options to calculate the 

297 # energy depending on kernel and/or averaging scheme. 

298 

299 # Option (1) - Spin-polarized form of kernel exists 

300 # e.g. rALDA, rAPBE. 

301 # Then, solve block diagonal form of Dyson 

302 # equation (dimensions (ns*nG) * (ns*nG)) 

303 # (note this does not necessarily mean that 

304 # the calculation is spin-polarized!) 

305 

306 if self.xcflags.spin_kernel: 

307 fv_GG = gcut.spin_cut(self.cache.handle(qi).read(), ns) 

308 

309 # the spin-polarized kernel constructed from wavevector average 

310 # is already multiplied by |q+G| |q+G'|/4pi, and doesn't require 

311 # special treatment of the head and wings. However not true for 

312 # density average: 

313 

314 if self.avg_scheme == 'density': 

315 # Create and modify a view: 

316 fv_sGsG = fv_GG.reshape(ns, nG, ns, nG) 

317 

318 for s1 in range(ns): 

319 for s2 in range(ns): 

320 fv_sGsG[s1, :, s2, :] *= ( 

321 G_G * G_G[:, np.newaxis] / (4 * np.pi)) 

322 

323 # XXX Gamma check changed cf. #1178 without 

324 # further testing. 

325 if np.prod(self.unit_cells) > 1 and qpd.optical_limit: 

326 fv_sGsG[s1, 0, s2, :] = 0.0 

327 fv_sGsG[s1, :, s2, 0] = 0.0 

328 fv_sGsG[s1, 0, s2, 0] = 1.0 

329 

330 else: 

331 fv_GG = np.eye(nG) 

332 

333 if qpd.optical_limit: 

334 G_G[0] = 1.0 

335 

336 energy_w = [] 

337 for chi0_sGG in np.swapaxes(chi0_swGG, 0, 1): 

338 chi0_sGG = gcut.cut(chi0_sGG, [1, 2]) 

339 

340 if self.xcflags.spin_kernel: 

341 chi0v_sGsG = get_chi0v_foreach_spin(chi0_sGG, G_G) 

342 else: 

343 chi0v_sGsG = get_chi0v_spinsum(chi0_sGG, G_G) 

344 

345 energy_w.append(self.calculate_energy_contribution( 

346 chi0v_sGsG, fv_GG, nG 

347 )) 

348 return np.array(energy_w) 

349 

350 

351class KernelIntegrator(ABC): 

352 def __init__(self, gs, xc, context, ibzq_qc, ecut): 

353 self.gs = gs 

354 self.xc = xc 

355 self.context = context 

356 

357 self.xcflags = XCFlags(xc) 

358 self.gd = gs.density.gd 

359 self.ibzq_qc = ibzq_qc 

360 self.ecut = ecut 

361 

362 def calculate_fhxc(self): 

363 self.context.print( 

364 f'Calculating {self.xc} kernel at {self.ecut} eV cutoff') 

365 fhxc_iterator = self._calculate_fhxc() 

366 

367 while True: 

368 with self.context.timer('FHXC'): 

369 try: 

370 yield next(fhxc_iterator) 

371 except StopIteration: 

372 return 

373 

374 @abstractmethod 

375 def _calculate_fhxc(self): 

376 """Perform computation and yield (iq, array) tuples.""" 

377 

378 

379class KernelWave(KernelIntegrator): 

380 def __init__(self, *, q_empty, **kwargs): 

381 super().__init__(**kwargs) 

382 

383 self.ns = self.gs.nspins 

384 self.q_empty = q_empty 

385 

386 # Density grid 

387 n_sg, finegd = self.gs.get_all_electron_density(gridrefinement=2) 

388 self.n_g = n_sg.sum(axis=0).flatten() 

389 

390 # For atoms with large vacuum regions 

391 # this apparently can take negative values! 

392 mindens = np.amin(self.n_g) 

393 

394 if mindens < 0: 

395 self.context.print('Negative densities found! (magnitude %s)' % 

396 np.abs(mindens), flush=False) 

397 self.context.print('These will be reset to 1E-12 elec/bohr^3)') 

398 self.n_g[np.where(self.n_g < 0.0)] = 1.0E-12 

399 

400 r_g = finegd.get_grid_point_coordinates() 

401 self.x_g = 1.0 * r_g[0].flatten() 

402 self.y_g = 1.0 * r_g[1].flatten() 

403 self.z_g = 1.0 * r_g[2].flatten() 

404 self.gridsize = len(self.x_g) 

405 assert len(self.n_g) == self.gridsize 

406 

407 # Enhancement factor for GGA 

408 if self.xcflags.is_apbe: 

409 nf_g = self.gs.hacky_all_electron_density(gridrefinement=4) 

410 gdf = self.gd.refine().refine() 

411 grad_v = [Gradient(gdf, v, n=1).apply for v in range(3)] 

412 gradnf_vg = gdf.empty(3) 

413 

414 for v in range(3): 

415 grad_v[v](nf_g, gradnf_vg[v]) 

416 

417 self.s2_g = np.sqrt(np.sum(gradnf_vg[:, ::2, ::2, ::2]**2.0, 

418 0)).flatten() # |\nabla\rho| 

419 self.s2_g *= 1.0 / (2.0 * (3.0 * np.pi**2.0)**(1.0 / 3.0) * 

420 self.n_g**(4.0 / 3.0)) 

421 # |\nabla\rho|/(2kF\rho) = s 

422 self.s2_g = self.s2_g**2 # s^2 

423 assert len(self.n_g) == len(self.s2_g) 

424 

425 # Now we find all the regions where the 

426 # APBE kernel wants to be positive, and hack s to = 0, 

427 # so that we are really using the ALDA kernel 

428 # at these points 

429 apbe_g = get_pbe_fxc(self.n_g, self.s2_g) 

430 poskern_ind = np.where(apbe_g >= 0.0) 

431 if len(poskern_ind[0]) > 0: 

432 self.context.print( 

433 'The APBE kernel takes positive values at ' 

434 + '%s grid points out of a total of %s (%3.2f%%).' 

435 % (len(poskern_ind[0]), self.gridsize, 100.0 * len( 

436 poskern_ind[0]) / self.gridsize), flush=False) 

437 self.context.print('The ALDA kernel will be used at these ' 

438 'points') 

439 self.s2_g[poskern_ind] = 0.0 

440 

441 def _calculate_fhxc(self): 

442 for iq, q_c in enumerate(self.ibzq_qc): 

443 if iq < self.q_empty: # don't recalculate q vectors 

444 continue 

445 

446 yield iq, self.calculate_one_qpoint(iq, q_c) 

447 

448 def calculate_one_qpoint(self, iq, q_c): 

449 qpd = SingleQPWDescriptor.from_q(q_c, self.ecut / Ha, self.gd) 

450 

451 nG = qpd.ngmax 

452 G_G = qpd.G2_qG[0]**0.5 # |G+q| 

453 Gv_G = qpd.get_reciprocal_vectors(q=0, add_q=False) 

454 # G as a vector (note we are at a specific q point here so set q=0) 

455 

456 # Distribute G vectors among processors 

457 # Later we calculate for iG' > iG, 

458 # so stagger allocation in order to balance load 

459 local_Gvec_grid_size = nG // self.context.comm.size 

460 my_Gints = (self.context.comm.rank + np.arange(0, 

461 local_Gvec_grid_size * self.context.comm.size, 

462 self.context.comm.size)) 

463 

464 if (self.context.comm.rank + 

465 (local_Gvec_grid_size) * self.context.comm.size) < nG: 

466 my_Gints = np.append(my_Gints, 

467 [self.context.comm.rank + 

468 local_Gvec_grid_size * 

469 self.context.comm.size]) 

470 

471 my_Gv_G = Gv_G[my_Gints] 

472 

473 # XXX Should this be if self.ns == 2 and self.xcflags.spin_kernel? 

474 calc_spincorr = (self.ns == 2) and (self.xc == 'rALDA' 

475 or self.xc == 'rAPBE') 

476 

477 if calc_spincorr: 

478 # Form spin-dependent kernel according to 

479 # PRB 88, 115131 (2013) equation 20 

480 # (note typo, should be \tilde{f^rALDA}) 

481 # spincorr is just the ALDA exchange kernel 

482 # with a step function (\equiv \tilde{f^rALDA}) 

483 # fHxc^{up up} = fHxc^{down down} = fv_nospin + fv_spincorr 

484 # fHxc^{up down} = fHxc^{down up} = fv_nospin - fv_spincorr 

485 fv_spincorr_GG = np.zeros((nG, nG), dtype=complex) 

486 

487 fv_nospin_GG = np.zeros((nG, nG), dtype=complex) 

488 

489 for iG, Gv in zip(my_Gints, my_Gv_G): # loop over G vecs 

490 

491 # For all kernels we 

492 # treat head and wings analytically 

493 if G_G[iG] > 1.0E-5: 

494 # Symmetrised |q+G||q+G'|, where iG' >= iG 

495 mod_Gpq = np.sqrt(G_G[iG] * G_G[iG:]) 

496 

497 # Phase factor \vec{G}-\vec{G'} 

498 deltaGv = Gv - Gv_G[iG:] 

499 

500 if self.xc == 'rALDA': 

501 # rALDA trick: the Hartree-XC kernel is exactly 

502 # zero for densities below rho_min = 

503 # min_Gpq^3/(24*pi^2), 

504 # so we don't need to include these contributions 

505 # in the Fourier transform 

506 

507 min_Gpq = np.amin(mod_Gpq) 

508 rho_min = min_Gpq**3.0 / (24.0 * np.pi**2.0) 

509 small_ind = np.where(self.n_g >= rho_min) 

510 elif self.xcflags.is_apbe: 

511 # rAPBE trick: the Hartree-XC kernel 

512 # is exactly zero at grid points where 

513 # min_Gpq > cutoff wavevector 

514 

515 min_Gpq = np.amin(mod_Gpq) 

516 small_ind = np.where(min_Gpq <= np.sqrt( 

517 -4.0 * np.pi / 

518 get_pbe_fxc(self.n_g, self.s2_g))) 

519 else: 

520 small_ind = np.arange(self.gridsize) 

521 

522 phase_Gpq = np.exp( 

523 -1.0j * 

524 (deltaGv[:, 0, np.newaxis] * self.x_g[small_ind] + 

525 deltaGv[:, 1, np.newaxis] * self.y_g[small_ind] + 

526 deltaGv[:, 2, np.newaxis] * self.z_g[small_ind])) 

527 

528 def scaled_fHxc(spincorr): 

529 return self.get_scaled_fHxc_q( 

530 q=mod_Gpq, 

531 sel_points=small_ind, 

532 Gphase=phase_Gpq, 

533 spincorr=spincorr) 

534 

535 fv_nospin_GG[iG, iG:] = scaled_fHxc(spincorr=False) 

536 

537 if calc_spincorr: 

538 fv_spincorr_GG[iG, iG:] = scaled_fHxc(spincorr=True) 

539 else: 

540 # head and wings of q=0 are dominated by 

541 # 1/q^2 divergence of scaled Coulomb interaction 

542 

543 assert iG == 0 

544 

545 # The [0, 0] element would ordinarily be set to 

546 # 'l' if we have nonlinear kernel (which we are 

547 # removing). Now l=1.0 always: 

548 fv_nospin_GG[0, 0] = 1.0 

549 fv_nospin_GG[0, 1:] = 0.0 

550 

551 if calc_spincorr: 

552 fv_spincorr_GG[0, :] = 0.0 

553 

554 # End loop over G vectors 

555 

556 self.context.comm.sum(fv_nospin_GG) 

557 

558 # We've only got half the matrix here, 

559 # so add the hermitian conjugate: 

560 fv_nospin_GG += np.conj(fv_nospin_GG.T) 

561 # but now the diagonal's been doubled, 

562 # so we multiply these elements by 0.5 

563 fv_nospin_GG[np.diag_indices(nG)] *= 0.5 

564 

565 # End of loop over coupling constant 

566 

567 if calc_spincorr: 

568 self.context.comm.sum(fv_spincorr_GG) 

569 fv_spincorr_GG += np.conj(fv_spincorr_GG.T) 

570 fv_spincorr_GG[np.diag_indices(nG)] *= 0.5 

571 

572 self.context.print('q point %s complete' % iq) 

573 

574 if self.context.comm.rank == 0: 

575 if calc_spincorr: 

576 # Form the block matrix kernel 

577 fv_full_2G2G = np.empty((2 * nG, 2 * nG), dtype=complex) 

578 fv_full_2G2G[:nG, :nG] = fv_nospin_GG + fv_spincorr_GG 

579 fv_full_2G2G[:nG, nG:] = fv_nospin_GG - fv_spincorr_GG 

580 fv_full_2G2G[nG:, :nG] = fv_nospin_GG - fv_spincorr_GG 

581 fv_full_2G2G[nG:, nG:] = fv_nospin_GG + fv_spincorr_GG 

582 fhxc_sGsG = fv_full_2G2G 

583 

584 else: 

585 fhxc_sGsG = fv_nospin_GG 

586 

587 return fhxc_sGsG 

588 else: 

589 return None 

590 

591 def get_scaled_fHxc_q(self, q, sel_points, Gphase, spincorr): 

592 # Given a coupling constant l, construct the Hartree-XC 

593 # kernel in q space a la Lein, Gross and Perdew, 

594 # Phys. Rev. B 61, 13431 (2000): 

595 # 

596 # f_{Hxc}^\lambda(q,\omega,r_s) = \frac{4\pi \lambda }{q^2} + 

597 # \frac{1}{\lambda} f_{xc}(q/\lambda,\omega/\lambda^2,\lambda r_s) 

598 # 

599 # divided by the unscaled Coulomb interaction!! 

600 # 

601 # i.e. this subroutine returns f_{Hxc}^\lambda(q,\omega,r_s) 

602 # * \frac{q^2}{4\pi} 

603 # = \lambda * [\frac{(q/lambda)^2}{4\pi} 

604 # f_{Hxc}(q/\lambda,\omega/\lambda^2,\lambda r_s)] 

605 # = \lambda * [1/scaled_coulomb * fHxc computed with scaled quantities] 

606 

607 # Apply scaling 

608 rho = self.n_g[sel_points] 

609 

610 # GGA enhancement factor s is lambda independent, 

611 # but we might want to truncate it 

612 if self.xcflags.is_apbe: 

613 s2_g = self.s2_g[sel_points] 

614 else: 

615 s2_g = None 

616 

617 l = 1.0 # Leftover from the age of non-linear kernels. 

618 # This would be an integration weight or something. 

619 scaled_q = q / l 

620 scaled_rho = rho / l**3.0 

621 scaled_rs = (3.0 / (4.0 * np.pi * scaled_rho))**(1.0 / 3.0 

622 ) # Wigner radius 

623 

624 if not spincorr: 

625 scaled_kernel = l * self.get_fHxc_q(scaled_rs, scaled_q, Gphase, 

626 s2_g) 

627 else: 

628 scaled_kernel = l * self.get_spinfHxc_q(scaled_rs, scaled_q, 

629 Gphase, s2_g) 

630 

631 return scaled_kernel 

632 

633 def get_fHxc_q(self, rs, q, Gphase, s2_g): 

634 # Construct fHxc(q,G,:), divided by scaled Coulomb interaction 

635 

636 heg = HEG(rs) 

637 qF = heg.qF 

638 

639 fHxc_Gr = get_fHxc_Gr(self.xcflags, rs, q, qF, s2_g) 

640 

641 # Integrate over r with phase 

642 fHxc_Gr *= Gphase 

643 fHxc_GG = np.sum(fHxc_Gr, 1) / self.gridsize 

644 return fHxc_GG 

645 

646 def get_spinfHxc_q(self, rs, q, Gphase, s2_g): 

647 qF = HEG(rs).qF 

648 

649 if self.xc == 'rALDA': 

650 fspinHxc_Gr = get_fspinHxc_Gr_rALDA(qF, q) 

651 

652 elif self.xc == 'rAPBE': 

653 fspinHxc_Gr = get_fspinHxc_Gr_rAPBE(rs, q, s2_g) 

654 

655 fspinHxc_Gr *= Gphase 

656 fspinHxc_GG = np.sum(fspinHxc_Gr, 1) / self.gridsize 

657 return fspinHxc_GG 

658 

659 

660class KernelDens(KernelIntegrator): 

661 def __init__(self, *, unit_cells, density_cut, **kwargs): 

662 super().__init__(**kwargs) 

663 

664 self.unit_cells = unit_cells 

665 self.density_cut = density_cut 

666 

667 self.A_x = -(3 / 4.) * (3 / np.pi)**(1 / 3.) 

668 

669 self.n_g = self.gs.hacky_all_electron_density(gridrefinement=1) 

670 

671 if self.xc[-3:] == 'PBE': 

672 nf_g = self.gs.hacky_all_electron_density(gridrefinement=2) 

673 gdf = self.gd.refine() 

674 grad_v = [Gradient(gdf, v, n=1).apply for v in range(3)] 

675 gradnf_vg = gdf.empty(3) 

676 for v in range(3): 

677 grad_v[v](nf_g, gradnf_vg[v]) 

678 self.gradn_vg = gradnf_vg[:, ::2, ::2, ::2] 

679 

680 qd = KPointDescriptor(self.ibzq_qc) 

681 self.pd = PWDescriptor(self.ecut / Ha, self.gd, complex, qd) 

682 

683 def _calculate_fhxc(self): 

684 if self.xc[0] == 'r': # wth? 

685 assert self.xcflags.spin_kernel 

686 yield from self.calculate_rkernel() 

687 else: 

688 assert self.xc[0] == 'A' # wth? 

689 assert self.xc == 'ALDA' 

690 yield from self.calculate_local_kernel() 

691 

692 def calculate_rkernel(self): 

693 gd = self.gd 

694 ng_c = gd.N_c 

695 cell_cv = gd.cell_cv 

696 icell_cv = 2 * np.pi * np.linalg.inv(cell_cv) 

697 vol = gd.volume 

698 

699 ns = self.gs.nspins 

700 n_g = self.n_g # density on rough grid 

701 

702 fx_g = ns * self.get_fxc_g(n_g) # local exchange kernel 

703 try: 

704 qc_g = (-4 * np.pi * ns / fx_g)**0.5 # cutoff functional 

705 except FloatingPointError as err: 

706 msg = ( 

707 'Kernel is negative yet we want its square root. ' 

708 'You probably should not rely on this feature at all. ', 

709 'See discussion https://gitlab.com/gpaw/gpaw/-/issues/723') 

710 raise RuntimeError(msg) from err 

711 flocal_g = qc_g**3 * fx_g / (6 * np.pi**2) # ren. x-kernel for r=r' 

712 Vlocal_g = 2 * qc_g / np.pi # ren. Hartree kernel for r=r' 

713 

714 ng = np.prod(ng_c) # number of grid points 

715 r_vg = gd.get_grid_point_coordinates() 

716 rx_g = r_vg[0].flatten() 

717 ry_g = r_vg[1].flatten() 

718 rz_g = r_vg[2].flatten() 

719 

720 self.context.print(' %d grid points and %d plane waves at the ' 

721 'Gamma point' % (ng, self.pd.ngmax), flush=False) 

722 

723 # Unit cells 

724 R_Rv = [] 

725 weight_R = [] 

726 nR_v = self.unit_cells 

727 nR = np.prod(nR_v) 

728 for i in range(-nR_v[0] + 1, nR_v[0]): 

729 for j in range(-nR_v[1] + 1, nR_v[1]): 

730 for h in range(-nR_v[2] + 1, nR_v[2]): 

731 R_Rv.append(i * cell_cv[0] + j * cell_cv[1] + 

732 h * cell_cv[2]) 

733 weight_R.append((nR_v[0] - abs(i)) * (nR_v[1] - abs(j)) * 

734 (nR_v[2] - abs(h)) / float(nR)) 

735 if nR > 1: 

736 # with more than one unit cell only the exchange kernel is 

737 # calculated on the grid. The bare Coulomb kernel is added 

738 # in PW basis and Vlocal_g only the exchange part 

739 dv = self.gs.density.gd.dv 

740 gc = (3 * dv / 4 / np.pi)**(1 / 3.) 

741 Vlocal_g -= 2 * np.pi * gc**2 / dv 

742 self.context.print( 

743 ' Lattice point sampling: (%s x %s x %s)^2 ' 

744 % (nR_v[0], nR_v[1], nR_v[2]) + ' Reduced to %s lattice points' 

745 % len(R_Rv), flush=False) 

746 

747 l_g_size = -(-ng // self.context.comm.size) 

748 l_g_range = range(self.context.comm.rank * l_g_size, 

749 min((self.context.comm.rank + 1) * l_g_size, ng)) 

750 

751 fhxc_qsGr = {} 

752 for iq in range(len(self.ibzq_qc)): 

753 fhxc_qsGr[iq] = np.zeros( 

754 (ns, len(self.pd.G2_qG[iq]), len(l_g_range)), dtype=complex) 

755 

756 inv_error = np.seterr() 

757 np.seterr(invalid='ignore') 

758 np.seterr(divide='ignore') 

759 

760 t0 = time() 

761 # Loop over Lattice points 

762 for i, R_v in enumerate(R_Rv): 

763 # Loop over r'. f_rr and V_rr are functions of r (dim. as r_vg[0]) 

764 if i == 1: 

765 self.context.print( 

766 ' Finished 1 cell in %s seconds' % int(time() - t0) + 

767 ' - estimated %s seconds left' % int((len(R_Rv) - 1) * 

768 (time() - t0))) 

769 if len(R_Rv) > 5: 

770 if (i + 1) % (len(R_Rv) / 5 + 1) == 0: 

771 self.context.print( 

772 ' Finished %s cells in %s seconds' 

773 % (i, int(time() - t0)) + ' - estimated ' 

774 '%s seconds left' % int((len(R_Rv) - i) * (time() - 

775 t0) / i)) 

776 for g in l_g_range: 

777 rx = rx_g[g] + R_v[0] 

778 ry = ry_g[g] + R_v[1] 

779 rz = rz_g[g] + R_v[2] 

780 

781 # |r-r'-R_i| 

782 rr = ((r_vg[0] - rx)**2 + (r_vg[1] - ry)**2 + 

783 (r_vg[2] - rz)**2)**0.5 

784 

785 n_av = (n_g + n_g.flatten()[g]) / 2. 

786 fx_g = ns * self.get_fxc_g(n_av, index=g) 

787 qc_g = (-4 * np.pi * ns / fx_g)**0.5 

788 x = qc_g * rr 

789 osc_x = np.sin(x) - x * np.cos(x) 

790 f_rr = fx_g * osc_x / (2 * np.pi**2 * rr**3) 

791 if nR > 1: # include only exchange part of the kernel here 

792 V_rr = (sici(x)[0] * 2 / np.pi - 1) / rr 

793 else: # include the full kernel (also hartree part) 

794 V_rr = (sici(x)[0] * 2 / np.pi) / rr 

795 

796 # Terms with r = r' 

797 if (np.abs(R_v) < 0.001).all(): 

798 tmp_flat = f_rr.flatten() 

799 tmp_flat[g] = flocal_g.flatten()[g] 

800 f_rr = tmp_flat.reshape(ng_c) 

801 tmp_flat = V_rr.flatten() 

802 tmp_flat[g] = Vlocal_g.flatten()[g] 

803 V_rr = tmp_flat.reshape(ng_c) 

804 del tmp_flat 

805 

806 f_rr[np.where(n_av < self.density_cut)] = 0.0 

807 V_rr[np.where(n_av < self.density_cut)] = 0.0 

808 

809 f_rr *= weight_R[i] 

810 V_rr *= weight_R[i] 

811 

812 # r-r'-R_i 

813 r_r = np.array([r_vg[0] - rx, r_vg[1] - ry, r_vg[2] - rz]) 

814 

815 # Fourier transform of r 

816 for iq, q in enumerate(self.ibzq_qc): 

817 q_v = np.dot(q, icell_cv) 

818 e_q = np.exp(-1j * gemmdot(q_v, r_r, beta=0.0)) 

819 f_q = self.pd.fft((f_rr + V_rr) * e_q, iq) * vol / ng 

820 fhxc_qsGr[iq][0, :, g - l_g_range[0]] += f_q 

821 if ns == 2: 

822 f_q = self.pd.fft(V_rr * e_q, iq) * vol / ng 

823 fhxc_qsGr[iq][1, :, g - l_g_range[0]] += f_q 

824 

825 self.context.comm.barrier() 

826 

827 np.seterr(**inv_error) 

828 

829 for iq, q in enumerate(self.ibzq_qc): 

830 npw = len(self.pd.G2_qG[iq]) 

831 fhxc_sGsG = np.zeros((ns * npw, ns * npw), complex) 

832 # parallelize over PW below 

833 l_pw_size = -(-npw // self.context.comm.size) 

834 l_pw_range = range(self.context.comm.rank * l_pw_size, 

835 min((self.context.comm.rank + 1) * l_pw_size, 

836 npw)) 

837 

838 if self.context.comm.size > 1: 

839 # redistribute grid and plane waves in fhxc_qsGr[iq] 

840 bg1 = BlacsGrid(self.context.comm, 1, self.context.comm.size) 

841 bg2 = BlacsGrid(self.context.comm, self.context.comm.size, 1) 

842 bd1 = bg1.new_descriptor(npw, ng, npw, 

843 -(-ng // self.context.comm.size)) 

844 bd2 = bg2.new_descriptor(npw, ng, 

845 -(-npw // self.context.comm.size), 

846 ng) 

847 

848 fhxc_Glr = np.zeros((len(l_pw_range), ng), dtype=complex) 

849 if ns == 2: 

850 Koff_Glr = np.zeros((len(l_pw_range), ng), dtype=complex) 

851 

852 r = Redistributor(bg1.comm, bd1, bd2) 

853 r.redistribute(fhxc_qsGr[iq][0], fhxc_Glr, npw, ng) 

854 if ns == 2: 

855 r.redistribute(fhxc_qsGr[iq][1], Koff_Glr, npw, ng) 

856 else: 

857 fhxc_Glr = fhxc_qsGr[iq][0] 

858 if ns == 2: 

859 Koff_Glr = fhxc_qsGr[iq][1] 

860 

861 # Fourier transform of r' 

862 for iG in range(len(l_pw_range)): 

863 f_g = fhxc_Glr[iG].reshape(ng_c) 

864 f_G = self.pd.fft(f_g.conj(), iq) * vol / ng 

865 fhxc_sGsG[l_pw_range[0] + iG, :npw] = f_G.conj() 

866 if ns == 2: 

867 v_g = Koff_Glr[iG].reshape(ng_c) 

868 v_G = self.pd.fft(v_g.conj(), iq) * vol / ng 

869 fhxc_sGsG[npw + l_pw_range[0] + iG, :npw] = v_G.conj() 

870 

871 if ns == 2: # f_00 = f_11 and f_01 = f_10 

872 fhxc_sGsG[:npw, npw:] = fhxc_sGsG[npw:, :npw] 

873 fhxc_sGsG[npw:, npw:] = fhxc_sGsG[:npw, :npw] 

874 

875 self.context.comm.sum(fhxc_sGsG) 

876 fhxc_sGsG /= vol 

877 

878 if self.context.comm.rank == 0: 

879 if nR > 1: # add Hartree kernel evaluated in PW basis 

880 Gq2_G = self.pd.G2_qG[iq] 

881 if (q == 0).all(): 

882 Gq2_G = Gq2_G.copy() 

883 Gq2_G[0] = 1. 

884 vq_G = 4 * np.pi / Gq2_G 

885 fhxc_sGsG += np.tile(np.eye(npw) * vq_G, (ns, ns)) 

886 yield iq, fhxc_sGsG 

887 else: 

888 yield iq, None 

889 

890 def calculate_local_kernel(self): 

891 # Standard ALDA exchange kernel 

892 # Use with care. Results are very difficult to converge 

893 # Sensitive to density_cut 

894 ns = self.gs.nspins 

895 gd = self.gd 

896 pd = self.pd 

897 cell_cv = gd.cell_cv 

898 icell_cv = 2 * np.pi * np.linalg.inv(cell_cv) 

899 vol = gd.volume 

900 

901 fxc_sg = ns * self.get_fxc_g(ns * self.n_g) 

902 fxc_sg[np.where(self.n_g < self.density_cut)] = 0.0 

903 

904 r_vg = gd.get_grid_point_coordinates() 

905 

906 for iq in range(len(self.ibzq_qc)): 

907 Gvec_Gc = np.dot(pd.get_reciprocal_vectors(q=iq, add_q=False), 

908 cell_cv / (2 * np.pi)) 

909 npw = len(Gvec_Gc) 

910 l_pw_size = -(-npw // self.context.comm.size) 

911 l_pw_range = range(self.context.comm.rank * l_pw_size, 

912 min((self.context.comm.rank + 1) * l_pw_size, 

913 npw)) 

914 fhxc_sGsG = np.zeros((ns * npw, ns * npw), dtype=complex) 

915 for s in range(ns): 

916 for iG in l_pw_range: 

917 for jG in range(npw): 

918 fxc = fxc_sg[s].copy() 

919 dG_c = Gvec_Gc[iG] - Gvec_Gc[jG] 

920 dG_v = np.dot(dG_c, icell_cv) 

921 dGr_g = gemmdot(dG_v, r_vg, beta=0.0) 

922 ft_fxc = gd.integrate(np.exp(-1j * dGr_g) * fxc) 

923 fhxc_sGsG[s * npw + iG, s * npw + jG] = ft_fxc 

924 

925 self.context.comm.sum(fhxc_sGsG) 

926 fhxc_sGsG /= vol 

927 

928 Gq2_G = self.pd.G2_qG[iq] 

929 if (self.ibzq_qc[iq] == 0).all(): 

930 Gq2_G[0] = 1. 

931 vq_G = 4 * np.pi / Gq2_G 

932 fhxc_sGsG += np.tile(np.eye(npw) * vq_G, (ns, ns)) 

933 

934 yield iq, fhxc_sGsG 

935 

936 def get_fxc_g(self, n_g, index=None): 

937 if self.xc[-3:] == 'LDA': 

938 return self.get_lda_g(n_g) 

939 elif self.xc[-3:] == 'PBE': 

940 return self.get_pbe_g(n_g, index=index) 

941 else: 

942 raise '%s kernel not recognized' % self.xc 

943 

944 def get_lda_g(self, n_g): 

945 return (4. / 9.) * self.A_x * n_g**(-2. / 3.) 

946 

947 def get_pbe_g(self, n_g, index=None): 

948 if index is None: 

949 gradn_vg = self.gradn_vg 

950 else: 

951 gradn_vg = self.gs.density.gd.empty(3) 

952 for v in range(3): 

953 gradn_vg[v] = (self.gradn_vg[v] + 

954 self.gradn_vg[v].flatten()[index]) / 2 

955 

956 kf_g = (3. * np.pi**2 * n_g)**(1 / 3.) 

957 s2_g = np.zeros_like(n_g) 

958 for v in range(3): 

959 axpy(1.0, gradn_vg[v]**2, s2_g) 

960 s2_g /= 4 * kf_g**2 * n_g**2 

961 

962 from gpaw.xc.fxc_kernels import get_pbe_fxc 

963 return get_pbe_fxc(n_g, s2_g) 

964 

965 

966class XCFlags: 

967 _accepted_flags = { 

968 'RPA', 

969 'rALDA', 

970 'rAPBE', 

971 'ALDA'} 

972 

973 _spin_kernels = {'rALDA', 'rAPBE', 'ALDA'} 

974 

975 def __init__(self, xc): 

976 if xc not in self._accepted_flags: 

977 raise RuntimeError('%s kernel not recognized' % self.xc) 

978 

979 self.xc = xc 

980 

981 @property 

982 def spin_kernel(self): 

983 # rALDA/rAPBE are the only kernels which have spin-dependent forms 

984 return self.xc in self._spin_kernels 

985 

986 @property 

987 def is_apbe(self): 

988 # If new GGA kernels are added, maybe there should be an 

989 # is_gga property. 

990 return self.xc in {'rAPBE'} 

991 

992 def choose_avg_scheme(self, avg_scheme=None): 

993 xc = self.xc 

994 

995 if self.spin_kernel: 

996 if avg_scheme is None: 

997 avg_scheme = 'density' 

998 # Two-point scheme default for rALDA and rAPBE 

999 

1000 if avg_scheme == 'density': 

1001 assert self.spin_kernel, ('Two-point density average ' 

1002 'only implemented for rALDA and rAPBE') 

1003 

1004 elif xc != 'RPA': 

1005 avg_scheme = 'wavevector' 

1006 else: 

1007 avg_scheme = None 

1008 

1009 return avg_scheme