dgl.khop_in_subgraph¶

dgl.khop_in_subgraph(graph, nodes, k, *, relabel_nodes=True, store_ids=True, output_device=None)[source]¶

Return the subgraph induced by k-hop in-neighborhood of the specified node(s).

We can expand a set of nodes by including the predecessors of them. From a specified node set, a k-hop in subgraph is obtained by first repeating the node set expansion for k times and then creating a node induced subgraph. In addition to extracting the subgraph, DGL also copies the features of the extracted nodes and edges to the resulting graph. The copy is lazy and incurs data movement only when needed.

If the graph is heterogeneous, DGL extracts a subgraph per relation and composes them as the resulting graph. Thus the resulting graph has the same set of relations as the input one.

Parameters

graph (DGLGraph) – The input graph.
nodes (nodes or dict[str, nodes]) –
The starting node(s) to expand. The allowed formats are:
- Int: ID of a single node.
- Int Tensor: Each element is a node ID. The tensor must have the same device type and ID data type as the graph’s.
- iterable[int]: Each element is a node ID.
If the graph is homogeneous, one can directly pass the above formats. Otherwise, the argument must be a dictionary with keys being node types and values being the node IDs in the above formats.
k (int) – The number of hops.
relabel_nodes (bool, optional) – If True, it will remove the isolated nodes and relabel the rest nodes in the extracted subgraph.
store_ids (bool, optional) – If True, it will store the raw IDs of the extracted edges in the edata of the resulting graph under name dgl.EID; if relabel_nodes is True, it will also store the raw IDs of the extracted nodes in the ndata of the resulting graph under name dgl.NID.
output_device (Framework-specific device context object, optional) – The output device. Default is the same as the input graph.

Returns

DGLGraph – The subgraph.
Tensor or dict[str, Tensor], optional – The new IDs of the input nodes after node relabeling. This is returned only when relabel_nodes is True. It is in the same form as nodes.

Notes

When k is 1, the result subgraph is different from the one obtained by dgl.in_subgraph(). The 1-hop in subgraph also includes the edges among the neighborhood.

Examples

The following example uses PyTorch backend.

>>> import dgl
>>> import torch

Extract a two-hop subgraph from a homogeneous graph.

>>> g = dgl.graph(([1, 1, 2, 3, 4], [0, 2, 0, 4, 2]))
>>> g.edata['w'] = torch.arange(10).view(5, 2)
>>> sg, inverse_indices = dgl.khop_in_subgraph(g, 0, k=2)
>>> sg
Graph(num_nodes=4, num_edges=4,
      ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)}
      edata_schemes={'w': Scheme(shape=(2,), dtype=torch.int64),
                     '_ID': Scheme(shape=(), dtype=torch.int64)})
>>> sg.edges()
(tensor([1, 1, 2, 3]), tensor([0, 2, 0, 2]))
>>> sg.edata[dgl.EID]  # original edge IDs
tensor([0, 1, 2, 4])
>>> sg.edata['w']  # also extract the features
tensor([[0, 1],
        [2, 3],
        [4, 5],
        [8, 9]])
>>> inverse_indices
tensor([0])

Extract a subgraph from a heterogeneous graph.

>>> g = dgl.heterograph({
...     ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1]),
...     ('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2])})
>>> sg, inverse_indices = dgl.khop_in_subgraph(g, {'game': 0}, k=2)
>>> sg
Graph(num_nodes={'game': 1, 'user': 2},
      num_edges={('user', 'follows', 'user'): 1, ('user', 'plays', 'game'): 2},
      metagraph=[('user', 'user', 'follows'), ('user', 'game', 'plays')])
>>> inverse_indices
{'game': tensor([0])}