Coverage for gpaw/dos.py: 95%

1from __future__ import annotations 

2from pathlib import Path 

3from typing import Union, List, Optional, Sequence, TYPE_CHECKING 

4 

5import numpy as np 

6from ase.dft.dos import linear_tetrahedron_integration as lti 

7 

8from gpaw.setup import Setup 

9from gpaw.spinorbit import soc_eigenstates, BZWaveFunctions 

10from gpaw.typing import Array1D, Array2D, Array3D, ArrayLike1D 

11 

12if TYPE_CHECKING: 

13 from gpaw.calculator import GPAW 

14 from gpaw.new.ase_interface import ASECalculator 

15 

16 

17class IBZWaveFunctions: 

18 """Container for eigenvalues and PAW projections (only IBZ).""" 

19 def __init__(self, calc: ASECalculator | GPAW): 

20 self.calc = calc 

21 self.fermi_level = self.calc.get_fermi_level() 

22 self.size = calc.wfs.kd.N_c 

23 self.bz2ibz_map = calc.wfs.kd.bz2ibz_k 

24 

25 def weights(self) -> Array1D: 

26 """Weigths of IBZ k-points (adds to 1.0).""" 

27 return self.calc.wfs.kd.weight_k 

28 

29 def eigenvalues(self) -> Array3D: 

30 """All eigenvalues.""" 

31 kd = self.calc.wfs.kd 

32 eigs = np.array([[self.calc.get_eigenvalues(kpt=k, spin=s) 

33 for k in range(kd.nibzkpts)] 

34 for s in range(kd.nspins)]) 

35 return eigs 

36 

37 def pdos_weights(self, 

38 a: int, 

39 indices: List[int] 

40 ) -> Array3D: 

41 """Projections for PDOS. 

42 

43 Returns (nibzkpts, nbands, nspins)-shaped ndarray 

44 of the square of absolute value of the projections. The *indices* 

45 list contains projector-indices. 

46 """ 

47 kd = self.calc.wfs.kd 

48 dos_kns = np.zeros((kd.nibzkpts, 

49 self.calc.wfs.bd.nbands, 

50 kd.nspins)) 

51 bands = self.calc.wfs.bd.get_slice() 

52 

53 for wf in self.calc.wfs.kpt_u: 

54 P_ani = wf.projections 

55 if a in P_ani: 

56 P_ni = P_ani[a][:, indices] 

57 dos_kns[wf.k, bands, wf.s] = (abs(P_ni)**2).sum(1) 

58 

59 self.calc.world.sum(dos_kns) 

60 return dos_kns 

61 

62 

63def get_projector_numbers(setup: Setup, ell: int) -> List[int]: 

64 """Find indices of bound-state PAW projector functions. 

65 

66 >>> from gpaw.setup import create_setup 

67 >>> setup = create_setup('Li') 

68 >>> get_projector_numbers(setup, 0) 

69 [0] 

70 >>> get_projector_numbers(setup, 1) 

71 [1, 2, 3] 

72 """ 

73 indices = [] 

74 i1 = 0 

75 for n, l in zip(setup.n_j, setup.l_j): 

76 i2 = i1 + 2 * l + 1 

77 if l == ell and n >= 0: 

78 indices += list(range(i1, i2)) 

79 i1 = i2 

80 return indices 

81 

82 

83def gaussian_dos(eig_kn, 

84 weight_kn, 

85 weight_k, 

86 energies, 

87 width: float) -> Array1D: 

88 """Simple broadening with a Gaussian.""" 

89 dos = np.zeros_like(energies) 

90 if weight_kn is None: 

91 for e_n, w in zip(eig_kn, weight_k): 

92 for e in e_n: 

93 dos += w * np.exp(-((energies - e) / width)**2) 

94 else: 

95 for e_n, w, w_n in zip(eig_kn, weight_k, weight_kn): 

96 for e, w2 in zip(e_n, w_n): 

97 dos += w * w2 * np.exp(-((energies - e) / width)**2) 

98 return dos / (np.pi**0.5 * width) 

99 

100 

101def linear_tetrahedron_dos(eig_kn, 

102 weight_kn, 

103 energies, 

104 cell, 

105 size, 

106 bz2ibz_map=None) -> Array1D: 

107 """Linear-tetrahedron method.""" 

108 if len(eig_kn) != np.prod(size): 

109 eig_kn = eig_kn[bz2ibz_map] 

110 if weight_kn is not None: 

111 weight_kn = weight_kn[bz2ibz_map] 

112 

113 shape = tuple(size) + (-1,) 

114 eig_kn = eig_kn.reshape(shape) 

115 if weight_kn is not None: 

116 weight_kn = weight_kn.reshape(shape) 

117 

118 dos = lti(cell, eig_kn, energies, weight_kn) 

119 return dos 

120 

121 

122class DOSCalculator: 

123 def __init__(self, 

124 wfs, 

125 setups=None, 

126 cell=None, 

127 shift_fermi_level=True): 

128 self.wfs = wfs 

129 self.setups = setups 

130 self.cell = cell 

131 

132 self.eig_skn = wfs.eigenvalues() 

133 self.fermi_level = wfs.fermi_level 

134 

135 if shift_fermi_level: 

136 self.eig_skn -= wfs.fermi_level 

137 

138 self.collinear = (self.eig_skn.ndim == 3) 

139 if self.collinear: 

140 self.degeneracy = 2 / len(self.eig_skn) 

141 else: 

142 self.eig_skn = np.array([self.eig_skn, self.eig_skn]) 

143 self.degeneracy = 0.5 

144 

145 self.nspins = len(self.eig_skn) 

146 self.weight_k = wfs.weights() 

147 

148 def get_energies(self, 

149 emin: Optional[float] = None, 

150 emax: Optional[float] = None, 

151 npoints: int = 100): 

152 emin = emin if emin is not None else self.eig_skn.min() 

153 emax = emax if emax is not None else self.eig_skn.max() 

154 return np.linspace(emin, emax, npoints) 

155 

156 @classmethod 

157 def from_calculator(cls, 

158 filename: ASECalculator | GPAW | Path | str, 

159 soc=False, theta=0.0, phi=0.0, 

160 shift_fermi_level=True) -> DOSCalculator: 

161 """Create DOSCalculator from a GPAW calculation. 

162 

163 filename: str 

164 Name of restart-file or GPAW calculator object. 

165 """ 

166 from gpaw.calculator import GPAW 

167 if not isinstance(filename, (str, Path)): 

168 calc = filename 

169 else: 

170 calc = GPAW(filename, txt=None) 

171 

172 wfs: Union[BZWaveFunctions, IBZWaveFunctions] 

173 if soc: 

174 wfs = soc_eigenstates(calc, theta=theta, phi=phi) 

175 else: 

176 wfs = IBZWaveFunctions(calc) 

177 

178 return DOSCalculator(wfs, 

179 calc.setups, calc.atoms.cell, 

180 shift_fermi_level) 

181 

182 def calculate(self, 

183 energies: ArrayLike1D, 

184 eig_kn: Array2D, 

185 weight_kn: Array2D = None, 

186 width: float = 0.1): 

187 energies = np.asarray(energies) 

188 if width > 0.0: 

189 return gaussian_dos(eig_kn, weight_kn, 

190 self.weight_k, energies, width) 

191 else: 

192 return linear_tetrahedron_dos( 

193 eig_kn, weight_kn, energies, 

194 self.cell, self.wfs.size, self.wfs.bz2ibz_map) 

195 

196 def raw_dos(self, 

197 energies: Sequence[float], 

198 spin: Optional[int] = None, 

199 width: float = 0.1) -> Array1D: 

200 """Calculate density of states. 

201 

202 width: float 

203 Width of Gaussians in eV. Use width=0.0 to use the 

204 linear-tetrahedron-interpolation method. 

205 """ 

206 if spin is None: 

207 dos = sum(self.calculate(energies, eig_kn, width=width) 

208 for eig_kn in self.eig_skn) 

209 dos *= self.degeneracy 

210 else: 

211 dos = self.calculate(energies, self.eig_skn[spin], width=width) 

212 

213 return dos 

214 

215 def raw_pdos(self, 

216 energies: Sequence[float], 

217 a: int, 

218 l: int, 

219 m: Optional[int] = None, 

220 spin: Optional[int] = None, 

221 width: float = 0.1) -> Array1D: 

222 """Calculate projected density of states. 

223 

224 a: 

225 Atom index. 

226 l: 

227 Angular momentum quantum number. 

228 m: 

229 Magnetic quantum number. Default is None meaning sum over all m. 

230 For p-orbitals, m=0,1,2 translates to y, z and x. 

231 For d-orbitals, m=0,1,2,3,4 translates to xy, yz, 3z2-r2, 

232 zx and x2-y2. 

233 spin: 

234 Must be 0, 1 or None meaning spin-up, down or total respectively. 

235 width: float 

236 Width of Gaussians in eV. Use width=0.0 to use the 

237 linear-tetrahedron-interpolation method. 

238 """ 

239 indices = get_projector_numbers(self.setups[a], l) 

240 if m is not None: 

241 indices = indices[m::(2 * l) + 1] 

242 weight_kns = self.wfs.pdos_weights(a, indices) 

243 

244 if spin is None: 

245 dos = sum(self.calculate(energies, 

246 eig_kn, 

247 weight_nk.T, 

248 width=width) 

249 for eig_kn, weight_nk 

250 in zip(self.eig_skn, weight_kns.T)) 

251 dos *= self.degeneracy 

252 else: 

253 dos = self.calculate(energies, 

254 self.eig_skn[spin], 

255 weight_kns[:, :, spin], 

256 width=width) 

257 

258 return dos