Coverage for eminus/atoms.py: 98.16%

1# SPDX-FileCopyrightText: 2021 The eminus developers

2# SPDX-License-Identifier: Apache-2.0

3"""Atoms class definition."""

5import math

6import numbers

8from scipy.fft import next_fast_len

10from . import backend as xp

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.

88 self.clear()

90 # ### Class properties ###

92 @property

93 def atom(self):

94 """Atom symbols."""

95 return self._atom

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 = xp.atleast_2d(xp.asarray(value, dtype=float))

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

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 if self.Natoms == 0:

131 log.warning("No atoms are placed in the unit cell.")

132 # The structure factor changes when changing pos

133 self.is_built = False

134

135 @property

136 def ecut(self):

137 """Cut-off energy."""

138 return self._ecut

139

140 @ecut.setter

141 def ecut(self, value):

142 self._ecut = value

143 # Calculate the sampling from the cut-off energy

144 s = xp.asarray(xp.linalg.norm(self.a, axis=1) / cutoff2gridspacing(value), dtype=int)

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

146 s = 2 * s + 1

147 # Calculate a fast length to optimize the FFT calculations

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

149 # The cell discretization changes when changing s or ecut

150 self.is_built = False

151

152 @property

153 def a(self):

154 """Cell/Vacuum size."""

155 return self._a

156

157 @a.setter

158 def a(self, value):

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

160 if xp.asarray(value).ndim <= 1:

161 self._a = xp.asarray(value, dtype=float) * xp.eye(3)

162 # Otherwise scaled cell vectors are given

163 else:

164 self._a = xp.asarray(value, dtype=float)

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

166 if hasattr(self, "ecut"):

167 self.ecut = self.ecut

168 # Calculate the unit cell volume

169 self._Omega = abs(xp.linalg.det(self._a))

170 if hasattr(self, "kpts"):

171 self.kpts.a = self._a

172 # The cell changes when changing a

173 self.is_built = False

174

175 @property

176 def spin(self):

177 """Number of unpaired electrons."""

178 return self.occ.spin

179

180 @spin.setter

181 def spin(self, value):

182 self.occ.spin = value

183

184 @property

185 def charge(self):

186 """System charge."""

187 return self.occ.charge

188

189 @charge.setter

190 def charge(self, value):

191 self.occ.charge = value

192

193 @property

194 def unrestricted(self):

195 """Enables unrestricted spin handling."""

196 return self.occ.Nspin == 2

197

198 @unrestricted.setter

199 def unrestricted(self, value):

200 if value is None:

201 self.occ.Nspin = value

202 else:

203 self.occ.Nspin = value + 1

204

205 @property

206 def center(self):

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

208 return self._center

209

210 @center.setter

211 def center(self, value):

212 if isinstance(value, str):

213 self._center = value.lower()

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

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

216 else:

217 self._center = value

218 # Do nothing when recentering

219 if self._center == "recentered":

220 return

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

222 # Rotate before shifting!

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

224 I = inertia_tensor(self.pos)

225 _, eigvecs = xp.linalg.eigh(I)

226 self.pos = (xp.linalg.inv(eigvecs) @ self.pos.T).T

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

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

229 com = center_of_mass(self.pos)

230 self.pos = self.pos - (com - xp.sum(self.a, axis=0) / 2)

231 # The structure factor changes when changing pos

232 self.is_built = False

233

234 @property

235 def verbose(self):

236 """Verbosity level."""

237 return self._verbose

238

239 @verbose.setter

240 def verbose(self, value):

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

242 if value is None:

243 value = log.verbose

244 self._verbose = get_level(value)

245 self._log.verbose = self._verbose

246

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

248

249 @property

250 def f(self):

251 """Occupation numbers per state."""

252 return self.occ.f

253

254 @f.setter

255 def f(self, value):

256 # Pass through to the Occupations object

257 self.occ.f = value

258

259 @property

260 def s(self):

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

262 return self._s

263

264 @s.setter

265 def s(self, value):

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

267 if isinstance(value, numbers.Integral):

268 value = value * xp.ones(3, dtype=int)

269 self._s = xp.asarray(value, dtype=int)

270 self._Ns = int(xp.prod(self._s))

271 # The cell discretization changes when changing s

272 self.is_built = False

273

274 @property

275 def Z(self):

276 """Valence charge per atom."""

277 return self._Z

278

279 @Z.setter

280 def Z(self, value):

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

282 if isinstance(value, numbers.Integral):

283 value = value * xp.ones(self.Natoms, dtype=int)

284 elif isinstance(value, dict):

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

286 # Get the valence charges from the GTH files

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

288 value = atom2charge(self.atom, value)

289 self._Z = xp.asarray(value, dtype=int)

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

291 msg = (

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

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

294 )

295 raise ValueError(msg)

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

297 self.occ.Nelec = xp.sum(self._Z) - self.charge

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

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

300 self.occ.bands = 0

301

302 # ### Read-only properties ###

303

304 @property

305 def Natoms(self):

306 """Number of atoms."""

307 return self._Natoms

308

309 @property

310 def Ns(self):

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

312 return self._Ns

313

314 @property

315 def Omega(self):

316 """Unit cell volume."""

317 return self._Omega

318

319 @property

320 def r(self):

321 """Real-space sampling points."""

322 return self._r

323

324 @property

325 def active(self):

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

327 return self._active

328

329 @property

330 def G(self):

331 """G-vectors."""

332 return self._G

333

334 @property

335 def G2(self):

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

337 return self._G2

338

339 @property

340 def G2c(self):

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

342 return self._G2c

343

344 @property

345 def Gk2(self):

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

347 return self._Gk2

348

349 @property

350 def Gk2c(self):

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

352 return self._Gk2c

353

354 @property

355 def Sf(self):

356 """Structure factor per atom."""

357 return self._Sf

358

359 @property

360 def dV(self):

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

362 return self.Omega / self._Ns

363

364 @property

365 def _atoms(self):

366 """The Atoms object itself."""

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

368 return self

369

370 # ### Class methods ###

371

372 def build(self):

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

374 self.kpts.build()

375 self._sample_unit_cell()

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

377 self.occ.fill()

378 self.is_built = True

379 return self

380

381 kernel = build

382

383 def recenter(self, center=None):

384 """Recenter the system inside the cell.

385

386 Keyword Args:

387 center: Point to center the system around.

388 """

389 com = center_of_mass(self.pos)

390 if center is None:

391 self.pos = self.pos - (com - xp.sum(self.a, axis=0) / 2)

392 else:

393 center = xp.asarray(center, dtype=float)

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

395 if self.Sf is not None:

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

397 self._Sf = xp.exp(1j * (self.G @ self.pos.T)).T

398 self._center = "recentered"

399 return self

400

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

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

403

404 Args:

405 k: k-point coordinates.

406

407 Keyword Args:

408 wk: k-point weights.

409 """

410 self.kpts.build()

411 self.kpts._k = xp.atleast_2d(xp.asarray(k, dtype=float))

412 if wk is None:

413 self.kpts._wk = xp.ones(len(self.kpts._k)) / len(self.kpts._k)

414 else:

415 self.kpts._wk = xp.asarray(wk, dtype=float)

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

417 self.kpts._kmesh = None

418 self.occ.wk = self.kpts.wk

419 self._sample_unit_cell()

420 return self

421

422 def clear(self):

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

424 self._r = None # Sample points in cell

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

426 self._G = None # G-vectors

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

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

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

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

431 self._Sf = None # Structure factor

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

433 return self

434

435 def _get_index_matrices(self):

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

437

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

439

440 Returns:

441 Index matrices.

442 """

443 # Build index matrix M

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

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

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

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

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

449 M = xp.permute_dims(xp.indices(self.s, dtype=float), (1, 2, 3, 0)).reshape(-1, 3)

450 # Build index matrix N

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

452 return M, N

453

454 def _sample_unit_cell(self):

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

456 # Calculate index matrices

457 M, N = self._get_index_matrices()

458 # Build the real-space sampling

459 self._r = M @ xp.linalg.inv(xp.diag(xp.astype(self.s, float))) @ self.a

460 # Build G-vectors

461 self._G = 2 * math.pi * N @ xp.linalg.inv(self.a.T)

462 # Calculate squared magnitudes of G-vectors

463 self._G2 = xp.linalg.norm(self.G, axis=1) ** 2

464 # Calculate the G2 restriction

465 self._active = [

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

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

468 ]

469 self._G2c = self._G2[xp.nonzero(2 * self.ecut >= self._G2)]

470 # Calculate G+k-vectors

471 self._Gk2 = xp.stack(

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

473 )

474 self._Gk2c = [self._Gk2[ik][self._active[ik]] for ik in range(self.kpts.Nk)]

475 # Calculate the structure factor per atom

476 self._Sf = xp.exp(1j * (self.G @ self.pos.T)).T

477

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

479 self._active.append(xp.nonzero(2 * self.ecut >= self._G2))

480 self._Gk2 = xp.vstack((self._Gk2, self._G2))

481 self._Gk2c.append(self._G2c)

482

483 O = operators.O

484 L = operators.L

485 Linv = operators.Linv

486 K = operators.K

487 T = operators.T

488 I = operators.I

489 J = operators.J

490 Idag = operators.Idag

491 Jdag = operators.Jdag

492

493 def __repr__(self):

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

495 out = "Atom Valence Position"

496 for i in range(self.Natoms):

497 out += (

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

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

500 )

501 return out