Coverage for gpaw/purepython.py: 90%

1"""Pure Python implementation of the _gpaw C-extension module. 

2 

3Used if GPAW_NO_C_EXTENSION=1. See also the gpaw.cgpaw module. 

4""" 

5import numpy as np 

6from scipy.interpolate import CubicSpline 

7from gpaw.gpu import cupy as cp, cupy_is_fake 

8from gpaw.typing import Array1D, ArrayND 

9 

10have_openmp = False 

11 

12 

13def gpaw_gpu_init(): 

14 pass 

15 

16 

17def get_num_threads(): 

18 return 1 

19 

20 

21class Spline: 

22 def __init__(self, l, rmax, f_g): 

23 self.spline = CubicSpline(np.linspace(0, rmax, len(f_g)), f_g, 

24 bc_type='clamped') 

25 self.l = l 

26 self.rmax = rmax 

27 

28 def __call__(self, r): 

29 return self.spline(r) 

30 

31 def get_angular_momentum_number(self): 

32 return self.l 

33 

34 def get_cutoff(self): 

35 return self.rmax 

36 

37 def map(self, r_g, out_g): 

38 out_g[:] = self.spline(r_g) 

39 out_g[r_g >= self.rmax] = 0.0 

40 

41 

42def hartree(l: int, 

43 nrdr: np.ndarray, 

44 r: np.ndarray, 

45 vr: np.ndarray) -> None: 

46 p = 0.0 

47 q = 0.0 

48 for g in range(len(r) - 1, 0, -1): 

49 R = r[g] 

50 rl = R**l 

51 dp = nrdr[g] / rl 

52 rlp1 = rl * R 

53 dq = nrdr[g] * rlp1 

54 vr[g] = (p + 0.5 * dp) * rlp1 - (q + 0.5 * dq) / rl 

55 p += dp 

56 q += dq 

57 vr[0] = 0.0 

58 f = 4.0 * np.pi / (2 * l + 1) 

59 vr[1:] += q / r[1:]**l 

60 vr[1:] *= f 

61 

62 

63def unpack(M, M2): 

64 n = len(M2) 

65 p = 0 

66 for r in range(n): 

67 for c in range(r, n): 

68 d = M[p] 

69 M2[r, c] = d 

70 M2[c, r] = d 

71 p += 1 

72 

73 

74def pack(M2): 

75 n = len(M2) 

76 M = np.empty(n * (n + 1) // 2) 

77 p = 0 

78 for r in range(n): 

79 M[p] = M2[r, r] 

80 p += 1 

81 for c in range(r + 1, n): 

82 M[p] = M2[r, c] + M2[c, r] 

83 p += 1 

84 return M 

85 

86 

87def add_to_density(f: float, 

88 psit_X: ArrayND, 

89 nt_X: ArrayND) -> None: 

90 nt_X += f * abs(psit_X)**2 

91 

92 

93def pw_precond(G2_G: Array1D, 

94 r_G: Array1D, 

95 ekin: float, 

96 o_G: Array1D) -> None: 

97 x = 1 / ekin / 3 * G2_G 

98 a = 27.0 + x * (18.0 + x * (12.0 + x * 8.0)) 

99 xx = x * x 

100 o_G[:] = -4.0 / 3 / ekin * a / (a + 16.0 * xx * xx) * r_G 

101 

102 

103def pw_insert(coef_G: Array1D, 

104 Q_G: Array1D, 

105 x: float, 

106 array_Q: Array1D) -> None: 

107 array_Q[:] = 0.0 

108 array_Q.ravel()[Q_G] = x * coef_G 

109 

110 

111def pw_insert_gpu(psit_nG, 

112 Q_G, 

113 scale, 

114 psit_bQ, 

115 nx, ny, nz): 

116 assert scale == 1.0 

117 psit_bQ[..., Q_G] = psit_nG 

118 if nx * ny * nz != psit_bQ.shape[-1]: 

119 n, m = nx // 2 - 1, ny // 2 - 1 

120 pw_amend_insert_realwf_gpu(psit_bQ.reshape((-1, nx, ny, nz // 2 + 1)), 

121 n, m) 

122 

123 

124def pwlfc_expand(f_Gs, Gk_Gv, pos_av, eikR_a, 

125 Y_GL, l_s, a_J, s_J, 

126 cc, f_GI, xp=np): 

127 emiGR_Ga = Gk_Gv @ pos_av.T 

128 emiGR_Ga = \ 

129 (xp.cos(emiGR_Ga) - 1j * xp.sin(emiGR_Ga)) * eikR_a 

130 real = np.issubdtype(f_GI.dtype, np.floating) 

131 I1 = 0 

132 for J, (a, s) in enumerate(zip(a_J, s_J)): 

133 l = l_s[s] 

134 I2 = I1 + 2 * l + 1 

135 f_Gi = (f_Gs[:, s] * 

136 emiGR_Ga[:, a] * 

137 Y_GL[:, l**2:(l + 1)**2].T * 

138 (-1.0j)**l).T 

139 if cc: 

140 np.conjugate(f_Gi, f_Gi) 

141 if real: 

142 f_GI[::2, I1:I2] = f_Gi.real 

143 f_GI[1::2, I1:I2] = f_Gi.imag 

144 else: 

145 f_GI[:, I1:I2] = f_Gi 

146 I1 = I2 

147 

148 

149def pwlfc_expand_gpu(f_Gs, Gk_Gv, pos_av, eikR_a, 

150 Y_GL, l_s, a_J, s_J, 

151 cc, f_GI, I_J): 

152 pwlfc_expand(f_Gs, Gk_Gv, pos_av, eikR_a, 

153 Y_GL, l_s, a_J, s_J, 

154 cc, f_GI, xp=cp) 

155 

156 

157def dH_aii_times_P_ani_gpu(dH_aii, ni_a, 

158 P_nI, out_nI): 

159 I1 = 0 

160 J1 = 0 

161 for ni in ni_a.get(): 

162 I2 = I1 + ni 

163 J2 = J1 + ni**2 

164 dH_ii = dH_aii[J1:J2].reshape((ni, ni)) 

165 out_nI[:, I1:I2] = P_nI[:, I1:I2] @ dH_ii 

166 I1 = I2 

167 J1 = J2 

168 

169 

170def pw_amend_insert_realwf_gpu(array_nQ, n, m): 

171 for array_Q in array_nQ: 

172 t = array_Q[:, :, 0] 

173 t[0, -m:] = t[0, m:0:-1].conj() 

174 t[n:0:-1, -m:] = t[-n:, m:0:-1].conj() 

175 t[-n:, -m:] = t[n:0:-1, m:0:-1].conj() 

176 t[-n:, 0] = t[n:0:-1, 0].conj() 

177 

178 

179def calculate_residuals_gpu(residual_nG, eps_n, wfs_nG): 

180 for residual_G, eps, wfs_G in zip(residual_nG, eps_n, wfs_nG): 

181 residual_G -= eps * wfs_G 

182 

183 

184def add_to_density_gpu(weight_n, psit_nR, nt_R): 

185 for weight, psit_R in zip(weight_n, psit_nR): 

186 nt_R += float(weight) * cp.abs(psit_R)**2 

187 

188 

189def symmetrize_ft(a_R, b_R, r_cc, t_c, offset_c): 

190 if (r_cc == np.eye(3, dtype=int)).all() and not t_c.any(): 

191 b_R[:] = a_R 

192 return 

193 raise NotImplementedError 

194 

195 

196def evaluate_lda_gpu(nt_sr, vxct_sr, e_r) -> None: 

197 if cupy_is_fake: 

198 from gpaw.xc.kernel import XCKernel 

199 XCKernel('LDA').calculate(e_r._data, nt_sr._data, vxct_sr._data) 

200 else: 

201 from _gpaw import evaluate_lda_gpu as evalf # type: ignore 

202 evalf(nt_sr, vxct_sr, e_r) 

203 

204 

205def evaluate_pbe_gpu(nt_sr, vxct_sr, e_r, sigma_xr, dedsigma_xr) -> None: 

206 if cupy_is_fake: 

207 from gpaw.xc.kernel import XCKernel 

208 XCKernel('PBE').calculate(e_r._data, nt_sr._data, vxct_sr._data, 

209 sigma_xr._data, dedsigma_xr._data) 

210 else: 

211 from _gpaw import evaluate_pbe_gpu as evalf # type: ignore 

212 evalf(nt_sr, vxct_sr, e_r, sigma_xr, dedsigma_xr) 

213 

214 

215def pw_norm_gpu(result_x, C_xG): 

216 if cupy_is_fake: 

217 result_x._data[:] = np.sum(np.abs(C_xG._data)**2, axis=1) 

218 else: 

219 result_x[:] = cp.sum(cp.abs(C_xG)**2, axis=1) 

220 

221 

222def pw_norm_kinetic_gpu(result_x, a_xG, kin_G): 

223 if cupy_is_fake: 

224 result_x._data[:] = np.sum( 

225 np.abs(a_xG._data)**2 * kin_G._data[None, :], 

226 axis=1) 

227 else: 

228 result_x[:] = cp.sum(cp.abs(a_xG)**2 * kin_G[None, :], axis=1)