from __future__ import annotations
import numpy as np
from pathlib import Path
from gpaw.mpi import world
from ..density_matrices.time import ConvolutionDensityMatrices
from ..utils import get_gaussian_pulse_values
[docs]
def calculate_pulserho_and_save(pulserho_fmt: str,
density_matrices: ConvolutionDensityMatrices):
""" Read density matrices in frequency space from disk, convolve with
Gaussian laser pulse, and inverse Fourier transform to get a time domain response.
Save the pulse density matrces for selected times in the simulation
Read density matrices in frequency space from disk.
Parameters
----------
pulserho_fmt
Formatting string for the density matrices saved to disk.
Example:
* pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`.
Accepts variables
* `{time}` - Time in units of as.
* `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`.
* `{pulsefreq}` - Pulse frequency in units of eV.
* `{pulsefwhm}` - Pulse FWHM in units of fs.
density_matrices
Object that gives the density matrix in the time domain.
"""
calc_comm = density_matrices.calc_comm
log = density_matrices.log
nlocaltot = len(density_matrices.local_work_plan)
if world.rank == 0:
log('Read frequency density matrix', flush=True)
tags_keys = [(tag, key) for s, (tag, key) in enumerate(
[('', 'rho_p'), ('-Iomega', 'drho_p'), ('-omega2', 'ddrho_p')]) if s in density_matrices.derivative_order_s]
# Do convolution time-by-time
for ndone, (work, dm) in enumerate(density_matrices, 1):
avg = log.elapsed('read')/ndone
estrem = avg * (nlocaltot - ndone)
log(f'Calculated t{work.time:09.1f} '
f'Avg: {avg:10.3f}s ETA: {estrem:10.3f}s', flush=True)
for tag, key in tags_keys:
fname_kw = dict(time=work.time, tag=tag, **get_gaussian_pulse_values(work.pulse))
fpath = Path(pulserho_fmt.format(**fname_kw))
rho_p = getattr(dm, key)
if calc_comm.rank == 0:
assert isinstance(rho_p, np.ndarray)
fpath.parent.mkdir(parents=True, exist_ok=True)
np.save(str(fpath), rho_p)