Coverage for gpaw/new/pw/poisson.py: 54%
240 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import warnings
2from functools import cached_property
3from math import pi
5import numpy as np
6from ase.units import Bohr, Ha
7from gpaw.core import PWArray, PWDesc, UGDesc
8from gpaw.new.poisson import PoissonSolver
9from scipy.sparse.linalg import LinearOperator, cg
10from scipy.special import erf
11from gpaw import get_scipy_version
13if get_scipy_version() >= [1, 14]:
14 RTOL = 'rtol'
15else:
16 RTOL = 'tol'
19def make_poisson_solver(pw: PWDesc,
20 grid: UGDesc,
21 charge: float,
22 environment=None,
23 strength: float = 1.0,
24 dipolelayer: bool = False,
25 **kwargs) -> PoissonSolver:
26 if charge != 0.0 and not grid.pbc_c.any():
27 return ChargedPWPoissonSolver(pw, grid, charge, strength, **kwargs)
29 ps = PWPoissonSolver(pw, charge, strength)
31 if dipolelayer:
32 return DipoleLayerPWPoissonSolver(ps, grid, **kwargs)
34 assert not kwargs
36 if hasattr(environment, 'dielectric'):
37 if 1:
38 return ConjugateGradientPoissonSolver(
39 pw, grid, environment.dielectric, zero_vacuum=True)
40 from gpaw.new.sjm import SJMPWPoissonSolver
41 return SJMPWPoissonSolver(pw, environment.dielectric, grid)
43 return ps
46class PWPoissonSolver(PoissonSolver):
47 def __init__(self,
48 pw: PWDesc,
49 charge: float = 0.0,
50 strength: float = 1.0):
51 self.pw = pw
52 self.charge = charge
53 self.strength = strength
55 self.ekin_g = pw.ekin_G.copy()
56 if pw.comm.rank == 0:
57 # Avoid division by zero:
58 self.ekin_g[0] = 1.0
60 def __str__(self) -> str:
61 txt = ('poisson solver:\n'
62 f' ecut: {self.pw.ecut * Ha} # eV\n')
63 if self.strength != 1.0:
64 txt += f' strength: {self.strength}\n'
65 if self.charge != 0.0:
66 txt += f' uniform background charge: {self.charge} # electrons\n'
67 return txt
69 def solve(self,
70 vHt_g: PWArray,
71 rhot_g: PWArray) -> float:
72 """Solve Poisson equation.
74 Places result in vHt_g ndarray.
75 """
76 epot = self._solve(vHt_g, rhot_g)
77 return epot
79 def _solve(self,
80 vHt_g,
81 rhot_g) -> float:
82 vHt_g.data[:] = 2 * pi * self.strength * rhot_g.data
83 if self.pw.comm.rank == 0:
84 # Use uniform background charge in case we have a charged system:
85 vHt_g.data[0] = 0.0
86 if not isinstance(self.ekin_g, vHt_g.xp.ndarray):
87 self.ekin_g = vHt_g.xp.array(self.ekin_g)
88 vHt_g.data /= self.ekin_g
89 epot = 0.5 * vHt_g.integrate(rhot_g)
90 return epot
93class ChargedPWPoissonSolver(PWPoissonSolver):
94 def __init__(self,
95 pw: PWDesc,
96 grid: UGDesc,
97 charge: float,
98 strength: float = 1.0,
99 alpha: float = None,
100 eps: float = 1e-5):
101 """Reciprocal-space Poisson solver for charged molecules.
103 * Add a compensating Gaussian-shaped charge to the density
104 in order to make the total charge neutral (placed in the
105 middle of the unit cell
107 * Solve Poisson equation.
109 * Correct potential so that it has the correct 1/r
110 asymptotic behavior
112 * Correct energy to remove the artificial interaction with
113 the compensation charge
115 Parameters
116 ----------
117 pw: PWDesc
118 grid: UGDesc
119 charge: float
120 strength: float
121 alpha: float
122 eps: float
124 Attributes
125 ----------
126 alpha : float
127 charge_g : np.ndarray
128 Gaussian-shaped charge in reciprocal space
129 potential_g : PWArray
130 Potential in reciprocal space created by charge_g
131 """
132 super().__init__(pw, charge, strength)
134 if alpha is None:
135 # Shortest distance from center to edge of cell:
136 rcut = 0.5 / (pw.icell**2).sum(axis=1).max()**0.5
138 # Make sure e^(-alpha*rcut^2)=eps:
139 alpha = -rcut**-2 * np.log(eps)
141 self.alpha = alpha
143 center_v = pw.cell_cv.sum(axis=0) / 2
144 G2_g = 2 * pw.ekin_G
145 G_gv = pw.G_plus_k_Gv
146 self.charge_g = np.exp(-1 / (4 * alpha) * G2_g +
147 1j * (G_gv @ center_v))
148 self.charge_g *= charge / pw.dv
150 R_Rv = grid.xyz()
151 d_R = ((R_Rv - center_v)**2).sum(axis=3)**0.5
152 potential_R = grid.empty()
154 # avoid division by 0
155 zero_indx = d_R == 0
156 d_R[zero_indx] = 1
157 potential_R.data[:] = charge * erf(alpha**0.5 * d_R) / d_R
158 # at zero we should have:
159 # erf(alpha**0.5 * d_R) / d_R = alpha**0.5 * 2 / sqrt(pi)
160 potential_R.data[zero_indx] = charge * alpha**0.5 * 2 / np.sqrt(pi)
161 self.potential_g = potential_R.fft(pw=pw)
163 def __str__(self) -> str:
164 txt, x, _ = super().__str__().rsplit('\n', 2)
165 assert x.startswith(' uniform background charge:')
166 txt += (
167 '\n # using Gaussian-shaped compensation charge: e^(-alpha r^2)\n'
168 f' alpha: {self.alpha} # bohr^-2')
169 return txt
171 def _solve(self,
172 vHt_g,
173 rhot_g) -> float:
174 neutral_g = rhot_g.copy()
175 neutral_g.data += self.charge_g
177 if neutral_g.desc.comm.rank == 0:
178 error = neutral_g.data[0] # * self.pd.gd.dv
179 assert error.imag == 0.0, error
180 assert abs(error.real) < 0.00001, error
181 neutral_g.data[0] = 0.0
183 vHt_g.data[:] = 2 * pi * neutral_g.data
184 vHt_g.data /= self.ekin_g
185 epot = 0.5 * vHt_g.integrate(neutral_g)
186 epot -= self.potential_g.integrate(rhot_g)
187 epot -= self.charge**2 * (self.alpha / 2 / pi)**0.5
188 vHt_g.data -= self.potential_g.data
189 return epot
192class DipoleLayerPWPoissonSolver(PoissonSolver):
193 def __init__(self,
194 ps: PWPoissonSolver,
195 grid: UGDesc,
196 width: float = 1.0, # Ångström
197 zero_vacuum=False):
198 self.ps = ps
199 self.grid = grid
200 self.width = width / Bohr
201 self.zero_vacuum = zero_vacuum
202 assert grid.pbc_c.sum() == 2
203 self.axis = np.where(~grid.pbc_c)[0][0]
204 self.correction = np.nan
205 self.pw = ps.pw
207 def solve(self,
208 vHt_g: PWArray,
209 rhot_g: PWArray) -> float:
210 epot = self.ps.solve(vHt_g, rhot_g)
211 dip_v = -rhot_g.moment()
212 c = self.axis
213 L = self.grid.cell_cv[c, c]
214 self.correction = 2 * np.pi * dip_v[c] * L / self.grid.volume
215 vHt_g.data -= 2 * self.correction * self.sawtooth_g.data
216 if self.zero_vacuum:
217 v0 = vHt_g.boundary_value(self.axis)
218 if vHt_g.desc.comm.rank == 0:
219 vHt_g.data[0] += self.correction - v0
220 return epot + 2 * np.pi * dip_v[c]**2 / self.grid.volume
222 def dipole_layer_correction(self) -> float:
223 return self.correction
225 @cached_property
226 def sawtooth_g(self) -> PWArray:
227 grid = self.grid
228 if grid.comm.rank == 0:
229 c = self.axis
230 L = grid.cell_cv[c, c]
231 w = self.width / 2
232 assert w < L / 2, (w, L, c)
233 gc = int(w / L * grid.size_c[c])
234 x = np.linspace(0, L, grid.size_c[c], endpoint=False)
235 sawtooth = x / L - 0.5
236 a = 1 / L - 0.75 / w
237 b = 0.25 / w**3
238 sawtooth[:gc] = x[:gc] * (a + b * x[:gc]**2)
239 sawtooth[-gc:] = -sawtooth[gc:0:-1]
240 sawtooth_r = grid.new(comm=None).empty()
241 shape = [1, 1, 1]
242 shape[c] = -1
243 sawtooth_r.data[:] = sawtooth.reshape(shape)
244 sawtooth_g = sawtooth_r.fft(pw=self.ps.pw.new(comm=None)).data
245 else:
246 sawtooth_g = None
248 result_g = self.ps.pw.empty()
249 result_g.scatter_from(sawtooth_g)
250 return result_g
253class ConjugateGradientPoissonSolver(PWPoissonSolver):
254 """Poisson solver using conjugate gradient method in reciprocal space.
255 """
257 def __init__(self,
258 pw: PWDesc,
259 grid,
260 dielectric,
261 charge: float = 0.0,
262 strength: float = 1.0,
263 eps=1e-4,
264 maxiter=15,
265 zero_vacuum=False):
266 """Initialize the conjugate gradient Poisson solver.
268 Parameters:
269 -----------
270 pw : PWDesc
271 Plane wave descriptor
272 charge : float, optional
273 Total charge of the system
274 strength : float, optional
275 Scaling factor for the potential
276 eps : float, optional
277 Convergence threshold for conjugate gradient algorithm
278 maxiter : int, optional
279 Maximum number of iterations for the conjugate gradient algorithm
280 """
281 super().__init__(pw, charge, strength)
282 self.dielectric = dielectric
283 self.grid = grid
284 self.pw0 = pw.new(comm=None)
285 self.grid0 = grid.new(comm=None)
286 if pw.comm.rank == 0:
287 self.ekin_g = self.pw0.ekin_G.copy()
288 self.ekin_g[0] = 1.0
290 self.eps = eps
291 self.maxiter = maxiter
293 self.eps0_R = None
295 self.drho_g = None
296 self.zero_vacuum = zero_vacuum
297 if zero_vacuum:
298 self.drho_g = dipole_layer(grid).fft(pw=pw)
300 def __str__(self) -> str:
301 txt = ('conjugate gradient poisson solver:\n'
302 f' ecut: {self.pw.ecut * Ha} # eV\n'
303 f' eps: {self.eps}\n'
304 f' maxiter: {self.maxiter}\n')
305 if self.strength != 1.0:
306 txt += f' strength: {self.strength}\n'
307 if self.charge != 0.0:
308 txt += f' uniform background charge: {self.charge} # electrons\n'
309 return txt
311 def get_description(self):
312 return 'Conjugate Gradient Poisson Solver'
314 def operator(self, phi_G):
315 """Apply the generalized Poisson operator in reciprocal space.
317 Parameters:
318 -----------
319 phi_q : ndarray
320 Input potential in reciprocal space
322 Returns:
323 --------
324 ndarray
325 Result of operator application
326 """
327 pw = self.pw0
328 grid = self.grid0
329 G_vG = pw.G_plus_k_Gv.T
331 ophi_G = np.zeros_like(phi_G)
332 for G_G in G_vG:
333 grad_G = pw.from_data(G_G * phi_G)
334 grad_R = grad_G.ifft(grid=grid)
335 grad_R.data *= self.eps0_R.data
336 ophi_G += grad_R.fft(pw=pw).data * G_G
338 return ophi_G
340 def _solve(self,
341 vHt_g,
342 rhot_g) -> float:
343 vHt_g.data[:] = 4 * np.pi * self.strength * rhot_g.data
344 eps_R = self.grid.from_data(self.dielectric.eps_gradeps[0])
345 self.eps0_R = eps_R.gather()
347 vHt0_g = vHt_g.gather()
349 if self.pw.comm.rank == 0:
350 vHt0_g.data[0] = 0.0
352 N = len(vHt0_g.data)
353 op = LinearOperator((N, N),
354 matvec=self.operator,
355 dtype=complex)
356 M = LinearOperator((N, N),
357 matvec=lambda x: 0.5 * x / self.ekin_g,
358 dtype=complex)
359 vHt0_g.data[:], info = cg(
360 op, vHt0_g.data, maxiter=self.maxiter, M=M, **{RTOL: self.eps})
361 if info != 0:
362 warnings.warn(
363 f'Conjugate gradient did not converge (info={info})')
365 vHt_g.scatter_from(vHt0_g)
367 if self.zero_vacuum:
368 self.zero_vacuum = False
369 dphi_g = self.pw.zeros()
370 self._solve(dphi_g, self.drho_g)
371 v0s, v1s = xy_average_at_boundary(dphi_g)
372 v0, v1 = xy_average_at_boundary(vHt_g)
373 vHt_g.data -= dphi_g.data * (v1 / v1s)
374 vHt_g.data[0] -= v0 - (v1 / v1s) * v0s
375 self.zero_vacuum = True
377 epot = 0.5 * vHt_g.integrate(rhot_g)
378 return epot
380 def correct_slope(self, vHt_g: PWArray):
381 from gpaw.new.sjm import modified_saw_tooth
382 eps_r = self.grid.from_data(self.dielectric.eps_gradeps[0])
383 eps0_r = eps_r.gather()
384 vHt0_g = vHt_g.gather()
385 if eps0_r is not None:
386 saw_tooth_z = modified_saw_tooth(eps0_r)
387 vHt0_r = vHt0_g.ifft(grid=self.grid.new(comm=None))
388 s1, s2 = saw_tooth_z[[2, 10]]
389 v1, v2 = vHt0_r.data[:, :, [2, 10]].mean(axis=(0, 1))
390 vHt0_r.data -= (v2 - v1) / (s2 - s1) * saw_tooth_z[np.newaxis,
391 np.newaxis]
392 vHt0_r.data -= vHt0_r.data[:, :, -1].mean()
393 vHt0_r.fft(out=vHt0_g)
394 vHt_g.scatter_from(vHt0_g)
397def dipole_layer(grid: UGDesc, z0: float = 2.0):
398 a_r = grid.empty()
399 z_r = grid.xyz()[:, :, :, 2]
400 n = z_r.shape[2]
401 h = grid.cell_cv[2, 2]
402 d = h / n
403 z0 = round(z0 / d) * d
404 z_r += h / 2 - z0
405 z_r %= h
406 z_r -= h / 2
407 alpha = 6.0 / z0**2
408 a_r.data[:] = 4 * alpha**1.5 / np.pi**0.5 * z_r * np.exp(-alpha * z_r**2)
409 return a_r
412def xy_average_at_boundary(f_G: PWArray) -> np.ndarray:
413 """Calculate average value and derivative at boundary of box."""
414 pw = f_G.desc
415 m0_G, m1_G = pw.indices_cG[:2, pw.ng1:pw.ng2] == 0
416 mask_G = m0_G & m1_G
417 f_q = f_G.data[mask_G]
418 value = f_q.real.sum() * 2
419 derivative = pw.G_plus_k_Gv[mask_G, 2] @ f_q.imag
420 if pw.comm.rank == 0:
421 value -= f_q[0].real
422 result = np.array([value, derivative])
423 pw.comm.sum(result)
424 return result