Coverage for gpaw/analyse/expandyl.py: 78%

1from math import pi, acos, sqrt 

2 

3import numpy as np 

4from ase.atoms import string2vector 

5from ase.units import Bohr, Hartree 

6from ase.utils import IOContext 

7 

8from gpaw.spherical_harmonics import Y 

9from gpaw.utilities.tools import coordinates 

10from gpaw.mpi import serial_comm 

11 

12 

13class AngularIntegral: 

14 

15 """Perform an angular integral on the grid. 

16 

17 center: 

18 the center (Ang) 

19 gd: 

20 grid_descriptor of the grids to expand 

21 Rmax: 

22 maximal radius of the expansion (Ang) 

23 dR: 

24 grid spacing in the radius (Ang) 

25 """ 

26 

27 def __init__(self, center, gd, Rmax=None, dR=None): 

28 assert gd.orthogonal 

29 center = Vector3d(center) / Bohr 

30 

31 self.center = center 

32 self.gd = gd 

33 

34 # set Rmax to the maximal radius possible 

35 # i.e. the corner distance 

36 if not Rmax: 

37 Rmax = 0 

38 extreme = gd.h_cv.diagonal() * gd.N_c 

39 for corner in ([0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], 

40 [0, 0, 1], [1, 0, 1], [0, 1, 1], [1, 1, 1]): 

41 Rmax = max(Rmax, 

42 self.center.distance(np.array(corner) * extreme)) 

43 else: 

44 Rmax /= Bohr 

45 self.Rmax = Rmax 

46 

47 if not dR: 

48 dR = min(gd.h_cv.diagonal()) 

49 else: 

50 dR /= Bohr 

51 self.dR = dR 

52 

53 self.initialize() 

54 

55 def initialize(self): 

56 """Initialize grids.""" 

57 

58 Rmax = self.Rmax 

59 dR = self.dR 

60 gd = self.gd 

61 

62 # initialize the ylm and Radial grids 

63 

64 # self.V_R will contain the volume of the R shell 

65 # self.R_g will contain the radial indicees corresponding to 

66 # each grid point 

67 # self.ball_g will contain the mask of the ball of radius Rmax 

68 # self.y_Lg will contain the YL values corresponding to 

69 # each grid point 

70 V_R = np.empty((int(Rmax / dR + 1),)) 

71 R_R = np.empty((int(Rmax / dR + 1),)) 

72 

73 r_cg, r2_g = coordinates(gd, self.center, tiny=1.e-78) 

74 r_g = np.sqrt(r2_g) 

75 rhat_cg = r_cg / r_g 

76 

77 ball_g = np.where(r_g < Rmax, 1, 0) 

78 self.R_g = np.where(r_g < Rmax, r_g / dR, -1).astype(int) 

79 

80 if hasattr(self, 'L_l'): # ExpandYl 

81 npY = np.vectorize(Y, (float,), 'spherical harmonic') 

82 nL = len(self.L_l) 

83 y_Lg = [] 

84 for L in range(nL): 

85 y_Lg.append(npY(L, rhat_cg[0], rhat_cg[1], rhat_cg[2])) 

86 self.y_Lg = y_Lg 

87 

88 for i, v in enumerate(V_R): 

89 R_g = np.where(self.R_g == i, 1., 0.) 

90 V_R[i] = gd.integrate(R_g) 

91 R_R[i] = gd.integrate(R_g * r_g) 

92 

93 self.ball_g = ball_g 

94 self.V_R = V_R 

95 self.nominalR_R = self.dR * (np.arange(len(self.V_R)) + .5) 

96 V_R = np.where(V_R > 0, V_R, -1) 

97 self.R_R = np.where(V_R > 0, R_R / V_R, self.nominalR_R) 

98 

99 def integrate(self, f_g): 

100 """Integrate a function on the grid over the angles. 

101 

102 Contains the weight 4*pi*R^2 with R in Angstrom.""" 

103 int_R = [] 

104 for i, dV in enumerate(self.V_R): 

105 # get the R shell 

106 R_g = np.where(self.R_g == i, 1, 0) 

107 int_R.append(self.gd.integrate(f_g * R_g) / self.dR) 

108 return np.array(int_R) * Bohr ** 2 

109 

110 def average(self, f_g): 

111 """Give the angular average of a function on the grid.""" 

112 V_R = np.where(self.V_R > 0, self.V_R, 1.e32) 

113 return self.integrate(f_g) * self.dR / V_R / Bohr ** 2 

114 

115 def radii(self, model='nominal'): 

116 """Return the radii of the radial shells in Angstrom""" 

117 if model == 'nominal': 

118 return self.nominalR_R * Bohr 

119 elif model == 'mean': 

120 return self.R_R * Bohr 

121 else: 

122 raise NotImplementedError 

123 

124 

125class ExpandYl(AngularIntegral): 

126 

127 """Expand the smooth wave functions in spherical harmonics 

128 relative to a given center. 

129 

130 center: 

131 the center for the expansion (Ang) 

132 gd: 

133 grid_descriptor of the grids to expand 

134 lmax: 

135 maximal angular momentum in the expansion (lmax<7) 

136 Rmax: 

137 maximal radius of the expansion (Ang) 

138 dR: 

139 grid spacing in the radius (Ang) 

140 """ 

141 

142 def __init__(self, center, gd, lmax=6, Rmax=None, dR=None): 

143 

144 self.lmax = lmax 

145 self.L_l = [] 

146 for l in range(lmax + 1): 

147 for m in range(2 * l + 1): 

148 self.L_l.append(l) 

149 

150 super().__init__(center, gd, Rmax, dR) 

151 

152 def expand(self, psit_g): 

153 """Expand a wave function""" 

154 

155 gamma_l = np.zeros(self.lmax + 1) 

156 nL = len(self.L_l) 

157 L_l = self.L_l 

158 

159 def abs2(z): 

160 return z.real**2 + z.imag**2 

161 

162 for i, dV in enumerate(self.V_R): 

163 # get the R shell and it's Volume 

164 R_g = np.where(self.R_g == i, 1, 0) 

165 if dV > 0: 

166 for L in range(nL): 

167 psit_LR = self.gd.integrate(psit_g * R_g * self.y_Lg[L]) 

168 gamma_l[L_l[L]] += 4 * pi / dV * abs2(psit_LR) 

169 # weight of the wave function inside the ball 

170 weight = self.gd.integrate((psit_g * psit_g.conj()).real * self.ball_g) 

171 

172 return gamma_l, weight 

173 

174 def to_file(self, calculator, 

175 filename='expandyl.dat', 

176 spins=None, 

177 kpoints=None, 

178 bands=None 

179 ): 

180 """Expand a range of wave functions and write the result 

181 to a file""" 

182 with IOContext() as io: 

183 fd = io.openfile(filename, comm=serial_comm) 

184 

185 if not spins: 

186 srange = range(calculator.wfs.nspins) 

187 else: 

188 srange = spins 

189 if not kpoints: 

190 krange = range(len(calculator.wfs.kd.ibzk_kc)) 

191 else: 

192 krange = kpoints 

193 if not bands: 

194 nrange = range(calculator.wfs.bd.nbands) 

195 else: 

196 nrange = bands 

197 

198 print('# Yl expansion', 'of smooth wave functions', file=fd) 

199 lu = 'Angstrom' 

200 print('# center =', self.center * Bohr, lu, file=fd) 

201 print('# Rmax =', self.Rmax * Bohr, lu, file=fd) 

202 print('# dR =', self.dR * Bohr, lu, file=fd) 

203 print('# lmax =', self.lmax, file=fd) 

204 print('# s k n', end=' ', file=fd) 

205 print('kpt-wght e[eV] occ', end=' ', file=fd) 

206 print(' norm sum weight', end=' ', file=fd) 

207 spdfghi = 's p d f g h i'.split() 

208 for l in range(self.lmax + 1): 

209 print(' %' + spdfghi[l], end=' ', file=fd) 

210 print(file=fd) 

211 

212 for s in srange: 

213 for k in krange: 

214 u = k * calculator.wfs.nspins + s 

215 for n in nrange: 

216 kpt = calculator.wfs.kpt_u[u] 

217 psit_G = kpt.psit_nG[n] 

218 norm = self.gd.integrate((psit_G * psit_G.conj()).real) 

219 

220 gl, weight = self.expand(psit_G) 

221 gsum = np.sum(gl) 

222 gl = 100 * gl / gsum 

223 

224 print(f'{s:2d} {k:5d} {n:5d}', end=' ', file=fd) 

225 print(f'{kpt.weight:6.4f} ' 

226 f'{(kpt.eps_n[n] * Hartree):10.4f} ' 

227 f'{kpt.f_n[n]:8.4f}', end=' ', file=fd) 

228 print(f'{norm:8.4f} {gsum:8.4f} {weight:8.4f}', 

229 end=' ', file=fd) 

230 

231 for g in gl: 

232 print(f'{g:8.2f}', end=' ', file=fd) 

233 print(file=fd) 

234 fd.flush() 

235 

236 

237class Vector3d(list): 

238 def __init__(self, vector=None): 

239 if vector is None: 

240 vector = [0, 0, 0] 

241 vector = string2vector(vector) 

242 super().__init__() 

243 for c in range(3): 

244 self.append(float(vector[c])) 

245 self.l = False 

246 

247 def __add__(self, other): 

248 result = self.copy() 

249 for c in range(3): 

250 result[c] += other[c] 

251 return result 

252 

253 def __truediv__(self, other): 

254 return Vector3d(np.array(self) / other) 

255 

256 __div__ = __truediv__ 

257 

258 def __mul__(self, x): 

259 if isinstance(x, type(self)): 

260 return np.dot(self, x) 

261 else: 

262 return Vector3d(x * np.array(self)) 

263 

264 def __rmul__(self, x): 

265 return self.__mul__(x) 

266 

267 def __lmul__(self, x): 

268 return self.__mul__(x) 

269 

270 def __neg__(self): 

271 return -1 * self 

272 

273 def __str__(self): 

274 return "(%g,%g,%g)" % tuple(self) 

275 

276 def __sub__(self, other): 

277 result = self.copy() 

278 for c in range(3): 

279 result[c] -= other[c] 

280 return result 

281 

282 def angle(self, other): 

283 """Return the angle between the directions of yourself and the 

284 other vector in radians.""" 

285 other = Vector3d(other) 

286 ll = self.length() * other.length() 

287 if not ll > 0: 

288 return None 

289 return acos((self * other) / ll) 

290 

291 def copy(self): 

292 return Vector3d(self) 

293 

294 def distance(self, vector): 

295 if not isinstance(vector, type(self)): 

296 vector = Vector3d(vector) 

297 return (self - vector).length() 

298 

299 def length(self, value=None): 

300 if value: 

301 fac = value / self.length() 

302 for c in range(3): 

303 self[c] *= fac 

304 self.l = False 

305 if not self.l: 

306 self.l = sqrt(self.norm()) 

307 return self.l 

308 

309 def norm(self): 

310 # return np.sum( self*self ) 

311 return self * self # XXX drop this class and use numpy arrays ... 

312 

313 def x(self): 

314 return self[0] 

315 

316 def y(self): 

317 return self[1] 

318 

319 def z(self): 

320 return self[2]