Coverage for gpaw/tddft/tdopers.py: 93%
274 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-12 00:18 +0000
1# Written by Lauri Lehtovaara, 2007
2"""This module implements classes for time-dependent variables and
3operators."""
5import numpy as np
7from gpaw.utilities import unpack_hermitian
8from gpaw.fd_operators import Laplace, Gradient
9from gpaw.overlap import Overlap
10from gpaw.wavefunctions.fd import FDWaveFunctions
13class TimeDependentHamiltonian:
14 """Time-dependent Hamiltonian, H(t)
16 This class contains information required to apply time-dependent
17 Hamiltonian to a wavefunction.
18 """
20 def __init__(self, wfs, spos_ac, hamiltonian, td_potential):
21 """Create the TimeDependentHamiltonian-object.
23 The time-dependent potential object must (be None or) have a member
24 function strength(self,time), which provides the strength of the
25 time-dependent external potential to x-direction at the given time.
27 Parameters
28 ----------
29 wfs: FDWaveFunctions
30 time-independent grid-based wavefunctions
31 hamiltonian: Hamiltonian
32 time-independent Hamiltonian
33 td_potential: TimeDependentPotential
34 time-dependent potential
35 """
37 self.wfs = wfs
38 self.hamiltonian = hamiltonian
39 self.td_potential = td_potential
40 self.time = self.old_time = 0
42 # internal smooth potential
43 self.vt_sG = hamiltonian.gd.zeros(hamiltonian.nspins)
45 # Increase the accuracy of Poisson solver
46 poisson = self.hamiltonian.poisson
47 if getattr(poisson, 'eps', None) and poisson.eps > 1e-12:
48 poisson.eps = 1e-12
50 # external potential
51 # if hamiltonian.vext_g is None:
52 # hamiltonian.vext_g = hamiltonian.finegd.zeros()
54 # self.ti_vext_g = hamiltonian.vext_g
55 # self.td_vext_g = hamiltonian.finegd.zeros(n=hamiltonian.nspins)
57 self.P = None
59 self.spos_ac = spos_ac
60 self.absorbing_boundary = None
62 def update(self, density, time):
63 """Updates the time-dependent Hamiltonian.
65 Parameters
66 ----------
67 density: Density
68 the density at the given time
69 (TimeDependentDensity.get_density())
70 time: float
71 the current time
73 """
75 self.old_time = self.time = time
76 self.hamiltonian.update(density)
78 def half_update(self, density, time):
79 """Updates the time-dependent Hamiltonian, in such a way, that a
80 half of the old Hamiltonian is kept and the other half is updated.
82 Parameters
83 ----------
84 density: Density
85 the density at the given time
86 (TimeDependentDensity.get_density())
87 time: float
88 the current time
90 """
92 self.old_time = self.time
93 self.time = time
95 # copy old
96 self.vt_sG[:] = self.hamiltonian.vt_sG
97 self.dH_asp = {}
98 for a, dH_sp in self.hamiltonian.dH_asp.items():
99 self.dH_asp[a] = dH_sp.copy()
100 # update
101 self.hamiltonian.update(density)
102 # average and difference
103 self.hamiltonian.vt_sG[:], self.vt_sG[:] = \
104 0.5 * (self.hamiltonian.vt_sG + self.vt_sG), \
105 self.hamiltonian.vt_sG - self.vt_sG
106 for a, dH_sp in self.hamiltonian.dH_asp.items():
107 dH_sp[:], self.dH_asp[a][:] = 0.5 * (dH_sp + self.dH_asp[a]), \
108 dH_sp - self.dH_asp[a] # pack/unpack is linear for real values
110 def half_apply_local_potential(self, psit_nG, Htpsit_nG, s):
111 """Apply the half-difference Hamiltonian operator to a set of vectors.
113 Parameters:
115 psit_nG: ndarray
116 set of vectors to which the overlap operator is applied.
117 psit_nG: ndarray, output
118 resulting H applied to psit_nG vectors.
119 s: int
120 spin index of k-point object defined in kpoint.py.
122 """
123 # Does exactly the same as Hamiltonian.apply_local_potential
124 # but uses the difference between vt_sG at time t and t+dt.
125 vt_G = self.vt_sG[s]
126 if psit_nG.ndim == 3:
127 Htpsit_nG += psit_nG * vt_G
128 else:
129 for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG):
130 Htpsit_G += psit_G * vt_G
132 def half_apply(self, kpt, psit, hpsit, calculate_P_ani=True):
133 """Applies the half-difference of the time-dependent Hamiltonian
134 to the wavefunction psit of the k-point kpt.
136 Parameters
137 ----------
138 kpt: Kpoint
139 the current k-point (kpt_u[index_of_k-point])
140 psit: List of coarse grid
141 the wavefuntions (on coarse grid)
142 (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
143 hpsit: List of coarse grid
144 the resulting "operated wavefunctions" (H psit)
145 calculate_P_ani: bool
146 When True, the integrals of projector times vectors
147 P_ni = <p_i | psit> are calculated.
148 When False, existing P_uni are used
150 """
152 hpsit.fill(0.0)
153 self.half_apply_local_potential(psit, hpsit, kpt.s)
155 # Does exactly the same as last part of Hamiltonian.apply but
156 # uses the difference between dH_asp at time t and t+dt.
157 shape = psit.shape[:-3]
158 P_axi = self.wfs.pt.dict(shape)
160 if calculate_P_ani:
161 self.wfs.pt.integrate(psit, P_axi, kpt.q)
162 else:
163 for a, P_ni in kpt.P_ani.items():
164 P_axi[a][:] = P_ni
166 for a, P_xi in P_axi.items():
167 dH_ii = unpack_hermitian(self.dH_asp[a][kpt.s])
168 P_axi[a][:] = np.dot(P_xi, dH_ii)
169 self.wfs.pt.add(hpsit, P_axi, kpt.q)
171 if self.td_potential is not None:
172 # FIXME: add half difference here... but maybe it's not important
173 # as this will be used only for getting initial guess. So, should
174 # not affect to the results, only to the speed of convergence.
175 # raise NotImplementedError
176 pass
178 def apply(self, kpt, psit, hpsit, calculate_P_ani=True):
179 """Applies the time-dependent Hamiltonian to the wavefunction psit of
180 the k-point kpt.
182 Parameters
183 ----------
184 kpt: Kpoint
185 the current k-point (kpt_u[index_of_k-point])
186 psit: List of coarse grid
187 the wavefuntions (on coarse grid)
188 (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
189 hpsit: List of coarse grid
190 the resulting "operated wavefunctions" (H psit)
191 calculate_P_ani: bool
192 When True, the integrals of projector times vectors
193 P_ni = <p_i | psit> are calculated.
194 When False, existing P_uni are used
196 """
198 self.hamiltonian.apply(psit, hpsit, self.wfs, kpt, calculate_P_ani)
200 # PAW correction
201 if self.P is not None:
202 self.P.add(psit, hpsit, self.wfs, kpt)
204 # Absorbing boundary conditions
206 # Imaginary potential
207 if self.absorbing_boundary is not None \
208 and self.absorbing_boundary.type == 'IPOT':
209 hpsit[:] += self.absorbing_boundary.get_potential_matrix() * psit
211 # Perfectly matched layers
212 if self.absorbing_boundary is not None \
213 and self.absorbing_boundary.type == 'PML':
214 # Perfectly matched layer is applied as potential Vpml = Tpml-T
215 # Where T = -0.5*\nabla^{2}\psi (Use latex for these equations)
216 # See abc.py for details
217 # This is probably not the most optimal approach and slows
218 # the propagation.
219 if self.lpsit is None:
220 self.lpsit = self.hamiltonian.gd.empty(n=len(psit),
221 dtype=complex)
222 self.laplace.apply(psit, self.lpsit, kpt.phase_cd)
223 hpsit[:] -= (.5 * (self.absorbing_boundary.get_G()**2 - 1.0) *
224 self.lpsit)
225 for i in range(3):
226 self.gradient[i].apply(psit, self.lpsit, kpt.phase_cd)
227 hpsit[:] -= (.5 * self.absorbing_boundary.get_G() *
228 self.absorbing_boundary.get_dG()[i] * self.lpsit)
230 # Time-dependent dipole field
231 if self.td_potential is not None:
232 # TODO on shaky ground here...
233 strength = self.td_potential.strength
234 add_linear_field(
235 self.wfs, self.spos_ac, psit, hpsit,
236 0.5 * strength(self.time) + 0.5 * strength(self.old_time), kpt)
238 def set_absorbing_boundary(self, absorbing_boundary):
239 """ Sets up the absorbing boundary.
240 Parameters:
241 absorbing_boundary: absorbing boundary object of any kind.
242 """
244 self.absorbing_boundary = absorbing_boundary
245 self.absorbing_boundary.set_up(self.hamiltonian.gd)
246 if self.absorbing_boundary.type == 'PML':
247 gd = self.hamiltonian.gd
248 self.laplace = Laplace(gd, n=2, dtype=complex)
249 self.gradient = np.array(
250 (Gradient(gd, 0, n=2,
251 dtype=complex), Gradient(gd, 1, n=2, dtype=complex),
252 Gradient(gd, 2, n=2, dtype=complex)))
253 self.lpsit = None
255 def calculate_paw_correction(self,
256 psit_nG,
257 hpsit,
258 wfs,
259 kpt,
260 v_atom,
261 calculate_P_ani=True):
262 """ Operates on psit_nG with the P-term that is present in PAW based
263 Ehrenfest dynamics
265 Parameters
266 ----------
267 psit_nG: list of coarse grid wavefunctions
269 hpsit: array to which P psit_nG is added
271 wfs: Wavefunctions object
273 kpt: Kpoint object
275 v_atom: atomic velocities
277 """
279 shape = psit_nG.shape[:-3]
280 P_axi = wfs.pt.dict(shape)
281 wfs.pt.integrate(psit_nG, P_axi, kpt.q)
283 # Coefficients for calculating P \psi_n
284 # P = -i sum_a v_a P^a, P^a = T^{\dagger} \nabla_{R_a} T
285 w_ani = wfs.pt.dict(wfs.bd.mynbands, zero=True)
286 # projector derivatives < nabla pt_i^a | psit_n >
287 dpt_aniv = wfs.pt.dict(wfs.bd.mynbands, derivative=True)
288 wfs.pt.derivative(psit_nG, dpt_aniv, kpt.q)
289 # wfs.calculate_forces(paw.hamiltonian, F_av)
290 for a in dpt_aniv.keys():
291 # ni = wfs.pt.get_function_count(a)
292 for c in range(3):
294 P_xi = P_axi[a]
295 # nabla_iiv contains terms < \phi_i1^a | d / d v phi_i2^a >
296 # - < phit_i1^a | d / dv phit_i2^a>, where v is either x,y or z
297 nabla_ii = wfs.setups[a].nabla_iiv[:, :, c]
298 dpt_ni = dpt_aniv[a][:, :, c]
299 dO_ii = wfs.setups[a].dO_ii
300 # dphi_aniv[a] = np.dot(P_xi, nabla_ii.transpose())
301 dphi_ni = np.dot(P_xi, nabla_ii.transpose())
302 pt_ni = np.dot(dpt_ni, dO_ii)
303 # pt_aniv[a] = np.dot(Dpt_ni, dO_ii)
304 w_ani[a] += (dphi_ni + pt_ni) * v_atom[a, c]
306 w_ani[a] *= complex(0, 1)
307 # dO_ani[a] *= complex(0,1)
309 # wfs.pt.add(ppsit, W_ani, kpt.q)
310 wfs.pt.add(hpsit, w_ani, kpt.q)
313# AbsorptionKickHamiltonian
314class AbsorptionKickHamiltonian:
315 """Absorption kick Hamiltonian, p.r
317 This class contains information required to apply absorption kick
318 Hamiltonian to a wavefunction.
319 """
321 def __init__(self, wfs, spos_ac, strength=[0.0, 0.0, 1e-3]):
322 """Create the AbsorptionKickHamiltonian-object.
324 Parameters
325 ----------
326 wfs: FDWaveFunctions
327 time-independent grid-based wavefunctions
328 spos_ac: ndarray
329 scaled positions
330 strength: float[3]
331 strength of the delta field to different directions
333 """
335 self.wfs = wfs
336 self.spos_ac = spos_ac
338 # magnitude
339 magnitude = np.sqrt(strength[0] * strength[0] +
340 strength[1] * strength[1] +
341 strength[2] * strength[2])
342 # iterations
343 self.iterations = int(round(magnitude / 1.0e-4))
344 if self.iterations < 1:
345 self.iterations = 1
346 # delta p
347 self.dp = strength / self.iterations
349 # hamiltonian
350 self.abs_hamiltonian = np.array([self.dp[0], self.dp[1], self.dp[2]])
352 def update(self, density, time):
353 """Dummy function = does nothing. Required to have correct interface.
355 Parameters
356 ----------
357 density: Density or None
358 the density at the given time or None (ignored)
359 time: Float or None
360 the current time (ignored)
362 """
363 pass
365 def half_update(self, density, time):
366 """Dummy function = does nothing. Required to have correct interface.
368 Parameters
369 ----------
370 density: Density or None
371 the density at the given time or None (ignored)
372 time: float or None
373 the current time (ignored)
375 """
376 pass
378 def apply(self, kpt, psit, hpsit, calculate_P_ani=True):
379 """Applies the absorption kick Hamiltonian to the wavefunction psit of
380 the k-point kpt.
382 Parameters
383 ----------
384 kpt: Kpoint
385 the current k-point (kpt_u[index_of_k-point])
386 psit: List of coarse grids
387 the wavefuntions (on coarse grid)
388 (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
389 hpsit: List of coarse grids
390 the resulting "operated wavefunctions" (H psit)
391 calculate_P_ani: bool
392 When True, the integrals of projector times vectors
393 P_ni = <p_i | psit> are calculated.
394 When False, existing P_uni are used
396 """
397 hpsit[:] = 0.0
399 # TODO on shaky ground here...
400 add_linear_field(self.wfs, self.spos_ac, psit, hpsit,
401 self.abs_hamiltonian, kpt)
404# Overlap
405class TimeDependentOverlap(Overlap):
406 """Time-dependent overlap operator S(t)
408 This class contains information required to apply time-dependent
409 overlap operator to a set of wavefunctions.
410 """
412 def __init__(self, timer):
413 """Creates the TimeDependentOverlap-object.
415 Parameters
416 ----------
417 XXX TODO
419 """
420 Overlap.__init__(self, timer)
422 def update_k_point_projections(self, wfs, kpt, psit=None):
423 """Updates the projector function overlap integrals
424 with the wavefunctions of a given k-point.
426 Parameters
427 ----------
428 wfs: TimeDependentWaveFunctions
429 time-independent grid-based wavefunctions
430 kpt: Kpoint
431 the current k-point (kpt_u[index_of_k-point])
432 psit: List of coarse grids (optional)
433 the wavefuntions (on coarse grid)
434 (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
436 """
437 if psit is not None:
438 wfs.pt.integrate(psit, kpt.P_ani, kpt.q)
439 else:
440 wfs.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
442 def update(self, wfs):
443 """Updates the time-dependent overlap operator.
445 Parameters
446 ----------
447 wfs: TimeDependentWaveFunctions
448 time-independent grid-based wavefunctions
450 """
451 for kpt in wfs.kpt_u:
452 self.update_k_point_projections(wfs, kpt)
454 def half_update(self, wfs):
455 """Updates the time-dependent overlap operator, in such a way,
456 that a half of the old overlap operator is kept and the other half
457 is updated. !Currently does nothing!
459 Parameters
460 ----------
461 wfs: TimeDependentWaveFunctions
462 time-independent grid-based wavefunctions
464 """
465 # for kpt in wfs.kpt_u:
466 # # copy old
467 # P_ani = {}
468 # for a,P_ni in kpt.P_ani.items():
469 # P_ani[a] = P_ni.copy()
470 # # update
471 # self.update_k_point_projections(wfs, kpt)
472 # # average
473 # for a,P_ni in P_ani.items():
474 # kpt.P_ani[a] += P_ni
475 # kpt.P_ani[a] *= .5
477 # !!! FIX ME !!! update overlap operator/projectors/...
478 pass
480 # def apply(self, psit, spsit, wfs, kpt, calculate_P_ani=True):
481 # """Apply the time-dependent overlap operator to the wavefunction
482 # psit of the k-point kpt.
483 #
484 # Parameters
485 # ----------
486 # psit: List of coarse grids
487 # the wavefuntions (on coarse grid)
488 # (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
489 # spsit: List of coarse grids
490 # the resulting "operated wavefunctions" (S psit)
491 # wfs: TimeDependentWaveFunctions
492 # time-independent grid-based wavefunctions
493 # kpt: Kpoint
494 # the current k-point (kpt_u[index_of_k-point])
495 # calculate_P_ani: bool
496 # When True, the integrals of projector times vectors
497 # P_ni = <p_i | psit> are calculated.
498 # When False, existing P_ani are used
499 #
500 # """
501 # self.overlap.apply(psit, spsit, wfs, kpt, calculate_P_ani)
502 #
503 def apply_inverse(self,
504 a_nG,
505 b_nG,
506 wfs,
507 kpt,
508 calculate_P_ani=True,
509 use_cg=True):
510 """Apply the approximative time-dependent inverse overlap operator
511 to the wavefunction psit of the k-point kpt.
513 Parameters
514 ----------
515 a_nG: List of coarse grids
516 the wavefuntions (on coarse grid)
517 (kpt_u[index_of_k-point].psit_nG[indices_of_wavefunc])
518 b_nG: List of coarse grids
519 the resulting "operated wavefunctions" (S^(-1) psit)
520 wfs: TimeDependentWaveFunctions
521 time-independent grid-based wavefunctions
522 kpt: Kpoint
523 the current k-point (kpt_u[index_of_k-point])
524 calculate_P_ani: bool
525 When True, the integrals of projector times vectors
526 P_ni = <p_i | psit> are calculated.
527 When False, existing P_uni are used
528 use_cg: bool
529 When True, use conjugate gradient method to solve for inverse.
531 """
532 if not use_cg:
533 self.timer.start('Apply approximate inverse overlap')
534 Overlap.apply_inverse(self, a_nG, b_nG, wfs, kpt, calculate_P_ani)
535 self.timer.stop('Apply approximate inverse overlap')
536 return
538 self.timer.start('Apply exact inverse overlap')
539 from gpaw.utilities.blas import axpy
541 # from gpaw.tddft.cscg import multi_zdotu, multi_scale, multi_zaxpy
542 # initialization
543 # Multivector dot product, a^T b, where ^T is transpose
545 def multi_zdotu(s, x, y, nvec):
546 for i in range(nvec):
547 s[i] = x[i].ravel().dot(y[i].ravel())
548 # s[i] = dotu(x[i],y[i])
549 wfs.gd.comm.sum(s)
550 return s
552 # Multivector ZAXPY: a x + y => y
553 def multi_zaxpy(a, x, y, nvec):
554 for i in range(nvec):
555 axpy(a[i] * (1 + 0J), x[i], y[i])
557 # Multiscale: a x => x
558 def multi_scale(a, x, nvec):
559 for i in range(nvec):
560 x[i] *= a[i]
562 nvec = len(a_nG)
563 r = wfs.gd.zeros(nvec, dtype=wfs.dtype)
564 z = wfs.gd.zeros((nvec, ), dtype=wfs.dtype)
565 p = wfs.gd.zeros(nvec, dtype=wfs.dtype)
566 q = wfs.gd.zeros(nvec, dtype=wfs.dtype)
567 alpha = np.zeros((nvec, ), dtype=wfs.dtype)
568 beta = np.zeros((nvec, ), dtype=wfs.dtype)
569 scale = np.zeros((nvec, ), dtype=wfs.dtype)
570 normr2 = np.zeros((nvec, ), dtype=wfs.dtype)
571 rho = np.zeros((nvec, ), dtype=wfs.dtype)
572 rho_prev = np.zeros((nvec, ), dtype=wfs.dtype)
573 rho_prev[:] = 1.0
574 tol_cg = 1e-14
575 multi_zdotu(scale, a_nG, a_nG, nvec)
576 scale = np.abs(scale)
578 x = b_nG # XXX TODO rename this
580 # x0 = S^-1_approx a_nG
581 self.apply_inverse(a_nG, x, wfs, kpt, calculate_P_ani, use_cg=False)
582 # r0 = a_nG - S x_0
583 self.apply(-x, r, wfs, kpt, calculate_P_ani)
584 r += a_nG
585 # print 'r.max() =', abs(r).max()
587 max_iter = 50
589 for i in range(max_iter):
591 # print 'iter =', i
593 self.apply_inverse(r, z, wfs, kpt, calculate_P_ani, use_cg=False)
595 multi_zdotu(rho, r, z, nvec)
597 beta = rho / rho_prev
598 multi_scale(beta, p, nvec)
599 p += z
601 self.apply(p, q, wfs, kpt, calculate_P_ani)
603 multi_zdotu(alpha, p, q, nvec)
604 alpha = rho / alpha
606 multi_zaxpy(alpha, p, x, nvec)
607 multi_zaxpy(-alpha, q, r, nvec)
609 multi_zdotu(normr2, r, r, nvec)
611 # rhoc = rho.copy()
612 rho_prev[:] = rho.copy()
613 # rho.copy()
614 # rho_prev = rho.copy()
616 # print '||r|| =', np.sqrt(np.abs(normr2/scale))
617 if ((np.sqrt(np.abs(normr2) / scale)) < tol_cg).all():
618 break
620 self.timer.stop('Apply exact inverse overlap')
623class TimeDependentWaveFunctions(FDWaveFunctions):
624 def __init__(self, stencil, parallel, initksl, gd, nvalence, collinear,
625 setups, bd, dtype, world, kd, kptband_comm, timer):
626 assert dtype == complex
627 FDWaveFunctions.__init__(self,
628 stencil,
629 parallel,
630 initksl,
631 gd,
632 nvalence,
633 setups,
634 bd,
635 dtype,
636 world,
637 kd,
638 kptband_comm,
639 collinear=collinear,
640 timer=timer)
641 self.overlap = self.make_overlap()
643 def make_overlap(self):
644 return TimeDependentOverlap(self.timer)
646 def calculate_forces(self, hamiltonian, F_av):
647 """ Calculate wavefunction forces with optional corrections for
648 Ehrenfest dynamics
649 """
651 # If td_correction is not none, we replace the overlap part of the
652 # force, sum_n f_n eps_n < psit_n | dO / dR_a | psit_n>, with
653 # sum_n f_n <psit_n | H S^-1 D^a + c.c. | psit_n >, with D^a
654 # defined as D^a = sum_{i1,i2}
655 # | pt_i1^a > [O_{i1,i2} < d pt_i2^a / dR_a |
656 # + (< phi_i1^a | d phi_i2^a / dR_a >
657 # - < phit_i1^a | d phit_i2^a / dR_a >) < pt_i1^a|].
658 # This is required in order to conserve the total energy also when
659 # electronic excitations start to play a significant role.
661 # TODO: move the corrections into the tddft directory
663 # Calculate force-contribution from k-points:
664 F_av.fill(0.0)
665 F_aniv = self.pt.dict(self.bd.mynbands, derivative=True)
666 # print 'self.dtype =', self.dtype
667 for kpt in self.kpt_u:
668 self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q)
670 # self.overlap.update_k_point_projections(kpt)
671 self.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
672 hpsit = self.gd.zeros(len(kpt.psit_nG), dtype=self.dtype)
673 # eps_psit = self.gd.zeros(len(kpt.psit_nG), dtype=self.dtype)
674 sinvhpsit = self.gd.zeros(len(kpt.psit_nG), dtype=self.dtype)
675 hamiltonian.apply(kpt.psit_nG,
676 hpsit,
677 self,
678 kpt,
679 calculate_P_ani=True)
680 self.overlap.apply_inverse(hpsit,
681 sinvhpsit,
682 self,
683 kpt,
684 calculate_P_ani=True)
686 G_axi = self.pt.dict(self.bd.mynbands)
687 self.pt.integrate(sinvhpsit, G_axi, kpt.q)
689 for a, F_niv in F_aniv.items():
690 F_niv = F_niv.conj()
691 F_niv *= kpt.f_n[:, np.newaxis, np.newaxis]
692 FdH1_niv = F_niv.copy()
693 dH_ii = unpack_hermitian(hamiltonian.dH_asp[a][kpt.s])
694 P_ni = kpt.P_ani[a]
695 dO_ii = hamiltonian.setups[a].dO_ii
696 F_vii = np.dot(np.dot(F_niv.transpose(), P_ni), dH_ii)
698 fP_ni = P_ni * kpt.f_n[:, np.newaxis]
699 G_ni = G_axi[a]
700 nabla_iiv = hamiltonian.setups[a].nabla_iiv
701 F_vii_sinvh_dpt = -np.dot(np.dot(FdH1_niv.transpose(), G_ni),
702 dO_ii)
703 F_vii_sinvh_dphi = -np.dot(
704 nabla_iiv.transpose(2, 0, 1),
705 np.dot(fP_ni.conj().transpose(), G_ni))
706 F_vii += F_vii_sinvh_dpt + F_vii_sinvh_dphi
708 # F_av_dO[a] += 2 * F_vii_dO.real.trace(0,1,2)
709 # F_av_dH[a] += 2 * F_vii_dH.real.trace(0,1,2)
710 F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
712 # Hack used in delta-scf calculations:
713 if hasattr(kpt, 'c_on'):
714 assert self.bd.comm.size == 1
715 self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) # XXX again
716 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
717 dtype=complex)
718 for ne, c_n in zip(kpt.ne_o, kpt.c_on):
719 d_nn += ne * np.outer(c_n.conj(), c_n)
720 for a, F_niv in F_aniv.items():
721 F_niv = F_niv.conj()
722 dH_ii = unpack_hermitian(hamiltonian.dH_asp[a][kpt.s])
723 Q_ni = np.dot(d_nn, kpt.P_ani[a])
724 F_vii = np.dot(np.dot(F_niv.transpose(), Q_ni), dH_ii)
725 F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis]
726 dO_ii = hamiltonian.setups[a].dO_ii
727 F_vii -= np.dot(np.dot(F_niv.transpose(), Q_ni), dO_ii)
728 F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
730 self.bd.comm.sum(F_av, 0)
732 if self.bd.comm.rank == 0:
733 self.kd.comm.sum(F_av, 0)
736# DummyDensity
737class DummyDensity:
738 """Implements dummy (= does nothing) density for AbsorptionKick."""
740 def __init__(self, wfs):
741 """Placeholder Density object for AbsorptionKick.
743 Parameters
744 ----------
745 wfs: FDWaveFunctions
746 time-independent grid-based wavefunctions
748 """
749 self.wfs = wfs
751 def update(self):
752 pass
754 def get_wavefunctions(self):
755 return self.wfs
757 def get_density(self):
758 return None
761# Density
762class TimeDependentDensity(DummyDensity):
763 """Time-dependent density rho(t)
765 This class contains information required to get the time-dependent
766 density.
767 """
769 def __init__(self, paw):
770 """Creates the TimeDependentDensity-object.
772 Parameters
773 ----------
774 paw: PAW
775 the PAW-object
776 """
777 DummyDensity.__init__(self, paw.wfs)
778 self.density = paw.density
780 def update(self):
781 """Updates the time-dependent density.
783 Parameters
784 ----------
785 None
787 """
788 # for kpt in self.wfs.kpt_u:
789 # self.wfs.pt.integrate(kpt.psit_nG, kpt.P_ani)
790 self.density.update(self.wfs)
792 def get_density(self):
793 """Returns the current density.
795 Parameters
796 ----------
797 None
799 """
800 return self.density
803def add_linear_field(wfs, spos_ac, a_nG, b_nG, strength, kpt):
804 """Adds (does NOT apply) linear field.
806 ::
808 f(x,y,z) = str_x * x + str_y * y + str_z * z to wavefunctions.
810 Parameters:
812 a_nG:
813 the wavefunctions
814 b_nG:
815 the result
816 strength: float[3]
817 strength of the linear field
818 kpt: KPoint
819 K-point
820 """
822 gd = wfs.gd
824 # apply local part of x to smooth wavefunctions psit_n
825 for i in range(gd.n_c[0]):
826 x = (i + gd.beg_c[0]) * gd.h_cv[0, 0]
827 b_nG[:, i, :, :] += (strength[0] * x) * a_nG[:, i, :, :]
829 # FIXME: combine y and z to one vectorized operation,
830 # i.e., make yz-array and take its product with a_nG
832 # apply local part of y to smooth wavefunctions psit_n
833 for i in range(gd.n_c[1]):
834 y = (i + gd.beg_c[1]) * gd.h_cv[1, 1]
835 b_nG[:, :, i, :] += (strength[1] * y) * a_nG[:, :, i, :]
837 # apply local part of z to smooth wavefunctions psit_n
838 for i in range(gd.n_c[2]):
839 z = (i + gd.beg_c[2]) * gd.h_cv[2, 2]
840 b_nG[:, :, :, i] += (strength[2] * z) * a_nG[:, :, :, i]
842 # apply the non-local part for each nucleus
844 # number of wavefunctions, psit_nG
845 n = len(a_nG)
846 P_ani = wfs.pt.dict(n)
847 wfs.pt.integrate(a_nG, P_ani, kpt.q)
849 coef_ani = {}
850 for a, P_ni in P_ani.items():
851 c0 = np.dot(spos_ac[a] * gd.cell_cv.diagonal(), strength)
852 cxyz = strength
853 # calculate coefficient
854 # ---------------------
855 #
856 # coeffs_ni =
857 # P_nj * c0 * 1_ij
858 # + P_nj * cx * x_ij
859 #
860 # where (see spherical_harmonics.py)
861 #
862 # 1_ij = sqrt(4pi) Delta_0ij
863 # y_ij = sqrt(4pi/3) Delta_1ij
864 # z_ij = sqrt(4pi/3) Delta_2ij
865 # x_ij = sqrt(4pi/3) Delta_3ij
866 # ...
868 Delta_iiL = wfs.setups[a].Delta_iiL
870 # 1_ij = sqrt(4pi) Delta_0ij
871 # y_ij = sqrt(4pi/3) Delta_1ij
872 # z_ij = sqrt(4pi/3) Delta_2ij
873 # x_ij = sqrt(4pi/3) Delta_3ij
874 oneij = np.sqrt(4 * np.pi) \
875 * np.dot(P_ni, Delta_iiL[:, :, 0])
876 yij = np.sqrt(4 * np.pi / 3) \
877 * np.dot(P_ni, Delta_iiL[:, :, 1])
878 zij = np.sqrt(4 * np.pi / 3) \
879 * np.dot(P_ni, Delta_iiL[:, :, 2])
880 xij = np.sqrt(4 * np.pi / 3) \
881 * np.dot(P_ni, Delta_iiL[:, :, 3])
883 # coefficients
884 # coefs_ni = sum_j ( <phi_i| f(x,y,z) | phi_j>
885 # - <phit_i| f(x,y,z) | phit_j> ) P_nj
886 coef_ani[a] = (c0 * oneij + cxyz[0] * xij + cxyz[1] * yij +
887 cxyz[2] * zij)
889 # add partial wave pt_nG to psit_nG with proper coefficient
890 wfs.pt.add(b_nG, coef_ani, kpt.q)