Is there a way to split the graph based on an edge feature

[ntquanghai]

I came across a tutorial which can randomly split the graph into training, testing and validation sets by using the RandomLinkSplit class. I was wondering if there was a similar class, similar to this which is able to split the data based on a condition, or a based on an edge feature (For example, split the data based on "year")? I opted to find a similar implementation since I do need the edge_label and edge_label_index produced by this. Below is the tutorial code:

import os.path as osp

import torch
from sklearn.metrics import roc_auc_score

import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv, RGCNConv
from torch_geometric.utils import negative_sampling

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = T.Compose([
    T.NormalizeFeatures(),
    T.ToDevice(device),
    T.RandomLinkSplit(num_val=0.05, num_test=0.1, is_undirected=True,
                      add_negative_train_samples=False),
])
dataset = Planetoid('.', name='Cora', transform=transform)

# After applying the `RandomLinkSplit` transform, the data is transformed from
# a data object to a list of tuples (train_data, val_data, test_data), with
# each element representing the corresponding split.
train_data, val_data, test_data = dataset[0]

[rusty1s]

Currently, this needs to be manually defined by the user. We don't offer such a functionality besides random splitting. This should be quite straightforward, as you can simply sort edges/nodes by time, and then split them manually by a pre-defined ratio, e.g.:

perm = edge_time.argsort()
train_perm = perm[:perm.numel() * 0.8]
val_perm = perm[perm.numel() * 0.8:perm.numel() * 0.9]
test_perm = perm[perm.numel() * 0.9:]

train_edge_index = edge_index[:, train_perm]
...

Hope this helps.