Coverage for dynasor/modes/mode_projector.py: 100%
211 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 07:00 +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, mode_coordinates, mode_forces,
16 mode_momenta, symmetrize_eigenvectors)
17from .qpoint import QPoint
18from .atoms import Prim, Supercell
19from ..qpoints.tools import get_commensurate_lattice_points
22class ModeProjector:
23 """
24 The :class:`ModeProjector` maps between real atomic displacements `u` and
25 complex mode coordinates `Q`.
27 Some special python methods are implemented. The `__str__` and `__repr__`
28 provides useful info. :class:`QPoint <dynasor.modes.qpoint.QPoint>` objects are
29 representations of a single q-point and associated information and can be accessed either by
30 call providing a reduced wavevector
32 >>> mp((1/2, 0, 0)) # doctest: +SKIP
34 or by index corresponding to reduced q-point accessible from
35 :attr:`~ModeProjector.q_reduced`
37 >>> mp[2] # doctest: +SKIP
39 In addition to mode coordinates `Q` the class can also map the atomic
40 velocities `v` to mode momenta `P` as well as atomic forces `f` to mode
41 forces `F`. The class can also map back. This mapping is done using
42 getters and setters. Internally only Q, P and F are stored
44 >>> Q = mp.get_Q() # doctest: +SKIP
45 >>> mp.set_u(u) # doctest: +SKIP
47 In addition, the forces corresponding to the harmonic forces can be
48 accessed by :meth:`~ModeProjector.get_f_harmonic()` and
49 :meth:`~ModeProjector.get_F_harmonic()`. For ASE Atoms objects the
50 displacements etc. can be updated and applied by
52 >>> mp.update_from_atoms(atoms) # doctest: +SKIP
53 >>> atoms = mp.get_atoms(harmonic_forces=False) # doctest: +SKIP
55 The shapes for each property uses the following variables
57 * `N`: Number of atoms in supercell
58 * `Nu`: unit cells (`N/Np`)
59 * `Np`: primitive basis atoms (`N/Nu`)
60 * `Nb`: bands (`Np*3`)
61 * `Nq`: q-points (`Nu`)
63 Please consult the documentation or the specific getters and setters to see
64 the exact transformations used.
66 Units
67 ^^^^^
68 The internal units in dynasor are Å, fs and eV. All frequencies are angular
69 (a.k.a. the "physicist's convention" with 2π included). These are the units
70 dynasor will expect and return. In, e.g., print functions conventional units
71 such fs, Å, THz, Da, meV are commonly used.
73 **Mass:**
74 The internal unit choice (eV, Å, fs) means that the mass unit is not Dalton
75 but rather 0.009648533290731906 Da.
76 We refer to this unit as the "dynasor mass unit" (dmu),
77 i.e., 1 Da = 103.64269572045423 dmu.
78 As a user you will only see this unit in the output of :class:`ModeProjector` objects.
79 Masses provided via, e.g., ASE Atoms objects are converted internally.
81 **Waves:**
82 dynasor reports and expects spatial (angular) frequencies in rad/Å and temporal (angular)
83 frequencies in rad/fs. This follows the often-used convention in physics to include the
84 factor of 2π in the wave vectors. For instance the wavelength is given by λ=2π/q.
86 **Mode amplitudes:**
87 Mode amplitudes are reported in Å√dmu = fs√eV.
89 **Velocities:**
90 For modes the momenta are reported in Å√dmu/fs or just √eV
91 while atomic velocities are reported in Å/fs.
93 **Mode forces:**
94 The force is defined as the derivative of the momenta with respect to time
95 so the unit used when reporting mode forces is Å√dmu/fs² (or √eV/fs).
97 Internal arrays
98 ^^^^^^^^^^^^^^^
99 For the curious, the internal data arrays are
101 * :attr:`primitive`, :attr:`supercell`, :attr:`force_constants` (input)
102 * :attr:`_q`, :attr:`q_minus` (reduced q-points and which q-points are related by inversion)
103 * :attr:`_D`, :attr:`_w2`, :attr:`_W`
104 (dynamical matrices, frequencies (ev/Ų/Da), polarization vectors)
105 * :attr:`_X` (eigenmodes which are mass weighted "polarization vectors" in the supercell)
106 """
107 def __init__(self, primitive: Atoms, supercell: Atoms, force_constants: NDArray[float]):
108 """The mode projector is initialized by a primitive cell and a
109 supercell as well as harmonic force constants.
111 The force constants are assumed to be in units of eV/Ų as returned
112 from phonopy. Be careful about the permutations when working with force
113 constants and atoms object from different codes.
115 Parameters
116 ----------
117 primitive
118 Primitive cell. Note that the masses are stored internally as
119 Dalton in ASE but will be converted to the internal dynasor
120 mass unit (dmu).
121 supercell
122 Ideal supercell corresponding to the force constants.
123 force_constants
124 Force constants for the supercell in eV/Ų as a `(N, N, 3, 3)` array
125 where `N` is `len(supercell)`.
126 """
127 if len(primitive) == len(supercell):
128 warnings.warn('Primitive and supercell have the same size')
129 elif len(primitive) > len(supercell):
130 raise ValueError('Primitive cell larger than supercell')
131 elif not (len(supercell) / len(primitive)).is_integer():
132 raise ValueError('supercell size is not multiple of primitive size')
134 if len(supercell) != len(force_constants):
135 raise ValueError('force constants shape is not compatible with supercell size')
137 if force_constants.shape != (len(supercell), len(supercell), 3, 3):
138 raise ValueError('force constants shape should be (N, N, 3, 3)')
140 self.primitive = Prim(primitive)
141 self.supercell = Supercell(supercell, primitive)
142 self.force_constants = force_constants
144 # Find q-points in reduced primitive cell coordinates
145 q_integer = get_commensurate_lattice_points(self.supercell.P.T)
146 q_reduced = np.dot(q_integer, self.supercell.P_inv.T)
147 self._q = np.array(sorted(tuple(q) for q in q_reduced))
149 # The equivalent q-point corresponding to -q
150 self._q_minus = [[tuple(q) for q in self._q].index(tuple((-q) % 1)) for q in self._q]
152 # Construct dynamical matrix and diagonalize at each q-point
153 self._D, self._w2, self._W = [], [], []
154 for qi, q in enumerate(self._q):
156 D = get_dynamical_matrix(
157 self.force_constants, self.supercell.offsets, self.supercell.indices,
158 q.astype(np.float64))
160 if qi == self.q_minus[qi]:
161 assert np.allclose(D.imag, 0)
162 D = D.real
164 D = np.einsum('ijab,i,j->ijab',
165 D, self.primitive.masses**-0.5, self.primitive.masses**-0.5)
166 D_matrix = D.transpose(0, 2, 1, 3).reshape(-1, self.primitive.n_atoms * 3)
167 assert np.allclose(D_matrix, D_matrix.T.conj())
168 w2, W = np.linalg.eigh(D_matrix)
169 W = W.T.reshape(-1, self.primitive.n_atoms, 3)
171 self._D.append(D)
172 self._w2.append(w2)
173 self._W.append(W)
175 self._D = np.array(self._D)
176 self._w2 = np.array(self._w2)
177 self._W = np.array(self._W)
179 # Post check basic symmetries, group eigenvalues and try to make degenerate modes nicer
180 for q, q_minus in enumerate(self.q_minus):
181 q_minus = self.q_minus[q]
183 assert np.allclose(self._D[q], self._D[q_minus].conj())
184 assert np.allclose(self._w2[q], self._w2[q_minus])
186 # tolerances for grouping and sorting eigenvalues and eigenvectors
187 group_decimals = 12
188 tolerance = 10**(-group_decimals)
189 # Keep sorting coarser than the convergence tolerance so tiny SVD
190 # roundoff cannot change the order of equivalent modes.
191 sort_decimals = 10
193 for group in group_eigvals(self._w2[q], tolerance**0.5):
194 W = symmetrize_eigenvectors(self._W[q, group])
196 # Try to order them
197 W_sort = W.copy().transpose(0, 2, 1).reshape(len(W), -1)
198 # abs is because we want to consider the magnitude
199 # - (minus) basically reverts the sort order to place largest first
200 # T is just because how lexsort works, we want to consider each
201 # atom and direction as a key for the bands
202 # -1 is because we want to make the x-direction of the first
203 # atom the most significant key
204 # At the end the first band should have the largest magnitude
205 # for the first atom in x
206 # Suppress numerical noise in the mode signatures before sorting.
207 # Keep this precision deliberately below the factor-analysis
208 # convergence noise so tiny SVD/BLAS differences cannot permute
209 # otherwise equivalent degenerate modes.
210 argsort = np.lexsort(np.round(-np.abs(W_sort).T[::-1], sort_decimals))
211 self._W[q, group] = W[argsort]
213 self._W[q_minus] = self._W[q].conj()
215 # Construct supercell projection matrix: q_ks = X_ksna u_na.
216 # phase and mass depend only on (q, index), so compute once and
217 # broadcast against a gather of self._W.
218 indices = self.supercell.indices
220 # self._q has dtype=object (Fraction); cast to float for the phase
221 q_float = self._q.astype(np.float64)
222 phase = np.exp(-1j * 2*np.pi * q_float @ self.supercell.offsets.T) # (n_q, n_super)
223 mass_sqrt = self.primitive.masses[indices]**0.5 # (n_super,)
224 # Build _X in place to avoid extra full-size temporaries (halves peak
225 # memory). astype pins the dtype (no-op unless _W is real, e.g. Gamma).
226 self._X = self._W[:, :, indices, :].astype(np.complex128, copy=False)
227 np.conjugate(self._X, out=self._X)
228 self._X *= phase[:, None, :, None]
229 self._X *= mass_sqrt[None, None, :, None]
230 self._X /= (self.supercell.n_atoms / self.primitive.n_atoms)**0.5
232 # Init arrays to hold Q, P and F
233 self._Q = np.zeros((len(self._q), self.primitive.n_atoms*3), dtype=np.complex128)
234 self._P = np.zeros_like(self._Q)
235 self._F = np.zeros_like(self._Q)
237 def __str__(self):
238 strings = ['### ModeProjector ###']
239 strings += [f'{self.supercell}']
240 strings += [f'{self.primitive}']
241 string = '\n'.join(strings)
243 # ASCII DOS!
244 width = 80
245 height = 24
246 dos = np.full((height, width), ' ')
248 THz = self.omegas * radians_per_fs_to_THz
250 hist, bins = np.histogram(THz.flat, bins=width)
252 for i, h in enumerate(hist):
253 dos[np.round(h * (height - 1) / hist.max()).astype(int), i] = '+' # '·' or 'x'
255 dos = dos[::-1]
256 dos[-1, dos[-1] == ' '] = '-'
257 dos = '\n'.join([''.join(d) for d in dos])
259 string += f'\n{dos}'
260 string += f'\n|{THz.min():<10.2f} THz' + ' '*(width - 26) + f'{THz.max():>10.2f}|'
262 return string
264 def __repr__(self):
265 return str(self)
267 def __getitem__(self, q) -> QPoint:
268 """Returns the q-point object based on its index"""
269 if q < 0 or q >= len(self._q):
270 raise IndexError
271 return QPoint(q, self)
273 def __call__(self, qpoint) -> QPoint:
274 """Tries to find a matching q-point based on reduced coordinate"""
275 qpoint = np.array(qpoint).astype(np.float64) % 1
276 for q, qpoint2 in enumerate(np.array(self._q).astype(np.float64)):
277 if np.allclose(qpoint, qpoint2):
278 return QPoint(q, self)
279 raise ValueError('qpoint not compatible, check mp.q_reduced')
281 # Getters ans setters for internal mutable arrays
282 def get_Q(self) -> NDArray[complex]:
283 """Return the mode coordinates as a ``(Nq, Nb)`` complex array in Å√dmu."""
284 return self._Q.copy()
286 def get_P(self) -> NDArray[complex]:
287 """Return the mode momenta as a ``(Nq, Nb)`` complex array in √eV."""
288 return self._P.copy()
290 def get_F(self) -> NDArray[complex]:
291 """Return the mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu."""
292 return self._F.copy()
294 def get_F_harmonic(self) -> NDArray[complex]:
295 r"""Return the harmonic mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu.
297 Computed as :math:`-\omega^2 Q^*`. The conjugate follows the convention
298 :math:`F = X^* f / m` used by :meth:`set_f`, in which :math:`F` is the force
299 conjugate to :math:`Q`, i.e. :math:`-\partial V/\partial Q`. Note that
300 :math:`Q` and :math:`Q^*` coincide only at self-conjugate q-points.
301 """
302 return -self._w2 * self.get_Q().conj()
304 def set_Q(self, Q: NDArray[complex]) -> None:
305 """Sets the internal mode coordinates :math:`Q`.
307 The function ensures :math:`Q(-q)=Q^*(q)`.
308 """
309 # This ensures that stuff like mp.set_Q(0) works while not updating the
310 # array until the assert
311 Q_new = self.get_Q()
312 Q_new[:] = Q
313 if not np.allclose(np.conjugate(Q_new), Q_new[self.q_minus]):
314 raise ValueError('Supplied Q does not fulfill Q(-q) = Q(q)*')
315 self._Q[:] = Q_new
317 def set_P(self, P: NDArray[complex]) -> None:
318 """Sets the internal mode momenta :math:`P`.
320 The function ensures :math:`P(-q)=P^*(q)`.
321 """
322 P_new = self.get_P()
323 P_new[:] = P
324 if not np.allclose(np.conjugate(P_new), P_new[self.q_minus]):
325 raise ValueError('Supplied P does not fulfill P(-q) = P(q)*')
326 self._P[:] = P_new
328 def set_F(self, F: NDArray[complex]) -> None:
329 """Sets the internal mode forces :math:`F`.
331 The function ensures :math:`F(-q)=F^*(q)`.
332 """
333 F_new = self.get_F()
334 F_new[:] = F
335 if not np.allclose(np.conjugate(F_new), F_new[self.q_minus]):
336 raise ValueError('Supplied F does not fulfill F(-q) = F(q)*')
337 self._F[:] = F_new
339 def get_u(self) -> NDArray[float]:
340 """Return the atomic displacements as a ``(N, 3)`` array in Å."""
341 u = np.einsum('ksna,ks,n->na', self._X.conj(), self._Q, 1 / self.supercell.masses)
342 assert np.allclose(u.imag, 0)
343 return u.real
345 def get_v(self) -> NDArray[float]:
346 """Return the atomic velocities as a ``(N, 3)`` array in Å/fs."""
347 v = np.einsum('ksna,ks,n->na', self._X, self._P, 1 / self.supercell.masses)
348 assert np.allclose(v.imag, 0)
349 return v.real
351 def get_f(self) -> NDArray[float]:
352 """Return the atomic forces as a ``(N, 3)`` array in eV/Å."""
353 f = np.einsum('ksna,ks->na', self._X, self._F)
354 assert np.allclose(f.imag, 0)
355 return f.real
357 def get_f_harmonic(self) -> NDArray[float]:
358 """Return the harmonic atomic forces for the current displacements
359 as a ``(N, 3)`` array in eV/Å.
360 """
361 F_harmonic = self.get_F_harmonic()
362 f_harmonic = np.einsum('ksna,ks->na', self._X, F_harmonic)
363 assert np.allclose(f_harmonic.imag, 0)
364 return f_harmonic.real
366 def set_u(self, u: NDArray[float]) -> None:
367 """Sets the internal mode coordinates :math:`Q` given the atomic displacements :math:`u`.
369 .. math::
371 Q = X u
373 Parameters
374 ----------
375 u
376 The atomic displacements in Å.
377 """
378 Q = mode_coordinates(self._X, u)
379 self.set_Q(Q)
381 def set_v(self, v: NDArray[float]) -> None:
382 """Sets the internal mode momenta :math:`P` given the atomic velocities :math:`v`.
384 .. math::
386 P = X^* * v
388 Parameters
389 ----------
390 v
391 The atomic velocities in Å/fs.
392 """
393 P = mode_momenta(self._X, v)
394 self.set_P(P)
396 def set_f(self, f: NDArray[float]) -> None:
397 """Sets the internal mode forces :math:`F` given the atomic forces :math:`f`.
399 .. math::
401 F = X^* * f / m
403 Parameters
404 ----------
405 f
406 The atomic forces in eV/Å.
407 """
408 F = mode_forces(self._X, f, self.supercell.masses)
409 self.set_F(F)
411 # Convenience functions to handle ASE Atoms objects
412 def get_atoms(self, harmonic_forces: Optional[bool] = False) -> Atoms:
413 r"""Returns ASE :class:`~ase.Atoms` object with displacement,
414 velocities, forces, and harmonic energies.
416 Parameters
417 ----------
418 harmonic_forces
419 Whether the forces should be taken from the internal `F` or via `-\omega^2 Q`.
420 """
421 atoms = self.supercell.to_ase()
422 atoms.positions += self.get_u()
423 atoms.set_velocities(self.get_v() / fs)
424 E = self.potential_energies.sum()
425 f = self.get_f_harmonic() if harmonic_forces else self.get_f()
427 atoms.calc = SinglePointCalculator(
428 energy=E, forces=f, stress=None, magmoms=None, atoms=atoms)
430 return atoms
432 def update_from_atoms(self, atoms: Atoms) -> None:
433 """Updates the :class:`ModeProjector` objects with displacements, velocities,
434 and forces from an ASE :class:`~ase.Atoms` object.
436 Checks for an attached calculator in the first place and next for a forces array.
438 If no data sets corresponding array to zeros.
440 The velocities are converted to dynasor units internally. Note that
441 :attr:`atoms`'s masses are not used; masses are fixed from the
442 ``primitive``/``supercell`` structures given at construction, since the
443 dynamical matrix and mode eigenvectors already depend on them.
444 """
446 u = get_displacements(atoms, self.supercell)
447 if np.max(np.abs(u)) > 2.0:
448 warnings.warn('Displacements larger than 2Å. Is the atoms object permuted?')
449 self.set_u(u)
450 self.set_v(atoms.get_velocities() * fs)
451 try:
452 self.set_f(atoms.get_forces())
453 except RuntimeError:
454 if 'forces' in atoms.arrays:
455 self.set_f(atoms.arrays['forces'])
456 else:
457 self.set_f(np.zeros_like(atoms.positions))
459 # properties
460 @property
461 def q_minus(self) -> NDArray[float]:
462 """The index of the corresponding counter-propagating mode (:math:`-q`)."""
463 return self._q_minus.copy()
465 @property
466 def q_reduced(self) -> NDArray[float]:
467 """The q-points in reduced coordinates.
469 For example a zone boundary mode would be (1/2, 0, 0)
470 """
471 return self._q.astype(float)
473 @property
474 def q_cartesian(self) -> NDArray[float]:
475 """The q-points in cartesian coordinates with unit of rad/Å (2π included)."""
476 return 2 * np.pi * self.q_reduced @ self.primitive.inv_cell
478 @property
479 def omegas(self) -> NDArray[float]:
480 """The frequencies of each mode in rad/fs.
482 Following convention, negative values indicate imaginary frequencies.
483 """
484 return np.sign(self._w2) * np.sqrt(np.abs(self._w2))
486 @property
487 def polarizations(self) -> NDArray[float]:
488 """The polarization vectors for each mode `(Nq, Nb, Np, 3)`."""
489 return self._W
491 @property
492 def eigenmodes(self) -> NDArray[float]:
493 """The eigenmodes in the supercell as `(Nq, Nb, N, 3)`-array
495 The eigenmodes include the masses such that :math:`Q = X u`
496 where :math:`u` are the supercell displacements.
497 """
498 return self._X
500 @property
501 def potential_energies(self) -> NDArray[float]:
502 r"""Potential energy per mode as `(Nq, Nb)`-array.
504 The potential energies are defined as :math:`1/2 \omega^2 Q Q^*` and should equal
505 :math:`1/2 k_B T` in equilibrium for a harmonic system.
506 """
507 return 1 / 2 * np.abs(self._Q) ** 2 * self._w2
509 @property
510 def kinetic_energies(self) -> NDArray[float]:
511 """Kinetic energy per mode as `(Nq, Nb)`-array.
513 The kinetic energies are defined as :math:`1/2 P P^*`. Should equal
514 :math:`1/2 k_B T` in equilibrium.
515 """
516 return 1 / 2 * np.abs(self._P)**2
518 @property
519 def virial_energies(self) -> NDArray[complex]:
520 r"""The virial energies per mode as `(Nq, Nb)`-array.
522 The virial energies are defined here as :math:`-1/2 Q F`, which should have an
523 expectation value of :math:`1/2 k_B T` per mode in equilibrium. For a harmonic
524 system this is simply equal to the potential energy. This means that
525 the virial energy can be used to monitor the anharmonicity or
526 define a measure of the potential energy.
528 Note that while :math:`Q F` is in general complex the virial energy
529 obeys :math:`E(-\boldsymbol{q}) = E^*(\boldsymbol{q})`, meaning the
530 imaginary parts cancel between :math:`\boldsymbol{q}` and
531 :math:`-\boldsymbol{q}` and the total is real. Take the real part to
532 obtain per-mode energies; this leaves the sum unchanged.
533 """
534 return -1 / 2 * self._Q * self._F
536 def write(self, file_name: str) -> None:
537 """Uses pickle to write mode projector to file."""
538 with open(file_name, 'wb') as f:
539 pickle.dump(self, f)
541 @classmethod
542 def read(cls, file_name: str) -> ModeProjector:
543 """Return :class:`ModeProjector` instance from pickle file
544 that was saved using :func:`~ModeProjector.write`."""
545 with open(file_name, 'rb') as f:
546 mp = pickle.load(f)
547 return mp