Source code for ufp.adapters.ase

"""
ASE calculator integration for UFP potentials.

Use this adapter when an ``ase.Atoms`` workflow should consume UFP energies,
forces, and optional per-atom outputs.
"""

from __future__ import annotations

from typing import Iterable, Optional, Sequence, Union

from ase.calculators.calculator import (
    Calculator,
    PropertyNotImplementedError,
    all_changes,
)

from ufp.core.potential import UFPotential
from ufp.neighbors._neighbors import NeighborListBackend




class UFPASECalculator(Calculator):
    """
    Wrap a UFP potential as an ASE calculator.

    Args:
        potential: Potential to evaluate.
        neighbor_backend: Backend used when the potential builds neighbor lists.
    """

    implemented_properties = [
        "energy",
        "free_energy",
        "forces",
        "stress",
        "energies",
    ]

    def __init__(
        self,
        potential: UFPotential,
        *,
        neighbor_backend: Union[str, NeighborListBackend] = NeighborListBackend.AUTO,
        **kwargs,
    ) -> None:
        """Initialize the ASE calculator with one wrapped UFP potential."""
        super().__init__(**kwargs)
        self.potential = potential
        self.neighbor_backend = NeighborListBackend(neighbor_backend)
        self.parameters["neighbor_backend"] = self.neighbor_backend.value



    def calculate(
        self,
        atoms=None,
        properties: Optional[Sequence[str]] = None,
        system_changes: Optional[Iterable[str]] = all_changes,
    ) -> None:
        """Run the wrapped potential and populate the ASE ``results`` dictionary."""
        super().calculate(atoms, properties, system_changes)

        if atoms is None:
            raise ValueError("`atoms` is required")

        prediction = self.potential.compute(
            atoms=atoms,
            backend=self.neighbor_backend,
            derive_forces=properties is None or "forces" in properties,
        )
        self.results = prediction.as_ase_results(n_atoms=len(atoms))

        if properties is None:
            return

        missing = [name for name in properties if name not in self.results]
        if missing:
            raise PropertyNotImplementedError(
                "the wrapped UFPotential did not provide the requested ASE "
                f"properties: {', '.join(missing)}"
            )