outline tasks and models

README.md

@@ -16,9 +16,16 @@ If you have pretrained MLIP models that you would like to contribute to the MLIP
 
 #### Molecular dynamics calculations
 
+- [ ] [MD17](http://www.sgdml.org/#datasets)
+- [ ] [MD22](http://www.sgdml.org/#datasets)
+
 
 #### Single-point density functional theory calculations
 
+- [ ] MPTrj
+- [ ] QM9
+- [ ] [Alexandria](https://alexandria.icams.rub.de/)
+
 ### Add new training datasets
 
 [Hugging Face Auto-Train](https://huggingface.co/docs/hub/webhooks-guide-auto-retrain)

atomind_mlip/models/README.md

@@ -0,0 +1,6 @@
+
+
+## Note on model registration
+
+1. Use `ast` to parse model classes from the uploaded script.
+2. Add the classes and their supported tasks to the model registry file `registry.yaml`.

atomind_mlip/models/__init__.py

@@ -0,0 +1,39 @@
+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

atomind_mlip/models/mace-mp.py

@@ -0,0 +1,43 @@
+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

atomind_mlip/models/registry.yaml

@@ -0,0 +1,7 @@
+
+MACE_MP_Medium:
+  username: cyrusyc
+  datetime: 2022-03-28T14:30:00 # TODO: Fake datetime
+  tasks:
+    - alexandria
+    - qmof

atomind_mlip/tasks/__init__.py

@@ -0,0 +1,3 @@
+
+
+REGISTRY_FILE = 'registry.yaml'

atomind_mlip/tasks/alexandria.py

@@ -0,0 +1,3 @@
+
+
+URL = "https://alexandria.icams.rub.de/"

atomind_mlip/tasks/registry.yaml

@@ -0,0 +1,2 @@
+- alexandria
+- qmof