Source code for rhodent.response.gpaw

from __future__ import annotations

from typing import Collection

import numpy as np
from pathlib import Path

from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition

from .base import BaseResponse
from ..perturbation import create_perturbation, PerturbationLike, NoPerturbation
from ..density_matrices.frequency import (FrequencyDensityMatrices,
                                          FrequencyDensityMatricesFromWaveFunctions)
from ..density_matrices.time import (ConvolutionDensityMatrices,
                                     ConvolutionDensityMatricesFromWaveFunctions,
                                     TimeDensityMatricesFromWaveFunctions)
from ..utils import Logger
from ..typing import Array1D


[docs] class ResponseFromWaveFunctions(BaseResponse): """ Response from time-dependent wave functions file written by GPAW. Parameters ---------- ksd KohnShamDecomposition object or file name. wfs_fname File name 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] | Array1D[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] | Array1D[np.float64], frequency_broadening: float = 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, real=real, imag=imag, calc_size=self.calc_size, log=log, stridet=self.stridet) return density_matrices