Coverage for gpaw/wavefunctions/fd.py: 85%

1import numpy as np 

2from ase.units import Bohr 

3 

4from gpaw.fd_operators import Laplace, Gradient 

5from gpaw.kpoint import KPoint 

6from gpaw.kpt_descriptor import KPointDescriptor 

7from gpaw.lfc import LocalizedFunctionsCollection as LFC 

8from gpaw.mpi import serial_comm 

9from gpaw.preconditioner import Preconditioner 

10from gpaw.projections import Projections 

11from gpaw.transformers import Transformer 

12from gpaw.utilities.blas import axpy 

13from gpaw.wavefunctions.arrays import UniformGridWaveFunctions 

14from gpaw.wavefunctions.fdpw import FDPWWaveFunctions 

15from gpaw.wavefunctions.mode import Mode 

16import gpaw.cgpaw as cgpaw 

17 

18 

19class FD(Mode): 

20 name = 'fd' 

21 

22 def __init__(self, nn=3, interpolation=3, force_complex_dtype=False): 

23 self.nn = nn 

24 self.interpolation = interpolation 

25 Mode.__init__(self, force_complex_dtype) 

26 

27 def __call__(self, *args, **kwargs): 

28 return FDWaveFunctions(self.nn, *args, **kwargs) 

29 

30 def todict(self): 

31 dct = Mode.todict(self) 

32 dct['nn'] = self.nn 

33 dct['interpolation'] = self.interpolation 

34 return dct 

35 

36 

37class FDWaveFunctions(FDPWWaveFunctions): 

38 mode = 'fd' 

39 

40 def __init__(self, stencil, parallel, initksl, 

41 gd, nvalence, setups, bd, 

42 dtype, world, kd, kptband_comm, timer, reuse_wfs_method=None, 

43 collinear=True): 

44 FDPWWaveFunctions.__init__(self, parallel, initksl, 

45 reuse_wfs_method=reuse_wfs_method, 

46 collinear=collinear, 

47 gd=gd, nvalence=nvalence, setups=setups, 

48 bd=bd, dtype=dtype, world=world, kd=kd, 

49 kptband_comm=kptband_comm, timer=timer) 

50 

51 # Kinetic energy operator: 

52 self.kin = Laplace(self.gd, -0.5, stencil, self.dtype) 

53 

54 self.taugrad_v = None # initialized by MGGA functional 

55 self.read_from_file_init_wfs_dm = False 

56 

57 def empty(self, n=(), global_array=False, realspace=False, q=-1): 

58 return self.gd.empty(n, self.dtype, global_array) 

59 

60 def integrate(self, a_xg, b_yg=None, global_integral=True): 

61 return self.gd.integrate(a_xg, b_yg, global_integral) 

62 

63 def bytes_per_wave_function(self): 

64 return self.gd.bytecount(self.dtype) 

65 

66 def set_setups(self, setups): 

67 self.pt = LFC(self.gd, [setup.pt_j for setup in setups], 

68 self.kd, dtype=self.dtype, forces=True) 

69 FDPWWaveFunctions.set_setups(self, setups) 

70 

71 def set_positions(self, spos_ac, atom_partition=None): 

72 FDPWWaveFunctions.set_positions(self, spos_ac, atom_partition) 

73 

74 def __str__(self): 

75 s = 'Wave functions: Uniform real-space grid\n' 

76 s += ' Kinetic energy operator: %s\n' % self.kin.description 

77 return s + FDPWWaveFunctions.__str__(self) 

78 

79 def make_preconditioner(self, block=1): 

80 return Preconditioner(self.gd, self.kin, self.dtype, block) 

81 

82 def apply_pseudo_hamiltonian(self, kpt, ham, psit_xG, Htpsit_xG): 

83 self.timer.start('Apply hamiltonian') 

84 self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) 

85 ham.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) 

86 ham.xc.apply_orbital_dependent_hamiltonian( 

87 kpt, psit_xG, Htpsit_xG, ham.dH_asp) 

88 self.timer.stop('Apply hamiltonian') 

89 

90 def get_pseudo_partial_waves(self): 

91 phit_aj = [setup.get_partial_waves_for_atomic_orbitals() 

92 for setup in self.setups] 

93 return LFC(self.gd, phit_aj, kd=self.kd, cut=True, dtype=self.dtype) 

94 

95 def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n): 

96 # Used in calculation of response part of GLLB-potential 

97 nt_G = nt_sG[kpt.s] 

98 for f, psit_G in zip(f_n, kpt.psit_nG): 

99 # Same as nt_G += f * abs(psit_G)**2, but much faster: 

100 cgpaw.add_to_density(f, psit_G, nt_G) 

101 

102 # Hack used in delta-scf calculations: 

103 if hasattr(kpt, 'c_on'): 

104 assert self.bd.comm.size == 1 

105 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands), 

106 dtype=complex) 

107 for ne, c_n in zip(kpt.ne_o, kpt.c_on): 

108 d_nn += ne * np.outer(c_n.conj(), c_n) 

109 for d_n, psi0_G in zip(d_nn, kpt.psit_nG): 

110 for d, psi_G in zip(d_n, kpt.psit_nG): 

111 if abs(d) > 1.e-12: 

112 nt_G += (psi0_G.conj() * d * psi_G).real 

113 

114 def calculate_kinetic_energy_density(self): 

115 if self.taugrad_v is None: 

116 self.taugrad_v = [ 

117 Gradient(self.gd, v, n=3, dtype=self.dtype).apply 

118 for v in range(3)] 

119 

120 assert not hasattr(self.kpt_u[0], 'c_on') 

121 if not isinstance(self.kpt_u[0].psit_nG, np.ndarray): 

122 return None 

123 

124 taut_sG = self.gd.zeros(self.nspins) 

125 dpsit_G = self.gd.empty(dtype=self.dtype) 

126 for kpt in self.kpt_u: 

127 for f, psit_G in zip(kpt.f_n, kpt.psit_nG): 

128 for v in range(3): 

129 self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd) 

130 axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s]) 

131 

132 self.kptband_comm.sum(taut_sG) 

133 for taut_G in taut_sG: 

134 self.kd.symmetry.symmetrize(taut_G, self.gd) 

135 return taut_sG 

136 

137 def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, 

138 Htpsit_xG, dH_asp, 

139 dedtaut_G): 

140 a_G = self.gd.empty(dtype=psit_xG.dtype) 

141 for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG): 

142 for v in range(3): 

143 self.taugrad_v[v](psit_G, a_G, kpt.phase_cd) 

144 self.taugrad_v[v](dedtaut_G * a_G, a_G, kpt.phase_cd) 

145 axpy(-0.5, a_G, Htpsit_G) 

146 

147 def ibz2bz(self, atoms): 

148 """Transform wave functions in IBZ to the full BZ.""" 

149 

150 assert self.kd.comm.size == 1 

151 

152 # New k-point descriptor for full BZ: 

153 kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins) 

154 kd.set_communicator(serial_comm) 

155 

156 self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups], 

157 kd, dtype=self.dtype) 

158 self.pt.set_positions(atoms.get_scaled_positions()) 

159 

160 self.initialize_wave_functions_from_restart_file() 

161 

162 weight = 2.0 / kd.nspins / kd.nbzkpts 

163 

164 # Build new list of k-points: 

165 kpt_qs = [] 

166 kpt_u = [] 

167 for k in range(kd.nbzkpts): 

168 kpt_s = [] 

169 for s in range(self.nspins): 

170 # Index of symmetry related point in the IBZ 

171 ik = self.kd.bz2ibz_k[k] 

172 r, q = self.kd.get_rank_and_index(ik) 

173 assert r == 0 

174 kpt = self.kpt_qs[q][s] 

175 

176 phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd * 

177 kd.bzk_kc[k, :, np.newaxis]) 

178 

179 # New k-point: 

180 kpt2 = KPoint(1.0 / kd.nbzkpts, weight, s, k, k, phase_cd) 

181 kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins 

182 kpt2.eps_n = kpt.eps_n.copy() 

183 

184 # Transform wave functions using symmetry operation: 

185 Psit_nG = self.gd.collect(kpt.psit_nG) 

186 if Psit_nG is not None: 

187 Psit_nG = Psit_nG.copy() 

188 for Psit_G in Psit_nG: 

189 Psit_G[:] = self.kd.transform_wave_function(Psit_G, k) 

190 kpt2.psit = UniformGridWaveFunctions( 

191 self.bd.nbands, self.gd, self.dtype, 

192 kpt=k, dist=(self.bd.comm, self.bd.comm.size), 

193 spin=kpt.s, collinear=True) 

194 self.gd.distribute(Psit_nG, kpt2.psit_nG) 

195 # Calculate PAW projections: 

196 nproj_a = [setup.ni for setup in self.setups] 

197 kpt2.projections = Projections( 

198 self.bd.nbands, nproj_a, 

199 kpt.projections.atom_partition, 

200 self.bd.comm, 

201 collinear=True, spin=s, dtype=self.dtype) 

202 

203 kpt2.psit.matrix_elements(self.pt, out=kpt2.projections) 

204 kpt_s.append(kpt2) 

205 kpt_u.append(kpt2) 

206 kpt_qs.append(kpt_s) 

207 

208 self.kd = kd 

209 self.kpt_qs = kpt_qs 

210 self.kpt_u = kpt_u 

211 

212 def _get_wave_function_array(self, u, n, realspace=True, periodic=False): 

213 assert realspace 

214 kpt = self.kpt_u[u] 

215 psit_G = kpt.psit_nG[n] 

216 if periodic and self.dtype == complex: 

217 k_c = self.kd.ibzk_kc[kpt.k] 

218 return self.gd.plane_wave(-k_c) * psit_G 

219 return psit_G 

220 

221 def write(self, writer, write_wave_functions=False): 

222 FDPWWaveFunctions.write(self, writer) 

223 

224 if not write_wave_functions: 

225 return 

226 

227 writer.add_array( 

228 'values', 

229 (self.nspins, self.kd.nibzkpts, self.bd.nbands) + 

230 tuple(self.gd.get_size_of_global_array()), 

231 self.dtype) 

232 

233 for s in range(self.nspins): 

234 for k in range(self.kd.nibzkpts): 

235 for n in range(self.bd.nbands): 

236 psit_G = self.get_wave_function_array(n, k, s) 

237 writer.fill(psit_G * Bohr**-1.5) 

238 

239 def read(self, reader): 

240 FDPWWaveFunctions.read(self, reader) 

241 

242 if 'values' in reader.wave_functions: 

243 name = 'values' 

244 elif 'coefficients' in reader.wave_functions: 

245 name = 'coefficients' 

246 else: 

247 return 

248 

249 c = reader.bohr**1.5 

250 if reader.version < 0: 

251 c = 1 # old gpw file 

252 

253 for kpt in self.kpt_u: 

254 # We may not be able to keep all the wave 

255 # functions in memory - so psit_nG will be a special type of 

256 # array that is really just a reference to a file: 

257 psit_nG = reader.wave_functions.proxy(name, kpt.s, kpt.k) 

258 psit_nG.scale = c 

259 

260 kpt.psit = UniformGridWaveFunctions( 

261 self.bd.nbands, self.gd, self.dtype, psit_nG, 

262 kpt=kpt.q, dist=(self.bd.comm, self.bd.comm.size), 

263 spin=kpt.s, collinear=True) 

264 

265 if self.world.size > 1: 

266 # Read to memory: 

267 for kpt in self.kpt_u: 

268 kpt.psit.read_from_file() 

269 self.read_from_file_init_wfs_dm = True 

270 

271 def initialize_from_lcao_coefficients(self, basis_functions): 

272 for kpt in self.kpt_u: 

273 kpt.psit = UniformGridWaveFunctions( 

274 self.bd.nbands, self.gd, self.dtype, kpt=kpt.q, 

275 dist=(self.bd.comm, self.bd.comm.size, 1), 

276 spin=kpt.s, collinear=True) 

277 kpt.psit_nG[:] = 0.0 

278 mynbands = len(kpt.C_nM) 

279 basis_functions.lcao_to_grid(kpt.C_nM, 

280 kpt.psit_nG[:mynbands], kpt.q) 

281 kpt.C_nM = None 

282 

283 def random_wave_functions(self, nao): 

284 """Generate random wave functions.""" 

285 

286 gpts = self.gd.N_c[0] * self.gd.N_c[1] * self.gd.N_c[2] 

287 rng = np.random.default_rng(4 + self.world.rank) 

288 

289 if self.bd.nbands < gpts / 64: 

290 gd1 = self.gd.coarsen() 

291 gd2 = gd1.coarsen() 

292 

293 psit_G1 = gd1.empty(dtype=self.dtype) 

294 psit_G2 = gd2.empty(dtype=self.dtype) 

295 

296 interpolate2 = Transformer(gd2, gd1, 1, self.dtype).apply 

297 interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply 

298 

299 shape = tuple(gd2.n_c) 

300 scale = np.sqrt(12 / gd2.volume) 

301 

302 for kpt in self.kpt_u: 

303 for psit_G in kpt.psit_nG[nao:]: 

304 if self.dtype == float: 

305 psit_G2[:] = (rng.random(shape) - 0.5) * scale 

306 else: 

307 psit_G2.real = (rng.random(shape) - 0.5) * scale 

308 psit_G2.imag = (rng.random(shape) - 0.5) * scale 

309 

310 interpolate2(psit_G2, psit_G1, kpt.phase_cd) 

311 interpolate1(psit_G1, psit_G, kpt.phase_cd) 

312 

313 elif gpts / 64 <= self.bd.nbands < gpts / 8: 

314 gd1 = self.gd.coarsen() 

315 

316 psit_G1 = gd1.empty(dtype=self.dtype) 

317 

318 interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply 

319 

320 shape = tuple(gd1.n_c) 

321 scale = np.sqrt(12 / gd1.volume) 

322 

323 for kpt in self.kpt_u: 

324 for psit_G in kpt.psit_nG[nao:]: 

325 if self.dtype == float: 

326 psit_G1[:] = (rng.random(shape) - 0.5) * scale 

327 else: 

328 psit_G1.real = (rng.random(shape) - 0.5) * scale 

329 psit_G1.imag = (rng.random(shape) - 0.5) * scale 

330 

331 interpolate1(psit_G1, psit_G, kpt.phase_cd) 

332 

333 else: 

334 shape = tuple(self.gd.n_c) 

335 scale = np.sqrt(12 / self.gd.volume) 

336 

337 for kpt in self.kpt_u: 

338 for psit_G in kpt.psit_nG[nao:]: 

339 if self.dtype == float: 

340 psit_G[:] = (rng.random(shape) - 0.5) * scale 

341 else: 

342 psit_G.real = (rng.random(shape) - 0.5) * scale 

343 psit_G.imag = (rng.random(shape) - 0.5) * scale 

344 

345 def estimate_memory(self, mem): 

346 FDPWWaveFunctions.estimate_memory(self, mem)