Understanding LinkNeighborLoader

[GianlucaDeStefano]

Hi everybody, I am trying to understand exactly how the LinkNeighborLoader works in practice with HeteroGraphs.

I have this minimal example:

import random

import torch
from torch_geometric.data import HeteroData
from torch_geometric.loader import LinkNeighborLoader
import torch_geometric.transforms as T

transforms = T.Compose([
    T.NormalizeFeatures(),
    T.AddSelfLoops(),
    T.ToUndirected(),
])

data = HeteroData()

data['A'].node_id = torch.arange(1000)
data['B'].node_id = torch.arange(1000)
data['C'].node_id = torch.arange(1000)

data['A', 'edge_1', 'B'].edge_index = torch.asarray(([random.randint(0, 999) for i in range(4000)], [random.randint(0, 999) for i in range(4000)]))
data['A', 'edge_1', 'B'].edge_attr = torch.arange(4000)


data['A', 'edge_2', 'C'].edge_index = torch.asarray(([random.randint(0, 999) for i in range(4000)], [random.randint(0, 999) for i in range(4000)]))
data['A', 'edge_2', 'C'].edge_attr = torch.arange(4000)

loader = LinkNeighborLoader(
    data=data,
    num_neighbors=[10,10],
    neg_sampling_ratio=1,
    edge_label_index=(("A", "edge_1", "B"), data['A', 'edge_1', 'B'].edge_index),
    edge_label=torch.ones(4000),
    batch_size=32,
    shuffle=True,
    drop_last=True
)

for batch in loader:
    print(batch)

Here a graph contains 3 distinct types of nodes('A','B','C') and 2 types of edges ('edge_1','edge_2') that connect respectively 'A' to 'B' and 'A' to 'C'.
And by sampling a batch using the LinkNeighbourLoader I would expect to receive a subgraph containing all 3 types of nodes.

However, the batches that the loader returns only contain the nodes 'A' and 'B':

HeteroData(
  A={ node_id=[239] },
  B={ node_id=[60] },
  C={ node_id=[0] },
  (A, edge_1, B)={
    edge_index=[2, 246],
    edge_attr=[246],
    input_id=[32],
    edge_label_index=[2, 64],
    edge_label=[64]
  },
  (A, edge_2, C)={
    edge_index=[2, 0],
    edge_attr=[0]
  }

Why does this happen? Furthermore, I noticed that if I reverse the direction of the second edge to make it point from 'C' to 'A'
I get a batch that this time contains also 'C' nodes. To me this makes even less sense, why should nodes in 'C' be considered in the neighbourhood of nodes 'A' only if there is an outgoing edge from 'C' to 'A' and not the contrary?

I am obviously misunderstanding something here. In my view, the LinkNeighborLoader should create a subgraph of the neighbours of both the root nodes of the sampled edge.

Thank you in advance for your help.

[GianlucaDeStefano]

Ok, this makes perfect sense as the LinkNeighborLoader samples only neighbours from which the label node is reachable and not vice versa. I was looking at it wrong. :)

[Dhano]

HI @GianlucaDeStefano , If i want a negative sampling only on ['A', 'edge_1', 'B'] edge , where as both edges should participate in message passing . I need to pass as edge_label_index = ['A', 'edge_1', 'B'] . right?
Also can you help with what should be edge label?