Link Prediction on Heterogeneous Graphs

[naarkhoo]

related to issue #3958 I have difficulties building a simple link prediction model on Heterogeneous data.

here is a simulated data where I have protein ligand interaction data (binary) and have a set of feature on both ligand and proteins;

import torch
import torch_geometric
from torch_geometric.data import HeteroData

# Set random seed for reproducibility
torch.manual_seed(1234)

# Define number of binding sites and ligands
num_bs = 10
num_ligands = 5

# Define random binary interaction matrix between binding sites and ligands
interaction = torch.randint(2, (num_bs, num_ligands))

# Define random features for binding sites
bs_features = torch.randn((num_bs, 3))

# Define random features for ligands
ligand_features = torch.randn((num_ligands, 2))

# Construct HeteroData object
data = HeteroData()

# Add binding sites as a node type
data['bs'].x = bs_features

# Add ligands as a node type
data['ligand'].x = ligand_features

# Add edges between binding sites and ligands
edge_list = []
for bs_idx in range(num_bs):
    for ligand_idx in range(num_ligands):
        if interaction[bs_idx][ligand_idx] == 1:
            edge_list.append((bs_idx, ligand_idx))

data['bs', 'binds', 'ligand'].edge_index = torch.tensor(edge_list).t()

print(data)

this is my link prediction model

import torch
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, negative_sampling
from sklearn.metrics import roc_auc_score

# Set random seed for reproducibility
torch.manual_seed(1234)

# Define the link prediction model
class LinkPredictionModel(MessagePassing):
    def __init__(self, hidden_channels):
        super().__init__(aggr='mean')

        # Define linear transformations for each node type
        self.bs_lin = torch.nn.Linear(3, hidden_channels)
        self.ligand_lin = torch.nn.Linear(2, hidden_channels)

        # Define the relation-specific transformations
        self.binds_lin = torch.nn.Linear(hidden_channels, hidden_channels)

    def forward(self, x_dict, edge_index_dict):
        # Get the node representations for each node type
        bs_x = self.bs_lin(x_dict['bs'])
        ligand_x = self.ligand_lin(x_dict['ligand'])

        # Get the edge indices for each edge type
        binds_edge_index = edge_index_dict[('bs', 'binds', 'ligand')]

        # Add self-loops for each node type
        bs_edge_index = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand')][0], num_nodes=bs_x.size(0))
        ligand_edge_index = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand')][1], num_nodes=ligand_x.size(0))

        # Update the node representations using message passing
        bs_x = self.propagate(bs_edge_index, x=bs_x)
        ligand_x = self.propagate(ligand_edge_index, x=ligand_x)

        # Get the edge representations for each edge type
        binds_e = self.binds_lin(bs_x[binds_edge_index[0]] * ligand_x[binds_edge_index[1]])

        # Compute the edge scores for each edge type
        pos_edge_scores = (bs_x[binds_edge_index[0]] * binds_e * ligand_x[binds_edge_index[1]]).sum(dim=1)
        neg_edge_scores = (bs_x[binds_edge_index[0]] * binds_e * ligand_x[negative_sampling(binds_edge_index, num_nodes=ligand_x.size(0))[1]]).sum(dim=1)

        return torch.cat([pos_edge_scores, neg_edge_scores])

    def message(self, x_j):
        # Pass the node features as messages
        return x_j

    def update(self, aggr_out):
        # Update the node representations using mean aggregation
        return aggr_out

and

# Define the training function
def train(model, optimizer, criterion, data, device):
    model.train()

    optimizer.zero_grad()

    # Forward pass
    out = model(data)

    # Compute loss
    loss = criterion(out.view(-1), data.y.view(-1))

    # Backward pass
    loss.backward()
    optimizer.step()

    return loss.item()


# Define the testing function
def test(model, criterion, test_loader, device):
    model.eval()

I have put the code in this colab notebook where you can play with it - I keep getting error ... I appreciate any hint

[rusty1s]

You need to call your model via model(data.x_dict, data.edge_index_dict).

[naarkhoo]

thanks but then I get "IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)" from

[rusty1s]

It is crashing in add_self_loops now, because the expected input is only one-dimensional. Try

# Add self-loops for each node type
bs_edge_index = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand')], num_nodes=bs_x.size(0))
ligand_edge_index = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand'), num_nodes=ligand_x.size(0))

[naarkhoo]

now got

"ValueError: MessagePassing.propagate only supports integer tensors of shape [2, num_messages], torch_sparse.SparseTensoror torch.sparse.Tensor for argument edge_index

I tried to fix the conditionality but got some device related error.

[rusty1s]

add_self_loops returns a tuple of edge indices and weights, so it should be

bs_edge_index, _ = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand')], num_nodes=bs_x.size(0))
ligand_edge_index, _ = add_self_loops(edge_index_dict[('bs', 'binds', 'ligand')], num_nodes=ligand_x.size(0))

Anyway, I think there are a few issues in your code.

1. Since you are working on bipartite graphs, add_self_loops isn't really well-defined. I suggest to drop it.
2. Your message passing should actually work on a tuple of node features to correctly exchange messages in bipartite graphs.

Here is an updated example, which at least gets you past message passing:

bs_edge_index = edge_index_dict[('bs', 'binds', 'ligand')]
ligand_edge_index = edge_index_dict[('bs', 'binds', 'ligand')].flip([0])

ligand_x = self.propagate(bs_edge_index, x=(bs_x, ligand_x))
bs_x = self.propagate(ligand_edge_index, x=(ligand_x, bs_x))