pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/datasets/hgb_dataset.py

@@ -123,8 +123,8 @@ class HGBDataset(InMemoryDataset):
             start = info.index('LINK\tSTART\tEND\tMEANING') + 1
             end = info[start:].index('')
             for key, row in enumerate(info[start:start + end]):
-                row = row.split('\t')[1:]
-                src, dst, rel = (v for v in row if v != '')
+                edge = row.split('\t')[1:]
+                src, dst, rel = (v for v in edge if v != '')
                 src, dst = n_types[int(src)], n_types[int(dst)]
                 rel = rel.split('-')[1]
                 e_types[key] = (src, rel, dst)

torch_geometric/datasets/hm.py

@@ -81,7 +81,7 @@ class HM(InMemoryDataset):
             xs.append(torch.from_numpy(x).to(torch.float))
 
         x = torch.from_numpy(df['age'].values).to(torch.float).view(-1, 1)
-        x = x.nan_to_num(nan=x.nanmean())
+        x = x.nan_to_num(nan=x.nanmean())  # type: ignore
         xs.append(x / x.max())
 
         data['customer'].x = torch.cat(xs, dim=-1)

torch_geometric/datasets/instruct_mol_dataset.py

@@ -0,0 +1,134 @@
+import json
+import sys
+from typing import Callable, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import Data, InMemoryDataset
+from torch_geometric.io import fs
+from torch_geometric.utils import one_hot
+
+
+class InstructMolDataset(InMemoryDataset):
+    r"""The dataset from the `"InstructMol: Multi-Modal Integration for
+    Building a Versatile and Reliable Molecular Assistant in Drug Discovery"
+    <https://arxiv.org/pdf/2311.16208>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    raw_url = 'https://huggingface.co/datasets/OpenMol/PubChemSFT/resolve/main'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['all_clean.json']
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['data.pt']
+
+    def download(self) -> None:
+        print('downloading dataset...')
+        fs.cp(f'{self.raw_url}/all_clean.json', self.raw_dir)
+
+    def process(self) -> None:
+        try:
+            from rdkit import Chem
+            from rdkit.Chem.rdchem import BondType as BT
+            WITH_RDKIT = True
+
+        except ImportError:
+            WITH_RDKIT = False
+
+        if not WITH_RDKIT:
+            print(("Using a pre-processed version of the dataset. Please "
+                   "install 'rdkit' to alternatively process the raw data."),
+                  file=sys.stderr)
+
+            data_list = fs.torch_load(self.raw_paths[0])
+            data_list = [Data(**data_dict) for data_dict in data_list]
+
+            if self.pre_filter is not None:
+                data_list = [d for d in data_list if self.pre_filter(d)]
+
+            if self.pre_transform is not None:
+                data_list = [self.pre_transform(d) for d in data_list]
+
+            self.save(data_list, self.processed_paths[0])
+            return
+
+        # types of atom and bond
+        types = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4, 'Unknow': 5}
+        bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3}
+
+        # load data
+        mols = json.load(open(f'{self.raw_dir}/all_clean.json'))
+
+        data_list = []
+        for smiles, qa_pairs in tqdm(mols.items(), total=len(mols)):
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                continue
+
+            x: torch.Tensor = torch.tensor([
+                types[atom.GetSymbol()] if atom.GetSymbol() in types else 5
+                for atom in mol.GetAtoms()
+            ])
+            x = one_hot(x, num_classes=len(types), dtype=torch.float)
+
+            rows, cols, edge_types = [], [], []
+            for bond in mol.GetBonds():
+                i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+                edge_types += [bonds[bond.GetBondType()]] * 2
+                rows += [i, j]
+                cols += [j, i]
+
+            edge_index = torch.tensor([rows, cols], dtype=torch.long)
+            edge_type = torch.tensor(edge_types, dtype=torch.long)
+            edge_attr = one_hot(edge_type, num_classes=len(bonds))
+
+            for question, answer in qa_pairs:
+                data = Data(
+                    x=x,
+                    edge_index=edge_index,
+                    edge_attr=edge_attr,
+                    smiles=smiles,
+                    instruction=question,
+                    y=answer,
+                )
+
+                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)
+
+        self.save(data_list, self.processed_paths[0])

torch_geometric/datasets/md17.py

@@ -57,7 +57,7 @@ class MD17(InMemoryDataset):
     +--------------------+--------------------+-------------------------------+-----------+
     | Uracil             | DFT                | :obj:`uracil`                 | 133,770   |
     +--------------------+--------------------+-------------------------------+-----------+
-    | Naphthalene        | DFT                | :obj:`napthalene`             | 326,250   |
+    | Naphthalene        | DFT                | :obj:`naphthalene`            | 326,250   |
     +--------------------+--------------------+-------------------------------+-----------+
     | Aspirin            | DFT                | :obj:`aspirin`                | 211,762   |
     +--------------------+--------------------+-------------------------------+-----------+
@@ -77,7 +77,7 @@ class MD17(InMemoryDataset):
     +--------------------+--------------------+-------------------------------+-----------+
     | Uracil (R)         | DFT (PBE/def2-SVP) | :obj:`revised uracil`         | 100,000   |
     +--------------------+--------------------+-------------------------------+-----------+
-    | Naphthalene (R)    | DFT (PBE/def2-SVP) | :obj:`revised napthalene`     | 100,000   |
+    | Naphthalene (R)    | DFT (PBE/def2-SVP) | :obj:`revised naphthalene`    | 100,000   |
     +--------------------+--------------------+-------------------------------+-----------+
     | Aspirin (R)        | DFT (PBE/def2-SVP) | :obj:`revised aspirin`        | 100,000   |
     +--------------------+--------------------+-------------------------------+-----------+
@@ -309,7 +309,7 @@ class MD17(InMemoryDataset):
     file_names = {
         'benzene': 'md17_benzene2017.npz',
         'uracil': 'md17_uracil.npz',
-        'naphtalene': 'md17_naphthalene.npz',
+        'naphthalene': 'md17_naphthalene.npz',
         'aspirin': 'md17_aspirin.npz',
         'salicylic acid': 'md17_salicylic.npz',
         'malonaldehyde': 'md17_malonaldehyde.npz',

torch_geometric/datasets/medshapenet.py

@@ -0,0 +1,145 @@
+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
+
+
+class MedShapeNet(InMemoryDataset):
+    r"""The MedShapeNet datasets from the `"MedShapeNet -- A Large-Scale
+    Dataset of 3D Medical Shapes for Computer Vision"
+    <https://arxiv.org/abs/2308.16139>`_ paper,
+    containing 8 different type of structures (classes).
+
+    .. note::
+
+        Data objects hold mesh faces instead of edge indices.
+        To convert the mesh to a graph, use the
+        :obj:`torch_geometric.transforms.FaceToEdge` as :obj:`pre_transform`.
+        To convert the mesh to a point cloud, use the
+        :obj:`torch_geometric.transforms.SamplePoints` as :obj:`transform` to
+        sample a fixed number of points on the mesh faces according to their
+        face area.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        size (int): Number of invividual 3D structures to download per
+            type (classes).
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        root: str,
+        size: int = 100,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+        self.size = size
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload=force_reload)
+
+        path = self.processed_paths[0]
+        self.load(path)
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [
+            '3DTeethSeg', 'CoronaryArteries', 'FLARE', 'KITS', 'PULMONARY',
+            'SurgicalInstruments', 'ThoracicAorta_Saitta', 'ToothFairy'
+        ]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['dataset.pt']
+
+    @property
+    def raw_paths(self) -> List[str]:
+        r"""The absolute filepaths that must be present in order to skip
+        downloading.
+        """
+        return [osp.join(self.raw_dir, f) for f in self.raw_file_names]
+
+    def process(self) -> None:
+        import urllib3
+        from MedShapeNet import MedShapeNet as msn
+
+        msn_instance = msn(timeout=120)
+
+        urllib3.HTTPConnectionPool("medshapenet.ddns.net", maxsize=50)
+
+        list_of_datasets = msn_instance.datasets(False)
+        list_of_datasets = list(
+            filter(
+                lambda x: x not in [
+                    'medshapenetcore/ASOCA', 'medshapenetcore/AVT',
+                    'medshapenetcore/AutoImplantCraniotomy',
+                    'medshapenetcore/FaceVR'
+                ], list_of_datasets))
+
+        subset = []
+        for dataset in list_of_datasets:
+            parts = dataset.split("/")
+            self.newpath = self.root + '/' + parts[1 if len(parts) > 1 else 0]
+            if not os.path.exists(self.newpath):
+                os.makedirs(self.newpath)
+            stl_files = msn_instance.dataset_files(dataset, '.stl')
+            subset.extend(stl_files[:self.size])
+
+            for stl_file in stl_files[:self.size]:
+                msn_instance.download_stl_as_numpy(bucket_name=dataset,
+                                                   stl_file=stl_file,
+                                                   output_dir=self.newpath,
+                                                   print_output=False)
+
+        class_mapping = {
+            '3DTeethSeg': 0,
+            'CoronaryArteries': 1,
+            'FLARE': 2,
+            'KITS': 3,
+            'PULMONARY': 4,
+            'SurgicalInstruments': 5,
+            'ThoracicAorta_Saitta': 6,
+            'ToothFairy': 7
+        }
+
+        for dataset, path in zip([subset], self.processed_paths):
+            data_list = []
+            for item in dataset:
+                class_name = item.split("/")[0]
+                item = item.split("stl")[0]
+                target = class_mapping[class_name]
+                file = osp.join(self.root, item + 'npz')
+
+                data = np.load(file)
+                pre_data_list = Data(
+                    pos=torch.tensor(data["vertices"], dtype=torch.float),
+                    face=torch.tensor(data["faces"],
+                                      dtype=torch.long).t().contiguous())
+                pre_data_list.y = torch.tensor([target], dtype=torch.long)
+                data_list.append(pre_data_list)
+
+            if self.pre_filter is not None:
+                data_list = [d for d in data_list if self.pre_filter(d)]
+
+            if self.pre_transform is not None:
+                data_list = [self.pre_transform(d) for d in data_list]
+
+            self.save(data_list, path)

torch_geometric/datasets/modelnet.py

@@ -79,7 +79,7 @@ class ModelNet(InMemoryDataset):
 
     urls = {
         '10':
-        'http://vision.princeton.edu/projects/2014/3DShapeNets/ModelNet10.zip',
+        'http://3dvision.princeton.edu/projects/2014/3DShapeNets/ModelNet10.zip',  # noqa
         '40': 'http://modelnet.cs.princeton.edu/ModelNet40.zip'
     }