Fixes + MD22

openqdc/__init__.py

@@ -40,6 +40,7 @@ def get_project_root():
     "PCQM_B3LYP": "openqdc.datasets.potential.pcqm",
     "PCQM_PM6": "openqdc.datasets.potential.pcqm",
     "RevMD17": "openqdc.datasets.potential.revmd17",
+    "MD22": "openqdc.datasets.potential.md22",
     "Transition1X": "openqdc.datasets.potential.transition1x",
     "MultixcQM9": "openqdc.datasets.potential.multixcqm9",
     "MultixcQM9_V2": "openqdc.datasets.potential.multixcqm9",
@@ -105,6 +106,7 @@ def __dir__():
     from .datasets.potential.gdml import GDML
     from .datasets.potential.geom import GEOM
     from .datasets.potential.iso_17 import ISO17
+    from .datasets.potential.md22 import MD22
     from .datasets.potential.molecule3d import Molecule3D
     from .datasets.potential.multixcqm9 import MultixcQM9, MultixcQM9_V2
     from .datasets.potential.nabladft import NablaDFT

openqdc/datasets/base.py

@@ -546,6 +546,48 @@ def wrapper(idx):
         datum["idxs"] = idxs
         return datum
 
+    @classmethod
+    def as_dataloader(
+        cls,
+        batch_size: int = 8,
+        energy_unit: Optional[str] = None,
+        distance_unit: Optional[str] = None,
+        array_format: str = "torch",
+        energy_type: str = "formation",
+        overwrite_local_cache: bool = False,
+        cache_dir: Optional[str] = None,
+        recompute_statistics: bool = False,
+        transform: Optional[Callable] = None,
+    ):
+        """
+        Return the dataset as a dataloader.
+
+        Parameters
+        ----------
+        batch_size : int, optional
+            Batch size, by default 8
+        For other parameters, see the __init__ method.
+        """
+        if not has_package("torch_geometric"):
+            raise ImportError("torch_geometric is required to use this method.")
+        assert array_format in ["torch", "jax"], f"Format {array_format} must be torch or jax."
+        from torch_geometric.data import Data
+        from torch_geometric.loader import DataLoader
+
+        return DataLoader(
+            cls(
+                energy_unit=energy_unit,
+                distance_unit=distance_unit,
+                array_format=array_format,
+                energy_type=energy_type,
+                overwrite_local_cache=overwrite_local_cache,
+                cache_dir=cache_dir,
+                recompute_statistics=recompute_statistics,
+                transform=lambda x: Data(**x) if transform is None else transform,
+            ),
+            batch_size=batch_size,
+        )
+
     def as_iter(self, atoms: bool = False, energy_method: int = 0):
         """
         Return the dataset as an iterator.

openqdc/datasets/potential/__init__.py

@@ -4,6 +4,7 @@
 from .gdml import GDML
 from .geom import GEOM
 from .iso_17 import ISO17
+from .md22 import MD22
 from .molecule3d import Molecule3D
 from .multixcqm9 import MultixcQM9, MultixcQM9_V2
 from .nabladft import NablaDFT
@@ -48,4 +49,5 @@
     "multixcqm9": MultixcQM9,
     "multixcqm9v2": MultixcQM9_V2,
     "revmd17": RevMD17,
+    "md22": MD22,
 }

openqdc/datasets/potential/dummy.py

@@ -94,6 +94,16 @@ def preprocess_path(self, overwrite_local_cache=False):
 
         return p_join(get_project_root(), "tests", "files", self.__name__, "preprocessed")
 
+    # override
+    @property
+    def data_types(self):
+        return {
+            "atomic_inputs": np.float32,
+            "position_idx_range": np.int32,
+            "energies": np.float32,
+            "forces": np.float32,
+        }
+
     def is_preprocessed(self):
         return True
 

openqdc/datasets/potential/iso_17.py

@@ -50,7 +50,7 @@ def __smiles_converter__(self, x):
         return "-".join(x.decode("ascii").split("_")[:-1])
 
     def read_raw_entries(self):
-        raw_path = p_join(self.root, "iso_17.h5")
+        raw_path = p_join(self.root, "iso_17.h5.gz")
         samples = read_qc_archive_h5(raw_path, "iso_17", self.energy_target_names, self.force_target_names)
 
         return samples

openqdc/datasets/potential/md22.py

@@ -0,0 +1,49 @@
+from os.path import join as p_join
+
+import numpy as np
+
+from openqdc.datasets.potential.revmd17 import RevMD17, shape_atom_inputs
+
+trajectories = [
+    "Ac-Ala3-NHMe",
+    "DHA",
+    "stachyose",
+    "AT-AT",
+    "AT-AT-CG-CG",
+    "double-walled_nanotube",
+    "buckyball-catcher",
+]
+
+
+def read_npz_entry(filename, root):
+    data = np.load(create_path(filename, root))
+    nuclear_charges, coords, energies, forces = (
+        data["z"],
+        data["R"],
+        data["E"],
+        data["F"],
+    )
+    frames = coords.shape[0]
+    res = dict(
+        name=np.array([filename] * frames),
+        subset=np.array([filename] * frames),
+        energies=energies.reshape(-1, 1).astype(np.float64),
+        forces=forces.reshape(-1, 3, 1).astype(np.float32),
+        atomic_inputs=shape_atom_inputs(coords, nuclear_charges),
+        n_atoms=np.array([len(nuclear_charges)] * frames, dtype=np.int32),
+    )
+    return res
+
+
+def create_path(filename, root):
+    return p_join(root, filename + ".npz")
+
+
+class MD22(RevMD17):
+    __name__ = "md22"
+
+    def read_raw_entries(self):
+        entries_list = []
+        for trajectory in trajectories:
+            entries_list.append(read_npz_entry(trajectory, self.root))
+        return entries_list

openqdc/datasets/potential/revmd17.py

@@ -79,6 +79,7 @@ class RevMD17(BaseDataset):
         PotentialMethod.PBE_DEF2_TZVP
         # "pbe/def2-tzvp",
     ]
+    __force_mask__ = [True]
 
     energy_target_names = [
         "PBE-TS Energy",

openqdc/datasets/potential/spice.py

@@ -56,7 +56,7 @@ class Spice(BaseDataset):
     """
 
     __name__ = "spice"
-    __energy_methods__ = [PotentialMethod.WB97M_D3BJ_DEF2_TZVPPD]  # "wb97m-d3bj/def2-tzvppd"]
+    __energy_methods__ = [PotentialMethod.WB97M_D3BJ_DEF2_TZVPPD]
     __force_mask__ = [True]
     __energy_unit__ = "hartree"
     __distance_unit__ = "bohr"

openqdc/raws/config_factory.py

@@ -354,6 +354,21 @@ class DataConfigFactory:
         dataset_name="revmd17",
         links={"revmd17.zip": "https://figshare.com/ndownloader/articles/12672038/versions/3"},
     )
+    md22 = dict(
+        dataset_name="md22",
+        links={
+            f"{x}.npz": f"http://www.quantum-machine.org/gdml/repo/datasets/md22_{x}.npz"
+            for x in [
+                "Ac-Ala3-NHMe",
+                "DHA",
+                "stachyose",
+                "AT-AT",
+                "AT-AT-CG-CG",
+                "double-walled_nanotube",
+                "buckyball-catcher",
+            ]
+        },
+    )
 
     x40 = dict(
         dataset_name="x40",