Coverage for ase / vibrations / data.py: 98.33%
180 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +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 Returns:
407 zero-point energy in eV
408 """
409 return self._calculate_zero_point_energy(self.get_energies())
411 @staticmethod
412 def _calculate_zero_point_energy(energies: Union[Sequence[complex],
413 np.ndarray]) -> float:
414 return 0.5 * np.asarray(energies).real.sum()
416 def tabulate(self, im_tol: float = 1e-8) -> str:
417 """Get a summary of the vibrational frequencies.
419 Args:
420 im_tol:
421 Tolerance for imaginary frequency in eV. If frequency has a
422 larger imaginary component than im_tol, the imaginary component
423 is shown in the summary table.
425 Returns:
426 Summary table as formatted text
427 """
429 energies = self.get_energies()
431 return ('\n'.join(self._tabulate_from_energies(energies,
432 im_tol=im_tol))
433 + '\n')
435 @classmethod
436 def _tabulate_from_energies(cls,
437 energies: Union[Sequence[complex], np.ndarray],
438 im_tol: float = 1e-8) -> List[str]:
439 summary_lines = ['---------------------',
440 ' # meV cm^-1',
441 '---------------------']
443 for n, e in enumerate(energies):
444 if abs(e.imag) > im_tol:
445 c = 'i'
446 e = e.imag
447 else:
448 c = ''
449 e = e.real
451 summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'
452 .format(index=n, mev=(e * 1e3),
453 cm=(e / units.invcm), im=c))
454 summary_lines.append('---------------------')
455 summary_lines.append('Zero-point energy: {:.3f} eV'.format(
456 cls._calculate_zero_point_energy(energies=energies)))
458 return summary_lines
460 def iter_animated_mode(self, mode_index: int,
461 temperature: float = units.kB * 300,
462 frames: int = 30) -> Iterator[Atoms]:
463 """Obtain animated mode as a series of Atoms
465 Args:
466 mode_index: Selection of mode to animate
467 temperature: In energy units - use units.kB * T_IN_KELVIN
468 frames: number of image frames in animation
470 Yields:
471 Displaced atoms following vibrational mode
473 """
475 mode = (self.get_modes(all_atoms=True)[mode_index]
476 * sqrt(temperature / abs(self.get_energies()[mode_index])))
478 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
479 atoms = self.get_atoms()
480 atoms.positions += sin(phase) * mode
482 yield atoms
484 def show_as_force(self,
485 mode: int,
486 scale: float = 0.2,
487 show: bool = True) -> Atoms:
488 """Illustrate mode as "forces" on atoms
490 Args:
491 mode: mode index
492 scale: scale factor
493 show: if True, open the ASE GUI and show atoms
495 Returns:
496 Atoms with scaled forces corresponding to mode eigenvectors (using
497 attached SinglePointCalculator).
499 """
501 atoms = self.get_atoms()
502 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
503 atoms.calc = SinglePointCalculator(atoms, forces=mode)
504 if show:
505 atoms.edit()
507 return atoms
509 def write_jmol(self,
510 filename: str = 'vib.xyz',
511 ir_intensities: Union[Sequence[float], np.ndarray] = None
512 ) -> None:
513 """Writes file for viewing of the modes with jmol.
515 This is an extended XYZ file with eigenvectors given as extra columns
516 and metadata given in the label/comment line for each image. The format
517 is not quite human-friendly, but has the advantage that it can be
518 imported back into ASE with ase.io.read.
520 Args:
521 filename: Path for output file
522 ir_intensities: If available, IR intensities can be included in the
523 header lines. This does not affect the visualisation, but may
524 be convenient when comparing to experimental data.
525 """
527 all_images = list(self._get_jmol_images(atoms=self.get_atoms(),
528 energies=self.get_energies(),
529 modes=self.get_modes(
530 all_atoms=True),
531 ir_intensities=ir_intensities))
532 ase.io.write(filename, all_images, format='extxyz')
534 @staticmethod
535 def _get_jmol_images(atoms: Atoms,
536 energies: np.ndarray,
537 modes: np.ndarray,
538 ir_intensities:
539 Union[Sequence[float], np.ndarray] = None
540 ) -> Iterator[Atoms]:
541 """Get vibrational modes as a series of Atoms with attached data
543 For each image (Atoms object):
545 - eigenvalues are attached to image.arrays['mode']
546 - "mode#" and "frequency_cm-1" are set in image.info
547 - "IR_intensity" is set if provided in ir_intensities
548 - "masses" is removed
550 This is intended to set up the object for JMOL-compatible export using
551 ase.io.extxyz.
554 Args:
555 atoms: The base atoms object; all images have the same positions
556 energies: Complex vibrational energies in eV
557 modes: Eigenvectors array corresponding to atoms and energies. This
558 should cover the full set of atoms (i.e. modes =
559 vib.get_modes(all_atoms=True)).
560 ir_intensities: If available, IR intensities can be included in the
561 header lines. This does not affect the visualisation, but may
562 be convenient when comparing to experimental data.
563 Returns:
564 Iterator of Atoms objects
566 """
567 for i, (energy, mode) in enumerate(zip(energies, modes)):
568 # write imaginary frequencies as negative numbers
569 if energy.imag > energy.real:
570 energy = float(-energy.imag)
571 else:
572 energy = energy.real
574 image = atoms.copy()
575 image.info.update({'mode#': str(i),
576 'frequency_cm-1': energy / units.invcm,
577 })
578 image.arrays['mode'] = mode
580 # Custom masses are quite useful in vibration analysis, but will
581 # show up in the xyz file unless we remove them
582 if image.has('masses'):
583 del image.arrays['masses']
585 if ir_intensities is not None:
586 image.info['IR_intensity'] = float(ir_intensities[i])
588 yield image
590 def get_dos(self) -> RawDOSData:
591 """Total phonon DOS"""
592 energies = self.get_energies()
593 return RawDOSData(energies, np.ones_like(energies))
595 def get_pdos(self) -> DOSCollection:
596 """Phonon DOS, including atomic contributions"""
597 energies = self.get_energies()
598 masses = self._atoms[self.get_mask()].get_masses()
600 # Get weights as N_moving_atoms x N_modes array
601 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5
602 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T
604 mask = self.get_mask()
605 all_info = [{'index': i, 'symbol': a.symbol}
606 for i, a in enumerate(self._atoms) if mask[i]]
608 return DOSCollection([RawDOSData(energies, weights, info=info)
609 for weights, info in zip(all_weights, all_info)])
611 def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray]
612 ) -> VD:
613 """Get a copy of vibrations with modified masses and the same Hessian
615 Args:
616 masses:
617 New sequence of masses corresponding to the atom order in
618 self.get_atoms()
619 Returns:
620 A copy of the data with new masses for the same Hessian
621 """
623 new_atoms = self.get_atoms()
624 new_atoms.set_masses(masses)
625 return self.__class__(new_atoms, self.get_hessian(),
626 indices=self.get_indices())