dgl.DGLGraph.filter_edges

DGLGraph.filter_edges(predicate, edges='__ALL__', etype=None)

返回满足给定谓词的具有给定边类型的边的ID。

Parameters:

• predicate (callable) – 一个签名为 func(edges) -> Tensor 的函数。 edges 是 dgl.EdgeBatch 对象。 它的输出张量应该是一个一维布尔张量， 每个元素指示批次中的相应边是否满足谓词。

• edges (edges) –

  用于发送和接收消息的边。允许的输入格式有：

  • int: 单个边ID。

  • Int Tensor: 每个元素都是一个边ID。张量必须具有与图相同的设备类型和ID数据类型。

  • iterable[int]: 每个元素都是一个边ID。

  • (Tensor, Tensor): 节点张量格式，其中两个张量的第i个元素指定一条边。

  • (iterable[int], iterable[int]): 类似于节点张量格式，但将边端点存储在python可迭代对象中。

  默认情况下，它考虑所有边。

• etype (str or (str, str, str), optional) –

  The type name of the edges. The allowed type name formats are:

  • (str, str, str) for source node type, edge type and destination node type.

  • or one str edge type name if the name can uniquely identify a triplet format in the graph.

  Can be omitted if the graph has only one type of edges.

Returns:

一个包含满足谓词的边的ID的一维张量。

Return type:

张量

示例

以下示例使用PyTorch后端。

>>> import dgl
>>> import torch

定义一个谓词函数。

>>> def edges_with_feature_one(edges):
...     # Whether an edge has feature 1
...     return (edges.data['h'] == 1.).squeeze(1)

为同构图过滤边。

>>> g = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
>>> g.edata['h'] = torch.tensor([[0.], [1.], [1.]])
>>> print(g.filter_edges(edges_with_feature_one))
tensor([1, 2])

过滤ID为0和1的边

>>> print(g.filter_edges(edges_with_feature_one, edges=torch.tensor([0, 1])))
tensor([1])

为异构图过滤边。

>>> g = dgl.heterograph({
...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
...                                 torch.tensor([0, 0, 1, 1])),
...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2]))})
>>> g.edges['plays'].data['h'] = torch.tensor([[0.], [1.], [1.], [0.]])
>>> # Filter for 'plays' nodes
>>> print(g.filter_edges(edges_with_feature_one, etype='plays'))
tensor([1, 2])