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

torch_geometric/datasets/hgb_dataset.py

@@ -99,7 +99,7 @@ class HGBDataset(InMemoryDataset):
         # node_types = {0: 'paper', 1, 'author', ...}
         # edge_types = {0: ('paper', 'cite', 'paper'), ...}
         if self.name in ['acm', 'dblp', 'imdb']:
-            with open(self.raw_paths[0], 'r') as f:  # `info.dat`
+            with open(self.raw_paths[0]) as f:  # `info.dat`
                 info = json.load(f)
             n_types = info['node.dat']['node type']
             n_types = {int(k): v for k, v in n_types.items()}
@@ -112,7 +112,7 @@ class HGBDataset(InMemoryDataset):
                 e_types[key] = (src, rel, dst)
             num_classes = len(info['label.dat']['node type']['0'])
         elif self.name in ['freebase']:
-            with open(self.raw_paths[0], 'r') as f:  # `info.dat`
+            with open(self.raw_paths[0]) as f:  # `info.dat`
                 info = f.read().split('\n')
             start = info.index('TYPE\tMEANING') + 1
             end = info[start:].index('')
@@ -124,7 +124,7 @@ class HGBDataset(InMemoryDataset):
             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 != '']
+                src, dst, rel = (v for v in row if v != '')
                 src, dst = n_types[int(src)], n_types[int(dst)]
                 rel = rel.split('-')[1]
                 e_types[key] = (src, rel, dst)
@@ -134,8 +134,8 @@ class HGBDataset(InMemoryDataset):
         # Extract node information:
         mapping_dict = {}  # Maps global node indices to local ones.
         x_dict = defaultdict(list)
-        num_nodes_dict: Dict[str, int] = defaultdict(lambda: 0)
-        with open(self.raw_paths[1], 'r') as f:  # `node.dat`
+        num_nodes_dict: Dict[str, int] = defaultdict(int)
+        with open(self.raw_paths[1]) as f:  # `node.dat`
             xs = [v.split('\t') for v in f.read().split('\n')[:-1]]
         for x in xs:
             n_id, n_type = int(x[0]), n_types[int(x[2])]
@@ -151,7 +151,7 @@ class HGBDataset(InMemoryDataset):
 
         edge_index_dict = defaultdict(list)
         edge_weight_dict = defaultdict(list)
-        with open(self.raw_paths[2], 'r') as f:  # `link.dat`
+        with open(self.raw_paths[2]) as f:  # `link.dat`
             edges = [v.split('\t') for v in f.read().split('\n')[:-1]]
         for src, dst, rel, weight in edges:
             e_type = e_types[int(rel)]
@@ -168,9 +168,9 @@ class HGBDataset(InMemoryDataset):
 
         # Node classification:
         if self.name in ['acm', 'dblp', 'freebase', 'imdb']:
-            with open(self.raw_paths[3], 'r') as f:  # `label.dat`
+            with open(self.raw_paths[3]) as f:  # `label.dat`
                 train_ys = [v.split('\t') for v in f.read().split('\n')[:-1]]
-            with open(self.raw_paths[4], 'r') as f:  # `label.dat.test`
+            with open(self.raw_paths[4]) as f:  # `label.dat.test`
                 test_ys = [v.split('\t') for v in f.read().split('\n')[:-1]]
             for y in train_ys:
                 n_id, n_type = mapping_dict[int(y[0])], n_types[int(y[2])]

torch_geometric/datasets/imdb.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 (
@@ -69,6 +68,8 @@ class IMDB(InMemoryDataset):
         os.remove(path)
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         data = HeteroData()
 
         node_types = ['movie', 'director', 'actor']

torch_geometric/datasets/karate.py

@@ -8,8 +8,9 @@ from torch_geometric.data import Data, InMemoryDataset
 class KarateClub(InMemoryDataset):
     r"""Zachary's karate club network from the `"An Information Flow Model for
     Conflict and Fission in Small Groups"
-    <http://www1.ind.ku.dk/complexLearning/zachary1977.pdf>`_ paper, containing
-    34 nodes, connected by 156 (undirected and unweighted) edges.
+    <https://www.journals.uchicago.edu/doi/abs/10.1086/jar.33.4.3629752>`_
+    paper, containing 34 nodes,
+    connected by 156 (undirected and unweighted) edges.
     Every node is labeled by one of four classes obtained via modularity-based
     clustering, following the `"Semi-supervised Classification with Graph
     Convolutional Networks" <https://arxiv.org/abs/1609.02907>`_ paper.

torch_geometric/datasets/last_fm.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 (
@@ -68,6 +67,8 @@ class LastFM(InMemoryDataset):
         os.remove(path)
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         data = HeteroData()
 
         node_type_idx = np.load(osp.join(self.raw_dir, 'node_types.npy'))

torch_geometric/datasets/linkx_dataset.py

@@ -5,6 +5,7 @@ import numpy as np
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_url
+from torch_geometric.io import fs
 from torch_geometric.utils import one_hot
 
 
@@ -115,9 +116,9 @@ class LINKXDataset(InMemoryDataset):
 
     def _process_wiki(self) -> Data:
         paths = {x.split('/')[-1]: x for x in self.raw_paths}
-        x = torch.load(paths['wiki_features2M.pt'])
-        edge_index = torch.load(paths['wiki_edges2M.pt']).t().contiguous()
-        y = torch.load(paths['wiki_views2M.pt'])
+        x = fs.torch_load(paths['wiki_features2M.pt'])
+        edge_index = fs.torch_load(paths['wiki_edges2M.pt']).t().contiguous()
+        y = fs.torch_load(paths['wiki_views2M.pt'])
 
         return Data(x=x, edge_index=edge_index, y=y)
 

torch_geometric/datasets/lrgb.py

@@ -188,9 +188,8 @@ class LRGBDataset(InMemoryDataset):
                         graphs = pickle.load(f)
                 elif self.name.split('-')[0] == 'peptides':
                     # Peptides-func and Peptides-struct
-                    with open(osp.join(self.raw_dir, f'{split}.pt'),
-                              'rb') as f:
-                        graphs = torch.load(f)
+                    graphs = fs.torch_load(
+                        osp.join(self.raw_dir, f'{split}.pt'))
 
                 data_list = []
                 for graph in tqdm(graphs, desc=f'Processing {split} dataset'):
@@ -260,8 +259,7 @@ class LRGBDataset(InMemoryDataset):
 
     def process_pcqm_contact(self) -> None:
         for split in ['train', 'val', 'test']:
-            with open(osp.join(self.raw_dir, f'{split}.pt'), 'rb') as f:
-                graphs = torch.load(f)
+            graphs = fs.torch_load(osp.join(self.raw_dir, f'{split}.pt'))
 
             data_list = []
             for graph in tqdm(graphs, desc=f'Processing {split} dataset'):

torch_geometric/datasets/malnet_tiny.py

@@ -11,6 +11,7 @@ from torch_geometric.data import (
     extract_tar,
     extract_zip,
 )
+from torch_geometric.io import fs
 
 
 class MalNetTiny(InMemoryDataset):
@@ -65,7 +66,7 @@ class MalNetTiny(InMemoryDataset):
         self.load(self.processed_paths[0])
 
         if split is not None:
-            split_slices = torch.load(self.processed_paths[1])
+            split_slices = fs.torch_load(self.processed_paths[1])
             if split == 'train':
                 self._indices = range(split_slices[0], split_slices[1])
             elif split == 'val':
@@ -98,7 +99,7 @@ class MalNetTiny(InMemoryDataset):
         split_slices = [0]
 
         for split in ['train', 'val', 'test']:
-            with open(osp.join(self.raw_paths[1], f'{split}.txt'), 'r') as f:
+            with open(osp.join(self.raw_paths[1], f'{split}.txt')) as f:
                 filenames = f.read().split('\n')[:-1]
                 split_slices.append(split_slices[-1] + len(filenames))
 
@@ -107,7 +108,7 @@ class MalNetTiny(InMemoryDataset):
                 malware_type = filename.split('/')[0]
                 y = y_map.setdefault(malware_type, len(y_map))
 
-                with open(path, 'r') as f:
+                with open(path) as f:
                     edges = f.read().split('\n')[5:-1]
 
                 edge_indices = [[int(s) for s in e.split()] for e in edges]

torch_geometric/datasets/mnist_superpixels.py

@@ -1,14 +1,13 @@
 import os
 from typing import Callable, List, Optional
 
-import torch
-
 from torch_geometric.data import (
     Data,
     InMemoryDataset,
     download_url,
     extract_zip,
 )
+from torch_geometric.io import fs
 
 
 class MNISTSuperpixels(InMemoryDataset):
@@ -85,7 +84,7 @@ class MNISTSuperpixels(InMemoryDataset):
         os.unlink(path)
 
     def process(self) -> None:
-        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]]