Source code for directupsampling.calculators.effqh

from ase.atoms import Atoms
from ase.calculators.calculator import Calculator
from hiphive import ForceConstants
from hiphive.calculators import ForceConstantCalculator

from directupsampling.effqh.utils import get_ref_atoms
from directupsampling.effqh.fcfit import ForceConstantFit




class EffQHCalculator(Calculator):
    """Quasiharmonic ForceConstantCalculator.

    Implemented as a wrapper around the hiphive ForceConstantCalculator.
    Uses a caching over volume for instantiated ForceConstantCalculators,
    which provides acceptable speed when the volume is constant over an MD
    run, which is the present typical use case. It is not suitable for e.g.
    NPT MD, where the volume changes at every step.

    Parameters
    ----------
    ideal_atoms: atoms object whose positions represent the ideal structure
        used as reference for displacements.

    fc_fit: force constant-fit coefficients.

    max_disp: Maximum displacement passed to hiphive ForceConstantCalculator.
    """

    implemented_properties = [
        "energy",
        "forces",
    ]

    def __init__(
        self,
        reference_atoms: Atoms | list[Atoms] = None,
        fc_fit: ForceConstantFit = None,
        max_disp: float = 3.0,
    ):
        super().__init__()
        self.reference_atoms = (
            Atoms(reference_atoms)
            if hasattr(reference_atoms, "get_potential_energy")
            else [Atoms(_) for _ in reference_atoms]
        )
        self.fc_fit = fc_fit
        self.max_disp = max_disp
        self._fc_calc_cache = {}



    def calculate(self, atoms=None, properties=["energy"], *args):
        super().calculate(atoms, properties, *args)
        for prop in properties:
            volume = self.atoms.get_volume() / len(self.atoms)
            if not volume in self._fc_calc_cache:
                force_constants = self.fc_fit(volume)
                reference_atoms = get_ref_atoms(self.reference_atoms, volume)
                fcs = ForceConstants.from_arrays(reference_atoms, force_constants)
                fc_calc = ForceConstantCalculator(fcs)
                self._fc_calc_cache[volume] = fc_calc
            else:
                fc_calc = self._fc_calc_cache[volume]

            result = fc_calc.get_property(prop, self.atoms)
            self.results[prop] = result