Remap edge attributes after train/test split

[saiden89]

Hi,
I have a generic dataset with node and edge features like so:

data = Data(edge_attr=[100], edge_index=[2, 100], x=[50, 10])

After using the train_test_split_edges(data) I get the added pos/neg val/test splits.
How can I remap the original edge_attr to the positive datasets? That is, how can I know the correct attribute of each true edge after the split?

[rusty1s]

Great question. train_test_split_edges currently can not handle edge features. Part of the reason why we ignore edge features here is that it's impossible for us to infer how one would create edge features for negative samples. I agree though that mapping edge_attr to positive edges is important. Please let me know if you have any interest in integrating this :)

[saiden89]

Sure thing! It would be my first contribution to an actual package and I don't mind to put some extra effort :)
I agree that negative samples features are very use-case specific.

I've already started working on adding some lines to propagate the mask and slices also through edge_attr, with edge_attr.shape = [num_edges, num_features]. Everything is fine for val and test since it only requires basic indexing, but train is being quite of a pain since to_undirected gets called and reorders the edgelist. Any idea on how to approach this issue?

[rusty1s]

You can pass edge features to to_undirected as well, see here.

[saiden89]

Oh... for good 5 minutes I was shocked that I missed something this big, but then I realized I was just looking at the old version without edge_attr 😃 . I have modified train_test_split_edges accordingly,as well as the docstrings to reflect the changes.
If you find it correct I can send a pull request to integrate it, bear in mind that I'm still quite new though 😄 . Do I have to open a new issue to propose the feature enhancement? For now I'm posting the updated code below:

I have not found more elegant solutions for taking into account the possibility of not having edge_attr than the series of conditional statements. It should get the job done though.

Edit: created pull request #2692