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 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 GPAW wave functions file Parameters ---------- ksd KohnShamDecomposition object or filename wfs_fname Filename of the GPAW 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, stride_opts=None, 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.stride_opts = stride_opts 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 if len(pulses) == 1 and list(pulses)[0] == self.perturbation and set(derivative_order_s) == set([0]): density_matrices = TimeDensityMatricesFromWaveFunctions( ksd=self.ksd, wfs_fname=self.wfs_fname, times=times, real=real, imag=imag, log=log, calc_size=self.calc_size, stride_opts=self.stride_opts, stridet=self.stridet) else: # Perform convolution assert not isinstance(self.perturbation, NoPerturbation) 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, stride_opts=self.stride_opts, stridet=self.stridet) return density_matrices def _get_frequency_density_matrices(self, frequencies: list[float] | NDArray[np.float64], derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> FrequencyDensityMatrices: assert not isinstance(self.perturbation, NoPerturbation) density_matrices = FrequencyDensityMatricesFromWaveFunctions( ksd=self.ksd, wfs_fname=self.wfs_fname, frequencies=frequencies, derivative_order_s=derivative_order_s, real=real, imag=imag, calc_size=self.calc_size, stride_opts=self.stride_opts, 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 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: assert isinstance(self.perturbation, DeltaKick) assert np.isclose(self.perturbation.strength, 1e-5) # TODO easily fixed density_matrices = ConvolutionDensityMatricesFromFrequency( ksd=self.ksd, frho_fmt=self.frho_fmt, 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], derivative_order_s: list[int] = [0], real: bool = True, imag: bool = True, log: Logger | None = None, ) -> FrequencyDensityMatrices: assert isinstance(self.perturbation, DeltaKick) density_matrices = FrequencyDensityMatricesFromDisk( ksd=self.ksd, frho_fmt=self.frho_fmt, frequencies=frequencies, derivative_order_s=derivative_order_s, real=real, imag=imag, kickstr=self.perturbation.strength, 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 as * {tag} - Derivative tag, '', '-Iomega', or '-omega2' * {pulsefreq} - Pulse frequency in eV * {pulsefwhm} - Pulse FWHM in 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): # TODO create composite object with responses to many pulses 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: # 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], 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')