Coverage for gpaw/occupations.py: 88%

1"""Smearing functions and occupation number calculators.""" 

2from __future__ import annotations 

3 

4import warnings 

5from math import inf, nan, pi 

6from typing import Any, Callable, Dict, List, NamedTuple, Tuple, cast 

7 

8import numpy as np 

9from ase.units import Ha 

10from scipy.special import erf 

11 

12from gpaw.band_descriptor import BandDescriptor 

13from gpaw.mpi import MPIComm, broadcast_float, serial_comm 

14from gpaw.typing import Array1D, Array2D 

15 

16 

17class ParallelLayout(NamedTuple): 

18 """Collection of parallel stuff.""" 

19 bd: BandDescriptor 

20 kpt_comm: MPIComm 

21 domain_comm: MPIComm 

22 

23 

24def fermi_dirac(eig: np.ndarray, 

25 fermi_level: float, 

26 width: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: 

27 """Fermi-Dirac distribution function. 

28 

29 >>> f, _, _ = fermi_dirac(np.array([-1.0, 0.0, 1.0]), 0.0, 0.01) 

30 >>> f.round(5) 

31 array([1. , 0.5, 0. ]) 

32 

33 """ 

34 x = (eig - fermi_level) / width 

35 x = np.clip(x, -100, 100) 

36 y = np.exp(x) 

37 z = y + 1.0 

38 f = 1.0 / z 

39 dfde = (f - f**2) / width 

40 y *= x 

41 y /= z 

42 y -= np.log(z) 

43 e_entropy = y * width 

44 return f, dfde, e_entropy 

45 

46 

47def marzari_vanderbilt(eig: np.ndarray, 

48 fermi_level: float, 

49 width: float) -> Tuple[np.ndarray, 

50 np.ndarray, 

51 np.ndarray]: 

52 """Marzari-Vanderbilt distribution (cold smearing). 

53 

54 See: :doi:`10.1103/PhysRevLett.82.3296` 

55 """ 

56 x = (eig - fermi_level) / width 

57 expterm = np.exp(-(x + (1 / np.sqrt(2)))**2) 

58 f = expterm / np.sqrt(2 * np.pi) + 0.5 * (1 - erf(1. / np.sqrt(2) + x)) 

59 dfde = expterm * (2 + np.sqrt(2) * x) / np.sqrt(np.pi) / width 

60 s = expterm * (1 + np.sqrt(2) * x) / (2 * np.sqrt(np.pi)) 

61 e_entropy = -s * width 

62 return f, dfde, e_entropy 

63 

64 

65def methfessel_paxton(eig: np.ndarray, 

66 fermi_level: float, 

67 width: float, 

68 order: int = 0) -> Tuple[np.ndarray, 

69 np.ndarray, 

70 np.ndarray]: 

71 """Methfessel-Paxton distribution.""" 

72 x = (eig - fermi_level) / width 

73 f = 0.5 * (1 - erf(x)) 

74 for i in range(order): 

75 f += (coff_function(i + 1) * 

76 hermite_poly(2 * i + 1, x) * np.exp(-x**2)) 

77 dfde = 1 / np.sqrt(pi) * np.exp(-x**2) 

78 for i in range(order): 

79 dfde += (coff_function(i + 1) * 

80 hermite_poly(2 * i + 2, x) * np.exp(-x**2)) 

81 dfde *= 1.0 / width 

82 e_entropy = (0.5 * coff_function(order) * 

83 hermite_poly(2 * order, x) * np.exp(-x**2)) 

84 e_entropy = -e_entropy * width 

85 return f, dfde, e_entropy 

86 

87 

88def coff_function(n): 

89 return (-1)**n / (np.prod(np.arange(1, n + 1)) * 

90 4**n * np.sqrt(np.pi)) 

91 

92 

93def hermite_poly(n, x): 

94 if n == 0: 

95 return 1 

96 elif n == 1: 

97 return 2 * x 

98 else: 

99 return (2 * x * hermite_poly(n - 1, x) - 

100 2 * (n - 1) * hermite_poly(n - 2, x)) 

101 

102 

103class OccupationNumberCalculator: 

104 """Base class for all occupation number calculators.""" 

105 name = 'unknown' 

106 extrapolate_factor: float 

107 

108 def __init__(self, 

109 parallel_layout: ParallelLayout = None): 

110 """Object for calculating fermi level(s) and occupation numbers. 

111 

112 If fixmagmom=True then the fixed_magmom_value attribute must be set 

113 and two fermi levels will be calculated. 

114 """ 

115 if parallel_layout is None: 

116 parallel_layout = ParallelLayout(BandDescriptor(1), 

117 serial_comm, 

118 serial_comm) 

119 self.bd = parallel_layout.bd 

120 self.kpt_comm = parallel_layout.kpt_comm 

121 self.domain_comm = parallel_layout.domain_comm 

122 

123 @property 

124 def parallel_layout(self) -> ParallelLayout: 

125 return ParallelLayout(self.bd, self.kpt_comm, self.domain_comm) 

126 

127 def todict(self): 

128 return {'name': self.name} 

129 

130 def copy(self, 

131 parallel_layout: ParallelLayout = None, 

132 bz2ibzmap: List[int] = None) -> 'OccupationNumberCalculator': 

133 return create_occ_calc( 

134 self.todict(), 

135 parallel_layout=parallel_layout or self.parallel_layout) 

136 

137 def calculate(self, 

138 nelectrons: float, 

139 eigenvalues: list[list[float]], 

140 weights: list[float], 

141 fermi_levels_guess: list[float] = None, 

142 fix_fermi_level: bool = False) -> tuple[Array2D, 

143 list[float], 

144 float]: 

145 """Calculate occupation numbers and fermi level(s) from eigenvalues. 

146 

147 nelectrons: 

148 Number of electrons. 

149 eigenvalues: ndarray, shape=(nibzkpts, nbands) 

150 Eigenvalues in Hartree. 

151 weights: ndarray, shape=(nibzkpts,) 

152 Weights of k-points in IBZ (must sum to 1). 

153 parallel: 

154 Parallel distribution of eigenvalues. 

155 fermi_level_guesses: 

156 Optional guess(es) at fermi level(s). 

157 

158 Returns a tuple containing: 

159 

160 * occupation numbers (in the range 0 to 1) 

161 * fermi-level in Hartree 

162 * entropy as -S*T in Hartree 

163 

164 >>> occ = ZeroWidth() 

165 >>> occ.calculate(1, [[0, 1]], [1]) 

166 (array([[1., 0.]]), [0.5], 0.0) 

167 """ 

168 

169 eig_qn = [np.asarray(eig_n) for eig_n in eigenvalues] 

170 weight_q = np.asarray(weights) 

171 

172 if fermi_levels_guess is None: 

173 fermi_levels_guess = [nan] 

174 

175 f_qn = np.empty((len(weight_q), len(eig_qn[0]))) 

176 

177 result = np.empty(2) 

178 

179 if self.domain_comm.rank == 0: 

180 # Let the master domain do the work and broadcast results: 

181 result[:] = self._calculate( 

182 nelectrons, eig_qn, weight_q, f_qn, 

183 fermi_levels_guess[0], fix_fermi_level) 

184 

185 self.domain_comm.broadcast(result, 0) 

186 self.domain_comm.broadcast(f_qn, 0) 

187 

188 fermi_level, e_entropy = result 

189 return f_qn, [float(fermi_level)], float(e_entropy) 

190 

191 def _calculate(self, 

192 nelectrons: float, 

193 eig_qn: List[Array1D], 

194 weight_q: Array1D, 

195 f_qn: Array2D, 

196 fermi_level_guess: float, 

197 fix_fermi_level: bool = False) -> tuple[float, float]: 

198 raise NotImplementedError 

199 

200 

201class FixMagneticMomentOccupationNumberCalculator(OccupationNumberCalculator): 

202 """Base class for all occupation number objects.""" 

203 name = 'fixmagmom' 

204 

205 def __init__(self, 

206 occ: OccupationNumberCalculator, 

207 magmom: float): 

208 """Object for calculating fermi level(s) and occupation numbers. 

209 

210 If fixmagmom=True then the fixed_magmom_value attribute must be set 

211 and two fermi levels will be calculated. 

212 """ 

213 self.occ = occ 

214 self.fixed_magmom_value = magmom 

215 self.extrapolate_factor = occ.extrapolate_factor 

216 

217 def __str__(self): 

218 return (f'Fixed magnetic moment: {self.fixed_magmom_value:.3f}\n' + 

219 str(self.occ)) 

220 

221 def todict(self): 

222 dct = self.occ.todict() 

223 dct['fixmagmom'] = True 

224 return dct 

225 

226 def calculate(self, 

227 nelectrons: float, 

228 eigenvalues: List[List[float]], 

229 weights: List[float], 

230 fermi_levels_guess: List[float] = None, 

231 fix_fermi_level: bool = False) -> tuple[Array2D, 

232 List[float], 

233 float]: 

234 

235 magmom = self.fixed_magmom_value 

236 

237 if fermi_levels_guess is None: 

238 fermi_levels_guess = [nan, nan] 

239 

240 f1_qn, fermi_levels1, e_entropy1 = self.occ.calculate( 

241 (nelectrons + magmom) / 2, 

242 eigenvalues[::2], 

243 weights[::2], 

244 fermi_levels_guess[:1], 

245 fix_fermi_level) 

246 

247 f2_qn, fermi_levels2, e_entropy2 = self.occ.calculate( 

248 (nelectrons - magmom) / 2, 

249 eigenvalues[1::2], 

250 weights[1::2], 

251 fermi_levels_guess[1:], 

252 fix_fermi_level) 

253 

254 f_qn = [] 

255 for f1_n, f2_n in zip(f1_qn, f2_qn): 

256 f_qn += [f1_n, f2_n] 

257 

258 return (np.array(f_qn), 

259 fermi_levels1 + fermi_levels2, 

260 e_entropy1 + e_entropy2) 

261 

262 

263class SmoothDistribution(OccupationNumberCalculator): 

264 """Base class for Fermi-Dirac and other smooth distributions.""" 

265 def __init__(self, width: float, parallel_layout: ParallelLayout = None): 

266 """Smooth distribution. 

267 

268 width: float 

269 Width of distribution in eV. 

270 fixmagmom: bool 

271 Fix spin moment calculations. A separate Fermi level for 

272 spin up and down electrons is found. 

273 """ 

274 

275 self._width = width 

276 OccupationNumberCalculator.__init__(self, parallel_layout) 

277 

278 def todict(self): 

279 return {'name': self.name, 'width': self._width} 

280 

281 def _calculate(self, 

282 nelectrons, 

283 eig_qn, 

284 weight_q, 

285 f_qn, 

286 fermi_level_guess, 

287 fix_fermi_level): 

288 # Guess can be nan or inf: 

289 if not np.isfinite(fermi_level_guess) or self._width == 0.0: 

290 zero = ZeroWidth(self.parallel_layout) 

291 fermi_level_guess, _ = zero._calculate( 

292 nelectrons, eig_qn, weight_q, f_qn) 

293 if self._width == 0.0 or np.isinf(fermi_level_guess): 

294 return fermi_level_guess, 0.0 

295 

296 x = fermi_level_guess 

297 

298 data = np.empty(3) 

299 

300 def func(x, data=data): 

301 """calculate excess electrons (and update occupation numbers).""" 

302 data[:] = 0.0 

303 for eig_n, weight, f_n in zip(eig_qn, weight_q, f_qn): 

304 f_n[:], dfde_n, e_entropy_n = self.distribution(eig_n, x) 

305 data += [weight * x_n.sum() 

306 for x_n in [f_n, dfde_n, e_entropy_n]] 

307 self.bd.comm.sum(data) 

308 self.kpt_comm.sum(data) 

309 f, dfde = data[:2] 

310 df = f - nelectrons 

311 return df, dfde 

312 

313 if fix_fermi_level: 

314 fermi_level = x 

315 func(fermi_level) 

316 else: 

317 fermi_level, niter = findroot(func, x) 

318 

319 e_entropy = data[2] 

320 

321 return fermi_level, e_entropy 

322 

323 

324class FermiDiracCalculator(SmoothDistribution): 

325 name = 'fermi-dirac' 

326 extrapolate_factor = -0.5 

327 

328 def distribution(self, 

329 eig_n: np.ndarray, 

330 fermi_level: float) -> Tuple[np.ndarray, 

331 np.ndarray, 

332 np.ndarray]: 

333 return fermi_dirac(eig_n, fermi_level, self._width) 

334 

335 def __str__(self): 

336 return f'Fermi-Dirac:\n width: {self._width:.4f} # eV\n' 

337 

338 

339class MarzariVanderbiltCalculator(SmoothDistribution): 

340 name = 'marzari-vanderbilt' 

341 # According to Nicola Marzari, one should not extrapolate M-V energies 

342 # https://lists.quantum-espresso.org/pipermail/users/2005-October/003170.html 

343 extrapolate_factor = 0.0 

344 

345 def distribution(self, eig_n, fermi_level): 

346 return marzari_vanderbilt(eig_n, fermi_level, self._width) 

347 

348 def __str__(self): 

349 return f'Marzari-Vanderbilt:\n width: {self._width:.4f} # eV\n' 

350 

351 

352class MethfesselPaxtonCalculator(SmoothDistribution): 

353 name = 'methfessel_paxton' 

354 

355 def __init__(self, width, order=0, parallel_layout: ParallelLayout = None): 

356 SmoothDistribution.__init__(self, width, parallel_layout) 

357 self.order = order 

358 self.extrapolate_factor = -1.0 / (self.order + 2) 

359 

360 def todict(self): 

361 dct = SmoothDistribution.todict(self) 

362 dct['order'] = self.order 

363 return dct 

364 

365 def __str__(self): 

366 return ('Methfessel-Paxton:\n' 

367 f' width: {self._width:.4f} # eV\n' 

368 f' order: {self.order}\n') 

369 

370 def distribution(self, eig_n, fermi_level): 

371 return methfessel_paxton(eig_n, fermi_level, self._width, self.order) 

372 

373 

374def findroot(func: Callable[[float], Tuple[float, float]], 

375 x: float, 

376 tol: float = 1e-10) -> Tuple[float, int]: 

377 """Function used for locating Fermi level. 

378 

379 The function should return a (value, derivative) tuple: 

380 

381 >>> x, _ = findroot(lambda x: (x, 1.0), 1.0) 

382 >>> assert abs(x) < 1e-10 

383 """ 

384 

385 assert np.isfinite(x), x 

386 

387 xmin = -np.inf 

388 xmax = np.inf 

389 

390 # Try 10 step using the gradient: 

391 niter = 0 

392 while True: 

393 f, dfdx = func(x) 

394 if abs(f) < tol: 

395 return x, niter 

396 if f < 0.0 and x > xmin: 

397 xmin = x 

398 elif f > 0.0 and x < xmax: 

399 xmax = x 

400 dx = -f / max(dfdx, 1e-18) 

401 if niter == 10 or abs(dx) > 0.3 or not (xmin < x + dx < xmax): 

402 break # try bisection 

403 x += dx 

404 niter += 1 

405 

406 # Bracket the solution: 

407 if not np.isfinite(xmin): 

408 xmin = x 

409 fmin = f 

410 step = 0.01 

411 while fmin > tol: 

412 xmin -= step 

413 fmin = func(xmin)[0] 

414 step *= 2 

415 

416 if not np.isfinite(xmax): 

417 xmax = x 

418 fmax = f 

419 step = 0.01 

420 while fmax < 0: 

421 xmax += step 

422 fmax = func(xmax)[0] 

423 step *= 2 

424 

425 # Bisect: 

426 while True: 

427 x = (xmin + xmax) / 2 

428 f = func(x)[0] 

429 if abs(f) < tol: 

430 return x, niter 

431 if f > 0: 

432 xmax = x 

433 else: 

434 xmin = x 

435 niter += 1 

436 assert niter < 1000 

437 

438 

439def collect_eigelvalues(eig_qn: np.ndarray, 

440 weight_q: np.ndarray, 

441 bd: BandDescriptor, 

442 kpt_comm: MPIComm) -> Tuple[np.ndarray, 

443 np.ndarray, 

444 np.ndarray]: 

445 """Collect eigenvalues from bd.comm and kpt_comm.""" 

446 nkpts_r = np.zeros(kpt_comm.size, int) 

447 nkpts_r[kpt_comm.rank] = len(weight_q) 

448 kpt_comm.sum(nkpts_r) 

449 nk = cast(int, nkpts_r.sum()) 

450 weight_k = np.zeros(nk) 

451 k1 = nkpts_r[:kpt_comm.rank].sum() 

452 k2 = k1 + len(weight_q) 

453 weight_k[k1:k2] = weight_q 

454 kpt_comm.sum(weight_k, 0) 

455 

456 eig_kn: Array2D = np.zeros((0, 0)) 

457 k = 0 

458 for rank, nkpts in enumerate(nkpts_r): 

459 for q in range(nkpts): 

460 if rank == kpt_comm.rank: 

461 eig_n = eig_qn[q] 

462 eig_n = bd.collect(eig_n) 

463 if bd.comm.rank == 0: 

464 if kpt_comm.rank == 0: 

465 if k == 0: 

466 eig_kn = np.empty((nk, len(eig_n))) 

467 if rank == 0: 

468 eig_kn[k] = eig_n 

469 else: 

470 kpt_comm.receive(eig_kn[k], rank) 

471 elif rank == kpt_comm.rank: 

472 kpt_comm.send(eig_n, 0) 

473 k += 1 

474 return eig_kn, weight_k, nkpts_r 

475 

476 

477def distribute_occupation_numbers(f_kn: np.ndarray, # input 

478 f_qn: np.ndarray, # output 

479 nkpts_r: np.ndarray, 

480 bd: BandDescriptor, 

481 kpt_comm: MPIComm) -> None: 

482 """Distribute occupation numbers over bd.comm and kpt_comm.""" 

483 k = 0 

484 for rank, nkpts in enumerate(nkpts_r): 

485 for q in range(nkpts): 

486 if kpt_comm.rank == 0: 

487 if rank == 0: 

488 if bd.comm.size == 1: 

489 f_qn[q] = f_kn[k] 

490 else: 

491 bd.distribute(None if f_kn is None else f_kn[k], 

492 f_qn[q]) 

493 elif f_kn is not None: 

494 kpt_comm.send(f_kn[k], rank) 

495 elif rank == kpt_comm.rank: 

496 if bd.comm.size == 1: 

497 kpt_comm.receive(f_qn[q], 0) 

498 else: 

499 if bd.comm.rank == 0: 

500 f_n = bd.empty(global_array=True) 

501 kpt_comm.receive(f_n, 0) 

502 else: 

503 f_n = None 

504 bd.distribute(f_n, f_qn[q]) 

505 k += 1 

506 

507 

508class ZeroWidth(OccupationNumberCalculator): 

509 name = 'zero-width' 

510 extrapolate_factor = 0.0 

511 

512 def todict(self): 

513 return {'width': 0.0} 

514 

515 def __str__(self): 

516 return '# Zero width' 

517 

518 def distribution(self, eig_n, fermi_level): 

519 f_n = np.zeros_like(eig_n) 

520 f_n[eig_n < fermi_level] = 1.0 

521 f_n[eig_n == fermi_level] = 0.5 

522 return f_n, np.zeros_like(eig_n), np.zeros_like(eig_n) 

523 

524 def _calculate(self, 

525 nelectrons, 

526 eig_qn, 

527 weight_q, 

528 f_qn, 

529 fermi_level_guess=nan, 

530 fix_fermi_level=False): 

531 eig_kn, weight_k, nkpts_r = collect_eigelvalues(eig_qn, weight_q, 

532 self.bd, self.kpt_comm) 

533 

534 if eig_kn.size != 0: 

535 # Try to use integer weights (avoid round-off errors): 

536 N = int(round(1 / min(weight_k))) 

537 w_k = (weight_k * N).round().astype(int) 

538 if abs(w_k - N * weight_k).max() > 1e-10: 

539 # Did not work. Use original fractional weights: 

540 w_k = weight_k 

541 N = 1 

542 

543 f_kn = np.zeros_like(eig_kn) 

544 f_m = f_kn.ravel() 

545 w_kn = np.empty_like(eig_kn, dtype=w_k.dtype) 

546 w_kn[:] = w_k[:, np.newaxis] 

547 eig_m = eig_kn.ravel() 

548 w_m = w_kn.ravel() 

549 m_i = eig_m.argsort() 

550 if fix_fermi_level: 

551 fermi_level = fermi_level_guess 

552 f_m[eig_m < fermi_level] = 1.0 

553 f_m[eig_m == fermi_level] = 0.5 

554 else: 

555 w_i = w_m[m_i] 

556 sum_i = np.add.accumulate(w_i) 

557 filled_i = (sum_i <= nelectrons * N) 

558 i = sum(filled_i) 

559 f_m[m_i[:i]] = 1.0 

560 if i == len(m_i): 

561 fermi_level = inf 

562 else: 

563 extra = nelectrons * N - (sum_i[i - 1] if i > 0 else 0.0) 

564 if extra > 0: 

565 assert extra <= w_i[i] 

566 f_m[m_i[i]] = extra / w_i[i] 

567 fermi_level = eig_m[m_i[i]] 

568 else: 

569 fermi_level = (eig_m[m_i[i]] + eig_m[m_i[i - 1]]) / 2 

570 else: 

571 f_kn = None 

572 fermi_level = nan 

573 

574 distribute_occupation_numbers(f_kn, f_qn, nkpts_r, 

575 self.bd, self.kpt_comm) 

576 

577 if self.kpt_comm.rank == 0: 

578 fermi_level = broadcast_float(fermi_level, self.bd.comm) 

579 fermi_level = broadcast_float(fermi_level, self.kpt_comm) 

580 

581 e_entropy = 0.0 

582 return fermi_level, e_entropy 

583 

584 

585class FixedOccupationNumbers(OccupationNumberCalculator): 

586 extrapolate_factor = 0.0 

587 

588 def __init__(self, numbers, parallel_layout: ParallelLayout = None): 

589 """Fixed occupation numbers. 

590 

591 f_sn: ndarray, shape=(nspins, nbands) 

592 Occupation numbers (in the range from 0 to 1) 

593 

594 Example (excited state with 4 electrons):: 

595 

596 occ = FixedOccupationNumbers([[1, 0, 1, 0], [1, 1, 0, 0]]) 

597 

598 """ 

599 OccupationNumberCalculator.__init__(self, parallel_layout) 

600 self.f_sn = np.array(numbers) 

601 

602 def _calculate(self, 

603 nelectrons, 

604 eig_qn, 

605 weight_q, 

606 f_qn, 

607 fermi_level_guess=nan, 

608 fix_fermi_level=False): 

609 

610 calc_fixed(self.bd, self.f_sn, f_qn) 

611 

612 return inf, 0.0 

613 

614 def todict(self): 

615 return {'name': 'fixed', 'numbers': self.f_sn} 

616 

617 

618class FixedOccupationNumbersUniform(OccupationNumberCalculator): 

619 

620 extrapolate_factor = 0.0 

621 name = 'fixed-uniform' 

622 

623 def __init__(self, nelectrons, nspins, magmom, nkpts, nbands, 

624 parallel_layout: ParallelLayout = None): 

625 """ 

626 Uniform distribution of occupation numbers: each 

627 k-point per spin has the same number of occupied states 

628 Magnetic moment defines difference between two spins occ. numb. 

629 """ 

630 OccupationNumberCalculator.__init__(self, parallel_layout) 

631 

632 def get_f(nelectrons, magmom, nkpts, nbands, spin): 

633 """ 

634 :param nelectrons: 

635 :param magmom: 

636 :param nkpts: 

637 :param nbands: 

638 :param spin: +1 or -1 

639 :return: occupation numbers per spin channel 

640 """ 

641 

642 f_qn = np.zeros(shape=(nkpts, nbands)) 

643 nelecps = (nelectrons + spin * magmom) / 2 

644 assert int(nelecps) < nbands, 'need more bands!' 

645 

646 f_qn[:, : int(nelecps)] = 1.0 

647 f_qn[:, int(nelecps)] = nelecps - int(nelecps) 

648 

649 return f_qn 

650 

651 if nspins == 2: 

652 f1_qn = get_f(nelectrons, magmom, nkpts, nbands, 1) 

653 f2_qn = get_f(nelectrons, magmom, nkpts, nbands, -1) 

654 f_qn = [] 

655 for f1_n, f2_n in zip(f1_qn, f2_qn): 

656 f_qn += [f1_n, f2_n] 

657 else: 

658 f_qn = get_f(nelectrons, magmom, nkpts, nbands, 0) 

659 

660 self.f_sn = np.array(f_qn) 

661 self.nspins = nspins 

662 self.magmom = magmom 

663 

664 def _calculate(self, 

665 nelectrons, 

666 eig_qn, 

667 weight_q, 

668 f_qn, 

669 fermi_level_guess=nan, 

670 fix_fermi_level=False): 

671 assert not fix_fermi_level 

672 

673 calc_fixed(self.bd, self.f_sn, f_qn) 

674 

675 eig_kn, weight_k, nkpts_r = collect_eigelvalues( 

676 eig_qn, weight_q, self.bd, self.kpt_comm) 

677 

678 def get_homo(eig_kn, nelectrons, deg, magmom, spin): 

679 nelecps = int((nelectrons * deg + spin * magmom) / 2) 

680 return np.max(eig_kn[:, np.maximum(nelecps - 1, 0)]) 

681 

682 def get_lumo(eig_kn, nelectrons, deg, magmom, spin): 

683 nelecps = int((nelectrons * deg + spin * magmom) / 2) 

684 return np.min(eig_kn[:, nelecps]) 

685 

686 deg = 3 - self.nspins 

687 mm = self.magmom 

688 if eig_kn.size != 0: 

689 if self.nspins == 2: 

690 hup = get_homo(eig_kn[::2], nelectrons, deg, mm, 1) 

691 hdown = get_homo(eig_kn[1::2], nelectrons, deg, mm, -1) 

692 homo = np.maximum(hup, hdown) 

693 

694 lup = get_lumo(eig_kn[::2], nelectrons, deg, mm, 1) 

695 ldown = get_lumo(eig_kn[1::2], nelectrons, deg, mm, -1) 

696 lumo = np.maximum(lup, ldown) 

697 

698 else: 

699 homo = get_homo(eig_kn, nelectrons, deg, mm, 0) 

700 lumo = get_lumo(eig_kn, nelectrons, deg, mm, 0) 

701 fermi_level = (homo + lumo) / 2 

702 else: 

703 fermi_level = nan 

704 

705 if self.kpt_comm.rank == 0: 

706 fermi_level = broadcast_float(fermi_level, self.bd.comm) 

707 fermi_level = broadcast_float(fermi_level, self.kpt_comm) 

708 

709 return fermi_level, 0.0 

710 

711 def todict(self): 

712 return {'name': 'fixed-uniform'} 

713 

714 def __str__(self): 

715 return '# Uniform distribution of occupation numbers' 

716 

717 

718def calc_fixed(bd, f_sn, f_qn): 

719 if bd.nbands == f_sn.shape[1]: 

720 for q, f_n in enumerate(f_qn): 

721 s = q % len(f_sn) 

722 bd.distribute(f_sn[s], f_n) 

723 else: 

724 # Non-collinear calculation: 

725 bd.distribute(f_sn.T.flatten().copy(), f_qn[0]) 

726 

727 

728def FixedOccupations(f_sn): 

729 warnings.warn( 

730 "Please use occupations={'name': 'fixed', 'numbers': ...} instead.") 

731 if len(f_sn) == 1: 

732 f_sn = np.array(f_sn) / 2 

733 return {'name': 'fixed', 'numbers': f_sn} 

734 

735 

736class ThomasFermiOccupations(OccupationNumberCalculator): 

737 name = 'orbital-free' 

738 extrapolate_factor = 0.0 

739 

740 def _calculate(self, 

741 nelectrons, 

742 eig_qn, 

743 weight_q, 

744 f_qn, 

745 fermi_level_guess=nan, 

746 fix_fermi_level=False): 

747 assert len(f_qn) == 1 

748 f_qn[0][:] = [nelectrons] 

749 return inf, 0.0 

750 

751 

752def create_occ_calc(dct: Dict[str, Any], 

753 *, 

754 parallel_layout: ParallelLayout = None, 

755 fixed_magmom_value=None, 

756 rcell=None, 

757 monkhorst_pack_size=None, 

758 bz2ibzmap=None, 

759 nspins=None, 

760 nelectrons=None, 

761 nkpts=None, 

762 nbands=None 

763 ) -> OccupationNumberCalculator: 

764 """Surprise: Create occupation-number object. 

765 

766 The unit of width is eV and name must be one of: 

767 

768 * 'fermi-dirac' 

769 * 'marzari-vanderbilt' 

770 * 'methfessel-paxton' 

771 * 'fixed' 

772 * 'tetrahedron-method' 

773 * 'improved-tetrahedron-method' 

774 * 'orbital-free' 

775 

776 >>> occ = create_occ_calc({'width': 0.0}) 

777 >>> occ.calculate(nelectrons=3, 

778 ... eigenvalues=[[0, 1, 2], [0, 2, 3]], 

779 ... weights=[1, 1]) 

780 (array([[1., 1., 0.], 

781 [1., 0., 0.]]), [1.5], 0.0) 

782 """ 

783 kwargs = dct.copy() 

784 fix_the_magnetic_moment = kwargs.pop('fixmagmom', False) 

785 name = kwargs.pop('name', '') 

786 kwargs['parallel_layout'] = parallel_layout 

787