Coverage for rhodent/response/gpaw.py: 94%

1from __future__ import annotations 

2 

3from typing import Collection 

4 

5import numpy as np 

6from pathlib import Path 

7 

8from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition 

9 

10from .base import BaseResponse 

11from ..perturbation import create_perturbation, PerturbationLike, NoPerturbation 

12from ..density_matrices.frequency import (FrequencyDensityMatrices, 

13 FrequencyDensityMatricesFromWaveFunctions) 

14from ..density_matrices.time import (ConvolutionDensityMatrices, 

15 ConvolutionDensityMatricesFromWaveFunctions, 

16 TimeDensityMatricesFromWaveFunctions) 

17from ..utils import Logger 

18from ..typing import Array1D 

19 

20 

21class ResponseFromWaveFunctions(BaseResponse): 

22 

23 """ Response from time-dependent wave functions file written by GPAW. 

24 

25 Parameters 

26 ---------- 

27 ksd 

28 KohnShamDecomposition object or file name. 

29 wfs_fname 

30 File name of the time-dependent wave functions file, written by :class:`WaveFunctionsWriter`. 

31 perturbation 

32 Perturbation that was present during time propagation. 

33 calc_size 

34 Size of the calculation communicator. 

35 """ 

36 

37 def __init__(self, 

38 wfs_fname: Path | str, # File name of wfs.ulm file 

39 ksd: KohnShamDecomposition | str, 

40 perturbation: PerturbationLike = None, 

41 calc_size: int = 1, 

42 stridet: int = 1): 

43 super().__init__(ksd=ksd, 

44 perturbation=perturbation, 

45 calc_size=calc_size) 

46 

47 self.wfs_fname = str(wfs_fname) 

48 

49 # Options for reading the wfs.ulm file 

50 self.stridet = stridet 

51 

52 def __str__(self) -> str: 

53 lines = [f'{self.__class__.__name__}'] 

54 lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}'] 

55 lines += [f' wfs_fname: {self.wfs_fname}'] 

56 lines += [' perturbation:'] 

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

58 return '\n'.join(lines) 

59 

60 def _get_time_density_matrices(self, 

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

62 pulses: Collection[PerturbationLike], 

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

64 real: bool = True, 

65 imag: bool = True, 

66 log: Logger | None = None, 

67 ) -> ConvolutionDensityMatrices: 

68 density_matrices: ConvolutionDensityMatrices 

69 

70 # Perform convolution if pulses differ from perturbation or higher derivatives needed 

71 if (all(create_perturbation(pulse) == self.perturbation for pulse in pulses) and 

72 all(derivative == 0 for derivative in derivative_order_s)): 

73 # No convolution needed 

74 density_matrices = TimeDensityMatricesFromWaveFunctions( 

75 ksd=self.ksd, 

76 wfs_fname=self.wfs_fname, 

77 times=times, 

78 real=real, 

79 imag=imag, 

80 log=log, 

81 calc_size=self.calc_size, 

82 stridet=self.stridet) 

83 else: 

84 # Perform convolution 

85 if isinstance(self.perturbation, NoPerturbation): 

86 raise ValueError('Perturbation must be given to perform pulse convolution.') 

87 density_matrices = ConvolutionDensityMatricesFromWaveFunctions( 

88 ksd=self.ksd, 

89 wfs_fname=self.wfs_fname, 

90 perturbation=self.perturbation, 

91 pulses=pulses, 

92 times=times, 

93 derivative_order_s=derivative_order_s, 

94 real=real, 

95 imag=imag, 

96 log=log, 

97 calc_size=self.calc_size, 

98 stridet=self.stridet) 

99 

100 return density_matrices 

101 

102 def _get_frequency_density_matrices(self, 

103 frequencies: list[float] | Array1D[np.float64], 

104 frequency_broadening: float = 0, 

105 real: bool = True, 

106 imag: bool = True, 

107 log: Logger | None = None, 

108 ) -> FrequencyDensityMatrices: 

109 if isinstance(self.perturbation, NoPerturbation): 

110 raise ValueError('Perturbation must be given to normalize Fourier transform.') 

111 density_matrices = FrequencyDensityMatricesFromWaveFunctions( 

112 ksd=self.ksd, 

113 wfs_fname=self.wfs_fname, 

114 perturbation=self.perturbation, 

115 frequencies=frequencies, 

116 frequency_broadening=frequency_broadening, 

117 real=real, 

118 imag=imag, 

119 calc_size=self.calc_size, 

120 log=log, 

121 stridet=self.stridet) 

122 

123 return density_matrices