Coverage for gpaw/lcaotddft/hamiltonian.py: 90%
203 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1from gpaw.mixer import DummyMixer
2from gpaw.xc import XC
4from gpaw.lcaotddft.laser import create_laser
5from gpaw.lcaotddft.utilities import read_uMM
6from gpaw.lcaotddft.utilities import read_wuMM
7from gpaw.lcaotddft.utilities import write_uMM
8from gpaw.lcaotddft.utilities import write_wuMM
11class TimeDependentPotential:
12 def __init__(self):
13 self.ext_i = []
14 self.laser_i = []
15 self.initialized = False
17 def add(self, ext, laser):
18 self.ext_i.append(ext)
19 self.laser_i.append(laser)
21 def initialize(self, paw):
22 if self.initialized:
23 return
25 self.ksl = paw.wfs.ksl
26 self.kd = paw.wfs.kd
27 self.kpt_u = paw.wfs.kpt_u
28 get_matrix = paw.wfs.eigensolver.calculate_hamiltonian_matrix
29 self.V_iuMM = []
30 for ext in self.ext_i:
31 V_uMM = []
32 hamiltonian = KickHamiltonian(paw.hamiltonian, paw.density, ext)
33 for kpt in paw.wfs.kpt_u:
34 V_MM = get_matrix(hamiltonian, paw.wfs, kpt,
35 add_kinetic=False, root=-1)
36 V_uMM.append(V_MM)
37 self.V_iuMM.append(V_uMM)
38 self.Ni = len(self.ext_i)
39 self.initialized = True
41 def get_MM(self, u, time):
42 V_MM = self.laser_i[0].strength(time) * self.V_iuMM[0][u]
43 for i in range(1, self.Ni):
44 V_MM += self.laser_i[i].strength(time) * self.V_iuMM[i][u]
45 return V_MM
47 def write(self, writer):
48 writer.write(Ni=self.Ni)
49 writer.write(laser_i=[laser.todict() for laser in self.laser_i])
50 write_wuMM(self.kd, self.ksl, writer, 'V_iuMM',
51 self.V_iuMM, range(self.Ni))
53 def read(self, reader):
54 self.Ni = reader.Ni
55 self.laser_i = [create_laser(**laser) for laser in reader.laser_i]
56 self.V_iuMM = read_wuMM(self.kpt_u, self.ksl, reader, 'V_iuMM',
57 range(self.Ni))
58 self.initialized = True
61class KickHamiltonian:
62 def __init__(self, ham, dens, ext):
63 nspins = ham.nspins
64 vext_g = ext.get_potential(ham.finegd)
65 vext_G = ham.restrict_and_collect(vext_g)
66 self.vt_sG = ham.gd.empty(nspins)
67 for s in range(nspins):
68 self.vt_sG[s] = vext_G
70 dH_asp = ham.setups.empty_atomic_matrix(nspins,
71 ham.atom_partition)
72 W_aL = dens.ghat.dict()
73 dens.ghat.integrate(vext_g, W_aL)
74 # XXX this is a quick hack to get the distribution right
75 # It'd be better to have these distributions abstracted elsewhere.
76 # The idea is (thanks Ask, see discussion in !910):
77 # * gd has D cores and coarse grid
78 # * aux_gd is the same grid as gd but with either D or
79 # K * B * D cores (with augment_grids = True)
80 # * finegd is then aux_gd.refine().
81 # In integrals like W_aL, the atomic quantities are distributed
82 # according to the domains of those atoms (so finegd for W_aL,
83 # but gd for dH_asp).
84 # And, some things (atomic XC corrections) are calculated evenly
85 # distributed among all cores.
86 dHaux_asp = ham.atomdist.to_aux(dH_asp)
87 for a, W_L in W_aL.items():
88 setup = dens.setups[a]
89 dH_p = setup.Delta_pL @ W_L
90 for s in range(nspins):
91 dHaux_asp[a][s] = dH_p
92 self.dH_asp = ham.atomdist.from_aux(dHaux_asp)
95class TimeDependentHamiltonian:
96 def __init__(self, fxc=None, td_potential=None, scale=None,
97 rremission=None):
98 assert fxc is None or isinstance(fxc, str)
99 self.fxc_name = fxc
100 if isinstance(td_potential, dict):
101 td_potential = [td_potential]
102 if isinstance(td_potential, list):
103 pot_i = td_potential
104 td_potential = TimeDependentPotential()
105 for pot in pot_i:
106 td_potential.add(**pot)
107 self.td_potential = td_potential
108 self.has_scale = scale is not None
109 if self.has_scale:
110 self.scale = scale
111 self.rremission = rremission
113 def write(self, writer):
114 if self.has_fxc:
115 self.write_fxc(writer.child('fxc'))
116 if self.has_scale:
117 self.write_scale(writer.child('scale'))
118 if self.td_potential is not None:
119 self.td_potential.write(writer.child('td_potential'))
120 if self.rremission is not None:
121 self.rremission.write(writer.child('rremission'))
123 def write_fxc(self, writer):
124 writer.write(name=self.fxc_name)
125 write_uMM(self.wfs.kd, self.wfs.ksl, writer, 'deltaXC_H_uMM',
126 self.deltaXC_H_uMM)
128 def write_scale(self, writer):
129 writer.write(scale=self.scale)
130 write_uMM(self.wfs.kd, self.wfs.ksl, writer, 'scale_H_uMM',
131 self.scale_H_uMM)
133 def read(self, reader):
134 if 'fxc' in reader:
135 self.read_fxc(reader.fxc)
136 if 'scale' in reader:
137 self.read_scale(reader.scale)
138 if 'td_potential' in reader:
139 assert self.td_potential is None
140 self.td_potential = TimeDependentPotential()
141 self.td_potential.ksl = self.wfs.ksl
142 self.td_potential.kd = self.wfs.kd
143 self.td_potential.kpt_u = self.wfs.kpt_u
144 self.td_potential.read(reader.td_potential)
145 if 'rremission' in reader:
146 assert self.rremission is None
147 from gpaw.lcaotddft.qed import RRemission
148 self.rremission = RRemission.from_reader(reader.rremission)
150 def read_fxc(self, reader):
151 assert self.fxc_name is None or self.fxc_name == reader.name
152 self.fxc_name = reader.name
153 self.deltaXC_H_uMM = read_uMM(self.wfs.kpt_u, self.wfs.ksl, reader,
154 'deltaXC_H_uMM')
156 def read_scale(self, reader):
157 assert not self.has_scale or self.scale == reader.scale
158 self.has_scale = True
159 self.scale = reader.scale
160 self.scale_H_uMM = read_uMM(self.wfs.kpt_u, self.wfs.ksl, reader,
161 'scale_H_uMM')
163 def initialize(self, paw):
164 self.timer = paw.timer
165 self.timer.start('Initialize TDDFT Hamiltonian')
166 self.wfs = paw.wfs
167 self.density = paw.density
168 self.hamiltonian = paw.hamiltonian
169 niter = paw.niter
170 if self.rremission is not None:
171 self.rremission.initialize(paw)
173 # Reset the density mixer
174 # XXX: density mixer is not written to the gpw file
175 # XXX: so we need to set it always
176 self.density.set_mixer(DummyMixer())
177 self.update()
179 # Initialize fxc
180 self.initialize_fxc(niter)
182 # Initialize scale
183 self.initialize_scale(niter)
185 # Initialize td_potential
186 if self.td_potential is not None:
187 self.td_potential.initialize(paw)
189 self.timer.stop('Initialize TDDFT Hamiltonian')
191 def initialize_fxc(self, niter):
192 self.has_fxc = self.fxc_name is not None
193 if not self.has_fxc:
194 return
195 self.timer.start('Initialize fxc')
196 # XXX: Similar functionality is available in
197 # paw.py: PAW.linearize_to_xc(self, newxc)
198 # See test/lcaotddft/fxc_vs_linearize.py
200 # Calculate deltaXC: 1. take current H_MM
201 if niter == 0:
202 def get_H_MM(kpt):
203 return self.get_hamiltonian_matrix(kpt, time=0.0,
204 addfxc=False, addpot=False,
205 scale=False)
206 self.deltaXC_H_uMM = []
207 for kpt in self.wfs.kpt_u:
208 self.deltaXC_H_uMM.append(get_H_MM(kpt))
210 # Update hamiltonian.xc
211 if self.fxc_name == 'RPA':
212 xc_name = 'null'
213 else:
214 xc_name = self.fxc_name
215 # XXX: xc is not written to the gpw file
216 # XXX: so we need to set it always
217 xc = XC(xc_name)
218 xc.initialize(self.density, self.hamiltonian, self.wfs)
219 xc.set_positions(self.hamiltonian.spos_ac)
220 self.hamiltonian.xc = xc
221 self.update()
223 # Calculate deltaXC: 2. update with new H_MM
224 if niter == 0:
225 for u, kpt in enumerate(self.wfs.kpt_u):
226 self.deltaXC_H_uMM[u] -= get_H_MM(kpt)
227 self.timer.stop('Initialize fxc')
229 def initialize_scale(self, niter):
230 if not self.has_scale:
231 return
232 self.timer.start('Initialize scale')
234 # Take current H_MM and multiply with scale
235 if niter == 0:
236 def get_H_MM(kpt):
237 return self.get_hamiltonian_matrix(kpt, time=0.0,
238 addfxc=False, addpot=False,
239 scale=False)
240 self.scale_H_uMM = []
241 for kpt in self.wfs.kpt_u:
242 self.scale_H_uMM.append((1 - self.scale) * get_H_MM(kpt))
243 self.timer.stop('Initialize scale')
245 def update_projectors(self):
246 self.timer.start('Update projectors')
247 for kpt in self.wfs.kpt_u:
248 self.wfs.atomic_correction.calculate_projections(self.wfs, kpt)
249 self.timer.stop('Update projectors')
251 def get_hamiltonian_matrix(self, kpt, time, addfxc=True, addpot=True,
252 scale=True):
253 self.timer.start('Calculate H_MM')
254 kpt_rank, q = self.wfs.kd.get_rank_and_index(kpt.k)
255 u = q * self.wfs.nspins + kpt.s
256 assert kpt_rank == self.wfs.kd.comm.rank
258 get_matrix = self.wfs.eigensolver.calculate_hamiltonian_matrix
259 H_MM = get_matrix(self.hamiltonian, self.wfs, kpt, root=-1)
260 if addfxc and self.has_fxc:
261 H_MM += self.deltaXC_H_uMM[u]
263 if self.rremission is not None:
264 H_MM += self.rremission.vradiationreaction(kpt, time)
266 if scale and self.has_scale:
267 H_MM *= self.scale
268 H_MM += self.scale_H_uMM[u]
269 if addpot and self.td_potential is not None:
270 H_MM += self.td_potential.get_MM(u, time)
271 self.timer.stop('Calculate H_MM')
272 if hasattr(kpt, 'A_MM'):
273 return H_MM + kpt.A_MM
274 else:
275 return H_MM
277 def update(self, mode='all'):
278 self.timer.start('Update TDDFT Hamiltonian')
279 if mode in ['all', 'density']:
280 self.update_projectors()
281 self.density.update(self.wfs)
282 if mode in ['all']:
283 self.hamiltonian.update(self.density)
284 self.timer.stop('Update TDDFT Hamiltonian')