Coverage for gpaw/utilities/gauss.py: 94%

1import numpy as np 

2from numpy import sqrt, pi, exp, abs 

3from scipy.special import erf 

4 

5import gpaw.cgpaw as cgpaw 

6from gpaw import debug 

7from gpaw.utilities.tools import coordinates 

8from gpaw.utilities import is_contiguous 

9 

10 

11def Y_L(L, x, y, z, r2): 

12 if L == 0: 

13 return 0.28209479177387814 

14 elif L == 1: 

15 return 0.48860251190291992 * y 

16 elif L == 2: 

17 return 0.48860251190291992 * z 

18 elif L == 3: 

19 return 0.48860251190291992 * x 

20 elif L == 4: 

21 return 1.0925484305920792 * x * y 

22 elif L == 5: 

23 return 1.0925484305920792 * y * z 

24 elif L == 6: 

25 return 0.31539156525252005 * (3 * z * z - r2) 

26 elif L == 7: 

27 return 1.0925484305920792 * x * z 

28 elif L == 8: 

29 return 0.54627421529603959 * (x * x - y * y) 

30 

31 

32def gauss_L(a, L, x, y, z, r2, exp_ar2): 

33 if L == 0: 

34 return sqrt(a**3 * 4) / pi * exp_ar2 

35 elif L == 1: 

36 return sqrt(a**5 * 5.333333333333333) / pi * y * exp_ar2 

37 elif L == 2: 

38 return sqrt(a**5 * 5.333333333333333) / pi * z * exp_ar2 

39 elif L == 3: 

40 return sqrt(a**5 * 5.333333333333333) / pi * x * exp_ar2 

41 elif L == 4: 

42 return sqrt(a**7 * 4.2666666666666666) / pi * x * y * exp_ar2 

43 elif L == 5: 

44 return sqrt(a**7 * 4.2666666666666666) / pi * y * z * exp_ar2 

45 elif L == 6: 

46 return sqrt( 

47 a**7 * 0.35555555555555557) / pi * (3 * z * z - r2) * exp_ar2 

48 elif L == 7: 

49 return sqrt(a**7 * 4.2666666666666666) / pi * x * z * exp_ar2 

50 elif L == 8: 

51 return sqrt(a**7 * 1.0666666666666667) / pi * (x * x - y * y) * exp_ar2 

52 

53 

54def gausspot_L(a, L, x, y, z, r, r2, erf_sar, exp_ar2): 

55 if L == 0: 

56 return 2.0 * 1.7724538509055159 * erf_sar / r 

57 elif L == 1: 

58 return 1.1547005383792515 * (1.7724538509055159 * erf_sar - 

59 2 * sqrt(a) * r * exp_ar2) / r / r2 * y 

60 elif L == 2: 

61 return 1.1547005383792515 * (1.7724538509055159 * erf_sar - 

62 2 * sqrt(a) * r * exp_ar2) / r / r2 * z 

63 elif L == 3: 

64 return 1.1547005383792515 * (1.7724538509055159 * erf_sar - 

65 2 * sqrt(a) * r * exp_ar2) / r / r2 * x 

66 elif L == 4: 

67 return 0.5163977794943222 * ( 

68 5.3173615527165481 * erf_sar - 

69 (6 + 4 * 

70 (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * x * y 

71 elif L == 5: 

72 return 0.5163977794943222 * ( 

73 5.3173615527165481 * erf_sar - 

74 (6 + 4 * 

75 (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * y * z 

76 elif L == 6: 

77 return 0.14907119849998599 * (5.3173615527165481 * erf_sar - 

78 (6 + 4 * 

79 (sqrt(a) * r)**2) * sqrt(a) * r * 

80 exp_ar2) / r / r2**2 * (3 * z * z - r2) 

81 elif L == 7: 

82 return 0.5163977794943222 * ( 

83 5.3173615527165481 * erf_sar - 

84 (6 + 4 * 

85 (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * x * z 

86 elif L == 8: 

87 return 0.2581988897471611 * (5.3173615527165481 * erf_sar - 

88 (6 + 4 * (sqrt(a) * r)**2) * sqrt(a) * r * 

89 exp_ar2) / r / r2**2 * (x * x - y * y) 

90 

91 

92# end of computer generated code 

93 

94 

95class Gaussian: 

96 r"""Class offering several utilities related to the generalized gaussians. 

97 

98 Generalized gaussians are defined by:: 

99 

100 _____ 2 

101 / 1 l! l+3/2 -a r l m 

102 g (x,y,z) = / ----- --------- (4 a) e r Y (x,y,z), 

103 L \/ 4 pi (2l + 1)! l 

104 

105 where a is the inverse width of the gaussian, and Y_l^m is a real 

106 spherical harmonic. 

107 The gaussians are centered in the middle of input grid-descriptor.""" 

108 def __init__(self, gd, a=19., center=None): 

109 self.gd = gd 

110 self.xyz, self.r2 = coordinates(gd, center) 

111 self.r = np.sqrt(self.r2) 

112 self.set_width(a, center) 

113 self.exp_ar2 = exp(-self.a * self.r2) 

114 self.erf_sar = erf(sqrt(self.a) * self.r) 

115 

116 def set_width(self, a, center): 

117 """Set exponent of exp-function to -a on the boundary.""" 

118 if center is None: 

119 self.a = 4 * a * (self.gd.icell_cv**2).sum(1).max() 

120 else: 

121 cell_center = self.gd.cell_cv.sum(1) / 2 

122 r_min = (cell_center - np.abs(center - cell_center)).min() 

123 self.a = a / r_min**2 

124 

125 def get_gauss(self, L): 

126 a = self.a 

127 x, y, z = tuple(self.xyz) 

128 r2 = self.r2 

129 exp_ar2 = self.exp_ar2 

130 return gauss_L(a, L, x, y, z, r2, exp_ar2) 

131 

132 def get_gauss_pot(self, L): 

133 a = self.a 

134 x, y, z = tuple(self.xyz) 

135 r2 = self.r2 

136 r = self.r 

137 erf_sar = self.erf_sar 

138 exp_ar2 = self.exp_ar2 

139 return gausspot_L(a, L, x, y, z, r, r2, erf_sar, exp_ar2) 

140 

141 def get_moment(self, n, L): 

142 r2 = self.r2 

143 x, y, z = tuple(self.xyz) 

144 return self.gd.integrate(n * Y_L(L, x, y, z, r2)) 

145 

146 def remove_moment(self, n, L, q=None): 

147 # Determine multipole moment 

148 if q is None: 

149 q = self.get_moment(n, L) 

150 

151 # Don't do anything if moment is less than the tolerance 

152 if abs(q) < 1e-7: 

153 return 0. 

154 

155 # Remove moment from input density 

156 n -= q * self.get_gauss(L) 

157 

158 # Return correction 

159 return q * self.get_gauss_pot(L) 

160 

161 

162def gaussian_wave(r_vG, r0_v, sigma, k_v=None, A=None, dtype=float, 

163 out_G=None): 

164 r"""Generates function values for atom-centered Gaussian waves. 

165 

166 :: 

167 

168 _ _ 

169 _ / -|r-r0|^2 \ _ _ 

170 f(r) = A * exp( ----------- ) * exp( i k.r ) 

171 \ 2 sigma^2 / 

172 

173 If the parameter A is not specified, the Gaussian wave is normalized:: 

174 

175 oo 

176 / ____ \ -3/2 / _ 2 2 

177 A = ( / ' ) => 4 pi | dr |f(r)| r = 1 

178 \ \/ pi sigma / / 

179 0 

180 

181 Parameters: 

182 

183 r_vG: ndarray 

184 Set of coordinates defining the grid positions. 

185 r0_v: ndarray 

186 Set of coordinates defining the center of the Gaussian envelope. 

187 sigma: float 

188 Specifies the spatial width of the Gaussian envelope. 

189 k_v: ndarray or None 

190 Set of reciprocal lattice coordinates defining the wave vector. 

191 An argument of None is interpreted as the gamma point i.e. k_v=0. 

192 A: float, complex or None 

193 Specifies the amplitude of the Gaussian wave. Normalizes if None. 

194 dtype: type, defaults to float 

195 Specifies the output data type. Only returns the real-part if float. 

196 out_G: ndarray or None 

197 Optional pre-allocated buffer to fill in values. Allocates if None. 

198 

199 """ 

200 if k_v is None: 

201 k_v = np.zeros(r0_v.shape) 

202 

203 if A is None: 

204 # 4*pi*int(exp(-r^2/(2*sigma^2))^2 * r^2, r=0...infinity) 

205 # = sigma^3*pi^(3/2) = 1/A^2 -> A = (sqrt(Pi)*sigma)^(-3/2) 

206 A = 1 / (sigma * np.pi**0.5)**1.5 

207 

208 if debug: 

209 assert is_contiguous(r_vG, float) 

210 assert is_contiguous(r0_v, float) 

211 assert is_contiguous(k_v, float) 

212 assert r_vG.ndim >= 2 and r_vG.shape[0] > 0 

213 assert r0_v.ndim == 1 and r0_v.shape[0] > 0 

214 assert k_v.ndim == 1 and k_v.shape[0] > 0 

215 assert (r_vG.shape[0], ) == r0_v.shape == k_v.shape 

216 assert sigma > 0 

217 

218 if out_G is None: 

219 out_G = np.empty(r_vG.shape[1:], dtype=dtype) 

220 elif debug: 

221 assert is_contiguous(out_G) 

222 assert out_G.shape == r_vG.shape[1:] 

223 

224 # slice_v2vG = [slice(None)] + [np.newaxis]*3 

225 # gw = lambda r_vG, r0_v, sigma, k_v, A=1/(sigma*np.pi**0.5)**1.5: \ 

226 # * np.exp(-np.sum((r_vG-r0_v[slice_v2vG])**2, axis=0)/(2*sigma**2)) \ 

227 # * np.exp(1j*np.sum(np.r_vG*k_v[slice_v2vG], axis=0)) * A 

228 cgpaw.utilities_gaussian_wave(A, r_vG, r0_v, sigma, k_v, out_G) 

229 return out_G