How to use virtual node with PyG? How many methods are there?

[StefanIsSmart]

Hello everyone, I have a question:

My graph data is fully connected, and they have node_features and edge_features (Not None).

Now I want to add graph_level_features, just like a virtual node.

How to do that with PyG?

I want to update node_features by graph_level_features ; and update graph_level_features by node_features.

The difference is the edge_features between graph_level_features and node_features is None.

How to do that?
Which method is more elegant and effective?

[devanshamin]

Hi @StefanIsSmart ,

Below is an example of adding a virtual node to a molecule graph,

import torch
from torch_geometric.data import Data
from torch_geometric.utils import from_smiles

if __name__ == '__main__':
    
    data: Data = from_smiles(smiles='F/C=C/F')

    # Creating molecule (graph) level features of the same dim
    # as the node level features
    mol_features = torch.randint(0, 9, size=(1, data.x.shape[-1]))
    # Add a extra node (aka virtual node) containing the molecule level features
    data.x = torch.cat((data.x, mol_features))
    
    n = data.num_nodes - 1 # Exclude the virtual node
    # Undirected graph 
    #   All nodes -> Virtual node
    d1 = torch.tensor([list(range(n)), [n] * n])
    #   Virtual node -> All nodes
    d2 = torch.flip(d1, dims=[0])
    #   Append to the `edge_index`
    data.edge_index = torch.cat((data.edge_index, d1, d2), dim=-1)

    # Add edge attr for the newly added edges
    dim1, dim2 = data.edge_attr.shape
    no_of_missing_edge_attr = data.edge_index.shape[1] - dim1
    edge_attr = torch.tensor([0] * dim2).unsqueeze(0).repeat_interleave(no_of_missing_edge_attr, dim=0)
    data.edge_attr = torch.cat((data.edge_attr, edge_attr))

    print(data.x)
    print(data.edge_attr)
    print(data.edge_index)
    print(data.edge_attr)

    # Sort the indices
    perm = (data.edge_index[0] * data.x.size(0) + data.edge_index[1]).argsort()
    data.edge_index, data.edge_attr = data.edge_index[:, perm], data.edge_attr[perm]

    print(data.edge_index)
    print(data.edge_attr)

You can now use any GNN model and perform the message passing all_nodes <-> virtual node.