pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/datasets/actor.py

@@ -19,17 +19,15 @@ class Actor(InMemoryDataset):
     actor's Wikipedia.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
+        root: Root directory where the dataset should be saved.
+        transform: 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`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before being saved to
+            disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/airfrans.py

@@ -2,14 +2,13 @@ import json
 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 AirfRANS(InMemoryDataset):
@@ -26,13 +25,13 @@ class AirfRANS(InMemoryDataset):
     features: the inlet velocity (two components in meter per second), the
     distance to the airfoil (one component in meter), and the normals (two
     components in meter, set to :obj:`0` if the point is not on the airfoil).
-    Each point is given a target of 4 components for the underyling regression
+    Each point is given a target of 4 components for the underlying regression
     task: the velocity (two components in meter per second), the pressure
     divided by the specific mass (one component in meter squared per second
     squared), the turbulent kinematic viscosity (one component in meter squared
     per second).
-    Finaly, a boolean is attached to each point to inform if this point lies on
-    the airfoil or not.
+    Finally, a boolean is attached to each point to inform if this point lies
+    on the airfoil or not.
 
     A library for manipulating simulations of the dataset is available `here
     <https://airfrans.readthedocs.io/en/latest/index.html>`_.
@@ -47,26 +46,24 @@ class AirfRANS(InMemoryDataset):
         :obj:`torch_geometric.transforms.RadiusGraph` transform.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        task (str): The task to study (:obj:`"full"`, :obj:`"scarce"`,
+        root: Root directory where the dataset should be saved.
+        task: The task to study (:obj:`"full"`, :obj:`"scarce"`,
             :obj:`"reynolds"`, :obj:`"aoa"`) that defines the utilized training
             and test splits.
-        train (bool, optional): If :obj:`True`, loads the training dataset,
-            otherwise the test dataset. (default: :obj:`True`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        train: If :obj:`True`, loads the training dataset, otherwise the test
+            dataset.
+        transform: A function/transform that takes in an
+            :class:`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
+        pre_transform: A function/transform that takes in an
+            :class:`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
+            being saved to disk.
+        pre_filter: 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`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 
@@ -129,7 +126,7 @@ class AirfRANS(InMemoryDataset):
         partial = set(manifest[f'{self.task}_{self.split}'])
 
         data_list = []
-        raw_data = torch.load(self.raw_paths[0])
+        raw_data = fs.torch_load(self.raw_paths[0])
         for k, s in enumerate(total):
             if s in partial:
                 data = Data(**raw_data[k])

torch_geometric/datasets/airports.py

@@ -14,22 +14,20 @@ class Airports(InMemoryDataset):
     and labels correspond to activity levels.
     Features are given by one-hot encoded node identifiers, as described in the
     `"GraLSP: Graph Neural Networks with Local Structural Patterns"
-    ` <https://arxiv.org/abs/1911.07675>`_ paper.
+    <https://arxiv.org/abs/1911.07675>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
             :obj:`"Europe"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in an
+            :class:`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
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     edge_url = ('https://github.com/leoribeiro/struc2vec/'
                 'raw/master/graph/{}-airports.edgelist')

torch_geometric/datasets/amazon.py

@@ -15,19 +15,16 @@ class Amazon(InMemoryDataset):
     map goods to their respective product category.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Computers"`,
-            :obj:`"Photo"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Computers"`, :obj:`"Photo"`).
+        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`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/amazon_book.py

@@ -14,17 +14,16 @@ class AmazonBook(InMemoryDataset):
     No labels or features are provided.
 
     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.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` 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.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     url = ('https://raw.githubusercontent.com/gusye1234/LightGCN-PyTorch/'
            'master/data/amazon-book')

torch_geometric/datasets/amazon_products.py

@@ -14,17 +14,15 @@ class AmazonProducts(InMemoryDataset):
     containing products and its categories.
 
     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
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`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
+        pre_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
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/aminer.py

@@ -24,17 +24,16 @@ class AMiner(InMemoryDataset):
     truth labels for a subset of nodes.
 
     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.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` 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.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
 
     url = 'https://www.dropbox.com/s/1bnz8r7mofx0osf/net_aminer.zip?dl=1'

torch_geometric/datasets/aqsol.py

@@ -30,25 +30,22 @@ class AQSOL(InMemoryDataset):
     the :class:`~torch_geometric.datasets.ZINC` dataset.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        split (str, optional): If :obj:`"train"`, loads the training dataset.
+        root: Root directory where the dataset should be saved.
+        split: If :obj:`"train"`, loads the training dataset.
             If :obj:`"val"`, loads the validation dataset.
             If :obj:`"test"`, loads the test dataset.
-            (default: :obj:`"train"`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        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`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_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
-            being saved to disk. (default: :obj:`None`)
+            being saved to disk.
         pre_filter (callable, optional): A function that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a boolean
+            :class:`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`)
+            the final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/attributed_graph_dataset.py

@@ -19,21 +19,19 @@ class AttributedGraphDataset(InMemoryDataset):
     <https://arxiv.org/abs/2009.00826>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`,
             :obj:`"CiteSeer"`, :obj:`"PubMed"`, :obj:`"BlogCatalog"`,
             :obj:`"PPI"`, :obj:`"Flickr"`, :obj:`"Facebook"`, :obj:`"Twitter"`,
             :obj:`"TWeibo"`, :obj:`"MAG"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        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`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_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
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/ba_multi_shapes.py

@@ -25,21 +25,19 @@ class BAMultiShapesDataset(InMemoryDataset):
     This dataset is pre-computed from the official implementation.
 
     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
+        root: Root directory where the dataset should be saved.
+        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`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_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
-            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
+            being saved to disk.
+        pre_filter: A function that takes in a
+            :class:`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`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

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/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/city.py

@@ -0,0 +1,157 @@
+import os.path as osp
+from typing import Callable, Optional
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_tar,
+)
+from torch_geometric.io import fs
+
+
+class CityNetwork(InMemoryDataset):
+    r"""The City-Networks are introduced in
+    `"Towards Quantifying Long-Range Interactions in Graph Machine Learning:
+    a Large Graph Dataset and a Measurement"
+    <https://arxiv.org/abs/2503.09008>`_ paper.
+    The dataset contains four city networks: `paris`, `shanghai`, `la`,
+    and `london`, where nodes represent junctions and edges represent
+    undirected road segments. The task is to predict each node's eccentricity
+    score, which is approximated based on its 16-hop neighborhood and naturally
+    requires long-range information. The score indicates how accessible one
+    node is in the network, and is mapped to 10 quantiles for transductive
+    classification. See the original
+    `source code <https://github.com/LeonResearch/City-Networks>`_ for more
+    details on the individual networks.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset (``"paris"``, ``"shanghai"``,
+            ``"la"``, ``"london"``).
+        augmented (bool, optional): Whether to use the augmented node features
+            from edge features.(default: :obj:`True`)
+        transform (callable, optional): A function/transform that takes in an
+            :class:`~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 :class:`~torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+
+    **STATS:**
+
+    .. list-table::
+        :widths: 10 10 10 10 10
+        :header-rows: 1
+
+        * - Name
+          - #nodes
+          - #edges
+          - #features
+          - #classes
+        * - paris
+          - 114,127
+          - 182,511
+          - 37
+          - 10
+        * - shanghai
+          - 183,917
+          - 262,092
+          - 37
+          - 10
+        * - la
+          - 240,587
+          - 341,523
+          - 37
+          - 10
+        * - london
+          - 568,795
+          - 756,502
+          - 37
+          - 10
+    """
+    url = "https://github.com/LeonResearch/City-Networks/raw/refs/heads/main/data/"  # noqa: E501
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        augmented: bool = True,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+        delete_raw: bool = False,
+    ) -> None:
+        self.name = name.lower()
+        assert self.name in ["paris", "shanghai", "la", "london"]
+        self.augmented = augmented
+        self.delete_raw = delete_raw
+        super().__init__(
+            root,
+            transform,
+            pre_transform,
+            force_reload=force_reload,
+        )
+        self.load(self.processed_paths[0])
+
+    @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, self.name, "processed")
+
+    @property
+    def raw_file_names(self) -> str:
+        return f"{self.name}.json"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self) -> None:
+        self.download_path = download_url(
+            self.url + f"{self.name}.tar.gz",
+            self.raw_dir,
+        )
+
+    def process(self) -> None:
+        extract_tar(self.download_path, self.raw_dir)
+        data_path = osp.join(self.raw_dir, self.name)
+        node_feat = fs.torch_load(
+            osp.join(
+                data_path,
+                f"node_features{'_augmented' if self.augmented else ''}.pt",
+            ))
+        edge_index = fs.torch_load(osp.join(data_path, "edge_indices.pt"))
+        label = fs.torch_load(
+            osp.join(data_path, "10-chunk_16-hop_node_labels.pt"))
+        train_mask = fs.torch_load(osp.join(data_path, "train_mask.pt"))
+        val_mask = fs.torch_load(osp.join(data_path, "valid_mask.pt"))
+        test_mask = fs.torch_load(osp.join(data_path, "test_mask.pt"))
+        data = Data(
+            x=node_feat,
+            edge_index=edge_index,
+            y=label,
+            train_mask=train_mask,
+            val_mask=val_mask,
+            test_mask=test_mask,
+        )
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
+
+        if self.delete_raw:
+            fs.rm(data_path)
+
+    def __repr__(self) -> str:
+        return (f"{self.__class__.__name__}("
+                f"root='{self.root}', "
+                f"name='{self.name}', "
+                f"augmented={self.augmented})")

torch_geometric/datasets/dbp15k.py

@@ -73,7 +73,7 @@ class DBP15K(InMemoryDataset):
     def process(self) -> None:
         embs = {}
         with open(osp.join(self.raw_dir, 'sub.glove.300d')) as f:
-            for i, line in enumerate(f):
+            for line in f:
                 info = line.strip().split(' ')
                 if len(info) > 300:
                     embs[info[0]] = torch.tensor([float(x) for x in info[1:]])

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]: