Coverage for ase / vibrations / data.py: 98.90%
181 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-21 15:52 +0000
1"""Storage and analysis for vibrational data"""
3import collections
4from functools import cached_property
5from math import pi, sin, sqrt
6from numbers import Integral, Real
7from typing import Any, Iterator, Sequence, TypeVar
9import numpy as np
11import ase.io
12import ase.units as units
13from ase.atoms import Atoms
14from ase.calculators.singlepoint import SinglePointCalculator
15from ase.constraints import FixAtoms, FixCartesian, constrained_indices
16from ase.spectrum.doscollection import DOSCollection
17from ase.spectrum.dosdata import RawDOSData
18from ase.utils import jsonable
20RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]]
21VD = TypeVar('VD', bound='VibrationsData')
24@jsonable('vibrationsdata')
25class VibrationsData:
26 """Class for storing and analyzing vibrational data (i.e. Atoms + Hessian)
28 This class is not responsible for calculating Hessians; the Hessian should
29 be computed by a Calculator or some other algorithm.
30 Once :class:`VibrationsData` has been constructed, this class provides some
31 common processing options; frequency calculation, mode animation, DOS etc.
33 If the Atoms object is a periodic supercell, :class:`VibrationsData` may be
34 converted to a PhononData using the VibrationsData.to_phonondata() method.
35 This provides access to q-point-dependent analyses such as phonon
36 dispersion plotting.
38 If the Atoms object has :class:`~ase.constraints.FixAtoms` or
39 :class:`~.ase.constraints.FixCartesian` constraints,
40 these will be respected and the Hessian should be sized accordingly.
42 Parameters
43 ----------
44 atoms : :class:`~ase.Atoms`
45 Equilibrium geometry of vibrating system.
46 This will be stored as a full copy.
48 hessian : np.ndarray
49 Second-derivative in energy with respect to Cartesian nuclear movements
50 as an (N, 3, N, 3) array.
52 indices : Sequence[int] | np.ndarray | None, default: None
53 Indices of atoms which are included in Hessian.
54 By default, it includes all freely moving atoms (i.e. not fixed ones).
55 Leave at None if constraints should be determined automatically from
56 the atoms object.
58 """
60 def __init__(
61 self,
62 atoms: Atoms,
63 hessian: RealSequence4D | np.ndarray,
64 indices: Sequence[int] | np.ndarray | None = None,
65 ) -> None:
66 if indices is None:
67 indices = np.asarray(
68 self.indices_from_constraints(atoms), dtype=int
69 )
71 self._indices = np.array(indices, dtype=int)
73 n_atoms = self._check_dimensions(
74 atoms, np.asarray(hessian), indices=self._indices
75 )
76 self._atoms = atoms.copy()
78 self._hessian2d = (
79 np.asarray(hessian).reshape(3 * n_atoms, 3 * n_atoms).copy()
80 )
82 _setter_error = (
83 'VibrationsData properties cannot be modified: construct '
84 'a new VibrationsData with consistent atoms, Hessian and '
85 '(optionally) indices/mask.'
86 )
88 @classmethod
89 def from_2d(
90 cls,
91 atoms: Atoms,
92 hessian_2d: Sequence[Sequence[Real]] | np.ndarray,
93 indices: Sequence[int] | None = None,
94 ) -> 'VibrationsData':
95 """Instantiate VibrationsData when the Hessian is in a 3Nx3N format
97 Parameters
98 ----------
99 atoms : :class:`~ase.Atoms`
100 Equilibrium geometry of vibrating system.
102 hessian : Sequence[Sequence[Real]] | np.ndarray
103 Second-derivative in energy with respect to Cartesian nuclear
104 movements as a (3N, 3N) array.
106 indices : Sequence[int] | None, default: None
107 Indices of (non-frozen) atoms included in Hessian.
109 """
110 if indices is None:
111 indices = range(len(atoms))
112 assert indices is not None # Show Mypy that indices is now a sequence
114 hessian_2d_array = np.asarray(hessian_2d)
115 n_atoms = cls._check_dimensions(
116 atoms, hessian_2d_array, indices=indices, two_d=True
117 )
119 return cls(
120 atoms,
121 hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3),
122 indices=indices,
123 )
125 @staticmethod
126 def indices_from_constraints(atoms: Atoms) -> list[int]:
127 """Indices corresponding to Atoms Constraints
129 Deduces the freely moving atoms from the constraints set on the
130 atoms object. :class:`VibrationsData` only supports:
132 - :class:`~ase.constraints.FixCartesian`
133 - :class:`~ase.constraints.FixAtoms`
135 All others are neglected.
137 Parameters
138 ----------
139 atoms : :class:`~ase.Atoms`
141 Returns
142 -------
143 list[int]
144 Indices of free atoms.
146 """
147 # Only fully fixed atoms supported by VibrationsData
148 const_indices = constrained_indices(
149 atoms, only_include=(FixCartesian, FixAtoms)
150 )
151 # Invert the selection to get free atoms
152 indices = np.setdiff1d(
153 np.array(range(len(atoms))), const_indices
154 ).astype(int)
155 # TODO: use numpy.typing to resolve this error.
156 return indices.tolist() # type: ignore[return-value]
158 @staticmethod
159 def indices_from_mask(
160 mask: Sequence[bool] | np.ndarray,
161 ) -> list[int]:
162 """Indices corresponding to boolean mask
164 This is provided as a convenience for instantiating
165 :class:`VibrationsData` with a boolean mask.
166 For example, if the Hessian data includes only the H
167 atoms in a structure::
169 h_mask = atoms.get_chemical_symbols() == 'H'
170 vib_data = VibrationsData(atoms, hessian,
171 VibrationsData.indices_from_mask(h_mask))
173 Take care to ensure that the length of the mask corresponds to the full
174 number of atoms; this function is only aware of the mask it has been
175 given.
177 Parameters
178 ----------
179 mask: Sequence[bool] | np.ndarray
180 Sequence of boolean values.
182 Returns
183 -------
184 list[int]
185 Indices of :obj:`True` elements.
187 """
188 # TODO: use numpy.typing to resolve this error.
189 return np.where(mask)[0].tolist() # type: ignore[return-value]
191 @staticmethod
192 def _check_dimensions(
193 atoms: Atoms,
194 hessian: np.ndarray,
195 indices: np.ndarray | Sequence[int],
196 two_d: bool = False,
197 ) -> int:
198 """Sanity check on array shapes from input data
200 Parameters
201 ----------
202 atoms : :class:`~ase.Atoms`
203 hessian : np.ndarray
204 Proposed Hessian array.
205 indices : np.ndarray | Sequence[int]
206 Indices of atoms used in Hessian.
208 Returns
209 -------
210 Number of atoms contributing to Hessian.
212 Raises
213 ------
214 ValueError
215 If Hessian dimensions are not (N, 3, N, 3).
217 """
219 n_atoms = len(atoms[indices])
221 if two_d:
222 ref_shape = [n_atoms * 3, n_atoms * 3]
223 ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3))
225 else:
226 ref_shape = [n_atoms, 3, n_atoms, 3]
227 ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms)
229 if isinstance(hessian, np.ndarray) and hessian.shape == tuple(
230 ref_shape
231 ):
232 return n_atoms
233 else:
234 raise ValueError(
235 'Hessian for these atoms should be a {} numpy array.'.format(
236 ref_shape_txt
237 )
238 )
240 def get_atoms(self) -> Atoms:
241 return self._atoms.copy()
243 def get_indices(self) -> np.ndarray:
244 return self._indices.copy()
246 def get_mask(self) -> np.ndarray:
247 """Boolean mask of atoms selected by indices"""
248 return self._mask_from_indices(self._atoms, self.get_indices())
250 @staticmethod
251 def _mask_from_indices(
252 atoms: Atoms,
253 indices: Sequence[int] | np.ndarray | None,
254 ) -> np.ndarray:
255 """Boolean mask of atoms selected by indices"""
256 natoms = len(atoms)
258 # Wrap indices to allow negative values
259 indices = np.asarray(indices) % natoms
261 mask = np.full(natoms, False, dtype=bool)
262 mask[indices] = True
263 return mask
265 def get_hessian(self) -> np.ndarray:
266 """The Hessian; second derivative of energy wrt positions
268 This format is preferred for iteration over atoms and when
269 addressing specific elements of the Hessian.
271 Returns
272 -------
273 np.ndarray
274 Array with shape (n_atoms, 3, n_atoms, 3) where
276 - the first and third indices identify atoms in self.get_atoms()
278 - the second and fourth indices cover the corresponding
279 Cartesian movements in x, y, z
281 e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on
282 atoms[1] in the x-direction in response to a movement in the
283 z-direction of atoms[0]
284 """
285 n_atoms = int(self._hessian2d.shape[0] / 3)
286 return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy()
288 def get_hessian_2d(self) -> np.ndarray:
289 """Get the Hessian as a 2-D array
291 This format may be preferred for use with standard linear algebra
292 functions
294 Returns
295 -------
296 np.ndarray
297 Array with shape (n_atoms * 3, n_atoms * 3) where the elements are
298 ordered by atom and Cartesian direction::
300 >> [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...],
301 >> [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...],
302 >> [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...],
303 >> [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...],
304 >> ...]
307 e.g. the element h[2, 3] gives a harmonic force exerted on
308 atoms[1] in the x-direction in response to a movement in the
309 z-direction of atoms[0]
310 """
311 return self._hessian2d.copy()
313 def todict(self) -> dict[str, Any]:
314 unconstrained_indices = self.indices_from_constraints(self._atoms)
315 if (
316 len(self._indices) == len(unconstrained_indices)
317 and (self._indices == unconstrained_indices).all()
318 ):
319 indices = None
320 else:
321 indices = self.get_indices()
323 return {
324 'atoms': self.get_atoms(),
325 'hessian': self.get_hessian(),
326 'indices': indices,
327 }
329 @classmethod
330 def fromdict(cls, data: dict[str, Any]) -> 'VibrationsData':
331 # mypy is understandably suspicious of data coming from a dict that
332 # holds mixed types, but it can see if we sanity-check with 'assert'
333 assert isinstance(data['atoms'], Atoms)
334 assert isinstance(
335 data['hessian'], (collections.abc.Sequence, np.ndarray)
336 )
337 if data['indices'] is not None:
338 assert isinstance(
339 data['indices'], (collections.abc.Sequence, np.ndarray)
340 )
341 for index in data['indices']:
342 assert isinstance(index, Integral)
344 return cls(data['atoms'], data['hessian'], indices=data['indices'])
346 @cached_property
347 def _energies_and_modes(self) -> tuple[np.ndarray, np.ndarray]:
348 """Diagonalise the Hessian to obtain harmonic modes
350 This method is an internal implementation of get_energies_and_modes(),
351 see the docstring of that method for more information.
353 """
354 active_atoms = self._atoms[self.get_mask()]
355 n_atoms = len(active_atoms)
356 masses = active_atoms.get_masses()
358 if not np.all(masses):
359 raise ValueError(
360 'Zero mass encountered in one or more of '
361 'the vibrated atoms. Use Atoms.set_masses()'
362 ' to set all masses to non-zero values.'
363 )
364 mass_weights = np.repeat(masses**-0.5, 3)
366 omega2, vectors = np.linalg.eigh(
367 mass_weights * self.get_hessian_2d() * mass_weights[:, np.newaxis]
368 )
370 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)
371 energies = unit_conversion * omega2.astype(complex) ** 0.5
373 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)
374 modes = modes * masses[np.newaxis, :, np.newaxis] ** -0.5
376 return (energies, modes)
378 def get_energies_and_modes(
379 self,
380 all_atoms: bool = False,
381 ) -> tuple[np.ndarray, np.ndarray]:
382 """Diagonalise the Hessian to obtain harmonic modes
384 Results are cached so diagonalization will only be performed once for
385 this object instance.
387 Parameters
388 ----------
389 all_atoms : bool
390 If True, return modes as (3N, [N + N_frozen], 3) array where
391 the second axis corresponds to the full list of atoms in the
392 attached atoms object. Atoms that were not included in the
393 Hessian will have displacement vectors of (0, 0, 0).
395 Returns
396 -------
397 tuple[np.ndarray, np.ndarray]
399 Tuple of (energies, modes).
401 Energies are given in units of eV. (To convert these to frequencies
402 in cm-1, divide by ase.units.invcm.)
404 Modes are given in Cartesian coordinates as a (3N, N, 3) array
405 where indices correspond to the (mode_index, atom, direction).
407 """
409 energies, modes_from_hessian = self._energies_and_modes
411 if all_atoms:
412 n_active_atoms = len(self.get_indices())
413 n_all_atoms = len(self._atoms)
414 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))
415 modes[:, self.get_mask(), :] = modes_from_hessian
416 else:
417 modes = modes_from_hessian.copy()
419 return (energies.copy(), modes)
421 def get_modes(self, all_atoms: bool = False) -> np.ndarray:
422 """Diagonalise the Hessian to obtain harmonic modes
424 Results are cached so diagonalization will only be performed once for
425 this object instance.
427 Parameters
428 ----------
429 all_atoms : bool
430 If True, return modes as (3N, [N + N_frozen], 3) array where
431 the second axis corresponds to the full list of atoms in the
432 attached atoms object. Atoms that were not included in the
433 Hessian will have displacement vectors of (0, 0, 0).
435 Returns
436 -------
437 np.ndarray
438 Modes in Cartesian coordinates as a (3N, N, 3) array where indices
439 correspond to the (mode_index, atom, direction).
441 """
442 return self.get_energies_and_modes(all_atoms=all_atoms)[1]
444 def get_energies(self) -> np.ndarray:
445 """Diagonalise the Hessian to obtain eigenvalues
447 Results are cached so diagonalization will only be performed once for
448 this object instance.
450 Returns
451 -------
452 np.ndarray
453 Harmonic mode energies in units of eV.
455 """
456 return self.get_energies_and_modes()[0]
458 def get_frequencies(self) -> np.ndarray:
459 """Diagonalise the Hessian to obtain frequencies in cm^-1
461 Results are cached so diagonalization will only be performed once for
462 this object instance.
464 Returns
465 -------
466 Harmonic mode frequencies in units of cm^-1.
468 """
470 return self.get_energies() / units.invcm
472 def get_zero_point_energy(self) -> float:
473 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy
475 Returns
476 -------
477 float
478 Zero-point energy in eV.
480 """
481 return self._calculate_zero_point_energy(self.get_energies())
483 @staticmethod
484 def _calculate_zero_point_energy(
485 energies: Sequence[complex] | np.ndarray,
486 ) -> float:
487 return 0.5 * np.asarray(energies).real.sum()
489 def tabulate(self, im_tol: float = 1e-8) -> str:
490 """Get a summary of the vibrational frequencies.
492 Parameters
493 ----------
494 im_tol : float
495 Tolerance for imaginary frequency in eV.
496 If frequency has a larger imaginary component than im_tol,
497 the imaginary component is shown in the summary table.
499 Returns
500 -------
501 str
502 Summary table as formatted text.
504 """
506 energies = self.get_energies()
508 return (
509 '\n'.join(self._tabulate_from_energies(energies, im_tol=im_tol))
510 + '\n'
511 )
513 @classmethod
514 def _tabulate_from_energies(
515 cls,
516 energies: Sequence[complex] | np.ndarray,
517 im_tol: float = 1e-8,
518 ) -> list[str]:
519 summary_lines = [
520 '---------------------',
521 ' # meV cm^-1',
522 '---------------------',
523 ]
525 for n, e in enumerate(energies):
526 if abs(e.imag) > im_tol:
527 c = 'i'
528 e = e.imag
529 else:
530 c = ''
531 e = e.real
533 summary_lines.append(
534 '{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'.format(
535 index=n, mev=(e * 1e3), cm=(e / units.invcm), im=c
536 )
537 )
538 summary_lines.append('---------------------')
539 summary_lines.append(
540 'Zero-point energy: {:.3f} eV'.format(
541 cls._calculate_zero_point_energy(energies=energies)
542 )
543 )
545 return summary_lines
547 def iter_animated_mode(
548 self,
549 mode_index: int,
550 temperature: float = units.kB * 300.0,
551 frames: int = 30,
552 ) -> Iterator[Atoms]:
553 """Obtain animated mode as a series of Atoms
555 Parameters
556 ----------
557 mode_index : int
558 Selection of mode to animate.
559 temperature : float
560 Temperature in energy units - use ``units.kB * T_IN_KELVIN``
561 frames : int
562 Number of image frames in animation.
564 Yields
565 ------
566 :class:`~ase.Atoms`
567 Displaced atoms following vibrational mode.
569 """
571 mode = self.get_modes(all_atoms=True)[mode_index] * sqrt(
572 temperature / abs(self.get_energies()[mode_index])
573 )
575 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
576 atoms = self.get_atoms()
577 atoms.positions += sin(phase) * mode
579 yield atoms
581 def show_as_force(
582 self,
583 mode: int,
584 scale: float = 0.2,
585 show: bool = True,
586 ) -> Atoms:
587 """Illustrate mode as "forces" on atoms
589 Parameters
590 ----------
591 mode : int
592 Mode index.
593 scale : float
594 Scale factor.
595 show : bool
596 If True, open the ASE GUI and show atoms.
598 Returns
599 -------
600 :class:`~ase.Atoms`
601 Atoms with scaled forces corresponding to mode eigenvectors (using
602 attached SinglePointCalculator).
604 """
606 atoms = self.get_atoms()
607 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
608 atoms.calc = SinglePointCalculator(atoms, forces=mode)
609 if show:
610 atoms.edit()
612 return atoms
614 def write_jmol(
615 self,
616 filename: str = 'vib.xyz',
617 ir_intensities: Sequence[float] | np.ndarray | None = None,
618 ) -> None:
619 """Writes file for viewing of the modes with jmol.
621 This is an extended XYZ file with eigenvectors given as extra columns
622 and metadata given in the label/comment line for each image. The format
623 is not quite human-friendly, but has the advantage that it can be
624 imported back into ASE with ase.io.read.
626 Parameters
627 ----------
628 filename : str
629 Path for output file.
630 ir_intensities :
631 If available, IR intensities can be included in the header lines.
632 This does not affect the visualisation but may be convenient when
633 comparing to experimental data.
635 """
637 all_images = list(
638 self._get_jmol_images(
639 atoms=self.get_atoms(),
640 energies=self.get_energies(),
641 modes=self.get_modes(all_atoms=True),
642 ir_intensities=ir_intensities,
643 )
644 )
645 ase.io.write(filename, all_images, format='extxyz')
647 @staticmethod
648 def _get_jmol_images(
649 atoms: Atoms,
650 energies: np.ndarray,
651 modes: np.ndarray,
652 ir_intensities: Sequence[float] | np.ndarray | None = None,
653 ) -> Iterator[Atoms]:
654 """Get vibrational modes as a series of Atoms with attached data
656 For each image (Atoms object):
658 - eigenvalues are attached to image.arrays['mode']
659 - "mode#" and "frequency_cm-1" are set in image.info
660 - "IR_intensity" is set if provided in ir_intensities
661 - "masses" is removed
663 This is intended to set up the object for JMOL-compatible export using
664 ase.io.extxyz.
666 Parameters
667 ----------
668 atoms : :class:`~ase.Atoms`
669 The base atoms object; all images have the same positions
670 energies : np.ndarray
671 Complex vibrational energies in eV.
672 modes : np.ndarray
673 Eigenvectors array corresponding to atoms and energies.
674 This should cover the full set of atoms (i.e. modes =
675 vib.get_modes(all_atoms=True)).
676 ir_intensities: Sequence[float] | np.ndarray | None = None
677 If available, IR intensities can be included in the header lines.
678 This does not affect the visualisation but may be convenient when
679 comparing to experimental data.
681 Returns
682 -------
683 Iterator[Atoms]
685 """
686 for i, (energy, mode) in enumerate(zip(energies, modes)):
687 # write imaginary frequencies as negative numbers
688 if energy.imag > energy.real:
689 energy = float(-energy.imag)
690 else:
691 energy = energy.real
693 image = atoms.copy()
694 image.info.update(
695 {
696 'mode#': str(i),
697 'frequency_cm-1': energy / units.invcm,
698 }
699 )
700 image.arrays['mode'] = mode
702 # Custom masses are quite useful in vibration analysis, but will
703 # show up in the xyz file unless we remove them
704 if image.has('masses'):
705 del image.arrays['masses']
707 if ir_intensities is not None:
708 image.info['IR_intensity'] = float(ir_intensities[i])
710 yield image
712 def get_dos(self) -> RawDOSData:
713 """Total phonon DOS"""
714 energies = self.get_energies()
715 return RawDOSData(energies, np.ones_like(energies))
717 def get_pdos(self) -> DOSCollection:
718 """Phonon DOS, including atomic contributions"""
719 energies = self.get_energies()
720 masses = self._atoms[self.get_mask()].get_masses()
722 # Get weights as N_moving_atoms x N_modes array
723 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis] ** -0.5
724 all_weights = (np.linalg.norm(vectors, axis=-1) ** 2).T
726 mask = self.get_mask()
727 all_info = [
728 {'index': i, 'symbol': a.symbol}
729 for i, a in enumerate(self._atoms)
730 if mask[i]
731 ]
733 return DOSCollection(
734 [
735 RawDOSData(energies, weights, info=info)
736 for weights, info in zip(all_weights, all_info)
737 ]
738 )
740 def with_new_masses(
741 self: VD,
742 masses: Sequence[float] | np.ndarray,
743 ) -> VD:
744 """Get a copy of vibrations with modified masses and the same Hessian
746 Parameters
747 ----------
748 masses: Sequence[float] | np.ndarray
749 New sequence of masses corresponding to the atom order in
750 ``self.get_atoms()``.
752 Returns
753 -------
754 :class:`VibrationalData`
755 A copy of the data with new masses for the same Hessian.
757 """
759 new_atoms = self.get_atoms()
760 new_atoms.set_masses(masses)
761 return self.__class__(
762 new_atoms, self.get_hessian(), indices=self.get_indices()
763 )