"""A class for computing vibrational modes"""
import sys
from collections import namedtuple
from math import log, pi, sqrt
from pathlib import Path
from typing import Iterator, Tuple
import numpy as np
import ase.io
import ase.units as units
from ase.atoms import Atoms
from ase.constraints import FixAtoms
from ase.parallel import paropen, world
from ase.utils.filecache import get_json_cache
from .data import VibrationsData
class AtomicDisplacements:
def _disp(self, a, i, step):
if isinstance(i, str): # XXX Simplify by removing this.
i = 'xyz'.index(i)
return Displacement(a, i, np.sign(step), abs(step), self)
def _eq_disp(self):
return self._disp(0, 0, 0)
@property
def ndof(self):
return 3 * len(self.indices)
class Displacement(namedtuple('Displacement', ['a', 'i', 'sign', 'ndisp',
'vib'])):
@property
def name(self):
if self.sign == 0:
return 'eq'
axisname = 'xyz'[self.i]
dispname = self.ndisp * ' +-'[self.sign]
return f'{self.a}{axisname}{dispname}'
@property
def _cached(self):
return self.vib.cache[self.name]
def forces(self):
return self._cached['forces'].copy()
@property
def step(self):
return self.ndisp * self.sign * self.vib.delta
# XXX dipole only valid for infrared
def dipole(self):
return self._cached['dipole'].copy()
# XXX below stuff only valid for TDDFT excitation stuff
def save_ov_nn(self, ov_nn):
np.save(Path(self.vib.exname) / (self.name + '.ov'), ov_nn)
def load_ov_nn(self):
return np.load(Path(self.vib.exname) / (self.name + '.ov.npy'))
@property
def _exname(self):
return Path(self.vib.exname) / f'ex.{self.name}{self.vib.exext}'
def calculate_and_save_static_polarizability(self, atoms):
exobj = self.vib._new_exobj()
excitation_data = exobj(atoms)
np.savetxt(self._exname, excitation_data)
def load_static_polarizability(self):
return np.loadtxt(self._exname)
def read_exobj(self):
# XXX each exobj should allow for self._exname as Path
return self.vib.read_exobj(str(self._exname))
def calculate_and_save_exlist(self, atoms):
# exo = self.vib._new_exobj()
excalc = self.vib._new_exobj()
exlist = excalc.calculate(atoms)
# XXX each exobj should allow for self._exname as Path
exlist.write(str(self._exname))
[docs]
class Vibrations(AtomicDisplacements):
"""Class for calculating vibrational modes using finite difference.
The vibrational modes are calculated from a finite difference
approximation of the Hessian matrix.
The *summary()*, *get_energies()* and *get_frequencies()* methods all take
an optional *method* keyword. Use method='Frederiksen' to use the method
described in:
T. Frederiksen, M. Paulsson, M. Brandbyge, A. P. Jauho:
"Inelastic transport theory from first-principles: methodology and
applications for nanoscale devices", Phys. Rev. B 75, 205413 (2007)
atoms: Atoms object
The atoms to work on.
indices: list of int
List of indices of atoms to vibrate. Default behavior is
to vibrate all atoms.
name: str
Name to use for files.
delta: float
Magnitude of displacements.
nfree: int
Number of displacements per atom and cartesian coordinate, 2 and 4 are
supported. Default is 2 which will displace each atom +delta and
-delta for each cartesian coordinate.
Example:
>>> from ase import Atoms
>>> from ase.calculators.emt import EMT
>>> from ase.optimize import BFGS
>>> from ase.vibrations import Vibrations
>>> n2 = Atoms('N2', [(0, 0, 0), (0, 0, 1.1)],
... calculator=EMT())
>>> BFGS(n2).run(fmax=0.01)
BFGS: 0 16:01:21 0.440339 3.2518
BFGS: 1 16:01:21 0.271928 0.8211
BFGS: 2 16:01:21 0.263278 0.1994
BFGS: 3 16:01:21 0.262777 0.0088
>>> vib = Vibrations(n2)
>>> vib.run()
>>> vib.summary()
---------------------
# meV cm^-1
---------------------
0 0.0 0.0
1 0.0 0.0
2 0.0 0.0
3 1.4 11.5
4 1.4 11.5
5 152.7 1231.3
---------------------
Zero-point energy: 0.078 eV
>>> vib.write_mode(-1) # write last mode to trajectory file
"""
def __init__(self, atoms, indices=None, name='vib', delta=0.01, nfree=2):
assert nfree in [2, 4]
self.atoms = atoms
self.calc = atoms.calc
if indices is None:
fixed_indices = []
for constr in atoms.constraints:
if isinstance(constr, FixAtoms):
fixed_indices.extend(constr.get_indices())
fixed_indices = list(set(fixed_indices))
indices = [i for i in range(len(atoms)) if i not in fixed_indices]
if len(indices) != len(set(indices)):
raise ValueError(
'one (or more) indices included more than once')
self.indices = np.asarray(indices)
self.delta = delta
self.nfree = nfree
self.H = None
self.ir = None
self._vibrations = None
self.cache = get_json_cache(name)
@property
def name(self):
return str(self.cache.directory)
[docs]
def run(self):
"""Run the vibration calculations.
This will calculate the forces for 6 displacements per atom +/-x,
+/-y, +/-z. Only those calculations that are not already done will be
started. Be aware that an interrupted calculation may produce an empty
file (ending with .json), which must be deleted before restarting the
job. Otherwise the forces will not be calculated for that
displacement.
Note that the calculations for the different displacements can be done
simultaneously by several independent processes. This feature relies
on the existence of files and the subsequent creation of the file in
case it is not found.
If the program you want to use does not have a calculator in ASE, use
``iterdisplace`` to get all displaced structures and calculate the
forces on your own.
"""
if not self.cache.writable:
raise RuntimeError(
'Cannot run calculation. '
'Cache must be removed or split in order '
'to have only one sort of data structure at a time.')
self._check_old_pickles()
for disp, atoms in self.iterdisplace(inplace=True):
with self.cache.lock(disp.name) as handle:
if handle is None:
continue
result = self.calculate(atoms, disp)
if world.rank == 0:
handle.save(result)
def _check_old_pickles(self):
from pathlib import Path
eq_pickle_path = Path(f'{self.name}.eq.pckl')
pickle2json_instructions = f"""\
Found old pickle files such as {eq_pickle_path}. \
Please remove them and recalculate or run \
"python -m ase.vibrations.pickle2json --help"."""
if len(self.cache) == 0 and eq_pickle_path.exists():
raise RuntimeError(pickle2json_instructions)
[docs]
def iterdisplace(self, inplace=False) -> \
Iterator[Tuple[Displacement, Atoms]]:
"""Iterate over initial and displaced structures.
Use this to export the structures for each single-point calculation
to an external program instead of using ``run()``. Then save the
calculated gradients to <name>.json and continue using this instance.
Parameters:
------------
inplace: bool
If True, the atoms object will be modified in-place. Otherwise a
copy will be made.
Yields:
--------
disp: Displacement
Displacement object with information about the displacement.
atoms: Atoms
Atoms object with the displaced structure.
"""
# XXX change of type of disp
atoms = self.atoms if inplace else self.atoms.copy()
displacements = self.displacements()
eq_disp = next(displacements)
assert eq_disp.name == 'eq'
yield eq_disp, atoms
for disp in displacements:
if not inplace:
atoms = self.atoms.copy()
pos0 = atoms.positions[disp.a, disp.i]
atoms.positions[disp.a, disp.i] += disp.step
yield disp, atoms
if inplace:
atoms.positions[disp.a, disp.i] = pos0
[docs]
def iterimages(self):
"""Yield initial and displaced structures."""
for name, atoms in self.iterdisplace():
yield atoms
def _iter_ai(self):
for a in self.indices:
for i in range(3):
yield a, i
def displacements(self):
yield self._eq_disp()
for a, i in self._iter_ai():
for sign in [-1, 1]:
for ndisp in range(1, self.nfree // 2 + 1):
yield self._disp(a, i, sign * ndisp)
def calculate(self, atoms, disp):
results = {}
results['forces'] = self.calc.get_forces(atoms)
if self.ir:
results['dipole'] = self.calc.get_dipole_moment(atoms)
return results
[docs]
def clean(self, empty_files=False, combined=True):
"""Remove json-files.
Use empty_files=True to remove only empty files and
combined=False to not remove the combined file.
"""
if world.rank != 0:
return 0
if empty_files:
return self.cache.strip_empties() # XXX Fails on combined cache
nfiles = self.cache.filecount()
self.cache.clear()
return nfiles
[docs]
def combine(self):
"""Combine json-files to one file ending with '.all.json'.
The other json-files will be removed in order to have only one sort
of data structure at a time.
"""
nelements_before = self.cache.filecount()
self.cache = self.cache.combine()
return nelements_before
[docs]
def split(self):
"""Split combined json-file.
The combined json-file will be removed in order to have only one
sort of data structure at a time.
"""
count = self.cache.filecount()
self.cache = self.cache.split()
return count
def read(self, method='standard', direction='central'):
self.method = method.lower()
self.direction = direction.lower()
assert self.method in ['standard', 'frederiksen']
assert self.direction in ['central', 'forward', 'backward']
n = 3 * len(self.indices)
H = np.empty((n, n))
r = 0
eq_disp = self._eq_disp()
if direction != 'central':
feq = eq_disp.forces()
for a, i in self._iter_ai():
disp_minus = self._disp(a, i, -1)
disp_plus = self._disp(a, i, 1)
fminus = disp_minus.forces()
fplus = disp_plus.forces()
if self.method == 'frederiksen':
fminus[a] -= fminus.sum(0)
fplus[a] -= fplus.sum(0)
if self.nfree == 4:
fminusminus = self._disp(a, i, -2).forces()
fplusplus = self._disp(a, i, 2).forces()
if self.method == 'frederiksen':
fminusminus[a] -= fminusminus.sum(0)
fplusplus[a] -= fplusplus.sum(0)
if self.direction == 'central':
if self.nfree == 2:
H[r] = .5 * (fminus - fplus)[self.indices].ravel()
else:
assert self.nfree == 4
H[r] = H[r] = (-fminusminus +
8 * fminus -
8 * fplus +
fplusplus)[self.indices].ravel() / 12.0
elif self.direction == 'forward':
H[r] = (feq - fplus)[self.indices].ravel()
else:
assert self.direction == 'backward'
H[r] = (fminus - feq)[self.indices].ravel()
H[r] /= 2 * self.delta
r += 1
H += H.copy().T
self.H = H
masses = self.atoms.get_masses()
if any(masses[self.indices] == 0):
raise RuntimeError('Zero mass encountered in one or more of '
'the vibrated atoms. Use Atoms.set_masses()'
' to set all masses to non-zero values.')
self.im = np.repeat(masses[self.indices]**-0.5, 3)
self._vibrations = self.get_vibrations(read_cache=False,
method=self.method,
direction=self.direction)
omega2, modes = np.linalg.eigh(self.im[:, None] * H * self.im)
self.modes = modes.T.copy()
# Conversion factor:
s = units._hbar * 1e10 / sqrt(units._e * units._amu)
self.hnu = s * omega2.astype(complex)**0.5
[docs]
def get_vibrations(self, method='standard', direction='central',
read_cache=True, **kw):
"""Get vibrations as VibrationsData object
If read() has not yet been called, this will be called to assemble data
from the outputs of run(). Most of the arguments to this function are
options to be passed to read() in this case.
Args:
method (str): Calculation method passed to read()
direction (str): Finite-difference scheme passed to read()
read_cache (bool): The VibrationsData object will be cached for
quick access. Set False to force regeneration of the cache with
the current atoms/Hessian/indices data.
**kw: Any remaining keyword arguments are passed to read()
Returns:
VibrationsData
"""
if read_cache and (self._vibrations is not None):
return self._vibrations
else:
if (self.H is None or method.lower() != self.method or
direction.lower() != self.direction):
self.read(method, direction, **kw)
return VibrationsData.from_2d(self.atoms, self.H,
indices=self.indices)
[docs]
def get_energies(self, method='standard', direction='central', **kw):
"""Get vibration energies in eV."""
return self.get_vibrations(method=method,
direction=direction, **kw).get_energies()
[docs]
def get_frequencies(self, method='standard', direction='central'):
"""Get vibration frequencies in cm^-1."""
return self.get_vibrations(method=method,
direction=direction).get_frequencies()
def summary(self, method='standard', direction='central', freq=None,
log=sys.stdout):
if freq is not None:
energies = freq * units.invcm
else:
energies = self.get_energies(method=method, direction=direction)
summary_lines = VibrationsData._tabulate_from_energies(energies)
log_text = '\n'.join(summary_lines) + '\n'
if isinstance(log, str):
with paropen(log, 'a') as log_file:
log_file.write(log_text)
else:
log.write(log_text)
def get_zero_point_energy(self, freq=None):
if freq:
raise NotImplementedError
return self.get_vibrations().get_zero_point_energy()
[docs]
def get_mode(self, n):
"""Get mode number ."""
return self.get_vibrations().get_modes(all_atoms=True)[n]
[docs]
def write_mode(self, n=None, kT=units.kB * 300, nimages=30):
"""Write mode number n to trajectory file. If n is not specified,
writes all non-zero modes."""
if n is None:
for index, energy in enumerate(self.get_energies()):
if abs(energy) > 1e-5:
self.write_mode(n=index, kT=kT, nimages=nimages)
return
else:
n %= len(self.get_energies())
with ase.io.Trajectory('%s.%d.traj' % (self.name, n), 'w') as traj:
for image in (self.get_vibrations()
.iter_animated_mode(n,
temperature=kT, frames=nimages)):
traj.write(image)
def show_as_force(self, n, scale=0.2, show=True):
return self.get_vibrations().show_as_force(n, scale=scale, show=show)
[docs]
def write_jmol(self):
"""Writes file for viewing of the modes with jmol."""
with open(self.name + '.xyz', 'w') as fd:
self._write_jmol(fd)
def _write_jmol(self, fd):
symbols = self.atoms.get_chemical_symbols()
freq = self.get_frequencies()
for n in range(3 * len(self.indices)):
fd.write('%6d\n' % len(self.atoms))
if freq[n].imag != 0:
c = 'i'
freq[n] = freq[n].imag
else:
freq[n] = freq[n].real
c = ' '
fd.write('Mode #%d, f = %.1f%s cm^-1'
% (n, float(freq[n].real), c))
if self.ir:
fd.write(', I = %.4f (D/Å)^2 amu^-1.\n' % self.intensities[n])
else:
fd.write('.\n')
mode = self.get_mode(n)
for i, pos in enumerate(self.atoms.positions):
fd.write('%2s %12.5f %12.5f %12.5f %12.5f %12.5f %12.5f\n' %
(symbols[i], pos[0], pos[1], pos[2],
mode[i, 0], mode[i, 1], mode[i, 2]))
[docs]
def fold(self, frequencies, intensities,
start=800.0, end=4000.0, npts=None, width=4.0,
type='Gaussian', normalize=False):
"""Fold frequencies and intensities within the given range
and folding method (Gaussian/Lorentzian).
The energy unit is cm^-1.
normalize=True ensures the integral over the peaks to give the
intensity.
"""
lctype = type.lower()
assert lctype in ['gaussian', 'lorentzian']
if not npts:
npts = int((end - start) / width * 10 + 1)
prefactor = 1
if lctype == 'lorentzian':
intensities = intensities * width * pi / 2.
if normalize:
prefactor = 2. / width / pi
else:
sigma = width / 2. / sqrt(2. * log(2.))
if normalize:
prefactor = 1. / sigma / sqrt(2 * pi)
# Make array with spectrum data
spectrum = np.empty(npts)
energies = np.linspace(start, end, npts)
for i, energy in enumerate(energies):
energies[i] = energy
if lctype == 'lorentzian':
spectrum[i] = (intensities * 0.5 * width / pi /
((frequencies - energy)**2 +
0.25 * width**2)).sum()
else:
spectrum[i] = (intensities *
np.exp(-(frequencies - energy)**2 /
2. / sigma**2)).sum()
return [energies, prefactor * spectrum]
[docs]
def write_dos(self, out='vib-dos.dat', start=800, end=4000,
npts=None, width=10,
type='Gaussian', method='standard', direction='central'):
"""Write out the vibrational density of states to file.
First column is the wavenumber in cm^-1, the second column the
folded vibrational density of states.
Start and end points, and width of the Gaussian/Lorentzian
should be given in cm^-1."""
frequencies = self.get_frequencies(method, direction).real
intensities = np.ones(len(frequencies))
energies, spectrum = self.fold(frequencies, intensities,
start, end, npts, width, type)
# Write out spectrum in file.
outdata = np.empty([len(energies), 2])
outdata.T[0] = energies
outdata.T[1] = spectrum
with open(out, 'w') as fd:
fd.write(f'# {type.title()} folded, width={width:g} cm^-1\n')
fd.write('# [cm^-1] arbitrary\n')
for row in outdata:
fd.write('%.3f %15.5e\n' %
(row[0], row[1]))