dgl.DGLGraph.filter_edges

DGLGraph.filter_edges(predicate, edges='__ALL__', etype=None)[源代码]

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

参数:

predicate (callable) – 一个函数，签名为 func(edges) -> Tensor。 edges 是 dgl.EdgeBatch 对象。其输出张量应是一个 1D 布尔张量，每个元素指示批处理中相应的边是否满足谓词。
edges (edges) –
要发送和接收消息的边。允许的输入格式有：
- int: 单个边 ID。
- Int Tensor: 每个元素是一个边 ID。张量必须与图具有相同的设备类型和 ID 数据类型。
- iterable[int]: 每个元素是一个边 ID。
- (Tensor, Tensor): 节点张量格式，其中两个张量的第 i 个元素指定一条边。
- (iterable[int], iterable[int]): 类似于节点张量格式，但在 python 可迭代对象中存储边的端点。
默认情况下，考虑所有边。
etype (str 或 (str, str, str), 可选) –
边类型名称。允许的类型名称格式有：
- (str, str, str) 表示源节点类型、边类型和目标节点类型。
- 或单个 str 边类型名称，如果该名称可以在图中唯一标识一个三元组格式。
如果图只有一种边类型，则可以省略。

返回:

一个 1D 张量，包含满足谓词的边 ID。

返回类型:

Tensor

示例

以下示例使用 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])