Source code for rhodent.response.gpaw

from __future__ import annotations

from typing import Collection

import numpy as np
from numpy.typing import NDArray
from pathlib import Path

from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition

from .base import BaseResponse
from ..perturbation import create_perturbation, DeltaKick, PerturbationLike, NoPerturbation
from ..density_matrices.frequency import (FrequencyDensityMatrices,
                                          FrequencyDensityMatricesFromDisk,
                                          FrequencyDensityMatricesFromWaveFunctions)
from ..density_matrices.time import (ConvolutionDensityMatricesFromDisk,
                                     ConvolutionDensityMatricesFromFrequency,
                                     ConvolutionDensityMatrices,
                                     ConvolutionDensityMatricesFromWaveFunctions,
                                     TimeDensityMatricesFromWaveFunctions)
from ..utils import Logger


[docs] class ResponseFromWaveFunctions(BaseResponse): """ Response from time-dependent wave functions file written by GPAW. Parameters ---------- ksd KohnShamDecomposition object or filename. wfs_fname Filename of the time-dependent wave functions file, written by ``WaveFunctionsWriter``. perturbation Perturbation that was present during time propagation. calc_size Size of the calculation communicator. """ def __init__(self, wfs_fname: Path | str, # File name of wfs.ulm file ksd: KohnShamDecomposition | str, perturbation: PerturbationLike = None, calc_size: int = 1, stridet: int = 1): super().__init__(ksd=ksd, perturbation=perturbation, calc_size=calc_size) self.wfs_fname = str(wfs_fname) # Options for reading the wfs.ulm file self.stridet = stridet def __str__(self) -> str: lines = [f'{self.__class__.__name__}'] lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}'] lines += [f' wfs_fname: {self.wfs_fname}'] lines += [' perturbation:'] lines += [' ' + line for line in str(self.perturbation).split('\n')] return '\n'.join(lines) def _get_time_density_matrices(self, times: list[float] | NDArray[np.float64], pulses: Collection[PerturbationLike], derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> ConvolutionDensityMatrices: density_matrices: ConvolutionDensityMatrices # Perform convolution if pulses differ from perturbation or higher derivatives needed if (all(create_perturbation(pulse) == self.perturbation for pulse in pulses) and all(derivative == 0 for derivative in derivative_order_s)): # No convolution needed density_matrices = TimeDensityMatricesFromWaveFunctions( ksd=self.ksd, wfs_fname=self.wfs_fname, times=times, real=real, imag=imag, log=log, calc_size=self.calc_size, stridet=self.stridet) else: # Perform convolution if isinstance(self.perturbation, NoPerturbation): raise ValueError('Perturbation must be given to perform pulse convolution.') density_matrices = ConvolutionDensityMatricesFromWaveFunctions( ksd=self.ksd, wfs_fname=self.wfs_fname, perturbation=self.perturbation, pulses=pulses, times=times, derivative_order_s=derivative_order_s, real=real, imag=imag, log=log, calc_size=self.calc_size, stridet=self.stridet) return density_matrices def _get_frequency_density_matrices(self, frequencies: list[float] | NDArray[np.float64], frequency_broadening: float = 0, derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> FrequencyDensityMatrices: if isinstance(self.perturbation, NoPerturbation): raise ValueError('Perturbation must be given to normalize Fourier transform.') density_matrices = FrequencyDensityMatricesFromWaveFunctions( ksd=self.ksd, wfs_fname=self.wfs_fname, perturbation=self.perturbation, frequencies=frequencies, frequency_broadening=frequency_broadening, derivative_order_s=derivative_order_s, real=real, imag=imag, calc_size=self.calc_size, log=log, stridet=self.stridet) return density_matrices
[docs] class ResponseFromFourierTransform(BaseResponse): """ Response from Fourier transform of density matrices save on disk. Parameters ---------- frho_fmt Formatting string for the density matrices in frequency space saved to disk. Example: * frho_fmt = `frho/w{freq:05.2f}-{reim}.npy`. Accepts variables: * `{freq}` - Frequency in units of eV. * `{reim}` - `'Re'` or `'Im'` for Fourier transform of real/imaginary part of density matrix. ksd KohnShamDecomposition object or filename. perturbation Perturbation that was present during time propagation. calc_size Size of the calculation communicator. """ def __init__(self, frho_fmt: str, ksd: KohnShamDecomposition | str, perturbation: PerturbationLike = None, calc_size: int = 1): super().__init__(ksd=ksd, perturbation=perturbation, calc_size=calc_size) if not isinstance(self.perturbation, DeltaKick): raise NotImplementedError('Only delta kick implemented') self.frho_fmt = frho_fmt def __str__(self) -> str: lines = [f'{self.__class__.__name__}'] lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}'] lines += [f' frho_fmt: {self.frho_fmt}'] lines += [' perturbation:'] lines += [' ' + line for line in str(self.perturbation).split('\n')] return '\n'.join(lines) def _get_time_density_matrices(self, times: list[float] | NDArray[np.float64], pulses: Collection[PerturbationLike], derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> ConvolutionDensityMatrices: density_matrices = ConvolutionDensityMatricesFromFrequency( ksd=self.ksd, frho_fmt=self.frho_fmt, perturbation=self.perturbation, pulses=pulses, times=times, derivative_order_s=derivative_order_s, real=real, imag=imag, log=log, calc_size=self.calc_size) return density_matrices def _get_frequency_density_matrices(self, frequencies: list[float] | NDArray[np.float64], frequency_broadening: float = 0, derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> FrequencyDensityMatrices: if frequency_broadening > 0: raise NotImplementedError(f'Frequency broadening not implemented for {self.__class__.__name__}') density_matrices = FrequencyDensityMatricesFromDisk( ksd=self.ksd, frho_fmt=self.frho_fmt, perturbation=self.perturbation, frequencies=frequencies, derivative_order_s=derivative_order_s, real=real, imag=imag, log=log, calc_size=self.calc_size) return density_matrices
[docs] class ResponseFromDensityMatrices(BaseResponse): """ Response from density matrices saved on 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. ksd KohnShamDecomposition object or filename. perturbation Perturbation that was present during time propagation. calc_size Size of the calculation communicator. """ def __init__(self, pulserho_fmt: str, ksd: KohnShamDecomposition | str, perturbation: PerturbationLike = None, calc_size: int = 1): super().__init__(ksd=ksd, perturbation=perturbation, calc_size=calc_size) self.pulserho_fmt = pulserho_fmt def __str__(self) -> str: lines = [f'{self.__class__.__name__}'] lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}'] lines += [f' pulserho_fmt: {self.pulserho_fmt}'] lines += [' perturbation:'] lines += [' ' + line for line in str(self.perturbation).split('\n')] return '\n'.join(lines) def _get_time_density_matrices(self, times: list[float] | NDArray[np.float64], pulses: Collection[PerturbationLike], derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> ConvolutionDensityMatrices: if (len(pulses) == 1 and list(pulses)[0] == self.perturbation) or isinstance(self.perturbation, NoPerturbation): # Yield the density matrices without performing convolution density_matrices = ConvolutionDensityMatricesFromDisk( ksd=self.ksd, pulserho_fmt=self.pulserho_fmt, times=times, pulses=pulses, derivative_order_s=derivative_order_s, real=real, imag=imag, log=log, calc_size=self.calc_size) else: raise NotImplementedError('Pulse convolution of density matrices on disk is not implemented') return density_matrices def _get_frequency_density_matrices(self, frequencies: list[float] | NDArray[np.float64], frequency_broadening: float = 0, derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> FrequencyDensityMatrices: raise NotImplementedError('Fourier transformation of density matrices on disk is not implemented')