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

49 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-08-01 16:57 +0000

1from __future__ import annotations 

2 

3from abc import abstractmethod, ABC 

4from typing import Generator, Collection 

5 

6import numpy as np 

7 

8from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition 

9 

10from ..perturbation import create_perturbation, Perturbation, PerturbationLike 

11from ..density_matrices.density_matrix import DensityMatrix 

12from ..density_matrices.frequency import (FrequencyDensityMatrixMetadata, 

13 FrequencyDensityMatrices) 

14from ..density_matrices.time import (ConvolutionDensityMatrixMetadata, 

15 ConvolutionDensityMatrices) 

16from ..typing import Array1D 

17from ..utils import add_fake_kpts, Logger 

18 

19 

20class BaseResponse(ABC): 

21 

22 """ Object describing response; obtained from :term:`TDDFT` calculation. 

23 

24 Parameters 

25 ---------- 

26 ksd 

27 KohnShamDecomposition object or file name. 

28 perturbation 

29 The perturbation that was present during the TDDFT calculation. 

30 None to mark it as an ne perturbation. 

31 calc_size 

32 Size of the calculation communicator. 

33 """ 

34 def __init__(self, 

35 ksd: KohnShamDecomposition | str, 

36 perturbation: PerturbationLike = None, 

37 calc_size: int = 1): 

38 if isinstance(ksd, KohnShamDecomposition): 

39 self._ksd = ksd 

40 else: 

41 self._ksd = KohnShamDecomposition(filename=ksd) 

42 add_fake_kpts(self._ksd) 

43 

44 self._perturbation = create_perturbation(perturbation) 

45 self.calc_size = calc_size 

46 

47 def __str__(self) -> str: 

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

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

50 lines += [' perturbation:'] 

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

52 return '\n'.join(lines) 

53 

54 @property 

55 def ksd(self) -> KohnShamDecomposition: 

56 """ Kohn-Sham decomposition object. """ 

57 return self._ksd 

58 

59 @property 

60 def perturbation(self) -> Perturbation: 

61 """ The perturbation that caused this response. """ 

62 return self._perturbation 

63 

64 @abstractmethod 

65 def _get_time_density_matrices(self, 

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

67 pulses: Collection[PerturbationLike], 

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

69 real: bool = True, 

70 imag: bool = True, 

71 log: Logger | None = None, 

72 ) -> ConvolutionDensityMatrices: 

73 raise NotImplementedError 

74 

75 @abstractmethod 

76 def _get_frequency_density_matrices(self, 

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

78 frequency_broadening: float = 0, 

79 real: bool = True, 

80 imag: bool = True, 

81 log: Logger | None = None, 

82 ) -> FrequencyDensityMatrices: 

83 raise NotImplementedError 

84 

85 def iterate_density_matrices_in_time(self, 

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

87 pulses: Collection[PerturbationLike], 

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

89 real: bool = True, 

90 imag: bool = True, 

91 log: Logger | None = None, 

92 ) -> Generator[tuple[ConvolutionDensityMatrixMetadata, 

93 DensityMatrix], None, None]: 

94 """ Obtain density matrices at the given times in response to the given pulses. 

95 

96 If the given pulse(s) differ from the perturbation that caused this response, 

97 then the pulse convolution trick is applied to obtain the response to the given 

98 pulse(s). 

99 

100 Parameters 

101 ---------- 

102 times 

103 Calculate density matrices for these times (or as close to them as possible). In units of as. 

104 pulses 

105 Calculate density matrices in response to these pulses. 

106 derivative_order_s 

107 Calculate density matrix derivatives of the following orders. 

108 ``0`` for plain density matrix and positive integers for derivatives. 

109 real 

110 Calculate the real part of density matrices. 

111 imag 

112 Calculate the imaginary part of density matrices. 

113 log 

114 Logger object. 

115 

116 Yields 

117 ------ 

118 Tuple (work, dm) on the root rank of the calculation communicator: 

119 

120 work 

121 An object representing the metadata (time and pulse) for the work done. 

122 dm 

123 Density matrix for this time and pulse. 

124 """ 

125 density_matrices = self._get_time_density_matrices( 

126 times, pulses, derivative_order_s, real, imag, log) 

127 yield from density_matrices 

128 

129 def iterate_density_matrices_in_frequency(self, 

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

131 frequency_broadening: float = 0, 

132 real: bool = True, 

133 imag: bool = True, 

134 log: Logger | None = None, 

135 ) -> Generator[tuple[FrequencyDensityMatrixMetadata, 

136 DensityMatrix], None, None]: 

137 """ Obtain density matrices at the given frequencies. 

138 

139 Parameters 

140 ---------- 

141 frequencies 

142 Compute density matrices for these frequencies (or as close to them as possible). In units of eV. 

143 frequency_broadening 

144 Gaussian broadening width in atomic units. Default (0) is no broadening. 

145 real 

146 Calculate the Fourier transform of the real part of the density matrix. 

147 imag 

148 Calculate the Fourier transform of the imaginary part of the density matrix. 

149 

150 Yields 

151 ------ 

152 Tuple (work, dm) on the root rank of the calculation communicator: 

153 

154 work 

155 An object representing the metadata (frequency) for the work done. 

156 dm 

157 Density matrix for this frequency. 

158 """ 

159 density_matrices = self._get_frequency_density_matrices( 

160 frequencies, frequency_broadening, real, imag, log) 

161 yield from density_matrices 

162 

163 def write_in_time(self, 

164 pulserho_fmt: str, 

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

166 pulses: Collection[PerturbationLike], 

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

168 real: bool = True, 

169 imag: bool = True, 

170 log: Logger | None = None): 

171 density_matrices = self._get_time_density_matrices( 

172 times, pulses, derivative_order_s, real, imag, log) 

173 density_matrices.write_to_disk(pulserho_fmt) 

174 

175 def write_in_frequency(self, 

176 frho_fmt: str, 

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

178 frequency_broadening: float = 0, 

179 real: bool = True, 

180 imag: bool = True, 

181 log: Logger | None = None): 

182 density_matrices = self._get_frequency_density_matrices( 

183 frequencies, frequency_broadening, real, imag, log) 

184 density_matrices.write_to_disk(frho_fmt)