Coverage for gpaw/pw/descriptor.py: 84%

1import numbers

2from math import pi

3import numpy as np

5import gpaw.cgpaw as cgpaw

6import gpaw.fftw as fftw

7from gpaw.utilities.blas import mmm, r2k, rk

8from gpaw.gpu import cupy as cp

11class PWDescriptor:

12 ndim = 1 # all 3d G-vectors are stored in a 1d ndarray

14 def __init__(self, ecut, gd, dtype=None, kd=None,

15 fftwflags=fftw.MEASURE, gammacentered=False,

16 _new=False):

18 assert gd.pbc_c.all()

20 self.gd = gd

21 self.fftwflags = fftwflags

23 N_c = gd.N_c

24 self.comm = gd.comm

26 ecut0 = 0.5 * pi**2 / (self.gd.h_cv**2).sum(1).max()

27 if ecut is None:

28 ecut = 0.9999 * ecut0

29 else:

30 assert ecut <= ecut0

32 self.ecut = ecut

34 if dtype is None:

35 if kd is None or kd.gamma:

36 dtype = float

37 else:

38 dtype = complex

39 self.dtype = dtype

40 self.gammacentered = gammacentered

42 if dtype == float:

43 Nr_c = N_c.copy()

44 Nr_c[2] = N_c[2] // 2 + 1

45 i_Qc = np.indices(Nr_c).transpose((1, 2, 3, 0))

46 i_Qc[..., :2] += N_c[:2] // 2

47 i_Qc[..., :2] %= N_c[:2]

48 i_Qc[..., :2] -= N_c[:2] // 2

49 self.tmp_Q = fftw.empty(Nr_c, complex)

50 self.tmp_R = self.tmp_Q.view(float)[:, :, :N_c[2]]

51 else:

52 i_Qc = np.indices(N_c).transpose((1, 2, 3, 0))

53 i_Qc += N_c // 2

54 i_Qc %= N_c

55 i_Qc -= N_c // 2

56 self.tmp_Q = fftw.empty(N_c, complex)

57 self.tmp_R = self.tmp_Q

59 self.fftplan = fftw.create_plan(self.tmp_R, self.tmp_Q, -1, fftwflags)

60 self.ifftplan = fftw.create_plan(self.tmp_Q, self.tmp_R, 1, fftwflags)

62 # Calculate reciprocal lattice vectors:

63 B_cv = 2.0 * pi * gd.icell_cv

64 i_Qc.shape = (-1, 3)

65 self.G_Qv = np.dot(i_Qc, B_cv)

67 self.kd = kd

68 if kd is None:

69 self.K_qv = np.zeros((1, 3))

70 self.only_one_k_point = True

71 else:

72 self.K_qv = np.dot(kd.ibzk_qc, B_cv)

73 self.only_one_k_point = (kd.nbzkpts == 1)

75 # Map from vectors inside sphere to fft grid:

76 Q_Q = np.arange(len(i_Qc), dtype=np.int32)

78 self.ng_q, self.Q_qG, G2_qG = \

79 self.setup_pw_grid(i_Qc, Q_Q)

81 self.ngmin = min(self.ng_q)

82 self.ngmax = max(self.ng_q)

84 if kd is not None:

85 self.ngmin = kd.comm.min_scalar(self.ngmin)

86 self.ngmax = kd.comm.max_scalar(self.ngmax)

88 # Distribute things:

89 S = gd.comm.size

90 self.maxmyng = (self.ngmax + S - 1) // S

91 ng1 = gd.comm.rank * self.maxmyng

92 ng2 = ng1 + self.maxmyng

94 self.G2_qG = []

95 self.myQ_qG = []

96 self.myng_q = []

97 self.maxmyng_q = []

98 for q, G2_G in enumerate(G2_qG):

99 if _new:

100 x = (self.ng_q[q] + S - 1) // S

101 self.maxmyng_q.append(x)

102 ng1 = gd.comm.rank * x

103 ng2 = ng1 + x

104 G2_G = G2_G[ng1:ng2].copy()

105 G2_G.flags.writeable = False

106 self.G2_qG.append(G2_G)

107 myQ_G = self.Q_qG[q][ng1:ng2]

108 self.myQ_qG.append(myQ_G)

109 self.myng_q.append(len(myQ_G))

110

111 if S > 1:

112 self.tmp_G = np.empty(self.maxmyng * S, complex)

113 else:

114 self.tmp_G = None

115

116 self._new = _new

117

118 def setup_pw_grid(self, i_Qc, Q_Q):

119 ng_q = []

120 Q_qG = []

121 G2_qG = []

122 for q, K_v in enumerate(self.K_qv):

123 G2_Q = ((self.G_Qv + K_v)**2).sum(axis=1)

124 if self.gammacentered:

125 mask_Q = ((self.G_Qv**2).sum(axis=1) <= 2 * self.ecut)

126 else:

127 mask_Q = (G2_Q <= 2 * self.ecut)

128

129 if self.dtype == float:

130 mask_Q &= ((i_Qc[:, 2] > 0) |

131 (i_Qc[:, 1] > 0) |

132 ((i_Qc[:, 0] >= 0) & (i_Qc[:, 1] == 0)))

133 Q_G = Q_Q[mask_Q]

134 Q_qG.append(Q_G)

135 G2_qG.append(G2_Q[Q_G])

136 ng = len(Q_G)

137 ng_q.append(ng)

138

139 return ng_q, Q_qG, G2_qG

140

141 def get_reciprocal_vectors(self, q=0, add_q=True):

142 """Returns reciprocal lattice vectors plus q, G + q,

143 in xyz coordinates."""

144

145 if add_q:

146 q_v = self.K_qv[q]

147 return self.G_Qv[self.myQ_qG[q]] + q_v

148 return self.G_Qv[self.myQ_qG[q]]

149

150 def __getstate__(self):

151 return (self.ecut, self.gd, self.dtype, self.kd, self.fftwflags)

152

153 def __setstate__(self, state):

154 self.__init__(*state)

155

156 def estimate_memory(self, mem):

157 nbytes = (self.tmp_R.nbytes +

158 self.G_Qv.nbytes +

159 len(self.K_qv) * (self.ngmax * 4 +

160 self.maxmyng * (8 + 4)))

161 mem.subnode('Arrays', nbytes)

162

163 def bytecount(self, dtype=float):

164 return self.ngmax * 16

165

166 def zeros(self, x=(), dtype=None, q=None, global_array=False):

167 """Return zeroed array.

168

169 The shape of the array will be x + (ng,) where ng is the number

170 of G-vectors for on this core. Different k-points will have

171 different values for ng. Therefore, the q index must be given,

172 unless we are describibg a real-valued function."""

173

174 a_xG = self.empty(x, dtype, q, global_array)

175 a_xG.fill(0.0)

176 return a_xG

177

178 def empty(self, x=(), dtype=None, q=None, global_array=False):

179 """Return empty array."""

180 if dtype is not None:

181 assert dtype == self.dtype

182 if isinstance(x, numbers.Integral):

183 x = (x,)

184 if q is None:

185 assert self.only_one_k_point

186 q = 0

187 if global_array:

188 shape = x + (self.ng_q[q],)

189 else:

190 shape = x + (self.myng_q[q],)

191 return np.empty(shape, complex)

192

193 def fft(self, f_R, q=None, Q_G=None, local=False):

194 """Fast Fourier transform.

195

196 Returns c(G) for G<Gc::

197

198 -- -iG.R

199 c(G) = > f(R) e

200 --

201 R

202

203 If local=True, all cores will do an FFT without any

204 collect/scatter.

205 """

206

207 if local:

208 self.tmp_R[:] = f_R

209 else:

210 self.gd.collect(f_R, self.tmp_R)

211

212 if self.gd.comm.rank == 0 or local:

213 self.fftplan.execute()

214 if Q_G is None:

215 q = q or 0

216 Q_G = self.Q_qG[q]

217 f_G = self.tmp_Q.ravel()[Q_G]

218 if local:

219 return f_G

220 else:

221 f_G = None

222

223 return self.scatter(f_G, q)

224

225 def ifft(self, c_G, q=None, local=False, safe=True, distribute=True):

226 """Inverse fast Fourier transform.

227

228 Returns::

229

230 1 -- iG.R

231 f(R) = - > c(G) e

232 N --

233 G

234

235 If local=True, all cores will do an iFFT without any

236 gather/distribute.

237 """

238 assert q is not None or self.only_one_k_point

239 q = q or 0

240 if not local:

241 c_G = self.gather(c_G, q)

242 comm = self.gd.comm

243 scale = 1.0 / self.tmp_R.size

244 if comm.rank == 0 or local:

245 # Same as:

246 #

247 # self.tmp_Q[:] = 0.0

248 # self.tmp_Q.ravel()[self.Q_qG[q]] = scale * c_G

249 #

250 # but much faster:

251 Q_G = self.Q_qG[q]

252 assert len(c_G) == len(Q_G)

253 cgpaw.pw_insert(c_G, Q_G, scale, self.tmp_Q)

254 if self.dtype == float:

255 t = self.tmp_Q[:, :, 0]

256 n, m = self.gd.N_c[:2] // 2 - 1

257 t[0, -m:] = t[0, m:0:-1].conj()

258 t[n:0:-1, -m:] = t[-n:, m:0:-1].conj()

259 t[-n:, -m:] = t[n:0:-1, m:0:-1].conj()

260 t[-n:, 0] = t[n:0:-1, 0].conj()

261 self.ifftplan.execute()

262 if comm.size == 1 or local or not distribute:

263 if safe:

264 return self.tmp_R.copy()

265 return self.tmp_R

266 return self.gd.distribute(self.tmp_R)

267

268 def scatter(self, a_G, q=None):

269 """Scatter coefficients from master to all cores."""

270 comm = self.gd.comm

271 if comm.size == 1:

272 return a_G

273

274 mya_G = np.empty(self.maxmyng, complex)

275 comm.scatter(pad(a_G, self.maxmyng * comm.size), mya_G, 0)

276 return mya_G[:self.myng_q[q or 0]]

277

278 def gather(self, a_G, q=None):

279 """Gather coefficients on master."""

280 comm = self.gd.comm

281

282 if comm.size == 1:

283 return a_G

284

285 mya_G = pad(a_G, self.maxmyng)

286 if comm.rank == 0:

287 a_G = self.tmp_G

288 else:

289 a_G = None

290 comm.gather(mya_G, 0, a_G)

291 if comm.rank == 0:

292 return a_G[:self.ng_q[q or 0]]

293

294 def alltoall1(self, a_rG, q):

295 """Gather coefficients from a_rG[r] on rank r.

296

297 On rank r, an array of all G-vector coefficients will be returned.

298 These will be gathered from a_rG[r] on all the cores.

299 """

300 comm = self.gd.comm

301 if comm.size == 1:

302 return a_rG[0]

303 N = len(a_rG)

304 ng = self.ng_q[q]

305 ssize_r = np.zeros(comm.size, int)

306 ssize_r[:N] = self.myng_q[q]

307 soffset_r = np.arange(comm.size) * self.myng_q[q]

308 soffset_r[N:] = 0

309 myng = self.maxmyng_q[q] if self._new else self.maxmyng

310 roffset_r = (np.arange(comm.size) * myng).clip(max=ng)

311 rsize_r = np.zeros(comm.size, int)

312 if comm.rank < N:

313 rsize_r[:-1] = roffset_r[1:] - roffset_r[:-1]

314 rsize_r[-1] = ng - roffset_r[-1]

315 b_G = self.tmp_G[:ng]

316 comm.alltoallv(a_rG, ssize_r, soffset_r, b_G, rsize_r, roffset_r)

317 if comm.rank < N:

318 return b_G

319

320 def alltoall2(self, a_G, q, b_rG):

321 """Scatter all coefs. from rank r to B_rG[r] on other cores."""

322 comm = self.gd.comm

323 if comm.size == 1:

324 b_rG[0] += a_G

325 return

326 N = len(b_rG)

327 ng = self.ng_q[q]

328 rsize_r = np.zeros(comm.size, int)

329 rsize_r[:N] = self.myng_q[q]

330 roffset_r = np.arange(comm.size) * self.myng_q[q]

331 roffset_r[N:] = 0

332 myng = self.maxmyng_q[q] if self._new else self.maxmyng

333 soffset_r = (np.arange(comm.size) * myng).clip(max=ng)

334 ssize_r = np.zeros(comm.size, int)

335 if comm.rank < N:

336 ssize_r[:-1] = soffset_r[1:] - soffset_r[:-1]

337 ssize_r[-1] = ng - soffset_r[-1]

338 tmp_rG = self.tmp_G[:b_rG.size].reshape(b_rG.shape)

339 comm.alltoallv(a_G, ssize_r, soffset_r, tmp_rG, rsize_r, roffset_r)

340 b_rG += tmp_rG

341

342 def integrate(self, a_xg, b_yg=None,

343 global_integral=True, hermitian=False):

344 """Integrate function(s) over domain.

345

346 a_xg: ndarray

347 Function(s) to be integrated.

348 b_yg: ndarray

349 If present, integrate a_xg.conj() * b_yg.

350 global_integral: bool

351 If the array(s) are distributed over several domains, then the

352 total sum will be returned. To get the local contribution

353 only, use global_integral=False.

354 hermitian: bool

355 Result is hermitian.

356 """

357

358 if b_yg is None:

359 # Only one array:

360 assert self.dtype == float and self.gd.comm.size == 1

361 return a_xg[..., 0].real * self.gd.dv

362

363 if a_xg.ndim == 1:

364 A_xg = a_xg.reshape((1, len(a_xg)))

365 else:

366 A_xg = a_xg

367 if b_yg.ndim == 1:

368 B_yg = b_yg.reshape((1, len(b_yg)))

369 else:

370 B_yg = b_yg

371

372 alpha = self.gd.dv / self.gd.N_c.prod()

373

374 if self.dtype == float:

375 alpha *= 2

376 A_xg = A_xg.view(float)

377 B_yg = B_yg.view(float)

378

379 result_yx = np.zeros((len(B_yg), len(A_xg)), self.dtype)

380

381 if a_xg is b_yg:

382 rk(alpha, A_xg, 0.0, result_yx)

383 elif hermitian:

384 r2k(0.5 * alpha, A_xg, B_yg, 0.0, result_yx)

385 else:

386 mmm(alpha, B_yg, 'N', A_xg, 'C', 0.0, result_yx)

387

388 if self.dtype == float and self.gd.comm.rank == 0:

389 correction_yx = np.outer(B_yg[:, 0], A_xg[:, 0])

390 if hermitian:

391 result_yx -= 0.25 * alpha * (correction_yx + correction_yx.T)

392 else:

393 result_yx -= 0.5 * alpha * correction_yx

394

395 xshape = a_xg.shape[:-1]

396 yshape = b_yg.shape[:-1]

397 result = result_yx.T.reshape(xshape + yshape)

398

399 if result.ndim == 0:

400 if global_integral:

401 return self.gd.comm.sum_scalar(result.item())

402 return result.item()

403 else:

404 assert global_integral or self.gd.comm.size == 1

405 self.gd.comm.sum(result.T)

406 return result

407

408 def interpolate(self, a_R, pd):

409 if (pd.gd.N_c <= self.gd.N_c).any():

410 raise ValueError('Too few points in target grid!')

411

412 self.gd.collect(a_R, self.tmp_R[:])

413

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

415 self.fftplan.execute()

416

417 a_Q = self.tmp_Q

418 b_Q = pd.tmp_Q

419

420 e0, e1, e2 = 1 - self.gd.N_c % 2 # even or odd size

421 a0, a1, a2 = pd.gd.N_c // 2 - self.gd.N_c // 2

422 b0, b1, b2 = self.gd.N_c + (a0, a1, a2)

423

424 if self.dtype == float:

425 b2 = (b2 - a2) // 2 + 1

426 a2 = 0

427 axes = (0, 1)

428 else:

429 axes = (0, 1, 2)

430

431 b_Q[:] = 0.0

432 b_Q[a0:b0, a1:b1, a2:b2] = np.fft.fftshift(a_Q, axes=axes)

433

434 if e0:

435 b_Q[a0, a1:b1, a2:b2] *= 0.5

436 b_Q[b0, a1:b1, a2:b2] = b_Q[a0, a1:b1, a2:b2]

437 b0 += 1

438 if e1:

439 b_Q[a0:b0, a1, a2:b2] *= 0.5

440 b_Q[a0:b0, b1, a2:b2] = b_Q[a0:b0, a1, a2:b2]

441 b1 += 1

442 if self.dtype == complex:

443 if e2:

444 b_Q[a0:b0, a1:b1, a2] *= 0.5

445 b_Q[a0:b0, a1:b1, b2] = b_Q[a0:b0, a1:b1, a2]

446 else:

447 if e2:

448 b_Q[a0:b0, a1:b1, b2 - 1] *= 0.5

449

450 b_Q[:] = np.fft.ifftshift(b_Q, axes=axes)

451 pd.ifftplan.execute()

452

453 a_G = a_Q.ravel()[self.Q_qG[0]]

454 else:

455 a_G = None

456

457 return (pd.gd.distribute(pd.tmp_R) * (1.0 / self.tmp_R.size),

458 self.scatter(a_G))

459

460 def restrict(self, a_R, pd):

461 self.gd.collect(a_R, self.tmp_R[:])

462

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

464 a_Q = pd.tmp_Q

465 b_Q = self.tmp_Q

466

467 e0, e1, e2 = 1 - pd.gd.N_c % 2 # even or odd size

468 a0, a1, a2 = self.gd.N_c // 2 - pd.gd.N_c // 2

469 b0, b1, b2 = pd.gd.N_c // 2 + self.gd.N_c // 2 + 1

470

471 if self.dtype == float:

472 b2 = pd.gd.N_c[2] // 2 + 1

473 a2 = 0

474 axes = (0, 1)

475 else:

476 axes = (0, 1, 2)

477

478 self.fftplan.execute()

479 b_Q[:] = np.fft.fftshift(b_Q, axes=axes)

480

481 if e0:

482 b_Q[a0, a1:b1, a2:b2] += b_Q[b0 - 1, a1:b1, a2:b2]

483 b_Q[a0, a1:b1, a2:b2] *= 0.5

484 b0 -= 1

485 if e1:

486 b_Q[a0:b0, a1, a2:b2] += b_Q[a0:b0, b1 - 1, a2:b2]

487 b_Q[a0:b0, a1, a2:b2] *= 0.5

488 b1 -= 1

489 if self.dtype == complex and e2:

490 b_Q[a0:b0, a1:b1, a2] += b_Q[a0:b0, a1:b1, b2 - 1]

491 b_Q[a0:b0, a1:b1, a2] *= 0.5

492 b2 -= 1

493

494 a_Q[:] = b_Q[a0:b0, a1:b1, a2:b2]

495 a_Q[:] = np.fft.ifftshift(a_Q, axes=axes)

496 a_G = a_Q.ravel()[pd.Q_qG[0]] / 8

497 pd.ifftplan.execute()

498 else:

499 a_G = None

500

501 return (pd.gd.distribute(pd.tmp_R) * (1.0 / self.tmp_R.size),

502 pd.scatter(a_G))

503

504

505class PWMapping:

506 def __init__(self, pd1, pd2):

507 """Mapping from pd1 to pd2."""

508 N_c = np.array(pd1.tmp_Q.shape)

509 N2_c = pd2.tmp_Q.shape

510 Q1_G = pd1.Q_qG[0]

511 Q1_Gc = np.empty((len(Q1_G), 3), int)

512 Q1_Gc[:, 0], r_G = divmod(Q1_G, N_c[1] * N_c[2])

513 Q1_Gc.T[1:] = divmod(r_G, N_c[2])

514 if pd1.dtype == float:

515 C = 2

516 else:

517 C = 3

518 Q1_Gc[:, :C] += N_c[:C] // 2

519 Q1_Gc[:, :C] %= N_c[:C]

520 Q1_Gc[:, :C] -= N_c[:C] // 2

521 Q1_Gc[:, :C] %= N2_c[:C]

522 Q2_G = Q1_Gc[:, 2] + N2_c[2] * (Q1_Gc[:, 1] + N2_c[1] * Q1_Gc[:, 0])

523 G2_Q = np.empty(N2_c, int).ravel()

524 G2_Q[:] = -1

525 G2_Q[pd2.myQ_qG[0]] = np.arange(pd2.myng_q[0])

526 G2_G1 = G2_Q[Q2_G]

527

528 mask_G1 = (G2_G1 != -1)

529 self.G2_G1 = G2_G1[mask_G1]

530 self.G1 = np.arange(pd1.ngmax)[mask_G1]

531

532 self.pd1 = pd1

533 self.pd2 = pd2

534

535 def add_to1(self, a_G1, b_G2):

536 """Do a += b * scale, where a is on pd1 and b on pd2."""

537 scale = self.pd1.tmp_R.size / self.pd2.tmp_R.size

538

539 if self.pd1.gd.comm.size == 1:

540 a_G1 += b_G2[self.G2_G1] * scale

541 return

542

543 b_G1 = self.pd1.tmp_G

544 b_G1[:] = 0.0

545 b_G1[self.G1] = b_G2[self.G2_G1]

546 self.pd1.gd.comm.sum(b_G1)

547 ng1 = self.pd1.gd.comm.rank * self.pd1.maxmyng

548 ng2 = ng1 + self.pd1.myng_q[0]

549 a_G1 += b_G1[ng1:ng2] * scale

550

551 def add_to2(self, a_G2, b_G1):

552 """Do a += b * scale, where a is on pd2 and b on pd1."""

553 myb_G1 = b_G1 * (self.pd2.tmp_R.size / self.pd1.tmp_R.size)

554 if self.pd1.gd.comm.size == 1:

555 a_G2[self.G2_G1] += myb_G1

556 return

557

558 b_G1 = self.pd1.tmp_G

559 self.pd1.gd.comm.all_gather(pad(myb_G1, self.pd1.maxmyng), b_G1)

560 a_G2[self.G2_G1] += b_G1[self.G1]

561

562

563def count_reciprocal_vectors(ecut, gd, q_c):

564 assert gd.comm.size == 1

565 N_c = gd.N_c

566 i_Qc = np.indices(N_c).transpose((1, 2, 3, 0))

567 i_Qc += N_c // 2

568 i_Qc %= N_c

569 i_Qc -= N_c // 2

570

571 B_cv = 2.0 * pi * gd.icell_cv

572 i_Qc.shape = (-1, 3)

573 Gpq_Qv = np.dot(i_Qc, B_cv) + np.dot(q_c, B_cv)

574

575 G2_Q = (Gpq_Qv**2).sum(axis=1)

576 return (G2_Q <= 2 * ecut).sum()

577

578

579def pad(array, N):

580 """Pad 1-d ndarray with zeros up to length N.

581

582 >>> a = np.ones(2, complex)

583 >>> pad(a, 3)

584 array([1.+0.j, 1.+0.j, 0.+0.j])

585 >>> pad(a, 2) is a

586 True

587 >>> pad(None, 7) is None

588 True

589 """

590 if array is None:

591 return None

592 n = len(array)

593 if n == N:

594 return array

595 if isinstance(array, np.ndarray):

596 b = np.empty(N, complex)

597 else:

598 b = cp.empty(N, complex)

599 b[:n] = array

600 b[n:] = 0

601 return b