dgl.edge_subgraphο
- dgl.edge_subgraph(graph, edges, *, relabel_nodes=True, store_ids=True, output_device=None)[source]ο
Return a subgraph induced on the given edges.
An edge-induced subgraph is equivalent to creating a new graph using the given edges. 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 graph to extract the subgraph from.
edges (edges or dict[(str, str, str), edges]) β
The edges to form the subgraph. The allowed edges formats are:
Int Tensor: Each element is an edge ID. The tensor must have the same device type and ID data type as the graphβs.
iterable[int]: Each element is an edge ID.
Bool Tensor: Each \(i^{th}\) element is a bool flag indicating whether edge \(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 edge types and values being the edge IDs in the above formats.
relabel_nodes (bool, optional) β If True, it will remove the isolated nodes and relabel the incident 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 namedgl.EID
; ifrelabel_nodes
isTrue
, it will also store the raw IDs of the incident nodes in thendata
of the resulting graph under namedgl.NID
.output_device (Framework-specific device context object, optional) β The output device. Default is the same as the input graph.
- Returns:
G β The subgraph.
- Return type:
Notes
This function discards the batch information. Please use
dgl.DGLGraph.set_batch_num_nodes()
anddgl.DGLGraph.set_batch_num_edges()
on the transformed graph to maintain the information.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.edge_subgraph(g, [0, 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, 1]), tensor([2, 0])) >>> sg.ndata[dgl.NID] # original node IDs tensor([0, 4, 1]) >>> sg.edata[dgl.EID] # original edge IDs tensor([0, 4])
Extract a subgraph without node relabeling.
>>> sg = dgl.edge_subgraph(g, [0, 4], relabel_nodes=False) >>> sg Graph(num_nodes=5, num_edges=2, ndata_schemes={} edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)}) >>> sg.edges() (tensor([0, 4]), tensor([1, 0]))
Specify edges using a boolean mask.
>>> nodes = torch.tensor([True, False, False, False, True]) # choose edges [0, 4] >>> dgl.edge_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.edge_subgraph(g, [0, 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[dgl.NID] tensor([0, 4, 1]) >>> sg.ndata['x'] tensor([[0, 1], [8, 9], [2, 3]])
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.edge_subgraph(g, {('user', 'follows', 'user'): [1, 2], ... ('user', 'plays', 'game'): [2]}) >>> print(sub_g) Graph(num_nodes={'game': 1, user': 2}, num_edges={('user', 'follows', 'user'): 2, ('user', 'plays', 'game'): 1}, metagraph=[('user', 'user', 'follows'), ('user', 'game', 'plays')])
See also