DeepSNAP

https://snap.stanford.edu/deepsnap/
Admin
DeepSNAP is a Python library designed for efficient deep learning on graphs, featuring flexible...
Tokens:103,122
Snippets:556
Trust Score:10
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install Snap.py backend

Source: https://snap.stanford.edu/deepsnap/_sources/notes/installation.rst.txt

Command to install the Snap.py (SnapX) backend for DeepSNAP.

```bash
$ pip install snap-stanford
```

--------------------------------

### Build DeepSNAP from source

Source: https://snap.stanford.edu/deepsnap/_sources/notes/installation.rst.txt

Commands to clone the repository and install the package from the local source directory.

```bash
$ git clone https://github.com/snap-stanford/deepsnap
$ cd deepsnap
$ pip install .
```

--------------------------------

### Install DeepSNAP via pip

Source: https://snap.stanford.edu/deepsnap/_sources/notes/installation.rst.txt

Standard installation command for the DeepSNAP package.

```bash
$ pip install deepsnap
```

--------------------------------

### Install DeepSNAP via pip

Source: https://snap.stanford.edu/deepsnap/notes/installation.html

Install DeepSNAP using pip. Requires PyTorch Geometric and NetworkX to be installed beforehand.

```bash
$ pip install deepsnap

```

--------------------------------

### Create DeepSNAP Graph from PyTorch Geometric

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Starts creating a DeepSNAP Graph from PyTorch Geometric's Planetoid dataset. The example is incomplete in the source, showing only the imports and the beginning of the code.

```python
from deepsnap.graph import Graph
from torch_geometric.datasets import Planetoid
```

--------------------------------

### Preprocess and filter graphs

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/dataset.html

Handles graph input validation, on-the-fly generator setup, and filters graphs based on the minimum_node_per_graph constraint.

```python
# graphs preprocessing
if graphs is None or len(graphs) == 0:
    if generator is None:
        raise ValueError("Graphs are None")
    else:
        # on-the-fly dataset
        self.generator = generator
        self.graphs = None
        self.otf_device = None
elif generator is not None:
    raise ValueError(
        "Both graphs and on-the-fly generator are "
        "provided (only one should be provided."
    )
else:
    # by default on-the-fly generator is not used
    # when generator is not provide
    self.generator = None

    # filter graphs that are too small
    if self.minimum_node_per_graph > 0:
        graphs_filter = []
        for idx, graph in enumerate(self.graphs):
            if isinstance(graph, Graph):
                if isinstance(graph, HeteroGraph):
                    if (
                        sum(graph.num_nodes().values())
                        >= self.minimum_node_per_graph
                    ):
                        graphs_filter.append(graph)
                    else:
                        warnings.warn(
                            f"the {idx}-th graph in self.graphs is "
                            "filtered out as it contains "
                            f"{sum(graph.num_nodes().values())} nodes,"
                            " which is less than "
                            "self.minimum_node_per_graph: "
                            f"{self.minimum_node_per_graph}."
                        )
                else:
                    if graph.num_nodes >= self.minimum_node_per_graph:
                        graphs_filter.append(graph)
                    else:
                        warnings.warn(
                            f"the {idx}-th graph in self.graphs is "
                            "filtered out as it contains "
                            f"{graph.num_nodes} nodes,"
                            " which is less than "
                            "self.minimum_node_per_graph: "
                            f"{self.minimum_node_per_graph}."
                        )
            else:
                raise TypeError(
                    "element in self.graphs of unexpected type"
                )
        self.graphs = graphs_filter

    # update graph in self.graphs with appropriate custom
    # split component
    for graph in self.graphs:
        if not hasattr(graph, "_custom_update_flag"):
            # assign task to graph
            graph.task = self.task
```

--------------------------------

### Custom Negative Sampling Method

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/hetero_graph.html

Configures negative sampling for heterogeneous graphs, allowing for optional resampling of positive examples.

```python
    def _custom_create_neg_sampling(
        self,
        negative_sampling_ratio: float,
        split_types: List[str] = None,
        resample: bool = False
    ):
        r"""
        Args:
            negative_sampling_ratio (float or int): ratio of negative sampling
                edges compared with the original edges.
            resample (boolean): whether should resample.
        """
        if split_types is None:
            split_types = self.message_types
        if not isinstance(split_types, list):
            raise TypeError("Split_types must be string or list of string.")

        # filter split_types
        split_types = (
            [
                message_type for message_type in split_types
                if message_type in self.edge_label_index
            ]
        )

        if resample and self._num_positive_examples is not None:
            for (
                message_type, edge_type_positive_num
            ) in self._num_positive_examples.items():
                self.edge_label_index[message_type] = (
                    self.edge_label_index[message_type][
                        :, :edge_type_positive_num
                    ]
                )
```

--------------------------------

### Initialize dataset parameters

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/dataset.html

Assigns input parameters to instance variables and initializes private state.

```python
# parameter initialization
self.graphs = graphs
self.task = task
self.custom_split_graphs = custom_split_graphs
self.edge_message_ratio = edge_message_ratio
self.edge_negative_sampling_ratio = edge_negative_sampling_ratio
self.edge_train_mode = edge_train_mode
self.edge_split_mode = edge_split_mode
self.minimum_node_per_graph = minimum_node_per_graph
self.resample_negatives = resample_negatives
self.resample_disjoint = resample_disjoint
self.resample_disjoint_period = resample_disjoint_period
self.negative_label_val = negative_label_val

# set private parameters
self._split_types = None
self._is_tensor = False
```

--------------------------------

### Initialize edge lists for splitting

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/hetero_graph.html

Prepares the edge lists for splitting, optionally shuffling them if required.

```python
        split_types_all_flag = set(split_types) == set(self.message_types)

        if self.G is not None:
            edges = list(self.G.edges(data=True))
            if shuffle:
                random.shuffle(edges)
        else:
            edges = {}
            for message_type in self.message_types:
                if shuffle:
                    edges[message_type] = (
                        torch.randperm(self.num_edges(message_type))
                    )
                else:
                    edges[message_type] = (
                        torch.arange(self.num_edges(message_type))
                    )
```

--------------------------------

### get_num_labels

Source: https://snap.stanford.edu/deepsnap/modules/graph.html

Gets the unique labels for a specified key.

```APIDOC
## get_num_labels

### Description
Gets the unique labels for a specified key.

### Method
get_num_labels

### Signature
get_num_labels(key: str)

### Parameters
- **key** (str) - Required - The chosen property.

### Returns
Unique labels in tensor format.

### Return Type
torch.Tensor
```

--------------------------------

### Create HeteroGraph from NetworkX

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Shows how to initialize a HeteroGraph object by defining nodes and edges with specific types and features in a NetworkX DiGraph.

```python
import torch
import networkx as nx
from deepsnap.hetero_graph import HeteroGraph

G = nx.DiGraph()
G.add_node(0, node_type='n1', node_label=1, node_feature=torch.Tensor([0.1, 0.2, 0.3]))
G.add_node(1, node_type='n1', node_label=0, node_feature=torch.Tensor([0.2, 0.3, 0.4]))
G.add_node(2, node_type='n2', node_label=1, node_feature=torch.Tensor([0.3, 0.4, 0.5]))
G.add_edge(0, 1, edge_type='e1')
G.add_edge(0, 2, edge_type='e1')
G.add_edge(1, 2, edge_type='e2')
H = HeteroGraph(G)
for hetero_feature in H:
    print(hetero_feature)

>>> ('G', <networkx.classes.digraph.DiGraph object at 0x103642370>)
    ('edge_index', {('n1', 'e1', 'n1'): tensor([[0],
            [1]]), ('n1', 'e1', 'n2'): tensor([[0],
            [0]]), ('n1', 'e2', 'n2'): tensor([[1],
            [0]])})
    ('edge_label_index', {('n1', 'e1', 'n1'): tensor([[0],
            [1]]), ('n1', 'e1', 'n2'): tensor([[0],
            [0]]), ('n1', 'e2', 'n2'): tensor([[1],
            [0]])})
    ('edge_to_graph_mapping', {('n1', 'e1', 'n1'): tensor([0]), ('n1', 'e1', 'n2'): tensor([1]), ('n1', 'e2', 'n2'): tensor([2])})
    ('edge_to_tensor_mapping', tensor([0, 0, 0]))
    ('edge_type', {('n1', 'e1', 'n1'): ['e1'], ('n1', 'e1', 'n2'): ['e1'], ('n1', 'e2', 'n2'): ['e2']})
    ('node_feature', {'n1': tensor([[0.1000, 0.2000, 0.3000],
            [0.2000, 0.3000, 0.4000]]), 'n2': tensor([[0.3000, 0.4000, 0.5000]])})
    ('node_label', {'n1': tensor([1, 0]), 'n2': tensor([1])})
    ('node_label_index', {'n1': tensor([0, 1]), 'n2': tensor([0])})
    ('node_to_graph_mapping', {'n1': tensor([0, 1]), 'n2': tensor([2])})
    ('node_to_tensor_mapping', tensor([0, 1, 0]))
    ('node_type', {'n1': ['n1', 'n1'], 'n2': ['n2']})
```

--------------------------------

### Graph.__getitem__

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/graph.html

Gets the data of the attribute with the given key. Returns None if the attribute does not exist.

```APIDOC
## Graph.__getitem__

### Description
Gets the data of the attribute with the given key. Returns None if the attribute does not exist.

### Method
Operator Overload

### Signature
__getitem__(key: str)

### Parameters
- **key** (str) - The name of the attribute to retrieve.

### Returns
- The value of the attribute, or None if not set.
```

--------------------------------

### Graph.__init__(G=None, netlib=None, **kwargs)

Source: https://snap.stanford.edu/deepsnap/modules/graph.html

Constructor for the Graph class. Initializes a graph object with optional NetworkX or SnapX graph and tensor backend attributes.

```APIDOC
## Graph.__init__(G=None, netlib=None, **kwargs)

### Description
Initializes a DeepSNAP Graph object. If G is not specified, the tensor backend is used. The graph backend can be NetworkX or SnapX (undirected homogeneous only).

### Method
Constructor

### Parameters
- **G** (Graph object, optional) - The NetworkX or SnapX graph object containing features and labels.
- **netlib** (types.ModuleType, optional) - The graph backend module. Default is NetworkX.
- ****kwargs** (optional) - Keyworded argument list with keys such as node_feature, node_label and values that are corresponding attributes. Features are required for the tensor backend.

### Returns
A Graph object.
```

--------------------------------

### HeteroConv initialization with aggregation and parallelization

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/hetero_gnn.html

Initializes the HeteroConv layer, validating the aggregation method and optionally creating CUDA streams for parallel message passing when CUDA is available.

```python
assert aggr in ["add", "mean", "max", "mul", "concat", None]
        self.aggr = aggr

        if parallelize and torch.cuda.is_available():
            self.streams = {key: torch.cuda.Stream() for key in convs.keys()}
        else:
            self.streams = None
```

--------------------------------

### Graph Attribute Accessors

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/graph.html

Methods for getting and setting graph attributes using dictionary-like syntax.

```python
    def __getitem__(self, key: str):
        r"""
        Gets the data of the attribute :obj:`key`.
        """
        return getattr(self, key, None)

    def __setitem__(self, key: str, value):
        """Sets the attribute :obj:`key` to :obj:`value`."""
        setattr(self, key, value)

    @property
    def keys(self):
        r"""
        Returns all names of the graph attributes.

```

--------------------------------

### Graph Class Definition and Initialization

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/graph.html

The Graph class constructor handles graph initialization from NetworkX/SnapX objects or raw tensors, setting up necessary attributes like features and labels.

```python
import re
import types
import random
import copy
import math
import pdb
import numpy as np
import torch
from torch_geometric.utils import to_undirected
from typing import (
    Dict,
    List,
    Union
)
import warnings
import deepsnap



[docs]class Graph(object):
    r"""
    A plain python object modeling a single graph with various
    (optional) attributes.

    Args:
        G (Graph object, optional): The NetworkX or SnapX graph 
            object which contains features and labels. If it is not 
            specified, :class:`Graph` will use the tensor backend. 
        netlib (types.ModuleType, optional): The graph backend module. 
                Currently DeepSNAP supports the NetworkX and SnapX (for 
                SnapX only the undirected homogeneous graph) as the graph 
                backend. Default graph backend is the NetworkX.
        **kwargs (optional): Keyworded argument list with keys such
            as :obj:`node_feature`, :obj:`node_label` and 
            values that are corresponding attributes. The features 
            are required for the tensor backend.
    """

    def __init__(self, G=None, netlib=None, **kwargs):
        self.G = G
        if netlib is not None:
            deepsnap._netlib = netlib
        keys = [
            "node_feature",
            "node_label",
            "edge_feature",
            "edge_label",
            "graph_feature",
            "graph_label",
            "edge_index",
            "edge_label_index",
            "node_label_index",
            "custom"
        ]
        for key in keys:
            self[key] = None

        self._is_train = False
        self._num_positive_examples = None

        for key, item in kwargs.items():
            self[key] = item

        if G is None and kwargs:
            if "directed" not in kwargs:
                self.directed = True
            if "edge_index" not in kwargs:
                raise ValueError(
                    "A tensor of edge_index is required by using "
                    "the tensor backend."
                )
            # check for undirected edge_index format
            if not self.directed:
                edge_index_length = self.edge_index.shape[1]
                edge_index_first_half, _ = (
                    torch.sort(self.edge_index[:, :int(edge_index_length / 2)])
                )
                edge_index_second_half, _ = (
                    torch.sort(self.edge_index[:, int(edge_index_length / 2):])
                )
                if not torch.equal(
                    edge_index_first_half,
                    torch.flip(edge_index_second_half, [0])
                ):
                    raise ValueError(
                        "In tensor backend mode with undirected graph, "
                        "the user provided edge_index should contain "
                        "undirected edges for both directions."
                        "the first half of edge_index should contain "
                        "unique edges in one direction and the second "
                        "half of edge_index should contain the same set "
                        "of unique edges of another direction."
                    )

        if G is not None or kwargs:
            # handle tensor backend + custom support
            if (
                ("edge_label_index" not in kwargs)
                and ("node_label_index" not in kwargs)
            ):
                self._update_tensors(init=True)
            else:
                self._update_tensors(init=False)
```

--------------------------------

### Graph.__cat_dim__

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/graph.html

Returns the dimension for which value of attribute key will get concatenated when creating batches. Internal use only.

```APIDOC
## Graph.__cat_dim__

### Description
Returns the dimension for which value of attribute key will get concatenated when creating batches. Internal use only.

### Method Signature
`__cat_dim__(self, key: str, value) -> int`

### Parameters
- **key** (str) - Required - The attribute name.
- **value** - Required - The attribute value.

### Returns
- **int**: The concatenation dimension. Returns -1 if 'index' is in key, otherwise 0.
```

--------------------------------

### Graph.keys

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/graph.html

Returns a list of attribute names present in the Graph object, excluding those with value None or starting with underscore.

```APIDOC
## Graph.keys

### Description
Returns a list of attribute names present in the Graph object, excluding those with value None or starting with underscore.

### Method Signature
`keys(self) -> list`

### Returns
- **list**: List of attributes in the Graph object.
```

--------------------------------

### Switch DeepSNAP backend to SnapX

Source: https://snap.stanford.edu/deepsnap/modules/deepsnap.html

Use deepsnap.use(sx) to set SnapX as the backend. Requires importing snap and snapx as sx. The default backend is NetworkX; SnapX is supported only for undirected homogeneous graphs.

```python
import snap
import snapx as sx
import deepsnap
deepsnap.use(sx)

```

--------------------------------

### HeteroGraph.__init__

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/hetero_graph.html

Initializes a HeteroGraph object from a NetworkX graph or keyword arguments. Validates edge_index format for undirected graphs in tensor backend mode.

```APIDOC
## HeteroGraph.__init__

### Description
Initializes a HeteroGraph object. Accepts a NetworkX graph or keyword arguments for tensor backend. Validates edge_index format for undirected graphs.

### Method
Constructor

### Signature
```python
HeteroGraph(G=None, **kwargs)
```

### Parameters
- **G** (networkx.classes.graph, optional) - A NetworkX graph object containing features and labels for each node type and edge type.
- **kwargs** (dict, optional) - Keyworded argument list with keys such as "node_feature", "node_label", and corresponding attributes.

### Raises
- **ValueError** - If edge_index is not provided in tensor backend mode, or if edge_index format is invalid for undirected graphs.

```

--------------------------------

### GraphDataset.num_dims_dict

Source: https://snap.stanford.edu/deepsnap/modules/dataset.html

Returns a dictionary containing the dimensions of all fields in the graphs. For example, if graphs have node_feature and graph_label attributes, the dictionary will have keys for each and values for their dimensions.

```APIDOC
## GraphDataset.num_dims_dict

### Description
Returns a dictionary containing the dimensions of all fields in the graphs. For example, if graphs have node_feature and graph_label attributes, the dictionary will have keys for each and values for their dimensions.

### Method
num_dims_dict

### Returns
Dimensions of all fields. For example, if graphs have two attributes the node_feature and the graph_label. The returned dictionary will have two keys, node_feature and graph_label, and two values, node feature dimension and graph label dimension.

### Return Type
dict
```

--------------------------------

### Create HeteroGraph from NetworkX and iterate over attributes

Source: https://snap.stanford.edu/deepsnap/notes/introduction.html

Constructs a DeepSNAP HeteroGraph from a NetworkX directed graph with node types, features, labels, and edge types. Iterating over the HeteroGraph yields tuples of attribute names and values, including edge_index, node_feature, node_label, and mapping tensors. The output shows the printed attributes for the example graph.

```python
import torch
import networkx as nx
from deepsnap.hetero_graph import HeteroGraph

G = nx.DiGraph()
G.add_node(0, node_type='n1', node_label=1, node_feature=torch.Tensor([0.1, 0.2, 0.3]))
G.add_node(1, node_type='n1', node_label=0, node_feature=torch.Tensor([0.2, 0.3, 0.4]))
G.add_node(2, node_type='n2', node_label=1, node_feature=torch.Tensor([0.3, 0.4, 0.5]))
G.add_edge(0, 1, edge_type='e1')
G.add_edge(0, 2, edge_type='e1')
G.add_edge(1, 2, edge_type='e2')
H = HeteroGraph(G)
for hetero_feature in H:
    print(hetero_feature)
```

```text
>>> ('G', <networkx.classes.digraph.DiGraph object at 0x103642370>)
    ('edge_index', {('n1', 'e1', 'n1'): tensor([[0],
            [1]]), ('n1', 'e1', 'n2'): tensor([[0],
            [0]]), ('n1', 'e2', 'n2'): tensor([[1],
            [0]])})
    ('edge_label_index', {('n1', 'e1', 'n1'): tensor([[0],
            [1]]), ('n1', 'e1', 'n2'): tensor([[0],
            [0]]), ('n1', 'e2', 'n2'): tensor([[1],
            [0]])})
    ('edge_to_graph_mapping', {('n1', 'e1', 'n1'): tensor([0]), ('n1', 'e1', 'n2'): tensor([1]), ('n1', 'e2', 'n2'): tensor([2])})
    ('edge_to_tensor_mapping', tensor([0, 0, 0]))
    ('edge_type', {('n1', 'e1', 'n1'): ['e1'], ('n1', 'e1', 'n2'): ['e1'], ('n1', 'e2', 'n2'): ['e2']})
    ('node_feature', {'n1': tensor([[0.1000, 0.2000, 0.3000],
            [0.2000, 0.3000, 0.4000]]), 'n2': tensor([[0.3000, 0.4000, 0.5000]])})
    ('node_label', {'n1': tensor([1, 0]), 'n2': tensor([1])})
    ('node_label_index', {'n1': tensor([0, 1]), 'n2': tensor([0])})
    ('node_to_graph_mapping', {'n1': tensor([0, 1]), 'n2': tensor([2])})
    ('node_to_tensor_mapping', tensor([0, 1, 0]))
    ('node_type', {'n1': ['n1', 'n1'], 'n2': ['n2']})
```

--------------------------------

### deepsnap.use(netlib=None)

Source: https://snap.stanford.edu/deepsnap/modules/deepsnap.html

Specifies which graph library to use as the DeepSNAP backend. The default is NetworkX, and SnapX is also supported for undirected homogeneous graphs.

```APIDOC
## deepsnap.use(netlib=None)

### Description
Specifies which graph library to use as the DeepSNAP backend. The default backend is NetworkX. Current DeepSNAP also supports Snap.py (SnapX) backend for undirected homogeneous graph.

### Parameters
- **netlib** (types.ModuleType, optional) - The graph backend module. Currently DeepSNAP supports the NetworkX and SnapX (for SnapX only the undirected homogeneous graph) as the graph backend. Default graph backend is the NetworkX.

### Example
```
import snap
import snapx as sx
import deepsnap
deepsnap.use(sx)
```
```

--------------------------------

### Split dataset for node classification with transductive=True

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Creates a node classification dataset with node features and labels, then splits it into train, validation, and test sets with a ratio of 0.8, 0.1, 0.1. With transductive=True, each graph's nodes are split, so each split contains all graphs.

```python
import torch
import networkx as nx
from deepsnap.graph import Graph
from deepsnap.dataset import GraphDataset

G = nx.complete_graph(100)
Graph.add_node_attr(G, 'node_feature', torch.zeros([100, 1]))
Graph.add_node_attr(G, 'node_label', torch.zeros([100, 1]))
H1 = Graph(G)
H2 = H1.clone()
dataset = GraphDataset(graphs=[H1, H2], task='node')

train, val, test = dataset.split(transductive=True, split_ratio=[0.8, 0.1, 0.1])
print(train, val, test)
>>> GraphDataset(2) GraphDataset(2) GraphDataset(2)
```

--------------------------------

### Process Homogeneous Graphs with DeepSNAP

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Demonstrates loading a PyG dataset, converting it to DeepSNAP graphs, applying a custom clustering transformation, and inspecting the resulting batch.

```python
import torch
import networkx as nx
from deepsnap.batch import Batch
from deepsnap.dataset import GraphDataset
from torch_geometric.datasets import TUDataset
from torch.utils.data import DataLoader

def clustering_func(graph):
    clustering = torch.tensor(list(nx.clustering(graph.G).values()))
    clustering = clustering.view(-1, 1)
    graph.node_feature = torch.cat([graph.node_feature, clustering], dim=1)

pyg_dataset = TUDataset('./enzymes', 'ENZYMES')
graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
dataset = GraphDataset(graphs, task='graph', minimum_node_per_graph=0)
train, val, test = dataset.split(
            transductive=False, split_ratio = [0.8, 0.1, 0.1])
train_loader = DataLoader(train, collate_fn=Batch.collate(), batch_size=10, shuffle=True)
batch = next(iter(train_loader))
batch = batch.apply_transform(clustering_func, update_graph=True, update_tensor=False)
print(batch)
>>> Batch(G=[10], batch=[411], edge_index=[2, 1378], edge_label_index=[2, 1378], graph_label=[10], node_feature=[411, 4], node_label_index=[411])
print(nx.get_node_attributes(batch.G[0], 'node_feature')[0].shape[0])
>>> 4
```

--------------------------------

### GraphDataset.__init__

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/dataset.html

Constructor for the GraphDataset class. Initializes the dataset with graphs, task type, and various configuration parameters for edge sampling, splitting, and generation.

```APIDOC
## GraphDataset.__init__

### Description
Initializes a GraphDataset instance for graph-based machine learning tasks. Accepts a list of graphs, specifies the task type, and configures edge sampling, splitting, and generation options.

### Method
Constructor

### Signature
```python
GraphDataset(
    graphs: List[Graph] = None,
    task: str = "node",
    custom_split_graphs: List[Graph] = None,
    edge_negative_sampling_ratio: float = 1,
    edge_message_ratio: float = 0.8,
    edge_train_mode: str = "all",
    edge_split_mode: str = "exact",
    minimum_node_per_graph: int = 5,
    generator=None,
    resample_negatives: bool = False,
    resample_disjoint: bool = False,
    resample_disjoint_period: int = 1,
    negative_label_val: int = None,
    netlib=None
)
```

### Parameters
- **graphs** (list, optional) - A list of deepsnap.graph.Graph objects.
- **task** (str) - The task type: 'node', 'edge', 'link_pred', or 'graph'.
- **custom_split_graphs** (list) - A list of 2 (train and val) or 3 (train, val, test) lists of split graphs for custom split in graph task.
- **edge_negative_sampling_ratio** (float) - Number of negative samples compared to positive edges. Default: 1.
- **edge_message_ratio** (float) - Number of message passing edges compared to training supervision edges. Default: 0.8.
- **edge_train_mode** (str) - 'all' or 'disjoint'. In 'all' mode, training supervision edges are same as message passing edges. In 'disjoint' mode, they are different.
- **edge_split_mode** (str) - 'exact' or 'approximate'. For heterogeneous graphs, 'exact' splits according to ratio and split type, 'approximate' splits regardless of split type.
- **minimum_node_per_graph** (int) - Graphs with fewer nodes than this value are filtered out. Default: 5.
- **generator** (Generator) - On-the-fly generator used if graphs is empty or None and generator has an overwritten generate method.
- **resample_negatives** (bool) - Whether to resample negative edges in each iteration of link_pred task. Default: False.
- **resample_disjoint** (bool) - Whether to resample disjoint training edges in disjoint link_pred task. Default: False.
- **resample_disjoint_period** (int) - Number of iterations after which disjoint training edges are resampled. Default: 1.
- **negative_label_val** (int, optional) - Value of negative edges generated in link_pred task. Required for tensor backend custom split.
- **netlib** (types.ModuleType, optional) - Graph backend module. Supports NetworkX and SnapX (undirected homogeneous only). Default: NetworkX.

### Raises
- **ValueError** - If task is not one of 'node', 'edge', 'link_pred', 'graph'.
- **ValueError** - If edge_train_mode is not 'all' or 'disjoint'.
- **ValueError** - If edge_split_mode is not 'exact' or 'approximate'.

### Notes
- If netlib is provided, it sets the global deepsnap._netlib.
- If graphs is not a list, it is converted to a list.
- Non-Graph objects in graphs are converted to Graph instances using the provided netlib.

```

--------------------------------

### Handle edge splitting for heterogeneous graphs

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/hetero_graph.html

This logic handles the partitioning of edges into training, validation, and test sets based on a provided split ratio. It includes edge cases to ensure each set contains at least one edge when the split ratio results in zero.

```python
for split_type in message_types_sorted:
    if split_type not in split_types:
        edges_train[split_type] = edges[split_type]
        edges_val[split_type] = edges[split_type]

elif len(split_ratio) == 3:
    num_edges_train = int(
        split_ratio[0] * edges_split_type_length
    )
    num_edges_val = int(
        split_ratio[1] * edges_split_type_length
    )
    num_edges_test = (
        edges_split_type_length
        - num_edges_train
        - num_edges_val
    )
    if (
        (num_edges_train == 0)
        or (num_edges_val == 0)
        or (num_edges_test == 0)
    ):
        num_edges_train = (
            1
            + int(
                split_ratio[0]
                * (edges_split_type_length - 3)
            )
        )
        num_edges_val = (
            1
            + int(
                split_ratio[1]
                * (edges_split_type_length - 3)
            )
        )

    edges_train_index = sorted(
        edge_index[:num_edges_train]
    )
    edges_val_index = sorted(
        edge_index[
            num_edges_train:num_edges_train + num_edges_val
        ]
    )
    edges_test_index = sorted(
        edge_index[num_edges_train + num_edges_val:]
    )
    edges_train, edges_val, edges_test = {}, {}, {}
    split_type_cnt = 0
    for index in edges_train_index:
        while not (
            cumulative_split_type_cnt[split_type_cnt]
            <= index
            < cumulative_split_type_cnt[split_type_cnt+1]
        ):
            split_type_cnt += 1
        split_type = split_types_sorted[split_type_cnt]
        index_type = (
            index
            - cumulative_split_type_cnt[split_type_cnt]
        )
        if split_type not in edges_train:
            edges_train[split_type] = []
        edges_train[split_type].append(index_type)

    split_type_cnt = 0
    for index in edges_val_index:
        while not (
            cumulative_split_type_cnt[split_type_cnt]
            <= index
            < cumulative_split_type_cnt[split_type_cnt+1]
        ):
            split_type_cnt += 1
        split_type = split_types_sorted[split_type_cnt]
        index_type = (
            index
            - cumulative_split_type_cnt[split_type_cnt]
        )
        if split_type not in edges_val:
            edges_val[split_type] = []
        edges_val[split_type].append(index_type)

    split_type_cnt = 0
    for index in edges_test_index:
        while not (
            cumulative_split_type_cnt[split_type_cnt]
            <= index
            < cumulative_split_type_cnt[split_type_cnt+1]
        ):
            split_type_cnt += 1
        split_type = split_types_sorted[split_type_cnt]
        index_type = (
            index
            - cumulative_split_type_cnt[split_type_cnt]
        )
        if split_type not in edges_test:
            edges_test[split_type] = []
        edges_test[split_type].append(index_type)
```

--------------------------------

### Collate dataset into batches with DataLoader and apply transform

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Creates a DataLoader with a batch size of 10 using Batch.collate() as the collate function. Splits a graph classification dataset, loads a batch, applies the clustering transform to the batch, and prints the batch, showing 10 graphs and the added node_clustering attribute.

```python
import networkx as nx
from deepsnap.batch import Batch
from deepsnap.dataset import GraphDataset
from torch_geometric.datasets import TUDataset
from torch.utils.data import DataLoader

def clustering_func(graph):
    clustering = list(nx.clustering(graph.G).values())
    graph['node_clustering'] = clustering

pyg_dataset = TUDataset('./enzymes', 'ENZYMES')
graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
dataset = GraphDataset(graphs, task='graph', minimum_node_per_graph=0)
train, val, test = dataset.split(
            transductive=False, split_ratio = [0.8, 0.1, 0.1])
train_loader = DataLoader(train, collate_fn=Batch.collate(), batch_size=10, shuffle=True)
batch = next(iter(train_loader))
batch = batch.apply_transform(clustering_func, update_graph=True, update_tensor=False)
print(batch)
>>> Batch(G=[10], batch=[266], edge_index=[2, 1064], edge_label_index=[2, 1064], graph_label=[10], node_clustering=[10], node_feature=[266, 3], node_label_index=[266])
```

--------------------------------

### Implement __repr__ for dataset representation

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/dataset.html

Provides a string representation of the dataset object, showing either the number of graphs or the generator class name.

```python
def __repr__(self) -> str:  # pragma: no cover
    descriptor = (
        len(self) if self.graphs is not None else self.generator.__class__
    )
    return f"{self.__class__.__name__}({descriptor})"
```

--------------------------------

### to(device, *keys)

Source: https://snap.stanford.edu/deepsnap/modules/graph.html

Transfers tensor attributes to a specified device. If *keys is not given, all present attributes are transferred.

```APIDOC
## to(device, *keys)

### Description
Transfers tensor attributes to specified device for all attributes specified in *keys. If *keys is not given, applies to all present attributes.

### Parameters
- **device** (str) - Required - Specified device name, such as 'cpu' or 'cuda'.
- ***keys** (str, optional) - Optional - Tensor attributes to transfer.

### Returns
None
```

--------------------------------

### raw_to_graph(data)

Source: https://snap.stanford.edu/deepsnap/modules/graph.html

Static method intended for users to import their own data format. Not implemented in the base class.

```APIDOC
## raw_to_graph(data)

### Description
Static method for users to import their own data format. Not implemented in the base class.

### Parameters
- **data** - Required - The raw data to convert.

### Returns
Not implemented.
```

--------------------------------

### Batch.from_data_list

Source: https://snap.stanford.edu/deepsnap/modules/batch.html

Constructs a Batch object from a list of graph objects. The assignment vector 'batch' is created on the fly, and additional assignment batch vectors are created for keys in follow_batch.

```APIDOC
## Batch.from_data_list

### Description
Constructs a deepsnap.batch.Batch object from a python list holding torch_geometric.data.Data objects. The assignment vector 'batch' is created on the fly. Additionally, creates assignment batch vectors for each key in follow_batch.

### Method
from_data_list (static)

### Signature
from_data_list(data_list, follow_batch=None, transform=None, **kwargs)

### Parameters
- **data_list** (list) - Required - A list of deepsnap.graph.Graph objects.
- **follow_batch** (list) - Optional - Creates assignment batch vectors for each key.
- **transform** (callable) - Optional - If it is not None, apply transform when batching.
- **kwargs** - Optional - Other parameters.

### Returns
A Batch object constructed from the given data list.
```

--------------------------------

### Batch class definition and initialization

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap/batch.html

The Batch class inherits from Graph and initializes with an optional batch assignment vector.

```python
import torch
from deepsnap.graph import Graph
from deepsnap.hetero_graph import HeteroGraph
from typing import (
    Callable,
    Dict,
    List
)



[docs]class Batch(Graph):
    r"""
    A plain old python object modeling a batch of
    :class:`deepsnap.graph.Graph` objects as one big (disconnected) graph,
    with :class:`torch_geometric.data.Data` being the
    base class that all its methods can also be used here.
    In addition, graphs can be reconstructed via the assignment vector
    :obj:`batch`, which maps each node to its respective graph identifier.

    .. note::

        For more detailed use of :class:`deepsnap.batch.Batch`, see the `examples 
        <https://github.com/snap-stanford/deepsnap/tree/master/examples>`_ folder.

    """
    def __init__(self, batch=None, **kwargs):
        super(Batch, self).__init__(**kwargs)

        self.batch = batch
        self.__data_class__ = Graph
        self.__slices__ = None

    @staticmethod
    def collate(follow_batch=[], transform=None, **kwargs):
        return lambda batch: Batch.from_data_list(
            batch, follow_batch, transform, **kwargs
        )


[docs]    @staticmethod
    def from_data_list(
        data_list: List[Graph],
        follow_batch: List = None,
        transform: Callable = None,
        **kwargs
    ):
        r"""
        Constructs A :class:`deepsnap.batch.Batch` object from a python list
        holding :class:`torch_geometric.data.Data` objects.
        The assignment vector :obj:`batch` is created on the fly.
        Additionally, creates assignment batch vectors for each key in
        :obj:`follow_batch`.

        Args:
            data_list (list): A list of :class:`deepsnap.graph.Graph` objects.
            follow_batch (list): Creates assignment batch vectors
                for each key.
            transform (callable): If it is not `None`, apply transform 
                when batching.
            **kwargs: Other parameters.
        """
        if follow_batch is None:
            follow_batch = []
        if transform is not None:
            data_list = [
                data.apply_transform(
                    transform,
                    deep_copy=True,
                    **kwargs,
                )
                for data in data_list
            ]
        keys = [set(data.keys) for data in data_list]
        keys = list(set.union(*keys))
        assert "batch" not in keys

        batch, cumsum = Batch._init_batch_fields(keys, follow_batch)
        batch.__data_class__ = data_list[0].__class__
        batch.batch = []
        for i, data in enumerate(data_list):
            # Note: in heterogeneous graph, __inc__ logic is different
            Batch._collate_dict(
                data, cumsum,
                batch.__slices__, batch,
                data, follow_batch, i=i
            )
            if isinstance(data, Graph):
                if isinstance(data, HeteroGraph):
                    num_nodes = sum(data.num_nodes().values())
                else:
                    num_nodes = data.num_nodes
            else:
                raise TypeError(
                    "element in self.graphs of unexpected type"
                )
            if num_nodes is not None:
                item = torch.full((num_nodes, ), i, dtype=torch.long)
                batch.batch.append(item)

        if num_nodes is None:
            batch.batch = None

        Batch._dict_list_to_tensor(batch, data_list[0])

        return batch.contiguous()


    @staticmethod
    def _init_batch_fields(keys, follow_batch):
        batch = Batch()
        batch.__slices__ = {key: [0] for key in keys}

        for key in keys:
            batch[key] = []

        for key in follow_batch:
            batch[f"{key}_batch"] = []

        cumsum = {key: 0 for key in keys}
        return batch, cumsum

    @staticmethod
    def _collate_dict(
        curr_dict,
        cumsum: Dict[str, int],
        slices,
        batched_dict,
        graph,
        follow_batch,
        i=None
    ):
        r""" Called in from_data_list to collate a dictionary.
        This can also be applied to Graph object, since it has support for
        keys and __getitem__().

```

--------------------------------

### Import deepsnap modules and define __version__

Source: https://snap.stanford.edu/deepsnap/_modules/deepsnap.html

Imports required modules and defines the package version.

```python
import types
import torch
import random
import numpy as np
import deepsnap.graph
import deepsnap.dataset
import deepsnap.batch
import deepsnap.hetero_graph
import deepsnap.hetero_gnn

import networkx as _netlib

__version__ = "0.2.0"
```

--------------------------------

### Create DeepSNAP Graph with node features from NetworkX

Source: https://snap.stanford.edu/deepsnap/_sources/notes/introduction.rst.txt

Builds a NetworkX graph with node_feature attributes and converts it to a DeepSNAP Graph. The node features are automatically loaded and printed as a tensor.

```python
import torch
import networkx as nx
rom deepsnap.graph import Graph

G = nx.Graph()
G.add_node(0, node_feature=torch.tensor([1,2,3]))
G.add_node(1, node_feature=torch.tensor([4,5,6]))
G.add_edge(0, 1)
H = Graph(G)
print(H.node_feature)
>>> tensor([[1, 2, 3],
            [4, 5, 6]])
```
Tokens:
Raw