`LineDigraph` transformation

[Flunzmas]

trafficstars

🚀 The feature, motivation and pitch

Recent scientific work [1] leveraging GNNs for routing in computer networks has used a line digraph representation of the computer network topology as input for their GNN policies. The line digraph is formed by taking the original graph's edges as new nodes, and drawing an edge between those new nodes that, as edges in the original graph, form a directed path of length two: E' = {(u, v), (w, x) | (u, v) in E; (w, x) in E; v = w}.

However, AFAIK there is no existing PyG implementation that provides this conversion (and no issue/MR requesting this). The LineGraph transformation is defined differently and produces a different output. Bernárdez et al., in their code, have implemented the conversion manually on a networkX.DiGraph object, and while this works fine I suspect a PyG transformation implementation working on a Data object can provide the functionality more efficiently.

Alternatives

As mentioned above, the existing LineGraph transformation does something slightly different and is therefore not an alternative. Also, while it's possible to manually construct the graph as Bernárdez et al. did in their code, I assume we can do it quicker with tensor ops in PyG.

I have implemented a sketch for the transformation in this gist, however somebody with more PyG experience might help me in optimizing away the for-loop (if possible).

Additional context

Besides, is it possible to apply the transformations of PyG to Batch objects?

---

[1]: Bernárdez et al. "Magnneto: A graph neural network-based multi-agent system for traffic engineering." IEEE Transactions on Cognitive Communications and Networking 9.2 (2023)

[rusty1s]

You can send a PR with your initial implementation. I can try to optimize it in a follow-up.

[petern48]

Hi @Flunzmas , first-timer here. I optimized the for loop out of your code. I also ran some small test cases to ensure it produces the same result as your original code, and it was already 10x faster using some small graphs. Let me know if there's anything else I can do to help. I would be happy to integrate this into the codebase and send a PR, but wanted to leave the option to you, in case you prefer to do it yourself.

def to_line_digraph_optim(data: Data) -> Data:  # self
    assert data.edge_index is not None
    assert data.is_directed()
    edge_index, edge_attr = data.edge_index, data.edge_attr
    N, E = data.num_nodes, data.num_edges
    edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes=data.num_nodes)
    row, col = edge_index

    # Broadcast to create a mask for matching row and col elements
    mask = row.unsqueeze(0) == col.unsqueeze(1)   # (num_edges, num_edges)
    # Generate new_row/col indices. 2D tensor of arange
    new_row = torch.arange(E).unsqueeze(1).expand(-1, mask.size(1))
    new_col = torch.arange(mask.size(1)).unsqueeze(0).expand(E, -1)
    # Apply the mask to get valid new_row and new_col indices
    new_row = new_row[mask]
    new_col = new_col[mask]
    new_edge_index = torch.stack([new_row, new_col], dim=0)

    data.edge_index = new_edge_index
    data.x = edge_attr
    data.num_nodes = E
    data.edge_attr = None
    return data

[Flunzmas]

Hi @Flunzmas , ...

Hey, cool stuff! I have pushed my template on a branch of my fork here.

Feel free to build on top of that and submit the PR yourself, this way you can be credited for the optimization 🙂.

[petern48]

Hi @Flunzmas , ...

Hey, cool stuff! I have pushed my template on a branch of my fork here.

Feel free to build on top of that and submit the PR yourself, this way you can be credited for the optimization 🙂.

Sounds great, thanks. Could you give me write access to your fork? I tried forking your fork but there doesn't seem to be a way to send a pull request from mine to the original repo. I haven't done a collaborative pull request before, so I'm not sure what standard practice is, but I think committing to your fork is the easiest way.

All my changes are ready. I cleaned up the file to pass the style guide, added my optimization, and made a small change to make it work with the API correctly.

[Flunzmas]

[...] Could you give me write access to your fork? [...]

I haven't done this either, so I also don't know whether there's a canonical way. I have added you as a collaborator.