diff options
Diffstat (limited to 'atom.py')
| -rw-r--r-- | atom.py | 158 |
1 files changed, 158 insertions, 0 deletions
@@ -0,0 +1,158 @@ +from math import * +from scipy.special import erf, sph_harm, eval_genlaguerre + +import numpy as np + +import element + + +# Atomic Coulomb potential, modelled using HGH pseudopotentials. +# See https://arxiv.org/abs/cond-mat/9803286 for details. +class Atom: + def __init__(self, s, symbol, pos): + self.symbol = symbol + self.params = element.Element(symbol) # get element data + self.position = pos + + # Set up local grid in spherical coordinates + self.r = np.zeros((s.Nres, s.Nres, s.Nres)) + self.theta = np.zeros((s.Nres, s.Nres, s.Nres)) + self.phi = np.zeros((s.Nres, s.Nres, s.Nres)) + for i in range(s.Nres): + for j in range(s.Nres): + for k in range(s.Nres): + x = s.rx[i, j, k] + y = s.ry[i, j, k] + z = s.rz[i, j, k] + + r = sqrt((x - pos[0])**2 + (y - pos[1])**2 + (z - pos[2])**2) + if r == 0: + theta = 0 + phi = 0 + else: + theta = acos((z - pos[2]) / r) + phi = atan2((y - pos[1]), (x - pos[0])) + + self.r[i, j, k] = r + self.theta[i, j, k] = theta + self.phi[i, j, k] = phi + + # Initialize and cache orbitals for nonlocal potential contribution + self.orbitals = {} + for l in range(0, 3): + self.orbitals[l] = {} + for i in range(1, 4): + p = self.get_projector(l, i) + self.orbitals[l][i] = {} + for m in range(-l, l + 1): + Y = self.cubic_harmonic(m, l) + self.orbitals[l][i][m] = p * Y + + + # Local potential representing short-range interactions. + # This is Eq.(1), but without the long-range erf-term: + def get_local_potential(self): + V = np.zeros_like(self.r) + + rr = self.r / self.params.rloc # relative r + for i in range(4): + V += self.params.Cloc[i] * rr**(2 * i) * np.exp(-0.5 * rr**2) + + return V + + + # The local erf-term decays as 1/r, which is problematic for periodicity. + # Instead, we replace it with a Hartree potential from a Gaussian charge. + # This is exact: multiply first term of Eq.(5) by g^2 (from Poisson). + def get_local_density(self): + alpha = 0.5 / self.params.rloc**2 + rho = 4 * alpha**(3/2) * np.exp(-alpha * self.r**2) / sqrt(pi) + + # Normalize in spherical coordinates: 4pi int rho(r) r^2 dr = Zion + rho *= self.params.Zion / (4 * pi) + return rho + + + # Given a psi in real space, apply nonlocal potential by integration + # Vpsi = sum_{lijm} |p_i> h_{ij} <p_j | psi> + def apply_nonlocal_potential(self, psi): + Vpsi = np.zeros_like(psi) + + # Note that these if-statements provide a significant speedup + for l in range(0, 3): + if self.params.rl[l] != 0: + for i in range(1, 4): + for j in range(1, 4): + if self.params.h[l][i][j] != 0: + for m in range(-l, l + 1): + # Inner product normalization is handled by caller + ip = np.vdot(self.orbitals[l][j][m], psi) + Vpsi += self.orbitals[l][i][m] * self.params.h[l][i][j] * ip + + return Vpsi + + + # Calculate projectors p_i^l(r) for nonlocal potential, see Eq.(3) + def get_projector(self, l, i): + p = np.zeros_like(self.r) + + if self.params.rl[l] != 0: + p = sqrt(2) * self.r**(l + 2 * (i - 1)) + p *= np.exp(-0.5 * (self.r / self.params.rl[l])**2) + p /= self.params.rl[l]**(l + 0.5 * (4 * i - 1)) + p /= sqrt(gamma(l + 0.5 * (4 * i - 1))) + + return p + + + # Guess this atom's electron density (to initialize main DFT loop) + def guess_density(self): + # We choose a spherical gaussian with standard deviation sigma + sigma = 1 + rho = (sqrt(2 / pi) / sigma**3) * np.exp(- 0.5 * self.r**2 / sigma**2) + + # Normalize in spherical coordinates: 4pi int rho(r) r^2 dr = Zion + rho /= self.params.Zion / (4 * pi) + return rho + + + # Guess this atom's eigenstates (to initialize eigensolver). + # We simply choose hydrogen orbitals, without any modifications: + def guess_orbitals(self, Norb): + psis = [] + for n in range(1, 4): + rho = 2 * self.r / n + for l in range(0, n): + for m in range(-l, l + 1): + psi = self.cubic_harmonic(m, l) + psi *= np.exp(-0.5 * rho) * rho**l + psi *= eval_genlaguerre(n - l - 1, 2 * l + 1, rho) + psis.append(psi) + + return psis[0 : Norb] + + + # Use cubic harmonics instead of spherical harmonics. This makes no + # physical difference, but it avoids confusion when visualizing. + def cubic_harmonic(self, m, l): + # Note that phi and theta are swapped according to SciPy + Yp = sph_harm( abs(m), l, self.phi, self.theta) + Yn = sph_harm(-abs(m), l, self.phi, self.theta) + + f = np.zeros_like(self.r) + if m == -3: + f = (1 / sqrt(2)) * (Yn - Yp) + if m == -2: + f = (1j / sqrt(2)) * (Yn - Yp) + if m == -1: + f = (1 / sqrt(2)) * (Yn - Yp) + if m == 0: + f = Yp + if m == 1: + f = (1j / sqrt(2)) * (Yn + Yp) + if m == 2: + f = (1 / sqrt(2)) * (Yn + Yp) + if m == 3: + f = (1j / sqrt(2)) * (Yn + Yp) + + return np.real(f) |