Coverage for gpaw/test/response/mpa_interpolation_scalar.py: 96%

1from __future__ import annotations 

2import numpy as np 

3from numpy.linalg import eigvals 

4from typing import List, Tuple 

5from gpaw.typing import Array1D, Array2D 

6 

7 

8null_pole_thr = 1e-5 

9pole_resolution = 1e-5 

10epsilon = 1e-8 

11 

12# ------------------------------------------------------------- 

13# Old reference code 

14# ------------------------------------------------------------- 

15 

16# 1 pole 

17 

18 

19def Xeval(Omega_GGp, residues_GGp, omega_w): 

20 X_GGpw = ( 

21 residues_GGp[..., :, np.newaxis] * 2 * Omega_GGp[..., :, np.newaxis] / 

22 (omega_w[None, None, None, :]**2 - Omega_GGp[..., :, np.newaxis]**2) 

23 ) 

24 

25 return np.sum(X_GGpw, axis=2) 

26 

27 

28def mpa_cond1(z: tuple[complex, complex] | Array1D, 

29 E2: tuple[complex] | Array1D) -> \ 

30 tuple[[complex], [float]] | Array2D: 

31 PPcond_rate = 0 

32 if abs(E2) < null_pole_thr: # need to check also NAN(abs(E)) 

33 PPcond_rate = 1 

34 elif np.real(E2) > 0.: 

35 pass 

36 else: 

37 PPcond_rate = 1 

38 

39 # DALV: since MPA uses complex poles we need to guarantee the time ordering 

40 E2 = complex(abs(E2.real), -abs(E2.imag)) 

41 E = np.emath.sqrt(E2) 

42 

43 return E, PPcond_rate 

44 

45 

46def mpa_E_1p_solver(z, x): 

47 E2 = (x[0] * z[0]**2 - x[1] * z[1]**2) / (x[0] - x[1]) 

48 E, PPcond_rate = mpa_cond1(z, E2) 

49 return E, PPcond_rate 

50 

51 

52def pole_is_out(i, wmax, thr, E): 

53 is_out = False 

54 

55 if np.real(E[i]) > wmax: 

56 is_out = True 

57 

58 j = 0 

59 while j < i and not is_out: 

60 if abs(np.real(E[i]) - np.real(E[j])) < thr: 

61 is_out = True 

62 if abs(np.real(E[j])) > max(abs(np.imag(E[j])), 

63 abs(np.imag(E[i]))): 

64 E[j] = np.mean([np.real(E[j]), np.real(E[i])]) - 1j * \ 

65 max(abs(np.imag(E[j])), abs(np.imag(E[i]))) 

66 else: 

67 E[j] = np.mean([np.real(E[j]), np.real(E[i])]) - 1j * \ 

68 min(abs(np.imag(E[j])), abs(np.imag(E[i]))) 

69 j = j + 1 

70 

71 return is_out 

72 

73 

74def mpa_cond(npols: int, z: List[complex], E) ->\ 

75 Tuple[int, List[bool], List[complex]]: 

76 PPcond = np.full(npols, False) 

77 npr = npols 

78 wmax = np.max(np.real(np.emath.sqrt(z))) * 1.5 

79 Eaux = np.emath.sqrt(E) 

80 

81 i = 0 

82 while i < npr: 

83 Eaux[i] = max(abs(np.real(Eaux[i])), abs(np.imag(Eaux[i]))) - 1j * \ 

84 min(abs(np.real(Eaux[i])), abs(np.imag(Eaux[i]))) 

85 is_out = pole_is_out(i, wmax, pole_resolution, Eaux) 

86 

87 if is_out: 

88 Eaux[i] = np.emath.sqrt(E[npr - 1]) 

89 Eaux[i] = max(abs(np.real(Eaux[i])), abs(np.imag(Eaux[i]))) - 1j \ 

90 * min(abs(np.real(Eaux[i])), abs(np.imag(Eaux[i]))) 

91 PPcond[npr - 1] = True 

92 npr = npr - 1 

93 else: 

94 i = i + 1 

95 

96 E[:npr] = Eaux[:npr] 

97 if npr < npols: 

98 E[npr:npols] = 2 * wmax - 0.01j 

99 PPcond[npr:npols] = True 

100 

101 return E, npr, PPcond 

102 

103 

104def mpa_R_1p_fit(npols, npr, w, x, E): 

105 # Transforming the problem into a 2* larger least square with real numbers: 

106 A = np.zeros((4, 2), dtype='complex64') 

107 b = np.zeros((4), dtype='complex64') 

108 for k in range(2): 

109 b[2 * k] = np.real(x[k]) 

110 b[2 * k + 1] = np.imag(x[k]) 

111 A[2 * k][0] = 2. * np.real(E / (w[k]**2 - E**2)) 

112 A[2 * k][1] = -2. * np.imag(E / (w[k]**2 - E**2)) 

113 A[2 * k + 1][0] = 2. * np.imag(E / (w[k]**2 - E**2)) 

114 A[2 * k + 1][1] = 2. * np.real(E / (w[k]**2 - E**2)) 

115 

116 Rri = np.linalg.lstsq(A, b, rcond=None)[0] 

117 R = Rri[0] + 1j * Rri[1] 

118 return R 

119 

120 

121def mpa_R_fit(npols, npr, w, x, E): 

122 # Transforming the problem into a 2* larger least square with real numbers: 

123 A = np.zeros((2 * npols * 2, npr * 2), dtype='complex64') 

124 b = np.zeros((2 * npols * 2), dtype='complex64') 

125 for k in range(2 * npols): 

126 b[2 * k] = np.real(x[k]) 

127 b[2 * k + 1] = np.imag(x[k]) 

128 for i in range(npr): 

129 A[2 * k][2 * i] = 2. * np.real(E[i] / (w[k]**2 - E[i]**2)) 

130 A[2 * k][2 * i + 1] = -2. * np.imag(E[i] / (w[k]**2 - E[i]**2)) 

131 A[2 * k + 1][2 * i] = 2. * np.imag(E[i] / (w[k]**2 - E[i]**2)) 

132 A[2 * k + 1][2 * i + 1] = 2. * np.real(E[i] / (w[k]**2 - E[i]**2)) 

133 

134 Rri = np.linalg.lstsq(A, b, rcond=None)[0] 

135 R = np.zeros(npols, dtype='complex64') 

136 R[:npr] = Rri[::2] + 1j * Rri[1::2] 

137 return R 

138 

139 

140def mpa_E_solver_Pade(npols, z, x): 

141 b_m1 = b = np.zeros(npols + 1, dtype='complex64') 

142 b_m1[0] = b[0] = 1 

143 c = np.copy(x) 

144 

145 for i in range(1, 2 * npols): 

146 

147 c_m1 = np.copy(c) 

148 c[i:] = (c_m1[i - 1] - c_m1[i:]) / ((z[i:] - z[i - 1]) * c_m1[i:]) 

149 

150 b_m2 = np.copy(b_m1) 

151 b_m1 = np.copy(b) 

152 

153 b = b_m1 - z[i - 1] * c[i] * b_m2 

154 b_m2[npols:0:-1] = c[i] * b_m2[npols - 1::-1] 

155 b[1:] = b[1:] + b_m2[1:] 

156 

157 Companion = np.polynomial.polynomial.polycompanion(b[:npols + 1]) 

158 # DALV: /b[npols] it is carried inside 

159 

160 E = eigvals(Companion) 

161 

162 # DALV: here we need to force real(E) to be positive. 

163 # This is because of the way the residue integral is performed, later. 

164 E, npr, PPcond = mpa_cond(npols, z, E) 

165 

166 return E, npr, PPcond 

167 

168 

169def mpa_RE_solver(npols, w, x): 

170 if npols == 1: # DALV: we could particularize the solution for 2-3 poles 

171 E, PPcond_rate = mpa_E_1p_solver(w, x) 

172 R = mpa_R_1p_fit(1, 1, w, x, E) 

173 # DALV: if PPcond_rate=0, R = x[1]*(z[1]**2-E**2)/(2*E) 

174 MPred = 0 

175 else: 

176 # DALV: Pade-Thiele solver (mpa_sol='PT') 

177 E, npr, PPcond = mpa_E_solver_Pade(npols, w**2, x) 

178 R = mpa_R_fit(npols, npr, w, x, E) 

179 

180 PPcond_rate = 1 

181 MPred = 1 

182 

183 # for later: MP_err = err_func_X(np, R, E, w, x) 

184 

185 return R, E, MPred, PPcond_rate # , MP_err