MLIP and MLIPEnum

API reference for the MLIP base class, MLIPEnum, and MLIPMap.

Overview

The mlip_arena.models module exposes three top-level objects that provide a unified interface to all registered machine learning interatomic potential (MLIP) models:

• REGISTRY — raw dict loaded from registry.yaml
• MLIPMap — dict mapping model names to their Python classes
• MLIPEnum — an Enum built from MLIPMap for safe, enumerable model references
• MLIP — the base class all native MLIP models inherit from

---

MLIPEnum

MLIPEnum is a Python Enum whose members are the successfully-imported MLIP model classes. It is built at import time from MLIPMap.

from mlip_arena.models import MLIPEnum

Iterating over all models

for model in MLIPEnum:
    print(model.name, model.value)  # e.g. "MACE-MP(M)", <class 'MACE_MP_Medium'>

Accessing a model by name

model_cls = MLIPEnum["MACE-MP(M)"].value
print(model_cls)  # <class 'mlip_arena.models.externals.mace-mp.MACE_MP_Medium'>

Members

Members are populated at runtime from the registry. Any model whose package is not installed is silently skipped with a warning. The full set of registered models is:

Member name	Class	Family
MACE-MP(M)	MACE_MP_Medium	mace-mp
CHGNet	CHGNet	chgnet
M3GNet	M3GNet	matgl
MatterSim	MatterSim	mattersim
ORBv2	ORBv2	orb
SevenNet	SevenNet	sevennet
eqV2(OMat)	eqV2	fairchem
MACE-MPA	MACE_MPA	mace-mp
eSEN	eSEN	fairchem
EquiformerV2(OC22)	EquiformerV2	equiformer
EquiformerV2(OC20)	EquiformerV2OC20	equiformer
eSCN(OC20)	eSCN	escn
MACE-OFF(M)	MACE_OFF_Medium	mace-off
ANI2x	ANI2x	ani
ALIGNN	ALIGNN	alignn
DeepMD	DeepMD	deepmd
ORB	ORB	orb

Only models whose Python packages are installed in the current environment will appear as members of MLIPEnum. Missing packages produce a warning log and are skipped.

---

MLIPMap

MLIPMap is the plain dict from which MLIPEnum is built. Keys are model name strings; values are the corresponding Python classes.

from mlip_arena.models import MLIPMap

print(MLIPMap.keys())   # dict_keys(['MACE-MP(M)', 'CHGNet', ...])
model_cls = MLIPMap["CHGNet"]

You can use MLIPMap directly when you need a dict interface (e.g. programmatic selection, serialization).

---

MLIP

from mlip_arena.models import MLIP

Inheritance

MLIP inherits from both torch.nn.Module and huggingface_hub.PyTorchModelHubMixin, and is registered with the HuggingFace Hub tags ["atomistic-simulation", "MLIP"].

MLIP
├── torch.nn.Module
└── huggingface_hub.PyTorchModelHubMixin
    └── tags: ["atomistic-simulation", "MLIP"]

Constructor

MLIP(model: nn.Module)

model · torch.nn.Module · required

The underlying PyTorch model to wrap. Stored as self.model.

from_pretrained

Class method inherited from PyTorchModelHubMixin. Downloads and instantiates a model from the HuggingFace Hub or a local path.

mlip = MLIP.from_pretrained("atomind/mace-mp")

pretrained_model_name_or_path · string | Path · required

HuggingFace Hub model ID (e.g. "atomind/mace-mp") or local directory path.

force_download · boolean · default: "false"

Re-download files even if they already exist in the cache.

resume_download · boolean | None · default: "None"

Resume an incomplete download. None uses the hub’s default behaviour.

proxies · dict | None · default: "None"

Dict of proxies for HTTP/HTTPS requests, passed to requests.

token · string | boolean | None · default: "None"

HuggingFace authentication token. Pass True to use the cached token from huggingface-cli login.

cache_dir · string | Path | None · default: "None"

Override the default HuggingFace cache directory.

local_files_only · boolean · default: "false"

If True, only use locally cached files and raise an error if none exist.

revision · string | None · default: "None"

Git revision (branch, tag, or commit hash) to pull from the Hub.

**model_kwargs · dict

Additional keyword arguments forwarded to the model’s __init__.

return · MLIP

An instantiated MLIP subclass loaded from the specified source.

forward

def forward(self, x) -> Any

Delegates to self.model(x). Subclasses override this to handle graph construction and model-specific preprocessing.

x · Any · required

Model input. For MLIPCalculator subclasses this is a batched graph data object created by collate_fn.

return · Any

Raw output from the underlying model. Subclasses typically return a dict with keys energy, forces, and stress.

---

Code examples

Check which models are available

from mlip_arena.models import MLIPEnum

for model in MLIPEnum:
    print(model.name)

Instantiate a model by name

from mlip_arena.models import MLIPEnum

calculator = MLIPEnum["CHGNet"].value()

Conditional dispatch based on model

from mlip_arena.models import MLIPEnum

def make_calculator(name: str):
    if name not in MLIPEnum.__members__:
        raise ValueError(f"Unknown model: {name}")
    return MLIPEnum[name].value()

calc = make_calculator("MACE-MP(M)")