How to structure data object for multilayer graphs

[saiden89]

Hi,
I'm working on a problem of link prediction using four different edge sets for the same nodes, in which each edge_index reflects a different kind of association between the same nodes. I've tried to embed all the different layers into the same data object, but after due train/test split and negative sampling I get a rather messy data object:

data.x
data.layer1_train_pos_edge_index
data.layer1_train_neg_edge_index
data.layer1_val_pos_edge_index
data.layer1_val_neg_edge_index
data.layer1_test_pos_edge_index
data.layer1_test_neg_edge_index
data.layer1_train_pos_edge_weight
data.layer1_val_pos_edge_weight
data.layer1_test_pos_edge_weight
data.layer2_train_pos_edge_index
...

The goal is to train a model that, given the node embeddings and the different edge sets, is capable of predicting layer-specific links. For example, predicting a type-1 association based on the information also present in other layers.
My network(s) isn't that large, I have ~20k nodes and only one layer has ~10M edges, the others are much smaller. Is there a more readable and understandable way to achieve a proper data structure to the feed to the model?

Thanks a lot for any insight provided!
p.s. special thanks to @rusty1s for his incredible effort in devising this library :)

[rusty1s]

I guess you can also put them into a list of tensors:

data.train_pos_edge_index = [layer1_train_pos_edge_index, layer2_train_pos_edge_index, ...]

Does that solve your issues regarding readability?

[saiden89]

Sure thing, I was thinking about that possibility, along putting them in a dictionary of some kind. I was just investigating whether there was already some kind of built-in abstraction for heterogeneous graphs.
Thank you for the answer :)