Coverage for gpaw/wavefunctions/pw.py: 87%

1from __future__ import annotations 

2from math import pi 

3import numbers 

4 

5from ase.units import Bohr, Ha 

6from ase.utils.timing import timer 

7import numpy as np 

8 

9from gpaw.band_descriptor import BandDescriptor 

10from gpaw.blacs import BlacsDescriptor, BlacsGrid, Redistributor 

11from gpaw.lfc import BasisFunctions 

12from gpaw.matrix_descriptor import MatrixDescriptor 

13from gpaw.pw.descriptor import PWDescriptor 

14from gpaw.pw.lfc import PWLFC 

15from gpaw.typing import Array2D 

16from gpaw.utilities import unpack_hermitian 

17from gpaw.utilities.blas import axpy 

18from gpaw.utilities.progressbar import ProgressBar 

19from gpaw.wavefunctions.arrays import PlaneWaveExpansionWaveFunctions 

20from gpaw.wavefunctions.fdpw import FDPWWaveFunctions 

21from gpaw.wavefunctions.mode import Mode 

22import gpaw 

23import gpaw.cgpaw as cgpaw 

24import gpaw.fftw as fftw 

25 

26 

27class PW(Mode): 

28 name = 'pw' 

29 

30 def __init__(self, 

31 ecut: float = 340, 

32 *, 

33 fftwflags: int = fftw.MEASURE, 

34 cell: Array2D | None = None, 

35 gammacentered=False, # Deprecated 

36 pulay_stress: float | None = None, 

37 dedecut: float | str | None = None, 

38 force_complex_dtype: bool = False, 

39 interpolation: str | int = 'fft'): 

40 """Plane-wave basis mode. 

41 

42 Only one of ``dedecut`` and ``pulay_stress`` can be used. 

43 

44 Parameters 

45 ========== 

46 ecut: float 

47 Plane-wave cutoff in eV. 

48 gammacentered: bool 

49 Center the grid of chosen plane waves around the 

50 gamma point or q/k-vector 

51 dedecut: 

52 Estimate of derivative of total energy with respect to 

53 plane-wave cutoff. Used to calculate the Pulay-stress. 

54 pulay_stress: float or None 

55 Pulay-stress correction. 

56 fftwflags: int 

57 Flags for making an FFTW plan. There are 4 possibilities 

58 (default is MEASURE):: 

59 

60 from gpaw.fftw import ESTIMATE, MEASURE, PATIENT, EXHAUSTIVE 

61 

62 cell: np.ndarray 

63 Use this unit cell to chose the plane-waves. 

64 interpolation : str or int 

65 Interpolation scheme to construct the density on the fine grid. 

66 Default is ``'fft'`` and alternatively a stencil (integer) can 

67 be given to perform an explicit real-space interpolation. 

68 """ 

69 

70 assert not gammacentered 

71 self.ecut = ecut / Ha 

72 # Don't do expensive planning in dry-run mode: 

73 self.fftwflags = fftwflags if not gpaw.dry_run else fftw.MEASURE 

74 self.dedecut = dedecut 

75 self.interpolation = interpolation 

76 self.pulay_stress = (None 

77 if pulay_stress is None 

78 else pulay_stress * Bohr**3 / Ha) 

79 

80 assert pulay_stress is None or dedecut is None 

81 

82 if cell is None: 

83 self.cell_cv = None 

84 else: 

85 self.cell_cv = cell / Bohr 

86 

87 Mode.__init__(self, force_complex_dtype) 

88 

89 def __call__(self, parallel, initksl, gd, **kwargs): 

90 dedepsilon = 0.0 

91 

92 if self.cell_cv is None: 

93 ecut = self.ecut 

94 else: 

95 volume0 = abs(np.linalg.det(self.cell_cv)) 

96 ecut = self.ecut * (volume0 / gd.volume)**(2 / 3.0) 

97 

98 if self.pulay_stress is not None: 

99 dedepsilon = self.pulay_stress * gd.volume 

100 elif self.dedecut is not None: 

101 if self.dedecut == 'estimate': 

102 dedepsilon = 'estimate' 

103 else: 

104 dedepsilon = self.dedecut * 2 / 3 * ecut 

105 

106 wfs = PWWaveFunctions(ecut, 

107 self.fftwflags, dedepsilon, 

108 parallel, initksl, gd=gd, 

109 **kwargs) 

110 

111 return wfs 

112 

113 def todict(self): 

114 dct = Mode.todict(self) 

115 dct['ecut'] = self.ecut * Ha 

116 

117 if self.cell_cv is not None: 

118 dct['cell'] = self.cell_cv * Bohr 

119 if self.pulay_stress is not None: 

120 dct['pulay_stress'] = self.pulay_stress * Ha / Bohr**3 

121 if self.dedecut is not None: 

122 dct['dedecut'] = self.dedecut 

123 if self.interpolation != 'fft': 

124 dct['interpolation'] = self.interpolation 

125 return dct 

126 

127 

128class Preconditioner: 

129 """Preconditioner for KS equation. 

130 

131 From: 

132 

133 Teter, Payne and Allen, Phys. Rev. B 40, 12255 (1989) 

134 

135 as modified by: 

136 

137 Kresse and Furthmüller, Phys. Rev. B 54, 11169 (1996) 

138 """ 

139 

140 def __init__(self, G2_qG, pd, _scale=1.0): 

141 self.G2_qG = G2_qG 

142 self.pd = pd 

143 self._scale = _scale 

144 

145 def calculate_kinetic_energy(self, psit_xG, kpt): 

146 if psit_xG.ndim == 1: 

147 return self.calculate_kinetic_energy(psit_xG[np.newaxis], kpt)[0] 

148 G2_G = self.G2_qG[kpt.q] 

149 return np.array([self.pd.integrate(0.5 * G2_G * psit_G, psit_G).real 

150 for psit_G in psit_xG]) * self._scale 

151 

152 def __call__(self, R_xG, kpt, ekin_x, out=None): 

153 if out is None: 

154 out = np.empty_like(R_xG) 

155 G2_G = self.G2_qG[kpt.q] 

156 if R_xG.ndim == 1: 

157 cgpaw.pw_precond(G2_G, R_xG, ekin_x, out) 

158 else: 

159 for PR_G, R_G, ekin in zip(out, R_xG, ekin_x): 

160 cgpaw.pw_precond(G2_G, R_G, ekin, PR_G) 

161 return out 

162 

163 

164class NonCollinearPreconditioner(Preconditioner): 

165 def calculate_kinetic_energy(self, psit_xsG, kpt): 

166 shape = psit_xsG.shape 

167 ekin_xs = Preconditioner.calculate_kinetic_energy( 

168 self, psit_xsG.reshape((-1, shape[-1])), kpt) 

169 return ekin_xs.reshape(shape[:-1]).sum(-1) 

170 

171 def __call__(self, R_sG, kpt, ekin, out=None): 

172 return Preconditioner.__call__(self, R_sG, kpt, [ekin, ekin], out) 

173 

174 

175class PWWaveFunctions(FDPWWaveFunctions): 

176 mode = 'pw' 

177 

178 def __init__(self, ecut, fftwflags, dedepsilon, 

179 parallel, initksl, 

180 reuse_wfs_method, collinear, 

181 gd, nvalence, setups, bd, dtype, 

182 world, kd, kptband_comm, timer): 

183 self.ecut = ecut 

184 self.fftwflags = fftwflags 

185 self.dedepsilon = dedepsilon # Pulay correction for stress tensor 

186 

187 self.ng_k = None # number of G-vectors for all IBZ k-points 

188 

189 FDPWWaveFunctions.__init__(self, parallel, initksl, 

190 reuse_wfs_method=reuse_wfs_method, 

191 collinear=collinear, 

192 gd=gd, nvalence=nvalence, setups=setups, 

193 bd=bd, dtype=dtype, world=world, kd=kd, 

194 kptband_comm=kptband_comm, timer=timer) 

195 self.read_from_file_init_wfs_dm = False 

196 

197 def empty(self, n=(), global_array=False, realspace=False, q=None): 

198 if isinstance(n, numbers.Integral): 

199 n = (n,) 

200 if realspace: 

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

202 elif global_array: 

203 return np.zeros(n + (self.pd.ngmax,), complex) 

204 elif q is None: 

205 return np.zeros(n + (self.pd.maxmyng,), complex) 

206 else: 

207 return self.pd.empty(n, self.dtype, q) 

208 

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

210 return self.pd.integrate(a_xg, b_yg, global_integral) 

211 

212 def bytes_per_wave_function(self): 

213 return 16 * self.pd.maxmyng 

214 

215 def set_setups(self, setups): 

216 self.timer.start('PWDescriptor') 

217 self.pd = PWDescriptor(self.ecut, self.gd, self.dtype, self.kd, 

218 self.fftwflags) 

219 self.timer.stop('PWDescriptor') 

220 

221 # Build array of number of plane wave coefficiants for all k-points 

222 # in the IBZ: 

223 self.ng_k = np.zeros(self.kd.nibzkpts, dtype=int) 

224 for kpt in self.kpt_u: 

225 if kpt.s != 1: # avoid double counting (only sum over s=0 or None) 

226 self.ng_k[kpt.k] = len(self.pd.Q_qG[kpt.q]) 

227 self.kd.comm.sum(self.ng_k) 

228 

229 self.pt = PWLFC([setup.pt_j for setup in setups], self.pd) 

230 

231 FDPWWaveFunctions.set_setups(self, setups) 

232 

233 if self.dedepsilon == 'estimate': 

234 dedecut = self.setups.estimate_dedecut(self.ecut) 

235 self.dedepsilon = dedecut * 2 / 3 * self.ecut 

236 

237 def get_pseudo_partial_waves(self): 

238 return PWLFC([setup.get_partial_waves_for_atomic_orbitals() 

239 for setup in self.setups], self.pd) 

240 

241 def __str__(self): 

242 s = 'Wave functions: Plane wave expansion\n' 

243 s += ' Cutoff energy: %.3f eV\n' % (self.pd.ecut * Ha) 

244 

245 if self.dtype == float: 

246 s += (' Number of coefficients: %d (reduced to %d)\n' % 

247 (self.pd.ngmax * 2 - 1, self.pd.ngmax)) 

248 else: 

249 s += (' Number of coefficients (min, max): %d, %d\n' % 

250 (self.pd.ngmin, self.pd.ngmax)) 

251 

252 stress = self.dedepsilon / self.gd.volume * Ha / Bohr**3 

253 dedecut = 1.5 * self.dedepsilon / self.ecut 

254 s += (' Pulay-stress correction: {:.6f} eV/Ang^3 ' 

255 '(de/decut={:.6f})\n'.format(stress, dedecut)) 

256 

257 if fftw.have_fftw(): 

258 s += ' Using FFTW library\n' 

259 else: 

260 s += " Using Numpy's FFT\n" 

261 return s + FDPWWaveFunctions.__str__(self) 

262 

263 def make_preconditioner(self, block=1): 

264 if self.collinear: 

265 return Preconditioner(self.pd.G2_qG, self.pd) 

266 return NonCollinearPreconditioner(self.pd.G2_qG, self.pd) 

267 

268 @timer('Apply H') 

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

270 """Apply the pseudo Hamiltonian i.e. without PAW corrections.""" 

271 if not self.collinear: 

272 self.apply_pseudo_hamiltonian_nc(kpt, ham, psit_xG, Htpsit_xG) 

273 return 

274 

275 N = len(psit_xG) 

276 S = self.gd.comm.size 

277 

278 vt_R = self.gd.collect(ham.vt_sG[kpt.s], broadcast=True) 

279 Q_G = self.pd.Q_qG[kpt.q] 

280 T_G = 0.5 * self.pd.G2_qG[kpt.q] 

281 

282 for n1 in range(0, N, S): 

283 n2 = min(n1 + S, N) 

284 psit_G = self.pd.alltoall1(psit_xG[n1:n2], kpt.q) 

285 with self.timer('HMM T'): 

286 np.multiply(T_G, psit_xG[n1:n2], Htpsit_xG[n1:n2]) 

287 if psit_G is not None: 

288 psit_R = self.pd.ifft(psit_G, kpt.q, local=True, safe=False) 

289 psit_R *= vt_R 

290 self.pd.fftplan.execute() 

291 vtpsit_G = self.pd.tmp_Q.ravel()[Q_G] 

292 else: 

293 vtpsit_G = self.pd.tmp_G 

294 self.pd.alltoall2(vtpsit_G, kpt.q, Htpsit_xG[n1:n2]) 

295 

296 ham.xc.apply_orbital_dependent_hamiltonian( 

297 kpt, psit_xG, Htpsit_xG, ham.dH_asp) 

298 

299 def apply_pseudo_hamiltonian_nc(self, kpt, ham, psit_xG, Htpsit_xG): 

300 Htpsit_xG[:] = 0.5 * self.pd.G2_qG[kpt.q] * psit_xG 

301 v, x, y, z = ham.vt_xG 

302 iy = y * 1j 

303 for psit_sG, Htpsit_sG in zip(psit_xG, Htpsit_xG): 

304 a = self.pd.ifft(psit_sG[0], kpt.q) 

305 b = self.pd.ifft(psit_sG[1], kpt.q) 

306 Htpsit_sG[0] += self.pd.fft(a * (v + z) + b * (x - iy), kpt.q) 

307 Htpsit_sG[1] += self.pd.fft(a * (x + iy) + b * (v - z), kpt.q) 

308 

309 def add_orbital_density(self, nt_G, kpt, n): 

310 axpy(1.0, abs(self.pd.ifft(kpt.psit_nG[n], kpt.q))**2, nt_G) 

311 

312 def add_to_density_from_k_point_with_occupation(self, nt_xR, kpt, f_n): 

313 if not self.collinear: 

314 self.add_to_density_from_k_point_with_occupation_nc( 

315 nt_xR, kpt, f_n) 

316 return 

317 

318 comm = self.gd.comm 

319 

320 nt_R = self.gd.zeros(global_array=True) 

321 

322 for n1 in range(0, self.bd.mynbands, comm.size): 

323 n2 = min(n1 + comm.size, self.bd.mynbands) 

324 psit_G = self.pd.alltoall1(kpt.psit.array[n1:n2], kpt.q) 

325 if psit_G is not None: 

326 f = f_n[n1 + comm.rank] 

327 psit_R = self.pd.ifft(psit_G, kpt.q, local=True, safe=False) 

328 # Same as nt_R += f * abs(psit_R)**2, but much faster: 

329 cgpaw.add_to_density(f, psit_R, nt_R) 

330 

331 comm.sum(nt_R) 

332 nt_R = self.gd.distribute(nt_R) 

333 nt_xR[kpt.s] += nt_R 

334 

335 def add_to_density_from_k_point_with_occupation_nc(self, nt_xR, kpt, f_n): 

336 for f, psit_sG in zip(f_n, kpt.psit.array): 

337 p1 = self.pd.ifft(psit_sG[0], kpt.q) 

338 p2 = self.pd.ifft(psit_sG[1], kpt.q) 

339 p11 = p1.real**2 + p1.imag**2 

340 p22 = p2.real**2 + p2.imag**2 

341 p12 = p1.conj() * p2 

342 nt_xR[0] += f * (p11 + p22) 

343 nt_xR[1] += 2 * f * p12.real 

344 nt_xR[2] += 2 * f * p12.imag 

345 nt_xR[3] += f * (p11 - p22) 

346 

347 def add_to_kinetic_energy_density_kpt(self, kpt, psit_xG, taut_xR): 

348 N = self.bd.mynbands 

349 S = self.gd.comm.size 

350 Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) 

351 taut_R = self.gd.zeros(global_array=True) 

352 G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) 

353 comm = self.gd.comm 

354 

355 for v in range(3): 

356 for n1 in range(0, N, S): 

357 n2 = min(n1 + S, N) 

358 dn = n2 - n1 

359 Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] 

360 Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) 

361 if Gpsit_G is not None: 

362 f = kpt.f_n[n1 + comm.rank] 

363 a_R = self.pd.ifft(Gpsit_G, kpt.q, local=True, safe=False) 

364 cgpaw.add_to_density(0.5 * f, a_R, taut_R) 

365 

366 comm.sum(taut_R) 

367 taut_R = self.gd.distribute(taut_R) 

368 taut_xR[kpt.s] += taut_R 

369 

370 def add_to_ke_crossterms_kpt(self, kpt, psit_xG, taut_xR): 

371 N = self.bd.mynbands 

372 S = self.gd.comm.size 

373 Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) 

374 taut_wR = self.gd.zeros((6,), global_array=True) 

375 G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) 

376 comm = self.gd.comm 

377 

378 for n1 in range(0, N, S): 

379 n2 = min(n1 + S, N) 

380 dn = n2 - n1 

381 a_vR = {} 

382 for v in range(3): 

383 Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] 

384 Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) 

385 if Gpsit_G is not None: 

386 f = kpt.f_n[n1 + comm.rank] 

387 a_vR[v] = self.pd.ifft(Gpsit_G, kpt.q, 

388 local=True, safe=True) 

389 if len(a_vR) == 3: 

390 f = kpt.f_n[n1 + comm.rank] 

391 w = 0 

392 for v1 in range(3): 

393 for v2 in range(v1, 3): 

394 # imaginary parts should cancel 

395 taut_wR[w] += f * (a_vR[v1].conj() 

396 * a_vR[v2]).real 

397 w += 1 

398 elif len(a_vR) == 0: 

399 pass 

400 else: 

401 raise RuntimeError('Parallelization issue') 

402 

403 comm.sum(taut_wR) 

404 taut_wR = self.gd.distribute(taut_wR) 

405 taut_xR[kpt.s, :] += taut_wR 

406 

407 def calculate_kinetic_energy_density(self): 

408 if self.kpt_u[0].f_n is None: 

409 return None 

410 

411 taut_sR = self.gd.zeros(self.nspins) 

412 for kpt in self.kpt_u: 

413 self.add_to_kinetic_energy_density_kpt(kpt, kpt.psit_nG, taut_sR) 

414 

415 self.kptband_comm.sum(taut_sR) 

416 for taut_R in taut_sR: 

417 self.kd.symmetry.symmetrize(taut_R, self.gd) 

418 return taut_sR 

419 

420 def calculate_kinetic_energy_density_crossterms(self): 

421 if self.kpt_u[0].f_n is None: 

422 return None 

423 

424 taut_svvR = self.gd.zeros((self.nspins, 6)) 

425 for kpt in self.kpt_u: 

426 self.add_to_ke_crossterms_kpt(kpt, kpt.psit_nG, taut_svvR) 

427 

428 self.kptband_comm.sum(taut_svvR) 

429 for taut_R in taut_svvR.reshape(-1, *taut_svvR.shape[-3:]): 

430 self.kd.symmetry.symmetrize(taut_R, self.gd) 

431 return taut_svvR 

432 

433 def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, 

434 Htpsit_xG, dH_asp, 

435 dedtaut_R): 

436 N = len(psit_xG) 

437 S = self.gd.comm.size 

438 Q_G = self.pd.Q_qG[kpt.q] 

439 Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) 

440 tmp_xG = np.empty((S,) + psit_xG.shape[1:], dtype=Htpsit_xG.dtype) 

441 G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) 

442 

443 dedtaut_R = self.gd.collect(dedtaut_R, broadcast=True) 

444 

445 for v in range(3): 

446 for n1 in range(0, N, S): 

447 n2 = min(n1 + S, N) 

448 dn = n2 - n1 

449 Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] 

450 tmp_xG[:] = 0 

451 Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) 

452 if Gpsit_G is not None: 

453 a_R = self.pd.ifft(Gpsit_G, kpt.q, local=True, safe=False) 

454 a_R *= dedtaut_R 

455 self.pd.fftplan.execute() 

456 a_R = self.pd.tmp_Q.ravel()[Q_G] 

457 else: 

458 a_R = self.pd.tmp_G 

459 self.pd.alltoall2(a_R, kpt.q, tmp_xG[:dn]) 

460 axpy(-0.5, (1j * G_Gv[:, v] * tmp_xG[:dn]).ravel(), 

461 Htpsit_xG[n1:n2].ravel()) 

462 

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

464 kpt = self.kpt_u[u] 

465 psit_G = kpt.psit_nG[n] 

466 

467 if realspace: 

468 if psit_G.ndim == 2: 

469 psit_R = np.array([self.pd.ifft(psits_G, kpt.q) 

470 for psits_G in psit_G]) 

471 else: 

472 psit_R = self.pd.ifft(psit_G, kpt.q) 

473 if self.kd.gamma or periodic: 

474 return psit_R 

475 

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

477 eikr_R = self.gd.plane_wave(k_c) 

478 return psit_R * eikr_R 

479 

480 return psit_G 

481 

482 def get_wave_function_array(self, n, k, s, realspace=True, 

483 cut=True, periodic=False): 

484 kpt_rank, q = self.kd.get_rank_and_index(k) 

485 u = q * self.nspins + s 

486 band_rank, myn = self.bd.who_has(n) 

487 

488 rank = self.world.rank 

489 if (self.kd.comm.rank == kpt_rank and 

490 self.bd.comm.rank == band_rank): 

491 psit_G = self._get_wave_function_array(u, myn, realspace, periodic) 

492 

493 if realspace: 

494 psit_G = self.gd.collect(psit_G) 

495 else: 

496 assert not cut 

497 tmp_G = self.pd.gather(psit_G, self.kpt_u[u].q) 

498 if tmp_G is not None: 

499 ng = self.pd.ngmax 

500 if self.collinear: 

501 psit_G = np.zeros(ng, complex) 

502 else: 

503 psit_G = np.zeros((2, ng), complex) 

504 psit_G[..., :tmp_G.shape[-1]] = tmp_G 

505 

506 if rank == 0: 

507 return psit_G 

508 

509 # Domain master send this to the global master 

510 if self.gd.comm.rank == 0: 

511 self.world.ssend(psit_G, 0, 1398) 

512 

513 if rank == 0: 

514 # allocate full wave function and receive 

515 shape = () if self.collinear else (2,) 

516 psit_G = self.empty(shape, global_array=True, 

517 realspace=realspace) 

518 # XXX this will fail when using non-standard nesting 

519 # of communicators. 

520 world_rank = (kpt_rank * self.gd.comm.size * 

521 self.bd.comm.size + 

522 band_rank * self.gd.comm.size) 

523 self.world.receive(psit_G, world_rank, 1398) 

524 return psit_G 

525 

526 # We return a number instead of None on all the slaves. Most of 

527 # the time the return value will be ignored on the slaves, but 

528 # in some cases it will be multiplied by some other number and 

529 # then ignored. Allowing for this will simplify some code here 

530 # and there. 

531 return np.nan 

532 

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

534 FDPWWaveFunctions.write(self, writer) 

535 

536 if not write_wave_functions: 

537 return 

538 

539 if self.collinear: 

540 shape = (self.nspins, 

541 self.kd.nibzkpts, self.bd.nbands, self.pd.ngmax) 

542 else: 

543 shape = (self.kd.nibzkpts, self.bd.nbands, 2, self.pd.ngmax) 

544 

545 writer.add_array('coefficients', shape, complex) 

546 

547 c = Bohr**-1.5 

548 for s in range(self.nspins): 

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

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

551 psit_G = self.get_wave_function_array(n, k, s, 

552 realspace=False, 

553 cut=False) 

554 writer.fill(psit_G * c) 

555 

556 writer.add_array('indices', (self.kd.nibzkpts, self.pd.ngmax), 

557 np.int32) 

558 

559 if self.bd.comm.rank > 0: 

560 return 

561 

562 Q_G = np.empty(self.pd.ngmax, np.int32) 

563 kk = 0 

564 for r in range(self.kd.comm.size): 

565 for q, k in enumerate(self.kd.get_indices(r)): 

566 ng = self.ng_k[k] 

567 if r == self.kd.comm.rank: 

568 Q_G[:ng] = self.pd.Q_qG[q] 

569 if r > 0: 

570 self.kd.comm.send(Q_G, 0) 

571 if self.kd.comm.rank == 0: 

572 if r > 0: 

573 self.kd.comm.receive(Q_G, r) 

574 Q_G[ng:] = -1 

575 writer.fill(Q_G) 

576 assert k == kk 

577 kk += 1 

578 

579 def read(self, reader): 

580 FDPWWaveFunctions.read(self, reader) 

581 

582 if 'coefficients' not in reader.wave_functions: 

583 return 

584 

585 Q_kG = reader.wave_functions.indices 

586 for kpt in self.kpt_u: 

587 if kpt.s == 0: 

588 Q_G = Q_kG[kpt.k] 

589 ng = self.ng_k[kpt.k] 

590 assert (Q_G[:ng] == self.pd.Q_qG[kpt.q]).all() 

591 assert (Q_G[ng:] == -1).all() 

592 

593 c = reader.bohr**1.5 

594 if reader.version < 0: 

595 c = 1 # old gpw file 

596 elif reader.version >= 4: 

597 c *= self.gd.N_c.prod() 

598 for kpt in self.kpt_u: 

599 ng = self.ng_k[kpt.k] 

600 index = (kpt.s, kpt.k) if self.collinear else (kpt.k,) 

601 psit_nG = reader.wave_functions.proxy('coefficients', *index) 

602 psit_nG.scale = c 

603 psit_nG.length_of_last_dimension = ng 

604 

605 kpt.psit = PlaneWaveExpansionWaveFunctions( 

606 self.bd.nbands, self.pd, self.dtype, psit_nG, 

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

608 spin=kpt.s, collinear=self.collinear) 

609 

610 if self.world.size > 1: 

611 # Read to memory: 

612 for kpt in self.kpt_u: 

613 kpt.psit.read_from_file() 

614 self.read_from_file_init_wfs_dm = True 

615 

616 def hs(self, ham, q=-1, s=0, md=None): 

617 npw = len(self.pd.Q_qG[q]) 

618 N = self.pd.tmp_R.size 

619 

620 if md is None: 

621 H_GG = np.zeros((npw, npw), complex) 

622 S_GG = np.zeros((npw, npw), complex) 

623 G1 = 0 

624 G2 = npw 

625 else: 

626 H_GG = md.zeros(dtype=complex) 

627 S_GG = md.zeros(dtype=complex) 

628 if S_GG.size == 0: 

629 return H_GG, S_GG 

630 G1, G2 = next(md.my_blocks(S_GG))[:2] 

631 

632 H_GG.ravel()[G1::npw + 1] = (0.5 * self.pd.gd.dv / N * 

633 self.pd.G2_qG[q][G1:G2]) 

634 for G in range(G1, G2): 

635 x_G = self.pd.zeros(q=q) 

636 x_G[G] = 1.0 

637 H_GG[G - G1] += (self.pd.gd.dv / N * 

638 self.pd.fft(ham.vt_sG[s] * 

639 self.pd.ifft(x_G, q), q)) 

640 

641 if ham.xc.type == 'MGGA': 

642 G_Gv = self.pd.get_reciprocal_vectors(q=q) 

643 for G in range(G1, G2): 

644 x_G = self.pd.zeros(q=q) 

645 x_G[G] = 1.0 

646 for v in range(3): 

647 a_R = self.pd.ifft(1j * G_Gv[:, v] * x_G, q)