AOTI MLIP

Demo of pytorch AOTInductor for uMLIP.

Installation

From Source

git clone https://github.com/abhijeetgangan/aoti_mlip.git
cd aoti_mlip
pip install .

Quick Start

1. Compile a Model

Compile a pretrained MatterSim model with aoti. This will store the model with .pt2 extension:

from aoti_mlip.utils.aoti_compile import compile_mattersim

# Compile the 1M parameter model for CUDA
package_path = compile_mattersim(
    checkpoint_name="mattersim-v1.0.0-1M.pth",
    cutoff=5.0,
    threebody_cutoff=4.0,
    compute_force=True,
    compute_stress=True,
    device="cuda"
)
# Package saved to: ~/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pt2

The compiled package is device-specific.

2. Run simulation with ASE

Use the compiled model with ASE for fast calculations:

import torch
from ase.build import bulk
from aoti_mlip.calculators.mattersim import MatterSimCalculator

# Create your atomic structure
atoms = bulk("Fe", "bcc", a=2.86, cubic=True).repeat((3, 3, 3))

# Load the compiled model
calc = MatterSimCalculator(
    model_path="~/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pt2",
    device="cuda"
)
atoms.calc = calc

# Calculate properties
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
stress = atoms.get_stress()

3. Batched simulation with TorchSim

Use the compiled model with TorchSim for batched MD and structural relaxation:

import torch
import torch_sim as ts
from ase.build import bulk

from aoti_mlip.calculators.torchsim import MatterSimTorchSimModel

# Load the compiled model as a TorchSim ModelInterface
model = MatterSimTorchSimModel(
    model_path="~/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pt2",
    device="cuda",
)

# Batch-relax multiple structures simultaneously
structures = [
    bulk("Si", "diamond", a=5.43, cubic=True),
    bulk("Cu", "fcc", a=3.61, cubic=True),
    bulk("Fe", "bcc", a=2.86, cubic=True),
]
relaxed = ts.optimize(
    system=structures,
    model=model,
    optimizer=ts.Optimizer.fire,
    convergence_fn=ts.generate_force_convergence_fn(force_tol=1e-3),
    init_kwargs={"cell_filter": ts.CellFilter.frechet},
)

Performance

Single-structure throughput

Benchmarks on NVIDIA A100 40GB GPU:

Run examples/benchmark.py to generate your own performance plots.

Batch relaxation

Relaxation of 1000 WBM structures with TorchSim on RTX 4070M:

Run examples/batch_relaxation.py to reproduce.

License and third‑party notices

• Project license: MIT LICENSE
• Third‑party: Mattersim v1.2.0 — MIT. See third_party/mattersim/LICENSE and upstream at microsoft/mattersim v1.2.0.
• Third‑party: NequIP — MIT. See third_party/nequip/LICENSE and upstream at mir-group/nequip.

If any attribution is missing or incorrect, please open an issue or PR.