Coverage for strongcoca / calculators / retarded_polarizability_calculator.py: 100%

1import logging 

2 

3import numpy as np 

4from numpy.linalg import solve 

5 

6from .. import CoupledSystem 

7from ..response.utilities import Broadening, GaussianBroadening 

8from ..types import Array 

9from ..units import eV_to_au, au_to_k 

10from .base_calculator import BaseCalculator 

11from ..compute_config import get_compute_config 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16class RetardedPolarizabilityCalculator(BaseCalculator): 

17 """Instances of this class enable the calculation of correlation energy 

18 and spectrum, including retardation effects, of a coupled system consisting 

19 of polarizable objects which all have an internal representation in the 

20 form of a polarizability tensor. 

21 

22 Spectra calculated with this calculator derive broadening from 

23 the underlying polarizable units. No additional or separate broadening 

24 is added. 

25 

26 Parameters 

27 ---------- 

28 coupled_system 

29 Coupled system for which to carry out calculations. 

30 imaginary_frequencies 

31 Frequency grid along imaginary axis used for correlation energy 

32 calculation; eV by default, optionally atomic units (see :attr:`units`). 

33 units 

34 `eV` to specify energies in eV or `au` to specify inputs in atomic units. 

35 

36 This parameter determines whether a unit conversion should be performed 

37 during initialization. It does not affect other instance methods or members. 

38 name 

39 Name of response. 

40 """ 

41 def __init__(self, 

42 coupled_system: CoupledSystem, 

43 imaginary_frequencies: Array, 

44 units: str = 'eV', 

45 name: str = 'RetardedPolarizabilityCalculator') -> None: 

46 super().__init__(coupled_system, broadening=Broadening(), name=name) 

47 

48 imaginary_frequencies = np.asarray(imaginary_frequencies) 

49 if units == 'eV': 

50 imaginary_frequencies = imaginary_frequencies * eV_to_au 

51 elif units != 'au': 

52 raise ValueError(f"units has to be 'eV' or 'au', not '{units}'") 

53 

54 self._ifreq_w = imaginary_frequencies 

55 

56 difreq_w = self._ifreq_w[1:] - self._ifreq_w[:-1] 

57 self._difreq = difreq_w[0] 

58 if not np.allclose(difreq_w, self._difreq): 

59 raise ValueError('Frequency grid needs to be equally spaced.') 

60 

61 def _build_coupling_matrix(self, k) -> np.ndarray: 

62 """Build retarded coupling matrix for wave number k. 

63 

64 Coupling matrix is formed of blocks: 

65 Each block is a 3x3 retarded dipole-dipole tensor and there are NxN 

66 blocks in total (N=number of units in coupled system). 

67 The 3x3 blocks on the diagonal are zero. 

68 """ 

69 

70 Ni = len(self._coupled_system) 

71 N3 = 3 * Ni 

72 

73 pos_iv = np.array([pu._position for pu in self._coupled_system], dtype=float) 

74 

75 # R[i,j] = r_j - r_i 

76 R_ijv = pos_iv[np.newaxis, :] - pos_iv[:, np.newaxis] 

77 r_ij = np.linalg.norm(R_ijv, axis=-1) 

78 

79 # avoid division by zero temporarily 

80 np.fill_diagonal(r_ij, 1.0) 

81 

82 kr_ij = k * r_ij 

83 exp_ikr_ij = np.exp(1j * kr_ij) 

84 

85 a_ij = exp_ikr_ij * (1.0 - 1j * kr_ij - kr_ij**2) / r_ij**3 

86 b_ij = -exp_ikr_ij * (3.0 - 3j * kr_ij - kr_ij**2) / r_ij**5 

87 

88 RR_ijvw = R_ijv[:, :, :, np.newaxis] * R_ijv[:, :, np.newaxis, :] 

89 I_vw = np.eye(3)[np.newaxis, np.newaxis, :, :] 

90 

91 # Dipole-dipole tensor with retardation (free space Green's dyadic) 

92 # for all pairs 

93 T_ijvw = ( 

94 a_ij[:, :, np.newaxis, np.newaxis] * I_vw 

95 + b_ij[:, :, np.newaxis, np.newaxis] * RR_ijvw 

96 ) 

97 

98 # no self term; "self" retardation should be included in polarizability 

99 diag = np.arange(Ni) 

100 T_ijvw[diag, diag] = 0.0 

101 

102 return T_ijvw.transpose(0, 2, 1, 3).reshape(N3, N3).astype(complex) # type: ignore 

103 

104 def _calculate_correlation_energy(self) -> float: 

105 """Returns the correlation energy of the coupled system in atomic units.""" 

106 Ni = len(self._coupled_system) 

107 if Ni == 0: 

108 return 0.0 

109 

110 gpu_backend = get_compute_config().backend 

111 if gpu_backend != 'none': # pragma: no cover 

112 raise NotImplementedError( 

113 'GPU backend is not supported in RetardedPolarizabilityCalculator.' 

114 ) 

115 

116 W = len(self._ifreq_w) 

117 N3 = 3 * Ni 

118 

119 dm_wNvv = np.empty((W, Ni, 3, 3), dtype=complex) 

120 for i, pu in enumerate(self._coupled_system): 

121 dm_wNvv[:, i] = pu._get_dynamic_polarizability_imaginary_frequency(self._ifreq_w) 

122 

123 trace_logD_w = np.empty(W, dtype=complex) 

124 trace_chiK_w = np.empty(W, dtype=complex) 

125 

126 I_nn = np.eye(N3, dtype=complex) 

127 

128 # Current implementation: loop over frequencies 

129 for w, ifreq in enumerate(self._ifreq_w): 

130 k = 1j * ifreq * au_to_k 

131 K_nn = self._build_coupling_matrix(k) 

132 K_block = K_nn.reshape(Ni, 3, N3) 

133 

134 chi_K_nn = (dm_wNvv[w] @ K_block).reshape(N3, N3) 

135 

136 D_nn = I_nn - chi_K_nn 

137 sign, logabsdet = np.linalg.slogdet(D_nn) 

138 trace_logD_w[w] = np.log(sign) + logabsdet 

139 trace_chiK_w[w] = np.trace(chi_K_nn) 

140 

141 integrand_w = trace_logD_w + trace_chiK_w.conj() 

142 integral = np.sum(integrand_w) * self._difreq 

143 energy: float = float(np.real(integral)) / (2 * np.pi) 

144 return energy 

145 

146 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray: 

147 if any(isinstance(pu.broadening, GaussianBroadening) 

148 for pu in self._coupled_system): 

149 raise NotImplementedError( 

150 'Gaussian broadening is not supported in RetardedPolarizabilityCalculator. ' 

151 ) 

152 

153 gpu_backend = get_compute_config().backend 

154 if gpu_backend != 'none': # pragma: no cover 

155 raise NotImplementedError( 

156 'GPU backend is not supported in RetardedPolarizabilityCalculator.' 

157 ) 

158 

159 freq_w = np.asarray(frequencies) 

160 Ni = len(self._coupled_system) 

161 W = len(freq_w) 

162 N3 = 3 * Ni 

163 

164 dm_wNvv = np.empty((W, Ni, 3, 3), dtype=complex) 

165 for i, pu in enumerate(self._coupled_system): 

166 dm_wNvv[:, i] = pu._get_dynamic_polarizability(freq_w) 

167 

168 dm_wvv = np.empty((W, 3, 3), dtype=complex) 

169 to_MiB = 1024 ** -2 

170 mem_limit = 10 * get_compute_config().max_solve_mem / to_MiB 

171 syssize = 3 ** 2 * Ni ** 2 * 8 # size of A matrix per frequency in bytes 

172 chunksize = max(1, int(mem_limit) // syssize) 

173 for indices in np.array_split(np.arange(W), (W + chunksize - 1) // chunksize): 

174 freq_c = freq_w[indices] 

175 k_c = freq_c * au_to_k 

176 K_cNN = np.empty((len(indices), N3, N3), dtype=complex) 

177 for j, k in enumerate(k_c): 

178 K_cNN[j] = self._build_coupling_matrix(k) 

179 

180 dm_wvv[indices] = self._get_dynamic_polarizability_chunk( 

181 dm_wNvv[indices], K_cNN 

182 ) 

183 

184 return dm_wvv 

185 

186 def _get_dynamic_polarizability_chunk( 

187 self, 

188 dm_wNvv: np.ndarray, 

189 K_wNN: np.ndarray) -> np.ndarray: 

190 """CPU path: solve for a prebuilt frequency chunk. 

191 

192 Parameters 

193 ---------- 

194 dm_wNvv 

195 Per-unit polarizabilities for this chunk, shape (W, Ni, 3, 3). 

196 K_wNN 

197 Retarded coupling matrices for this chunk, shape (W, 3*Ni, 3*Ni). 

198 """ 

199 W, Ni = dm_wNvv.shape[:2] 

200 N3 = 3 * Ni 

201 

202 K_block = K_wNN.reshape(W, Ni, 3, N3) 

203 A_wnn = np.eye(N3, dtype=complex)[None] + (dm_wNvv @ K_block).reshape(W, N3, N3) 

204 rhs_wnv = np.ascontiguousarray(dm_wNvv.reshape(W, N3, 3)) 

205 red_wnv = solve(A_wnn, rhs_wnv) 

206 return red_wnv.reshape(W, Ni, 3, 3).sum(axis=1) # type: ignore 

207 

208 def _get_dynamic_polarizability_imaginary_frequency( 

209 self, frequencies: Array) -> np.ndarray: 

210 raise NotImplementedError()