Coverage for gpaw/test/ri/test_spherical_hse_kernel.py: 78%

1""" 

2 

3Test the implementation of spherical harmonic expansion of screened 

4Coulomb kernel. Based on 

5 

6 János G Ángyán et al 2006 J. Phys. A: Math. Gen. 39 8613 

7 

8 

9""" 

10 

11 

12from gpaw.xc.ri.spherical_hse_kernel import Phi as phi 

13from scipy.special import erfc 

14import numpy as np 

15from gpaw.sphere.lebedev import weight_n, Y_nL, R_nv 

16from gpaw.spherical_harmonics import Y 

17 

18 

19# [23, 31, 16] are indices to a triangle in the 50-point Lebedev grid. 

20# This is used to align the two angular grids nicely to avoid divergences. 

21Rdir_v = np.mean(R_nv[[23, 31, 16], :], axis=0) 

22# Obtain the tangent and bitangent vectors for a full basis 

23Q_vv, _ = np.linalg.qr(np.array([Rdir_v]).T, 'complete') 

24 

25# Get the rotated second angular integration grid 

26R2_nv = R_nv @ Q_vv 

27 

28 

29def phiold(n, mu, R, r): 

30 """ 

31 

32 Explicit implementation of spherical harmonic expansion up to l=2 

33 as given by the article Eqs. A1, A2 and A3. These are compared to 

34 the official implementation in `xc/ri/spherical_hse_kernel.py`. 

35 

36 """ 

37 Rg = np.maximum.reduce([R, r]) 

38 Rl = np.minimum.reduce([R, r]) 

39 Xi = mu * Rg 

40 xi = mu * Rl 

41 

42 if n == 0: 

43 prefactor = -1 / (2 * np.pi**0.5 * xi * Xi) 

44 A = np.exp(-(xi + Xi)**2) - np.exp(-(xi - Xi)**2) 

45 B = -np.pi**0.5 * ((xi - Xi) * erfc(Xi - xi) + 

46 (Xi + xi) * erfc(Xi + xi)) 

47 return mu * prefactor * (A + B) 

48 if n == 1: 

49 prefactor = -1 / (2 * np.pi**0.5 * xi**2 * Xi**2) 

50 A = 1 / 2 * ((np.exp(-(xi + Xi)**2) - np.exp(-(xi - Xi)**2)) * (2 * xi**2 + 2 * xi * Xi - (1 - 2 * Xi**2)) - 4 * xi * Xi * np.exp(-(xi + Xi)**2)) - np.pi**0.5 * ((xi**3 - Xi**3) * erfc(Xi - xi) + (xi**3 + Xi**3) * erfc(xi + Xi)) # noqa: E501 

51 return mu * prefactor * A 

52 if n == 2: 

53 prefactor = -1 / (2 * np.pi**0.5 * xi**3 * Xi**3) 

54 A = 1 / 4 * ((np.exp(-(xi + Xi)**2) - np.exp(-(xi - Xi)**2)) * (4 * (xi**4 + xi**3 * Xi + Xi**4) - 2 * xi**2 * (1 - 2 * Xi**2) + (1 - 2 * xi * Xi) * (3 - 2 * Xi**2)) - 4 * np.exp(-(xi + Xi)**2) * xi * Xi * (2 * xi**2 - (3 - 2 * Xi**2))) - np.pi**0.5 * ((xi**5 - Xi**5) * erfc(Xi - xi) + (xi**5 + Xi**5) * erfc(xi + Xi)) # noqa: E501 

55 return mu * prefactor * A 

56 raise NotImplementedError 

57 

58 

59def phi_lebedev(n, mu, R_x, r_x): 

60 # Target spherical harmonic, primary grid 

61 Y1_n = Y_nL[:, n**2] 

62 # Target spherical harmonic, secondary grid 

63 Y2_n = Y(n**2, *R2_nv.T) 

64 

65 V_x = np.zeros_like(R_x) 

66 for x, (R, r) in enumerate(zip(R_x, r_x)): 

67 C1_nv = R * R_nv 

68 C2_nv = r * R2_nv 

69 

70 D_nn = np.sum((C1_nv[:, None, :] - C2_nv[None, :, :])**2, axis=2)**0.5 

71 V_nn = erfc(D_nn * mu) / D_nn 

72 

73 V_x[x] = np.einsum('n,m,nm,n,m', weight_n, weight_n, V_nn, Y1_n, Y2_n) 

74 

75 return V_x * (4 * np.pi) * (2 * n + 1) 

76 

77 

78def test_old_vs_new_spherical_kernel(): 

79 """Test explicitely hard coded implementation against generic 

80 implementation. 

81 """ 

82 rng = np.random.RandomState(42) 

83 for n in range(3): 

84 R = rng.random(100) * 10 

85 r = rng.random(100) * 10 

86 params = (n, 0.11, R, r) 

87 new, old = phi(*params), phiold(*params) 

88 print(old, new, abs(old - new).max()) 

89 

90 assert np.allclose(new, old, atol=1.5e-6) 

91 

92 

93def test_wrt_lebedev_integrated_kernel(plot=False): 

94 """ 

95 Test a double angular numerically integrated kernel against 

96 the generic implementation. 

97 """ 

98 if plot: 

99 import matplotlib.pyplot as plt 

100 s = 25 

101 for n in range(5): 

102 for RR in [1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2]: 

103 R = RR * np.ones((125,)) 

104 r = np.logspace(-5, 3, 125 + 1)[1:] 

105 params = (n, 0.11, R.ravel(), r.ravel()) 

106 new, old = phi(*params), phi_lebedev(*params) 

107 if plot: 

108 plt.loglog(r, np.abs(old), '-r') 

109 plt.loglog(r, np.abs(new), '--b') 

110 plt.loglog(r, np.abs(old - new), '-k') 

111 plt.ylim([1e-7, 1e7]) 

112 err = np.abs(new - old) 

113 

114 # The angular integration error (due to only 50 point grid) is 

115 # too large on small separations. Therefore, they should not 

116 # be compared directly. 

117 err = np.where(R - r < 0.3, 0 * err, err) 

118 assert np.allclose(err, 0, atol=1e-2, rtol=1e-2) 

119 if plot: 

120 plt.show() 

121 

122 if plot: 

123 import matplotlib.pyplot as plt 

124 for n in range(5): 

125 t = np.logspace(-5, 5, s) 

126 R, r = np.meshgrid(t, t) 

127 params = (n, 0.11, R.ravel(), r.ravel()) 

128 new, old = phi(*params), phi_lebedev(*params) 

129 plt.contourf(np.log10(r), np.log10(R), 

130 np.reshape(np.log10(np.abs(old - new) + 1e-7), 

131 (s, s))) 

132 plt.colorbar() 

133 plt.show() 

134 

135 

136if __name__ == '__main__': 

137 test_wrt_lebedev_integrated_kernel(plot=True)