def test_topk_edges():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(14))
feat0 = F.randn((g0.number_of_edges(), 10))
g0.edata['x'] = feat0
# to test the case where k > number of edges.
dgl.topk_edges(g0, 'x', 30, idx=-1)
# test correctness
val, indices = dgl.topk_edges(g0, 'x', 7, idx=-1)
ground_truth = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 9, 10), 0, True)[:7], (7,))
assert F.allclose(ground_truth, indices)
g0.edata.pop('x')
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(12))
feat1 = F.randn((g1.number_of_edges(), 10))
bg = dgl.batch([g0, g1])
bg.edata['x'] = F.cat([feat0, feat1], 0)
# to test the case where k > number of edges.
dgl.topk_edges(bg, 'x', 33, idx=1)
# test correctness
val, indices = dgl.topk_edges(bg, 'x', 4, descending=False, idx=0)
ground_truth_0 = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 0, 1), 0, False)[:4], (4,))
ground_truth_1 = F.reshape(
F.argsort(F.slice_axis(feat1, -1, 0, 1), 0, False)[:4], (4,))
ground_truth = F.stack([ground_truth_0, ground_truth_1], 0)
assert F.allclose(ground_truth, indices)
# test idx=None
val, indices = dgl.topk_edges(bg, 'x', 6, descending=True)
assert F.allclose(val, F.stack([F.topk(feat0, 6, 0), F.topk(feat1, 6, 0)], 0))
def forward(self, dgl_data):
topknodes, _ = dgl.topk_nodes(dgl_data, 'h', 4)
topkedges, _ = dgl.topk_edges(dgl_data, 'h', 4)
meannode = torch.mean(topknodes, 1)
maxnode, _ = torch.max(topknodes, 1)
meanedge = torch.mean(topkedges, 1)
maxedge, _ = torch.max(topkedges, 1)
dgl_feat = torch.cat([meannode, maxnode, meanedge, maxedge], -1)
dgl_predict = self.activate(self.weight(dgl_feat))
return dgl_predict
def test_topk(g, idtype, descending):
g = g.astype(idtype).to(F.ctx())
g.ndata['x'] = F.randn((g.number_of_nodes(), 3))
# Test.1: to test the case where k > number of nodes.
dgl.topk_nodes(g, 'x', 100, sortby=-1)
# Test.2: test correctness
min_nnodes = F.asnumpy(g.batch_num_nodes()).min()
if min_nnodes <= 1:
return
k = min_nnodes - 1
val, indices = dgl.topk_nodes(g, 'x', k, descending=descending, sortby=-1)
print(k)
print(g.ndata['x'])
print('val', val)
print('indices', indices)
subg = dgl.unbatch(g)
subval, subidx = [], []
for sg in subg:
subx = F.asnumpy(sg.ndata['x'])
ai = np.argsort(subx[:, -1:].flatten())
if descending:
ai = np.ascontiguousarray(ai[::-1])
subx = np.expand_dims(subx[ai[:k]], 0)
subval.append(F.tensor(subx))
subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
print(F.cat(subval, dim=0))
assert F.allclose(val, F.cat(subval, dim=0))
assert F.allclose(indices, F.cat(subidx, dim=0))
# Test.3: sorby=None
dgl.topk_nodes(g, 'x', k, sortby=None)
g.edata['x'] = F.randn((g.number_of_edges(), 3))
# Test.4: topk edges where k > number of edges.
dgl.topk_edges(g, 'x', 100, sortby=-1)
# Test.5: topk edges test correctness
min_nedges = F.asnumpy(g.batch_num_edges()).min()
if min_nedges <= 1:
return
k = min_nedges - 1
val, indices = dgl.topk_edges(g, 'x', k, descending=descending, sortby=-1)
print(k)
print(g.edata['x'])
print('val', val)
print('indices', indices)
subg = dgl.unbatch(g)
subval, subidx = [], []
for sg in subg:
subx = F.asnumpy(sg.edata['x'])
ai = np.argsort(subx[:, -1:].flatten())
if descending:
ai = np.ascontiguousarray(ai[::-1])
subx = np.expand_dims(subx[ai[:k]], 0)
subval.append(F.tensor(subx))
subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
print(F.cat(subval, dim=0))
assert F.allclose(val, F.cat(subval, dim=0))
assert F.allclose(indices, F.cat(subidx, dim=0))