"""Storage and analysis for vibrational data"""
import collections
from functools import cached_property
from math import pi, sin, sqrt
from numbers import Integral, Real
from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union
import numpy as np
import ase.io
import ase.units as units
from ase.atoms import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
from ase.constraints import FixAtoms, FixCartesian, constrained_indices
from ase.spectrum.doscollection import DOSCollection
from ase.spectrum.dosdata import RawDOSData
from ase.utils import jsonable
RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]]
VD = TypeVar('VD', bound='VibrationsData')
[docs]
@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.
If the Atoms object has FixedAtoms or FixedCartesian constraints, these
will be respected and the Hessian should be sized accordingly.
Args:
atoms:
Equilibrium geometry of vibrating system. This will be stored as a
full copy.
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 freely
moving atoms (i.e. not fixed ones). Leave at None if
constraints should be determined automatically from the
atoms object.
"""
def __init__(self,
atoms: Atoms,
hessian: Union[RealSequence4D, np.ndarray],
indices: Union[Sequence[int], np.ndarray] = None,
) -> None:
if indices is None:
indices = np.asarray(self.indices_from_constraints(atoms),
dtype=int)
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.")
[docs]
@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)
[docs]
@staticmethod
def indices_from_constraints(atoms: Atoms) -> List[int]:
"""Indices corresponding to Atoms Constraints
Deduces the freely moving atoms from the constraints set on the
atoms object. VibrationsData only supports FixCartesian and
FixAtoms. All others are neglected.
Args:
atoms: Atoms object.
Retruns:
indices of free atoms.
"""
# Only fully fixed atoms supported by VibrationsData
const_indices = constrained_indices(
atoms, only_include=(FixCartesian, FixAtoms))
# Invert the selection to get free atoms
indices = np.setdiff1d(
np.array(
range(
len(atoms))),
const_indices).astype(int)
return indices.tolist()
[docs]
@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: Union[np.ndarray, 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()
[docs]
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
[docs]
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()
[docs]
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'])
@cached_property
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)
[docs]
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)
[docs]
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]
[docs]
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]
[docs]
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
[docs]
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()
[docs]
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
[docs]
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
[docs]
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
[docs]
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
[docs]
def get_dos(self) -> RawDOSData:
"""Total phonon DOS"""
energies = self.get_energies()
return RawDOSData(energies, np.ones_like(energies))
[docs]
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)])
[docs]
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())