Upload folder using huggingface_hub

README.md

@@ -47,13 +47,12 @@ model.to("cuda")
 
 # random batch of data
 n_atoms = 10
+elements = torch.tensor([1, 6, 7, 8]) # H, C, N, O
 data = TGData(
     pos=torch.randn(n_atoms, 3),
     batch=torch.zeros(n_atoms),
-    one_hot=torch.randint(0, 4, (n_atoms, 4)),
+    z=elements[torch.randint(0, 4, (n_atoms,))],
     natoms=n_atoms,
-    # just needs be a placeholder that decides the output energy shape
-    energy=torch.randn(1), 
 )
 data = Batch.from_data_list([data])
 # dataloader = TGDataLoader(dataset, batch_size=1, shuffle=False)

equiformer_ase_calculator.py

@@ -19,7 +19,10 @@ from torch_scatter import scatter_mean
 
 from ocpmodels.datasets import data_list_collater
 from ocpmodels.preprocessing import AtomsToGraphs
-from ocpmodels.common.relaxation.ase_utils import batch_to_atoms, ase_atoms_to_torch_geometric
+from ocpmodels.common.relaxation.ase_utils import (
+    batch_to_atoms,
+    ase_atoms_to_torch_geometric,
+)
 
 from nets.equiformer_v2.equiformer_v2_oc20 import EquiformerV2_OC20
 from nets.prediction_utils import compute_extra_props
@@ -114,8 +117,8 @@ class EquiformerASECalculator(Calculator):
         """
         # Call base class to set atoms attribute
         Calculator.calculate(self, atoms)
-        
-        # ocpmodels/common/relaxation/ase_utils.py 
+
+        # ocpmodels/common/relaxation/ase_utils.py
         # data_object = self.a2g.convert(atoms)
         # batch = data_list_collater([data_object], otf_graph=True)
 

example_inference.py

@@ -105,6 +105,7 @@ if __name__ == "__main__":
 
     # Example 1: load a dataset file and predict the first batch
     from ocpmodels.ff_lmdb import LmdbDataset
+
     dataset_path = os.path.join(project_root, "data/sample_100.lmdb")
     dataset = LmdbDataset(dataset_path)
     # either use the dataset directly or use a dataloader
@@ -120,10 +121,9 @@ if __name__ == "__main__":
 
     # Example 2: create a random data object with random positions and predict
     n_atoms = 10
-    elements = torch.tensor([1, 6, 7, 8]) # H, C, N, O
+    elements = torch.tensor([1, 6, 7, 8])  # H, C, N, O
     data = TGData(
         pos=torch.randn(n_atoms, 3),
-        batch=torch.zeros(n_atoms),
         z=elements[torch.randint(0, 4, (n_atoms,))],
         natoms=n_atoms,
     )

ocpmodels/common/relaxation/ase_utils.py

@@ -209,6 +209,7 @@ class OCPCalculator(Calculator):
         elif self.trainer.name == "is2re":
             self.results["energy"] = predictions["energy"].item()
 
+
 # by Andreas. We do not care about the graph, because connectivity is computed on the fly?
 def ase_atoms_to_torch_geometric(atoms):
     """
@@ -232,4 +233,4 @@ def ase_atoms_to_torch_geometric(atoms):
     )
     data = Batch.from_data_list([data])
 
-    return data
+    return data

ocpmodels/units.py

@@ -139,4 +139,4 @@ ELEMENT_TO_ATOMIC_NUMBER = {
     "Og": 118,
 }
 
-ATOMIC_NUMBER_TO_ELEMENT = {v: k for k, v in ELEMENT_TO_ATOMIC_NUMBER.items()}
+ATOMIC_NUMBER_TO_ELEMENT = {v: k for k, v in ELEMENT_TO_ATOMIC_NUMBER.items()}