dgl.node_subgraph¶
-
dgl.
node_subgraph
(graph, nodes)[source]¶ Return a subgraph induced on the given nodes.
A node-induced subgraph is a subset of the nodes of a graph together with any edges whose endpoints are both in this subset. In addition to extracting the subgraph, DGL conducts the following:
Relabel the extracted nodes to IDs starting from zero.
Copy the features of the extracted nodes and edges to the resulting graph. The copy is lazy and incurs data movement only when needed.
Store the IDs of the extracted nodes and edges in the
ndata
andedata
of the resulting graph under namedgl.NID
anddgl.EID
, respectively.
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 graph to extract subgraphs from.
nodes (nodes or dict[str, nodes]) –
The nodes to form the subgraph. The allowed nodes formats are:
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.
Bool Tensor: Each \(i^{th}\) element is a bool flag indicating whether node \(i\) is in the subgraph.
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 nodes.
- Returns
G – The subgraph.
- Return type
DGLGraph
Examples
The following example uses PyTorch backend.
>>> import dgl >>> import torch
Extract a subgraph from a homogeneous graph.
>>> g = dgl.graph(([0, 1, 2, 3, 4], [1, 2, 3, 4, 0])) # 5-node cycle >>> sg = dgl.node_subgraph(g, [0, 1, 4]) >>> sg Graph(num_nodes=3, num_edges=2, ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)} edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)}) >>> sg.edges() (tensor([0, 2]), tensor([1, 0])) >>> sg.ndata[dgl.NID] # original node IDs tensor([0, 1, 4]) >>> sg.edata[dgl.EID] # original edge IDs tensor([0, 4])
Specify nodes using a boolean mask.
>>> nodes = torch.tensor([True, True, False, False, True]) # choose nodes [0, 1, 4] >>> dgl.node_subgraph(g, nodes) Graph(num_nodes=3, num_edges=2, ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)} edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
The resulting subgraph also copies features from the parent graph.
>>> g.ndata['x'] = torch.arange(10).view(5, 2) >>> sg = dgl.node_subgraph(g, [0, 1, 4]) >>> sg Graph(num_nodes=3, num_edges=2, ndata_schemes={'x': Scheme(shape=(2,), dtype=torch.int64), '_ID': Scheme(shape=(), dtype=torch.int64)} edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)}) >>> sg.ndata['x'] tensor([[0, 1], [2, 3], [8, 9]])
Extract a subgraph from a hetergeneous graph.
>>> g = dgl.heterograph({ >>> ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1]), >>> ('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2]) >>> }) >>> sub_g = dgl.node_subgraph(g, {'user': [1, 2]}) >>> sub_g Graph(num_nodes={'user': 2, 'game': 0}, num_edges={('user', 'plays', 'game'): 0, ('user', 'follows', 'user'): 2}, metagraph=[('user', 'game'), ('user', 'user')])
See also