Coverage for gpaw/tddft/propagators.py: 73%

1# Initially written by Lauri Lehtovaara, 2007

2"""This module implements time propagators for time-dependent density

3functional theory calculations."""

4from abc import ABC, abstractmethod

5from collections import namedtuple

7import numpy as np

9from ase.utils.timing import timer

11from gpaw.utilities.blas import axpy

13from gpaw.tddft.utils import MultiBlas

14from gpaw.tddft.tdopers import DummyDensity

17def create_propagator(name, **kwargs):

18 if isinstance(name, BasePropagator):

19 return name

20 elif isinstance(name, dict):

21 kwargs.update(name)

22 return create_propagator(**kwargs)

23 elif name == 'ECN':

24 return ExplicitCrankNicolson(**kwargs)

25 elif name == 'SICN':

26 return SemiImplicitCrankNicolson(**kwargs)

27 elif name == 'EFSICN':

28 return EhrenfestPAWSICN(**kwargs)

29 elif name == 'EFSICN_HGH':

30 return EhrenfestHGHSICN(**kwargs)

31 elif name == 'ETRSCN':

32 return EnforcedTimeReversalSymmetryCrankNicolson(**kwargs)

33 elif name == 'SITE':

34 return SemiImplicitTaylorExponential(**kwargs)

35 elif name == 'SIKE':

36 return SemiImplicitKrylovExponential(**kwargs)

37 elif name.startswith('SITE') or name.startswith('SIKE'):

38 raise DeprecationWarning(

39 'Use dictionary to specify degree.')

40 else:

41 raise ValueError('Unknown propagator: %s' % name)

44def allocate_wavefunction_arrays(wfs):

45 """Allocate wavefunction arrays."""

46 DummyKPoint = namedtuple('DummyKPoint', ['psit_nG'])

47 new_kpt_u = []

48 for kpt in wfs.kpt_u:

49 psit_nG = wfs.gd.empty(n=kpt.psit_nG.shape[0], dtype=complex)

50 new_kpt_u.append(DummyKPoint(psit_nG))

51 return new_kpt_u

54class BasePropagator(ABC):

55 """Abstract base class for time propagators."""

57 def todict(self):

58 return {'name': self.__class__.__name__}

60 def initialize(self, td_density, td_hamiltonian, td_overlap, solver,

61 preconditioner, gd, timer):

62 """Initialize propagator using runtime objects.

64 Parameters

65 ----------

66 td_density: TimeDependentDensity

67 the time-dependent density

68 td_hamiltonian: TimeDependentHamiltonian

69 the time-dependent hamiltonian

70 td_overlap: TimeDependentOverlap

71 the time-dependent overlap operator

72 solver: LinearSolver

73 solver for linear equations

74 preconditioner: Preconditioner

75 preconditioner for linear equations

76 gd: GridDescriptor

77 coarse (/wavefunction) grid descriptor

78 timer: Timer

79 timer

80 """

81 self.td_density = td_density

82 self.td_hamiltonian = td_hamiltonian

83 self.td_overlap = td_overlap

85 self.wfs = td_density.get_wavefunctions()

87 self.solver = solver

88 self.preconditioner = preconditioner

89 self.gd = gd

90 self.timer = timer

92 self.mblas = MultiBlas(gd)

94 # Solve M psin = psi

95 def apply_preconditioner(self, psi, psin):

96 """Solves preconditioner equation.

98 Parameters

99 ----------

100 psi: List of coarse grids

101 the known wavefunctions

102 psin: List of coarse grids

103 the result

104

105 """

106 self.timer.start('Solve TDDFT preconditioner')

107 if self.preconditioner is not None:

108 self.preconditioner.apply(self.kpt, psi, psin)

109 else:

110 psin[:] = psi

111 self.timer.stop('Solve TDDFT preconditioner')

112

113 @timer('Update time-dependent operators')

114 def update_time_dependent_operators(self, time):

115 """Update overlap, density, and Hamiltonian.

116

117 Parameters

118 ----------

119 time: float

120 the current time

121 """

122 # Update overlap S(t) of kpt.psit_nG in kpt.P_ani.

123 self.td_overlap.update(self.wfs)

124

125 # Calculate density rho(t) based on the wavefunctions psit_nG

126 # in kpt_u for t = time. Updates wfs.D_asp based on kpt.P_ani.

127 self.td_density.update()

128

129 # Update Hamiltonian H(t) to reflect density rho(t)

130 self.td_hamiltonian.update(self.td_density.get_density(), time)

131

132 @timer('Update time-dependent operators')

133 def half_update_time_dependent_operators(self, time):

134 """Half-update overlap, density, and Hamiltonian.

135

136 Parameters

137 ----------

138 time: float

139 the time passed to hamiltonian.half_update()

140 """

141 # Update overlap S(t+dt) of kpt.psit_nG in kpt.P_ani.

142 self.td_overlap.update(self.wfs)

143

144 # Calculate density rho(t+dt) based on the wavefunctions psit_nG in

145 # kpt_u for t = time+time_step. Updates wfs.D_asp based on kpt.P_ani.

146 self.td_density.update()

147

148 # Estimate Hamiltonian H(t+dt/2) by averaging H(t) and H(t+dt)

149 # and retain the difference for a half-way Hamiltonian dH(t+dt/2).

150 self.td_hamiltonian.half_update(self.td_density.get_density(), time)

151

152 # Estimate overlap S(t+dt/2) by averaging S(t) and S(t+dt) #TODO!!!

153 # XXX this doesn't do anything, see TimeDependentOverlap.half_update()

154 self.td_overlap.half_update(self.wfs)

155

156 @abstractmethod

157 def propagate(self, time, time_step):

158 """Propagate wavefunctions once.

159

160 Parameters

161 ----------

162 time: float

163 the current time

164 time_step: float

165 the time step

166

167 """

168 raise NotImplementedError()

169

170

171class ExplicitCrankNicolson(BasePropagator):

172 """Explicit Crank-Nicolson propagator

173

174 Crank-Nicolson propagator, which approximates the time-dependent

175 Hamiltonian to be unchanged during one iteration step.

176

177 (S(t) + .5j dt H(t) / hbar) psi(t+dt) = (S(t) - .5j dt H(t) / hbar) psi(t)

178

179 """

180 def __init__(self):

181 """Create ExplicitCrankNicolson-object."""

182 BasePropagator.__init__(self)

183 self.tmp_kpt_u = None

184 self.hpsit = None

185 self.spsit = None

186 self.sinvhpsit = None

187

188 def todict(self):

189 return {'name': 'ECN'}

190

191 def initialize(self, *args, **kwargs):

192 BasePropagator.initialize(self, *args, **kwargs)

193

194 # Allocate temporary wavefunctions

195 self.tmp_kpt_u = allocate_wavefunction_arrays(self.wfs)

196

197 # Allocate memory for Crank-Nicolson stuff

198 nvec = len(self.wfs.kpt_u[0].psit_nG)

199 self.hpsit = self.gd.zeros(nvec, dtype=complex)

200 self.spsit = self.gd.zeros(nvec, dtype=complex)

201

202 # ( S + i H dt/2 ) psit(t+dt) = ( S - i H dt/2 ) psit(t)

203 def propagate(self, time, time_step):

204 """Propagate wavefunctions.

205

206 Parameters

207 ----------

208 time: float

209 the current time

210 time_step: float

211 time step

212

213 """

214 self.niter = 0

215

216 # Copy current wavefunctions psit_nG to work wavefunction arrays

217 for u, kpt in enumerate(self.wfs.kpt_u):

218 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

219

220 # Euler step

221 # Overwrite psit_nG in tmp_kpt_u by (1 - i S^(-1)(t) H(t) dt) psit_nG

222 # from corresponding kpt_u in a Euler step before predicting psit(t+dt)

223 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.tmp_kpt_u):

224 self.solve_propagation_equation(kpt,

225 rhs_kpt,

226 time_step,

227 guess=True)

228

229 self.update_time_dependent_operators(time + time_step)

230

231 return self.niter

232

233 # ( S + i H dt/2 ) psit(t+dt) = ( S - i H dt/2 ) psit(t)

234 def solve_propagation_equation(self, kpt, rhs_kpt, time_step, guess=False):

235

236 # kpt is guess, rhs_kpt is used to calculate rhs and is overwritten

237 nvec = len(rhs_kpt.psit_nG)

238

239 assert kpt != rhs_kpt, 'Data race condition detected'

240 assert len(kpt.psit_nG) == nvec, 'Incompatible lhs/rhs vectors'

241

242 self.timer.start('Apply time-dependent operators')

243 # Store H psi(t) as hpsit and S psit(t) as spsit

244 self.td_overlap.update_k_point_projections(self.wfs, kpt,

245 rhs_kpt.psit_nG)

246 self.td_hamiltonian.apply(kpt,

247 rhs_kpt.psit_nG,

248 self.hpsit,

249 calculate_P_ani=False)

250 self.td_overlap.apply(rhs_kpt.psit_nG,

251 self.spsit,

252 self.wfs,

253 kpt,

254 calculate_P_ani=False)

255 self.timer.stop('Apply time-dependent operators')

256

257 # Update rhs_kpt.psit_nG to reflect ( S - i H dt/2 ) psit(t)

258 # rhs_kpt.psit_nG[:] = self.spsit - .5J * self.hpsit * time_step

259 rhs_kpt.psit_nG[:] = self.spsit

260 self.mblas.multi_zaxpy(-.5j * time_step, self.hpsit, rhs_kpt.psit_nG,

261 nvec)

262

263 if guess:

264 if self.sinvhpsit is None:

265 self.sinvhpsit = self.gd.zeros(len(kpt.psit_nG), dtype=complex)

266

267 # Update estimate of psit(t+dt) to ( 1 - i S^(-1) H dt ) psit(t)

268 self.td_overlap.apply_inverse(self.hpsit,

269 self.sinvhpsit,

270 self.wfs,

271 kpt,

272 use_cg=False)

273 self.mblas.multi_zaxpy(-1.0j * time_step, self.sinvhpsit,

274 kpt.psit_nG, nvec)

275

276 # Information needed by solver.solve -> self.dot

277 self.kpt = kpt

278 self.time_step = time_step

279

280 # Solve A x = b where A is (S + i H dt/2) and b = rhs_kpt.psit_nG

281 self.niter += self.solver.solve(self, kpt.psit_nG, rhs_kpt.psit_nG)

282

283 # ( S + i H dt/2 ) psi

284 def dot(self, psi, psin):

285 """Applies the propagator matrix to the given wavefunctions.

286

287 Parameters

288 ----------

289 psi: List of coarse grids

290 the known wavefunctions

291 psin: List of coarse grids

292 the result ( S + i H dt/2 ) psi

293

294 """

295 self.timer.start('Apply time-dependent operators')

296 self.td_overlap.update_k_point_projections(self.wfs, self.kpt, psi)

297 self.td_hamiltonian.apply(self.kpt,

298 psi,

299 self.hpsit,

300 calculate_P_ani=False)

301 self.td_overlap.apply(psi,

302 self.spsit,

303 self.wfs,

304 self.kpt,

305 calculate_P_ani=False)

306 self.timer.stop('Apply time-dependent operators')

307

308 # psin[:] = self.spsit + .5J * self.time_step * self.hpsit

309 psin[:] = self.spsit

310 self.mblas.multi_zaxpy(.5j * self.time_step, self.hpsit, psin,

311 len(psi))

312

313

314class SemiImplicitCrankNicolson(ExplicitCrankNicolson):

315 """Semi-implicit Crank-Nicolson propagator

316

317 Crank-Nicolson propagator, which first approximates the time-dependent

318 Hamiltonian to be unchanged during one iteration step to predict future

319 wavefunctions. Then the approximations for the future wavefunctions are

320 used to approximate the Hamiltonian at the middle of the time step.

321

322 (S(t) + .5j dt H(t) / hbar) psi(t+dt) = (S(t) - .5j dt H(t) / hbar) psi(t)

323 (S(t) + .5j dt H(t+dt/2) / hbar) psi(t+dt)

324 = (S(t) - .5j dt H(t+dt/2) / hbar) psi(t)

325

326 """

327 def __init__(self):

328 """Create SemiImplicitCrankNicolson-object."""

329 ExplicitCrankNicolson.__init__(self)

330 self.old_kpt_u = None

331

332 def todict(self):

333 return {'name': 'SICN'}

334

335 def initialize(self, *args, **kwargs):

336 ExplicitCrankNicolson.initialize(self, *args, **kwargs)

337

338 # Allocate old wavefunctions

339 self.old_kpt_u = allocate_wavefunction_arrays(self.wfs)

340

341 def propagate(self, time, time_step):

342 """Propagate wavefunctions once.

343

344 Parameters

345 ----------

346 time: float

347 the current time

348 time_step: float

349 time step

350 """

351

352 self.niter = 0

353 nvec = len(self.wfs.kpt_u[0].psit_nG)

354

355 # Copy current wavefunctions to work and old wavefunction arrays

356 for u, kpt in enumerate(self.wfs.kpt_u):

357 self.old_kpt_u[u].psit_nG[:] = kpt.psit_nG

358 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

359

360 # Predictor step

361 # Overwrite psit_nG in tmp_kpt_u by (1 - i S^(-1)(t) H(t) dt) psit_nG

362 # from corresponding kpt_u in a Euler step before predicting psit(t+dt)

363 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.tmp_kpt_u):

364 self.solve_propagation_equation(kpt,

365 rhs_kpt,

366 time_step,

367 guess=True)

368

369 self.half_update_time_dependent_operators(time + time_step)

370

371 # Corrector step

372 # Use predicted psit_nG in kpt_u as an initial guess, whereas the old

373 # wavefunction in old_kpt_u are used to calculate rhs based on psit(t)

374 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

375 # Average of psit(t) and predicted psit(t+dt)

376 mean_psit_nG = 0.5 * (kpt.psit_nG + rhs_kpt.psit_nG)

377 self.td_hamiltonian.half_apply(kpt, mean_psit_nG, self.hpsit)

378 self.td_overlap.apply_inverse(self.hpsit,

379 self.sinvhpsit,

380 self.wfs,

381 kpt,

382 use_cg=False)

383

384 # Update kpt.psit_nG to reflect

385 # psit(t+dt) - i S^(-1) dH(t+dt/2) dt/2 psit(t+dt/2)

386 kpt.psit_nG[:] = kpt.psit_nG - .5J * self.sinvhpsit * time_step

387 self.mblas.multi_zaxpy(-.5j * time_step, self.sinvhpsit,

388 kpt.psit_nG, nvec)

389

390 self.solve_propagation_equation(kpt, rhs_kpt, time_step)

391

392 self.update_time_dependent_operators(time + time_step)

393

394 return self.niter

395

396

397class EhrenfestPAWSICN(SemiImplicitCrankNicolson):

398 """Semi-implicit Crank-Nicolson propagator for Ehrenfest dynamics

399 TODO: merge this with the ordinary SICN

400 """

401 def __init__(self,

402 corrector_guess=True,

403 predictor_guess=(True, False),

404 use_cg=(False, False)):

405 """Create SemiImplicitCrankNicolson-object.

406

407 Parameters

408 ----------

409 corrector_guess: Bool

410 use initial guess for the corrector step (default is True)

411 predictor_guess: (Bool, Bool)

412 use (first, second) order initial guesses for the predictor step

413 default is (True, False)

414 use_cg: (Bool, Bool)

415 use CG for calculating the inverse overlap (predictor, corrector)

416 default is (False, False)

417

418 """

419 SemiImplicitCrankNicolson.__init__(self)

420 self.old_kpt_u = None

421 self.corrector_guess = corrector_guess

422 self.predictor_guess = predictor_guess

423 self.use_cg = use_cg

424

425 # self.hsinvhpsit = None

426 self.sinvh2psit = None

427

428 def todict(self):

429 return {'name': 'EFSICN'}

430

431 def update_velocities(self, v_at_new, v_at_old=None):

432 self.v_at = v_at_new.copy()

433 if (v_at_old is not None):

434 self.v_at_old = v_at_old.copy()

435

436 def propagate(self, time, time_step, v_a):

437 """Propagate wavefunctions once.

438

439 Parameters

440 ----------

441 time: float

442 the current time

443 time_step: float

444 time step

445 """

446

447 self.niter = 0

448 nvec = len(self.wfs.kpt_u[0].psit_nG)

449

450 # update the atomic velocities which are required

451 # for calculating the P term

452 self.update_velocities(v_a)

453

454 # Copy current wavefunctions to work and old wavefunction arrays

455 for u, kpt in enumerate(self.wfs.kpt_u):

456 self.old_kpt_u[u].psit_nG[:] = kpt.psit_nG

457 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

458

459 # Predictor step

460 # Overwrite psit_nG in tmp_kpt_u by (1 - i S^(-1)(t) H(t) dt) psit_nG

461 # from corresponding kpt_u in a Euler step before predicting psit(t+dt)

462 # self.v_at = self.v_at_old.copy() #v(t) for predictor step

463 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.tmp_kpt_u):

464 self.solve_propagation_equation(kpt,

465 rhs_kpt,

466 time_step,

467 guess=self.predictor_guess[0])

468

469 self.half_update_time_dependent_operators(time + time_step)

470

471 # Corrector step

472 # Use predicted psit_nG in kpt_u as an initial guess, whereas the old

473 # wavefunction in old_kpt_u are used to calculate rhs based on psit(t)

474 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

475 # Average of psit(t) and predicted psit(t+dt)

476 if (self.corrector_guess):

477 mean_psit_nG = 0.5 * (kpt.psit_nG + rhs_kpt.psit_nG)

478 self.td_hamiltonian.half_apply(kpt, mean_psit_nG, self.hpsit)

479 self.td_overlap.apply_inverse(self.hpsit,

480 self.sinvhpsit,

481 self.wfs,

482 kpt,

483 use_cg=self.use_cg[1])

484

485 # Update kpt.psit_nG to reflect

486 # psit(t+dt) - i S^(-1) dH(t+dt/2) dt/2 psit(t+dt/2)

487 kpt.psit_nG[:] = kpt.psit_nG - .5J * self.sinvhpsit * time_step

488 self.mblas.multi_zaxpy(-.5j * time_step, self.sinvhpsit,

489 kpt.psit_nG, nvec)

490

491 self.solve_propagation_equation(kpt, rhs_kpt, time_step)

492

493 self.update_time_dependent_operators(time + time_step)

494

495 return self.niter

496

497 # ( S + i H dt/2 ) psit(t+dt) = ( S - i H dt/2 ) psit(t)

498 def solve_propagation_equation(self,

499 kpt,

500 rhs_kpt,

501 time_step,

502 calculate_P_ani=False,

503 guess=False):

504

505 # kpt is guess, rhs_kpt is used to calculate rhs and is overwritten

506 nvec = len(rhs_kpt.psit_nG)

507

508 assert kpt != rhs_kpt, 'Data race condition detected'

509 assert len(kpt.psit_nG) == nvec, 'Incompatible lhs/rhs vectors'

510

511 self.timer.start('Apply time-dependent operators')

512 # Store H psi(t) as hpsit and S psit(t) as spsit

513 self.td_overlap.update_k_point_projections(self.wfs, kpt,

514 rhs_kpt.psit_nG)

515 self.td_hamiltonian.apply(kpt,

516 rhs_kpt.psit_nG,

517 self.hpsit,

518 calculate_P_ani=False)

519 self.td_hamiltonian.calculate_paw_correction(rhs_kpt.psit_nG,

520 self.hpsit,

521 self.wfs,

522 kpt,

523 self.v_at,

524 calculate_P_ani=False)

525

526 self.td_overlap.apply(rhs_kpt.psit_nG,

527 self.spsit,

528 self.wfs,

529 kpt,

530 calculate_P_ani=False)

531 self.timer.stop('Apply time-dependent operators')

532

533 # Update rhs_kpt.psit_nG to reflect ( S - i H dt/2 ) psit(t)

534 # rhs_kpt.psit_nG[:] = self.spsit - .5J * self.hpsit * time_step

535 rhs_kpt.psit_nG[:] = self.spsit

536 self.mblas.multi_zaxpy(-.5j * time_step, self.hpsit, rhs_kpt.psit_nG,

537 nvec)

538

539 if guess:

540 if self.sinvhpsit is None:

541 self.sinvhpsit = self.gd.zeros(len(kpt.psit_nG), dtype=complex)

542

543 if self.predictor_guess[1]:

544 if self.sinvh2psit is None:

545 self.sinvh2psit = self.gd.zeros(len(kpt.psit_nG),

546 dtype=complex)

547

548 # Update estimate of psit(t+dt) to ( 1 - i S^(-1) H dt ) psit(t)

549 self.td_overlap.apply_inverse(self.hpsit,

550 self.sinvhpsit,

551 self.wfs,

552 kpt,

553 use_cg=self.use_cg[0])

554

555 self.mblas.multi_zaxpy(-1.0j * time_step, self.sinvhpsit,

556 kpt.psit_nG, nvec)

557 if (self.predictor_guess[1]):

558 self.td_hamiltonian.apply(kpt,

559 self.sinvhpsit,

560 self.sinvh2psit,

561 calculate_P_ani=False)

562 self.td_hamiltonian.calculate_paw_correction(

563 self.sinvhpsit,

564 self.sinvh2psit,

565 self.wfs,

566 kpt,

567 self.v_at,

568 calculate_P_ani=False)

569 self.td_overlap.apply_inverse(self.sinvh2psit,

570 self.sinvh2psit,

571 self.wfs,

572 kpt,

573 use_cg=self.use_cg[0])

574 self.mblas.multi_zaxpy(-.5 * time_step * time_step,

575 self.sinvh2psit, kpt.psit_nG, nvec)

576

577 # Information needed by solver.solve -> self.dot

578 self.kpt = kpt

579 self.time_step = time_step

580

581 # Solve A x = b where A is (S + i H dt/2) and b = rhs_kpt.psit_nG

582 self.niter += self.solver.solve(self, kpt.psit_nG, rhs_kpt.psit_nG)

583

584 # ( S + i H dt/2 ) psi

585 def dot(self, psi, psin):

586 """Applies the propagator matrix to the given wavefunctions.

587

588 Parameters

589 ----------

590 psi: List of coarse grids

591 the known wavefunctions

592 psin: List of coarse grids

593 the result ( S + i H dt/2 ) psi

594

595 """

596 self.timer.start('Apply time-dependent operators')

597 self.td_overlap.update_k_point_projections(self.wfs, self.kpt, psi)

598 self.td_hamiltonian.apply(self.kpt,

599 psi,

600 self.hpsit,

601 calculate_P_ani=False)

602 self.td_hamiltonian.calculate_paw_correction(psi,

603 self.hpsit,

604 self.wfs,

605 self.kpt,

606 self.v_at,

607 calculate_P_ani=False)

608 self.td_overlap.apply(psi,

609 self.spsit,

610 self.wfs,

611 self.kpt,

612 calculate_P_ani=False)

613 self.timer.stop('Apply time-dependent operators')

614

615 # psin[:] = self.spsit + .5J * self.time_step * self.hpsit

616 psin[:] = self.spsit

617 self.mblas.multi_zaxpy(.5j * self.time_step, self.hpsit, psin,

618 len(psi))

619

620

621class EhrenfestHGHSICN(SemiImplicitCrankNicolson):

622 """Semi-implicit Crank-Nicolson propagator for Ehrenfest dynamics

623 using HGH pseudopotentials

624

625 """

626 def __init__(self):

627 """Create SemiImplicitCrankNicolson-object."""

628 SemiImplicitCrankNicolson.__init__(self)

629

630 def todict(self):

631 return {'name': 'EFSICN_HGH'}

632

633 def propagate(self, time, time_step):

634 """Propagate wavefunctions once.

635

636 Parameters

637 ----------

638 time: float

639 the current time

640 time_step: float

641 time step

642 """

643

644 self.niter = 0

645

646 # Copy current wavefunctions to work and old wavefunction arrays

647 for u, kpt in enumerate(self.wfs.kpt_u):

648 self.old_kpt_u[u].psit_nG[:] = kpt.psit_nG

649 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

650

651 # Predictor step

652 # Overwrite psit_nG in tmp_kpt_u by (1 - i S^(-1)(t) H(t) dt) psit_nG

653 # from corresponding kpt_u in a Euler step before predicting psit(t+dt)

654 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.tmp_kpt_u):

655 self.solve_propagation_equation(kpt,

656 rhs_kpt,

657 time_step,

658 guess=False)

659

660 self.half_update_time_dependent_operators(time + time_step)

661

662 # Corrector step

663 # Use predicted psit_nG in kpt_u as an initial guess, whereas the old

664 # wavefunction in old_kpt_u are used to calculate rhs based on psit(t)

665 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

666

667 self.solve_propagation_equation(kpt, rhs_kpt, time_step)

668

669 self.update_time_dependent_operators(time + time_step)

670

671 return self.niter

672

673 # ( S + i H dt/2 ) psit(t+dt) = ( S - i H dt/2 ) psit(t)

674 def solve_propagation_equation(self,

675 kpt,

676 rhs_kpt,

677 time_step,

678 calculate_P_ani=False,

679 guess=False):

680

681 # kpt is guess, rhs_kpt is used to calculate rhs and is overwritten

682 nvec = len(rhs_kpt.psit_nG)

683

684 assert kpt != rhs_kpt, 'Data race condition detected'

685 assert len(kpt.psit_nG) == nvec, 'Incompatible lhs/rhs vectors'

686

687 self.timer.start('Apply time-dependent operators')

688 # Store H psi(t) as hpsit and S psit(t) as spsit

689 self.td_overlap.update_k_point_projections(self.wfs, kpt,

690 rhs_kpt.psit_nG)

691 self.td_hamiltonian.apply(kpt,

692 rhs_kpt.psit_nG,

693 self.hpsit,

694 calculate_P_ani=False)

695 self.td_overlap.apply(rhs_kpt.psit_nG,

696 self.spsit,

697 self.wfs,

698 kpt,

699 calculate_P_ani=False)

700 self.timer.stop('Apply time-dependent operators')

701

702 # Update rhs_kpt.psit_nG to reflect ( S - i H dt/2 ) psit(t)

703 # rhs_kpt.psit_nG[:] = self.spsit - .5J * self.hpsit * time_step

704 rhs_kpt.psit_nG[:] = self.spsit

705 self.mblas.multi_zaxpy(-.5j * time_step, self.hpsit, rhs_kpt.psit_nG,

706 nvec)

707

708 # Information needed by solver.solve -> self.dot

709 self.kpt = kpt

710 self.time_step = time_step

711

712 # Solve A x = b where A is (S + i H dt/2) and b = rhs_kpt.psit_nG

713 self.niter += self.solver.solve(self, kpt.psit_nG, rhs_kpt.psit_nG)

714

715

716class EnforcedTimeReversalSymmetryCrankNicolson(SemiImplicitCrankNicolson):

717 """Enforced time-reversal symmetry Crank-Nicolson propagator

718

719 Crank-Nicolson propagator, which first approximates the time-dependent

720 Hamiltonian to be unchanged during one iteration step to predict future

721 wavefunctions. Then the approximations for the future wavefunctions are

722 used to approximate the Hamiltonian in the future.

723

724 (S(t) + .5j dt H(t) / hbar) psi(t+dt) = (S(t) - .5j dt H(t) / hbar) psi(t)

725 (S(t) + .5j dt H(t+dt) / hbar) psi(t+dt)

726 = (S(t) - .5j dt H(t) / hbar) psi(t)

727

728 """

729 def __init__(self):

730 SemiImplicitCrankNicolson.__init__(self)

731

732 def todict(self):

733 return {'name': 'ETRSCN'}

734

735 def propagate(self, time, time_step, update_callback=None):

736 """Propagate wavefunctions once.

737

738 Parameters

739 ----------

740 time: float

741 the current time

742 time_step: float

743 time step

744 """

745

746 self.niter = 0

747

748 # Copy current wavefunctions to work and old wavefunction arrays

749 for u, kpt in enumerate(self.wfs.kpt_u):

750 self.old_kpt_u[u].psit_nG[:] = kpt.psit_nG

751 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

752

753 # Predictor step

754 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

755 self.create_rhs(rhs_kpt, kpt, time_step)

756

757 if update_callback is not None:

758 update_callback()

759

760 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

761 self.solve_propagation_equation(kpt, rhs_kpt, time_step)

762

763 # XXX why here is full update and not half update?

764 # (compare to other propagators)

765 self.update_time_dependent_operators(time + time_step)

766

767 # Corrector step

768 # Use predicted psit_nG in kpt_u as an initial guess, whereas the old

769 # wavefunction in old_kpt_u are used to calculate rhs based on psit(t)

770 for [kpt, rhs_kpt] in zip(self.wfs.kpt_u, self.old_kpt_u):

771 self.solve_propagation_equation(kpt, rhs_kpt, time_step)

772

773 self.update_time_dependent_operators(time + time_step)

774

775 return self.niter

776

777 def create_rhs(self, rhs_kpt, kpt, time_step):

778 # kpt is guess, rhs_kpt is used to calculate rhs and is overwritten

779 nvec = len(rhs_kpt.psit_nG)

780

781 assert kpt != rhs_kpt, 'Data race condition detected'

782 assert len(kpt.psit_nG) == nvec, 'Incompatible lhs/rhs vectors'

783

784 self.timer.start('Apply time-dependent operators')

785 # Store H psi(t) as hpsit and S psit(t) as spsit

786 self.td_overlap.update_k_point_projections(self.wfs, kpt,

787 rhs_kpt.psit_nG)

788 self.td_hamiltonian.apply(kpt,

789 rhs_kpt.psit_nG,

790 self.hpsit,

791 calculate_P_ani=False)

792 self.td_overlap.apply(rhs_kpt.psit_nG,

793 self.spsit,

794 self.wfs,

795 kpt,

796 calculate_P_ani=False)

797 self.timer.stop('Apply time-dependent operators')

798

799 # Update rhs_kpt.psit_nG to reflect ( S - i H dt/2 ) psit(t)

800 # rhs_kpt.psit_nG[:] = self.spsit - .5J * self.hpsit * time_step

801 rhs_kpt.psit_nG[:] = self.spsit

802 self.mblas.multi_zaxpy(-.5j * time_step, self.hpsit, rhs_kpt.psit_nG,

803 nvec)

804

805 # ( S + i H(t+dt) dt/2 ) psit(t+dt) = ( S - i H(t) dt/2 ) psit(t)

806 # rhs_kpt = ( S - i H(t) dt/2 ) psit(t)

807 def solve_propagation_equation(self, kpt, rhs_kpt, time_step, guess=False):

808 # Information needed by solver.solve -> self.dot

809 self.kpt = kpt

810 self.time_step = time_step

811

812 # Solve A x = b where A is (S + i H dt/2) and b = rhs_kpt.psit_nG

813 self.niter += self.solver.solve(self, kpt.psit_nG, rhs_kpt.psit_nG)

814

815

816class AbsorptionKick:

817 """Absorption kick propagator

818

819 Absorption kick propagator::

820

821 (S(t) + .5j dt p.r / hbar) psi(0+) = (S(t) - .5j dt p.r / hbar) psi(0-)

822

823 where ``|p| = (eps e / hbar)``, and eps is field strength, e is elementary

824 charge.

825

826 """

827 def __init__(self, wfs, abs_kick_hamiltonian, td_overlap, solver,

828 preconditioner, gd, timer):

829 """Create AbsorptionKick-object.

830

831 Parameters

832 ----------

833 wfs: FDWaveFunctions

834 time-independent grid-based wavefunctions

835 abs_kick_hamiltonian: AbsorptionKickHamiltonian

836 the absorption kick hamiltonian

837 td_overlap: TimeDependentOverlap

838 the time-dependent overlap operator

839 solver: LinearSolver

840 solver for linear equations

841 preconditioner: Preconditioner

842 preconditioner for linear equations

843 gd: GridDescriptor

844 coarse (wavefunction) grid descriptor

845 timer: Timer

846 timer

847

848 """

849 self.propagator = ExplicitCrankNicolson()

850 self.propagator.initialize(DummyDensity(wfs),

851 abs_kick_hamiltonian, td_overlap,

852 solver, preconditioner, gd, timer)

853

854 def kick(self):

855 """Excite all possible frequencies.

856 """

857 for l in range(self.propagator.td_hamiltonian.iterations):

858 self.propagator.propagate(0, 1.0)

859

860

861class SemiImplicitTaylorExponential(BasePropagator):

862 """Semi-implicit Taylor exponential propagator

863 exp(-i S^-1 H t) = 1 - i S^-1 H t + (1/2) (-i S^-1 H t)^2 + ...

864

865 """

866 def __init__(self, degree=4):

867 """Create SemiImplicitTaylorExponential-object.

868

869 Parameters

870 ----------

871 degree: integer

872 Degree of the Taylor polynomial (default is 4)

873

874 """

875 raise RuntimeError('SITE propagator is unstable')

876 BasePropagator.__init__(self)

877 self.degree = degree

878 self.tmp_kpt_u = None

879 self.psin = None

880 self.hpsit = None

881

882 def todict(self):

883 return {'name': 'SITE',

884 'degree': self.degree}

885

886 def initialize(self, *args, **kwargs):

887 BasePropagator.initialize(self, *args, **kwargs)

888

889 # Allocate temporary wavefunctions

890 self.tmp_kpt_u = allocate_wavefunction_arrays(self.wfs)

891

892 # Allocate memory for Taylor exponential stuff

893 nvec = len(self.wfs.kpt_u[0].psit_nG)

894 self.psin = self.gd.zeros(nvec, dtype=complex)

895 self.hpsit = self.gd.zeros(nvec, dtype=complex)

896

897 def propagate(self, time, time_step):

898 """Propagate wavefunctions once.

899

900 Parameters

901 ----------

902 time: float

903 the current time

904 time_step: float

905 time step

906 """

907

908 self.niter = 0

909

910 # copy current wavefunctions to temporary variable

911 for u, kpt in enumerate(self.wfs.kpt_u):

912 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

913

914 # predict for each k-point

915 for kpt in self.wfs.kpt_u:

916 self.solve_propagation_equation(kpt, time_step)

917

918 self.half_update_time_dependent_operators(time + time_step)

919

920 # propagate psit(t), not psit(t+dt), in correct

921 for u, kpt in enumerate(self.wfs.kpt_u):

922 kpt.psit_nG[:] = self.tmp_kpt_u[u].psit_nG

923

924 # correct for each k-point

925 for kpt in self.wfs.kpt_u:

926 self.solve_propagation_equation(kpt, time_step)

927

928 self.update_time_dependent_operators(time + time_step)

929

930 return self.niter

931

932 # psi(t) = exp(-i t S^-1 H) psi(0)

933 # psi(t) = 1 + (-i S^-1 H t) (1 + (1/2) (-i S^-1 H t) (1 + ... ) )

934 def solve_propagation_equation(self, kpt, time_step):

935

936 nvec = len(kpt.psit_nG)

937

938 # Information needed by solver.solve -> self.dot

939 self.kpt = kpt

940 self.time_step = time_step

941

942 # psin = psi(0)

943 self.psin[:] = kpt.psit_nG

944 for k in range(self.degree, 0, -1):

945 # psin = psi(0) + (1/k) (-i S^-1 H t) psin

946 self.td_hamiltonian.apply(kpt, self.psin, self.hpsit)

947 # S psin = H psin

948 self.psin[:] = self.hpsit

949 self.niter += self.solver.solve(self, self.psin, self.hpsit)

950 # psin = psi(0) + (-it/k) S^-1 H psin

951 self.mblas.multi_scale(-1.0j * time_step / k, self.psin, nvec)

952 self.mblas.multi_zaxpy(1.0, kpt.psit_nG, self.psin, nvec)

953

954 kpt.psit_nG[:] = self.psin

955

956 def dot(self, psit, spsit):

957 self.td_overlap.apply(psit, spsit, self.wfs, self.kpt)

958

959

960class SemiImplicitKrylovExponential(BasePropagator):

961 """Semi-implicit Krylov exponential propagator

962

963

964 """

965 def __init__(self, degree=4):

966 """Create SemiImplicitKrylovExponential-object.

967

968 Parameters

969 ----------

970 degree: integer

971 Degree of the Krylov subspace (default is 4)

972 """

973 BasePropagator.__init__(self)

974 self.degree = degree

975 self.kdim = degree + 1

976 self.tmp_kpt_u = None

977 self.lm = None

978 self.em = None

979 self.hm = None

980 self.sm = None

981 self.xm = None

982 self.qm = None

983 self.Hqm = None

984 self.Sqm = None

985 self.rqm = None

986

987 def todict(self):

988 return {'name': 'SIKE',

989 'degree': self.degree}

990

991 def initialize(self, *args, **kwargs):

992 BasePropagator.initialize(self, *args, **kwargs)

993

994 # Allocate temporary wavefunctions

995 self.tmp_kpt_u = allocate_wavefunction_arrays(self.wfs)

996

997 # Allocate memory for Krylov subspace stuff

998 nvec = len(self.wfs.kpt_u[0].psit_nG)

999 self.em = np.zeros((nvec, self.kdim), dtype=float)

1000 self.lm = np.zeros((nvec, ), dtype=complex)

1001 self.hm = np.zeros((nvec, self.kdim, self.kdim), dtype=complex)

1002 self.sm = np.zeros((nvec, self.kdim, self.kdim), dtype=complex)

1003 self.xm = np.zeros((nvec, self.kdim, self.kdim), dtype=complex)

1004 self.qm = self.gd.zeros((self.kdim, nvec), dtype=complex)

1005 self.Hqm = self.gd.zeros((self.kdim, nvec), dtype=complex)

1006 self.Sqm = self.gd.zeros((self.kdim, nvec), dtype=complex)

1007 self.rqm = self.gd.zeros((nvec, ), dtype=complex)

1008

1009 def propagate(self, time, time_step):

1010 """Propagate wavefunctions once.

1011

1012 Parameters

1013 ----------

1014 time: float

1015 the current time

1016 time_step: float

1017 time step

1018 """

1019

1020 self.niter = 0

1021

1022 self.time_step = time_step

1023

1024 # copy current wavefunctions to temporary variable

1025 for u, kpt in enumerate(self.wfs.kpt_u):

1026 self.tmp_kpt_u[u].psit_nG[:] = kpt.psit_nG

1027

1028 # predict for each k-point

1029 for kpt in self.wfs.kpt_u:

1030 self.solve_propagation_equation(kpt, time_step)

1031

1032 self.half_update_time_dependent_operators(time + time_step)

1033

1034 # propagate psit(t), not psit(t+dt), in correct

1035 for u, kpt in enumerate(self.wfs.kpt_u):

1036 kpt.psit_nG[:] = self.tmp_kpt_u[u].psit_nG

1037

1038 # correct for each k-point

1039 for kpt in self.wfs.kpt_u:

1040 self.solve_propagation_equation(kpt, time_step)

1041

1042 self.update_time_dependent_operators(time + time_step)

1043

1044 return self.niter

1045

1046 # psi(t) = exp(-i t S^-1 H) psi(0)

1047 def solve_propagation_equation(self, kpt, time_step):

1048

1049 nvec = len(kpt.psit_nG)

1050 tmp = np.zeros((nvec, ), complex)

1051 xm_tmp = np.zeros((nvec, self.kdim), complex)

1052

1053 qm = self.qm

1054 Hqm = self.Hqm

1055 Sqm = self.Sqm

1056

1057 # Information needed by solver.solve -> self.dot

1058 self.kpt = kpt

1059

1060 scale = self.create_krylov_subspace(kpt, self.td_hamiltonian,

1061 self.td_overlap, self.qm, self.Hqm,

1062 self.Sqm)

1063

1064 # Calculate hm and sm

1065 for i in range(self.kdim):

1066 for j in range(self.kdim):

1067 self.mblas.multi_zdotc(tmp, qm[i], Hqm[j], nvec)

1068 tmp *= self.gd.dv

1069 for k in range(nvec):

1070 self.hm[k][i][j] = tmp[k]

1071 self.mblas.multi_zdotc(tmp, qm[i], Sqm[j], nvec)

1072 tmp *= self.gd.dv

1073 for k in range(nvec):

1074 self.sm[k][i][j] = tmp[k]

1075

1076 # Diagonalize

1077 # Propagate

1078 # psi(t) = Qm Xm exp(-i Em t) Xm^H Sm e_1

1079 # = Qm Xm exp(-i Em t) Sm Qm^H S psi(0) ???

1080 # = Qm Xm exp(-i Em t) y

1081 # = Qm Xm z

1082 # y = Sm Qm^H S psi(0) = Xm^H Sm e_1

1083 # if Sm = I then y is the first row of Xm^*

1084 # and z = exp(-i Em t) y

1085 for k in range(nvec):

1086 (self.em[k], self.xm[k]) = np.linalg.eigh(self.hm[k])

1087

1088 self.em = np.exp(self.em * (-1.0j * time_step))

1089 for k in range(nvec):

1090 z = self.em[k] * np.conj(self.xm[k, 0])

1091 xm_tmp[k][:] = np.dot(self.xm[k], z)

1092 kpt.psit_nG[:] = 0.0

1093 for k in range(nvec):

1094 for i in range(self.kdim):

1095 axpy(xm_tmp[k][i] / scale[k], self.qm[i][k], kpt.psit_nG[k])

1096

1097 # Create Krylov subspace

1098 # K_v = { psi, S^-1 H psi, (S^-1 H)^2 psi, ... }

1099

1100 def create_krylov_subspace(self, kpt, h, s, qm, Hqm, Sqm):

1101 nvec = len(kpt.psit_nG)

1102 tmp = np.zeros((nvec, ), complex)

1103 scale = np.zeros((nvec, ), complex)

1104 scale[:] = 0.0

1105 rqm = self.rqm

1106

1107 # q_0 = psi

1108 rqm[:] = kpt.psit_nG

1109

1110 for i in range(self.kdim):

1111 qm[i][:] = rqm

1112

1113 # S orthogonalize

1114 # q_i = q_i - sum_j<i <q_j|S|q_i> q_j

1115 for j in range(i):

1116 self.mblas.multi_zdotc(tmp, qm[i], Sqm[j], nvec)

1117 tmp *= self.gd.dv

1118 tmp = np.conj(tmp)

1119 self.mblas.multi_zaxpy(-tmp, qm[j], qm[i], nvec)

1120

1121 # S q_i

1122 s.apply(qm[i], Sqm[i], self.wfs, kpt)

1123 self.mblas.multi_zdotc(tmp, qm[i], Sqm[i], nvec)

1124 tmp *= self.gd.dv

1125 self.mblas.multi_scale(1. / np.sqrt(tmp), qm[i], nvec)

1126 self.mblas.multi_scale(1. / np.sqrt(tmp), Sqm[i], nvec)

1127 if i == 0:

1128 scale[:] = 1 / np.sqrt(tmp)

1129

1130 # H q_i

1131 h.apply(kpt, qm[i], Hqm[i])

1132

1133 # S r = H q_i, (if stuff, to save one inversion)

1134 if i + 1 < self.kdim:

1135 rqm[:] = Hqm[i]

1136 self.solver.solve(self, rqm, Hqm[i])

1137

1138 return scale

1139

1140 def dot(self, psit, spsit):

1141 self.td_overlap.apply(psit, spsit, self.wfs, self.kpt)

1142

1143 # Below this, just for testing & debug

1144 def Sdot(self, psit, spsit):

1145 self.apply_preconditioner(psit, self.tmp)

1146 self.td_overlap.apply(self.tmp, spsit, self.wfs, self.kpt)

1147

1148 def Hdot(self, psit, spsit):

1149 self.apply_preconditioner(psit, self.tmp)

1150 self.td_hamiltonian.apply(self.kpt, self.tmp, spsit)

1151

1152 def inverse_overlap(self, kpt_u, degree):

1153 self.dot = self.Sdot

1154 self.kpt = kpt_u[0]

1155 nvec = len(self.kpt.psit_nG)

1156 nrm2 = np.zeros(nvec, dtype=complex)

1157 self.tmp = self.gd.zeros(n=nvec, dtype=complex)

1158

1159 for i in range(10):

1160 self.solver.solve(self, self.kpt.psit_nG, self.kpt.psit_nG)

1161 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.kpt.psit_nG,

1162 nvec)

1163 nrm2 *= self.gd.dv

1164 self.mblas.multi_scale(1 / np.sqrt(nrm2), self.kpt.psit_nG, nvec)

1165 self.td_overlap.apply(self.kpt.psit_nG, self.tmp, self.wfs, self.kpt)

1166 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.tmp, nvec)

1167 nrm2 *= self.gd.dv

1168 print('S min eig = ', nrm2)

1169

1170 def overlap(self, kpt_u, degree):

1171 self.dot = self.Sdot

1172 self.kpt = kpt_u[0]

1173 nvec = len(self.kpt.psit_nG)

1174 nrm2 = np.zeros(nvec, dtype=complex)

1175 self.tmp = self.gd.zeros(n=nvec, dtype=complex)

1176

1177 for i in range(100):

1178 self.tmp[:] = self.kpt.psit_nG

1179 self.td_overlap.apply(self.tmp, self.kpt.psit_nG, self.wfs,

1180 self.kpt)

1181 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.kpt.psit_nG,

1182 nvec)

1183 nrm2 *= self.gd.dv

1184 self.mblas.multi_scale(1 / np.sqrt(nrm2), self.kpt.psit_nG, nvec)

1185 self.td_overlap.apply(self.kpt.psit_nG, self.tmp, self.wfs, self.kpt)

1186 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.tmp, nvec)

1187 nrm2 *= self.gd.dv

1188 print('S max eig = ', nrm2)

1189

1190 def inverse_hamiltonian(self, kpt_u, degree):

1191 self.dot = self.Hdot

1192 self.kpt = kpt_u[0]

1193 nvec = len(self.kpt.psit_nG)

1194 nrm2 = np.zeros(nvec, dtype=complex)

1195 self.tmp = self.gd.zeros(n=nvec, dtype=complex)

1196

1197 for i in range(10):

1198 self.solver.solve(self, self.kpt.psit_nG, self.kpt.psit_nG)

1199 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.kpt.psit_nG,

1200 nvec)

1201 nrm2 *= self.gd.dv

1202 self.mblas.multi_scale(1 / np.sqrt(nrm2), self.kpt.psit_nG, nvec)

1203 self.td_hamiltonian.apply(self.kpt, self.kpt.psit_nG, self.tmp)

1204 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.tmp, nvec)

1205 nrm2 *= self.gd.dv

1206 print('H min eig = ', nrm2)

1207

1208 def hamiltonian(self, kpt_u, degree):

1209 self.dot = self.Hdot

1210 self.kpt = kpt_u[0]

1211 nvec = len(self.kpt.psit_nG)

1212 nrm2 = np.zeros(nvec, dtype=complex)

1213 self.tmp = self.gd.zeros(n=nvec, dtype=complex)

1214

1215 for i in range(100):

1216 self.tmp[:] = self.kpt.psit_nG

1217 self.td_hamiltonian.apply(self.kpt, self.tmp, self.kpt.psit_nG)

1218 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.kpt.psit_nG,

1219 nvec)

1220 nrm2 *= self.gd.dv

1221 self.mblas.multi_scale(1 / np.sqrt(nrm2), self.kpt.psit_nG, nvec)

1222 self.td_hamiltonian.apply(self.kpt, self.kpt.psit_nG, self.tmp)

1223 self.mblas.multi_zdotc(nrm2, self.kpt.psit_nG, self.tmp, nvec)

1224 nrm2 *= self.gd.dv

1225 print('H max eig = ', nrm2)