dgl.sampling.select_topk

dgl.sampling.select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False, copy_ndata=True, copy_edata=True, output_device=None)[source]

Select the neighboring edges with k-largest (or k-smallest) weights of the given nodes and return the induced subgraph.

For each node, a number of inbound (or outbound when edge_dir == 'out') edges with the largest (or smallest when ascending == True) weights will be chosen. The graph returned will then contain all the nodes in the original graph, but only the sampled edges.

Node/edge features are not preserved. The original IDs of the sampled edges are stored as the dgl.EID feature in the returned graph.

Parameters

  • g (DGLGraph) – The graph. Must be on CPU.

  • k (int or dict[etype, int]) –

    The number of edges to be selected for each node on each edge type.

    This argument can take a single int or a dictionary of edge types and ints. If a single int is given, DGL will select this number of edges for each node for every edge type.

    If -1 is given for a single edge type, all the neighboring edges with that edge type will be selected.

  • weight (str) – Feature name of the weights associated with each edge. The feature should have only one element for each edge. The feature can be either int32/64 or float32/64.

  • nodes (tensor or dict, optional) –

    Node IDs to sample neighbors from.

    This argument can take a single ID tensor or a dictionary of node types and ID tensors. If a single tensor is given, the graph must only have one type of nodes.

    If None, DGL will select the edges for all nodes.

  • edge_dir (str, optional) –

    Determines whether to sample inbound or outbound edges.

    Can take either in for inbound edges or out for outbound edges.

  • ascending (bool, optional) – If True, DGL will return edges with k-smallest weights instead of k-largest weights.

  • copy_ndata (bool, optional) –

    If True, the node features of the new graph are copied from the original graph. If False, the new graph will not have any node features.

    (Default: True)

  • copy_edata (bool, optional) –

    If True, the edge features of the new graph are copied from the original graph. If False, the new graph will not have any edge features.

    (Default: True)

  • output_device (Framework-specific device context object, optional) – The output device. Default is the same as the input graph.

Returns

A sampled subgraph containing only the sampled neighboring edges. It is on CPU.

Return type

DGLGraph

Notes

If copy_ndata or copy_edata is True, same tensors are used as the node or edge features of the original graph and the new graph. As a result, users should avoid performing in-place operations on the node features of the new graph to avoid feature corruption.

Examples

>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))
>>> g.edata['weight'] = torch.FloatTensor([0, 1, 0, 1, 0, 1])
>>> sg = dgl.sampling.select_topk(g, 1, 'weight')
>>> sg.edges(order='eid')
(tensor([2, 1, 0]), tensor([0, 1, 2]))