Coverage for gpaw/lcaotddft/propagators.py: 83%
338 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-19 00:19 +0000
1import numpy as np
3from numpy.linalg import inv, solve
5from ase.utils.timing import timer
7from gpaw.lcaotddft.hamiltonian import KickHamiltonian
8from gpaw import debug
9from gpaw.tddft.units import au_to_as
10from gpaw.utilities.scalapack import (pblas_simple_hemm, pblas_simple_gemm,
11 scalapack_inverse, scalapack_solve,
12 scalapack_tri2full)
15def create_propagator(name, **kwargs):
16 if name is None:
17 return create_propagator('sicn')
18 elif isinstance(name, Propagator):
19 return name
20 elif isinstance(name, dict):
21 kwargs.update(name)
22 return create_propagator(**kwargs)
23 elif name == 'sicn':
24 return SICNPropagator(**kwargs)
25 elif name == 'scpc':
26 return SelfConsistentPropagator(**kwargs)
27 elif name == 'ecn':
28 return ECNPropagator(**kwargs)
29 elif name.endswith('.ulm'):
30 return ReplayPropagator(name, **kwargs)
31 else:
32 raise ValueError('Unknown propagator: %s' % name)
35def equal(a, b, eps=1e-8):
36 return abs(a - b) < eps
39class Propagator:
41 def __init__(self):
42 object.__init__(self)
44 def initialize(self, paw):
45 self.timer = paw.timer
46 self.log = paw.log
48 def kick(self, ext, time):
49 raise NotImplementedError()
51 def propagate(self, time, time_step):
52 raise NotImplementedError()
54 def control_paw(self, paw):
55 raise NotImplementedError()
57 def todict(self):
58 raise NotImplementedError()
60 def get_description(self):
61 return '%s' % self.__class__.__name__
64class LCAOPropagator(Propagator):
66 def __init__(self):
67 super().__init__()
69 def initialize(self, paw):
70 super().initialize(paw)
71 self.wfs = paw.wfs
72 self.density = paw.density
73 self.hamiltonian = paw.td_hamiltonian
76class ReplayPropagator(LCAOPropagator):
78 def __init__(self, filename, update='all'):
79 from gpaw.lcaotddft.wfwriter import WaveFunctionReader
80 super().__init__()
81 self.filename = filename
82 self.update_mode = update
83 self.reader = WaveFunctionReader(self.filename)
84 self.read_index = 1
85 self.read_count = len(self.reader)
87 def _align_read_index(self, time):
88 while self.read_index < self.read_count:
89 r = self.reader[self.read_index]
90 if equal(r.time, time):
91 break
92 self.read_index += 1
93 if self.read_index == self.read_count:
94 raise RuntimeError('Time not found: %f' % time)
96 def _read(self):
97 reader = self.reader[self.read_index]
98 r = reader.wave_functions
99 self.wfs.read_wave_functions(r)
100 self.wfs.read_occupations(r)
101 self.read_index += 1
103 def kick(self, ext, time):
104 self._align_read_index(time)
105 # Check that this is the step after kick
106 assert not equal(self.reader[self.read_index].time,
107 self.reader[self.read_index + 1].time)
108 self._read()
109 self.hamiltonian.update(self.update_mode)
111 def propagate(self, time, time_step):
112 next_time = time + time_step
113 self._align_read_index(next_time)
114 self._read()
115 self.hamiltonian.update(self.update_mode)
116 return next_time
118 def control_paw(self, paw):
119 # Read the initial state
120 index = 1
121 r = self.reader[index]
122 assert r.action == 'init'
123 assert equal(r.time, paw.time)
124 self.read_index = index
125 self._read()
126 index += 1
127 # Read the rest
128 while index < self.read_count:
129 r = self.reader[index]
130 if r.action == 'init':
131 index += 1
132 elif r.action == 'kick':
133 assert equal(r.time, paw.time)
134 paw.absorption_kick(r.kick_strength)
135 assert equal(r.time, paw.time)
136 index += 1
137 elif r.action == 'propagate':
138 # Skip earlier times
139 if r.time < paw.time or equal(r.time, paw.time):
140 index += 1
141 continue
142 # Count the number of steps with the same time step
143 time = paw.time
144 time_step = r.time - time
145 iterations = 0
146 while index < self.read_count:
147 r = self.reader[index]
148 if (r.action != 'propagate' or
149 not equal(r.time - time, time_step)):
150 break
151 iterations += 1
152 time = r.time
153 index += 1
154 # Propagate
155 paw.propagate(time_step * au_to_as, iterations)
156 assert equal(time, paw.time)
157 else:
158 raise RuntimeError('Unknown action: %s' % r.action)
160 def __del__(self):
161 self.reader.close()
163 def todict(self):
164 return {'name': self.filename,
165 'update': self.update_mode}
167 def get_description(self):
168 lines = [self.__class__.__name__]
169 lines += [' File: %s' % (self.filename)]
170 lines += [' Update: %s' % (self.update_mode)]
171 return '\n'.join(lines)
174class ECNPropagator(LCAOPropagator):
176 def __init__(self):
177 super().__init__()
178 self.have_velocity_operator_matrix = False
180 def initialize(self, paw, hamiltonian=None):
181 super().initialize(paw)
182 if hamiltonian is not None:
183 self.hamiltonian = hamiltonian
185 ksl = self.wfs.ksl
186 using_blacs = ksl.using_blacs
187 if using_blacs:
188 from gpaw.blacs import Redistributor
189 self.log('BLACS Parallelization')
191 # Propagator function
192 self.propagate_wfs = self.propagate_wfs_blacs
194 # Parallel grid descriptors
195 grid = ksl.blockgrid
196 assert grid.nprow * grid.npcol == ksl.block_comm.size
197 self.mm_block_descriptor = ksl.mmdescriptor
198 self.Cnm_block_descriptor = grid.new_descriptor(ksl.bd.nbands,
199 ksl.nao,
200 ksl.blocksize,
201 ksl.blocksize)
202 self.CnM_unique_descriptor = ksl.nM_unique_descriptor
204 # Redistributors
205 self.Cnm2nM = Redistributor(ksl.block_comm,
206 self.Cnm_block_descriptor,
207 self.CnM_unique_descriptor)
208 self.CnM2nm = Redistributor(ksl.block_comm,
209 self.CnM_unique_descriptor,
210 self.Cnm_block_descriptor)
212 for kpt in self.wfs.kpt_u:
213 scalapack_tri2full(self.mm_block_descriptor, kpt.S_MM)
214 scalapack_tri2full(self.mm_block_descriptor, kpt.T_MM)
216 if self.density.gd.comm.rank != 0:
217 # This is a (0, 0) dummy array that is needed for
218 # redistributing between nM and nm block descriptor.
219 # See propagate_wfs() and also
220 # BlacsOrbitalLayouts.calculate_blocked_density_matrix()
221 self.dummy_C_nM = \
222 self.CnM_unique_descriptor.zeros(dtype=complex)
224 else:
225 # Propagator function
226 self.propagate_wfs = self.propagate_wfs_numpy
228 if debug and using_blacs:
229 nao = ksl.nao
230 self.MM_descriptor = grid.new_descriptor(nao, nao, nao, nao)
231 self.mm2MM = Redistributor(ksl.block_comm,
232 self.mm_block_descriptor,
233 self.MM_descriptor)
234 self.MM2mm = Redistributor(ksl.block_comm,
235 self.MM_descriptor,
236 self.mm_block_descriptor)
238 def calculate_velocity_operator_matrix(self):
239 if getattr(self, 'have_velocity_operator_matrix', False):
240 return
241 ksl = self.wfs.ksl
243 gcomm = self.wfs.gd.comm
244 manytci = self.wfs.manytci
245 Vkick_qvmM = manytci.O_qMM_T_qMM(gcomm,
246 ksl.Mstart,
247 ksl.Mstop,
248 ignore_upper=ksl.using_blacs,
249 derivative=True)[0] * (-1j)
251 my_atoms = self.wfs.atom_partition.my_indices
252 dnabla_vaii = {v: {a: -self.wfs.setups[a].nabla_iiv[:, :, v] * (-1j)
253 for a in my_atoms} for v in range(3)}
254 for kpt in self.wfs.kpt_u:
255 assert kpt.k == 0
257 for v in range(3):
258 self.wfs.atomic_correction.calculate(0, dnabla_vaii[v],
259 Vkick_qvmM[kpt.q][v],
260 ksl.Mstart, ksl.Mstop)
262 if ksl.using_blacs:
263 for Vkick_vmM in Vkick_qvmM:
264 for Vkick_mM in Vkick_vmM:
265 scalapack_tri2full(ksl.mMdescriptor, Vkick_mM)
267 q = 0
268 if ksl.using_blacs:
269 Vkick_vmm = self.wfs.ksl.distribute_overlap_matrix(
270 Vkick_qvmM[kpt.q]
271 )
272 else:
273 gcomm.sum(Vkick_qvmM[q])
274 Vkick_vmm = Vkick_qvmM[q]
276 for kpt in self.wfs.kpt_u:
277 assert kpt.q == 0
278 kpt.Vkick_vmm = Vkick_vmm
280 self.have_velocity_operator_matrix = True
282 def velocity_gauge_kick(self, magnitude, direction, time):
283 self.calculate_velocity_operator_matrix()
284 for kpt in self.wfs.kpt_u:
285 kpt.A_MM = (
286 -magnitude * np.einsum('v,vMN->MN', direction, kpt.Vkick_vmm)
287 )
289 # Update Hamiltonian (and density)
290 self.hamiltonian.update()
292 def kick(self, ext, time):
293 # Propagate
294 get_matrix = self.wfs.eigensolver.calculate_hamiltonian_matrix
295 kick_hamiltonian = KickHamiltonian(self.hamiltonian.hamiltonian,
296 self.density, ext)
297 for kpt in self.wfs.kpt_u:
298 Vkick_MM = get_matrix(kick_hamiltonian, self.wfs, kpt,
299 add_kinetic=False, root=-1)
300 for i in range(10):
301 self.propagate_wfs(kpt.C_nM, kpt.C_nM, kpt.S_MM, Vkick_MM, 0.1)
303 # Update Hamiltonian (and density)
304 self.hamiltonian.update()
306 def propagate(self, time, time_step):
307 get_H_MM = self.hamiltonian.get_hamiltonian_matrix
308 for kpt in self.wfs.kpt_u:
309 H_MM = get_H_MM(kpt, time)
310 self.propagate_wfs(kpt.C_nM, kpt.C_nM, kpt.S_MM, H_MM, time_step)
311 self.hamiltonian.update()
312 return time + time_step
314 @timer('Linear solve')
315 def propagate_wfs_blacs(self, source_C_nM, target_C_nM, S_mm, H_mm, dt):
316 # XXX, Preallocate?
317 target_C_nm = self.Cnm_block_descriptor.empty(dtype=complex)
318 source_C_nm = self.Cnm_block_descriptor.empty(dtype=complex)
320 # C_nM is duplicated over all ranks in gd.comm.
321 # Master rank will provide the actual data and other
322 # ranks use a dummy array in redistribute().
323 if self.density.gd.comm.rank != 0:
324 source = self.dummy_C_nM
325 else:
326 source = source_C_nM
327 self.CnM2nm.redistribute(source, source_C_nm)
329 # Note: tri2full for S_mm and T_mm is done already in initialize().
330 # H_mm seems to be a full matrix as we are working with complex
331 # dtype, so no need to do tri2full here again XXX
332 # scalapack_tri2full(self.mm_block_descriptor, H_mm)
333 SjH_mm = S_mm + (0.5j * dt) * H_mm
335 # 1. target = (S - 0.5j*H*dt) * source
336 pblas_simple_gemm(self.Cnm_block_descriptor,
337 self.mm_block_descriptor,
338 self.Cnm_block_descriptor,
339 source_C_nm,
340 SjH_mm,
341 target_C_nm,
342 transb='C')
344 # 2. target = (S + 0.5j*H*dt)^-1 * target
345 scalapack_solve(self.mm_block_descriptor,
346 self.Cnm_block_descriptor,
347 SjH_mm,
348 target_C_nm)
350 # C_nM is duplicated over all ranks in gd.comm.
351 # Master rank will receive the data and other
352 # ranks use a dummy array in redistribute()
353 if self.density.gd.comm.rank != 0:
354 target = self.dummy_C_nM
355 else:
356 target = target_C_nM
357 self.Cnm2nM.redistribute(target_C_nm, target)
359 # Broadcast the new C_nM to all ranks in gd.comm
360 self.density.gd.comm.broadcast(target_C_nM, 0)
362 @timer('Linear solve')
363 def propagate_wfs_numpy(self, source_C_nM, target_C_nM, S_MM, H_MM, dt):
364 SjH_MM = S_MM + (0.5j * dt) * H_MM
365 target_C_nM[:] = np.dot(source_C_nM, SjH_MM.conj().T)
366 target_C_nM[:] = solve(SjH_MM.T, target_C_nM.T).T
368 def blacs_mm_to_global(self, H_mm):
369 if not debug:
370 raise RuntimeError('Use debug mode')
371 # Someone could verify that this works and remove the error.
372 raise NotImplementedError('Method untested and thus unreliable')
373 target = self.MM_descriptor.empty(dtype=complex)
374 self.mm2MM.redistribute(H_mm, target)
375 self.wfs.world.barrier()
376 return target
378 def blacs_nm_to_global(self, C_nm):
379 # Someone could verify that this works and remove the error.
380 raise NotImplementedError('Method untested and thus unreliable')
381 target = self.CnM_unique_descriptor.empty(dtype=complex)
382 self.Cnm2nM.redistribute(C_nm, target)
383 self.wfs.world.barrier()
384 return target
386 def todict(self):
387 return {'name': 'ecn'}
390class SICNPropagator(ECNPropagator):
392 def __init__(self):
393 super().__init__()
395 def initialize(self, paw):
396 super().initialize(paw)
397 # Allocate kpt.C2_nM arrays
398 for kpt in self.wfs.kpt_u:
399 kpt.C2_nM = np.empty_like(kpt.C_nM)
401 def propagate(self, time, time_step):
402 get_H_MM = self.hamiltonian.get_hamiltonian_matrix
403 # --------------
404 # Predictor step
405 # --------------
406 # 1. Store current C_nM
407 self.save_wfs() # kpt.C2_nM = kpt.C_nM
408 for kpt in self.wfs.kpt_u:
409 # H_MM(t) = <M|H(t)|M>
410 kpt.H0_MM = get_H_MM(kpt, time)
411 # 2. Solve Psi(t+dt) from
412 # (S_MM - 0.5j*H_MM(t)*dt) Psi(t+dt)
413 # = (S_MM + 0.5j*H_MM(t)*dt) Psi(t)
414 self.propagate_wfs(kpt.C_nM, kpt.C_nM, kpt.S_MM, kpt.H0_MM,
415 time_step)
416 # ---------------
417 # Propagator step
418 # ---------------
419 # 1. Calculate H(t+dt)
420 self.hamiltonian.update()
421 for kpt in self.wfs.kpt_u:
422 # 2. Estimate H(t+0.5*dt) ~ 0.5 * [ H(t) + H(t+dt) ]
423 kpt.H0_MM += get_H_MM(kpt, time + time_step)
424 kpt.H0_MM *= 0.5
425 # 3. Solve Psi(t+dt) from
426 # (S_MM - 0.5j*H_MM(t+0.5*dt)*dt) Psi(t+dt)
427 # = (S_MM + 0.5j*H_MM(t+0.5*dt)*dt) Psi(t)
428 self.propagate_wfs(kpt.C2_nM, kpt.C_nM, kpt.S_MM, kpt.H0_MM,
429 time_step)
430 kpt.H0_MM = None
431 # 4. Calculate new Hamiltonian (and density)
432 self.hamiltonian.update()
433 return time + time_step
435 def save_wfs(self):
436 for kpt in self.wfs.kpt_u:
437 kpt.C2_nM[:] = kpt.C_nM
439 def todict(self):
440 return {'name': 'sicn'}
443# ToDo: Should there be an abstract baseclass for SelfConsistentPropagator
444# and SICNPropagator instead of inheriting from ECNPropagator?
445class SelfConsistentPropagator(SICNPropagator):
446 """
447 This is an actual Predictor-Corrector propagator that uses the SICN
448 and combines it with an actual Corrector step. This is identical to
449 SICN for a very low tolerance of e.g. 1e-2. The higher the tolerance,
450 the better energy etc will be preserved by the propagator. Notice,
451 that the standard SICN accuracy is often sufficient, but some
452 routines (like the 3rd time-derivative in the RRemission class)
453 require higher accuracy for reliable predictions. The PC update
454 become especially important for large time-steps. Try for instance
455 dt=100 and propagte for a few thousand step, compare SICN vs SCPC.
456 You will notice an artifical exponential decay of the SICN dipole
457 after kick while SCPC will preserve the dipole oscillations.
458 """
459 def __init__(self, tolerance=1e-8, max_pc_iterations=20):
460 super().__init__()
461 self.tolerance = tolerance
462 self.max_pc_iterations = max_pc_iterations
463 self.last_pc_iterations = 0
465 def propagate(self, time, time_step):
466 """
467 Since the propagate + update call will change the result
468 for H0 at time t, we have to somehow safe the previous H0
469 in order to estimate the intermediate H0 at t+dt/2.
470 """
471 prevH0 = []
472 get_H_MM = self.hamiltonian.get_hamiltonian_matrix
473 # --------------
474 # Predictor step
475 # --------------
476 # 1. Store current C_nM
477 self.save_wfs() # kpt.C2_nM = kpt.C_nM
478 for kpt in self.wfs.kpt_u:
479 # H_MM(t) = <M|H(t)|M>
480 kpt.H0_MM = get_H_MM(kpt, time)
481 prevH0.append(kpt.H0_MM)
482 # 2. Solve Psi(t+dt) from
483 # (S_MM - 0.5j*H_MM(t)*dt) Psi(t+dt)
484 # = (S_MM + 0.5j*H_MM(t)*dt) Psi(t)
485 self.propagate_wfs(kpt.C_nM, kpt.C_nM, kpt.S_MM, kpt.H0_MM,
486 time_step)
487 self.hamiltonian.update()
488 for last_pc_iterations in range(self.max_pc_iterations):
489 self.last_pc_iterations = last_pc_iterations
490 # ---------------
491 # Propagator step
492 # ---------------
493 # 1. Calculate H(t+dt)
494 itkpt = - 1
495 for kpt in self.wfs.kpt_u:
496 # 2. Estimate H(t+0.5*dt) ~ 0.5 * [ H(t) + H(t+dt) ]
497 itkpt += 1
498 kpt.H0_MM = prevH0[itkpt] + get_H_MM(kpt, time + time_step)
499 kpt.H0_MM *= 0.5
500 # 3. Solve Psi(t+dt) from
501 # (S_MM - 0.5j*H_MM(t+0.5*dt)*dt) Psi(t+dt)
502 # = (S_MM + 0.5j*H_MM(t+0.5*dt)*dt) Psi(t)
503 self.propagate_wfs(kpt.C2_nM, kpt.C_nM, kpt.S_MM, kpt.H0_MM,
504 time_step)
505 kpt.H0_MM = None
507 prev_dipole_v = self.density.calculate_dipole_moment()
508 # 4. Calculate new Hamiltonian (and density)
509 self.hamiltonian.update()
510 dipole_v = self.density.calculate_dipole_moment()
511 if np.sum(np.abs(dipole_v - prev_dipole_v)) < self.tolerance:
512 break
513 if last_pc_iterations == self.max_pc_iterations - 1:
514 raise RuntimeError('The SCPC propagator required too ',
515 'many iterations to reach the ',
516 'demanded accuracy.')
517 return time + time_step
519 def todict(self):
520 return {'name': 'scpc', 'tolerance': self.tolerance,
521 'max_pc_iterations': self.max_pc_iterations}
524class TaylorPropagator(Propagator):
526 def __init__(self):
527 super().__init__()
528 raise NotImplementedError('TaylorPropagator not implemented')
530 def initialize(self, paw):
531 if 1:
532 # XXX to propagator class
533 if self.propagator == 'taylor' and self.blacs:
534 # cholS_mm = self.mm_block_descriptor.empty(dtype=complex)
535 for kpt in self.wfs.kpt_u:
536 kpt.invS_MM = kpt.S_MM.copy()
537 scalapack_inverse(self.mm_block_descriptor,
538 kpt.invS_MM, 'L')
539 if self.propagator == 'taylor' and not self.blacs:
540 tmp = inv(self.wfs.kpt_u[0].S_MM)
541 self.wfs.kpt_u[0].invS = tmp
543 def taylor_propagator(self, sourceC_nM, targetC_nM, S_MM, H_MM, dt):
544 self.timer.start('Taylor propagator')
546 if self.blacs:
547 # XXX, Preallocate
548 target_blockC_nm = self.Cnm_block_descriptor.empty(dtype=complex)
549 if self.density.gd.comm.rank != 0:
550 # XXX Fake blacks nbands, nao, nbands, nao grid because some
551 # weird asserts
552 # (these are 0,x or x,0 arrays)
553 sourceC_nM = self.CnM_unique_descriptor.zeros(dtype=complex)
555 # Zeroth order taylor to target
556 self.CnM2nm.redistribute(sourceC_nM, target_blockC_nm)
558 # XXX, preallocate, optimize use of temporal arrays
559 temp_blockC_nm = target_blockC_nm.copy()
560 temp2_blockC_nm = target_blockC_nm.copy()
562 order = 4
563 assert self.wfs.kd.comm.size == 1
564 for n in range(order):
565 # Multiply with hamiltonian
566 pblas_simple_hemm(self.mm_block_descriptor,
567 self.Cnm_block_descriptor,
568 self.Cnm_block_descriptor,
569 H_MM,
570 temp_blockC_nm,
571 temp2_blockC_nm, side='R')
572 # XXX: replace with not simple gemm
573 temp2_blockC_nm *= -1j * dt / (n + 1)
574 # Multiply with inverse overlap
575 pblas_simple_hemm(self.mm_block_descriptor,
576 self.Cnm_block_descriptor,
577 self.Cnm_block_descriptor,
578 self.wfs.kpt_u[0].invS_MM, # XXX
579 temp2_blockC_nm,
580 temp_blockC_nm, side='R')
581 target_blockC_nm += temp_blockC_nm
582 if self.density.gd.comm.rank != 0: # Todo: Change to gd.rank
583 # XXX Fake blacks nbands, nao, nbands, nao grid because
584 # some weird asserts
585 # (these are 0,x or x,0 arrays)
586 target = self.CnM_unique_descriptor.zeros(dtype=complex)
587 else:
588 target = targetC_nM
589 self.Cnm2nM.redistribute(target_blockC_nm, target)
591 self.density.gd.comm.broadcast(targetC_nM, 0)
592 else:
593 assert self.wfs.kd.comm.size == 1
594 if self.density.gd.comm.rank == 0:
595 targetC_nM[:] = sourceC_nM[:]
596 tempC_nM = sourceC_nM.copy()
597 order = 4
598 for n in range(order):
599 tempC_nM[:] = \
600 np.dot(self.wfs.kpt_u[0].invS,
601 np.dot(H_MM, 1j * dt / (n + 1) *
602 tempC_nM.T.conjugate())).T.conjugate()
603 targetC_nM += tempC_nM
604 self.density.gd.comm.broadcast(targetC_nM, 0)
606 self.timer.stop('Taylor propagator')