Inductive train/test splitting graph

[ogawayuto]

Hi, I have an idea for new features about experiments setting.
When working on inductive settings for a single graph, it's crucial to establish train/test masks for nodes/links to ensure that test samples don't appear in the training phase. While you can use masks and the subgraph() function to implement inductive node/link tasks, there don't appear to be explicit functions for splitting data into training and testing sets for inductive settings. Moreover, in the case of inductive link prediction, it's essential to identify which edges are designated for training and testing, but the subgraph() function doesn't provide information about test edges.

To address this, I've created an inductive_train_test_split() function that facilitates the splitting of a graph into a train graph and a test graph. This function allows you to specify which edges should be used for training and testing, ensuring a clear separation of data for inductive tasks.

I'm considering sharing and committing this function on GitHub. Would there be any interest or demand for such a contribution?

from typing import List, Union, Tuple
from torch_geometric.typing import OptTensor
from torch_geometric.data import Data
from torch_geometric.utils import subgraph
from torch_geometric.utils.mask import index_to_mask
from torch_geometric.utils.num_nodes import maybe_num_nodes
import torch
from torch import Tensor

def split_graph(
    data: Data,
    subset1: Union[Tensor, List[int]],
    subset2: Union[Tensor, List[int]],
    bridge: bool = True,
) -> Tuple[Tensor, Tensor, OptTensor, OptTensor, OptTensor, OptTensor]:

    device = data.edge_index.device

    if isinstance(subset1, (list, tuple)):
        subset1 = torch.tensor(subset1, dtype=torch.long, device=device)

    if isinstance(subset2, (list, tuple)):
        subset2 = torch.tensor(subset2, dtype=torch.long, device=device)

    if subset1.dtype != torch.bool:
        num_nodes = maybe_num_nodes(data.edge_index, data.num_nodes)
        node_mask1 = index_to_mask(subset1, size=num_nodes)
    else:
        num_nodes = subset1.size(0)
        node_mask1 = subset1
        node_mask1 = node_mask1.nonzero().view(-1)

    if subset2.dtype != torch.bool:
        num_nodes = maybe_num_nodes(data.edge_index, data.num_nodes)
        node_mask2 = index_to_mask(subset2, size=num_nodes)
    else:
        num_nodes = subset2.size(0)
        node_mask2 = subset2
        node_mask2 = node_mask2.nonzero().view(-1)

    edge_index1, edge_attr1 = subgraph(
        node_mask1, data.edge_index, data.edge_attr, return_edge_mask=False
    )
    edge_index2, edge_attr2 = subgraph(
        node_mask2, data.edge_index, data.edge_attr, return_edge_mask=False
    )
    if bridge:
        bridge_mask = (node_mask1[data.edge_index[0]] & node_mask2[data.edge_index[1]]) | (
            node_mask2[data.edge_index[0]] & node_mask1[data.edge_index[1]]
        )
        bridge_edge_index = data.edge_index[:, bridge_mask]
        bridge_edge_attr = data.edge_attr[bridge_mask]
    else:
        bridge_edge_index = None
        bridge_edge_attr = None

    return edge_index1, edge_index2, bridge_edge_index, edge_attr1, edge_attr2, bridge_edge_attr


def inductive_train_test_split(data, train_node, test_node, bridge=True) -> Tuple[Data, Data]:

    device = data.edge_index.device

    if isinstance(train_node, (list, tuple)):
        train_node = torch.tensor(train_node, dtype=torch.long, device=device)

    if isinstance(test_node, (list, tuple)):
        test_node = torch.tensor(test_node, dtype=torch.long, device=device)

    if train_node.dtype != torch.bool:
        num_nodes = maybe_num_nodes(data.edge_index, data.num_nodes)
        node_mask1 = index_to_mask(train_node, size=num_nodes)
    else:
        num_nodes = train_node.size(0)
        node_mask1 = train_node
        node_mask1 = node_mask1.nonzero().view(-1)

    if test_node.dtype != torch.bool:
        num_nodes = maybe_num_nodes(data.edge_index, data.num_nodes)
        node_mask2 = index_to_mask(test_node, size=num_nodes)
    else:
        num_nodes = test_node.size(0)
        node_mask2 = test_node
        node_mask2 = node_mask2.nonzero().view(-1)

    if (node_mask1 & node_mask2).any():
        raise Exception("train and test node must be disjoint, it has intersection")
    if not (node_mask1 | node_mask2).all():
        raise Exception("train and test node must cover all nodes, some nodes are not included")

    (
        train_edge_index,
        test_edge_index,
        bridge_edge_index,
        train_edge_attr,
        test_edge_attr,
        bridge_edge_attr,
    ) = split_graph(data, train_node, test_node, bridge)

    train_data = Data(
        edge_index=train_edge_index,
        edge_attr=train_edge_attr,
    )

    if bridge:
        test_data = Data(
            edge_index=test_edge_index,
            edge_attr=test_edge_attr,
            bridge_edge_index=bridge_edge_index,
            bridge_edge_attr=bridge_edge_attr,
        )
    else:
        test_data = Data(
            edge_index=test_edge_index,
            edge_attr=test_edge_attr,
        )
    if data.x is not None:
        train_data.x = data.x[train_node]
        test_data.x = data.x[test_node]
    if data.y is not None:
        train_data.y = data.y[train_node]
        test_data.y = data.y[test_node]

    return train_data, test_data

[rusty1s]

Sounds wonderful. Please go ahead with your contribution, looking forward to it :)