Coverage for gpaw/wavefunctions/base.py: 79%

1import numpy as np 

2from ase.units import Ha 

3 

4from gpaw.projections import Projections 

5from gpaw.utilities import pack_density 

6from gpaw.utilities.blas import axpy, mmm 

7from gpaw.utilities.partition import AtomPartition 

8 

9 

10class WaveFunctions: 

11 """... 

12 

13 setups: 

14 List of setup objects. 

15 symmetry: 

16 Symmetry object. 

17 kpt_u: 

18 List of **k**-point objects. 

19 nbands: int 

20 Number of bands. 

21 nspins: int 

22 Number of spins. 

23 dtype: dtype 

24 Data type of wave functions (float or complex). 

25 bzk_kc: ndarray 

26 Scaled **k**-points used for sampling the whole 

27 Brillouin zone - values scaled to [-0.5, 0.5). 

28 ibzk_kc: ndarray 

29 Scaled **k**-points in the irreducible part of the 

30 Brillouin zone. 

31 weight_k: ndarray 

32 Weights of the **k**-points in the irreducible part 

33 of the Brillouin zone (summing up to 1). 

34 kpt_comm: 

35 MPI-communicator for parallelization over **k**-points. 

36 """ 

37 

38 def __init__(self, gd, nvalence, setups, bd, dtype, collinear, 

39 world, kd, kptband_comm, timer): 

40 self.gd = gd 

41 self.nspins = kd.nspins 

42 self.collinear = collinear 

43 self.nvalence = nvalence 

44 self.bd = bd 

45 self.dtype = dtype 

46 assert dtype == float or dtype == complex 

47 self.world = world 

48 self.kd = kd 

49 self.kptband_comm = kptband_comm 

50 self.timer = timer 

51 self.atom_partition = None 

52 

53 self.kpt_qs = kd.create_k_points(self.gd.sdisp_cd, collinear) 

54 self.kpt_u = [kpt for kpt_s in self.kpt_qs for kpt in kpt_s] 

55 

56 self.occupations = None 

57 self.fermi_levels = None 

58 

59 self.eigensolver = None 

60 self.positions_set = False 

61 self.spos_ac = None 

62 

63 self.set_setups(setups) 

64 

65 @property 

66 def fermi_level(self): 

67 assert len(self.fermi_levels) == 1 

68 return self.fermi_levels[0] 

69 

70 def summary(self, log): 

71 log(eigenvalue_string(self)) 

72 

73 func = None 

74 if hasattr(self.eigensolver, 'dm_helper'): 

75 func = getattr(self.eigensolver.dm_helper.func, 'name', None) 

76 elif hasattr(self.eigensolver, 'odd'): 

77 func = getattr(self.eigensolver.odd, 'name', None) 

78 if func is None: 

79 pass 

80 elif 'SIC' in func: 

81 self.summary_func(log) 

82 

83 if self.fermi_levels is None: 

84 return 

85 

86 if len(self.fermi_levels) == 1: 

87 log(f'Fermi level: {self.fermi_levels[0] * Ha:.5f}\n') 

88 else: 

89 f1, f2 = (f * Ha for f in self.fermi_levels) 

90 log(f'Fermi levels: {f1:.5f}, {f2:.5f}\n') 

91 

92 def set_setups(self, setups): 

93 self.setups = setups 

94 

95 def set_eigensolver(self, eigensolver): 

96 self.eigensolver = eigensolver 

97 

98 def add_realspace_orbital_to_density(self, nt_G, psit_G): 

99 if psit_G.dtype == float: 

100 axpy(1.0, psit_G**2, nt_G) 

101 else: 

102 assert psit_G.dtype == complex 

103 axpy(1.0, psit_G.real**2, nt_G) 

104 axpy(1.0, psit_G.imag**2, nt_G) 

105 

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

107 self.add_realspace_orbital_to_density(nt_G, kpt.psit_nG[n]) 

108 

109 def calculate_band_energy(self): 

110 e_band = 0.0 

111 for kpt in self.kpt_u: 

112 e_band += np.dot(kpt.f_n, kpt.eps_n) 

113 

114 try: # DCSF needs this ... 

115 e_band += self.occupations.calculate_band_energy(self) 

116 except AttributeError: 

117 pass 

118 

119 return self.kptband_comm.sum_scalar(e_band) 

120 

121 def calculate_density_contribution(self, nt_sG): 

122 """Calculate contribution to pseudo density from wave functions. 

123 

124 Array entries are written to (not added to).""" 

125 nt_sG.fill(0.0) 

126 for kpt in self.kpt_u: 

127 self.add_to_density_from_k_point(nt_sG, kpt) 

128 self.kptband_comm.sum(nt_sG) 

129 

130 self.timer.start('Symmetrize density') 

131 for nt_G in nt_sG: 

132 self.kd.symmetry.symmetrize(nt_G, self.gd) 

133 self.timer.stop('Symmetrize density') 

134 

135 def add_to_density_from_k_point(self, nt_sG, kpt): 

136 self.add_to_density_from_k_point_with_occupation(nt_sG, kpt, kpt.f_n) 

137 

138 def get_orbital_density_matrix(self, a, kpt, n): 

139 """Add the nth band density from kpt to density matrix D_sp""" 

140 ni = self.setups[a].ni 

141 D_sii = np.zeros((self.nspins, ni, ni)) 

142 P_i = kpt.P_ani[a][n] 

143 D_sii[kpt.s] += np.outer(P_i.conj(), P_i).real 

144 D_sp = [pack_density(D_ii) for D_ii in D_sii] 

145 return D_sp 

146 

147 def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n): 

148 if kpt.rho_MM is not None: 

149 P_Mi = self.P_aqMi[a][kpt.q] 

150 rhoP_Mi = np.zeros_like(P_Mi) 

151 D_ii = np.zeros(D_sii[kpt.s].shape, kpt.rho_MM.dtype) 

152 mmm(1.0, kpt.rho_MM, 'N', P_Mi, 'N', 0.0, rhoP_Mi) 

153 mmm(1.0, P_Mi, 'C', rhoP_Mi, 'N', 0.0, D_ii) 

154 D_sii[kpt.s] += D_ii.real 

155 else: 

156 if self.collinear: 

157 P_ni = kpt.projections[a] 

158 D_sii[kpt.s] += np.dot(P_ni.T.conj() * f_n, P_ni).real 

159 else: 

160 P_nsi = kpt.projections[a] 

161 D_ssii = np.einsum('nsi,n,nzj->szij', 

162 P_nsi.conj(), f_n, P_nsi) 

163 D_sii[0] += (D_ssii[0, 0] + D_ssii[1, 1]).real 

164 D_sii[1] += 2 * D_ssii[0, 1].real 

165 D_sii[2] += 2 * D_ssii[0, 1].imag 

166 D_sii[3] += (D_ssii[0, 0] - D_ssii[1, 1]).real 

167 

168 if hasattr(kpt, 'c_on'): 

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

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

171 D_sii[kpt.s] += np.dot(P_ni.T.conj(), np.dot(d_nn, P_ni)).real 

172 

173 def calculate_atomic_density_matrices(self, D_asp): 

174 """Calculate atomic density matrices from projections.""" 

175 f_un = [kpt.f_n for kpt in self.kpt_u] 

176 self.calculate_atomic_density_matrices_with_occupation(D_asp, f_un) 

177 

178 def calculate_atomic_density_matrices_with_occupation(self, D_asp, f_un): 

179 """Calculate atomic density matrices from projections with 

180 custom occupation f_un.""" 

181 

182 # Parameter check (if user accidentally passes f_n instead of f_un) 

183 if f_un[0] is not None: # special case for transport calculations... 

184 assert isinstance(f_un[0], np.ndarray) 

185 # Varying f_n used in calculation of response part of GLLB-potential 

186 for a, D_sp in D_asp.items(): 

187 ni = self.setups[a].ni 

188 D_sii = np.zeros((len(D_sp), ni, ni)) 

189 for f_n, kpt in zip(f_un, self.kpt_u): 

190 self.calculate_atomic_density_matrices_k_point(D_sii, kpt, a, 

191 f_n) 

192 D_sp[:] = [pack_density(D_ii) for D_ii in D_sii] 

193 self.kptband_comm.sum(D_sp) 

194 

195 self.symmetrize_atomic_density_matrices(D_asp) 

196 

197 def symmetrize_atomic_density_matrices(self, D_asp): 

198 if len(self.kd.symmetry.op_scc) == 0: 

199 return 

200 

201 D_asp.redistribute(self.atom_partition.as_serial()) 

202 self.setups.atomrotations.symmetrize_atomic_density_matrices( 

203 D_asp, a_sa=self.kd.symmetry.a_sa) 

204 D_asp.redistribute(self.atom_partition) 

205 

206 def calculate_occupation_numbers(self, fix_fermi_level=False): 

207 if self.collinear and self.nspins == 1: 

208 degeneracy = 2 

209 else: 

210 degeneracy = 1 

211 

212 f_qn, fermi_levels, e_entropy = self.occupations.calculate( 

213 nelectrons=self.nvalence / degeneracy, 

214 eigenvalues=[kpt.eps_n * Ha for kpt in self.kpt_u], 

215 weights=[kpt.weightk for kpt in self.kpt_u], 

216 fermi_levels_guess=(self.fermi_levels * Ha 

217 if self.fermi_levels is not None else None), 

218 fix_fermi_level=fix_fermi_level) 

219 

220 if not fix_fermi_level or self.fermi_levels is None: 

221 self.fermi_levels = np.array(fermi_levels) / Ha 

222 

223 for f_n, kpt in zip(f_qn, self.kpt_u): 

224 kpt.f_n = f_n * (kpt.weightk * degeneracy) 

225 

226 return e_entropy * degeneracy / Ha 

227 

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

229 self.positions_set = False 

230 # rank_a = self.gd.get_ranks_from_positions(spos_ac) 

231 # atom_partition = AtomPartition(self.gd.comm, rank_a) 

232 # XXX pass AtomPartition around instead of spos_ac? 

233 # All the classes passing around spos_ac end up needing the ranks 

234 # anyway. 

235 

236 if atom_partition is None: 

237 rank_a = self.gd.get_ranks_from_positions(spos_ac) 

238 atom_partition = AtomPartition(self.gd.comm, rank_a) 

239 

240 if self.atom_partition is not None and self.kpt_u[0].P_ani is not None: 

241 with self.timer('Redistribute'): 

242 for kpt in self.kpt_u: 

243 P = kpt.projections 

244 assert self.atom_partition == P.atom_partition 

245 kpt.projections = P.redist(atom_partition) 

246 assert atom_partition == kpt.projections.atom_partition 

247 

248 self.atom_partition = atom_partition 

249 self.kd.symmetry.check(spos_ac) 

250 self.spos_ac = spos_ac 

251 

252 def allocate_arrays_for_projections(self, my_atom_indices): # XXX unused 

253 if not self.positions_set and self.kpt_u[0]._projections is not None: 

254 # Projections have been read from file - don't delete them! 

255 pass 

256 else: 

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

258 for kpt in self.kpt_u: 

259 kpt.projections = Projections( 

260 self.bd.nbands, nproj_a, 

261 self.atom_partition, 

262 self.bd.comm, 

263 collinear=self.collinear, spin=kpt.s, dtype=self.dtype) 

264 

265 def collect_eigenvalues(self, k, s): 

266 return self.collect_array('eps_n', k, s) 

267 

268 def collect_occupations(self, k, s): 

269 return self.collect_array('f_n', k, s) 

270 

271 def collect_array(self, name, k, s, subset=None): 

272 """Helper method for collect_eigenvalues and collect_occupations. 

273 

274 For the parallel case find the rank in kpt_comm that contains 

275 the (k,s) pair, for this rank, collect on the corresponding 

276 domain a full array on the domain master and send this to the 

277 global master.""" 

278 

279 kpt_qs = self.kpt_qs 

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

281 if self.kd.comm.rank == kpt_rank: 

282 a_nx = getattr(kpt_qs[q][s], name) 

283 

284 if subset is not None: 

285 a_nx = a_nx[subset] 

286 

287 # Domain master send this to the global master 

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

289 if self.bd.comm.size == 1: 

290 if kpt_rank == 0: 

291 return a_nx 

292 else: 

293 self.kd.comm.ssend(a_nx, 0, 1301) 

294 else: 

295 b_nx = self.bd.collect(a_nx) 

296 if self.bd.comm.rank == 0: 

297 if kpt_rank == 0: 

298 return b_nx 

299 else: 

300 self.kd.comm.ssend(b_nx, 0, 1301) 

301 

302 elif self.world.rank == 0 and kpt_rank != 0: 

303 # Only used to determine shape and dtype of receiving buffer: 

304 a_nx = getattr(kpt_qs[0][0], name) 

305 

306 if subset is not None: 

307 a_nx = a_nx[subset] 

308 

309 b_nx = np.zeros((self.bd.nbands,) + a_nx.shape[1:], 

310 dtype=a_nx.dtype) 

311 self.kd.comm.receive(b_nx, kpt_rank, 1301) 

312 return b_nx 

313 

314 return np.zeros(0) # see comment in get_wave_function_array() method 

315 

316 def collect_auxiliary(self, value, k, s, shape=1, dtype=float): 

317 """Helper method for collecting band-independent scalars/arrays. 

318 

319 For the parallel case find the rank in kpt_comm that contains 

320 the (k,s) pair, for this rank, collect on the corresponding 

321 domain a full array on the domain master and send this to the 

322 global master.""" 

323 

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

325 

326 if self.kd.comm.rank == kpt_rank: 

327 if isinstance(value, str): 

328 a_o = getattr(self.kpt_qs[q][s], value) 

329 else: 

330 u = q * self.nspins + s 

331 a_o = value[u] # assumed list 

332 

333 # Make sure data is a mutable object 

334 a_o = np.asarray(a_o) 

335 

336 if a_o.dtype is not dtype: 

337 a_o = a_o.astype(dtype) 

338 

339 # Domain master send this to the global master 

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

341 if kpt_rank == 0: 

342 return a_o 

343 else: 

344 self.kd.comm.send(a_o, 0, 1302) 

345 

346 elif self.world.rank == 0 and kpt_rank != 0: 

347 b_o = np.zeros(shape, dtype=dtype) 

348 self.kd.comm.receive(b_o, kpt_rank, 1302) 

349 return b_o 

350 

351 def collect_projections(self, k, s): 

352 """Helper method for collecting projector overlaps across domains. 

353 

354 For the parallel case find the rank in kpt_comm that contains 

355 the (k,s) pair, for this rank, send to the global master.""" 

356 

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

358 

359 if self.kd.comm.rank == kpt_rank: 

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

361 P_nI = kpt.projections.collect() 

362 if self.world.rank == 0: 

363 return P_nI 

364 if P_nI is not None: 

365 self.kd.comm.send(np.ascontiguousarray(P_nI), 0, tag=117) 

366 if self.world.rank == 0: 

367 nproj = sum(setup.ni for setup in self.setups) 

368 if not self.collinear: 

369 nproj *= 2 

370 P_nI = np.empty((self.bd.nbands, nproj), self.dtype) 

371 self.kd.comm.receive(P_nI, kpt_rank, tag=117) 

372 return P_nI 

373 

374 def get_wave_function_array(self, n, k, s, realspace=True, periodic=False): 

375 """Return pseudo-wave-function array on master. 

376 

377 n: int 

378 Global band index. 

379 k: int 

380 Global IBZ k-point index. 

381 s: int 

382 Spin index (0 or 1). 

383 realspace: bool 

384 Transform plane wave or LCAO expansion coefficients to real-space. 

385 

386 For the parallel case find the ranks in kd.comm and bd.comm 

387 that contains to (n, k, s), and collect on the corresponding 

388 domain a full array on the domain master and send this to the 

389 global master.""" 

390 

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

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

393 

394 rank = self.world.rank 

395 

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

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

398 u = q * self.nspins + s 

399 psit_G = self._get_wave_function_array(u, myn, 

400 realspace, periodic) 

401 

402 if realspace: 

403 psit_G = self.gd.collect(psit_G) 

404 

405 if rank == 0: 

406 return psit_G 

407 

408 # Domain master send this to the global master 

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

410 psit_G = np.ascontiguousarray(psit_G) 

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

412 

413 if rank == 0: 

414 # allocate full wave function and receive 

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

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

417 realspace=realspace) 

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

419 # of communicators. 

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

421 self.bd.comm.size + 

422 band_rank * self.gd.comm.size) 

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

424 return psit_G 

425 

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

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

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

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

430 # and there. 

431 return np.nan 

432 

433 def get_homo_lumo(self, spin=None, _gllb=False): 

434 """Return HOMO and LUMO eigenvalues.""" 

435 if spin is None: 

436 if self.nspins == 1: 

437 return self.get_homo_lumo(0) 

438 h0, l0 = self.get_homo_lumo(0) 

439 h1, l1 = self.get_homo_lumo(1) 

440 return np.array([max(h0, h1), min(l0, l1)]) 

441 

442 if _gllb: 

443 # Backwards compatibility (see test/gllb/test_metallic.py test) 

444 n = self.nvalence // 2 

445 else: 

446 nocc = 0.0 

447 for kpt in self.kpt_u: 

448 if kpt.s == spin: 

449 nocc += kpt.f_n.sum() 

450 nocc = self.kptband_comm.sum_scalar(nocc) * self.nspins / 2 

451 n = int(round(nocc)) 

452 

453 band_rank, myn = self.bd.who_has(n - 1) 

454 homo = -np.inf 

455 if self.bd.comm.rank == band_rank: 

456 for kpt in self.kpt_u: 

457 if kpt.s == spin: 

458 homo = max(kpt.eps_n[myn], homo) 

459 homo = self.world.max_scalar(homo) 

460 

461 lumo = np.inf 

462 if n < self.bd.nbands: # there are not enough bands for LUMO 

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

464 if self.bd.comm.rank == band_rank: 

465 for kpt in self.kpt_u: 

466 if kpt.s == spin: 

467 lumo = min(kpt.eps_n[myn], lumo) 

468 lumo = self.world.min_scalar(lumo) 

469 

470 return np.array([homo, lumo]) 

471 

472 def write(self, writer): 

473 writer.write(version=2, ha=Ha) 

474 writer.write(fermi_levels=self.fermi_levels * Ha) 

475 writer.write(kpts=self.kd) 

476 self.write_projections(writer) 

477 self.write_eigenvalues(writer) 

478 self.write_occupations(writer) 

479 

480 def write_projections(self, writer): 

481 nproj = sum(setup.ni for setup in self.setups) 

482 

483 if self.collinear: 

484 shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands, nproj) 

485 else: 

486 shape = (self.kd.nibzkpts, self.bd.nbands, 2, nproj) 

487 

488 writer.add_array('projections', shape, self.dtype) 

489 

490 for s in range(self.nspins): 

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

492 P_nI = self.collect_projections(k, s) 

493 if not self.collinear and P_nI is not None: 

494 P_nI.shape = (self.bd.nbands, 2, nproj) 

495 writer.fill(P_nI) 

496 

497 def write_eigenvalues(self, writer): 

498 if self.collinear: 

499 shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands) 

500 else: 

501 shape = (self.kd.nibzkpts, self.bd.nbands) 

502 

503 writer.add_array('eigenvalues', shape) 

504 for s in range(self.nspins): 

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

506 writer.fill(self.collect_eigenvalues(k, s) * Ha) 

507 

508 def write_occupations(self, writer): 

509 

510 if self.collinear: 

511 shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands) 

512 else: 

513 shape = (self.kd.nibzkpts, self.bd.nbands) 

514 

515 writer.add_array('occupations', shape) 

516 for s in range(self.nspins): 

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

518 # Scale occupation numbers when writing: 

519 # XXX fix this in the code also ... 

520 weight = self.kd.weight_k[k] * 2 / self.nspins 

521 writer.fill(self.collect_occupations(k, s) / weight) 

522 

523 def read(self, reader): 

524 wfs_reader = reader.wave_functions 

525 # Backward compatibility: 

526 # Take parameters from main reader 

527 if 'ha' not in wfs_reader: 

528 wfs_reader.ha = reader.ha 

529 if 'version' not in wfs_reader: 

530 wfs_reader.version = reader.version 

531 

532 if reader.version >= 3: 

533 self.fermi_levels = wfs_reader.fermi_levels / wfs_reader.ha 

534 else: 

535 o = reader.occupations 

536 self.fermi_levels = np.array( 

537 [o.fermilevel + o.split / 2, 

538 o.fermilevel - o.split / 2]) / wfs_reader.ha 

539 if self.occupations.name != 'fixmagmom': 

540 assert o.split == 0.0 

541 self.fermi_levels = self.fermi_levels[:1] 

542 

543 if reader.version >= 2: 

544 kpts = wfs_reader.kpts 

545 assert np.allclose(kpts.ibzkpts, self.kd.ibzk_kc) 

546 assert np.allclose(kpts.bzkpts, self.kd.bzk_kc) 

547 assert (kpts.bz2ibz == self.kd.bz2ibz_k).all() 

548 assert np.allclose(kpts.weights, self.kd.weight_k) 

549 

550 if 'projections' in wfs_reader: 

551 self.read_projections(wfs_reader) 

552 self.read_eigenvalues(wfs_reader, wfs_reader.version <= 0) 

553 self.read_occupations(wfs_reader, wfs_reader.version <= 0) 

554 

555 def read_projections(self, reader): 

556 nslice = self.bd.get_slice() 

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

558 atom_partition = AtomPartition(self.gd.comm, 

559 np.zeros(len(nproj_a), int)) 

560 for u, kpt in enumerate(self.kpt_u): 

561 if self.collinear: 

562 index = (kpt.s, kpt.k) 

563 else: 

564 index = (kpt.k,) 

565 kpt.projections = Projections( 

566 self.bd.nbands, nproj_a, 

567 atom_partition, self.bd.comm, 

568 collinear=self.collinear, spin=kpt.s, dtype=self.dtype) 

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

570 P_nI = reader.proxy('projections', *index)[nslice] 

571 if not self.collinear: 

572 P_nI.shape = (self.bd.mynbands, -1) 

573 kpt.projections.matrix.array[:] = P_nI 

574 

575 def read_eigenvalues(self, reader, old=False): 

576 nslice = self.bd.get_slice() 

577 for u, kpt in enumerate(self.kpt_u): 

578 if self.collinear: 

579 index = (kpt.s, kpt.k) 

580 else: 

581 index = (kpt.k,) 

582 eps_n = reader.proxy('eigenvalues', *index)[nslice] 

583 x = self.bd.mynbands - len(eps_n) # missing bands? 

584 if x > 0: 

585 # Working on a real fix to this parallelization problem ... 

586 eps_n = np.pad(eps_n, (0, x), 'constant') 

587 if not old: # skip for old tar-files gpw's 

588 eps_n /= reader.ha 

589 kpt.eps_n = eps_n 

590 

591 def read_occupations(self, reader, old=False): 

592 nslice = self.bd.get_slice() 

593 for u, kpt in enumerate(self.kpt_u): 

594 if self.collinear: 

595 index = (kpt.s, kpt.k) 

596 else: 

597 index = (kpt.k,) 

598 f_n = reader.proxy('occupations', *index)[nslice] 

599 x = self.bd.mynbands - len(f_n) # missing bands? 

600 if x > 0: 

601 # Working on a real fix to this parallelization problem ... 

602 f_n = np.pad(f_n, (0, x), 'constant') 

603 if not old: # skip for old tar-files gpw's 

604 f_n *= kpt.weight 

605 kpt.f_n = f_n 

606 

607 def summary_func(self, log): 

608 

609 pot = None 

610 if hasattr(self.eigensolver, 'dm_helper'): 

611 pot = self.eigensolver.dm_helper.func 

612 elif hasattr(self.eigensolver, 'odd'): 

613 pot = self.eigensolver.odd 

614 

615 f_sn = {} 

616 for kpt in self.kpt_u: 

617 u = kpt.s * self.kd.nibzkpts + kpt.q 

618 f_sn[u] = kpt.f_n 

619 

620 log("Diagonal elements of Lagrange matrix:") 

621 if self.nspins == 1: 

622 header = " Band L_ii " \ 

623 "Occupancy" 

624 log(header) 

625 

626 lagr = pot.lagr_diag_s[0] 

627 for i, x in enumerate(lagr): 

628 log('%5d %11.5f %9.5f' % ( 

629 i, Ha * x, f_sn[0][i])) 

630 

631 if self.nspins == 2: 

632 if self.kd.comm.size > 1: 

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

634 # occupation numbers 

635 size = np.array([0]) 

636 self.kd.comm.receive(size, 1) 

637 f_2n = np.zeros(shape=(int(size[0]))) 

638 self.kd.comm.receive(f_2n, 1) 

639 f_sn[1] = f_2n 

640 

641 # orbital energies 

642 size = np.array([0]) 

643 self.kd.comm.receive(size, 1) 

644 lagr_1 = np.zeros(shape=(int(size[0]))) 

645 self.kd.comm.receive(lagr_1, 1) 

646 

647 else: 

648 # occupations 

649 size = np.array([f_sn[1].shape[0]]) 

650 self.kd.comm.send(size, 0) 

651 self.kd.comm.send(f_sn[1], 0) 

652 

653 # orbital energies 

654 a = pot.lagr_diag_s[1].copy() 

655 size = np.array([a.shape[0]]) 

656 self.kd.comm.send(size, 0) 

657 self.kd.comm.send(a, 0) 

658 

659 del a 

660 

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

662 log(' Up ' 

663 ' Down') 

664 log(' Band L_ii Occupancy ' 

665 ' Band L_ii Occupancy') 

666 

667 lagr_0 = np.sort(pot.lagr_diag_s[0]) 

668 lagr_labeled_0 = {} 

669 for c, x in enumerate(pot.lagr_diag_s[0]): 

670 lagr_labeled_0[str(round(x, 12))] = c 

671 

672 if self.kd.comm.size == 1: 

673 lagr_1 = np.sort(pot.lagr_diag_s[1]) 

674 lagr_labeled_1 = {} 

675 for c, x in enumerate( 

676 pot.lagr_diag_s[1]): 

677 lagr_labeled_1[str(round(x, 12))] = c 

678 else: 

679 lagr_labeled_1 = {} 

680 for c, x in enumerate(lagr_1): 

681 lagr_labeled_1[str(round(x, 12))] = c 

682 lagr_1 = np.sort(lagr_1) 

683 

684 for x, y in zip(lagr_0, lagr_1): 

685 i0 = lagr_labeled_0[str(round(x, 12))] 

686 i1 = lagr_labeled_1[str(round(y, 12))] 

687 

688 log('%5d %11.5f %9.5f' 

689 '%5d %11.5f %9.5f' % 

690 (i0, Ha * x, 

691 f_sn[0][i0], 

692 i1, 

693 Ha * y,