Coverage for eminus/occupations.py: 98.28%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Determine occupations for atomic systems from simple inputs.""" 

4 

5import dataclasses 

6import numbers 

7 

8from . import backend as xp 

9from .logger import log 

10from .tools import fermi_distribution, get_Efermi 

11 

12 

13@dataclasses.dataclass 

14class Occupations: 

15 """Occupations class to save electronic state information in one place. 

16 

17 The attribute Nelec has to be given first after instantiation. 

18 """ 

19 

20 # Set the private variable for the attributes that are properties. 

21 _Nelec: int = 0 #: Number of electrons. 

22 _Nspin: int = 0 #: Number of spin states. 

23 _spin: float = 0 #: Number of unpaired electrons. 

24 _charge: int = 0 #: System charge. 

25 _Nstate: int = 0 #: Number of states. 

26 _Nempty: int = 0 #: Number of empty states. 

27 _Nk: int = 1 #: Number of k-points. 

28 _bands: int = 0 #: Number of bands. 

29 _smearing: float = 0 #: Smearing width in Hartree. 

30 is_filled: bool = False #: Determines the Occupations object fill status. 

31 

32 # ### Class properties ### 

33 

34 @property 

35 def Nelec(self): 

36 """Number of electrons.""" 

37 return self._Nelec 

38 

39 @Nelec.setter 

40 def Nelec(self, value): 

41 # Only update Nelec if it actually gets updated 

42 if self._Nelec != int(value): 

43 self._Nelec = int(value) 

44 self.is_filled = False 

45 

46 @property 

47 def Nspin(self): 

48 """Number of spin states.""" 

49 return self._Nspin 

50 

51 @Nspin.setter 

52 def Nspin(self, value): 

53 # Use a spin-paired calculation for an even number of electrons 

54 if value is None: 

55 if self.Nelec % 2 == 0 and self.spin == 0: 

56 value = 1 

57 else: 

58 value = 2 

59 # Only update if needed 

60 if self._Nspin != int(value): 

61 self._Nspin = int(value) 

62 self.spin = self.spin 

63 self.is_filled = False 

64 if hasattr(self, "f"): 

65 log.warning("Reset previously set fillings.") 

66 

67 @property 

68 def spin(self): 

69 """Number of unpaired electrons.""" 

70 return self._spin 

71 

72 @spin.setter 

73 def spin(self, value): 

74 # If no spin is given try to determine it from the number of electrons 

75 if value is None: 

76 if self.Nelec % 2 == 0: 

77 value = 0 

78 else: 

79 value = 1 

80 if self._spin != value: 

81 self.is_filled = False 

82 # We have no spin in the spin-paired case 

83 if self.Nspin == 1: 

84 self._spin = 0 

85 else: 

86 self._spin = value 

87 

88 @property 

89 def charge(self): 

90 """System charge.""" 

91 return self._charge 

92 

93 @charge.setter 

94 def charge(self, value): 

95 # If we set a charge using this setter update the number of electrons 

96 if hasattr(self, "Nelec"): 

97 self.Nelec += self._charge - int(value) 

98 if self._charge != int(value): 

99 self._charge = int(value) 

100 self.is_filled = False 

101 

102 @property 

103 def f(self): 

104 """Occupation numbers per state.""" 

105 return self._f 

106 

107 @f.setter 

108 def f(self, value): 

109 # Make sure the occupations are in a two-dimensional array 

110 if isinstance(value, (list, tuple)) or xp.is_array(value): 

111 value = xp.atleast_2d(xp.asarray(value, dtype=float)) 

112 self.is_filled = False 

113 # This setter will only be called when explicitly setting f 

114 # Call the fill function in that case 

115 self.fill(value) 

116 

117 @property 

118 def Nk(self): 

119 """Number of k-points.""" 

120 return self._Nk 

121 

122 @Nk.setter 

123 def Nk(self, value): 

124 self._Nk = int(value) 

125 self.is_filled = False 

126 

127 @property 

128 def wk(self): 

129 """k-point weights.""" 

130 return self._wk 

131 

132 @wk.setter 

133 def wk(self, value): 

134 self._wk = xp.asarray(value, dtype=float) 

135 self._Nk = len(self._wk) 

136 self.is_filled = False 

137 

138 @property 

139 def bands(self): 

140 """Total number of bands.""" 

141 return self._bands 

142 

143 @bands.setter 

144 def bands(self, value): 

145 self._bands = int(value) 

146 self.is_filled = False 

147 

148 @property 

149 def smearing(self): 

150 """Smearing width in Hartree.""" 

151 return self._smearing 

152 

153 @smearing.setter 

154 def smearing(self, value): 

155 if value < 0: 

156 log.error("The smearing width can not be negative.") 

157 self._smearing = value 

158 self.is_filled = False 

159 if self.Nempty > 0: 

160 log.warning("Empty states with smearing enabled found.") 

161 

162 @property 

163 def magnetization(self): 

164 """Magnetization from occupation numbers.""" 

165 # There is no magnetization in the spin-paired case 

166 if self.Nspin == 1: 

167 return 0 

168 if not self.is_filled: 

169 log.warning("Can not calculate magnetization for unfilled occupations.") 

170 return 0 

171 return xp.sum(self.wk * (self.f[:, 0] - self.f[:, 1])) / xp.sum(self.wk * self.f) 

172 

173 @magnetization.setter 

174 def magnetization(self, value): 

175 if value is not None: 

176 if self.is_filled: 

177 log.warning("Reset previously set fillings.") 

178 self.is_filled = False 

179 self.fill(None, value) 

180 

181 # ### Read-only properties ### 

182 

183 @property 

184 def multiplicity(self): 

185 """Multiplicity, i.e., 2S+1.""" 

186 return self.spin + 1 

187 

188 @property 

189 def Nstate(self): 

190 """Number of states.""" 

191 return self._Nstate 

192 

193 @property 

194 def Nempty(self): 

195 """Number of empty states.""" 

196 return self._Nempty 

197 

198 @property 

199 def F(self): 

200 """Diagonal matrices of f per k-point and spin.""" 

201 return [[xp.diag(f) for f in f_spin] for f_spin in self.f] 

202 

203 # ### Class methods ### 

204 

205 def fill(self, f=None, magnetization=None): 

206 """Fill the states of the object. 

207 

208 Keyword Args: 

209 f: Fillings. 

210 magnetization: Magnetization. 

211 """ 

212 # Do nothing if the object is already filled 

213 if self.is_filled: 

214 return self 

215 # If no f is given just use the standard fillings: 2 for restricted and 1 for unrestricted 

216 if f is None: 

217 f = 2 / self.Nspin 

218 self._update_from_fillings(f, magnetization) 

219 self.is_filled = True 

220 return self 

221 

222 kernel = fill 

223 

224 def _update_from_fillings(self, value, magnetization): 

225 """Update fillings. 

226 

227 Args: 

228 value: Fillings. 

229 magnetization: Magnetization. 

230 """ 

231 # Do not use the setter methods in this place to not trigger the setter effects 

232 # If f is a number use this occupation for all states 

233 if isinstance(value, numbers.Real) or isinstance(magnetization, numbers.Real): 

234 # Do not leave the states array empty when no electrons are present 

235 if self.Nelec <= 0: 

236 self._Nstate = 1 

237 self._f = xp.zeros((self.Nk, self.Nspin, 1)) 

238 # Always use the fractional fillings method if a magnetization is given 

239 elif isinstance(magnetization, numbers.Real): 

240 self._fractional_fillings(value, magnetization) 

241 elif self.Nspin == 1 or self.Nelec % 2 == self.spin % 2: 

242 self._integer_fillings(value) 

243 else: 

244 log.warning("Using fractional occupations.") 

245 self._fractional_fillings(value) 

246 # If f is an array redetermine all attributes from it 

247 else: 

248 self._f = value 

249 self._Nspin = len(value) 

250 self._Nstate = len(value[0]) 

251 self._charge += self.Nelec - int(xp.sum(value)) 

252 self.Nelec = int(xp.sum(value)) 

253 if self.Nspin == 1: 

254 self.spin = 0 

255 else: 

256 self.spin = abs(xp.sum(self.f[0] - self.f[1])) 

257 

258 def _integer_fillings(self, f): 

259 """Update fillings while maintaining integer occupations numbers. 

260 

261 Args: 

262 f: Fillings. 

263 """ 

264 # Determine the electrons per spin channel 

265 Nup = self.Nelec // self.Nspin + self.spin // self.Nspin + self.Nelec % self.Nspin 

266 Ndw = self.Nelec // self.Nspin - self.spin // self.Nspin 

267 elecs = xp.asarray([Nup, Ndw]) 

268 

269 # Get the number of states 

270 Nstate = int(xp.ceil(max(elecs / f))) 

271 # If no bands are set, set them now 

272 if self.bands == 0: 

273 self.bands = Nstate 

274 self._Nstate = Nstate 

275 if self.bands < Nstate: 

276 log.error("Number of bands is smaller than the number of valence electrons.") 

277 if self.bands == Nstate and self.smearing > 0: 

278 log.warning("Smearing has been enabled but no extra bands have been set.") 

279 # Set the number of empty states if no smearing is used 

280 if self.smearing == 0: 

281 self._Nstate = Nstate 

282 self._Nempty = self.bands - Nstate 

283 

284 # Simply build the occupations array 

285 self._f = f * xp.ones((self.Nspin, Nstate), dtype=int) 

286 # If we have filled too much correct it in the second spin channel 

287 # The following procedure ensures that we have no negative fillings 

288 rest = xp.sum(self._f) - self.Nelec 

289 i = 1 

290 while rest > 0: 

291 # If the occupation is greater than the rest we are done with removing electrons 

292 if self._f[-1, -i] >= rest: 

293 self._f[-1, -i] -= rest 

294 break 

295 # Otherwise zero out the state, update the rest, and move to the next state 

296 rest -= self._f[-1, -i] 

297 self._f[-1, -i] = 0 

298 i += 1 

299 

300 if self.smearing > 0: 

301 # Append extra states 

302 self._f = xp.hstack((self._f, xp.zeros((self.Nspin, self.bands - Nstate)))) 

303 self._Nstate = self.bands 

304 self._f = xp.vstack([xp.stack([self._f])] * self.Nk) 

305 

306 def _fractional_fillings(self, f, magnetization=None): 

307 """Update fillings while allowing fractional occupation numbers. 

308 

309 Args: 

310 f: Fillings. 

311 

312 Keyword Args: 

313 magnetization: Magnetization. 

314 """ 

315 # Determine the electrons per spin channel 

316 if magnetization is None: 

317 Nup = self.Nelec / self.Nspin + self.spin / self.Nspin 

318 Ndw = self.Nelec / self.Nspin - self.spin / self.Nspin 

319 else: 

320 # M = (Nup - Ndw) / N = (Nup - N + Nup) / N 

321 Nup = (magnetization * self.Nelec + self.Nelec) / 2 

322 Ndw = self.Nelec - Nup 

323 self.spin = abs(Nup - Ndw) 

324 elecs = xp.asarray([Nup, Ndw]) 

325 

326 # Get the number of states 

327 Nstate = int(xp.ceil(max(elecs / f))) 

328 # If no bands are set, set them now 

329 if self.bands == 0: 

330 self.bands = Nstate 

331 self._Nstate = Nstate 

332 if self.bands < Nstate: 

333 log.error("Number of bands is smaller than the number of valence electrons.") 

334 if self.bands == Nstate and self.smearing > 0: 

335 log.warning("Smearing has been enabled but no extra bands have been set.") 

336 # Set the number of empty states if no smearing is used 

337 if self.smearing == 0: 

338 self._Nstate = Nstate 

339 self._Nempty = self.bands - Nstate 

340 

341 # Simply build the occupations array 

342 self._f = f * xp.ones((self.Nspin, Nstate)) 

343 # If we have filled too much correct it in both spin channels 

344 # The following procedure ensures that we have no negative fillings 

345 rest = xp.sum(self._f, axis=1) - elecs 

346 for spin in range(self.Nspin): 

347 i = 1 

348 while rest[spin] > 0: 

349 # If the occupation is greater than the rest we are done with removing electrons 

350 if self._f[spin, -i] >= rest[spin]: 

351 self._f[spin, -i] -= rest[spin] 

352 break 

353 # Otherwise zero out the state, update the rest, and move to the next state 

354 rest[spin] -= self._f[spin, -i] 

355 self._f[spin, -i] = 0 

356 i += 1 

357 

358 if self.smearing > 0: 

359 # Append extra states 

360 self._f = xp.hstack((self._f, xp.zeros((self.Nspin, self.bands - Nstate)))) 

361 self._Nstate = self.bands 

362 self._f = xp.vstack([xp.stack([self._f])] * self.Nk) 

363 

364 def __repr__(self): 

365 """Print the parameters stored in the Occupations object.""" 

366 return ( 

367 f"Spin handling: {'un' if self.Nspin == 2 else ''}restricted\n" 

368 f"Number of electrons: {self.Nelec}\n" 

369 f"Spin: {self.spin}\n" 

370 f"Charge: {self.charge}\n" 

371 f"Number of bands: {self.bands}\n" 

372 f"Number of states: {self.Nstate}\n" 

373 f"Number of empty states: {self.Nempty}\n" 

374 f"Number of k-points: {self.Nk}\n" 

375 f"Smearing width: {self.smearing} Eh\n" 

376 f"Fillings: \n{self.f if self.is_filled else 'Not filled'}" 

377 ) 

378 

379 def smear(self, epsilon): 

380 """Update fillings according to a Fermi distribution. 

381 

382 Args: 

383 epsilon: Eigenenergies. 

384 

385 Returns: 

386 Efermi: Fermi energy. 

387 """ 

388 if self.smearing == 0: 

389 log.info("Smearing is set to zero, nothing to do.") 

390 return 0 

391 

392 Efermi = get_Efermi(self, epsilon) 

393 self._f = fermi_distribution(epsilon, Efermi, self.smearing) * 2 / self.Nspin 

394 return Efermi