Coverage for gpaw/hamiltonian.py: 98%

1"""This module defines a Hamiltonian.""" 

2 

3import functools 

4 

5import numpy as np 

6from ase.units import Ha 

7 

8from gpaw.arraydict import ArrayDict 

9from gpaw.external import create_external_potential 

10from gpaw.lfc import LFC 

11from gpaw.poisson import PoissonSolver 

12from gpaw.spinorbit import soc 

13from gpaw.transformers import Transformer 

14from gpaw.utilities import (pack_atomic_matrices, pack_hermitian, 

15 unpack_atomic_matrices, unpack_density, 

16 unpack_hermitian) 

17from gpaw.utilities.partition import AtomPartition 

18 

19ENERGY_NAMES = ['e_kinetic', 'e_coulomb', 'e_zero', 'e_external', 'e_xc', 

20 'e_entropy', 'e_total_free', 'e_total_extrapolated'] 

21 

22 

23def apply_non_local_hamilton(dH_asp, collinear, P, out=None): 

24 if out is None: 

25 out = P.new() 

26 for a, I1, I2 in P.indices: 

27 if collinear: 

28 dH_ii = unpack_hermitian(dH_asp[a][P.spin]) 

29 out.array[:, I1:I2] = np.dot(P.array[:, I1:I2], dH_ii) 

30 else: 

31 dH_xp = dH_asp[a] 

32 # We need the transpose because 

33 # we are dotting from the left 

34 dH_ii = unpack_hermitian(dH_xp[0]).T 

35 dH_vii = [unpack_hermitian(dH_p).T for dH_p in dH_xp[1:]] 

36 out.array[:, 0, I1:I2] = (np.dot(P.array[:, 0, I1:I2], 

37 dH_ii + dH_vii[2]) + 

38 np.dot(P.array[:, 1, I1:I2], 

39 dH_vii[0] - 1j * dH_vii[1])) 

40 out.array[:, 1, I1:I2] = (np.dot(P.array[:, 1, I1:I2], 

41 dH_ii - dH_vii[2]) + 

42 np.dot(P.array[:, 0, I1:I2], 

43 dH_vii[0] + 1j * dH_vii[1])) 

44 return out 

45 

46 

47# from gpaw.utilities.debug import frozen 

48# @frozen 

49class Hamiltonian: 

50 

51 def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world, 

52 redistributor, vext=None): 

53 self.gd = gd 

54 self.finegd = finegd 

55 self.nspins = nspins 

56 self.collinear = collinear 

57 self.setups = setups 

58 self.timer = timer 

59 self.xc = xc 

60 self.world = world 

61 self.redistributor = redistributor 

62 

63 self.ncomponents = self.nspins if self.collinear else 1 + 3 

64 

65 self.atomdist = None 

66 self.dH_asp = None 

67 self.vt_xG = None 

68 self.vt_sG = None 

69 self.vt_vG = None 

70 self.vHt_g = None 

71 self.vt_xg = None 

72 self.vt_sg = None 

73 self.vt_vg = None 

74 self.atom_partition = None 

75 

76 # Energy contributioons that sum up to e_total_free: 

77 self.e_kinetic = None 

78 self.e_coulomb = None 

79 self.e_zero = None 

80 self.e_external = None 

81 self.e_xc = None 

82 self.e_entropy = None 

83 self.e_band = None 

84 self.e_sic = None 

85 

86 self.e_total_free = None 

87 self.e_total_extrapolated = None 

88 self.e_kinetic0 = None 

89 

90 self.ref_vt_sG = None 

91 self.ref_dH_asp = None 

92 

93 if isinstance(vext, dict): 

94 vext = create_external_potential(**vext) 

95 self.vext = vext # external potential 

96 

97 self.positions_set = False 

98 self.spos_ac = None 

99 self.soc = False 

100 

101 @property 

102 def dH(self): 

103 return functools.partial(apply_non_local_hamilton, 

104 self.dH_asp, self.collinear) 

105 

106 def update_atomic_hamiltonians(self, value): 

107 if isinstance(value, dict): 

108 dtype = complex if self.soc else float 

109 tmp = self.setups.empty_atomic_matrix(self.ncomponents, 

110 self.atom_partition, dtype) 

111 tmp.update(value) 

112 value = tmp 

113 assert isinstance(value, ArrayDict) or value is None, type(value) 

114 if value is not None: 

115 value.check_consistency() 

116 self.dH_asp = value 

117 

118 def __str__(self): 

119 s = 'Hamiltonian:\n' 

120 s += (' XC and Coulomb potentials evaluated on a {}*{}*{} grid\n' 

121 .format(*self.finegd.N_c)) 

122 s += ' Using the %s Exchange-Correlation functional\n' % self.xc.name 

123 # We would get the description of the XC functional here, 

124 # except the thing has probably not been fully initialized yet. 

125 if self.vext is not None: 

126 s += f' External potential:\n {self.vext}\n' 

127 return s 

128 

129 def summary(self, wfs, log): 

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

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

132 

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

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

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

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

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

138 ('Local: ', self.e_zero), 

139 ('SIC: ', self.e_sic)] 

140 

141 for name, e in energies: 

142 if name == 'SIC: ' and e is None: 

143 continue 

144 else: 

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

146 

147 log('--------------------------') 

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

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

150 log() 

151 self.xc.summary(log) 

152 

153 try: 

154 workfunctions = self.get_workfunctions(wfs) 

155 except ValueError: 

156 pass 

157 else: 

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

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

160 log() 

161 

162 def get_workfunctions(self, wfs): 

163 """ 

164 Returns the work functions, in Hartree, for a dipole-corrected 

165 simulation. Returns None if no dipole correction is present. 

166 (wfs can be obtained from calc.wfs) 

167 """ 

168 try: 

169 dipole_correction = self.poisson.correction 

170 except AttributeError: 

171 raise ValueError( 

172 'Work function not defined if no field-free region. Consider ' 

173 'using a dipole correction if you are looking for a ' 

174 'work function.') 

175 c = self.poisson.c # index of axis perpendicular to dipole-layer 

176 if not self.gd.pbc_c[c]: 

177 # zero boundary conditions 

178 vacuum = 0.0 

179 else: 

180 v_q = self.pd3.gather(self.vHt_q) 

181 if self.pd3.comm.rank == 0: 

182 axes = (c, (c + 1) % 3, (c + 2) % 3) 

183 v_g = self.pd3.ifft(v_q, local=True).transpose(axes) 

184 vacuum = v_g[0].mean() 

185 else: 

186 vacuum = np.nan 

187 

188 fermilevel = wfs.fermi_level 

189 

190 wf1 = vacuum - fermilevel + dipole_correction 

191 wf2 = vacuum - fermilevel - dipole_correction 

192 return np.array([wf1, wf2]) 

193 

194 def set_positions_without_ruining_everything(self, spos_ac, 

195 atom_partition): 

196 self.spos_ac = spos_ac 

197 rank_a = atom_partition.rank_a 

198 

199 # If both old and new atomic ranks are present, start a blank dict if 

200 # it previously didn't exist but it will needed for the new atoms. 

201 # XXX what purpose does this serve? In what case does it happen? 

202 # How would one even go about figuring it out? Why does it all have 

203 # to be so unreadable? -Ask 

204 # 

205 if (self.atom_partition is not None and 

206 self.dH_asp is None and (rank_a == self.gd.comm.rank).any()): 

207 self.update_atomic_hamiltonians({}) 

208 

209 if self.atom_partition is not None and self.dH_asp is not None: 

210 self.timer.start('Redistribute') 

211 self.dH_asp.redistribute(atom_partition) 

212 self.timer.stop('Redistribute') 

213 

214 self.atom_partition = atom_partition 

215 self.atomdist = self.redistributor.get_atom_distributions(spos_ac) 

216 

217 def set_positions(self, spos_ac, atom_partition): 

218 self.vbar.set_positions(spos_ac, atom_partition) 

219 self.xc.set_positions(spos_ac) 

220 self.set_positions_without_ruining_everything(spos_ac, atom_partition) 

221 self.positions_set = True 

222 

223 def initialize(self): 

224 self.vt_xg = self.finegd.empty(self.ncomponents) 

225 self.vt_sg = self.vt_xg[:self.nspins] 

226 self.vt_vg = self.vt_xg[self.nspins:] 

227 self.vHt_g = self.finegd.zeros() 

228 self.vt_xG = self.gd.empty(self.ncomponents) 

229 self.vt_sG = self.vt_xG[:self.nspins] 

230 self.vt_vG = self.vt_xG[self.nspins:] 

231 

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

233 """Calculate effective potential. 

234 

235 The XC-potential and the Hartree potential are evaluated on 

236 the fine grid, and the sum is then restricted to the coarse 

237 grid.""" 

238 

239 self.timer.start('Hamiltonian') 

240 if self.vt_sg is None: 

241 with self.timer('Initialize Hamiltonian'): 

242 self.initialize() 

243 

244 finegrid_energies = self.update_pseudo_potential(density) 

245 coarsegrid_e_kinetic = self.calculate_kinetic_energy(density) 

246 

247 with self.timer('Calculate atomic Hamiltonians'): 

248 W_aL = self.calculate_atomic_hamiltonians(density) 

249 

250 atomic_energies = self.update_corrections(density, W_aL) 

251 

252 # Make energy contributions summable over world: 

253 finegrid_energies *= self.finegd.comm.size / self.world.size 

254 coarsegrid_e_kinetic *= self.gd.comm.size / self.world.size 

255 # (careful with array orderings/contents) 

256 

257 if 0: 

258 print('kinetic', coarsegrid_e_kinetic, atomic_energies[0]) 

259 print('coulomb', finegrid_energies[0], atomic_energies[1]) 

260 print('zero', finegrid_energies[1], atomic_energies[2]) 

261 print('xc', finegrid_energies[3], atomic_energies[4]) 

262 print('external', finegrid_energies[2], atomic_energies[3]) 

263 

264 energies = atomic_energies # kinetic, coulomb, zero, external, xc 

265 energies[1:] += finegrid_energies # coulomb, zero, external, xc 

266 energies[0] += coarsegrid_e_kinetic # kinetic 

267 

268 with self.timer('Communicate'): # time possible load imbalance 

269 self.world.sum(energies) 

270 if not kin_en_using_band: 

271 assert wfs is not None 

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

273 energies[0] = \ 

274 self.calculate_kinetic_energy_directly(density, 

275 wfs) 

276 (self.e_kinetic0, self.e_coulomb, self.e_zero, 

277 self.e_external, self.e_xc) = energies 

278 

279 self.timer.stop('Hamiltonian') 

280 

281 def update_corrections(self, dens, W_aL): 

282 self.timer.start('Atomic') 

283 self.update_atomic_hamiltonians(None) # XXXX 

284 

285 e_kinetic = 0.0 

286 e_coulomb = 0.0 

287 e_zero = 0.0 

288 e_external = 0.0 

289 e_xc = 0.0 

290 

291 D_asp = self.atomdist.to_work(dens.D_asp) 

292 dtype = complex if self.soc else float 

293 dH_asp = self.setups.empty_atomic_matrix(self.ncomponents, 

294 D_asp.partition, dtype) 

295 

296 for a, D_sp in D_asp.items(): 

297 W_L = W_aL[a] 

298 setup = self.setups[a] 

299 

300 if self.nspins == 2: 

301 D_p = D_sp.sum(0) 

302 else: 

303 D_p = D_sp[0] 

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

305 setup.MB_p + 2.0 * np.dot(setup.M_pp, D_p) + 

306 np.dot(setup.Delta_pL, W_L)) 

307 e_kinetic += np.dot(setup.K_p, D_p) + setup.Kc 

308 e_zero += setup.MB + np.dot(setup.MB_p, D_p) 

309 e_coulomb += setup.M + np.dot(D_p, (setup.M_p + 

310 np.dot(setup.M_pp, D_p))) 

311 

312 if self.soc: 

313 dH_vii = soc(setup, self.xc, D_sp) 

314 dH_sp = np.zeros_like(D_sp, dtype=complex) 

315 dH_sp[1:] = pack_hermitian(dH_vii) 

316 else: 

317 dH_sp = np.zeros_like(D_sp) 

318 

319 if setup.hubbard_u is not None: 

320 eU, dHU_sii = setup.hubbard_u.calculate(setup, 

321 unpack_density(D_sp)) 

322 e_xc += eU 

323 dH_sp += pack_hermitian(dHU_sii) 

324 

325 dH_sp[:self.nspins] += dH_p 

326 

327 if self.vext is not None: 

328 self.vext.paw_correction(setup.Delta_pL[:, 0], dH_sp) 

329 

330 if self.vext and self.vext.get_name() == 'CDFTPotential': 

331 # cDFT atomic hamiltonian, eq. 25 

332 # energy correction added in cDFT main 

333 h_cdft_a, h_cdft_b = self.vext.get_atomic_hamiltonians( 

334 setups=setup.Delta_pL[:, 0], atom=a) 

335 

336 dH_sp[0] += h_cdft_a 

337 dH_sp[1] += h_cdft_b 

338 

339 if self.ref_dH_asp: 

340 assert self.collinear 

341 dH_sp += self.ref_dH_asp[a] 

342 

343 dH_asp[a] = dH_sp 

344 

345 self.timer.start('XC Correction') 

346 for a, D_sp in D_asp.items(): 

347 e_xc += self.xc.calculate_paw_correction(self.setups[a], D_sp, 

348 dH_asp[a], a=a) 

349 self.timer.stop('XC Correction') 

350 for a, D_sp in D_asp.items(): 

351 e_kinetic -= (D_sp * dH_asp[a]).sum().real 

352 

353 self.update_atomic_hamiltonians(self.atomdist.from_work(dH_asp)) 

354 self.timer.stop('Atomic') 

355 

356 # Make corrections due to non-local xc: 

357 # self.Enlxc = 0.0 # XXXxcfunc.get_non_local_energy() 

358 e_kinetic += self.xc.get_kinetic_energy_correction() / self.world.size 

359 

360 return np.array([e_kinetic, e_coulomb, e_zero, e_external, e_xc]) 

361 

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

363 """Sum up all eigenvalues weighted with occupation numbers""" 

364 self.e_band = wfs.calculate_band_energy() 

365 if kin_en_using_band: 

366 self.e_kinetic = self.e_kinetic0 + self.e_band 

367 else: 

368 self.e_kinetic = self.e_kinetic0 

369 self.e_entropy = e_entropy 

370 if 0: 

371 print(self.e_kinetic0, 

372 self.e_band, 

373 self.e_coulomb, 

374 self.e_external, 

375 self.e_zero, 

376 self.e_xc, 

377 self.e_entropy) 

378 self.e_total_free = (self.e_kinetic + self.e_coulomb + 

379 self.e_external + self.e_zero + self.e_xc + 

380 self.e_entropy) 

381 

382 if e_sic is not None: 

383 self.e_sic = e_sic 

384 self.e_total_free += e_sic 

385 

386 self.e_total_extrapolated = ( 

387 self.e_total_free + 

388 wfs.occupations.extrapolate_factor * e_entropy) 

389 

390 return self.e_total_free 

391 

392 def linearize_to_xc(self, new_xc, density): 

393 # Store old hamiltonian 

394 ref_vt_sG = self.vt_sG.copy() 

395 ref_dH_asp = self.dH_asp.copy() 

396 self.xc = new_xc 

397 self.xc.set_positions(self.spos_ac) 

398 self.update(density) 

399 

400 ref_vt_sG -= self.vt_sG 

401 for a, dH_sp in self.dH_asp.items(): 

402 ref_dH_asp[a] -= dH_sp 

403 self.ref_vt_sG = ref_vt_sG 

404 self.ref_dH_asp = self.atomdist.to_work(ref_dH_asp) 

405 

406 def calculate_forces(self, dens, F_av): 

407 ghat_aLv = dens.ghat.dict(derivative=True) 

408 nct_av = dens.nct.dict(derivative=True) 

409 vbar_av = self.vbar.dict(derivative=True) 

410 

411 self.calculate_forces2(dens, ghat_aLv, nct_av, vbar_av) 

412 F_coarsegrid_av = np.zeros_like(F_av) 

413 

414 # Force from compensation charges: 

415 _Q, Q_aL = dens.calculate_multipole_moments() 

416 for a, dF_Lv in ghat_aLv.items(): 

417 F_av[a] += np.dot(Q_aL[a], dF_Lv) 

418 

419 # Force from smooth core charge: 

420 for a, dF_v in nct_av.items(): 

421 F_coarsegrid_av[a] += dF_v[0] 

422 

423 # Force from zero potential: 

424 for a, dF_v in vbar_av.items(): 

425 F_av[a] += dF_v[0] 

426 

427 self.xc.add_forces(F_av) 

428 self.gd.comm.sum(F_coarsegrid_av, 0) 

429 self.finegd.comm.sum(F_av, 0) 

430 if self.vext: 

431 if self.vext.get_name() == 'CDFTPotential': 

432 F_av += self.vext.get_cdft_forces() 

433 F_av += F_coarsegrid_av 

434 

435 def apply_local_potential(self, psit_nG, Htpsit_nG, s): 

436 vt_G = self.vt_sG[s] 

437 if psit_nG.ndim == 3: 

438 Htpsit_nG += psit_nG * vt_G 

439 else: 

440 for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG): 

441 Htpsit_G += psit_G * vt_G 

442 

443 def apply(self, a_xG, b_xG, wfs, kpt, calculate_P_ani=True): 

444 """Apply the Hamiltonian operator to a set of vectors. 

445 

446 Parameters: 

447 

448 a_nG: ndarray 

449 Set of vectors to which the overlap operator is applied. 

450 b_nG: ndarray, output 

451 Resulting S times a_nG vectors. 

452 wfs: WaveFunctions 

453 Wave-function object defined in wavefunctions.py 

454 kpt: KPoint object 

455 k-point object defined in kpoint.py. 

456 calculate_P_ani: bool 

457 When True, the integrals of projector times vectors 

458 P_ni = <p_i | a_nG> are calculated. 

459 When False, existing P_ani are used 

460 

461 """ 

462 

463 wfs.kin.apply(a_xG, b_xG, kpt.phase_cd) 

464 self.apply_local_potential(a_xG, b_xG, kpt.s) 

465 shape = a_xG.shape[:-3] 

466 P_axi = wfs.pt.dict(shape) 

467 

468 if calculate_P_ani: # TODO calculate_P_ani=False is experimental 

469 wfs.pt.integrate(a_xG, P_axi, kpt.q) 

470 else: 

471 for a, P_ni in kpt.P_ani.items(): 

472 P_axi[a][:] = P_ni 

473 

474 for a, P_xi in P_axi.items(): 

475 dH_ii = unpack_hermitian(self.dH_asp[a][kpt.s]) 

476 P_axi[a] = np.dot(P_xi, dH_ii) 

477 wfs.pt.add(b_xG, P_axi, kpt.q) 

478 

479 def get_xc_difference(self, xc, density): 

480 """Calculate non-selfconsistent XC-energy difference.""" 

481 if density.nt_sg is None: 

482 density.interpolate_pseudo_density() 

483 nt_sg = density.nt_sg 

484 if hasattr(xc, 'hybrid'): 

485 xc.calculate_exx() 

486 finegd_e_xc = xc.calculate(density.finegd, nt_sg) 

487 D_asp = self.atomdist.to_work(density.D_asp) 

488 atomic_e_xc = 0.0 

489 for a, D_sp in D_asp.items(): 

490 setup = self.setups[a] 

491 atomic_e_xc += xc.calculate_paw_correction(setup, D_sp, a=a) 

492 e_xc = finegd_e_xc + self.world.sum_scalar(atomic_e_xc) 

493 return e_xc - self.e_xc 

494 

495 def estimate_memory(self, mem): 

496 nbytes = self.gd.bytecount() 

497 nfinebytes = self.finegd.bytecount() 

498 arrays = mem.subnode('Arrays', 0) 

499 arrays.subnode('vHt_g', nfinebytes) 

500 arrays.subnode('vt_sG', self.nspins * nbytes) 

501 arrays.subnode('vt_sg', self.nspins * nfinebytes) 

502 self.xc.estimate_memory(mem.subnode('XC')) 

503 self.poisson.estimate_memory(mem.subnode('Poisson')) 

504 self.vbar.estimate_memory(mem.subnode('vbar')) 

505 

506 def write(self, writer): 

507 # Write all eneriges: 

508 for name in ENERGY_NAMES: 

509 energy = getattr(self, name) 

510 if energy is not None: 

511 energy *= Ha 

512 writer.write(name, energy) 

513 

514 writer.write( 

515 potential=self.gd.collect(self.vt_xG) * Ha, 

516 atomic_hamiltonian_matrices=pack_atomic_matrices(self.dH_asp) * Ha) 

517 

518 self.xc.write(writer.child('xc')) 

519 

520 if hasattr(self.poisson, 'write'): 

521 self.poisson.write(writer.child('poisson')) 

522 

523 def read(self, reader): 

524 h = reader.hamiltonian 

525 

526 # Read all energies: 

527 for name in ENERGY_NAMES: 

528 energy = h.get(name) 

529 if energy is not None: 

530 energy /= reader.ha 

531 setattr(self, name, energy) 

532 

533 # Read pseudo potential on the coarse grid 

534 # and broadcast on kpt/band comm: 

535 self.initialize() 

536 self.gd.distribute(h.potential / reader.ha, self.vt_xG) 

537 

538 self.atom_partition = AtomPartition(self.gd.comm, 

539 np.zeros(len(self.setups), int), 

540 name='hamiltonian-init-serial') 

541 

542 # Read non-local part of hamiltonian 

543 self.update_atomic_hamiltonians({}) 

544 dH_sP = h.atomic_hamiltonian_matrices / reader.ha 

545 

546 if self.gd.comm.rank == 0: 

547 self.update_atomic_hamiltonians( 

548 unpack_atomic_matrices(dH_sP, self.setups, 

549 new=reader.version >= 4)) 

550 

551 if hasattr(self.poisson, 'read'): 

552 self.poisson.read(reader) 

553 self.poisson.set_grid_descriptor(self.finegd) 

554 

555 def calculate_kinetic_energy_directly(self, density, wfs): 

556 

557 """ 

558 Calculate kinetic energy as 1/2 (nable psi)^2 

559 it gives better estimate of kinetic energy during the SCF. 

560 Important for direct min. 

561 

562 'calculate_kinetic_energy' method gives a correct 

563 value of kinetic energy only at self-consistent solution. 

564 

565 :param density: 

566 :param wfs: 

567 :return: total kinetic energy 

568 """ 

569 # pseudo-part 

570 if wfs.mode == 'lcao': 

571 return self.calculate_kinetic_energy_using_kin_en_matrix( 

572 density, wfs) 

573 elif wfs.mode == 'pw': 

574 e_kin = 0.0 

575 for kpt in wfs.kpt_u: 

576 for f, psit_G in zip(kpt.f_n, kpt.psit_nG): 

577 if f > 1.0e-10: 

578 G2_G = wfs.pd.G2_qG[kpt.q] 

579 e_kin += f * wfs.pd.integrate( 

580 0.5 * G2_G * psit_G, psit_G).real 

581 else: 

582 e_kin = 0.0 

583 

584 def Lapl(psit_G, kpt): 

585 Lpsit_G = np.zeros_like(psit_G) 

586 wfs.kin.apply(psit_G, Lpsit_G, kpt.phase_cd) 

587 return Lpsit_G 

588 

589 for kpt in wfs.kpt_u: 

590 for f, psit_G in zip(kpt.f_n, kpt.psit_nG): 

591 if f > 1.0e-10: 

592 e_kin += f * wfs.integrate( 

593 Lapl(psit_G, kpt), psit_G, False) 

594 e_kin = e_kin.real 

595 e_kin = wfs.gd.comm.sum_scalar(e_kin) 

596 

597 e_kin = wfs.kd.comm.sum_scalar(e_kin) # ? 

598 # paw corrections 

599 e_kin_paw = 0.0 

600 for a, D_sp in density.D_asp.items(): 

601 setup = wfs.setups[a] 

602 D_p = D_sp.sum(0) 

603 e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc 

604 e_kin_paw = density.gd.comm.sum_scalar(e_kin_paw) 

605 return e_kin + e_kin_paw 

606 

607 def calculate_kinetic_energy_using_kin_en_matrix(self, density, 

608 wfs): 

609 """ 

610 E_k = sum_{M'M} rho_MM' T_M'M 

611 better agreement between gradients of energy and 

612 the total energy during the direct minimisation. 

613 This is important when the line search is used. 

614 Also avoids using the eigenvalues which are 

615 not calculated during the direct minimisation. 

616 

617 'calculate_kinetic_energy' method gives a correct 

618 value of kinetic energy only at self-consistent solution. 

619 

620 :param density: 

621 :param wfs: 

622 :return: total kinetic energy 

623 """ 

624 # pseudo-part 

625 e_kinetic = 0.0 

626 e_kin_paw = 0.0 

627 

628 for kpt in wfs.kpt_u: 

629 # calculation of the density matrix directly 

630 # can be expansive for a large scale 

631 # as there are lot of empty states 

632 # when the exponential transformation is used 

633 # (n_bands=n_basis_functions.) 

634 # 

635 # rho_MM = \ 

636 # wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM) 

637 # e_kinetic += np.einsum('ij,ji->', kpt.T_MM, rho_MM) 

638 # 

639 # the code below is faster 

640 self.timer.start('Pseudo part') 

641 occ = kpt.f_n > 1e-10 

642 x_nn = np.dot(kpt.C_nM[occ], 

643 np.dot(kpt.T_MM, 

644 kpt.C_nM[occ].T.conj())).real 

645 e_kinetic += np.einsum('i,ii->', kpt.f_n[occ], x_nn) 

646 self.timer.stop('Pseudo part') 

647 # del rho_MM 

648 

649 e_kinetic = wfs.kd.comm.sum_scalar(e_kinetic) 

650 # paw corrections 

651 for a, D_sp in density.D_asp.items(): 

652 setup = wfs.setups[a] 

653 D_p = D_sp.sum(0) 

654 e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc 

655 

656 e_kin_paw = self.gd.comm.sum_scalar(e_kin_paw) 

657 

658 return e_kinetic.real + e_kin_paw 

659 

660 

661class RealSpaceHamiltonian(Hamiltonian): 

662 def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world, 

663 vext=None, 

664 psolver=None, stencil=3, redistributor=None, 

665 charge: float = 0.0): 

666 Hamiltonian.__init__(self, gd, finegd, nspins, collinear, 

667 setups, timer, xc, 

668 world, vext=vext, 

669 redistributor=redistributor) 

670 

671 # Solver for the Poisson equation: 

672 if psolver is None: 

673 psolver = {} 

674 if isinstance(psolver, dict): 

675 psolver = PoissonSolver(**psolver) 

676 self.poisson = psolver 

677 self.poisson.set_grid_descriptor(self.finegd) 

678 

679 # Restrictor function for the potential: 

680 self.restrictor = Transformer(self.finegd, self.redistributor.aux_gd, 

681 stencil) 

682 self.restrict = self.restrictor.apply 

683 

684 self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups], 

685 forces=True) 

686 

687 self.vbar_g = None 

688 self.npoisson = None 

689 

690 def restrict_and_collect(self, a_xg, b_xg=None, phases=None): 

691 if self.redistributor.enabled: 

692 tmp_xg = self.restrictor.apply(a_xg, output=None, phases=phases) 

693 b_xg = self.redistributor.collect(tmp_xg, b_xg) 

694 else: 

695 b_xg = self.restrictor.apply(a_xg, output=b_xg, phases=phases) 

696 return b_xg 

697 

698 def __str__(self): 

699 s = Hamiltonian.__str__(self) 

700 

701 degree = self.restrictor.nn * 2 - 1 

702 name = ['linear', 'cubic', 'quintic', 'heptic'][degree // 2] 

703 s += (' Interpolation: tri-%s ' % name + 

704 '(%d. degree polynomial)\n' % degree) 

705 s += ' Poisson solver: %s' % self.poisson.get_description() 

706 return s 

707 

708 def set_positions(self, spos_ac, rank_a): 

709 Hamiltonian.set_positions(self, spos_ac, rank_a) 

710 if self.vbar_g is None: 

711 self.vbar_g = self.finegd.empty() 

712 self.vbar_g[:] = 0.0 

713 self.vbar.add(self.vbar_g) 

714 

715 def update_pseudo_potential(self, dens): 

716 self.timer.start('vbar') 

717 e_zero = self.finegd.integrate(self.vbar_g, dens.nt_g, 

718 global_integral=False) 

719 

720 vt_g = self.vt_sg[0] 

721 vt_g[:] = self.vbar_g 

722 self.timer.stop('vbar') 

723 

724 e_external = 0.0 

725 if self.vext is not None: 

726 if self.vext.get_name() == 'CDFTPotential': 

727 vext_g = self.vext.get_potential(self.finegd).copy() 

728 e_external += self.vext.get_cdft_external_energy(