Coverage for gpaw/poisson_extravacuum.py: 99%
174 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import numpy as np
2from ase.geometry import cell_to_cellpar
3from ase.units import Bohr
5from gpaw.fd_operators import Laplace
6from gpaw.poisson import _PoissonSolver, create_poisson_solver
7from gpaw.transformers import Transformer
8from gpaw.utilities.extend_grid import (deextend_array, extend_array,
9 extended_grid_descriptor)
12def ext_gd(gd, **kwargs):
13 egd, _, _ = extended_grid_descriptor(gd, **kwargs)
14 return egd
17class ExtraVacuumPoissonSolver(_PoissonSolver):
18 """Wrapper around PoissonSolver extending the vacuum size.
20 Poisson equation is solved on the large grid defined by `gpts`
21 using `poissonsolver_large`.
23 If `coarses` is given, then the large grid is first coarsed
24 `coarses` times from to the original grid. Then, the coarse
25 potential is used to correct the boundary conditions
26 of the potential calculated on the original (small, fine)
27 grid by `poissonsolver_small`.
29 The parameters `nn_*` control the finite difference stencils
30 used in the coarsening, refining, and Laplace.
32 If the parameter `use_aux_grid` is `True`, an auxiliary
33 medium-sized grid is used between the large and small grids.
34 The parameter does not affect the result but can be used to
35 achieve speed-up depending on the grid sizes.
36 """
38 def __init__(self, gpts, poissonsolver_large,
39 poissonsolver_small=None, coarses=0,
40 nn_coarse=3, nn_refine=3, nn_laplace=3,
41 use_aux_grid=True):
42 # TODO: Alternative options: vacuum size and h
43 self.N_large_fine_c = np.array(gpts, dtype=int)
44 self.Ncoar = coarses # coar == coarse
45 if self.Ncoar > 0:
46 self.use_coarse = True
47 else:
48 self.use_coarse = False
49 self.ps_large_coar = create_poisson_solver(poissonsolver_large)
50 if self.use_coarse:
51 self.ps_small_fine = create_poisson_solver(poissonsolver_small)
52 else:
53 assert poissonsolver_small is None
54 self.nn_coarse = nn_coarse
55 self.nn_refine = nn_refine
56 self.nn_laplace = nn_laplace
57 self.use_aux_grid = use_aux_grid
58 self._initialized = False
60 def set_grid_descriptor(self, gd):
61 # If non-periodic boundary conditions is used,
62 # there is problems with auxiliary grid.
63 # Maybe with use_aux_grid=False it would work?
64 if gd.pbc_c.any():
65 raise NotImplementedError('Only non-periodic boundary '
66 'conditions are tested')
68 self.gd_small_fine = gd
69 assert np.all(self.gd_small_fine.N_c <= self.N_large_fine_c), \
70 'extended grid has to be larger than the original one'
72 if self.use_coarse:
73 # 1.1. Construct coarse chain on the small grid
74 self.coarser_i = []
75 gd = self.gd_small_fine
76 N_c = self.N_large_fine_c.copy()
77 for i in range(self.Ncoar):
78 gd2 = gd.coarsen()
79 self.coarser_i.append(Transformer(gd, gd2, self.nn_coarse))
80 N_c //= 2
81 gd = gd2
82 self.gd_small_coar = gd
83 else:
84 self.gd_small_coar = self.gd_small_fine
85 N_c = self.N_large_fine_c
87 # 1.2. Construct coarse extended grid
88 self.gd_large_coar = ext_gd(self.gd_small_coar, N_c=N_c)
90 # Initialize poissonsolvers
91 self.ps_large_coar.set_grid_descriptor(self.gd_large_coar)
92 if not self.use_coarse:
93 return
94 self.ps_small_fine.set_grid_descriptor(self.gd_small_fine)
96 if self.use_aux_grid:
97 # 2.1. Construct an auxiliary grid that is the small grid plus
98 # a buffer region allowing Laplace and refining
99 # with the used stencils
100 buf = self.nn_refine
101 for i in range(self.Ncoar):
102 buf = 2 * buf + self.nn_refine
103 buf += self.nn_laplace
104 div = 2**self.Ncoar
105 if buf % div != 0:
106 buf += div - buf % div
107 N_c = self.gd_small_fine.N_c + 2 * buf
108 if np.any(N_c > self.N_large_fine_c):
109 self.use_aux_grid = False
110 N_c = self.N_large_fine_c
111 self.gd_aux_fine = ext_gd(self.gd_small_fine, N_c=N_c)
112 else:
113 self.gd_aux_fine = ext_gd(self.gd_small_fine,
114 N_c=self.N_large_fine_c)
116 # 2.2 Construct Laplace on the aux grid
117 self.laplace_aux_fine = Laplace(self.gd_aux_fine, - 0.25 / np.pi,
118 self.nn_laplace)
120 # 2.3 Construct refine chain
121 self.refiner_i = []
122 gd = self.gd_aux_fine
123 N_c = gd.N_c.copy()
124 for i in range(self.Ncoar):
125 gd2 = gd.coarsen()
126 self.refiner_i.append(Transformer(gd2, gd, self.nn_refine))
127 N_c //= 2
128 gd = gd2
129 self.refiner_i = self.refiner_i[::-1]
130 self.gd_aux_coar = gd
132 if self.use_aux_grid:
133 # 2.4 Construct large coarse grid from aux grid
134 self.gd_large_coar_from_aux = ext_gd(self.gd_aux_coar,
135 N_c=self.gd_large_coar.N_c)
136 # Check the consistency of the grids
137 gd1 = self.gd_large_coar
138 gd2 = self.gd_large_coar_from_aux
139 assert np.all(gd1.N_c == gd2.N_c)
140 assert np.allclose(gd1.h_cv, gd2.h_cv, rtol=0, atol=1e-12)
142 self._initialized = False
144 def _init(self):
145 if self._initialized:
146 return
147 # Allocate arrays
148 self.phi_large_coar_g = self.gd_large_coar.zeros()
149 self._initialized = True
151 # Initialize poissonsolvers
152 # self.ps_large_coar._init()
153 # if not self.use_coarse:
154 # return
155 # self.ps_small_fine._init()
157 def solve(self, phi, rho, **kwargs):
158 self._init()
159 phi_small_fine_g = phi
160 rho_small_fine_g = rho.copy()
162 if self.use_coarse:
163 # 1.1. Coarse rho on the small grid
164 tmp_g = rho_small_fine_g
165 for coarser in self.coarser_i:
166 tmp_g = coarser.apply(tmp_g)
167 rho_small_coar_g = tmp_g
168 else:
169 rho_small_coar_g = rho_small_fine_g
171 # 1.2. Extend rho to the large grid
172 rho_large_coar_g = self.gd_large_coar.zeros()
173 extend_array(self.gd_small_coar, self.gd_large_coar,
174 rho_small_coar_g, rho_large_coar_g)
176 # 1.3 Solve potential on the large coarse grid
177 niter_large = self.ps_large_coar.solve(self.phi_large_coar_g,
178 rho_large_coar_g, **kwargs)
179 rho_large_coar_g = None
181 if not self.use_coarse:
182 deextend_array(self.gd_small_fine, self.gd_large_coar,
183 phi_small_fine_g, self.phi_large_coar_g)
184 return niter_large
186 if self.use_aux_grid:
187 # 2.1 De-extend the potential to the auxiliary grid
188 phi_aux_coar_g = self.gd_aux_coar.empty()
189 deextend_array(self.gd_aux_coar, self.gd_large_coar_from_aux,
190 phi_aux_coar_g, self.phi_large_coar_g)
191 else:
192 phi_aux_coar_g = self.phi_large_coar_g
194 # 3.1 Refine the potential
195 tmp_g = phi_aux_coar_g
196 for refiner in self.refiner_i:
197 tmp_g = refiner.apply(tmp_g)
198 phi_aux_coar_g = None
199 phi_aux_fine_g = tmp_g
201 # 3.2 Calculate the corresponding density with Laplace
202 # (the refined coarse density would not accurately match with
203 # the potential)
204 rho_aux_fine_g = self.gd_aux_fine.empty()
205 self.laplace_aux_fine.apply(phi_aux_fine_g, rho_aux_fine_g)
207 # 3.3 De-extend the potential and density to the small grid
208 cor_phi_small_fine_g = self.gd_small_fine.empty()
209 deextend_array(self.gd_small_fine, self.gd_aux_fine,
210 cor_phi_small_fine_g, phi_aux_fine_g)
211 phi_aux_fine_g = None
212 cor_rho_small_fine_g = self.gd_small_fine.empty()
213 deextend_array(self.gd_small_fine, self.gd_aux_fine,
214 cor_rho_small_fine_g, rho_aux_fine_g)
215 rho_aux_fine_g = None
217 # 3.4 Remove the correcting density and potential
218 rho_small_fine_g -= cor_rho_small_fine_g
219 phi_small_fine_g -= cor_phi_small_fine_g
221 # 3.5 Solve potential on the small grid
222 niter_small = self.ps_small_fine.solve(phi_small_fine_g,
223 rho_small_fine_g, **kwargs)
225 # 3.6 Correct potential and density
226 phi_small_fine_g += cor_phi_small_fine_g
227 # rho_small_fine_g += cor_rho_small_fine_g
229 return (niter_large, niter_small)
231 def estimate_memory(self, mem):
232 self.ps_large_coar.estimate_memory(mem.subnode('Large grid Poisson'))
233 if self.use_coarse:
234 ps = self.ps_small_fine
235 ps.estimate_memory(mem.subnode('Small grid Poisson'))
236 mem.subnode('Large coarse phi', self.gd_large_coar.bytecount())
237 tmp = max(self.gd_small_fine.bytecount(),
238 self.gd_large_coar.bytecount())
239 if self.use_coarse:
240 tmp = max(tmp,
241 self.gd_aux_coar.bytecount(),
242 self.gd_aux_fine.bytecount() * 2 +
243 self.gd_small_fine.bytecount(),
244 self.gd_aux_fine.bytecount() +
245 self.gd_small_fine.bytecount() * 2)
246 mem.subnode('Temporary arrays', tmp)
248 def get_description(self):
249 line = '%s with ' % self.__class__.__name__
250 if self.use_coarse:
251 line += 'large and small grids'
252 else:
253 line += 'large grid'
254 lines = [line]
256 def add_line(line, pad=0):
257 lines.extend(['{}{}'.format(' ' * pad, line)])
259 def get_cell(ps):
260 descr = ps.get_description().replace('\n', '\n%s' % (' ' * 8))
261 add_line('Poisson solver: %s' % descr, 8)
262 if hasattr(ps, 'gd'):
263 gd = ps.gd
264 par = cell_to_cellpar(gd.cell_cv * Bohr)
265 h_eff = gd.dv**(1.0 / 3.0) * Bohr
266 l1 = '{:8d} x {:8d} x {:8d} points'.format(*gd.N_c)
267 l2 = '{:8.2f} x {:8.2f} x {:8.2f} AA'.format(*par[:3])
268 l3 = f'Effective grid spacing dv^(1/3) = {h_eff:.4f}'
269 add_line('Grid: %s' % l1, 8)
270 add_line(' %s' % l2, 8)
271 add_line(' %s' % l3, 8)
273 add_line('Large grid:', 4)
274 get_cell(self.ps_large_coar)
276 if self.use_coarse:
277 add_line('Small grid:', 4)
278 get_cell(self.ps_small_fine)
280 return '\n'.join(lines)
282 def todict(self):
283 d = {'name': self.__class__.__name__}
284 d['gpts'] = self.N_large_fine_c
285 d['coarses'] = self.Ncoar
286 d['nn_coarse'] = self.nn_coarse
287 d['nn_refine'] = self.nn_refine
288 d['nn_laplace'] = self.nn_laplace
289 d['use_aux_grid'] = self.use_aux_grid
290 d['poissonsolver_large'] = self.ps_large_coar.todict()
291 if self.use_coarse:
292 d['poissonsolver_small'] = self.ps_small_fine.todict()
293 else:
294 d['poissonsolver_small'] = None
295 return d