Refactor Interaction and Better Testing

openqdc/__init__.py

@@ -14,6 +14,7 @@ def get_project_root():
 _lazy_imports_obj = {
     "__version__": "openqdc._version",
     "BaseDataset": "openqdc.datasets.base",
+    # POTENTIAL
     "ANI1": "openqdc.datasets.potential.ani",
     "ANI1CCX": "openqdc.datasets.potential.ani",
     "ANI1X": "openqdc.datasets.potential.ani",
@@ -32,12 +33,23 @@ def get_project_root():
     "SolvatedPeptides": "openqdc.datasets.potential.solvated_peptides",
     "WaterClusters": "openqdc.datasets.potential.waterclusters3_30",
     "TMQM": "openqdc.datasets.potential.tmqm",
-    "Dummy": "openqdc.datasets.potential.dummy",
     "PCQM_B3LYP": "openqdc.datasets.potential.pcqm",
     "PCQM_PM6": "openqdc.datasets.potential.pcqm",
     "RevMD17": "openqdc.datasets.potential.revmd17",
     "Transition1X": "openqdc.datasets.potential.transition1x",
     "MultixcQM9": "openqdc.datasets.potential.multixcqm9",
+    # INTERACTION
+    "DES5M": "openqdc.datasets.interaction.des",
+    "DES370K": "openqdc.datasets.interaction.des",
+    "DESS66": "openqdc.datasets.interaction.des",
+    "DESS66x8": "openqdc.datasets.interaction.des",
+    "L7": "openqdc.datasets.interaction.l7",
+    "X40": "openqdc.datasets.interaction.x40",
+    "Metcalf": "openqdc.datasets.interaction.metcalf",
+    "Splinter": "openqdc.datasets.interaction.splinter",
+    # DEBUG
+    "Dummy": "openqdc.datasets.potential.dummy",
+    # ALL
     "AVAILABLE_DATASETS": "openqdc.datasets",
     "AVAILABLE_POTENTIAL_DATASETS": "openqdc.datasets.potential",
     "AVAILABLE_INTERACTION_DATASETS": "openqdc.datasets.interaction",
@@ -75,6 +87,13 @@ def __dir__():
     from ._version import __version__  # noqa
     from .datasets import AVAILABLE_DATASETS  # noqa
     from .datasets.base import BaseDataset  # noqa
+
+    # INTERACTION
+    from .datasets.interaction.des import DES5M, DES370K, DESS66, DESS66x8  # noqa
+    from .datasets.interaction.l7 import L7  # noqa
+    from .datasets.interaction.metcalf import Metcalf  # noqa
+    from .datasets.interaction.splinter import Splinter  # noqa
+    from .datasets.interaction.x40 import X40  # noqa
     from .datasets.potential.ani import ANI1, ANI1CCX, ANI1X  # noqa
     from .datasets.potential.comp6 import COMP6  # noqa
     from .datasets.potential.dummy import Dummy  # noqa

openqdc/datasets/base.py

@@ -49,11 +49,15 @@
 
 @requires_package("torch")
 def to_torch(x: np.ndarray):
+    if isinstance(x, torch.Tensor):
+        return x
     return torch.from_numpy(x)
 
 
 @requires_package("jax")
 def to_jax(x: np.ndarray):
+    if isinstance(x, jnp.ndarray):
+        return x
     return jnp.array(x)
 
 
@@ -166,6 +170,7 @@ def _precompute_statistics(self, overwrite_local_cache: bool = False):
             PerAtomFormationEnergyStats,
         )
         self.statistics.run_calculators()  # run the calculators
+        self._compute_average_nb_atoms()
 
     @classmethod
     def no_init(cls):
@@ -243,6 +248,14 @@ def data_keys(self):
             keys.remove("forces")
         return keys
 
+    @property
+    def pkl_data_keys(self):
+        return list(self.pkl_data_types.keys())
+
+    @property
+    def pkl_data_types(self):
+        return {"name": str, "subset": str, "n_atoms": np.int32}
+
     @property
     def data_types(self):
         return {
@@ -257,8 +270,8 @@ def data_shapes(self):
         return {
             "atomic_inputs": (-1, NB_ATOMIC_FEATURES),
             "position_idx_range": (-1, 2),
-            "energies": (-1, len(self.energy_target_names)),
-            "forces": (-1, 3, len(self.force_target_names)),
+            "energies": (-1, len(self.energy_methods)),
+            "forces": (-1, 3, len(self.force_methods)),
         }
 
     def _set_units(self, en, ds):
@@ -332,8 +345,14 @@ def save_preprocess(self, data_dict):
 
         # save smiles and subset
         local_path = p_join(self.preprocess_path, "props.pkl")
-        for key in ["name", "subset"]:
-            data_dict[key] = np.unique(data_dict[key], return_inverse=True)
+
+        # assert that (required) pkl keys are present in data_dict
+        assert all([key in data_dict.keys() for key in self.pkl_data_keys])
+
+        # store unique and inverse indices for str-based pkl keys
+        for key in self.pkl_data_keys:
+            if self.pkl_data_types[key] == str:
+                data_dict[key] = np.unique(data_dict[key], return_inverse=True)
 
         with open(local_path, "wb") as f:
             pkl.dump(data_dict, f)
@@ -369,7 +388,10 @@ def read_preprocess(self, overwrite_local_cache=False):
         pull_locally(filename, overwrite=overwrite_local_cache)
         with open(filename, "rb") as f:
             tmp = pkl.load(f)
-            for key in ["name", "subset", "n_atoms"]:
+            all_pkl_keys = set(tmp.keys()) - set(self.data_keys)
+            # assert required pkl_keys are present in all_pkl_keys
+            assert all([key in all_pkl_keys for key in self.pkl_data_keys])
+            for key in all_pkl_keys:
                 x = tmp.pop(key)
                 if len(x) == 2:
                     self.data[key] = x[0][x[1]]

openqdc/datasets/interaction/__init__.py

@@ -1,12 +1,9 @@
 from .base import BaseInteractionDataset  # noqa
-from .des5m import DES5M
-from .des370k import DES370K
-from .dess66 import DESS66
-from .dess66x8 import DESS66x8
-from .L7 import L7
+from .des import DES5M, DES370K, DESS66, DESS66x8
+from .l7 import L7
 from .metcalf import Metcalf
 from .splinter import Splinter
-from .X40 import X40
+from .x40 import X40
 
 AVAILABLE_INTERACTION_DATASETS = {
     "des5m": DES5M,