Coverage for gpaw/lfc.py: 80%

1from math import pi 

2 

3import numpy as np 

4from ase.units import Bohr 

5 

6import gpaw.cgpaw as cgpaw 

7from gpaw import debug 

8from gpaw.grid_descriptor import GridDescriptor, GridBoundsError 

9from gpaw.gpu import cupy_is_fake 

10from gpaw.new import trace 

11from gpaw.utilities import smallest_safe_grid_spacing 

12 

13""" 

14 

15=== ================================================= 

16 M Global localized function number. 

17 W Global volume number. 

18 G Global grid point number. 

19 g Local (inside sphere) grid point number. 

20 i Index into list of current spheres for current G. 

21=== ================================================= 

22 

23l 

24m 

25b 

26w 

27 

28Global grid point number (*G*) for a 7*6 grid:: 

29 ------------- 

30 |5 . . . . . .| 

31 |4 . . . . . .| 

32 |3 9 . . . . .| 

33 |2 8 . . . . .| 

34 |1 7 . . . . .| 

35 |0 6 . . . . .| 

36 ------------- 

37 

38For this example *G* runs from 0 to 41. 

39 

40Here is a sphere inside the box with grid points (*g*) numbered from 0 

41to 7:: 

42 

43 ------------- 

44 |. . . . . . .| 

45 |. . . . 5 . .| 

46 |. . . 1 4 7 .| 

47 |. . . 0 3 6 .| 

48 |. . . . 2 . .| 

49 |. . . . . . .| 

50 ------------- 

51 

52~ _ ^ ~ ~ 

53p v g n F 

54 i L c M 

55 

56i 

57d d d d d 

58s s 

59 s s 

60""" 

61 

62 

63class Sphere: 

64 def __init__(self, spline_j): 

65 self.spline_j = spline_j 

66 self.spos_c = None 

67 self.rank = None # Rank corresponding to center 

68 self.ranks = None # Ranks with at least some overlap 

69 self.Mmax = None 

70 self.A_wgm = None 

71 self.G_wb = None 

72 self.M_w = None 

73 self.sdisp_wc = None 

74 self.normalized = False 

75 self.I_M = None 

76 

77 def set_position(self, spos_c, gd, cut): 

78 if self.spos_c is not None and not (self.spos_c - spos_c).any(): 

79 return False 

80 

81 self.A_wgm = [] 

82 self.G_wb = [] 

83 self.M_w = [] 

84 self.sdisp_wc = [] 

85 

86 ng = 0 

87 M = 0 

88 for spline in self.spline_j: 

89 rcut = spline.get_cutoff() 

90 l = spline.get_angular_momentum_number() 

91 for beg_c, end_c, sdisp_c in gd.get_boxes(spos_c, rcut, cut): 

92 pos_v = np.dot(spos_c - sdisp_c, gd.cell_cv) 

93 A_gm, G_b = cgpaw.spline_to_grid( 

94 spline.spline, 

95 beg_c, end_c, pos_v, 

96 np.ascontiguousarray(gd.h_cv), 

97 gd.n_c, gd.beg_c) 

98 if len(G_b) > 0: 

99 self.A_wgm.append(A_gm) 

100 self.G_wb.append(G_b) 

101 self.M_w.append(M) 

102 self.sdisp_wc.append(sdisp_c) 

103 ng += A_gm.shape[0] 

104 assert A_gm.shape[0] > 0 

105 M += 2 * l + 1 

106 

107 self.Mmax = M 

108 

109 self.rank = gd.get_rank_from_position(spos_c) 

110 if ng == 0: 

111 self.ranks = None # What about making empty lists instead? 

112 self.A_wgm = None 

113 self.G_wb = None 

114 self.M_w = None 

115 self.sdisp_wc = None 

116 

117 self.spos_c = spos_c.copy() 

118 self.normalized = False 

119 return True 

120 

121 def get_function_count(self): 

122 return sum([2 * spline.get_angular_momentum_number() + 1 

123 for spline in self.spline_j]) 

124 

125 def normalize(self, integral, a, dv, comm): 

126 """Normalize localized functions.""" 

127 if self.normalized or integral < 1e-15: 

128 self.normalized = True 

129 yield None 

130 yield None 

131 yield None 

132 return 

133 

134 I_M = np.zeros(self.Mmax) 

135 

136 nw = len(self.A_wgm) // len(self.spline_j) 

137 assert nw * len(self.spline_j) == len(self.A_wgm) 

138 

139 for M, A_gm in zip(self.M_w, self.A_wgm): 

140 I_m = A_gm.sum(axis=0) 

141 I_M[M:M + len(I_m)] += I_m * dv 

142 

143 requests = [] 

144 if len(self.ranks) > 0: 

145 I_rM = np.empty((len(self.ranks), self.Mmax)) 

146 for r, J_M in zip(self.ranks, I_rM): 

147 requests.append(comm.receive(J_M, r, a, False)) 

148 if self.rank != comm.rank: 

149 requests.append(comm.send(I_M, self.rank, a, False)) 

150 

151 yield None 

152 

153 for request in requests: 

154 comm.wait(request) 

155 

156 requests = [] 

157 if len(self.ranks) > 0: 

158 I_M += I_rM.sum(axis=0) 

159 for r in self.ranks: 

160 requests.append(comm.send(I_M, r, a, False)) 

161 if self.rank != comm.rank: 

162 requests.append(comm.receive(I_M, self.rank, a, False)) 

163 

164 yield None 

165 

166 for request in requests: 

167 comm.wait(request) 

168 

169 w = 0 

170 for M, A_gm in zip(self.M_w, self.A_wgm): 

171 if M == 0: 

172 A_gm *= integral / I_M[0] 

173 else: 

174 A_gm -= (I_M[M:M + A_gm.shape[1]] / integral * 

175 self.A_wgm[w % nw]) 

176 w += 1 

177 self.normalized = True 

178 self.I_M = I_M 

179 yield None 

180 

181 def estimate_gridpointcount(self, gd): 

182 points = 0.0 

183 for spline in self.spline_j: 

184 l = spline.get_angular_momentum_number() 

185 rc = spline.get_cutoff() 

186 points += 4.0 * np.pi / 3.0 * rc**3 / gd.dv * (2 * l + 1) 

187 return points 

188 

189 

190# Quick hack: base class to share basic functionality across LFC classes 

191class BaseLFC: 

192 def dict(self, shape=(), derivative=False, zero=False): 

193 if isinstance(shape, int): 

194 shape = (shape,) 

195 if derivative: 

196 assert not zero 

197 c_axiv = {} 

198 for a in self.my_atom_indices: 

199 ni = self.get_function_count(a) 

200 c_axiv[a] = np.empty(shape + (ni, 3), self.dtype) 

201 return c_axiv 

202 else: 

203 c_axi = {} 

204 for a in self.my_atom_indices: 

205 ni = self.get_function_count(a) 

206 c_axi[a] = np.empty(shape + (ni,), self.dtype) 

207 if zero: 

208 c_axi[a][:] = 0.0 

209 return c_axi 

210 

211 def estimate_memory(self, mem): 

212 points = 0 

213 for sphere in self.sphere_a: 

214 points += sphere.estimate_gridpointcount(self.gd) 

215 nbytes = points * mem.floatsize 

216 mem.setsize(nbytes / self.gd.comm.size) # Assume equal distribution 

217 

218 

219class LocalizedFunctionsCollection(BaseLFC): 

220 """LocalizedFunctionsCollection 

221 

222 Utilizes that localized functions can be stored on a spherical subset of 

223 the uniform grid, as opposed to LocalizedFunctionsCollection which is just 

224 a wrapper around the old localized_functions which use rectangular grids. 

225 

226 """ 

227 def __init__(self, gd, spline_aj, kd=None, cut=False, dtype=float, 

228 integral=None, forces=None, xp=np): 

229 self.gd = gd 

230 self.kd = kd 

231 self.sphere_a = [Sphere(spline_j) for spline_j in spline_aj] 

232 self.cut = cut 

233 self.dtype = dtype 

234 self.Mmax = None 

235 self.xp = xp 

236 

237 if kd is None: 

238 self.ibzk_qc = np.zeros((1, 3)) 

239 self.gamma = True 

240 else: 

241 self.ibzk_qc = kd.ibzk_qc 

242 self.gamma = kd.gamma 

243 

244 # Global or local M-indices? 

245 self.use_global_indices = False 

246 

247 if integral is not None: 

248 if isinstance(integral, (float, int)): 

249 self.integral_a = np.empty(len(spline_aj)) 

250 self.integral_a[:] = integral 

251 else: 

252 self.integral_a = np.array(integral) 

253 else: 

254 self.integral_a = None 

255 

256 self.my_atom_indices = None 

257 self.lfc = None 

258 

259 def set_positions(self, spos_ac, atom_partition=None) -> bool: 

260 """Set positions and return True if any atoms have migrated to 

261 another rank. 

262 """ 

263 assert len(spos_ac) == len(self.sphere_a) 

264 spos_ac = np.asarray(spos_ac) 

265 movement = False 

266 old_ranks = [sphere.rank for sphere in self.sphere_a] 

267 for a, (spos_c, sphere) in enumerate(zip(spos_ac, self.sphere_a)): 

268 try: 

269 movement |= sphere.set_position(spos_c, self.gd, self.cut) 

270 except GridBoundsError as e: 

271 e.args = (f'Atom {a} too close to edge: {e}',) 

272 raise 

273 

274 if self.my_atom_indices is None: 

275 self._update(spos_ac) 

276 return False 

277 

278 if movement: 

279 self._update(spos_ac) 

280 for rank, sphere in zip(old_ranks, self.sphere_a): 

281 if rank != sphere.rank: 

282 return True 

283 return False 

284 

285 def _update(self, spos_ac): 

286 nB = 0 

287 nW = 0 

288 self.my_atom_indices = [] 

289 self.atom_indices = [] 

290 M = 0 

291 self.M_a = [] 

292 for a, sphere in enumerate(self.sphere_a): 

293 self.M_a.append(M) 

294 if sphere.rank == self.gd.comm.rank: 

295 self.my_atom_indices.append(a) 

296 G_wb = sphere.G_wb 

297 if G_wb: 

298 nB += sum([len(G_b) for G_b in G_wb]) 

299 nW += len(G_wb) 

300 self.atom_indices.append(a) 

301 

302 if not self.use_global_indices: 

303 M += sphere.Mmax 

304 

305 if self.use_global_indices: 

306 M += sphere.Mmax 

307 

308 self.Mmax = M 

309 

310 natoms = len(spos_ac) 

311 # Holm-Nielsen check: 

312 if (self.gd.comm.sum_scalar(float(sum(self.my_atom_indices))) != 

313 natoms * (natoms - 1) // 2): 

314 raise ValueError('Holm-Nielsen check failed. Grid might be ' 

315 'too coarse. Use h < %.3f' 

316 % (smallest_safe_grid_spacing * Bohr)) 

317 

318 self.M_W = np.empty(nW, np.intc) 

319 self.G_B = np.empty(nB, np.intc) 

320 self.W_B = np.empty(nB, np.intc) 

321 self.A_Wgm = [] 

322 sdisp_Wc = np.empty((nW, 3), int) 

323 self.pos_Wv = np.empty((nW, 3)) 

324 

325 B1 = 0 

326 W = 0 

327 for a in self.atom_indices: 

328 sphere = self.sphere_a[a] 

329 self.A_Wgm.extend(sphere.A_wgm) 

330 nw = len(sphere.M_w) 

331 self.M_W[W:W + nw] = self.M_a[a] + np.array(sphere.M_w) 

332 sdisp_Wc[W:W + nw] = sphere.sdisp_wc 

333 self.pos_Wv[W:W + nw] = np.dot(spos_ac[a] - 

334 np.array(sphere.sdisp_wc), 

335 self.gd.cell_cv) 

336 for G_b in sphere.G_wb: 

337 B2 = B1 + len(G_b) 

338 self.G_B[B1:B2] = G_b 

339 self.W_B[B1:B2:2] = W 

340 self.W_B[B1 + 1:B2 + 1:2] = -W - 1 

341 B1 = B2 

342 W += 1 

343 assert B1 == nB 

344 

345 if self.gamma: 

346 if self.dtype == float: 

347 self.phase_qW = np.empty((0, nW), complex) 

348 else: 

349 # TDDFT calculation: 

350 self.phase_qW = np.ones((1, nW), complex) 

351 else: 

352 self.phase_qW = np.exp(2j * pi * np.inner(self.ibzk_qc, sdisp_Wc)) 

353 

354 indices = np.argsort(self.G_B, kind='mergesort') 

355 self.G_B = self.G_B[indices] 

356 self.W_B = self.W_B[indices] 

357 

358 # Find out which ranks have a piece of the 

359 # localized functions: 

360 x_a = np.zeros(natoms, bool) 

361 x_a[self.atom_indices] = True 

362 x_a[self.my_atom_indices] = False 

363 x_ra = np.empty((self.gd.comm.size, natoms), bool) 

364 self.gd.comm.all_gather(x_a, x_ra) 

365 for a in self.atom_indices: 

366 sphere = self.sphere_a[a] 

367 if sphere.rank == self.gd.comm.rank: 

368 sphere.ranks = x_ra[:, a].nonzero()[0] 

369 else: 

370 sphere.ranks = [] 

371 

372 if self.integral_a is not None: 

373 iterators = [] 

374 for a in self.atom_indices: 

375 iterator = self.sphere_a[a].normalize(self.integral_a[a], a, 

376 self.gd.dv, 

377 self.gd.comm) 

378 iterators.append(iterator) 

379 for i in range(3): 

380 for iterator in iterators: 

381 next(iterator) 

382 

383 self.lfc = cgpaw.LFC(self.A_Wgm, self.M_W, self.G_B, self.W_B, 

384 self.gd.dv, self.phase_qW, self.xp is not np) 

385 

386 return sdisp_Wc 

387 

388 def M_to_ai(self, src_xM, dst_axi): 

389 xshape = src_xM.shape[:-1] 

390 src_xM = src_xM.reshape(np.prod(xshape), self.Mmax) 

391 for a in self.my_atom_indices: 

392 M1 = self.M_a[a] 

393 M2 = M1 + self.sphere_a[a].Mmax 

394 dst_axi[a] = src_xM[:, M1:M2].copy() 

395 

396 def ai_to_M(self, src_axi, dst_xM): 

397 xshape = dst_xM.shape[:-1] 

398 dst_xM = dst_xM.reshape(np.prod(xshape), self.Mmax) 

399 for a in self.my_atom_indices: 

400 M1 = self.M_a[a] 

401 M2 = M1 + self.sphere_a[a].Mmax 

402 dst_xM[:, M1:M2] = src_axi[a] 

403 

404 def add(self, a_xG, c_axi=1.0, q=-1): 

405 """Add localized functions to extended arrays. 

406 

407 :: 

408 

409 -- a a 

410 a (G) += > c Phi (G) 

411 x -- xi i 

412 a,i 

413 """ 

414 

415 assert not self.use_global_indices 

416 if q == -1: 

417 assert self.dtype == float 

418 

419 if isinstance(c_axi, float): 

420 assert q == -1 and a_xG.ndim == 3 

421 c_xM = self.xp.empty(self.Mmax) 

422 c_xM.fill(c_axi) 

423 if self.xp is np: 

424 self.lfc.add(c_xM, a_xG, q) 

425 elif cupy_is_fake: 

426 self.lfc.add(c_xM._data, a_xG._data, q) 

427 else: 

428 self.lfc.add_gpu(c_xM.data.ptr, 

429 c_xM.shape, 

430 a_xG.data.ptr, 

431 a_xG.shape, q) 

432 return 

433 

434 dtype = a_xG.dtype 

435 

436 if debug: 

437 assert a_xG.ndim >= 3 

438 assert sorted(c_axi.keys()) == self.my_atom_indices 

439 for c_xi in c_axi.values(): 

440 assert c_xi.dtype == dtype 

441 

442 comm = self.gd.comm 

443 xshape = a_xG.shape[:-3] 

444 assert len(xshape) <= 1 

445 requests = [] 

446 M1 = 0 

447 b_axi = {} 

448 for a in self.atom_indices: 

449 c_xi = c_axi.get(a) 

450 sphere = self.sphere_a[a] 

451 M2 = M1 + sphere.Mmax 

452 if c_xi is None: 

453 c_xi = self.xp.empty(xshape + (sphere.Mmax,), dtype) 

454 b_axi[a] = c_xi 

455 requests.append(comm.receive(c_xi, sphere.rank, a, False)) 

456 else: 

457 for r in sphere.ranks: 

458 requests.append(comm.send(c_xi.copy(), r, a, False)) 

459 

460 M1 = M2 

461 

462 for request in requests: 

463 comm.wait(request) 

464 

465 c_xM = self.xp.empty(xshape + (self.Mmax,), dtype) 

466 M1 = 0 

467 for a in self.atom_indices: 

468 c_xi = c_axi.get(a) 

469 sphere = self.sphere_a[a] 

470 M2 = M1 + sphere.Mmax 

471 if c_xi is None: 

472 c_xi = b_axi[a] 

473 c_xM[..., M1:M2] = c_xi 

474 M1 = M2 

475 

476 if self.xp is np: 

477 self.lfc.add(c_xM, a_xG, q) 

478 elif cupy_is_fake: 

479 self.lfc.add(c_xM._data, a_xG._data, q) 

480 else: 

481 self.lfc.add_gpu(c_xM.data.ptr, 

482 c_xM.shape, 

483 a_xG.data.ptr, 

484 a_xG.shape, q) 

485 

486 def add_derivative(self, a, v, a_xG, c_axi=1.0, q=-1): 

487 """Add derivative of localized functions on atom to extended arrays. 

488 

489 Parameters: 

490 

491 a: int 

492 Atomic index of the derivative 

493 v: int 

494 Cartesian coordinate of the derivative (0, 1 or 2) 

495 

496 This function adds the following sum to the extended arrays:: 

497 

498 -- a a 

499 a (G) += > c dPhi (G) 

500 x -- xi iv 

501 i 

502 

503 where:: 

504 

505 a d a 

506 dPhi (G) = -- Phi (g) 

507 iv dv i 

508 

509 is the derivative of the Phi^a and v is either x, y, or z. 

510 

511 """ 

512 

513 assert v in [0, 1, 2] 

514 assert not self.use_global_indices 

515 

516 if q == -1: 

517 assert self.dtype == float 

518 

519 if isinstance(c_axi, float): 

520 assert q == -1 

521 c_xM = np.empty(self.Mmax) 

522 c_xM.fill(c_axi) 

523 self.lfc.add(c_xM, a_xG, q) 

524 return 

525 

526 dtype = a_xG.dtype 

527 

528 if debug: 

529 assert a_xG.ndim >= 3 

530 assert dtype == self.dtype 

531 if isinstance(c_axi, dict): 

532 assert sorted(c_axi.keys()) == self.my_atom_indices 

533 for c_xi in c_axi.values(): 

534 assert c_xi.dtype == dtype 

535 

536 cspline_M = [] 

537 for a_ in self.atom_indices: 

538 for spline in self.sphere_a[a_].spline_j: 

539 nm = 2 * spline.get_angular_momentum_number() + 1 

540 cspline_M.extend([spline.spline] * nm) 

541 

542 # Temp solution - set all coefficient to zero except for those at 

543 # atom a 

544 

545 # Coefficient indices for atom a 

546 M1 = self.M_a[a] 

547 M2 = M1 + self.sphere_a[a].Mmax 

548 

549 if isinstance(c_axi, float): 

550 assert q == -1 

551 c_xM = np.zeros(self.Mmax) 

552 c_xM[..., M1:M2] = c_axi 

553 else: 

554 xshape = a_xG.shape[:-3] 

555 c_xM = np.zeros(xshape + (self.Mmax,), dtype) 

556 c_xM[..., M1:M2] = c_axi[a] 

557 

558 gd = self.gd 

559 

560 self.lfc.add_derivative(c_xM, a_xG, np.ascontiguousarray(gd.h_cv), 

561 gd.n_c, cspline_M, 

562 gd.beg_c, self.pos_Wv, a, v, q) 

563 

564 def integrate(self, a_xG, c_axi, q=-1, add_to=False): 

565 """Calculate integrals of arrays times localized functions. 

566 

567 :: 

568 

569 / a* 

570 c_axi = | dG a (G) Phi (G) 

571 / x i 

572 

573 """ 

574 assert not add_to 

575 assert not self.use_global_indices 

576 if q == -1: 

577 assert self.dtype == float 

578 

579 xshape = a_xG.shape[:-3] 

580 dtype = a_xG.dtype 

581 

582 if debug: 

583 assert self.dtype == dtype 

584 assert a_xG.ndim >= 3 

585 assert sorted(c_axi.keys()) == self.my_atom_indices 

586 for c_xi in c_axi.values(): 

587 assert c_xi.shape[:-1] == xshape 

588 

589 comm = self.gd.comm 

590 

591 c_xM = self.xp.zeros(xshape + (self.Mmax,), dtype) 

592 if self.xp is np: 

593 self.lfc.integrate(a_xG, c_xM, q) 

594 elif cupy_is_fake: 

595 self.lfc.integrate(a_xG._data, c_xM._data, q) 

596 else: 

597 self.lfc.integrate_gpu(a_xG.data.ptr, 

598 a_xG.shape, 

599 c_xM.data.ptr, 

600 c_xM.shape, q) 

601 c_xM *= self.gd.dv 

602 

603 rank = comm.rank 

604 srequests = [] 

605 rrequests = [] 

606 c_arxi = {} 

607 b_axi = {} 

608 M1 = 0 

609 for a in self.atom_indices: 

610 sphere = self.sphere_a[a] 

611 M2 = M1 + sphere.Mmax 

612 if sphere.rank != rank: 

613 c_xi = c_xM[..., M1:M2].copy() 

614 b_axi[a] = c_xi 

615 srequests.append(comm.send(c_xi, 

616 sphere.rank, a, False)) 

617 else: 

618 if len(sphere.ranks) > 0: 

619 c_rxi = self.xp.empty( 

620 sphere.ranks.shape + xshape + (M2 - M1,), 

621 dtype) 

622 c_arxi[a] = c_rxi 

623 for r, b_xi in zip(sphere.ranks, c_rxi): 

624 rrequests.append(comm.receive(b_xi, r, a, False)) 

625 M1 = M2 

626 

627 for request in rrequests: 

628 comm.wait(request) 

629 

630 M1 = 0 

631 for a in self.atom_indices: 

632 c_xi = c_axi.get(a) 

633 sphere = self.sphere_a[a] 

634 M2 = M1 + sphere.Mmax 

635 if c_xi is not None: 

636 if len(sphere.ranks) > 0: 

637 if c_xM.shape[-1] > M1: 

638 c_xi[:] = c_xM[..., M1:M2] + c_arxi[a].sum(axis=0) 

639 elif c_xM.shape[-1] > M1: 

640 c_xi[:] = c_xM[..., M1:M2] 

641 M1 = M2 

642 

643 for request in srequests: 

644 comm.wait(request) 

645 

646 def derivative(self, a_xG, c_axiv, q=-1): 

647 """Calculate x-, y-, and z-derivatives of localized function integrals. 

648 

649 :: 

650 

651 / a* 

652 c_axiv = | dG a (G) dPhi (G) 

653 / x iv 

654 

655 where:: 

656 

657 

658 a d a 

659 dPhi (G) = -- Phi (g) 

660 iv dv i 

661 

662 

663 and v is either x, y, or z, and R^a_v is the center of Phi^a. 

664 

665 Notice that d Phi^a_i / dR^a_v == - d Phi^a_i / d v. 

666 

667 """ 

668 

669 assert not self.use_global_indices 

670 

671 if debug: 

672 assert a_xG.ndim >= 3 

673 assert sorted(c_axiv.keys()) == self.my_atom_indices 

674 

675 if self.integral_a is not None: 

676 # assert q == -1 

677 assert a_xG.ndim == 3 

678 assert a_xG.dtype == float 

679 self._normalized_derivative(a_xG, c_axiv) 

680 return 

681 

682 dtype = a_xG.dtype 

683 

684 xshape = a_xG.shape[:-3] 

685 c_xMv = np.zeros(xshape + (self.Mmax, 3), dtype) 

686 cspline_M = [] 

687 for a in self.atom_indices: 

688 for spline in self.sphere_a[a].spline_j: 

689 nm = 2 * spline.get_angular_momentum_number() + 1 

690 cspline_M.extend([spline.spline] * nm) 

691 

692 gd = self.gd 

693 self.lfc.derivative(a_xG, c_xMv, np.ascontiguousarray(gd.h_cv), 

694 gd.n_c, cspline_M, 

695 gd.beg_c, self.pos_Wv, q) 

696 

697 comm = self.gd.comm 

698 rank = comm.rank 

699 srequests = [] 

700 rrequests = [] 

701 c_arxiv = {} # see also https://arXiv.org 

702 b_axiv = {} 

703 M1 = 0 

704 for a in self.atom_indices: 

705 sphere = self.sphere_a[a] 

706 M2 = M1 + sphere.Mmax 

707 if sphere.rank != rank: 

708 c_xiv = c_xMv[..., M1:M2, :].copy() 

709 b_axiv[a] = c_xiv 

710 srequests.append(comm.send(c_xiv, 

711 sphere.rank, a, False)) 

712 else: 

713 if len(sphere.ranks) > 0: 

714 c_rxiv = np.empty(sphere.ranks.shape + xshape + 

715 (M2 - M1, 3), dtype) 

716 c_arxiv[a] = c_rxiv 

717 for r, b_xiv in zip(sphere.ranks, c_rxiv): 

718 rrequests.append(comm.receive(b_xiv, r, a, False)) 

719 M1 = M2 

720 

721 for request in rrequests: 

722 comm.wait(request) 

723 

724 M1 = 0 

725 for a in self.atom_indices: 

726 c_xiv = c_axiv.get(a) 

727 sphere = self.sphere_a[a] 

728 M2 = M1 + sphere.Mmax 

729 if c_xiv is not None: 

730 if len(sphere.ranks) > 0: 

731 c_xiv[:] = c_xMv[..., M1:M2, :] + c_arxiv[a].sum(axis=0) 

732 else: 

733 c_xiv[:] = c_xMv[..., M1:M2, :] 

734 M1 = M2 

735 

736 for request in srequests: 

737 comm.wait(request) 

738 

739 def _normalized_derivative(self, a_G, c_aiv): 

740 """Calculate x-, y-, and z-derivatives of localized function integrals. 

741 

742 Calculates the derivatives of this integral:: 

743 

744 a / _ _ a - _a 

745 A = | dr a(r) f (r - R ), 

746 lm / lm 

747