Coverage for rhodent/density_matrices/time.py: 89%

378 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-08-01 16:57 +0000

1from __future__ import annotations 

2 

3from typing import Generator, Collection 

4 

5from pathlib import Path 

6import numpy as np 

7from numpy.typing import NDArray 

8 

9from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition 

10from gpaw.tddft.units import as_to_au, au_to_as, eV_to_au, au_to_eV 

11from gpaw.mpi import world 

12 

13from .buffer import DensityMatrixBuffer 

14from .distributed.pulse import PulseConvolver 

15from .readers.gpaw import KohnShamRhoWfsReader 

16from .readers.numpy import FrequencyDensityMatrixReader, TimeDensityMatrixReader 

17from .density_matrix import DensityMatrix 

18from .base import BaseDensityMatrices, WorkMetadata 

19from ..perturbation import create_perturbation, Perturbation, DeltaKick, PerturbationLike, PulsePerturbation 

20from ..utils import Logger, two_communicator_sizes, get_gaussian_pulse_values 

21from ..utils.logging import format_times 

22from ..utils.memory import MemoryEstimate 

23from ..typing import Array1D, Communicator 

24 

25 

26class ConvolutionDensityMatrixMetadata(WorkMetadata): 

27 """ Metadata to the density matrices. 

28 

29 Properties 

30 ---------- 

31 density_matrices 

32 Parent of this object. 

33 globalt 

34 Time index. 

35 localt 

36 Time index on this rank. 

37 globalp 

38 Pulse index. 

39 localp 

40 Pulse index on this rank. 

41 """ 

42 density_matrices: ConvolutionDensityMatrices 

43 globalt: int 

44 globalp: int 

45 localt: int 

46 localp: int 

47 

48 def __new__(cls, 

49 density_matrices: ConvolutionDensityMatrices, 

50 globalt: int, 

51 globalp: int, 

52 localt: int, 

53 localp: int): 

54 self = WorkMetadata.__new__(cls, density_matrices=density_matrices) 

55 self.globalt = globalt 

56 self.globalp = globalp 

57 self.localt = localt 

58 self.localp = localp 

59 return self 

60 

61 @property 

62 def global_indices(self): 

63 return (self.globalp, self.globalt) 

64 

65 @property 

66 def time(self) -> float: 

67 """ Simulation time in units of as. """ 

68 return self.density_matrices.times[self.globalt] 

69 

70 @property 

71 def pulse(self) -> Perturbation: 

72 """ Pulse. """ 

73 return self.density_matrices.pulses[self.globalp] 

74 

75 @property 

76 def desc(self) -> str: 

77 timestr = f'{self.time:.1f}as' 

78 if len(self.density_matrices.pulses) == 1: 

79 return timestr 

80 

81 d = self.pulse.todict() 

82 if d['name'] == 'GaussianPulse': 

83 pulsestr = f'{d["frequency"]:.1f}eV' 

84 else: 

85 pulsestr = f'#{self.globalp}' 

86 return f'{timestr} @ Pulse {pulsestr}' 

87 

88 

89class ConvolutionDensityMatrices(BaseDensityMatrices[ConvolutionDensityMatrixMetadata]): 

90 

91 """ 

92 Collection of density matrices in the Kohn-Sham basis for different times, 

93 after convolution with various pulses. 

94 

95 Plain density matrices and/or derivatives thereof may be represented. 

96 

97 Parameters 

98 ---------- 

99 ksd 

100 KohnShamDecomposition object or file name. 

101 pulses 

102 Convolute the density matrices with these pulses. 

103 times 

104 Compute density matrices for these times (or as close to them as possible). In units of as. 

105 derivative_order_s 

106 Compute density matrix derivatives of the following orders. 

107 0 for plain density matrix and positive integers for derivatives. 

108 real 

109 Calculate the real part of density matrices. 

110 imag 

111 Calculate the imaginary part of density matrices. 

112 calc_size 

113 Size of the calculation communicator. 

114 """ 

115 

116 _derivative_order_s: list[int] 

117 

118 def __init__(self, 

119 ksd: KohnShamDecomposition | str, 

120 pulses: Collection[PerturbationLike], 

121 times: list[float] | NDArray[np.float64], 

122 derivative_order_s: list[int] = [0], 

123 real: bool = True, 

124 imag: bool = True, 

125 calc_size: int = 1, 

126 log: Logger | None = None): 

127 self._time_t = np.array(times) 

128 self._pulses = [create_perturbation(pulse) for pulse in pulses] 

129 self._derivative_order_s = derivative_order_s 

130 

131 super().__init__(ksd=ksd, 

132 real=real, imag=imag, 

133 calc_size=calc_size, log=log) 

134 

135 tag_s = ['', '-Iomega', '-omega2'] 

136 return_keys_by_order = {0: 'pulserho', 1: 'pulsedrho', 2: 'pulseddrho'} 

137 self.keys_by_order = {0: 'rho_ia', 1: 'drho_ia', 2: 'ddrho_ia'} 

138 if 1 not in derivative_order_s: 

139 tag_s.remove('-Iomega') 

140 return_keys_by_order.pop(1) 

141 if 2 not in derivative_order_s: 

142 tag_s.remove('-omega2') 

143 return_keys_by_order.pop(2) 

144 self.tag_s = tag_s 

145 self.return_keys_by_order = return_keys_by_order 

146 

147 @property 

148 def times(self) -> Array1D[np.float64]: 

149 """ Simulation time in units of as. """ 

150 return self._time_t # type: ignore 

151 

152 @property 

153 def nt(self) -> int: 

154 """ Number of times. """ 

155 return len(self.times) 

156 

157 @property 

158 def pulses(self) -> list[Perturbation]: 

159 """ Pulses with which density matrices are convoluted. """ 

160 return self._pulses 

161 

162 @property 

163 def derivative_order_s(self) -> list[int]: 

164 """ 

165 List of orders of the density matrix derivatives to be computed. 

166 `0` for plain density matrix and positive integers for derivatives. 

167 """ 

168 return self._derivative_order_s 

169 

170 def work_loop(self, 

171 rank: int) -> Generator[ConvolutionDensityMatrixMetadata | None, None, None]: 

172 nt = len(self.times) 

173 ntperrank = (nt + self.loop_comm.size - 1) // self.loop_comm.size 

174 

175 # Do convolution pulse-by-pulse, time-by-time 

176 for p, pulse in enumerate(self.pulses): 

177 # Determine which times to compute on this loop_comm rank for good load balancing 

178 shift = (p * nt + rank) % self.loop_comm.size 

179 for localt in range(ntperrank): 

180 globalt = shift + localt * self.loop_comm.size 

181 if globalt < nt: 

182 yield ConvolutionDensityMatrixMetadata(density_matrices=self, globalt=globalt, 

183 localt=localt, globalp=p, localp=p) 

184 else: 

185 yield None 

186 

187 def write_to_disk(self, 

188 pulserho_fmt: str): 

189 """ Calculate the density matrices and save to disk. 

190 

191 Parameters 

192 ---------- 

193 pulserho_fmt 

194 Formatting string for the density matrices saved to disk. 

195 

196 The formatting string should be a plain string containing variable 

197 placeholders within curly brackets `{}`. It should not be confused with 

198 a formatted string literal (f-string). 

199 

200 Example: 

201 

202 * pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`. 

203 

204 Accepts variables 

205 

206 * `{time}` - Time in units of as. 

207 * `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`. 

208 * `{pulsefreq}` - Pulse frequency in units of eV. 

209 * `{pulsefwhm}` - Pulse FWHM in units of fs. 

210 """ 

211 nlocaltot = len(self.local_work_plan) 

212 

213 tags_keys = [(tag, key) for s, (tag, key) in enumerate( 

214 [('', 'rho_ia'), ('-Iomega', 'drho_ia'), ('-omega2', 'ddrho_ia')]) if s in self.derivative_order_s] 

215 

216 # Iterate over density matrices on all ranks 

217 for ndone, (work, dm) in enumerate(self, 1): 

218 avg = self.log.elapsed('read')/ndone 

219 estrem = avg * (nlocaltot - ndone) 

220 self.log(f'Obtained density matrix {ndone:4.0f}/{nlocaltot:4.0f} on this rank. ' 

221 f'Avg: {avg:10.3f}s, estimated remaining: {estrem:10.3f}s', who='Writer', flush=True) 

222 for tag, key in tags_keys: 

223 fname_kw = dict(time=work.time, tag=tag, **get_gaussian_pulse_values(work.pulse)) 

224 fpath = Path(pulserho_fmt.format(**fname_kw)) 

225 rho_ia = getattr(dm, key) 

226 if self.calc_comm.rank == 0: 

227 assert isinstance(rho_ia, np.ndarray) 

228 fpath.parent.mkdir(parents=True, exist_ok=True) 

229 if fpath.suffix == '.npz': 

230 # Write numpy archive 

231 np.savez(str(fpath), rho_ia=rho_ia) 

232 else: 

233 # Save numpy binary file 

234 np.save(str(fpath), rho_ia) 

235 self.log(f'Written {fpath}', who='Writer', flush=True) 

236 world.barrier() 

237 

238 

239class TimeDensityMatricesFromWaveFunctions(ConvolutionDensityMatrices): 

240 

241 """ 

242 Collection of density matrices in the Kohn-Sham basis for different times, 

243 obtained by reading the time-dependent wave functions file. 

244 

245 Parameters 

246 ---------- 

247 ksd 

248 KohnShamDecomposition object or file name. 

249 wfs_fname 

250 File name of the time-dependent wave functions file. 

251 times 

252 Compute density matrices for these times (or as close to them as possible). In units of as. 

253 real 

254 Calculate the real part of density matrices. 

255 imag 

256 Calculate the imaginary part of density matrices. 

257 calc_size 

258 Size of the calculation communicator. 

259 stridet 

260 Skip this many steps when reading the time-dependent wave functions file. 

261 """ 

262 

263 def __init__(self, 

264 ksd: KohnShamDecomposition | str, 

265 wfs_fname: str, 

266 times: list[float] | Array1D[np.float64], 

267 real: bool = True, 

268 imag: bool = True, 

269 calc_size: int = 1, 

270 log: Logger | None = None, 

271 stridet: int = 1): 

272 rho_nn_direct = KohnShamRhoWfsReader(wfs_fname=wfs_fname, ksd=ksd, 

273 yield_re=real, yield_im=imag, 

274 filter_times=np.array(times) * as_to_au, # type: ignore 

275 stridet=stridet, log=log) 

276 self.rho_nn_direct = rho_nn_direct 

277 

278 super().__init__(ksd=rho_nn_direct.ksd, 

279 times=rho_nn_direct.time_t * au_to_as, 

280 real=real, 

281 imag=imag, 

282 pulses=[None], 

283 calc_size=1, 

284 log=rho_nn_direct.log) 

285 

286 imin, imax, amin, amax = self.ksd.ialims() 

287 

288 # Read density matrices corresponding to ksd ialims 

289 self._n1slice = slice(imin, imax + 1) 

290 self._n2slice = slice(amin, amax + 1) 

291 

292 def __str__(self) -> str: 

293 lines = ['Response from time-dependent wave functions'] 

294 

295 lines.append('') 

296 lines += [' ' + line for line in str(self.rho_nn_direct).split('\n')] 

297 shape = (self._n1slice.stop - self._n1slice.start, 

298 self._n2slice.stop - self._n2slice.start,) 

299 if 'Re' in self.reim and 'Im' in self.reim: 

300 reim_desc = 'real and imaginary parts' 

301 elif 'Re' in self.reim: 

302 reim_desc = 'real parts' 

303 else: 

304 reim_desc = 'imaginary parts' 

305 

306 lines.append('') 

307 lines.append(f'Density matrix shape {shape}') 

308 lines.append(f'reading {reim_desc}') 

309 

310 return '\n'.join(lines) 

311 

312 def get_memory_estimate(self) -> MemoryEstimate: 

313 shape = (self._n1slice.stop - self._n1slice.start, 

314 self._n2slice.stop - self._n2slice.start,) 

315 narrays = 2 if 'Re' in self.reim and 'Im' in self.reim else 1 

316 for t_r in self.rho_nn_direct.work_loop_by_ranks(): 

317 nreading = len(t_r) 

318 break 

319 memory_estimate = MemoryEstimate() 

320 memory_estimate.add_child('Time-dependent wave functions reader', 

321 self.rho_nn_direct.get_memory_estimate()) 

322 memory_estimate.add_key('Density matrix', (shape) + (narrays, ), 

323 on_num_ranks=nreading) 

324 

325 return memory_estimate 

326 

327 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]: 

328 assert self.calc_comm.size == 1 # TODO 

329 

330 for work, dm_buffer in zip(self.work_loop(self.loop_comm.rank), 

331 self.rho_nn_direct.iread(0, 0, self._n1slice, self._n2slice)): 

332 assert work is not None 

333 if 'Re' in self.reim: 

334 Rerho_ia = dm_buffer._get_real(0) 

335 if 'Im' in self.reim: 

336 Imrho_ia = dm_buffer._get_imag(0) 

337 if 'Re' in self.reim and 'Im' in self.reim: 

338 # Complex result 

339 # Compared to numpy, we use another convention, hence the minus sign 

340 rho_ia = Rerho_ia - 1j * Imrho_ia 

341 elif 'Re' in self.reim: 

342 # Real result 

343 rho_ia = Rerho_ia 

344 else: 

345 rho_ia = -1j * Imrho_ia 

346 matrices: dict[int, NDArray[np.complex128] | None] = {0: rho_ia} 

347 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm) 

348 

349 yield work, dm 

350 

351 def parallel_prepare(self): 

352 self.rho_nn_direct.parallel_prepare() 

353 

354 

355class ConvolutionDensityMatricesFromDisk(ConvolutionDensityMatrices): 

356 

357 """ 

358 Collection of density matrices in the Kohn-Sham basis after convolution with 

359 one or several pulses, for different times. Read from disk. 

360 

361 Parameters 

362 ---------- 

363 ksd 

364 KohnShamDecomposition object or file name. 

365 pulserho_fmt 

366 Formatting string for the density matrices saved to disk. 

367 

368 The formatting string should be a plain string containing variable 

369 placeholders within curly brackets `{}`. It should not be confused with 

370 a formatted string literal (f-string). 

371 

372 Example: 

373 

374 * pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`. 

375 

376 Accepts variables 

377 

378 * `{time}` - Time in units of as. 

379 * `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`. 

380 * `{pulsefreq}` - Pulse frequency in units of eV. 

381 * `{pulsefwhm}` - Pulse FWHM in units of fs. 

382 pulses 

383 Convolute the density matrices with these pulses. 

384 times 

385 Compute density matrices for these times (or as close to them as possible). In units of as. 

386 derivative_order_s 

387 Compute density matrix derivatives of the following orders. 

388 `0` for plain density matrix and positive integers for derivatives. 

389 real 

390 Calculate real part of density matrices. 

391 imag 

392 Calculate imaginary part of density matrices. 

393 calc_size 

394 Size of the calculation communicator. 

395 """ 

396 

397 def __init__(self, 

398 ksd: KohnShamDecomposition | str, 

399 pulserho_fmt: str, 

400 pulses: Collection[PerturbationLike], 

401 times: list[float] | Array1D[np.float64], 

402 derivative_order_s: list[int] = [0], 

403 log: Logger | None = None, 

404 real: bool = True, 

405 imag: bool = True, 

406 calc_size: int = 1): 

407 reader = TimeDensityMatrixReader(pulserho_fmt=pulserho_fmt, 

408 ksd=ksd, 

409 pulses=pulses, 

410 filter_times=times, 

411 derivative_order_s=derivative_order_s, 

412 log=log) 

413 self.reader = reader 

414 super().__init__(ksd=reader.ksd, pulses=reader.pulses, times=reader.times, 

415 real=real, imag=imag, 

416 derivative_order_s=derivative_order_s, calc_size=calc_size, log=reader.log) 

417 self.ksd.distribute(self.calc_comm) 

418 self.pulserho_fmt = pulserho_fmt 

419 

420 def __str__(self) -> str: 

421 return str(self.reader) 

422 

423 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]: 

424 for work in self.local_work_plan: 

425 self.log.start('read') 

426 matrices: dict[int, NDArray[np.complex128] | None] = dict() 

427 for derivative in self.derivative_order_s: 

428 if self.calc_comm.rank > 0: 

429 # Don't read on non calc comm root ranks 

430 matrices[derivative] = None 

431 continue 

432 

433 matrices[derivative] = self.reader.read(work.time, work.pulse, derivative) 

434 

435 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm) 

436 

437 yield work, dm 

438 

439 

440class ConvolutionDensityMatricesFromFrequency(ConvolutionDensityMatrices): 

441 

442 """ 

443 Collection of density matrices in the Kohn-Sham basis after convolution with 

444 one or several pulses, for different times. Obtained from the the density 

445 matrices in frequency space, which are read from disk. 

446 

447 Parameters 

448 ---------- 

449 ksd 

450 KohnShamDecomposition object or file name. 

451 frho_fmt 

452 Formatting string for the density matrices 

453 in frequency space saved to disk. 

454 

455 The formatting string should be a plain string containing variable 

456 placeholders within curly brackets `{}`. It should not be confused with 

457 a formatted string literal (f-string). 

458 

459 Example: 

460 

461 * frho_fmt = `frho/w{freq:05.2f}-{reim}.npy'`. 

462 

463 Accepts variables: 

464 

465 * `{freq}` - Frequency in units of eV. 

466 * `{reim}` - `'Re'` or `'Im'` for Fourier transform of real/imaginary 

467 part of density matrix. 

468 perturbation 

469 Perturbation that was present during time propagation. 

470 pulses 

471 Convolute the density matrices with these pulses. 

472 times 

473 Compute density matrices for these times (or as close to them as possible). In units of as. 

474 derivative_order_s 

475 Compute density matrix derivatives of the following orders. 

476 `0` for plain density matrix and positive integers for derivatives. 

477 real 

478 Calculate real part of density matrices. 

479 imag 

480 Calculate imaginary part of density matrices. 

481 calc_size 

482 Size of the calculation communicator. 

483 """ 

484 

485 def __init__(self, 

486 ksd: KohnShamDecomposition | str, 

487 frho_fmt: str, 

488 perturbation: PerturbationLike, 

489 pulses: Collection[PerturbationLike], 

490 times: list[float] | NDArray[np.float64], 

491 derivative_order_s: list[int] = [0], 

492 log: Logger | None = None, 

493 real: bool = True, 

494 imag: bool = True, 

495 calc_size: int = 1): 

496 super().__init__(ksd=ksd, pulses=pulses, times=times, 

497 real=real, imag=imag, 

498 derivative_order_s=derivative_order_s, calc_size=calc_size, log=log) 

499 # Frequency grid for convolution 

500 if not all(pulse.todict()['name'] == 'GaussianPulse' for pulse in self.pulses): 

501 raise NotImplementedError('Only Gaussian pulses implemented for ResponseFromFourierTransform') 

502 

503 freq_spacing = 0.05 # 50meV 

504 freq_min = min((pulse.pulse.omega0 - 4 * pulse.pulse.sigma) * au_to_eV # type: ignore 

505 for pulse in self.pulses) 

506 freq_min = max(freq_min, freq_spacing) 

507 freq_max = min((pulse.pulse.omega0 + 4 * pulse.pulse.sigma) * au_to_eV # type: ignore 

508 for pulse in self.pulses) 

509 

510 transformer = ExactFourierTransformer.from_file( 

511 frho_fmt=frho_fmt, ksd=ksd, perturbation=perturbation, 

512 real=real, imag=imag, log=self.log, comm=self.calc_comm, 

513 freq_spacing=freq_spacing, freq_min=freq_min, freq_max=freq_max) 

514 self.transformer = transformer 

515 

516 def __str__(self) -> str: 

517 lines = ['Response from Fourier transform of density matrices on disk'] 

518 lines.append('') 

519 lines.append(f'Calculating response for {self.nt} times and {len(self.pulses)} pulses') 

520 lines.append(f' times: {format_times(self.times)}') 

521 

522 return '\n'.join(lines) 

523 

524 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]: 

525 self.transformer.parallel_prepare() 

526 

527 for work in self.local_work_plan: 

528 matrices: dict[int, NDArray[np.complex128] | None] = dict() 

529 

530 for derivative in self.derivative_order_s: 

531 buffer = self.transformer.convolve(work.time, work.pulse, derivative) 

532 if self.calc_comm.rank > 0: 

533 matrices[derivative] = None 

534 continue 

535 assert buffer is not None 

536 # Buffer shape is i, a, pulses, times 

537 if 'Re' in self.reim: 

538 Rerho_ia = buffer._get_real(derivative) 

539 if 'Im' in self.reim: 

540 Imrho_ia = buffer._get_imag(derivative) 

541 if 'Re' in self.reim and 'Im' in self.reim: 

542 # Complex result 

543 rho_ia = Rerho_ia + 1j * Imrho_ia 

544 elif 'Re' in self.reim: 

545 # Real result 

546 rho_ia = Rerho_ia 

547 else: 

548 rho_ia = 1j * Imrho_ia 

549 matrices[derivative] = rho_ia 

550 

551 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm) 

552 

553 yield work, dm 

554 

555 

556class ConvolutionDensityMatricesFromWaveFunctions(ConvolutionDensityMatrices): 

557 

558 """ 

559 Collection of density matrices in the Kohn-Sham basis after convolution with 

560 one or several pulses, for different times. Obtained from the time-dependent wave functions file, 

561 which is read from disk. 

562 

563 Parameters 

564 ---------- 

565 ksd 

566 KohnShamDecomposition object or file name. 

567 wfs_fname 

568 File name of the time-dependent wave functions file. 

569 perturbation 

570 Perturbation that was present during time propagation. 

571 pulses 

572 Convolute the density matrices with these pulses. 

573 times 

574 Compute density matrices for these times (or as close to them as possible). In units of as. 

575 derivative_order_s 

576 Compute density matrix derivatives of the following orders. 

577 `0` for plain density matrix and positive integers for derivatives. 

578 real 

579 Calculate real part of density matrices. 

580 imag 

581 Calculate imaginary part of density matrices. 

582 calc_size 

583 Size of the calculation communicator. 

584 stridet 

585 Skip this many steps when reading the time-dependent wave functions file. 

586 """ 

587 

588 def __init__(self, 

589 ksd: KohnShamDecomposition | str, 

590 wfs_fname: str, 

591 perturbation: PerturbationLike, 

592 pulses: Collection[PerturbationLike], 

593 times: list[float] | Array1D[np.float64], 

594 derivative_order_s: list[int] = [0], 

595 real: bool = True, 

596 imag: bool = True, 

597 calc_size: int = 1, 

598 log: Logger | None = None, 

599 stridet: int = 1): 

600 _, calc_size = two_communicator_sizes(-1, calc_size) 

601 # The calc_comm rank 0's are world ranks 0, with a spacing of calc_size 

602 result_on_ranks = list(range(0, world.size, calc_size)) 

603 

604 rho_nn_conv = PulseConvolver.from_parameters( 

605 wfs_fname, ksd, 

606 pulses=pulses, 

607 perturbation=perturbation, 

608 yield_re=real, yield_im=imag, 

609 filter_times=np.array(times) * as_to_au, 

610 derivative_order_s=list(sorted(derivative_order_s)), 

611 stridet=stridet, 

612 result_on_ranks=result_on_ranks, 

613 log=log) 

614 self.rho_nn_conv = rho_nn_conv 

615 

616 super().__init__(ksd=rho_nn_conv.ksd, pulses=pulses, times=rho_nn_conv.time_t * au_to_as, 

617 real=real, imag=imag, 

618 derivative_order_s=derivative_order_s, calc_size=calc_size, 

619 log=rho_nn_conv.log) 

620 

621 def __str__(self) -> str: 

622 lines = [] 

623 lines.append('Response from time-dependent wave functions ') 

624 lines.append('') 

625 lines += ['' + line for line in str(self.rho_nn_conv.rho_nn_reader.rho_wfs_reader).split('\n')] 

626 lines.append('') 

627 lines += ['' + line for line in str(self.rho_nn_conv.rho_nn_reader).split('\n')] 

628 lines.append('') 

629 lines += ['' + line for line in str(self.rho_nn_conv).split('\n')] 

630 

631 return '\n'.join(lines) 

632 

633 def get_memory_estimate(self) -> MemoryEstimate: 

634 return self.rho_nn_conv.get_memory_estimate() 

635 

636 @property 

637 def myt(self) -> list[int]: 

638 """ List of indices corresponding to the time indices on held on this rank. """ 

639 return self.rho_nn_conv.my_work() 

640 

641 def work_loop(self, 

642 rank: int) -> Generator[ConvolutionDensityMatrixMetadata | None, None, None]: 

643 nt = len(self.times) 

644 ntperrank = (nt + self.loop_comm.size - 1) // self.loop_comm.size 

645 

646 for p, pulse in enumerate(self.pulses): 

647 for localt in range(ntperrank): 

648 globalt = rank + localt * self.loop_comm.size 

649 if globalt < nt: 

650 yield ConvolutionDensityMatrixMetadata(density_matrices=self, globalt=globalt, 

651 localt=localt, globalp=p, localp=p) 

652 else: 

653 yield None 

654 

655 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]: 

656 parameters = self.rho_nn_conv.rho_nn_reader._parameters 

657 flt = (slice(parameters.n1size), slice(parameters.n2size)) 

658 

659 dist_buffer = self.rho_nn_conv.dist_buffer # Perform the redistribution 

660 self.ksd.distribute(self.calc_comm) 

661 

662 if self.calc_comm.rank != 0: 

663 assert len(self.myt) == 0, self.myt 

664 

665 for work in self.local_work_plan: 

666 if self.calc_comm.rank == 0: 

667 assert self.myt[work.localt] == work.globalt 

668 localflt = flt + (work.localp, work.localt) 

669 

670 matrices: dict[int, NDArray[np.complex128] | None] = dict() 

671 for derivative in self.derivative_order_s: 

672 if self.calc_comm.rank > 0: 

673 matrices[derivative] = None 

674 continue 

675 # Buffer shape is i, a, pulses, times 

676 if 'Re' in self.reim: 

677 Rerho_ia = dist_buffer._get_real(derivative)[localflt] 

678 if 'Im' in self.reim: 

679 Imrho_ia = dist_buffer._get_imag(derivative)[localflt] 

680 if 'Re' in self.reim and 'Im' in self.reim: 

681 # Complex result 

682 # Compared to numpy, we use another convention, hence the minus sign 

683 rho_ia = Rerho_ia - 1j * Imrho_ia 

684 elif 'Re' in self.reim: 

685 # Real result 

686 rho_ia = Rerho_ia 

687 else: 

688 rho_ia = -1j * Imrho_ia 

689 matrices[derivative] = rho_ia 

690 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm) 

691 

692 yield work, dm 

693 

694 def parallel_prepare(self): 

695 self.rho_nn_conv.dist_buffer # Perform the redistribution 

696 

697 

698class ExactFourierTransformer: 

699 

700 def __init__(self, 

701 reader: FrequencyDensityMatrixReader, 

702 perturbation: PerturbationLike, 

703 real: bool = True, 

704 imag: bool = True, 

705 comm: Communicator | None = None, 

706 log: Logger | None = None): 

707 if comm is None: 

708 comm = world 

709 self._comm = comm 

710 self.reader = reader 

711 if log is None: 

712 log = Logger() 

713 self.log = log 

714 

715 self.real = real 

716 self.imag = imag 

717 

718 self.frequency_buffer = DensityMatrixBuffer(self.nnshape, (self.reader.nw, ), np.complex128) 

719 self._ready = False 

720 

721 omega_w = self.reader.frequencies * eV_to_au 

722 

723 self.perturbation = create_perturbation(perturbation) 

724 assert isinstance(self.perturbation, DeltaKick) 

725 

726 # Need a uniform spacing for Simpson's integration rule 

727 dw = omega_w[1] - omega_w[0] 

728 nw = len(omega_w) 

729 if not np.allclose(omega_w[1:] - dw, omega_w[:-1]): 

730 raise ValueError('Variable frequency spacing.') 

731 # integration weights from Simpson's integration rule 

732 self._weight_w = dw / 3 * np.array([1] + [4, 2] * int((nw - 2) / 2) 

733 + [4] * (nw % 2) + [1]) 

734 

735 self._weight_w *= 2 / (2 * np.pi) 

736 self._weight_w /= self.perturbation.strength 

737 

738 @property 

739 def comm(self) -> Communicator: 

740 """ MPI commonicator. 

741 

742 During convolution, the density matrices are distributed with 

743 different chunks on different ranks of this communicator. 

744 """ 

745 return self._comm 

746 

747 @property 

748 def full_nnshape(self) -> tuple[int, int]: 

749 imin, imax, amin, amax = [int(i) for i in self.reader.ksd.ialims()] 

750 return (imax - imin + 1, amax - amin + 1) 

751 

752 @property 

753 def nnshape(self) -> tuple[int, int]: 

754 stridei = (self.full_nnshape[0] + self.comm.size - 1) // self.comm.size 

755 return (stridei, self.full_nnshape[1]) 

756 

757 @property 

758 def myslice(self) -> slice: 

759 return self.slice_of_rank(self.comm.rank) 

760 

761 def slice_of_rank(self, rank: int) -> slice: 

762 stridei = (self.full_nnshape[0] + self.comm.size - 1) // self.comm.size 

763 return slice(rank * stridei, (rank + 1) * stridei) 

764 

765 def parallel_prepare(self): 

766 # Read all the Fourier transform of density matrices 

767 self.frequency_buffer.zero_buffers(real=self.real, imag=self.imag, derivative_order_s=[0]) 

768 for w, freq in enumerate(self.reader.frequencies): 

769 if self.real: 

770 rho_ia = self.reader.read(frequency=freq, real=True)[self.myslice] 

771 self.frequency_buffer[w].real[:len(rho_ia)] = rho_ia 

772 if self.imag: 

773 rho_ia = self.reader.read(frequency=freq, real=False)[self.myslice] 

774 self.frequency_buffer[w].imag[:len(rho_ia)] = rho_ia 

775 self._ready = True 

776 

777 def convolve(self, 

778 time: float, 

779 pulse: Perturbation, 

780 derivative: int) -> DensityMatrixBuffer | None: 

781 if not self._ready: 

782 self.parallel_prepare() 

783 

784 assert isinstance(pulse, PulsePerturbation) 

785 omega_w = self.reader.frequencies * eV_to_au 

786 pulse_w = pulse.pulse.fourier(omega_w) 

787 

788 buffer = DensityMatrixBuffer(self.frequency_buffer.nnshape, (), dtype=np.float64) 

789 

790 self.log.start('convolve') 

791 exp_w = np.exp(-1j * omega_w * time * as_to_au) 

792 exp_w *= self._weight_w * pulse_w * (-1j * omega_w) ** derivative 

793 

794 if self.real: 

795 # Real part 

796 rho_iaw = self.frequency_buffer.real[:] 

797 conv_rho_ia = rho_iaw.real @ exp_w.real 

798 conv_rho_ia -= rho_iaw.imag @ exp_w.imag 

799 

800 buffer.store(True, derivative, conv_rho_ia) 

801 

802 if self.imag: 

803 # Imaginary part 

804 rho_iaw = self.frequency_buffer.imag[:] 

805 conv_rho_ia = rho_iaw.real @ exp_w.real 

806 conv_rho_ia -= rho_iaw.imag @ exp_w.imag 

807 

808 buffer.store(False, derivative, conv_rho_ia) 

809 

810 if self.comm.rank == 0: 

811 full_buffer = DensityMatrixBuffer(self.frequency_buffer.nnshape, (), dtype=np.float64) 

812 full_buffer.zero_buffers(real=self.real, imag=self.imag, derivative_order_s=[derivative]) 

813 

814 # Send to root 

815 for rank in range(self.comm.size): 

816 if rank == self.comm.rank: 

817 # Send own work to root 

818 buffer.send_arrays(self.comm, 0) 

819 

820 if self.comm.rank == 0: 

821 # Receive on root, fill buffer 

822 buffer.recv_arrays(self.comm, rank) 

823 for full_rho_ia, part_rho_ia in zip(full_buffer._iter_buffers(), 

824 buffer._iter_buffers()): 

825 full_slice_ia = full_rho_ia[self.slice_of_rank(rank)] 

826 full_slice_ia[:] = part_rho_ia[:len(full_slice_ia)] 

827 

828 if self.comm.rank == 0: 

829 self.log(f'Convolved density matrices in {self.log.elapsed("convolve"):.1f} s', 

830 who='Response', flush=True) 

831 

832 if self.comm.rank == 0: 

833 return full_buffer 

834 

835 return None 

836 

837 @classmethod 

838 def from_file(cls, 

839 frho_fmt: str, 

840 ksd: KohnShamDecomposition | str, 

841 perturbation: PerturbationLike, 

842 freq_min: float, 

843 freq_spacing: float, 

844 freq_max: float, 

845 real: bool = True, 

846 imag: bool = True, 

847 comm: Communicator | None = None, 

848 log: Logger | None = None) -> ExactFourierTransformer: 

849 frequencies = np.arange(freq_min, freq_max + freq_spacing * 0.1, freq_spacing) 

850 

851 reader = FrequencyDensityMatrixReader(frho_fmt=frho_fmt, 

852 ksd=ksd, 

853 filter_frequencies=frequencies, 

854 real=real, imag=imag, 

855 comm=None, # Look for files on world 

856 log=log) 

857 if reader.frequencies[0] > freq_min: 

858 raise ValueError(f'Could not satisfy lower bound for frequency grid with files matching {frho_fmt}. ' 

859 f'Found {reader.frequencies[0]:.2f}eV, reqested {freq_min:.2f}eV.') 

860 

861 if reader.frequencies[-1] > freq_max: 

862 raise ValueError(f'Could not satisfy upper bound for frequency grid with files matching {frho_fmt}. ' 

863 f'Found {reader.frequencies[-1]:.2f}eV, reqested {freq_max:.2f}eV.') 

864 

865 dw = reader.frequencies[1] - reader.frequencies[0] 

866 if dw > freq_spacing + 1e-8: 

867 raise ValueError(f'Could not construct sufficiently dense frequency grid with files matching {frho_fmt}. ' 

868 f'Found spacing of {dw:.4f}eV, reqested {freq_spacing:.4f}eV.') 

869 

870 return cls(reader=reader, perturbation=perturbation, real=real, imag=imag, comm=comm, log=log)