Unify model naming convention

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

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

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