Coverage for gpaw/new/eigensolver.py: 49%

1from __future__ import annotations 

2 

3import numpy as np 

4 

5from ase.units import Ha 

6 

7from gpaw.new.density import Density 

8from gpaw.new.hamiltonian import Hamiltonian 

9from gpaw.new.potential import Potential 

10from gpaw.new.energies import DFTEnergies 

11from gpaw.new.ibzwfs import IBZWaveFunctions 

12from gpaw.new.pot_calc import PotentialCalculator 

13from gpaw.mpi import broadcast_float 

14from gpaw.typing import Array1D 

15 

16 

17class Eigensolver: 

18 direct = False 

19 

20 def iterate(self, 

21 ibzwfs: IBZWaveFunctions, 

22 density: Density, 

23 potential: Potential, 

24 hamiltonian: Hamiltonian, 

25 pot_calc: PotentialCalculator, 

26 energies: DFTEnergies) -> tuple[float, float, DFTEnergies]: 

27 raise NotImplementedError 

28 

29 def postprocess(self, ibzwfs, density, potential, hamiltonian): 

30 pass 

31 

32 def iterate_kpt(self, wfs, weight_n, iter_func, **fkwargs): 

33 had_eigs_and_occs = wfs.has_eigs and wfs.has_occs 

34 if had_eigs_and_occs: 

35 eig_old = wfs.myeig_n 

36 eigs_error = iter_func(wfs=wfs, weight_n=weight_n, **fkwargs) 

37 if had_eigs_and_occs: 

38 eig_error = np.max(weight_n * np.abs(eig_old - wfs.myeig_n), 

39 initial=0) 

40 else: # no eigenvalues to compare with 

41 eig_error = np.inf 

42 return eigs_error, eig_error 

43 

44 

45def calculate_weights(converge_bands: int | str, 

46 ibzwfs: IBZWaveFunctions) -> list[Array1D | None]: 

47 """Calculate convergence weights for all eigenstates.""" 

48 weight_un = [] 

49 nu = len(ibzwfs.wfs_qs) * ibzwfs.nspins 

50 nbands = ibzwfs.nbands 

51 

52 if converge_bands == 'occupied': 

53 # Converge occupied bands: 

54 for wfs in ibzwfs: 

55 if wfs.has_occs: 

56 # Methfessel-Paxton or cold-smearing distributions can give 

57 # negative occupation numbers - so we take the absolute value: 

58 weight_n = np.abs(wfs.myocc_n) 

59 else: 

60 # No eigenvalues yet: 

61 return [None] * nu 

62 weight_un.append(weight_n) 

63 return weight_un 

64 

65 if converge_bands == 'all': 

66 converge_bands = nbands 

67 

68 if not isinstance(converge_bands, str): 

69 # Converge fixed number of bands: 

70 n = converge_bands 

71 if n < 0: 

72 n += nbands 

73 assert n >= 0 

74 for wfs in ibzwfs: 

75 weight_n = np.zeros(wfs.n2 - wfs.n1) 

76 m = max(wfs.n1, min(n, wfs.n2)) - wfs.n1 

77 weight_n[:m] = 1.0 

78 weight_un.append(weight_n) 

79 return weight_un 

80 

81 # Converge states with energy up to CBM + delta: 

82 assert converge_bands.startswith('CBM+') 

83 delta = float(converge_bands[4:]) / Ha 

84 

85 if ibzwfs.fermi_levels is None: 

86 return [None] * nu 

87 

88 efermi = np.mean(ibzwfs.fermi_levels) 

89 

90 # Find CBM: 

91 cbm = np.inf 

92 nocc_u = np.empty(nu, int) 

93 for u, wfs in enumerate(ibzwfs): 

94 n = (wfs.eig_n < efermi).sum() # number of occupied bands 

95 nocc_u[u] = n 

96 if n < nbands: 

97 cbm = min(cbm, wfs.eig_n[n]) 

98 

99 # If all k-points don't have the same number of occupied bands, 

100 # then it's a metal: 

101 n0 = int(broadcast_float(float(nocc_u[0]), ibzwfs.kpt_comm)) 

102 metal = bool(ibzwfs.kpt_comm.sum_scalar(float((nocc_u != n0).any()))) 

103 if metal: 

104 cbm = efermi 

105 else: 

106 cbm = ibzwfs.kpt_comm.min_scalar(cbm) 

107 

108 ecut = cbm + delta 

109 

110 for wfs in ibzwfs: 

111 weight_n = (wfs.myeig_n < ecut).astype(float) 

112 if wfs.eig_n[-1] < ecut: 

113 # We don't have enough bands! 

114 weight_n[:] = np.inf 

115 weight_un.append(weight_n) 

116 

117 return weight_un