Coverage for gpaw/new/rttddft.py: 31%
148 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from __future__ import annotations
3from typing import Generator, NamedTuple
5import numpy as np
6from numpy.linalg import solve
8from ase.units import Bohr, Hartree
10from gpaw.external import ExternalPotential, ConstantElectricField
11from gpaw.typing import Vector
12from gpaw.mpi import world
13from gpaw.new.ase_interface import ASECalculator
14from gpaw.new.calculation import DFTState, DFTCalculation
15from gpaw.new.lcao.hamiltonian import HamiltonianMatrixCalculator
16from gpaw.new.lcao.wave_functions import LCAOWaveFunctions
17from gpaw.new.gpw import read_gpw
18from gpaw.new.pot_calc import PotentialCalculator
19from gpaw.tddft.units import asetime_to_autime, autime_to_asetime, au_to_eA
20from gpaw.utilities.timing import nulltimer
23class TDAlgorithm:
25 def kick(self,
26 state: DFTState,
27 pot_calc: PotentialCalculator,
28 dm_calc: HamiltonianMatrixCalculator):
29 raise NotImplementedError()
31 def propagate(self,
32 time_step: float,
33 state: DFTState,
34 pot_calc: PotentialCalculator,
35 ham_calc: HamiltonianMatrixCalculator):
36 raise NotImplementedError()
38 def get_description(self):
39 return self.__class__.__name__
42def propagate_wave_functions_numpy(source_C_nM: np.ndarray,
43 target_C_nM: np.ndarray,
44 S_MM: np.ndarray,
45 H_MM: np.ndarray,
46 dt: float):
47 SjH_MM = S_MM + (0.5j * dt) * H_MM
48 target_C_nM[:] = source_C_nM @ SjH_MM.conj().T
49 target_C_nM[:] = solve(SjH_MM.T, target_C_nM.T).T
52class ECNAlgorithm(TDAlgorithm):
54 def kick(self,
55 state: DFTState,
56 pot_calc: PotentialCalculator,
57 dm_calc: HamiltonianMatrixCalculator):
58 """Propagate wavefunctions by delta-kick.
60 ::
62 0+
63 ^ -1 / ^ -1
64 U(0+, 0) = T exp[-iS | δ(τ) V (r) dτ ] = T exp[-iS V (r)]
65 / ext ext
66 0
68 (1) Calculate propagator U(0+, 0)
69 (2) Update wavefunctions ψ_n(0+) ← U(0+, 0) ψ_n(0)
70 (3) Update density and hamiltonian H(0+)
72 Parameters
73 ----------
74 state
75 Current state of the wave functions, that is to be updated
76 pot_calc
77 The potential calculator
78 dm_calc
79 Dipole moment operator calculator, which contains the dipole moment
80 operator
81 """
82 for wfs in state.ibzwfs:
83 assert isinstance(wfs, LCAOWaveFunctions)
84 V_MM = dm_calc.calculate_matrix(wfs)
86 # Phi_n <- U(0+, 0) Phi_n
87 nkicks = 10
88 for i in range(nkicks):
89 propagate_wave_functions_numpy(wfs.C_nM.data, wfs.C_nM.data,
90 wfs.S_MM.data,
91 V_MM.data, 1 / nkicks)
92 # Update density
93 state.density.update(state.ibzwfs)
95 # Calculate Hamiltonian H(t+dt) = H[n[Phi_n]]
96 state.potential, state.energies, _ = pot_calc.calculate(
97 state.density, state.potential.vHt_x)
99 def propagate(self,
100 time_step: float,
101 state: DFTState,
102 pot_calc: PotentialCalculator,
103 ham_calc: HamiltonianMatrixCalculator):
104 """ One explicit Crank-Nicolson propagation step, i.e.
106 (1) Calculate propagator U[H(t)]
107 (2) Update wavefunctions ψ_n(t+dt) ← U[H(t)] ψ_n(t)
108 (3) Update density and hamiltonian H(t+dt)
109 """
110 for wfs in state.ibzwfs:
111 assert isinstance(wfs, LCAOWaveFunctions)
112 H_MM = ham_calc.calculate_matrix(wfs)
114 # Phi_n <- U[H(t)] Phi_n
115 propagate_wave_functions_numpy(wfs.C_nM.data, wfs.C_nM.data,
116 wfs.S_MM.data,
117 H_MM.data, time_step)
118 # Update density
119 state.density.update(state.ibzwfs)
121 # Calculate Hamiltonian H(t+dt) = H[n[Phi_n]]
122 state.potential, state.energies, _ = pot_calc.calculate(
123 state.density, state.potential.vHt_x)
126class RTTDDFTHistory:
128 kick_strength: Vector | None # Kick strength in atomic units
129 niter: int # Number of propagation steps
130 time: float # Simulation time in atomic units
132 def __init__(self):
133 """Object that keeps track of the RT-TDDFT history, that is
135 - Has a kick been performed?
136 - The number of propagation states performed
137 """
138 self.kick_strength = None
139 self.niter = 0
140 self.time = 0.0
142 def absorption_kick(self,
143 kick_strength: Vector):
144 """ Store the kick strength in history
146 At most one kick can be done, and it must happen before any
147 propagation steps
149 Parameters
150 ----------
151 kick_strength
152 Strength of the kick in atomic units
153 """
154 assert self.niter == 0, 'Cannot kick if already propagated'
155 assert self.kick_strength is None, 'Cannot kick if already kicked'
156 self.kick_strength = np.array(kick_strength, dtype=float)
158 def propagate(self,
159 time_step: float) -> float:
160 """ Increment the number of propagation steps and simulation time
161 in history
163 Parameters
164 ----------
165 time_step
166 Time step in atomic units
168 Returns
169 -------
170 The new simulation time in atomic units
171 """
172 self.niter += 1
173 self.time += time_step
175 return self.time
177 def todict(self):
178 absorption_kick = self.absorption_kick
179 if absorption_kick is not None:
180 absorption_kick = absorption_kick.tolist()
181 return {'niter': self.niter, 'time': self.time,
182 'absorption_kick': absorption_kick}
185class RTTDDFTResult(NamedTuple):
187 """ Results are stored in atomic units, but displayed to the user in
188 ASE units
189 """
191 time: float # Time in atomic units
192 dipolemoment: Vector # Dipole moment in atomic units
194 def __repr__(self):
195 timestr = f'{self.time * autime_to_asetime:.3f} Å√(u/eV)'
196 dmstr = ', '.join([f'{dm * au_to_eA:10.4g}'
197 for dm in self.dipolemoment])
198 dmstr = f'[{dmstr}]'
200 return (f'{self.__class__.__name__}: '
201 f'(time: {timestr}, dipolemoment: {dmstr} eÅ)')
204class RTTDDFT:
205 def __init__(self,
206 state: DFTState,
207 pot_calc: PotentialCalculator,
208 hamiltonian,
209 history: RTTDDFTHistory,
210 propagator: TDAlgorithm | None = None):
211 if propagator is None:
212 propagator = ECNAlgorithm()
214 self.state = state
215 self.pot_calc = pot_calc
216 self.propagator = propagator
217 self.hamiltonian = hamiltonian
218 self.history = history
220 self.kick_ext: ExternalPotential | None = None
222 # Dipole moment operators in each Cartesian direction
223 self.dm_operator_c: list[HamiltonianMatrixCalculator] | None = None
225 self.timer = nulltimer
226 self.log = print
228 self.ham_calc = hamiltonian.create_hamiltonian_matrix_calculator(state)
230 @classmethod
231 def from_dft_calculation(cls,
232 calc: ASECalculator | DFTCalculation,
233 propagator: TDAlgorithm | None = None):
235 if isinstance(calc, DFTCalculation):
236 dft = calc
237 else:
238 assert calc.dft is not None
239 dft = calc.dft
241 state = dft.get_state()
242 pot_calc = dft.pot_calc
243 hamiltonian = dft.scf_loop.hamiltonian
244 history = RTTDDFTHistory()
246 return cls(state, pot_calc, hamiltonian, propagator=propagator,
247 history=history)
249 @classmethod
250 def from_dft_file(cls,
251 filepath: str,
252 propagator: TDAlgorithm | None = None):
253 _, dft, params, builder = read_gpw(filepath,
254 log='-',
255 comm=world,
256 dtype=complex)
258 state = dft.get_state()
259 pot_calc = dft.pot_calc
260 hamiltonian = builder.create_hamiltonian_operator()
261 history = RTTDDFTHistory()
263 return cls(state, pot_calc, hamiltonian, propagator=propagator,
264 history=history)
266 def absorption_kick(self,
267 kick_strength: Vector):
268 """Kick with a weak electric field.
270 Parameters
271 ----------
272 kick_strength
273 Strength of the kick in atomic units
274 """
275 with self.timer('Kick'):
276 kick_strength = np.array(kick_strength, dtype=float)
277 self.history.absorption_kick(kick_strength)
279 magnitude = np.sqrt(np.sum(kick_strength**2))
280 direction = kick_strength / magnitude
281 dirstr = [f'{d:.4f}' for d in direction]
283 self.log('---- Applying absorption kick')
284 self.log(f'---- Magnitude: {magnitude:.8f} Hartree/Bohr')
285 self.log(f'---- Direction: {dirstr}')
287 # Create Hamiltonian object for absorption kick
288 cef = ConstantElectricField(magnitude * Hartree / Bohr, direction)
290 # Propagate kick
291 self.kick(cef)
293 def kick(self,
294 ext: ExternalPotential):
295 """Kick with any external potential.
297 Note that unless this function is called by absorption_kick, the kick
298 is not logged in history
300 Parameters
301 ----------
302 ext
303 External potential
304 """
305 with self.timer('Kick'):
306 self.log('---- Applying kick')
307 self.log(f'---- {ext}')
309 dm_operator_calc = self.hamiltonian.create_kick_matrix_calculator(
310 self.state, ext, self.pot_calc)
312 self.kick_ext = ext
314 # Propagate kick
315 self.propagator.kick(state=self.state,
316 pot_calc=self.pot_calc,
317 dm_calc=dm_operator_calc)
319 def ipropagate(self,
320 time_step: float = 10.0,
321 maxiter: int = 2000,
322 ) -> Generator[RTTDDFTResult, None, None]:
323 """Propagate the electronic system.
325 Parameters
326 ----------
327 time_step
328 Time step in ASE time units Å√(u/eV)
329 iterations
330 Number of propagation steps
331 """
333 time_step = time_step * asetime_to_autime
335 for iteration in range(maxiter):
336 self.propagator.propagate(time_step,
337 state=self.state,
338 pot_calc=self.pot_calc,
339 ham_calc=self.ham_calc)
340 time = self.history.propagate(time_step)
341 # TODO This seems to be broken
342 # dipolemoment = self.state.density.calculate_dipole_moment(
343 # self.pot_calc.relpos_ac)
344 dipolemoment_xv = [
345 self.calculate_dipole_moment(wfs) # type: ignore
346 for wfs in self.state.ibzwfs]
347 dipolemoment_v = np.sum(dipolemoment_xv, axis=0)
348 result = RTTDDFTResult(time=time, dipolemoment=dipolemoment_v)
349 yield result
351 def calculate_dipole_moment(self,
352 wfs: LCAOWaveFunctions) -> np.ndarray:
353 """ Calculates the dipole moment
355 The dipole moment is calculated as the expectation value of the
356 dipole moment operator, i.e. the trace of it times the density matrix::
358 d = - Σ ρ d
359 μν μν νμ
361 """
362 if self.dm_operator_c is None:
363 self.dm_operator_c = []
365 # Create external potentials in each direction
366 ext_c = [ConstantElectricField(Hartree / Bohr, dir)
367 for dir in np.eye(3)]
368 dm_operator_c = [self.hamiltonian.create_kick_matrix_calculator(
369 self.state, ext, self.pot_calc) for ext in ext_c]
370 self.dm_operator_c = dm_operator_c
372 dm_v = np.zeros(3)
373 for c, dm_operator in enumerate(self.dm_operator_c):
374 rho_MM = wfs.calculate_density_matrix()
375 dm_MM = dm_operator.calculate_matrix(wfs)
376 dm = - np.einsum('MN,NM->', rho_MM, dm_MM.data)
377 assert np.abs(dm.imag) < 1e-20
378 dm_v[c] = dm.real
380 return dm_v