torch_geometric.transforms.AddRandomWalkPE doesn't support HeteroData()

[Jerome-Cong]

Most transformations like AddSelfLoops() works for HeteroData objects as well, but AddRandomWalkPE doesn't, it requires the attribute edge_index and cannot deal with edge_index_dict like others. For example:

import torch
import torch_geometric.transforms as T
from torch_geometric.data import HeteroData, Data

transform = T.AddRandomWalkPE(walk_length=4, attr_name='pe')
nodes = torch.empty((5,6)).float()
edge_idx = torch.Tensor([[0,0,1,1,2,2,3,3],[1,2,0,3,0,3,1,2]]).long()
edge_idx_aug = torch.Tensor([[0,1,2,3], [0,0,0,0]]).long()

hetero_data=HeteroData()
hetero_data['place'].x = nodes[:-1,:]
hetero_data['room'].x = nodes[-1,:].view(-1, nodes_aug.shape[1])
hetero_data['place', 'with', 'place'].edge_index = edge_idx
hetero_data['place', 'connects', 'room'].edge_index = edge_idx_aug
post_data = transform(hetero_data)

You will get:

AttributeError: 'HeteroData' has no attribute 'edge_index'

One possible solution is to create a Data() subgraph for place nodes can compute the pe attribute with transform, then assign it to the dict.

place_data = Data(x=nodes[:-1,:], edge_index=edge_idx)
post_data = transform(place_data)
hetero_data['place'].pe = post_data.pe

Is there any better way to solve this?

[wsad1]

You are right AddRandomWalkPE only supports homogenous graphs. The paper it is based on only discusses homogenous graphs.
Approaches to tackle this are

1. The approach you suggested looks good. you can use HeteroData.node_type_subgraph to get a subgraph that contains only place nodes.
2. Use MetaPath2Vec to generate embeddings for heterogenous graphs.

[wsad1]

Looking up literature on heterogenous graph embedding yielded https://www.amazon.science/publications/panrep-universal-node-embeddings-for-heterogeneous-graphs and https://arxiv.org/abs/2011.14867. It would be great if we could some of these into pyg.