Skip to content

Commit dffc10f

Browse files
authored
Upstream Metatomic Model (#523)
1 parent 9a17923 commit dffc10f

4 files changed

Lines changed: 28 additions & 357 deletions

File tree

examples/tutorials/metatomic_tutorial.py

Lines changed: 0 additions & 69 deletions
This file was deleted.

pyproject.toml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ dependencies = [
3030
"h5py>=3.15.1",
3131
"numpy>=1.26,<3; python_version < '3.13'",
3232
"numpy>=2.3.2,<3; python_version >= '3.13'",
33-
"nvalchemi-toolkit-ops>=0.3.0",
33+
"nvalchemi-toolkit-ops[torch]>=0.3.0",
3434
"tables>=3.11.1",
3535
"torch>=2",
3636
"tqdm>=4.67",
@@ -52,8 +52,8 @@ io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2025.6.14"]
5252
symmetry = ["moyopy>=0.7.8"]
5353
mace = ["mace-torch>=0.3.15"]
5454
mattersim = ["mattersim>=0.1.2"]
55-
metatomic = ["metatomic-torch>=0.1.3", "metatrain[pet]>=2025.12"]
56-
orb = ["orb-models>=0.6.0"]
55+
metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0", "upet>=0.2.0"]
56+
orb = ["orb-models>=0.6.2"]
5757
sevenn = ["sevenn[torchsim]>=0.12.1"]
5858
graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"]
5959
nequip = ["nequip>=0.17.0"]

tests/models/test_metatomic.py

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,9 @@
1212

1313

1414
try:
15-
from metatomic.torch import ase_calculator
16-
from metatrain.utils.io import load_model
15+
from metatomic.torch import AtomisticModel
16+
from metatomic_ase import MetatomicCalculator
17+
from upet import get_upet
1718

1819
from torch_sim.models.metatomic import MetatomicModel
1920
except ImportError:
@@ -24,26 +25,22 @@
2425

2526

2627
@pytest.fixture
27-
def metatomic_calculator():
28-
"""Load a pretrained metatomic model for testing."""
29-
model_url = "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt"
30-
return ase_calculator.MetatomicCalculator(
31-
model=load_model(model_url).export(), device=DEVICE
32-
)
28+
def metatomic_module() -> AtomisticModel:
29+
return get_upet(model="pet-mad")
3330

3431

3532
@pytest.fixture
36-
def metatomic_model() -> MetatomicModel:
37-
"""Create an MetatomicModel wrapper for the pretrained model."""
38-
return MetatomicModel(model="pet-mad", device=DEVICE)
33+
def metatomic_calculator(metatomic_module: AtomisticModel) -> MetatomicCalculator:
34+
return MetatomicCalculator(model=metatomic_module, device=DEVICE)
35+
36+
37+
@pytest.fixture
38+
def metatomic_model(metatomic_module: AtomisticModel) -> MetatomicModel:
39+
return MetatomicModel(model=metatomic_module, device=DEVICE)
3940

4041

4142
def test_metatomic_initialization() -> None:
42-
"""Test that the metatomic model initializes correctly."""
43-
model = MetatomicModel(
44-
model="pet-mad",
45-
device=DEVICE,
46-
)
43+
model = MetatomicModel(model=get_upet(model="pet-mad"), device=DEVICE)
4744
assert model.device == DEVICE
4845
assert model.dtype == torch.float32
4946

0 commit comments

Comments
 (0)