common_neighbors(edge_index: torch.Tensor, num_nodes: int) → Dict[int, Set[int]]

Create a dictionary of neighbors for each node.

Args:

edge_index (Tensor): Edge indices, shape (2, num_edges). num_nodes (int): Total number of nodes.

Returns:

dict: A dictionary mapping each node to its set of neighbors.

generate_negative_samples(edge_index, num_nodes, num_neg_samples, max_attempts=1000)

Generate negative samples for link prediction.

Args:

edge_index (Tensor): Edge indices, shape (2, num_edges). num_nodes (int): Total number of nodes. num_neg_samples (int): Number of negative samples. max_attempts (int, optional): Max attempts to find samples (default 1000).

Returns:

list: Negative samples, each a pair of nodes [node1, node2].

Example:

generate_negative_samples(torch.tensor([[0, 1], [1, 2]]), 3, 2) # Returns [[0, 2], [1, 0]]