Spaces:
Running
Running
File size: 1,607 Bytes
d390139 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import torch
from huggingface_hub import hf_hub_download
from torch_geometric.data import Data
from ase import Atoms
from ase.calculators.calculator import all_changes
from atomind_mlip.models import MLIP
class MACE_MP_Medium(MLIP):
def __init__(self):
super().__init__()
self.name = "MACE-MP-0 (medium)"
self.version = "1.0.0"
fpath = hf_hub_download(repo_id="cyrusyc/mace-universal", subfolder="pretrained", filename="2023-12-12-mace-128-L1_epoch-199.model")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = torch.load(fpath, map_location="cpu")
self.model.to(self.device)
self.implemented_properties = [
"energy",
"forces",
"stress",
]
def calculate(self, atoms: Atoms, properties: list[str], system_changes: dict = all_changes):
"""Calculate energies and forces for the given Atoms object"""
super().calculate(atoms, properties, system_changes)
output = self.forward(atoms)
self.results = {}
if "energy" in properties:
self.results["energy"] = output["energy"].item()
if "forces" in properties:
self.results["forces"] = output["forces"].cpu().detach().numpy()
if "stress" in properties:
self.results["stress"] = output["stress"].cpu().detach().numpy()
def forward(self, x: Data | Atoms) -> dict[str, torch.Tensor]:
"""Implement data conversion, graph creation, and model forward pass"""
# TODO
raise NotImplementedError
|