Example #1
0
def get_radial_potential(calc, a, ai):
    """Calculates dV/dr / r for the effective potential.
    Below, f_g denotes dV/dr = minus the radial force"""

    rgd = a.xc_correction.rgd
    r_g = rgd.r_g
    r_g[0] = 1.0e-12
    dr_g = rgd.dr_g
    Ns = calc.wfs.nspins

    D_sp = calc.density.D_asp[ai]
    B_pq = a.xc_correction.B_pqL[:, :, 0]
    n_qg = a.xc_correction.n_qg
    D_sq = np.dot(D_sp, B_pq)
    n_sg = np.dot(D_sq, n_qg) / (4 * np.pi)**0.5
    n_sg[:] += a.xc_correction.nc_g / Ns

    # Coulomb force from nucleus
    fc_g = a.Z / r_g**2

    # Hartree force
    rho_g = 4 * np.pi * r_g**2 * dr_g * np.sum(n_sg, axis=0)
    fh_g = -np.array([np.sum(rho_g[:ig]) for ig in range(len(r_g))]) / r_g**2

    # xc force
    xc = XC(calc.get_xc_functional())
    v_sg = np.zeros_like(n_sg)
    xc.calculate_spherical(a.xc_correction.rgd, n_sg, v_sg)
    fxc_sg = np.array([a.xc_correction.rgd.derivative(v_sg[s])
                       for s in range(Ns)])
    fxc_g = np.sum(fxc_sg, axis=0) / Ns

    # f_sg = np.tile(fc_g, (Ns, 1)) + np.tile(fh_g, (Ns, 1)) + fxc_sg
    f_sg = np.tile(fc_g + fh_g + fxc_g, (Ns, 1))

    return f_sg[:] / r_g
Example #2
0
class AllElectronAtom:
    def __init__(self, symbol, xc='LDA', spinpol=False, dirac=False,
                 configuration=None,
                 log=None):
        """All-electron calculation for spherically symmetric atom.

        symbol: str (or int)
            Chemical symbol (or atomic number).
        xc: str
            Name of XC-functional.
        spinpol: bool
            If true, do spin-polarized calculation.  Default is spin-paired.
        dirac: bool
            Solve Dirac equation instead of Schrödinger equation.
        configuration: list
            Electronic configuration for symbol, format as in
            gpaw.atom.configurations
        log: stream
            Text output."""

        if isinstance(symbol, int):
            symbol = chemical_symbols[symbol]
        self.symbol = symbol
        self.Z = atomic_numbers[symbol]

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

        self.dirac = bool(dirac)

        if configuration is not None:
            self.configuration = copy.deepcopy(configuration)
        else:
            self.configuration = None

        self.scalar_relativistic = False

        if isinstance(xc, str):
            self.xc = XC(xc)
        else:
            self.xc = xc

        self.fd = log or sys.stdout

        self.vr_sg = None  # potential * r
        self.n_sg = 0.0  # density
        self.rgd = None  # radial grid descriptor

        # Energies:
        self.ekin = None
        self.eeig = None
        self.eH = None
        self.eZ = None

        self.channels = None

        self.initialize_configuration(self.configuration)

        self.log('Z:              ', self.Z)
        self.log('Name:           ', atomic_names[self.Z])
        self.log('Symbol:         ', symbol)
        self.log('XC-functional:  ', self.xc.name)
        self.log('Equation:       ', ['Schrödinger', 'Dirac'][self.dirac])

        self.method = 'Gaussian basis-set'

    def log(self, *args, **kwargs):
        print(file=self.fd, *args, **kwargs)

    def initialize_configuration(self, configuration=None):
        self.f_lsn = {}

        if configuration is None:
            configuration = configurations[self.symbol][1]

        for n, l, f, e in configuration:

            if l not in self.f_lsn:
                self.f_lsn[l] = [[] for s in range(self.nspins)]
            if self.nspins == 1:
                self.f_lsn[l][0].append(f)
            else:
                # Use Hund's rule:
                f0 = min(f, 2 * l + 1)
                self.f_lsn[l][0].append(f0)
                self.f_lsn[l][1].append(f - f0)

        if 0:
            n = 2 + len(self.f_lsn[2][0])
            if self.f_lsn[0][0][n] == 2:
                self.f_lsn[0][0][n] = 1
                self.f_lsn[2][0][n - 3] += 1

    def add(self, n, l, df=+1, s=None):
        """Add (remove) electrons."""
        if s is None:
            if self.nspins == 1:
                s = 0
            else:
                self.add(n, l, 0.5 * df, 0)
                self.add(n, l, 0.5 * df, 1)
                return

        if l not in self.f_lsn:
            self.f_lsn[l] = [[] for x in range(self.nspins)]

        f_n = self.f_lsn[l][s]
        if len(f_n) < n - l:
            f_n.extend([0] * (n - l - len(f_n)))
        f_n[n - l - 1] += df

    def initialize(self, ngpts=2000, rcut=50.0,
                   alpha1=0.01, alpha2=None, ngauss=50,
                   eps=1.0e-7):
        """Initialize basis sets and radial grid.

        ngpts: int
            Number of grid points for radial grid.
        rcut: float
            Cutoff for radial grid.
        alpha1: float
            Smallest exponent for gaussian.
        alpha2: float
            Largest exponent for gaussian.
        ngauss: int
            Number of gaussians.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        if alpha2 is None:
            alpha2 = 50.0 * self.Z**2

        # Use grid with r(0)=0, r(1)=a and r(ngpts)=rcut:
        a = 1 / alpha2**0.5 / 20
        b = (rcut - a * ngpts) / (rcut * ngpts)
        b = 1 / round(1 / b)
        self.rgd = AERadialGridDescriptor(a, b, ngpts)

        self.log('Grid points:     %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %
                 ((self.rgd.N,) + tuple(self.rgd.r_g[[0, 1, 2, -2, -1]])))

        # Distribute exponents between alpha1 and alpha2:
        alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)
        self.log('Exponents:       %d (%.3f, %.3f, ..., %.3f, %.3f)' %
                 ((ngauss,) + tuple(alpha_B[[0, 1, -2, -1]])))

        # Maximum l value:
        lmax = max(self.f_lsn.keys())

        self.channels = []
        nb_l = []
        if not self.dirac:
            for l in range(lmax + 1):
                basis = GaussianBasis(l, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for s in range(self.nspins):
                    self.channels.append(Channel(l, s, self.f_lsn[l][s],
                                                 basis))
        else:
            for K in range(1, lmax + 2):
                leff = (K**2 - (self.Z / c)**2)**0.5 - 1
                basis = GaussianBasis(leff, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for k, l in [(-K, K - 1), (K, K)]:
                    if l > lmax:
                        continue
                    f_n = self.f_lsn[l][0]
                    j = abs(k) - 0.5
                    f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)
                    self.channels.append(DiracChannel(k, f_n, basis))

        self.log('Basis functions: %s (%s)' %
                 (', '.join([str(nb) for nb in nb_l]),
                  ', '.join('spdf'[:lmax + 1])))

        self.vr_sg = self.rgd.zeros(self.nspins)
        self.vr_sg[:] = -self.Z

    def solve(self):
        """Diagonalize Schrödinger equation."""
        self.eeig = 0.0
        for channel in self.channels:
            if self.method == 'Gaussian basis-set':
                channel.solve(self.vr_sg[channel.s])
            else:
                channel.solve2(self.vr_sg[channel.s], self.scalar_relativistic,
                               self.Z)

            self.eeig += channel.get_eigenvalue_sum()

    def calculate_density(self):
        """Calculate elctron density and kinetic energy."""
        self.n_sg = self.rgd.zeros(self.nspins)
        for channel in self.channels:
            self.n_sg[channel.s] += channel.calculate_density()

    def calculate_electrostatic_potential(self):
        """Calculate electrostatic potential and energy."""
        n_g = self.n_sg.sum(0)
        self.vHr_g = self.rgd.poisson(n_g)
        self.eH = 0.5 * self.rgd.integrate(n_g * self.vHr_g, -1)
        self.eZ = -self.Z * self.rgd.integrate(n_g, -1)

    def calculate_xc_potential(self):
        self.vxc_sg = self.rgd.zeros(self.nspins)
        self.exc = self.xc.calculate_spherical(self.rgd, self.n_sg,
                                               self.vxc_sg)

    def step(self):
        self.solve()
        self.calculate_density()
        self.calculate_electrostatic_potential()
        self.calculate_xc_potential()
        self.vr_sg = self.vxc_sg * self.rgd.r_g
        self.vr_sg += self.vHr_g
        self.vr_sg -= self.Z
        self.ekin = (self.eeig -
                     self.rgd.integrate((self.vr_sg * self.n_sg).sum(0), -1))

    def run(self, mix=0.4, maxiter=117, dnmax=1e-9):
        if self.channels is None:
            self.initialize()

        if self.dirac:
            equation = 'Dirac'
        elif self.scalar_relativistic:
            equation = 'scalar-relativistic Schrödinger'
        else:
            equation = 'non-relativistic Schrödinger'
        self.log('\nSolving %s equation using %s:' % (equation, self.method))

        dn = self.Z

        vr_old_sg = None
        n_old_sg = None
        for iter in range(maxiter):
            self.log('.', end='')
            self.fd.flush()
            if iter > 0:
                self.vr_sg *= mix
                self.vr_sg += (1 - mix) * vr_old_sg
                dn = self.rgd.integrate(abs(self.n_sg - n_old_sg).sum(0))
                if dn <= dnmax:
                    self.log('\nConverged in', iter, 'steps')
                    break

            vr_old_sg = self.vr_sg
            n_old_sg = self.n_sg
            self.step()

        self.summary()

        if self.method != 'Gaussian basis-set':
            for channel in self.channels:
                assert channel.solve2ok

        if dn > dnmax:
            raise RuntimeError('Did not converge!')

    def refine(self):
        self.method = 'finite difference'
        self.run(dnmax=1e-6, mix=0.14, maxiter=200)

    def summary(self):
        self.write_states()
        self.write_energies()

    def write_states(self):
        self.log('\n state  occupation         eigenvalue          <r>')
        if self.dirac:
            self.log(' nl(j)               [Hartree]        [eV]    [Bohr]')
        else:
            self.log(' nl                  [Hartree]        [eV]    [Bohr]')
        self.log('-----------------------------------------------------')
        states = []
        for ch in self.channels:
            for n, f in enumerate(ch.f_n):
                states.append((ch.e_n[n], ch, n))
        states.sort()
        for e, ch, n in states:
            name = str(n + ch.l + 1) + ch.name
            if self.nspins == 2:
                name += '(%s)' % '+-'[ch.s]
            n_g = ch.calculate_density(n)
            rave = self.rgd.integrate(n_g, 1)
            self.log(' %-7s  %6.3f %13.6f  %13.5f %6.3f' %
                     (name, ch.f_n[n], e, e * units.Hartree, rave))

    def write_energies(self):
        self.log('\nEnergies:          [Hartree]           [eV]')
        self.log('--------------------------------------------')
        for text, e in [('kinetic      ', self.ekin),
                        ('coulomb (e-e)', self.eH),
                        ('coulomb (e-n)', self.eZ),
                        ('xc           ', self.exc),
                        ('total        ',
                         self.ekin + self.eH + self.eZ + self.exc)]:
            self.log(' %s %+13.6f  %+13.5f' % (text, e, units.Hartree * e))

        self.calculate_exx()
        self.log('\nExact exchange energy: %.6f Hartree, %.5f eV' %
                 (self.exx, self.exx * units.Hartree))

    def get_channel(self, l=None, s=0, k=None):
        if self.dirac:
            for channel in self.channels:
                if channel.k == k:
                    return channel
        else:
            for channel in self.channels:
                if channel.l == l and channel.s == s:
                    return channel
        raise ValueError

    def get_orbital(self, n, l=None, s=0, k=None):
        channel = self.get_channel(l, s, k)
        return channel.basis.expand(channel.C_nb[n])

    def plot_wave_functions(self, rc=4.0):
        import matplotlib.pyplot as plt
        for ch in self.channels:
            for n in range(len(ch.f_n)):
                fr_g = ch.basis.expand(ch.C_nb[n]) * self.rgd.r_g
                # fr_g = ch.phi_ng[n]
                name = str(n + ch.l + 1) + ch.name
                lw = 2
                if self.nspins == 2:
                    name += '(%s)' % '+-'[ch.s]
                    if ch.s == 1:
                        lw = 1
                if self.dirac and ch.k > 0:
                    lw = 1
                ls = ['-', '--', '-.', ':'][ch.l]
                n_g = ch.calculate_density(n)
                rave = self.rgd.integrate(n_g, 1)
                gave = self.rgd.round(rave)
                fr_g *= np.sign(fr_g[gave], 0)
                plt.plot(self.rgd.r_g, fr_g,
                         ls=ls, lw=lw, color=colors[n + ch.l], label=name)

        plt.legend(loc='best')
        plt.xlabel('r [Bohr]')
        plt.ylabel('$r\\phi(r)$')
        plt.axis(xmax=rc)
        plt.show()

    def logarithmic_derivative(self, l, energies, rcut):
        ch = Channel(l)
        gcut = self.rgd.round(rcut)
        u_g = self.rgd.empty()
        logderivs = []
        d0 = 42.0
        offset = 0
        for e in energies:
            dudr = ch.integrate_outwards(u_g, self.rgd, self.vr_sg[0],
                                         gcut, e, self.scalar_relativistic,
                                         self.Z)[0]
            d1 = np.arctan(dudr / u_g[gcut]) / pi + offset
            if d1 > d0:
                offset -= 1
                d1 -= 1
            logderivs.append(d1)
            d0 = d1

        return np.array(logderivs)

    def calculate_exx(self, s=None):
        if s is None:
            self.exx = sum(self.calculate_exx(s)
                           for s in range(self.nspins)) / self.nspins
            return self.exx

        states = []
        lmax = 0
        for ch in self.channels:
            l = ch.l
            for n, phi_g in enumerate(ch.phi_ng):
                f = ch.f_n[n]
                if f > 0 and ch.s == s:
                    states.append((l, f * self.nspins / 2.0 / (2 * l + 1),
                                   phi_g))
                    if l > lmax:
                        lmax = l

        G_LLL = gaunt(lmax)

        exx = 0.0
        j1 = 0
        for l1, f1, phi1_g in states:
            f = 1.0
            for l2, f2, phi2_g in states[j1:]:
                n_g = phi1_g * phi2_g
                for l in range((l1 + l2) % 2, l1 + l2 + 1, 2):
                    G = (G_LLL[l1**2:(l1 + 1)**2,
                               l2**2:(l2 + 1)**2,
                               l**2:(l + 1)**2]**2).sum()
                    vr_g = self.rgd.poisson(n_g, l)
                    e = f * self.rgd.integrate(vr_g * n_g, -1) / 4 / pi
                    exx -= e * G * f1 * f2
                f = 2.0
            j1 += 1

        return exx
Example #3
0
from gpaw.atom.radialgd import EquidistantRadialGridDescriptor
from gpaw.grid_descriptor import GridDescriptor
from gpaw.xc import XC
import numpy as np
from gpaw.test import equal

rgd = EquidistantRadialGridDescriptor(0.01, 100)

for name in ['LDA', 'PBE']:
    xc = XC(name)
    for nspins in [1, 2]:
        n = rgd.zeros(nspins)
        v = rgd.zeros(nspins)
        n[:] = np.exp(-rgd.r_g**2)
        n[-1] *= 2
        E = xc.calculate_spherical(rgd, n, v)
        i = 23
        x = v[-1, i] * rgd.dv_g[i]
        n[-1, i] += 0.000001
        Ep = xc.calculate_spherical(rgd, n, v)
        n[-1, i] -= 0.000002
        Em = xc.calculate_spherical(rgd, n, v)
        x2 = (Ep - Em) / 0.000002
        print(name, nspins, E, x, x2, x - x2)
        equal(x, x2, 1e-9)
        n[-1, i] += 0.000001
        if nspins == 1:
            ns = rgd.empty(2)
            ns[:] = n / 2
            Es = xc.calculate_spherical(rgd, ns, 0 * ns)
            equal(E, Es, 1e-13)
Example #4
0
class C_GLLBScr(Contribution):
    def __init__(self, nlfunc, weight, functional='GGA_X_B88', metallic=False):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional
        self.old_coeffs = None
        self.iter = 0
        self.metallic = metallic
        
    def get_name(self):
        return 'SCREENING'

    def get_desc(self):
        return '(' + self.functional + ')'
        
    # Initialize GLLBScr functional
    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional)
        self.v_g = np.zeros(self.ae.N)
        self.e_g = np.zeros(self.ae.N)

    # Calcualte the GLLB potential and energy 1d
    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.n.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        v_g += 2 * self.weight * self.e_g / (self.ae.n + 1e-10)
        Exc = self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)
        return Exc

    def initialize(self):
        self.occupations = self.nlfunc.occupations
        self.xc = XC(self.functional)

        # Always 1 spin, no matter what calculation nspins is
        self.vt_sg = self.nlfunc.finegd.empty(1) 
        self.e_g = self.nlfunc.finegd.empty()#.ravel()

    def get_coefficient_calculator(self):
        return self

    def f(self, f):
        return sqrt(f)
    
    def get_coefficients(self, e_j, f_j):
        homo_e = max( [ np.where(f>1e-3, e, -1000) for f,e in zip(f_j, e_j)] ) 
        return [ f * K_G * self.f( max(0, homo_e - e)) for e,f in zip(e_j, f_j) ]

    def get_coefficients_1d(self, smooth=False, lumo_perturbation = False):
        homo_e = max( [ np.where(f>1e-3, e, -1000) for f,e in zip(self.ae.f_j, self.ae.e_j)]) 
        if not smooth:
            if lumo_perturbation:
                lumo_e = min( [ np.where(f<1e-3, e, 1000) for f,e in zip(self.ae.f_j, self.ae.e_j)])
                return np.array([ f * K_G * (self.f( max(0, lumo_e - e)) - self.f(max(0, homo_e -e)))
                                        for e,f in zip(self.ae.e_j, self.ae.f_j) ])
            else:
                return np.array([ f * K_G * (self.f( max(0, homo_e - e)))
                                   for e,f in zip(self.ae.e_j, self.ae.f_j) ])
        else:
            return [ [ f * K_G * self.f( max(0, homo_e - e))
                    for e,f in zip(e_n, f_n) ]
                     for e_n, f_n in zip(self.ae.e_ln, self.ae.f_ln) ]
        

    def get_coefficients_by_kpt(self, kpt_u, lumo_perturbation=False, homolumo=None, nspins=1):
        if not hasattr(kpt_u[0],'orbitals_ready'):
            kpt_u[0].orbitals_ready = True
            return None
        #if kpt_u[0].psit_nG is None or isinstance(kpt_u[0].psit_nG,
        #                                          TarFileReference): 
        #    if kpt_u[0].C_nM==None:
        #        return None

        if homolumo == None:
            if self.metallic:
                # For metallic systems, the calculated fermi level represents 
                # the most accurate estimate for reference-energy
                eref_lumo_s = eref_s = nspins * [ self.occupations.get_fermi_level() ]
            else:
                # Find h**o and lumo levels for each spin
                eref_s = []
                eref_lumo_s = []
                for s in range(nspins):
                    h**o, lumo = self.occupations.get_homo_lumo_by_spin(self.nlfunc.wfs, s)
                    eref_s.append(h**o)
                    eref_lumo_s.append(lumo)
        else:
            eref_s, eref_lumo_s = homolumo
            if not isinstance(eref_s, (list, tuple)):
                eref_s = [ eref_s ]
                eref_lumo_s = [ eref_lumo_s ]

        # The parameter ee might sometimes be set to small thereshold value to
        # achieve convergence on small systems with degenerate H**O.
        if len(kpt_u) > nspins:
            ee = 0.0
        else:
            ee = 0.05 / 27.21

        if lumo_perturbation:
            return [np.array([
                f * K_G * (self.f( np.where(eref_lumo_s[kpt.s] - e>ee, eref_lumo_s[kpt.s]-e,0))
                         -self.f( np.where(eref_s[kpt.s]      - e>ee, eref_s[kpt.s]-e,0)))
                     for e, f in zip(kpt.eps_n, kpt.f_n) ])
                     for kpt in kpt_u ]
            
            
        else:
            coeff = [ np.array([ f * K_G * self.f( np.where(eref_s[kpt.s] - e>ee, eref_s[kpt.s]-e,0))
                     for e, f in zip(kpt.eps_n, kpt.f_n) ])
                     for kpt in kpt_u ]
            return coeff
        

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        self.e_g[:] = np.where(n_g<1e-10, 0, self.e_g)
        v_g += self.weight * 2 * self.e_g / (n_g + 1e-10)
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, n_sg, v_sg):
	# Calculate spinpolarized exchange screening as two spin-paired calculations n=2*n_s
        for n, v in [ (n_sg[0], v_sg[0]), (n_sg[1], v_sg[1]) ]:
		self.e_g[:] = 0.0
	        self.vt_sg[:] = 0.0
	        self.xc.calculate(self.nlfunc.finegd, 2*n[None, ...], self.vt_sg, self.e_g)
	        self.e_g[:] = np.where(n<1e-10, 0, self.e_g)
	        v += self.weight * 2 * self.e_g / (2 * n + 1e-9)
	        e_g += self.weight * self.e_g / 2

    def calculate_energy_and_derivatives(self, setup, D_sp, H_sp, a, addcoredensity=True):
        # Get the XC-correction instance
        c = setup.xc_correction
	nspins = self.nlfunc.nspins

        E = 0
	for D_p, dEdD_p in zip(D_sp, H_sp):
	        D_Lq = np.dot(c.B_pqL.T, nspins*D_p)
	        n_Lg = np.dot(D_Lq, c.n_qg)
                if addcoredensity:
                     n_Lg[0] += c.nc_g * sqrt(4 * pi)
	        nt_Lg = np.dot(D_Lq, c.nt_qg)
                if addcoredensity:
                     nt_Lg[0] += c.nct_g * sqrt(4 * pi)
	        dndr_Lg = np.zeros((c.Lmax, c.ng))
	        dntdr_Lg = np.zeros((c.Lmax, c.ng))
	        for L in range(c.Lmax):
	            c.rgd.derivative(n_Lg[L], dndr_Lg[L])
	            c.rgd.derivative(nt_Lg[L], dntdr_Lg[L])
	        vt_g = np.zeros(c.ng)
	        v_g = np.zeros(c.ng)
	        e_g = np.zeros(c.ng)
	        deda2_g = np.zeros(c.ng)
	        for y, (w, Y_L) in enumerate(zip(weight_n, c.Y_nL)):
	            # Cut gradient releated coefficient to match the setup's Lmax
	            A_Li = rnablaY_nLv[y, :c.Lmax]
	
	            # Expand pseudo density
	            nt_g = np.dot(Y_L, nt_Lg)
	
	            # Expand pseudo density gradient
	            a1x_g = np.dot(A_Li[:, 0], nt_Lg)
	            a1y_g = np.dot(A_Li[:, 1], nt_Lg)
	            a1z_g = np.dot(A_Li[:, 2], nt_Lg)
	            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
	            a2_g[1:] /= c.rgd.r_g[1:]**2
	            a2_g[0] = a2_g[1]
	            a1_g = np.dot(Y_L, dntdr_Lg)
	            a2_g += a1_g**2
	            
	            vt_g[:] = 0.0
	            e_g[:] = 0.0
	            # Calculate pseudo GGA energy density (potential is discarded)
	            self.xc.kernel.calculate(e_g, nt_g.reshape((1, -1)),
	                                     vt_g.reshape((1, -1)),
	                                     a2_g.reshape((1, -1)),
	                                     deda2_g.reshape((1, -1)))
	
	            # Calculate pseudo GLLB-potential from GGA-energy density
	            vt_g[:] = 2 * e_g / (nt_g + 1e-10)
	
	            dEdD_p -= self.weight * w * np.dot(np.dot(c.B_pqL, Y_L),
	                                  np.dot(c.nt_qg, vt_g * c.rgd.dv_g))
	
	            E -= w * np.dot(e_g, c.rgd.dv_g) / nspins
            
	            # Expand density
	            n_g = np.dot(Y_L, n_Lg)
	
	            # Expand density gradient
	            a1x_g = np.dot(A_Li[:, 0], n_Lg)
	            a1y_g = np.dot(A_Li[:, 1], n_Lg)
	            a1z_g = np.dot(A_Li[:, 2], n_Lg)
	            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
	            a2_g[1:] /= c.rgd.r_g[1:]**2
	            a2_g[0] = a2_g[1]
	            a1_g = np.dot(Y_L, dndr_Lg)
	            a2_g += a1_g**2
	            
	            v_g[:] = 0.0
	            e_g[:] = 0.0
	            # Calculate GGA energy density (potential is discarded)
	            self.xc.kernel.calculate(e_g, n_g.reshape((1, -1)),
	                                     v_g.reshape((1, -1)),
	                                     a2_g.reshape((1, -1)),
	                                     deda2_g.reshape((1, -1)))
	
	            # Calculate GLLB-potential from GGA-energy density
	            v_g[:] = 2 * e_g / (n_g + 1e-10)
	            
	            dEdD_p += self.weight * w * np.dot(np.dot(c.B_pqL, Y_L),
	                                  np.dot(c.n_qg, v_g * c.rgd.dv_g))
	            E += w * np.dot(e_g, c.rgd.dv_g) / nspins
            
        return E * self.weight

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.nt.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        vt_g += 2 * self.weight * self.e_g / (self.ae.nt + 1e-10)
        return self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)

    def initialize_from_atomic_orbitals(self, basis_functions):
        # GLLBScr needs only density which is already initialized
        pass
        
    def add_extra_setup_data(self, dict):
        # GLLBScr has not any special data
        pass

    def read(self, reader):
        # GLLBScr has no special data to be read
        pass

    def write(self, writer, natoms):
        # GLLBScr has no special data to be written
        pass
Example #5
0
class C_XC(Contribution):
    def __init__(self, nlfunc, weight, functional = 'LDA'):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional

    def get_name(self):
        return 'XC'

    def get_desc(self):
        return "("+self.functional+")"
        
    def initialize(self):
        self.xc = XC(self.functional)
        self.vt_sg = self.nlfunc.finegd.empty(self.nlfunc.nspins)
        self.e_g = self.nlfunc.finegd.empty()

    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional) 
        self.v_g = np.zeros(self.ae.N)

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        v_g += self.weight * self.vt_sg[0]
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, n_sg, v_sg):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_sg, self.vt_sg, self.e_g)
        #self.xc.get_energy_and_potential(na_g, self.vt_sg[0], nb_g, self.vt_sg[1], e_g=self.e_g)
        v_sg[0] += self.weight * self.vt_sg[0]
        v_sg[1] += self.weight * self.vt_sg[1]
        e_g += self.weight * self.e_g

    def calculate_energy_and_derivatives(self, setup, D_sp, H_sp, a, addcoredensity=True):
        E = self.xc.calculate_paw_correction(setup, D_sp, H_sp, True, a)
        E += setup.xc_correction.Exc0
        print("E", E)
        return E

    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.n.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        v_g += self.weight * self.v_g
        return self.weight * Exc

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.nt.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        vt_g += self.weight * self.v_g
        return self.weight * Exc

    def initialize_from_atomic_orbitals(self, basis_functions):
        # LDA needs only density, which is already initialized
        pass

    def add_extra_setup_data(self, dict):
        # LDA has not any special data
        pass

    def write(self, writer, natoms):
        # LDA has not any special data to be written
        pass

    def read(self, reader):
        # LDA has not any special data to be read
        pass
Example #6
0
class C_XC(Contribution):
    def __init__(self, nlfunc, weight, functional='LDA'):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional

    def get_name(self):
        return 'XC'

    def get_desc(self):
        return "(" + self.functional + ")"

    def initialize(self):
        self.xc = XC(self.functional)
        self.vt_sg = self.nlfunc.finegd.empty(self.nlfunc.nspins)
        self.e_g = self.nlfunc.finegd.empty()

    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional)
        self.v_g = np.zeros(self.ae.N)

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        v_g += self.weight * self.vt_sg[0]
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, n_sg, v_sg):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_sg, self.vt_sg, self.e_g)
        #self.xc.get_energy_and_potential(na_g, self.vt_sg[0], nb_g, self.vt_sg[1], e_g=self.e_g)
        v_sg[0] += self.weight * self.vt_sg[0]
        v_sg[1] += self.weight * self.vt_sg[1]
        e_g += self.weight * self.e_g

    def calculate_energy_and_derivatives(self,
                                         setup,
                                         D_sp,
                                         H_sp,
                                         a,
                                         addcoredensity=True):
        E = self.xc.calculate_paw_correction(setup, D_sp, H_sp, True, a)
        E += setup.xc_correction.Exc0
        print("E", E)
        return E

    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.n.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        v_g += self.weight * self.v_g
        return self.weight * Exc

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.nt.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        vt_g += self.weight * self.v_g
        return self.weight * Exc

    def initialize_from_atomic_orbitals(self, basis_functions):
        # LDA needs only density, which is already initialized
        pass

    def add_extra_setup_data(self, dict):
        # LDA has not any special data
        pass

    def write(self, writer, natoms):
        # LDA has not any special data to be written
        pass

    def read(self, reader):
        # LDA has not any special data to be read
        pass
Example #7
0
class C_GLLBScr(Contribution):
    def __init__(self,
                 nlfunc,
                 weight,
                 functional='GGA_X_B88',
                 width=None,
                 eps=0.05,
                 damp=1e-10):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional
        self.old_coeffs = None
        self.iter = 0
        self.damp = damp
        if width is not None:
            width = width / 27.21
        self.eps = eps / 27.21
        self.width = width

    def get_name(self):
        return 'SCREENING'

    def set_damp(self, damp):
        self.damp = damp

    def get_desc(self):
        return '(' + self.functional + ')'

    # Initialize GLLBScr functional
    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional)
        self.v_g = np.zeros(self.ae.N)
        self.e_g = np.zeros(self.ae.N)

    # Calcualte the GLLB potential and energy 1d
    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.n.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        v_g += 2 * self.weight * self.e_g / (self.ae.n + self.damp)
        Exc = self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)
        return Exc

    def initialize(self):
        self.occupations = self.nlfunc.occupations
        self.xc = XC(self.functional)

        # Always 1 spin, no matter what calculation nspins is
        self.vt_sg = self.nlfunc.finegd.empty(1)
        self.e_g = self.nlfunc.finegd.empty()  #.ravel()

    def get_coefficient_calculator(self):
        return self

    def f(self, f):
        if self.width is None:
            if f > self.eps:
                return sqrt(f)
            else:
                return 0.0
        else:
            dEH = -f
            w = self.width
            if dEH / w < -100:
                return sqrt(f)
            Knew = -0.5 * erf(sqrt((max(0.0,dEH)-dEH)/w)) * \
                    sqrt(w*pi) * exp(-dEH/w)
            Knew += 0.5 * sqrt(w * pi) * exp(-dEH / w)
            Knew += sqrt(max(0.0, dEH) - dEH) * exp(max(0.0, dEH) / w)
            #print dEH, w, dEH/w, Knew, f**0.5
            return Knew

    def get_coefficients(self, e_j, f_j):
        homo_e = max([np.where(f > 1e-3, e, -1000) for f, e in zip(f_j, e_j)])
        return [f * K_G * self.f(homo_e - e) for e, f in zip(e_j, f_j)]

    def get_coefficients_1d(self, smooth=False, lumo_perturbation=False):
        homo_e = max([
            np.where(f > 1e-3, e, -1000)
            for f, e in zip(self.ae.f_j, self.ae.e_j)
        ])
        if not smooth:
            if lumo_perturbation:
                lumo_e = min([
                    np.where(f < 1e-3, e, 1000)
                    for f, e in zip(self.ae.f_j, self.ae.e_j)
                ])
                return np.array([
                    f * K_G * (self.f(lumo_e - e) - self.f(homo_e - e))
                    for e, f in zip(self.ae.e_j, self.ae.f_j)
                ])
            else:
                return np.array([
                    f * K_G * (self.f(homo_e - e))
                    for e, f in zip(self.ae.e_j, self.ae.f_j)
                ])
        else:
            return [[f * K_G * self.f(homo_e - e) for e, f in zip(e_n, f_n)]
                    for e_n, f_n in zip(self.ae.e_ln, self.ae.f_ln)]

    def get_coefficients_by_kpt(self,
                                kpt_u,
                                lumo_perturbation=False,
                                homolumo=None,
                                nspins=1):
        #if not hasattr(kpt_u[0],'orbitals_ready'):
        #    kpt_u[0].orbitals_ready = True
        #    return None
        #if not hasattr(self.occupations, 'nvalence'):
        #    print "occupations not ready"
        #    return None
        #if self.occupations.nvalence is None:
        #    return None
        #if kpt_u[0].psit_nG is None or isinstance(kpt_u[0].psit_nG,
        #                                          TarFileReference):
        #    if kpt_u[0].C_nM is None:
        #        return None

        if homolumo is None:
            # Find h**o and lumo levels for each spin
            eref_s = []
            eref_lumo_s = []
            for s in range(nspins):
                h**o, lumo = self.nlfunc.wfs.get_homo_lumo(s)
                eref_s.append(h**o)
                eref_lumo_s.append(lumo)
        else:
            eref_s, eref_lumo_s = homolumo
            if not isinstance(eref_s, (list, tuple)):
                eref_s = [eref_s]
                eref_lumo_s = [eref_lumo_s]

        # The parameter ee might sometimes be set to small thereshold value to
        # achieve convergence on small systems with degenerate H**O.
        if len(kpt_u) > nspins:
            ee = 0.0
        else:
            ee = 0.05 / 27.21

        if lumo_perturbation:
            return [
                np.array([
                    f * K_G * (self.f(eref_lumo_s[kpt.s] - e) -
                               self.f(eref_s[kpt.s] - e))
                    for e, f in zip(kpt.eps_n, kpt.f_n)
                ]) for kpt in kpt_u
            ]

        else:
            coeff = [
                np.array([
                    f * K_G * self.f(eref_s[kpt.s] - e)
                    for e, f in zip(kpt.eps_n, kpt.f_n)
                ]) for kpt in kpt_u
            ]
            #print coeff
            return coeff

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        self.e_g[:] = np.where(n_g < self.damp, 0, self.e_g)
        v_g += self.weight * 2 * self.e_g / (n_g + self.damp)
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, n_sg, v_sg):
        # Calculate spinpolarized exchange screening as two spin-paired calculations n=2*n_s
        for n, v in [(n_sg[0], v_sg[0]), (n_sg[1], v_sg[1])]:
            self.e_g[:] = 0.0
            self.vt_sg[:] = 0.0
            self.xc.calculate(self.nlfunc.finegd, 2 * n[None, ...], self.vt_sg,
                              self.e_g)
            self.e_g[:] = np.where(n < self.damp, 0, self.e_g)
            v += self.weight * 2 * self.e_g / (2 * n + self.damp)
            e_g += self.weight * self.e_g / 2

    def calculate_energy_and_derivatives(self,
                                         setup,
                                         D_sp,
                                         H_sp,
                                         a,
                                         addcoredensity=True):
        # Get the XC-correction instance
        c = setup.xc_correction
        nspins = self.nlfunc.nspins

        E = 0
        for D_p, dEdD_p in zip(D_sp, H_sp):
            D_Lq = np.dot(c.B_pqL.T, nspins * D_p)
            n_Lg = np.dot(D_Lq, c.n_qg)
            if addcoredensity:
                n_Lg[0] += c.nc_g * sqrt(4 * pi)
            nt_Lg = np.dot(D_Lq, c.nt_qg)
            if addcoredensity:
                nt_Lg[0] += c.nct_g * sqrt(4 * pi)
            dndr_Lg = np.zeros((c.Lmax, c.ng))
            dntdr_Lg = np.zeros((c.Lmax, c.ng))
            for L in range(c.Lmax):
                c.rgd.derivative(n_Lg[L], dndr_Lg[L])
                c.rgd.derivative(nt_Lg[L], dntdr_Lg[L])
            vt_g = np.zeros(c.ng)
            v_g = np.zeros(c.ng)
            e_g = np.zeros(c.ng)
            deda2_g = np.zeros(c.ng)
            for y, (w, Y_L) in enumerate(zip(weight_n, c.Y_nL)):
                # Cut gradient releated coefficient to match the setup's Lmax
                A_Li = rnablaY_nLv[y, :c.Lmax]

                # Expand pseudo density
                nt_g = np.dot(Y_L, nt_Lg)

                # Expand pseudo density gradient
                a1x_g = np.dot(A_Li[:, 0], nt_Lg)
                a1y_g = np.dot(A_Li[:, 1], nt_Lg)
                a1z_g = np.dot(A_Li[:, 2], nt_Lg)
                a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
                a2_g[1:] /= c.rgd.r_g[1:]**2
                a2_g[0] = a2_g[1]
                a1_g = np.dot(Y_L, dntdr_Lg)
                a2_g += a1_g**2

                vt_g[:] = 0.0
                e_g[:] = 0.0
                # Calculate pseudo GGA energy density (potential is discarded)
                self.xc.kernel.calculate(e_g, nt_g.reshape((1, -1)),
                                         vt_g.reshape((1, -1)),
                                         a2_g.reshape((1, -1)),
                                         deda2_g.reshape((1, -1)))

                # Calculate pseudo GLLB-potential from GGA-energy density
                vt_g[:] = 2 * e_g / (nt_g + self.damp)

                dEdD_p -= self.weight * w * np.dot(
                    np.dot(c.B_pqL, Y_L), np.dot(c.nt_qg, vt_g * c.rgd.dv_g))

                E -= w * np.dot(e_g, c.rgd.dv_g) / nspins

                # Expand density
                n_g = np.dot(Y_L, n_Lg)

                # Expand density gradient
                a1x_g = np.dot(A_Li[:, 0], n_Lg)
                a1y_g = np.dot(A_Li[:, 1], n_Lg)
                a1z_g = np.dot(A_Li[:, 2], n_Lg)
                a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
                a2_g[1:] /= c.rgd.r_g[1:]**2
                a2_g[0] = a2_g[1]
                a1_g = np.dot(Y_L, dndr_Lg)
                a2_g += a1_g**2

                v_g[:] = 0.0
                e_g[:] = 0.0
                # Calculate GGA energy density (potential is discarded)
                self.xc.kernel.calculate(e_g, n_g.reshape((1, -1)),
                                         v_g.reshape((1, -1)),
                                         a2_g.reshape((1, -1)),
                                         deda2_g.reshape((1, -1)))

                # Calculate GLLB-potential from GGA-energy density
                v_g[:] = 2 * e_g / (n_g + self.damp)

                dEdD_p += self.weight * w * np.dot(
                    np.dot(c.B_pqL, Y_L), np.dot(c.n_qg, v_g * c.rgd.dv_g))
                E += w * np.dot(e_g, c.rgd.dv_g) / nspins

        return E * self.weight

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.nt.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        vt_g += 2 * self.weight * self.e_g / (self.ae.nt + self.damp)
        return self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)

    def initialize_from_atomic_orbitals(self, basis_functions):
        # GLLBScr needs only density which is already initialized
        pass

    def add_extra_setup_data(self, dict):
        # GLLBScr has not any special data
        pass

    def read(self, reader):
        # GLLBScr has no special data to be read
        pass

    def write(self, writer):
        # GLLBScr has no special data to be written
        pass
Example #8
0
class AllElectronAtom:
    def __init__(self, symbol, xc='LDA', spinpol=False, dirac=False,
                 log=sys.stdout):
        """All-electron calculation for spherically symmetric atom.

        symbol: str (or int)
            Chemical symbol (or atomic number).
        xc: str
            Name of XC-functional.
        spinpol: bool
            If true, do spin-polarized calculation.  Default is spin-paired.
        dirac: bool
            Solve Dirac equation instead of Schrödinger equation.
        log: stream
            Text output."""

        if isinstance(symbol, int):
            symbol = chemical_symbols[symbol]
        self.symbol = symbol
        self.Z = atomic_numbers[symbol]

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

        self.dirac = bool(dirac)
        
        self.scalar_relativistic = False

        if isinstance(xc, str):
            self.xc = XC(xc)
        else:
            self.xc = xc

        if log is None:
            log = devnull
        self.fd = log

        self.vr_sg = None  # potential * r
        self.n_sg = 0.0    # density
        self.rgd = None     # radial grid descriptor

        # Energies:
        self.ekin = None
        self.eeig = None
        self.eH = None
        self.eZ = None

        self.channels = None

        self.initialize_configuration()

        self.log('Z:              ', self.Z)
        self.log('Name:           ', atomic_names[self.Z])
        self.log('Symbol:         ', symbol)
        self.log('XC-functional:  ', self.xc.name)
        self.log('Equation:       ', ['Schrödinger', 'Dirac'][self.dirac])

        self.method = 'Gaussian basis-set'

    def log(self, *args, **kwargs):
        prnt(file=self.fd, *args, **kwargs)

    def initialize_configuration(self):
        self.f_lsn = {}
        for n, l, f, e in configurations[self.symbol][1]:
            
            if l not in self.f_lsn:
                self.f_lsn[l] = [[] for s in range(self.nspins)]
            if self.nspins == 1:
                self.f_lsn[l][0].append(f)
            else:
                # Use Hund's rule:
                f0 = min(f, 2 * l + 1)
                self.f_lsn[l][0].append(f0)
                self.f_lsn[l][1].append(f - f0)

    def add(self, n, l, df=+1, s=None):
        """Add (remove) electrons."""
        if s is None:
            if self.nspins == 1:
                s = 0
            else:
                self.add(n, l, 0.5 * df, 0)
                self.add(n, l, 0.5 * df, 1)
                return
            
        if l not in self.f_lsn:
            self.f_lsn[l] = [[] for x in range(self.nspins)]
            
        f_n = self.f_lsn[l][s]
        if len(f_n) < n - l:
            f_n.extend([0] * (n - l - len(f_n)))
        f_n[n - l - 1] += df

    def initialize(self, ngpts=2000, rcut=50.0,
                   alpha1=0.01, alpha2=None, ngauss=50,
                   eps=1.0e-7):
        """Initialize basis sets and radial grid.

        ngpts: int
            Number of grid points for radial grid.
        rcut: float
            Cutoff for radial grid.
        alpha1: float
            Smallest exponent for gaussian.
        alpha2: float
            Largest exponent for gaussian.
        ngauss: int
            Number of gaussians.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        if alpha2 is None:
            alpha2 = 50.0 * self.Z**2

        # Use grid with r(0)=0, r(1)=a and r(ngpts)=rcut:
        a = 1 / alpha2**0.5 / 20
        b = (rcut - a * ngpts) / (rcut * ngpts)
        b = 1 / round(1 / b)
        self.rgd = AERadialGridDescriptor(a, b, ngpts)
        
        self.log('Grid points:     %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %
                 ((self.rgd.N,) + tuple(self.rgd.r_g[[0, 1, 2, -2, -1]])))

        # Distribute exponents between alpha1 and alpha2:
        alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)
        self.log('Exponents:       %d (%.3f, %.3f, ..., %.3f, %.3f)' %
                 ((ngauss,) + tuple(alpha_B[[0, 1, -2, -1]])))

        # Maximum l value:
        lmax = max(self.f_lsn.keys())

        self.channels = []
        nb_l = []
        if not self.dirac:
            for l in range(lmax + 1):
                basis = GaussianBasis(l, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for s in range(self.nspins):
                    self.channels.append(Channel(l, s, self.f_lsn[l][s],
                                                 basis))
        else:
            for K in range(1, lmax + 2):
                leff = (K**2 - (self.Z / c)**2)**0.5 - 1
                basis = GaussianBasis(leff, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for k, l in [(-K, K - 1), (K, K)]:
                    if l > lmax:
                        continue
                    f_n = self.f_lsn[l][0]
                    j = abs(k) - 0.5
                    f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)
                    self.channels.append(DiracChannel(k, f_n, basis))

        self.log('Basis functions: %s (%s)' %
                 (', '.join([str(nb) for nb in nb_l]),
                  ', '.join('spdf'[:lmax + 1])))

        self.vr_sg = self.rgd.zeros(self.nspins)
        self.vr_sg[:] = -self.Z

    def solve(self):
        """Diagonalize Schrödinger equation."""
        self.eeig = 0.0
        for channel in self.channels:
            if self.method == 'Gaussian basis-set':
                channel.solve(self.vr_sg[channel.s])
            else:
                channel.solve2(self.vr_sg[channel.s], self.scalar_relativistic)
            self.eeig += channel.get_eigenvalue_sum()

    def calculate_density(self):
        """Calculate elctron density and kinetic energy."""
        self.n_sg = self.rgd.zeros(self.nspins)
        for channel in self.channels:
            self.n_sg[channel.s] += channel.calculate_density()

    def calculate_electrostatic_potential(self):
        """Calculate electrostatic potential and energy."""
        n_g = self.n_sg.sum(0)
        self.vHr_g = self.rgd.poisson(n_g)        
        self.eH = 0.5 * self.rgd.integrate(n_g * self.vHr_g, -1)
        self.eZ = -self.Z * self.rgd.integrate(n_g, -1)
        
    def calculate_xc_potential(self):
        self.vxc_sg = self.rgd.zeros(self.nspins)
        self.exc = self.xc.calculate_spherical(self.rgd, self.n_sg, self.vxc_sg)

    def step(self):
        self.solve()
        self.calculate_density()
        self.calculate_electrostatic_potential()
        self.calculate_xc_potential()
        self.vr_sg = self.vxc_sg * self.rgd.r_g
        self.vr_sg += self.vHr_g
        self.vr_sg -= self.Z
        self.ekin = (self.eeig -
                     self.rgd.integrate((self.vr_sg * self.n_sg).sum(0), -1))
        
    def run(self, mix=0.4, maxiter=117, dnmax=1e-9):
        if self.channels is None:
            self.initialize()

        if self.dirac:
            equation = 'Dirac'
        elif self.scalar_relativistic:
            equation = 'scalar-relativistic Schrödinger'
        else:
            equation = 'non-relativistic Schrödinger'
        self.log('\nSolving %s equation using %s:' % (equation, self.method))

        dn = self.Z
        
        for iter in range(maxiter):
            self.log('.', end='')
            self.fd.flush()
            if iter > 0:
                self.vr_sg *= mix
                self.vr_sg += (1 - mix) * vr_old_sg
                dn = self.rgd.integrate(abs(self.n_sg - n_old_sg).sum(0))
                if dn <= dnmax:
                    self.log('\nConverged in', iter, 'steps')
                    break

            vr_old_sg = self.vr_sg
            n_old_sg = self.n_sg
            self.step()

        self.summary()
        if dn > dnmax:
            raise RuntimeError('Did not converge!')

    def refine(self):
        self.method = 'finite difference'
        self.run(dnmax=1e-6, mix=0.14, maxiter=200)
        
    def summary(self):
        self.write_states()
        self.write_energies()
            
    def write_states(self):
        self.log('\n state  occupation         eigenvalue          <r>')
        if self.dirac:
            self.log(' nl(j)               [Hartree]        [eV]    [Bohr]')
        else:
            self.log(' nl                  [Hartree]        [eV]    [Bohr]')
        self.log('-----------------------------------------------------')
        states = []
        for ch in self.channels:
            for n, f in enumerate(ch.f_n):
                states.append((ch.e_n[n], ch, n))
        states.sort()
        for e, ch, n in states:
            name = str(n + ch.l + 1) + ch.name
            if self.nspins == 2:
                name += '(%s)' % '+-'[ch.s]    
            n_g = ch.calculate_density(n)
            rave = self.rgd.integrate(n_g, 1)
            self.log(' %-7s  %6.3f %13.6f  %13.5f %6.3f' %
                     (name, ch.f_n[n], e, e * units.Hartree, rave))

    def write_energies(self):
        self.log('\nEnergies:          [Hartree]           [eV]')
        self.log('--------------------------------------------')
        for text, e in [('kinetic      ', self.ekin),
                        ('coulomb (e-e)', self.eH),
                        ('coulomb (e-n)', self.eZ),
                        ('xc           ', self.exc),
                        ('total        ',
                         self.ekin + self.eH + self.eZ + self.exc)]:
            self.log(' %s %+13.6f  %+13.5f' % (text, e, units.Hartree * e))

        self.calculate_exx()
        self.log('\nExact exchange energy: %.6f Hartree, %.5f eV' %
                 (self.exx, self.exx * units.Hartree))

    def get_channel(self, l=None, s=0, k=None):
        if self.dirac:
            for channel in self.channels:
                if channel.k == k:
                    return channel
        else:
            for channel in self.channels:
                if channel.l == l and channel.s == s:
                    return channel
        raise ValueError

    def get_orbital(self, n, l=None, s=0, k=None):
        channel = self.get_channel(l, s, k)
        return channel.basis.expand(channel.C_nb[n])

    def plot_wave_functions(self, rc=4.0):
        import matplotlib.pyplot as plt
        for ch in self.channels:
            for n in range(len(ch.f_n)):
                fr_g = ch.basis.expand(ch.C_nb[n]) * self.rgd.r_g
                name = str(n + ch.l + 1) + ch.name
                lw = 2
                if self.nspins == 2:
                    name += '(%s)' % '+-'[ch.s]    
                    if ch.s == 1:
                        lw = 1
                if self.dirac and ch.k > 0:
                    lw = 1
                ls = ['-', '--', '-.', ':'][ch.l]
                n_g = ch.calculate_density(n)
                rave = self.rgd.integrate(n_g, 1)
                gave = self.rgd.round(rave)
                fr_g *= cmp(fr_g[gave], 0)
                plt.plot(self.rgd.r_g, fr_g,
                         ls=ls, lw=lw, color=colors[n + ch.l], label=name)
        plt.legend(loc='best')
        plt.xlabel('r [Bohr]')
        plt.ylabel('$r\\phi(r)$')
        plt.axis(xmax=rc)
        plt.show()

    def logarithmic_derivative(self, l, energies, rcut):
        ch = Channel(l)
        gcut = self.rgd.round(rcut)
        u_g = self.rgd.empty()
        logderivs = []
        for e in energies:
            dudr = ch.integrate_outwards(u_g, self.rgd, self.vr_sg[0],
                                         gcut, e, self.scalar_relativistic)
            logderivs.append(dudr / u_g[gcut])
        return logderivs
            
    def calculate_exx(self, s=None):
        if s is None:
            self.exx = sum(self.calculate_exx(s)
                           for s in range(self.nspins)) / self.nspins
            return self.exx

        states = []
        lmax = 0
        for ch in self.channels:
            l = ch.l
            for n, phi_g in enumerate(ch.phi_ng):
                f = ch.f_n[n]
                if f > 0 and ch.s == s:
                    states.append((l, f * self.nspins / 2.0 / (2 * l + 1),
                                   phi_g))
                    if l > lmax:
                        lmax = l

        G_LLL = make_gaunt(lmax)

        exx = 0.0
        j1 = 0
        for l1, f1, phi1_g in states:
            f = 1.0
            for l2, f2, phi2_g in states[j1:]:
                n_g = phi1_g * phi2_g
                for l in range((l1 + l2) % 2, l1 + l2 + 1, 2):
                    G = (G_LLL[l1**2:(l1 + 1)**2,
                               l2**2:(l2 + 1)**2,
                               l**2:(l + 1)**2]**2).sum()
                    vr_g = self.rgd.poisson(n_g, l)
                    e = f * self.rgd.integrate(vr_g * n_g, -1) / 4 / pi
                    exx -= e * G * f1 * f2
                f = 2.0
            j1 += 1

        return exx
Example #9
0
from gpaw.xc import XC
import numpy as np
from gpaw.test import equal

r = 0.01 * np.arange(100)
dr = 0.01 * np.ones(100)
rgd = RadialGridDescriptor(r, dr)

for name in ['LDA', 'PBE']:
    xc = XC(name)
    for nspins in [1, 2]:
        n = rgd.zeros(nspins)
        v = rgd.zeros(nspins)
        n[:] = np.exp(-r**2)
        n[-1] *= 2
        E = xc.calculate_spherical(rgd, n, v)
        i = 23
        x = v[-1, i] * rgd.dv_g[i]
        n[-1, i] += 0.000001
        Ep = xc.calculate_spherical(rgd, n, v)
        n[-1, i] -= 0.000002
        Em = xc.calculate_spherical(rgd, n, v)
        x2 = (Ep - Em) / 0.000002
        print name, nspins, E, x, x2, x - x2
        equal(x, x2, 1e-9)
        n[-1, i] += 0.000001
        if nspins == 1:
            ns = rgd.empty(2)
            ns[:] = n / 2
            Es = xc.calculate_spherical(rgd, ns, 0 * ns)
            equal(E, Es, 1e-13)
Example #10
0
class AllElectron:
    """Object for doing an atomic DFT calculation."""
    def __init__(self,
                 symbol,
                 xcname='LDA',
                 scalarrel=True,
                 corehole=None,
                 configuration=None,
                 nofiles=True,
                 txt='-',
                 gpernode=150,
                 orbital_free=False,
                 tw_coeff=1.):
        """Do an atomic DFT calculation.

        Example::

          a = AllElectron('Fe')
          a.run()
        """

        if txt is None:
            txt = devnull
        elif txt == '-':
            txt = sys.stdout
        elif isinstance(txt, str):
            txt = open(txt, 'w')
        self.txt = txt

        self.symbol = symbol
        self.xcname = xcname
        self.scalarrel = scalarrel
        self.nofiles = nofiles

        # Get reference state:
        self.Z, nlfe_j = configurations[symbol]

        # Collect principal quantum numbers, angular momentum quantum
        # numbers, occupation numbers and eigenvalues (j is a combined
        # index for n and l):
        self.n_j = [n for n, l, f, e in nlfe_j]
        self.l_j = [l for n, l, f, e in nlfe_j]
        self.f_j = [f for n, l, f, e in nlfe_j]
        self.e_j = [e for n, l, f, e in nlfe_j]

        if configuration is not None:
            j = 0
            for conf in configuration.split(','):
                if conf[0].isdigit():
                    n = int(conf[0])
                    l = 'spdf'.find(conf[1])
                    if len(conf) == 2:
                        f = 1.0
                    else:
                        f = float(conf[2:])
                    #try:
                    assert n == self.n_j[j]
                    assert l == self.l_j[j]
                    self.f_j[j] = f
                    #except IndexError:
                    #    self.n_j.append(n)
                    #    self.l_j.append(l)
                    #    self.f_j.append(f)
                    #    self.e_j.append(self.e_j[-1])
                    j += 1
                else:
                    j += {'He': 1, 'Ne': 3, 'Ar': 5, 'Kr': 8, 'Xe': 11}[conf]

        maxnodes = max([n - l - 1 for n, l in zip(self.n_j, self.l_j)])
        self.N = (maxnodes + 1) * gpernode
        self.beta = 0.4

        self.orbital_free = orbital_free
        self.tw_coeff = tw_coeff

        if self.orbital_free:
            self.n_j = [1]
            self.l_j = [0]
            self.f_j = [self.Z]
            self.e_j = [self.e_j[-1]]

        t = self.text
        t()
        if scalarrel:
            t('Scalar-relativistic atomic ', end='')
        else:
            t('Atomic ', end='')
        t('%s calculation for %s (%s, Z=%d)' %
          (xcname, symbol, atomic_names[self.Z], self.Z))

        if corehole is not None:
            self.ncorehole, self.lcorehole, self.fcorehole = corehole

            # Find j for core hole and adjust occupation:
            for j in range(len(self.f_j)):
                if (self.n_j[j] == self.ncorehole
                        and self.l_j[j] == self.lcorehole):
                    assert self.f_j[j] == 2 * (2 * self.lcorehole + 1)
                    self.f_j[j] -= self.fcorehole
                    self.jcorehole = j
                    break

            coreholestate = '%d%s' % (self.ncorehole, 'spdf'[self.lcorehole])
            t('Core hole in %s state (%s occupation: %.1f)' %
              (coreholestate, coreholestate, self.f_j[self.jcorehole]))
        else:
            self.jcorehole = None
            self.fcorehole = 0

    def text(self, *args, **kwargs):
        self.txt.write(
            kwargs.get('sep', ' ').join([str(arg) for arg in args]) +
            kwargs.get('end', '\n'))

    def initialize_wave_functions(self):
        r = self.r
        dr = self.dr
        # Initialize with Slater function:
        for l, e, u in zip(self.l_j, self.e_j, self.u_j):
            if self.symbol in ['Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au']:
                a = sqrt(-4.0 * e)
            else:
                a = sqrt(-2.0 * e)

            u[:] = r**(1 + l) * np.exp(-a * r)
            norm = np.dot(u**2, dr)
            u *= 1.0 / sqrt(norm)

    def run(self, use_restart_file=True):
        #     beta g
        # r = ------, g = 0, 1, ..., N - 1
        #     N - g
        #
        #        rN
        # g = --------
        #     beta + r

        t = self.text
        N = self.N
        beta = self.beta
        t(N, 'radial gridpoints.')
        self.rgd = AERadialGridDescriptor(beta / N, 1.0 / N, N)
        g = np.arange(N, dtype=float)
        self.r = self.rgd.r_g
        self.dr = self.rgd.dr_g
        self.d2gdr2 = self.rgd.d2gdr2()

        # Number of orbitals:
        nj = len(self.n_j)

        # Radial wave functions multiplied by radius:
        self.u_j = np.zeros((nj, self.N))

        # Effective potential multiplied by radius:
        self.vr = np.zeros(N)

        # Electron density:
        self.n = np.zeros(N)

        # Always spinpaired nspins=1
        self.xc = XC(self.xcname)

        # Initialize for non-local functionals
        if self.xc.type == 'GLLB':
            self.xc.pass_stuff_1d(self)
            self.xc.initialize_1d()

        n_j = self.n_j
        l_j = self.l_j
        f_j = self.f_j
        e_j = self.e_j

        Z = self.Z  # nuclear charge
        r = self.r  # radial coordinate
        dr = self.dr  # dr/dg
        n = self.n  # electron density
        vr = self.vr  # effective potential multiplied by r

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

        restartfile = '%s/%s.restart' % (tempdir, self.symbol)
        if self.xc.type == 'GLLB' or not use_restart_file:
            # Do not start from initial guess when doing
            # non local XC!
            # This is because we need wavefunctions as well
            # on the first iteration.
            fd = None
        else:
            try:
                fd = open(restartfile, 'rb')
            except IOError:
                fd = None
            else:
                try:
                    n[:] = pickle.load(fd)
                except (ValueError, IndexError):
                    fd = None
                else:
                    norm = np.dot(n * r**2, dr) * 4 * pi
                    if abs(norm - sum(f_j)) > 0.01:
                        fd = None
                    else:
                        t('Using old density for initial guess.')
                        n *= sum(f_j) / norm

        if fd is None:
            self.initialize_wave_functions()
            n[:] = self.calculate_density()

        bar = '|------------------------------------------------|'
        t(bar)

        niter = 0
        nitermax = 117
        qOK = log(1e-10)
        mix = 0.4

        # orbital_free needs more iterations and coefficient
        if self.orbital_free:
            mix = 0.01
            nitermax = 2000
            e_j[0] /= self.tw_coeff
            if Z > 10:  #help convergence for third row elements
                mix = 0.002
                nitermax = 10000

        vrold = None

        while True:
            # calculate hartree potential
            hartree(0, n * r * dr, r, vHr)

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

            # calculated exchange correlation potential and energy
            self.vXC[:] = 0.0

            if self.xc.type == 'GLLB':
                # Update the potential to self.vXC an the energy to self.Exc
                Exc = self.xc.get_xc_potential_and_energy_1d(self.vXC)
            else:
                Exc = self.xc.calculate_spherical(self.rgd, n.reshape((1, -1)),
                                                  self.vXC.reshape((1, -1)))

            # calculate new total Kohn-Sham effective potential and
            # admix with old version

            vr[:] = (vHr + self.vXC * r)

            if self.orbital_free:
                vr /= self.tw_coeff

            if niter > 0:
                vr[:] = mix * vr + (1 - mix) * vrold
            vrold = vr.copy()

            # solve Kohn-Sham equation and determine the density change
            self.solve()
            dn = self.calculate_density() - n
            n += dn

            # estimate error from the square of the density change integrated
            q = log(np.sum((r * dn)**2))

            # print progress bar
            if niter == 0:
                q0 = q
                b0 = 0
            else:
                b = int((q0 - q) / (q0 - qOK) * 50)
                if b > b0:
                    self.txt.write(bar[b0:min(b, 50)])
                    self.txt.flush()
                    b0 = b

            # check if converged and break loop if so
            if q < qOK:
                self.txt.write(bar[b0:])
                self.txt.flush()
                break

            niter += 1
            if niter > nitermax:
                raise RuntimeError('Did not converge!')

        tau = self.calculate_kinetic_energy_density()

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

        try:
            fd = open(restartfile, 'wb')
        except IOError:
            pass
        else:
            pickle.dump(n, fd)
            try:
                os.chmod(restartfile, 0o666)
            except OSError:
                pass

        Ekin = 0

        for f, e in zip(f_j, e_j):
            Ekin += f * e

        e_coulomb = 2 * pi * np.dot(n * r * (vHr - Z), dr)
        Ekin += -4 * pi * np.dot(n * vr * r, dr)

        if self.orbital_free:
            # e and vr are not scaled back
            # instead Ekin is scaled for total energy (printed and inside setup)
            Ekin *= self.tw_coeff
            t()
            t('Lambda:{0}'.format(self.tw_coeff))
            t('Correct eigenvalue:{0}'.format(e_j[0] * self.tw_coeff))
            t()

        t()
        t('Energy contributions:')
        t('-------------------------')
        t('Kinetic:   %+13.6f' % Ekin)
        t('XC:        %+13.6f' % Exc)
        t('Potential: %+13.6f' % e_coulomb)
        t('-------------------------')
        t('Total:     %+13.6f' % (Ekin + Exc + e_coulomb))
        self.ETotal = Ekin + Exc + e_coulomb
        t()

        t('state      eigenvalue         ekin         rmax')
        t('-----------------------------------------------')
        for m, l, f, e, u in zip(n_j, l_j, f_j, e_j, self.u_j):
            # Find kinetic energy:
            k = e - np.sum((
                np.where(abs(u) < 1e-160, 0, u)**2 *  # XXXNumeric!
                vr * dr)[1:] / r[1:])

            # Find outermost maximum:
            g = self.N - 4
            while u[g - 1] >= u[g]:
                g -= 1
            x = r[g - 1:g + 2]
            y = u[g - 1:g + 2]
            A = np.transpose(np.array([x**i for i in range(3)]))
            c, b, a = np.linalg.solve(A, y)
            assert a < 0.0
            rmax = -0.5 * b / a

            s = 'spdf'[l]
            t('%d%s^%-4.1f: %12.6f %12.6f %12.3f' % (m, s, f, e, k, rmax))
        t('-----------------------------------------------')
        t('(units: Bohr and Hartree)')

        for m, l, u in zip(n_j, l_j, self.u_j):
            self.write(u, 'ae', n=m, l=l)

        self.write(n, 'n')
        self.write(vr, 'vr')
        self.write(vHr, 'vHr')
        self.write(self.vXC, 'vXC')
        self.write(tau, 'tau')

        self.Ekin = Ekin
        self.e_coulomb = e_coulomb
        self.Exc = Exc

    def write(self, array, name=None, n=None, l=None):
        if self.nofiles:
            return

        if name:
            name = self.symbol + '.' + name
        else:
            name = self.symbol

        if l is not None:
            assert n is not None
            if n > 0:
                # Bound state:
                name += '.%d%s' % (n, 'spdf'[l])
            else:
                name += '.x%d%s' % (-n, 'spdf'[l])

        f = open(name, 'w')
        for r, a in zip(self.r, array):
            print(r, a, file=f)

    def calculate_density(self):
        """Return the electron charge density divided by 4 pi"""
        n = np.dot(self.f_j,
                   np.where(abs(self.u_j) < 1e-160, 0, self.u_j)**2) / (4 * pi)
        n[1:] /= self.r[1:]**2
        n[0] = n[1]
        return n

    def calculate_kinetic_energy_density(self):
        """Return the kinetic energy density"""
        return self.radial_kinetic_energy_density(self.f_j, self.l_j, self.u_j)

    def radial_kinetic_energy_density(self, f_j, l_j, u_j):
        """Kinetic energy density from a restricted set of wf's
        """
        shape = np.shape(u_j[0])
        dudr = np.zeros(shape)
        tau = np.zeros(shape)
        for f, l, u in zip(f_j, l_j, u_j):
            self.rgd.derivative(u, dudr)
            # contribution from angular derivatives
            if l > 0:
                tau += f * l * (l + 1) * np.where(abs(u) < 1e-160, 0, u)**2
            # contribution from radial derivatives
            dudr = u - self.r * dudr
            tau += f * np.where(abs(dudr) < 1e-160, 0, dudr)**2
        tau[1:] /= self.r[1:]**4
        tau[0] = tau[1]

        return 0.5 * tau / (4 * pi)

    def calculate_kinetic_energy_density2(self):
        """Return the kinetic energy density
        calculation over R(r)=u(r)/r
        slower convergence with # of radial grid points for
        Ekin of H than radial_kinetic_energy_density
        """

        shape = self.u_j.shape[1]
        R = np.zeros(shape)
        dRdr = np.zeros(shape)
        tau = np.zeros(shape)
        for f, l, u in zip(self.f_j, self.l_j, self.u_j):
            R[1:] = u[1:] / self.r[1:]
            if l == 0:
                # estimate value at origin by Taylor series to first order
                d1 = self.r[1]
                d2 = self.r[2]
                R[0] = .5 * (R[1] + R[2] + (R[1] - R[2]) * (d1 + d2) /
                             (d2 - d1))
            else:
                R[0] = 0
            self.rgd.derivative(R, dRdr)
            # contribution from radial derivatives
            tau += f * np.where(abs(dRdr) < 1e-160, 0, dRdr)**2
            # contribution from angular derivatives
            if l > 0:
                R[1:] = R[1:] / self.r[1:]
                if l == 1:
                    R[0] = R[1]
                else:
                    R[0] = 0
                tau += f * l * (l + 1) * np.where(abs(R) < 1e-160, 0, R)**2

        return 0.5 * tau / (4 * pi)

    def solve(self):
        """Solve the Schrodinger equation

        ::

             2
            d u     1  dv  du   u     l(l + 1)
          - --- - ---- -- (-- - -) + [-------- + 2M(v - e)] u = 0,
              2      2 dr  dr   r         2
            dr    2Mc                    r


        where the relativistic mass::

                   1
          M = 1 - --- (v - e)
                    2
                  2c

        and the fine-structure constant alpha = 1/c = 1/137.036
        is set to zero for non-scalar-relativistic calculations.

        On the logaritmic radial grids defined by::

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

                 rN
          g = --------, r = [0; oo[
              beta + r

        the Schrodinger equation becomes::

           2
          d u      du
          --- c  + -- c  + u c  = 0
            2  2   dg  1      0
          dg

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

                 2 dg 2
          c  = -r (--)
           2       dr

                  2         2
                 d g  2    r   dg dv
          c  = - --- r  - ---- -- --
           1       2         2 dr dr
                 dr       2Mc

                                    2    r   dv
          c  = l(l + 1) + 2M(v - e)r  + ---- --
           0                               2 dr
                                        2Mc
        """
        r = self.r
        dr = self.dr
        vr = self.vr

        c2 = -(r / dr)**2
        c10 = -self.d2gdr2 * r**2  # first part of c1 vector

        if self.scalarrel:
            self.r2dvdr = np.zeros(self.N)
            self.rgd.derivative(vr, self.r2dvdr)
            self.r2dvdr *= r
            self.r2dvdr -= vr
        else:
            self.r2dvdr = None

        # solve for each quantum state separately
        for j, (n, l, e,
                u) in enumerate(zip(self.n_j, self.l_j, self.e_j, self.u_j)):
            nodes = n - l - 1  # analytically expected number of nodes
            delta = -0.2 * e
            nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                          self.scalarrel)
            # adjust eigenenergy until u has the correct number of nodes
            while nn != nodes:
                diff = np.sign(nn - nodes)
                while diff == np.sign(nn - nodes):
                    e -= diff * delta
                    nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                                  self.scalarrel)
                delta /= 2

            # adjust eigenenergy until u is smooth at the turning point
            de = 1.0
            while abs(de) > 1e-9:
                norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)
                u *= 1.0 / sqrt(norm)
                de = 0.5 * A / norm
                x = abs(de / e)
                if x > 0.1:
                    de *= 0.1 / x
                e -= de
                assert e < 0.0
                nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                              self.scalarrel)
            self.e_j[j] = e
            u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr))

    def solve_confined(self, j, rc, vconf=None):
        """Solve the Schroedinger equation in a confinement potential.

        Solves the Schroedinger equation like the solve method, but with a
        number of differences.  Before invoking this method, run solve() to
        get initial guesses.

        Parameters:
            j: solves only for the state given by j
            rc: solution cutoff. Solution will be zero outside this.
            vconf: added to the potential (use this as confinement potential)

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

        Unlike the solve method, this method will not alter any attributes of
        this object.
        """
        r = self.r
        dr = self.dr
        vr = self.vr.copy()
        if vconf is not None:
            vr += vconf * r

        c2 = -(r / dr)**2
        c10 = -self.d2gdr2 * r**2  # first part of c1 vector

        if j is None:
            n, l, e, u = 3, 2, -0.15, self.u_j[-1].copy()
        else:
            n = self.n_j[j]
            l = self.l_j[j]
            e = self.e_j[j]
            u = self.u_j[j].copy()

        nn, A = shoot_confined(u,
                               l,
                               vr,
                               e,
                               self.r2dvdr,
                               r,
                               dr,
                               c10,
                               c2,
                               self.scalarrel,
                               rc=rc,
                               beta=self.beta)
        assert nn == n - l - 1  # run() should have been called already

        # adjust eigenenergy until u is smooth at the turning point
        de = 1.0
        while abs(de) > 1e-9:
            norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)
            u *= 1.0 / sqrt(norm)
            de = 0.5 * A / norm
            x = abs(de / e)
            if x > 0.1:
                de *= 0.1 / x
            e -= de
            assert e < 0.0

            nn, A = shoot_confined(u,
                                   l,
                                   vr,
                                   e,
                                   self.r2dvdr,
                                   r,
                                   dr,
                                   c10,
                                   c2,
                                   self.scalarrel,
                                   rc=rc,
                                   beta=self.beta)
        u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr))
        return u, e

    def kin(self, l, u, e=None):  # XXX move to Generator
        r = self.r[1:]
        dr = self.dr[1:]

        c0 = 0.5 * l * (l + 1) / r**2
        c1 = -0.5 * self.d2gdr2[1:]
        c2 = -0.5 * dr**-2

        if e is not None and self.scalarrel:
            x = 0.5 * alpha**2
            Mr = r * (1.0 + x * e) - x * self.vr[1:]
            c0 += (
                (Mr - r) *
                (self.vr[1:] - e * r) + 0.5 * x * self.r2dvdr[1:] / Mr) / r**2
            c1 -= 0.5 * x * self.r2dvdr[1:] / (Mr * dr * r)

        fp = c2 + 0.5 * c1
        fm = c2 - 0.5 * c1
        f0 = c0 - 2 * c2
        kr = np.zeros(self.N)
        kr[1:] = f0 * u[1:] + fm * u[:-1]
        kr[1:-1] += fp[:-1] * u[2:]
        kr[0] = 0.0
        return kr

    def r2g(self, r):
        """Convert radius to index of the radial grid."""
        return int(r * self.N / (self.beta + r))

    def get_confinement_potential(self, alpha, ri, rc):
        r"""Create a smooth confinement potential.

        Returns a (potential) function which is zero inside the radius ri
        and goes to infinity smoothly at rc, after which point it is nan.
        The potential is given by::

                   alpha         /   rc - ri \
          V(r) = --------   exp ( - --------- )   for   ri < r < rc
                  rc - r         \    r - ri /

        """
        i_ri = self.r2g(ri)
        i_rc = self.r2g(rc)
        if self.r[i_rc] == rc:
            # Avoid division by zero in the odd case that rc coincides
            # exactly with a grid point (which actually happens sometimes)
            i_rc -= 1

        potential = np.zeros(np.shape(self.r))
        r = self.r[i_ri + 1:i_rc + 1]
        exponent = -(rc - ri) / (r - ri)
        denom = rc - r
        value = np.exp(exponent) / denom
        potential[i_ri + 1:i_rc + 1] = value
        potential[i_rc + 1:] = np.inf

        return alpha * potential
class AllElectron:
    """Object for doing an atomic DFT calculation."""

    def __init__(self, symbol, xcname='LDA', scalarrel=False,
                 corehole=None, configuration=None, nofiles=True,
                 txt='-', gpernode=150, orbital_free=False, tf_coeff=1.):
        """Do an atomic DFT calculation.

        Example::

          a = AllElectron('Fe')
          a.run()
        """
        
        if txt is None:
            txt = devnull
        elif txt == '-':
            txt = sys.stdout
        elif isinstance(txt, str):
            txt = open(txt, 'w')
        self.txt = txt

        self.symbol = symbol
        self.xcname = xcname
        self.scalarrel = scalarrel
        self.nofiles = nofiles

        # Get reference state:
        self.Z, nlfe_j = configurations[symbol]

        # Collect principal quantum numbers, angular momentum quantum
        # numbers, occupation numbers and eigenvalues (j is a combined
        # index for n and l):
        self.n_j = [n for n, l, f, e in nlfe_j]
        self.l_j = [l for n, l, f, e in nlfe_j]
        self.f_j = [f for n, l, f, e in nlfe_j]
        self.e_j = [e for n, l, f, e in nlfe_j]

        if configuration is not None:
            j = 0
            for conf in configuration.split(','):
                if conf[0].isdigit():
                    n = int(conf[0])
                    l = 'spdf'.find(conf[1])
                    if len(conf) == 2:
                        f = 1.0
                    else:
                        f = float(conf[2:])
                    #try:
                    assert n == self.n_j[j]
                    assert l == self.l_j[j]
                    self.f_j[j] = f
                    #except IndexError:
                    #    self.n_j.append(n)
                    #    self.l_j.append(l)
                    #    self.f_j.append(f)
                    #    self.e_j.append(self.e_j[-1])
                    j += 1
                else:
                    j += {'He': 1,
                          'Ne': 3,
                          'Ar': 5,
                          'Kr': 8,
                          'Xe': 11}[conf]

        maxnodes = max([n - l - 1 for n, l in zip(self.n_j, self.l_j)])
        self.N = (maxnodes + 1) * gpernode
        self.beta = 0.4

        self.orbital_free = orbital_free
        self.tf_coeff = tf_coeff

        if self.orbital_free:
            self.n_j = [1]
            self.l_j = [0]
            self.f_j = [self.Z]
            self.e_j = [self.e_j[-1]]
            
        t = self.text
        t()
        if scalarrel:
            t('Scalar-relativistic atomic ', end='')
        else:
            t('Atomic ', end='')
        t('%s calculation for %s (%s, Z=%d)' % (xcname, symbol,
                                                atomic_names[self.Z], self.Z))

        if corehole is not None:
            self.ncorehole, self.lcorehole, self.fcorehole = corehole

            # Find j for core hole and adjust occupation:
            for j in range(len(self.f_j)):
                if (self.n_j[j] == self.ncorehole and
                    self.l_j[j] == self.lcorehole):
                    assert self.f_j[j] == 2 * (2 * self.lcorehole + 1)
                    self.f_j[j] -= self.fcorehole
                    self.jcorehole = j
                    break

            coreholestate = '%d%s' % (self.ncorehole, 'spdf'[self.lcorehole])
            t('Core hole in %s state (%s occupation: %.1f)' % (
                coreholestate, coreholestate, self.f_j[self.jcorehole]))
        else:
            self.jcorehole = None
            self.fcorehole = 0

    def text(self, *args, **kwargs):
        self.txt.write(kwargs.get('sep', ' ').join([str(arg)
                                                    for arg in args]) +
                       kwargs.get('end', '\n'))

    def initialize_wave_functions(self):
        r = self.r
        dr = self.dr
        # Initialize with Slater function:
        for l, e, u in zip(self.l_j, self.e_j, self.u_j):
            if self.symbol in ['Hf', 'Ta', 'W', 'Re', 'Os',
                               'Ir', 'Pt', 'Au']:
                a = sqrt(-4.0 * e)
            else:
                a = sqrt(-2.0 * e)

            u[:] = r**(1 + l) * np.exp(-a * r)
            norm = np.dot(u**2, dr)
            u *= 1.0 / sqrt(norm)
            
    def run(self, use_restart_file=True):
        #     beta g
        # r = ------, g = 0, 1, ..., N - 1
        #     N - g
        #
        #        rN
        # g = --------
        #     beta + r

        t = self.text
        N = self.N
        beta = self.beta
        t(N, 'radial gridpoints.')
        self.rgd = AERadialGridDescriptor(beta / N, 1.0 / N, N)
        g = np.arange(N, dtype=float)
        self.r = self.rgd.r_g
        self.dr = self.rgd.dr_g
        self.d2gdr2 = self.rgd.d2gdr2()

        # Number of orbitals:
        nj = len(self.n_j)

        # Radial wave functions multiplied by radius:
        self.u_j = np.zeros((nj, self.N))

        # Effective potential multiplied by radius:
        self.vr = np.zeros(N)

        # Electron density:
        self.n = np.zeros(N)

        # Always spinpaired nspins=1
        self.xc = XC(self.xcname)

        # Initialize for non-local functionals
        if self.xc.type == 'GLLB':
            self.xc.pass_stuff_1d(self)
            self.xc.initialize_1d()
            
        n_j = self.n_j
        l_j = self.l_j
        f_j = self.f_j
        e_j = self.e_j
        
        Z = self.Z    # nuclear charge
        r = self.r    # radial coordinate
        dr = self.dr  # dr/dg
        n = self.n    # electron density
        vr = self.vr  # effective potential multiplied by r

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

        restartfile = '%s/%s.restart' % (tempdir, self.symbol)
        if self.xc.type == 'GLLB' or not use_restart_file:
            # Do not start from initial guess when doing
            # non local XC!
            # This is because we need wavefunctions as well
            # on the first iteration.
            fd = None
        else:
            try:
                fd = open(restartfile, 'r')
            except IOError:
                fd = None
            else:
                try:
                    n[:] = pickle.load(fd)
                except (ValueError, IndexError):
                    fd = None
                else:
                    norm = np.dot(n * r**2, dr) * 4 * pi
                    if abs(norm - sum(f_j)) > 0.01:
                        fd = None
                    else:
                        t('Using old density for initial guess.')
                        n *= sum(f_j) / norm

        if fd is None:
            self.initialize_wave_functions()
            n[:] = self.calculate_density()

        bar = '|------------------------------------------------|'
        t(bar)
        
        niter = 0
        nitermax = 117
        qOK = log(1e-10)
        mix = 0.4
        
        # orbital_free needs more iterations and coefficient
        if self.orbital_free:
            #qOK = log(1e-14)
            e_j[0] /= self.tf_coeff
            mix = 0.01
            nitermax = 1000
            
        vrold = None
        
        while True:
            # calculate hartree potential
            hartree(0, n * r * dr, r, vHr)

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

            # calculated exchange correlation potential and energy
            self.vXC[:] = 0.0

            if self.xc.type == 'GLLB':
                # Update the potential to self.vXC an the energy to self.Exc
                Exc = self.xc.get_xc_potential_and_energy_1d(self.vXC)
            else:
                Exc = self.xc.calculate_spherical(self.rgd,
                                                  n.reshape((1, -1)),
                                                  self.vXC.reshape((1, -1)))

            # calculate new total Kohn-Sham effective potential and
            # admix with old version
            vr[:] = (vHr + self.vXC * r) / self.tf_coeff

            if niter > 0:
                vr[:] = mix * vr + (1 - mix) * vrold
            vrold = vr.copy()

            # solve Kohn-Sham equation and determine the density change
            self.solve()
            dn = self.calculate_density() - n
            n += dn

            # estimate error from the square of the density change integrated
            q = log(np.sum((r * dn)**2))

            # print progress bar
            if niter == 0:
                q0 = q
                b0 = 0
            else:
                b = int((q0 - q) / (q0 - qOK) * 50)
                if b > b0:
                    self.txt.write(bar[b0:min(b, 50)])
                    self.txt.flush()
                    b0 = b

            # check if converged and break loop if so
            if q < qOK:
                self.txt.write(bar[b0:])
                self.txt.flush()
                break

            niter += 1
            if niter > nitermax:
                raise RuntimeError('Did not converge!')

        tau = self.calculate_kinetic_energy_density()

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

        try:
            fd = open(restartfile, 'w')
        except IOError:
            pass
        else:
            pickle.dump(n, fd)
            try:
                os.chmod(restartfile, 0666)
            except OSError:
                pass

        Ekin = 0
        if self.orbital_free:
            e_j[0] *= self.tf_coeff
            vr *= self.tf_coeff
        
        for f, e in zip(f_j, e_j):
            Ekin += f * e

        Epot = 2 * pi * np.dot(n * r * (vHr - Z), dr)
        Ekin += -4 * pi * np.dot(n * vr * r, dr)

        t()
        t('Energy contributions:')
        t('-------------------------')
        t('Kinetic:   %+13.6f' % Ekin)
        t('XC:        %+13.6f' % Exc)
        t('Potential: %+13.6f' % Epot)
        t('-------------------------')
        t('Total:     %+13.6f' % (Ekin + Exc + Epot))
        self.ETotal = Ekin + Exc + Epot
        t()

        t('state      eigenvalue         ekin         rmax')
        t('-----------------------------------------------')
        for m, l, f, e, u in zip(n_j, l_j, f_j, e_j, self.u_j):
            # Find kinetic energy:
            k = e - np.sum((np.where(abs(u) < 1e-160, 0, u)**2 *  # XXXNumeric!
                            vr * dr)[1:] / r[1:])

            # Find outermost maximum:
            g = self.N - 4
            while u[g - 1] >= u[g]:
                g -= 1
            x = r[g - 1:g + 2]
            y = u[g - 1:g + 2]
            A = np.transpose(np.array([x**i for i in range(3)]))
            c, b, a = np.linalg.solve(A, y)
            assert a < 0.0
            rmax = -0.5 * b / a

            s = 'spdf'[l]
            t('%d%s^%-4.1f: %12.6f %12.6f %12.3f' % (m, s, f, e, k, rmax))
        t('-----------------------------------------------')
        t('(units: Bohr and Hartree)')

        for m, l, u in zip(n_j, l_j, self.u_j):
            self.write(u, 'ae', n=m, l=l)

        self.write(n, 'n')
        self.write(vr, 'vr')
        self.write(vHr, 'vHr')
        self.write(self.vXC, 'vXC')
        self.write(tau, 'tau')

        self.Ekin = Ekin
        self.Epot = Epot
        self.Exc = Exc

    def write(self, array, name=None, n=None, l=None):
        if self.nofiles:
            return

        if name:
            name = self.symbol + '.' + name
        else:
            name = self.symbol

        if l is not None:
            assert n is not None
            if n > 0:
                # Bound state:
                name += '.%d%s' % (n, 'spdf'[l])
            else:
                name += '.x%d%s' % (-n, 'spdf'[l])

        f = open(name, 'w')
        for r, a in zip(self.r, array):
            print(r, a, file=f)

    def calculate_density(self):
        """Return the electron charge density divided by 4 pi"""
        n = np.dot(self.f_j,
                   np.where(abs(self.u_j) < 1e-160, 0,
                            self.u_j)**2) / (4 * pi)
        n[1:] /= self.r[1:]**2
        n[0] = n[1]
        return n

    def calculate_kinetic_energy_density(self):
        """Return the kinetic energy density"""
        return self.radial_kinetic_energy_density(self.f_j, self.l_j, self.u_j)

    def radial_kinetic_energy_density(self, f_j, l_j, u_j):
        """Kinetic energy density from a restricted set of wf's
        """
        shape = np.shape(u_j[0])
        dudr = np.zeros(shape)
        tau = np.zeros(shape)
        for f, l, u in zip(f_j, l_j, u_j):
            self.rgd.derivative(u, dudr)
            # contribution from angular derivatives
            if l > 0:
                tau += f * l * (l + 1) * np.where(abs(u) < 1e-160, 0, u)**2
            # contribution from radial derivatives
            dudr = u - self.r * dudr
            tau += f * np.where(abs(dudr) < 1e-160, 0, dudr)**2
        tau[1:] /= self.r[1:]**4
        tau[0] = tau[1]

        return 0.5 * tau / (4 * pi)

    def calculate_kinetic_energy_density2(self):
        """Return the kinetic energy density
        calculation over R(r)=u(r)/r
        slower convergence with # of radial grid points for
        Ekin of H than radial_kinetic_energy_density
        """

        shape = self.u_j.shape[1]
        R = np.zeros(shape)
        dRdr = np.zeros(shape)
        tau = np.zeros(shape)
        for f, l, u in zip(self.f_j, self.l_j, self.u_j):
            R[1:] = u[1:] / self.r[1:]
            if l == 0:
                # estimate value at origin by Taylor series to first order
                d1 = self.r[1]
                d2 = self.r[2]
                R[0] = .5 * (R[1] + R[2] + (R[1] - R[2]) *
                             (d1 + d2) / (d2 - d1))
            else:
                R[0] = 0
            self.rgd.derivative(R, dRdr)
            # contribution from radial derivatives
            tau += f * np.where(abs(dRdr) < 1e-160, 0, dRdr)**2
            # contribution from angular derivatives
            if l > 0:
                R[1:] = R[1:] / self.r[1:]
                if l == 1:
                    R[0] = R[1]
                else:
                    R[0] = 0
                tau += f * l * (l + 1) * np.where(abs(R) < 1e-160, 0, R)**2

        return 0.5 * tau / (4 * pi)

    def solve(self):
        """Solve the Schrodinger equation

        ::

             2
            d u     1  dv  du   u     l(l + 1)
          - --- - ---- -- (-- - -) + [-------- + 2M(v - e)] u = 0,
              2      2 dr  dr   r         2
            dr    2Mc                    r


        where the relativistic mass::

                   1
          M = 1 - --- (v - e)
                    2
                  2c

        and the fine-structure constant alpha = 1/c = 1/137.036
        is set to zero for non-scalar-relativistic calculations.

        On the logaritmic radial grids defined by::

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

                 rN
          g = --------, r = [0; oo[
              beta + r

        the Schrodinger equation becomes::

           2
          d u      du
          --- c  + -- c  + u c  = 0
            2  2   dg  1      0
          dg

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

                 2 dg 2
          c  = -r (--)
           2       dr

                  2         2
                 d g  2    r   dg dv
          c  = - --- r  - ---- -- --
           1       2         2 dr dr
                 dr       2Mc

                                    2    r   dv
          c  = l(l + 1) + 2M(v - e)r  + ---- --
           0                               2 dr
                                        2Mc
        """
        r = self.r
        dr = self.dr
        vr = self.vr

        c2 = -(r / dr)**2
        c10 = -self.d2gdr2 * r**2  # first part of c1 vector

        if self.scalarrel:
            self.r2dvdr = np.zeros(self.N)
            self.rgd.derivative(vr, self.r2dvdr)
            self.r2dvdr *= r
            self.r2dvdr -= vr
        else:
            self.r2dvdr = None

        # solve for each quantum state separately
        for j, (n, l, e, u) in enumerate(zip(self.n_j, self.l_j,
                                             self.e_j, self.u_j)):
            nodes = n - l - 1  # analytically expected number of nodes
            delta = -0.2 * e
            nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                          self.scalarrel)
            # adjust eigenenergy until u has the correct number of nodes
            while nn != nodes:
                diff = cmp(nn, nodes)
                while diff == cmp(nn, nodes):
                    e -= diff * delta
                    nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                                  self.scalarrel)
                delta /= 2

            # adjust eigenenergy until u is smooth at the turning point
            de = 1.0
            while abs(de) > 1e-9:
                norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)
                u *= 1.0 / sqrt(norm)
                de = 0.5 * A / norm
                x = abs(de / e)
                if x > 0.1:
                    de *= 0.1 / x
                e -= de
                assert e < 0.0
                nn, A = shoot(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                              self.scalarrel)
            self.e_j[j] = e
            u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr))

    def solve_confined(self, j, rc, vconf=None):
        """Solve the Schroedinger equation in a confinement potential.
        
        Solves the Schroedinger equation like the solve method, but with a
        number of differences.  Before invoking this method, run solve() to
        get initial guesses.

        Parameters:
            j: solves only for the state given by j
            rc: solution cutoff. Solution will be zero outside this.
            vconf: added to the potential (use this as confinement potential)

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

        Unlike the solve method, this method will not alter any attributes of
        this object.
        """
        r = self.r
        dr = self.dr
        vr = self.vr.copy()
        if vconf is not None:
            vr += vconf * r

        c2 = -(r / dr)**2
        c10 = -self.d2gdr2 * r**2  # first part of c1 vector

        if j is None:
            n, l, e, u = 3, 2, -0.15, self.u_j[-1].copy()
        else:
            n = self.n_j[j]
            l = self.l_j[j]
            e = self.e_j[j]
            u = self.u_j[j].copy()
            
        nn, A = shoot_confined(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                               self.scalarrel, rc=rc, beta=self.beta)
        assert nn == n - l - 1  # run() should have been called already
        
        # adjust eigenenergy until u is smooth at the turning point
        de = 1.0
        while abs(de) > 1e-9:
            norm = np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr)
            u *= 1.0 / sqrt(norm)
            de = 0.5 * A / norm
            x = abs(de / e)
            if x > 0.1:
                de *= 0.1 / x
            e -= de
            assert e < 0.0

            nn, A = shoot_confined(u, l, vr, e, self.r2dvdr, r, dr, c10, c2,
                                   self.scalarrel, rc=rc, beta=self.beta)
        u *= 1.0 / sqrt(np.dot(np.where(abs(u) < 1e-160, 0, u)**2, dr))
        return u, e

    def kin(self, l, u, e=None):  # XXX move to Generator
        r = self.r[1:]
        dr = self.dr[1:]

        c0 = 0.5 * l * (l + 1) / r**2
        c1 = -0.5 * self.d2gdr2[1:]
        c2 = -0.5 * dr**-2

        if e is not None and self.scalarrel:
            x = 0.5 * alpha**2
            Mr = r * (1.0 + x * e) - x * self.vr[1:]
            c0 += ((Mr - r) * (self.vr[1:] - e * r) +
                   0.5 * x * self.r2dvdr[1:] / Mr) / r**2
            c1 -= 0.5 * x * self.r2dvdr[1:] / (Mr * dr * r)

        fp = c2 + 0.5 * c1
        fm = c2 - 0.5 * c1
        f0 = c0 - 2 * c2
        kr = np.zeros(self.N)
        kr[1:] = f0 * u[1:] + fm * u[:-1]
        kr[1:-1] += fp[:-1] * u[2:]
        kr[0] = 0.0
        return kr

    def r2g(self, r):
        """Convert radius to index of the radial grid."""
        return int(r * self.N / (self.beta + r))

    def get_confinement_potential(self, alpha, ri, rc):
        """Create a smooth confinement potential.
        
        Returns a (potential) function which is zero inside the radius ri
        and goes to infinity smoothly at rc, after which point it is nan.
        The potential is given by::

                   alpha         /   rc - ri \
          V(r) = --------   exp ( - --------- )   for   ri < r < rc
                  rc - r         \    r - ri /

        """
        i_ri = self.r2g(ri)
        i_rc = self.r2g(rc)
        if self.r[i_rc] == rc:
            # Avoid division by zero in the odd case that rc coincides
            # exactly with a grid point (which actually happens sometimes)
            i_rc -= 1

        potential = np.zeros(np.shape(self.r))
        r = self.r[i_ri + 1:i_rc + 1]
        exponent = - (rc - ri) / (r - ri)
        denom = rc - r
        value = np.exp(exponent) / denom
        potential[i_ri + 1:i_rc + 1] = value
        potential[i_rc + 1:] = np.inf

        return alpha * potential
Example #12
0
class AllElectronAtom:
    def __init__(self,
                 symbol,
                 xc='LDA',
                 spinpol=False,
                 dirac=False,
                 log=sys.stdout):
        """All-electron calculation for spherically symmetric atom.

        symbol: str (or int)
            Chemical symbol (or atomic number).
        xc: str
            Name of XC-functional.
        spinpol: bool
            If true, do spin-polarized calculation.  Default is spin-paired.
        dirac: bool
            Solve Dirac equation instead of Schrödinger equation.
        log: stream
            Text output."""

        if isinstance(symbol, int):
            symbol = chemical_symbols[symbol]
        self.symbol = symbol
        self.Z = atomic_numbers[symbol]

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

        self.dirac = bool(dirac)

        if isinstance(xc, str):
            self.xc = XC(xc)
        else:
            self.xc = xc

        if log is None:
            log = devnull
        self.fd = log

        self.vr_sg = None  # potential * r
        self.n_sg = 0.0  # density
        self.gd = None  # radial grid descriptor

        # Energies:
        self.ekin = None
        self.eeig = None
        self.eH = None
        self.eZ = None

        self.channels = None

        self.initialize_configuration()

        self.log('Z:              ', self.Z)
        self.log('Name:           ', atomic_names[self.Z])
        self.log('Symbol:         ', symbol)
        self.log('XC-functional:  ', self.xc.name)
        self.log('Equation:       ', ['Schrödinger', 'Dirac'][self.dirac])

    def log(self, *args, **kwargs):
        self.fd.write(
            kwargs.get('sep', ' ').join([str(arg) for arg in args]) +
            kwargs.get('end', '\n'))

    def initialize_configuration(self):
        self.f_lsn = {}
        for n, l, f, e in configurations[self.symbol][1]:

            if l not in self.f_lsn:
                self.f_lsn[l] = [[] for s in range(self.nspins)]
            if self.nspins == 1:
                self.f_lsn[l][0].append(f)
            else:
                # Use Hund's rule:
                f0 = min(f, 2 * l + 1)
                self.f_lsn[l][0].append(f0)
                self.f_lsn[l][1].append(f - f0)

    def add(self, n, l, df=+1, s=None):
        """Add (remove) electrons."""
        if s is None:
            if self.nspins == 1:
                s = 0
            else:
                self.add(n, l, 0.5 * df, 0)
                self.add(n, l, 0.5 * df, 1)
                return

        if l not in self.f_lsn:
            self.f_lsn[l] = [[] for x in range(self.nspins)]

        f_n = self.f_lsn[l][s]
        if len(f_n) < n - l:
            f_n.extend([0] * (n - l - len(f_n)))
        f_n[n - l - 1] += df

    def initialize(self,
                   ngpts=1000,
                   rcut=50.0,
                   alpha1=0.01,
                   alpha2=None,
                   ngauss=50,
                   eps=1.0e-7):
        """Initialize basis sets and radial grid.

        ngpts: int
            Number of grid points for radial grid.
        rcut: float
            Cutoff for radial grid.
        alpha1: float
            Smallest exponent for gaussian.
        alpha2: float
            Largest exponent for gaussian.
        ngauss: int
            Number of gaussians.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        if alpha2 is None:
            alpha2 = 50.0 * self.Z**2

        self.gd = GridDescriptor(r1=1 / alpha2**0.5 / 50, rN=rcut, N=ngpts)
        self.log('Grid points:     %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %
                 ((self.gd.N, ) + tuple(self.gd.r_g[[0, 1, 2, -2, -1]])))

        # Distribute exponents between alpha1 and alpha2:
        alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)
        self.log('Exponents:       %d (%.3f, %.3f, ..., %.3f, %.3f)' %
                 ((ngauss, ) + tuple(alpha_B[[0, 1, -2, -1]])))

        # Maximum l value:
        lmax = max(self.f_lsn.keys())

        self.channels = []
        nb_l = []
        if not self.dirac:
            for l in range(lmax + 1):
                basis = GaussianBasis(l, alpha_B, self.gd, eps)
                nb_l.append(len(basis))
                for s in range(self.nspins):
                    self.channels.append(Channel(l, s, self.f_lsn[l][s],
                                                 basis))
        else:
            for K in range(1, lmax + 2):
                leff = (K**2 - (self.Z / c)**2)**0.5 - 1
                basis = GaussianBasis(leff, alpha_B, self.gd, eps)
                nb_l.append(len(basis))
                for k, l in [(-K, K - 1), (K, K)]:
                    if l > lmax:
                        continue
                    f_n = self.f_lsn[l][0]
                    j = abs(k) - 0.5
                    f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)
                    self.channels.append(DiracChannel(k, f_n, basis))

        self.log('Basis functions: %s (%s)' %
                 (', '.join([str(nb)
                             for nb in nb_l]), ', '.join('spdf'[:lmax + 1])))

        self.vr_sg = self.gd.zeros(self.nspins)
        self.vr_sg[:] = -self.Z

    def solve(self):
        """Diagonalize Schrödinger equation."""
        self.eeig = 0.0
        for channel in self.channels:
            channel.solve(self.vr_sg[channel.s])
            self.eeig += channel.get_eigenvalue_sum()

    def calculate_density(self):
        """Calculate elctron density and kinetic energy."""
        self.n_sg = self.gd.zeros(self.nspins)
        for channel in self.channels:
            self.n_sg[channel.s] += channel.calculate_density()

    def calculate_electrostatic_potential(self):
        """Calculate electrostatic potential and energy."""
        n_g = self.n_sg.sum(0)
        self.vHr_g = self.gd.poisson(n_g)
        self.eH = 0.5 * self.gd.integrate(n_g * self.vHr_g, -1)
        self.eZ = -self.Z * self.gd.integrate(n_g, -1)

    def calculate_xc_potential(self):
        self.vxc_sg = self.gd.zeros(self.nspins)
        self.exc = self.xc.calculate_spherical(self.gd, self.n_sg, self.vxc_sg)

    def step(self):
        self.solve()
        self.calculate_density()
        self.calculate_electrostatic_potential()
        self.calculate_xc_potential()
        self.vr_sg = self.vxc_sg * self.gd.r_g
        self.vr_sg += self.vHr_g
        self.vr_sg -= self.Z
        self.ekin = (self.eeig - self.gd.integrate(
            (self.vr_sg * self.n_sg).sum(0), -1))

    def run(self, mix=0.4, maxiter=117, dnmax=1e-9):
        if self.channels is None:
            self.initialize()

        dn = self.Z
        pb = ProgressBar(log(dnmax / dn), 0, 53, self.fd)
        self.log()

        for iter in range(maxiter):
            if iter > 1:
                self.vr_sg *= mix
                self.vr_sg += (1 - mix) * vr_old_sg
                dn = self.gd.integrate(abs(self.n_sg - n_old_sg).sum(0))
                pb(log(dnmax / dn))
                if dn <= dnmax:
                    break

            vr_old_sg = self.vr_sg
            n_old_sg = self.n_sg
            self.step()

        self.summary()
        if dn > dnmax:
            raise RuntimeError('Did not converge!')

    def summary(self):
        self.write_states()
        self.write_energies()

    def write_states(self):
        self.log('\n state  occupation         eigenvalue          <r>')
        if self.dirac:
            self.log(' nl(j)               [Hartree]        [eV]    [Bohr]')
        else:
            self.log(' nl                  [Hartree]        [eV]    [Bohr]')
        self.log('=====================================================')
        states = []
        for ch in self.channels:
            for n, f in enumerate(ch.f_n):
                states.append((ch.e_n[n], ch, n))
        states.sort()
        for e, ch, n in states:
            name = str(n + ch.l + 1) + ch.name
            if self.nspins == 2:
                name += '(%s)' % '+-'[ch.s]
            n_g = ch.calculate_density(n)
            rave = self.gd.integrate(n_g, 1)
            self.log(' %-7s  %6.3f %13.6f  %13.5f %6.3f' %
                     (name, ch.f_n[n], e, e * units.Hartree, rave))
        self.log('=====================================================')

    def write_energies(self):
        self.log('\nEnergies:          [Hartree]           [eV]')
        self.log('============================================')
        for text, e in [
            ('kinetic      ', self.ekin), ('coulomb (e-e)', self.eH),
            ('coulomb (e-n)', self.eZ), ('xc           ', self.exc),
            ('total        ', self.ekin + self.eH + self.eZ + self.exc)
        ]:
            self.log(' %s %+13.6f  %+13.5f' % (text, e, units.Hartree * e))
        self.log('============================================')

    def get_channel(self, l=None, s=0, k=None):
        if self.dirac:
            for channel in self.channels:
                if channel.k == k:
                    return channel
        else:
            for channel in self.channels:
                if channel.l == l and channel.s == s:
                    return channel
        raise ValueError

    def get_orbital(self, n, l=None, s=0, k=None):
        channel = self.get_channel(l, s, k)
        return channel.basis.expand(channel.C_nb[n])

    def plot_wave_functions(self, rc=4.0):
        import matplotlib.pyplot as plt
        colors = 'krgbycm'
        for ch in self.channels:
            for n in range(len(ch.f_n)):
                fr_g = ch.basis.expand(ch.C_nb[n]) * self.gd.r_g
                name = str(n + ch.l + 1) + ch.name
                lw = 2
                if self.nspins == 2:
                    name += '(%s)' % '+-'[ch.s]
                    if ch.s == 1:
                        lw = 1
                if self.dirac and ch.k > 0:
                    lw = 1
                ls = ['-', '--', '-.', ':'][ch.l]
                n_g = ch.calculate_density(n)
                rave = self.gd.integrate(n_g, 1)
                gave = self.gd.get_index(rave)
                fr_g *= cmp(fr_g[gave], 0)
                plt.plot(self.gd.r_g,
                         fr_g,
                         ls=ls,
                         lw=lw,
                         color=colors[n + ch.l],
                         label=name)
        plt.legend(loc='best')
        plt.axis(xmax=rc)
        plt.show()

    def logarithmic_derivative(self, l, energies, rcut):
        vr = splrep(self.gd.r_g, self.vr_sg[0])

        def v(r):
            return splev(r, vr) / r

        def f(y, r, e):
            if r == 0:
                return [y[1], -2.0]
            return [y[1], 2 * (v(r) - e) * y[0]]

        logderivs = []
        for e in energies:
            u, dudr = odeint(f, [0, 1], [0, rcut], (e, ))[1, :]
            logderivs.append(dudr / u)
        return logderivs
Example #13
0
class AllElectronAtom:
    def __init__(self, symbol, xc='LDA', spinpol=False, dirac=False,
                 log=sys.stdout):
        """All-electron calculation for spherically symmetric atom.

        symbol: str (or int)
            Chemical symbol (or atomic number).
        xc: str
            Name of XC-functional.
        spinpol: bool
            If true, do spin-polarized calculation.  Default is spin-paired.
        dirac: bool
            Solve Dirac equation instead of Schrödinger equation.
        log: stream
            Text output."""

        if isinstance(symbol, int):
            symbol = chemical_symbols[symbol]
        self.symbol = symbol
        self.Z = atomic_numbers[symbol]

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

        self.dirac = bool(dirac)

        if isinstance(xc, str):
            self.xc = XC(xc)
        else:
            self.xc = xc

        if log is None:
            log = devnull
        self.fd = log

        self.vr_sg = None  # potential * r
        self.n_sg = 0.0    # density
        self.gd = None     # radial grid descriptor

        # Energies:
        self.ekin = None
        self.eeig = None
        self.eH = None
        self.eZ = None

        self.channels = None

        self.initialize_configuration()

        self.log('Z:              ', self.Z)
        self.log('Name:           ', atomic_names[self.Z])
        self.log('Symbol:         ', symbol)
        self.log('XC-functional:  ', self.xc.name)
        self.log('Equation:       ', ['Schrödinger', 'Dirac'][self.dirac])

    def log(self, *args, **kwargs):
        self.fd.write(kwargs.get('sep', ' ').join([str(arg) for arg in args]) +
                      kwargs.get('end', '\n'))

    def initialize_configuration(self):
        self.f_lsn = {}
        for n, l, f, e in configurations[self.symbol][1]:
            
            if l not in self.f_lsn:
                self.f_lsn[l] = [[] for s in range(self.nspins)]
            if self.nspins == 1:
                self.f_lsn[l][0].append(f)
            else:
                # Use Hund's rule:
                f0 = min(f, 2 * l + 1)
                self.f_lsn[l][0].append(f0)
                self.f_lsn[l][1].append(f - f0)

    def add(self, n, l, df=+1, s=None):
        """Add (remove) electrons."""
        if s is None:
            if self.nspins == 1:
                s = 0
            else:
                self.add(n, l, 0.5 * df, 0)
                self.add(n, l, 0.5 * df, 1)
                return
            
        if l not in self.f_lsn:
            self.f_lsn[l] = [[] for x in range(self.nspins)]
            
        f_n = self.f_lsn[l][s]
        if len(f_n) < n - l:
            f_n.extend([0] * (n - l - len(f_n)))
        f_n[n - l - 1] += df

    def initialize(self, ngpts=1000, rcut=50.0,
                   alpha1=0.01, alpha2=None, ngauss=50,
                   eps=1.0e-7):
        """Initialize basis sets and radial grid.

        ngpts: int
            Number of grid points for radial grid.
        rcut: float
            Cutoff for radial grid.
        alpha1: float
            Smallest exponent for gaussian.
        alpha2: float
            Largest exponent for gaussian.
        ngauss: int
            Number of gaussians.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        if alpha2 is None:
            alpha2 = 50.0 * self.Z**2

        self.gd = GridDescriptor(r1=1 / alpha2**0.5 / 50, rN=rcut, N=ngpts)
        self.log('Grid points:     %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %
                 ((self.gd.N,) + tuple(self.gd.r_g[[0, 1, 2, -2, -1]])))

        # Distribute exponents between alpha1 and alpha2:
        alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)
        self.log('Exponents:       %d (%.3f, %.3f, ..., %.3f, %.3f)' %
                 ((ngauss,) + tuple(alpha_B[[0, 1, -2, -1]])))

        # Maximum l value:
        lmax = max(self.f_lsn.keys())

        self.channels = []
        nb_l = []
        if not self.dirac:
            for l in range(lmax + 1):
                basis = GaussianBasis(l, alpha_B, self.gd, eps)
                nb_l.append(len(basis))
                for s in range(self.nspins):
                    self.channels.append(Channel(l, s, self.f_lsn[l][s],
                                                 basis))
        else:
            for K in range(1, lmax + 2):
                leff = (K**2 - (self.Z / c)**2)**0.5 - 1
                basis = GaussianBasis(leff, alpha_B, self.gd, eps)
                nb_l.append(len(basis))
                for k, l in [(-K, K - 1), (K, K)]:
                    if l > lmax:
                        continue
                    f_n = self.f_lsn[l][0]
                    j = abs(k) - 0.5
                    f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)
                    self.channels.append(DiracChannel(k, f_n, basis))

        self.log('Basis functions: %s (%s)' %
                 (', '.join([str(nb) for nb in nb_l]),
                  ', '.join('spdf'[:lmax + 1])))

        self.vr_sg = self.gd.zeros(self.nspins)
        self.vr_sg[:] = -self.Z

    def solve(self):
        """Diagonalize Schrödinger equation."""
        self.eeig = 0.0
        for channel in self.channels:
            channel.solve(self.vr_sg[channel.s])
            self.eeig += channel.get_eigenvalue_sum()

    def calculate_density(self):
        """Calculate elctron density and kinetic energy."""
        self.n_sg = self.gd.zeros(self.nspins)
        for channel in self.channels:
            self.n_sg[channel.s] += channel.calculate_density()

    def calculate_electrostatic_potential(self):
        """Calculate electrostatic potential and energy."""
        n_g = self.n_sg.sum(0)
        self.vHr_g = self.gd.poisson(n_g)        
        self.eH = 0.5 * self.gd.integrate(n_g * self.vHr_g, -1)
        self.eZ = -self.Z * self.gd.integrate(n_g, -1)
        
    def calculate_xc_potential(self):
        self.vxc_sg = self.gd.zeros(self.nspins)
        self.exc = self.xc.calculate_spherical(self.gd, self.n_sg, self.vxc_sg)

    def step(self):
        self.solve()
        self.calculate_density()
        self.calculate_electrostatic_potential()
        self.calculate_xc_potential()
        self.vr_sg = self.vxc_sg * self.gd.r_g
        self.vr_sg += self.vHr_g
        self.vr_sg -= self.Z
        self.ekin = (self.eeig -
                     self.gd.integrate((self.vr_sg * self.n_sg).sum(0), -1))
        
    def run(self, mix=0.4, maxiter=117, dnmax=1e-9):
        if self.channels is None:
            self.initialize()

        dn = self.Z
        pb = ProgressBar(log(dnmax / dn), 0, 53, self.fd)
        self.log()
        
        for iter in range(maxiter):
            if iter > 1:
                self.vr_sg *= mix
                self.vr_sg += (1 - mix) * vr_old_sg
                dn = self.gd.integrate(abs(self.n_sg - n_old_sg).sum(0))
                pb(log(dnmax / dn))
                if dn <= dnmax:
                    break

            vr_old_sg = self.vr_sg
            n_old_sg = self.n_sg
            self.step()

        self.summary()
        if dn > dnmax:
            raise RuntimeError('Did not converge!')

    def summary(self):
        self.write_states()
        self.write_energies()

    def write_states(self):
        self.log('\n state  occupation         eigenvalue          <r>')
        if self.dirac:
            self.log(' nl(j)               [Hartree]        [eV]    [Bohr]')
        else:
            self.log(' nl                  [Hartree]        [eV]    [Bohr]')
        self.log('=====================================================')
        states = []
        for ch in self.channels:
            for n, f in enumerate(ch.f_n):
                states.append((ch.e_n[n], ch, n))
        states.sort()
        for e, ch, n in states:
            name = str(n + ch.l + 1) + ch.name
            if self.nspins == 2:
                name += '(%s)' % '+-'[ch.s]    
            n_g = ch.calculate_density(n)
            rave = self.gd.integrate(n_g, 1)
            self.log(' %-7s  %6.3f %13.6f  %13.5f %6.3f' %
                     (name, ch.f_n[n], e, e * units.Hartree, rave))
        self.log('=====================================================')

    def write_energies(self):
        self.log('\nEnergies:          [Hartree]           [eV]')
        self.log('============================================')
        for text, e in [('kinetic      ', self.ekin),
                        ('coulomb (e-e)', self.eH),
                        ('coulomb (e-n)', self.eZ),
                        ('xc           ', self.exc),
                        ('total        ',
                         self.ekin + self.eH + self.eZ + self.exc)]:
            self.log(' %s %+13.6f  %+13.5f' % (text, e, units.Hartree * e))
        self.log('============================================')

    def get_channel(self, l=None, s=0, k=None):
        if self.dirac:
            for channel in self.channels:
                if channel.k == k:
                    return channel
        else:
            for channel in self.channels:
                if channel.l == l and channel.s == s:
                    return channel
        raise ValueError

    def get_orbital(self, n, l=None, s=0, k=None):
        channel = self.get_channel(l, s, k)
        return channel.basis.expand(channel.C_nb[n])

    def plot_wave_functions(self, rc=4.0):
        import matplotlib.pyplot as plt
        colors = 'krgbycm'
        for ch in self.channels:
            for n in range(len(ch.f_n)):
                fr_g = ch.basis.expand(ch.C_nb[n]) * self.gd.r_g
                name = str(n + ch.l + 1) + ch.name
                lw = 2
                if self.nspins == 2:
                    name += '(%s)' % '+-'[ch.s]    
                    if ch.s == 1:
                        lw = 1
                if self.dirac and ch.k > 0:
                    lw = 1
                ls = ['-', '--', '-.', ':'][ch.l]
                n_g = ch.calculate_density(n)
                rave = self.gd.integrate(n_g, 1)
                gave = self.gd.get_index(rave)
                fr_g *= cmp(fr_g[gave], 0)
                plt.plot(self.gd.r_g, fr_g,
                         ls=ls, lw=lw, color=colors[n + ch.l], label=name)
        plt.legend(loc='best')
        plt.axis(xmax=rc)
        plt.show()

    def logarithmic_derivative(self, l, energies, rcut):
        vr = splrep(self.gd.r_g, self.vr_sg[0])

        def v(r):
            return splev(r, vr) / r
        
        def f(y, r, e):
            if r == 0:
                return [y[1], -2.0]
            return [y[1], 2 * (v(r) - e) * y[0]]

        logderivs = []
        for e in energies:
            u, dudr = odeint(f, [0, 1], [0, rcut], (e,))[1, :]
            logderivs.append(dudr / u)
        return logderivs
Example #14
0
class C_XC(Contribution):
    def __init__(self, nlfunc, weight, functional = 'LDA'):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional

    def get_name(self):
        return 'XC'

    def get_desc(self):
        return "("+self.functional+")"
        
    def initialize(self):
        self.xc = XC(self.functional)
        self.vt_sg = self.nlfunc.finegd.empty(self.nlfunc.nspins)
        self.e_g = self.nlfunc.finegd.empty()

    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional) 
        self.v_g = np.zeros(self.ae.N)

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        v_g += self.weight * self.vt_sg[0]
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, na_g, va_g, nb_g, vb_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.get_energy_and_potential(na_g, self.vt_sg[0], nb_g, self.vt_sg[1], e_g=self.e_g)
        va_g += self.weight * self.vt_sg[0]
        vb_g += self.weight * self.vt_sg[1]
        e_g += (self.weight * self.e_g).ravel()

    def calculate_energy_and_derivatives(self, D_sp, H_sp, a):
        # Get the XC-correction instance
        c = self.nlfunc.setups[a].xc_correction

        assert self.nlfunc.nspins == 1
        D_p = D_sp[0]
        dEdD_p = H_sp[0][:]
        D_Lq = dot3(c.B_pqL.T, D_p)
        n_Lg = np.dot(D_Lq, c.n_qg)
        n_Lg[0] += c.nc_g * sqrt(4 * pi)
        nt_Lg = np.dot(D_Lq, c.nt_qg)
        nt_Lg[0] += c.nct_g * sqrt(4 * pi)
        dndr_Lg = np.zeros((c.Lmax, c.ng))
        dntdr_Lg = np.zeros((c.Lmax, c.ng))
        for L in range(c.Lmax):
            c.rgd.derivative(n_Lg[L], dndr_Lg[L])
            c.rgd.derivative(nt_Lg[L], dntdr_Lg[L])
                                                            
        E = 0
        vt_g = np.zeros(c.ng)
        v_g = np.zeros(c.ng)
        e_g = np.zeros(c.ng)
        y = 0
        for w, Y_L in zip(weight_n, c.Y_nL):
            A_Li = rnablaY_nLv[y, :c.Lmax]
            a1x_g = np.dot(A_Li[:, 0], n_Lg)
            a1y_g = np.dot(A_Li[:, 1], n_Lg)
            a1z_g = np.dot(A_Li[:, 2], n_Lg)
            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
            a2_g[1:] /= c.rgd.r_g[1:]**2
            a2_g[0] = a2_g[1]
            a1_g = np.dot(Y_L, dndr_Lg)
            a2_g += a1_g**2
            deda2_g = np.zeros(c.ng)  

            v_g[:] = 0.0
            e_g[:] = 0.0
            n_g = np.dot(Y_L, n_Lg)
            self.xc.kernel.calculate(e_g, n_g.reshape((1, -1)),
                                     v_g.reshape((1, -1)),
                                     a2_g.reshape((1, -1)),
                                     deda2_g.reshape((1, -1)))
            
            E += w * np.dot(e_g, c.rgd.dv_g)
            x_g = -2.0 * deda2_g * c.rgd.dv_g * a1_g
            c.rgd.derivative2(x_g, x_g)
            x_g += v_g * c.rgd.dv_g
            dEdD_p += self.weight * w * np.dot(dot3(c.B_pqL, Y_L),
                                  np.dot(c.n_qg, x_g))
            x_g = 8.0 * pi * deda2_g * c.rgd.dr_g
            dEdD_p += w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 0]),
                                  np.dot(c.n_qg, x_g * a1x_g))
            dEdD_p += w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 1]),
                                  np.dot(c.n_qg, x_g * a1y_g))
            dEdD_p += w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 2]),
                                  np.dot(c.n_qg, x_g * a1z_g))

            n_g = np.dot(Y_L, nt_Lg)
            a1x_g = np.dot(A_Li[:, 0], nt_Lg)
            a1y_g = np.dot(A_Li[:, 1], nt_Lg)
            a1z_g = np.dot(A_Li[:, 2], nt_Lg)
            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
            a2_g[1:] /= c.rgd.r_g[1:]**2
            a2_g[0] = a2_g[1]
            a1_g = np.dot(Y_L, dntdr_Lg)
            a2_g += a1_g**2
            v_g = np.zeros(c.ng)
            e_g = np.zeros(c.ng)
            deda2_g = np.zeros(c.ng)

            v_g[:] = 0.0
            e_g[:] = 0.0
            self.xc.kernel.calculate(e_g, n_g.reshape((1, -1)),
                                     v_g.reshape((1, -1)),
                                     a2_g.reshape((1, -1)),
                                     deda2_g.reshape((1, -1)))

            E -= w * np.dot(e_g, c.dv_g)
            x_g = -2.0 * deda2_g * c.dv_g * a1_g
            c.rgd.derivative2(x_g, x_g)
            x_g += v_g * c.dv_g

            B_Lqp = c.B_pqL.T
            dEdD_p -= w * np.dot(dot3(c.B_pqL, Y_L),
                                  np.dot(c.nt_qg, x_g))
            x_g = 8.0 * pi * deda2_g * c.rgd.dr_g
            dEdD_p -= w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 0]),
                                  np.dot(c.nt_qg, x_g * a1x_g))
            dEdD_p -= w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 1]),
                                  np.dot(c.nt_qg, x_g * a1y_g))
            
            dEdD_p -= w * np.dot(dot3(c.B_pqL,
                                       A_Li[:, 2]),
                                  np.dot(c.nt_qg, x_g * a1z_g))
            
            y += 1
        
        return (E) * self.weight

    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.n.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        v_g += self.weight * self.v_g
        return self.weight * Exc

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        Exc = self.xc.calculate_spherical(self.ae.rgd,
                                          self.ae.nt.reshape((1, -1)),
                                          self.v_g.reshape((1, -1)))
        vt_g += self.weight * self.v_g
        return self.weight * Exc

    def initialize_from_atomic_orbitals(self, basis_functions):
        # LDA needs only density, which is already initialized
        pass

    def add_extra_setup_data(self, dict):
        # LDA has not any special data
        pass

    def write(self, writer, natoms):
        # LDA has not any special data to be written
        pass

    def read(self, reader):
        # LDA has not any special data to be read
        pass
Example #15
0
class C_GLLBScr(Contribution):
    def __init__(self, nlfunc, weight, functional='GGA_X_B88'):
        Contribution.__init__(self, nlfunc, weight)
        self.functional = functional
        self.old_coeffs = None
        self.iter = 0
        
    def get_name(self):
        return 'SCREENING'

    def get_desc(self):
        return '(' + self.functional + ')'
        
    # Initialize GLLBScr functional
    def initialize_1d(self):
        self.ae = self.nlfunc.ae
        self.xc = XC(self.functional)
        self.v_g = np.zeros(self.ae.N)
        self.e_g = np.zeros(self.ae.N)

    # Calcualte the GLLB potential and energy 1d
    def add_xc_potential_and_energy_1d(self, v_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.n.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        v_g += 2 * self.weight * self.e_g / (self.ae.n + 1e-10)
        Exc = self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)
        return Exc

    def initialize(self):
        self.occupations = self.nlfunc.occupations
        self.xc = XC(self.functional)
        self.vt_sg = self.nlfunc.finegd.empty(self.nlfunc.nspins)
        self.e_g = self.nlfunc.finegd.empty()#.ravel()

    def get_coefficient_calculator(self):
        return self

    def f(self, f):
        return sqrt(f)
    
    def get_coefficients_1d(self, smooth=False, lumo_perturbation = False):
        homo_e = max( [ np.where(f>1e-3, e, -1000) for f,e in zip(self.ae.f_j, self.ae.e_j)]) 
        if not smooth:
            if lumo_perturbation:
                lumo_e = min( [ np.where(f<1e-3, e, 1000) for f,e in zip(self.ae.f_j, self.ae.e_j)])
                return np.array([ f * K_G * (self.f( max(0, lumo_e - e)) - self.f(max(0, homo_e -e)))
                                        for e,f in zip(self.ae.e_j, self.ae.f_j) ])
            else:
                return np.array([ f * K_G * (self.f( max(0, homo_e - e)))
                                   for e,f in zip(self.ae.e_j, self.ae.f_j) ])
        else:
            return [ [ f * K_G * self.f( max(0, homo_e - e))
                    for e,f in zip(e_n, f_n) ]
                     for e_n, f_n in zip(self.ae.e_ln, self.ae.f_ln) ]
        

    def get_coefficients_by_kpt(self, kpt_u, lumo_perturbation=False, homolumo=None):
        if kpt_u[0].psit_nG is None or isinstance(kpt_u[0].psit_nG,
                                                  TarFileReference): 
            return None

        if homolumo == None:
            e_ref, e_ref_lumo = self.occupations.get_homo_lumo(self.nlfunc.wfs)
        else:
            e_ref, e_ref_lumo = homolumo

        # The parameter ee might sometimes be set to small thereshold value to
        # achieve convergence on systems with degenerate H**O.
        if len(kpt_u) > 1:
            ee = 0.0
        else:
            ee = 0.1 / 27.21

        if lumo_perturbation:
            return [np.array([
                f * K_G * (self.f( np.where(e_ref_lumo - e>ee, e_ref_lumo-e,0))
                         -self.f( np.where(e_ref      - e>ee, e_ref-e,0)))
                     for e, f in zip(kpt.eps_n, kpt.f_n) ])
                     for kpt in kpt_u ]
            
            
        else:
            coeff = [ np.array([ f * K_G * self.f( np.where(e_ref - e>ee, e_ref-e,0))
                     for e, f in zip(kpt.eps_n, kpt.f_n) ])
                     for kpt in kpt_u ]
            if self.old_coeffs is None:
                self.old_coeffs = coeff
            else:
                # Mix the coefficients with 25%
                mix = 0.25
                self.old_coeffs = [ (1-mix) * old + mix * new for new, old in zip(coeff, self.old_coeffs) ]
            return self.old_coeffs
        

    def calculate_spinpaired(self, e_g, n_g, v_g):
        self.e_g[:] = 0.0
        self.vt_sg[:] = 0.0
        self.xc.calculate(self.nlfunc.finegd, n_g[None, ...], self.vt_sg,
                          self.e_g)
        v_g += self.weight * 2 * self.e_g / (n_g + 1e-10)
        e_g += self.weight * self.e_g

    def calculate_spinpolarized(self, e_g, na_g, va_g, nb_g, vb_g, 
                                a2_g=None, aa2_g=None, ab2_g=None, deda2_g=None,
                                dedaa2_g=None, dedab2_g=None):
        raise NotImplementedError

    def calculate_energy_and_derivatives(self, D_sp, H_sp, a):
        # Get the XC-correction instance
        c = self.nlfunc.setups[a].xc_correction

        assert self.nlfunc.nspins == 1

        D_p = D_sp[0]
        dEdD_p = H_sp[0][:]
        D_Lq = np.dot(c.B_pqL.T, D_p)
        n_Lg = np.dot(D_Lq, c.n_qg)
        n_Lg[0] += c.nc_g * sqrt(4 * pi)
        nt_Lg = np.dot(D_Lq, c.nt_qg)
        nt_Lg[0] += c.nct_g * sqrt(4 * pi)
        dndr_Lg = np.zeros((c.Lmax, c.ng))
        dntdr_Lg = np.zeros((c.Lmax, c.ng))
        for L in range(c.Lmax):
            c.rgd.derivative(n_Lg[L], dndr_Lg[L])
            c.rgd.derivative(nt_Lg[L], dntdr_Lg[L])
        E = 0
        vt_g = np.zeros(c.ng)
        v_g = np.zeros(c.ng)
        e_g = np.zeros(c.ng)
        deda2_g = np.zeros(c.ng)
        for y, (w, Y_L) in enumerate(zip(weight_n, c.Y_nL)):
            # Cut gradient releated coefficient to match the setup's Lmax
            A_Li = rnablaY_nLv[y, :c.Lmax]

            # Expand pseudo density
            nt_g = np.dot(Y_L, nt_Lg)

            # Expand pseudo density gradient
            a1x_g = np.dot(A_Li[:, 0], nt_Lg)
            a1y_g = np.dot(A_Li[:, 1], nt_Lg)
            a1z_g = np.dot(A_Li[:, 2], nt_Lg)
            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
            a2_g[1:] /= c.rgd.r_g[1:]**2
            a2_g[0] = a2_g[1]
            a1_g = np.dot(Y_L, dntdr_Lg)
            a2_g += a1_g**2
            
            vt_g[:] = 0.0
            e_g[:] = 0.0
            # Calculate pseudo GGA energy density (potential is discarded)
            self.xc.kernel.calculate(e_g, nt_g.reshape((1, -1)),
                                     vt_g.reshape((1, -1)),
                                     a2_g.reshape((1, -1)),
                                     deda2_g.reshape((1, -1)))

            # Calculate pseudo GLLB-potential from GGA-energy density
            vt_g[:] = 2 * e_g / (nt_g + 1e-10)

            
            dEdD_p -= self.weight * w * np.dot(np.dot(c.B_pqL, Y_L),
                                  np.dot(c.nt_qg, vt_g * c.rgd.dv_g))

            E -= w * np.dot(e_g, c.rgd.dv_g)
            
            # Expand density
            n_g = np.dot(Y_L, n_Lg)

            # Expand density gradient
            a1x_g = np.dot(A_Li[:, 0], n_Lg)
            a1y_g = np.dot(A_Li[:, 1], n_Lg)
            a1z_g = np.dot(A_Li[:, 2], n_Lg)
            a2_g = a1x_g**2 + a1y_g**2 + a1z_g**2
            a2_g[1:] /= c.rgd.r_g[1:]**2
            a2_g[0] = a2_g[1]
            a1_g = np.dot(Y_L, dndr_Lg)
            a2_g += a1_g**2
            
            v_g[:] = 0.0
            e_g[:] = 0.0
            # Calculate GGA energy density (potential is discarded)
            self.xc.kernel.calculate(e_g, n_g.reshape((1, -1)),
                                     v_g.reshape((1, -1)),
                                     a2_g.reshape((1, -1)),
                                     deda2_g.reshape((1, -1)))

            # Calculate GLLB-potential from GGA-energy density
            v_g[:] = 2 * e_g / (n_g + 1e-10)
            
            dEdD_p += self.weight * w * np.dot(np.dot(c.B_pqL, Y_L),
                                  np.dot(c.n_qg, v_g * c.rgd.dv_g))
            E += w * np.dot(e_g, c.rgd.dv_g)
            
        return (E) * self.weight

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        self.v_g[:] = 0.0
        self.e_g[:] = 0.0
        self.xc.calculate_spherical(self.ae.rgd, self.ae.nt.reshape((1, -1)),
                                    self.v_g.reshape((1, -1)), self.e_g)
        vt_g += 2 * self.weight * self.e_g / (self.ae.nt + 1e-10)
        return self.weight * np.sum(self.e_g * self.ae.rgd.dv_g)

    def initialize_from_atomic_orbitals(self, basis_functions):
        # GLLBScr needs only density which is already initialized
        pass
        
    def add_extra_setup_data(self, dict):
        # GLLBScr has not any special data
        pass

    def read(self, reader):
        # GLLBScr has no special data to be read
        pass

    def write(self, writer, natoms):
        # GLLBScr has no special data to be written
        pass