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
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1"""
2Atomic Density Functional Theory
3"""
5from __future__ import annotations
6from dataclasses import dataclass
7from functools import cached_property
8from math import log, pi, sqrt
10import numpy as np
11from ase.data import atomic_names
12from ase.utils import IOContext
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
21# fine-structure constant
22alpha = 1 / 137.036
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
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)
37 # add potential from nuclear point charge (v = -Z / r)
38 vHr -= Z
39 return vHr
42def calculate_xc(rgd, xc, n_g):
43 vxc_g = np.zeros(rgd.N)
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
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
60 if tw_coeff is not None:
61 vr_g /= tw_coeff
63 return vr_g, vHr_g, vxc_g, Exc
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
73class AllElectron(IOContext):
74 """Object for doing an atomic DFT calculation."""
76 default_gpernode = 150
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.
84 Example::
86 a = AllElectron('Fe')
87 a.run()
88 """
90 self.txt = self.openfile(txt, comm=serial_comm)
92 self.symbol = symbol
93 self.xcname = xcname
94 self.scalarrel = scalarrel
95 self.nofiles = nofiles
97 # Get reference state:
98 self.Z, nlfe_j = configurations[symbol]
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]
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]
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)
135 self.orbital_free = orbital_free
136 self.tw_coeff = tw_coeff
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]]
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))
153 if corehole is not None:
154 self.ncorehole, self.lcorehole, self.fcorehole = corehole
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
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
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'))
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)
188 u[:] = r**(1 + l) * np.exp(-a * r)
189 norm = np.dot(u**2, dr)
190 u *= 1.0 / sqrt(norm)
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
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()
209 # Number of orbitals:
210 nj = len(self.n_j)
212 # Radial wave functions multiplied by radius:
213 self.u_j = np.zeros((nj, self.N))
215 # Effective potential multiplied by radius:
216 self.vr = np.zeros(N)
218 # Electron density:
219 self.n = np.zeros(N)
221 # Always spinpaired nspins=1
222 self.xc = XC(self.xcname)
224 # Initialize for non-local functionals
225 if self.xc.type == 'GLLB':
226 self.xc.initialize_1d(self)
228 n_j = self.n_j
229 l_j = self.l_j
230 f_j = self.f_j
231 e_j = self.e_j
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
239 vHr = np.zeros(self.N)
240 self.vXC = np.zeros(self.N)
242 self.initialize_wave_functions()
243 n[:] = self.calculate_density()
245 bar = '|------------------------------------------------|'
246 t(bar)
248 niter = 0
249 nitermax = 117
250 qOK = log(1e-10)
251 mix = 0.4
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
262 vrold = None
264 while True:
265 tw_coeff = self.tw_coeff if self.orbital_free else None
267 vr[:], vHr[:], self.vXC[:], Exc = calculate_potentials(
268 rgd=self.rgd, xc=self.xc, n_g=n, Z=Z, tw_coeff=tw_coeff)
270 # calculate new total Kohn-Sham effective potential and
271 # admix with old version
273 if niter > 0:
274 vr[:] = mix * vr + (1 - mix) * vrold
275 vrold = vr.copy()
277 # solve Kohn-Sham equation and determine the density change
278 self.solve()
279 dn = self.calculate_density() - n
280 n += dn
282 # estimate error from the square of the density change integrated
283 q = log(np.sum((r * dn)**2))
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
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
302 niter += 1
303 if niter > nitermax:
304 raise RuntimeError('Did not converge!')
306 tau = self.calculate_kinetic_energy_density()
308 t()
309 t('Converged in %d iteration%s.' % (niter, 's'[:niter != 1]))
311 Ekin = 0
313 for f, e in zip(f_j, e_j):
314 Ekin += f * e
316 e_coulomb = 2 * pi * np.dot(n * r * (vHr - Z), dr)
317 Ekin += -4 * pi * np.dot(n * vr * r, dr)
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()
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()
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:])
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
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)')
363 for m, l, u in zip(n_j, l_j, self.u_j):
364 self.write(u, 'ae', n=m, l=l)
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')
372 self.Ekin = Ekin
373 self.e_coulomb = e_coulomb
374 self.Exc = Exc
376 def write(self, array, name=None, n=None, l=None):
377 if self.nofiles:
378 return
380 if name:
381 name = self.symbol + '.' + name
382 else:
383 name = self.symbol
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])
393 f = open(name, 'w')
394 for r, a in zip(self.r, array):
395 print(r, a, file=f)
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)
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)
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]
422 return 0.5 * tau / (4 * pi)
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 """
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
457 return 0.5 * tau / (4 * pi)
459 def solve(self):
460 """Solve the Schrodinger equation
462 ::
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
471 where the relativistic mass::
473 1
474 M = 1 - --- (v - e)
475 2
476 2c
478 and the fine-structure constant alpha = 1/c = 1/137.036
479 is set to zero for non-scalar-relativistic calculations.
481 On the logaritmic radial grids defined by::
483 beta g
484 r = ------, g = 0, 1, ..., N - 1
485 N - g
487 rN
488 g = --------, r = [0; oo[
489 beta + r
491 the Schrodinger equation becomes::
493 2
494 d u du
495 --- c + -- c + u c = 0
496 2 2 dg 1 0
497 dg
499 with the vectors c0, c1, and c2 defined by::
501 2 dg 2
502 c = -r (--)
503 2 dr
505 2 2
506 d g 2 r dg dv
507 c = - --- r - ---- -- --
508 1 2 2 dr dr
509 dr 2Mc
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
520 c2 = -(r / dr)**2
521 c10 = -self.d2gdr2 * r**2 # first part of c1 vector
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
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
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))
561 def kin(self, l, u, e=None): # XXX move to Generator
562 r = self.r[1:]
563 dr = self.dr[1:]
565 c0 = 0.5 * l * (l + 1) / r**2
566 c1 = -0.5 * self.d2gdr2[1:]
567 c2 = -0.5 * dr**-2
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)
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
585 def __del__(self):
586 self.close()
589def shoot(u, l, vr, e, r2dvdr, r, dr, c10, c2, scalarrel, gmax=None):
590 """n, A = shoot(u, l, vr, e, ...)
592 For guessed trial eigenenergy e, integrate the radial Schrodinger
593 equation::
595 2
596 d u du
597 --- c + -- c + u c = 0
598 2 2 dg 1 0
599 dg
601 2 dg 2
602 c = -r (--)
603 2 dr
605 2 2
606 d g 2 r dg dv
607 c = - --- r - ---- -- --
608 1 2 2 dr dr
609 dr 2Mc
611 2 r dv
612 c = l(l + 1) + 2M(v - e)r + ---- --
613 0 2 dr
614 2Mc
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."""
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!')
632 c1 = c10
633 if scalarrel:
634 c0 += x * r2dvdr / Mr
635 c1 = c10 - x * r * r2dvdr / (Mr * dr)
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
642 if gmax is None:
643 # set boundary conditions at r -> oo (u(oo) = 0 is implicit)
644 u[-1] = 1.0
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
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
670 # set boundary conditions at r -> 0
671 u[0] = 0.0
672 u[1] = 1.0
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
687 # scale first part of wavefunction, such that it is continuous at gtp
688 u[:gtp + 2] *= utp / u[gtp]
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
694 return nodes, A
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
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)
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
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.
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]
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
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
756 # set boundary conditions at r -> 0
757 u[0] = 0.0
758 u[1] = 1.0
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
773 # scale first part of wavefunction, such that it is continuous at gtp
774 u[:gtp + 2] *= utp / u[gtp]
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
780 return nodes, A
783@dataclass
784class ValenceData:
785 symbol: str
786 xcname: str
788 rgd: AERadialGridDescriptor
790 n_j: list[int]
791 l_j: list[int]
792 e_j: list[float]
793 f_j: list[float]
795 scalarrel: bool
797 phi_jg: list[np.ndarray]
798 phit_jg: list[np.ndarray]
799 pt_jg: list[np.ndarray]
800 rcut_j: list[float]
802 vr_g: np.ndarray
803 r2dvdr_g: np.ndarray | None # Actually: None means not scalarrel
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()
809 for attr in jattributes:
810 thing_j = getattr(self, attr)
811 assert len(thing_j) == self.nj, (attr, len(thing_j), self.nj)
813 err = abs(self.rgd.beta / len(self.rgd.r_g) - self.rgd.a)
814 assert err < 1e-15, f'Inconsistent rgd spacing, {err=}'
816 @property
817 def nj(self):
818 return len(self.l_j)
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)
826 # XXX GLLB needs special initialization I think.
827 assert 'GLLB' not in setupdata.setupname
829 if setupdata.orbital_free:
830 raise RuntimeError('Setup is orbital-free')
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
839 vr_g = calculate_potentials(rgd, xc, n_g, setupdata.Z,
840 tw_coeff=None)[0]
841 r2dvdr_g = get_r2dvdr(rgd, vr_g)
843 assert setupdata.type in {'scalar-relativistic', 'non-relativistic'}
844 scalarrel = setupdata.type == 'scalar-relativistic'
845 return vr_g, r2dvdr_g, scalarrel
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)
856 @classmethod
857 def from_setupdata_and_potentials(cls, setupdata, *, vr_g, r2dvdr_g,
858 scalarrel):
860 assert len(setupdata.phi_jg) == len(setupdata.l_j)
862 def multiply_r(array_jg):
863 return array_jg * setupdata.rgd.r_g[None, :]
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 )
882 @property
883 def N(self):
884 return len(self.rgd.r_g)
886 @cached_property
887 def d2gdr2_g(self):
888 return self.rgd.d2gdr2()
890 def solve_confined(self, j, rc, vconf=None):
891 """Solve the Schroedinger equation in a confinement potential.
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.
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)
902 Returns: a tuple containing the solution u and its energy e.
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
913 c2 = -(r / dr)**2
914 c10 = -self.d2gdr2_g * r**2 # first part of c1 vector
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()
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
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
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
945 def get_confinement_potential(self, alpha, ri, rc):
946 r"""Create a smooth confinement potential.
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::
952 alpha / rc - ri \
953 V(r) = -------- exp ( - --------- ) for ri < r < rc
954 rc - r \ r - ri /
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
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
972 return alpha * potential
975if __name__ == '__main__':
976 a = AllElectron('Cu', scalarrel=True)
977 a.run()