class PoissonSolver:
def __init__(self, nn=3, relax='J', eps=2e-10):
self.relax = relax
self.nn = nn
self.eps = eps
self.charged_periodic_correction = None
self.maxiter = 1000
# 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)
def get_method(self):
return ['Gauss-Seidel', 'Jacobi'][self.relax_method - 1]
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, allocate=False)]
self.B = LaplaceB(gd, allocate=False)
else:
self.operators = [Laplace(gd, scale, self.nn, allocate=False)]
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 < 4:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1, allocate=False))
self.interpolators.append(Transformer(gd2, gd, allocate=False))
self.restrictors.append(Transformer(gd, gd2, allocate=False))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
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
for obj in self.operators + self.interpolators + self.restrictors:
obj.allocate()
if self.B is not None:
self.B.allocate()
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 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
if self.charged_periodic_correction is None:
print "+-----------------------------------------------------+"
print "| Calculating charged periodic correction using the |"
print "| Ewald potential from a lattice of probe charges in |"
print "| a homogenous background density |"
print "+-----------------------------------------------------+"
self.charged_periodic_correction = madelung(self.gd.cell_cv)
print "Potential shift will be ", \
self.charged_periodic_correction , "Ha."
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
phi -= charge * self.charged_periodic_correction
iters = self.solve_neutral(phi, rho - background, eps=eps)
phi += charge * self.charged_periodic_correction
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
raise NotImplementedError
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:
charge = np.sum(rho.ravel()) * self.dv
print 'CHARGE, eps:', charge, eps
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)
for i, obj in enumerate(self.restrictors + self.interpolators):
obj.estimate_memory(mem.subnode('Transformer %d' % i))
for i, operator in enumerate(self.operators):
name = operator.__class__.__name__
operator.estimate_memory(
mem.subnode('Operator %d [%s]' % (i, name)))
if self.B is not None:
name = self.B.__class__.__name__
self.B.estimate_memory(mem.subnode('B [%s]' % name))
def __repr__(self):
template = 'PoissonSolver(relax=\'%s\', nn=%s, eps=%e)'
representation = template % (self.relax, repr(self.nn), self.eps)
return representation
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
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
class PoissonSolver:
def __init__(self, nn=3, relax='J', eps=2e-10):
self.relax = relax
self.nn = nn
self.eps = eps
self.charged_periodic_correction = None
self.maxiter = 1000
# 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)
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, allocate=False)]
self.B = LaplaceB(gd, allocate=False)
else:
self.operators = [Laplace(gd, scale, self.nn, allocate=False)]
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 < 4:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1, allocate=False))
self.interpolators.append(Transformer(gd2, gd, allocate=False))
self.restrictors.append(Transformer(gd, gd2, allocate=False))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
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
for obj in self.operators + self.interpolators + self.restrictors:
obj.allocate()
if self.B is not None:
self.B.allocate()
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 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
if self.charged_periodic_correction is None:
print "+-----------------------------------------------------+"
print "| Calculating charged periodic correction using the |"
print "| Ewald potential from a lattice of probe charges in |"
print "| a homogenous background density |"
print "+-----------------------------------------------------+"
self.charged_periodic_correction = madelung(self.gd.cell_cv)
print "Potential shift will be ", \
self.charged_periodic_correction , "Ha."
# Set initial guess for potential
if zero_initial_phi:
phi[:] = 0.0
else:
phi -= charge * self.charged_periodic_correction
iters = self.solve_neutral(phi, rho - background, eps=eps)
phi += charge * self.charged_periodic_correction
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
raise NotImplementedError
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:
charge = np.sum(rho.ravel()) * self.dv
print 'CHARGE, eps:', charge, eps
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)
for i, obj in enumerate(self.restrictors + self.interpolators):
obj.estimate_memory(mem.subnode('Transformer %d' % i))
for i, operator in enumerate(self.operators):
name = operator.__class__.__name__
operator.estimate_memory(mem.subnode('Operator %d [%s]' % (i,
name)))
if self.B is not None:
name = self.B.__class__.__name__
self.B.estimate_memory(mem.subnode('B [%s]' % name))
def __repr__(self):
template = 'PoissonSolver(relax=\'%s\', nn=%s, eps=%e)'
representation = template % (self.relax, repr(self.nn), self.eps)
return representation
class FDPoissonSolver(BasePoissonSolver):
def __init__(self,
nn=3,
relax='J',
eps=2e-10,
maxiter=1000,
remove_moment=None,
use_charge_center=False):
super(FDPoissonSolver,
self).__init__(eps=eps,
remove_moment=remove_moment,
use_charge_center=use_charge_center)
self.relax = relax
self.nn = nn
self.charged_periodic_correction = None
self.maxiter = maxiter
# 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
self._initialized = False
def todict(self):
d = super(FDPoissonSolver, self).todict()
d.update({'name': 'fd', 'nn': self.nn, 'relax': self.relax})
return d
def get_stencil(self):
return self.nn
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
"""Instantiate and return a Laplace operator
Allows subclasses to change the Laplace operator
"""
return Laplace(gd, scale, n, dtype)
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 = [self.create_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(self.create_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)
self._initialized = False
# The Gaussians depend on the grid as well so we have to 'unload' them
if hasattr(self, 'rho_gauss'):
del self.rho_gauss
del self.phi_gauss
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,
' Max iterations: %d' % self.maxiter
])
lines.append(super(FDPoissonSolver, self).get_description())
return '\n'.join(lines)
def _init(self):
if self._initialized:
return
# 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
self._initialized = True
def solve_neutral(self, phi, rho, eps=2e-10, timer=None):
self._init()
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"""
self._init()
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 = 'FDPoissonSolver(relax=\'%s\', nn=%s, eps=%e)'
representation = template % (self.relax, repr(self.nn), self.eps)
return representation