Coverage for gpaw/test/solvation/test_pw_poisson.py: 100%
153 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1import numpy as np
2import pytest
3import time
4from ase.units import Bohr
5from gpaw.core import PWDesc, UGDesc, PWArray
6from gpaw.new.pw.poisson import ConjugateGradientPoissonSolver
7from ase.parallel import parprint
9nn = 3
10accuracy = 2e-10
11BOX = 12.0
14def generate_grid(box, h=0.4, pbc=[True, True, True], dtype=float):
15 """
16 Create a grid for testing.
18 Parameters:
19 ----------
20 box : float
21 Box size in Ångström
22 h : float
23 Grid spacing in Ångström
24 pbc : list
25 Periodic boundary conditions
26 dtype : type
27 Data type for the grid (float or complex)
29 Returns:
30 -------
31 grid : UGDesc
32 Grid descriptor
33 """
34 diag = np.array([box] * 3)
35 cell = np.diag(diag)
36 grid_shape = tuple((diag / h * 2).astype(int))
38 return UGDesc(cell=cell / Bohr, size=grid_shape, pbc=pbc, dtype=dtype)
41def generate_dipole_charges(grid, pw, box):
42 """
43 Generate a dipole charge distribution (positive and negative charges).
45 Parameters:
46 -----------
47 grid : UGDesc
48 Grid descriptor
49 pw : PWDesc
50 Plane wave descriptor
51 box : float
52 Box size
54 Returns:
55 ---------
56 rho : PWArray
57 Charge density for testing
58 """
59 xyz = grid.xyz()
60 x = xyz[:, :, :, 0]
61 y = xyz[:, :, :, 1]
62 z = xyz[:, :, :, 2]
63 center = box / (2 * Bohr)
65 radius = 1.0 / Bohr
66 pos_center = np.array([center - 2.0, center, center])
67 neg_center = np.array([center + 2.0, center, center])
69 r_pos = np.sqrt(
70 (x - pos_center[0])**2 +
71 (y - pos_center[1])**2 +
72 (z - pos_center[2])**2)
73 r_neg = np.sqrt(
74 (x - neg_center[0])**2 +
75 (y - neg_center[1])**2 +
76 (z - neg_center[2])**2)
78 shape = tuple(grid.size_c)
79 charge_array = np.zeros(shape)
81 sigma = 0.5 / Bohr
82 pos_charge = 0.5 * (1.0 - np.tanh((r_pos - radius) / sigma))
83 neg_charge = -0.5 * (1.0 - np.tanh((r_neg - radius) / sigma))
85 charge_array = pos_charge + neg_charge
87 total_charge = np.sum(charge_array) * grid.dv
88 charge_array -= total_charge / (grid.dv * np.prod(shape))
90 rho_r = grid.zeros()
91 rho_r.data[:] = charge_array
93 rho = rho_r.fft(pw=pw)
95 return rho
98def generate_quadrupole_charges(grid, pw, box, seed=None):
99 """
100 Generate a quadrupole-like charge distribution with four charges
101 arranged in a tetrahedral pattern.
103 Parameters:
104 -----------
105 grid : UGDesc
106 Grid descriptor
107 pw : PWDesc
108 Plane wave descriptor
109 box : float
110 Box size
111 seed : int, optional
112 Seed for random number generator
114 Returns:
115 ---------
116 rho : PWArray
117 Charge density for testing
118 """
119 if seed is not None:
120 np.random.seed(seed)
122 xyz = grid.xyz()
123 x = xyz[:, :, :, 0]
124 y = xyz[:, :, :, 1]
125 z = xyz[:, :, :, 2]
126 center = box / (2 * Bohr)
128 radius = box / (4 * Bohr)
129 charge_radius = 1.0 / Bohr
131 vertices = np.array([
132 [1, 1, 1],
133 [1, -1, -1],
134 [-1, 1, -1],
135 [-1, -1, 1]
136 ]) / np.sqrt(3)
138 centers = vertices * radius
139 centers += np.array([center, center, center])
141 charges = [1.0, -1.0, 1.0, -1.0]
143 shape = tuple(grid.size_c)
144 charge_array = np.zeros(shape)
145 sigma = 0.5 / Bohr
147 for q, pos in zip(charges, centers):
148 r = np.sqrt((x - pos[0])**2 + (y - pos[1])**2 + (z - pos[2])**2)
149 charge = q * 0.5 * (1.0 - np.tanh((r - charge_radius) / sigma))
150 charge_array += charge
152 total_charge = np.sum(charge_array) * grid.dv
153 charge_array -= total_charge / (grid.dv * np.prod(shape))
155 rho_r = grid.zeros()
156 rho_r.data[:] = charge_array
158 rho = rho_r.fft(pw=pw)
160 return rho
163def spherical_dielectric_function(grid, box, epsinf=80.0, eps1=1.0,
164 radius=None):
165 """
166 Generate a spherical dielectric function with high dielectric inside
167 a sphere and low dielectric outside, with a smooth transition at the
168 boundary.
170 Parameters:
171 ----------
172 grid : UGDesc
173 Grid descriptor
174 box : float
175 Box size in Ångström
176 epsinf : float
177 Dielectric constant inside the sphere (typically water = 80)
178 eps1 : float
179 Dielectric constant outside the sphere (typically vacuum = 1)
180 radius : float, optional
181 Radius of the dielectric sphere in Ångström. Defaults to 1/3
182 of box size.
184 Returns:
185 --------
186 eps : UGArray
187 Dielectric function for testing
188 eps_gradeps : list
189 List containing eps and its gradient as UGArray objects
190 """
191 xyz = grid.xyz()
192 x = xyz[:, :, :, 0]
193 y = xyz[:, :, :, 1]
194 z = xyz[:, :, :, 2]
195 center = box / (2 * Bohr)
197 if radius is None:
198 radius = box / 3.0
199 radius_bohr = radius / Bohr
201 r = np.sqrt((x - center)**2 + (y - center)**2 + (z - center)**2)
203 width = 0.5 / Bohr
204 transition = 0.5 * (1.0 - np.tanh((r - radius_bohr) / width))
206 eps_np = eps1 + (epsinf - eps1) * transition
208 deps_dr_np = -(epsinf - eps1) * 0.5 * (
209 1.0 - np.tanh((r - radius_bohr) / width)**2) / width
211 grad_x_np = np.zeros_like(r)
212 grad_y_np = np.zeros_like(r)
213 grad_z_np = np.zeros_like(r)
215 mask = r > 1e-10
216 grad_x_np[mask] = deps_dr_np[mask] * (x[mask] - center) / r[mask]
217 grad_y_np[mask] = deps_dr_np[mask] * (y[mask] - center) / r[mask]
218 grad_z_np[mask] = deps_dr_np[mask] * (z[mask] - center) / r[mask]
220 eps = grid.zeros()
221 eps.data[:] = eps_np
223 grad_x = grid.zeros()
224 grad_x.data[:] = grad_x_np
226 grad_y = grid.zeros()
227 grad_y.data[:] = grad_y_np
229 grad_z = grid.zeros()
230 grad_z.data[:] = grad_z_np
232 eps_gradeps = [eps, grad_x, grad_y, grad_z]
234 return eps, eps_gradeps
237@pytest.mark.parametrize("density_generator", [
238 generate_dipole_charges,
239 lambda grid, pw, box: generate_quadrupole_charges(grid, pw, box, seed=0),
240 lambda grid, pw, box: generate_quadrupole_charges(grid, pw, box, seed=1),
241])
242def test_cg_poisson_solver_constant_dielectric(density_generator):
243 """Test the conjugate gradient Poisson solver with constant dielectric."""
244 box = BOX
245 grid = generate_grid(box, dtype=complex)
247 pw = PWDesc(ecut=20.0, cell=grid.cell, kpt=grid.kpt, comm=grid.comm,
248 dtype=complex)
250 epsinf = 80.0
252 shape = tuple(grid.size_c)
253 eps_np = np.ones(shape) * epsinf
255 eps = grid.zeros()
256 eps.data[:] = eps_np
258 grad_zeros_np = np.zeros(shape)
260 grad_x = grid.zeros()
261 grad_x.data[:] = grad_zeros_np
262 grad_y = grid.zeros()
263 grad_y.data[:] = grad_zeros_np
264 grad_z = grid.zeros()
265 grad_z.data[:] = grad_zeros_np
267 eps_gradeps = [eps, grad_x, grad_y, grad_z]
269 rho = density_generator(grid, pw, box)
271 class MockDielectric:
272 def __init__(self, eps_gradeps):
273 self.eps_gradeps = [e.data for e in eps_gradeps]
275 solver = ConjugateGradientPoissonSolver(
276 pw, grid, MockDielectric(eps_gradeps),
277 eps=1e-6, maxiter=100)
279 phi = PWArray(pw=pw)
280 start = time.time()
281 solver.solve(phi, rho)
282 end = time.time()
283 parprint(f'CG Solver timing: {end - start}')
285 laplacian_phi = PWArray(pw=pw)
286 G2_q = pw.ekin_G.copy()
287 laplacian_phi.data[:] = G2_q * phi.data.copy()
289 expected_laplacian = PWArray(pw=pw)
290 factor = 2 * np.pi / float(epsinf)
291 expected_laplacian.data[:] = factor * rho.data.copy()
293 laplacian_phi_r = laplacian_phi.ifft(grid=grid).data
294 expected_laplacian_r = expected_laplacian.ifft(grid=grid).data
296 assert laplacian_phi_r == pytest.approx(expected_laplacian_r, abs=1e-3)
299@pytest.mark.parametrize("density_generator", [
300 generate_dipole_charges,
301 lambda grid, pw, box: generate_quadrupole_charges(grid, pw, box, seed=0),
302])
303def test_cg_poisson_solver_variable_dielectric(density_generator):
304 """Test the conjugate gradient Poisson solver with variable dielectric."""
305 box = BOX
306 grid = generate_grid(box, dtype=complex)
308 pw = PWDesc(ecut=20.0, cell=grid.cell, kpt=grid.kpt, comm=grid.comm,
309 dtype=complex)
311 eps, eps_gradeps = spherical_dielectric_function(grid, box,
312 epsinf=80.0, eps1=1.0)
314 rho = density_generator(grid, pw, box)
316 class MockDielectric:
317 def __init__(self, eps_gradeps):
318 self.eps_gradeps = [e.data for e in eps_gradeps]
320 solver = ConjugateGradientPoissonSolver(
321 pw=pw, grid=grid,
322 dielectric=MockDielectric(eps_gradeps),
323 eps=1e-6, maxiter=100)
325 phi = PWArray(pw=pw)
326 start = time.time()
327 solver.solve(phi, rho)
328 end = time.time()
329 parprint(f'CG Solver timing: {end - start}')
331 phi_operator = solver.operator(phi.data.copy())
333 expected = rho.data.copy()
334 expected *= 4 * np.pi
336 assert np.allclose(phi_operator, expected, atol=1e-3)