Coverage for gpaw/inducedfield/inducedfield_base.py: 80%
291 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import numpy as np
2from ase.units import Bohr, Hartree
4import gpaw.mpi as mpi
5from gpaw.tddft.units import eV_to_aufrequency
6from gpaw.poisson import PoissonSolver
7from gpaw.fd_operators import Gradient
8from gpaw.grid_descriptor import GridDescriptor
9from gpaw.utilities.extend_grid import extended_grid_descriptor, \
10 extend_array, deextend_array, move_atoms
13def sendreceive_dict(comm, a_i, dest, b_i, src_i, iitems):
14 """
15 Send iitems of dictionary b_i distributed according to src_i
16 to dictionary a_i in dest.
17 """
18 requests = []
19 for i in iitems:
20 # Send and receive
21 if comm.rank == dest:
22 # Check if dest has it already
23 if dest == src_i[i]:
24 a_i[i] = b_i[i].copy()
25 else:
26 requests.append(comm.receive(a_i[i], src_i[i], tag=i,
27 block=False))
28 elif comm.rank == src_i[i]:
29 requests.append(comm.send(b_i[i], dest, tag=i, block=False))
30 comm.waitall(requests)
33class BaseInducedField:
34 """Partially virtual base class for induced field calculations.
36 Attributes:
37 -----------
38 omega_w: ndarray
39 Frequencies for Fourier transform in atomic units
40 folding: string
41 Folding type ('Gauss' or 'Lorentz' or None)
42 width: float
43 Width parameter for folding
44 Frho_wg: ndarray (complex)
45 Fourier transform of induced charge density
46 Fphi_wg: ndarray (complex)
47 Fourier transform of induced electric potential
48 Fef_wvg: ndarray (complex)
49 Fourier transform of induced electric field
50 Ffe_wg: ndarray (float)
51 Fourier transform of field enhancement
52 Fbgef_v: ndarray (float)
53 Fourier transform of background electric field
54 """
56 def __init__(self, filename=None, paw=None,
57 frequencies=None, folding='Gauss', width=0.08,
58 readmode=''):
59 """
60 Parameters:
61 -----------
62 filename: string
63 Filename of a previous ``InducedField`` to be loaded.
64 Setting filename disables parameters ``frequencies``,
65 ``folding`` and ``width``.
66 paw: PAW object
67 PAW object for InducedField
68 frequencies: ndarray or list of floats
69 Frequencies in eV for Fourier transforms.
70 This parameter is neglected if ``filename`` is given.
71 folding: string
72 Folding type: 'Gauss' or 'Lorentz' or None:
73 Gaussian or Lorentzian folding or no folding.
74 This parameter is neglected if ``filename`` is given.
75 width: float
76 Width in eV for the Gaussian (sigma) or Lorentzian (eta) folding
77 This parameter is neglected if ``filename`` is given.
78 """
79 self.dtype = complex
81 # These are allocated when calculated
82 self.fieldgd = None
83 self.Frho_wg = None
84 self.Fphi_wg = None
85 self.Fef_wvg = None
86 self.Ffe_wg = None
87 self.Fbgef_v = None
89 # has variables
90 self.has_paw = False
91 self.has_field = False
93 if not hasattr(self, 'readwritemode_str_to_list'):
94 self.readwritemode_str_to_list = \
95 {'': ['Frho', 'atoms'],
96 'all': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms'],
97 'field': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms']}
99 if filename is not None:
100 self.initialize(paw, allocate=False)
101 self.read(filename, mode=readmode)
102 return
104 self.folding = folding
105 self.width = width * eV_to_aufrequency
106 self.set_folding(folding, width * eV_to_aufrequency)
108 self.omega_w = np.asarray(frequencies) * eV_to_aufrequency
109 self.nw = len(self.omega_w)
111 self.nv = 3 # dimensionality of the space
113 self.initialize(paw, allocate=True)
115 def initialize(self, paw, allocate=True):
116 self.allocated = False
117 self.has_paw = paw is not None
119 if self.has_paw:
120 # If no paw is given, then the variables
121 # marked with "# !" are created in _read().
122 # Other variables are not accessible without
123 # paw (TODO: could be implemented...).
124 self.paw = paw
125 self.world = paw.wfs.world # !
126 self.domain_comm = paw.wfs.gd.comm # !
127 self.band_comm = paw.wfs.bd.comm # !
128 self.kpt_comm = paw.wfs.kd.comm # !
129 self.rank_a = paw.wfs.atom_partition.rank_a
130 self.nspins = paw.density.nspins # !
131 self.setups = paw.wfs.setups
132 self.density = paw.density
133 self.atoms = paw.atoms # !
134 self.na = len(self.atoms.get_atomic_numbers()) # !
135 self.gd = self.density.gd # !
136 self.stencil = self.density.stencil
137 self.log = paw.log
139 if allocate:
140 self.allocate()
142 def allocate(self):
143 self.allocated = True
145 def deallocate(self):
146 self.fieldgd = None
147 self.Frho_wg = None
148 self.Fphi_wg = None
149 self.Fef_wvg = None
150 self.Ffe_wg = None
151 self.Fbgef_v = None
153 def set_folding(self, folding, width):
154 """
155 width: float
156 Width in atomic units
157 """
158 if width is None:
159 folding = None
161 self.folding = folding
162 if self.folding is None:
163 self.width = None
164 else:
165 self.width = width
167 def get_induced_density(self, from_density, gridrefinement):
168 raise RuntimeError('Virtual member function called')
170 def calculate_induced_field(self, from_density='comp',
171 gridrefinement=2,
172 extend_N_cd=None,
173 deextend=False,
174 poissonsolver=None,
175 gradient_n=3):
176 if self.has_field and \
177 from_density == self.field_from_density and \
178 self.Frho_wg is not None:
179 Frho_wg = self.Frho_wg
180 gd = self.fieldgd
181 else:
182 Frho_wg, gd = self.get_induced_density(from_density,
183 gridrefinement)
185 # Always extend a bit to get field without jumps
186 if extend_N_cd is None:
187 extend_N = max(8, 2**int(np.ceil(np.log(gradient_n) / np.log(2.))))
188 extend_N_cd = extend_N * np.ones(shape=(3, 2), dtype=int)
189 deextend = True
191 # Extend grid
192 oldgd = gd
193 egd, cell_cv, move_c = \
194 extended_grid_descriptor(gd, extend_N_cd=extend_N_cd)
195 Frho_we = egd.zeros((self.nw,), dtype=self.dtype)
196 for w in range(self.nw):
197 extend_array(gd, egd, Frho_wg[w], Frho_we[w])
198 Frho_wg = Frho_we
199 gd = egd
200 if not deextend:
201 # TODO: this will make atoms unusable with original grid
202 self.atoms.set_cell(cell_cv, scale_atoms=False)
203 move_atoms(self.atoms, move_c)
205 # Allocate arrays
206 Fphi_wg = gd.zeros((self.nw,), dtype=self.dtype)
207 Fef_wvg = gd.zeros((self.nw, self.nv,), dtype=self.dtype)
208 Ffe_wg = gd.zeros((self.nw,), dtype=float)
210 # Poissonsolver
211 if poissonsolver is None:
212 poissonsolver = PoissonSolver(name='fd', eps=1e-20)
213 poissonsolver.set_grid_descriptor(gd)
215 for w in range(self.nw):
216 # TODO: better output of progress
217 # parprint('%d' % w)
218 calculate_field(gd, Frho_wg[w], self.Fbgef_v,
219 Fphi_wg[w], Fef_wvg[w], Ffe_wg[w],
220 poissonsolver=poissonsolver,
221 nv=self.nv,
222 gradient_n=gradient_n)
224 # De-extend grid
225 if deextend:
226 Frho_wo = oldgd.zeros((self.nw,), dtype=self.dtype)
227 Fphi_wo = oldgd.zeros((self.nw,), dtype=self.dtype)
228 Fef_wvo = oldgd.zeros((self.nw, self.nv,), dtype=self.dtype)
229 Ffe_wo = oldgd.zeros((self.nw,), dtype=float)
230 for w in range(self.nw):
231 deextend_array(oldgd, gd, Frho_wo[w], Frho_wg[w])
232 deextend_array(oldgd, gd, Fphi_wo[w], Fphi_wg[w])
233 deextend_array(oldgd, gd, Ffe_wo[w], Ffe_wg[w])
234 for v in range(self.nv):
235 deextend_array(oldgd, gd, Fef_wvo[w][v], Fef_wvg[w][v])
236 Frho_wg = Frho_wo
237 Fphi_wg = Fphi_wo
238 Fef_wvg = Fef_wvo
239 Ffe_wg = Ffe_wo
240 gd = oldgd
242 # Store results
243 self.has_field = True
244 self.field_from_density = from_density
245 self.fieldgd = gd
246 self.Frho_wg = Frho_wg
247 self.Fphi_wg = Fphi_wg
248 self.Fef_wvg = Fef_wvg
249 self.Ffe_wg = Ffe_wg
251 def _parse_readwritemode(self, mode):
252 if isinstance(mode, str):
253 try:
254 readwrites = self.readwritemode_str_to_list[mode]
255 except KeyError:
256 raise OSError('unknown readwrite mode string')
257 elif isinstance(mode, list):
258 readwrites = mode
259 else:
260 raise OSError('unknown readwrite mode type')
262 if any(k in readwrites for k in ['Frho', 'Fphi', 'Fef', 'Ffe']):
263 readwrites.append('field')
265 return readwrites
267 def read(self, filename, mode='', idiotproof=True):
268 if idiotproof and not filename.endswith('.ind'):
269 raise OSError('Filename must end with `.ind`.')
271 reads = self._parse_readwritemode(mode)
273 # Open reader
274 from gpaw.io import Reader
275 reader = Reader(filename)
276 self._read(reader, reads)
277 reader.close()
278 self.world.barrier()
280 def _read(self, reader, reads):
281 r = reader
283 # Test data type
284 dtype = {'float': float, 'complex': complex}[r.dtype]
285 if dtype != self.dtype:
286 raise OSError('Data is an incompatible type.')
288 # Read dimensions
289 na = r.na
290 self.nv = r.nv
291 nspins = r.nspins
292 ng = r.ng
294 # Background electric field
295 self.Fbgef_v = r.Fbgef_v
297 if self.has_paw:
298 # Test dimensions
299 if na != self.na:
300 raise OSError('natoms is incompatible with calculator')
301 if nspins != self.nspins:
302 raise OSError('nspins is incompatible with calculator')
303 if (ng != self.gd.get_size_of_global_array()).any():
304 raise OSError('grid is incompatible with calculator')
305 if (self.Fbgef_v != self.paw.kick_strength).any():
306 raise OSError('kick is incompatible with calculator')
307 else:
308 # Construct objects / assign values without paw
309 self.na = na
310 self.nspins = nspins
312 from ase.io.trajectory import read_atoms
313 self.atoms = read_atoms(r.atoms)
315 self.world = mpi.world
316 self.gd = GridDescriptor(ng + 1, self.atoms.get_cell() / Bohr,
317 pbc_c=False, comm=self.world)
318 self.domain_comm = self.gd.comm
319 self.band_comm = mpi.SerialCommunicator()
320 self.kpt_comm = mpi.SerialCommunicator()
322 # Folding
323 folding = r.folding
324 width = r.width
325 self.set_folding(folding, width)
327 # Frequencies
328 self.omega_w = r.omega_w
329 self.nw = len(self.omega_w)
331 # Read field
332 if 'field' in reads:
333 nfieldg = r.nfieldg
334 self.has_field = True
335 self.field_from_density = r.field_from_density
336 self.fieldgd = self.gd.new_descriptor(N_c=nfieldg + 1)
338 def readarray(name, shape, dtype):
339 if name.split('_')[0] in reads:
340 setattr(self, name, self.fieldgd.empty(shape, dtype=dtype))
341 self.fieldgd.distribute(r.get(name), getattr(self, name))
343 readarray('Frho_wg', (self.nw,), self.dtype)
344 readarray('Fphi_wg', (self.nw,), self.dtype)
345 readarray('Fef_wvg', (self.nw, self.nv), self.dtype)
346 readarray('Ffe_wg', (self.nw,), float)
348 def write(self, filename, mode='', idiotproof=True):
349 """
350 Parameters
351 ----------
352 mode: string or list of strings
355 """
356 if idiotproof and not filename.endswith('.ind'):
357 raise OSError('Filename must end with `.ind`.')
359 writes = self._parse_readwritemode(mode)
361 if 'field' in writes and self.fieldgd is None:
362 raise OSError('field variables cannot be written ' +
363 'before they are calculated')
365 from gpaw.io import Writer
366 writer = Writer(filename, self.world, tag='INDUCEDFIELD')
367 # Actual write
368 self._write(writer, writes)
369 # Make sure slaves don't return before master is done
370 writer.close()
371 self.world.barrier()
373 def _write(self, writer, writes):
374 # Write parameters/dimensions
375 writer.write(dtype={float: 'float', complex: 'complex'}[self.dtype],
376 folding=self.folding,
377 width=self.width,
378 na=self.na,
379 nv=self.nv,
380 nspins=self.nspins,
381 ng=self.gd.get_size_of_global_array())
383 # Write field grid
384 if 'field' in writes:
385 writer.write(field_from_density=self.field_from_density,
386 nfieldg=self.fieldgd.get_size_of_global_array())
388 # Write frequencies
389 writer.write(omega_w=self.omega_w)
391 # Write background electric field
392 writer.write(Fbgef_v=self.Fbgef_v)
394 from ase.io.trajectory import write_atoms
395 write_atoms(writer.child('atoms'), self.atoms)
397 if 'field' in writes:
398 def writearray(name, shape, dtype):
399 if name.split('_')[0] in writes:
400 writer.add_array(name, shape, dtype)
401 a_wxg = getattr(self, name)
402 for w in range(self.nw):
403 writer.fill(self.fieldgd.collect(a_wxg[w]))
405 shape = (self.nw,) + tuple(self.fieldgd.get_size_of_global_array())
406 writearray('Frho_wg', shape, self.dtype)
407 writearray('Fphi_wg', shape, self.dtype)
408 writearray('Ffe_wg', shape, float)
410 shape3 = shape[:1] + (self.nv,) + shape[1:]
411 writearray('Fef_wvg', shape3, self.dtype)
414def calculate_field(gd, rho_g, bgef_v,
415 phi_g, ef_vg, fe_g, # preallocated numpy arrays
416 poissonsolver,
417 nv=3,
418 gradient_n=3):
420 dtype = rho_g.dtype
421 yes_complex = dtype == complex
423 phi_g[:] = 0.0
424 ef_vg[:] = 0.0
425 fe_g[:] = 0.0
426 tmp_g = gd.zeros(dtype=float)
428 # Potential, real part
429 poissonsolver.solve(tmp_g, rho_g.real.copy())
430 phi_g += tmp_g
431 # Potential, imag part
432 if yes_complex:
433 tmp_g[:] = 0.0
434 poissonsolver.solve(tmp_g, rho_g.imag.copy())
435 phi_g += 1.0j * tmp_g
437 # Gradient
438 gradient = [Gradient(gd, v, scale=1.0, n=gradient_n)
439 for v in range(nv)]
440 for v in range(nv):
441 # Electric field, real part
442 gradient[v].apply(-phi_g.real, tmp_g)
443 ef_vg[v] += tmp_g
444 # Electric field, imag part
445 if yes_complex:
446 gradient[v].apply(-phi_g.imag, tmp_g)
447 ef_vg[v] += 1.0j * tmp_g
449 # Electric field enhancement
450 tmp_g[:] = 0.0 # total electric field norm
451 bgefnorm = 0.0 # background electric field norm
452 for v in range(nv):
453 tmp_g += np.absolute(bgef_v[v] + ef_vg[v])**2
454 bgefnorm += np.absolute(bgef_v[v])**2
456 tmp_g = np.sqrt(tmp_g)
457 bgefnorm = np.sqrt(bgefnorm)
459 fe_g[:] = tmp_g / bgefnorm
462def zero_pad(a):
463 """Add zeros to both sides of all axis on array a.
465 Not parallel safe.
466 """
468 z_shape = np.zeros(shape=len(a.shape))
469 z_shape[-3:] = 2
470 b_shape = np.array(a.shape) + z_shape
472 b = np.zeros(shape=b_shape, dtype=a.dtype)
474 b[..., 1:-1, 1:-1, 1:-1] = a
475 return b
478# TOOD: remove/edit this function
479def calculate_oscstr(Fn_wg, omega_w, box, kick):
480 omega_w = np.array(omega_w)
481 omega_w *= eV_to_aufrequency # to a.u.
482 box = box.copy() / Bohr # to a.u
483 volume = box.prod()
484 ng = np.array(Fn_wg[0].shape)
485 dV = volume / (ng - 1).prod() # Note: data is zeropadded to both sides
487 Fn_wg_im = Fn_wg.imag
488 if kick[0] != 0.0:
489 osc_w = -2 * omega_w / np.pi / kick[0] * \
490 ((Fn_wg_im.sum(axis=2)).sum(axis=2) *
491 np.linspace(0, box[0], ng[0])).sum(axis=1) * dV
492 elif kick[1] != 0.0:
493 osc_w = -2 * omega_w / np.pi / kick[1] * \
494 ((Fn_wg_im.sum(axis=1)).sum(axis=2) *
495 np.linspace(0, box[1], ng[1])).sum(axis=1) * dV
496 elif kick[2] != 0.0:
497 osc_w = -2 * omega_w / np.pi / kick[2] * \
498 ((Fn_wg_im.sum(axis=1)).sum(axis=1) *
499 np.linspace(0, box[2], ng[2])).sum(axis=1) * dV
500 return osc_w / Hartree # to 1/eV
503# TOOD: remove/edit this function
504def calculate_polarizability(Fn_wg, box, kick):
505 box = box.copy() / Bohr # to a.u
506 volume = box.prod()
507 ng = np.array(Fn_wg[0].shape)
508 dV = volume / (ng - 1).prod() # Note: data is zeropadded to both sides
510 if kick[0] != 0.0:
511 pol_w = -1.0 / kick[0] * \
512 ((Fn_wg.sum(axis=2)).sum(axis=2) *
513 np.linspace(0, box[0], ng[0])).sum(axis=1) * dV
514 elif kick[1] != 0.0:
515 pol_w = -1.0 / kick[1] * \
516 ((Fn_wg.sum(axis=1)).sum(axis=2) *
517 np.linspace(0, box[1], ng[1])).sum(axis=1) * dV
518 elif kick[2] != 0.0:
519 pol_w = -1.0 / kick[2] * \
520 ((Fn_wg.sum(axis=1)).sum(axis=1) *
521 np.linspace(0, box[2], ng[2])).sum(axis=1) * dV
522 return pol_w / Hartree**2 # to 1/eV**2