# Source code for gpaw.atom.basis

```"""This module is used to generate atomic orbital basis sets."""
import sys
from io import StringIO

import numpy as np
from ase.units import Hartree
from gpaw.utilities import devnull

from gpaw import __version__ as version
from gpaw.utilities import divrl
from gpaw.atom.generator import Generator
from gpaw.atom.all_electron import AllElectron
from gpaw.atom.configurations import parameters
from gpaw.basis_data import Basis, BasisFunction, get_basis_name

def make_split_valence_basis_function(r_g, psi_g, l, gcut):
"""Get polynomial which joins psi smoothly at rcut.

Returns an array of function values f(r) * r, where::

l           2
f(r) = r  * (a - b r ),  r < rcut
f(r) = psi(r),           r >= rcut

where a and b are determined such that f(r) is continuous and
differentiable at rcut.  The parameter psi should be an atomic
orbital.
"""
r1 = r_g[gcut]  # ensure that rcut is moved to a grid point
r2 = r_g[gcut + 1]
y1 = psi_g[gcut] / r_g[gcut]
y2 = psi_g[gcut + 1] / r_g[gcut + 1]
b = - (y2 / r2**l - y1 / r1**l) / (r2**2 - r1**2)
a = (y1 / r1**l + b * r1**2)
psi_g2 = r_g**(l + 1) * (a - b * r_g**2)
psi_g2[gcut:] = psi_g[gcut:]
return psi_g2

def rsplit_by_norm(rgd, l, u, tailnorm_squared, txt):
"""Find radius outside which remaining tail has a particular norm."""
norm_squared = np.dot(rgd.dr_g, u * u)
partial_norm_squared = 0.
i = len(u) - 1
absolute_tailnorm_squared = tailnorm_squared * norm_squared
while partial_norm_squared < absolute_tailnorm_squared:
# Integrate backwards.  This is important since the pseudo
# wave functions have strange behaviour near the core.
partial_norm_squared += rgd.dr_g[i] * u[i]**2
i -= 1
rsplit = rgd.r_g[i + 1]
msg = ('Tail norm %.03f :: rsplit=%.02f Bohr' %
((partial_norm_squared / norm_squared)**0.5, rsplit))
print(msg, file=txt)
gsplit = rgd.floor(rsplit)
splitwave = make_split_valence_basis_function(rgd.r_g, u, l, gsplit)
return rsplit, partial_norm_squared, splitwave

[docs]class BasisMaker:
"""Class for creating atomic basis functions."""
def __init__(self, generator, name=None, run=True, gtxt='-',
non_relativistic_guess=False, xc='PBE',
save_setup=False):

if isinstance(generator, str):  # treat 'generator' as symbol
generator = Generator(generator, scalarrel=True,
xcname=xc, txt=gtxt,
nofiles=True)
generator.N *= 4
self.generator = generator
1.0 / generator.N, generator.N,
default_spline_points=100)
self.name = name

if run:
if non_relativistic_guess:
ae0 = AllElectron(generator.symbol, scalarrel=False,
nofiles=False, txt=gtxt, xcname=xc)
ae0.N = generator.N
ae0.beta = generator.beta
ae0.run()
# Now files will be stored such that they can
# automagically be used by the next run()
setup = generator.run(write_xml=False,
name=name,
**parameters[generator.symbol])

if save_setup:
setup.write_xml()
else:
if save_setup:
raise ValueError('cannot save setup here because setup '
'was already generated before basis '
'generation.')

def smoothify(self, psi_mg, l):
r"""Generate pseudo wave functions from all-electron ones.

The pseudo wave function is::

___
~                   \    /   ~             \    ~     ~
| psi  > = | psi  > +   )  | | phi > - | phi > ) < p  | psi > ,
m          m      /__  \     i         i /     i      m
i

where the scalar products are found by solving::

___
~             \     ~             ~     ~
< p | psi  > =   )  < p  | phi  > < p  | psi  > .
i     m      /__    i      j      j      m
j

In order to ensure smoothness close to the core, the
all-electron wave function and partial wave are then
multiplied by a radial function which approaches 0 near the
core, such that the pseudo wave function approaches::

___
~        \      ~       ~     ~
| psi  > =   )  | phi >  < p  | psi >    (for r << rcut),
m      /__      i      i      m
i

which is exact if the projectors/pseudo partial waves are complete.
"""
if psi_mg.ndim == 1:
return self.smoothify(psi_mg[None], l)[0]

g = self.generator
u_ng = g.u_ln[l]
q_ng = g.q_ln[l]
s_ng = g.s_ln[l]

Pi_nn = np.dot(g.r * q_ng, u_ng.T)
Q_nm = np.dot(g.r * q_ng, psi_mg.T)
Qt_nm = np.linalg.solve(Pi_nn, Q_nm)

# Weight-function for truncating all-electron parts smoothly near core
gmerge = g.r2g(g.rcut_l[l])
w_g = np.ones(g.r.shape)
w_g[0:gmerge] = (g.r[0:gmerge] / g.r[gmerge])**2.
w_g = w_g[None]

psit_mg = psi_mg * w_g + np.dot(Qt_nm.T, s_ng - u_ng * w_g)
return psit_mg

def rcut_by_energy(self, j, esplit=.1, tolerance=.1, rguess=6.,
vconf_args=None):
"""Find confinement cutoff corresponding to given orbital energy shift.

Creates a confinement potential for the orbital given by j,
such that the confined-orbital energy is (emin to emax) eV larger
than the free-orbital energy."""
g = self.generator
e_base = g.e_j[j]
rc = rguess

if vconf_args is None:
vconf = None
else:
amplitude, ri_rel = vconf_args
vconf = g.get_confinement_potential(amplitude, ri_rel * rc, rc)

psi_g, e = g.solve_confined(j, rc, vconf)
de_min, de_max = esplit / Hartree, (esplit + tolerance) / Hartree

rmin = 0.
rmax = g.r[-1]

de = e - e_base
while de < de_min or de > de_max:
if de < de_min:  # Move rc left -> smaller cutoff, higher energy
rmax = rc
rc = (rc + rmin) / 2.
else:  # Move rc right
rmin = rc
rc = (rc + rmax) / 2.
if vconf is not None:
vconf = g.get_confinement_potential(amplitude, ri_rel * rc, rc)
psi_g, e = g.solve_confined(j, rc, vconf)
de = e - e_base
if g.r2g(rmax) - g.r2g(rmin) <= 1:  # adjacent points
break  # cannot meet tolerance due to grid resolution

return psi_g, e, de, vconf, rc

def generate(self, zetacount=2, polarizationcount=1,
tailnorm=(0.16, 0.3, 0.6), energysplit=0.1, tolerance=1.0e-3,
rcutpol_rel=1.0,
rcutmax=20.0,
rcharpol_rel=None,
vconf_args=(12.0, 0.6), txt='-',
include_energy_derivatives=False,
jvalues=None,
l_pol=None):
"""Generate an entire basis set.

This is a high-level method which will return a basis set
consisting of several different basis vector types.

Parameters:

===================== =================================================
``zetacount``         Number of basis functions per occupied orbital
``polarizationcount`` Number of polarization functions
``tailnorm``          List of tail norms for split-valence scheme
``energysplit``       Energy increase defining confinement radius (eV)
``tolerance``         Tolerance of energy split (eV)
``rcutpol_rel``       Polarization rcut relative to largest other rcut
``rcutmax``           No cutoff will be greater than this value
``vconf_args``        Parameters (alpha, ri/rc) for conf. potential
``txt``               Log filename or '-' for stdout
===================== =================================================

Returns a fully initialized Basis object.
"""

if txt == '-':
txt = sys.stdout
elif txt is None:
txt = devnull

if isinstance(tailnorm, float):
tailnorm = (tailnorm,)
if 1 + len(tailnorm) < max(polarizationcount, zetacount):
raise ValueError(
'Needs %d tail norm values, but only %d are specified' %
(max(polarizationcount, zetacount) - 1, len(tailnorm)))

textbuffer = StringIO()

class TeeStream:  # quick hack to both write and save output
def __init__(self, out1, out2):
self.out1 = out1
self.out2 = out2

def write(self, string):
self.out1.write(string)
self.out2.write(string)

txt = TeeStream(txt, textbuffer)

if vconf_args is not None:
amplitude, ri_rel = vconf_args

g = self.generator
rgd = self.rgd

njcore = g.njcore
n_j = g.n_j[njcore:]
l_j = g.l_j[njcore:]
f_j = g.f_j[njcore:]

if jvalues is None:
jvalues = []
sortkeys = []
for j in range(len(n_j)):
if f_j[j] == 0 and l_j[j] != 0:
continue
jvalues.append(j)
sortkeys.append(l_j[j])

# Now order jvalues by l
#
# Use a stable sort so the energy ordering within each
# angular momentum is guaranteed to be preserved
args = np.argsort(sortkeys, kind='mergesort')
jvalues = np.array(jvalues)[args]

fulljvalues = [njcore + j for j in jvalues]

if isinstance(energysplit, float):
energysplit = [energysplit] * len(jvalues)

title = f'{g.xcname} Basis functions for {g.symbol}'
print(title, file=txt)
print('=' * len(title), file=txt)

singlezetas = []
energy_derivative_functions = []
multizetas = [[] for i in range(zetacount - 1)]
polarization_functions = []

splitvalencedescr = 'split-valence wave, fixed tail norm'
derivativedescr = 'derivative of sz wrt. (ri/rc) of potential'

for vj, fullj, esplit in zip(jvalues, fulljvalues, energysplit):
l = l_j[vj]
n = n_j[vj]
assert n > 0
orbitaltype = str(n) + 'spdf'[l]
msg = 'Basis functions for l=%d, n=%d' % (l, n)
print(file=txt)
print(msg + '\n', '-' * len(msg), file=txt)
print(file=txt)
if vconf_args is None:
else:
print('Zeta 1: %s confined pseudo wave,' % adverb, end=' ',
file=txt)

u, e, de, vconf, rc = self.rcut_by_energy(fullj, esplit,
tolerance,
vconf_args=vconf_args)
if rc > rcutmax:
rc = rcutmax  # scale things down
if vconf is not None:
vconf = g.get_confinement_potential(amplitude, ri_rel * rc,
rc)
u, e = g.solve_confined(fullj, rc, vconf)
print('using maximum cutoff', file=txt)
print('rc=%.02f Bohr' % rc, file=txt)
else:
print('fixed energy shift', file=txt)
print('DE=%.03f eV :: rc=%.02f Bohr'
% (de * Hartree, rc), file=txt)

if vconf is not None:
print('Potential amp=%.02f :: ri/rc=%.02f' %
(amplitude, ri_rel), file=txt)
phit_g = self.smoothify(u, l)
bf = BasisFunction(n, l, rc, phit_g,
'%s-sz confined orbital' % orbitaltype)
norm = np.dot(g.dr, phit_g * phit_g)**.5
print('Norm=%.03f' % norm, file=txt)
singlezetas.append(bf)

zetacounter = iter(range(2, zetacount + 1))

if include_energy_derivatives:
assert zetacount > 1
zeta = next(zetacounter)
print('\nZeta %d: %s' % (zeta, derivativedescr), file=txt)
vconf2 = g.get_confinement_potential(amplitude,
ri_rel * rc * .99, rc)
u2, e2 = g.solve_confined(fullj, rc, vconf2)

phit2_g = self.smoothify(u2, l)
dphit_g = phit2_g - phit_g
dphit_norm = np.dot(rgd.dr_g, dphit_g * dphit_g) ** .5
dphit_g /= dphit_norm
descr = '%s-dz E-derivative of sz' % orbitaltype
bf = BasisFunction(None, l, rc, dphit_g, descr)
energy_derivative_functions.append(bf)

for i, zeta in enumerate(zetacounter):
print('\nZeta %d: %s' % (zeta, splitvalencedescr), file=txt)
# Unresolved issue:  how does the lack of normalization
# of the first function impact the tail norm scheme?
# Presumably not much, since most interesting stuff happens
# close to the core.
rsplit, norm, splitwave = rsplit_by_norm(rgd, l, phit_g,
tailnorm[i]**2.0,
txt)
zetastring = '0sdtq56789'[zeta]
descr = f'{orbitaltype}-{zetastring}z split-valence wave'
bf = BasisFunction(None, l, rsplit, phit_g - splitwave, descr)
multizetas[i].append(bf)

if polarizationcount > 0 or l_pol is not None:
if l_pol is None:
# Now make up some properties for the polarization orbital
# We just use the cutoffs from the previous one times a factor
# Find 'missing' values in lvalues
lvalues = [l_j[vj] for vj in jvalues]
for i in range(max(lvalues) + 1):
if list(lvalues).count(i) == 0:
l_pol = i
break
else:
l_pol = max(lvalues) + 1

# Find the last state with l=l_pol - 1, which will be the state we
# base the polarization function on
for vj, fullj, bf in zip(jvalues[::-1], fulljvalues[::-1],
singlezetas[::-1]):
if bf.l == l_pol - 1:
fullj_pol = fullj
rcut = bf.rc * rcutpol_rel
break
else:
raise ValueError('The requested value l_pol=%d requires l=%d '
'among valence states' % (l_pol, l_pol - 1))
rcut = min(rcut, rcutmax)
msg = 'Polarization function: l=%d, rc=%.02f' % (l_pol, rcut)
print('\n' + msg, file=txt)
print('-' * len(msg), file=txt)
# Make a single Gaussian for polarization function.
#
# It is known that for given l, the sz cutoff defined
# by some fixed energy is strongly correlated to the
# value of the characteristic radius which best reproduces
# the wave function found by interpolation.
#
# We know that for e.g. d orbitals:
#   rchar ~= .37 rcut[sz](.3eV)
# Since we don't want to spend a lot of time finding
# these value for other energies, we just find the energy
# shift at .3 eV now

u, e, de, vconf, rc_fixed = self.rcut_by_energy(fullj_pol,
.3, 1e-2,
6., (12., .6))

default_rchar_rel = .25
# Defaults for each l.  Actually we don't care right now
rchar_rels = {}

if rcharpol_rel is None:
rcharpol_rel = rchar_rels.get(l_pol, default_rchar_rel)
rchar = rcharpol_rel * rc_fixed
gaussian = QuasiGaussian(1.0 / rchar**2, rcut)
psi_pol = gaussian(rgd.r_g) * rgd.r_g**(l_pol + 1)
norm = np.dot(rgd.dr_g, psi_pol * psi_pol) ** .5
psi_pol /= norm
print('Single quasi Gaussian', file=txt)
msg = f'Rchar = {rcharpol_rel:.3f}*rcut = {rchar:.3f} Bohr'
print(msg, file=txt)
typestring = 'spdfg'[l_pol]
bf_pol = BasisFunction(None, l_pol, rcut, psi_pol, type)

polarization_functions.append(bf_pol)
for i in range(polarizationcount - 1):
npol = i + 2
levelstring = ['Secondary', 'Tertiary', 'Quaternary',
'Quintary', 'Sextary', 'Septenary'][i]
msg = f'\n{levelstring}: {splitvalencedescr}'
print(msg, file=txt)
rsplit, norm, splitwave = rsplit_by_norm(rgd, l_pol, psi_pol,
tailnorm[i], txt)
descr = ('%s-type split-valence polarization %d'
% ('spdfg'[l_pol], npol))
bf_pol = BasisFunction(None, l_pol, rsplit,
psi_pol - splitwave,
descr)
polarization_functions.append(bf_pol)

bf_j = []
bf_j.extend(singlezetas)
bf_j.extend(energy_derivative_functions)
for multizeta_list in multizetas:
bf_j.extend(multizeta_list)
bf_j.extend(polarization_functions)

rcmax = max([bf.rc for bf in bf_j])

# The non-equidistant grids are really only suited for AE WFs
d = 1.0 / 64
equidistant_grid = np.arange(0.0, rcmax + d, d)
ng = len(equidistant_grid)

for bf in bf_j:
# We have been storing phit_g * r, but we just want phit_g
bf.phit_g = divrl(bf.phit_g, 1, rgd.r_g)

gcut = min(int(1 + bf.rc / d), ng - 1)

assert equidistant_grid[gcut] >= bf.rc
assert equidistant_grid[gcut - 1] <= bf.rc

bf.rc = equidistant_grid[gcut]
# Note: bf.rc *must* correspond to a grid point (spline issues)
bf.ng = gcut + 1
# XXX all this should be done while building the basis vectors,
# not here

# Quick hack to change to equidistant coordinates
spline = rgd.spline(bf.phit_g, rgd.r_g[rgd.floor(bf.rc)], bf.l,
points=100)
bf.phit_g = np.array([spline(r) * r**bf.l
for r in equidistant_grid[:bf.ng]])
bf.phit_g[-1] = 0.

basistype = get_basis_name(zetacount, polarizationcount)
if self.name is None:
compound_name = basistype
else:
compound_name = f'{self.name}.{basistype}'

basis = Basis(g.symbol, compound_name, False,
basis.bf_j = bf_j
basis.generatordata = textbuffer.getvalue().strip()
basis.generatorattrs = {'version': version}
textbuffer.close()

return basis

class QuasiGaussian:
"""Gaussian-like functions for expansion of orbitals.

Implements f(r) = A [G(r) - P(r)] where::

G(r) = exp{- alpha r^2}
P(r) = a - b r^2

with (a, b) such that f(rcut) == f'(rcut) == 0.
"""
def __init__(self, alpha, rcut, A=1.):
self.alpha = alpha
self.rcut = rcut
expmar2 = np.exp(-alpha * rcut**2)
a = (1 + alpha * rcut**2) * expmar2
b = alpha * expmar2
self.a = a
self.b = b
self.A = A

def __call__(self, r):
"""Evaluate function values at r, which is a numpy array."""
condition = (r < self.rcut) & (self.alpha * r**2 < 700.)
r2 = np.where(condition, r**2., 0.)  # prevent overflow
g = np.exp(-self.alpha * r2)
p = (self.a - self.b * r2)
y = np.where(condition, g - p, 0.)
return self.A * y
```