Merge pull request #192 from guillemsimeon/main

Inclusion of MD22 dataset

Author: guillemsimeon

torchmdnet/datasets/__init__.py

@@ -4,6 +4,7 @@
 from .custom import Custom
 from .hdf import HDF5
 from .md17 import MD17
+from .md22 import MD22
 from .qm9 import QM9
 from .qm9q import QM9q
 from .spice import SPICE
@@ -21,6 +22,7 @@
     "GDB10to13",
     "HDF5",
     "MD17",
+    "MD22",
     "QM9",
     "QM9q",
     "S66X8",

torchmdnet/datasets/md22.py

@@ -0,0 +1,93 @@
+import os
+import os.path as osp
+from typing import Callable, List, Optional
+import numpy as np
+import torch
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_tar,
+    extract_zip,
+)
+
+class MD22(InMemoryDataset):
+
+    gdml_url = 'http://quantum-machine.org/gdml/data/npz'
+
+    file_names = {
+        'AT-AT-CG-CG': 'md22_AT-AT-CG-CG.npz',
+        'AT-AT': 'md22_AT-AT.npz',
+        'Ac-Ala3-NHMe': 'md22_Ac-Ala3-NHMe.npz',
+        'DHA': 'md22_DHA.npz',
+        'buckyball-catcher': 'md22_buckyball-catcher.npz',
+        'dw-nanotube': 'md22_dw_nanotube.npz',
+        'stachyose': 'md22_stachyose.npz',        
+    }
+
+    def __init__(
+        self,
+        root: str,
+        molecules: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+    ):
+        name = molecules
+        if name not in self.file_names:
+            raise ValueError(f"Unknown dataset name '{name}'")
+
+        self.name = name
+
+        super().__init__(root, transform, pre_transform, pre_filter)
+
+        idx = 0
+        self.data, self.slices = torch.load(self.processed_paths[idx])
+
+    def mean(self) -> float:
+        return float(self._data.energy.mean())
+
+    @property
+    def raw_dir(self) -> str:
+        return osp.join(self.root, self.name, 'raw')
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, 'processed', self.name)
+
+    @property
+    def raw_file_names(self) -> str:
+        name = self.file_names[self.name]
+        return name
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['data.pt']
+
+    def download(self):
+        url = f'{self.gdml_url}/{self.file_names[self.name]}'
+        path = download_url(url, self.raw_dir)
+
+    def process(self):
+        it = zip(self.raw_paths, self.processed_paths)
+        for raw_path, processed_path in it:
+            raw_data = np.load(raw_path)
+            
+            z = torch.from_numpy(raw_data['z']).long()
+            pos = torch.from_numpy(raw_data['R']).float()
+            energy = torch.from_numpy(raw_data['E']).float()
+            force = torch.from_numpy(raw_data['F']).float()
+
+            data_list = []
+            for i in range(pos.size(0)):
+                data = Data(z=z, pos=pos[i], y=energy[i].unsqueeze(-1), neg_dy=force[i])
+                if self.pre_filter is not None and not self.pre_filter(data):
+                    continue
+                if self.pre_transform is not None:
+                    data = self.pre_transform(data)
+                data_list.append(data)
+
+            torch.save(self.collate(data_list), processed_path)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({len(self)}, name='{self.name}')"