from collections.abc import Iterable
from typing import Optional
import numpy as np
from ase import Atoms
from ase.calculators.calculator import BaseCalculator, all_properties
[docs]
class FiniteDifferenceCalculator(BaseCalculator):
"""Wrapper calculator using the finite-difference method.
The forces and the stress are computed using the finite-difference method.
.. versionadded:: 3.24.0
"""
implemented_properties = all_properties
def __init__(
self,
calc: BaseCalculator,
eps_disp: float = 1e-6,
eps_strain: float = 1e-6,
) -> None:
"""
Parameters
----------
calc : :class:`~ase.calculators.calculator.BaseCalculator`
ASE Calculator object to be wrapped.
eps_disp : float, default 1e-6
Displacement used for computing forces.
eps_strain : float, default 1e-6
Strain used for computing stress.
"""
super().__init__()
self.calc = calc
self.eps_disp = eps_disp
self.eps_strain = eps_strain
def calculate(self, atoms: Atoms, properties, system_changes) -> None:
atoms = atoms.copy() # copy to not mess up original `atoms`
atoms.calc = self.calc
self.results = {
'energy': atoms.get_potential_energy(),
'forces': calculate_numerical_forces(atoms, eps=self.eps_disp),
'stress': calculate_numerical_stress(atoms, eps=self.eps_strain),
}
self.results['free_energy'] = self.results['energy']
def _numeric_force(
atoms: Atoms,
iatom: int,
icart: int,
eps: float = 1e-6,
) -> float:
"""Calculate numerical force on a specific atom along a specific direction.
Parameters
----------
atoms : :class:`~ase.Atoms`
ASE :class:`~ase.Atoms` object.
iatom : int
Index of atoms.
icart : {0, 1, 2}
Index of Cartesian component.
eps : float, default 1e-6
Displacement.
"""
p0 = atoms.get_positions()
p = p0.copy()
p[iatom, icart] = p0[iatom, icart] + eps
atoms.set_positions(p, apply_constraint=False)
eplus = atoms.get_potential_energy()
p[iatom, icart] = p0[iatom, icart] - eps
atoms.set_positions(p, apply_constraint=False)
eminus = atoms.get_potential_energy()
atoms.set_positions(p0, apply_constraint=False)
return (eminus - eplus) / (2 * eps)
[docs]
def calculate_numerical_forces(
atoms: Atoms,
eps: float = 1e-6,
iatoms: Optional[Iterable[int]] = None,
icarts: Optional[Iterable[int]] = None,
) -> np.ndarray:
"""Calculate forces numerically based on the finite-difference method.
Parameters
----------
atoms : :class:`~ase.Atoms`
ASE :class:`~ase.Atoms` object.
eps : float, default 1e-6
Displacement.
iatoms : Optional[Iterable[int]]
Indices of atoms for which forces are computed.
By default, all atoms are considered.
icarts : Optional[Iterable[int]]
Indices of Cartesian coordinates for which forces are computed.
By default, all three coordinates are considered.
Returns
-------
forces : np.ndarray
Forces computed numerically based on the finite-difference method.
"""
if iatoms is None:
iatoms = range(len(atoms))
if icarts is None:
icarts = [0, 1, 2]
return np.array(
[
[_numeric_force(atoms, iatom, icart, eps) for icart in icarts]
for iatom in iatoms
]
)
[docs]
def calculate_numerical_stress(
atoms: Atoms,
eps: float = 1e-6,
voigt: bool = True,
) -> np.ndarray:
"""Calculate stress numerically based on the finite-difference method.
Parameters
----------
atoms : :class:`~ase.Atoms`
ASE :class:`~ase.Atoms` object.
eps : float, default 1e-6
Strain in the Voigt notation.
voigt : bool, default True
If True, the stress is returned in the Voigt notation.
Returns
-------
stress : np.ndarray
Stress computed numerically based on the finite-difference method.
"""
stress = np.zeros((3, 3), dtype=float)
cell = atoms.cell.copy()
volume = atoms.get_volume()
for i in range(3):
x = np.eye(3)
x[i, i] = 1.0 + eps
atoms.set_cell(cell @ x, scale_atoms=True)
eplus = atoms.get_potential_energy(force_consistent=True)
x[i, i] = 1.0 - eps
atoms.set_cell(cell @ x, scale_atoms=True)
eminus = atoms.get_potential_energy(force_consistent=True)
stress[i, i] = (eplus - eminus) / (2 * eps * volume)
x[i, i] = 1.0
j = i - 2
x[i, j] = x[j, i] = +0.5 * eps
atoms.set_cell(cell @ x, scale_atoms=True)
eplus = atoms.get_potential_energy(force_consistent=True)
x[i, j] = x[j, i] = -0.5 * eps
atoms.set_cell(cell @ x, scale_atoms=True)
eminus = atoms.get_potential_energy(force_consistent=True)
stress[i, j] = stress[j, i] = (eplus - eminus) / (2 * eps * volume)
atoms.set_cell(cell, scale_atoms=True)
return stress.flat[[0, 4, 8, 5, 2, 1]] if voigt else stress