Spaces:
Running
Running
File size: 1,424 Bytes
1547ed2 |
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 |
import torch
from torch import nn
from torch_geometric.data import Data
from ase import Atoms
from ase.calculators.calculator import Calculator, all_changes
class MLIP(Calculator):
def __init__(self):
super().__init__()
self.name: str = "MLIP"
self.version: str = None
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model: nn.Module = None
self.implemented_properties = ["energy", "forces"]
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()
def forward(self, x: Data | Atoms) -> dict[str, torch.Tensor]:
"""Implement data conversion, graph creation, and model forward pass
Example implementation:
1. Use `ase.neighborlist.NeighborList` to get neighbor list
2. Create `torch_geometric.data.Data` object and copy the data
3. Pass the `Data` object to the model and return the output
"""
raise NotImplementedError
|