gammagl.utils.k_hop_subgraph

class k_hop_subgraph(node_idx, num_hops, edge_index, relabel_nodes=False, num_nodes=None, reverse=False)[source]

Computes the induced subgraph of edge_index around all nodes in node_idx reachable within \(k\) hops. The flow argument denotes the direction of edges for finding \(k\)-hop neighbors. If set to "source_to_target", then the method will find all neighbors that point to the initial set of seed nodes in node_idx. This mimics the natural flow of message passing in Graph Neural Networks. The method returns (1) the nodes involved in the subgraph, (2) the filtered edge_index connectivity, (3) the mapping from node indices in node_idx to their new location, and (4) the edge mask indicating which edges were preserved.

Parameters:

  • node_idx (int, list, tuple, tensor) – The central seed node(s).

  • num_hops (int) – The number of hops \(k\).

  • edge_index (tensor) – The edge indices.

  • relabel_nodes (bool, optional) – If set to True, the resulting edge_index will be relabeled to hold consecutive indices starting from zero. (default: False)

  • num_nodes (int, optional) – The number of nodes, i.e. max_val + 1 of edge_index. (default: None)

  • reverse (bool, optional) – The flow direction of \(k\)-hop, False for “source to target” or vice versa.