Coverage for gpaw/atom/basis.py: 94%
292 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1"""This module is used to generate atomic orbital basis sets."""
2import sys
3from io import StringIO
5import numpy as np
6from ase.units import Hartree
7from gpaw.utilities import devnull
9from gpaw import __version__ as version
10from gpaw.utilities import divrl
11from gpaw.atom.generator import Generator
12from gpaw.atom.all_electron import ValenceData
13from gpaw.atom.configurations import parameters
14from gpaw.basis_data import Basis, BasisFunction, get_basis_name
15from gpaw.atom.radialgd import (AERadialGridDescriptor,
16 EquidistantRadialGridDescriptor)
19def make_split_valence_basis_function(r_g, psi_g, l, gcut):
20 """Get polynomial which joins psi smoothly at rcut.
22 Returns an array of function values f(r) * r, where::
24 l 2
25 f(r) = r * (a - b r ), r < rcut
26 f(r) = psi(r), r >= rcut
28 where a and b are determined such that f(r) is continuous and
29 differentiable at rcut. The parameter psi should be an atomic
30 orbital.
31 """
32 r1 = r_g[gcut] # ensure that rcut is moved to a grid point
33 r2 = r_g[gcut + 1]
34 y1 = psi_g[gcut] / r_g[gcut]
35 y2 = psi_g[gcut + 1] / r_g[gcut + 1]
36 b = - (y2 / r2**l - y1 / r1**l) / (r2**2 - r1**2)
37 a = (y1 / r1**l + b * r1**2)
38 psi_g2 = r_g**(l + 1) * (a - b * r_g**2)
39 psi_g2[gcut:] = psi_g[gcut:]
40 return psi_g2
43def rsplit_by_norm(rgd, l, u, tailnorm_squared, txt):
44 """Find radius outside which remaining tail has a particular norm."""
45 norm_squared = np.dot(rgd.dr_g, u * u)
46 partial_norm_squared = 0.
47 i = len(u) - 1
48 absolute_tailnorm_squared = tailnorm_squared * norm_squared
49 while partial_norm_squared < absolute_tailnorm_squared:
50 # Integrate backwards. This is important since the pseudo
51 # wave functions have strange behaviour near the core.
52 partial_norm_squared += rgd.dr_g[i] * u[i]**2
53 i -= 1
54 rsplit = rgd.r_g[i + 1]
55 msg = ('Tail norm %.03f :: rsplit=%.02f Bohr' %
56 ((partial_norm_squared / norm_squared)**0.5, rsplit))
57 print(msg, file=txt)
58 gsplit = rgd.floor(rsplit)
59 splitwave = make_split_valence_basis_function(rgd.r_g, u, l, gsplit)
60 return rsplit, partial_norm_squared, splitwave
63class BasisMaker:
64 """Class for creating atomic basis functions."""
65 def __init__(self, valence_data, name=None, run=True, gtxt='-',
66 non_relativistic_guess=False, xc='PBE',
67 save_setup=False):
69 if not isinstance(valence_data, ValenceData):
70 raise TypeError('BasisMaker no longer accepts Generator or str. '
71 'Please provide ValenceData or use one of the '
72 'helper classmethods.')
74 self.name = name
75 self.valence_data = valence_data
76 rgd = valence_data.rgd
77 rgd = AERadialGridDescriptor(rgd.a,
78 rgd.b,
79 len(rgd.r_g),
80 default_spline_points=100)
81 self.rgd = rgd
83 @classmethod
84 def from_setup_and_generator(cls, setup, generator, **kwargs):
85 valdata = ValenceData.from_setupdata_and_potentials(
86 setup, vr_g=generator.vr, r2dvdr_g=generator.r2dvdr,
87 scalarrel=generator.scalarrel)
88 return cls(valdata, **kwargs)
90 @classmethod
91 def from_symbol(cls, symbol, gtxt='-', xc='PBE', name=None,
92 save_setup=False,
93 generator_run_kwargs=None, **kwargs):
94 generator = Generator(symbol, scalarrel=True,
95 xcname=xc, txt=gtxt,
96 nofiles=True,
97 gpernode=Generator.default_gpernode * 4)
99 run_kwargs = {
100 'write_xml': False,
101 'name': name,
102 **parameters[generator.symbol]}
104 if generator_run_kwargs is not None:
105 run_kwargs.update(generator_run_kwargs)
107 setup = generator.run(**run_kwargs)
109 if save_setup:
110 setup.write_xml()
112 return cls.from_setup_and_generator(setup, generator, name=name,
113 **kwargs)
115 def smoothify(self, psi_mg, j):
116 r"""Generate pseudo wave functions from all-electron ones.
118 The pseudo wave function is::
120 ___
121 ~ \ / ~ \ ~ ~
122 | psi > = | psi > + ) | | phi > - | phi > ) < p | psi > ,
123 m m /__ \ i i / i m
124 i
126 where the scalar products are found by solving::
128 ___
129 ~ \ ~ ~ ~
130 < p | psi > = ) < p | phi > < p | psi > .
131 i m /__ i j j m
132 j
134 In order to ensure smoothness close to the core, the
135 all-electron wave function and partial wave are then
136 multiplied by a radial function which approaches 0 near the
137 core, such that the pseudo wave function approaches::
139 ___
140 ~ \ ~ ~ ~
141 | psi > = ) | phi > < p | psi > (for r << rcut),
142 m /__ i i m
143 i
145 which is exact if the projectors/pseudo partial waves are complete.
146 """
148 if psi_mg.ndim == 1:
149 return self.smoothify(psi_mg[None], j)[0]
151 valdata = self.valence_data
152 l = valdata.l_j[j]
154 def angular_buddies(array_j):
155 return np.array([array_j[j] for j, l1 in enumerate(valdata.l_j)
156 if l1 == l])
158 u_ng = angular_buddies(valdata.phi_jg)
159 q_ng = angular_buddies(valdata.pt_jg)
160 s_ng = angular_buddies(valdata.phit_jg)
162 r_g = valdata.rgd.r_g
164 Pi_nn = np.dot(r_g * q_ng, u_ng.T)
165 Q_nm = np.dot(r_g * q_ng, psi_mg.T)
166 Qt_nm = np.linalg.solve(Pi_nn, Q_nm)
168 # Weight-function for truncating all-electron parts smoothly near core
169 gmerge = valdata.rgd.floor(valdata.rcut_j[j])
170 w_g = np.ones_like(r_g)
171 w_g[0:gmerge] = (r_g[0:gmerge] / r_g[gmerge])**2.
172 w_g = w_g[None]
174 psit_mg = psi_mg * w_g + np.dot(Qt_nm.T, s_ng - u_ng * w_g)
175 return psit_mg
177 def rcut_by_energy(self, j, esplit=.1, tolerance=.1, rguess=6.,
178 vconf_args=None):
179 """Find confinement cutoff corresponding to given orbital energy shift.
181 Creates a confinement potential for the orbital given by j,
182 such that the confined-orbital energy is (emin to emax) eV larger
183 than the free-orbital energy."""
184 valdata = self.valence_data
185 e_base = valdata.e_j[j]
186 rc = rguess
188 if vconf_args is None:
189 vconf = None
190 else:
191 amplitude, ri_rel = vconf_args
192 vconf = valdata.get_confinement_potential(
193 amplitude, ri_rel * rc, rc)
195 psi_g, e = valdata.solve_confined(j, rc, vconf)
196 de_min, de_max = esplit / Hartree, (esplit + tolerance) / Hartree
198 rmin = 0.
199 rmax = valdata.rgd.r_g[-1]
201 de = e - e_base
202 while de < de_min or de > de_max:
203 if de < de_min: # Move rc left -> smaller cutoff, higher energy
204 rmax = rc
205 rc = (rc + rmin) / 2.
206 else: # Move rc right
207 rmin = rc
208 rc = (rc + rmax) / 2.
209 if vconf is not None:
210 vconf = valdata.get_confinement_potential(
211 amplitude, ri_rel * rc, rc)
212 psi_g, e = valdata.solve_confined(j, rc, vconf)
213 de = e - e_base
214 if valdata.rgd.floor(rmax) - valdata.rgd.floor(rmin) <= 1:
215 # Adjacent points, cannot meet tolerance due to grid resolution
216 break
218 return psi_g, e, de, vconf, rc
220 def generate(self, zetacount=2, polarizationcount=1,
221 tailnorm=(0.16, 0.3, 0.6), energysplit=0.1, tolerance=1.0e-3,
222 rcutpol_rel=1.0,
223 rcutmax=20.0,
224 rcharpol_rel=None,
225 vconf_args=(12.0, 0.6), txt='-',
226 include_energy_derivatives=False,
227 jvalues=None,
228 l_pol=None):
229 """Generate an entire basis set.
231 This is a high-level method which will return a basis set
232 consisting of several different basis vector types.
234 Parameters:
236 ===================== =================================================
237 ``zetacount`` Number of basis functions per occupied orbital
238 ``polarizationcount`` Number of polarization functions
239 ``tailnorm`` List of tail norms for split-valence scheme
240 ``energysplit`` Energy increase defining confinement radius (eV)
241 ``tolerance`` Tolerance of energy split (eV)
242 ``rcutpol_rel`` Polarization rcut relative to largest other rcut
243 ``rcutmax`` No cutoff will be greater than this value
244 ``vconf_args`` Parameters (alpha, ri/rc) for conf. potential
245 ``txt`` Log filename or '-' for stdout
246 ===================== =================================================
248 Returns a fully initialized Basis object.
249 """
251 if txt == '-':
252 txt = sys.stdout
253 elif txt is None:
254 txt = devnull
256 if isinstance(tailnorm, float):
257 tailnorm = (tailnorm,)
258 if 1 + len(tailnorm) < max(polarizationcount, zetacount):
259 raise ValueError(
260 'Needs %d tail norm values, but only %d are specified' %
261 (max(polarizationcount, zetacount) - 1, len(tailnorm)))
263 textbuffer = StringIO()
265 class TeeStream: # quick hack to both write and save output
266 def __init__(self, out1, out2):
267 self.out1 = out1
268 self.out2 = out2
270 def write(self, string):
271 self.out1.write(string)
272 self.out2.write(string)
274 txt = TeeStream(txt, textbuffer)
276 if vconf_args is not None:
277 amplitude, ri_rel = vconf_args
279 valdata = self.valence_data
280 rgd = self.rgd
282 n_j = valdata.n_j
283 l_j = valdata.l_j
284 f_j = valdata.f_j
286 if jvalues is None:
287 jvalues = []
288 sortkeys = []
289 for j in range(len(n_j)):
290 if f_j[j] == 0 and l_j[j] != 0:
291 continue
292 jvalues.append(j)
293 sortkeys.append(l_j[j])
295 # Now order jvalues by l
296 #
297 # Use a stable sort so the energy ordering within each
298 # angular momentum is guaranteed to be preserved
299 args = np.argsort(sortkeys, kind='mergesort')
300 jvalues = np.array(jvalues)[args]
302 if isinstance(energysplit, float):
303 energysplit = [energysplit] * len(jvalues)
305 title = f'{valdata.xcname} Basis functions for {valdata.symbol}'
306 print(title, file=txt)
307 print('=' * len(title), file=txt)
309 singlezetas = []
310 energy_derivative_functions = []
311 multizetas = [[] for i in range(zetacount - 1)]
312 polarization_functions = []
314 splitvalencedescr = 'split-valence wave, fixed tail norm'
315 derivativedescr = 'derivative of sz wrt. (ri/rc) of potential'
317 jvalues = [j for j in jvalues if n_j[j] > 0]
319 for vj, esplit in zip(jvalues, energysplit):
320 l = l_j[vj]
321 n = n_j[vj]
322 assert n > 0
323 orbitaltype = str(n) + 'spdf'[l]
324 msg = 'Basis functions for l=%d, n=%d' % (l, n)
325 print(file=txt)
326 print(msg + '\n', '-' * len(msg), file=txt)
327 print(file=txt)
328 if vconf_args is None:
329 adverb = 'sharply'
330 else:
331 adverb = 'softly'
332 print('Zeta 1: %s confined pseudo wave,' % adverb, end=' ',
333 file=txt)
335 u, e, de, vconf, rc = self.rcut_by_energy(vj, esplit,
336 tolerance,
337 vconf_args=vconf_args)
338 if rc > rcutmax:
339 rc = rcutmax # scale things down
340 if vconf is not None:
341 vconf = valdata.get_confinement_potential(
342 amplitude, ri_rel * rc, rc)
343 u, e = valdata.solve_confined(vj, rc, vconf)
344 print('using maximum cutoff', file=txt)
345 print('rc=%.02f Bohr' % rc, file=txt)
346 else:
347 print('fixed energy shift', file=txt)
348 print('DE=%.03f eV :: rc=%.02f Bohr'
349 % (de * Hartree, rc), file=txt)
351 if vconf is not None:
352 print('Potential amp=%.02f :: ri/rc=%.02f' %
353 (amplitude, ri_rel), file=txt)
354 phit_g = self.smoothify(u, vj)
355 bf = BasisFunction(n, l, rc, phit_g,
356 '%s-sz confined orbital' % orbitaltype)
357 norm = np.dot(valdata.rgd.dr_g, phit_g * phit_g)**.5
358 print('Norm=%.03f' % norm, file=txt)
359 singlezetas.append(bf)
361 zetacounter = iter(range(2, zetacount + 1))
363 if include_energy_derivatives:
364 assert zetacount > 1
365 zeta = next(zetacounter)
366 print('\nZeta %d: %s' % (zeta, derivativedescr), file=txt)
367 vconf2 = valdata.get_confinement_potential(
368 amplitude, ri_rel * rc * .99, rc)
369 u2, e2 = valdata.solve_confined(vj, rc, vconf2)
371 phit2_g = self.smoothify(u2, vj)
372 dphit_g = phit2_g - phit_g
373 dphit_norm = np.dot(rgd.dr_g, dphit_g * dphit_g) ** .5
374 dphit_g /= dphit_norm
375 descr = '%s-dz E-derivative of sz' % orbitaltype
376 bf = BasisFunction(None, l, rc, dphit_g, descr)
377 energy_derivative_functions.append(bf)
379 for i, zeta in enumerate(zetacounter):
380 print('\nZeta %d: %s' % (zeta, splitvalencedescr), file=txt)
381 # Unresolved issue: how does the lack of normalization
382 # of the first function impact the tail norm scheme?
383 # Presumably not much, since most interesting stuff happens
384 # close to the core.
385 rsplit, norm, splitwave = rsplit_by_norm(rgd, l, phit_g,
386 tailnorm[i]**2.0,
387 txt)
388 zetastring = '0sdtq56789'[zeta]
389 descr = f'{orbitaltype}-{zetastring}z split-valence wave'
390 bf = BasisFunction(None, l, rsplit, phit_g - splitwave, descr)
391 multizetas[i].append(bf)
393 if polarizationcount > 0 or l_pol is not None:
394 if l_pol is None:
395 # Now make up some properties for the polarization orbital
396 # We just use the cutoffs from the previous one times a factor
397 # Find 'missing' values in lvalues
398 lvalues = [l_j[vj] for vj in jvalues]
399 for i in range(max(lvalues) + 1):
400 if list(lvalues).count(i) == 0:
401 l_pol = i
402 break
403 else:
404 l_pol = max(lvalues) + 1
406 # Find the last state with l=l_pol - 1, which will be the state we
407 # base the polarization function on
408 for vj, bf in zip(jvalues[::-1], singlezetas[::-1]):
409 if bf.l == l_pol - 1:
410 j_pol = vj
411 rcut = bf.rc * rcutpol_rel
412 break
413 else:
414 raise ValueError('The requested value l_pol=%d requires l=%d '
415 'among valence states' % (l_pol, l_pol - 1))
416 rcut = min(rcut, rcutmax)
417 msg = 'Polarization function: l=%d, rc=%.02f' % (l_pol, rcut)
418 print('\n' + msg, file=txt)
419 print('-' * len(msg), file=txt)
420 # Make a single Gaussian for polarization function.
421 #
422 # It is known that for given l, the sz cutoff defined
423 # by some fixed energy is strongly correlated to the
424 # value of the characteristic radius which best reproduces
425 # the wave function found by interpolation.
426 #
427 # We know that for e.g. d orbitals:
428 # rchar ~= .37 rcut[sz](.3eV)
429 # Since we don't want to spend a lot of time finding
430 # these value for other energies, we just find the energy
431 # shift at .3 eV now
433 u, e, de, vconf, rc_fixed = self.rcut_by_energy(j_pol,
434 .3, 1e-2,
435 6., (12., .6))
437 default_rchar_rel = .25
438 # Defaults for each l. Actually we don't care right now
439 rchar_rels = {}
441 if rcharpol_rel is None:
442 rcharpol_rel = rchar_rels.get(l_pol, default_rchar_rel)
443 rchar = rcharpol_rel * rc_fixed
444 gaussian = QuasiGaussian(1.0 / rchar**2, rcut)
445 psi_pol = gaussian(rgd.r_g) * rgd.r_g**(l_pol + 1)
446 norm = np.dot(rgd.dr_g, psi_pol * psi_pol) ** .5
447 psi_pol /= norm
448 print('Single quasi Gaussian', file=txt)
449 msg = f'Rchar = {rcharpol_rel:.3f}*rcut = {rchar:.3f} Bohr'
450 adjective = 'Gaussian'
451 print(msg, file=txt)
452 typestring = 'spdfg'[l_pol]
453 type = f'{typestring}-type {adjective} polarization'
454 bf_pol = BasisFunction(None, l_pol, rcut, psi_pol, type)
456 polarization_functions.append(bf_pol)
457 for i in range(polarizationcount - 1):
458 npol = i + 2
459 levelstring = ['Secondary', 'Tertiary', 'Quaternary',
460 'Quintary', 'Sextary', 'Septenary'][i]
461 msg = f'\n{levelstring}: {splitvalencedescr}'
462 print(msg, file=txt)
463 rsplit, norm, splitwave = rsplit_by_norm(rgd, l_pol, psi_pol,
464 tailnorm[i], txt)
465 descr = ('%s-type split-valence polarization %d'
466 % ('spdfg'[l_pol], npol))
467 bf_pol = BasisFunction(None, l_pol, rsplit,
468 psi_pol - splitwave,
469 descr)
470 polarization_functions.append(bf_pol)
472 bf_j = []
473 bf_j.extend(singlezetas)
474 bf_j.extend(energy_derivative_functions)
475 for multizeta_list in multizetas:
476 bf_j.extend(multizeta_list)
477 bf_j.extend(polarization_functions)
479 rcmax = max([bf.rc for bf in bf_j])
481 # The non-equidistant grids are really only suited for AE WFs
482 d = 1.0 / 64
483 equidistant_grid = np.arange(0.0, rcmax + d, d)
484 ng = len(equidistant_grid)
486 for bf in bf_j:
487 # We have been storing phit_g * r, but we just want phit_g
488 bf.phit_g = divrl(bf.phit_g, 1, rgd.r_g)
490 gcut = min(int(1 + bf.rc / d), ng - 1)
492 assert equidistant_grid[gcut] >= bf.rc
493 assert equidistant_grid[gcut - 1] <= bf.rc
495 bf.rc = equidistant_grid[gcut]
496 # Note: bf.rc *must* correspond to a grid point (spline issues)
497 bf.ng = gcut + 1
498 # XXX all this should be done while building the basis vectors,
499 # not here
501 # Quick hack to change to equidistant coordinates
502 spline = rgd.spline(bf.phit_g, rgd.r_g[rgd.floor(bf.rc)], bf.l,
503 points=100)
504 bf.phit_g = np.array([spline(r) * r**bf.l
505 for r in equidistant_grid[:bf.ng]])
506 bf.phit_g[-1] = 0.
508 basistype = get_basis_name(zetacount, polarizationcount)
509 if self.name is None:
510 compound_name = basistype
511 else:
512 compound_name = f'{self.name}.{basistype}'
514 basis = Basis(
515 valdata.symbol, compound_name,
516 rgd=EquidistantRadialGridDescriptor(d, ng),
517 bf_j=bf_j,
518 generatordata=textbuffer.getvalue().strip(),
519 generatorattrs={'version': version})
520 textbuffer.close()
522 return basis
525class QuasiGaussian:
526 """Gaussian-like functions for expansion of orbitals.
528 Implements f(r) = A [G(r) - P(r)] where::
530 G(r) = exp{- alpha r^2}
531 P(r) = a - b r^2
533 with (a, b) such that f(rcut) == f'(rcut) == 0.
534 """
535 def __init__(self, alpha, rcut, A=1.):
536 self.alpha = alpha
537 self.rcut = rcut
538 expmar2 = np.exp(-alpha * rcut**2)
539 a = (1 + alpha * rcut**2) * expmar2
540 b = alpha * expmar2
541 self.a = a
542 self.b = b
543 self.A = A
545 def __call__(self, r):
546 """Evaluate function values at r, which is a numpy array."""
547 condition = (r < self.rcut) & (self.alpha * r**2 < 700.)
548 r2 = np.where(condition, r**2., 0.) # prevent overflow
549 g = np.exp(-self.alpha * r2)
550 p = (self.a - self.b * r2)
551 y = np.where(condition, g - p, 0.)
552 return self.A * y