ase-calculator-kit 0.3.2__py3-none-any.whl

ase_calculator_kit/__init__.py

@@ -0,0 +1,50 @@
+"""ase-calculator-kit: a unified ASE calculator factory.
+
+Call supported MLIP and DFT ASE calculators from one public factory::
+
+    from ase_calculator_kit import get_calculator
+    atoms.calc = get_calculator("uma", task="omat")
+    atoms.calc = get_calculator("vasp", config="examples/dft/vasp_pbe_static.yaml")
+    energy = atoms.get_potential_energy()
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .config import resolve_calculator_config
+from .dispersion import DispersionPolicy, get_dispersion_policy
+from .errors import CalculatorKitError, DispersionError, MissingDependencyError
+from .factory import (
+    attach_calculator,
+    available_calculators,
+    available_dft_calculators,
+    available_mlip_models,
+    available_models,
+    get_calculator,
+    get_dft_calculator,
+    get_mlip_calculator,
+)
+
+try:
+    __version__ = version("ase-calculator-kit")
+except PackageNotFoundError:  # not installed (e.g. running from a source tree)
+    __version__ = "0.0.0"
+
+__all__ = [
+    "get_calculator",
+    "get_mlip_calculator",
+    "get_dft_calculator",
+    "attach_calculator",
+    "available_calculators",
+    "available_models",
+    "available_mlip_models",
+    "available_dft_calculators",
+    "resolve_calculator_config",
+    "DispersionPolicy",
+    "get_dispersion_policy",
+    "CalculatorKitError",
+    "MissingDependencyError",
+    "DispersionError",
+    "__version__",
+]

ase_calculator_kit/backends/__init__.py

@@ -0,0 +1,24 @@
+"""Calculator backends."""
+
+from __future__ import annotations
+
+from .base import BaseBackend
+from .dft import EspressoBackend, VaspBackend
+from .mlip import (
+    CHGNetBackend,
+    FairChemBackend,
+    MatterSimBackend,
+    NequIPBackend,
+    SevenNetBackend,
+)
+
+__all__ = [
+    "BaseBackend",
+    "CHGNetBackend",
+    "EspressoBackend",
+    "FairChemBackend",
+    "MatterSimBackend",
+    "NequIPBackend",
+    "SevenNetBackend",
+    "VaspBackend",
+]

ase_calculator_kit/backends/base.py

@@ -0,0 +1,25 @@
+"""Backend base class."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from ase.calculators.calculator import Calculator
+
+
+class BaseBackend(ABC):
+    """Abstract base for a calculator backend.
+
+    Subclasses lazily import their underlying package inside
+    :meth:`create_calculator` and raise
+    :class:`ase_calculator_kit.errors.MissingDependencyError` if it is not
+    installed.
+    """
+
+    #: Canonical backend name (e.g. ``"chgnet"`` or ``"vasp"``).
+    name: str
+
+    @abstractmethod
+    def create_calculator(self, **kwargs) -> Calculator:
+        """Build and return a fresh ASE calculator."""
+        raise NotImplementedError

ase_calculator_kit/backends/dft/__init__.py

@@ -0,0 +1,11 @@
+"""External DFT calculator backends."""
+
+from __future__ import annotations
+
+from .espresso import EspressoBackend
+from .vasp import VaspBackend
+
+__all__ = [
+    "EspressoBackend",
+    "VaspBackend",
+]

ase_calculator_kit/backends/dft/espresso.py

@@ -0,0 +1,56 @@
+"""Quantum ESPRESSO backend using ASE's Espresso calculator."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from ase.calculators.calculator import Calculator
+
+from ...config import resolve_calculator_config, write_resolved_config_file
+from ..base import BaseBackend
+
+
+class EspressoBackend(BaseBackend):
+    name = "qe"
+
+    def create_calculator(
+        self,
+        *,
+        config: str | Path | dict[str, Any],
+        overrides: dict[str, Any] | None = None,
+        write_resolved_config: bool = False,
+    ) -> Calculator:
+        """Create an ASE :class:`ase.calculators.espresso.Espresso` from config."""
+        from ase.calculators.espresso import Espresso, EspressoProfile
+
+        resolved = resolve_calculator_config(
+            "qe",
+            config=config,
+            overrides=overrides,
+        )
+        profile_cfg = resolved.get("profile", {})
+        if "command" not in profile_cfg:
+            raise ValueError("QE config requires profile.command.")
+        if "pseudo_dir" not in profile_cfg:
+            raise ValueError("QE config requires profile.pseudo_dir.")
+
+        pseudopotentials = resolved.get("pseudopotentials")
+        if not pseudopotentials:
+            raise ValueError("QE config requires pseudopotentials.")
+
+        parameters = resolved.get("parameters", {})
+        directory = resolved.get("directory", ".")
+        if write_resolved_config:
+            write_resolved_config_file(resolved, directory)
+
+        profile = EspressoProfile(
+            command=profile_cfg["command"],
+            pseudo_dir=profile_cfg["pseudo_dir"],
+        )
+        return Espresso(
+            profile=profile,
+            directory=directory,
+            pseudopotentials=pseudopotentials,
+            **parameters,
+        )

ase_calculator_kit/backends/dft/vasp.py

@@ -0,0 +1,46 @@
+"""VASP backend using ASE's Vasp calculator."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from ase.calculators.calculator import Calculator
+
+from ...config import resolve_calculator_config, write_resolved_config_file
+from ..base import BaseBackend
+
+
+class VaspBackend(BaseBackend):
+    name = "vasp"
+
+    def create_calculator(
+        self,
+        *,
+        config: str | Path | dict[str, Any],
+        overrides: dict[str, Any] | None = None,
+        write_resolved_config: bool = False,
+    ) -> Calculator:
+        """Create an ASE :class:`ase.calculators.vasp.Vasp` from config."""
+        from ase.calculators.vasp import Vasp
+
+        resolved = resolve_calculator_config(
+            "vasp",
+            config=config,
+            overrides=overrides,
+        )
+        profile = resolved.get("profile", {})
+        if "command" not in profile:
+            raise ValueError("VASP config requires profile.command.")
+
+        parameters = resolved.get("parameters", {})
+        directory = resolved.get("directory", ".")
+        if write_resolved_config:
+            write_resolved_config_file(resolved, directory)
+
+        return Vasp(
+            command=profile["command"],
+            directory=directory,
+            txt=profile.get("txt", "vasp.out"),
+            **parameters,
+        )

ase_calculator_kit/backends/mlip/__init__.py

@@ -0,0 +1,17 @@
+"""Machine-learning interatomic potential backends."""
+
+from __future__ import annotations
+
+from .chgnet import CHGNetBackend
+from .fairchem import FairChemBackend
+from .mattersim import MatterSimBackend
+from .nequip import NequIPBackend
+from .sevennet import SevenNetBackend
+
+__all__ = [
+    "CHGNetBackend",
+    "FairChemBackend",
+    "MatterSimBackend",
+    "NequIPBackend",
+    "SevenNetBackend",
+]

ase_calculator_kit/backends/mlip/chgnet.py

@@ -0,0 +1,73 @@
+"""CHGNet backend (https://github.com/CederGroupHub/chgnet)."""
+
+from __future__ import annotations
+
+from ase.calculators.calculator import Calculator
+
+from ...device import resolve_device
+from ...dispersion import precheck_dispersion_xc, wrap_with_d3
+from ...errors import MissingDependencyError
+from ..base import BaseBackend
+
+
+class CHGNetBackend(BaseBackend):
+    name = "chgnet"
+
+    def create_calculator(
+        self,
+        *,
+        device: str = "auto",
+        model: str | None = None,
+        checkpoint: str | None = None,
+        dispersion: bool = False,
+        dispersion_xc: str | None = None,
+        **kwargs,
+    ) -> Calculator:
+        """Create a :class:`chgnet.model.dynamics.CHGNetCalculator`.
+
+        CHGNet is best for inorganic crystalline materials and fast
+        pre-relaxation. Be careful with isolated atoms and chemistry far from
+        its training domain.
+
+        Parameters
+        ----------
+        device:
+            ``"auto"`` (cuda > mps > cpu), or ``"cuda"`` / ``"mps"`` / ``"cpu"``.
+            CHGNet supports Apple Silicon ``"mps"``.
+        model:
+            Optional pretrained model name passed to ``CHGNet.load(model_name=...)``.
+            When omitted, CHGNet's bundled default model is used.
+        checkpoint:
+            Optional path to a ``.pth`` checkpoint, loaded via
+            ``CHGNetCalculator.from_file``.
+        dispersion, dispersion_xc:
+            Add a Grimme-D3(BJ) correction (CHGNet is PBE+U, so ``xc="pbe"`` by
+            default). See ``docs/models.md`` for the per-model policy.
+        """
+        # Validate the dispersion policy before loading the model (fail fast).
+        d3_xc = precheck_dispersion_xc(
+            self.name, model or "default",
+            dispersion=dispersion, dispersion_xc=dispersion_xc,
+        )
+        use_device = resolve_device(device, allow_mps=True)
+
+        try:
+            from chgnet.model.dynamics import CHGNetCalculator
+        except ImportError as exc:  # pragma: no cover - exercised via tests with mocks
+            raise MissingDependencyError("CHGNet") from exc
+
+        if checkpoint is not None:
+            bare = CHGNetCalculator.from_file(
+                checkpoint, use_device=use_device, **kwargs
+            )
+        elif model is not None:
+            from chgnet.model.model import CHGNet
+
+            loaded = CHGNet.load(model_name=model)
+            bare = CHGNetCalculator(model=loaded, use_device=use_device, **kwargs)
+        else:
+            bare = CHGNetCalculator(use_device=use_device, **kwargs)
+
+        if d3_xc is not None:
+            return wrap_with_d3(bare, xc=d3_xc, device=use_device)
+        return bare

ase_calculator_kit/backends/mlip/fairchem.py

@@ -0,0 +1,95 @@
+"""fairchem / UMA backend (https://github.com/facebookresearch/fairchem)."""
+
+from __future__ import annotations
+
+from ase.calculators.calculator import Calculator
+
+from ...device import resolve_device
+from ...dispersion import precheck_dispersion_xc, wrap_with_d3
+from ...errors import MissingDependencyError
+from ..base import BaseBackend
+
+
+class FairChemBackend(BaseBackend):
+    name = "uma"
+
+    def create_calculator(
+        self,
+        *,
+        device: str = "auto",
+        model: str = "uma-s-1p2",
+        task: str = "omat",
+        dispersion: bool = False,
+        dispersion_xc: str | None = None,
+        **kwargs,
+    ) -> Calculator:
+        """Create a :class:`fairchem.core.FAIRChemCalculator` for a UMA model.
+
+        UMA checkpoints are gated on Hugging Face. If creation fails with an
+        authorization error, request access to the model repository and run
+        ``huggingface-cli login``.
+
+        Parameters
+        ----------
+        device:
+            ``"auto"`` (cuda > cpu) or explicit ``"cuda"`` / ``"cpu"``. Apple
+            Silicon ``"mps"`` is not supported: fairchem-core's predict unit
+            asserts ``device in {"cpu", "cuda"}``, so ``"mps"`` is rejected
+            before this wrapper runs.
+        model:
+            UMA model name. Defaults to ``"uma-s-1p2"``.
+        task:
+            The ``task_name`` selecting the domain-specific head. A single UMA
+            model serves many domains; pick the task matching your system:
+
+            ======= ==================================================
+            ``task`` Use for
+            ======= ==================================================
+            ``omat`` Inorganic bulk/materials, stress, cell optimization
+            ``omol`` Molecules and polymers
+            ``oc20`` Catalyst surfaces and adsorption
+            ``oc22`` Oxide catalysis
+            ``oc25`` Electrochemistry / solid-liquid interfaces
+            ``odac`` MOFs and direct air capture
+            ``omc``  Molecular crystals
+            ======= ==================================================
+
+            For the molecular task (``omol``), set ``atoms.info["charge"]`` (total
+            charge) and ``atoms.info["spin"]`` (spin multiplicity, ``2S+1``)
+            *before* computing::
+
+                atoms.info["charge"] = -1
+                atoms.info["spin"] = 2
+                atoms.calc = get_calculator("uma", task="omol")
+
+            This is not optional in practice, only in form: fairchem does **not**
+            raise when they are missing. It logs a warning, writes
+            ``charge=0`` / ``spin=1`` into ``atoms.info`` (mutating the object you
+            passed in), and returns a neutral closed-shell result. An anion or a
+            radical therefore comes back silently wrong unless both keys are set.
+            ``charge`` and ``spin`` are read only by the ``omol`` head; other
+            tasks ignore them.
+        dispersion, dispersion_xc:
+            Add a Grimme-D3(BJ) correction. The D3 ``xc`` depends on the task's
+            DFT level (e.g. ``omat``→pbe, ``oc20``→rpbe). Rejected for the tasks
+            whose reference data already accounts for dispersion: ``oc25``
+            (RPBE+D3), ``omol`` (ωB97M-V nonlocal), ``odac`` (PBE-D3) and ``omc``
+            (PBE+D3). See ``docs/models.md``.
+        """
+        # Validate the dispersion policy before loading the model (fail fast).
+        d3_xc = precheck_dispersion_xc(
+            self.name, task, dispersion=dispersion, dispersion_xc=dispersion_xc,
+        )
+        resolved_device = resolve_device(device)
+
+        try:
+            from fairchem.core import FAIRChemCalculator, pretrained_mlip
+        except ImportError as exc:  # pragma: no cover - exercised via tests with mocks
+            raise MissingDependencyError("fairchem-core") from exc
+
+        predictor = pretrained_mlip.get_predict_unit(model, device=resolved_device)
+        bare = FAIRChemCalculator(predictor, task_name=task, **kwargs)
+
+        if d3_xc is not None:
+            return wrap_with_d3(bare, xc=d3_xc, device=resolved_device)
+        return bare

ase_calculator_kit/backends/mlip/mattersim.py

@@ -0,0 +1,86 @@
+"""MatterSim backend (https://github.com/microsoft/mattersim)."""
+
+from __future__ import annotations
+
+from ase.calculators.calculator import Calculator
+
+from ...device import resolve_device
+from ...dispersion import precheck_dispersion_xc, wrap_with_d3
+from ...errors import MissingDependencyError
+from ..base import BaseBackend
+
+#: Friendly model keys -> checkpoint file names shipped with MatterSim.
+_MODEL_TO_LOAD_PATH = {
+    "1M": "MatterSim-v1.0.0-1M.pth",
+    "5M": "MatterSim-v1.0.0-5M.pth",
+    "MatterSim-v1.0.0-1M": "MatterSim-v1.0.0-1M.pth",
+    "MatterSim-v1.0.0-5M": "MatterSim-v1.0.0-5M.pth",
+}
+_DEFAULT_MODELS = {"1M", "MatterSim-v1.0.0-1M"}
+
+
+class MatterSimBackend(BaseBackend):
+    name = "mattersim"
+
+    def create_calculator(
+        self,
+        *,
+        device: str = "auto",
+        model: str = "1M",
+        load_path: str | None = None,
+        dispersion: bool = False,
+        dispersion_xc: str | None = None,
+        **kwargs,
+    ) -> Calculator:
+        """Create a :class:`mattersim.forcefield.MatterSimCalculator`.
+
+        Parameters
+        ----------
+        device:
+            ``"auto"`` (cuda > mps > cpu), or explicit ``"cuda"`` / ``"mps"`` /
+            ``"cpu"``. MatterSim supports Apple Silicon ``"mps"``.
+        model:
+            ``"1M"`` (default, fast screening) or ``"5M"`` (more accurate).
+            Keep the checkpoint fixed across a campaign.
+        load_path:
+            Explicit checkpoint path; overrides ``model``. When ``model="1M"``
+            and no ``load_path`` is given, MatterSim's default 1M checkpoint is
+            used (no ``load_path`` passed).
+        dispersion, dispersion_xc:
+            Add a Grimme-D3(BJ) correction (MatterSim is PBE, so ``xc="pbe"`` by
+            default). See ``docs/models.md``.
+        """
+        if model in {"1M", "MatterSim-v1.0.0-1M"}:
+            key = "1M"
+        elif model in {"5M", "MatterSim-v1.0.0-5M"}:
+            key = "5M"
+        else:
+            key = "default"
+        # Validate the dispersion policy before loading the model (fail fast).
+        d3_xc = precheck_dispersion_xc(
+            self.name, key, dispersion=dispersion, dispersion_xc=dispersion_xc,
+        )
+        resolved_device = resolve_device(device, allow_mps=True)
+
+        try:
+            from mattersim.forcefield import MatterSimCalculator
+        except ImportError as exc:  # pragma: no cover - exercised via tests with mocks
+            raise MissingDependencyError("MatterSim") from exc
+
+        params: dict = {"device": resolved_device}
+
+        resolved = load_path or _MODEL_TO_LOAD_PATH.get(model)
+        if load_path is not None or model not in _DEFAULT_MODELS:
+            if resolved is None:
+                raise ValueError(
+                    f"Unknown MatterSim model '{model}'. Use '1M', '5M', or pass "
+                    "an explicit load_path."
+                )
+            params["load_path"] = resolved
+
+        params.update(kwargs)
+        bare = MatterSimCalculator(**params)
+
+        if d3_xc is not None:
+            return wrap_with_d3(bare, xc=d3_xc, device=resolved_device)
+        return bare

ase_calculator_kit/backends/mlip/nequip.py

@@ -0,0 +1,110 @@
+"""NequIP OAM backend (https://www.nequip.net/)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from ase.calculators.calculator import Calculator
+
+from ...device import resolve_device
+from ...dispersion import precheck_dispersion_xc, wrap_with_d3
+from ...errors import MissingDependencyError
+from ..base import BaseBackend
+
+_MODEL_IDS = {
+    "S": "mir-group/NequIP-OAM-S:0.1",
+    "M": "mir-group/NequIP-OAM-M:0.1",
+    "L": "mir-group/NequIP-OAM-L:0.1",
+    "XL": "mir-group/NequIP-OAM-XL:0.1",
+}
+
+
+def _normalized_model_size(model: str) -> str:
+    normalized = model.upper()
+    if normalized not in _MODEL_IDS:
+        valid = ", ".join(_MODEL_IDS)
+        raise ValueError(f"Unknown NequIP OAM model '{model}'. Use one of: {valid}.")
+    return normalized
+
+
+def _load_target(model: str, model_path: str | Path | None) -> str:
+    """Return either an explicit local model or NequIP's pinned OAM identifier."""
+    return str(model_path) if model_path is not None else f"nequip.net:{_MODEL_IDS[model]}"
+
+
+class NequIPBackend(BaseBackend):
+    name = "nequip"
+
+    def create_calculator(
+        self,
+        *,
+        device: str = "auto",
+        model: str = "L",
+        model_path: str | Path | None = None,
+        chemical_species_to_atom_type_map: dict[str, str] | bool | None = True,
+        compile_mode: str = "eager",
+        model_name: str = "sole_model",
+        neighborlist_backend: str = "matscipy",
+        allow_tf32: bool = False,
+        dispersion: bool = False,
+        dispersion_xc: str | None = None,
+        **kwargs,
+    ) -> Calculator:
+        """Create a NequIP ASE calculator for one of the OAM foundation models.
+
+        Parameters
+        ----------
+        device:
+            ``"auto"`` (cuda > cpu) or explicit ``"cuda"`` / ``"cpu"``. Apple
+            Silicon ``"mps"`` is intentionally not enabled for NequIP OAM:
+            local testing fails with ``Cannot convert a MPS Tensor to float64``
+            because the packaged OAM models use float64 buffers and PyTorch MPS
+            does not support float64.
+        model:
+            OAM model size: ``"S"``, ``"M"``, ``"L"`` (default), or ``"XL"``.
+            Case-insensitive.
+        model_path:
+            Optional local ``.nequip.zip`` or checkpoint path. When omitted, the
+            selected OAM model is loaded through NequIP's ``nequip.net:`` loader
+            and cached by NequIP.
+        chemical_species_to_atom_type_map:
+            Passed to NequIP's ASE integration. ``True`` selects the identity
+            map and silences the upstream fallback warning, which is correct for
+            OAM model type names.
+        compile_mode, model_name, neighborlist_backend, allow_tf32:
+            Forwarded to NequIP's saved-model ASE loader.
+        dispersion, dispersion_xc:
+            Optionally add Grimme-D3(BJ). OAM uses PBE-level reference data, so
+            the default D3 parameters are PBE; an explicit ``dispersion_xc``
+            overrides that verified default.
+        """
+        normalized_model = _normalized_model_size(model)
+
+        # Validate the dispersion policy before loading the model (fail fast).
+        d3_xc = precheck_dispersion_xc(
+            self.name,
+            normalized_model,
+            dispersion=dispersion,
+            dispersion_xc=dispersion_xc,
+        )
+        resolved_device = resolve_device(device)
+
+        try:
+            from nequip.integrations.ase import NequIPCalculator
+        except ImportError as exc:  # pragma: no cover - exercised via tests with mocks
+            raise MissingDependencyError("NequIP") from exc
+
+        bare = NequIPCalculator._from_saved_model(
+            _load_target(normalized_model, model_path),
+            device=resolved_device,
+            chemical_species_to_atom_type_map=chemical_species_to_atom_type_map,
+            allow_tf32=allow_tf32,
+            model_name=model_name,
+            compile_mode=compile_mode,
+            neighborlist_backend=neighborlist_backend,
+            **kwargs,
+        )
+
+        if d3_xc is not None:
+            return wrap_with_d3(bare, xc=d3_xc, device=resolved_device)
+        return bare