Coverage for eminus/atoms.py: 98.15%

270 statements  

« prev     ^ index     » next       coverage.py v7.8.2, created at 2025-06-02 10:16 +0000

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Atoms class definition.""" 

4 

5import numbers 

6 

7import numpy as np 

8from scipy.fft import next_fast_len 

9from scipy.linalg import det, eigh, inv, norm 

10 

11from . import operators 

12from .kpoints import KPoints 

13from .logger import create_logger, get_level, log 

14from .occupations import Occupations 

15from .tools import center_of_mass, cutoff2gridspacing, inertia_tensor 

16from .utils import atom2charge, BaseObject, molecule2list 

17 

18 

19class Atoms(BaseObject): 

20 """Atoms object that holds all system and cell parameters. 

21 

22 Args: 

23 atom: Atom symbols. 

24 

25 A string can be given, e.g., with :code:`CH4` that will be parsed to 

26 :code:`["C", "H", "H", "H", "H"]`. When calculating atoms one can directly provide the 

27 charge, e.g., with :code:`Li-q3`. 

28 pos: Atom positions. 

29 

30 Keyword Args: 

31 ecut: Cut-off energy. 

32 

33 Defaults to 30 Eh (ca. 816 eV). 

34 a: Cell size or vacuum size. 

35 

36 Floats will create a cubic unit cell. Defaults to a 20 a0 (ca. 10.5 A) cubic cell. 

37 Scaled lattice vectors can be given to build a custom cell. 

38 spin: Number of unpaired electrons. 

39 

40 This is the difference between the number of up and down electrons. 

41 charge: Charge of the system. 

42 unrestricted: Handling of spin. 

43 

44 :code:`False` for restricted, :code:`True` for unrestricted, and :code:`None` for 

45 automatic detection. 

46 center: Center the system inside the cell. 

47 

48 Aligns the geometric center of mass with the center of the call and rotates the system, 

49 such that its geometric moment of inertia aligns with the coordinate axes. Can be one of 

50 bool, "shift", and "rotate". 

51 verbose: Level of output. 

52 

53 Can be one of "critical", "error", "warning", "info" (default), or "debug". An integer 

54 value can be used as well, where larger numbers mean more output, starting from 0. 

55 None will use the global logger verbosity value. 

56 """ 

57 

58 def __init__( 

59 self, 

60 atom, 

61 pos, 

62 ecut=30, 

63 a=20, 

64 spin=None, 

65 charge=0, 

66 unrestricted=None, 

67 center=False, 

68 verbose=None, 

69 ): 

70 """Initialize the Atoms object.""" 

71 # Set the input parameters (the ordering is important) 

72 self._log = create_logger(self) #: Logger object. 

73 self.verbose = verbose #: Verbosity level. 

74 self.occ = Occupations() #: Occupations object. 

75 self.atom = atom #: Atom symbols. 

76 self.pos = pos #: Atom positions. 

77 self.a = a #: Cell/Vacuum size. 

78 self.ecut = ecut #: Cut-off energy. 

79 self.center = center #: Enables centering the system in the cell. 

80 self.charge = charge #: System charge. 

81 self.spin = spin #: Number of unpaired electrons. 

82 self.unrestricted = unrestricted #: Enables unrestricted spin handling. 

83 self.kpts = KPoints("sc", self.a) #: KPoints object. 

84 

85 # Initialize other attributes 

86 self.occ.fill() #: Fill states from the given input. 

87 self.is_built = False #: Determines the Atoms object build status. 

88 self.clear() 

89 

90 # ### Class properties ### 

91 

92 @property 

93 def atom(self): 

94 """Atom symbols.""" 

95 return self._atom 

96 

97 @atom.setter 

98 def atom(self, value): 

99 # Quick option to set the charge for single atoms 

100 if isinstance(value, str) and "-q" in value: 

101 atom, Z = value.split("-q") 

102 self._atom = [atom] 

103 self._Natoms = 1 

104 self.Z = int(Z) 

105 else: 

106 # If a string, i.e., chemical formula is given convert it to a list of chemical symbols 

107 if isinstance(value, str): 

108 self._atom = molecule2list(value) 

109 else: 

110 self._atom = value 

111 # Get the number of atoms and determine the charges 

112 self._Natoms = len(self._atom) 

113 self.Z = None 

114 

115 @property 

116 def pos(self): 

117 """Atom positions.""" 

118 return self._pos 

119 

120 @pos.setter 

121 def pos(self, value): 

122 # We need atom positions as a two-dimensional array 

123 self._pos = np.atleast_2d(value) 

124 if self.Natoms != len(self._pos): 

125 msg = ( 

126 f"Mismatch between number of atoms ({self.Natoms}) and number of " 

127 f"coordinates ({len(self._pos)})." 

128 ) 

129 raise ValueError(msg) 

130 # The structure factor changes when changing pos 

131 self.is_built = False 

132 

133 @property 

134 def ecut(self): 

135 """Cut-off energy.""" 

136 return self._ecut 

137 

138 @ecut.setter 

139 def ecut(self, value): 

140 self._ecut = value 

141 # Calculate the sampling from the cut-off energy 

142 s = np.int64(norm(self.a, axis=1) / cutoff2gridspacing(value)) 

143 # Multiply by two and add one to match PWDFT.jl 

144 s = 2 * s + 1 

145 # Calculate a fast length to optimize the FFT calculations 

146 self.s = [next_fast_len(i) for i in s] 

147 # The cell discretization changes when changing s or ecut 

148 self.is_built = False 

149 

150 @property 

151 def a(self): 

152 """Cell/Vacuum size.""" 

153 return self._a 

154 

155 @a.setter 

156 def a(self, value): 

157 # Build a cubic cell if a number or 1d-array is given 

158 if np.asarray(value).ndim <= 1: 

159 self._a = value * np.eye(3) 

160 # Otherwise scaled cell vectors are given 

161 else: 

162 self._a = np.asarray(value) 

163 # Update ecut and s if it has been set before 

164 if hasattr(self, "ecut"): 

165 self.ecut = self.ecut 

166 # Calculate the unit cell volume 

167 self._Omega = abs(det(self._a)) 

168 if hasattr(self, "kpts"): 

169 self.kpts.a = self._a 

170 # The cell changes when changing a 

171 self.is_built = False 

172 

173 @property 

174 def spin(self): 

175 """Number of unpaired electrons.""" 

176 return self.occ.spin 

177 

178 @spin.setter 

179 def spin(self, value): 

180 self.occ.spin = value 

181 

182 @property 

183 def charge(self): 

184 """System charge.""" 

185 return self.occ.charge 

186 

187 @charge.setter 

188 def charge(self, value): 

189 self.occ.charge = value 

190 

191 @property 

192 def unrestricted(self): 

193 """Enables unrestricted spin handling.""" 

194 return self.occ.Nspin == 2 

195 

196 @unrestricted.setter 

197 def unrestricted(self, value): 

198 if value is None: 

199 self.occ.Nspin = value 

200 else: 

201 self.occ.Nspin = value + 1 

202 

203 @property 

204 def center(self): 

205 """Enables centering the system in the cell.""" 

206 return self._center 

207 

208 @center.setter 

209 def center(self, value): 

210 if isinstance(value, str): 

211 self._center = value.lower() 

212 if self._center not in {"rotate", "shift", "recentered"}: 

213 log.error(f"{self._center} is not a recognized center method.") 

214 else: 

215 self._center = value 

216 # Do nothing when recentering 

217 if self._center == "recentered": 

218 return 

219 # Center system such that the geometric inertia tensor will be diagonal 

220 # Rotate before shifting! 

221 if self._center is True or self._center == "rotate": 

222 I = inertia_tensor(self.pos) 

223 _, eigvecs = eigh(I) 

224 self.pos = (inv(eigvecs) @ self.pos.T).T 

225 # Shift system such that its geometric center of mass is in the center of the cell 

226 if self._center is True or self._center == "shift": 

227 com = center_of_mass(self.pos) 

228 self.pos = self.pos - (com - np.sum(self.a, axis=0) / 2) 

229 # The structure factor changes when changing pos 

230 self.is_built = False 

231 

232 @property 

233 def verbose(self): 

234 """Verbosity level.""" 

235 return self._verbose 

236 

237 @verbose.setter 

238 def verbose(self, value): 

239 # If no verbosity is given use the global verbosity level 

240 if value is None: 

241 value = log.verbose 

242 self._verbose = get_level(value) 

243 self._log.verbose = self._verbose 

244 

245 # ### Class properties with a setter outside of the init method ### 

246 

247 @property 

248 def f(self): 

249 """Occupation numbers per state.""" 

250 return self.occ.f 

251 

252 @f.setter 

253 def f(self, value): 

254 # Pass through to the Occupations object 

255 self.occ.f = value 

256 

257 @property 

258 def s(self): 

259 """Real-space sampling of the cell.""" 

260 return self._s 

261 

262 @s.setter 

263 def s(self, value): 

264 # Choose the same sampling for every direction if an integer is given 

265 if isinstance(value, numbers.Integral): 

266 value = value * np.ones(3, dtype=int) 

267 self._s = np.asarray(value) 

268 self._Ns = int(np.prod(self._s)) 

269 # The cell discretization changes when changing s 

270 self.is_built = False 

271 

272 @property 

273 def Z(self): 

274 """Valence charge per atom.""" 

275 return self._Z 

276 

277 @Z.setter 

278 def Z(self, value): 

279 # Assume same charges for all atoms if an integer is given 

280 if isinstance(value, numbers.Integral): 

281 value = value * np.ones(self.Natoms, dtype=int) 

282 elif isinstance(value, dict): 

283 value = [value[ia] for ia in self.atom] 

284 # Get the valence charges from the GTH files 

285 elif value is None or isinstance(value, str): 

286 value = atom2charge(self.atom, value) 

287 self._Z = np.asarray(value) 

288 if self.Natoms != len(self._Z): 

289 msg = ( 

290 f"Mismatch between number of atoms ({self.Natoms}) and number of " 

291 f"charges ({len(self._Z)})." 

292 ) 

293 raise ValueError(msg) 

294 # Get the number of calculated electrons and pass it to occ 

295 self.occ.Nelec = np.sum(self._Z) - self.charge 

296 if self.occ.Nspin and self.occ.bands < self.occ.Nelec * self.occ.Nspin // 2: 

297 log.warning("The number of bands is too small, reset to the minimally needed amount.") 

298 self.occ.bands = 0 

299 

300 # ### Read-only properties ### 

301 

302 @property 

303 def Natoms(self): 

304 """Number of atoms.""" 

305 return self._Natoms 

306 

307 @property 

308 def Ns(self): 

309 """Number of real-space grid points.""" 

310 return self._Ns 

311 

312 @property 

313 def Omega(self): 

314 """Unit cell volume.""" 

315 return self._Omega 

316 

317 @property 

318 def r(self): 

319 """Real-space sampling points.""" 

320 return self._r 

321 

322 @property 

323 def active(self): 

324 """Mask for active G-vectors.""" 

325 return self._active 

326 

327 @property 

328 def G(self): 

329 """G-vectors.""" 

330 return self._G 

331 

332 @property 

333 def G2(self): 

334 """Squared magnitudes of G-vectors.""" 

335 return self._G2 

336 

337 @property 

338 def G2c(self): 

339 """Truncated squared magnitudes of G-vectors.""" 

340 return self._G2c 

341 

342 @property 

343 def Gk2(self): 

344 """Squared magnitudes of G+k-vectors.""" 

345 return self._Gk2 

346 

347 @property 

348 def Gk2c(self): 

349 """Truncated squared magnitudes of G+k-vectors.""" 

350 return self._Gk2c 

351 

352 @property 

353 def Sf(self): 

354 """Structure factor per atom.""" 

355 return self._Sf 

356 

357 @property 

358 def dV(self): 

359 """Volume element to multiply when integrating field properties.""" 

360 return self.Omega / self._Ns 

361 

362 @property 

363 def _atoms(self): 

364 """Return the Atoms object itself.""" 

365 # This way we can access the object from Atoms and SCF classes with the same code 

366 return self 

367 

368 # ### Class methods ### 

369 

370 def build(self): 

371 """Build all parameters of the Atoms object.""" 

372 self.kpts.build() 

373 self._sample_unit_cell() 

374 self.occ.wk = self.kpts.wk # Pass the weights of k-points to the Occupations object 

375 self.occ.fill() 

376 self.is_built = True 

377 return self 

378 

379 kernel = build 

380 

381 def recenter(self, center=None): 

382 """Recenter the system inside the cell. 

383 

384 Keyword Args: 

385 center: Point to center the system around. 

386 """ 

387 com = center_of_mass(self.pos) 

388 if center is None: 

389 self.pos = self.pos - (com - np.sum(self.a, axis=0) / 2) 

390 else: 

391 center = np.asarray(center) 

392 self.pos = self.pos - (com - center) 

393 if self.Sf is not None: 

394 # Recalculate the structure factor since it depends on the atom positions 

395 self._Sf = np.exp(1j * self.G @ self.pos.T).T 

396 self._center = "recentered" 

397 return self 

398 

399 def set_k(self, k, wk=None): 

400 """Interface to set custom k-points. 

401 

402 Args: 

403 k: k-point coordinates. 

404 

405 Keyword Args: 

406 wk: k-point weights. 

407 """ 

408 self.kpts.build() 

409 self.kpts._k = np.atleast_2d(k) 

410 if wk is None: 

411 self.kpts._wk = np.ones(len(self.kpts._k)) / len(self.kpts._k) 

412 else: 

413 self.kpts._wk = np.asarray(wk) 

414 self.kpts._Nk = len(self.kpts._wk) 

415 self.kpts._kmesh = None 

416 self.occ.wk = self.kpts.wk 

417 self._sample_unit_cell() 

418 return self 

419 

420 def clear(self): 

421 """Initialize or clear parameters that will be built out of the inputs.""" 

422 self._r = None # Sample points in cell 

423 self._active = None # Mask for active G-vectors 

424 self._G = None # G-vectors 

425 self._G2 = None # Squared magnitudes of G-vectors 

426 self._G2c = None # Truncated squared magnitudes of G-vectors 

427 self._Gk2 = None # Squared magnitudes of G+k-vectors 

428 self._Gk2c = None # Truncated squared magnitudes of G+k-vectors 

429 self._Sf = None # Structure factor 

430 self.is_built = False # Flag to determine if the object was built or not 

431 return self 

432 

433 def _get_index_matrices(self): 

434 """Build index matrices M and N to build the real and reciprocal space samplings. 

435 

436 The matrices are using C ordering (the last index is the fastest). 

437 

438 Returns: 

439 Index matrices. 

440 """ 

441 # Build index matrix M 

442 # ms = np.arange(self._Ns) 

443 # m1 = np.floor(ms / (self.s[2] * self.s[1])) % self.s[0] 

444 # m2 = np.floor(ms / self.s[2]) % self.s[1] 

445 # m3 = ms % self.s[2] 

446 # M = np.column_stack((m1, m2, m3)) 

447 M = np.indices(self.s).transpose((1, 2, 3, 0)).reshape((-1, 3)) 

448 # Build index matrix N 

449 N = M - (self.s / 2 < M) * self.s 

450 return M, N 

451 

452 def _sample_unit_cell(self): 

453 """Build the real-space sampling and all G-vector parameters.""" 

454 # Calculate index matrices 

455 M, N = self._get_index_matrices() 

456 # Build the real-space sampling 

457 self._r = M @ inv(np.diag(self.s)) @ self.a 

458 # Build G-vectors 

459 self._G = 2 * np.pi * N @ inv(self.a.T) 

460 # Calculate squared magnitudes of G-vectors 

461 self._G2 = norm(self.G, axis=1) ** 2 

462 # Calculate the G2 restriction 

463 self._active = [ 

464 np.nonzero(2 * self.ecut >= norm(self.G + self.kpts.k[ik], axis=1) ** 2) 

465 for ik in range(self.kpts.Nk) 

466 ] 

467 self._G2c = self._G2[np.nonzero(2 * self.ecut >= self._G2)] 

468 # Calculate G+k-vectors 

469 self._Gk2 = np.asarray( 

470 [norm(self.G + self.kpts.k[ik], axis=1) ** 2 for ik in range(self.kpts.Nk)] 

471 ) 

472 self._Gk2c = [self.Gk2[ik][self._active[ik]] for ik in range(self.kpts.Nk)] 

473 # Calculate the structure factor per atom 

474 self._Sf = np.exp(1j * self.G @ self.pos.T).T 

475 

476 # Create the grid used for the non-wave function fields and append it to the end 

477 self._active.append(np.nonzero(2 * self.ecut >= self._G2)) 

478 self._Gk2 = np.vstack((self._Gk2, self._G2)) 

479 self._Gk2c.append(self._G2c) 

480 

481 O = operators.O 

482 L = operators.L 

483 Linv = operators.Linv 

484 K = operators.K 

485 T = operators.T 

486 I = operators.I 

487 J = operators.J 

488 Idag = operators.Idag 

489 Jdag = operators.Jdag 

490 

491 def __repr__(self): 

492 """Print the parameters stored in the Atoms object.""" 

493 out = "Atom Valence Position" 

494 for i in range(self.Natoms): 

495 out += ( 

496 f"\n{self.atom[i]:>3} {self.Z[i]:>6} " 

497 f"{self.pos[i, 0]:10.5f} {self.pos[i, 1]:10.5f} {self.pos[i, 2]:10.5f}" 

498 ) 

499 return out