diff --git a/src/matcalc/_neb.py b/src/matcalc/_neb.py index 3d0ae5e..9219e87 100644 --- a/src/matcalc/_neb.py +++ b/src/matcalc/_neb.py @@ -8,11 +8,7 @@ import numpy as np from ase.io import Trajectory - -try: - from ase.mep import NEB, NEBTools -except ImportError: - from ase.neb import NEB, NEBTools # type: ignore[no-redef] +from ase.mep import NEB, NEBTools from ase.utils.forcecurve import fit_images from pymatgen.core import Lattice, Structure from pymatgen.core.periodic_table import Species diff --git a/src/matcalc/utils.py b/src/matcalc/utils.py index 5757033..7ea6f68 100644 --- a/src/matcalc/utils.py +++ b/src/matcalc/utils.py @@ -32,7 +32,6 @@ "SevenNet", "TensorNet", "GRACE", - "TensorPotential", "ORB", "PBE", "r2SCAN", @@ -53,16 +52,33 @@ except Exception: # noqa: BLE001 warnings.warn("Unable to get pre-trained MatGL universal calculators.", stacklevel=1) -# Provide simple aliases for some common models. The key in MODEL_ALIASES must be lower case. -MODEL_ALIASES = { - "tensornet": "TensorNet-MatPES-PBE-v2025.1-PES", - "m3gnet": "M3GNet-MatPES-PBE-v2025.1-PES", - "chgnet": "CHGNet-MatPES-PBE-2025.2.10-2.7M-PES", - "pbe": "TensorNet-MatPES-PBE-v2025.1-PES", - "r2scan": "TensorNet-MatPES-r2SCAN-v2025.1-PES", + +# Common aliases and abbreviations will load the most advanced or widely used model. +ALIAS_TO_ID = { + ("tensornet", "tensornet-pbe", "pbe"): "TensorNet-MatPES-PBE-v2025.1-S", + ("tensornet-r2scan", "r2scan"): "TensorNet-MatPES-r2SCAN-v2025.1-S", + ("m3gnet", "m3gnet-pbe"): "M3GNet-MatPES-PBE-v2025.1-S", + ("m3gnet-r2scan",): "M3GNet-MatPES-r2SCAN-v2025.1-S", + ("chgnet", "chgnet-pbe"): "CHGNet-MatPES-PBE-v2025.2-M", + ("chgnet-r2scan",): "CHGNet-MatPES-r2SCAN-v2025.2-M", + ("mace-mp-0", "mace-mp-0-m", "mace-mp-pbe-0"): "MACE-MP-PBE-0-M", + ("mace", "mace-mpa-0", "mace-mpa-0-m", "mace-mpa-pbe-0"): "MACE-MPA-PBE-0-M", + ("mace-omat-0", "mace-omat-0-m", "mace-omat-pbe-0"): "MACE-OMAT-PBE-0-M", + ("mace-matpes-pbe", "mace-matpes-pbe-0"): "MACE-MatPES-PBE-0-M", + ("mace-matpes-r2scan", "mace-matpes-r2scan-0"): "MACE-MatPES-r2SCAN-0-M", + ("grace-mp",): "GRACE-MP-PBE-0-L", + ("grace", "grace-oam"): "GRACE-OAM-PBE-0-L", + ("grace-omat",): "GRACE-OMAT-PBE-0-L", +} + +ALIAS_HASH_TABLE: dict[str, str] = { + alias.lower(): canonical for aliases, canonical in ALIAS_TO_ID.items() for alias in aliases } +MODEL_ID_PARTS = 5 # Architecture-Dataset-Functional-Version-Size + + UNIVERSAL_CALCULATORS = Enum("UNIVERSAL_CALCULATORS", {k: k for k in _universal_calculators}) # type: ignore[misc] @@ -134,7 +150,6 @@ def calculate( """ from ase.calculators.calculator import all_changes, all_properties from maml.apps.pes import EnergyForceStress - from pymatgen.io.ase import AseAtomsAdaptor properties = properties or all_properties system_changes = system_changes or all_changes @@ -366,45 +381,48 @@ def load_universal(name: str | Calculator, **kwargs: Any) -> Calculator: # noqa """ result: Calculator - if not isinstance(name, str): # e.g. already an ase Calculator instance - result = name + if not isinstance(name, str): + return name - elif any(name.lower().startswith(m) for m in ("m3gnet", "chgnet", "tensornet", "pbe", "r2scan")): - name = MODEL_ALIASES.get(name.lower(), name) - result = PESCalculator.load_matgl(name, **kwargs) + model_id = _resolve_model(name) + backend = _infer_backend(model_id) - elif name.lower() == "mace": + if backend == "matgl": + model_name = _parse_matgl_model_id(model_id) + result = PESCalculator.load_matgl(model_name, **kwargs) + + elif backend == "mace": from mace.calculators import mace_mp - result = mace_mp(**kwargs) + model_name = _parse_mace_model_id(model_id) + result = mace_mp(model=model_name, **kwargs) - elif name.lower() == "sevennet": - from sevenn.calculator import SevenNetCalculator + elif backend == "grace": + from tensorpotential.calculator.foundation_models import grace_fm - result = SevenNetCalculator(**kwargs) + model_name = _parse_grace_model_id(model_id) + result = grace_fm(model=model_name, **kwargs) - elif name.lower() == "grace" or name.lower() == "tensorpotential": - from tensorpotential.calculator.foundation_models import grace_fm + elif backend == "sevennet": + from sevenn.calculator import SevenNetCalculator - kwargs.setdefault("model", "GRACE-2L-OAM") - result = grace_fm(**kwargs) + result = SevenNetCalculator(**kwargs) - elif name.lower() == "orb": + elif backend == "orb": from orb_models.forcefield.calculator import ORBCalculator from orb_models.forcefield.pretrained import ORB_PRETRAINED_MODELS model = kwargs.pop("model", "orb-v2") device = kwargs.get("device", "cpu") - orbff = ORB_PRETRAINED_MODELS[model](device=device) result = ORBCalculator(orbff, **kwargs) - elif name.lower() == "mattersim": # pragma: no cover + elif backend == "mattersim": # pragma: no cover from mattersim.forcefield import MatterSimCalculator result = MatterSimCalculator(**kwargs) - elif name.lower() == "fairchem": # pragma: no cover + elif backend == "fairchem": # pragma: no cover from fairchem.core import FAIRChemCalculator, pretrained_mlip device = kwargs.pop("device", "cpu") @@ -413,12 +431,12 @@ def load_universal(name: str | Calculator, **kwargs: Any) -> Calculator: # noqa predictor = pretrained_mlip.get_predict_unit(model, device=device) result = FAIRChemCalculator(predictor, task_name=task_name, **kwargs) - elif name.lower() == "petmad": # pragma: no cover + elif backend == "petmad": # pragma: no cover from pet_mad.calculator import PETMADCalculator result = PETMADCalculator(**kwargs) - elif name.lower().startswith("deepmd"): # pragma: no cover + elif backend == "deepmd": # pragma: no cover from pathlib import Path from deepmd.calculator import DP @@ -487,3 +505,91 @@ def to_pmg_molecule(structure: Atoms | Structure | Molecule | IMolecule) -> IMol structure = AseAtomsAdaptor.get_molecule(structure) return Molecule.from_sites(structure) # type: ignore[return-value] + + +def _resolve_model(name: str) -> str: + key = name.lower() + + return ALIAS_HASH_TABLE.get(key, name) + + +def _infer_backend(model_id: str) -> str: + model_id = model_id.lower() + backend = "unknown" + + if model_id.startswith(("tensornet", "m3gnet", "chgnet")): + backend = "matgl" + elif model_id.startswith("mace"): + backend = "mace" + elif model_id.startswith("grace"): + backend = "grace" + elif model_id == "sevennet": + backend = "sevennet" + elif model_id == "orb": + backend = "orb" + elif model_id == "mattersim": + backend = "mattersim" + elif model_id == "fairchem": + backend = "fairchem" + elif model_id == "petmad": + backend = "petmad" + elif model_id.startswith("deepmd"): + backend = "deepmd" + + return backend + + +def _parse_mace_model_id(model_id: str) -> str: + parts = model_id.split("-") + result = model_id + + if len(parts) == MODEL_ID_PARTS and parts[0] == "MACE": + _, dataset, functional, version, size = parts + + size_map = {"S": "small", "M": "medium", "L": "large"} + mapped_size = size_map.get(size.upper()) + + if mapped_size is not None: + version = version.lower() + functional = functional.lower() + + if dataset == "MP": + result = mapped_size if version == "0" else f"{mapped_size}-{version}" + elif dataset == "MPA": + result = f"{mapped_size}-mpa-{version}" + elif dataset == "OMAT": + result = f"{mapped_size}-omat-{version}" + elif dataset == "MatPES": + result = f"mace-matpes-{functional}-{version}" + + return result + + +def _parse_matgl_model_id(model_id: str) -> str: + parts = model_id.split("-") + if len(parts) == MODEL_ID_PARTS and parts[0] in {"TensorNet", "M3GNet", "CHGNet"} and parts[1] == "MatPES": + return "-".join([*parts[:-1], "PES"]) + + return model_id + + +def _parse_grace_model_id(model_id: str) -> str: + parts = model_id.split("-") + result = model_id + + if len(parts) == MODEL_ID_PARTS and parts[0] == "GRACE": + _, dataset, functional, version, size = parts + size = size.upper() + + if functional == "PBE" and version == "0": + if dataset == "MP": + if size == "L": + result = "GRACE-2L-MP-r6" + elif dataset in {"OMAT", "OAM"}: + if size == "S": + result = "GRACE-2L-OMAT" if dataset == "OMAT" else "GRACE-2L-OAM" + elif size in {"M", "L"}: + suffix = "E" if dataset == "OMAT" else "AM" + result = f"GRACE-2L-OMAT-medium-ft-{suffix}" if size == "M" else f"GRACE-2L-OMAT-large-ft-{suffix}" + + return result diff --git a/tests/test_utils.py b/tests/test_utils.py index 2470a7e..02e15ea 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -11,9 +11,13 @@ from matcalc import RelaxCalc from matcalc.utils import ( - MODEL_ALIASES, + ALIAS_TO_ID, UNIVERSAL_CALCULATORS, PESCalculator, + _parse_grace_model_id, + _parse_mace_model_id, + _parse_matgl_model_id, + _resolve_model, ) DIR = Path(__file__).parent.absolute() @@ -176,10 +180,50 @@ def test_pescalculator_calculate() -> None: assert list(stresses.shape) == [6] -def test_aliases() -> None: - # Ensures that model aliases always point to valid models. - names = [u.name for u in UNIVERSAL_CALCULATORS] - for v in MODEL_ALIASES.values(): - # We are not testing DGL based models. - if "M3GNet" not in v and "CHGNet" not in v: - assert v in names +def test_allias_to_id(): + for aliases, model_id in ALIAS_TO_ID.items(): + for alias in aliases: + assert _resolve_model(alias) == model_id + + +ID_TO_NAME = { + "TensorNet-MatPES-PBE-v2025.1-S": "TensorNet-MatPES-PBE-v2025.1-PES", + "TensorNet-MatPES-r2SCAN-v2025.1-S": "TensorNet-MatPES-r2SCAN-v2025.1-PES", + "M3GNet-MatPES-PBE-v2025.1-S": "M3GNet-MatPES-PBE-v2025.1-PES", + "M3GNet-MatPES-r2SCAN-v2025.1-S": "M3GNet-MatPES-r2SCAN-v2025.1-PES", + "MACE-MP-PBE-0-S": "small", + "MACE-MP-PBE-0-M": "medium", + "MACE-MP-PBE-0-L": "large", + "MACE-MP-PBE-0b-S": "small-0b", + "MACE-MP-PBE-0b-M": "medium-0b", + "MACE-MP-PBE-0b2-S": "small-0b2", + "MACE-MP-PBE-0b2-M": "medium-0b2", + "MACE-MP-PBE-0b2-L": "large-0b2", + "MACE-MP-PBE-0b3-M": "medium-0b3", + "MACE-MPA-PBE-0-M": "medium-mpa-0", + "MACE-OMAT-PBE-0-S": "small-omat-0", + "MACE-OMAT-PBE-0-M": "medium-omat-0", + "MACE-MatPES-PBE-0-M": "mace-matpes-pbe-0", + "MACE-MatPES-r2SCAN-0-M": "mace-matpes-r2scan-0", + "GRACE-MP-PBE-0-L": "GRACE-2L-MP-r6", + "GRACE-OAM-PBE-0-S": "GRACE-2L-OAM", + "GRACE-OAM-PBE-0-M": "GRACE-2L-OMAT-medium-ft-AM", + "GRACE-OAM-PBE-0-L": "GRACE-2L-OMAT-large-ft-AM", + "GRACE-OMAT-PBE-0-S": "GRACE-2L-OMAT", + "GRACE-OMAT-PBE-0-M": "GRACE-2L-OMAT-medium-ft-E", + "GRACE-OMAT-PBE-0-L": "GRACE-2L-OMAT-large-ft-E", +} + + +@pytest.mark.parametrize("model_id, expected", ID_TO_NAME.items()) +def test_id_to_name(model_id: str, expected: str) -> None: + if model_id.startswith("MACE"): + parsed = _parse_mace_model_id(model_id) + elif model_id.startswith(("TensorNet", "M3GNet", "CHGNet")): + parsed = _parse_matgl_model_id(model_id) + elif model_id.startswith("GRACE"): + parsed = _parse_grace_model_id(model_id) + else: + pytest.skip(f"No parser for {model_id}") + + assert parsed == expected