Coverage for gpaw/atom/all_electron.py: 90%

536 statements  

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

1""" 

2Atomic Density Functional Theory 

3""" 

4 

5from __future__ import annotations 

6from dataclasses import dataclass 

7from functools import cached_property 

8from math import log, pi, sqrt 

9 

10import numpy as np 

11from ase.data import atomic_names 

12from ase.utils import IOContext 

13 

14from gpaw import ConvergenceError 

15from gpaw.atom.configurations import configurations 

16from gpaw.atom.radialgd import AERadialGridDescriptor 

17from gpaw.utilities import hartree 

18from gpaw.xc import XC 

19from gpaw.mpi import serial_comm 

20 

21# fine-structure constant 

22alpha = 1 / 137.036 

23 

24 

25def get_r2dvdr(rgd, vr_g): 

26 r2dvdr_g = np.zeros(rgd.N) 

27 rgd.derivative(vr_g, r2dvdr_g) 

28 r2dvdr_g *= rgd.r_g 

29 r2dvdr_g -= vr_g 

30 return r2dvdr_g 

31 

32 

33def calculate_hartree(rgd, n, Z): 

34 vHr = np.zeros(rgd.N) 

35 hartree(0, n * rgd.r_g * rgd.dr_g, rgd.r_g, vHr) 

36 

37 # add potential from nuclear point charge (v = -Z / r) 

38 vHr -= Z 

39 return vHr 

40 

41 

42def calculate_xc(rgd, xc, n_g): 

43 vxc_g = np.zeros(rgd.N) 

44 

45 if xc.type == 'GLLB': 

46 # Wait, the xc functional necessarily needs to see the density! 

47 # ☠☠☠ XXX From where does it have the density? XXX ☠☠☠ 

48 Exc = xc.get_xc_potential_and_energy_1d(vxc_g) 

49 else: 

50 Exc = xc.calculate_spherical( 

51 rgd, n_g.reshape((1, -1)), vxc_g.reshape((1, -1))) 

52 return vxc_g, Exc 

53 

54 

55def calculate_potentials(rgd, xc, n_g, Z, tw_coeff=None): 

56 vHr_g = calculate_hartree(rgd, n_g, Z) 

57 vxc_g, Exc = calculate_xc(rgd, xc, n_g) 

58 vr_g = vHr_g + vxc_g * rgd.r_g 

59 

60 if tw_coeff is not None: 

61 vr_g /= tw_coeff 

62 

63 return vr_g, vHr_g, vxc_g, Exc 

64 

65 

66def calculate_density(f_j, u_jg, r_g): 

67 n_g = np.dot(f_j, np.where(abs(u_jg) < 1e-160, 0, u_jg)**2) / (4 * pi) 

68 n_g[1:] /= r_g[1:]**2 

69 n_g[0] = n_g[1] 

70 return n_g 

71 

72 

73class AllElectron(IOContext): 

74 """Object for doing an atomic DFT calculation.""" 

75 

76 default_gpernode = 150 

77 

78 def __init__(self, symbol, xcname='LDA', scalarrel=True, 

79 corehole=None, configuration=None, nofiles=True, 

80 txt='-', gpernode=default_gpernode, 

81 orbital_free=False, tw_coeff=1.): 

82 """Do an atomic DFT calculation. 

83 

84 Example:: 

85 

86 a = AllElectron('Fe') 

87 a.run() 

88 """ 

89 

90 self.txt = self.openfile(txt, comm=serial_comm) 

91 

92 self.symbol = symbol 

93 self.xcname = xcname 

94 self.scalarrel = scalarrel 

95 self.nofiles = nofiles 

96 

97 # Get reference state: 

98 self.Z, nlfe_j = configurations[symbol] 

99 

100 # Collect principal quantum numbers, angular momentum quantum 

101 # numbers, occupation numbers and eigenvalues (j is a combined 

102 # index for n and l): 

103 self.n_j = [n for n, l, f, e in nlfe_j] 

104 self.l_j = [l for n, l, f, e in nlfe_j] 

105 self.f_j = [f for n, l, f, e in nlfe_j] 

106 self.e_j = [e for n, l, f, e in nlfe_j] 

107 

108 if configuration is not None: 

109 j = 0 

110 for conf in configuration.split(','): 

111 if conf[0].isdigit(): 

112 n = int(conf[0]) 

113 l = 'spdf'.find(conf[1]) 

114 if len(conf) == 2: 

115 f = 1.0 

116 else: 

117 f = float(conf[2:]) 

118 assert n == self.n_j[j] 

119 assert l == self.l_j[j] 

120 self.f_j[j] = f 

121 j += 1 

122 else: 

123 j += {'He': 1, 

124 'Ne': 3, 

125 'Ar': 5, 

126 'Kr': 8, 

127 'Xe': 11}[conf] 

128 

129 maxnodes = max([n - l - 1 for n, l in zip(self.n_j, self.l_j)]) 

130 self.N = (maxnodes + 1) * gpernode 

131 self.beta = 0.4 

132 self.rgd = AERadialGridDescriptor(self.beta / self.N, 1.0 / self.N, 

133 self.N) 

134 

135 self.orbital_free = orbital_free 

136 self.tw_coeff = tw_coeff 

137 

138 if self.orbital_free: 

139 self.n_j = [1] 

140 self.l_j = [0] 

141 self.f_j = [self.Z] 

142 self.e_j = [self.e_j[-1]] 

143 

144 t = self.text 

145 t() 

146 if scalarrel: 

147 t('Scalar-relativistic atomic ', end='') 

148 else: 

149 t('Atomic ', end='') 

150 t('%s calculation for %s (%s, Z=%d)' % (xcname, symbol, 

151 atomic_names[self.Z], self.Z)) 

152 

153 if corehole is not None: 

154 self.ncorehole, self.lcorehole, self.fcorehole = corehole 

155 

156 # Find j for core hole and adjust occupation: 

157 for j in range(len(self.f_j)): 

158 if (self.n_j[j] == self.ncorehole and 

159 self.l_j[j] == self.lcorehole): 

160 assert self.f_j[j] == 2 * (2 * self.lcorehole + 1) 

161 self.f_j[j] -= self.fcorehole 

162 self.jcorehole = j 

163 break 

164 

165 coreholestate = '%d%s' % (self.ncorehole, 'spdf'[self.lcorehole]) 

166 t('Core hole in {} state ({} occupation: {:.1f})'.format( 

167 coreholestate, coreholestate, self.f_j[self.jcorehole])) 

168 else: 

169 self.jcorehole = None 

170 self.fcorehole = 0 

171 

172 def text(self, *args, **kwargs): 

173 self.txt.write(kwargs.get('sep', ' ').join([str(arg) 

174 for arg in args]) + 

175 kwargs.get('end', '\n')) 

176 

177 def initialize_wave_functions(self): 

178 r = self.r 

179 dr = self.dr 

180 # Initialize with Slater function: 

181 for l, e, u in zip(self.l_j, self.e_j, self.u_j): 

182 if self.symbol in ['Hf', 'Ta', 'W', 'Re', 'Os', 

183 'Ir', 'Pt', 'Au']: 

184 a = sqrt(-4.0 * e) 

185 else: 

186 a = sqrt(-2.0 * e) 

187 

188 u[:] = r**(1 + l) * np.exp(-a * r) 

189 norm = np.dot(u**2, dr) 

190 u *= 1.0 / sqrt(norm) 

191 

192 def run(self): 

193 # beta g 

194 # r = ------, g = 0, 1, ..., N - 1 

195 # N - g 

196 # 

197 # rN 

198 # g = -------- 

199 # beta + r 

200 

201 t = self.text 

202 N = self.N 

203 t(N, 'radial gridpoints.') 

204 g = np.arange(N, dtype=float) 

205 self.r = self.rgd.r_g 

206 self.dr = self.rgd.dr_g 

207 self.d2gdr2 = self.rgd.d2gdr2() 

208 

209 # Number of orbitals: 

210 nj = len(self.n_j) 

211 

212 # Radial wave functions multiplied by radius: 

213 self.u_j = np.zeros((nj, self.N)) 

214 

215 # Effective potential multiplied by radius: 

216 self.vr = np.zeros(N) 

217 

218 # Electron density: 

219 self.n = np.zeros(N) 

220 

221 # Always spinpaired nspins=1 

222 self.xc = XC(self.xcname) 

223 

224 # Initialize for non-local functionals 

225 if self.xc.type == 'GLLB': 

226 self.xc.initialize_1d(self) 

227 

228 n_j = self.n_j 

229 l_j = self.l_j 

230 f_j = self.f_j 

231 e_j = self.e_j 

232 

233 Z = self.Z # nuclear charge 

234 r = self.r # radial coordinate 

235 dr = self.dr # dr/dg 

236 n = self.n # electron density 

237 vr = self.vr # effective potential multiplied by r 

238 

239 vHr = np.zeros(self.N) 

240 self.vXC = np.zeros(self.N) 

241 

242 self.initialize_wave_functions() 

243 n[:] = self.calculate_density() 

244 

245 bar = '|------------------------------------------------|' 

246 t(bar) 

247 

248 niter = 0 

249 nitermax = 117 

250 qOK = log(1e-10) 

251 mix = 0.4 

252 

253 # orbital_free needs more iterations and coefficient 

254 if self.orbital_free: 

255 mix = 0.01 

256 nitermax = 2000 

257 e_j[0] /= self.tw_coeff 

258 if Z > 10: # help convergence for third row elements 

259 mix = 0.002 

260 nitermax = 10000 

261 

262 vrold = None 

263 

264 while True: 

265 tw_coeff = self.tw_coeff if self.orbital_free else None 

266 

267 vr[:], vHr[:], self.vXC[:], Exc = calculate_potentials( 

268 rgd=self.rgd, xc=self.xc, n_g=n, Z=Z, tw_coeff=tw_coeff) 

269 

270 # calculate new total Kohn-Sham effective potential and 

271 # admix with old version 

272 

273 if niter > 0: 

274 vr[:] = mix * vr + (1 - mix) * vrold 

275 vrold = vr.copy() 

276 

277 # solve Kohn-Sham equation and determine the density change 

278 self.solve() 

279 dn = self.calculate_density() - n 

280 n += dn 

281 

282 # estimate error from the square of the density change integrated 

283 q = log(np.sum((r * dn)**2)) 

284 

285 # print progress bar 

286 if niter == 0: 

287 q0 = q 

288 b0 = 0 

289 else: 

290 b = int((q0 - q) / (q0 - qOK) * 50) 

291 if b > b0: 

292 self.txt.write(bar[b0:min(b, 50)]) 

293 self.txt.flush() 

294 b0 = b 

295 

296 # check if converged and break loop if so 

297 if q < qOK: 

298 self.txt.write(bar[b0:]) 

299 self.txt.flush() 

300 break 

301 

302 niter += 1 

303 if niter > nitermax: 

304 raise RuntimeError('Did not converge!') 

305 

306 tau = self.calculate_kinetic_energy_density() 

307 

308 t() 

309 t('Converged in %d iteration%s.' % (niter, 's'[:niter != 1])) 

310 

311 Ekin = 0 

312 

313 for f, e in zip(f_j, e_j): 

314 Ekin += f * e 

315 

316 e_coulomb = 2 * pi * np.dot(n * r * (vHr - Z), dr) 

317 Ekin += -4 * pi * np.dot(n * vr * r, dr) 

318 

319 if self.orbital_free: 

320 # e and vr are not scaled back 

321 # instead Ekin is scaled for total energy 

322 # (printed and inside setup) 

323 Ekin *= self.tw_coeff 

324 t() 

325 t(f'Lambda:{self.tw_coeff}') 

326 t(f'Correct eigenvalue:{e_j[0] * self.tw_coeff}') 

327 t() 

328 

329 t() 

330 t('Energy contributions:') 

331 t('-------------------------') 

332 t('Kinetic: %+13.6f' % Ekin) 

333 t('XC: %+13.6f' % Exc) 

334 t('Potential: %+13.6f' % e_coulomb) 

335 t('-------------------------') 

336 t('Total: %+13.6f' % (Ekin + Exc + e_coulomb)) 

337 self.ETotal = Ekin + Exc + e_coulomb 

338 t() 

339 

340 t('state eigenvalue ekin rmax') 

341 t('-----------------------------------------------') 

342 for m, l, f, e, u in zip(n_j, l_j, f_j, e_j, self.u_j): 

343 # Find kinetic energy: 

344 k = e - np.sum((np.where(abs(u) < 1e-160, 0, u)**2 * # XXXNumeric! 

345 vr * dr)[1:] / r[1:]) 

346 

347 # Find outermost maximum: 

348 g = self.N - 4 

349 while u[g - 1] >= u[g]: 

350 g -= 1 

351 x = r[g - 1:g + 2] 

352 y = u[g - 1:g + 2] 

353 A = np.transpose(np.array([x**i for i in range(3)])) 

354 c, b, a = np.linalg.solve(A, y) 

355 assert a < 0.0 

356 rmax = -0.5 * b / a 

357 

358 s = 'spdf'[l] 

359 t('%d%s^%-4.1f: %12.6f %12.6f %12.3f' % (m, s, f, e, k, rmax)) 

360 t('-----------------------------------------------') 

361 t('(units: Bohr and Hartree)') 

362 

363 for m, l, u in zip(n_j, l_j, self.u_j): 

364 self.write(u, 'ae', n=m, l=l) 

365 

366 self.write(n, 'n') 

367 self.write(vr, 'vr') 

368 self.write(vHr, 'vHr') 

369 self.write(self.vXC, 'vXC') 

370 self.write(tau, 'tau') 

371 

372 self.Ekin = Ekin 

373 self.e_coulomb = e_coulomb 

374 self.Exc = Exc 

375 

376 def write(self, array, name=None, n=None, l=None): 

377 if self.nofiles: 

378 return 

379 

380 if name: 

381 name = self.symbol + '.' + name 

382 else: 

383 name = self.symbol 

384 

385 if l is not None: 

386 assert n is not None 

387 if n > 0: 

388 # Bound state: 

389 name += '.%d%s' % (n, 'spdf'[l]) 

390 else: 

391 name += '.x%d%s' % (-n, 'spdf'[l]) 

392 

393 f = open(name, 'w') 

394 for r, a in zip(self.r, array): 

395 print(r, a, file=f) 

396 

397 def calculate_density(self): 

398 """Return the electron charge density divided by 4 pi.""" 

399 return calculate_density(self.f_j, self.u_j, self.r) 

400 

401 def calculate_kinetic_energy_density(self): 

402 """Return the kinetic energy density""" 

403 return self.radial_kinetic_energy_density(self.f_j, self.l_j, self.u_j) 

404 

405 def radial_kinetic_energy_density(self, f_j, l_j, u_j): 

406 """Kinetic energy density from a restricted set of wf's 

407 """ 

408 shape = np.shape(u_j[0]) 

409 dudr = np.zeros(shape) 

410 tau = np.zeros(shape) 

411 for f, l, u in zip(f_j, l_j, u_j): 

412 self.rgd.derivative(u, dudr) 

413 # contribution from angular derivatives 

414 if l > 0: 

415 tau += f * l * (l + 1) * np.where(abs(u) < 1e-160, 0, u)**2 

416 # contribution from radial derivatives 

417 dudr = u - self.r * dudr 

418 tau += f * np.where(abs(dudr) < 1e-160, 0, dudr)**2 

419 tau[1:] /= self.r[1:]**4 

420 tau[0] = tau[1] 

421 

422 return 0.5 * tau / (4 * pi) 

423 

424 def calculate_kinetic_energy_density2(self): 

425 """Return the kinetic energy density 

426 calculation over R(r)=u(r)/r 

427 slower convergence with # of radial grid points for 

428 Ekin of H than radial_kinetic_energy_density 

429 """ 

430 

431 shape = self.u_j.shape[1] 

432 R = np.zeros(shape) 

433 dRdr = np.zeros(shape) 

434 tau = np.zeros(shape) 

435 for f, l, u in zip(self.f_j, self.l_j, self.u_j): 

436 R[1:] = u[1:] / self.r[1:] 

437 if l == 0: 

438 # estimate value at origin by Taylor series to first order 

439 d1 = self.r[1] 

440 d2 = self.r[2] 

441 R[0] = .5 * (R[1] + R[2] + (R[1] - R[2]) * 

442 (d1 + d2) / (d2 - d1)) 

443 else: 

444 R[0] = 0 

445 self.rgd.derivative(R, dRdr) 

446 # contribution from radial derivatives 

447 tau += f * np.where(abs(dRdr) < 1e-160, 0, dRdr)**2 

448 # contribution from angular derivatives 

449 if l > 0: 

450 R[1:] = R[1:] / self.r[1:] 

451 if l == 1: 

452 R[0] = R[1] 

453 else: 

454 R[0] = 0 

455 tau += f * l * (l + 1) * np.where(abs(R) < 1e-160, 0, R)**2 

456 

457 return 0.5 * tau / (4 * pi) 

458 

459 def solve(self): 

460 """Solve the Schrodinger equation 

461 

462 :: 

463 

464 2 

465 d u 1 dv du u l(l + 1) 

466 - --- - ---- -- (-- - -) + [-------- + 2M(v - e)] u = 0, 

467 2 2 dr dr r 2 

468 dr 2Mc r 

469 

470 

471 where the relativistic mass:: 

472 

473 1 

474 M = 1 - --- (v - e) 

475 2 

476 2c 

477 

478 and the fine-structure constant alpha = 1/c = 1/137.036 

479 is set to zero for non-scalar-relativistic calculations. 

480 

481 On the logaritmic radial grids defined by:: 

482 

483 beta g 

484 r = ------, g = 0, 1, ..., N - 1 

485 N - g 

486 

487 rN 

488 g = --------, r = [0; oo[ 

489 beta + r 

490 

491 the Schrodinger equation becomes:: 

492 

493 2 

494 d u du 

495 --- c + -- c + u c = 0 

496 2 2 dg 1 0 

497 dg 

498 

499 with the vectors c0, c1, and c2 defined by:: 

500 

501 2 dg 2 

502 c = -r (--) 

503 2 dr 

504 

505 2 2 

506 d g 2 r dg dv 

507 c = - --- r - ---- -- -- 

508 1 2 2 dr dr 

509 dr 2Mc 

510 

511 2 r dv 

512 c = l(l + 1) + 2M(v - e)r + ---- -- 

513 0 2 dr 

514 2Mc 

515 """ 

516 r = self.r 

517 dr = self.dr 

518 vr = self.vr 

519 

520 c2 = -(r / dr)**2 

521 c10 = -self.d2gdr2 * r**2 # first part of c1 vector 

522 

523 if self.scalarrel: 

524 assert self.N == self.rgd.N 

525 self.r2dvdr = get_r2dvdr(self.rgd, vr) 

526 else: 

527 self.r2dvdr = None 

528 

529 # solve for each quantum state separately 

530 for j, (n, l, e, u) in enumerate(zip(self.n_j, self.l_j, 

531 self.e_j, self.u_j)): 

532 nodes = n - l - 1 # analytically expected number of nodes 

533 delta = -0.2 * e 

534 nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2, 

535 self.scalarrel) 

536 # adjust eigenenergy until u has the correct number of nodes 

537 while nn != nodes: 

538 diff = np.sign(nn - nodes) 

539 while diff == np.sign(nn - nodes): 

540 e -= diff * delta 

541 nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2, 

542 self.scalarrel) 

543 delta /= 2 

544 

545 # adjust eigenenergy until u is smooth at the turning point 

546 de = 1.0 

547 while abs(de) > 1e-9: 

548 norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr) 

549 u *= 1.0 / sqrt(norm) 

550 de = 0.5 * A / norm 

551 x = abs(de / e) 

552 if x > 0.1: 

553 de *= 0.1 / x 

554 e -= de 

555 assert e < 0.0 

556 nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2, 

557 self.scalarrel) 

558 self.e_j[j] = e 

559 u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)) 

560 

561 def kin(self, l, u, e=None): # XXX move to Generator 

562 r = self.r[1:] 

563 dr = self.dr[1:] 

564 

565 c0 = 0.5 * l * (l + 1) / r**2 

566 c1 = -0.5 * self.d2gdr2[1:] 

567 c2 = -0.5 * dr**-2 

568 

569 if e is not None and self.scalarrel: 

570 x = 0.5 * alpha**2 

571 Mr = r * (1.0 + x * e) - x * self.vr[1:] 

572 c0 += ((Mr - r) * (self.vr[1:] - e * r) + 

573 0.5 * x * self.r2dvdr[1:] / Mr) / r**2 

574 c1 -= 0.5 * x * self.r2dvdr[1:] / (Mr * dr * r) 

575 

576 fp = c2 + 0.5 * c1 

577 fm = c2 - 0.5 * c1 

578 f0 = c0 - 2 * c2 

579 kr = np.zeros(self.N) 

580 kr[1:] = f0 * u[1:] + fm * u[:-1] 

581 kr[1:-1] += fp[:-1] * u[2:] 

582 kr[0] = 0.0 

583 return kr 

584 

585 def __del__(self): 

586 self.close() 

587 

588 

589def shoot(u, l, vr, e, r2dvdr, r, dr, c10, c2, scalarrel, gmax=None): 

590 """n, A = shoot(u, l, vr, e, ...) 

591 

592 For guessed trial eigenenergy e, integrate the radial Schrodinger 

593 equation:: 

594 

595 2 

596 d u du 

597 --- c + -- c + u c = 0 

598 2 2 dg 1 0 

599 dg 

600 

601 2 dg 2 

602 c = -r (--) 

603 2 dr 

604 

605 2 2 

606 d g 2 r dg dv 

607 c = - --- r - ---- -- -- 

608 1 2 2 dr dr 

609 dr 2Mc 

610 

611 2 r dv 

612 c = l(l + 1) + 2M(v - e)r + ---- -- 

613 0 2 dr 

614 2Mc 

615 

616 The resulting wavefunction is returned in input vector u. 

617 The number of nodes of u is returned in attribute n. 

618 Returned attribute A, is a measure of the size of the derivative 

619 discontinuity at the classical turning point. 

620 The trial energy e is correct if A is zero and n is the correct number 

621 of nodes.""" 

622 

623 if scalarrel: 

624 x = 0.5 * alpha**2 # x = 1 / (2c^2) 

625 Mr = r * (1.0 + x * e) - x * vr 

626 else: 

627 Mr = r 

628 c0 = l * (l + 1) + 2 * Mr * (vr - e * r) 

629 if gmax is None and (c0 > 0).all(): 

630 raise ConvergenceError('Bad initial electron density guess!') 

631 

632 c1 = c10 

633 if scalarrel: 

634 c0 += x * r2dvdr / Mr 

635 c1 = c10 - x * r * r2dvdr / (Mr * dr) 

636 

637 # vectors needed for numeric integration of diff. equation 

638 fm = 0.5 * c1 - c2 

639 fp = 0.5 * c1 + c2 

640 f0 = c0 - 2 * c2 

641 

642 if gmax is None: 

643 # set boundary conditions at r -> oo (u(oo) = 0 is implicit) 

644 u[-1] = 1.0 

645 

646 # perform backwards integration from infinity to the turning point 

647 g = len(u) - 2 

648 u[-2] = u[-1] * f0[-1] / fm[-1] 

649 while c0[g] > 0.0: # this defines the classical turning point 

650 u[g - 1] = (f0[g] * u[g] + fp[g] * u[g + 1]) / fm[g] 

651 if u[g - 1] < 0.0: 

652 # There should't be a node here! Use a more negative 

653 # eigenvalue: 

654 print('!!!!!!', end=' ') 

655 return 100, None 

656 if u[g - 1] > 1e100: 

657 u *= 1e-100 

658 g -= 1 

659 

660 # stored values of the wavefunction and the first derivative 

661 # at the turning point 

662 gtp = g + 1 

663 utp = u[gtp] 

664 if gtp == len(u) - 1: 

665 return 100, 0.0 

666 dudrplus = 0.5 * (u[gtp + 1] - u[gtp - 1]) / dr[gtp] 

667 else: 

668 gtp = gmax 

669 

670 # set boundary conditions at r -> 0 

671 u[0] = 0.0 

672 u[1] = 1.0 

673 

674 # perform forward integration from zero to the turning point 

675 g = 1 

676 nodes = 0 

677 # integrate one step further than gtp 

678 # (such that dudr is defined in gtp) 

679 while g <= gtp: 

680 u[g + 1] = (fm[g] * u[g - 1] - f0[g] * u[g]) / fp[g] 

681 if u[g + 1] * u[g] < 0: 

682 nodes += 1 

683 g += 1 

684 if gmax is not None: 

685 return 

686 

687 # scale first part of wavefunction, such that it is continuous at gtp 

688 u[:gtp + 2] *= utp / u[gtp] 

689 

690 # determine size of the derivative discontinuity at gtp 

691 dudrminus = 0.5 * (u[gtp + 1] - u[gtp - 1]) / dr[gtp] 

692 A = (dudrplus - dudrminus) * utp 

693 

694 return nodes, A 

695 

696 

697def shoot_confined(u, l, vr, e, r2dvdr, r, dr, c10, c2, scalarrel, 

698 gmax=None, rc=10., beta=7.): 

699 """This method is used by the solve_confined method.""" 

700 # XXX much of this is pasted from the ordinary shoot method 

701 assert l < 3 

702 

703 if scalarrel: 

704 x = 0.5 * alpha**2 # x = 1 / (2c^2) 

705 Mr = r * (1.0 + x * e) - x * vr 

706 else: 

707 Mr = r 

708 c0 = l * (l + 1) + 2 * Mr * (vr - e * r) 

709 if gmax is None and (c0 > 0).all(): 

710 print(""" 

711Problem with initial electron density guess! Try to run the program 

712with the '-n' option (non-scalar-relativistic calculation) and then 

713try again without the '-n' option (this will generate a good initial 

714guess for the density). 

715""") 

716 raise SystemExit 

717 c1 = c10 

718 if scalarrel: 

719 c0 += x * r2dvdr / Mr 

720 c1 = c10 - x * r * r2dvdr / (Mr * dr) 

721 

722 # vectors needed for numeric integration of diff. equation 

723 fm = 0.5 * c1 - c2 

724 fp = 0.5 * c1 + c2 

725 f0 = c0 - 2 * c2 

726 

727 if gmax is None: 

728 gcut = int(rc * len(r) / (beta + rc)) 

729 # set boundary conditions at r -> oo (u(oo) = 0 is implicit) 

730 u[gcut - 1] = 1. 

731 u[gcut:] = 0. 

732 

733 # perform backwards integration from infinity to the turning point 

734 g = gcut - 2 

735 u[g] = u[g + 1] * f0[g + 1] / fm[g + 1] 

736 

737 while c0[g] > 0.0: # this defines the classical turning point 

738 u[g - 1] = (f0[g] * u[g] + fp[g] * u[g + 1]) / fm[g] 

739 if u[g - 1] < 0.0: 

740 # There should't be a node here! Use a more negative 

741 # eigenvalue: 

742 print('!!!!!!', end=' ') 

743 return 100, None 

744 if u[g - 1] > 1e100: 

745 u *= 1e-100 

746 g -= 1 

747 

748 # stored values of the wavefunction and the first derivative 

749 # at the turning point 

750 gtp = g + 1 

751 utp = u[gtp] 

752 dudrplus = 0.5 * (u[gtp + 1] - u[gtp - 1]) / dr[gtp] 

753 else: 

754 gtp = gmax 

755 

756 # set boundary conditions at r -> 0 

757 u[0] = 0.0 

758 u[1] = 1.0 

759 

760 # perform forward integration from zero to the turning point 

761 g = 1 

762 nodes = 0 

763 # integrate one step further than gtp 

764 # (such that dudr is defined in gtp) 

765 while g <= gtp: 

766 u[g + 1] = (fm[g] * u[g - 1] - f0[g] * u[g]) / fp[g] 

767 if u[g + 1] * u[g] < 0: 

768 nodes += 1 

769 g += 1 

770 if gmax is not None: 

771 return 

772 

773 # scale first part of wavefunction, such that it is continuous at gtp 

774 u[:gtp + 2] *= utp / u[gtp] 

775 

776 # determine size of the derivative discontinuity at gtp 

777 dudrminus = 0.5 * (u[gtp + 1] - u[gtp - 1]) / dr[gtp] 

778 A = (dudrplus - dudrminus) * utp 

779 

780 return nodes, A 

781 

782 

783@dataclass 

784class ValenceData: 

785 symbol: str 

786 xcname: str 

787 

788 rgd: AERadialGridDescriptor 

789 

790 n_j: list[int] 

791 l_j: list[int] 

792 e_j: list[float] 

793 f_j: list[float] 

794 

795 scalarrel: bool 

796 

797 phi_jg: list[np.ndarray] 

798 phit_jg: list[np.ndarray] 

799 pt_jg: list[np.ndarray] 

800 rcut_j: list[float] 

801 

802 vr_g: np.ndarray 

803 r2dvdr_g: np.ndarray | None # Actually: None means not scalarrel 

804 

805 def __post_init__(self): 

806 assert self.scalarrel == (self.r2dvdr_g is not None) 

807 jattributes = 'n_j l_j e_j f_j rcut_j'.split() 

808 

809 for attr in jattributes: 

810 thing_j = getattr(self, attr) 

811 assert len(thing_j) == self.nj, (attr, len(thing_j), self.nj) 

812 

813 err = abs(self.rgd.beta / len(self.rgd.r_g) - self.rgd.a) 

814 assert err < 1e-15, f'Inconsistent rgd spacing, {err=}' 

815 

816 @property 

817 def nj(self): 

818 return len(self.l_j) 

819 

820 @classmethod 

821 def calculate_potential_data(cls, setupdata): 

822 sqrt4pi = np.sqrt(4.0 * np.pi) 

823 rgd = setupdata.rgd 

824 xc = XC(setupdata.setupname) 

825 

826 # XXX GLLB needs special initialization I think. 

827 assert 'GLLB' not in setupdata.setupname 

828 

829 if setupdata.orbital_free: 

830 raise RuntimeError('Setup is orbital-free') 

831 

832 # f_j includes only valence states, so we need to add also 

833 # all-electron core density. 

834 n_g = calculate_density(setupdata.f_j, 

835 setupdata.phi_jg * rgd.r_g[None, :], 

836 rgd.r_g) 

837 n_g += setupdata.nc_g / sqrt4pi 

838 

839 vr_g = calculate_potentials(rgd, xc, n_g, setupdata.Z, 

840 tw_coeff=None)[0] 

841 r2dvdr_g = get_r2dvdr(rgd, vr_g) 

842 

843 assert setupdata.type in {'scalar-relativistic', 'non-relativistic'} 

844 scalarrel = setupdata.type == 'scalar-relativistic' 

845 return vr_g, r2dvdr_g, scalarrel 

846 

847 @classmethod 

848 def from_setupdata_onthefly_potentials(cls, setupdata): 

849 vr_g, r2dvdr_g, scalarrel = cls.calculate_potential_data(setupdata) 

850 return cls.from_setupdata_and_potentials( 

851 setupdata, vr_g=vr_g, 

852 # To comply with the assertion in `.__post_init__()` 

853 r2dvdr_g=r2dvdr_g if scalarrel else None, 

854 scalarrel=scalarrel) 

855 

856 @classmethod 

857 def from_setupdata_and_potentials(cls, setupdata, *, vr_g, r2dvdr_g, 

858 scalarrel): 

859 

860 assert len(setupdata.phi_jg) == len(setupdata.l_j) 

861 

862 def multiply_r(array_jg): 

863 return array_jg * setupdata.rgd.r_g[None, :] 

864 

865 return cls( 

866 xcname=setupdata.setupname, 

867 symbol=setupdata.symbol, 

868 rgd=setupdata.rgd, 

869 n_j=setupdata.n_j, 

870 l_j=setupdata.l_j, 

871 e_j=setupdata.eps_j, 

872 f_j=setupdata.f_j, 

873 rcut_j=setupdata.rcut_j, 

874 phi_jg=multiply_r(setupdata.phi_jg), 

875 phit_jg=multiply_r(setupdata.phit_jg), 

876 pt_jg=multiply_r(setupdata.pt_jg), 

877 vr_g=vr_g, 

878 r2dvdr_g=r2dvdr_g, 

879 scalarrel=scalarrel, 

880 ) 

881 

882 @property 

883 def N(self): 

884 return len(self.rgd.r_g) 

885 

886 @cached_property 

887 def d2gdr2_g(self): 

888 return self.rgd.d2gdr2() 

889 

890 def solve_confined(self, j, rc, vconf=None): 

891 """Solve the Schroedinger equation in a confinement potential. 

892 

893 Solves the Schroedinger equation like the solve method, but with a 

894 number of differences. Before invoking this method, run solve() to 

895 get initial guesses. 

896 

897 Parameters: 

898 j: solves only for the state given by j 

899 rc: solution cutoff. Solution will be zero outside this. 

900 vconf: added to the potential (use this as confinement potential) 

901 

902 Returns: a tuple containing the solution u and its energy e. 

903 

904 Unlike the solve method, this method will not alter any attributes of 

905 this object. 

906 """ 

907 r = self.rgd.r_g 

908 dr = self.rgd.dr_g 

909 vr = self.vr_g.copy() 

910 if vconf is not None: 

911 vr += vconf * r 

912 

913 c2 = -(r / dr)**2 

914 c10 = -self.d2gdr2_g * r**2 # first part of c1 vector 

915 

916 if j is None: 

917 n, l, e, u = 3, 2, -0.15, self.phi_jg[-1].copy() 

918 else: 

919 n = self.n_j[j] 

920 l = self.l_j[j] 

921 e = self.e_j[j] 

922 u = self.phi_jg[j].copy() 

923 

924 nn, A = shoot_confined(u, l, vr, e, self.r2dvdr_g, r, dr, c10, c2, 

925 self.scalarrel, rc=rc, beta=self.rgd.beta) 

926 assert nn == n - l - 1 # run() should have been called already 

927 

928 # adjust eigenenergy until u is smooth at the turning point 

929 de = 1.0 

930 while abs(de) > 1e-9: 

931 norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr) 

932 u *= 1.0 / sqrt(norm) 

933 de = 0.5 * A / norm 

934 x = abs(de / e) 

935 if x > 0.1: 

936 de *= 0.1 / x 

937 e -= de 

938 assert e < 0.0 

939 

940 nn, A = shoot_confined(u, l, vr, e, self.r2dvdr_g, r, dr, c10, c2, 

941 self.scalarrel, rc=rc, beta=self.rgd.beta) 

942 u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)) 

943 return u, e 

944 

945 def get_confinement_potential(self, alpha, ri, rc): 

946 r"""Create a smooth confinement potential. 

947 

948 Returns a (potential) function which is zero inside the radius ri 

949 and goes to infinity smoothly at rc, after which point it is nan. 

950 The potential is given by:: 

951 

952 alpha / rc - ri \ 

953 V(r) = -------- exp ( - --------- ) for ri < r < rc 

954 rc - r \ r - ri / 

955 

956 """ 

957 i_ri = self.rgd.floor(ri) 

958 i_rc = self.rgd.floor(rc) 

959 if self.rgd.r_g[i_rc] == rc: 

960 # Avoid division by zero in the odd case that rc coincides 

961 # exactly with a grid point (which actually happens sometimes) 

962 i_rc -= 1 

963 

964 potential = np.zeros(np.shape(self.rgd.r_g)) 

965 r = self.rgd.r_g[i_ri + 1:i_rc + 1] 

966 exponent = - (rc - ri) / (r - ri) 

967 denom = rc - r 

968 value = np.exp(exponent) / denom 

969 potential[i_ri + 1:i_rc + 1] = value 

970 potential[i_rc + 1:] = np.inf 

971 

972 return alpha * potential 

973 

974 

975if __name__ == '__main__': 

976 a = AllElectron('Cu', scalarrel=True) 

977 a.run()