"""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, Iterator, Sequence, TypeVar
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 :class:`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, :class:`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 :class:`~ase.constraints.FixAtoms` or
:class:`~.ase.constraints.FixCartesian` constraints,
these will be respected and the Hessian should be sized accordingly.
Parameters
----------
atoms : :class:`~ase.Atoms`
Equilibrium geometry of vibrating system.
This will be stored as a full copy.
hessian : np.ndarray
Second-derivative in energy with respect to Cartesian nuclear movements
as an (N, 3, N, 3) array.
indices : Sequence[int] | np.ndarray | None, default: None
Indices of atoms which are included in Hessian.
By default, it 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: RealSequence4D | np.ndarray,
indices: Sequence[int] | np.ndarray | None = 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: Sequence[Sequence[Real]] | np.ndarray,
indices: Sequence[int] | None = None,
) -> 'VibrationsData':
"""Instantiate VibrationsData when the Hessian is in a 3Nx3N format
Parameters
----------
atoms : :class:`~ase.Atoms`
Equilibrium geometry of vibrating system.
hessian : Sequence[Sequence[Real]] | np.ndarray
Second-derivative in energy with respect to Cartesian nuclear
movements as a (3N, 3N) array.
indices : Sequence[int] | None, default: None
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. :class:`VibrationsData` only supports:
- :class:`~ase.constraints.FixCartesian`
- :class:`~ase.constraints.FixAtoms`
All others are neglected.
Parameters
----------
atoms : :class:`~ase.Atoms`
Returns
-------
list[int]
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)
# TODO: use numpy.typing to resolve this error.
return indices.tolist() # type: ignore[return-value]
[docs]
@staticmethod
def indices_from_mask(
mask: Sequence[bool] | np.ndarray,
) -> list[int]:
"""Indices corresponding to boolean mask
This is provided as a convenience for instantiating
:class:`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.
Parameters
----------
mask: Sequence[bool] | np.ndarray
Sequence of boolean values.
Returns
-------
list[int]
Indices of :obj:`True` elements.
"""
# TODO: use numpy.typing to resolve this error.
return np.where(mask)[0].tolist() # type: ignore[return-value]
@staticmethod
def _check_dimensions(
atoms: Atoms,
hessian: np.ndarray,
indices: np.ndarray | Sequence[int],
two_d: bool = False,
) -> int:
"""Sanity check on array shapes from input data
Parameters
----------
atoms : :class:`~ase.Atoms`
hessian : np.ndarray
Proposed Hessian array.
indices : np.ndarray | Sequence[int]
Indices of atoms used in Hessian.
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: Sequence[int] | np.ndarray | None,
) -> 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
-------
np.ndarray
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
-------
np.ndarray
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.
Parameters
----------
all_atoms : bool
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[np.ndarray, np.ndarray]
Tuple of (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.
Parameters
----------
all_atoms : bool
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
-------
np.ndarray
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
-------
np.ndarray
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
Returns
-------
float
Zero-point energy in eV.
"""
return self._calculate_zero_point_energy(self.get_energies())
@staticmethod
def _calculate_zero_point_energy(
energies: 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.
Parameters
----------
im_tol : float
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
-------
str
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: 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.0,
frames: int = 30,
) -> Iterator[Atoms]:
"""Obtain animated mode as a series of Atoms
Parameters
----------
mode_index : int
Selection of mode to animate.
temperature : float
Temperature in energy units - use ``units.kB * T_IN_KELVIN``
frames : int
Number of image frames in animation.
Yields
------
:class:`~ase.Atoms`
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
Parameters
----------
mode : int
Mode index.
scale : float
Scale factor.
show : bool
If True, open the ASE GUI and show atoms.
Returns
-------
:class:`~ase.Atoms`
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: Sequence[float] | np.ndarray | None = 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.
Parameters
----------
filename : str
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: Sequence[float] | np.ndarray | None = 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.
Parameters
----------
atoms : :class:`~ase.Atoms`
The base atoms object; all images have the same positions
energies : np.ndarray
Complex vibrational energies in eV.
modes : np.ndarray
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: Sequence[float] | np.ndarray | None = None
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[Atoms]
"""
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: Sequence[float] | np.ndarray,
) -> VD:
"""Get a copy of vibrations with modified masses and the same Hessian
Parameters
----------
masses: Sequence[float] | np.ndarray
New sequence of masses corresponding to the atom order in
``self.get_atoms()``.
Returns
-------
:class:`VibrationalData`
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()
)
