utils.py

Utility Functions and Helper Transforms

class rings.utils.Shuffle(shuffle_edges=False, shuffle_features=False, generator=None)

A transform that shuffles node features and/or edges in a graph.

This transform can randomize the relationship between node features and graph structure by shuffling feature vectors among nodes, shuffling edges between nodes, or both. It preserves the feature and edge distributions while destroying their correlations.

Parameters:

• shuffle_edges (bool, default=False) – If True, shuffle the order of edges and reassign them to random nodes.

• shuffle_features (bool, default=False) – If True, shuffle the order of node features among nodes.

• generator (torch.Generator, optional) – Random number generator for reproducibility. If None, a new generator is created.

Examples

>>> import torch
>>> from torch_geometric.data import Data
>>> from rings.utils import Shuffle
>>>
>>> # Create a simple graph with features and edges
>>> x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float)
>>> edge_index = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
>>> data = Data(x=x, edge_index=edge_index)
>>>
>>> # Shuffle only features
>>> transform1 = Shuffle(shuffle_features=True)
>>> t1_data = transform1(data.clone())
>>> # Feature vectors are preserved but reassigned among nodes
>>> original_features = set(tuple(f.tolist()) for f in data.x)
>>> shuffled_features = set(tuple(f.tolist()) for f in t1_data.x)
>>> assert original_features == shuffled_features
>>> assert t1_data.edge_index.equal(data.edge_index)
>>>
>>> # Shuffle only edges
>>> transform2 = Shuffle(shuffle_edges=True)
>>> t2_data = transform2(data.clone())
>>> # Edge structure is changed, features unchanged
>>> assert t2_data.x.equal(data.x)
>>> assert not t2_data.edge_index.equal(data.edge_index)

__init__(shuffle_edges=False, shuffle_features=False, generator=None)

Initialize the Shuffle transform.

Parameters:

• shuffle_edges (bool, default=False) – If True, shuffle the order of edges and reassign them to random nodes.

• shuffle_features (bool, default=False) – If True, shuffle the order of node features among nodes.

• generator (torch.Generator, optional) – Random number generator for reproducibility. If None, a new generator is created.

rings.utils.is_connected(data)

Check if a graph is connected.

A connected graph has a path between every pair of nodes. This function converts the PyTorch Geometric graph to NetworkX format and uses NetworkX’s is_connected function to check connectivity.

Parameters:

data (torch_geometric.data.Data) – Graph data object.

Returns:

True if the graph is connected, False otherwise. Empty graphs (0 nodes) are considered connected by convention.

Return type:

bool

Examples

>>> import torch
>>> from torch_geometric.data import Data
>>> from rings.utils import is_connected
>>>
>>> # Create a connected graph: 0-1-2
>>> edge_index1 = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
>>> data1 = Data(edge_index=edge_index1, num_nodes=3)
>>> assert is_connected(data1) == True
>>>
>>> # Create a disconnected graph: 0-1  2
>>> edge_index2 = torch.tensor([[0], [1]], dtype=torch.long)
>>> data2 = Data(edge_index=edge_index2, num_nodes=3)
>>> assert is_connected(data2) == False
>>>
>>> # Empty graph
>>> data3 = Data(edge_index=torch.zeros((2, 0), dtype=torch.long), num_nodes=0)
>>> assert is_connected(data3) == True  # Connected by convention

rings.utils.ensure_no_self_loops(source_nodes, target_nodes, num_nodes, generator)

Ensure no self-loops exist in the shuffled edges.

If any self-loops are detected after shuffling, they are replaced with random edges (by assigning a random target node) to maintain the original number of edges.

Parameters:

• source_nodes (torch.Tensor) – Source nodes of the edges.

• target_nodes (torch.Tensor) – Target nodes of the edges.

• num_nodes (int) – Number of nodes in the graph.

Returns:

Updated (source_nodes, target_nodes) without self-loops.

Return type:

tuple of torch.Tensor

Notes

This method iteratively replaces target nodes that would create self-loops until all self-loops are eliminated. The source nodes are kept intact.