def test_topk_nodes():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(14))
feat0 = F.randn((g0.number_of_nodes(), 10))
g0.ndata['x'] = feat0
# to test the case where k > number of nodes.
dgl.topk_nodes(g0, 'x', 20, idx=-1)
# test correctness
val, indices = dgl.topk_nodes(g0, 'x', 5, idx=-1)
ground_truth = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 9, 10), 0, True)[:5], (5,))
assert F.allclose(ground_truth, indices)
g0.ndata.pop('x')
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(12))
feat1 = F.randn((g1.number_of_nodes(), 10))
bg = dgl.batch([g0, g1])
bg.ndata['x'] = F.cat([feat0, feat1], 0)
# to test the case where k > number of nodes.
dgl.topk_nodes(bg, 'x', 16, idx=1)
# test correctness
val, indices = dgl.topk_nodes(bg, 'x', 6, descending=False, idx=0)
ground_truth_0 = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 0, 1), 0, False)[:6], (6,))
ground_truth_1 = F.reshape(
F.argsort(F.slice_axis(feat1, -1, 0, 1), 0, False)[:6], (6,))
ground_truth = F.stack([ground_truth_0, ground_truth_1], 0)
assert F.allclose(ground_truth, indices)
# test idx=None
val, indices = dgl.topk_nodes(bg, 'x', 6, descending=True)
assert F.allclose(val, F.stack([F.topk(feat0, 6, 0), F.topk(feat1, 6, 0)], 0))
def test_random_walk():
g1 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (2, 0)]
})
g2 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (1, 3), (2, 0), (3, 0)]
})
g3 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (2, 0)],
('user', 'view', 'item'): [(0, 0), (1, 1), (2, 2)],
('item', 'viewed-by', 'user'): [(0, 0), (1, 1), (2, 2)]
})
g4 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (1, 3), (2, 0), (3, 0)],
('user', 'view', 'item'): [(0, 0), (0, 1), (1, 1), (2, 2), (3, 2),
(3, 1)],
('item', 'viewed-by', 'user'): [(0, 0), (1, 0), (1, 1), (2, 2), (2, 3),
(1, 3)]
})
g2.edata['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g2.edata['p2'] = F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32)
g4.edges['follow'].data['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g4.edges['viewed-by'].data['p'] = F.tensor([1, 1, 1, 1, 1, 1],
dtype=F.float32)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2], length=4)
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=4,
restart_prob=0.)
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=4,
restart_prob=F.zeros((4, ),
F.float32,
F.cpu()))
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=5,
restart_prob=F.tensor(
[0, 0, 0, 0, 1],
dtype=F.float32))
check_random_walk(g1, ['follow'] * 4, F.slice_axis(traces, 1, 0, 5),
F.slice_axis(ntypes, 0, 0, 5))
assert (F.asnumpy(traces)[:, 5] == -1).all()
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4)
check_random_walk(g2, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4,
prob='p')
check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p')
try:
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4,
prob='p2')
fail = False
except dgl.DGLError:
fail = True
assert fail
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g3, [0, 1, 2, 0, 1, 2],
metapath=metapath)
check_random_walk(g3, metapath, traces, ntypes)
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath)
check_random_walk(g4, metapath, traces, ntypes)
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p')
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p',
restart_prob=0.)
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p',
restart_prob=F.zeros((6, ),
F.float32,
F.cpu()))
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath + ['follow'],
prob='p',
restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32))
check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p')
assert (F.asnumpy(traces[:, 7]) == -1).all()
