dgl.node_type_subgraph

dgl.node_type_subgraph(graph, ntypes, output_device=None)[source]

Return the subgraph induced on given node types.

A node-type-induced subgraph contains all the nodes of the given subset of the node types of a graph and any edges whose endpoints are both in this subset. 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.

Parameters:

graph (DGLGraph) – The graph to extract subgraphs from.
ntypes (list[str]) – The type names of the nodes in the subgraph.
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:

DGLGraph

Notes

This function discards the batch information. Please use dgl.DGLGraph.set_batch_num_nodes() and dgl.DGLGraph.set_batch_num_edges() on the transformed graph to maintain the information.

Examples

The following example uses PyTorch backend.

>>> import dgl
>>> import torch

Instantiate a heterograph.

>>> g = dgl.heterograph({
>>>     ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1]),
>>>     ('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2])
>>> })
>>> # Set node features
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])

Get subgraphs.

>>> sub_g = g.node_type_subgraph(['user'])
>>> print(sub_g)
Graph(num_nodes=3, num_edges=3,
      ndata_schemes={'h': Scheme(shape=(1,), dtype=torch.float32)}
      edata_schemes={})

Get the extracted node features.

>>> sub_g.nodes['user'].data['h']
tensor([[0.],
        [1.],
        [2.]])