Coverage for eminus/atoms.py: 98.12%

1# SPDX-FileCopyrightText: 2021 The eminus developers

2# SPDX-License-Identifier: Apache-2.0

3"""Atoms class definition."""

5import numbers

7import numpy as np

8from scipy.fft import next_fast_len

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

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

19class Atoms(BaseObject):

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

22 Args:

23 atom: Atom symbols.

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.

30 Keyword Args:

31 ecut: Cut-off energy.

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

34 a: Cell size or vacuum size.

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.

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

41 charge: Charge of the system.

42 unrestricted: Handling of spin.

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

45 automatic detection.

46 center: Center the system inside the cell.

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.

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 """

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.

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.

89 # ### Class properties ###

91 @property

92 def atom(self):

93 """Atom symbols."""

94 return self._atom

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=1) / 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 if self.occ.Nspin and self.occ.bands < self.occ.Nelec * self.occ.Nspin // 2:

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

297 self.occ.bands = 0

298

299 # ### Read-only properties ###

300

301 @property

302 def Natoms(self):

303 """Number of atoms."""

304 return self._Natoms

305

306 @property

307 def Ns(self):

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

309 return self._Ns

310

311 @property

312 def Omega(self):

313 """Unit cell volume."""

314 return self._Omega

315

316 @property

317 def r(self):

318 """Real-space sampling points."""

319 return self._r

320

321 @property

322 def active(self):

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

324 return self._active

325

326 @property

327 def G(self):

328 """G-vectors."""

329 return self._G

330

331 @property

332 def G2(self):

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

334 return self._G2

335

336 @property

337 def G2c(self):

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

339 return self._G2c

340

341 @property

342 def Gk2(self):

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

344 return self._Gk2

345

346 @property

347 def Gk2c(self):

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

349 return self._Gk2c

350

351 @property

352 def Sf(self):

353 """Structure factor per atom."""

354 return self._Sf

355

356 @property

357 def dV(self):

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

359 return self.Omega / self._Ns

360

361 @property

362 def _atoms(self):

363 """Return the Atoms object itself."""

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

365 return self

366

367 # ### Class methods ###

368

369 def build(self):

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

371 self.kpts.build()

372 self._sample_unit_cell()

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

374 self.occ.fill()

375 self.is_built = True

376 return self

377

378 kernel = build

379

380 def recenter(self, center=None):

381 """Recenter the system inside the cell.

382

383 Keyword Args:

384 center: Point to center the system around.

385 """

386 com = center_of_mass(self.pos)

387 if center is None:

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

389 else:

390 center = np.asarray(center)

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

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

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

394 self._center = "recentered"

395 return self

396

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

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

399

400 Args:

401 k: k-point coordinates.

402

403 Keyword Args:

404 wk: k-point weights.

405 """

406 self.kpts.build()

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

408 if wk is None:

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

410 else:

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

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

413 self.kpts._kmesh = None

414 self.occ.wk = self.kpts.wk

415 self._sample_unit_cell()

416 return self

417

418 def clear(self):

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

420 self._r = None # Sample points in cell

421 self._G = None # G-vectors

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

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

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

425 self._Sf = None # Structure factor

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

427 return self

428

429 def _get_index_matrices(self):

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

431

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

433

434 Returns:

435 Index matrices.

436 """

437 # Build index matrix M

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

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

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

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

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

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

444 # Build index matrix N

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

446 return M, N

447

448 def _sample_unit_cell(self):

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

450 # Calculate index matrices

451 M, N = self._get_index_matrices()

452 # Build the real-space sampling

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

454 # Build G-vectors

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

456 # Calculate squared magnitudes of G-vectors

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

458 # Calculate the G2 restriction

459 self._active = [

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

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

462 ]

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

464 # Calculate G+k-vectors

465 self._Gk2 = np.asarray(

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

467 )

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

469 # Calculate the structure factor per atom

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

471

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

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

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

475 self._Gk2c.append(self._G2c)

476

477 O = operators.O

478 L = operators.L

479 Linv = operators.Linv

480 K = operators.K

481 T = operators.T

482 I = operators.I

483 J = operators.J

484 Idag = operators.Idag

485 Jdag = operators.Jdag

486

487 def __repr__(self):

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

489 out = "Atom Valence Position"

490 for i in range(self.Natoms):

491 out += (

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

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

494 )

495 return out