"""Relaxation step of ASSYST."""
from dataclasses import dataclass
from typing import Literal, Iterable, Iterator
from .calculators import AseCalculatorConfig
from .utils import update_uuid
from ase import Atoms
from ase.calculators.calculator import Calculator
from ase.calculators.singlepoint import SinglePointCalculator
from ase.constraints import FixAtoms, FixSymmetry
from ase.filters import FrechetCellFilter
from ase.optimize import BFGS, FIRE, LBFGS, CellAwareBFGS
import numpy as np
[docs]
@dataclass(frozen=True, eq=True)
class Relax:
"""Minimize energy with respect to internal positions.
Also used as a base class for all other relaxation classes."""
max_steps: int = 100
force_tolerance: float = 1e-3
algorithm: Literal["LBFGS", "BFGS", "FIRE"] = "LBFGS"
[docs]
def apply_filter_and_constraints(self, structure: Atoms):
"""Hook to allow subclasses to add filters and constraints."""
return structure
[docs]
def relax(self, structure: Atoms) -> Atoms:
"""Relax a structure and return result.
Structure must have a calculator attached.
Returned structure will have a SinglePointCalculator with the final energy, forces and stresses attached.
Args:
structure (:class:`ase.Atoms`): structure to relax
Returns:
:class:`ase.Atoms`: relaxed structure with attached single point calculator.
"""
calc = structure.calc
structure = structure.copy()
update_uuid(structure)
structure.calc = calc
optimizer_cls = {"LBFGS": LBFGS, "BFGS": BFGS, "FIRE": FIRE}[self.algorithm]
optimizer = optimizer_cls(self.apply_filter_and_constraints(structure), logfile="/dev/null")
optimizer.run(fmax=self.force_tolerance, steps=self.max_steps)
structure.calc = None
structure.calc = SinglePointCalculator(
structure,
energy=calc.get_potential_energy(),
forces=calc.get_forces(),
stress=calc.get_stress(),
)
structure.constraints.clear()
return structure
[docs]
@dataclass(frozen=True, eq=True)
class CellRelax(Relax):
"""Minimize energy while keeping relative positions and volume constant."""
[docs]
def apply_filter_and_constraints(self, structure: Atoms):
structure.set_constraint(FixAtoms(np.ones(len(structure), dtype=bool)))
return FrechetCellFilter(structure, constant_volume=True)
[docs]
@dataclass(frozen=True, eq=True)
class VolumeRelax(Relax):
"""Minimize energy while keeping relative positions and cell shape constant."""
pressure: float = 0.0
[docs]
def apply_filter_and_constraints(self, structure: Atoms):
structure.set_constraint(FixAtoms(np.ones(len(structure), dtype=bool)))
return FrechetCellFilter(
structure, hydrostatic_strain=True, scalar_pressure=self.pressure
)
[docs]
@dataclass(frozen=True, eq=True)
class SymmetryRelax(Relax):
"""Minimize energy with respect to internal positions and cell, while keeping space group fixed."""
pressure: float = 0.0
[docs]
def apply_filter_and_constraints(self, structure: Atoms):
structure.set_constraint(FixSymmetry(structure))
return FrechetCellFilter(structure, scalar_pressure=self.pressure)
[docs]
@dataclass(frozen=True, eq=True)
class FullRelax(Relax):
"""Minimize energy with respect to internal positions and cell without constraints."""
pressure: float = 0.0
[docs]
def apply_filter_and_constraints(self, structure: Atoms):
return FrechetCellFilter(structure, scalar_pressure=self.pressure)
[docs]
def relax(
structures: Iterable[Atoms],
settings: Relax,
calculator: AseCalculatorConfig | Calculator,
) -> Iterator[Atoms]:
"""Relax structures according the given relaxation settings.
Output structures have the final energy and force attached as ase's SinglePointCalculator.
Args:
structures (:class:`collections.abc.Iterable` of :class:`ase.Atoms`): the structures to minimize
settings (:class:`.Relax`): the kind of relaxation to perform (position, volume, etc.)
calculator (:class:`.AseCalculatorConfig` or :class:`ase.calculators.calculator.Calculator`): the energy/force engine to use
Yields:
:class:`ase.Atoms`: the corresponding relaxed configuration to each input structure
"""
for s in structures:
s = s.copy()
if isinstance(calculator, AseCalculatorConfig):
s.calc = calculator.get_calculator()
else:
s.calc = calculator
yield settings.relax(s)
__all__ = [
"Relax",
"CellRelax",
"VolumeRelax",
"SymmetryRelax",
"FullRelax",
"relax",
]