"""Storage and analysis for vibrational data""" import collections from math import sin, pi, sqrt from numbers import Real, Integral from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union import numpy as np from ase.atoms import Atoms import ase.units as units import ase.io from ase.utils import jsonable, lazymethod from ase.calculators.singlepoint import SinglePointCalculator from ase.spectrum.dosdata import RawDOSData from ase.spectrum.doscollection import DOSCollection RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]] VD = TypeVar('VD', bound='VibrationsData') @jsonable('vibrationsdata') class VibrationsData: """Class for storing and analyzing vibrational data (i.e. Atoms + Hessian) This class is not responsible for calculating Hessians; the Hessian should be computed by a Calculator or some other algorithm. Once the VibrationsData has been constructed, this class provides some common processing options; frequency calculation, mode animation, DOS etc. If the Atoms object is a periodic supercell, VibrationsData may be converted to a PhononData using the VibrationsData.to_phonondata() method. This provides access to q-point-dependent analyses such as phonon dispersion plotting. Args: atoms: Equilibrium geometry of vibrating system. This will be stored as a lightweight copy with just positions, masses, unit cell. hessian: Second-derivative in energy with respect to Cartesian nuclear movements as an (N, 3, N, 3) array. indices: indices of atoms which are included in Hessian. Default value (None) includes all atoms. """ def __init__(self, atoms: Atoms, hessian: Union[RealSequence4D, np.ndarray], indices: Union[Sequence[int], np.ndarray] = None, ) -> None: if indices is None: self._indices = np.arange(len(atoms), dtype=int) else: self._indices = np.array(indices, dtype=int) n_atoms = self._check_dimensions(atoms, np.asarray(hessian), indices=self._indices) self._atoms = atoms.copy() self._hessian2d = (np.asarray(hessian) .reshape(3 * n_atoms, 3 * n_atoms).copy()) _setter_error = ("VibrationsData properties cannot be modified: construct " "a new VibrationsData with consistent atoms, Hessian and " "(optionally) indices/mask.") @classmethod def from_2d(cls, atoms: Atoms, hessian_2d: Union[Sequence[Sequence[Real]], np.ndarray], indices: Sequence[int] = None) -> 'VibrationsData': """Instantiate VibrationsData when the Hessian is in a 3Nx3N format Args: atoms: Equilibrium geometry of vibrating system hessian: Second-derivative in energy with respect to Cartesian nuclear movements as a (3N, 3N) array. indices: Indices of (non-frozen) atoms included in Hessian """ if indices is None: indices = range(len(atoms)) assert indices is not None # Show Mypy that indices is now a sequence hessian_2d_array = np.asarray(hessian_2d) n_atoms = cls._check_dimensions(atoms, hessian_2d_array, indices=indices, two_d=True) return cls(atoms, hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3), indices=indices) @staticmethod def indices_from_mask(mask: Union[Sequence[bool], np.ndarray] ) -> List[int]: """Indices corresponding to boolean mask This is provided as a convenience for instantiating VibrationsData with a boolean mask. For example, if the Hessian data includes only the H atoms in a structure:: h_mask = atoms.get_chemical_symbols() == 'H' vib_data = VibrationsData(atoms, hessian, VibrationsData.indices_from_mask(h_mask)) Take care to ensure that the length of the mask corresponds to the full number of atoms; this function is only aware of the mask it has been given. Args: mask: a sequence of True, False values Returns: indices of True elements """ return np.where(mask)[0].tolist() @staticmethod def _check_dimensions(atoms: Atoms, hessian: np.ndarray, indices: Sequence[int], two_d: bool = False) -> int: """Sanity check on array shapes from input data Args: atoms: Structure indices: Indices of atoms used in Hessian hessian: Proposed Hessian array Returns: Number of atoms contributing to Hessian Raises: ValueError if Hessian dimensions are not (N, 3, N, 3) """ n_atoms = len(atoms[indices]) if two_d: ref_shape = [n_atoms * 3, n_atoms * 3] ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3)) else: ref_shape = [n_atoms, 3, n_atoms, 3] ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms) if (isinstance(hessian, np.ndarray) and hessian.shape == tuple(ref_shape)): return n_atoms else: raise ValueError("Hessian for these atoms should be a " "{} numpy array.".format(ref_shape_txt)) def get_atoms(self) -> Atoms: return self._atoms.copy() def get_indices(self) -> np.ndarray: return self._indices.copy() def get_mask(self) -> np.ndarray: """Boolean mask of atoms selected by indices""" return self._mask_from_indices(self._atoms, self.get_indices()) @staticmethod def _mask_from_indices(atoms: Atoms, indices: Union[None, Sequence[int], np.ndarray] ) -> np.ndarray: """Boolean mask of atoms selected by indices""" natoms = len(atoms) # Wrap indices to allow negative values indices = np.asarray(indices) % natoms mask = np.full(natoms, False, dtype=bool) mask[indices] = True return mask def get_hessian(self) -> np.ndarray: """The Hessian; second derivative of energy wrt positions This format is preferred for iteration over atoms and when addressing specific elements of the Hessian. Returns: array with shape (n_atoms, 3, n_atoms, 3) where - the first and third indices identify atoms in self.get_atoms() - the second and fourth indices cover the corresponding Cartesian movements in x, y, z e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on atoms[1] in the x-direction in response to a movement in the z-direction of atoms[0] """ n_atoms = int(self._hessian2d.shape[0] / 3) return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy() def get_hessian_2d(self) -> np.ndarray: """Get the Hessian as a 2-D array This format may be preferred for use with standard linear algebra functions Returns: array with shape (n_atoms * 3, n_atoms * 3) where the elements are ordered by atom and Cartesian direction [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...], [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...], [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...], [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...], ...] e.g. the element h[2, 3] gives a harmonic force exerted on atoms[1] in the x-direction in response to a movement in the z-direction of atoms[0] """ return self._hessian2d.copy() def todict(self) -> Dict[str, Any]: if np.allclose(self._indices, range(len(self._atoms))): indices = None else: indices = self.get_indices() return {'atoms': self.get_atoms(), 'hessian': self.get_hessian(), 'indices': indices} @classmethod def fromdict(cls, data: Dict[str, Any]) -> 'VibrationsData': # mypy is understandably suspicious of data coming from a dict that # holds mixed types, but it can see if we sanity-check with 'assert' assert isinstance(data['atoms'], Atoms) assert isinstance(data['hessian'], (collections.abc.Sequence, np.ndarray)) if data['indices'] is not None: assert isinstance(data['indices'], (collections.abc.Sequence, np.ndarray)) for index in data['indices']: assert isinstance(index, Integral) return cls(data['atoms'], data['hessian'], indices=data['indices']) @lazymethod def _energies_and_modes(self) -> Tuple[np.ndarray, np.ndarray]: """Diagonalise the Hessian to obtain harmonic modes This method is an internal implementation of get_energies_and_modes(), see the docstring of that method for more information. """ active_atoms = self._atoms[self.get_mask()] n_atoms = len(active_atoms) masses = active_atoms.get_masses() if not np.all(masses): raise ValueError('Zero mass encountered in one or more of ' 'the vibrated atoms. Use Atoms.set_masses()' ' to set all masses to non-zero values.') mass_weights = np.repeat(masses**-0.5, 3) omega2, vectors = np.linalg.eigh(mass_weights * self.get_hessian_2d() * mass_weights[:, np.newaxis]) unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu) energies = unit_conversion * omega2.astype(complex)**0.5 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3) modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5 return (energies, modes) def get_energies_and_modes(self, all_atoms: bool = False ) -> Tuple[np.ndarray, np.ndarray]: """Diagonalise the Hessian to obtain harmonic modes Results are cached so diagonalization will only be performed once for this object instance. Args: all_atoms: If True, return modes as (3N, [N + N_frozen], 3) array where the second axis corresponds to the full list of atoms in the attached atoms object. Atoms that were not included in the Hessian will have displacement vectors of (0, 0, 0). Returns: tuple (energies, modes) Energies are given in units of eV. (To convert these to frequencies in cm-1, divide by ase.units.invcm.) Modes are given in Cartesian coordinates as a (3N, N, 3) array where indices correspond to the (mode_index, atom, direction). """ energies, modes_from_hessian = self._energies_and_modes() if all_atoms: n_active_atoms = len(self.get_indices()) n_all_atoms = len(self._atoms) modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3)) modes[:, self.get_mask(), :] = modes_from_hessian else: modes = modes_from_hessian.copy() return (energies.copy(), modes) def get_modes(self, all_atoms: bool = False) -> np.ndarray: """Diagonalise the Hessian to obtain harmonic modes Results are cached so diagonalization will only be performed once for this object instance. all_atoms: If True, return modes as (3N, [N + N_frozen], 3) array where the second axis corresponds to the full list of atoms in the attached atoms object. Atoms that were not included in the Hessian will have displacement vectors of (0, 0, 0). Returns: Modes in Cartesian coordinates as a (3N, N, 3) array where indices correspond to the (mode_index, atom, direction). """ return self.get_energies_and_modes(all_atoms=all_atoms)[1] def get_energies(self) -> np.ndarray: """Diagonalise the Hessian to obtain eigenvalues Results are cached so diagonalization will only be performed once for this object instance. Returns: Harmonic mode energies in units of eV """ return self.get_energies_and_modes()[0] def get_frequencies(self) -> np.ndarray: """Diagonalise the Hessian to obtain frequencies in cm^-1 Results are cached so diagonalization will only be performed once for this object instance. Returns: Harmonic mode frequencies in units of cm^-1 """ return self.get_energies() / units.invcm def get_zero_point_energy(self) -> float: """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy Args: energies: Pre-computed energy eigenvalues. Use if available to avoid re-calculating these from the Hessian. Returns: zero-point energy in eV """ return self._calculate_zero_point_energy(self.get_energies()) @staticmethod def _calculate_zero_point_energy(energies: Union[Sequence[complex], np.ndarray]) -> float: return 0.5 * np.asarray(energies).real.sum() def tabulate(self, im_tol: float = 1e-8) -> str: """Get a summary of the vibrational frequencies. Args: im_tol: Tolerance for imaginary frequency in eV. If frequency has a larger imaginary component than im_tol, the imaginary component is shown in the summary table. Returns: Summary table as formatted text """ energies = self.get_energies() return ('\n'.join(self._tabulate_from_energies(energies, im_tol=im_tol)) + '\n') @classmethod def _tabulate_from_energies(cls, energies: Union[Sequence[complex], np.ndarray], im_tol: float = 1e-8) -> List[str]: summary_lines = ['---------------------', ' # meV cm^-1', '---------------------'] for n, e in enumerate(energies): if abs(e.imag) > im_tol: c = 'i' e = e.imag else: c = '' e = e.real summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}' .format(index=n, mev=(e * 1e3), cm=(e / units.invcm), im=c)) summary_lines.append('---------------------') summary_lines.append('Zero-point energy: {:.3f} eV'.format( cls._calculate_zero_point_energy(energies=energies))) return summary_lines def iter_animated_mode(self, mode_index: int, temperature: float = units.kB * 300, frames: int = 30) -> Iterator[Atoms]: """Obtain animated mode as a series of Atoms Args: mode_index: Selection of mode to animate temperature: In energy units - use units.kB * T_IN_KELVIN frames: number of image frames in animation Yields: Displaced atoms following vibrational mode """ mode = (self.get_modes(all_atoms=True)[mode_index] * sqrt(temperature / abs(self.get_energies()[mode_index]))) for phase in np.linspace(0, 2 * pi, frames, endpoint=False): atoms = self.get_atoms() atoms.positions += sin(phase) * mode yield atoms def show_as_force(self, mode: int, scale: float = 0.2, show: bool = True) -> Atoms: """Illustrate mode as "forces" on atoms Args: mode: mode index scale: scale factor show: if True, open the ASE GUI and show atoms Returns: Atoms with scaled forces corresponding to mode eigenvectors (using attached SinglePointCalculator). """ atoms = self.get_atoms() mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale atoms.calc = SinglePointCalculator(atoms, forces=mode) if show: atoms.edit() return atoms def write_jmol(self, filename: str = 'vib.xyz', ir_intensities: Union[Sequence[float], np.ndarray] = None ) -> None: """Writes file for viewing of the modes with jmol. This is an extended XYZ file with eigenvectors given as extra columns and metadata given in the label/comment line for each image. The format is not quite human-friendly, but has the advantage that it can be imported back into ASE with ase.io.read. Args: filename: Path for output file ir_intensities: If available, IR intensities can be included in the header lines. This does not affect the visualisation, but may be convenient when comparing to experimental data. """ all_images = list(self._get_jmol_images(atoms=self.get_atoms(), energies=self.get_energies(), modes=self.get_modes(all_atoms=True), ir_intensities=ir_intensities)) ase.io.write(filename, all_images, format='extxyz') @staticmethod def _get_jmol_images(atoms: Atoms, energies: np.ndarray, modes: np.ndarray, ir_intensities: Union[Sequence[float], np.ndarray] = None ) -> Iterator[Atoms]: """Get vibrational modes as a series of Atoms with attached data For each image (Atoms object): - eigenvalues are attached to image.arrays['mode'] - "mode#" and "frequency_cm-1" are set in image.info - "IR_intensity" is set if provided in ir_intensities - "masses" is removed This is intended to set up the object for JMOL-compatible export using ase.io.extxyz. Args: atoms: The base atoms object; all images have the same positions energies: Complex vibrational energies in eV modes: Eigenvectors array corresponding to atoms and energies. This should cover the full set of atoms (i.e. modes = vib.get_modes(all_atoms=True)). ir_intensities: If available, IR intensities can be included in the header lines. This does not affect the visualisation, but may be convenient when comparing to experimental data. Returns: Iterator of Atoms objects """ for i, (energy, mode) in enumerate(zip(energies, modes)): # write imaginary frequencies as negative numbers if energy.imag > energy.real: energy = float(-energy.imag) else: energy = energy.real image = atoms.copy() image.info.update({'mode#': str(i), 'frequency_cm-1': energy / units.invcm, }) image.arrays['mode'] = mode # Custom masses are quite useful in vibration analysis, but will # show up in the xyz file unless we remove them if image.has('masses'): del image.arrays['masses'] if ir_intensities is not None: image.info['IR_intensity'] = float(ir_intensities[i]) yield image def get_dos(self) -> RawDOSData: """Total phonon DOS""" energies = self.get_energies() return RawDOSData(energies, np.ones_like(energies)) def get_pdos(self) -> DOSCollection: """Phonon DOS, including atomic contributions""" energies = self.get_energies() masses = self._atoms[self.get_mask()].get_masses() # Get weights as N_moving_atoms x N_modes array vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T mask = self.get_mask() all_info = [{'index': i, 'symbol': a.symbol} for i, a in enumerate(self._atoms) if mask[i]] return DOSCollection([RawDOSData(energies, weights, info=info) for weights, info in zip(all_weights, all_info)]) def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray] ) -> VD: """Get a copy of vibrations with modified masses and the same Hessian Args: masses: New sequence of masses corresponding to the atom order in self.get_atoms() Returns: A copy of the data with new masses for the same Hessian """ new_atoms = self.get_atoms() new_atoms.set_masses(masses) return self.__class__(new_atoms, self.get_hessian(), indices=self.get_indices())