Coverage for dynasor/modes/mode_projector.py: 100%
211 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
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 <dynasor.modes.qpoint.QPoint>` objects are
28 representations of a 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.
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.
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')
133 if len(supercell) != len(force_constants):
134 raise ValueError('force constants shape is not compatible with supercell size')
136 if force_constants.shape != (len(supercell), len(supercell), 3, 3):
137 raise ValueError('force constants shape should be (N, N, 3, 3)')
139 self.primitive = Prim(primitive)
140 self.supercell = Supercell(supercell, primitive)
141 self.force_constants = force_constants
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))
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]
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):
155 D = get_dynamical_matrix(
156 self.force_constants, self.supercell.offsets, self.supercell.indices,
157 q.astype(np.float64))
159 if qi == self.q_minus[qi]:
160 assert np.allclose(D.imag, 0)
161 D = D.real
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)
170 self._D.append(D)
171 self._w2.append(w2)
172 self._W.append(W)
174 self._D = np.array(self._D)
175 self._w2 = np.array(self._w2)
176 self._W = np.array(self._W)
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]
182 assert np.allclose(self._D[q], self._D[q_minus].conj())
183 assert np.allclose(self._w2[q], self._w2[q_minus])
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
192 for group in group_eigvals(self._w2[q], tolerance**0.5):
193 W = symmetrize_eigenvectors(self._W[q, group])
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]
212 self._W[q_minus] = self._W[q].conj()
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
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
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)
236 def __str__(self):
237 strings = ['### ModeProjector ###']
238 strings += [f'{self.supercell}']
239 strings += [f'{self.primitive}']
240 string = '\n'.join(strings)
242 # ASCII DOS!
243 width = 80
244 height = 24
245 dos = np.full((height, width), ' ')
247 THz = self.omegas * radians_per_fs_to_THz
249 hist, bins = np.histogram(THz.flat, bins=width)
251 for i, h in enumerate(hist):
252 dos[np.round(h * (height - 1) / hist.max()).astype(int), i] = '+' # '·' or 'x'
254 dos = dos[::-1]
255 dos[-1, dos[-1] == ' '] = '-'
256 dos = '\n'.join([''.join(d) for d in dos])
258 string += f'\n{dos}'
259 string += f'\n|{THz.min():<10.2f} THz' + ' '*(width - 26) + f'{THz.max():>10.2f}|'
261 return string
263 def __repr__(self):
264 return str(self)
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)
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')
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()
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()
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()
293 def get_F_harmonic(self) -> NDArray[complex]:
294 r"""Return the harmonic mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu.
296 Computed as :math:`-\omega^2 Q^*`. The conjugate follows the convention
297 :math:`F = X^* f / m` used by :meth:`set_f`, in which :math:`F` is the force
298 conjugate to :math:`Q`, i.e. :math:`-\partial V/\partial Q`. Note that
299 :math:`Q` and :math:`Q^*` coincide only at self-conjugate q-points.
300 """
301 return -self._w2 * self.get_Q().conj()
303 def set_Q(self, Q: NDArray[complex]) -> None:
304 """Sets the internal mode coordinates :math:`Q`.
306 The function ensures :math:`Q(-q)=Q^*(q)`.
307 """
308 # This ensures that stuff like mp.set_Q(0) works while not updating the
309 # array until the assert
310 Q_new = self.get_Q()
311 Q_new[:] = Q
312 if not np.allclose(np.conjugate(Q_new), Q_new[self.q_minus]):
313 raise ValueError('Supplied Q does not fulfill Q(-q) = Q(q)*')
314 self._Q[:] = Q_new
316 def set_P(self, P: NDArray[complex]) -> None:
317 """Sets the internal mode momenta :math:`P`.
319 The function ensures :math:`P(-q)=P^*(q)`.
320 """
321 P_new = self.get_P()
322 P_new[:] = P
323 if not np.allclose(np.conjugate(P_new), P_new[self.q_minus]):
324 raise ValueError('Supplied P does not fulfill P(-q) = P(q)*')
325 self._P[:] = P_new
327 def set_F(self, F: NDArray[complex]) -> None:
328 """Sets the internal mode forces :math:`F`.
330 The function ensures :math:`F(-q)=F^*(q)`.
331 """
332 F_new = self.get_F()
333 F_new[:] = F
334 if not np.allclose(np.conjugate(F_new), F_new[self.q_minus]):
335 raise ValueError('Supplied F does not fulfill F(-q) = F(q)*')
336 self._F[:] = F_new
338 def get_u(self) -> NDArray[float]:
339 """Return the atomic displacements as a ``(N, 3)`` array in Å."""
340 u = np.einsum('ksna,ks,n->na', self._X.conj(), self._Q, 1 / self.supercell.masses)
341 assert np.allclose(u.imag, 0)
342 return u.real
344 def get_v(self) -> NDArray[float]:
345 """Return the atomic velocities as a ``(N, 3)`` array in Å/fs."""
346 v = np.einsum('ksna,ks,n->na', self._X, self._P, 1 / self.supercell.masses)
347 assert np.allclose(v.imag, 0)
348 return v.real
350 def get_f(self) -> NDArray[float]:
351 """Return the atomic forces as a ``(N, 3)`` array in eV/Å."""
352 f = np.einsum('ksna,ks->na', self._X, self._F)
353 assert np.allclose(f.imag, 0)
354 return f.real
356 def get_f_harmonic(self) -> NDArray[float]:
357 """Return the harmonic atomic forces for the current displacements
358 as a ``(N, 3)`` array in eV/Å.
359 """
360 F_harmonic = self.get_F_harmonic()
361 f_harmonic = np.einsum('ksna,ks->na', self._X, F_harmonic)
362 assert np.allclose(f_harmonic.imag, 0)
363 return f_harmonic.real
365 def set_u(self, u: NDArray[float]) -> None:
366 """Sets the internal mode coordinates :math:`Q` given the atomic displacements :math:`u`.
368 .. math::
370 Q = X u
372 Parameters
373 ----------
374 u
375 The atomic displacements in Å.
376 """
377 Q = np.einsum('ksna,na->ks', self._X, u)
378 self.set_Q(Q)
380 def set_v(self, v: NDArray[float]) -> None:
381 """Sets the internal mode momenta :math:`P` given the atomic velocities :math:`v`.
383 .. math::
385 P = X^* * v
387 Parameters
388 ----------
389 v
390 The atomic velocities in Å/fs.
391 """
392 P = np.einsum('ksna,na->ks', self._X.conj(), v)
393 self.set_P(P)
395 def set_f(self, f: NDArray[float]) -> None:
396 """Sets the internal mode forces :math:`F` given the atomic forces :math:`f`.
398 .. math::
400 F = X^* * f / m
402 Parameters
403 ----------
404 f
405 The atomic forces in eV/Å.
406 """
407 F = np.einsum('ksna,na,n->ks', self._X.conj(), f, 1 / self.supercell.masses)
408 self.set_F(F)
410 # Convenience functions to handle ASE Atoms objects
411 def get_atoms(self, harmonic_forces: Optional[bool] = False) -> Atoms:
412 r"""Returns ASE :class:`~ase.Atoms` object with displacement,
413 velocities, forces, and harmonic energies.
415 Parameters
416 ----------
417 harmonic_forces
418 Whether the forces should be taken from the internal `F` or via `-\omega^2 Q`.
419 """
420 atoms = self.supercell.to_ase()
421 atoms.positions += self.get_u()
422 atoms.set_velocities(self.get_v() / fs)
423 E = self.potential_energies.sum()
424 f = self.get_f_harmonic() if harmonic_forces else self.get_f()
426 atoms.calc = SinglePointCalculator(
427 energy=E, forces=f, stress=None, magmoms=None, atoms=atoms)
429 return atoms
431 def update_from_atoms(self, atoms: Atoms) -> None:
432 """Updates the :class:`ModeProjector` objects with displacements, velocities,
433 and forces from an ASE :class:`~ase.Atoms` object.
435 Checks for an attached calculator in the first place and next for a forces array.
437 If no data sets corresponding array to zeros.
439 The velocities are converted to dynasor units internally. Note that
440 :attr:`atoms`'s masses are not used; masses are fixed from the
441 ``primitive``/``supercell`` structures given at construction, since the
442 dynamical matrix and mode eigenvectors already depend on them.
443 """
445 u = get_displacements(atoms, self.supercell)
446 if np.max(np.abs(u)) > 2.0:
447 warnings.warn('Displacements larger than 2Å. Is the atoms object permuted?')
448 self.set_u(u)
449 self.set_v(atoms.get_velocities() * fs)
450 try:
451 self.set_f(atoms.get_forces())
452 except RuntimeError:
453 if 'forces' in atoms.arrays:
454 self.set_f(atoms.arrays['forces'])
455 else:
456 self.set_f(np.zeros_like(atoms.positions))
458 # properties
459 @property
460 def q_minus(self) -> NDArray[float]:
461 """The index of the corresponding counter-propagating mode (:math:`-q`)."""
462 return self._q_minus.copy()
464 @property
465 def q_reduced(self) -> NDArray[float]:
466 """The q-points in reduced coordinates.
468 For example a zone boundary mode would be (1/2, 0, 0)
469 """
470 return self._q.astype(float)
472 @property
473 def q_cartesian(self) -> NDArray[float]:
474 """The q-points in cartesian coordinates with unit of rad/Å (2π included)."""
475 return 2 * np.pi * self.q_reduced @ self.primitive.inv_cell
477 @property
478 def omegas(self) -> NDArray[float]:
479 """The frequencies of each mode in rad/fs.
481 Following convention, negative values indicate imaginary frequencies.
482 """
483 return np.sign(self._w2) * np.sqrt(np.abs(self._w2))
485 @property
486 def polarizations(self) -> NDArray[float]:
487 """The polarization vectors for each mode `(Nq, Nb, Np, 3)`."""
488 return self._W
490 @property
491 def eigenmodes(self) -> NDArray[float]:
492 """The eigenmodes in the supercell as `(Nq, Nb, N, 3)`-array
494 The eigenmodes include the masses such that :math:`Q = X u`
495 where :math:`u` are the supercell displacements.
496 """
497 return self._X
499 @property
500 def potential_energies(self) -> NDArray[float]:
501 r"""Potential energy per mode as `(Nq, Nb)`-array.
503 The potential energies are defined as :math:`1/2 \omega^2 Q Q^*` and should equal
504 :math:`1/2 k_B T` in equilibrium for a harmonic system.
505 """
506 return 1 / 2 * np.abs(self._Q) ** 2 * self._w2
508 @property
509 def kinetic_energies(self) -> NDArray[float]:
510 """Kinetic energy per mode as `(Nq, Nb)`-array.
512 The kinetic energies are defined as :math:`1/2 P P^*`. Should equal
513 :math:`1/2 k_B T` in equilibrium.
514 """
515 return 1 / 2 * np.abs(self._P)**2
517 @property
518 def virial_energies(self) -> NDArray[complex]:
519 r"""The virial energies per mode as `(Nq, Nb)`-array.
521 The virial energies are defined here as :math:`-1/2 Q F`, which should have an
522 expectation value of :math:`1/2 k_B T` per mode in equilibrium. For a harmonic
523 system this is simply equal to the potential energy. This means that
524 the virial energy can be used to monitor the anharmonicity or
525 define a measure of the potential energy.
527 Note that while :math:`Q F` is in general complex the virial energy
528 obeys :math:`E(-\boldsymbol{q}) = E^*(\boldsymbol{q})`, meaning the
529 imaginary parts cancel between :math:`\boldsymbol{q}` and
530 :math:`-\boldsymbol{q}` and the total is real. Take the real part to
531 obtain per-mode energies; this leaves the sum unchanged.
532 """
533 return -1 / 2 * self._Q * self._F
535 def write(self, file_name: str) -> None:
536 """Uses pickle to write mode projector to file."""
537 with open(file_name, 'wb') as f:
538 pickle.dump(self, f)
540 @classmethod
541 def read(cls, file_name: str) -> ModeProjector:
542 """Return :class:`ModeProjector` instance from pickle file
543 that was saved using :func:`~ModeProjector.write`."""
544 with open(file_name, 'rb') as f:
545 mp = pickle.load(f)
546 return mp