Skip to content

GitLab

Unverified Commit dc8ca88e authored by Jinjing Zhou's avatar Jinjing Zhou Committed by GitHub
Browse files

[Refactor] Explicit dtype for HeteroGraph (#1467) 



* 111

* 111

* lint

* lint

* lint

* lint

* fix

* lint

* try

* fix

* lint

* lint

* test

* fix

* ttt

* test

* fix

* fix

* fix

* mxnet

* 111

* fix 64bits computation

* pylint

* roll back

* fix

* lint

* fix hetero_from_relations

* remove index_dtype in to_homo and to_hetero

* fix

* fix

* fix

* fix

* remove default

* fix

* lint

* fix

* fix error message

* fix error

* lint

* macro dispatch

* try

* lint

* remove nbits

* error message

* fix

* fix

* lint

* lint

* lint

* fix

* lint

* fix

* fix random walk

* lint

* lint

* fix

* fix

* fix

* lint

* fix

* lint
Co-authored-by: default avatarMinjie Wang <wmjlyjemaine@gmail.com>
parent de34e15a
Hide whitespace changes
Inline Side-by-side
Showing with 246 additions and 161 deletions
tests/compute/test_heterograph.py
View file @ dc8ca88e
...@@ -8,8 +8,9 @@ import backend as F ...@@ -8,8 +8,9 @@ import backend as F
import networkx as nx import networkx as nx
import unittest, pytest import unittest, pytest
from dgl import DGLError from dgl import DGLError
from utils import parametrize_dtype
def create_test_heterograph(): def create_test_heterograph(index_dtype):
# test heterograph from the docstring, plus a user -- wishes -- game relation # test heterograph from the docstring, plus a user -- wishes -- game relation
# 3 users, 2 games, 2 developers # 3 users, 2 games, 2 developers
# metagraph: # metagraph:
...@@ -25,14 +26,19 @@ def create_test_heterograph(): ...@@ -25,14 +26,19 @@ def create_test_heterograph():
wishes_nx.add_edge('u0', 'g1', id=0) wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1) wishes_nx.add_edge('u2', 'g0', id=1)
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows') follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', index_dtype=index_dtype)
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game') plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game', index_dtype=index_dtype)
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game') wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game', index_dtype=index_dtype)
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game') develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game', index_dtype=index_dtype)
assert follows_g._idtype_str == index_dtype
assert plays_g._idtype_str == index_dtype
assert wishes_g._idtype_str == index_dtype
assert develops_g._idtype_str == index_dtype
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g]) g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
assert g._idtype_str == index_dtype
return g return g
def create_test_heterograph1(): def create_test_heterograph1(index_dtype):
edges = [] edges = []
edges.extend([(0,1), (1,2)]) # follows edges.extend([(0,1), (1,2)]) # follows
edges.extend([(0,3), (1,3), (2,4), (1,4)]) # plays edges.extend([(0,3), (1,3), (2,4), (1,4)]) # plays
...@@ -45,7 +51,7 @@ def create_test_heterograph1(): ...@@ -45,7 +51,7 @@ def create_test_heterograph1():
g0.edata[dgl.ETYPE] = etypes g0.edata[dgl.ETYPE] = etypes
return dgl.to_hetero(g0, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops']) return dgl.to_hetero(g0, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'])
def create_test_heterograph2(): def create_test_heterograph2(index_dtype):
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1]))) plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph() wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0) wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
...@@ -62,7 +68,7 @@ def create_test_heterograph2(): ...@@ -62,7 +68,7 @@ def create_test_heterograph2():
}) })
return g return g
def create_test_heterograph3(): def create_test_heterograph3(index_dtype):
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1]))) plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph() wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0) wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
...@@ -82,10 +88,11 @@ def create_test_heterograph3(): ...@@ -82,10 +88,11 @@ def create_test_heterograph3():
def get_redfn(name): def get_redfn(name):
return getattr(F, name) return getattr(F, name)
def test_create(): @parametrize_dtype
g0 = create_test_heterograph() def test_create(index_dtype):
g1 = create_test_heterograph1() g0 = create_test_heterograph(index_dtype)
g2 = create_test_heterograph2() g1 = create_test_heterograph1(index_dtype)
g2 = create_test_heterograph2(index_dtype)
assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes) assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
assert set(g0.canonical_etypes) == set(g1.canonical_etypes) == set(g2.canonical_etypes) assert set(g0.canonical_etypes) == set(g1.canonical_etypes) == set(g2.canonical_etypes)
...@@ -142,8 +149,9 @@ def test_create(): ...@@ -142,8 +149,9 @@ def test_create():
_test_validate_bipartite((3, 3)) _test_validate_bipartite((3, 3))
_test_validate_bipartite((2, 4)) _test_validate_bipartite((2, 4))
def test_query(): @parametrize_dtype
g = create_test_heterograph() def test_query(index_dtype):
g = create_test_heterograph(index_dtype)
ntypes = ['user', 'game', 'developer'] ntypes = ['user', 'game', 'developer']
canonical_etypes = [ canonical_etypes = [
...@@ -265,11 +273,11 @@ def test_query(): ...@@ -265,11 +273,11 @@ def test_query():
'wishes': ([0, 1], [0, 1]), 'wishes': ([0, 1], [0, 1]),
'develops': ([0, 1], [1, 0]), 'develops': ([0, 1], [1, 0]),
} }
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
_test(g) _test(g)
g = create_test_heterograph1() g = create_test_heterograph1(index_dtype)
_test(g) _test(g)
g = create_test_heterograph3() g = create_test_heterograph3(index_dtype)
_test(g) _test(g)
etypes = canonical_etypes etypes = canonical_etypes
...@@ -286,11 +294,11 @@ def test_query(): ...@@ -286,11 +294,11 @@ def test_query():
('user', 'wishes', 'game'): ([0, 1], [0, 1]), ('user', 'wishes', 'game'): ([0, 1], [0, 1]),
('developer', 'develops', 'game'): ([0, 1], [1, 0]), ('developer', 'develops', 'game'): ([0, 1], [1, 0]),
} }
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
_test(g) _test(g)
g = create_test_heterograph1() g = create_test_heterograph1(index_dtype)
_test(g) _test(g)
g = create_test_heterograph3() g = create_test_heterograph3(index_dtype)
_test(g) _test(g)
# test repr # test repr
...@@ -387,8 +395,9 @@ def test_edge_ids(): ...@@ -387,8 +395,9 @@ def test_edge_ids():
with pytest.raises(AssertionError): with pytest.raises(AssertionError):
eid = g2.edge_id(0, 1, etype='follows') eid = g2.edge_id(0, 1, etype='follows')
def test_adj(): @parametrize_dtype
g = create_test_heterograph() def test_adj(index_dtype):
g = create_test_heterograph(index_dtype)
adj = F.sparse_to_numpy(g.adj(etype='follows')) adj = F.sparse_to_numpy(g.adj(etype='follows'))
assert np.allclose( assert np.allclose(
adj, adj,
...@@ -442,8 +451,9 @@ def test_adj(): ...@@ -442,8 +451,9 @@ def test_adj():
[1., 0., 0.], [1., 0., 0.],
[0., 1., 0.]])) [0., 1., 0.]]))
def test_inc(): @parametrize_dtype
g = create_test_heterograph() def test_inc(index_dtype):
g = create_test_heterograph(index_dtype)
#follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows') #follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
adj = F.sparse_to_numpy(g['follows'].inc('in')) adj = F.sparse_to_numpy(g['follows'].inc('in'))
assert np.allclose( assert np.allclose(
...@@ -480,10 +490,11 @@ def test_inc(): ...@@ -480,10 +490,11 @@ def test_inc():
np.array([[-1., 0.], np.array([[-1., 0.],
[1., -1.], [1., -1.],
[0., 1.]])) [0., 1.]]))
def test_view(): @parametrize_dtype
def test_view(index_dtype):
# test data view # test data view
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
f1 = F.randn((3, 6)) f1 = F.randn((3, 6))
g.nodes['user'].data['h'] = f1 # ok g.nodes['user'].data['h'] = f1 # ok
...@@ -499,9 +510,10 @@ def test_view(): ...@@ -499,9 +510,10 @@ def test_view():
assert F.array_equal(f3, f5) assert F.array_equal(f3, f5)
assert F.array_equal(F.tensor(g.edges(etype='follows', form='eid')), F.arange(0, 2)) assert F.array_equal(F.tensor(g.edges(etype='follows', form='eid')), F.arange(0, 2))
def test_view1(): @parametrize_dtype
def test_view1(index_dtype):
# test relation view # test relation view
HG = create_test_heterograph() HG = create_test_heterograph(index_dtype)
ntypes = ['user', 'game', 'developer'] ntypes = ['user', 'game', 'developer']
canonical_etypes = [ canonical_etypes = [
('user', 'follows', 'user'), ('user', 'follows', 'user'),
...@@ -642,7 +654,8 @@ def test_view1(): ...@@ -642,7 +654,8 @@ def test_view1():
fail = True fail = True
assert fail assert fail
def test_flatten(): @parametrize_dtype
def test_flatten(index_dtype):
def check_mapping(g, fg): def check_mapping(g, fg):
if len(fg.ntypes) == 1: if len(fg.ntypes) == 1:
SRC = DST = fg.ntypes[0] SRC = DST = fg.ntypes[0]
...@@ -658,15 +671,15 @@ def test_flatten(): ...@@ -658,15 +671,15 @@ def test_flatten():
src_fg, dst_fg = fg.find_edges([i]) src_fg, dst_fg = fg.find_edges([i])
# TODO(gq): I feel this code is quite redundant; can we just add new members (like # TODO(gq): I feel this code is quite redundant; can we just add new members (like
# "induced_srcid") to returned heterograph object and not store them as features? # "induced_srcid") to returned heterograph object and not store them as features?
assert src_g == F.gather_row(fg.nodes[SRC].data[dgl.NID], src_fg)[0] assert F.asnumpy(src_g) == F.asnumpy(F.gather_row(fg.nodes[SRC].data[dgl.NID], src_fg)[0])
tid = F.asnumpy(F.gather_row(fg.nodes[SRC].data[dgl.NTYPE], src_fg)).item() tid = F.asnumpy(F.gather_row(fg.nodes[SRC].data[dgl.NTYPE], src_fg)).item()
assert g.canonical_etypes[etype][0] == g.ntypes[tid] assert g.canonical_etypes[etype][0] == g.ntypes[tid]
assert dst_g == F.gather_row(fg.nodes[DST].data[dgl.NID], dst_fg)[0] assert F.asnumpy(dst_g) == F.asnumpy(F.gather_row(fg.nodes[DST].data[dgl.NID], dst_fg)[0])
tid = F.asnumpy(F.gather_row(fg.nodes[DST].data[dgl.NTYPE], dst_fg)).item() tid = F.asnumpy(F.gather_row(fg.nodes[DST].data[dgl.NTYPE], dst_fg)).item()
assert g.canonical_etypes[etype][2] == g.ntypes[tid] assert g.canonical_etypes[etype][2] == g.ntypes[tid]
# check for wildcard slices # check for wildcard slices
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
g.nodes['user'].data['h'] = F.ones((3, 5)) g.nodes['user'].data['h'] = F.ones((3, 5))
g.nodes['game'].data['i'] = F.ones((2, 5)) g.nodes['game'].data['i'] = F.ones((2, 5))
g.edges['plays'].data['e'] = F.ones((4, 4)) g.edges['plays'].data['e'] = F.ones((4, 4))
...@@ -713,8 +726,8 @@ def test_flatten(): ...@@ -713,8 +726,8 @@ def test_flatten():
check_mapping(g, fg) check_mapping(g, fg)
# Test another heterograph # Test another heterograph
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows') g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows', index_dtype=index_dtype)
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows') g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', index_dtype=index_dtype)
g_x.nodes['user'].data['h'] = F.randn((4, 3)) g_x.nodes['user'].data['h'] = F.randn((4, 3))
g_x.edges['follows'].data['w'] = F.randn((3, 2)) g_x.edges['follows'].data['w'] = F.randn((3, 2))
g_y.nodes['user'].data['hh'] = F.randn((4, 5)) g_y.nodes['user'].data['hh'] = F.randn((4, 5))
...@@ -737,23 +750,25 @@ def test_flatten(): ...@@ -737,23 +750,25 @@ def test_flatten():
check_mapping(g, fg) check_mapping(g, fg)
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test") @unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
def test_to_device(): @parametrize_dtype
hg = create_test_heterograph() def test_to_device(index_dtype):
hg = create_test_heterograph(index_dtype)
if F.is_cuda_available(): if F.is_cuda_available():
hg = hg.to(F.cuda()) hg = hg.to(F.cuda())
assert hg is not None assert hg is not None
def test_convert_bound(): @parametrize_dtype
def test_convert_bound(index_dtype):
def _test_bipartite_bound(data, card): def _test_bipartite_bound(data, card):
try: try:
dgl.bipartite(data, num_nodes=card) dgl.bipartite(data, num_nodes=card, index_dtype=index_dtype)
except dgl.DGLError: except dgl.DGLError:
return return
assert False, 'bipartite bound test with wrong uid failed' assert False, 'bipartite bound test with wrong uid failed'
def _test_graph_bound(data, card): def _test_graph_bound(data, card):
try: try:
dgl.graph(data, num_nodes=card) dgl.graph(data, num_nodes=card, index_dtype=index_dtype)
except dgl.DGLError: except dgl.DGLError:
return return
assert False, 'graph bound test with wrong uid failed' assert False, 'graph bound test with wrong uid failed'
...@@ -764,8 +779,9 @@ def test_convert_bound(): ...@@ -764,8 +779,9 @@ def test_convert_bound():
_test_graph_bound(([0,1],[1,3]),3) _test_graph_bound(([0,1],[1,3]),3)
def test_convert(): @parametrize_dtype
hg = create_test_heterograph() def test_convert(index_dtype):
hg = create_test_heterograph(index_dtype)
hs = [] hs = []
for ntype in hg.ntypes: for ntype in hg.ntypes:
h = F.randn((hg.number_of_nodes(ntype), 5)) h = F.randn((hg.number_of_nodes(ntype), 5))
...@@ -780,6 +796,7 @@ def test_convert(): ...@@ -780,6 +796,7 @@ def test_convert():
hg.edges['plays'].data['x'] = F.randn((4, 3)) hg.edges['plays'].data['x'] = F.randn((4, 3))
g = dgl.to_homo(hg) g = dgl.to_homo(hg)
assert g._idtype_str == index_dtype
assert F.array_equal(F.cat(hs, dim=0), g.ndata['h']) assert F.array_equal(F.cat(hs, dim=0), g.ndata['h'])
assert 'x' not in g.ndata assert 'x' not in g.ndata
assert F.array_equal(F.cat(ws, dim=0), g.edata['w']) assert F.array_equal(F.cat(ws, dim=0), g.edata['w'])
...@@ -821,10 +838,11 @@ def test_convert(): ...@@ -821,10 +838,11 @@ def test_convert():
assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w']) assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w'])
# hetero_from_homo test case 2 # hetero_from_homo test case 2
g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)]) g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)], index_dtype=index_dtype)
g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2]) g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2]) g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2']) hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2'])
assert hg._idtype_str == index_dtype
assert set(hg.canonical_etypes) == set( assert set(hg.canonical_etypes) == set(
[('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')]) [('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')])
assert hg.number_of_nodes('l0') == 2 assert hg.number_of_nodes('l0') == 2
...@@ -838,11 +856,12 @@ def test_convert(): ...@@ -838,11 +856,12 @@ def test_convert():
mg = nx.MultiDiGraph([ mg = nx.MultiDiGraph([
('user', 'movie', 'watches'), ('user', 'movie', 'watches'),
('user', 'TV', 'watches')]) ('user', 'TV', 'watches')])
g = dgl.graph([(0, 1), (0, 2)]) g = dgl.graph([(0, 1), (0, 2)], index_dtype=index_dtype)
g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2]) g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0]) g.edata[dgl.ETYPE] = F.tensor([0, 0])
for _mg in [None, mg]: for _mg in [None, mg]:
hg = dgl.to_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg) hg = dgl.to_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg)
assert hg._idtype_str == index_dtype
assert set(hg.canonical_etypes) == set( assert set(hg.canonical_etypes) == set(
[('user', 'watches', 'movie'), ('user', 'watches', 'TV')]) [('user', 'watches', 'movie'), ('user', 'watches', 'TV')])
assert hg.number_of_nodes('user') == 1 assert hg.number_of_nodes('user') == 1
...@@ -857,25 +876,27 @@ def test_convert(): ...@@ -857,25 +876,27 @@ def test_convert():
g = dgl.to_homo(hg) g = dgl.to_homo(hg)
assert g.number_of_nodes() == 5 assert g.number_of_nodes() == 5
def test_transform(): @parametrize_dtype
g = create_test_heterograph() def test_transform(index_dtype):
g = create_test_heterograph(index_dtype)
x = F.randn((3, 5)) x = F.randn((3, 5))
g.nodes['user'].data['h'] = x g.nodes['user'].data['h'] = x
new_g = dgl.metapath_reachable_graph(g, ['follows', 'plays']) new_g = dgl.metapath_reachable_graph(g, ['follows', 'plays'])
assert new_g._idtype_str == index_dtype
assert new_g.ntypes == ['user', 'game'] assert new_g.ntypes == ['user', 'game']
assert new_g.number_of_edges() == 3 assert new_g.number_of_edges() == 3
assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all() assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all()
new_g = dgl.metapath_reachable_graph(g, ['follows']) new_g = dgl.metapath_reachable_graph(g, ['follows'])
assert new_g._idtype_str == index_dtype
assert new_g.ntypes == ['user'] assert new_g.ntypes == ['user']
assert new_g.number_of_edges() == 2 assert new_g.number_of_edges() == 2
assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all() assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all()
def test_subgraph(): @parametrize_dtype
g = create_test_heterograph() def test_subgraph(index_dtype):
g = create_test_heterograph(index_dtype)
x = F.randn((3, 5)) x = F.randn((3, 5))
y = F.randn((2, 4)) y = F.randn((2, 4))
g.nodes['user'].data['h'] = x g.nodes['user'].data['h'] = x
...@@ -952,13 +973,14 @@ def test_subgraph(): ...@@ -952,13 +973,14 @@ def test_subgraph():
assert np.array_equal(src, np.array([1])) assert np.array_equal(src, np.array([1]))
assert np.array_equal(dst, np.array([0])) assert np.array_equal(dst, np.array([0]))
def test_apply(): @parametrize_dtype
def test_apply(index_dtype):
def node_udf(nodes): def node_udf(nodes):
return {'h': nodes.data['h'] * 2} return {'h': nodes.data['h'] * 2}
def edge_udf(edges): def edge_udf(edges):
return {'h': edges.data['h'] * 2 + edges.src['h']} return {'h': edges.data['h'] * 2 + edges.src['h']}
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
g.nodes['user'].data['h'] = F.ones((3, 5)) g.nodes['user'].data['h'] = F.ones((3, 5))
g.apply_nodes(node_udf, ntype='user') g.apply_nodes(node_udf, ntype='user')
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 2) assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 2)
...@@ -990,14 +1012,15 @@ def test_apply(): ...@@ -990,14 +1012,15 @@ def test_apply():
fail = True fail = True
assert fail assert fail
def test_level1(): @parametrize_dtype
def test_level1(index_dtype):
#edges = { #edges = {
# 'follows': ([0, 1], [1, 2]), # 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]), # 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]), # 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]), # 'develops': ([0, 1], [0, 1]),
#} #}
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
def rfunc(nodes): def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)} return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes): def rfunc2(nodes):
...@@ -1082,14 +1105,15 @@ def test_level1(): ...@@ -1082,14 +1105,15 @@ def test_level1():
assert fail assert fail
def test_level2(): @parametrize_dtype
def test_level2(index_dtype):
#edges = { #edges = {
# 'follows': ([0, 1], [1, 2]), # 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]), # 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]), # 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]), # 'develops': ([0, 1], [0, 1]),
#} #}
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
def rfunc(nodes): def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)} return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes): def rfunc2(nodes):
...@@ -1310,14 +1334,15 @@ def test_level2(): ...@@ -1310,14 +1334,15 @@ def test_level2():
g.nodes['game'].data.clear() g.nodes['game'].data.clear()
def test_updates(): @parametrize_dtype
def test_updates(index_dtype):
def msg_func(edges): def msg_func(edges):
return {'m': edges.src['h']} return {'m': edges.src['h']}
def reduce_func(nodes): def reduce_func(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)} return {'y': F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes): def apply_func(nodes):
return {'y': nodes.data['y'] * 2} return {'y': nodes.data['y'] * 2}
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
x = F.randn((3, 5)) x = F.randn((3, 5))
g.nodes['user'].data['h'] = x g.nodes['user'].data['h'] = x
...@@ -1358,8 +1383,10 @@ def test_updates(): ...@@ -1358,8 +1383,10 @@ def test_updates():
assert F.array_equal(y[0], x[0] * multiplier) assert F.array_equal(y[0], x[0] * multiplier)
del g.nodes['game'].data['y'] del g.nodes['game'].data['y']
def test_backward():
g = create_test_heterograph() @parametrize_dtype
def test_backward(index_dtype):
g = create_test_heterograph(index_dtype)
x = F.randn((3, 5)) x = F.randn((3, 5))
F.attach_grad(x) F.attach_grad(x)
g.nodes['user'].data['h'] = x g.nodes['user'].data['h'] = x
...@@ -1375,7 +1402,9 @@ def test_backward(): ...@@ -1375,7 +1402,9 @@ def test_backward():
[2., 2., 2., 2., 2.], [2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.]])) [2., 2., 2., 2., 2.]]))
def test_empty_heterograph():
@parametrize_dtype
def test_empty_heterograph(index_dtype):
def assert_empty(g): def assert_empty(g):
assert g.number_of_nodes('user') == 0 assert g.number_of_nodes('user') == 0
assert g.number_of_edges('plays') == 0 assert g.number_of_edges('plays') == 0
...@@ -1390,12 +1419,15 @@ def test_empty_heterograph(): ...@@ -1390,12 +1419,15 @@ def test_empty_heterograph():
# empty networkx graph # empty networkx graph
assert_empty(dgl.heterograph({('user', 'plays', 'game'): nx.DiGraph()})) assert_empty(dgl.heterograph({('user', 'plays', 'game'): nx.DiGraph()}))
g = dgl.heterograph({('user', 'follows', 'user'): []}) g = dgl.heterograph({('user', 'follows', 'user'): []}, index_dtype=index_dtype)
assert g._idtype_str == index_dtype
assert g.number_of_nodes('user') == 0 assert g.number_of_nodes('user') == 0
assert g.number_of_edges('follows') == 0 assert g.number_of_edges('follows') == 0
# empty relation graph with others # empty relation graph with others
g = dgl.heterograph({('user', 'plays', 'game'): [], ('developer', 'develops', 'game'): [(0, 0), (1, 1)]}) g = dgl.heterograph({('user', 'plays', 'game'): [], ('developer', 'develops', 'game'): [
(0, 0), (1, 1)]}, index_dtype=index_dtype)
assert g._idtype_str == index_dtype
assert g.number_of_nodes('user') == 0 assert g.number_of_nodes('user') == 0
assert g.number_of_edges('plays') == 0 assert g.number_of_edges('plays') == 0
assert g.number_of_nodes('game') == 2 assert g.number_of_nodes('game') == 2
...@@ -1460,14 +1492,15 @@ def test_types_in_function(): ...@@ -1460,14 +1492,15 @@ def test_types_in_function():
g.filter_nodes(filter_nodes2, ntype='game') g.filter_nodes(filter_nodes2, ntype='game')
g.filter_edges(filter_edges2) g.filter_edges(filter_edges2)
def test_stack_reduce(): @parametrize_dtype
def test_stack_reduce(index_dtype):
#edges = { #edges = {
# 'follows': ([0, 1], [1, 2]), # 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]), # 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]), # 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]), # 'develops': ([0, 1], [0, 1]),
#} #}
g = create_test_heterograph() g = create_test_heterograph(index_dtype)
g.nodes['user'].data['h'] = F.randn((3, 200)) g.nodes['user'].data['h'] = F.randn((3, 200))
def rfunc(nodes): def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)} return {'y': F.sum(nodes.mailbox['m'], 1)}
...@@ -1486,7 +1519,8 @@ def test_stack_reduce(): ...@@ -1486,7 +1519,8 @@ def test_stack_reduce():
'stack') 'stack')
assert g.nodes['game'].data['y'].shape == (g.number_of_nodes('game'), 1, 200) assert g.nodes['game'].data['y'].shape == (g.number_of_nodes('game'), 1, 200)
def test_isolated_ntype(): @parametrize_dtype
def test_isolated_ntype(index_dtype):
g = dgl.heterograph({ g = dgl.heterograph({
('A', 'AB', 'B'): [(0, 1), (1, 2), (2, 3)]}, ('A', 'AB', 'B'): [(0, 1), (1, 2), (2, 3)]},
num_nodes_dict={'A': 3, 'B': 4, 'C': 4}) num_nodes_dict={'A': 3, 'B': 4, 'C': 4})
...@@ -1511,14 +1545,20 @@ def test_isolated_ntype(): ...@@ -1511,14 +1545,20 @@ def test_isolated_ntype():
assert g.number_of_nodes('B') == 4 assert g.number_of_nodes('B') == 4
assert g.number_of_nodes('C') == 4 assert g.number_of_nodes('C') == 4
def test_ismultigraph():
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5), (2, 5)], 'A', 'AB', 'B', num_nodes=(6, 6)) @parametrize_dtype
def test_ismultigraph(index_dtype):
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5), (2, 5)], 'A',
'AB', 'B', num_nodes=(6, 6), index_dtype=index_dtype)
assert g1.is_multigraph == False assert g1.is_multigraph == False
g2 = dgl.bipartite([(0, 1), (0, 1), (0, 2), (1, 5)], 'A', 'AC', 'C', num_nodes=(6, 6)) g2 = dgl.bipartite([(0, 1), (0, 1), (0, 2), (1, 5)], 'A',
'AC', 'C', num_nodes=(6, 6), index_dtype=index_dtype)
assert g2.is_multigraph == True assert g2.is_multigraph == True
g3 = dgl.graph([(0, 1), (1, 2)], 'A', 'AA', num_nodes=6) g3 = dgl.graph([(0, 1), (1, 2)], 'A', 'AA',
num_nodes=6, index_dtype=index_dtype)
assert g3.is_multigraph == False assert g3.is_multigraph == False
g4 = dgl.graph([(0, 1), (0, 1), (1, 2)], 'A', 'AA', num_nodes=6) g4 = dgl.graph([(0, 1), (0, 1), (1, 2)], 'A', 'AA',
num_nodes=6, index_dtype=index_dtype)
assert g4.is_multigraph == True assert g4.is_multigraph == True
g = dgl.hetero_from_relations([g1, g3]) g = dgl.hetero_from_relations([g1, g3])
assert g.is_multigraph == False assert g.is_multigraph == False
...@@ -1529,8 +1569,9 @@ def test_ismultigraph(): ...@@ -1529,8 +1569,9 @@ def test_ismultigraph():
g = dgl.hetero_from_relations([g2, g4]) g = dgl.hetero_from_relations([g2, g4])
assert g.is_multigraph == True assert g.is_multigraph == True
def test_bipartite(): @parametrize_dtype
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5)], 'A', 'AB', 'B') def test_bipartite(index_dtype):
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5)], 'A', 'AB', 'B', index_dtype=index_dtype)
assert g1.is_unibipartite assert g1.is_unibipartite
assert len(g1.ntypes) == 2 assert len(g1.ntypes) == 2
assert g1.etypes == ['AB'] assert g1.etypes == ['AB']
...@@ -1553,7 +1594,7 @@ def test_bipartite(): ...@@ -1553,7 +1594,7 @@ def test_bipartite():
assert F.array_equal(g1.nodes['DST/B'].data['h'], g1.dstdata['h']) assert F.array_equal(g1.nodes['DST/B'].data['h'], g1.dstdata['h'])
# more complicated bipartite # more complicated bipartite
g2 = dgl.bipartite([(1, 0), (0, 0)], 'A', 'AC', 'C') g2 = dgl.bipartite([(1, 0), (0, 0)], 'A', 'AC', 'C', index_dtype=index_dtype)
g3 = dgl.hetero_from_relations([g1, g2]) g3 = dgl.hetero_from_relations([g1, g2])
assert g3.is_unibipartite assert g3.is_unibipartite
assert g3.srctypes == ['A'] assert g3.srctypes == ['A']
...@@ -1570,31 +1611,49 @@ def test_bipartite(): ...@@ -1570,31 +1611,49 @@ def test_bipartite():
assert F.array_equal(g3.nodes['A'].data['h'], g3.srcdata['h']) assert F.array_equal(g3.nodes['A'].data['h'], g3.srcdata['h'])
assert F.array_equal(g3.nodes['SRC/A'].data['h'], g3.srcdata['h']) assert F.array_equal(g3.nodes['SRC/A'].data['h'], g3.srcdata['h'])
g4 = dgl.graph([(0, 0), (1, 1)], 'A', 'AA') g4 = dgl.graph([(0, 0), (1, 1)], 'A', 'AA', index_dtype=index_dtype)
g5 = dgl.hetero_from_relations([g1, g2, g4]) g5 = dgl.hetero_from_relations([g1, g2, g4])
assert not g5.is_unibipartite assert not g5.is_unibipartite
@parametrize_dtype
def test_dtype_cast(index_dtype):
g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], index_dtype=index_dtype)
assert g._idtype_str == index_dtype
g.ndata["feat"] = F.tensor([3, 4, 5])
g.edata["h"] = F.tensor([3, 4, 5, 6])
if index_dtype == "int32":
g_cast = g.long()
assert g_cast._idtype_str == 'int64'
else:
g_cast = g.int()
assert g_cast._idtype_str == 'int32'
assert "feat" in g_cast.ndata
assert "h" in g_cast.edata
assert F.array_equal(g.ndata["feat"], g_cast.ndata["feat"])
assert F.array_equal(g.edata["h"], g_cast.edata["h"])
if __name__ == '__main__': if __name__ == '__main__':
test_create() # test_create()
test_query() # test_query()
test_hypersparse() # test_hypersparse()
test_adj() # test_adj("int32")
test_inc() # test_inc()
test_view() # test_view()
test_view1() # test_view1("int32")
test_flatten() # test_flatten()
test_convert_bound() # test_convert_bound()
test_convert() # test_convert()
test_to_device() # test_to_device()
test_transform() # test_transform("int32")
test_subgraph() # test_subgraph()
test_apply() # test_apply()
test_level1() # test_level1()
test_level2() # test_level2()
test_updates() # test_updates()
test_backward() # test_backward()
test_empty_heterograph() # test_empty_heterograph()
test_types_in_function() # test_types_in_function()
test_stack_reduce() # test_stack_reduce()
test_isolated_ntype() # test_isolated_ntype()
test_bipartite() # test_bipartite()
test_dtype_cast()
tests/compute/test_randomwalk.py
View file @ dc8ca88e
...@@ -2,6 +2,8 @@ import dgl ...@@ -2,6 +2,8 @@ import dgl
from dgl import utils from dgl import utils
import backend as F import backend as F
import numpy as np import numpy as np
from utils import parametrize_dtype
import pytest
def test_random_walk(): def test_random_walk():
edge_list = [(0, 1), (1, 2), (2, 3), (3, 4), edge_list = [(0, 1), (1, 2), (2, 3), (3, 4),
...@@ -58,9 +60,10 @@ def test_random_walk_with_restart(): ...@@ -58,9 +60,10 @@ def test_random_walk_with_restart():
trace_diff = np.diff(F.zerocopy_to_numpy(t), axis=-1) trace_diff = np.diff(F.zerocopy_to_numpy(t), axis=-1)
assert (trace_diff % 2 == 0).all() assert (trace_diff % 2 == 0).all()
def test_metapath_random_walk(): @parametrize_dtype
g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b') def test_metapath_random_walk(index_dtype):
g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a') g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b', index_dtype=index_dtype)
g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a', index_dtype=index_dtype)
G = dgl.hetero_from_relations([g1, g2]) G = dgl.hetero_from_relations([g1, g2])
seeds = [0, 1] seeds = [0, 1]
traces = dgl.contrib.sampling.metapath_random_walk(G, ['ab', 'ba'] * 4, seeds, 3) traces = dgl.contrib.sampling.metapath_random_walk(G, ['ab', 'ba'] * 4, seeds, 3)
...@@ -74,5 +77,5 @@ def test_metapath_random_walk(): ...@@ -74,5 +77,5 @@ def test_metapath_random_walk():
assert g2.has_edge_between(trace[2 * i + 1], trace[2 * i + 2]) assert g2.has_edge_between(trace[2 * i + 1], trace[2 * i + 2])
if __name__ == '__main__': if __name__ == '__main__':
test_random_walk() # test_random_walk()
test_metapath_random_walk() test_metapath_random_walk("int32")
tests/compute/test_transform.py
View file @ dc8ca88e
...@@ -7,6 +7,7 @@ import dgl.function as fn ...@@ -7,6 +7,7 @@ import dgl.function as fn
import backend as F import backend as F
from dgl.graph_index import from_scipy_sparse_matrix from dgl.graph_index import from_scipy_sparse_matrix
import unittest import unittest
from utils import parametrize_dtype
D = 5 D = 5
...@@ -271,13 +272,15 @@ def test_metis_partition(): ...@@ -271,13 +272,15 @@ def test_metis_partition():
print(g.number_of_edges() - num_inner_edges) print(g.number_of_edges() - num_inner_edges)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_in_subgraph(): @parametrize_dtype
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow') def test_in_subgraph(index_dtype):
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game') g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user') g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin') g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
hg = dgl.hetero_from_relations([g1, g2, g3, g4]) hg = dgl.hetero_from_relations([g1, g2, g3, g4])
subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0}) subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0})
assert subg._idtype_str == index_dtype
assert len(subg.ntypes) == 3 assert len(subg.ntypes) == 3
assert len(subg.etypes) == 4 assert len(subg.etypes) == 4
u, v = subg['follow'].edges() u, v = subg['follow'].edges()
...@@ -295,13 +298,15 @@ def test_in_subgraph(): ...@@ -295,13 +298,15 @@ def test_in_subgraph():
assert subg['flips'].number_of_edges() == 0 assert subg['flips'].number_of_edges() == 0
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_out_subgraph(): @parametrize_dtype
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow') def test_out_subgraph(index_dtype):
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game') g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user') g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin') g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
hg = dgl.hetero_from_relations([g1, g2, g3, g4]) hg = dgl.hetero_from_relations([g1, g2, g3, g4])
subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0}) subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0})
assert subg._idtype_str == index_dtype
assert len(subg.ntypes) == 3 assert len(subg.ntypes) == 3
assert len(subg.etypes) == 4 assert len(subg.etypes) == 4
u, v = subg['follow'].edges() u, v = subg['follow'].edges()
...@@ -322,20 +327,21 @@ def test_out_subgraph(): ...@@ -322,20 +327,21 @@ def test_out_subgraph():
assert F.array_equal(hg['flips'].edge_ids(u, v), subg['flips'].edata[dgl.EID]) assert F.array_equal(hg['flips'].edge_ids(u, v), subg['flips'].edata[dgl.EID])
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
def test_compact(): @parametrize_dtype
def test_compact(index_dtype):
g1 = dgl.heterograph({ g1 = dgl.heterograph({
('user', 'follow', 'user'): [(1, 3), (3, 5)], ('user', 'follow', 'user'): [(1, 3), (3, 5)],
('user', 'plays', 'game'): [(2, 4), (3, 4), (2, 5)], ('user', 'plays', 'game'): [(2, 4), (3, 4), (2, 5)],
('game', 'wished-by', 'user'): [(6, 7), (5, 7)]}, ('game', 'wished-by', 'user'): [(6, 7), (5, 7)]},
{'user': 20, 'game': 10}) {'user': 20, 'game': 10}, index_dtype=index_dtype)
g2 = dgl.heterograph({ g2 = dgl.heterograph({
('game', 'clicked-by', 'user'): [(3, 1)], ('game', 'clicked-by', 'user'): [(3, 1)],
('user', 'likes', 'user'): [(1, 8), (8, 9)]}, ('user', 'likes', 'user'): [(1, 8), (8, 9)]},
{'user': 20, 'game': 10}) {'user': 20, 'game': 10}, index_dtype=index_dtype)
g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user') g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user', index_dtype=index_dtype)
g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user') g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user', index_dtype=index_dtype)
def _check(g, new_g, induced_nodes): def _check(g, new_g, induced_nodes):
assert g.ntypes == new_g.ntypes assert g.ntypes == new_g.ntypes
...@@ -358,13 +364,15 @@ def test_compact(): ...@@ -358,13 +364,15 @@ def test_compact():
new_g1 = dgl.compact_graphs(g1) new_g1 = dgl.compact_graphs(g1)
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes} induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7]) assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
assert set(induced_nodes['game']) == set([4, 5, 6]) assert set(induced_nodes['game']) == set([4, 5, 6])
_check(g1, new_g1, induced_nodes) _check(g1, new_g1, induced_nodes)
# Test with always_preserve given a dict # Test with always_preserve given a dict
new_g1 = dgl.compact_graphs( new_g1 = dgl.compact_graphs(
g1, always_preserve={'game': F.tensor([4, 7], dtype=F.int64)}) g1, always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
assert new_g1._idtype_str == index_dtype
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes} induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7]) assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
...@@ -373,9 +381,11 @@ def test_compact(): ...@@ -373,9 +381,11 @@ def test_compact():
# Test with always_preserve given a tensor # Test with always_preserve given a tensor
new_g3 = dgl.compact_graphs( new_g3 = dgl.compact_graphs(
g3, always_preserve=F.tensor([1, 7], dtype=F.int64)) g3, always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes} induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g3._idtype_str == index_dtype
assert set(induced_nodes['user']) == set([0, 1, 2, 7]) assert set(induced_nodes['user']) == set([0, 1, 2, 7])
_check(g3, new_g3, induced_nodes) _check(g3, new_g3, induced_nodes)
...@@ -383,6 +393,8 @@ def test_compact(): ...@@ -383,6 +393,8 @@ def test_compact():
new_g1, new_g2 = dgl.compact_graphs([g1, g2]) new_g1, new_g2 = dgl.compact_graphs([g1, g2])
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes} induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert new_g2._idtype_str == index_dtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9]) assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
assert set(induced_nodes['game']) == set([3, 4, 5, 6]) assert set(induced_nodes['game']) == set([3, 4, 5, 6])
_check(g1, new_g1, induced_nodes) _check(g1, new_g1, induced_nodes)
...@@ -390,9 +402,11 @@ def test_compact(): ...@@ -390,9 +402,11 @@ def test_compact():
# Test multiple graphs with always_preserve given a dict # Test multiple graphs with always_preserve given a dict
new_g1, new_g2 = dgl.compact_graphs( new_g1, new_g2 = dgl.compact_graphs(
[g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=F.int64)}) [g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes} induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert new_g2._idtype_str == index_dtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9]) assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
assert set(induced_nodes['game']) == set([3, 4, 5, 6, 7]) assert set(induced_nodes['game']) == set([3, 4, 5, 6, 7])
_check(g1, new_g1, induced_nodes) _check(g1, new_g1, induced_nodes)
...@@ -400,18 +414,21 @@ def test_compact(): ...@@ -400,18 +414,21 @@ def test_compact():
# Test multiple graphs with always_preserve given a tensor # Test multiple graphs with always_preserve given a tensor
new_g3, new_g4 = dgl.compact_graphs( new_g3, new_g4 = dgl.compact_graphs(
[g3, g4], always_preserve=F.tensor([1, 7], dtype=F.int64)) [g3, g4], always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes} induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()} induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g3._idtype_str == index_dtype
assert new_g4._idtype_str == index_dtype
assert set(induced_nodes['user']) == set([0, 1, 2, 3, 5, 7]) assert set(induced_nodes['user']) == set([0, 1, 2, 3, 5, 7])
_check(g3, new_g3, induced_nodes) _check(g3, new_g3, induced_nodes)
_check(g4, new_g4, induced_nodes) _check(g4, new_g4, induced_nodes)
@parametrize_dtype
def test_to_simple(): def test_to_simple(index_dtype):
g = dgl.heterograph({ g = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)], ('user', 'follow', 'user'): [(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)],
('user', 'plays', 'game'): [(3, 5), (2, 3), (1, 4), (1, 4), (3, 5), (2, 3), (2, 3)]}) ('user', 'plays', 'game'): [(3, 5), (2, 3), (1, 4), (1, 4), (3, 5), (2, 3), (2, 3)]}, index_dtype=index_dtype)
sg = dgl.to_simple(g, return_counts='weights', writeback_mapping='new_eid') sg = dgl.to_simple(g, return_counts='weights', writeback_mapping='new_eid')
for etype in g.canonical_etypes: for etype in g.canonical_etypes:
...@@ -434,7 +451,8 @@ def test_to_simple(): ...@@ -434,7 +451,8 @@ def test_to_simple():
assert eid_map[i] == suv.index(e) assert eid_map[i] == suv.index(e)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
def test_to_block(): @parametrize_dtype
def test_to_block(index_dtype):
def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True): def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
if dst_nodes is not None: if dst_nodes is not None:
assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes) assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
...@@ -471,7 +489,7 @@ def test_to_block(): ...@@ -471,7 +489,7 @@ def test_to_block():
g = dgl.heterograph({ g = dgl.heterograph({
('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)], ('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)], ('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
('B', 'BA', 'A'): [(2, 3), (3, 2)]}) ('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
g_a = g['AA'] g_a = g['AA']
bg = dgl.to_block(g_a) bg = dgl.to_block(g_a)
...@@ -484,39 +502,41 @@ def test_to_block(): ...@@ -484,39 +502,41 @@ def test_to_block():
assert bg.number_of_src_nodes() == 4 assert bg.number_of_src_nodes() == 4
assert bg.number_of_dst_nodes() == 4 assert bg.number_of_dst_nodes() == 4
dst_nodes = F.tensor([3, 4], dtype=F.int64) dst_nodes = F.tensor([3, 4], dtype=getattr(F, index_dtype))
bg = dgl.to_block(g_a, dst_nodes) bg = dgl.to_block(g_a, dst_nodes)
check(g_a, bg, 'A', 'AA', dst_nodes) check(g_a, bg, 'A', 'AA', dst_nodes)
dst_nodes = F.tensor([4, 3, 2, 1], dtype=F.int64) dst_nodes = F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype))
bg = dgl.to_block(g_a, dst_nodes) bg = dgl.to_block(g_a, dst_nodes)
check(g_a, bg, 'A', 'AA', dst_nodes) check(g_a, bg, 'A', 'AA', dst_nodes)
g_ab = g['AB'] g_ab = g['AB']
bg = dgl.to_block(g_ab) bg = dgl.to_block(g_ab)
assert bg._idtype_str == index_dtype
assert bg.number_of_nodes('SRC/B') == 4 assert bg.number_of_nodes('SRC/B') == 4
assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID]) assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
assert bg.number_of_nodes('DST/A') == 0 assert bg.number_of_nodes('DST/A') == 0
checkall(g_ab, bg, None) checkall(g_ab, bg, None)
dst_nodes = {'B': F.tensor([5, 6], dtype=F.int64)} dst_nodes = {'B': F.tensor([5, 6], dtype=getattr(F, index_dtype))}
bg = dgl.to_block(g, dst_nodes) bg = dgl.to_block(g, dst_nodes)
assert bg.number_of_nodes('SRC/B') == 2 assert bg.number_of_nodes('SRC/B') == 2
assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID]) assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
assert bg.number_of_nodes('DST/A') == 0 assert bg.number_of_nodes('DST/A') == 0
checkall(g, bg, dst_nodes) checkall(g, bg, dst_nodes)
dst_nodes = {'A': F.tensor([3, 4], dtype=F.int64), 'B': F.tensor([5, 6], dtype=F.int64)} dst_nodes = {'A': F.tensor([3, 4], dtype=getattr(F, index_dtype)), 'B': F.tensor([5, 6], dtype=getattr(F, index_dtype))}
bg = dgl.to_block(g, dst_nodes) bg = dgl.to_block(g, dst_nodes)
checkall(g, bg, dst_nodes) checkall(g, bg, dst_nodes)
dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=F.int64), 'B': F.tensor([3, 5, 6, 1], dtype=F.int64)} dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype)), 'B': F.tensor([3, 5, 6, 1], dtype=getattr(F, index_dtype))}
bg = dgl.to_block(g, dst_nodes=dst_nodes) bg = dgl.to_block(g, dst_nodes=dst_nodes)
checkall(g, bg, dst_nodes) checkall(g, bg, dst_nodes)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_remove_edges(): @parametrize_dtype
def test_remove_edges(index_dtype):
def check(g1, etype, g, edges_removed): def check(g1, etype, g, edges_removed):
src, dst, eid = g.edges(etype=etype, form='all') src, dst, eid = g.edges(etype=etype, form='all')
src1, dst1 = g1.edges(etype=etype, order='eid') src1, dst1 = g1.edges(etype=etype, order='eid')
...@@ -535,34 +555,35 @@ def test_remove_edges(): ...@@ -535,34 +555,35 @@ def test_remove_edges():
for s, d, e in zip(src1, dst1, eid1): for s, d, e in zip(src1, dst1, eid1):
assert (s, d, e) in sde_set assert (s, d, e) in sde_set
assert not np.isin(edges_removed, eid1).any() assert not np.isin(edges_removed, eid1).any()
assert g1.idtype == g.idtype
for fmt in ['coo', 'csr', 'csc']: for fmt in ['coo', 'csr', 'csc']:
for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]: for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]:
g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], restrict_format=fmt) g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], restrict_format=fmt, index_dtype=index_dtype)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove)) g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
check(g1, None, g, edges_to_remove) check(g1, None, g, edges_to_remove)
g = dgl.graph( g = dgl.graph(
spsp.csr_matrix(([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])), shape=(5, 5)), spsp.csr_matrix(([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])), shape=(5, 5)),
restrict_format=fmt) restrict_format=fmt, index_dtype=index_dtype)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove)) g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
check(g1, None, g, edges_to_remove) check(g1, None, g, edges_to_remove)
g = dgl.heterograph({ g = dgl.heterograph({
('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)], ('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)], ('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
('B', 'BA', 'A'): [(2, 3), (3, 2)]}) ('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
g2 = dgl.remove_edges(g, {'AA': F.tensor([2]), 'AB': F.tensor([3]), 'BA': F.tensor([1])}) g2 = dgl.remove_edges(g, {'AA': F.tensor([2], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
check(g2, 'AA', g, [2]) check(g2, 'AA', g, [2])
check(g2, 'AB', g, [3]) check(g2, 'AB', g, [3])
check(g2, 'BA', g, [1]) check(g2, 'BA', g, [1])
g3 = dgl.remove_edges(g, {'AA': F.tensor([]), 'AB': F.tensor([3]), 'BA': F.tensor([1])}) g3 = dgl.remove_edges(g, {'AA': F.tensor([], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
check(g3, 'AA', g, []) check(g3, 'AA', g, [])
check(g3, 'AB', g, [3]) check(g3, 'AB', g, [3])
check(g3, 'BA', g, [1]) check(g3, 'BA', g, [1])
g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0])}) g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0], getattr(F, index_dtype))})
check(g4, 'AA', g, []) check(g4, 'AA', g, [])
check(g4, 'AB', g, [3, 1, 2, 0]) check(g4, 'AB', g, [3, 1, 2, 0])
check(g4, 'BA', g, []) check(g4, 'BA', g, [])
...@@ -594,22 +615,22 @@ def test_cast(): ...@@ -594,22 +615,22 @@ def test_cast():
assert F.array_equal(g2dst, gdst) assert F.array_equal(g2dst, gdst)
if __name__ == '__main__': if __name__ == '__main__':
test_line_graph() # test_line_graph()
test_no_backtracking() # test_no_backtracking()
test_reverse() # test_reverse()
test_reverse_shared_frames() # test_reverse_shared_frames()
test_simple_graph() # test_simple_graph()
test_bidirected_graph() # test_bidirected_graph()
test_khop_adj() # test_khop_adj()
test_khop_graph() # test_khop_graph()
test_laplacian_lambda_max() # test_laplacian_lambda_max()
test_remove_self_loop() # test_remove_self_loop()
test_add_self_loop() # test_add_self_loop()
test_partition_with_halo() # test_partition_with_halo()
test_metis_partition() # test_metis_partition()
test_compact() # test_compact()
test_to_simple() # test_to_simple()
test_in_subgraph() # test_in_subgraph("int32")
test_out_subgraph() # test_out_subgraph()
test_to_block() test_to_block("int32")
test_remove_edges() # test_remove_edges()
tests/compute/utils.py
View file @ dc8ca88e
import pytest
parametrize_dtype = pytest.mark.parametrize("index_dtype", ['int32', 'int64'])
def check_fail(fn, *args, **kwargs): def check_fail(fn, *args, **kwargs):
try: try:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment