dgl.edge_subgraph

dgl.edge_subgraph(graph, edges, preserve_nodes=False)[source]

Return a subgraph induced on the given edges.

An edge-induced subgraph is equivalent to creating a new graph with the same number of nodes using the given edges. In addition to extracting the subgraph, DGL conducts the following:

  • Relabel the incident nodes to IDs starting from zero. Isolated nodes are removed.

  • 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 and edata of the resulting graph under name dgl.NID and dgl.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 the subgraph from.

  • edges (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 nodes.

  • preserve_nodes (bool, optional) – If true, do not relabel the incident nodes and remove the isolated nodes in the extracted subgraph. (Default: False)

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.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], preserve_nodes=True)
>>> sg
Graph(num_nodes=5, num_edges=2,
      ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)}
      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={'user': 2, 'game': 1},
      num_edges={('user', 'plays', 'game'): 1, ('user', 'follows', 'user'): 2},
      metagraph=[('user', 'game'), ('user', 'user')])

See also

node_subgraph()