Coverage for ase / dft / dos.py: 79.41%

136 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-03-30 08:22 +0000

1# fmt: off 

2 

3from math import pi, sqrt 

4 

5import numpy as np 

6 

7from ase.dft.kpoints import get_monkhorst_pack_size_and_offset 

8from ase.parallel import world 

9from ase.utils.cext import cextension 

10 

11 

12class DOS: 

13 def __init__(self, calc, width=0.1, window=None, npts=401, comm=world): 

14 """Electronic Density Of States object. 

15 

16 calc: calculator object 

17 Any ASE compliant calculator object. 

18 width: float 

19 Width of guassian smearing. Use width=0.0 for linear tetrahedron 

20 interpolation. 

21 window: tuple of two float 

22 Use ``window=(emin, emax)``. If not specified, a window 

23 big enough to hold all the eigenvalues will be used. 

24 npts: int 

25 Number of points. 

26 comm: communicator object 

27 MPI communicator for lti_dos 

28 

29 """ 

30 

31 self.comm = comm 

32 self.npts = npts 

33 self.width = width 

34 self.w_k = calc.get_k_point_weights() 

35 self.nspins = calc.get_number_of_spins() 

36 self.e_skn = np.array([[calc.get_eigenvalues(kpt=k, spin=s) 

37 for k in range(len(self.w_k))] 

38 for s in range(self.nspins)]) 

39 try: # two Fermi levels 

40 for i, eF in enumerate(calc.get_fermi_level()): 

41 self.e_skn[i] -= eF 

42 except TypeError: # a single Fermi level 

43 self.e_skn -= calc.get_fermi_level() 

44 

45 if window is None: 

46 emin = None 

47 emax = None 

48 else: 

49 emin, emax = window 

50 

51 if emin is None: 

52 emin = self.e_skn.min() - 5 * self.width 

53 if emax is None: 

54 emax = self.e_skn.max() + 5 * self.width 

55 

56 self.energies = np.linspace(emin, emax, npts) 

57 

58 if width == 0.0: 

59 bzkpts = calc.get_bz_k_points() 

60 size, _offset = get_monkhorst_pack_size_and_offset(bzkpts) 

61 bz2ibz = calc.get_bz_to_ibz_map() 

62 shape = (self.nspins,) + tuple(size) + (-1,) 

63 self.e_skn = self.e_skn[:, bz2ibz].reshape(shape) 

64 self.cell = calc.atoms.cell 

65 

66 def get_energies(self): 

67 """Return the array of energies used to sample the DOS. 

68 

69 The energies are reported relative to the Fermi level. 

70 """ 

71 return self.energies 

72 

73 def delta(self, energy): 

74 """Return a delta-function centered at 'energy'.""" 

75 x = -((self.energies - energy) / self.width)**2 

76 return np.exp(x) / (sqrt(pi) * self.width) 

77 

78 def get_dos(self, spin=None): 

79 """Get array of DOS values. 

80 

81 The *spin* argument can be 0 or 1 (spin up or down) - if not 

82 specified, the total DOS is returned. 

83 """ 

84 

85 if spin is None: 

86 if self.nspins == 2: 

87 # Return the total DOS 

88 return self.get_dos(spin=0) + self.get_dos(spin=1) 

89 else: 

90 return 2 * self.get_dos(spin=0) 

91 elif spin == 1 and self.nspins == 1: 

92 # For an unpolarized calculation, spin up and down are equivalent 

93 spin = 0 

94 

95 if self.width == 0.0: 

96 dos = linear_tetrahedron_integration(self.cell, self.e_skn[spin], 

97 self.energies, comm=self.comm) 

98 return dos 

99 

100 dos = np.zeros(self.npts) 

101 for w, e_n in zip(self.w_k, self.e_skn[spin]): 

102 for e in e_n: 

103 dos += w * self.delta(e) 

104 return dos 

105 

106 

107def linear_tetrahedron_integration(cell, eigs, energies, 

108 weights=None, comm=world): 

109 """DOS from linear tetrahedron interpolation. 

110 

111 cell: 3x3 ndarray-like 

112 Unit cell. 

113 eigs: (n1, n2, n3, nbands)-shaped ndarray 

114 Eigenvalues on a Monkhorst-Pack grid (not reduced). 

115 energies: 1-d array-like 

116 Energies where the DOS is calculated (must be a uniform grid). 

117 weights: ndarray of shape (n1, n2, n3, nbands) or (n1, n2, n3, nbands, nw) 

118 Weights. Defaults to a (n1, n2, n3, nbands)-shaped ndarray 

119 filled with ones. Can also have an extra dimednsion if there are 

120 nw weights. 

121 comm: communicator object 

122 MPI communicator for lti_dos 

123 

124 Returns 

125 ------- 

126 

127 DOS as an ndarray of same length as energies or as an 

128 ndarray of shape (nw, len(energies)). 

129 

130 See: 

131 

132 Extensions of the tetrahedron method for evaluating 

133 spectral properties of solids, 

134 A. H. MacDonald, S. H. Vosko and P. T. Coleridge, 

135 1979 J. Phys. C: Solid State Phys. 12 2991, 

136 :doi:`10.1088/0022-3719/12/15/008` 

137 """ 

138 

139 from scipy.spatial import Delaunay 

140 

141 # Find the 6 tetrahedra: 

142 size = eigs.shape[:3] 

143 B = (np.linalg.inv(cell) / size).T 

144 indices = np.array([[i, j, k] 

145 for i in [0, 1] for j in [0, 1] for k in [0, 1]]) 

146 dt = Delaunay(np.dot(indices, B)) 

147 

148 if weights is None: 

149 weights = np.ones_like(eigs) 

150 

151 if weights.ndim == 4: 

152 extra_dimension_added = True 

153 weights = weights[:, :, :, :, np.newaxis] 

154 else: 

155 extra_dimension_added = False 

156 

157 nweights = weights.shape[4] 

158 dos = np.empty((nweights, len(energies))) 

159 

160 lti_dos(indices[dt.simplices], eigs, weights, energies, dos, comm) 

161 

162 dos /= np.prod(size) 

163 

164 if extra_dimension_added: 

165 return dos[0] 

166 return dos 

167 

168 

169@cextension 

170def lti_dos(simplices, eigs, weights, energies, dos, world): 

171 shape = eigs.shape[:3] 

172 nweights = weights.shape[-1] 

173 dos[:] = 0.0 

174 n = -1 

175 for index in np.indices(shape).reshape((3, -1)).T: 

176 n += 1 

177 if n % world.size != world.rank: 

178 continue 

179 i = ((index + simplices) % shape).T 

180 E = eigs[i[0], i[1], i[2]].reshape((4, -1)) 

181 W = weights[i[0], i[1], i[2]].reshape((4, -1, nweights)) 

182 for e, w in zip(E.T, W.transpose((1, 0, 2))): 

183 lti_dos1(e, w, energies, dos) 

184 

185 dos /= 6.0 

186 world.sum(dos) 

187 

188 

189def lti_dos1(e, w, energies, dos): 

190 i = e.argsort() 

191 e0, e1, e2, e3 = en = e[i] 

192 w = w[i] 

193 

194 zero = energies[0] 

195 if len(energies) > 1: 

196 de = energies[1] - zero 

197 nn = (np.floor((en - zero) / de).astype(int) + 1).clip(0, 

198 len(energies)) 

199 else: 

200 nn = (en > zero).astype(int) 

201 

202 n0, n1, n2, n3 = nn 

203 

204 if n1 > n0: 

205 s = slice(n0, n1) 

206 x = energies[s] - e0 

207 f10 = x / (e1 - e0) 

208 f20 = x / (e2 - e0) 

209 f30 = x / (e3 - e0) 

210 f01 = 1 - f10 

211 f02 = 1 - f20 

212 f03 = 1 - f30 

213 g = f20 * f30 / (e1 - e0) 

214 dos[:, s] += w.T.dot([f01 + f02 + f03, 

215 f10, 

216 f20, 

217 f30]) * g 

218 if n2 > n1: 

219 delta = e3 - e0 

220 s = slice(n1, n2) 

221 x = energies[s] 

222 f20 = (x - e0) / (e2 - e0) 

223 f30 = (x - e0) / (e3 - e0) 

224 f21 = (x - e1) / (e2 - e1) 

225 f31 = (x - e1) / (e3 - e1) 

226 f02 = 1 - f20 

227 f03 = 1 - f30 

228 f12 = 1 - f21 

229 f13 = 1 - f31 

230 g = 3 / delta * (f12 * f20 + f21 * f13) 

231 dos[:, s] += w.T.dot([g * f03 / 3 + f02 * f20 * f12 / delta, 

232 g * f12 / 3 + f13 * f13 * f21 / delta, 

233 g * f21 / 3 + f20 * f20 * f12 / delta, 

234 g * f30 / 3 + f31 * f13 * f21 / delta]) 

235 if n3 > n2: 

236 s = slice(n2, n3) 

237 x = energies[s] - e3 

238 f03 = x / (e0 - e3) 

239 f13 = x / (e1 - e3) 

240 f23 = x / (e2 - e3) 

241 f30 = 1 - f03 

242 f31 = 1 - f13 

243 f32 = 1 - f23 

244 g = f03 * f13 / (e3 - e2) 

245 dos[:, s] += w.T.dot([f03, 

246 f13, 

247 f23, 

248 f30 + f31 + f32]) * g