Using k_hop_subgraph on graphs with periodic boundary conditions

[fedeotto]

I'm wondering whether it is possible to use k_hop_subgraph from torch_geometric.utils when dealing with a graph with periodic boundary conditions. In particular I'm talking about a graph of a crystal structure obtained using StructureGraph from pymatgen. In general, the obtained crystal graph will be a multi-graph with multiple connections between the nodes that signify bonds across periodic images of the original unit cell. For example considering a crystal graph with two nodes we may observe a connectivity like edge_index = [[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]] and we have an additional attribute,to_jimages tensor like

tensor([[ 0, -1,  1],
        [-1, -1, -1],
        [ 0,  0,  0],
        [ 0,  1,  0],
        [-1,  0,  0],
        [ 0,  1,  0]])

which tells based on which integer translation of the original unit cell the bond/edge must be considered. I guess that running directly k_hop_subgraph wouldn't make sense in this context? My goal would be to efficiently obtain k-hop neighbours for each node in the original crystal graph.

[rusty1s]

Sorry for late response. k_hop_subgraph will maintain multi edges in a multi graph that are reachable within k hops. Is that what you want?

[fedeotto]

hi @rusty1s many thanks for your answer. Actually, I'm not sure whether k_hop_subgraph does what I'm looking for. The point was being able to dynamically modify the graph during the training stage by connecting the original nodes with k-hop neighbors during message passing and thus facilitating the exchange of long-range information. I'm not sure that this can be achieved directly though. As I have been reading recent works that try to tackle this kind of problem like this, and a general interest towards long-range interactions as showed in the recently proposed long-range graph benchmark. Many thanks

[rusty1s]

Got it, yeah. This is not what k_hop_subgraph does. If you want to rewire your graph, this would require a specialized implementation.