update default data shapes for interaction datasets

valence-labs · Mar 7, 2024 · 3c84ee9 · 3c84ee9
1 parent f046ea9
commit 3c84ee9
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/src/openqdc/datasets/interaction/base.py b/src/openqdc/datasets/interaction/base.py
@@ -10,6 +10,9 @@
     set_cache_dir,
 )
 from openqdc.datasets.potential.base import BaseDataset
+from openqdc.utils.constants import (
+    NB_ATOMIC_FEATURES
+)
 
 from loguru import logger
 
@@ -43,3 +46,12 @@ def collate_list(self, list_entries: List[Dict]):
         res["position_idx_range"] = x
 
         return res
+
+    @property
+    def data_shapes(self):
+        return {
+            "atomic_inputs": (-1, NB_ATOMIC_FEATURES),
+            "position_idx_range": (-1, 2),
+            "energies": (-1, len(self.__energy_methods__)),
+            "forces": (-1, 3, len(self.force_target_names)),
+        }