def _test2(p, replace): # fanout > #neighbors
subg = dgl.sampling.sample_neighbors(g, [0, 2],
-1,
prob=p,
replace=replace,
edge_dir='out')
assert subg.number_of_nodes() == g.number_of_nodes()
u, v = subg.edges()
u_ans, v_ans = subg.out_edges([0, 2])
uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
assert uv == uv_ans
for i in range(10):
subg = dgl.sampling.sample_neighbors(g, [0, 2],
2,
prob=p,
replace=replace,
edge_dir='out')
assert subg.number_of_nodes() == g.number_of_nodes()
num_edges = 4 if replace else 3
assert subg.number_of_edges() == num_edges
u, v = subg.edges()
assert set(F.asnumpy(F.unique(u))) == {0, 2}
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64),
F.ones((num_edges, ), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
# check no duplication
assert len(edge_set) == num_edges
if p is not None:
assert not (0, 3) in edge_set
def _test1(p, replace):
subg = dgl.sampling.sample_neighbors(g, [0, 1],
-1,
prob=p,
replace=replace)
assert subg.number_of_nodes() == g.number_of_nodes()
u, v = subg.edges()
u_ans, v_ans = subg.in_edges([0, 1])
uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
assert uv == uv_ans
for i in range(10):
subg = dgl.sampling.sample_neighbors(g, [0, 1],
2,
prob=p,
replace=replace)
assert subg.number_of_nodes() == g.number_of_nodes()
assert subg.number_of_edges() == 4
u, v = subg.edges()
assert set(F.asnumpy(F.unique(v))) == {0, 1}
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64),
F.ones((4, ), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
# check no duplication
assert len(edge_set) == 4
if p is not None:
assert not (3, 0) in edge_set
assert not (3, 1) in edge_set
def test_batch_recv():
# basic recv test
g = generate_graph()
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
reduce_msg_shapes.clear()
g.send((u, v), message_func)
g.recv(F.unique(v), reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
def test_batch_recv(index_dtype):
# basic recv test
g = generate_graph(index_dtype=index_dtype)
u = F.tensor([0, 0, 0, 4, 5, 6], dtype=F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 9, 9, 9], dtype=F.data_type_dict[index_dtype])
reduce_msg_shapes.clear()
g.send((u, v), message_func)
g.recv(F.astype(F.unique(v), F.data_type_dict[index_dtype]), reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
def _test1(p, replace):
for i in range(10):
subg = dgl.sampling.sample_neighbors(g, [0, 1], 2, prob=p, replace=replace, edge_dir='out')
assert subg.number_of_nodes() == g.number_of_nodes()
assert subg.number_of_edges() == 4
u, v = subg.edges()
assert set(F.asnumpy(F.unique(u))) == {0, 1}
assert F.array_equal(g.has_edges_between(u, v), F.ones((4,), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
# check no duplication
assert len(edge_set) == 4
if p is not None:
assert not (0, 3) in edge_set
assert not (1, 3) in edge_set
def verify_hetero_graph(g, parts):
num_nodes = {ntype: 0 for ntype in g.ntypes}
num_edges = {etype: 0 for etype in g.etypes}
for part in parts:
assert len(g.ntypes) == len(F.unique(part.ndata[dgl.NTYPE]))
assert len(g.etypes) == len(F.unique(part.edata[dgl.ETYPE]))
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
inner_node_mask = _get_inner_node_mask(part, ntype_id)
num_inner_nodes = F.sum(F.astype(inner_node_mask, F.int64), 0)
num_nodes[ntype] += num_inner_nodes
for etype in g.etypes:
etype_id = g.get_etype_id(etype)
inner_edge_mask = _get_inner_edge_mask(part, etype_id)
num_inner_edges = F.sum(F.astype(inner_edge_mask, F.int64), 0)
num_edges[etype] += num_inner_edges
# Verify the number of nodes are correct.
for ntype in g.ntypes:
print('node {}: {}, {}'.format(ntype, g.number_of_nodes(ntype),
num_nodes[ntype]))
assert g.number_of_nodes(ntype) == num_nodes[ntype]
# Verify the number of edges are correct.
for etype in g.etypes:
print('edge {}: {}, {}'.format(etype, g.number_of_edges(etype),
num_edges[etype]))
assert g.number_of_edges(etype) == num_edges[etype]
nids = {ntype: [] for ntype in g.ntypes}
eids = {etype: [] for etype in g.etypes}
for part in parts:
src, dst, eid = part.edges(form='all')
orig_src = F.gather_row(part.ndata['orig_id'], src)
orig_dst = F.gather_row(part.ndata['orig_id'], dst)
orig_eid = F.gather_row(part.edata['orig_id'], eid)
etype_arr = F.gather_row(part.edata[dgl.ETYPE], eid)
eid_type = F.gather_row(part.edata[dgl.EID], eid)
for etype in g.etypes:
etype_id = g.get_etype_id(etype)
src1 = F.boolean_mask(orig_src, etype_arr == etype_id)
dst1 = F.boolean_mask(orig_dst, etype_arr == etype_id)
eid1 = F.boolean_mask(orig_eid, etype_arr == etype_id)
exist = g.has_edges_between(src1, dst1, etype=etype)
assert np.all(F.asnumpy(exist))
eid2 = g.edge_ids(src1, dst1, etype=etype)
assert np.all(F.asnumpy(eid1 == eid2))
eids[etype].append(F.boolean_mask(eid_type, etype_arr == etype_id))
# Make sure edge Ids fall into a range.
inner_edge_mask = _get_inner_edge_mask(part, etype_id)
inner_eids = np.sort(
F.asnumpy(F.boolean_mask(part.edata[dgl.EID],
inner_edge_mask)))
assert np.all(
inner_eids == np.arange(inner_eids[0], inner_eids[-1] + 1))
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
# Make sure inner nodes have Ids fall into a range.
inner_node_mask = _get_inner_node_mask(part, ntype_id)
inner_nids = F.boolean_mask(part.ndata[dgl.NID], inner_node_mask)
assert np.all(
F.asnumpy(
inner_nids == F.arange(F.as_scalar(inner_nids[0]),
F.as_scalar(inner_nids[-1]) + 1)))
nids[ntype].append(inner_nids)
for ntype in nids:
nids_type = F.cat(nids[ntype], 0)
uniq_ids = F.unique(nids_type)
# We should get all nodes.
assert len(uniq_ids) == g.number_of_nodes(ntype)
for etype in eids:
eids_type = F.cat(eids[etype], 0)
uniq_ids = F.unique(eids_type)
assert len(uniq_ids) == g.number_of_edges(etype)