Coverage for dynasor / modes / mode_projector.py: 88%
209 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-16 12:31 +0000
1from __future__ import annotations
2from typing import Optional
4import numpy as np
5import warnings
6import pickle
7import itertools
9from ase.calculators.singlepoint import SinglePointCalculator
10from ase.units import fs
11from ase import Atoms
12from numpy.typing import NDArray
14from dynasor.units import radians_per_fs_to_THz
15from dynasor.tools.structures import get_displacements
16from .tools import get_dynamical_matrix, group_eigvals, 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` objects are representations of a
29 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): 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 raise ValueError('Primitive cell larger than supercell')
131 elif not (len(supercell) / len(primitive)).is_integer(): 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true
132 raise ValueError('supercell size is not multiple of primitive size')
134 if len(supercell) != len(force_constants): 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true
135 raise ValueError('force constants shape is not compatible with supercell size')
137 if force_constants.shape != (len(supercell), len(supercell), 3, 3): 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
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 round_decimals = 12
188 tolerance = 10**(-round_decimals)
190 for group in group_eigvals(self._w2[q], tolerance**0.5):
191 W = symmetrize_eigenvectors(self._W[q, group])
193 # Try to order them
194 W_sort = W.copy().transpose(0, 2, 1).reshape(len(W), -1)
195 # abs is because we want to consider the magnitude
196 # - (minus) basically reverts the sort order to place largest first
197 # T is just because how lexsort works, we want to consider each
198 # atom and direction as a key for the bands
199 # -1 is because we want to make the x-direction of the first
200 # atom the mist significant key
201 # At the end the first band should have the largest magnitude
202 # for the first atom in x
203 argsort = np.lexsort(np.round(-np.abs(W_sort).T[::-1], round_decimals))
204 self._W[q, group] = W[argsort]
206 self._W[q_minus] = self._W[q].conj()
208 # Construct supercell projection matrix
209 # q_ks = X_ksna u_na
210 self._X = np.zeros((len(self._q), self.primitive.n_atoms * 3, self.supercell.n_atoms, 3),
211 dtype=np.complex128)
213 for index in range(self.supercell.n_atoms):
214 i, N = self.supercell.indices[index], self.supercell.offsets[index]
215 for q, s, a in itertools.product(
216 range(len(self._q)), range(self.primitive.n_atoms * 3), range(3)):
217 phase = np.exp(-1j * 2*np.pi * self._q[q] @ N)
218 self._X[q, s, index, a] = (
219 self.primitive.masses[i]**0.5 * phase * self._W[q, s, i, a].conj())
220 self._X /= (self.supercell.n_atoms / self.primitive.n_atoms)**0.5
222 # Init arrays to hold Q, P and F
223 self._Q = np.zeros((len(self._q), self.primitive.n_atoms*3), dtype=np.complex128)
224 self._P = np.zeros_like(self._Q)
225 self._F = np.zeros_like(self._Q)
227 def __str__(self):
228 strings = ['### ModeProjector ###']
229 strings += [f'{self.supercell}']
230 strings += [f'{self.primitive}']
231 string = '\n'.join(strings)
233 # ASCII DOS!
234 width = 80
235 height = 24
236 dos = np.full((height, width), ' ')
238 THz = self.omegas * radians_per_fs_to_THz
240 hist, bins = np.histogram(THz.flat, bins=width)
242 for i, h in enumerate(hist):
243 dos[np.round(h * (height - 1) / hist.max()).astype(int), i] = '+' # '·' or 'x'
245 dos = dos[::-1]
246 dos[-1, dos[-1] == ' '] = '-'
247 dos = '\n'.join([''.join(d) for d in dos])
249 string += f'\n{dos}'
250 string += f'\n|{THz.min():<10.2f} THz' + ' '*(width - 26) + f'{THz.max():>10.2f}|'
252 return string
254 def __repr__(self):
255 return str(self)
257 def __getitem__(self, q) -> QPoint:
258 """Returns the q-point object based on its index"""
259 if q < 0 or q >= len(self._q):
260 raise IndexError
261 return QPoint(q, self)
263 def __call__(self, qpoint) -> QPoint:
264 """Tries to find a matching q-point based on reduced coordinate"""
265 qpoint = np.array(qpoint).astype(np.float64) % 1
266 for q, qpoint2 in enumerate(np.array(self._q).astype(np.float64)): 266 ↛ 269line 266 didn't jump to line 269 because the loop on line 266 didn't complete
267 if np.allclose(qpoint, qpoint2):
268 return QPoint(q, self)
269 raise ValueError('qpoint not compatible, check mp.q_reduced')
271 # Getters ans setters for internal mutable arrays
272 def get_Q(self) -> NDArray[float]:
273 """Return the mode coordinates as a ``(Nq, Nb)`` complex array in Å√dmu."""
274 return self._Q.copy()
276 def get_P(self) -> NDArray[float]:
277 """Return the mode momenta as a ``(Nq, Nb)`` complex array in √eV."""
278 return self._P.copy()
280 def get_F(self) -> NDArray[float]:
281 """Return the mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu."""
282 return self._F.copy()
284 def get_F_harmonic(self) -> NDArray[float]:
285 r"""Return the harmonic mode forces as a ``(Nq, Nb)`` complex array in eV/Å√dmu.
287 Computed as :math:`-\omega^2 Q`.
288 """
289 return -self._w2 * self.get_Q()
291 def set_Q(self, Q: NDArray[float]) -> None:
292 """Sets the internal mode coordinates :math:`Q`.
294 The function ensures :math:`Q(-q)=Q^*(q)`.
295 """
296 # This ensures that stuff like mp.set_Q(0) works while not updating the
297 # array until the assert
298 Q_new = self.get_Q()
299 Q_new[:] = Q
300 if not np.allclose(np.conjugate(Q_new), Q_new[self.q_minus]): 300 ↛ 301line 300 didn't jump to line 301 because the condition on line 300 was never true
301 raise ValueError('Supplied Q does not fulfill Q(-q) = Q(q)*')
302 self._Q[:] = Q_new
304 def set_P(self, P: NDArray[float]) -> None:
305 """Sets the internal mode momenta :math:`P`.
307 The function ensures :math:`P(-q)=P^*(q)`.
308 """
309 P_new = self.get_P()
310 P_new[:] = P
311 if not np.allclose(np.conjugate(P_new), P_new[self.q_minus]): 311 ↛ 312line 311 didn't jump to line 312 because the condition on line 311 was never true
312 raise ValueError('Supplied P does not fulfill P(-q) = P(q)*')
313 self._P[:] = P_new
315 def set_F(self, F: NDArray[float]) -> None:
316 """Sets the internal mode forces :math:`F`.
318 The function ensures :math:`F(-q)=F^*(q)`.
319 """
320 F_new = self.get_F()
321 F_new[:] = F
322 if not np.allclose(np.conjugate(F_new), F_new[self.q_minus]): 322 ↛ 323line 322 didn't jump to line 323 because the condition on line 322 was never true
323 raise ValueError('Supplied F does not fulfill F(-q) = F(q)*')
324 self._F[:] = F_new
326 def get_u(self) -> NDArray[float]:
327 """Return the atomic displacements as a ``(N, 3)`` array in Å."""
328 u = np.einsum('ksna,ks,n->na', self._X.conj(), self._Q, 1 / self.supercell.masses)
329 assert np.allclose(u.imag, 0)
330 return u.real
332 def get_v(self) -> NDArray[float]:
333 """Return the atomic velocities as a ``(N, 3)`` array in Å/fs."""
334 v = np.einsum('ksna,ks,n->na', self._X, self._P, 1 / self.supercell.masses)
335 assert np.allclose(v.imag, 0)
336 return v.real
338 def get_f(self) -> NDArray[float]:
339 """Return the atomic forces as a ``(N, 3)`` array in eV/Å."""
340 f = np.einsum('ksna,ks->na', self._X, self._F)
341 assert np.allclose(f.imag, 0)
342 return f.real
344 def get_f_harmonic(self) -> NDArray[float]:
345 """Return the harmonic atomic forces for the current displacements
346 as a ``(N, 3)`` array in eV/Å.
347 """
348 F_harmonic = self.get_F_harmonic()
349 f_harmonic = np.einsum('ksna,ks->na', self._X, F_harmonic)
350 assert np.allclose(f_harmonic.imag, 0)
351 return f_harmonic.real
353 def set_u(self, u: NDArray[float]) -> None:
354 """Sets the internal mode coordinates :math:`Q` given the atomic displacements :math:`u`.
356 .. math::
358 Q = X u
360 Parameters
361 ----------
362 u
363 The atomic displacements in Å.
364 """
365 Q = np.einsum('ksna,na->ks', self._X, u)
366 self.set_Q(Q)
368 def set_v(self, v: NDArray[float]) -> None:
369 """Sets the internal mode momenta :math:`P` given the atomic velocities :math:`v`.
371 .. math::
373 P = X^* * v
375 Parameters
376 ----------
377 v
378 The atomic velocities in Å/fs.
379 """
380 P = np.einsum('ksna,na->ks', self._X.conj(), v)
381 self.set_P(P)
383 def set_f(self, f: NDArray[float]) -> None:
384 """Sets the internal mode forces :math:`F` given the atomic forces :math:`f`.
386 .. math::
388 F = X^* * f / m
390 Parameters
391 ----------
392 f
393 The atomic forces in eV/Å.
394 """
395 F = np.einsum('ksna,na,n->ks', self._X.conj(), f, 1 / self.supercell.masses)
396 self.set_F(F)
398 # Convenience functions to handle ASE Atoms objects
399 def get_atoms(self, harmonic_forces: Optional[bool] = False) -> Atoms:
400 r"""Returns ASE :class:`Atoms` object with displacement,
401 velocities, forces, and harmonic energies.
403 Parameters
404 ----------
405 harmonic_forces
406 Whether the forces should be taken from the internal `F` or via `-\omega^2 Q`.
407 """
408 atoms = self.supercell.to_ase()
409 atoms.positions += self.get_u()
410 atoms.set_velocities(self.get_v() / fs)
411 E = self.potential_energies.sum()
412 f = self.get_f_harmonic() if harmonic_forces else self.get_f()
414 atoms.calc = SinglePointCalculator(
415 energy=E, forces=f, stress=None, magmoms=None, atoms=atoms)
417 return atoms
419 def update_from_atoms(self, atoms: Atoms) -> None:
420 """Updates the :class:`ModeProjector` objects with displacements, velocities,
421 and forces from an ASE :class:`Atoms` object.
423 Checks for an attached calculator in the first place and next for a forces array.
425 If no data sets corresponding array to zeros.
427 The masses and velocities are converted to dynasor units internally.
428 """
430 u = get_displacements(atoms, self.supercell)
431 if np.max(np.abs(u)) > 2.0: 431 ↛ 432line 431 didn't jump to line 432 because the condition on line 431 was never true
432 warnings.warn('Displacements larger than 2Å. Is the atoms object permuted?')
433 self.set_u(u)
434 self.set_v(atoms.get_velocities() * fs)
435 try:
436 self.set_f(atoms.get_forces())
437 except RuntimeError:
438 if 'forces' in atoms.arrays: 438 ↛ 441line 438 didn't jump to line 441 because the condition on line 438 was always true
439 self.set_f(atoms.arrays['forces'])
440 else:
441 self.set_f(np.zeros_like(atoms.positions))
443 # properties
444 @property
445 def q_minus(self) -> NDArray[float]:
446 """The index of the corresponding counter-propagating mode (:math:`-q`)."""
447 return self._q_minus.copy()
449 @property
450 def q_reduced(self) -> NDArray[float]:
451 """The q-points in reduced coordinates.
453 For example a zone boundary mode would be (1/2, 0, 0)
454 """
455 return self._q.astype(float)
457 @property
458 def q_cartesian(self) -> NDArray[float]:
459 """The q-points in cartesian coordinates with unit of rad/Å (2π included)."""
460 return 2 * np.pi * self.q_reduced @ self.primitive.inv_cell
462 @property
463 def omegas(self) -> NDArray[float]:
464 """The frequencies of each mode in rad/fs.
466 Following convention, negative values indicate imaginary frequencies.
467 """
468 return np.sign(self._w2) * np.sqrt(np.abs(self._w2))
470 @property
471 def polarizations(self) -> NDArray[float]:
472 """The polarization vectors for each mode `(Nq, Nb, Np, 3)`."""
473 return self._W
475 @property
476 def eigenmodes(self) -> NDArray[float]:
477 """The eigenmodes in the supercell as `(Nq, Nb, N, 3)`-array
479 The eigenmodes include the masses such that :math:`Q = X u`
480 where :math:`u` are the supercell displacements.
481 """
482 return self._X
484 @property
485 def potential_energies(self) -> NDArray[float]:
486 """Potential energy per mode as `(Nq, Nb)`-array.
488 The potential energies are defined as :math:`1/2 Q Q^*` and should equal
489 :math:`1/2 k_B T` in equilibrium for a harmonic system.
490 """
491 return 1 / 2 * np.abs(self._Q) ** 2 * self._w2
493 @property
494 def kinetic_energies(self) -> NDArray[float]:
495 """Kinetic energy per mode as `(Nq, Nb)`-array.
497 The kinetic energies are defined as :math:`1/2 P P^*`. Should equal
498 :math:`1/2 k_B T` in equilibrium.`
499 """
500 return 1 / 2 * np.abs(self._P)**2
502 @property
503 def virial_energies(self) -> NDArray[float]:
504 """The virial energies per mode as `(Nq, Nb)`-array.
506 The virial energies are defined here as :math:`-1/2 Q F`, which should have an
507 expectation value of :math:`1/2 k_B T` per mode in equilibrium. For a harmonic
508 system this is simply equal to the potential energy. This means that
509 the virial energy can be used to monitor the anharmonicity or
510 define a measure of the potential energy.
511 """
512 return -1 / 2 * self._Q * self._F
514 def write(self, file_name: str) -> None:
515 """Uses pickle to write mode projector to file."""
516 with open(file_name, 'wb') as f:
517 pickle.dump(self, f)
519 @classmethod
520 def read(cls, file_name: str) -> ModeProjector:
521 """Return :class:`ModeProjector` instance from pickle file
522 that was saved using :func:`~ModeProjector.write`."""
523 with open(file_name, 'rb') as f:
524 mp = pickle.load(f)
525 return mp