Coverage for /builds/debichem-team/python-ase/ase/vibrations/data.py: 98.33%
180 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +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, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union
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. Once the
30 VibrationsData has been constructed, this class provides some common
31 processing options; frequency calculation, mode animation, DOS etc.
33 If the Atoms object is a periodic supercell, 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 FixedAtoms or FixedCartesian constraints, these
39 will be respected and the Hessian should be sized accordingly.
41 Args:
42 atoms:
43 Equilibrium geometry of vibrating system. This will be stored as a
44 full copy.
46 hessian: Second-derivative in energy with respect to
47 Cartesian nuclear movements as an (N, 3, N, 3) array.
49 indices: indices of atoms which are included
50 in Hessian. Default value (None) includes all freely
51 moving atoms (i.e. not fixed ones). Leave at None if
52 constraints should be determined automatically from the
53 atoms object.
55 """
57 def __init__(self,
58 atoms: Atoms,
59 hessian: Union[RealSequence4D, np.ndarray],
60 indices: Union[Sequence[int], np.ndarray] = None,
61 ) -> None:
63 if indices is None:
64 indices = np.asarray(self.indices_from_constraints(atoms),
65 dtype=int)
67 self._indices = np.array(indices, dtype=int)
69 n_atoms = self._check_dimensions(atoms, np.asarray(hessian),
70 indices=self._indices)
71 self._atoms = atoms.copy()
73 self._hessian2d = (np.asarray(hessian)
74 .reshape(3 * n_atoms, 3 * n_atoms).copy())
76 _setter_error = ("VibrationsData properties cannot be modified: construct "
77 "a new VibrationsData with consistent atoms, Hessian and "
78 "(optionally) indices/mask.")
80 @classmethod
81 def from_2d(cls, atoms: Atoms,
82 hessian_2d: Union[Sequence[Sequence[Real]], np.ndarray],
83 indices: Sequence[int] = None) -> 'VibrationsData':
84 """Instantiate VibrationsData when the Hessian is in a 3Nx3N format
86 Args:
87 atoms: Equilibrium geometry of vibrating system
89 hessian: Second-derivative in energy with respect to
90 Cartesian nuclear movements as a (3N, 3N) array.
92 indices: Indices of (non-frozen) atoms included in Hessian
94 """
95 if indices is None:
96 indices = range(len(atoms))
97 assert indices is not None # Show Mypy that indices is now a sequence
99 hessian_2d_array = np.asarray(hessian_2d)
100 n_atoms = cls._check_dimensions(atoms, hessian_2d_array,
101 indices=indices, two_d=True)
103 return cls(atoms, hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3),
104 indices=indices)
106 @staticmethod
107 def indices_from_constraints(atoms: Atoms) -> List[int]:
108 """Indices corresponding to Atoms Constraints
110 Deduces the freely moving atoms from the constraints set on the
111 atoms object. VibrationsData only supports FixCartesian and
112 FixAtoms. All others are neglected.
114 Args:
115 atoms: Atoms object.
117 Retruns:
118 indices of free atoms.
120 """
121 # Only fully fixed atoms supported by VibrationsData
122 const_indices = constrained_indices(
123 atoms, only_include=(FixCartesian, FixAtoms))
124 # Invert the selection to get free atoms
125 indices = np.setdiff1d(
126 np.array(
127 range(
128 len(atoms))),
129 const_indices).astype(int)
130 return indices.tolist()
132 @staticmethod
133 def indices_from_mask(mask: Union[Sequence[bool], np.ndarray]
134 ) -> List[int]:
135 """Indices corresponding to boolean mask
137 This is provided as a convenience for instantiating VibrationsData with
138 a boolean mask. For example, if the Hessian data includes only the H
139 atoms in a structure::
141 h_mask = atoms.get_chemical_symbols() == 'H'
142 vib_data = VibrationsData(atoms, hessian,
143 VibrationsData.indices_from_mask(h_mask))
145 Take care to ensure that the length of the mask corresponds to the full
146 number of atoms; this function is only aware of the mask it has been
147 given.
149 Args:
150 mask: a sequence of True, False values
152 Returns:
153 indices of True elements
155 """
156 return np.where(mask)[0].tolist()
158 @staticmethod
159 def _check_dimensions(atoms: Atoms,
160 hessian: np.ndarray,
161 indices: Union[np.ndarray, Sequence[int]],
162 two_d: bool = False) -> int:
163 """Sanity check on array shapes from input data
165 Args:
166 atoms: Structure
167 indices: Indices of atoms used in Hessian
168 hessian: Proposed Hessian array
170 Returns:
171 Number of atoms contributing to Hessian
173 Raises:
174 ValueError if Hessian dimensions are not (N, 3, N, 3)
176 """
178 n_atoms = len(atoms[indices])
180 if two_d:
181 ref_shape = [n_atoms * 3, n_atoms * 3]
182 ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3))
184 else:
185 ref_shape = [n_atoms, 3, n_atoms, 3]
186 ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms)
188 if (isinstance(hessian, np.ndarray)
189 and hessian.shape == tuple(ref_shape)):
190 return n_atoms
191 else:
192 raise ValueError("Hessian for these atoms should be a "
193 "{} numpy array.".format(ref_shape_txt))
195 def get_atoms(self) -> Atoms:
196 return self._atoms.copy()
198 def get_indices(self) -> np.ndarray:
199 return self._indices.copy()
201 def get_mask(self) -> np.ndarray:
202 """Boolean mask of atoms selected by indices"""
203 return self._mask_from_indices(self._atoms, self.get_indices())
205 @staticmethod
206 def _mask_from_indices(atoms: Atoms,
207 indices: Union[None, Sequence[int], np.ndarray]
208 ) -> np.ndarray:
209 """Boolean mask of atoms selected by indices"""
210 natoms = len(atoms)
212 # Wrap indices to allow negative values
213 indices = np.asarray(indices) % natoms
215 mask = np.full(natoms, False, dtype=bool)
216 mask[indices] = True
217 return mask
219 def get_hessian(self) -> np.ndarray:
220 """The Hessian; second derivative of energy wrt positions
222 This format is preferred for iteration over atoms and when
223 addressing specific elements of the Hessian.
225 Returns:
226 array with shape (n_atoms, 3, n_atoms, 3) where
228 - the first and third indices identify atoms in self.get_atoms()
230 - the second and fourth indices cover the corresponding
231 Cartesian movements in x, y, z
233 e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on
234 atoms[1] in the x-direction in response to a movement in the
235 z-direction of atoms[0]
236 """
237 n_atoms = int(self._hessian2d.shape[0] / 3)
238 return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy()
240 def get_hessian_2d(self) -> np.ndarray:
241 """Get the Hessian as a 2-D array
243 This format may be preferred for use with standard linear algebra
244 functions
246 Returns:
247 array with shape (n_atoms * 3, n_atoms * 3) where the elements are
248 ordered by atom and Cartesian direction::
250 >> [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...],
251 >> [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...],
252 >> [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...],
253 >> [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...],
254 >> ...]
257 e.g. the element h[2, 3] gives a harmonic force exerted on
258 atoms[1] in the x-direction in response to a movement in the
259 z-direction of atoms[0]
260 """
261 return self._hessian2d.copy()
263 def todict(self) -> Dict[str, Any]:
264 if np.allclose(self._indices, range(len(self._atoms))):
265 indices = None
266 else:
267 indices = self.get_indices()
269 return {'atoms': self.get_atoms(),
270 'hessian': self.get_hessian(),
271 'indices': indices}
273 @classmethod
274 def fromdict(cls, data: Dict[str, Any]) -> 'VibrationsData':
275 # mypy is understandably suspicious of data coming from a dict that
276 # holds mixed types, but it can see if we sanity-check with 'assert'
277 assert isinstance(data['atoms'], Atoms)
278 assert isinstance(data['hessian'], (collections.abc.Sequence,
279 np.ndarray))
280 if data['indices'] is not None:
281 assert isinstance(data['indices'], (collections.abc.Sequence,
282 np.ndarray))
283 for index in data['indices']:
284 assert isinstance(index, Integral)
286 return cls(data['atoms'], data['hessian'], indices=data['indices'])
288 @cached_property
289 def _energies_and_modes(self) -> Tuple[np.ndarray, np.ndarray]:
290 """Diagonalise the Hessian to obtain harmonic modes
292 This method is an internal implementation of get_energies_and_modes(),
293 see the docstring of that method for more information.
295 """
296 active_atoms = self._atoms[self.get_mask()]
297 n_atoms = len(active_atoms)
298 masses = active_atoms.get_masses()
300 if not np.all(masses):
301 raise ValueError('Zero mass encountered in one or more of '
302 'the vibrated atoms. Use Atoms.set_masses()'
303 ' to set all masses to non-zero values.')
304 mass_weights = np.repeat(masses**-0.5, 3)
306 omega2, vectors = np.linalg.eigh(mass_weights
307 * self.get_hessian_2d()
308 * mass_weights[:, np.newaxis])
310 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)
311 energies = unit_conversion * omega2.astype(complex)**0.5
313 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)
314 modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5
316 return (energies, modes)
318 def get_energies_and_modes(self, all_atoms: bool = False
319 ) -> Tuple[np.ndarray, np.ndarray]:
320 """Diagonalise the Hessian to obtain harmonic modes
322 Results are cached so diagonalization will only be performed once for
323 this object instance.
325 Args:
326 all_atoms:
327 If True, return modes as (3N, [N + N_frozen], 3) array where
328 the second axis corresponds to the full list of atoms in the
329 attached atoms object. Atoms that were not included in the
330 Hessian will have displacement vectors of (0, 0, 0).
332 Returns:
333 tuple (energies, modes)
335 Energies are given in units of eV. (To convert these to frequencies
336 in cm-1, divide by ase.units.invcm.)
338 Modes are given in Cartesian coordinates as a (3N, N, 3) array
339 where indices correspond to the (mode_index, atom, direction).
341 """
343 energies, modes_from_hessian = self._energies_and_modes
345 if all_atoms:
346 n_active_atoms = len(self.get_indices())
347 n_all_atoms = len(self._atoms)
348 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))
349 modes[:, self.get_mask(), :] = modes_from_hessian
350 else:
351 modes = modes_from_hessian.copy()
353 return (energies.copy(), modes)
355 def get_modes(self, all_atoms: bool = False) -> np.ndarray:
356 """Diagonalise the Hessian to obtain harmonic modes
358 Results are cached so diagonalization will only be performed once for
359 this object instance.
361 all_atoms:
362 If True, return modes as (3N, [N + N_frozen], 3) array where
363 the second axis corresponds to the full list of atoms in the
364 attached atoms object. Atoms that were not included in the
365 Hessian will have displacement vectors of (0, 0, 0).
367 Returns:
368 Modes in Cartesian coordinates as a (3N, N, 3) array where indices
369 correspond to the (mode_index, atom, direction).
371 """
372 return self.get_energies_and_modes(all_atoms=all_atoms)[1]
374 def get_energies(self) -> np.ndarray:
375 """Diagonalise the Hessian to obtain eigenvalues
377 Results are cached so diagonalization will only be performed once for
378 this object instance.
380 Returns:
381 Harmonic mode energies in units of eV
383 """
384 return self.get_energies_and_modes()[0]
386 def get_frequencies(self) -> np.ndarray:
387 """Diagonalise the Hessian to obtain frequencies in cm^-1
389 Results are cached so diagonalization will only be performed once for
390 this object instance.
392 Returns:
393 Harmonic mode frequencies in units of cm^-1
395 """
397 return self.get_energies() / units.invcm
399 def get_zero_point_energy(self) -> float:
400 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy
402 Args:
403 energies:
404 Pre-computed energy eigenvalues. Use if available to avoid
405 re-calculating these from the Hessian.
407 Returns:
408 zero-point energy in eV
409 """
410 return self._calculate_zero_point_energy(self.get_energies())
412 @staticmethod
413 def _calculate_zero_point_energy(energies: Union[Sequence[complex],
414 np.ndarray]) -> float:
415 return 0.5 * np.asarray(energies).real.sum()
417 def tabulate(self, im_tol: float = 1e-8) -> str:
418 """Get a summary of the vibrational frequencies.
420 Args:
421 im_tol:
422 Tolerance for imaginary frequency in eV. If frequency has a
423 larger imaginary component than im_tol, the imaginary component
424 is shown in the summary table.
426 Returns:
427 Summary table as formatted text
428 """
430 energies = self.get_energies()
432 return ('\n'.join(self._tabulate_from_energies(energies,
433 im_tol=im_tol))
434 + '\n')
436 @classmethod
437 def _tabulate_from_energies(cls,
438 energies: Union[Sequence[complex], np.ndarray],
439 im_tol: float = 1e-8) -> List[str]:
440 summary_lines = ['---------------------',
441 ' # meV cm^-1',
442 '---------------------']
444 for n, e in enumerate(energies):
445 if abs(e.imag) > im_tol:
446 c = 'i'
447 e = e.imag
448 else:
449 c = ''
450 e = e.real
452 summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'
453 .format(index=n, mev=(e * 1e3),
454 cm=(e / units.invcm), im=c))
455 summary_lines.append('---------------------')
456 summary_lines.append('Zero-point energy: {:.3f} eV'.format(
457 cls._calculate_zero_point_energy(energies=energies)))
459 return summary_lines
461 def iter_animated_mode(self, mode_index: int,
462 temperature: float = units.kB * 300,
463 frames: int = 30) -> Iterator[Atoms]:
464 """Obtain animated mode as a series of Atoms
466 Args:
467 mode_index: Selection of mode to animate
468 temperature: In energy units - use units.kB * T_IN_KELVIN
469 frames: number of image frames in animation
471 Yields:
472 Displaced atoms following vibrational mode
474 """
476 mode = (self.get_modes(all_atoms=True)[mode_index]
477 * sqrt(temperature / abs(self.get_energies()[mode_index])))
479 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
480 atoms = self.get_atoms()
481 atoms.positions += sin(phase) * mode
483 yield atoms
485 def show_as_force(self,
486 mode: int,
487 scale: float = 0.2,
488 show: bool = True) -> Atoms:
489 """Illustrate mode as "forces" on atoms
491 Args:
492 mode: mode index
493 scale: scale factor
494 show: if True, open the ASE GUI and show atoms
496 Returns:
497 Atoms with scaled forces corresponding to mode eigenvectors (using
498 attached SinglePointCalculator).
500 """
502 atoms = self.get_atoms()
503 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
504 atoms.calc = SinglePointCalculator(atoms, forces=mode)
505 if show:
506 atoms.edit()
508 return atoms
510 def write_jmol(self,
511 filename: str = 'vib.xyz',
512 ir_intensities: Union[Sequence[float], np.ndarray] = None
513 ) -> None:
514 """Writes file for viewing of the modes with jmol.
516 This is an extended XYZ file with eigenvectors given as extra columns
517 and metadata given in the label/comment line for each image. The format
518 is not quite human-friendly, but has the advantage that it can be
519 imported back into ASE with ase.io.read.
521 Args:
522 filename: Path for output file
523 ir_intensities: If available, IR intensities can be included in the
524 header lines. This does not affect the visualisation, but may
525 be convenient when comparing to experimental data.
526 """
528 all_images = list(self._get_jmol_images(atoms=self.get_atoms(),
529 energies=self.get_energies(),
530 modes=self.get_modes(
531 all_atoms=True),
532 ir_intensities=ir_intensities))
533 ase.io.write(filename, all_images, format='extxyz')
535 @staticmethod
536 def _get_jmol_images(atoms: Atoms,
537 energies: np.ndarray,
538 modes: np.ndarray,
539 ir_intensities:
540 Union[Sequence[float], np.ndarray] = None
541 ) -> Iterator[Atoms]:
542 """Get vibrational modes as a series of Atoms with attached data
544 For each image (Atoms object):
546 - eigenvalues are attached to image.arrays['mode']
547 - "mode#" and "frequency_cm-1" are set in image.info
548 - "IR_intensity" is set if provided in ir_intensities
549 - "masses" is removed
551 This is intended to set up the object for JMOL-compatible export using
552 ase.io.extxyz.
555 Args:
556 atoms: The base atoms object; all images have the same positions
557 energies: Complex vibrational energies in eV
558 modes: Eigenvectors array corresponding to atoms and energies. This
559 should cover the full set of atoms (i.e. modes =
560 vib.get_modes(all_atoms=True)).
561 ir_intensities: If available, IR intensities can be included in the
562 header lines. This does not affect the visualisation, but may
563 be convenient when comparing to experimental data.
564 Returns:
565 Iterator of Atoms objects
567 """
568 for i, (energy, mode) in enumerate(zip(energies, modes)):
569 # write imaginary frequencies as negative numbers
570 if energy.imag > energy.real:
571 energy = float(-energy.imag)
572 else:
573 energy = energy.real
575 image = atoms.copy()
576 image.info.update({'mode#': str(i),
577 'frequency_cm-1': energy / units.invcm,
578 })
579 image.arrays['mode'] = mode
581 # Custom masses are quite useful in vibration analysis, but will
582 # show up in the xyz file unless we remove them
583 if image.has('masses'):
584 del image.arrays['masses']
586 if ir_intensities is not None:
587 image.info['IR_intensity'] = float(ir_intensities[i])
589 yield image
591 def get_dos(self) -> RawDOSData:
592 """Total phonon DOS"""
593 energies = self.get_energies()
594 return RawDOSData(energies, np.ones_like(energies))
596 def get_pdos(self) -> DOSCollection:
597 """Phonon DOS, including atomic contributions"""
598 energies = self.get_energies()
599 masses = self._atoms[self.get_mask()].get_masses()
601 # Get weights as N_moving_atoms x N_modes array
602 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5
603 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T
605 mask = self.get_mask()
606 all_info = [{'index': i, 'symbol': a.symbol}
607 for i, a in enumerate(self._atoms) if mask[i]]
609 return DOSCollection([RawDOSData(energies, weights, info=info)
610 for weights, info in zip(all_weights, all_info)])
612 def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray]
613 ) -> VD:
614 """Get a copy of vibrations with modified masses and the same Hessian
616 Args:
617 masses:
618 New sequence of masses corresponding to the atom order in
619 self.get_atoms()
620 Returns:
621 A copy of the data with new masses for the same Hessian
622 """
624 new_atoms = self.get_atoms()
625 new_atoms.set_masses(masses)
626 return self.__class__(new_atoms, self.get_hessian(),
627 indices=self.get_indices())