Coverage for /builds/debichem-team/python-ase/ase/vibrations/data.py : 98.27%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Storage and analysis for vibrational data"""
3import collections
4from math import sin, pi, sqrt
5from numbers import Real, Integral
6from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union
8import numpy as np
10from ase.atoms import Atoms
11import ase.units as units
12import ase.io
13from ase.utils import jsonable, lazymethod
15from ase.calculators.singlepoint import SinglePointCalculator
16from ase.spectrum.dosdata import RawDOSData
17from ase.spectrum.doscollection import DOSCollection
19RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]]
20VD = TypeVar('VD', bound='VibrationsData')
23@jsonable('vibrationsdata')
24class VibrationsData:
25 """Class for storing and analyzing vibrational data (i.e. Atoms + Hessian)
27 This class is not responsible for calculating Hessians; the Hessian should
28 be computed by a Calculator or some other algorithm. Once the
29 VibrationsData has been constructed, this class provides some common
30 processing options; frequency calculation, mode animation, DOS etc.
32 If the Atoms object is a periodic supercell, VibrationsData may be
33 converted to a PhononData using the VibrationsData.to_phonondata() method.
34 This provides access to q-point-dependent analyses such as phonon
35 dispersion plotting.
37 Args:
38 atoms:
39 Equilibrium geometry of vibrating system. This will be stored as a
40 lightweight copy with just positions, masses, unit cell.
42 hessian: Second-derivative in energy with respect to
43 Cartesian nuclear movements as an (N, 3, N, 3) array.
44 indices: indices of atoms which are included
45 in Hessian. Default value (None) includes all atoms.
47 """
48 def __init__(self,
49 atoms: Atoms,
50 hessian: Union[RealSequence4D, np.ndarray],
51 indices: Union[Sequence[int], np.ndarray] = None,
52 ) -> None:
54 if indices is None:
55 self._indices = np.arange(len(atoms), dtype=int)
56 else:
57 self._indices = np.array(indices, dtype=int)
59 n_atoms = self._check_dimensions(atoms, np.asarray(hessian),
60 indices=self._indices)
61 self._atoms = atoms.copy()
63 self._hessian2d = (np.asarray(hessian)
64 .reshape(3 * n_atoms, 3 * n_atoms).copy())
66 _setter_error = ("VibrationsData properties cannot be modified: construct "
67 "a new VibrationsData with consistent atoms, Hessian and "
68 "(optionally) indices/mask.")
70 @classmethod
71 def from_2d(cls, atoms: Atoms,
72 hessian_2d: Union[Sequence[Sequence[Real]], np.ndarray],
73 indices: Sequence[int] = None) -> 'VibrationsData':
74 """Instantiate VibrationsData when the Hessian is in a 3Nx3N format
76 Args:
77 atoms: Equilibrium geometry of vibrating system
79 hessian: Second-derivative in energy with respect to
80 Cartesian nuclear movements as a (3N, 3N) array.
82 indices: Indices of (non-frozen) atoms included in Hessian
84 """
85 if indices is None:
86 indices = range(len(atoms))
87 assert indices is not None # Show Mypy that indices is now a sequence
89 hessian_2d_array = np.asarray(hessian_2d)
90 n_atoms = cls._check_dimensions(atoms, hessian_2d_array,
91 indices=indices, two_d=True)
93 return cls(atoms, hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3),
94 indices=indices)
96 @staticmethod
97 def indices_from_mask(mask: Union[Sequence[bool], np.ndarray]
98 ) -> List[int]:
99 """Indices corresponding to boolean mask
101 This is provided as a convenience for instantiating VibrationsData with
102 a boolean mask. For example, if the Hessian data includes only the H
103 atoms in a structure::
105 h_mask = atoms.get_chemical_symbols() == 'H'
106 vib_data = VibrationsData(atoms, hessian,
107 VibrationsData.indices_from_mask(h_mask))
109 Take care to ensure that the length of the mask corresponds to the full
110 number of atoms; this function is only aware of the mask it has been
111 given.
113 Args:
114 mask: a sequence of True, False values
116 Returns:
117 indices of True elements
119 """
120 return np.where(mask)[0].tolist()
122 @staticmethod
123 def _check_dimensions(atoms: Atoms,
124 hessian: np.ndarray,
125 indices: Sequence[int],
126 two_d: bool = False) -> int:
127 """Sanity check on array shapes from input data
129 Args:
130 atoms: Structure
131 indices: Indices of atoms used in Hessian
132 hessian: Proposed Hessian array
134 Returns:
135 Number of atoms contributing to Hessian
137 Raises:
138 ValueError if Hessian dimensions are not (N, 3, N, 3)
140 """
142 n_atoms = len(atoms[indices])
144 if two_d:
145 ref_shape = [n_atoms * 3, n_atoms * 3]
146 ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3))
148 else:
149 ref_shape = [n_atoms, 3, n_atoms, 3]
150 ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms)
152 if (isinstance(hessian, np.ndarray)
153 and hessian.shape == tuple(ref_shape)):
154 return n_atoms
155 else:
156 raise ValueError("Hessian for these atoms should be a "
157 "{} numpy array.".format(ref_shape_txt))
159 def get_atoms(self) -> Atoms:
160 return self._atoms.copy()
162 def get_indices(self) -> np.ndarray:
163 return self._indices.copy()
165 def get_mask(self) -> np.ndarray:
166 """Boolean mask of atoms selected by indices"""
167 return self._mask_from_indices(self._atoms, self.get_indices())
169 @staticmethod
170 def _mask_from_indices(atoms: Atoms,
171 indices: Union[None, Sequence[int], np.ndarray]
172 ) -> np.ndarray:
173 """Boolean mask of atoms selected by indices"""
174 natoms = len(atoms)
176 # Wrap indices to allow negative values
177 indices = np.asarray(indices) % natoms
179 mask = np.full(natoms, False, dtype=bool)
180 mask[indices] = True
181 return mask
183 def get_hessian(self) -> np.ndarray:
184 """The Hessian; second derivative of energy wrt positions
186 This format is preferred for iteration over atoms and when
187 addressing specific elements of the Hessian.
189 Returns:
190 array with shape (n_atoms, 3, n_atoms, 3) where
191 - the first and third indices identify atoms in self.get_atoms()
192 - the second and fourth indices cover the corresponding Cartesian
193 movements in x, y, z
195 e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on
196 atoms[1] in the x-direction in response to a movement in the
197 z-direction of atoms[0]
199 """
200 n_atoms = int(self._hessian2d.shape[0] / 3)
201 return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy()
203 def get_hessian_2d(self) -> np.ndarray:
204 """Get the Hessian as a 2-D array
206 This format may be preferred for use with standard linear algebra
207 functions
209 Returns:
210 array with shape (n_atoms * 3, n_atoms * 3) where the elements are
211 ordered by atom and Cartesian direction
213 [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...],
214 [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...],
215 [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...],
216 [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...],
217 ...]
219 e.g. the element h[2, 3] gives a harmonic force exerted on
220 atoms[1] in the x-direction in response to a movement in the
221 z-direction of atoms[0]
223 """
224 return self._hessian2d.copy()
226 def todict(self) -> Dict[str, Any]:
227 if np.allclose(self._indices, range(len(self._atoms))):
228 indices = None
229 else:
230 indices = self.get_indices()
232 return {'atoms': self.get_atoms(),
233 'hessian': self.get_hessian(),
234 'indices': indices}
236 @classmethod
237 def fromdict(cls, data: Dict[str, Any]) -> 'VibrationsData':
238 # mypy is understandably suspicious of data coming from a dict that
239 # holds mixed types, but it can see if we sanity-check with 'assert'
240 assert isinstance(data['atoms'], Atoms)
241 assert isinstance(data['hessian'], (collections.abc.Sequence,
242 np.ndarray))
243 if data['indices'] is not None:
244 assert isinstance(data['indices'], (collections.abc.Sequence,
245 np.ndarray))
246 for index in data['indices']:
247 assert isinstance(index, Integral)
249 return cls(data['atoms'], data['hessian'], indices=data['indices'])
251 @lazymethod
252 def _energies_and_modes(self) -> Tuple[np.ndarray, np.ndarray]:
253 """Diagonalise the Hessian to obtain harmonic modes
255 This method is an internal implementation of get_energies_and_modes(),
256 see the docstring of that method for more information.
258 """
259 active_atoms = self._atoms[self.get_mask()]
260 n_atoms = len(active_atoms)
261 masses = active_atoms.get_masses()
263 if not np.all(masses):
264 raise ValueError('Zero mass encountered in one or more of '
265 'the vibrated atoms. Use Atoms.set_masses()'
266 ' to set all masses to non-zero values.')
267 mass_weights = np.repeat(masses**-0.5, 3)
269 omega2, vectors = np.linalg.eigh(mass_weights
270 * self.get_hessian_2d()
271 * mass_weights[:, np.newaxis])
273 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)
274 energies = unit_conversion * omega2.astype(complex)**0.5
276 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)
277 modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5
279 return (energies, modes)
281 def get_energies_and_modes(self, all_atoms: bool = False
282 ) -> Tuple[np.ndarray, np.ndarray]:
283 """Diagonalise the Hessian to obtain harmonic modes
285 Results are cached so diagonalization will only be performed once for
286 this object instance.
288 Args:
289 all_atoms:
290 If True, return modes as (3N, [N + N_frozen], 3) array where
291 the second axis corresponds to the full list of atoms in the
292 attached atoms object. Atoms that were not included in the
293 Hessian will have displacement vectors of (0, 0, 0).
295 Returns:
296 tuple (energies, modes)
298 Energies are given in units of eV. (To convert these to frequencies
299 in cm-1, divide by ase.units.invcm.)
301 Modes are given in Cartesian coordinates as a (3N, N, 3) array
302 where indices correspond to the (mode_index, atom, direction).
304 """
306 energies, modes_from_hessian = self._energies_and_modes()
308 if all_atoms:
309 n_active_atoms = len(self.get_indices())
310 n_all_atoms = len(self._atoms)
311 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))
312 modes[:, self.get_mask(), :] = modes_from_hessian
313 else:
314 modes = modes_from_hessian.copy()
316 return (energies.copy(), modes)
318 def get_modes(self, all_atoms: bool = False) -> np.ndarray:
319 """Diagonalise the Hessian to obtain harmonic modes
321 Results are cached so diagonalization will only be performed once for
322 this object instance.
324 all_atoms:
325 If True, return modes as (3N, [N + N_frozen], 3) array where
326 the second axis corresponds to the full list of atoms in the
327 attached atoms object. Atoms that were not included in the
328 Hessian will have displacement vectors of (0, 0, 0).
330 Returns:
331 Modes in Cartesian coordinates as a (3N, N, 3) array where indices
332 correspond to the (mode_index, atom, direction).
334 """
335 return self.get_energies_and_modes(all_atoms=all_atoms)[1]
337 def get_energies(self) -> np.ndarray:
338 """Diagonalise the Hessian to obtain eigenvalues
340 Results are cached so diagonalization will only be performed once for
341 this object instance.
343 Returns:
344 Harmonic mode energies in units of eV
346 """
347 return self.get_energies_and_modes()[0]
349 def get_frequencies(self) -> np.ndarray:
350 """Diagonalise the Hessian to obtain frequencies in cm^-1
352 Results are cached so diagonalization will only be performed once for
353 this object instance.
355 Returns:
356 Harmonic mode frequencies in units of cm^-1
358 """
360 return self.get_energies() / units.invcm
362 def get_zero_point_energy(self) -> float:
363 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy
365 Args:
366 energies:
367 Pre-computed energy eigenvalues. Use if available to avoid
368 re-calculating these from the Hessian.
370 Returns:
371 zero-point energy in eV
372 """
373 return self._calculate_zero_point_energy(self.get_energies())
375 @staticmethod
376 def _calculate_zero_point_energy(energies: Union[Sequence[complex],
377 np.ndarray]) -> float:
378 return 0.5 * np.asarray(energies).real.sum()
380 def tabulate(self, im_tol: float = 1e-8) -> str:
381 """Get a summary of the vibrational frequencies.
383 Args:
384 im_tol:
385 Tolerance for imaginary frequency in eV. If frequency has a
386 larger imaginary component than im_tol, the imaginary component
387 is shown in the summary table.
389 Returns:
390 Summary table as formatted text
391 """
393 energies = self.get_energies()
395 return ('\n'.join(self._tabulate_from_energies(energies,
396 im_tol=im_tol))
397 + '\n')
399 @classmethod
400 def _tabulate_from_energies(cls,
401 energies: Union[Sequence[complex], np.ndarray],
402 im_tol: float = 1e-8) -> List[str]:
403 summary_lines = ['---------------------',
404 ' # meV cm^-1',
405 '---------------------']
407 for n, e in enumerate(energies):
408 if abs(e.imag) > im_tol:
409 c = 'i'
410 e = e.imag
411 else:
412 c = ''
413 e = e.real
415 summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'
416 .format(index=n, mev=(e * 1e3),
417 cm=(e / units.invcm), im=c))
418 summary_lines.append('---------------------')
419 summary_lines.append('Zero-point energy: {:.3f} eV'.format(
420 cls._calculate_zero_point_energy(energies=energies)))
422 return summary_lines
424 def iter_animated_mode(self, mode_index: int,
425 temperature: float = units.kB * 300,
426 frames: int = 30) -> Iterator[Atoms]:
427 """Obtain animated mode as a series of Atoms
429 Args:
430 mode_index: Selection of mode to animate
431 temperature: In energy units - use units.kB * T_IN_KELVIN
432 frames: number of image frames in animation
434 Yields:
435 Displaced atoms following vibrational mode
437 """
439 mode = (self.get_modes(all_atoms=True)[mode_index]
440 * sqrt(temperature / abs(self.get_energies()[mode_index])))
442 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):
443 atoms = self.get_atoms()
444 atoms.positions += sin(phase) * mode
446 yield atoms
448 def show_as_force(self,
449 mode: int,
450 scale: float = 0.2,
451 show: bool = True) -> Atoms:
452 """Illustrate mode as "forces" on atoms
454 Args:
455 mode: mode index
456 scale: scale factor
457 show: if True, open the ASE GUI and show atoms
459 Returns:
460 Atoms with scaled forces corresponding to mode eigenvectors (using
461 attached SinglePointCalculator).
463 """
465 atoms = self.get_atoms()
466 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale
467 atoms.calc = SinglePointCalculator(atoms, forces=mode)
468 if show:
469 atoms.edit()
471 return atoms
473 def write_jmol(self,
474 filename: str = 'vib.xyz',
475 ir_intensities: Union[Sequence[float], np.ndarray] = None
476 ) -> None:
477 """Writes file for viewing of the modes with jmol.
479 This is an extended XYZ file with eigenvectors given as extra columns
480 and metadata given in the label/comment line for each image. The format
481 is not quite human-friendly, but has the advantage that it can be
482 imported back into ASE with ase.io.read.
484 Args:
485 filename: Path for output file
486 ir_intensities: If available, IR intensities can be included in the
487 header lines. This does not affect the visualisation, but may
488 be convenient when comparing to experimental data.
489 """
491 all_images = list(self._get_jmol_images(atoms=self.get_atoms(),
492 energies=self.get_energies(),
493 modes=self.get_modes(all_atoms=True),
494 ir_intensities=ir_intensities))
495 ase.io.write(filename, all_images, format='extxyz')
497 @staticmethod
498 def _get_jmol_images(atoms: Atoms,
499 energies: np.ndarray,
500 modes: np.ndarray,
501 ir_intensities:
502 Union[Sequence[float], np.ndarray] = None
503 ) -> Iterator[Atoms]:
504 """Get vibrational modes as a series of Atoms with attached data
506 For each image (Atoms object):
508 - eigenvalues are attached to image.arrays['mode']
509 - "mode#" and "frequency_cm-1" are set in image.info
510 - "IR_intensity" is set if provided in ir_intensities
511 - "masses" is removed
513 This is intended to set up the object for JMOL-compatible export using
514 ase.io.extxyz.
517 Args:
518 atoms: The base atoms object; all images have the same positions
519 energies: Complex vibrational energies in eV
520 modes: Eigenvectors array corresponding to atoms and energies. This
521 should cover the full set of atoms (i.e. modes =
522 vib.get_modes(all_atoms=True)).
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 Returns:
527 Iterator of Atoms objects
529 """
530 for i, (energy, mode) in enumerate(zip(energies, modes)):
531 # write imaginary frequencies as negative numbers
532 if energy.imag > energy.real:
533 energy = float(-energy.imag)
534 else:
535 energy = energy.real
537 image = atoms.copy()
538 image.info.update({'mode#': str(i),
539 'frequency_cm-1': energy / units.invcm,
540 })
541 image.arrays['mode'] = mode
543 # Custom masses are quite useful in vibration analysis, but will
544 # show up in the xyz file unless we remove them
545 if image.has('masses'):
546 del image.arrays['masses']
548 if ir_intensities is not None:
549 image.info['IR_intensity'] = float(ir_intensities[i])
551 yield image
553 def get_dos(self) -> RawDOSData:
554 """Total phonon DOS"""
555 energies = self.get_energies()
556 return RawDOSData(energies, np.ones_like(energies))
558 def get_pdos(self) -> DOSCollection:
559 """Phonon DOS, including atomic contributions"""
560 energies = self.get_energies()
561 masses = self._atoms[self.get_mask()].get_masses()
563 # Get weights as N_moving_atoms x N_modes array
564 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5
565 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T
567 mask = self.get_mask()
568 all_info = [{'index': i, 'symbol': a.symbol}
569 for i, a in enumerate(self._atoms) if mask[i]]
571 return DOSCollection([RawDOSData(energies, weights, info=info)
572 for weights, info in zip(all_weights, all_info)])
574 def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray]
575 ) -> VD:
576 """Get a copy of vibrations with modified masses and the same Hessian
578 Args:
579 masses:
580 New sequence of masses corresponding to the atom order in
581 self.get_atoms()
582 Returns:
583 A copy of the data with new masses for the same Hessian
584 """
586 new_atoms = self.get_atoms()
587 new_atoms.set_masses(masses)
588 return self.__class__(new_atoms, self.get_hessian(),
589 indices=self.get_indices())