torch_geometric.utils.k_hop_subgraph

主要记录一下k_hop_subgraph在有向图里关于子图提取的需要注意的地方。
函数参数包括有k_hop_subgraph(
node_idx,
num_hops,
edge_index,
relabel_nodes=False,
num_nodes=None,
flow=‘source_to_target’,
)，每个参数这里就不细讲了。
文档里该函数的描述为：Computes the :math:k-hop subgraph of :obj:edge_index around node
:attr:node_idx.
It returns (1) the nodes involved in the subgraph, (2) the filtered
:obj:edge_index connectivity, (3) the mapping from node indices in
:obj:node_idx to their new location, and (4) the edge mask indicating
which edges were preserved.
通过一些实验个人的理解是：对于给定节点或节点集node_idx，将其作为目标节点（如果为有向图），会先将所有能够通过num_hops跳到达该目标节点的节点（如果为有向图则为源节点）添加进subset（即第一个返回值），然后将连接这些节点的所有边（是所有边，无论是出边还是入边）都添加进子图的边集（第二个返回值）里。由于在处理有向图时关于边的提取有不太清楚的地方，特地做了一些实验，结果记录如下。

先给定一个图并可视化。

from torch_geometric.data import Data
from torch_geometric.utils import to_networkx, k_hop_subgraph
import networkx as nx
import matplotlib.pyplot as plt

edge_index_original = torch.LongTensor([[0, 0, 1, 1, 2, 3, 3, 4, 4],
                               [2, 4, 0, 2, 3, 2, 4, 3, 0]])
x_original = torch.ones(5, 1)
g = Data(edge_index=edge_index_original, x=x_original)
g = to_networkx(g)
nx.draw(g, with_labels=g.nodes)
plt.show()

得到图如下：
用k_hop_subgraph提取节点0的1跳子图：

k_hop_subgraph(node_idx=[0], num_hops=1, edge_index=g.edge_index)
"""
(tensor([0, 1, 4]),
 tensor([[0, 1, 4],
         [4, 0, 0]]),
 tensor([0]),
 tensor([False,  True,  True, False, False, False, False, False,  True]))
"""

输出的第一个参数为得到的子图节点，可以看到是不包含节点2的，因为节点2只是节点0的一个目标节点，不是源节点。另外，输出的第二个参数为得到的子图上的边集，可以看到是包含该子图上三个节点（0， 1， 4）的所有边的，即不仅包括了4指向0的边，也有0指向4的边。
同理，以下关于不同节点的子图提取的结果就很好理解了：

k_hop_subgraph(node_idx=[1], num_hops=1, edge_index=g.edge_index)
"""
(tensor([1]),
 tensor([], size=(2, 0), dtype=torch.int64),
 tensor([0]),
 tensor([False, False, False, False, False, False, False, False, False]))
 没有指向节点1的节点，故子图里只有自身。
"""
k_hop_subgraph(node_idx=[3], num_hops=1, edge_index=g.edge_index)
"""
(tensor([2, 3, 4]),
 tensor([[2, 3, 3, 4],
         [3, 2, 4, 3]]),
 tensor([1]),
 tensor([False, False, False, False,  True,  True,  True,  True, False]))
 输出的子图的边集里有连接节点2，3，4的所有边，不分方向。
 """