Coverage for dynasor/modes/mode

1from __future__ import annotations

2from typing import Optional

4import numpy as np

5import warnings

6import pickle

8from ase.calculators.singlepoint import SinglePointCalculator

9from ase.units import fs

10from ase import Atoms

11from numpy.typing import NDArray

13from dynasor.units import radians_per_fs_to_THz

14from dynasor.tools.structures import get_displacements

15from .tools import get_dynamical_matrix, group_eigvals, symmetrize_eigenvectors

16from .qpoint import QPoint

17from .atoms import Prim, Supercell

18from ..qpoints.tools import get_commensurate_lattice_points

21class ModeProjector:

22 """

23 The :class:`ModeProjector` maps between real atomic displacements `u` and

24 complex mode coordinates `Q`.

26 Some special python methods are implemented. The `__str__` and `__repr__`

27 provides useful info. :class:`QPoint` objects are representations of a

28 single q-point and associated information and can be accessed either by

29 call providing a reduced wavevector

31 >>> mp((1/2, 0, 0)) # doctest: +SKIP

33 or by index corresponding to reduced q-point accessible from

34 :attr:`~ModeProjector.q_reduced`

36 >>> mp[2] # doctest: +SKIP

38 In addition to mode coordinates `Q` the class can also map the atomic

39 velocities `v` to mode momenta `P` as well as atomic forces `f` to mode

40 forces `F`. The class can also map back. This mapping is done using

41 getters and setters. Internally only Q, P and F are stored

43 >>> Q = mp.get_Q() # doctest: +SKIP

44 >>> mp.set_u(u) # doctest: +SKIP

46 In addition, the forces corresponding to the harmonic forces can be

47 accessed by :meth:`~ModeProjector.get_f_harmonic()` and

48 :meth:`~ModeProjector.get_F_harmonic()`. For ASE Atoms objects the

49 displacements etc. can be updated and applied by

51 >>> mp.update_from_atoms(atoms) # doctest: +SKIP

52 >>> atoms = mp.get_atoms(harmonic_forces=False) # doctest: +SKIP

54 The shapes for each property uses the following variables

56 * `N`: Number of atoms in supercell

57 * `Nu`: unit cells (`N/Np`)

58 * `Np`: primitive basis atoms (`N/Nu`)

59 * `Nb`: bands (`Np*3`)

60 * `Nq`: q-points (`Nu`)

62 Please consult the documentation or the specific getters and setters to see

63 the exact transformations used.

65 Units

66 ^^^^^

67 The internal units in dynasor are Å, fs and eV. All frequencies are angular

68 (a.k.a. the "physicist's convention" with 2π included). These are the units

69 dynasor will expect and return. In, e.g., print functions conventional units

70 such fs, Å, THz, Da, meV are commonly used.

72 **Mass:**

73 The internal unit choice (eV, Å, fs) means that the mass unit is not Dalton

74 but rather 0.009648533290731906 Da.

75 We refer to this unit as the "dynasor mass unit" (dmu),

76 i.e., 1 Da = 103.64269572045423 dmu.

77 As a user you will only see this unit in the output of :class:`ModeProjector` objects.

78 Masses provided via, e.g., ASE Atoms objects are converted internally.

80 **Waves:**

81 dynasor reports and expects spatial (angular) frequencies in rad/Å and temporal (angular)

82 frequencies in rad/fs. This follows the often-used convention in physics to include the

83 factor of 2π in the wave vectors. For instance the wavelength is given by λ=2π/q.

85 **Mode amplitudes:**

86 Mode amplitudes are reported in Å√dmu = fs√eV.

88 **Velocities:**

89 For modes the momenta are reported in Å√dmu/fs or just √eV

90 while atomic velocities are reported in Å/fs.

92 **Mode forces:**

93 The force is defined as the derivative of the momenta with respect to time

94 so the unit used when reporting mode forces is Å√dmu/fs² (or √eV/fs).

96 Internal arrays

97 ^^^^^^^^^^^^^^^

98 For the curious, the internal data arrays are

100 * :attr:`primitive`, :attr:`supercell`, :attr:`force_constants` (input)

101 * :attr:`_q`, :attr:`q_minus` (reduced q-points and which q-points are related by inversion)

102 * :attr:`_D`, :attr:`_w2`, :attr:`_W`

103 (dynamical matrices, frequencies (ev/Å²/Da), polarization vectors)

104 * :attr:`_X` (eigenmodes which are mass weighted "polarization vectors" in the supercell)

105 """

106 def __init__(self, primitive: Atoms, supercell: Atoms, force_constants: NDArray[float]):

107 """The mode projector is initialized by a primitive cell and a

108 supercell as well as harmonic force constants.

109

110 The force constants are assumed to be in units of eV/Å² as returned

111 from phonopy. Be careful about the permutations when working with force

112 constants and atoms object from different codes.

113

114 Parameters

115 ----------

116 primitive

117 Primitive cell. Note that the masses are stored internally as

118 Dalton in ASE but will be converted to the internal dynasor

119 mass unit (dmu).

120 supercell

121 Ideal supercell corresponding to the force constants.

122 force_constants

123 Force constants for the supercell in eV/Å² as a `(N, N, 3, 3)` array

124 where `N` is `len(supercell)`.

125 """

126 if len(primitive) == len(supercell):

127 warnings.warn('Primitive and supercell have the same size')

128 elif len(primitive) > len(supercell):

129 raise ValueError('Primitive cell larger than supercell')

130 elif not (len(supercell) / len(primitive)).is_integer():

131 raise ValueError('supercell size is not multiple of primitive size')

132

133 if len(supercell) != len(force_constants):

134 raise ValueError('force constants shape is not compatible with supercell size')

135

136 if force_constants.shape != (len(supercell), len(supercell), 3, 3):

137 raise ValueError('force constants shape should be (N, N, 3, 3)')

138

139 self.primitive = Prim(primitive)

140 self.supercell = Supercell(supercell, primitive)

141 self.force_constants = force_constants

142

143 # Find q-points in reduced primitive cell coordinates

144 q_integer = get_commensurate_lattice_points(self.supercell.P.T)

145 q_reduced = np.dot(q_integer, self.supercell.P_inv.T)

146 self._q = np.array(sorted(tuple(q) for q in q_reduced))

147

148 # The equivalent q-point corresponding to -q

149 self._q_minus = [[tuple(q) for q in self._q].index(tuple((-q) % 1)) for q in self._q]

150

151 # Construct dynamical matrix and diagonalize at each q-point

152 self._D, self._w2, self._W = [], [], []

153 for qi, q in enumerate(self._q):

154

155 D = get_dynamical_matrix(

156 self.force_constants, self.supercell.offsets, self.supercell.indices,

157 q.astype(np.float64))

158

159 if qi == self.q_minus[qi]:

160 assert np.allclose(D.imag, 0)

161 D = D.real

162

163 D = np.einsum('ijab,i,j->ijab',

164 D, self.primitive.masses**-0.5, self.primitive.masses**-0.5)

165 D_matrix = D.transpose(0, 2, 1, 3).reshape(-1, self.primitive.n_atoms * 3)

166 assert np.allclose(D_matrix, D_matrix.T.conj())

167 w2, W = np.linalg.eigh(D_matrix)

168 W = W.T.reshape(-1, self.primitive.n_atoms, 3)

169

170 self._D.append(D)

171 self._w2.append(w2)

172 self._W.append(W)

173

174 self._D = np.array(self._D)

175 self._w2 = np.array(self._w2)

176 self._W = np.array(self._W)

177

178 # Post check basic symmetries, group eigenvalues and try to make degenerate modes nicer

179 for q, q_minus in enumerate(self.q_minus):

180 q_minus = self.q_minus[q]

181

182 assert np.allclose(self._D[q], self._D[q_minus].conj())

183 assert np.allclose(self._w2[q], self._w2[q_minus])

184

185 # tolerances for grouping and sorting eigenvalues and eigenvectors

186 group_decimals = 12

187 tolerance = 10**(-group_decimals)

188 # Keep sorting coarser than the convergence tolerance so tiny SVD

189 # roundoff cannot change the order of equivalent modes.

190 sort_decimals = 10

191

192 for group in group_eigvals(self._w2[q], tolerance**0.5):

193 W = symmetrize_eigenvectors(self._W[q, group])

194

195 # Try to order them

196 W_sort = W.copy().transpose(0, 2, 1).reshape(len(W), -1)

197 # abs is because we want to consider the magnitude

198 # - (minus) basically reverts the sort order to place largest first

199 # T is just because how lexsort works, we want to consider each

200 # atom and direction as a key for the bands

201 # -1 is because we want to make the x-direction of the first

202 # atom the most significant key

203 # At the end the first band should have the largest magnitude

204 # for the first atom in x

205 # Suppress numerical noise in the mode signatures before sorting.

206 # Keep this precision deliberately below the factor-analysis

207 # convergence noise so tiny SVD/BLAS differences cannot permute

208 # otherwise equivalent degenerate modes.

209 argsort = np.lexsort(np.round(-np.abs(W_sort).T[::-1], sort_decimals))

210 self._W[q, group] = W[argsort]

211

212 self._W[q_minus] = self._W[q].conj()

213

214 # Construct supercell projection matrix: q_ks = X_ksna u_na.

215 # phase and mass depend only on (q, index), so compute once and

216 # broadcast against a gather of self._W.

217 indices = self.supercell.indices

218

219 # self._q has dtype=object (Fraction); cast to float for the phase

220 q_float = self._q.astype(np.float64)

221 phase = np.exp(-1j * 2*np.pi * q_float @ self.supercell.offsets.T) # (n_q, n_super)

222 mass_sqrt = self.primitive.masses[indices]**0.5 # (n_super,)

223 # Build _X in place to avoid extra full-size temporaries (halves peak

224 # memory). astype pins the dtype (no-op unless _W is real, e.g. Gamma).

225 self._X = self._W[:, :, indices, :].astype(np.complex128, copy=False)

226 np.conjugate(self._X, out=self._X)

227 self._X *= phase[:, None, :, None]

228 self._X *= mass_sqrt[None, None, :, None]

229 self._X /= (self.supercell.n_atoms / self.primitive.n_atoms)**0.5

230

231 # Init arrays to hold Q, P and F

232 self._Q = np.zeros((len(self._q), self.primitive.n_atoms*3), dtype=np.complex128)

233 self._P = np.zeros_like(self._Q)

234 self._F = np.zeros_like(self._Q)

235

236 def __str__(self):

237 strings = ['### ModeProjector ###']

238 strings += [f'{self.supercell}']

239 strings += [f'{self.primitive}']

240 string = '\n'.join(strings)

241

242 # ASCII DOS!

243 width = 80

244 height = 24

245 dos = np.full((height, width), ' ')

246

247 THz = self.omegas * radians_per_fs_to_THz

248

249 hist, bins = np.histogram(THz.flat, bins=width)

250

251 for i, h in enumerate(hist):

252 dos[np.round(h * (height - 1) / hist.max()).astype(int), i] = '+' # '·' or 'x'

253

254 dos = dos[::-1]

255 dos[-1, dos[-1] == ' '] = '-'

256 dos = '\n'.join([''.join(d) for d in dos])

257

258 string += f'\n{dos}'

259 string += f'\n|{THz.min():<10.2f} THz' + ' '*(width - 26) + f'{THz.max():>10.2f}|'

260

261 return string

262

263 def __repr__(self):

264 return str(self)

265

266 def __getitem__(self, q) -> QPoint:

267 """Returns the q-point object based on its index"""

268 if q < 0 or q >= len(self._q):

269 raise IndexError

270 return QPoint(q, self)

271

272 def __call__(self, qpoint) -> QPoint:

273 """Tries to find a matching q-point based on reduced coordinate"""

274 qpoint = np.array(qpoint).astype(np.float64) % 1

275 for q, qpoint2 in enumerate(np.array(self._q).astype(np.float64)):

276 if np.allclose(qpoint, qpoint2):

277 return QPoint(q, self)

278 raise ValueError('qpoint not compatible, check mp.q_reduced')

279

280 # Getters ans setters for internal mutable arrays

281 def get_Q(self) -> NDArray[complex]:

282 """Return the mode coordinates as a ``(Nq, Nb)`` complex array in Å√dmu."""

283 return self._Q.copy()

284

285 def get_P(self) -> NDArray[complex]:

286 """Return the mode momenta as a ``(Nq, Nb)`` complex array in √eV."""

287 return self._P.copy()

288

289 def get_F(self) -> NDArray[complex]:

290 """Return the mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu."""

291 return self._F.copy()

292

293 def get_F_harmonic(self) -> NDArray[complex]:

294 r"""Return the harmonic mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu.

295

296 Computed as :math:`-\omega^2 Q`.

297 """

298 return -self._w2 * self.get_Q()

299

300 def set_Q(self, Q: NDArray[complex]) -> None:

301 """Sets the internal mode coordinates :math:`Q`.

302

303 The function ensures :math:`Q(-q)=Q^*(q)`.

304 """

305 # This ensures that stuff like mp.set_Q(0) works while not updating the

306 # array until the assert

307 Q_new = self.get_Q()

308 Q_new[:] = Q

309 if not np.allclose(np.conjugate(Q_new), Q_new[self.q_minus]):

310 raise ValueError('Supplied Q does not fulfill Q(-q) = Q(q)*')

311 self._Q[:] = Q_new

312

313 def set_P(self, P: NDArray[complex]) -> None:

314 """Sets the internal mode momenta :math:`P`.

315

316 The function ensures :math:`P(-q)=P^*(q)`.

317 """

318 P_new = self.get_P()

319 P_new[:] = P

320 if not np.allclose(np.conjugate(P_new), P_new[self.q_minus]):

321 raise ValueError('Supplied P does not fulfill P(-q) = P(q)*')

322 self._P[:] = P_new

323

324 def set_F(self, F: NDArray[complex]) -> None:

325 """Sets the internal mode forces :math:`F`.

326

327 The function ensures :math:`F(-q)=F^*(q)`.

328 """

329 F_new = self.get_F()

330 F_new[:] = F

331 if not np.allclose(np.conjugate(F_new), F_new[self.q_minus]):

332 raise ValueError('Supplied F does not fulfill F(-q) = F(q)*')

333 self._F[:] = F_new

334

335 def get_u(self) -> NDArray[float]:

336 """Return the atomic displacements as a ``(N, 3)`` array in Å."""

337 u = np.einsum('ksna,ks,n->na', self._X.conj(), self._Q, 1 / self.supercell.masses)

338 assert np.allclose(u.imag, 0)

339 return u.real

340

341 def get_v(self) -> NDArray[float]:

342 """Return the atomic velocities as a ``(N, 3)`` array in Å/fs."""

343 v = np.einsum('ksna,ks,n->na', self._X, self._P, 1 / self.supercell.masses)

344 assert np.allclose(v.imag, 0)

345 return v.real

346

347 def get_f(self) -> NDArray[float]:

348 """Return the atomic forces as a ``(N, 3)`` array in eV/Å."""

349 f = np.einsum('ksna,ks->na', self._X, self._F)

350 assert np.allclose(f.imag, 0)

351 return f.real

352

353 def get_f_harmonic(self) -> NDArray[float]:

354 """Return the harmonic atomic forces for the current displacements

355 as a ``(N, 3)`` array in eV/Å.

356 """

357 F_harmonic = self.get_F_harmonic()

358 f_harmonic = np.einsum('ksna,ks->na', self._X, F_harmonic)

359 assert np.allclose(f_harmonic.imag, 0)

360 return f_harmonic.real

361

362 def set_u(self, u: NDArray[float]) -> None:

363 """Sets the internal mode coordinates :math:`Q` given the atomic displacements :math:`u`.

364

365 .. math::

366

367 Q = X u

368

369 Parameters

370 ----------

371 u

372 The atomic displacements in Å.

373 """

374 Q = np.einsum('ksna,na->ks', self._X, u)

375 self.set_Q(Q)

376

377 def set_v(self, v: NDArray[float]) -> None:

378 """Sets the internal mode momenta :math:`P` given the atomic velocities :math:`v`.

379

380 .. math::

381

382 P = X^* * v

383

384 Parameters

385 ----------

386 v

387 The atomic velocities in Å/fs.

388 """

389 P = np.einsum('ksna,na->ks', self._X.conj(), v)

390 self.set_P(P)

391

392 def set_f(self, f: NDArray[float]) -> None:

393 """Sets the internal mode forces :math:`F` given the atomic forces :math:`f`.

394

395 .. math::

396

397 F = X^* * f / m

398

399 Parameters

400 ----------

401 f

402 The atomic forces in eV/Å.

403 """

404 F = np.einsum('ksna,na,n->ks', self._X.conj(), f, 1 / self.supercell.masses)

405 self.set_F(F)

406

407 # Convenience functions to handle ASE Atoms objects

408 def get_atoms(self, harmonic_forces: Optional[bool] = False) -> Atoms:

409 r"""Returns ASE :class:`Atoms` object with displacement,

410 velocities, forces, and harmonic energies.

411

412 Parameters

413 ----------

414 harmonic_forces

415 Whether the forces should be taken from the internal `F` or via `-\omega^2 Q`.

416 """

417 atoms = self.supercell.to_ase()

418 atoms.positions += self.get_u()

419 atoms.set_velocities(self.get_v() / fs)

420 E = self.potential_energies.sum()

421 f = self.get_f_harmonic() if harmonic_forces else self.get_f()

422

423 atoms.calc = SinglePointCalculator(

424 energy=E, forces=f, stress=None, magmoms=None, atoms=atoms)

425

426 return atoms

427

428 def update_from_atoms(self, atoms: Atoms) -> None:

429 """Updates the :class:`ModeProjector` objects with displacements, velocities,

430 and forces from an ASE :class:`Atoms` object.

431

432 Checks for an attached calculator in the first place and next for a forces array.

433

434 If no data sets corresponding array to zeros.

435

436 The velocities are converted to dynasor units internally. Note that

437 :attr:`atoms`'s masses are not used; masses are fixed from the

438 ``primitive``/``supercell`` structures given at construction, since the

439 dynamical matrix and mode eigenvectors already depend on them.

440 """

441

442 u = get_displacements(atoms, self.supercell)

443 if np.max(np.abs(u)) > 2.0:

444 warnings.warn('Displacements larger than 2Å. Is the atoms object permuted?')

445 self.set_u(u)

446 self.set_v(atoms.get_velocities() * fs)

447 try:

448 self.set_f(atoms.get_forces())

449 except RuntimeError:

450 if 'forces' in atoms.arrays:

451 self.set_f(atoms.arrays['forces'])

452 else:

453 self.set_f(np.zeros_like(atoms.positions))

454

455 # properties

456 @property

457 def q_minus(self) -> NDArray[float]:

458 """The index of the corresponding counter-propagating mode (:math:`-q`)."""

459 return self._q_minus.copy()

460

461 @property

462 def q_reduced(self) -> NDArray[float]:

463 """The q-points in reduced coordinates.

464

465 For example a zone boundary mode would be (1/2, 0, 0)

466 """

467 return self._q.astype(float)

468

469 @property

470 def q_cartesian(self) -> NDArray[float]:

471 """The q-points in cartesian coordinates with unit of rad/Å (2π included)."""

472 return 2 * np.pi * self.q_reduced @ self.primitive.inv_cell

473

474 @property

475 def omegas(self) -> NDArray[float]:

476 """The frequencies of each mode in rad/fs.

477

478 Following convention, negative values indicate imaginary frequencies.

479 """

480 return np.sign(self._w2) * np.sqrt(np.abs(self._w2))

481

482 @property

483 def polarizations(self) -> NDArray[float]:

484 """The polarization vectors for each mode `(Nq, Nb, Np, 3)`."""

485 return self._W

486

487 @property

488 def eigenmodes(self) -> NDArray[float]:

489 """The eigenmodes in the supercell as `(Nq, Nb, N, 3)`-array

490

491 The eigenmodes include the masses such that :math:`Q = X u`

492 where :math:`u` are the supercell displacements.

493 """

494 return self._X

495

496 @property

497 def potential_energies(self) -> NDArray[float]:

498 r"""Potential energy per mode as `(Nq, Nb)`-array.

499

500 The potential energies are defined as :math:`1/2 \omega^2 Q Q^*` and should equal

501 :math:`1/2 k_B T` in equilibrium for a harmonic system.

502 """

503 return 1 / 2 * np.abs(self._Q) ** 2 * self._w2

504

505 @property

506 def kinetic_energies(self) -> NDArray[float]:

507 """Kinetic energy per mode as `(Nq, Nb)`-array.

508

509 The kinetic energies are defined as :math:`1/2 P P^*`. Should equal

510 :math:`1/2 k_B T` in equilibrium.

511 """

512 return 1 / 2 * np.abs(self._P)**2

513

514 @property

515 def virial_energies(self) -> NDArray[complex]:

516 r"""The virial energies per mode as `(Nq, Nb)`-array.

517

518 The virial energies are defined here as :math:`-1/2 Q F`, which should have an

519 expectation value of :math:`1/2 k_B T` per mode in equilibrium. For a harmonic

520 system this is simply equal to the potential energy. This means that

521 the virial energy can be used to monitor the anharmonicity or

522 define a measure of the potential energy.

523

524 Note that while :math:`Q F` is in general complex the virial energy

525 obeys :math:`E(-\boldsymbol{q}) = E^*(\boldsymbol{q})`, meaning the

526 imaginary parts cancel between :math:`\boldsymbol{q}` and

527 :math:`-\boldsymbol{q}` and the total is real. Take the real part to

528 obtain per-mode energies; this leaves the sum unchanged.

529 """

530 return -1 / 2 * self._Q * self._F

531

532 def write(self, file_name: str) -> None:

533 """Uses pickle to write mode projector to file."""

534 with open(file_name, 'wb') as f:

535 pickle.dump(self, f)

536

537 @classmethod

538 def read(cls, file_name: str) -> ModeProjector:

539 """Return :class:`ModeProjector` instance from pickle file

540 that was saved using :func:`~ModeProjector.write`."""

541 with open(file_name, 'rb') as f:

542 mp = pickle.load(f)

543 return mp

Coverage for dynasor/modes/mode_projector.py: 100%

211 statements