Coverage for gpaw/solvation/poisson.py: 91%
160 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import warnings
3import numpy as np
4from scipy.special import erf
6from gpaw.poisson import BasePoissonSolver, PoissonSolver, FDPoissonSolver
7from gpaw.fd_operators import Laplace, Gradient
8from gpaw.wfd_operators import WeightedFDOperator
9from gpaw.utilities.gauss import Gaussian
10from gpaw import PoissonConvergenceError
11from gpaw.utilities.timing import NullTimer
14class SolvationPoissonSolver(FDPoissonSolver):
15 """Base class for Poisson solvers with spatially varying dielectric.
17 The Poisson equation
18 div(epsilon(r) grad phi(r)) = -4 pi rho(r)
19 is solved.
20 """
22 def __init__(self, nn=3, relax='J', eps=2e-10, maxiter=1000,
23 remove_moment=None, use_charge_center=False):
24 if remove_moment is not None:
25 raise NotImplementedError(
26 'Removing arbitrary multipole moments '
27 'is not implemented for SolvationPoissonSolver!')
28 if nn == 'M':
29 raise NotImplementedError(
30 'Mehrstellen stencil is not implemented '
31 'for SolvationPoissonSolver!')
32 FDPoissonSolver.__init__(self, nn, relax, eps, maxiter, remove_moment,
33 use_charge_center=use_charge_center)
35 def set_dielectric(self, dielectric):
36 """Set the dielectric.
38 Arguments:
39 dielectric -- A Dielectric instance.
40 """
41 self.dielectric = dielectric
43 def load_gauss(self, center=None):
44 """Load compensating charge distribution for charged systems.
46 See Appendix B of
47 A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014).
48 """
49 # XXX Check if update is needed (dielectric changed)?
50 epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
51 gauss = Gaussian(self.gd, center=center)
52 rho_g = gauss.get_gauss(0)
53 phi_g = gauss.get_gauss_pot(0)
54 x, y, z = gauss.xyz
55 fac = 2. * np.sqrt(gauss.a) * np.exp(-gauss.a * gauss.r2)
56 fac /= np.sqrt(np.pi) * gauss.r2
57 fac -= erf(np.sqrt(gauss.a) * gauss.r) / (gauss.r2 * gauss.r)
58 fac *= 2.0 * 1.7724538509055159
59 dx_phi_g = fac * x
60 dy_phi_g = fac * y
61 dz_phi_g = fac * z
62 sp = dx_phi_g * dx_epsr + dy_phi_g * dy_epsr + dz_phi_g * dz_epsr
63 rho = epsr * rho_g - 1. / (4. * np.pi) * sp
64 invnorm = np.sqrt(4. * np.pi) / self.gd.integrate(rho)
65 self.phi_gauss = phi_g * invnorm
66 self.rho_gauss = rho * invnorm
69class WeightedFDPoissonSolver(SolvationPoissonSolver):
70 """Weighted finite difference Poisson solver with dielectric.
72 Following V. M. Sanchez, M. Sued and D. A. Scherlis,
73 J. Chem. Phys. 131, 174108 (2009).
74 """
76 def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
77 operators = [Laplace(gd, scale, n, dtype)]
78 operators += [Gradient(gd, j, scale, n, dtype) for j in (0, 1, 2)]
79 return WeightedFDOperator(operators)
81 def solve(self, phi, rho, charge=None,
82 maxcharge=1e-6,
83 zero_initial_phi=False, timer=None):
84 self._init()
85 if self.gd.pbc_c.all():
86 actual_charge = self.gd.integrate(rho)
87 if abs(actual_charge) > maxcharge:
88 raise NotImplementedError(
89 'charged periodic systems are not implemented')
90 self.restrict_op_weights()
91 ret = FDPoissonSolver.solve(self, phi, rho, charge, maxcharge,
92 zero_initial_phi)
93 return ret
95 def restrict_op_weights(self):
96 """Restric operator weights to coarse grids."""
97 self._init()
98 weights = [self.dielectric.eps_gradeps] + self.op_coarse_weights
99 for i, res in enumerate(self.restrictors):
100 for j in range(4):
101 res.apply(weights[i][j], weights[i + 1][j])
102 self.step = 0.66666666 / self.operators[0].get_diagonal_element()
104 def get_description(self):
105 description = SolvationPoissonSolver.get_description(self)
106 return description.replace(
107 'solver with',
108 'weighted FD solver with dielectric and')
110 def _init(self):
111 if self._initialized:
112 return
113 self.operators[0].set_weights(self.dielectric.eps_gradeps)
114 self.op_coarse_weights = []
115 for operator in self.operators[1:]:
116 weights = [gd.empty() for gd in (operator.gd, ) * 4]
117 self.op_coarse_weights.append(weights)
118 operator.set_weights(weights)
119 return SolvationPoissonSolver._init(self)
122class PolarizationPoissonSolver(BasePoissonSolver):
123 """Poisson solver with dielectric.
125 Calculates the polarization charges first using only the
126 vacuum poisson equation, then solves the vacuum equation
127 with polarization charges.
128 """
130 def __init__(self, nn=3, relax='J', eps=2e-10, maxiter=1000,
131 remove_moment=None, use_charge_center=False,
132 gas_phase_poisson='fast'):
133 self.nn = nn
134 self.eps = eps
135 self.maxiter = maxiter
136 self.phi_tilde = None
138 self.gas_phase_poisson = PoissonSolver(
139 name=gas_phase_poisson, nn=nn, eps=eps,
140 remove_moment=remove_moment,
141 use_charge_center=use_charge_center)
143 def set_dielectric(self, dielectric):
144 """Set the dielectric.
146 Arguments:
147 dielectric -- A Dielectric instance.
148 """
149 self.dielectric = dielectric
151 def get_description(self):
152 return ('PolarizationPoissonSolver based on '
153 + self.gas_phase_poisson.get_description())
155 def set_grid_descriptor(self, gd):
156 self.gd = gd
157 self.gas_phase_poisson.set_grid_descriptor(gd)
159 def solve(self, phi, rho, charge=None,
160 maxcharge=1e-6,
161 zero_initial_phi=False, timer=NullTimer()):
162 # get initial meaningful phi -> only do this if necessary
163 if zero_initial_phi:
164 niter = self.gas_phase_poisson.solve(
165 phi, rho, charge=None, maxcharge=maxcharge,
166 zero_initial_phi=zero_initial_phi, timer=timer)
167 else:
168 niter = 0
170 phi_old = phi.copy()
171 while niter < self.maxiter:
172 rho_mod = self.rho_with_polarization_charge(phi, rho)
173 niter += self.gas_phase_poisson.solve(
174 phi, rho_mod, charge=None, maxcharge=maxcharge,
175 zero_initial_phi=zero_initial_phi, timer=timer)
176 residual = phi - phi_old
177 error = self.gd.comm.sum_scalar(
178 np.dot(residual.ravel(), residual.ravel())) * self.gd.dv
179 if error < self.eps:
180 return niter
181 phi_old = phi.copy()
183 raise PoissonConvergenceError(
184 'PolarizationPoisson solver did not converge in '
185 + f'{niter} iterations!')
187 def rho_with_polarization_charge(self, phi, rho):
188 epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
189 dx_phi = self.gd.empty()
190 dy_phi = self.gd.empty()
191 dz_phi = self.gd.empty()
192 Gradient(self.gd, 0, 1.0, self.nn).apply(phi, dx_phi)
193 Gradient(self.gd, 1, 1.0, self.nn).apply(phi, dy_phi)
194 Gradient(self.gd, 2, 1.0, self.nn).apply(phi, dz_phi)
196 scalar_product = (
197 dx_epsr * dx_phi +
198 dy_epsr * dy_phi +
199 dz_epsr * dz_phi)
201 return (rho + scalar_product / (4. * np.pi)) / epsr
203 def estimate_memory(self, mem):
204 # XXX estimate your own contribution
205 return self.gas_phase_poisson.estimate_memory(mem)
207 def todict(self):
208 my_dict = self.gas_phase_poisson.todict()
209 my_dict['name'] = f"polarization-{my_dict['name']}"
210 return my_dict
213class ADM12PoissonSolver(SolvationPoissonSolver):
214 """Poisson solver with dielectric.
216 Following O. Andreussi, I. Dabo, and N. Marzari,
217 J. Chem. Phys. 136, 064102 (2012).
219 Warning: Not intended for production use, as it is not tested
220 thouroughly!
222 XXX TODO : * Correction for charged systems???
223 * Check: Can the polarization charge introduce a monopole?
224 * Convergence problems depending on eta. Apparently this
225 method works best with FFT as in the original Paper.
226 * Optimize numerics.
227 """
229 def __init__(self, nn=3, relax='J', eps=2e-10, maxiter=1000,
230 remove_moment=None, eta=.6, use_charge_center=False):
231 """Constructor for ADM12PoissonSolver.
233 Additional arguments not present in SolvationPoissonSolver:
234 eta -- linear mixing parameter
235 """
236 adm12_warning = UserWarning(
237 'ADM12PoissonSolver is not tested thoroughly'
238 ' and therefore not recommended for production code!')
239 warnings.warn(adm12_warning)
240 self.eta = eta
241 SolvationPoissonSolver.__init__(
242 self, nn, relax, eps, maxiter, remove_moment,
243 use_charge_center=use_charge_center)
245 def set_grid_descriptor(self, gd):
246 SolvationPoissonSolver.set_grid_descriptor(self, gd)
247 self.gradx = Gradient(gd, 0, 1.0, self.nn)
248 self.grady = Gradient(gd, 1, 1.0, self.nn)
249 self.gradz = Gradient(gd, 2, 1.0, self.nn)
251 def get_description(self):
252 if len(self.operators) == 0:
253 return 'uninitialized ADM12PoissonSolver'
254 else:
255 description = SolvationPoissonSolver.get_description(self)
256 return description.replace(
257 'solver with',
258 'ADM12 solver with dielectric and')
260 def _init(self):
261 if self._initialized:
262 return
263 self.rho_iter = self.gd.zeros()
264 self.d_phi = self.gd.empty()
265 return SolvationPoissonSolver._init(self)
267 def solve(self, phi, rho, charge=None,
268 maxcharge=1e-6,
269 zero_initial_phi=False, timer=None):
270 self._init()
271 if self.gd.pbc_c.all():
272 actual_charge = self.gd.integrate(rho)
273 if abs(actual_charge) > maxcharge:
274 raise NotImplementedError(
275 'charged periodic systems are not implemented')
276 return FDPoissonSolver.solve(
277 self, phi, rho, charge, maxcharge, zero_initial_phi,
278 timer=timer)
280 def solve_neutral(self, phi, rho, timer=None):
281 self._init()
282 self.rho = rho
283 return SolvationPoissonSolver.solve_neutral(self, phi, rho)
285 def iterate2(self, step, level=0):
286 self._init()
287 if level == 0:
288 epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
289 self.gradx.apply(self.phis[0], self.d_phi)
290 sp = dx_epsr * self.d_phi
291 self.grady.apply(self.phis[0], self.d_phi)
292 sp += dy_epsr * self.d_phi
293 self.gradz.apply(self.phis[0], self.d_phi)
294 sp += dz_epsr * self.d_phi
295 self.rho_iter = self.eta / (4. * np.pi) * sp + \
296 (1. - self.eta) * self.rho_iter
297 self.rhos[0][:] = (self.rho_iter + self.rho) / epsr
298 return SolvationPoissonSolver.iterate2(self, step, level)