Coverage for /builds/ase/ase/ase/vibrations/data.py: 98.33%
180 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3"""Storage and analysis for vibrational data"""
5import collections
6from functools import cached_property
7from math import pi, sin, sqrt
8from numbers import Integral, Real
9from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union
11import numpy as np
13import ase.io
14import ase.units as units
15from ase.atoms import Atoms
16from ase.calculators.singlepoint import SinglePointCalculator
17from ase.constraints import FixAtoms, FixCartesian, constrained_indices
18from ase.spectrum.doscollection import DOSCollection
19from ase.spectrum.dosdata import RawDOSData
20from ase.utils import jsonable
22RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]]
23VD = TypeVar('VD', bound='VibrationsData')
26@jsonable('vibrationsdata')
27class VibrationsData:
28 """Class for storing and analyzing vibrational data (i.e. Atoms + Hessian)
30 This class is not responsible for calculating Hessians; the Hessian should
31 be computed by a Calculator or some other algorithm. Once the
32 VibrationsData has been constructed, this class provides some common
33 processing options; frequency calculation, mode animation, DOS etc.
35 If the Atoms object is a periodic supercell, VibrationsData may be
36 converted to a PhononData using the VibrationsData.to_phonondata() method.
37 This provides access to q-point-dependent analyses such as phonon
38 dispersion plotting.
40 If the Atoms object has FixedAtoms or FixedCartesian constraints, these
41 will be respected and the Hessian should be sized accordingly.
43 Args:
44 atoms:
45 Equilibrium geometry of vibrating system. This will be stored as a
46 full copy.
48 hessian: Second-derivative in energy with respect to
49 Cartesian nuclear movements as an (N, 3, N, 3) array.
51 indices: indices of atoms which are included
52 in Hessian. Default value (None) includes all freely
53 moving atoms (i.e. not fixed ones). Leave at None if
54 constraints should be determined automatically from the
55 atoms object.
57 """
59 def __init__(self,
60 atoms: Atoms,
61 hessian: Union[RealSequence4D, np.ndarray],
62 indices: Union[Sequence[int], np.ndarray] = None,
63 ) -> None:
65 if indices is None:
66 indices = np.asarray(self.indices_from_constraints(atoms),
67 dtype=int)
69 self._indices = np.array(indices, dtype=int)
71 n_atoms = self._check_dimensions(atoms, np.asarray(hessian),
72 indices=self._indices)
73 self._atoms = atoms.copy()
75 self._hessian2d = (np.asarray(hessian)
76 .reshape(3 * n_atoms, 3 * n_atoms).copy())
78 _setter_error = ("VibrationsData properties cannot be modified: construct "
79 "a new VibrationsData with consistent atoms, Hessian and "
80 "(optionally) indices/mask.")
82 @classmethod
83 def from_2d(cls, atoms: Atoms,
84 hessian_2d: Union[Sequence[Sequence[Real]], np.ndarray],
85 indices: Sequence[int] = None) -> 'VibrationsData':
86 """Instantiate VibrationsData when the Hessian is in a 3Nx3N format
88 Args:
89 atoms: Equilibrium geometry of vibrating system
91 hessian: Second-derivative in energy with respect to
92 Cartesian nuclear movements as a (3N, 3N) array.
94 indices: Indices of (non-frozen) atoms included in Hessian
96 """
97 if indices is None:
98 indices = range(len(atoms))
99 assert indices is not None # Show Mypy that indices is now a sequence
101 hessian_2d_array = np.asarray(hessian_2d)
102 n_atoms = cls._check_dimensions(atoms, hessian_2d_array,
103 indices=indices, two_d=True)
105 return cls(atoms, hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3),
106 indices=indices)
108 @staticmethod
109 def indices_from_constraints(atoms: Atoms) -> List[int]:
110 """Indices corresponding to Atoms Constraints
112 Deduces the freely moving atoms from the constraints set on the
113 atoms object. VibrationsData only supports FixCartesian and
114 FixAtoms. All others are neglected.
116 Args:
117 atoms: Atoms object.
119 Retruns:
120 indices of free atoms.
122 """
123 # Only fully fixed atoms supported by VibrationsData
124 const_indices = constrained_indices(
125 atoms, only_include=(FixCartesian, FixAtoms))
126 # Invert the selection to get free atoms
127 indices = np.setdiff1d(
128 np.array(
129 range(
130 len(atoms))),
131 const_indices).astype(int)
132 # TODO: use numpy.typing to resolve this error.
133 return indices.tolist() # type: ignore[return-value]
135 @staticmethod
136 def indices_from_mask(mask: Union[Sequence[bool], np.ndarray]
137 ) -> List[int]:
138 """Indices corresponding to boolean mask
140 This is provided as a convenience for instantiating VibrationsData with
141 a boolean mask. For example, if the Hessian data includes only the H
142 atoms in a structure::
144 h_mask = atoms.get_chemical_symbols() == 'H'
145 vib_data = VibrationsData(atoms, hessian,
146 VibrationsData.indices_from_mask(h_mask))
148 Take care to ensure that the length of the mask corresponds to the full
149 number of atoms; this function is only aware of the mask it has been
150 given.
152 Args:
153 mask: a sequence of True, False values
155 Returns:
156 indices of True elements
158 """
159 # TODO: use numpy.typing to resolve this error.
160 return np.where(mask)[0].tolist() # type: ignore[return-value]
162 @staticmethod
163 def _check_dimensions(atoms: Atoms,
164 hessian: np.ndarray,
165 indices: Union[np.ndarray, Sequence[int]],
166 two_d: bool = False) -> int:
167 """Sanity check on array shapes from input data
169 Args:
170 atoms: Structure
171 indices: Indices of atoms used in Hessian
172 hessian: Proposed Hessian array
174 Returns:
175 Number of atoms contributing to Hessian
177 Raises:
178 ValueError if Hessian dimensions are not (N, 3, N, 3)
180 """
182 n_atoms = len(atoms[indices])
184 if two_d:
185 ref_shape = [n_atoms * 3, n_atoms * 3]
186 ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3))
188 else:
189 ref_shape = [n_atoms, 3, n_atoms, 3]
190 ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms)
192 if (isinstance(hessian, np.ndarray)
193 and hessian.shape == tuple(ref_shape)):
194 return n_atoms
195 else:
196 raise ValueError("Hessian for these atoms should be a "
197 "{} numpy array.".format(ref_shape_txt))
199 def get_atoms(self) -> Atoms:
200 return self._atoms.copy()
202 def get_indices(self) -> np.ndarray:
203 return self._indices.copy()
205 def get_mask(self) -> np.ndarray:
206 """Boolean mask of atoms selected by indices"""
207 return self._mask_from_indices(self._atoms, self.get_indices())
209 @staticmethod
210 def _mask_from_indices(atoms: Atoms,
211 indices: Union[None, Sequence[int], np.ndarray]
212 ) -> np.ndarray:
213 """Boolean mask of atoms selected by indices"""
214 natoms = len(atoms)
216 # Wrap indices to allow negative values
217 indices = np.asarray(indices) % natoms
219 mask = np.full(natoms, False, dtype=bool)
220 mask[indices] = True
221 return mask
223 def get_hessian(self) -> np.ndarray:
224 """The Hessian; second derivative of energy wrt positions
226 This format is preferred for iteration over atoms and when
227 addressing specific elements of the Hessian.
229 Returns:
230 array with shape (n_atoms, 3, n_atoms, 3) where
232 - the first and third indices identify atoms in self.get_atoms()
234 - the second and fourth indices cover the corresponding
235 Cartesian movements in x, y, z
237 e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on
238 atoms[1] in the x-direction in response to a movement in the
239 z-direction of atoms[0]
240 """
241 n_atoms = int(self._hessian2d.shape[0] / 3)
242 return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy()
244 def get_hessian_2d(self) -> np.ndarray:
245 """Get the Hessian as a 2-D array
247 This format may be preferred for use with standard linear algebra
248 functions
250 Returns:
251 array with shape (n_atoms * 3, n_atoms * 3) where the elements are
252 ordered by atom and Cartesian direction::
254 >> [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...],
255 >> [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...],
256 >> [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...],
257 >> [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...],
258 >> ...]
261 e.g. the element h[2, 3] gives a harmonic force exerted on
262 atoms[1] in the x-direction in response to a movement in the
263 z-direction of atoms[0]
264 """
265 return self._hessian2d.copy()
267 def todict(self) -> Dict[str, Any]:
268 if np.allclose(self._indices, range(len(self._atoms))):
269 indices = None
270 else:
271 indices = self.get_indices()
273 return {'atoms': self.get_atoms(),
274 'hessian': self.get_hessian(),
275 'indices': indices}
277 @classmethod
278 def fromdict(cls, data: Dict[str, Any]) -> 'VibrationsData':
279 # mypy is understandably suspicious of data coming from a dict that
280 # holds mixed types, but it can see if we sanity-check with 'assert'
281 assert isinstance(data['atoms'], Atoms)
282 assert isinstance(data['hessian'], (collections.abc.Sequence,
283 np.ndarray))
284 if data['indices'] is not None:
285 assert isinstance(data['indices'], (collections.abc.Sequence,
286 np.ndarray))
287 for index in data['indices']:
288 assert isinstance(index, Integral)
290 return cls(data['atoms'], data['hessian'], indices=data['indices'])
292 @cached_property
293 def _energies_and_modes(self) -> Tuple[np.ndarray, np.ndarray]:
294 """Diagonalise the Hessian to obtain harmonic modes
296 This method is an internal implementation of get_energies_and_modes(),
297 see the docstring of that method for more information.
299 """
300 active_atoms = self._atoms[self.get_mask()]
301 n_atoms = len(active_atoms)
302 masses = active_atoms.get_masses()
304 if not np.all(masses):
305 raise ValueError('Zero mass encountered in one or more of '
306 'the vibrated atoms. Use Atoms.set_masses()'
307 ' to set all masses to non-zero values.')
308 mass_weights = np.repeat(masses**-0.5, 3)
310 omega2, vectors = np.linalg.eigh(mass_weights
311 * self.get_hessian_2d()
312 * mass_weights[:, np.newaxis])
314 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)
315 energies = unit_conversion * omega2.astype(complex)**0.5
317 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)
318 modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5
320 return (energies, modes)
322 def get_energies_and_modes(self, all_atoms: bool = False
323 ) -> Tuple[np.ndarray, np.ndarray]:
324 """Diagonalise the Hessian to obtain harmonic modes
326 Results are cached so diagonalization will only be performed once for
327 this object instance.
329 Args:
330 all_atoms:
331 If True, return modes as (3N, [N + N_frozen], 3) array where
332 the second axis corresponds to the full list of atoms in the
333 attached atoms object. Atoms that were not included in the
334 Hessian will have displacement vectors of (0, 0, 0).
336 Returns:
337 tuple (energies, modes)
339 Energies are given in units of eV. (To convert these to frequencies
340 in cm-1, divide by ase.units.invcm.)
342 Modes are given in Cartesian coordinates as a (3N, N, 3) array
343 where indices correspond to the (mode_index, atom, direction).
345 """
347 energies, modes_from_hessian = self._energies_and_modes
349 if all_atoms:
350 n_active_atoms = len(self.get_indices())
351 n_all_atoms = len(self._atoms)
352 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))
353 modes[:, self.get_mask(), :] = modes_from_hessian
354 else:
355 modes = modes_from_hessian.copy()
357 return (energies.copy(), modes)
359 def get_modes(self, all_atoms: bool = False) -> np.ndarray:
360 """Diagonalise the Hessian to obtain harmonic modes
362 Results are cached so diagonalization will only be performed once for
363 this object instance.
365 all_atoms:
366 If True, return modes as (3N, [N + N_frozen], 3) array where
367 the second axis corresponds to the full list of atoms in the
368 attached atoms object. Atoms that were not included in the
369 Hessian will have displacement vectors of (0, 0, 0).
371 Returns:
372 Modes in Cartesian coordinates as a (3N, N, 3) array where indices
373 correspond to the (mode_index, atom, direction).
375 """
376 return self.get_energies_and_modes(all_atoms=all_atoms)[1]
378 def get_energies(self) -> np.ndarray:
379 """Diagonalise the Hessian to obtain eigenvalues
381 Results are cached so diagonalization will only be performed once for
382 this object instance.
384 Returns:
385 Harmonic mode energies in units of eV
387 """
388 return self.get_energies_and_modes()[0]
390 def get_frequencies(self) -> np.ndarray:
391 """Diagonalise the Hessian to obtain frequencies in cm^-1
393 Results are cached so diagonalization will only be performed once for
394 this object instance.
396 Returns:
397 Harmonic mode frequencies in units of cm^-1
399 """
401 return self.get_energies() / units.invcm
403 def get_zero_point_energy(self) -> float:
404 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy
406 Args:
407 energies:
408 Pre-computed energy eigenvalues. Use if available to avoid
409 re-calculating these from the Hessian.
411 Returns:
412 zero-point energy in eV
413 """
414 return self._calculate_zero_point_energy(self.get_energies())
416 @staticmethod
417 def _calculate_zero_point_energy(energies: Union[Sequence[complex],
418 np.ndarray]) -> float:
419 return 0.5 * np.asarray(energies).real.sum()
421 def tabulate(self, im_tol: float = 1e-8) -> str:
422 """Get a summary of the vibrational frequencies.
424 Args:
425 im_tol:
426 Tolerance for imaginary frequency in eV. If frequency has a
427 larger imaginary component than im_tol, the imaginary component
428 is shown in the summary table.
430 Returns:
431 Summary table as formatted text
432 """
434 energies = self.get_energies()
436 return ('\n'.join(self._tabulate_from_energies(energies,
437 im_tol=im_tol))
438 + '\n')
440 @classmethod
441 def _tabulate_from_energies(cls,
442 energies: Union[Sequence[complex], np.ndarray],
443 im_tol: float = 1e-8) -> List[str]:
444 summary_lines = ['---------------------',
445 ' # meV cm^-1',
446 '---------------------']
448 for n, e in enumerate(energies):
449 if abs(e.imag) > im_tol:
450 c = 'i'
451 e = e.imag
452 else:
453 c = ''
454 e = e.real
456 summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'
457 .format(index=n, mev=(e * 1e3),
458 cm=(e / units.invcm), im=c))
459 summary_lines.append('---------------------')
460 summary_lines.append('Zero-point energy: {:.3f} eV'.format(
461 cls._calculate_zero_point_energy(energies=energies)))
463 return summary_lines
465 def iter_animated_mode(self, mode_index: int,
466 temperature: float = units.kB * 300,
467 frames: int = 30) -> Iterator[Atoms]:
468 """Obtain animated mode as a series of Atoms
470 Args:
471 mode_index: Selection of mode to animate
472 temperature: In energy units - use units.kB * T_IN_KELVIN
473 frames: number of image frames in animation
475 Yields:
476 Displaced atoms following vibrational mode
478 """
480 mode = (self.get_modes(all_atoms=True)[mode_index]
481 * sqrt(temperature / abs(self.get_energies()[mode_index])))
483 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
484 atoms = self.get_atoms()
485 atoms.positions += sin(phase) * mode
487 yield atoms
489 def show_as_force(self,
490 mode: int,
491 scale: float = 0.2,
492 show: bool = True) -> Atoms:
493 """Illustrate mode as "forces" on atoms
495 Args:
496 mode: mode index
497 scale: scale factor
498 show: if True, open the ASE GUI and show atoms
500 Returns:
501 Atoms with scaled forces corresponding to mode eigenvectors (using
502 attached SinglePointCalculator).
504 """
506 atoms = self.get_atoms()
507 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
508 atoms.calc = SinglePointCalculator(atoms, forces=mode)
509 if show:
510 atoms.edit()
512 return atoms
514 def write_jmol(self,
515 filename: str = 'vib.xyz',
516 ir_intensities: Union[Sequence[float], np.ndarray] = None
517 ) -> None:
518 """Writes file for viewing of the modes with jmol.
520 This is an extended XYZ file with eigenvectors given as extra columns
521 and metadata given in the label/comment line for each image. The format
522 is not quite human-friendly, but has the advantage that it can be
523 imported back into ASE with ase.io.read.
525 Args:
526 filename: Path for output file
527 ir_intensities: If available, IR intensities can be included in the
528 header lines. This does not affect the visualisation, but may
529 be convenient when comparing to experimental data.
530 """
532 all_images = list(self._get_jmol_images(atoms=self.get_atoms(),
533 energies=self.get_energies(),
534 modes=self.get_modes(
535 all_atoms=True),
536 ir_intensities=ir_intensities))
537 ase.io.write(filename, all_images, format='extxyz')
539 @staticmethod
540 def _get_jmol_images(atoms: Atoms,
541 energies: np.ndarray,
542 modes: np.ndarray,
543 ir_intensities:
544 Union[Sequence[float], np.ndarray] = None
545 ) -> Iterator[Atoms]:
546 """Get vibrational modes as a series of Atoms with attached data
548 For each image (Atoms object):
550 - eigenvalues are attached to image.arrays['mode']
551 - "mode#" and "frequency_cm-1" are set in image.info
552 - "IR_intensity" is set if provided in ir_intensities
553 - "masses" is removed
555 This is intended to set up the object for JMOL-compatible export using
556 ase.io.extxyz.
559 Args:
560 atoms: The base atoms object; all images have the same positions
561 energies: Complex vibrational energies in eV
562 modes: Eigenvectors array corresponding to atoms and energies. This
563 should cover the full set of atoms (i.e. modes =
564 vib.get_modes(all_atoms=True)).
565 ir_intensities: If available, IR intensities can be included in the
566 header lines. This does not affect the visualisation, but may
567 be convenient when comparing to experimental data.
568 Returns:
569 Iterator of Atoms objects
571 """
572 for i, (energy, mode) in enumerate(zip(energies, modes)):
573 # write imaginary frequencies as negative numbers
574 if energy.imag > energy.real:
575 energy = float(-energy.imag)
576 else:
577 energy = energy.real
579 image = atoms.copy()
580 image.info.update({'mode#': str(i),
581 'frequency_cm-1': energy / units.invcm,
582 })
583 image.arrays['mode'] = mode
585 # Custom masses are quite useful in vibration analysis, but will
586 # show up in the xyz file unless we remove them
587 if image.has('masses'):
588 del image.arrays['masses']
590 if ir_intensities is not None:
591 image.info['IR_intensity'] = float(ir_intensities[i])
593 yield image
595 def get_dos(self) -> RawDOSData:
596 """Total phonon DOS"""
597 energies = self.get_energies()
598 return RawDOSData(energies, np.ones_like(energies))
600 def get_pdos(self) -> DOSCollection:
601 """Phonon DOS, including atomic contributions"""
602 energies = self.get_energies()
603 masses = self._atoms[self.get_mask()].get_masses()
605 # Get weights as N_moving_atoms x N_modes array
606 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5
607 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T
609 mask = self.get_mask()
610 all_info = [{'index': i, 'symbol': a.symbol}
611 for i, a in enumerate(self._atoms) if mask[i]]
613 return DOSCollection([RawDOSData(energies, weights, info=info)
614 for weights, info in zip(all_weights, all_info)])
616 def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray]
617 ) -> VD:
618 """Get a copy of vibrations with modified masses and the same Hessian
620 Args:
621 masses:
622 New sequence of masses corresponding to the atom order in
623 self.get_atoms()
624 Returns:
625 A copy of the data with new masses for the same Hessian
626 """
628 new_atoms = self.get_atoms()
629 new_atoms.set_masses(masses)
630 return self.__class__(new_atoms, self.get_hessian(),
631 indices=self.get_indices())