Coverage for gpaw/atom/aeatom.py: 81%

1import copy

2import sys

3from itertools import cycle

4from math import pi

6import numpy as np

7from numpy.linalg import eigh

8from scipy.special import gamma

9import ase.units as units

10from ase.data import atomic_numbers, atomic_names, chemical_symbols

11from ase.utils import seterr

13import gpaw.cgpaw as cgpaw

14from gpaw.xc import XC

15from gpaw.gaunt import gaunt

16from gpaw.atom.configurations import configurations

17from gpaw.atom.radialgd import (AERadialGridDescriptor,

18 AbinitRadialGridDescriptor)

21# Velocity of light in atomic units:

22c = 2 * units._hplanck / (units._mu0 * units._c * units._e**2)

24# Colors for s, p, d, f, g:

27class _Colors:

28 def __init__(self, items):

29 self.items = items

31 def __getitem__(self, i):

32 return self.items[i % len(self.items)]

34 def __iter__(self):

35 yield from cycle(self.items)

38colors = _Colors('krgbycm')

41class GaussianBasis:

42 def __init__(self, l, alpha_B, rgd, eps=1.0e-7):

43 """Guassian basis set for spherically symmetric atom.

45 l: int

46 Angular momentum quantum number.

47 alpha_B: ndarray

48 Exponents.

49 rgd: GridDescriptor

50 Grid descriptor.

51 eps: float

52 Cutoff for eigenvalues of overlap matrix."""

54 self.l = l

55 self.alpha_B = alpha_B

56 self.rgd = rgd

57 self.eps = eps

59 A_BB = np.add.outer(alpha_B, alpha_B)

60 M_BB = np.multiply.outer(alpha_B, alpha_B)

62 # Overlap matrix:

63 S_BB = (2 * M_BB**0.5 / A_BB)**(l + 1.5)

65 # Kinetic energy matrix:

66 T_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (

67 gamma(l + 2.5) * M_BB / A_BB**(l + 2.5) -

68 0.5 * (l + 1) * gamma(l + 1.5) / A_BB**(l + 0.5) +

69 0.25 * (l + 1) * (2 * l + 1) * gamma(l + 0.5) / A_BB**(l + 0.5))

71 # Derivative matrix:

72 D_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (

73 0.5 * (l + 1) * gamma(l + 1) / A_BB**(l + 1) -

74 gamma(l + 2) * alpha_B / A_BB**(l + 2))

76 # 1/r matrix:

77 K_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (

78 0.5 * gamma(l + 1) / A_BB**(l + 1))

80 # Find set of linearly independent functions.

81 # We have len(alpha_B) gaussians (index B) and self.nbasis

82 # linearly independent functions (index b).

83 s_B, U_BB = eigh(S_BB)

84 self.nbasis = int((s_B > eps).sum())

86 Q_Bb = np.dot(U_BB[:, -self.nbasis:],

87 np.diag(s_B[-self.nbasis:]**-0.5))

89 self.T_bb = np.dot(np.dot(Q_Bb.T, T_BB), Q_Bb)

90 self.D_bb = np.dot(np.dot(Q_Bb.T, D_BB), Q_Bb)

91 self.K_bb = np.dot(np.dot(Q_Bb.T, K_BB), Q_Bb)

93 r_g = rgd.r_g

94 prefactors_B = (2 * (2 * alpha_B)**(l + 1.5) / gamma(l + 1.5))**0.5

95 with seterr(divide='ignore', invalid='ignore'):

96 # Avoid errors division by zero when l<0.0:

97 gaussians_Bg = np.exp(-np.outer(alpha_B, r_g**2)) * r_g**l

98 self.basis_bg = Q_Bb.T @ (prefactors_B[:, None] * gaussians_Bg)

100 def __len__(self):

101 return self.nbasis

102

103 def expand(self, C_xb):

104 return np.dot(C_xb, self.basis_bg)

105

106 def calculate_potential_matrix(self, vr_g):

107 vr2dr_g = vr_g * self.rgd.r_g * self.rgd.dr_g

108 V_bb = np.inner(self.basis_bg[:, 1:],

109 self.basis_bg[:, 1:] * vr2dr_g[1:])

110 return V_bb

111

112

113def coefs(rgd, l, vr_g, e, scalar_relativistic=True, Z=None):

114 r_g = rgd.r_g

115

116 x0_g = 2 * (e * r_g - vr_g) * r_g

117 x1_g = 2 * (l + 1) * r_g / rgd.dr_g + r_g**2 * rgd.d2gdr2()

118 x2_g = r_g**2 / rgd.dr_g**2

119 x = 1.0

120

121 if scalar_relativistic:

122 x = (1 + l + l**2 - (Z / c)**2)**0.5 - l

123 r_g = r_g.copy()

124 r_g[0] = 1.0

125 v_g = vr_g / r_g

126 M_g = 1 + (e - v_g) / (2 * c**2)

127 kappa_g = (rgd.derivative(vr_g) - v_g) / r_g / (2 * c**2 * M_g)

128 x0_g = (2 * M_g * (e * r_g - vr_g) * r_g +

129 (l + x - 1) * kappa_g * r_g +

130 (l + x) * (l + x - 1) - l * (l + 1))

131 x1_g = (2 * (l + x) * r_g / rgd.dr_g +

132 r_g**2 * rgd.d2gdr2() +

133 r_g**2 * kappa_g / rgd.dr_g)

134

135 cm1_g = x2_g - x1_g / 2

136 c0_g = x0_g - 2 * x2_g

137 cp1_g = x2_g + x1_g / 2

138

139 return cm1_g, c0_g, cp1_g, x

140

141

142class Channel:

143 def __init__(self, l, s=0, f_n=(), basis=None):

144 self.l = l

145 self.s = s

146 self.basis = basis

147

148 self.C_nb = None # eigenvectors

149 self.e_n = None # eigenvalues

150 self.f_n = np.array(f_n, dtype=float) # occupation numbers

151 self.phi_ng = None # wave functions

152

153 self.name = 'spdfg'[l]

154 self.solve2ok = False

155

156 def solve(self, vr_g):

157 """Diagonalize Schrödinger equation in basis set."""

158 H_bb = self.basis.calculate_potential_matrix(vr_g)

159 H_bb += self.basis.T_bb

160 self.e_n, C_bn = eigh(H_bb)

161 self.C_nb = C_bn.T

162 self.phi_ng = self.basis.expand(self.C_nb[:len(self.f_n)])

163

164 def solve2(self, vr_g, scalar_relativistic=True, Z=None, rgd=None):

165 rgd = rgd or self.basis.rgd

166 r_g = rgd.r_g

167 l = self.l

168 u_g = rgd.empty()

169 self.solve2ok = True

170 for n in range(len(self.f_n)):

171 e = self.e_n[n]

172

173 if e > 0.0:

174 # Skip orbitals with positive energies

175 self.e_n[n] = 42

176 self.phi_ng[n] = 0.0

177 continue

178

179 # Find classical turning point:

180 x = vr_g * r_g + 0.5 * l * (l + 1) - e * r_g**2

181 g0 = rgd.round(4.0)

182 while x[g0] > 0:

183 g0 -= 1

184

185 iter = 0

186 ok = False

187 while True:

188 du1dr, a = self.integrate_outwards(u_g, rgd, vr_g, g0, e,

189 scalar_relativistic, Z)

190 u1 = u_g[g0]

191 du2dr = self.integrate_inwards(u_g, rgd, vr_g, g0, e,

192 scalar_relativistic, Z)

193 u2 = u_g[g0]

194 A = du1dr / u1 - du2dr / u2

195 u_g[g0:] *= u1 / u2

196 norm = rgd.integrate(u_g**2, -2) / (4 * pi)

197 u_g /= norm**0.5

198 a /= norm**0.5

199

200 nodes = (u_g[:-1] * u_g[1:] < 0).sum()

201

202 if abs(A) < 1e-7 and nodes == n:

203 ok = True

204 break

205

206 if nodes > n:

207 e *= 1.2

208 elif nodes < n:

209 e *= 0.8

210 else:

211 e += 0.5 * A * u_g[g0]**2

212 if e > 0:

213 break

214

215 iter += 1

216 assert iter < 400, (n, l, e)

217

218 if ok:

219 self.e_n[n] = e

220 self.phi_ng[n, 1:] = u_g[1:] / r_g[1:]

221 self.phi_ng[n, 0] = a

222 else:

223 self.solve2ok = False

224

225 def calculate_density(self, n=None):

226 """Calculate density."""

227 if n is None:

228 n_g = 0.0

229 for n, f in enumerate(self.f_n):

230 n_g += f * self.calculate_density(n)

231 else:

232 n_g = self.phi_ng[n]**2 / (4 * pi)

233 return n_g

234

235 def calculate_kinetic_energy_density(self, n):

236 """Calculate kinetic energy density."""

237 phi_g = self.phi_ng[n]

238 rgd = self.basis.rgd

239 tau_g = rgd.derivative(phi_g)**2 / (8 * pi)

240 if self.l > 0:

241 tau_g[1:] += (self.l * (self.l + 1) *

242 (phi_g[1:] / rgd.r_g[1:])**2 / (8 * pi))

243 return tau_g

244

245 def get_eigenvalue_sum(self):

246 f_n = self.f_n

247 return np.dot(f_n, self.e_n[:len(f_n)])

248

249 def integrate_outwards(self, u_g, rgd, vr_g, g0, e,

250 scalar_relativistic=True, Z=None, pt_g=None):

251 l = self.l

252 r_g = rgd.r_g

253

254 cm1_g, c0_g, cp1_g, x = coefs(rgd, l, vr_g, e, scalar_relativistic, Z)

255

256 b_g = rgd.zeros()

257

258 if Z is not None:

259 a0 = 1.0

260 if scalar_relativistic:

261 dadr = 2 * ((l + x - 1) * c**2 / Z - Z) / (1 + 2 * (l + x))

262 else:

263 dadr = -Z / (l + 1)

264 a1 = a0 + dadr * r_g[1]

265 else:

266 assert not scalar_relativistic

267 if pt_g is None:

268 a0 = 1.0

269 else:

270 b_g[1:] = 2 * pt_g[1:] * r_g[1:]**(2 - l)

271 a0 = pt_g[1] / r_g[1]**l / (vr_g[1] / r_g[1] - e)

272 a1 = a0

273

274 if 0:

275 u_g[0] = 0.0

276 g = 1

277 agm1 = a0

278 ag = a1

279 while True:

280 u_g[g] = ag * r_g[g]**(l + x)

281 agp1 = -(agm1 * cm1_g[g] + ag * c0_g[g] + b_g[g]) / cp1_g[g]

282 if g == g0:

283 break

284 g += 1

285 agm1 = ag

286 ag = agp1

287 else:

288 a_g = np.zeros_like(u_g)

289 a_g[:2] = (a0, a1)

290 cgpaw.integrate_outwards(g0, cm1_g, c0_g, cp1_g, b_g, a_g)

291 u_g[:g0 + 1] = a_g[:g0 + 1] * r_g[:g0 + 1]**(l + x)

292 g = g0

293 agm1, ag, agp1 = a_g[g - 1:g + 2]

294

295 r = r_g[g0]

296 dr = rgd.dr_g[g0]

297 da = 0.5 * (agp1 - agm1)

298 dudr = (l + x) * r**(l + x - 1) * ag + r**(l + x) * da / dr

299

300 if l - 1 + x < 0:

301 phi0 = a0 * (r_g[1] * 0.1)**(l - 1 + x)

302 else:

303 phi0 = a0 * 0.0**(l - 1 + x)

304

305 return dudr, phi0

306

307 def integrate_inwards(self, u_g, rgd, vr_g, g0, e,

308 scalar_relativistic=True, Z=None, gmax=None):

309 l = self.l

310 r_g = rgd.r_g

311

312 cm1_g, c0_g, cp1_g, x = coefs(rgd, l, vr_g, e, scalar_relativistic, Z)

313

314 cm1_g[:g0] = 1.0 # prevent division by zero

315 c0_g /= -cm1_g

316 cp1_g /= -cm1_g

317

318 if gmax is None:

319 gmax = len(u_g)

320

321 g = gmax - 2

322 agp1 = 1.0

323 u_g[gmax - 1] = agp1 * r_g[gmax - 1]**(l + x)

324 with np.errstate(over='raise'):

325 try:

326 ag = np.exp(-(-2 * e)**0.5 * (r_g[gmax - 2] - r_g[gmax - 1]))

327 except FloatingPointError:

328 ag = 2e50

329

330 if 0:

331 while True:

332 u_g[g] = ag * r_g[g]**(l + x)

333 if ag > 1e50:

334 u_g[g:] /= 1e50

335 ag = ag / 1e50

336 agp1 = agp1 / 1e50

337 agm1 = agp1 * cp1_g[g] + ag * c0_g[g]

338 if g == g0:

339 break

340 g -= 1

341 agp1 = ag

342 ag = agm1

343 else:

344 a_g = np.zeros_like(u_g)

345 a_g[g:g + 2] = (ag, agp1)

346 cgpaw.integrate_inwards(g, g0, c0_g, cp1_g, a_g)

347 u_g[g0:g + 2] = a_g[g0:g + 2] * r_g[g0:g + 2]**(l + x)

348 g = g0

349 agm1, ag, agp1 = a_g[g - 1:g + 2]

350

351 # print(agm1, ag, agp1, u_g[g - 1:]);asdfg

352 r = r_g[g]

353 dr = rgd.dr_g[g]

354 da = 0.5 * (agp1 - agm1)

355 dudr = (l + x) * r**(l + x - 1) * ag + r**(l + x) * da / dr

356

357 return dudr

358

359

360class DiracChannel(Channel):

361 def __init__(self, k, f_n, basis):

362 l = (abs(2 * k + 1) - 1) // 2

363 super().__init__(l, 0, f_n, basis)

364 self.k = k

365 self.j = abs(k) - 0.5

366 self.c_nb = None # eigenvectors (small component)

367

368 self.name += '(%d/2)' % (2 * self.j)

369

370 def solve(self, vr_g):

371 """Solve Dirac equation in basis set."""

372 nb = len(self.basis)

373 V_bb = self.basis.calculate_potential_matrix(vr_g)

374 H_bb = np.zeros((2 * nb, 2 * nb))

375 H_bb[:nb, :nb] = V_bb

376 H_bb[nb:, nb:] = V_bb - 2 * c**2 * np.eye(nb)

377 H_bb[nb:, :nb] = -c * (-self.basis.D_bb.T + self.k * self.basis.K_bb)

378 e_n, C_bn = eigh(H_bb)

379 if self.k < 0:

380 n0 = nb

381 else:

382 n0 = nb + 1

383 self.e_n = e_n[n0:].copy()

384 self.C_nb = C_bn[:nb, n0:].T.copy() # large component

385 self.c_nb = C_bn[nb:, n0:].T.copy() # small component

386

387 def calculate_density(self, n=None):

388 """Calculate density."""

389 if n is None:

390 n_g = Channel.calculate_density(self)

391 else:

392 n_g = (self.basis.expand(self.C_nb[n])**2 +

393 self.basis.expand(self.c_nb[n])**2) / (4 * pi)

394 if self.basis.l < 0:

395 n_g[0] = n_g[1]

396 return n_g

397

398

399class AllElectronAtom:

400 def __init__(self, symbol, xc='LDA', spinpol=False, dirac=False,

401 configuration=None,

402 ee_interaction=True,

403 Z=None,

404 log=None,

405 scalar_relativistic=True):

406 """All-electron calculation for spherically symmetric atom.

407

408 symbol: str (or int)

409 Chemical symbol (or atomic number).

410 xc: str

411 Name of XC-functional.

412 spinpol: bool

413 If true, do spin-polarized calculation. Default is spin-paired.

414 dirac: bool

415 Solve Dirac equation instead of Schrödinger equation.

416 configuration: list

417 Electronic configuration for symbol, format as in

418 gpaw.atom.configurations

419 ee_interaction: bool

420 Use ee_interaction=False to turn off electron-electron interaction.

421 Default is True.

422 log: stream

423 Text output."""

424

425 if isinstance(symbol, int):

426 symbol = chemical_symbols[symbol]

427 self.symbol = symbol

428 self.Z = Z or atomic_numbers[symbol]

429

430 self.nspins = 1 + int(bool(spinpol))

431

432 self.dirac = bool(dirac)

433

434 self.ee_interaction = ee_interaction

435

436 if configuration is not None:

437 self.configuration = copy.deepcopy(configuration)

438 else:

439 self.configuration = None

440

441 self.scalar_relativistic = bool(scalar_relativistic)

442

443 if isinstance(xc, str):

444 self.xc = XC(xc)

445 else:

446 self.xc = xc

447

448 self.fd = log or sys.stdout

449

450 self.vr_sg = None # potential * r

451 self.n_sg = 0.0 # density

452 self.rgd = None # radial grid descriptor

453

454 # Energies:

455 self.ekin = None

456 self.eeig = None

457 self.eH = None

458 self.eZ = None

459

460 self.channels = None

461

462 self.initialize_configuration(self.configuration)

463

464 self.log('Z: ', self.Z)

465 self.log('Name: ', atomic_names[atomic_numbers[symbol]])

466 self.log('Symbol: ', symbol)

467 self.log('XC-functional: ', self.xc.name)

468 self.log('Equation: ', ['Schrödinger', 'Dirac'][self.dirac])

469

470 self.method = 'Gaussian basis-set'

471

472 def log(self, *args, **kwargs):

473 print(file=self.fd, *args, **kwargs)

474

475 def initialize_configuration(self, configuration=None):

476 self.f_lsn = {}

477

478 if configuration is None:

479 configs = configurations[self.symbol][1]

480 elif isinstance(configuration, str):

481 configs = []

482 if configuration[0] == '[':

483 symbol, configuration = configuration[1:].split(']')

484 configs = configurations[symbol][1]

485 for nlf in configuration.split(','):

486 configs.append((int(nlf[0]),

487 'spdfg'.index(nlf[1]),

488 int(nlf[2:]),

489 -1.0))

490 else:

491 configs = configuration

492

493 for n, l, f, e in configs:

494

495 if l not in self.f_lsn:

496 self.f_lsn[l] = [[] for s in range(self.nspins)]

497 if self.nspins == 1:

498 self.f_lsn[l][0].append(f)

499 else:

500 # Use Hund's rule:

501 f0 = min(f, 2 * l + 1)

502 self.f_lsn[l][0].append(f0)

503 self.f_lsn[l][1].append(f - f0)

504

505 if 0:

506 n = 2 + len(self.f_lsn[2][0])

507 if self.f_lsn[0][0][n] == 2:

508 self.f_lsn[0][0][n] = 1

509 self.f_lsn[2][0][n - 3] += 1

510

511 def add(self, n, l, df=+1, s=None):

512 """Add (remove) electrons."""

513 if s is None:

514 if self.nspins == 1:

515 s = 0

516 else:

517 self.add(n, l, 0.5 * df, 0)

518 self.add(n, l, 0.5 * df, 1)

519 return

520

521 if l not in self.f_lsn:

522 self.f_lsn[l] = [[] for x in range(self.nspins)]

523

524 f_n = self.f_lsn[l][s]

525 if len(f_n) < n - l:

526 f_n.extend([0] * (n - l - len(f_n)))

527 f_n[n - l - 1] += df

528

529 def initialize(self, ngpts=2000, rcut=50.0,

530 alpha1=0.01, alpha2=None, ngauss=50,

531 eps=1.0e-7):

532 """Initialize basis sets and radial grid.

533

534 ngpts: int

535 Number of grid points for radial grid.

536 rcut: float

537 Cutoff for radial grid.

538 alpha1: float

539 Smallest exponent for gaussian.

540 alpha2: float

541 Largest exponent for gaussian.

542 ngauss: int

543 Number of gaussians.

544 eps: float

545 Cutoff for eigenvalues of overlap matrix."""

546

547 if alpha2 is None:

548 alpha2 = 50.0 * self.Z**2

549

550 if 1:

551 # Use grid with r(0)=0, r(1)=a and r(ngpts)=rcut:

552 a = 1 / alpha2**0.5 / 20

553 b = (rcut - a * ngpts) / (rcut * ngpts)

554 b = 1 / round(1 / b)

555 self.rgd = AERadialGridDescriptor(a, b, ngpts)

556 else:

557 from scipy.optimize import root

558 rT = self.Z / 137**2

559 r1 = rT / 10 / 5

560 sol = root(lambda d: r1 / d * (np.exp(d * (ngpts - 1)) - 1) - rcut,

561 0.1)

562 d = sol.x[0]

563 a = r1 / d

564 self.rgd = AbinitRadialGridDescriptor(a, d, ngpts)

565

566 self.log('Grid points: %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %

567 ((self.rgd.N,) + tuple(self.rgd.r_g[[0, 1, 2, -2, -1]])))

568

569 # Distribute exponents between alpha1 and alpha2:

570 alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)

571 self.log('Exponents: %d (%.3f, %.3f, ..., %.3f, %.3f)' %

572 ((ngauss,) + tuple(alpha_B[[0, 1, -2, -1]])))

573

574 # Maximum l value:

575 lmax = max(self.f_lsn.keys())

576

577 self.channels = []

578 nb_l = []

579 if not self.dirac:

580 for l in range(lmax + 1):

581 basis = GaussianBasis(l, alpha_B, self.rgd, eps)

582 nb_l.append(len(basis))

583 for s in range(self.nspins):

584 self.channels.append(Channel(l, s, self.f_lsn[l][s],

585 basis))

586 else:

587 for K in range(1, lmax + 2):

588 leff = (K**2 - (self.Z / c)**2)**0.5 - 1

589 basis = GaussianBasis(leff, alpha_B, self.rgd, eps)

590 nb_l.append(len(basis))

591 for k, l in [(-K, K - 1), (K, K)]:

592 if l > lmax:

593 continue

594 f_n = self.f_lsn[l][0]

595 j = abs(k) - 0.5

596 f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)

597 self.channels.append(DiracChannel(k, f_n, basis))

598

599 self.log('Basis functions: %s (%s)' %

600 (', '.join([str(nb) for nb in nb_l]),

601 ', '.join('spdf'[:lmax + 1])))

602

603 self.vr_sg = self.rgd.zeros(self.nspins)

604 self.vr_sg[:] = -self.Z

605

606 def solve(self):

607 """Diagonalize Schrödinger equation."""

608 self.eeig = 0.0

609 for channel in self.channels:

610 if self.method == 'Gaussian basis-set':

611 channel.solve(self.vr_sg[channel.s])

612 else:

613 channel.solve2(self.vr_sg[channel.s], self.scalar_relativistic,

614 self.Z)

615

616 self.eeig += channel.get_eigenvalue_sum()

617

618 def calculate_density(self):

619 """Calculate elctron density and kinetic energy."""

620 self.n_sg = self.rgd.zeros(self.nspins)

621 for channel in self.channels:

622 self.n_sg[channel.s] += channel.calculate_density()

623

624 def calculate_electrostatic_potential(self):

625 """Calculate electrostatic potential and energy."""

626 n_g = self.n_sg.sum(0)

627 if self.ee_interaction:

628 self.vHr_g = self.rgd.poisson(n_g)

629 else:

630 self.vHr_g = self.rgd.zeros()

631 self.eH = 0.5 * self.rgd.integrate(n_g * self.vHr_g, -1)

632 self.eZ = -self.Z * self.rgd.integrate(n_g, -1)

633

634 def calculate_xc_potential(self):

635 self.vxc_sg = self.rgd.zeros(self.nspins)

636 self.exc = self.xc.calculate_spherical(self.rgd, self.n_sg,

637 self.vxc_sg)

638

639 def step(self):

640 self.solve()

641 self.calculate_density()

642 self.calculate_electrostatic_potential()

643 self.calculate_xc_potential()

644 self.vr_sg = self.vxc_sg * self.rgd.r_g

645 self.vr_sg += self.vHr_g

646 self.vr_sg -= self.Z

647 self.ekin = (self.eeig -

648 self.rgd.integrate((self.vr_sg * self.n_sg).sum(0), -1))

649

650 def run(self, mix=0.4, maxiter=117, dnmax=1e-9):

651 if self.channels is None:

652 self.initialize()

653

654 if self.dirac:

655 equation = 'Dirac'

656 elif self.scalar_relativistic:

657 equation = 'scalar-relativistic Schrödinger'

658 else:

659 equation = 'non-relativistic Schrödinger'

660 self.log(f'\nSolving {equation} equation using {self.method}:')

661

662 dn = self.Z

663

664 vr_old_sg = None

665 n_old_sg = None

666 for iter in range(maxiter):

667 self.log('.', end='')

668 self.fd.flush()

669 if iter > 0:

670 self.vr_sg *= mix

671 self.vr_sg += (1 - mix) * vr_old_sg

672 dn = self.rgd.integrate(abs(self.n_sg - n_old_sg).sum(0))

673 if dn <= dnmax:

674 self.log('\nConverged in', iter, 'steps')

675 break

676

677 vr_old_sg = self.vr_sg

678 n_old_sg = self.n_sg

679 self.step()

680

681 self.summary()

682

683 if self.method != 'Gaussian basis-set':

684 for channel in self.channels:

685 assert channel.solve2ok

686

687 if dn > dnmax:

688 raise RuntimeError('Did not converge!')

689

690 def refine(self):

691 self.method = 'finite difference'

692 self.run(dnmax=1e-6, mix=0.14, maxiter=200)

693

694 def summary(self):

695 self.write_states()

696 self.write_energies()

697

698 def write_states(self):

699 self.log('\n state occupation eigenvalue <r>')

700 if self.dirac:

701 self.log(' nl(j) [Hartree] [eV] [Bohr]')

702 else:

703 self.log(' nl [Hartree] [eV] [Bohr]')

704 self.log('-----------------------------------------------------')

705 states = []

706 for ch in self.channels:

707 for n, f in enumerate(ch.f_n):

708 states.append((ch.e_n[n], n, ch.s, ch))

709 states.sort()

710 for e, n, s, ch in states:

711 name = str(n + ch.l + 1) + ch.name

712 if self.nspins == 2:

713 name += '(%s)' % '+-'[s]

714 n_g = ch.calculate_density(n)

715 rave = self.rgd.integrate(n_g, 1)

716 self.log(' %-7s %6.3f %13.6f %13.5f %6.3f' %

717 (name, ch.f_n[n], e, e * units.Hartree, rave))

718

719 def write_energies(self):

720 self.log('\nEnergies: [Hartree] [eV]')

721 self.log('--------------------------------------------')

722 for text, e in [('kinetic ', self.ekin),

723 ('coulomb (e-e)', self.eH),

724 ('coulomb (e-n)', self.eZ),

725 ('xc ', self.exc),

726 ('total ',

727 self.ekin + self.eH + self.eZ + self.exc)]:

728 self.log(f' {text} {e:+13.6f} {units.Hartree * e:+13.5f}')

729

730 if not self.dirac:

731 self.calculate_exx()

732 self.log(f'\nExact exchange energy: {self.exx:.6f} Hartree, '

733 f'{self.exx * units.Hartree:.5f} eV')

734

735 def get_channel(self, l=None, s=0, k=None):

736 if self.dirac:

737 for channel in self.channels:

738 if channel.k == k:

739 return channel

740 else:

741 for channel in self.channels:

742 if channel.l == l and channel.s == s:

743 return channel

744 raise ValueError

745

746 def get_orbital(self, n, l=None, s=0, k=None):

747 channel = self.get_channel(l, s, k)

748 return channel.basis.expand(channel.C_nb[n])

749

750 def plot_wave_functions(self, rc=4.0, show=True):

751 import matplotlib.pyplot as plt

752 for ch in self.channels:

753 for n in range(len(ch.f_n)):

754 fr_g = ch.basis.expand(ch.C_nb[n]) * self.rgd.r_g

755 # fr_g = ch.phi_ng[n]

756 name = str(n + ch.l + 1) + ch.name

757 lw = 2

758 if self.nspins == 2:

759 name += '(%s)' % '+-'[ch.s]

760 if ch.s == 1:

761 lw = 1

762 if self.dirac and ch.k > 0:

763 lw = 1

764 ls = ['-', '--', '-.', ':'][ch.l]

765 n_g = ch.calculate_density(n)

766 rave = self.rgd.integrate(n_g, 1)

767 gave = self.rgd.round(rave)

768 fr_g *= np.sign(fr_g[gave])

769 plt.plot(self.rgd.r_g, fr_g,

770 ls=ls, lw=lw, color=colors[n + ch.l], label=name)

771

772 plt.legend(loc='best')

773 plt.xlabel('r [Bohr]')

774 plt.ylabel('$r\\phi(r)$')

775 plt.axis(xmin=0, xmax=rc)

776 if show:

777 plt.show()

778

779 def logarithmic_derivative(self, l, energies, rcut):

780 ch = Channel(l)

781 gcut = self.rgd.round(rcut)

782 u_g = self.rgd.empty()

783 logderivs = []

784 d0 = 42.0

785 offset = 0

786 for e in energies:

787 dudr = ch.integrate_outwards(u_g, self.rgd, self.vr_sg[0],

788 gcut, e, self.scalar_relativistic,

789 self.Z)[0]

790 d1 = np.arctan(dudr / u_g[gcut]) / pi + offset

791 if d1 > d0:

792 offset -= 1

793 d1 -= 1

794 logderivs.append(d1)

795 d0 = d1

796

797 return np.array(logderivs)

798

799 def calculate_exx(self, s=None):

800 if s is None:

801 self.exx = sum(self.calculate_exx(s)

802 for s in range(self.nspins)) / self.nspins

803 return self.exx

804

805 states = []

806 lmax = 0

807 for ch in self.channels:

808 l = ch.l

809 for n, phi_g in enumerate(ch.phi_ng):

810 f = ch.f_n[n]

811 if f > 0 and ch.s == s:

812 states.append((l, f * self.nspins / 2.0 / (2 * l + 1),

813 phi_g))

814 if l > lmax:

815 lmax = l

816

817 G_LLL = gaunt(lmax)

818

819 exx = 0.0

820 j1 = 0

821 for l1, f1, phi1_g in states:

822 f = 1.0

823 for l2, f2, phi2_g in states[j1:]:

824 n_g = phi1_g * phi2_g

825 for l in range((l1 + l2) % 2, l1 + l2 + 1, 2):

826 G = (G_LLL[l1**2:(l1 + 1)**2,

827 l2**2:(l2 + 1)**2,

828 l**2:(l + 1)**2]**2).sum()

829 vr_g = self.rgd.poisson(n_g, l)

830 e = f * self.rgd.integrate(vr_g * n_g, -1) / 4 / pi

831 exx -= e * G * f1 * f2

832 f = 2.0

833 j1 += 1

834

835 return exx

836

837

838class CLICommand:

839 """Solve radial equation for an atom.

840

841 Example:

842

843 gpaw atom Li -f PBE -p # plot wave functions for a lithium atom

844 """

845

846 @staticmethod

847 def add_arguments(parser):

848 add = parser.add_argument

849 add('symbol')

850 add('-f', '--xc-functional', type=str, default='LDA',

851 help='Exchange-Correlation functional ' +

852 '(default value LDA)',

853 metavar='<XC>')

854 add('-a', '--add', metavar='states',

855 help='Add electron(s). Use "1s0.5a" to add 0.5 1s ' +

856 'electrons to the alpha-spin channel (use "b" for ' +

857 'beta-spin). The number of electrons defaults to ' +

858 'one. Examples: "1s", "2p2b", "4f0.1b,3d-0.1a".')

859 add('--spin-polarized', action='store_true')

860 add('-d', '--dirac', action='store_true',

861 help='Solve Dirac equation.')

862 add('-p', '--plot', action='store_true')

863 add('-e', '--exponents',

864 help='Exponents a: exp(-a*r^2). Use "-e 0.1:20.0:30" ' +

865 'to get 30 exponents from 0.1 to 20.0.')

866 add('-l', '--logarithmic-derivatives',

867 metavar='spdfg,e1:e2:de,radius',

868 help='Plot logarithmic derivatives. ' +

869 'Example: -l spdf,-1:1:0.05,1.3. ' +

870 'Energy range and/or radius can be left out.')

871 add('-n', '--ngrid', help='Specify number of grid points.')

872 add('-R', '--rcut', help='Radial cutoff.')

873 add('-r', '--refine', action='store_true')

874 add('-s', '--scalar-relativistic',

875 action='store_const', const=True,

876 help='Use scalar-relativistic setups (default).')

877 add('-S', '--non-relativistic',

878 action='store_const', const=False, dest='scalar_relativistic',

879 help='Don\'t use non-relativistic setups.')

880 add('--no-ee-interaction', action='store_true',

881 help='Turn off electron-electron interaction.')

882

883 @staticmethod

884 def run(args):

885 main(args)

886

887

888def parse_ld_str(s, energies=None, r=2.0):

889 parts = s.split(',')

890 lvalues = ['spdfg'.find(x) for x in parts.pop(0)]

891 if parts:

892 e1, e2, de = (float(x) for x in parts.pop(0).split(':'))

893 else:

894 e1, e2, de = energies

895 if parts:

896 r = float(parts.pop())

897 energies = np.linspace(e1, e2, int((e2 - e1) / de) + 1)

898 return lvalues, energies, r

899

900

901def main(args):

902 symbol = args.symbol

903

904 nlfs = []

905 if args.add:

906 for x in args.add.split(','):

907 n = int(x[0])

908 l = 'spdfg'.find(x[1])

909 x = x[2:]

910 if x and x[-1] in 'ab':

911 s = int(x[-1] == 'b')

912 args.spin_polarized = True

913 x = x[:-1]

914 else:

915 s = None

916 if x:

917 f = float(x)

918 else:

919 f = 1

920 nlfs.append((n, l, f, s))

921

922 aea_kwargs = dict(xc=args.xc_functional,

923 spinpol=args.spin_polarized,

924 dirac=args.dirac,

925 ee_interaction=not args.no_ee_interaction)

926 if args.scalar_relativistic is not None:

927 aea_kwargs['scalar_relativistic'] = args.scalar_relativistic

928 aea = AllElectronAtom(symbol, **aea_kwargs)

929

930 kwargs = {}

931 if args.exponents:

932 parts = args.exponents.split(':')

933 kwargs['alpha1'] = float(parts[0])

934 if len(parts) > 1:

935 kwargs['alpha2'] = float(parts[1])

936 if len(parts) > 2:

937 kwargs['ngauss'] = int(parts[2])

938

939 for n, l, f, s in nlfs:

940 aea.add(n, l, f, s)

941

942 if args.ngrid:

943 kwargs['ngpts'] = int(args.ngrid)

944 if args.rcut:

945 kwargs['rcut'] = float(args.rcut)

946

947 aea.initialize(**kwargs)

948 aea.run()

949

950 if args.refine or args.scalar_relativistic:

951 aea.refine()

952

953 if args.logarithmic_derivatives:

954 lvalues, energies, r = parse_ld_str(args.logarithmic_derivatives,

955 (-1, 1, 0.05))

956 import matplotlib.pyplot as plt

957 for l in lvalues:

958 ld = aea.logarithmic_derivative(l, energies, r)

959 plt.plot(energies, ld, colors[l])

960 plt.show()

961

962 if args.plot:

963 aea.plot_wave_functions()