Coverage for ase / vibrations / data.py: 98.33%
180 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +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 if np.allclose(self._indices, range(len(self._atoms))):
315 indices = None
316 else:
317 indices = self.get_indices()
319 return {
320 'atoms': self.get_atoms(),
321 'hessian': self.get_hessian(),
322 'indices': indices,
323 }
325 @classmethod
326 def fromdict(cls, data: dict[str, Any]) -> 'VibrationsData':
327 # mypy is understandably suspicious of data coming from a dict that
328 # holds mixed types, but it can see if we sanity-check with 'assert'
329 assert isinstance(data['atoms'], Atoms)
330 assert isinstance(
331 data['hessian'], (collections.abc.Sequence, np.ndarray)
332 )
333 if data['indices'] is not None:
334 assert isinstance(
335 data['indices'], (collections.abc.Sequence, np.ndarray)
336 )
337 for index in data['indices']:
338 assert isinstance(index, Integral)
340 return cls(data['atoms'], data['hessian'], indices=data['indices'])
342 @cached_property
343 def _energies_and_modes(self) -> tuple[np.ndarray, np.ndarray]:
344 """Diagonalise the Hessian to obtain harmonic modes
346 This method is an internal implementation of get_energies_and_modes(),
347 see the docstring of that method for more information.
349 """
350 active_atoms = self._atoms[self.get_mask()]
351 n_atoms = len(active_atoms)
352 masses = active_atoms.get_masses()
354 if not np.all(masses):
355 raise ValueError(
356 'Zero mass encountered in one or more of '
357 'the vibrated atoms. Use Atoms.set_masses()'
358 ' to set all masses to non-zero values.'
359 )
360 mass_weights = np.repeat(masses**-0.5, 3)
362 omega2, vectors = np.linalg.eigh(
363 mass_weights * self.get_hessian_2d() * mass_weights[:, np.newaxis]
364 )
366 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)
367 energies = unit_conversion * omega2.astype(complex) ** 0.5
369 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)
370 modes = modes * masses[np.newaxis, :, np.newaxis] ** -0.5
372 return (energies, modes)
374 def get_energies_and_modes(
375 self,
376 all_atoms: bool = False,
377 ) -> tuple[np.ndarray, np.ndarray]:
378 """Diagonalise the Hessian to obtain harmonic modes
380 Results are cached so diagonalization will only be performed once for
381 this object instance.
383 Parameters
384 ----------
385 all_atoms : bool
386 If True, return modes as (3N, [N + N_frozen], 3) array where
387 the second axis corresponds to the full list of atoms in the
388 attached atoms object. Atoms that were not included in the
389 Hessian will have displacement vectors of (0, 0, 0).
391 Returns
392 -------
393 tuple[np.ndarray, np.ndarray]
395 Tuple of (energies, modes).
397 Energies are given in units of eV. (To convert these to frequencies
398 in cm-1, divide by ase.units.invcm.)
400 Modes are given in Cartesian coordinates as a (3N, N, 3) array
401 where indices correspond to the (mode_index, atom, direction).
403 """
405 energies, modes_from_hessian = self._energies_and_modes
407 if all_atoms:
408 n_active_atoms = len(self.get_indices())
409 n_all_atoms = len(self._atoms)
410 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))
411 modes[:, self.get_mask(), :] = modes_from_hessian
412 else:
413 modes = modes_from_hessian.copy()
415 return (energies.copy(), modes)
417 def get_modes(self, all_atoms: bool = False) -> np.ndarray:
418 """Diagonalise the Hessian to obtain harmonic modes
420 Results are cached so diagonalization will only be performed once for
421 this object instance.
423 Parameters
424 ----------
425 all_atoms : bool
426 If True, return modes as (3N, [N + N_frozen], 3) array where
427 the second axis corresponds to the full list of atoms in the
428 attached atoms object. Atoms that were not included in the
429 Hessian will have displacement vectors of (0, 0, 0).
431 Returns
432 -------
433 np.ndarray
434 Modes in Cartesian coordinates as a (3N, N, 3) array where indices
435 correspond to the (mode_index, atom, direction).
437 """
438 return self.get_energies_and_modes(all_atoms=all_atoms)[1]
440 def get_energies(self) -> np.ndarray:
441 """Diagonalise the Hessian to obtain eigenvalues
443 Results are cached so diagonalization will only be performed once for
444 this object instance.
446 Returns
447 -------
448 np.ndarray
449 Harmonic mode energies in units of eV.
451 """
452 return self.get_energies_and_modes()[0]
454 def get_frequencies(self) -> np.ndarray:
455 """Diagonalise the Hessian to obtain frequencies in cm^-1
457 Results are cached so diagonalization will only be performed once for
458 this object instance.
460 Returns
461 -------
462 Harmonic mode frequencies in units of cm^-1.
464 """
466 return self.get_energies() / units.invcm
468 def get_zero_point_energy(self) -> float:
469 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy
471 Returns
472 -------
473 float
474 Zero-point energy in eV.
476 """
477 return self._calculate_zero_point_energy(self.get_energies())
479 @staticmethod
480 def _calculate_zero_point_energy(
481 energies: Sequence[complex] | np.ndarray,
482 ) -> float:
483 return 0.5 * np.asarray(energies).real.sum()
485 def tabulate(self, im_tol: float = 1e-8) -> str:
486 """Get a summary of the vibrational frequencies.
488 Parameters
489 ----------
490 im_tol : float
491 Tolerance for imaginary frequency in eV.
492 If frequency has a larger imaginary component than im_tol,
493 the imaginary component is shown in the summary table.
495 Returns
496 -------
497 str
498 Summary table as formatted text.
500 """
502 energies = self.get_energies()
504 return (
505 '\n'.join(self._tabulate_from_energies(energies, im_tol=im_tol))
506 + '\n'
507 )
509 @classmethod
510 def _tabulate_from_energies(
511 cls,
512 energies: Sequence[complex] | np.ndarray,
513 im_tol: float = 1e-8,
514 ) -> list[str]:
515 summary_lines = [
516 '---------------------',
517 ' # meV cm^-1',
518 '---------------------',
519 ]
521 for n, e in enumerate(energies):
522 if abs(e.imag) > im_tol:
523 c = 'i'
524 e = e.imag
525 else:
526 c = ''
527 e = e.real
529 summary_lines.append(
530 '{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'.format(
531 index=n, mev=(e * 1e3), cm=(e / units.invcm), im=c
532 )
533 )
534 summary_lines.append('---------------------')
535 summary_lines.append(
536 'Zero-point energy: {:.3f} eV'.format(
537 cls._calculate_zero_point_energy(energies=energies)
538 )
539 )
541 return summary_lines
543 def iter_animated_mode(
544 self,
545 mode_index: int,
546 temperature: float = units.kB * 300.0,
547 frames: int = 30,
548 ) -> Iterator[Atoms]:
549 """Obtain animated mode as a series of Atoms
551 Parameters
552 ----------
553 mode_index : int
554 Selection of mode to animate.
555 temperature : float
556 Temperature in energy units - use ``units.kB * T_IN_KELVIN``
557 frames : int
558 Number of image frames in animation.
560 Yields
561 ------
562 :class:`~ase.Atoms`
563 Displaced atoms following vibrational mode.
565 """
567 mode = self.get_modes(all_atoms=True)[mode_index] * sqrt(
568 temperature / abs(self.get_energies()[mode_index])
569 )
571 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
572 atoms = self.get_atoms()
573 atoms.positions += sin(phase) * mode
575 yield atoms
577 def show_as_force(
578 self,
579 mode: int,
580 scale: float = 0.2,
581 show: bool = True,
582 ) -> Atoms:
583 """Illustrate mode as "forces" on atoms
585 Parameters
586 ----------
587 mode : int
588 Mode index.
589 scale : float
590 Scale factor.
591 show : bool
592 If True, open the ASE GUI and show atoms.
594 Returns
595 -------
596 :class:`~ase.Atoms`
597 Atoms with scaled forces corresponding to mode eigenvectors (using
598 attached SinglePointCalculator).
600 """
602 atoms = self.get_atoms()
603 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
604 atoms.calc = SinglePointCalculator(atoms, forces=mode)
605 if show:
606 atoms.edit()
608 return atoms
610 def write_jmol(
611 self,
612 filename: str = 'vib.xyz',
613 ir_intensities: Sequence[float] | np.ndarray | None = None,
614 ) -> None:
615 """Writes file for viewing of the modes with jmol.
617 This is an extended XYZ file with eigenvectors given as extra columns
618 and metadata given in the label/comment line for each image. The format
619 is not quite human-friendly, but has the advantage that it can be
620 imported back into ASE with ase.io.read.
622 Parameters
623 ----------
624 filename : str
625 Path for output file.
626 ir_intensities :
627 If available, IR intensities can be included in the header lines.
628 This does not affect the visualisation but may be convenient when
629 comparing to experimental data.
631 """
633 all_images = list(
634 self._get_jmol_images(
635 atoms=self.get_atoms(),
636 energies=self.get_energies(),
637 modes=self.get_modes(all_atoms=True),
638 ir_intensities=ir_intensities,
639 )
640 )
641 ase.io.write(filename, all_images, format='extxyz')
643 @staticmethod
644 def _get_jmol_images(
645 atoms: Atoms,
646 energies: np.ndarray,
647 modes: np.ndarray,
648 ir_intensities: Sequence[float] | np.ndarray | None = None,
649 ) -> Iterator[Atoms]:
650 """Get vibrational modes as a series of Atoms with attached data
652 For each image (Atoms object):
654 - eigenvalues are attached to image.arrays['mode']
655 - "mode#" and "frequency_cm-1" are set in image.info
656 - "IR_intensity" is set if provided in ir_intensities
657 - "masses" is removed
659 This is intended to set up the object for JMOL-compatible export using
660 ase.io.extxyz.
662 Parameters
663 ----------
664 atoms : :class:`~ase.Atoms`
665 The base atoms object; all images have the same positions
666 energies : np.ndarray
667 Complex vibrational energies in eV.
668 modes : np.ndarray
669 Eigenvectors array corresponding to atoms and energies.
670 This should cover the full set of atoms (i.e. modes =
671 vib.get_modes(all_atoms=True)).
672 ir_intensities: Sequence[float] | np.ndarray | None = None
673 If available, IR intensities can be included in the header lines.
674 This does not affect the visualisation but may be convenient when
675 comparing to experimental data.
677 Returns
678 -------
679 Iterator[Atoms]
681 """
682 for i, (energy, mode) in enumerate(zip(energies, modes)):
683 # write imaginary frequencies as negative numbers
684 if energy.imag > energy.real:
685 energy = float(-energy.imag)
686 else:
687 energy = energy.real
689 image = atoms.copy()
690 image.info.update(
691 {
692 'mode#': str(i),
693 'frequency_cm-1': energy / units.invcm,
694 }
695 )
696 image.arrays['mode'] = mode
698 # Custom masses are quite useful in vibration analysis, but will
699 # show up in the xyz file unless we remove them
700 if image.has('masses'):
701 del image.arrays['masses']
703 if ir_intensities is not None:
704 image.info['IR_intensity'] = float(ir_intensities[i])
706 yield image
708 def get_dos(self) -> RawDOSData:
709 """Total phonon DOS"""
710 energies = self.get_energies()
711 return RawDOSData(energies, np.ones_like(energies))
713 def get_pdos(self) -> DOSCollection:
714 """Phonon DOS, including atomic contributions"""
715 energies = self.get_energies()
716 masses = self._atoms[self.get_mask()].get_masses()
718 # Get weights as N_moving_atoms x N_modes array
719 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis] ** -0.5
720 all_weights = (np.linalg.norm(vectors, axis=-1) ** 2).T
722 mask = self.get_mask()
723 all_info = [
724 {'index': i, 'symbol': a.symbol}
725 for i, a in enumerate(self._atoms)
726 if mask[i]
727 ]
729 return DOSCollection(
730 [
731 RawDOSData(energies, weights, info=info)
732 for weights, info in zip(all_weights, all_info)
733 ]
734 )
736 def with_new_masses(
737 self: VD,
738 masses: Sequence[float] | np.ndarray,
739 ) -> VD:
740 """Get a copy of vibrations with modified masses and the same Hessian
742 Parameters
743 ----------
744 masses: Sequence[float] | np.ndarray
745 New sequence of masses corresponding to the atom order in
746 ``self.get_atoms()``.
748 Returns
749 -------
750 :class:`VibrationalData`
751 A copy of the data with new masses for the same Hessian.
753 """
755 new_atoms = self.get_atoms()
756 new_atoms.set_masses(masses)
757 return self.__class__(
758 new_atoms, self.get_hessian(), indices=self.get_indices()
759 )