Coverage for eminus/atoms.py: 98.10%

263 statements  

« prev     ^ index     » next       coverage.py v7.6.4, created at 2024-11-08 12:59 +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 

89 # ### Class properties ### 

90 

91 @property 

92 def atom(self): 

93 """Atom symbols.""" 

94 return self._atom 

95 

96 @atom.setter 

97 def atom(self, value): 

98 # Quick option to set the charge for single atoms 

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

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

101 self._atom = [atom] 

102 self._Natoms = 1 

103 self.Z = int(Z) 

104 else: 

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

106 if isinstance(value, str): 

107 self._atom = molecule2list(value) 

108 else: 

109 self._atom = value 

110 # Get the number of atoms and determine the charges 

111 self._Natoms = len(self._atom) 

112 self.Z = None 

113 

114 @property 

115 def pos(self): 

116 """Atom positions.""" 

117 return self._pos 

118 

119 @pos.setter 

120 def pos(self, value): 

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

122 self._pos = np.atleast_2d(value) 

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

124 msg = ( 

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

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

127 ) 

128 raise ValueError(msg) 

129 # The structure factor changes when changing pos 

130 self.is_built = False 

131 

132 @property 

133 def ecut(self): 

134 """Cut-off energy.""" 

135 return self._ecut 

136 

137 @ecut.setter 

138 def ecut(self, value): 

139 self._ecut = value 

140 # Calculate the sampling from the cut-off energy 

141 s = np.int64(norm(self.a, axis=0) / cutoff2gridspacing(value)) 

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

143 s = 2 * s + 1 

144 # Calculate a fast length to optimize the FFT calculations 

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

146 # The cell discretization changes when changing s or ecut 

147 self.is_built = False 

148 

149 @property 

150 def a(self): 

151 """Cell/Vacuum size.""" 

152 return self._a 

153 

154 @a.setter 

155 def a(self, value): 

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

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

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

159 # Otherwise scaled cell vectors are given 

160 else: 

161 self._a = np.asarray(value) 

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

163 if hasattr(self, "ecut"): 

164 self.ecut = self.ecut 

165 # Calculate the unit cell volume 

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

167 if hasattr(self, "kpts"): 

168 self.kpts.a = self._a 

169 # The cell changes when changing a 

170 self.is_built = False 

171 

172 @property 

173 def spin(self): 

174 """Number of unpaired electrons.""" 

175 return self.occ.spin 

176 

177 @spin.setter 

178 def spin(self, value): 

179 self.occ.spin = value 

180 

181 @property 

182 def charge(self): 

183 """System charge.""" 

184 return self.occ.charge 

185 

186 @charge.setter 

187 def charge(self, value): 

188 self.occ.charge = value 

189 

190 @property 

191 def unrestricted(self): 

192 """Enables unrestricted spin handling.""" 

193 return self.occ.Nspin == 2 

194 

195 @unrestricted.setter 

196 def unrestricted(self, value): 

197 if value is None: 

198 self.occ.Nspin = value 

199 else: 

200 self.occ.Nspin = value + 1 

201 

202 @property 

203 def center(self): 

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

205 return self._center 

206 

207 @center.setter 

208 def center(self, value): 

209 if isinstance(value, str): 

210 self._center = value.lower() 

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

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

213 else: 

214 self._center = value 

215 # Do nothing when recentering 

216 if self._center == "recentered": 

217 return 

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

219 # Rotate before shifting! 

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

221 I = inertia_tensor(self.pos) 

222 _, eigvecs = eigh(I) 

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

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

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

226 com = center_of_mass(self.pos) 

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

228 # The structure factor changes when changing pos 

229 self.is_built = False 

230 

231 @property 

232 def verbose(self): 

233 """Verbosity level.""" 

234 return self._verbose 

235 

236 @verbose.setter 

237 def verbose(self, value): 

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

239 if value is None: 

240 value = log.verbose 

241 self._verbose = get_level(value) 

242 self._log.verbose = self._verbose 

243 

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

245 

246 @property 

247 def f(self): 

248 """Occupation numbers per state.""" 

249 return self.occ.f 

250 

251 @f.setter 

252 def f(self, value): 

253 # Pass through to the Occupations object 

254 self.occ.f = value 

255 

256 @property 

257 def s(self): 

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

259 return self._s 

260 

261 @s.setter 

262 def s(self, value): 

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

264 if isinstance(value, numbers.Integral): 

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

266 self._s = np.asarray(value) 

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

268 # The cell discretization changes when changing s 

269 self.is_built = False 

270 

271 @property 

272 def Z(self): 

273 """Valence charge per atom.""" 

274 return self._Z 

275 

276 @Z.setter 

277 def Z(self, value): 

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

279 if isinstance(value, numbers.Integral): 

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

281 elif isinstance(value, dict): 

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

283 # Get the valence charges from the GTH files 

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

285 value = atom2charge(self.atom, value) 

286 self._Z = np.asarray(value) 

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

288 msg = ( 

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

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

291 ) 

292 raise ValueError(msg) 

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

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

295 

296 # ### Read-only properties ### 

297 

298 @property 

299 def Natoms(self): 

300 """Number of atoms.""" 

301 return self._Natoms 

302 

303 @property 

304 def Ns(self): 

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

306 return self._Ns 

307 

308 @property 

309 def Omega(self): 

310 """Unit cell volume.""" 

311 return self._Omega 

312 

313 @property 

314 def r(self): 

315 """Real-space sampling points.""" 

316 return self._r 

317 

318 @property 

319 def active(self): 

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

321 return self._active 

322 

323 @property 

324 def G(self): 

325 """G-vectors.""" 

326 return self._G 

327 

328 @property 

329 def G2(self): 

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

331 return self._G2 

332 

333 @property 

334 def G2c(self): 

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

336 return self._G2c 

337 

338 @property 

339 def Gk2(self): 

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

341 return self._Gk2 

342 

343 @property 

344 def Gk2c(self): 

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

346 return self._Gk2c 

347 

348 @property 

349 def Sf(self): 

350 """Structure factor per atom.""" 

351 return self._Sf 

352 

353 @property 

354 def dV(self): 

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

356 return self.Omega / self._Ns 

357 

358 @property 

359 def _atoms(self): 

360 """Return the Atoms object itself.""" 

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

362 return self 

363 

364 # ### Class methods ### 

365 

366 def build(self): 

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

368 self.kpts.build() 

369 self._sample_unit_cell() 

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

371 self.occ.fill() 

372 self.is_built = True 

373 return self 

374 

375 kernel = build 

376 

377 def recenter(self, center=None): 

378 """Recenter the system inside the cell. 

379 

380 Keyword Args: 

381 center: Point to center the system around. 

382 """ 

383 com = center_of_mass(self.pos) 

384 if center is None: 

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

386 else: 

387 center = np.asarray(center) 

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

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

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

391 self._center = "recentered" 

392 return self 

393 

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

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

396 

397 Args: 

398 k: k-point coordinates. 

399 

400 Keyword Args: 

401 wk: k-point weights. 

402 """ 

403 self.kpts.build() 

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

405 if wk is None: 

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

407 else: 

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

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

410 self.kpts._kmesh = None 

411 self.occ.wk = self.kpts.wk 

412 self._sample_unit_cell() 

413 return self 

414 

415 def clear(self): 

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

417 self._r = None # Sample points in cell 

418 self._G = None # G-vectors 

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

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

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

422 self._Sf = None # Structure factor 

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

424 return self 

425 

426 def _get_index_matrices(self): 

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

428 

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

430 

431 Returns: 

432 Index matrices. 

433 """ 

434 # Build index matrix M 

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

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

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

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

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

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

441 # Build index matrix N 

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

443 return M, N 

444 

445 def _sample_unit_cell(self): 

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

447 # Calculate index matrices 

448 M, N = self._get_index_matrices() 

449 # Build the real-space sampling 

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

451 # Build G-vectors 

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

453 # Calculate squared magnitudes of G-vectors 

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

455 # Calculate the G2 restriction 

456 self._active = [ 

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

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

459 ] 

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

461 # Calculate G+k-vectors 

462 self._Gk2 = np.asarray( 

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

464 ) 

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

466 # Calculate the structure factor per atom 

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

468 

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

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

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

472 self._Gk2c.append(self._G2c) 

473 

474 O = operators.O 

475 L = operators.L 

476 Linv = operators.Linv 

477 K = operators.K 

478 T = operators.T 

479 I = operators.I 

480 J = operators.J 

481 Idag = operators.Idag 

482 Jdag = operators.Jdag 

483 

484 def __repr__(self): 

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

486 out = "Atom Valence Position" 

487 for i in range(self.Natoms): 

488 out += ( 

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

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

491 ) 

492 return out