pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/datasets/ba_shapes.py

@@ -30,15 +30,14 @@ class BAShapes(InMemoryDataset):
         :class:`torch_geometric.datasets.graph_generator.BAGraph` instead.
 
     Args:
-        connection_distribution (str, optional): Specifies how the houses
-            and the BA graph get connected. Valid inputs are :obj:`"random"`
+        connection_distribution: Specifies how the houses and the BA graph get
+            connected. Valid inputs are :obj:`"random"`
             (random BA graph nodes are selected for connection to the houses),
             and :obj:`"uniform"` (uniformly distributed BA graph nodes are
-            selected for connection to the houses). (default: :obj:`"random"`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+            selected for connection to the houses).
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
     """
     def __init__(
         self,

torch_geometric/datasets/bitcoin_otc.py

@@ -87,7 +87,7 @@ class BitcoinOTC(InMemoryDataset):
         os.unlink(path)
 
     def process(self) -> None:
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             lines = [[x for x in line.split(',')]
                      for line in f.read().split('\n')[:-1]]
 

torch_geometric/datasets/brca_tgca.py

@@ -94,7 +94,7 @@ class BrcaTcga(InMemoryDataset):
         graph_feat = torch.from_numpy(graph_feat).to(torch.float)
         graph_labels = np.loadtxt(self.raw_paths[1], delimiter=',')
         graph_label = torch.from_numpy(graph_labels).to(torch.float)
-        edge_index = torch.load(self.raw_paths[2])
+        edge_index = fs.torch_load(self.raw_paths[2])
 
         data_list = []
         for x, y in zip(graph_feat, graph_label):

torch_geometric/datasets/cornell.py

@@ -0,0 +1,145 @@
+import os.path as osp
+from typing import Callable, List, Optional
+
+import torch
+
+from torch_geometric.data import InMemoryDataset, download_url
+from torch_geometric.data.hypergraph_data import HyperGraphData
+
+
+class CornellTemporalHyperGraphDataset(InMemoryDataset):
+    r"""A collection of temporal higher-order network datasets from the
+    `"Simplicial Closure and higher-order link prediction"
+    <https://arxiv.org/abs/1802.06916>`_ paper.
+    Each of the datasets is a timestamped sequence of simplices, where a
+    simplex is a set of :math:`k` nodes.
+
+    See the original `datasets page
+    <https://www.cs.cornell.edu/~arb/data/>`_ for more details about
+    individual datasets.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset.
+            (default: :obj:`"train"`)
+        setting (str, optional): If :obj:`"transductive"`, loads the dataset
+            for transductive training.
+            If :obj:`"inductive"`, loads the dataset for inductive training.
+            (default: :obj:`"transductive"`)
+        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`)
+    """
+    names = [
+        'email-Eu',
+        'email-Enron',
+        'NDC-classes',
+        'tags-math-sx',
+        'email-Eu-25',
+        'NDC-substances',
+        'congress-bills',
+        'tags-ask-ubuntu',
+        'email-Enron-25',
+        'NDC-classes-25',
+        'threads-ask-ubuntu',
+        'contact-high-school',
+        'NDC-substances-25',
+        'congress-bills-25',
+        'contact-primary-school',
+    ]
+    url = ('https://huggingface.co/datasets/SauravMaheshkar/{}/raw/main/'
+           'processed/{}/{}')
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        split: str = 'train',
+        setting: str = 'transductive',
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+        assert name in self.names
+        assert setting in ['transductive', 'inductive']
+
+        self.name = name
+        self.setting = setting
+
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload)
+
+        if split == 'train':
+            path = self.processed_paths[0]
+        elif split == 'val':
+            path = self.processed_paths[1]
+        elif split == 'test':
+            path = self.processed_paths[2]
+        else:
+            raise ValueError(f"Split '{split}' not found")
+
+        self.load(path)
+
+    @property
+    def raw_dir(self) -> str:
+        return osp.join(self.root, self.name, self.setting, 'raw')
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['train_df.csv', 'val_df.csv', 'test_df.csv']
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, self.name, self.setting, 'processed')
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['train_data.pt', 'val_data.pt', 'test_data.pt']
+
+    def download(self) -> None:
+        for filename in self.raw_file_names:
+            url = self.url.format(self.name, self.setting, filename)
+            download_url(url, self.raw_dir)
+
+    def process(self) -> None:
+        import pandas as pd
+
+        for raw_path, path in zip(self.raw_paths, self.processed_paths):
+            df = pd.read_csv(raw_path)
+
+            data_list = []
+            for i, row in df.iterrows():
+                edge_indices: List[List[int]] = [[], []]
+                for node in eval(row['nodes']):  # str(list) -> list:
+                    edge_indices[0].append(node)
+                    edge_indices[1].append(i)  # Use `i` as hyper-edge index.
+
+                x = torch.tensor([[row['timestamp']]], dtype=torch.float)
+                edge_index = torch.tensor(edge_indices)
+
+                data = HyperGraphData(x=x, edge_index=edge_index)
+
+                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, path)

torch_geometric/datasets/dblp.py

@@ -4,7 +4,6 @@ from itertools import product
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import (
@@ -110,6 +109,8 @@ class DBLP(InMemoryDataset):
         os.remove(path)
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         data = HeteroData()
 
         node_types = ['author', 'paper', 'term', 'conference']

torch_geometric/datasets/dbp15k.py

@@ -72,7 +72,7 @@ class DBP15K(InMemoryDataset):
 
     def process(self) -> None:
         embs = {}
-        with open(osp.join(self.raw_dir, 'sub.glove.300d'), 'r') as f:
+        with open(osp.join(self.raw_dir, 'sub.glove.300d')) as f:
             for i, line in enumerate(f):
                 info = line.strip().split(' ')
                 if len(info) > 300:
@@ -112,7 +112,7 @@ class DBP15K(InMemoryDataset):
         subj, rel, obj = g1.t()
 
         x_dict = {}
-        with open(feature_path, 'r') as f:
+        with open(feature_path) as f:
             for line in f:
                 info = line.strip().split('\t')
                 info = info if len(info) == 2 else info + ['**UNK**']

torch_geometric/datasets/fake.py

@@ -170,7 +170,7 @@ class FakeHeteroDataset(InMemoryDataset):
         random.shuffle(edge_types)
 
         self.edge_types: List[Tuple[str, str, str]] = []
-        count: Dict[Tuple[str, str], int] = defaultdict(lambda: 0)
+        count: Dict[Tuple[str, str], int] = defaultdict(int)
         for edge_type in edge_types[:max(num_edge_types, 1)]:
             rel = f'e{count[edge_type]}'
             count[edge_type] += 1
@@ -222,8 +222,6 @@ class FakeHeteroDataset(InMemoryDataset):
             elif self.edge_dim == 1:
                 store.edge_weight = torch.rand(store.num_edges)
 
-            pass
-
         if self._num_classes > 0 and self.task == 'graph':
             data.y = torch.tensor([random.randint(0, self._num_classes - 1)])
 

torch_geometric/datasets/flickr.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_google_url
@@ -73,6 +72,8 @@ class Flickr(InMemoryDataset):
         download_google_url(self.role_id, self.raw_dir, 'role.json')
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         f = np.load(osp.join(self.raw_dir, 'adj_full.npz'))
         adj = sp.csr_matrix((f['data'], f['indices'], f['indptr']), f['shape'])
         adj = adj.tocoo()

torch_geometric/datasets/freebase.py

@@ -75,7 +75,7 @@ class FB15k_237(InMemoryDataset):
         rel_dict: Dict[str, int] = {}
 
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 lines = [x.split('\t') for x in f.read().split('\n')[:-1]]
 
             edge_index = torch.empty((2, len(lines)), dtype=torch.long)

torch_geometric/datasets/gdelt_lite.py

@@ -9,6 +9,7 @@ from torch_geometric.data import (
     download_url,
     extract_zip,
 )
+from torch_geometric.io import fs
 
 
 class GDELTLite(InMemoryDataset):
@@ -80,9 +81,9 @@ class GDELTLite(InMemoryDataset):
     def process(self) -> None:
         import pandas as pd
 
-        x = torch.load(self.raw_paths[0])
+        x = fs.torch_load(self.raw_paths[0])
         df = pd.read_csv(self.raw_paths[1])
-        edge_attr = torch.load(self.raw_paths[2])
+        edge_attr = fs.torch_load(self.raw_paths[2])
 
         row = torch.from_numpy(df['src'].values)
         col = torch.from_numpy(df['dst'].values)

torch_geometric/datasets/ged_dataset.py

@@ -13,6 +13,7 @@ from torch_geometric.data import (
     extract_tar,
     extract_zip,
 )
+from torch_geometric.io import fs
 from torch_geometric.utils import one_hot, to_undirected
 
 
@@ -145,9 +146,9 @@ class GEDDataset(InMemoryDataset):
         path = self.processed_paths[0] if train else self.processed_paths[1]
         self.load(path)
         path = osp.join(self.processed_dir, f'{self.name}_ged.pt')
-        self.ged = torch.load(path)
+        self.ged = fs.torch_load(path)
         path = osp.join(self.processed_dir, f'{self.name}_norm_ged.pt')
-        self.norm_ged = torch.load(path)
+        self.norm_ged = fs.torch_load(path)
 
     @property
     def raw_file_names(self) -> List[str]:

torch_geometric/datasets/git_mol_dataset.py

@@ -0,0 +1,263 @@
+import sys
+from typing import Any, Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_google_url,
+    extract_zip,
+)
+from torch_geometric.io import fs
+
+
+def safe_index(lst: List[Any], e: int) -> int:
+    return lst.index(e) if e in lst else len(lst) - 1
+
+
+class GitMolDataset(InMemoryDataset):
+    r"""The dataset from the `"GIT-Mol: A Multi-modal Large Language Model
+    for Molecular Science with Graph, Image, and Text"
+    <https://arxiv.org/pdf/2308.06911>`_ 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`)
+        split (int, optional): Datasets split, train/valid/test=0/1/2.
+            (default: :obj:`0`)
+    """
+
+    raw_url_id = '1loBXabD6ncAFY-vanRsVtRUSFkEtBweg'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+        split: int = 0,
+    ):
+        from torchvision import transforms
+
+        self.split = split
+
+        if self.split == 0:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((224, 224)),
+                transforms.RandomRotation(15),
+                transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+        else:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((224, 224)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+
+        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 ['train_3500.pkl', 'valid_450.pkl', 'test_450.pkl']
+
+    @property
+    def processed_file_names(self) -> str:
+        return ['train.pt', 'valid.pt', 'test.pt'][self.split]
+
+    def download(self) -> None:
+        file_path = download_google_url(
+            self.raw_url_id,
+            self.raw_dir,
+            'gitmol.zip',
+        )
+        extract_zip(file_path, self.raw_dir)
+
+    def process(self) -> None:
+        import pandas as pd
+        from PIL import Image
+
+        try:
+            from rdkit import Chem, RDLogger
+            RDLogger.DisableLog('rdApp.*')  # type: ignore
+            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
+
+        allowable_features: Dict[str, List[Any]] = {
+            'possible_atomic_num_list':
+            list(range(1, 119)) + ['misc'],
+            'possible_formal_charge_list':
+            [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 'misc'],
+            'possible_chirality_list': [
+                Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
+                Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
+                Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
+                Chem.rdchem.ChiralType.CHI_OTHER
+            ],
+            'possible_hybridization_list': [
+                Chem.rdchem.HybridizationType.SP,
+                Chem.rdchem.HybridizationType.SP2,
+                Chem.rdchem.HybridizationType.SP3,
+                Chem.rdchem.HybridizationType.SP3D,
+                Chem.rdchem.HybridizationType.SP3D2,
+                Chem.rdchem.HybridizationType.UNSPECIFIED, 'misc'
+            ],
+            'possible_numH_list': [0, 1, 2, 3, 4, 5, 6, 7, 8, 'misc'],
+            'possible_implicit_valence_list': [0, 1, 2, 3, 4, 5, 6],
+            'possible_degree_list': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 'misc'],
+            'possible_number_radical_e_list': [0, 1, 2, 3, 4, 'misc'],
+            'possible_is_aromatic_list': [False, True],
+            'possible_is_in_ring_list': [False, True],
+            'possible_bond_type_list': [
+                Chem.rdchem.BondType.SINGLE, Chem.rdchem.BondType.DOUBLE,
+                Chem.rdchem.BondType.TRIPLE, Chem.rdchem.BondType.AROMATIC,
+                Chem.rdchem.BondType.ZERO
+            ],
+            'possible_bond_dirs': [  # only for double bond stereo information
+                Chem.rdchem.BondDir.NONE, Chem.rdchem.BondDir.ENDUPRIGHT,
+                Chem.rdchem.BondDir.ENDDOWNRIGHT
+            ],
+            'possible_bond_stereo_list': [
+                Chem.rdchem.BondStereo.STEREONONE,
+                Chem.rdchem.BondStereo.STEREOZ,
+                Chem.rdchem.BondStereo.STEREOE,
+                Chem.rdchem.BondStereo.STEREOCIS,
+                Chem.rdchem.BondStereo.STEREOTRANS,
+                Chem.rdchem.BondStereo.STEREOANY,
+            ],
+            'possible_is_conjugated_list': [False, True]
+        }
+
+        data = pd.read_pickle(
+            f'{self.raw_dir}/igcdata_toy/{self.raw_file_names[self.split]}')
+
+        data_list = []
+        for _, r in tqdm(data.iterrows(), total=data.shape[0]):
+            smiles = r['isosmiles']
+            mol = Chem.MolFromSmiles(smiles.strip('\n'))
+            if mol is not None:
+                # text
+                summary = r['summary']
+                # image
+                cid = r['cid']
+                img_file = f'{self.raw_dir}/igcdata_toy/imgs/CID_{cid}.png'
+                img = Image.open(img_file).convert('RGB')
+                img = self.img_transform(img).unsqueeze(0)
+                # graph
+                atom_features_list = []
+                for atom in mol.GetAtoms():
+                    atom_feature = [
+                        safe_index(
+                            allowable_features['possible_atomic_num_list'],
+                            atom.GetAtomicNum()),
+                        allowable_features['possible_chirality_list'].index(
+                            atom.GetChiralTag()),
+                        safe_index(allowable_features['possible_degree_list'],
+                                   atom.GetTotalDegree()),
+                        safe_index(
+                            allowable_features['possible_formal_charge_list'],
+                            atom.GetFormalCharge()),
+                        safe_index(allowable_features['possible_numH_list'],
+                                   atom.GetTotalNumHs()),
+                        safe_index(
+                            allowable_features[
+                                'possible_number_radical_e_list'],
+                            atom.GetNumRadicalElectrons()),
+                        safe_index(
+                            allowable_features['possible_hybridization_list'],
+                            atom.GetHybridization()),
+                        allowable_features['possible_is_aromatic_list'].index(
+                            atom.GetIsAromatic()),
+                        allowable_features['possible_is_in_ring_list'].index(
+                            atom.IsInRing()),
+                    ]
+                    atom_features_list.append(atom_feature)
+                x = torch.tensor(np.array(atom_features_list),
+                                 dtype=torch.long)
+
+                edges_list = []
+                edge_features_list = []
+                for bond in mol.GetBonds():
+                    i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+                    edge_feature = [
+                        safe_index(
+                            allowable_features['possible_bond_type_list'],
+                            bond.GetBondType()),
+                        allowable_features['possible_bond_stereo_list'].index(
+                            bond.GetStereo()),
+                        allowable_features['possible_is_conjugated_list'].
+                        index(bond.GetIsConjugated()),
+                    ]
+                    edges_list.append((i, j))
+                    edge_features_list.append(edge_feature)
+                    edges_list.append((j, i))
+                    edge_features_list.append(edge_feature)
+
+                edge_index = torch.tensor(
+                    np.array(edges_list).T,
+                    dtype=torch.long,
+                )
+                edge_attr = torch.tensor(
+                    np.array(edge_features_list),
+                    dtype=torch.long,
+                )
+
+                data = Data(
+                    x=x,
+                    edge_index=edge_index,
+                    smiles=smiles,
+                    edge_attr=edge_attr,
+                    image=img,
+                    caption=summary,
+                )
+
+                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/gnn_benchmark_dataset.py

@@ -12,6 +12,7 @@ from torch_geometric.data import (
     download_url,
     extract_zip,
 )
+from torch_geometric.io import fs
 from torch_geometric.utils import remove_self_loops
 
 
@@ -61,31 +62,31 @@ class GNNBenchmarkDataset(InMemoryDataset):
           - #features
           - #classes
         * - PATTERN
-          - 10,000
+          - 14,000
           - ~118.9
           - ~6,098.9
           - 3
           - 2
         * - CLUSTER
-          - 10,000
+          - 12,000
           - ~117.2
           - ~4,303.9
           - 7
           - 6
         * - MNIST
-          - 55,000
+          - 70,000
           - ~70.6
           - ~564.5
           - 3
           - 10
         * - CIFAR10
-          - 45,000
+          - 60,000
           - ~117.6
           - ~941.2
           - 5
           - 10
         * - TSP
-          - 10,000
+          - 12,000
           - ~275.4
           - ~6,885.0
           - 2
@@ -126,9 +127,9 @@ class GNNBenchmarkDataset(InMemoryDataset):
         if self.name == 'CSL' and split != 'train':
             split = 'train'
             logging.warning(
-                ("Dataset 'CSL' does not provide a standardized splitting. "
-                 "Instead, it is recommended to perform 5-fold cross "
-                 "validation with stratifed sampling"))
+                "Dataset 'CSL' does not provide a standardized splitting. "
+                "Instead, it is recommended to perform 5-fold cross "
+                "validation with stratifed sampling")
 
         super().__init__(root, transform, pre_transform, pre_filter,
                          force_reload=force_reload)
@@ -181,7 +182,7 @@ class GNNBenchmarkDataset(InMemoryDataset):
             data_list = self.process_CSL()
             self.save(data_list, self.processed_paths[0])
         else:
-            inputs = torch.load(self.raw_paths[0])
+            inputs = fs.torch_load(self.raw_paths[0])
             for i in range(len(inputs)):
                 data_list = [Data(**data_dict) for data_dict in inputs[i]]
 
@@ -197,7 +198,7 @@ class GNNBenchmarkDataset(InMemoryDataset):
         with open(self.raw_paths[0], 'rb') as f:
             adjs = pickle.load(f)
 
-        ys = torch.load(self.raw_paths[1]).tolist()
+        ys = fs.torch_load(self.raw_paths[1]).tolist()
 
         data_list = []
         for adj, y in zip(adjs, ys):