def load(self, method):
"""Make sure all necessary attributes have been initialized"""
assert method in ('real', 'recip_gauss', 'recip_ewald'),\
str(method) + ' is an invalid method name,\n' +\
'use either real, recip_gauss, or recip_ewald'
if method.startswith('recip'):
if self.gd.comm.size > 1:
raise RuntimeError("Cannot do parallel FFT, use method='real'")
if not hasattr(self, 'k2'):
self.k2, self.N3 = construct_reciprocal(self.gd)
if method.endswith('ewald') and not hasattr(self, 'ewald'):
# cutoff radius
assert self.gd.orthogonal
rc = 0.5 * np.average(self.gd.cell_cv.diagonal())
# ewald potential: 1 - cos(k rc)
self.ewald = (np.ones(self.gd.n_c) -
np.cos(np.sqrt(self.k2) * rc))
# lim k -> 0 ewald / k2
self.ewald[0, 0, 0] = 0.5 * rc**2
elif method.endswith('gauss') and not hasattr(self, 'ng'):
gauss = Gaussian(self.gd)
self.ng = gauss.get_gauss(0) / sqrt(4 * pi)
self.vg = gauss.get_gauss_pot(0) / sqrt(4 * pi)
else: # method == 'real'
if not hasattr(self, 'solve'):
if self.poisson is not None:
self.solve = self.poisson.solve
else:
solver = PoissonSolver(nn=2)
solver.set_grid_descriptor(self.gd)
solver.initialize(load_gauss=True)
self.solve = solver.solve
def initialize(self, gd, load_gauss=False):
# XXX this won't work now, but supposedly this class will be deprecated
# in favour of FFTPoissonSolver, no?
self.gd = gd
if self.gd.comm.size > 1:
raise RuntimeError('Cannot do parallel FFT.')
self.k2, self.N3 = construct_reciprocal(self.gd)
if load_gauss:
gauss = Gaussian(self.gd)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def initialize(self, gd, load_gauss=False):
# XXX this won't work now, but supposedly this class will be deprecated
# in favour of FFTPoissonSolver, no?
self.gd = gd
if self.gd.comm.size > 1:
raise RuntimeError('Cannot do parallel FFT.')
self.k2, self.N3 = construct_reciprocal(self.gd)
if load_gauss:
gauss = Gaussian(self.gd)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def load_moment_corrections_gauss(self):
if not hasattr(self, 'gauss_i'):
self.gauss_i = []
mask_ir = []
r_ir = []
self.mom_ij = []
for rmom in self.moment_corrections:
if rmom['center'] is None:
center = None
else:
center = np.array(rmom['center'])
mom_j = rmom['moms']
gauss = Gaussian(self.gd, center=center)
self.gauss_i.append(gauss)
r_ir.append(gauss.r.ravel())
mask_ir.append(self.gd.zeros(dtype=int).ravel())
self.mom_ij.append(mom_j)
r_ir = np.array(r_ir)
mask_ir = np.array(mask_ir)
Ni = r_ir.shape[0]
Nr = r_ir.shape[1]
for r in range(Nr):
i = np.argmin(r_ir[:, r])
mask_ir[i, r] = 1
self.mask_ig = []
for i in range(Ni):
mask_r = mask_ir[i]
mask_g = mask_r.reshape(self.gd.n_c)
self.mask_ig.append(mask_g)
def load_gauss(self, center=None):
"""Load compensating charge distribution for charged systems.
See Appendix B of
A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014).
"""
# XXX Check if update is needed (dielectric changed)?
epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
gauss = Gaussian(self.gd, center=center)
rho_g = gauss.get_gauss(0)
phi_g = gauss.get_gauss_pot(0)
x, y, z = gauss.xyz
fac = 2. * np.sqrt(gauss.a) * np.exp(-gauss.a * gauss.r2)
fac /= np.sqrt(np.pi) * gauss.r2
fac -= erf(np.sqrt(gauss.a) * gauss.r) / (gauss.r2 * gauss.r)
fac *= 2.0 * 1.7724538509055159
dx_phi_g = fac * x
dy_phi_g = fac * y
dz_phi_g = fac * z
sp = dx_phi_g * dx_epsr + dy_phi_g * dy_epsr + dz_phi_g * dz_epsr
rho = epsr * rho_g - 1. / (4. * np.pi) * sp
invnorm = np.sqrt(4. * np.pi) / self.gd.integrate(rho)
self.phi_gauss = phi_g * invnorm
self.rho_gauss = rho * invnorm
def load_gauss(self):
if not hasattr(self, 'rho_gauss'):
gauss = Gaussian(self.gd)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
class FDPoissonSolver:
def __init__(self, nn=3, relax='J', eps=2e-10, maxiter=1000,
remove_moment=None, use_charge_center=False):
self.relax = relax
self.nn = nn
self.eps = eps
self.charged_periodic_correction = None
self.maxiter = maxiter
self.remove_moment = remove_moment
self.use_charge_center = use_charge_center
# Relaxation method
if relax == 'GS':
# Gauss-Seidel
self.relax_method = 1
elif relax == 'J':
# Jacobi
self.relax_method = 2
else:
raise NotImplementedError('Relaxation method %s' % relax)
self.description = None
def todict(self):
return {'name': 'fd', 'nn': self.nn, 'relax': self.relax,
'eps': self.eps, 'remove_moment': self.remove_moment}
def get_stencil(self):
return self.nn
def set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale)]
self.B = LaplaceB(gd)
else:
self.operators = [Laplace(gd, scale, self.nn)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 8:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1))
self.interpolators.append(Transformer(gd2, gd))
self.restrictors.append(Transformer(gd, gd2))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
if self.operators[-1].gd.N_c.max() > 36:
# Try to warn exactly once no matter how one uses the solver.
if gd.comm.parent is None:
warn = (gd.comm.rank == 0)
else:
warn = (gd.comm.parent.rank == 0)
if warn:
warntxt = '\n'.join([POISSON_GRID_WARNING, '',
self.get_description()])
else:
warntxt = ('Poisson warning from domain rank %d'
% self.gd.comm.rank)
# Warn from all ranks to avoid deadlocks.
warnings.warn(warntxt, stacklevel=2)
def get_description(self):
name = {1: 'Gauss-Seidel', 2: 'Jacobi'}[self.relax_method]
coarsest_grid = self.operators[-1].gd.N_c
coarsest_grid_string = ' x '.join([str(N) for N in coarsest_grid])
assert self.levels + 1 == len(self.operators)
lines = ['%s solver with %d multi-grid levels'
% (name, self.levels + 1),
' Coarsest grid: %s points' % coarsest_grid_string]
if coarsest_grid.max() > 24:
# This friendly warning has lower threshold than the big long
# one that we print when things are really bad.
lines.extend([' Warning: Coarse grid has more than 24 points.',
' More multi-grid levels recommended.'])
lines.extend([' Stencil: %s' % self.operators[0].description,
' Tolerance: %e' % self.eps,
' Max iterations: %d' % self.maxiter])
if self.remove_moment is not None:
lines.append(' Remove moments up to L=%d' % self.remove_moment)
if self.use_charge_center:
lines.append(' Compensate for charged system using center of '
'majority charge')
return '\n'.join(lines)
def initialize(self, load_gauss=False):
# Should probably be renamed allocate
gd = self.gd
self.rhos = [gd.empty()]
self.phis = [None]
self.residuals = [gd.empty()]
for level in range(self.levels):
gd2 = gd.coarsen()
self.phis.append(gd2.empty())
self.rhos.append(gd2.empty())
self.residuals.append(gd2.empty())
gd = gd2
assert len(self.phis) == len(self.rhos)
level += 1
assert level == self.levels
self.step = 0.66666666 / self.operators[0].get_diagonal_element()
self.presmooths[level] = 8
self.postsmooths[level] = 8
if load_gauss:
self.load_gauss()
def load_gauss(self, center=None):
if not hasattr(self, 'rho_gauss') or center is not None:
gauss = Gaussian(self.gd, center=center)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def solve(self, phi, rho, charge=None, eps=None, maxcharge=1e-6,
zero_initial_phi=False):
assert np.all(phi.shape == self.gd.n_c)
assert np.all(rho.shape == self.gd.n_c)
if eps is None:
eps = self.eps
actual_charge = self.gd.integrate(rho)
background = (actual_charge / self.gd.dv /
self.gd.get_size_of_global_array().prod())
if self.remove_moment:
assert not self.gd.pbc_c.any()
if not hasattr(self, 'gauss'):
self.gauss = Gaussian(self.gd)
rho_neutral = rho.copy()
phi_cor_L = []
for L in range(self.remove_moment):
phi_cor_L.append(self.gauss.remove_moment(rho_neutral, L))
# Remove multipoles for better initial guess
for phi_cor in phi_cor_L:
phi -= phi_cor
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing multipoles
for phi_cor in phi_cor_L:
phi += phi_cor
return niter
if charge is None:
charge = actual_charge
if abs(charge) <= maxcharge:
# System is charge neutral. Use standard solver
return self.solve_neutral(phi, rho - background, eps=eps)
elif abs(charge) > maxcharge and self.gd.pbc_c.all():
# System is charged and periodic. Subtract a homogeneous
# background charge
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
iters = self.solve_neutral(phi, rho - background, eps=eps)
return iters
elif abs(charge) > maxcharge and not self.gd.pbc_c.any():
# The system is charged and in a non-periodic unit cell.
# Determine the potential by 1) subtract a gaussian from the
# density, 2) determine potential from the neutralized density
# and 3) add the potential from the gaussian density.
# Load necessary attributes
# use_charge_center: The monopole will be removed at the
# center of the majority charge, which prevents artificial
# dipoles.
# Due to the shape of the Gaussian and it's Fourier-Transform,
# the Gaussian representing the charge should stay at least
# 7 gpts from the borders - see:
# https://listserv.fysik.dtu.dk/pipermail/gpaw-developers/2015-July/005806.html
if self.use_charge_center:
charge_sign = actual_charge / abs(actual_charge)
rho_sign = rho * charge_sign
rho_sign[np.where(rho_sign < 0)] = 0
absolute_charge = self.gd.integrate(rho_sign)
center = - self.gd.calculate_dipole_moment(rho_sign) \
/ absolute_charge
border_offset = np.inner(self.gd.h_cv, np.array((7, 7, 7)))
borders = np.inner(self.gd.h_cv, self.gd.N_c)
borders -= border_offset
if np.any(center > borders) or np.any(center < border_offset):
raise RuntimeError(
'Poisson solver: center of charge outside' + \
' borders - please increase box')
center[np.where(center > borders)] = borders
self.load_gauss(center=center)
else:
self.load_gauss()
# Remove monopole moment
q = actual_charge / np.sqrt(4 * pi) # Monopole moment
rho_neutral = rho - q * self.rho_gauss # neutralized density
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
axpy(-q, self.phi_gauss, phi) # phi -= q * self.phi_gauss
# Determine potential from neutral density using standard solver
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing monopole
axpy(q, self.phi_gauss, phi) # phi += q * self.phi_gauss
return niter
else:
# System is charged with mixed boundaryconditions
msg = ('Charged systems with mixed periodic/zero'
' boundary conditions')
raise NotImplementedError(msg)
def solve_neutral(self, phi, rho, eps=2e-10):
self.phis[0] = phi
if self.B is None:
self.rhos[0][:] = rho
else:
self.B.apply(rho, self.rhos[0])
niter = 1
maxiter = self.maxiter
while self.iterate2(self.step) > eps and niter < maxiter:
niter += 1
if niter == maxiter:
msg = 'Poisson solver did not converge in %d iterations!' % maxiter
raise PoissonConvergenceError(msg)
# Set the average potential to zero in periodic systems
if np.alltrue(self.gd.pbc_c):
phi_ave = self.gd.comm.sum(np.sum(phi.ravel()))
N_c = self.gd.get_size_of_global_array()
phi_ave /= np.product(N_c)
phi -= phi_ave
return niter
def iterate2(self, step, level=0):
"""Smooths the solution in every multigrid level"""
residual = self.residuals[level]
if level < self.levels:
self.operators[level].relax(self.relax_method,
self.phis[level],
self.rhos[level],
self.presmooths[level],
self.weights[level])
self.operators[level].apply(self.phis[level], residual)
residual -= self.rhos[level]
self.restrictors[level].apply(residual,
self.rhos[level + 1])
self.phis[level + 1][:] = 0.0
self.iterate2(4.0 * step, level + 1)
self.interpolators[level].apply(self.phis[level + 1], residual)
self.phis[level] -= residual
self.operators[level].relax(self.relax_method,
self.phis[level],
self.rhos[level],
self.postsmooths[level],
self.weights[level])
if level == 0:
self.operators[level].apply(self.phis[level], residual)
residual -= self.rhos[level]
error = self.gd.comm.sum(np.dot(residual.ravel(),
residual.ravel())) * self.gd.dv
# How about this instead:
# error = self.gd.comm.max(abs(residual).max())
return error
def estimate_memory(self, mem):
# XXX Memory estimate works only for J and GS, not FFT solver
# Poisson solver appears to use same amount of memory regardless
# of whether it's J or GS, which is a bit strange
gdbytes = self.gd.bytecount()
nbytes = -gdbytes # No phi on finest grid, compensate ahead
for level in range(self.levels):
nbytes += 3 * gdbytes # Arrays: rho, phi, residual
gdbytes //= 8
mem.subnode('rho, phi, residual [%d levels]' % self.levels, nbytes)
def __repr__(self):
template = 'PoissonSolver(relax=\'%s\', nn=%s, eps=%e)'
representation = template % (self.relax, repr(self.nn), self.eps)
return representation
# Initialize classes
a = 20 # Size of cell
N = 48 # Number of grid points
Nc = (N, N, N) # Number of grid points along each axis
gd = GridDescriptor(Nc, (a, a, a), 0) # Grid-descriptor object
solver = PoissonSolver(nn=3) # Numerical poisson solver
solver.set_grid_descriptor(gd)
solver.initialize()
solve = solver.solve
xyz, r2 = coordinates(gd) # Matrix with the square of the radial coordinate
print(r2.shape)
r = np.sqrt(r2) # Matrix with the values of the radial coordinate
nH = np.exp(-2 * r) / pi # Density of the hydrogen atom
gauss = Gaussian(gd) # An instance of Gaussian
# /------------------------------------------------\
# | Check if Gaussian densities are made correctly |
# \------------------------------------------------/
for gL in range(2, 9):
g = gauss.get_gauss(gL) # a gaussian of gL'th order
print("\nGaussian of order", gL)
for mL in range(9):
m = gauss.get_moment(g, mL) # the mL'th moment of g
print(" %s'th moment = %2.6f" % (mL, m))
equal(m, gL == mL, 1e-4)
# Check the moments of the constructed 1s density
print("\nDensity of Hydrogen atom")
for L in range(4):
def load_gauss(self, center=None):
if not hasattr(self, 'rho_gauss') or center is not None:
gauss = Gaussian(self.gd, center=center)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def solve(self, phi, rho, charge=None, eps=None, maxcharge=1e-6,
zero_initial_phi=False):
assert np.all(phi.shape == self.gd.n_c)
assert np.all(rho.shape == self.gd.n_c)
if eps is None:
eps = self.eps
actual_charge = self.gd.integrate(rho)
background = (actual_charge / self.gd.dv /
self.gd.get_size_of_global_array().prod())
if self.remove_moment:
assert not self.gd.pbc_c.any()
if not hasattr(self, 'gauss'):
self.gauss = Gaussian(self.gd)
rho_neutral = rho.copy()
phi_cor_L = []
for L in range(self.remove_moment):
phi_cor_L.append(self.gauss.remove_moment(rho_neutral, L))
# Remove multipoles for better initial guess
for phi_cor in phi_cor_L:
phi -= phi_cor
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing multipoles
for phi_cor in phi_cor_L:
phi += phi_cor
return niter
if charge is None:
charge = actual_charge
if abs(charge) <= maxcharge:
# System is charge neutral. Use standard solver
return self.solve_neutral(phi, rho - background, eps=eps)
elif abs(charge) > maxcharge and self.gd.pbc_c.all():
# System is charged and periodic. Subtract a homogeneous
# background charge
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
iters = self.solve_neutral(phi, rho - background, eps=eps)
return iters
elif abs(charge) > maxcharge and not self.gd.pbc_c.any():
# The system is charged and in a non-periodic unit cell.
# Determine the potential by 1) subtract a gaussian from the
# density, 2) determine potential from the neutralized density
# and 3) add the potential from the gaussian density.
# Load necessary attributes
# use_charge_center: The monopole will be removed at the
# center of the majority charge, which prevents artificial
# dipoles.
# Due to the shape of the Gaussian and it's Fourier-Transform,
# the Gaussian representing the charge should stay at least
# 7 gpts from the borders - see:
# https://listserv.fysik.dtu.dk/pipermail/gpaw-developers/2015-July/005806.html
if self.use_charge_center:
charge_sign = actual_charge / abs(actual_charge)
rho_sign = rho * charge_sign
rho_sign[np.where(rho_sign < 0)] = 0
absolute_charge = self.gd.integrate(rho_sign)
center = - self.gd.calculate_dipole_moment(rho_sign) \
/ absolute_charge
border_offset = np.inner(self.gd.h_cv, np.array((7, 7, 7)))
borders = np.inner(self.gd.h_cv, self.gd.N_c)
borders -= border_offset
if np.any(center > borders) or np.any(center < border_offset):
raise RuntimeError(
'Poisson solver: center of charge outside' + \
' borders - please increase box')
center[np.where(center > borders)] = borders
self.load_gauss(center=center)
else:
self.load_gauss()
# Remove monopole moment
q = actual_charge / np.sqrt(4 * pi) # Monopole moment
rho_neutral = rho - q * self.rho_gauss # neutralized density
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
axpy(-q, self.phi_gauss, phi) # phi -= q * self.phi_gauss
# Determine potential from neutral density using standard solver
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing monopole
axpy(q, self.phi_gauss, phi) # phi += q * self.phi_gauss
return niter
else:
# System is charged with mixed boundaryconditions
msg = ('Charged systems with mixed periodic/zero'
' boundary conditions')
raise NotImplementedError(msg)
class BasePoissonSolver(_PoissonSolver):
def __init__(self, eps=None, remove_moment=None, use_charge_center=False):
self.gd = None
self.remove_moment = remove_moment
self.use_charge_center = use_charge_center
self.eps = eps
def todict(self):
d = {'name': 'basepoisson'}
if self.eps is not None:
d['eps'] = self.eps
if self.remove_moment:
d['remove_moment'] = self.remove_moment
if self.use_charge_center:
d['use_charge_center'] = self.use_charge_center
return d
def get_description(self):
# The idea is that the subclass writes a header and main parameters,
# then adds the below string.
lines = []
if self.eps is not None:
lines.append(' Tolerance: %e' % self.eps),
if self.remove_moment is not None:
lines.append(' Remove moments up to L=%d' % self.remove_moment)
if self.use_charge_center:
lines.append(' Compensate for charged system using center of '
'majority charge')
return '\n'.join(lines)
def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False,
timer=NullTimer()):
self._init()
assert np.all(phi.shape == self.gd.n_c)
assert np.all(rho.shape == self.gd.n_c)
if eps is None:
eps = self.eps
actual_charge = self.gd.integrate(rho)
background = (actual_charge / self.gd.dv /
self.gd.get_size_of_global_array().prod())
if self.remove_moment:
assert not self.gd.pbc_c.any()
if not hasattr(self, 'gauss'):
self.gauss = Gaussian(self.gd)
rho_neutral = rho.copy()
phi_cor_L = []
for L in range(self.remove_moment):
phi_cor_L.append(self.gauss.remove_moment(rho_neutral, L))
# Remove multipoles for better initial guess
for phi_cor in phi_cor_L:
phi -= phi_cor
niter = self.solve_neutral(phi, rho_neutral, eps=eps, timer=timer)
# correct error introduced by removing multipoles
for phi_cor in phi_cor_L:
phi += phi_cor
return niter
if charge is None:
charge = actual_charge
if abs(charge) <= maxcharge:
# System is charge neutral. Use standard solver
return self.solve_neutral(phi,
rho - background,
eps=eps,
timer=timer)
elif abs(charge) > maxcharge and self.gd.pbc_c.all():
# System is charged and periodic. Subtract a homogeneous
# background charge
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
iters = self.solve_neutral(phi,
rho - background,
eps=eps,
timer=timer)
return iters
elif abs(charge) > maxcharge and not self.gd.pbc_c.any():
# The system is charged and in a non-periodic unit cell.
# Determine the potential by 1) subtract a gaussian from the
# density, 2) determine potential from the neutralized density
# and 3) add the potential from the gaussian density.
# Load necessary attributes
# use_charge_center: The monopole will be removed at the
# center of the majority charge, which prevents artificial
# dipoles.
# Due to the shape of the Gaussian and it's Fourier-Transform,
# the Gaussian representing the charge should stay at least
# 7 gpts from the borders - see:
# https://listserv.fysik.dtu.dk/pipermail/gpaw-developers/2015-July/005806.html
if self.use_charge_center:
charge_sign = actual_charge / abs(actual_charge)
rho_sign = rho * charge_sign
rho_sign[np.where(rho_sign < 0)] = 0
absolute_charge = self.gd.integrate(rho_sign)
center = -(self.gd.calculate_dipole_moment(rho_sign) /
absolute_charge)
border_offset = np.inner(self.gd.h_cv, np.array((7, 7, 7)))
borders = np.inner(self.gd.h_cv, self.gd.N_c)
borders -= border_offset
if np.any(center > borders) or np.any(center < border_offset):
raise RuntimeError('Poisson solver: '
'center of charge outside borders '
'- please increase box')
center[np.where(center > borders)] = borders
self.load_gauss(center=center)
else:
self.load_gauss()
# Remove monopole moment
q = actual_charge / np.sqrt(4 * pi) # Monopole moment
rho_neutral = rho - q * self.rho_gauss # neutralized density
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
axpy(-q, self.phi_gauss, phi) # phi -= q * self.phi_gauss
# Determine potential from neutral density using standard solver
niter = self.solve_neutral(phi, rho_neutral, eps=eps, timer=timer)
# correct error introduced by removing monopole
axpy(q, self.phi_gauss, phi) # phi += q * self.phi_gauss
return niter
else:
# System is charged with mixed boundaryconditions
msg = ('Charged systems with mixed periodic/zero'
' boundary conditions')
raise NotImplementedError(msg)
def load_gauss(self, center=None):
if not hasattr(self, 'rho_gauss') or center is not None:
gauss = Gaussian(self.gd, center=center)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False):
if eps is None:
eps = self.eps
actual_charge = self.gd.integrate(rho)
background = (actual_charge / self.gd.dv /
self.gd.get_size_of_global_array().prod())
if self.remove_moment:
assert not self.gd.pbc_c.any()
if not hasattr(self, 'gauss'):
self.gauss = Gaussian(self.gd)
rho_neutral = rho.copy()
phi_cor_L = []
for L in range(self.remove_moment):
phi_cor_L.append(self.gauss.remove_moment(rho_neutral, L))
# Remove multipoles for better initial guess
for phi_cor in phi_cor_L:
phi -= phi_cor
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing multipoles
for phi_cor in phi_cor_L:
phi += phi_cor
return niter
if charge is None:
charge = actual_charge
if abs(charge) <= maxcharge:
# System is charge neutral. Use standard solver
return self.solve_neutral(phi, rho - background, eps=eps)
elif abs(charge) > maxcharge and self.gd.pbc_c.all():
# System is charged and periodic. Subtract a homogeneous
# background charge
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
iters = self.solve_neutral(phi, rho - background, eps=eps)
return iters
elif abs(charge) > maxcharge and not self.gd.pbc_c.any():
# The system is charged and in a non-periodic unit cell.
# Determine the potential by 1) subtract a gaussian from the
# density, 2) determine potential from the neutralized density
# and 3) add the potential from the gaussian density.
# Load necessary attributes
self.load_gauss()
# Remove monopole moment
q = actual_charge / np.sqrt(4 * pi) # Monopole moment
rho_neutral = rho - q * self.rho_gauss # neutralized density
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
axpy(-q, self.phi_gauss, phi) # phi -= q * self.phi_gauss
# Determine potential from neutral density using standard solver
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing monopole
axpy(q, self.phi_gauss, phi) # phi += q * self.phi_gauss
return niter
else:
# System is charged with mixed boundaryconditions
msg = 'Charged systems with mixed periodic/zero'
msg += ' boundary conditions'
raise NotImplementedError(msg)
class PoissonSolver:
def __init__(self,
nn=3,
relax='J',
eps=2e-10,
maxiter=1000,
remove_moment=None):
self.relax = relax
self.nn = nn
self.eps = eps
self.charged_periodic_correction = None
self.maxiter = maxiter
self.remove_moment = remove_moment
# Relaxation method
if relax == 'GS':
# Gauss-Seidel
self.relax_method = 1
elif relax == 'J':
# Jacobi
self.relax_method = 2
else:
raise NotImplementedError('Relaxation method %s' % relax)
self.description = None
def get_stencil(self):
return self.nn
def set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
self.dv = gd.dv
gd = self.gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale)]
self.B = LaplaceB(gd)
else:
self.operators = [Laplace(gd, scale, self.nn)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 8:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1))
self.interpolators.append(Transformer(gd2, gd))
self.restrictors.append(Transformer(gd, gd2))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
def get_description(self):
name = {1: 'Gauss-Seidel', 2: 'Jacobi'}[self.relax_method]
coarsest_grid = self.operators[-1].gd.N_c
coarsest_grid_string = ' x '.join([str(N) for N in coarsest_grid])
assert self.levels + 1 == len(self.operators)
lines = [
'%s solver with %d multi-grid levels' % (name, self.levels + 1),
' Coarsest grid: %s points' % coarsest_grid_string
]
if coarsest_grid.max() > 24:
lines.extend([
' Warning: Coarse grid has more than 24 points.',
' More multi-grid levels recommended.'
])
lines.extend([
' Stencil: %s' % self.operators[0].description,
' Tolerance: %e' % self.eps,
' Max iterations: %d' % self.maxiter
])
return '\n'.join(lines)
def initialize(self, load_gauss=False):
# Should probably be renamed allocate
gd = self.gd
self.rhos = [gd.empty()]
self.phis = [None]
self.residuals = [gd.empty()]
for level in range(self.levels):
gd2 = gd.coarsen()
self.phis.append(gd2.empty())
self.rhos.append(gd2.empty())
self.residuals.append(gd2.empty())
gd = gd2
assert len(self.phis) == len(self.rhos)
level += 1
assert level == self.levels
self.step = 0.66666666 / self.operators[0].get_diagonal_element()
self.presmooths[level] = 8
self.postsmooths[level] = 8
if load_gauss:
self.load_gauss()
def load_gauss(self):
if not hasattr(self, 'rho_gauss'):
gauss = Gaussian(self.gd)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False):
if eps is None:
eps = self.eps
actual_charge = self.gd.integrate(rho)
background = (actual_charge / self.gd.dv /
self.gd.get_size_of_global_array().prod())
if self.remove_moment:
assert not self.gd.pbc_c.any()
if not hasattr(self, 'gauss'):
self.gauss = Gaussian(self.gd)
rho_neutral = rho.copy()
phi_cor_L = []
for L in range(self.remove_moment):
phi_cor_L.append(self.gauss.remove_moment(rho_neutral, L))
# Remove multipoles for better initial guess
for phi_cor in phi_cor_L:
phi -= phi_cor
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing multipoles
for phi_cor in phi_cor_L:
phi += phi_cor
return niter
if charge is None:
charge = actual_charge
if abs(charge) <= maxcharge:
# System is charge neutral. Use standard solver
return self.solve_neutral(phi, rho - background, eps=eps)
elif abs(charge) > maxcharge and self.gd.pbc_c.all():
# System is charged and periodic. Subtract a homogeneous
# background charge
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
iters = self.solve_neutral(phi, rho - background, eps=eps)
return iters
elif abs(charge) > maxcharge and not self.gd.pbc_c.any():
# The system is charged and in a non-periodic unit cell.
# Determine the potential by 1) subtract a gaussian from the
# density, 2) determine potential from the neutralized density
# and 3) add the potential from the gaussian density.
# Load necessary attributes
self.load_gauss()
# Remove monopole moment
q = actual_charge / np.sqrt(4 * pi) # Monopole moment
rho_neutral = rho - q * self.rho_gauss # neutralized density
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
axpy(-q, self.phi_gauss, phi) # phi -= q * self.phi_gauss
# Determine potential from neutral density using standard solver
niter = self.solve_neutral(phi, rho_neutral, eps=eps)
# correct error introduced by removing monopole
axpy(q, self.phi_gauss, phi) # phi += q * self.phi_gauss
return niter
else:
# System is charged with mixed boundaryconditions
msg = 'Charged systems with mixed periodic/zero'
msg += ' boundary conditions'
raise NotImplementedError(msg)
def solve_neutral(self, phi, rho, eps=2e-10):
self.phis[0] = phi
if self.B is None:
self.rhos[0][:] = rho
else:
self.B.apply(rho, self.rhos[0])
niter = 1
maxiter = self.maxiter
while self.iterate2(self.step) > eps and niter < maxiter:
niter += 1
if niter == maxiter:
msg = 'Poisson solver did not converge in %d iterations!' % maxiter
raise PoissonConvergenceError(msg)
# Set the average potential to zero in periodic systems
if np.alltrue(self.gd.pbc_c):
phi_ave = self.gd.comm.sum(np.sum(phi.ravel()))
N_c = self.gd.get_size_of_global_array()
phi_ave /= np.product(N_c)
phi -= phi_ave
return niter
def iterate(self, step, level=0):
residual = self.residuals[level]
niter = 0
while True:
niter += 1
if level > 0 and niter == 1:
residual[:] = -self.rhos[level]
else:
self.operators[level].apply(self.phis[level], residual)
residual -= self.rhos[level]
error = self.gd.comm.sum(np.vdot(residual, residual))
if niter == 1 and level < self.levels:
self.restrictors[level].apply(residual, self.rhos[level + 1])
self.phis[level + 1][:] = 0.0
self.iterate(4.0 * step, level + 1)
self.interpolators[level].apply(self.phis[level + 1], residual)
self.phis[level] -= residual
continue
residual *= step
self.phis[level] -= residual
if niter == 2:
break
return error
def iterate2(self, step, level=0):
"""Smooths the solution in every multigrid level"""
residual = self.residuals[level]
if level < self.levels:
self.operators[level].relax(self.relax_method, self.phis[level],
self.rhos[level],
self.presmooths[level],
self.weights[level])
self.operators[level].apply(self.phis[level], residual)
residual -= self.rhos[level]
self.restrictors[level].apply(residual, self.rhos[level + 1])
self.phis[level + 1][:] = 0.0
self.iterate2(4.0 * step, level + 1)
self.interpolators[level].apply(self.phis[level + 1], residual)
self.phis[level] -= residual
self.operators[level].relax(self.relax_method, self.phis[level],
self.rhos[level], self.postsmooths[level],
self.weights[level])
if level == 0:
self.operators[level].apply(self.phis[level], residual)
residual -= self.rhos[level]
error = self.gd.comm.sum(np.dot(residual.ravel(),
residual.ravel())) * self.dv
return error
def estimate_memory(self, mem):
# XXX Memory estimate works only for J and GS, not FFT solver
# Poisson solver appears to use same amount of memory regardless
# of whether it's J or GS, which is a bit strange
gdbytes = self.gd.bytecount()
nbytes = -gdbytes # No phi on finest grid, compensate ahead
for level in range(self.levels):
nbytes += 3 * gdbytes # Arrays: rho, phi, residual
gdbytes //= 8
mem.subnode('rho, phi, residual [%d levels]' % self.levels, nbytes)
def __repr__(self):
template = 'PoissonSolver(relax=\'%s\', nn=%s, eps=%e)'
representation = template % (self.relax, repr(self.nn), self.eps)
return representation
def load_gauss(self):
if not hasattr(self, 'rho_gauss'):
gauss = Gaussian(self.gd)
self.rho_gauss = gauss.get_gauss(0)
self.phi_gauss = gauss.get_gauss_pot(0)
psolvers = (WeightedFDPoissonSolver, ADM12PoissonSolver,
PolarizationPoissonSolver)
# test neutral system with constant permittivity
parprint('neutral, constant permittivity')
epsinf = 80.
eps = gd.zeros()
eps.fill(epsinf)
qs = (-1., 1.)
shifts = (-1., 1.)
rho = gd.zeros()
phi_expected = gd.zeros()
for q, shift in zip(qs, shifts):
gauss_norm = q / np.sqrt(4 * np.pi)
gauss = Gaussian(gd, center=(box / 2. + shift) * np.ones(3) / Bohr)
rho += gauss_norm * gauss.get_gauss(0)
phi_expected += gauss_norm * gauss.get_gauss_pot(0) / epsinf
for ps in psolvers:
phi = solve(ps, eps, rho)
parprint(ps, np.abs(phi - phi_expected).max())
equal(phi, phi_expected, 1e-3)
# test charged system with constant permittivity
parprint('charged, constant permittivity')
epsinf = 80.
eps = gd.zeros()
eps.fill(epsinf)
q = -2.
gauss_norm = q / np.sqrt(4 * np.pi)
from gpaw.utilities.gauss import Gaussian from gpaw.grid_descriptor import GridDescriptor from gpaw.poisson import PoissonSolver # Initialize classes a = 20.0 # Size of cell inv_width = 19 # inverse width of the gaussian N = 48 # Number of grid points center_of_charge = (a / 2, a / 2, 3 * a / 4) # off center charge Nc = (N, N, N) # Number of grid points along each axis gd = GridDescriptor(Nc, (a, a, a), 0) # Grid-descriptor object solver = PoissonSolver(nn=3, use_charge_center=True) solver.set_grid_descriptor(gd) solver.initialize() gauss = Gaussian(gd, a=inv_width, center=center_of_charge) test_poisson = Gaussian(gd, a=inv_width, center=center_of_charge) # /-------------------------------------------------\ # | Check if Gaussian potentials are made correctly | # \-------------------------------------------------/ # Array for storing the potential pot = gd.zeros(dtype=float, global_array=False) solver.load_gauss() vg = test_poisson.get_gauss_pot(0) # Get analytic functions ng = gauss.get_gauss(0) # vg = solver.phi_gauss # Solve potential numerically niter = solver.solve(pot, ng, charge=1.0, zero_initial_phi=False)
from gpaw.helmholtz import HelmholtzSolver, ScreenedPoissonGaussian # Initialize classes a = 20 # Size of cell inv_width = 31 # inverse width of the gaussian N = 48 # Number of grid points coupling = -0.4 # dampening Nc = (N, N, N) # Number of grid points along each axis gd = GridDescriptor(Nc, (a, a, a), 0) # Grid-descriptor object solver = HelmholtzSolver(k2=coupling, nn=3) # Numerical poisson solver # solver = PoissonSolver(nn=3) # Numerical poisson solver # solver = HelmholtzSolver(0.16) # Numerical poisson solver solver.set_grid_descriptor(gd) solver.initialize() xyz, r2 = coordinates(gd) # Matrix with the square of the radial coordinate gauss = Gaussian(gd, a=inv_width) # An instance of Gaussian test_screened_poisson = ScreenedPoissonGaussian(gd, a=inv_width) # /-------------------------------------------------\ # | Check if Gaussian potentials are made correctly | # \-------------------------------------------------/ # Array for storing the potential pot = gd.zeros(dtype=float, global_array=False) solver.load_gauss() vg = test_screened_poisson.get_phi(-coupling) # esp. for dampening # Get analytic functions ng = gauss.get_gauss(0) # vg = solver.phi_gauss # Solve potential numerically niter = solver.solve(pot, ng, charge=None, zero_initial_phi=True)
return np.sqrt(np.sum(a.ravel()**2)) / len(a.ravel())
# Initialize classes
a = 20 # Size of cell
N = 48 # Number of grid points
Nc = (N, N, N) # Number of grid points along each axis
gd = GridDescriptor(Nc, (a, a, a), 0) # Grid-descriptor object
solver = PoissonSolver(nn=3) # Numerical poisson solver
solver.set_grid_descriptor(gd)
solve = solver.solve
xyz, r2 = coordinates(gd) # Matrix with the square of the radial coordinate
print(r2.shape)
r = np.sqrt(r2) # Matrix with the values of the radial coordinate
nH = np.exp(-2 * r) / pi # Density of the hydrogen atom
gauss = Gaussian(gd) # An instance of Gaussian
# /------------------------------------------------\
# | Check if Gaussian densities are made correctly |
# \------------------------------------------------/
for gL in range(2, 9):
g = gauss.get_gauss(gL) # a gaussian of gL'th order
print('\nGaussian of order', gL)
for mL in range(9):
m = gauss.get_moment(g, mL) # the mL'th moment of g
print(' %s\'th moment = %2.6f' % (mL, m))
equal(m, gL == mL, 1e-4)
# Check the moments of the constructed 1s density
print('\nDensity of Hydrogen atom')
for L in range(4):
from gpaw.helmholtz import HelmholtzSolver, ScreenedPoissonGaussian # Initialize classes a = 20 # Size of cell inv_width = 31 # inverse width of the gaussian N = 48 # Number of grid points coupling = -0.4 # dampening Nc = (N, N, N) # Number of grid points along each axis gd = GridDescriptor(Nc, (a,a,a), 0) # Grid-descriptor object solver = HelmholtzSolver(k2=coupling, nn=3) # Numerical poisson solver # solver = PoissonSolver(nn=3) # Numerical poisson solver # solver = HelmholtzSolver(0.16) # Numerical poisson solver solver.set_grid_descriptor(gd) solver.initialize() xyz, r2 = coordinates(gd) # Matrix with the square of the radial coordinate gauss = Gaussian(gd, a=inv_width) # An instance of Gaussian test_screened_poisson = ScreenedPoissonGaussian(gd, a=inv_width) # /-------------------------------------------------\ # | Check if Gaussian potentials are made correctly | # \-------------------------------------------------/ # Array for storing the potential pot = gd.zeros(dtype=float, global_array=False) solver.load_gauss() vg = test_screened_poisson.get_phi(-coupling) # esp. for dampening # Get analytic functions ng = gauss.get_gauss(0) # vg = solver.phi_gauss # Solve potential numerically niter = solver.solve(pot, ng, charge=None, zero_initial_phi=True)
