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 smear(self, epsilon): 

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

226 

227 Args: 

228 epsilon: Eigenenergies. 

229 

230 Returns: 

231 Efermi: Fermi energy. 

232 """ 

233 if self.smearing == 0: 

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

235 return 0 

236 

237 Efermi = get_Efermi(self, epsilon) 

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

239 return Efermi 

240 

241 def _update_from_fillings(self, value, magnetization): 

242 """Update fillings. 

243 

244 Args: 

245 value: Fillings. 

246 magnetization: Magnetization. 

247 """ 

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

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

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

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

252 if self.Nelec <= 0: 

253 self._Nstate = 1 

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

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

256 elif isinstance(magnetization, numbers.Real): 

257 self._fractional_fillings(value, magnetization) 

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

259 self._integer_fillings(value) 

260 else: 

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

262 self._fractional_fillings(value) 

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

264 else: 

265 self._f = value 

266 self._Nspin = len(value) 

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

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

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

270 if self.Nspin == 1: 

271 self.spin = 0 

272 else: 

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

274 

275 def _integer_fillings(self, f): 

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

277 

278 Args: 

279 f: Fillings. 

280 """ 

281 # Determine the electrons per spin channel 

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

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

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

285 

286 # Get the number of states 

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

288 # If no bands are set, set them now 

289 if self.bands == 0: 

290 self.bands = Nstate 

291 self._Nstate = Nstate 

292 if self.bands < Nstate: 

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

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

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

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

297 if self.smearing == 0: 

298 self._Nstate = Nstate 

299 self._Nempty = self.bands - Nstate 

300 

301 # Simply build the occupations array 

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

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

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

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

306 i = 1 

307 while rest > 0: 

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

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

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

311 break 

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

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

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

315 i += 1 

316 

317 if self.smearing > 0: 

318 # Append extra states 

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

320 self._Nstate = self.bands 

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

322 

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

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

325 

326 Args: 

327 f: Fillings. 

328 

329 Keyword Args: 

330 magnetization: Magnetization. 

331 """ 

332 # Determine the electrons per spin channel 

333 if magnetization is None: 

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

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

336 else: 

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

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

339 Ndw = self.Nelec - Nup 

340 self.spin = abs(Nup - Ndw) 

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

342 

343 # Get the number of states 

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

345 # If no bands are set, set them now 

346 if self.bands == 0: 

347 self.bands = Nstate 

348 self._Nstate = Nstate 

349 if self.bands < Nstate: 

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

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

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

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

354 if self.smearing == 0: 

355 self._Nstate = Nstate 

356 self._Nempty = self.bands - Nstate 

357 

358 # Simply build the occupations array 

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

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

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

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

363 for spin in range(self.Nspin): 

364 i = 1 

365 while rest[spin] > 0: 

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

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

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

369 break 

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

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

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

373 i += 1 

374 

375 if self.smearing > 0: 

376 # Append extra states 

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

378 self._Nstate = self.bands 

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

380 

381 def __repr__(self): 

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

383 return ( 

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

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

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

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

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

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

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

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

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

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

394 )