dmlc--dgl
2974 行
116 KiB
Python
2974 行
116 KiB
Python
import dgl
|
|
import dgl.function as fn
|
|
from collections import Counter
|
|
import numpy as np
|
|
import scipy.sparse as ssp
|
|
import itertools
|
|
import backend as F
|
|
import networkx as nx
|
|
import unittest, pytest
|
|
from dgl import DGLError
|
|
import test_utils
|
|
from test_utils import parametrize_idtype, get_cases
|
|
from utils import assert_is_identical_hetero
|
|
from scipy.sparse import rand
|
|
import multiprocessing as mp
|
|
|
|
def create_test_heterograph(idtype):
|
|
# test heterograph from the docstring, plus a user -- wishes -- game relation
|
|
# 3 users, 2 games, 2 developers
|
|
# metagraph:
|
|
# ('user', 'follows', 'user'),
|
|
# ('user', 'plays', 'game'),
|
|
# ('user', 'wishes', 'game'),
|
|
# ('developer', 'develops', 'game')])
|
|
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2]),
|
|
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
|
|
('developer', 'develops', 'game'): ([0, 1], [0, 1])
|
|
}, idtype=idtype, device=F.ctx())
|
|
assert g.idtype == idtype
|
|
assert g.device == F.ctx()
|
|
return g
|
|
|
|
def create_test_heterograph1(idtype):
|
|
edges = []
|
|
edges.extend([(0, 1), (1, 2)]) # follows
|
|
edges.extend([(0, 3), (1, 3), (2, 4), (1, 4)]) # plays
|
|
edges.extend([(0, 4), (2, 3)]) # wishes
|
|
edges.extend([(5, 3), (6, 4)]) # develops
|
|
edges = tuple(zip(*edges))
|
|
ntypes = F.tensor([0, 0, 0, 1, 1, 2, 2])
|
|
etypes = F.tensor([0, 0, 1, 1, 1, 1, 2, 2, 3, 3])
|
|
g0 = dgl.graph(edges, idtype=idtype, device=F.ctx())
|
|
g0.ndata[dgl.NTYPE] = ntypes
|
|
g0.edata[dgl.ETYPE] = etypes
|
|
return dgl.to_heterogeneous(g0, ['user', 'game', 'developer'],
|
|
['follows', 'plays', 'wishes', 'develops'])
|
|
|
|
def create_test_heterograph2(idtype):
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2]),
|
|
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
('user', 'wishes', 'game'): ('csr', ([0, 1, 1, 2], [1, 0], [])),
|
|
('developer', 'develops', 'game'): ('csc', ([0, 1, 2], [0, 1], [0, 1])),
|
|
}, idtype=idtype, device=F.ctx())
|
|
assert g.idtype == idtype
|
|
assert g.device == F.ctx()
|
|
return g
|
|
|
|
def create_test_heterograph3(idtype):
|
|
g = dgl.heterograph({
|
|
('user', 'plays', 'game'): (F.tensor([0, 1, 1, 2], dtype=idtype),
|
|
F.tensor([0, 0, 1, 1], dtype=idtype)),
|
|
('developer', 'develops', 'game'): (F.tensor([0, 1], dtype=idtype),
|
|
F.tensor([0, 1], dtype=idtype))},
|
|
idtype=idtype, device=F.ctx())
|
|
|
|
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['developer'].data['h'] = F.copy_to(F.tensor([3, 3], dtype=idtype), ctx=F.ctx())
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
return g
|
|
|
|
def create_test_heterograph4(idtype):
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([0, 1, 1, 2, 2, 2], dtype=idtype),
|
|
F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype)),
|
|
('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
|
|
F.tensor([0, 1], dtype=idtype))},
|
|
idtype=idtype, device=F.ctx())
|
|
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
|
|
g.edges['follows'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4, 5, 6], dtype=idtype), ctx=F.ctx())
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
|
|
return g
|
|
|
|
def create_test_heterograph5(idtype):
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([1, 2], dtype=idtype),
|
|
F.tensor([0, 1], dtype=idtype)),
|
|
('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
|
|
F.tensor([0, 1], dtype=idtype))},
|
|
idtype=idtype, device=F.ctx())
|
|
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
|
|
g.edges['follows'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
|
|
return g
|
|
|
|
def get_redfn(name):
|
|
return getattr(F, name)
|
|
|
|
@parametrize_idtype
|
|
def test_create(idtype):
|
|
device = F.ctx()
|
|
g0 = create_test_heterograph(idtype)
|
|
g1 = create_test_heterograph1(idtype)
|
|
g2 = create_test_heterograph2(idtype)
|
|
assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
|
|
assert set(g0.canonical_etypes) == set(g1.canonical_etypes) == set(g2.canonical_etypes)
|
|
|
|
# Create a bipartite graph from a SciPy matrix
|
|
src_ids = np.array([2, 3, 4])
|
|
dst_ids = np.array([1, 2, 3])
|
|
eweight = np.array([0.2, 0.3, 0.5])
|
|
sp_mat = ssp.coo_matrix((eweight, (src_ids, dst_ids)))
|
|
g = dgl.bipartite_from_scipy(sp_mat, utype='user', etype='plays',
|
|
vtype='game', idtype=idtype, device=device)
|
|
assert g.idtype == idtype
|
|
assert g.device == device
|
|
assert g.num_src_nodes() == 5
|
|
assert g.num_dst_nodes() == 4
|
|
assert g.num_edges() == 3
|
|
src, dst = g.edges()
|
|
assert F.allclose(src, F.tensor([2, 3, 4], dtype=idtype))
|
|
assert F.allclose(dst, F.tensor([1, 2, 3], dtype=idtype))
|
|
g = dgl.bipartite_from_scipy(sp_mat, utype='_U', etype='_E', vtype='_V',
|
|
eweight_name='w', idtype=idtype, device=device)
|
|
assert F.allclose(g.edata['w'], F.tensor(eweight))
|
|
|
|
# Create a bipartite graph from a NetworkX graph
|
|
nx_g = nx.DiGraph()
|
|
nx_g.add_nodes_from([1, 3], bipartite=0, feat1=np.zeros((2)), feat2=np.ones((2)))
|
|
nx_g.add_nodes_from([2, 4, 5], bipartite=1, feat3=np.zeros((3)))
|
|
nx_g.add_edge(1, 4, weight=np.ones((1)), eid=np.array([1]))
|
|
nx_g.add_edge(3, 5, weight=np.ones((1)), eid=np.array([0]))
|
|
g = dgl.bipartite_from_networkx(nx_g, utype='user', etype='plays',
|
|
vtype='game', idtype=idtype, device=device)
|
|
assert g.idtype == idtype
|
|
assert g.device == device
|
|
assert g.num_src_nodes() == 2
|
|
assert g.num_dst_nodes() == 3
|
|
assert g.num_edges() == 2
|
|
src, dst = g.edges()
|
|
assert F.allclose(src, F.tensor([0, 1], dtype=idtype))
|
|
assert F.allclose(dst, F.tensor([1, 2], dtype=idtype))
|
|
g = dgl.bipartite_from_networkx(nx_g, utype='_U', etype='_E', vtype='V',
|
|
u_attrs=['feat1', 'feat2'],
|
|
e_attrs = ['weight'], v_attrs = ['feat3'])
|
|
assert F.allclose(g.srcdata['feat1'], F.tensor(np.zeros((2, 2))))
|
|
assert F.allclose(g.srcdata['feat2'], F.tensor(np.ones((2, 2))))
|
|
assert F.allclose(g.dstdata['feat3'], F.tensor(np.zeros((3, 3))))
|
|
assert F.allclose(g.edata['weight'], F.tensor(np.ones((2, 1))))
|
|
g = dgl.bipartite_from_networkx(nx_g, utype='_U', etype='_E', vtype='V',
|
|
edge_id_attr_name='eid', idtype=idtype, device=device)
|
|
src, dst = g.edges()
|
|
assert F.allclose(src, F.tensor([1, 0], dtype=idtype))
|
|
assert F.allclose(dst, F.tensor([2, 1], dtype=idtype))
|
|
|
|
# create from scipy
|
|
spmat = ssp.coo_matrix(([1,1,1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
|
|
g = dgl.from_scipy(spmat, idtype=idtype, device=device)
|
|
assert g.num_nodes() == 4
|
|
assert g.num_edges() == 3
|
|
assert g.idtype == idtype
|
|
assert g.device == device
|
|
|
|
# test inferring number of nodes for heterograph
|
|
g = dgl.heterograph({
|
|
('l0', 'e0', 'l1'): ([0, 0], [1, 2]),
|
|
('l0', 'e1', 'l2'): ([2], [2]),
|
|
('l2', 'e2', 'l2'): ([1, 3], [1, 3])
|
|
}, idtype=idtype, device=device)
|
|
assert g.num_nodes('l0') == 3
|
|
assert g.num_nodes('l1') == 3
|
|
assert g.num_nodes('l2') == 4
|
|
assert g.idtype == idtype
|
|
assert g.device == device
|
|
|
|
# test if validate flag works
|
|
# homo graph
|
|
with pytest.raises(DGLError):
|
|
g = dgl.graph(
|
|
([0, 0, 0, 1, 1, 2], [0, 1, 2, 0, 1, 2]),
|
|
num_nodes=2,
|
|
idtype=idtype, device=device
|
|
)
|
|
# bipartite graph
|
|
def _test_validate_bipartite(card):
|
|
with pytest.raises(DGLError):
|
|
g = dgl.heterograph({
|
|
('_U', '_E', '_V'): ([0, 0, 1, 1, 2], [1, 1, 2, 2, 3])
|
|
}, {'_U': card[0], '_V': card[1]}, idtype=idtype, device=device)
|
|
|
|
_test_validate_bipartite((3, 3))
|
|
_test_validate_bipartite((2, 4))
|
|
|
|
# test from_scipy
|
|
num_nodes = 10
|
|
density = 0.25
|
|
for fmt in ['csr', 'coo', 'csc']:
|
|
adj = rand(num_nodes, num_nodes, density=density, format=fmt)
|
|
g = dgl.from_scipy(adj, eweight_name='w', idtype=idtype)
|
|
assert g.idtype == idtype
|
|
assert g.device == F.cpu()
|
|
assert F.array_equal(g.edata['w'], F.copy_to(F.tensor(adj.data), F.cpu()))
|
|
|
|
def test_create2():
|
|
mat = ssp.random(20, 30, 0.1)
|
|
|
|
# coo
|
|
mat = mat.tocoo()
|
|
row = F.tensor(mat.row, dtype=F.int64)
|
|
col = F.tensor(mat.col, dtype=F.int64)
|
|
g = dgl.heterograph(
|
|
{('A', 'AB', 'B'): ('coo', (row, col))}, num_nodes_dict={'A': 20, 'B': 30})
|
|
|
|
# csr
|
|
mat = mat.tocsr()
|
|
indptr = F.tensor(mat.indptr, dtype=F.int64)
|
|
indices = F.tensor(mat.indices, dtype=F.int64)
|
|
data = F.tensor([], dtype=F.int64)
|
|
g = dgl.heterograph(
|
|
{('A', 'AB', 'B'): ('csr', (indptr, indices, data))}, num_nodes_dict={'A': 20, 'B': 30})
|
|
|
|
# csc
|
|
mat = mat.tocsc()
|
|
indptr = F.tensor(mat.indptr, dtype=F.int64)
|
|
indices = F.tensor(mat.indices, dtype=F.int64)
|
|
data = F.tensor([], dtype=F.int64)
|
|
g = dgl.heterograph(
|
|
{('A', 'AB', 'B'): ('csc', (indptr, indices, data))}, num_nodes_dict={'A': 20, 'B': 30})
|
|
|
|
@parametrize_idtype
|
|
def test_query(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
|
|
ntypes = ['user', 'game', 'developer']
|
|
canonical_etypes = [
|
|
('user', 'follows', 'user'),
|
|
('user', 'plays', 'game'),
|
|
('user', 'wishes', 'game'),
|
|
('developer', 'develops', 'game')]
|
|
etypes = ['follows', 'plays', 'wishes', 'develops']
|
|
|
|
# node & edge types
|
|
assert set(ntypes) == set(g.ntypes)
|
|
assert set(etypes) == set(g.etypes)
|
|
assert set(canonical_etypes) == set(g.canonical_etypes)
|
|
|
|
# metagraph
|
|
mg = g.metagraph()
|
|
assert set(g.ntypes) == set(mg.nodes)
|
|
etype_triplets = [(u, v, e) for u, v, e in mg.edges(keys=True)]
|
|
assert set([
|
|
('user', 'user', 'follows'),
|
|
('user', 'game', 'plays'),
|
|
('user', 'game', 'wishes'),
|
|
('developer', 'game', 'develops')]) == set(etype_triplets)
|
|
for i in range(len(etypes)):
|
|
assert g.to_canonical_etype(etypes[i]) == canonical_etypes[i]
|
|
|
|
def _test(g):
|
|
# number of nodes
|
|
assert [g.num_nodes(ntype) for ntype in ntypes] == [3, 2, 2]
|
|
|
|
# number of edges
|
|
assert [g.num_edges(etype) for etype in etypes] == [2, 4, 2, 2]
|
|
|
|
# has_node & has_nodes
|
|
for ntype in ntypes:
|
|
n = g.number_of_nodes(ntype)
|
|
for i in range(n):
|
|
assert g.has_node(i, ntype)
|
|
assert not g.has_node(n, ntype)
|
|
assert np.array_equal(
|
|
F.asnumpy(g.has_nodes([0, n], ntype)).astype('int32'), [1, 0])
|
|
|
|
assert not g.is_multigraph
|
|
|
|
for etype in etypes:
|
|
srcs, dsts = edges[etype]
|
|
for src, dst in zip(srcs, dsts):
|
|
assert g.has_edges_between(src, dst, etype)
|
|
assert F.asnumpy(g.has_edges_between(srcs, dsts, etype)).all()
|
|
|
|
srcs, dsts = negative_edges[etype]
|
|
for src, dst in zip(srcs, dsts):
|
|
assert not g.has_edges_between(src, dst, etype)
|
|
assert not F.asnumpy(g.has_edges_between(srcs, dsts, etype)).any()
|
|
|
|
srcs, dsts = edges[etype]
|
|
n_edges = len(srcs)
|
|
|
|
# predecessors & in_edges & in_degree
|
|
pred = [s for s, d in zip(srcs, dsts) if d == 0]
|
|
assert set(F.asnumpy(g.predecessors(0, etype)).tolist()) == set(pred)
|
|
u, v = g.in_edges([0], etype=etype)
|
|
assert F.asnumpy(v).tolist() == [0] * len(pred)
|
|
assert set(F.asnumpy(u).tolist()) == set(pred)
|
|
assert g.in_degrees(0, etype) == len(pred)
|
|
|
|
# successors & out_edges & out_degree
|
|
succ = [d for s, d in zip(srcs, dsts) if s == 0]
|
|
assert set(F.asnumpy(g.successors(0, etype)).tolist()) == set(succ)
|
|
u, v = g.out_edges([0], etype=etype)
|
|
assert F.asnumpy(u).tolist() == [0] * len(succ)
|
|
assert set(F.asnumpy(v).tolist()) == set(succ)
|
|
assert g.out_degrees(0, etype) == len(succ)
|
|
|
|
# edge_id & edge_ids
|
|
for i, (src, dst) in enumerate(zip(srcs, dsts)):
|
|
assert g.edge_ids(src, dst, etype=etype) == i
|
|
_, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
|
|
assert eid == i
|
|
assert F.asnumpy(g.edge_ids(srcs, dsts, etype=etype)).tolist() == list(range(n_edges))
|
|
u, v, e = g.edge_ids(srcs, dsts, etype=etype, return_uv=True)
|
|
u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
|
|
assert u[e].tolist() == srcs
|
|
assert v[e].tolist() == dsts
|
|
|
|
# find_edges
|
|
for eid in [list(range(n_edges)), np.arange(n_edges), F.astype(F.arange(0, n_edges), g.idtype)]:
|
|
u, v = g.find_edges(eid, etype)
|
|
assert F.asnumpy(u).tolist() == srcs
|
|
assert F.asnumpy(v).tolist() == dsts
|
|
|
|
# all_edges.
|
|
for order in ['eid']:
|
|
u, v, e = g.edges('all', order, etype)
|
|
assert F.asnumpy(u).tolist() == srcs
|
|
assert F.asnumpy(v).tolist() == dsts
|
|
assert F.asnumpy(e).tolist() == list(range(n_edges))
|
|
|
|
# in_degrees & out_degrees
|
|
in_degrees = F.asnumpy(g.in_degrees(etype=etype))
|
|
out_degrees = F.asnumpy(g.out_degrees(etype=etype))
|
|
src_count = Counter(srcs)
|
|
dst_count = Counter(dsts)
|
|
utype, _, vtype = g.to_canonical_etype(etype)
|
|
for i in range(g.number_of_nodes(utype)):
|
|
assert out_degrees[i] == src_count[i]
|
|
for i in range(g.number_of_nodes(vtype)):
|
|
assert in_degrees[i] == dst_count[i]
|
|
|
|
edges = {
|
|
'follows': ([0, 1], [1, 2]),
|
|
'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
'wishes': ([0, 2], [1, 0]),
|
|
'develops': ([0, 1], [0, 1]),
|
|
}
|
|
# edges that does not exist in the graph
|
|
negative_edges = {
|
|
'follows': ([0, 1], [0, 1]),
|
|
'plays': ([0, 2], [1, 0]),
|
|
'wishes': ([0, 1], [0, 1]),
|
|
'develops': ([0, 1], [1, 0]),
|
|
}
|
|
g = create_test_heterograph(idtype)
|
|
_test(g)
|
|
g = create_test_heterograph1(idtype)
|
|
_test(g)
|
|
if F._default_context_str != 'gpu':
|
|
# XXX: CUDA COO operators have not been live yet.
|
|
g = create_test_heterograph2(idtype)
|
|
_test(g)
|
|
|
|
etypes = canonical_etypes
|
|
edges = {
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2]),
|
|
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
|
|
('developer', 'develops', 'game'): ([0, 1], [0, 1]),
|
|
}
|
|
# edges that does not exist in the graph
|
|
negative_edges = {
|
|
('user', 'follows', 'user'): ([0, 1], [0, 1]),
|
|
('user', 'plays', 'game'): ([0, 2], [1, 0]),
|
|
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
|
|
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
|
|
}
|
|
g = create_test_heterograph(idtype)
|
|
_test(g)
|
|
g = create_test_heterograph1(idtype)
|
|
_test(g)
|
|
if F._default_context_str != 'gpu':
|
|
# XXX: CUDA COO operators have not been live yet.
|
|
g = create_test_heterograph2(idtype)
|
|
_test(g)
|
|
|
|
# test repr
|
|
print(g)
|
|
|
|
@parametrize_idtype
|
|
def test_empty_query(idtype):
|
|
g = dgl.graph(([1, 2, 3], [0, 4, 5]), idtype=idtype, device=F.ctx())
|
|
g.add_nodes(0)
|
|
g.add_edges([], [])
|
|
g.remove_edges([])
|
|
g.remove_nodes([])
|
|
assert F.shape(g.has_nodes([])) == (0,)
|
|
assert F.shape(g.has_edges_between([], [])) == (0,)
|
|
g.edge_ids([], [])
|
|
g.edge_ids([], [], return_uv=True)
|
|
g.find_edges([])
|
|
|
|
assert F.shape(g.in_edges([], form='eid')) == (0,)
|
|
u, v = g.in_edges([], form='uv')
|
|
assert F.shape(u) == (0,)
|
|
assert F.shape(v) == (0,)
|
|
u, v, e = g.in_edges([], form='all')
|
|
assert F.shape(u) == (0,)
|
|
assert F.shape(v) == (0,)
|
|
assert F.shape(e) == (0,)
|
|
|
|
assert F.shape(g.out_edges([], form='eid')) == (0,)
|
|
u, v = g.out_edges([], form='uv')
|
|
assert F.shape(u) == (0,)
|
|
assert F.shape(v) == (0,)
|
|
u, v, e = g.out_edges([], form='all')
|
|
assert F.shape(u) == (0,)
|
|
assert F.shape(v) == (0,)
|
|
assert F.shape(e) == (0,)
|
|
|
|
assert F.shape(g.in_degrees([])) == (0,)
|
|
assert F.shape(g.out_degrees([])) == (0,)
|
|
|
|
g = dgl.graph(([], []), idtype=idtype, device=F.ctx())
|
|
error_thrown = True
|
|
try:
|
|
g.in_degrees([0])
|
|
fail = False
|
|
except:
|
|
pass
|
|
assert error_thrown
|
|
error_thrown = True
|
|
try:
|
|
g.out_degrees([0])
|
|
fail = False
|
|
except:
|
|
pass
|
|
assert error_thrown
|
|
|
|
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU does not have COO impl.")
|
|
def _test_hypersparse():
|
|
N1 = 1 << 50 # should crash if allocated a CSR
|
|
N2 = 1 << 48
|
|
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([0], F.int64), F.tensor([1], F.int64)),
|
|
('user', 'plays', 'game'): (F.tensor([0], F.int64), F.tensor([N2], F.int64))},
|
|
{'user': N1, 'game': N1},
|
|
device=F.ctx())
|
|
assert g.number_of_nodes('user') == N1
|
|
assert g.number_of_nodes('game') == N1
|
|
assert g.number_of_edges('follows') == 1
|
|
assert g.number_of_edges('plays') == 1
|
|
|
|
assert g.has_edges_between(0, 1, 'follows')
|
|
assert not g.has_edges_between(0, 0, 'follows')
|
|
mask = F.asnumpy(g.has_edges_between([0, 0], [0, 1], 'follows')).tolist()
|
|
assert mask == [0, 1]
|
|
|
|
assert g.has_edges_between(0, N2, 'plays')
|
|
assert not g.has_edges_between(0, 0, 'plays')
|
|
mask = F.asnumpy(g.has_edges_between([0, 0], [0, N2], 'plays')).tolist()
|
|
assert mask == [0, 1]
|
|
|
|
assert F.asnumpy(g.predecessors(0, 'follows')).tolist() == []
|
|
assert F.asnumpy(g.successors(0, 'follows')).tolist() == [1]
|
|
assert F.asnumpy(g.predecessors(1, 'follows')).tolist() == [0]
|
|
assert F.asnumpy(g.successors(1, 'follows')).tolist() == []
|
|
|
|
assert F.asnumpy(g.predecessors(0, 'plays')).tolist() == []
|
|
assert F.asnumpy(g.successors(0, 'plays')).tolist() == [N2]
|
|
assert F.asnumpy(g.predecessors(N2, 'plays')).tolist() == [0]
|
|
assert F.asnumpy(g.successors(N2, 'plays')).tolist() == []
|
|
|
|
assert g.edge_ids(0, 1, etype='follows') == 0
|
|
assert g.edge_ids(0, N2, etype='plays') == 0
|
|
|
|
u, v = g.find_edges([0], 'follows')
|
|
assert F.asnumpy(u).tolist() == [0]
|
|
assert F.asnumpy(v).tolist() == [1]
|
|
u, v = g.find_edges([0], 'plays')
|
|
assert F.asnumpy(u).tolist() == [0]
|
|
assert F.asnumpy(v).tolist() == [N2]
|
|
u, v, e = g.all_edges('all', 'eid', 'follows')
|
|
assert F.asnumpy(u).tolist() == [0]
|
|
assert F.asnumpy(v).tolist() == [1]
|
|
assert F.asnumpy(e).tolist() == [0]
|
|
u, v, e = g.all_edges('all', 'eid', 'plays')
|
|
assert F.asnumpy(u).tolist() == [0]
|
|
assert F.asnumpy(v).tolist() == [N2]
|
|
assert F.asnumpy(e).tolist() == [0]
|
|
|
|
assert g.in_degrees(0, 'follows') == 0
|
|
assert g.in_degrees(1, 'follows') == 1
|
|
assert F.asnumpy(g.in_degrees([0, 1], 'follows')).tolist() == [0, 1]
|
|
assert g.in_degrees(0, 'plays') == 0
|
|
assert g.in_degrees(N2, 'plays') == 1
|
|
assert F.asnumpy(g.in_degrees([0, N2], 'plays')).tolist() == [0, 1]
|
|
assert g.out_degrees(0, 'follows') == 1
|
|
assert g.out_degrees(1, 'follows') == 0
|
|
assert F.asnumpy(g.out_degrees([0, 1], 'follows')).tolist() == [1, 0]
|
|
assert g.out_degrees(0, 'plays') == 1
|
|
assert g.out_degrees(N2, 'plays') == 0
|
|
assert F.asnumpy(g.out_degrees([0, N2], 'plays')).tolist() == [1, 0]
|
|
|
|
def _test_edge_ids():
|
|
N1 = 1 << 50 # should crash if allocated a CSR
|
|
N2 = 1 << 48
|
|
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([0], F.int64), F.tensor([1], F.int64)),
|
|
('user', 'plays', 'game'): (F.tensor([0], F.int64), F.tensor([N2], F.int64))},
|
|
{'user': N1, 'game': N1})
|
|
with pytest.raises(DGLError):
|
|
eid = g.edge_ids(0, 0, etype='follows')
|
|
|
|
g2 = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([0, 0], F.int64), F.tensor([1, 1], F.int64)),
|
|
('user', 'plays', 'game'): (F.tensor([0], F.int64), F.tensor([N2], F.int64))},
|
|
{'user': N1, 'game': N1}, device=F.cpu())
|
|
|
|
eid = g2.edge_ids(0, 1, etype='follows')
|
|
assert eid == 0
|
|
|
|
@parametrize_idtype
|
|
def test_adj(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
adj = F.sparse_to_numpy(g.adj(transpose=True, etype='follows'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[0., 0., 0.],
|
|
[1., 0., 0.],
|
|
[0., 1., 0.]]))
|
|
adj = F.sparse_to_numpy(g.adj(transpose=False, etype='follows'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[0., 1., 0.],
|
|
[0., 0., 1.],
|
|
[0., 0., 0.]]))
|
|
adj = F.sparse_to_numpy(g.adj(transpose=True, etype='plays'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[1., 1., 0.],
|
|
[0., 1., 1.]]))
|
|
adj = F.sparse_to_numpy(g.adj(transpose=False, etype='plays'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[1., 0.],
|
|
[1., 1.],
|
|
[0., 1.]]))
|
|
|
|
adj = g.adj(transpose=True, scipy_fmt='csr', etype='follows')
|
|
assert np.allclose(
|
|
adj.todense(),
|
|
np.array([[0., 0., 0.],
|
|
[1., 0., 0.],
|
|
[0., 1., 0.]]))
|
|
adj = g.adj(transpose=True, scipy_fmt='coo', etype='follows')
|
|
assert np.allclose(
|
|
adj.todense(),
|
|
np.array([[0., 0., 0.],
|
|
[1., 0., 0.],
|
|
[0., 1., 0.]]))
|
|
adj = g.adj(transpose=True, scipy_fmt='csr', etype='plays')
|
|
assert np.allclose(
|
|
adj.todense(),
|
|
np.array([[1., 1., 0.],
|
|
[0., 1., 1.]]))
|
|
adj = g.adj(transpose=True, scipy_fmt='coo', etype='plays')
|
|
assert np.allclose(
|
|
adj.todense(),
|
|
np.array([[1., 1., 0.],
|
|
[0., 1., 1.]]))
|
|
adj = F.sparse_to_numpy(g['follows'].adj(transpose=True))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[0., 0., 0.],
|
|
[1., 0., 0.],
|
|
[0., 1., 0.]]))
|
|
|
|
@parametrize_idtype
|
|
def test_inc(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
adj = F.sparse_to_numpy(g['follows'].inc('in'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[0., 0.],
|
|
[1., 0.],
|
|
[0., 1.]]))
|
|
adj = F.sparse_to_numpy(g['follows'].inc('out'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[1., 0.],
|
|
[0., 1.],
|
|
[0., 0.]]))
|
|
adj = F.sparse_to_numpy(g['follows'].inc('both'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[-1., 0.],
|
|
[1., -1.],
|
|
[0., 1.]]))
|
|
adj = F.sparse_to_numpy(g.inc('in', etype='plays'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[1., 1., 0., 0.],
|
|
[0., 0., 1., 1.]]))
|
|
adj = F.sparse_to_numpy(g.inc('out', etype='plays'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[1., 0., 0., 0.],
|
|
[0., 1., 0., 1.],
|
|
[0., 0., 1., 0.]]))
|
|
adj = F.sparse_to_numpy(g.inc('both', etype='follows'))
|
|
assert np.allclose(
|
|
adj,
|
|
np.array([[-1., 0.],
|
|
[1., -1.],
|
|
[0., 1.]]))
|
|
|
|
@parametrize_idtype
|
|
def test_view(idtype):
|
|
# test single node type
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2])
|
|
}, idtype=idtype, device=F.ctx())
|
|
f1 = F.randn((3, 6))
|
|
g.ndata['h'] = f1
|
|
f2 = g.nodes['user'].data['h']
|
|
assert F.array_equal(f1, f2)
|
|
fail = False
|
|
try:
|
|
g.ndata['h'] = {'user' : f1}
|
|
except Exception:
|
|
fail = True
|
|
assert fail
|
|
|
|
# test single edge type
|
|
f3 = F.randn((2, 4))
|
|
g.edata['h'] = f3
|
|
f4 = g.edges['follows'].data['h']
|
|
assert F.array_equal(f3, f4)
|
|
fail = False
|
|
try:
|
|
g.edata['h'] = {'follows' : f3}
|
|
except Exception:
|
|
fail = True
|
|
assert fail
|
|
|
|
# test data view
|
|
g = create_test_heterograph(idtype)
|
|
|
|
f1 = F.randn((3, 6))
|
|
g.nodes['user'].data['h'] = f1 # ok
|
|
f2 = g.nodes['user'].data['h']
|
|
assert F.array_equal(f1, f2)
|
|
assert F.array_equal(g.nodes('user'), F.arange(0, 3, idtype))
|
|
g.nodes['user'].data.pop('h')
|
|
|
|
# multi type ndata
|
|
f1 = F.randn((3, 6))
|
|
f2 = F.randn((2, 6))
|
|
fail = False
|
|
try:
|
|
g.ndata['h'] = f1
|
|
except Exception:
|
|
fail = True
|
|
assert fail
|
|
|
|
f3 = F.randn((2, 4))
|
|
g.edges['user', 'follows', 'user'].data['h'] = f3
|
|
f4 = g.edges['user', 'follows', 'user'].data['h']
|
|
f5 = g.edges['follows'].data['h']
|
|
assert F.array_equal(f3, f4)
|
|
assert F.array_equal(f3, f5)
|
|
assert F.array_equal(g.edges(etype='follows', form='eid'), F.arange(0, 2, idtype))
|
|
g.edges['follows'].data.pop('h')
|
|
|
|
f3 = F.randn((2, 4))
|
|
fail = False
|
|
try:
|
|
g.edata['h'] = f3
|
|
except Exception:
|
|
fail = True
|
|
assert fail
|
|
|
|
# test srcdata
|
|
f1 = F.randn((3, 6))
|
|
g.srcnodes['user'].data['h'] = f1 # ok
|
|
f2 = g.srcnodes['user'].data['h']
|
|
assert F.array_equal(f1, f2)
|
|
assert F.array_equal(g.srcnodes('user'), F.arange(0, 3, idtype))
|
|
g.srcnodes['user'].data.pop('h')
|
|
|
|
# test dstdata
|
|
f1 = F.randn((3, 6))
|
|
g.dstnodes['user'].data['h'] = f1 # ok
|
|
f2 = g.dstnodes['user'].data['h']
|
|
assert F.array_equal(f1, f2)
|
|
assert F.array_equal(g.dstnodes('user'), F.arange(0, 3, idtype))
|
|
g.dstnodes['user'].data.pop('h')
|
|
|
|
@parametrize_idtype
|
|
def test_view1(idtype):
|
|
# test relation view
|
|
HG = create_test_heterograph(idtype)
|
|
ntypes = ['user', 'game', 'developer']
|
|
canonical_etypes = [
|
|
('user', 'follows', 'user'),
|
|
('user', 'plays', 'game'),
|
|
('user', 'wishes', 'game'),
|
|
('developer', 'develops', 'game')]
|
|
etypes = ['follows', 'plays', 'wishes', 'develops']
|
|
|
|
def _test_query():
|
|
for etype in etypes:
|
|
utype, _, vtype = HG.to_canonical_etype(etype)
|
|
g = HG[etype]
|
|
srcs, dsts = edges[etype]
|
|
for src, dst in zip(srcs, dsts):
|
|
assert g.has_edges_between(src, dst)
|
|
assert F.asnumpy(g.has_edges_between(srcs, dsts)).all()
|
|
|
|
srcs, dsts = negative_edges[etype]
|
|
for src, dst in zip(srcs, dsts):
|
|
assert not g.has_edges_between(src, dst)
|
|
assert not F.asnumpy(g.has_edges_between(srcs, dsts)).any()
|
|
|
|
srcs, dsts = edges[etype]
|
|
n_edges = len(srcs)
|
|
|
|
# predecessors & in_edges & in_degree
|
|
pred = [s for s, d in zip(srcs, dsts) if d == 0]
|
|
assert set(F.asnumpy(g.predecessors(0)).tolist()) == set(pred)
|
|
u, v = g.in_edges([0])
|
|
assert F.asnumpy(v).tolist() == [0] * len(pred)
|
|
assert set(F.asnumpy(u).tolist()) == set(pred)
|
|
assert g.in_degrees(0) == len(pred)
|
|
|
|
# successors & out_edges & out_degree
|
|
succ = [d for s, d in zip(srcs, dsts) if s == 0]
|
|
assert set(F.asnumpy(g.successors(0)).tolist()) == set(succ)
|
|
u, v = g.out_edges([0])
|
|
assert F.asnumpy(u).tolist() == [0] * len(succ)
|
|
assert set(F.asnumpy(v).tolist()) == set(succ)
|
|
assert g.out_degrees(0) == len(succ)
|
|
|
|
# edge_id & edge_ids
|
|
for i, (src, dst) in enumerate(zip(srcs, dsts)):
|
|
assert g.edge_ids(src, dst, etype=etype) == i
|
|
_, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
|
|
assert eid == i
|
|
assert F.asnumpy(g.edge_ids(srcs, dsts)).tolist() == list(range(n_edges))
|
|
u, v, e = g.edge_ids(srcs, dsts, return_uv=True)
|
|
u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
|
|
assert u[e].tolist() == srcs
|
|
assert v[e].tolist() == dsts
|
|
|
|
# find_edges
|
|
u, v = g.find_edges(list(range(n_edges)))
|
|
assert F.asnumpy(u).tolist() == srcs
|
|
assert F.asnumpy(v).tolist() == dsts
|
|
|
|
# all_edges.
|
|
for order in ['eid']:
|
|
u, v, e = g.all_edges(form='all', order=order)
|
|
assert F.asnumpy(u).tolist() == srcs
|
|
assert F.asnumpy(v).tolist() == dsts
|
|
assert F.asnumpy(e).tolist() == list(range(n_edges))
|
|
|
|
# in_degrees & out_degrees
|
|
in_degrees = F.asnumpy(g.in_degrees())
|
|
out_degrees = F.asnumpy(g.out_degrees())
|
|
src_count = Counter(srcs)
|
|
dst_count = Counter(dsts)
|
|
for i in range(g.number_of_nodes(utype)):
|
|
assert out_degrees[i] == src_count[i]
|
|
for i in range(g.number_of_nodes(vtype)):
|
|
assert in_degrees[i] == dst_count[i]
|
|
|
|
edges = {
|
|
'follows': ([0, 1], [1, 2]),
|
|
'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
'wishes': ([0, 2], [1, 0]),
|
|
'develops': ([0, 1], [0, 1]),
|
|
}
|
|
# edges that does not exist in the graph
|
|
negative_edges = {
|
|
'follows': ([0, 1], [0, 1]),
|
|
'plays': ([0, 2], [1, 0]),
|
|
'wishes': ([0, 1], [0, 1]),
|
|
'develops': ([0, 1], [1, 0]),
|
|
}
|
|
_test_query()
|
|
etypes = canonical_etypes
|
|
edges = {
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2]),
|
|
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
|
|
('developer', 'develops', 'game'): ([0, 1], [0, 1]),
|
|
}
|
|
# edges that does not exist in the graph
|
|
negative_edges = {
|
|
('user', 'follows', 'user'): ([0, 1], [0, 1]),
|
|
('user', 'plays', 'game'): ([0, 2], [1, 0]),
|
|
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
|
|
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
|
|
}
|
|
_test_query()
|
|
|
|
# test features
|
|
HG.nodes['user'].data['h'] = F.ones((HG.number_of_nodes('user'), 5))
|
|
HG.nodes['game'].data['m'] = F.ones((HG.number_of_nodes('game'), 3)) * 2
|
|
|
|
# test only one node type
|
|
g = HG['follows']
|
|
assert g.number_of_nodes() == 3
|
|
|
|
# test ndata and edata
|
|
f1 = F.randn((3, 6))
|
|
g.ndata['h'] = f1 # ok
|
|
f2 = HG.nodes['user'].data['h']
|
|
assert F.array_equal(f1, f2)
|
|
assert F.array_equal(g.nodes(), F.arange(0, 3, g.idtype))
|
|
|
|
f3 = F.randn((2, 4))
|
|
g.edata['h'] = f3
|
|
f4 = HG.edges['follows'].data['h']
|
|
assert F.array_equal(f3, f4)
|
|
assert F.array_equal(g.edges(form='eid'), F.arange(0, 2, g.idtype))
|
|
|
|
@parametrize_idtype
|
|
def test_flatten(idtype):
|
|
def check_mapping(g, fg):
|
|
if len(fg.ntypes) == 1:
|
|
SRC = DST = fg.ntypes[0]
|
|
else:
|
|
SRC = fg.ntypes[0]
|
|
DST = fg.ntypes[1]
|
|
|
|
etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
|
|
eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
|
|
|
|
for i, (etype, eid) in enumerate(zip(etypes, eids)):
|
|
src_g, dst_g = g.find_edges([eid], g.canonical_etypes[etype])
|
|
src_fg, dst_fg = fg.find_edges([i])
|
|
# 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?
|
|
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()
|
|
assert g.canonical_etypes[etype][0] == g.ntypes[tid]
|
|
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()
|
|
assert g.canonical_etypes[etype][2] == g.ntypes[tid]
|
|
|
|
# check for wildcard slices
|
|
g = create_test_heterograph(idtype)
|
|
g.nodes['user'].data['h'] = F.ones((3, 5))
|
|
g.nodes['game'].data['i'] = F.ones((2, 5))
|
|
g.edges['plays'].data['e'] = F.ones((4, 4))
|
|
g.edges['wishes'].data['e'] = F.ones((2, 4))
|
|
g.edges['wishes'].data['f'] = F.ones((2, 4))
|
|
|
|
fg = g['user', :, 'game'] # user--plays->game and user--wishes->game
|
|
assert len(fg.ntypes) == 2
|
|
assert fg.ntypes == ['user', 'game']
|
|
assert fg.etypes == ['plays+wishes']
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
etype = fg.etypes[0]
|
|
assert fg[etype] is not None # Issue #2166
|
|
|
|
assert F.array_equal(fg.nodes['user'].data['h'], F.ones((3, 5)))
|
|
assert F.array_equal(fg.nodes['game'].data['i'], F.ones((2, 5)))
|
|
assert F.array_equal(fg.edata['e'], F.ones((6, 4)))
|
|
assert 'f' not in fg.edata
|
|
|
|
etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
|
|
eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
|
|
assert set(zip(etypes, eids)) == set([(3, 0), (3, 1), (2, 1), (2, 0), (2, 3), (2, 2)])
|
|
|
|
check_mapping(g, fg)
|
|
|
|
fg = g['user', :, 'user']
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
# NOTE(gq): The node/edge types from the parent graph is returned if there is only one
|
|
# node/edge type. This differs from the behavior above.
|
|
assert fg.ntypes == ['user']
|
|
assert fg.etypes == ['follows']
|
|
u1, v1 = g.edges(etype='follows', order='eid')
|
|
u2, v2 = fg.edges(etype='follows', order='eid')
|
|
assert F.array_equal(u1, u2)
|
|
assert F.array_equal(v1, v2)
|
|
|
|
fg = g['developer', :, 'game']
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
assert fg.ntypes == ['developer', 'game']
|
|
assert fg.etypes == ['develops']
|
|
u1, v1 = g.edges(etype='develops', order='eid')
|
|
u2, v2 = fg.edges(etype='develops', order='eid')
|
|
assert F.array_equal(u1, u2)
|
|
assert F.array_equal(v1, v2)
|
|
|
|
fg = g[:, :, :]
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
assert fg.ntypes == ['developer+user', 'game+user']
|
|
assert fg.etypes == ['develops+follows+plays+wishes']
|
|
check_mapping(g, fg)
|
|
|
|
# Test another heterograph
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1, 2], [1, 2, 3]),
|
|
('user', 'knows', 'user'): ([0, 2], [2, 3])
|
|
}, idtype=idtype, device=F.ctx())
|
|
g.nodes['user'].data['h'] = F.randn((4, 3))
|
|
g.edges['follows'].data['w'] = F.randn((3, 2))
|
|
g.nodes['user'].data['hh'] = F.randn((4, 5))
|
|
g.edges['knows'].data['ww'] = F.randn((2, 10))
|
|
|
|
fg = g['user', :, 'user']
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
assert fg.ntypes == ['user']
|
|
assert fg.etypes == ['follows+knows']
|
|
check_mapping(g, fg)
|
|
|
|
fg = g['user', :, :]
|
|
assert fg.idtype == g.idtype
|
|
assert fg.device == g.device
|
|
assert fg.ntypes == ['user']
|
|
assert fg.etypes == ['follows+knows']
|
|
check_mapping(g, fg)
|
|
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
|
|
@parametrize_idtype
|
|
def test_to_device(idtype):
|
|
# TODO: rewrite this test case to accept different graphs so we
|
|
# can test reverse graph and batched graph
|
|
g = create_test_heterograph(idtype)
|
|
g.nodes['user'].data['h'] = F.ones((3, 5))
|
|
g.nodes['game'].data['i'] = F.ones((2, 5))
|
|
g.edges['plays'].data['e'] = F.ones((4, 4))
|
|
assert g.device == F.ctx()
|
|
g = g.to(F.cpu())
|
|
assert g.device == F.cpu()
|
|
assert F.context(g.nodes['user'].data['h']) == F.cpu()
|
|
assert F.context(g.nodes['game'].data['i']) == F.cpu()
|
|
assert F.context(g.edges['plays'].data['e']) == F.cpu()
|
|
for ntype in g.ntypes:
|
|
assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
|
|
for etype in g.canonical_etypes:
|
|
assert F.context(g.batch_num_edges(etype)) == F.cpu()
|
|
|
|
if F.is_cuda_available():
|
|
g1 = g.to(F.cuda())
|
|
assert g1.device == F.cuda()
|
|
assert F.context(g1.nodes['user'].data['h']) == F.cuda()
|
|
assert F.context(g1.nodes['game'].data['i']) == F.cuda()
|
|
assert F.context(g1.edges['plays'].data['e']) == F.cuda()
|
|
for ntype in g1.ntypes:
|
|
assert F.context(g1.batch_num_nodes(ntype)) == F.cuda()
|
|
for etype in g1.canonical_etypes:
|
|
assert F.context(g1.batch_num_edges(etype)) == F.cuda()
|
|
assert F.context(g.nodes['user'].data['h']) == F.cpu()
|
|
assert F.context(g.nodes['game'].data['i']) == F.cpu()
|
|
assert F.context(g.edges['plays'].data['e']) == F.cpu()
|
|
for ntype in g.ntypes:
|
|
assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
|
|
for etype in g.canonical_etypes:
|
|
assert F.context(g.batch_num_edges(etype)) == F.cpu()
|
|
with pytest.raises(DGLError):
|
|
g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cpu())
|
|
with pytest.raises(DGLError):
|
|
g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cpu())
|
|
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize('g', get_cases(['block']))
|
|
def test_to_device2(g, idtype):
|
|
g = g.astype(idtype)
|
|
g = g.to(F.cpu())
|
|
assert g.device == F.cpu()
|
|
if F.is_cuda_available():
|
|
g1 = g.to(F.cuda())
|
|
assert g1.device == F.cuda()
|
|
assert g1.ntypes == g.ntypes
|
|
assert g1.etypes == g.etypes
|
|
assert g1.canonical_etypes == g.canonical_etypes
|
|
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
|
|
@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="Pinning graph inplace only supported for PyTorch")
|
|
@parametrize_idtype
|
|
def test_pin_memory_(idtype):
|
|
# TODO: rewrite this test case to accept different graphs so we
|
|
# can test reverse graph and batched graph
|
|
g = create_test_heterograph(idtype)
|
|
g.nodes['user'].data['h'] = F.ones((3, 5))
|
|
g.nodes['game'].data['i'] = F.ones((2, 5))
|
|
g.edges['plays'].data['e'] = F.ones((4, 4))
|
|
g = g.to(F.cpu())
|
|
assert not g.is_pinned()
|
|
|
|
# unpin an unpinned CPU graph, directly return
|
|
g.unpin_memory_()
|
|
assert not g.is_pinned()
|
|
assert g.device == F.cpu()
|
|
|
|
# pin a CPU graph
|
|
g.pin_memory_()
|
|
assert g.is_pinned()
|
|
assert g.device == F.cpu()
|
|
assert g.nodes['user'].data['h'].is_pinned()
|
|
assert g.nodes['game'].data['i'].is_pinned()
|
|
assert g.edges['plays'].data['e'].is_pinned()
|
|
assert F.context(g.nodes['user'].data['h']) == F.cpu()
|
|
assert F.context(g.nodes['game'].data['i']) == F.cpu()
|
|
assert F.context(g.edges['plays'].data['e']) == F.cpu()
|
|
for ntype in g.ntypes:
|
|
assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
|
|
for etype in g.canonical_etypes:
|
|
assert F.context(g.batch_num_edges(etype)) == F.cpu()
|
|
|
|
# it's fine to clone with new formats, but new graphs are not pinned
|
|
# >>> g.formats()
|
|
# {'created': ['coo'], 'not created': ['csr', 'csc']}
|
|
assert not g.formats('csc').is_pinned()
|
|
assert not g.formats('csr').is_pinned()
|
|
# 'coo' formats is already created and thus not cloned
|
|
assert g.formats('coo').is_pinned()
|
|
|
|
# pin a pinned graph, directly return
|
|
g.pin_memory_()
|
|
assert g.is_pinned()
|
|
assert g.device == F.cpu()
|
|
|
|
# unpin a pinned graph
|
|
g.unpin_memory_()
|
|
assert not g.is_pinned()
|
|
assert g.device == F.cpu()
|
|
|
|
g1 = g.to(F.cuda())
|
|
|
|
# unpin an unpinned GPU graph, directly return
|
|
g1.unpin_memory_()
|
|
assert not g1.is_pinned()
|
|
assert g1.device == F.cuda()
|
|
|
|
# error pinning a GPU graph
|
|
with pytest.raises(DGLError):
|
|
g1.pin_memory_()
|
|
|
|
# test pin empty homograph
|
|
g2 = dgl.graph(([], []))
|
|
g2.pin_memory_()
|
|
assert g2.is_pinned()
|
|
g2.unpin_memory_()
|
|
assert not g2.is_pinned()
|
|
|
|
# test pin heterograph with 0 edge of one relation type
|
|
g3 = dgl.heterograph({
|
|
('a','b','c'): ([0, 1], [1, 2]),
|
|
('c','d','c'): ([], [])}).astype(idtype)
|
|
g3.pin_memory_()
|
|
assert g3.is_pinned()
|
|
g3.unpin_memory_()
|
|
assert not g3.is_pinned()
|
|
|
|
@parametrize_idtype
|
|
def test_convert_bound(idtype):
|
|
def _test_bipartite_bound(data, card):
|
|
with pytest.raises(DGLError):
|
|
dgl.heterograph({
|
|
('_U', '_E', '_V'): data
|
|
}, {'_U': card[0], '_V': card[1]}, idtype=idtype, device=F.ctx())
|
|
|
|
def _test_graph_bound(data, card):
|
|
with pytest.raises(DGLError):
|
|
dgl.graph(data, num_nodes=card, idtype=idtype, device=F.ctx())
|
|
|
|
_test_bipartite_bound(([1, 2], [1, 2]), (2, 3))
|
|
_test_bipartite_bound(([0, 1], [1, 4]), (2, 3))
|
|
_test_graph_bound(([1, 3], [1, 2]), 3)
|
|
_test_graph_bound(([0, 1], [1, 3]), 3)
|
|
|
|
|
|
@parametrize_idtype
|
|
def test_convert(idtype):
|
|
hg = create_test_heterograph(idtype)
|
|
hs = []
|
|
for ntype in hg.ntypes:
|
|
h = F.randn((hg.number_of_nodes(ntype), 5))
|
|
hg.nodes[ntype].data['h'] = h
|
|
hs.append(h)
|
|
hg.nodes['user'].data['x'] = F.randn((3, 3))
|
|
ws = []
|
|
for etype in hg.canonical_etypes:
|
|
w = F.randn((hg.number_of_edges(etype), 5))
|
|
hg.edges[etype].data['w'] = w
|
|
ws.append(w)
|
|
hg.edges['plays'].data['x'] = F.randn((4, 3))
|
|
|
|
g = dgl.to_homogeneous(hg, ndata=['h'], edata=['w'])
|
|
assert g.idtype == idtype
|
|
assert g.device == hg.device
|
|
assert F.array_equal(F.cat(hs, dim=0), g.ndata['h'])
|
|
assert 'x' not in g.ndata
|
|
assert F.array_equal(F.cat(ws, dim=0), g.edata['w'])
|
|
assert 'x' not in g.edata
|
|
|
|
src, dst = g.all_edges(order='eid')
|
|
src = F.asnumpy(src)
|
|
dst = F.asnumpy(dst)
|
|
etype_id, eid = F.asnumpy(g.edata[dgl.ETYPE]), F.asnumpy(g.edata[dgl.EID])
|
|
ntype_id, nid = F.asnumpy(g.ndata[dgl.NTYPE]), F.asnumpy(g.ndata[dgl.NID])
|
|
for i in range(g.number_of_edges()):
|
|
srctype = hg.ntypes[ntype_id[src[i]]]
|
|
dsttype = hg.ntypes[ntype_id[dst[i]]]
|
|
etype = hg.etypes[etype_id[i]]
|
|
src_i, dst_i = hg.find_edges([eid[i]], (srctype, etype, dsttype))
|
|
assert np.asscalar(F.asnumpy(src_i)) == nid[src[i]]
|
|
assert np.asscalar(F.asnumpy(dst_i)) == nid[dst[i]]
|
|
|
|
mg = nx.MultiDiGraph([
|
|
('user', 'user', 'follows'),
|
|
('user', 'game', 'plays'),
|
|
('user', 'game', 'wishes'),
|
|
('developer', 'game', 'develops')])
|
|
|
|
for _mg in [None, mg]:
|
|
hg2 = dgl.to_heterogeneous(
|
|
g, hg.ntypes, hg.etypes,
|
|
ntype_field=dgl.NTYPE, etype_field=dgl.ETYPE, metagraph=_mg)
|
|
assert hg2.idtype == hg.idtype
|
|
assert hg2.device == hg.device
|
|
assert set(hg.ntypes) == set(hg2.ntypes)
|
|
assert set(hg.canonical_etypes) == set(hg2.canonical_etypes)
|
|
for ntype in hg.ntypes:
|
|
assert hg.number_of_nodes(ntype) == hg2.number_of_nodes(ntype)
|
|
assert F.array_equal(hg.nodes[ntype].data['h'], hg2.nodes[ntype].data['h'])
|
|
for canonical_etype in hg.canonical_etypes:
|
|
src, dst = hg.all_edges(etype=canonical_etype, order='eid')
|
|
src2, dst2 = hg2.all_edges(etype=canonical_etype, order='eid')
|
|
assert F.array_equal(src, src2)
|
|
assert F.array_equal(dst, dst2)
|
|
assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w'])
|
|
|
|
# hetero_from_homo test case 2
|
|
g = dgl.graph(([0, 1, 2, 0], [2, 2, 3, 3]), idtype=idtype, device=F.ctx())
|
|
g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
|
|
g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
|
|
hg = dgl.to_heterogeneous(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2'])
|
|
assert hg.idtype == idtype
|
|
assert hg.device == g.device
|
|
assert set(hg.canonical_etypes) == set(
|
|
[('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')])
|
|
assert hg.number_of_nodes('l0') == 2
|
|
assert hg.number_of_nodes('l1') == 1
|
|
assert hg.number_of_nodes('l2') == 1
|
|
assert hg.number_of_edges('e0') == 2
|
|
assert hg.number_of_edges('e1') == 1
|
|
assert hg.number_of_edges('e2') == 1
|
|
assert F.array_equal(hg.ndata[dgl.NID]['l0'], F.tensor([0, 1], F.int64))
|
|
assert F.array_equal(hg.ndata[dgl.NID]['l1'], F.tensor([2], F.int64))
|
|
assert F.array_equal(hg.ndata[dgl.NID]['l2'], F.tensor([3], F.int64))
|
|
assert F.array_equal(hg.edata[dgl.EID][('l0', 'e0', 'l1')], F.tensor([0, 1], F.int64))
|
|
assert F.array_equal(hg.edata[dgl.EID][('l0', 'e2', 'l2')], F.tensor([3], F.int64))
|
|
assert F.array_equal(hg.edata[dgl.EID][('l1', 'e1', 'l2')], F.tensor([2], F.int64))
|
|
|
|
# hetero_from_homo test case 3
|
|
mg = nx.MultiDiGraph([
|
|
('user', 'movie', 'watches'),
|
|
('user', 'TV', 'watches')])
|
|
g = dgl.graph(((0, 0), (1, 2)), idtype=idtype, device=F.ctx())
|
|
g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
|
|
g.edata[dgl.ETYPE] = F.tensor([0, 0])
|
|
for _mg in [None, mg]:
|
|
hg = dgl.to_heterogeneous(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg)
|
|
assert hg.idtype == g.idtype
|
|
assert hg.device == g.device
|
|
assert set(hg.canonical_etypes) == set(
|
|
[('user', 'watches', 'movie'), ('user', 'watches', 'TV')])
|
|
assert hg.number_of_nodes('user') == 1
|
|
assert hg.number_of_nodes('TV') == 1
|
|
assert hg.number_of_nodes('movie') == 1
|
|
assert hg.number_of_edges(('user', 'watches', 'TV')) == 1
|
|
assert hg.number_of_edges(('user', 'watches', 'movie')) == 1
|
|
assert len(hg.etypes) == 2
|
|
|
|
# hetero_to_homo test case 2
|
|
hg = dgl.heterograph({
|
|
('_U', '_E', '_V'): ([0, 1], [0, 1])
|
|
}, {'_U': 2, '_V': 3}, idtype=idtype, device=F.ctx())
|
|
g = dgl.to_homogeneous(hg)
|
|
assert hg.idtype == g.idtype
|
|
assert hg.device == g.device
|
|
assert g.number_of_nodes() == 5
|
|
|
|
# hetero_to_subgraph_to_homo
|
|
hg = dgl.heterograph({
|
|
('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1]),
|
|
('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2])
|
|
}, idtype=idtype, device=F.ctx())
|
|
hg.nodes['user'].data['h'] = F.copy_to(
|
|
F.tensor([[1, 0], [0, 1], [1, 1]], dtype=idtype), ctx=F.ctx())
|
|
sg = dgl.node_subgraph(hg, {'user': [1, 2]})
|
|
assert len(sg.ntypes) == 2
|
|
assert len(sg.etypes) == 2
|
|
assert sg.num_nodes('user') == 2
|
|
assert sg.num_nodes('game') == 0
|
|
g = dgl.to_homogeneous(sg, ndata=['h'])
|
|
assert 'h' in g.ndata.keys()
|
|
assert g.num_nodes() == 2
|
|
|
|
@unittest.skipIf(F._default_context_str == 'gpu', reason="Test on cpu is enough")
|
|
@parametrize_idtype
|
|
def test_to_homo_zero_nodes(idtype):
|
|
# Fix gihub issue #2870
|
|
g = dgl.heterograph({
|
|
('A', 'AB', 'B'): (np.random.randint(0, 200, (1000,)), np.random.randint(0, 200, (1000,))),
|
|
('B', 'BA', 'A'): (np.random.randint(0, 200, (1000,)), np.random.randint(0, 200, (1000,))),
|
|
}, num_nodes_dict={'A': 200, 'B': 200, 'C': 0}, idtype=idtype)
|
|
g.nodes['A'].data['x'] = F.randn((200, 3))
|
|
g.nodes['B'].data['x'] = F.randn((200, 3))
|
|
gg = dgl.to_homogeneous(g, ['x'])
|
|
assert 'x' in gg.ndata
|
|
|
|
@parametrize_idtype
|
|
def test_to_homo2(idtype):
|
|
# test the result homogeneous graph has nodes and edges sorted by their types
|
|
hg = create_test_heterograph(idtype)
|
|
g = dgl.to_homogeneous(hg)
|
|
ntypes = F.asnumpy(g.ndata[dgl.NTYPE])
|
|
etypes = F.asnumpy(g.edata[dgl.ETYPE])
|
|
p = 0
|
|
for tid, ntype in enumerate(hg.ntypes):
|
|
num_nodes = hg.num_nodes(ntype)
|
|
for i in range(p, p + num_nodes):
|
|
assert ntypes[i] == tid
|
|
p += num_nodes
|
|
p = 0
|
|
for tid, etype in enumerate(hg.canonical_etypes):
|
|
num_edges = hg.num_edges(etype)
|
|
for i in range(p, p + num_edges):
|
|
assert etypes[i] == tid
|
|
p += num_edges
|
|
# test store_type=False
|
|
g = dgl.to_homogeneous(hg, store_type=False)
|
|
assert dgl.NTYPE not in g.ndata
|
|
assert dgl.ETYPE not in g.edata
|
|
# test return_count=True
|
|
g, ntype_count, etype_count = dgl.to_homogeneous(hg, return_count=True)
|
|
for i, count in enumerate(ntype_count):
|
|
assert count == hg.num_nodes(hg.ntypes[i])
|
|
for i, count in enumerate(etype_count):
|
|
assert count == hg.num_edges(hg.canonical_etypes[i])
|
|
|
|
@parametrize_idtype
|
|
def test_invertible_conversion(idtype):
|
|
# Test whether to_homogeneous and to_heterogeneous are invertible
|
|
hg = create_test_heterograph(idtype)
|
|
g = dgl.to_homogeneous(hg)
|
|
hg2 = dgl.to_heterogeneous(g, hg.ntypes, hg.etypes)
|
|
assert_is_identical_hetero(hg, hg2, True)
|
|
|
|
@parametrize_idtype
|
|
def test_metagraph_reachable(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
x = F.randn((3, 5))
|
|
g.nodes['user'].data['h'] = x
|
|
|
|
new_g = dgl.metapath_reachable_graph(g, ['follows', 'plays'])
|
|
assert new_g.idtype == idtype
|
|
assert new_g.ntypes == ['game', 'user']
|
|
assert new_g.number_of_edges() == 3
|
|
assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all()
|
|
|
|
new_g = dgl.metapath_reachable_graph(g, ['follows'])
|
|
assert new_g.idtype == idtype
|
|
assert new_g.ntypes == ['user']
|
|
assert new_g.number_of_edges() == 2
|
|
assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all()
|
|
|
|
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support bool tensor")
|
|
@parametrize_idtype
|
|
def test_subgraph_mask(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
g_graph = g['follows']
|
|
g_bipartite = g['plays']
|
|
|
|
x = F.randn((3, 5))
|
|
y = F.randn((2, 4))
|
|
g.nodes['user'].data['h'] = x
|
|
g.edges['follows'].data['h'] = y
|
|
|
|
def _check_subgraph(g, sg):
|
|
assert sg.idtype == g.idtype
|
|
assert sg.device == g.device
|
|
assert sg.ntypes == g.ntypes
|
|
assert sg.etypes == g.etypes
|
|
assert sg.canonical_etypes == g.canonical_etypes
|
|
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
|
|
F.tensor([1, 2], idtype))
|
|
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
|
|
F.tensor([0], idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
|
|
F.tensor([1], idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
|
|
F.tensor([1], idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
|
|
F.tensor([1], idtype))
|
|
assert sg.number_of_nodes('developer') == 0
|
|
assert sg.number_of_edges('develops') == 0
|
|
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
|
|
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'][1:2])
|
|
|
|
sg1 = g.subgraph({'user': F.tensor([False, True, True], dtype=F.bool),
|
|
'game': F.tensor([True, False, False, False], dtype=F.bool)})
|
|
_check_subgraph(g, sg1)
|
|
if F._default_context_str != 'gpu':
|
|
# TODO(minjie): enable this later
|
|
sg2 = g.edge_subgraph({'follows': F.tensor([False, True], dtype=F.bool),
|
|
'plays': F.tensor([False, True, False, False], dtype=F.bool),
|
|
'wishes': F.tensor([False, True], dtype=F.bool)})
|
|
_check_subgraph(g, sg2)
|
|
|
|
@parametrize_idtype
|
|
def test_subgraph(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
g_graph = g['follows']
|
|
g_bipartite = g['plays']
|
|
|
|
x = F.randn((3, 5))
|
|
y = F.randn((2, 4))
|
|
g.nodes['user'].data['h'] = x
|
|
g.edges['follows'].data['h'] = y
|
|
|
|
def _check_subgraph(g, sg):
|
|
assert sg.idtype == g.idtype
|
|
assert sg.device == g.device
|
|
assert sg.ntypes == g.ntypes
|
|
assert sg.etypes == g.etypes
|
|
assert sg.canonical_etypes == g.canonical_etypes
|
|
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
|
|
F.tensor([1, 2], g.idtype))
|
|
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
|
|
F.tensor([0], g.idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
|
|
F.tensor([1], g.idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
|
|
F.tensor([1], g.idtype))
|
|
assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
|
|
F.tensor([1], g.idtype))
|
|
assert sg.number_of_nodes('developer') == 0
|
|
assert sg.number_of_edges('develops') == 0
|
|
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
|
|
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'][1:2])
|
|
|
|
sg1 = g.subgraph({'user': [1, 2], 'game': [0]})
|
|
_check_subgraph(g, sg1)
|
|
if F._default_context_str != 'gpu':
|
|
# TODO(minjie): enable this later
|
|
sg2 = g.edge_subgraph({'follows': [1], 'plays': [1], 'wishes': [1]})
|
|
_check_subgraph(g, sg2)
|
|
|
|
# backend tensor input
|
|
sg1 = g.subgraph({'user': F.tensor([1, 2], dtype=idtype),
|
|
'game': F.tensor([0], dtype=idtype)})
|
|
_check_subgraph(g, sg1)
|
|
if F._default_context_str != 'gpu':
|
|
# TODO(minjie): enable this later
|
|
sg2 = g.edge_subgraph({'follows': F.tensor([1], dtype=idtype),
|
|
'plays': F.tensor([1], dtype=idtype),
|
|
'wishes': F.tensor([1], dtype=idtype)})
|
|
_check_subgraph(g, sg2)
|
|
|
|
# numpy input
|
|
sg1 = g.subgraph({'user': np.array([1, 2]),
|
|
'game': np.array([0])})
|
|
_check_subgraph(g, sg1)
|
|
if F._default_context_str != 'gpu':
|
|
# TODO(minjie): enable this later
|
|
sg2 = g.edge_subgraph({'follows': np.array([1]),
|
|
'plays': np.array([1]),
|
|
'wishes': np.array([1])})
|
|
_check_subgraph(g, sg2)
|
|
|
|
def _check_subgraph_single_ntype(g, sg, preserve_nodes=False):
|
|
assert sg.idtype == g.idtype
|
|
assert sg.device == g.device
|
|
assert sg.ntypes == g.ntypes
|
|
assert sg.etypes == g.etypes
|
|
assert sg.canonical_etypes == g.canonical_etypes
|
|
|
|
if not preserve_nodes:
|
|
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
|
|
F.tensor([1, 2], g.idtype))
|
|
else:
|
|
for ntype in sg.ntypes:
|
|
assert g.number_of_nodes(ntype) == sg.number_of_nodes(ntype)
|
|
|
|
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
|
|
F.tensor([1], g.idtype))
|
|
|
|
if not preserve_nodes:
|
|
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
|
|
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'][1:2])
|
|
|
|
def _check_subgraph_single_etype(g, sg, preserve_nodes=False):
|
|
assert sg.ntypes == g.ntypes
|
|
assert sg.etypes == g.etypes
|
|
assert sg.canonical_etypes == g.canonical_etypes
|
|
|
|
if not preserve_nodes:
|
|
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
|
|
F.tensor([0, 1], g.idtype))
|
|
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
|
|
F.tensor([0], g.idtype))
|
|
else:
|
|
for ntype in sg.ntypes:
|
|
assert g.number_of_nodes(ntype) == sg.number_of_nodes(ntype)
|
|
|
|
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
|
|
F.tensor([0, 1], g.idtype))
|
|
|
|
sg1_graph = g_graph.subgraph([1, 2])
|
|
_check_subgraph_single_ntype(g_graph, sg1_graph)
|
|
if F._default_context_str != 'gpu':
|
|
# TODO(minjie): enable this later
|
|
sg1_graph = g_graph.edge_subgraph([1])
|
|
_check_subgraph_single_ntype(g_graph, sg1_graph)
|
|
sg1_graph = g_graph.edge_subgraph([1], relabel_nodes=False)
|
|
_check_subgraph_single_ntype(g_graph, sg1_graph, True)
|
|
sg2_bipartite = g_bipartite.edge_subgraph([0, 1])
|
|
_check_subgraph_single_etype(g_bipartite, sg2_bipartite)
|
|
sg2_bipartite = g_bipartite.edge_subgraph([0, 1], relabel_nodes=False)
|
|
_check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)
|
|
|
|
def _check_typed_subgraph1(g, sg):
|
|
assert g.idtype == sg.idtype
|
|
assert g.device == sg.device
|
|
assert set(sg.ntypes) == {'user', 'game'}
|
|
assert set(sg.etypes) == {'follows', 'plays', 'wishes'}
|
|
for ntype in sg.ntypes:
|
|
assert sg.number_of_nodes(ntype) == g.number_of_nodes(ntype)
|
|
for etype in sg.etypes:
|
|
src_sg, dst_sg = sg.all_edges(etype=etype, order='eid')
|
|
src_g, dst_g = g.all_edges(etype=etype, order='eid')
|
|
assert F.array_equal(src_sg, src_g)
|
|
assert F.array_equal(dst_sg, dst_g)
|
|
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
|
|
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'])
|
|
g.nodes['user'].data['h'] = F.scatter_row(g.nodes['user'].data['h'], F.tensor([2]), F.randn((1, 5)))
|
|
g.edges['follows'].data['h'] = F.scatter_row(g.edges['follows'].data['h'], F.tensor([1]), F.randn((1, 4)))
|
|
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
|
|
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'])
|
|
|
|
def _check_typed_subgraph2(g, sg):
|
|
assert set(sg.ntypes) == {'developer', 'game'}
|
|
assert set(sg.etypes) == {'develops'}
|
|
for ntype in sg.ntypes:
|
|
assert sg.number_of_nodes(ntype) == g.number_of_nodes(ntype)
|
|
for etype in sg.etypes:
|
|
src_sg, dst_sg = sg.all_edges(etype=etype, order='eid')
|
|
src_g, dst_g = g.all_edges(etype=etype, order='eid')
|
|
assert F.array_equal(src_sg, src_g)
|
|
assert F.array_equal(dst_sg, dst_g)
|
|
|
|
sg3 = g.node_type_subgraph(['user', 'game'])
|
|
_check_typed_subgraph1(g, sg3)
|
|
sg4 = g.edge_type_subgraph(['develops'])
|
|
_check_typed_subgraph2(g, sg4)
|
|
sg5 = g.edge_type_subgraph(['follows', 'plays', 'wishes'])
|
|
_check_typed_subgraph1(g, sg5)
|
|
|
|
@parametrize_idtype
|
|
def test_apply(idtype):
|
|
def node_udf(nodes):
|
|
return {'h': nodes.data['h'] * 2}
|
|
def node_udf2(nodes):
|
|
return {'h': F.sum(nodes.data['h'], dim=1, keepdims=True)}
|
|
def edge_udf(edges):
|
|
return {'h': edges.data['h'] * 2 + edges.src['h']}
|
|
|
|
g = create_test_heterograph(idtype)
|
|
g.nodes['user'].data['h'] = F.ones((3, 5))
|
|
g.apply_nodes(node_udf, ntype='user')
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 2)
|
|
|
|
g['plays'].edata['h'] = F.ones((4, 5))
|
|
g.apply_edges(edge_udf, etype=('user', 'plays', 'game'))
|
|
assert F.array_equal(g['plays'].edata['h'], F.ones((4, 5)) * 4)
|
|
|
|
# test apply on graph with only one type
|
|
g['follows'].apply_nodes(node_udf)
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 4)
|
|
|
|
g['plays'].apply_edges(edge_udf)
|
|
assert F.array_equal(g['plays'].edata['h'], F.ones((4, 5)) * 12)
|
|
|
|
# Test the case that feature size changes
|
|
g.nodes['user'].data['h'] = F.ones((3, 5))
|
|
g.apply_nodes(node_udf2, ntype='user')
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 1)) * 5)
|
|
|
|
# test fail case
|
|
# fail due to multiple types
|
|
with pytest.raises(DGLError):
|
|
g.apply_nodes(node_udf)
|
|
|
|
with pytest.raises(DGLError):
|
|
g.apply_edges(edge_udf)
|
|
|
|
@parametrize_idtype
|
|
def test_level2(idtype):
|
|
#edges = {
|
|
# 'follows': ([0, 1], [1, 2]),
|
|
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
# 'wishes': ([0, 2], [1, 0]),
|
|
# 'develops': ([0, 1], [0, 1]),
|
|
#}
|
|
g = create_test_heterograph(idtype)
|
|
def rfunc(nodes):
|
|
return {'y': F.sum(nodes.mailbox['m'], 1)}
|
|
def rfunc2(nodes):
|
|
return {'y': F.max(nodes.mailbox['m'], 1)}
|
|
def mfunc(edges):
|
|
return {'m': edges.src['h']}
|
|
def afunc(nodes):
|
|
return {'y' : nodes.data['y'] + 1}
|
|
|
|
#############################################################
|
|
# send_and_recv
|
|
#############################################################
|
|
|
|
g.nodes['user'].data['h'] = F.ones((3, 2))
|
|
g.send_and_recv([2, 3], mfunc, rfunc, etype='plays')
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
|
|
|
|
# only one type
|
|
g['plays'].send_and_recv([2, 3], mfunc, rfunc)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
|
|
|
|
# test fail case
|
|
# fail due to multiple types
|
|
with pytest.raises(DGLError):
|
|
g.send_and_recv([2, 3], mfunc, rfunc)
|
|
|
|
g.nodes['game'].data.clear()
|
|
|
|
#############################################################
|
|
# pull
|
|
#############################################################
|
|
|
|
g.nodes['user'].data['h'] = F.ones((3, 2))
|
|
g.pull(1, mfunc, rfunc, etype='plays')
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
|
|
|
|
# only one type
|
|
g['plays'].pull(1, mfunc, rfunc)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
|
|
|
|
# test fail case
|
|
with pytest.raises(DGLError):
|
|
g.pull(1, mfunc, rfunc)
|
|
|
|
g.nodes['game'].data.clear()
|
|
|
|
#############################################################
|
|
# update_all
|
|
#############################################################
|
|
|
|
g.nodes['user'].data['h'] = F.ones((3, 2))
|
|
g.update_all(mfunc, rfunc, etype='plays')
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[2., 2.], [2., 2.]]))
|
|
|
|
# only one type
|
|
g['plays'].update_all(mfunc, rfunc)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y, F.tensor([[2., 2.], [2., 2.]]))
|
|
|
|
# test fail case
|
|
# fail due to multiple types
|
|
with pytest.raises(DGLError):
|
|
g.update_all(mfunc, rfunc)
|
|
|
|
# test multi
|
|
g.multi_update_all(
|
|
{'plays' : (mfunc, rfunc),
|
|
('user', 'wishes', 'game'): (mfunc, rfunc2)},
|
|
'sum')
|
|
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))
|
|
|
|
# test multi
|
|
g.multi_update_all(
|
|
{'plays' : (mfunc, rfunc, afunc),
|
|
('user', 'wishes', 'game'): (mfunc, rfunc2)},
|
|
'sum', afunc)
|
|
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))
|
|
|
|
# test cross reducer
|
|
g.nodes['user'].data['h'] = F.randn((3, 2))
|
|
for cred in ['sum', 'max', 'min', 'mean', 'stack']:
|
|
g.multi_update_all(
|
|
{'plays' : (mfunc, rfunc, afunc),
|
|
'wishes': (mfunc, rfunc2)},
|
|
cred, afunc)
|
|
y = g.nodes['game'].data['y']
|
|
g['plays'].update_all(mfunc, rfunc, afunc)
|
|
y1 = g.nodes['game'].data['y']
|
|
g['wishes'].update_all(mfunc, rfunc2)
|
|
y2 = g.nodes['game'].data['y']
|
|
if cred == 'stack':
|
|
# stack has an internal order by edge type id
|
|
yy = F.stack([y1, y2], 1)
|
|
yy = yy + 1 # final afunc
|
|
assert F.array_equal(y, yy)
|
|
else:
|
|
yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
|
|
yy = yy + 1 # final afunc
|
|
assert F.array_equal(y, yy)
|
|
|
|
# test fail case
|
|
# fail because cannot infer ntype
|
|
with pytest.raises(DGLError):
|
|
g.update_all(
|
|
{'plays' : (mfunc, rfunc),
|
|
'follows': (mfunc, rfunc2)},
|
|
'sum')
|
|
|
|
g.nodes['game'].data.clear()
|
|
|
|
@parametrize_idtype
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
|
|
def test_more_nnz(idtype):
|
|
g = dgl.graph(([0, 0, 0, 0, 0], [1, 1, 1, 1, 1]), idtype=idtype, device=F.ctx())
|
|
g.ndata['x'] = F.copy_to(F.ones((2, 5)), ctx=F.ctx())
|
|
g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))
|
|
y = g.ndata['y']
|
|
ans = np.zeros((2, 5))
|
|
ans[1] = 5
|
|
ans = F.copy_to(F.tensor(ans, dtype=F.dtype(y)), ctx=F.ctx())
|
|
assert F.array_equal(y, ans)
|
|
|
|
@parametrize_idtype
|
|
def test_updates(idtype):
|
|
def msg_func(edges):
|
|
return {'m': edges.src['h']}
|
|
def reduce_func(nodes):
|
|
return {'y': F.sum(nodes.mailbox['m'], 1)}
|
|
def apply_func(nodes):
|
|
return {'y': nodes.data['y'] * 2}
|
|
g = create_test_heterograph(idtype)
|
|
x = F.randn((3, 5))
|
|
g.nodes['user'].data['h'] = x
|
|
|
|
for msg, red, apply in itertools.product(
|
|
[fn.copy_u('h', 'm'), msg_func], [fn.sum('m', 'y'), reduce_func],
|
|
[None, apply_func]):
|
|
multiplier = 1 if apply is None else 2
|
|
|
|
g['user', 'plays', 'game'].update_all(msg, red, apply)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
|
|
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
|
|
del g.nodes['game'].data['y']
|
|
|
|
g['user', 'plays', 'game'].send_and_recv(([0, 1, 2], [0, 1, 1]), msg, red, apply)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y[0], x[0] * multiplier)
|
|
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
|
|
del g.nodes['game'].data['y']
|
|
|
|
# pulls from destination (game) node 0
|
|
g['user', 'plays', 'game'].pull(0, msg, red, apply)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
|
|
del g.nodes['game'].data['y']
|
|
|
|
# pushes from source (user) node 0
|
|
g['user', 'plays', 'game'].push(0, msg, red, apply)
|
|
y = g.nodes['game'].data['y']
|
|
assert F.array_equal(y[0], x[0] * multiplier)
|
|
del g.nodes['game'].data['y']
|
|
|
|
|
|
@parametrize_idtype
|
|
def test_backward(idtype):
|
|
g = create_test_heterograph(idtype)
|
|
x = F.randn((3, 5))
|
|
F.attach_grad(x)
|
|
g.nodes['user'].data['h'] = x
|
|
with F.record_grad():
|
|
g.multi_update_all(
|
|
{'plays' : (fn.copy_u('h', 'm'), fn.sum('m', 'y')),
|
|
'wishes': (fn.copy_u('h', 'm'), fn.sum('m', 'y'))},
|
|
'sum')
|
|
y = g.nodes['game'].data['y']
|
|
F.backward(y, F.ones(y.shape))
|
|
print(F.grad(x))
|
|
assert F.array_equal(F.grad(x), F.tensor([[2., 2., 2., 2., 2.],
|
|
[2., 2., 2., 2., 2.],
|
|
[2., 2., 2., 2., 2.]]))
|
|
|
|
|
|
@parametrize_idtype
|
|
def test_empty_heterograph(idtype):
|
|
def assert_empty(g):
|
|
assert g.number_of_nodes('user') == 0
|
|
assert g.number_of_edges('plays') == 0
|
|
assert g.number_of_nodes('game') == 0
|
|
|
|
# empty src-dst pair
|
|
assert_empty(dgl.heterograph({('user', 'plays', 'game'): ([], [])}))
|
|
|
|
g = dgl.heterograph({('user', 'follows', 'user'): ([], [])}, idtype=idtype, device=F.ctx())
|
|
assert g.idtype == idtype
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 0
|
|
assert g.number_of_edges('follows') == 0
|
|
|
|
# empty relation graph with others
|
|
g = dgl.heterograph({('user', 'plays', 'game'): ([], []), ('developer', 'develops', 'game'):
|
|
([0, 1], [0, 1])}, idtype=idtype, device=F.ctx())
|
|
assert g.idtype == idtype
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 0
|
|
assert g.number_of_edges('plays') == 0
|
|
assert g.number_of_nodes('game') == 2
|
|
assert g.number_of_edges('develops') == 2
|
|
assert g.number_of_nodes('developer') == 2
|
|
|
|
@parametrize_idtype
|
|
def test_types_in_function(idtype):
|
|
def mfunc1(edges):
|
|
assert edges.canonical_etype == ('user', 'follow', 'user')
|
|
return {}
|
|
|
|
def rfunc1(nodes):
|
|
assert nodes.ntype == 'user'
|
|
return {}
|
|
|
|
def filter_nodes1(nodes):
|
|
assert nodes.ntype == 'user'
|
|
return F.zeros((3,))
|
|
|
|
def filter_edges1(edges):
|
|
assert edges.canonical_etype == ('user', 'follow', 'user')
|
|
return F.zeros((2,))
|
|
|
|
def mfunc2(edges):
|
|
assert edges.canonical_etype == ('user', 'plays', 'game')
|
|
return {}
|
|
|
|
def rfunc2(nodes):
|
|
assert nodes.ntype == 'game'
|
|
return {}
|
|
|
|
def filter_nodes2(nodes):
|
|
assert nodes.ntype == 'game'
|
|
return F.zeros((3,))
|
|
|
|
def filter_edges2(edges):
|
|
assert edges.canonical_etype == ('user', 'plays', 'game')
|
|
return F.zeros((2,))
|
|
|
|
g = dgl.heterograph({('user', 'follow', 'user'): ((0, 1), (1, 2))},
|
|
idtype=idtype, device=F.ctx())
|
|
g.apply_nodes(rfunc1)
|
|
g.apply_edges(mfunc1)
|
|
g.update_all(mfunc1, rfunc1)
|
|
g.send_and_recv([0, 1], mfunc1, rfunc1)
|
|
g.push([0], mfunc1, rfunc1)
|
|
g.pull([1], mfunc1, rfunc1)
|
|
g.filter_nodes(filter_nodes1)
|
|
g.filter_edges(filter_edges1)
|
|
|
|
g = dgl.heterograph({('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
g.apply_nodes(rfunc2, ntype='game')
|
|
g.apply_edges(mfunc2)
|
|
g.update_all(mfunc2, rfunc2)
|
|
g.send_and_recv([0, 1], mfunc2, rfunc2)
|
|
g.push([0], mfunc2, rfunc2)
|
|
g.pull([1], mfunc2, rfunc2)
|
|
g.filter_nodes(filter_nodes2, ntype='game')
|
|
g.filter_edges(filter_edges2)
|
|
|
|
@parametrize_idtype
|
|
def test_stack_reduce(idtype):
|
|
#edges = {
|
|
# 'follows': ([0, 1], [1, 2]),
|
|
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
|
|
# 'wishes': ([0, 2], [1, 0]),
|
|
# 'develops': ([0, 1], [0, 1]),
|
|
#}
|
|
g = create_test_heterograph(idtype)
|
|
g.nodes['user'].data['h'] = F.randn((3, 200))
|
|
def rfunc(nodes):
|
|
return {'y': F.sum(nodes.mailbox['m'], 1)}
|
|
def rfunc2(nodes):
|
|
return {'y': F.max(nodes.mailbox['m'], 1)}
|
|
def mfunc(edges):
|
|
return {'m': edges.src['h']}
|
|
g.multi_update_all(
|
|
{'plays' : (mfunc, rfunc),
|
|
'wishes': (mfunc, rfunc2)},
|
|
'stack')
|
|
assert g.nodes['game'].data['y'].shape == (g.number_of_nodes('game'), 2, 200)
|
|
# only one type-wise update_all, stack still adds one dimension
|
|
g.multi_update_all(
|
|
{'plays' : (mfunc, rfunc)},
|
|
'stack')
|
|
assert g.nodes['game'].data['y'].shape == (g.number_of_nodes('game'), 1, 200)
|
|
|
|
@parametrize_idtype
|
|
def test_isolated_ntype(idtype):
|
|
g = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 1, 2], [1, 2, 3])},
|
|
num_nodes_dict={'A': 3, 'B': 4, 'C': 4},
|
|
idtype=idtype, device=F.ctx())
|
|
assert g.number_of_nodes('A') == 3
|
|
assert g.number_of_nodes('B') == 4
|
|
assert g.number_of_nodes('C') == 4
|
|
|
|
g = dgl.heterograph({
|
|
('A', 'AC', 'C'): ([0, 1, 2], [1, 2, 3])},
|
|
num_nodes_dict={'A': 3, 'B': 4, 'C': 4},
|
|
idtype=idtype, device=F.ctx())
|
|
assert g.number_of_nodes('A') == 3
|
|
assert g.number_of_nodes('B') == 4
|
|
assert g.number_of_nodes('C') == 4
|
|
|
|
G = dgl.graph(([0, 1, 2], [4, 5, 6]), num_nodes=11, idtype=idtype, device=F.ctx())
|
|
G.ndata[dgl.NTYPE] = F.tensor([0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=F.int64)
|
|
G.edata[dgl.ETYPE] = F.tensor([0, 0, 0], dtype=F.int64)
|
|
g = dgl.to_heterogeneous(G, ['A', 'B', 'C'], ['AB'])
|
|
assert g.number_of_nodes('A') == 3
|
|
assert g.number_of_nodes('B') == 4
|
|
assert g.number_of_nodes('C') == 4
|
|
|
|
|
|
@parametrize_idtype
|
|
def test_ismultigraph(idtype):
|
|
g1 = dgl.heterograph({('A', 'AB', 'B'): ([0, 0, 1, 2], [1, 2, 5, 5])},
|
|
{'A': 6, 'B': 6}, idtype=idtype, device=F.ctx())
|
|
assert g1.is_multigraph == False
|
|
g2 = dgl.heterograph({('A', 'AC', 'C'): ([0, 0, 0, 1], [1, 1, 2, 5])},
|
|
{'A': 6, 'C': 6}, idtype=idtype, device=F.ctx())
|
|
assert g2.is_multigraph == True
|
|
g3 = dgl.graph(((0, 1), (1, 2)), num_nodes=6, idtype=idtype, device=F.ctx())
|
|
assert g3.is_multigraph == False
|
|
g4 = dgl.graph(([0, 0, 1], [1, 1, 2]), num_nodes=6, idtype=idtype, device=F.ctx())
|
|
assert g4.is_multigraph == True
|
|
g = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1, 2], [1, 2, 5, 5]),
|
|
('A', 'AA', 'A'): ([0, 1], [1, 2])},
|
|
{'A': 6, 'B': 6}, idtype=idtype, device=F.ctx())
|
|
assert g.is_multigraph == False
|
|
g = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1, 2], [1, 2, 5, 5]),
|
|
('A', 'AC', 'C'): ([0, 0, 0, 1], [1, 1, 2, 5])},
|
|
{'A': 6, 'B': 6, 'C': 6}, idtype=idtype, device=F.ctx())
|
|
assert g.is_multigraph == True
|
|
g = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1, 2], [1, 2, 5, 5]),
|
|
('A', 'AA', 'A'): ([0, 0, 1], [1, 1, 2])},
|
|
{'A': 6, 'B': 6}, idtype=idtype, device=F.ctx())
|
|
assert g.is_multigraph == True
|
|
g = dgl.heterograph({
|
|
('A', 'AC', 'C'): ([0, 0, 0, 1], [1, 1, 2, 5]),
|
|
('A', 'AA', 'A'): ([0, 1], [1, 2])},
|
|
{'A': 6, 'C': 6}, idtype=idtype, device=F.ctx())
|
|
assert g.is_multigraph == True
|
|
|
|
@parametrize_idtype
|
|
def test_bipartite(idtype):
|
|
g1 = dgl.heterograph({('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5])},
|
|
idtype=idtype, device=F.ctx())
|
|
assert g1.is_unibipartite
|
|
assert len(g1.ntypes) == 2
|
|
assert g1.etypes == ['AB']
|
|
assert g1.srctypes == ['A']
|
|
assert g1.dsttypes == ['B']
|
|
assert g1.number_of_nodes('A') == 2
|
|
assert g1.number_of_nodes('B') == 6
|
|
assert g1.number_of_src_nodes('A') == 2
|
|
assert g1.number_of_src_nodes() == 2
|
|
assert g1.number_of_dst_nodes('B') == 6
|
|
assert g1.number_of_dst_nodes() == 6
|
|
assert g1.number_of_edges() == 3
|
|
g1.srcdata['h'] = F.randn((2, 5))
|
|
assert F.array_equal(g1.srcnodes['A'].data['h'], g1.srcdata['h'])
|
|
assert F.array_equal(g1.nodes['A'].data['h'], g1.srcdata['h'])
|
|
assert F.array_equal(g1.nodes['SRC/A'].data['h'], g1.srcdata['h'])
|
|
g1.dstdata['h'] = F.randn((6, 3))
|
|
assert F.array_equal(g1.dstnodes['B'].data['h'], g1.dstdata['h'])
|
|
assert F.array_equal(g1.nodes['B'].data['h'], g1.dstdata['h'])
|
|
assert F.array_equal(g1.nodes['DST/B'].data['h'], g1.dstdata['h'])
|
|
|
|
# more complicated bipartite
|
|
g2 = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
|
|
('A', 'AC', 'C'): ([1, 0], [0, 0])
|
|
}, idtype=idtype, device=F.ctx())
|
|
|
|
assert g2.is_unibipartite
|
|
assert g2.srctypes == ['A']
|
|
assert set(g2.dsttypes) == {'B', 'C'}
|
|
assert g2.number_of_nodes('A') == 2
|
|
assert g2.number_of_nodes('B') == 6
|
|
assert g2.number_of_nodes('C') == 1
|
|
assert g2.number_of_src_nodes('A') == 2
|
|
assert g2.number_of_src_nodes() == 2
|
|
assert g2.number_of_dst_nodes('B') == 6
|
|
assert g2.number_of_dst_nodes('C') == 1
|
|
g2.srcdata['h'] = F.randn((2, 5))
|
|
assert F.array_equal(g2.srcnodes['A'].data['h'], g2.srcdata['h'])
|
|
assert F.array_equal(g2.nodes['A'].data['h'], g2.srcdata['h'])
|
|
assert F.array_equal(g2.nodes['SRC/A'].data['h'], g2.srcdata['h'])
|
|
|
|
g3 = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
|
|
('A', 'AC', 'C'): ([1, 0], [0, 0]),
|
|
('A', 'AA', 'A'): ([0, 1], [0, 1])
|
|
}, idtype=idtype, device=F.ctx())
|
|
assert not g3.is_unibipartite
|
|
|
|
g4 = dgl.heterograph({
|
|
('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
|
|
('C', 'CA', 'A'): ([1, 0], [0, 0])
|
|
}, idtype=idtype, device=F.ctx())
|
|
|
|
assert not g4.is_unibipartite
|
|
|
|
@parametrize_idtype
|
|
def test_dtype_cast(idtype):
|
|
g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
|
|
assert g.idtype == idtype
|
|
g.ndata["feat"] = F.tensor([3, 4, 5])
|
|
g.edata["h"] = F.tensor([3, 4, 5, 6])
|
|
if idtype == "int32":
|
|
g_cast = g.long()
|
|
assert g_cast.idtype == F.int64
|
|
else:
|
|
g_cast = g.int()
|
|
assert g_cast.idtype == F.int32
|
|
test_utils.check_graph_equal(g, g_cast, check_idtype=False)
|
|
|
|
def test_float_cast():
|
|
for t in [F.float16, F.float32, F.float64]:
|
|
idtype = F.int32
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): (F.tensor([0, 1, 1, 2, 2, 3], dtype=idtype),
|
|
F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype)),
|
|
('user', 'plays', 'game'): (F.tensor([0, 1, 1], dtype=idtype),
|
|
F.tensor([0, 0, 1], dtype=idtype))},
|
|
idtype=idtype, device=F.ctx())
|
|
uvalues = [1, 2, 3, 4]
|
|
gvalues = [5, 6]
|
|
fvalues = [7, 8, 9, 10, 11, 12]
|
|
pvalues = [13, 14, 15]
|
|
dataNamesTypes = [
|
|
('a',F.float16),
|
|
('b',F.float32),
|
|
('c',F.float64),
|
|
('d',F.int32),
|
|
('e',F.int64)]
|
|
for name,type in dataNamesTypes:
|
|
g.nodes['user'].data[name] = F.copy_to(F.tensor(uvalues, dtype=type), ctx=F.ctx())
|
|
for name,type in dataNamesTypes:
|
|
g.nodes['game'].data[name] = F.copy_to(F.tensor(gvalues, dtype=type), ctx=F.ctx())
|
|
for name,type in dataNamesTypes:
|
|
g.edges['follows'].data[name] = F.copy_to(F.tensor(fvalues, dtype=type), ctx=F.ctx())
|
|
for name,type in dataNamesTypes:
|
|
g.edges['plays'].data[name] = F.copy_to(F.tensor(pvalues, dtype=type), ctx=F.ctx())
|
|
|
|
if t == F.float16:
|
|
g = dgl.transforms.functional.to_half(g)
|
|
if t == F.float32:
|
|
g = dgl.transforms.functional.to_float(g)
|
|
if t == F.float64:
|
|
g = dgl.transforms.functional.to_double(g)
|
|
|
|
for name,origType in dataNamesTypes:
|
|
# integer tensors shouldn't be converted
|
|
reqType = t if (origType in [F.float16,F.float32,F.float64]) else origType
|
|
|
|
values = g.nodes['user'].data[name]
|
|
assert values.dtype == reqType
|
|
assert len(values) == len(uvalues)
|
|
assert F.allclose(values, F.tensor(uvalues), 0, 0)
|
|
|
|
values = g.nodes['game'].data[name]
|
|
assert values.dtype == reqType
|
|
assert len(values) == len(gvalues)
|
|
assert F.allclose(values, F.tensor(gvalues), 0, 0)
|
|
|
|
values = g.edges['follows'].data[name]
|
|
assert values.dtype == reqType
|
|
assert len(values) == len(fvalues)
|
|
assert F.allclose(values, F.tensor(fvalues), 0, 0)
|
|
|
|
values = g.edges['plays'].data[name]
|
|
assert values.dtype == reqType
|
|
assert len(values) == len(pvalues)
|
|
assert F.allclose(values, F.tensor(pvalues), 0, 0)
|
|
|
|
@parametrize_idtype
|
|
def test_format(idtype):
|
|
# single relation
|
|
g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
|
|
assert g.formats()['created'] == ['coo']
|
|
g1 = g.formats(['coo', 'csr', 'csc'])
|
|
assert len(g1.formats()['created']) + len(g1.formats()['not created']) == 3
|
|
g1.create_formats_()
|
|
assert len(g1.formats()['created']) == 3
|
|
assert g.formats()['created'] == ['coo']
|
|
|
|
# multiple relation
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1], [1, 2]),
|
|
('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1]),
|
|
('developer', 'develops', 'game'): ([0, 1], [0, 1])
|
|
}, idtype=idtype, device=F.ctx())
|
|
user_feat = F.randn((g['follows'].number_of_src_nodes(), 5))
|
|
g['follows'].srcdata['h'] = user_feat
|
|
g1 = g.formats('csc')
|
|
# test frame
|
|
assert F.array_equal(g1['follows'].srcdata['h'], user_feat)
|
|
# test each relation graph
|
|
assert g1.formats()['created'] == ['csc']
|
|
assert len(g1.formats()['not created']) == 0
|
|
|
|
# in_degrees
|
|
g = dgl.rand_graph(100, 2340).to(F.ctx())
|
|
ind_arr = []
|
|
for vid in range(0, 100):
|
|
ind_arr.append(g.in_degrees(vid))
|
|
in_degrees = g.in_degrees()
|
|
g = g.formats('coo')
|
|
for vid in range(0, 100):
|
|
assert g.in_degrees(vid) == ind_arr[vid]
|
|
assert F.array_equal(in_degrees, g.in_degrees())
|
|
|
|
@parametrize_idtype
|
|
def test_edges_order(idtype):
|
|
# (0, 2), (1, 2), (0, 1), (0, 1), (2, 1)
|
|
g = dgl.graph((
|
|
np.array([0, 1, 0, 0, 2]),
|
|
np.array([2, 2, 1, 1, 1])
|
|
), idtype=idtype, device=F.ctx())
|
|
|
|
print(g.formats())
|
|
src, dst = g.all_edges(order='srcdst')
|
|
assert F.array_equal(src, F.tensor([0, 0, 0, 1, 2], dtype=idtype))
|
|
assert F.array_equal(dst, F.tensor([1, 1, 2, 2, 1], dtype=idtype))
|
|
|
|
@parametrize_idtype
|
|
def test_reverse(idtype):
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
|
|
}, idtype=idtype, device=F.ctx())
|
|
gidx = g._graph
|
|
r_gidx = gidx.reverse()
|
|
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
# force to start with 'csr'
|
|
gidx = gidx.formats('csr')
|
|
gidx = gidx.formats(['coo', 'csr', 'csc'])
|
|
r_gidx = gidx.reverse()
|
|
assert 'csr' in gidx.formats()['created']
|
|
assert 'csc' in r_gidx.formats()['created']
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
# force to start with 'csc'
|
|
gidx = gidx.formats('csc')
|
|
gidx = gidx.formats(['coo', 'csr', 'csc'])
|
|
r_gidx = gidx.reverse()
|
|
assert 'csc' in gidx.formats()['created']
|
|
assert 'csr' in r_gidx.formats()['created']
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
g = dgl.heterograph({
|
|
('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
|
|
('user', 'plays', 'game'): ([0, 0, 2, 3, 3, 4, 1], [1, 0, 1, 0, 1, 0, 0]),
|
|
('developer', 'develops', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1]),
|
|
}, idtype=idtype, device=F.ctx())
|
|
gidx = g._graph
|
|
r_gidx = gidx.reverse()
|
|
|
|
# metagraph
|
|
mg = gidx.metagraph
|
|
r_mg = r_gidx.metagraph
|
|
for etype in range(3):
|
|
assert mg.find_edge(etype) == r_mg.find_edge(etype)[::-1]
|
|
|
|
# three node types and three edge types
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
|
|
assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
|
|
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(1)
|
|
rg_s, rg_d, _ = r_gidx.edges(1)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(2)
|
|
rg_s, rg_d, _ = r_gidx.edges(2)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
# force to start with 'csr'
|
|
gidx = gidx.formats('csr')
|
|
gidx = gidx.formats(['coo', 'csr', 'csc'])
|
|
r_gidx = gidx.reverse()
|
|
# three node types and three edge types
|
|
assert 'csr' in gidx.formats()['created']
|
|
assert 'csc' in r_gidx.formats()['created']
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
|
|
assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
|
|
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(1)
|
|
rg_s, rg_d, _ = r_gidx.edges(1)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(2)
|
|
rg_s, rg_d, _ = r_gidx.edges(2)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
# force to start with 'csc'
|
|
gidx = gidx.formats('csc')
|
|
gidx = gidx.formats(['coo', 'csr', 'csc'])
|
|
r_gidx = gidx.reverse()
|
|
# three node types and three edge types
|
|
assert 'csc' in gidx.formats()['created']
|
|
assert 'csr' in r_gidx.formats()['created']
|
|
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
|
|
assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
|
|
assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
|
|
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
|
|
assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
|
|
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
|
|
g_s, g_d, _ = gidx.edges(0)
|
|
rg_s, rg_d, _ = r_gidx.edges(0)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(1)
|
|
rg_s, rg_d, _ = r_gidx.edges(1)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
g_s, g_d, _ = gidx.edges(2)
|
|
rg_s, rg_d, _ = r_gidx.edges(2)
|
|
assert F.array_equal(g_s, rg_d)
|
|
assert F.array_equal(g_d, rg_s)
|
|
|
|
@parametrize_idtype
|
|
def test_clone(idtype):
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
|
|
|
|
new_g = g.clone()
|
|
assert g.number_of_nodes() == new_g.number_of_nodes()
|
|
assert g.number_of_edges() == new_g.number_of_edges()
|
|
assert g.device == new_g.device
|
|
assert g.idtype == new_g.idtype
|
|
assert F.array_equal(g.ndata['h'], new_g.ndata['h'])
|
|
assert F.array_equal(g.edata['h'], new_g.edata['h'])
|
|
# data change
|
|
new_g.ndata['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
|
|
assert (F.array_equal(g.ndata['h'], new_g.ndata['h']) == False)
|
|
g.edata['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
|
|
assert (F.array_equal(g.edata['h'], new_g.edata['h']) == False)
|
|
# graph structure change
|
|
g.add_nodes(1)
|
|
assert g.number_of_nodes() != new_g.number_of_nodes()
|
|
new_g.add_edges(1, 1)
|
|
assert g.number_of_edges() != new_g.number_of_edges()
|
|
|
|
# zero data graph
|
|
g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
|
|
new_g = g.clone()
|
|
assert g.number_of_nodes() == new_g.number_of_nodes()
|
|
assert g.number_of_edges() == new_g.number_of_edges()
|
|
|
|
# heterograph
|
|
g = create_test_heterograph3(idtype)
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
|
|
new_g = g.clone()
|
|
assert g.number_of_nodes('user') == new_g.number_of_nodes('user')
|
|
assert g.number_of_nodes('game') == new_g.number_of_nodes('game')
|
|
assert g.number_of_nodes('developer') == new_g.number_of_nodes('developer')
|
|
assert g.number_of_edges('plays') == new_g.number_of_edges('plays')
|
|
assert g.number_of_edges('develops') == new_g.number_of_edges('develops')
|
|
assert F.array_equal(g.nodes['user'].data['h'], new_g.nodes['user'].data['h'])
|
|
assert F.array_equal(g.nodes['game'].data['h'], new_g.nodes['game'].data['h'])
|
|
assert F.array_equal(g.edges['plays'].data['h'], new_g.edges['plays'].data['h'])
|
|
assert g.device == new_g.device
|
|
assert g.idtype == new_g.idtype
|
|
u, v = g.edges(form='uv', order='eid', etype='plays')
|
|
nu, nv = new_g.edges(form='uv', order='eid', etype='plays')
|
|
assert F.array_equal(u, nu)
|
|
assert F.array_equal(v, nv)
|
|
# graph structure change
|
|
u = F.tensor([0, 4], dtype=idtype)
|
|
v = F.tensor([2, 6], dtype=idtype)
|
|
g.add_edges(u, v, etype='plays')
|
|
u, v = g.edges(form='uv', order='eid', etype='plays')
|
|
assert u.shape[0] != nu.shape[0]
|
|
assert v.shape[0] != nv.shape[0]
|
|
assert g.nodes['user'].data['h'].shape[0] != new_g.nodes['user'].data['h'].shape[0]
|
|
assert g.nodes['game'].data['h'].shape[0] != new_g.nodes['game'].data['h'].shape[0]
|
|
assert g.edges['plays'].data['h'].shape[0] != new_g.edges['plays'].data['h'].shape[0]
|
|
|
|
|
|
@parametrize_idtype
|
|
def test_add_edges(idtype):
|
|
# homogeneous graph
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
u = 0
|
|
v = 1
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes() == 3
|
|
assert g.number_of_edges() == 3
|
|
u = [0]
|
|
v = [1]
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes() == 3
|
|
assert g.number_of_edges() == 4
|
|
u = F.tensor(u, dtype=idtype)
|
|
v = F.tensor(v, dtype=idtype)
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes() == 3
|
|
assert g.number_of_edges() == 5
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
|
|
|
|
# node id larger than current max node id
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
u = F.tensor([0, 1], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
g.add_edges(u, v)
|
|
assert g.number_of_nodes() == 4
|
|
assert g.number_of_edges() == 4
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
|
|
|
|
# has data
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
|
|
u = F.tensor([0, 1], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
|
|
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
|
|
g.add_edges(u, v, e_feat)
|
|
assert g.number_of_nodes() == 4
|
|
assert g.number_of_edges() == 4
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
|
|
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
|
|
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
|
|
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
|
|
|
|
# zero data graph
|
|
g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
|
|
u = F.tensor([0, 1], dtype=idtype)
|
|
v = F.tensor([2, 2], dtype=idtype)
|
|
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
|
|
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
|
|
g.add_edges(u, v, e_feat)
|
|
assert g.number_of_nodes() == 3
|
|
assert g.number_of_edges() == 2
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
|
|
assert F.array_equal(g.edata['h'], F.tensor([2, 2], dtype=idtype))
|
|
assert F.array_equal(g.edata['hh'], F.tensor([2, 2], dtype=idtype))
|
|
|
|
# bipartite graph
|
|
g = dgl.heterograph({('user', 'plays', 'game'): ([0, 1], [1, 2])},
|
|
idtype=idtype, device=F.ctx())
|
|
u = 0
|
|
v = 1
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 2
|
|
assert g.number_of_nodes('game') == 3
|
|
assert g.number_of_edges() == 3
|
|
u = [0]
|
|
v = [1]
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 2
|
|
assert g.number_of_nodes('game') == 3
|
|
assert g.number_of_edges() == 4
|
|
u = F.tensor(u, dtype=idtype)
|
|
v = F.tensor(v, dtype=idtype)
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 2
|
|
assert g.number_of_nodes('game') == 3
|
|
assert g.number_of_edges() == 5
|
|
u, v = g.edges(form='uv')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
|
|
|
|
# node id larger than current max node id
|
|
g = dgl.heterograph({('user', 'plays', 'game'): ([0, 1], [1, 2])},
|
|
idtype=idtype, device=F.ctx())
|
|
u = F.tensor([0, 2], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
g.add_edges(u, v)
|
|
assert g.device == F.ctx()
|
|
assert g.number_of_nodes('user') == 3
|
|
assert g.number_of_nodes('game') == 4
|
|
assert g.number_of_edges() == 4
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
|
|
|
|
# has data
|
|
g = dgl.heterograph({
|
|
('user', 'plays', 'game'): ([0, 1], [1, 2])
|
|
}, idtype=idtype, device=F.ctx())
|
|
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
|
|
u = F.tensor([0, 2], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
|
|
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
|
|
g.add_edges(u, v, e_feat)
|
|
assert g.number_of_nodes('user') == 3
|
|
assert g.number_of_nodes('game') == 4
|
|
assert g.number_of_edges() == 4
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 0], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 0], dtype=idtype))
|
|
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
|
|
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
|
|
|
|
# heterogeneous graph
|
|
g = create_test_heterograph3(idtype)
|
|
u = F.tensor([0, 2], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
g.add_edges(u, v, etype='plays')
|
|
assert g.number_of_nodes('user') == 3
|
|
assert g.number_of_nodes('game') == 4
|
|
assert g.number_of_nodes('developer') == 2
|
|
assert g.number_of_edges('plays') == 6
|
|
assert g.number_of_edges('develops') == 2
|
|
u, v = g.edges(form='uv', order='eid', etype='plays')
|
|
assert F.array_equal(u, F.tensor([0, 1, 1, 2, 0, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0, 0, 1, 1, 2, 3], dtype=idtype))
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 0, 0], dtype=idtype))
|
|
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 1, 1, 1, 0, 0], dtype=idtype))
|
|
|
|
# add with feature
|
|
e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
|
|
u = F.tensor([0, 2], dtype=idtype)
|
|
v = F.tensor([2, 3], dtype=idtype)
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.add_edges(u, v, data=e_feat, etype='develops')
|
|
assert g.number_of_nodes('user') == 3
|
|
assert g.number_of_nodes('game') == 4
|
|
assert g.number_of_nodes('developer') == 3
|
|
assert g.number_of_edges('plays') == 6
|
|
assert g.number_of_edges('develops') == 4
|
|
u, v = g.edges(form='uv', order='eid', etype='develops')
|
|
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0, 1, 2, 3], dtype=idtype))
|
|
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3, 0], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 1, 1], dtype=idtype))
|
|
assert F.array_equal(g.edges['develops'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
|
|
|
|
@parametrize_idtype
|
|
def test_add_nodes(idtype):
|
|
# homogeneous Graphs
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
|
|
g.add_nodes(1)
|
|
assert g.number_of_nodes() == 4
|
|
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
|
|
|
|
# zero node graph
|
|
g = dgl.graph(([], []), num_nodes=3, idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
|
|
g.add_nodes(1, data={'h' : F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())})
|
|
assert g.number_of_nodes() == 4
|
|
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 2], dtype=idtype))
|
|
|
|
# bipartite graph
|
|
g = dgl.heterograph({('user', 'plays', 'game'): ([0, 1], [1, 2])},
|
|
idtype=idtype, device=F.ctx())
|
|
g.add_nodes(2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='user')
|
|
assert g.number_of_nodes('user') == 4
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
|
|
g.add_nodes(2, ntype='game')
|
|
assert g.number_of_nodes('game') == 5
|
|
|
|
# heterogeneous graph
|
|
g = create_test_heterograph3(idtype)
|
|
g.add_nodes(1, ntype='user')
|
|
g.add_nodes(2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='game')
|
|
g.add_nodes(0, ntype='developer')
|
|
assert g.number_of_nodes('user') == 4
|
|
assert g.number_of_nodes('game') == 4
|
|
assert g.number_of_nodes('developer') == 2
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 2], dtype=idtype))
|
|
|
|
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet has error with (0,) shape tensor.")
|
|
@parametrize_idtype
|
|
def test_remove_edges(idtype):
|
|
# homogeneous Graphs
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
e = 0
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2], dtype=idtype))
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
e = [0]
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2], dtype=idtype))
|
|
e = F.tensor([0], dtype=idtype)
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 0
|
|
|
|
# has node data
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
|
|
g.remove_edges(1)
|
|
assert g.number_of_edges() == 1
|
|
assert F.array_equal(g.ndata['h'], F.tensor([1, 2, 3], dtype=idtype))
|
|
|
|
# has edge data
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
|
|
g.remove_edges(0)
|
|
assert g.number_of_edges() == 1
|
|
assert F.array_equal(g.edata['h'], F.tensor([2], dtype=idtype))
|
|
|
|
# invalid eid
|
|
assert_fail = False
|
|
try:
|
|
g.remove_edges(1)
|
|
except:
|
|
assert_fail = True
|
|
assert assert_fail
|
|
|
|
# bipartite graph
|
|
g = dgl.heterograph({
|
|
('user', 'plays', 'game'): ([0, 1], [1, 2])
|
|
}, idtype=idtype, device=F.ctx())
|
|
e = 0
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2], dtype=idtype))
|
|
g = dgl.heterograph(
|
|
{('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
e = [0]
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2], dtype=idtype))
|
|
e = F.tensor([0], dtype=idtype)
|
|
g.remove_edges(e)
|
|
assert g.number_of_edges() == 0
|
|
|
|
# has data
|
|
g = dgl.heterograph(
|
|
{('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
|
|
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
|
|
g.remove_edges(1)
|
|
assert g.number_of_edges() == 1
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2], dtype=idtype))
|
|
assert F.array_equal(g.edata['h'], F.tensor([1], dtype=idtype))
|
|
|
|
# heterogeneous graph
|
|
g = create_test_heterograph3(idtype)
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
|
|
g.remove_edges(1, etype='plays')
|
|
assert g.number_of_edges('plays') == 3
|
|
u, v = g.edges(form='uv', order='eid', etype='plays')
|
|
assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
|
|
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 3, 4], dtype=idtype))
|
|
# remove all edges of 'develops'
|
|
g.remove_edges([0, 1], etype='develops')
|
|
assert g.number_of_edges('develops') == 0
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2], dtype=idtype))
|
|
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
|
|
|
|
@parametrize_idtype
|
|
def test_remove_nodes(idtype):
|
|
# homogeneous Graphs
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
n = 0
|
|
g.remove_nodes(n)
|
|
assert g.number_of_nodes() == 2
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1], dtype=idtype))
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
n = [1]
|
|
g.remove_nodes(n)
|
|
assert g.number_of_nodes() == 2
|
|
assert g.number_of_edges() == 0
|
|
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
|
|
n = F.tensor([2], dtype=idtype)
|
|
g.remove_nodes(n)
|
|
assert g.number_of_nodes() == 2
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1], dtype=idtype))
|
|
|
|
# invalid nid
|
|
assert_fail = False
|
|
try:
|
|
g.remove_nodes(3)
|
|
except:
|
|
assert_fail = True
|
|
assert assert_fail
|
|
|
|
# has node and edge data
|
|
g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
|
|
g.ndata['hv'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
|
|
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
|
|
g.remove_nodes(F.tensor([0], dtype=idtype))
|
|
assert g.number_of_nodes() == 2
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(g.ndata['hv'], F.tensor([2, 3], dtype=idtype))
|
|
assert F.array_equal(g.edata['he'], F.tensor([3], dtype=idtype))
|
|
|
|
# node id larger than current max node id
|
|
g = dgl.heterograph(
|
|
{('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
n = 0
|
|
g.remove_nodes(n, ntype='user')
|
|
assert g.number_of_nodes('user') == 1
|
|
assert g.number_of_nodes('game') == 3
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([2], dtype=idtype))
|
|
g = dgl.heterograph(
|
|
{('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
n = [1]
|
|
g.remove_nodes(n, ntype='user')
|
|
assert g.number_of_nodes('user') == 1
|
|
assert g.number_of_nodes('game') == 3
|
|
assert g.number_of_edges() == 1
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([1], dtype=idtype))
|
|
g = dgl.heterograph(
|
|
{('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
|
|
n = F.tensor([0], dtype=idtype)
|
|
g.remove_nodes(n, ntype='game')
|
|
assert g.number_of_nodes('user') == 2
|
|
assert g.number_of_nodes('game') == 2
|
|
assert g.number_of_edges() == 2
|
|
u, v = g.edges(form='uv', order='eid')
|
|
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0 ,1], dtype=idtype))
|
|
|
|
# heterogeneous graph
|
|
g = create_test_heterograph3(idtype)
|
|
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
|
|
g.remove_nodes(0, ntype='game')
|
|
assert g.number_of_nodes('user') == 3
|
|
assert g.number_of_nodes('game') == 1
|
|
assert g.number_of_nodes('developer') == 2
|
|
assert g.number_of_edges('plays') == 2
|
|
assert g.number_of_edges('develops') == 1
|
|
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
|
|
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2], dtype=idtype))
|
|
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
|
|
u, v = g.edges(form='uv', order='eid', etype='plays')
|
|
assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
|
|
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([3, 4], dtype=idtype))
|
|
u, v = g.edges(form='uv', order='eid', etype='develops')
|
|
assert F.array_equal(u, F.tensor([1], dtype=idtype))
|
|
assert F.array_equal(v, F.tensor([0], dtype=idtype))
|
|
|
|
@parametrize_idtype
|
|
def test_frame(idtype):
|
|
g = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
|
|
g.ndata['h'] = F.copy_to(F.tensor([0, 1, 2, 3], dtype=idtype), ctx=F.ctx())
|
|
g.edata['h'] = F.copy_to(F.tensor([0, 1, 2], dtype=idtype), ctx=F.ctx())
|
|
|
|
# remove nodes
|
|
sg = dgl.remove_nodes(g, [3])
|
|
# check for lazy update
|
|
assert F.array_equal(sg._node_frames[0]._columns['h'].storage, g.ndata['h'])
|
|
assert F.array_equal(sg._edge_frames[0]._columns['h'].storage, g.edata['h'])
|
|
assert sg.ndata['h'].shape[0] == 3
|
|
assert sg.edata['h'].shape[0] == 2
|
|
# update after read
|
|
assert F.array_equal(sg._node_frames[0]._columns['h'].storage, F.tensor([0, 1, 2], dtype=idtype))
|
|
assert F.array_equal(sg._edge_frames[0]._columns['h'].storage, F.tensor([0, 1], dtype=idtype))
|
|
|
|
ng = dgl.add_nodes(sg, 1)
|
|
assert ng.ndata['h'].shape[0] == 4
|
|
assert F.array_equal(ng._node_frames[0]._columns['h'].storage, F.tensor([0, 1, 2, 0], dtype=idtype))
|
|
ng = dgl.add_edges(ng, [3], [1])
|
|
assert ng.edata['h'].shape[0] == 3
|
|
assert F.array_equal(ng._edge_frames[0]._columns['h'].storage, F.tensor([0, 1, 0], dtype=idtype))
|
|
|
|
# multi level lazy update
|
|
sg = dgl.remove_nodes(g, [3])
|
|
assert F.array_equal(sg._node_frames[0]._columns['h'].storage, g.ndata['h'])
|
|
assert F.array_equal(sg._edge_frames[0]._columns['h'].storage, g.edata['h'])
|
|
ssg = dgl.remove_nodes(sg, [1])
|
|
assert F.array_equal(ssg._node_frames[0]._columns['h'].storage, g.ndata['h'])
|
|
assert F.array_equal(ssg._edge_frames[0]._columns['h'].storage, g.edata['h'])
|
|
# ssg is changed
|
|
assert ssg.ndata['h'].shape[0] == 2
|
|
assert ssg.edata['h'].shape[0] == 0
|
|
assert F.array_equal(ssg._node_frames[0]._columns['h'].storage, F.tensor([0, 2], dtype=idtype))
|
|
# sg still in lazy model
|
|
assert F.array_equal(sg._node_frames[0]._columns['h'].storage, g.ndata['h'])
|
|
assert F.array_equal(sg._edge_frames[0]._columns['h'].storage, g.edata['h'])
|
|
|
|
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TensorFlow always create a new tensor")
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="cpu do not have context change problem")
|
|
@parametrize_idtype
|
|
def test_frame_device(idtype):
|
|
g = dgl.graph(([0,1,2], [2,3,1]))
|
|
g.ndata['h'] = F.copy_to(F.tensor([1,1,1,2], dtype=idtype), ctx=F.cpu())
|
|
g.ndata['hh'] = F.copy_to(F.ones((4,3), dtype=idtype), ctx=F.cpu())
|
|
g.edata['h'] = F.copy_to(F.tensor([1,2,3], dtype=idtype), ctx=F.cpu())
|
|
|
|
g = g.to(F.ctx())
|
|
# lazy device copy
|
|
assert F.context(g._node_frames[0]._columns['h'].storage) == F.cpu()
|
|
assert F.context(g._node_frames[0]._columns['hh'].storage) == F.cpu()
|
|
print(g.ndata['h'])
|
|
assert F.context(g._node_frames[0]._columns['h'].storage) == F.ctx()
|
|
assert F.context(g._node_frames[0]._columns['hh'].storage) == F.cpu()
|
|
assert F.context(g._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
|
|
# lazy device copy in subgraph
|
|
sg = dgl.node_subgraph(g, [0,1,2])
|
|
assert F.context(sg._node_frames[0]._columns['h'].storage) == F.ctx()
|
|
assert F.context(sg._node_frames[0]._columns['hh'].storage) == F.cpu()
|
|
assert F.context(sg._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
print(sg.ndata['hh'])
|
|
assert F.context(sg._node_frames[0]._columns['hh'].storage) == F.ctx()
|
|
assert F.context(sg._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
|
|
# back to cpu
|
|
sg = sg.to(F.cpu())
|
|
assert F.context(sg._node_frames[0]._columns['h'].storage) == F.ctx()
|
|
assert F.context(sg._node_frames[0]._columns['hh'].storage) == F.ctx()
|
|
assert F.context(sg._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
print(sg.ndata['h'])
|
|
print(sg.ndata['hh'])
|
|
print(sg.edata['h'])
|
|
assert F.context(sg._node_frames[0]._columns['h'].storage) == F.cpu()
|
|
assert F.context(sg._node_frames[0]._columns['hh'].storage) == F.cpu()
|
|
assert F.context(sg._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
|
|
# set some field
|
|
sg = sg.to(F.ctx())
|
|
assert F.context(sg._node_frames[0]._columns['h'].storage) == F.cpu()
|
|
sg.ndata['h'][0] = 5
|
|
assert F.context(sg._node_frames[0]._columns['h'].storage) == F.ctx()
|
|
assert F.context(sg._node_frames[0]._columns['hh'].storage) == F.cpu()
|
|
assert F.context(sg._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
|
|
# add nodes
|
|
ng = dgl.add_nodes(sg, 3)
|
|
assert F.context(ng._node_frames[0]._columns['h'].storage) == F.ctx()
|
|
assert F.context(ng._node_frames[0]._columns['hh'].storage) == F.ctx()
|
|
assert F.context(ng._edge_frames[0]._columns['h'].storage) == F.cpu()
|
|
|
|
@parametrize_idtype
|
|
def test_create_block(idtype):
|
|
block = dgl.create_block(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes() == 3
|
|
assert block.num_dst_nodes() == 4
|
|
assert block.num_edges() == 3
|
|
|
|
block = dgl.create_block(([], []), idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes() == 0
|
|
assert block.num_dst_nodes() == 0
|
|
assert block.num_edges() == 0
|
|
|
|
block = dgl.create_block(([], []), 3, 4, idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes() == 3
|
|
assert block.num_dst_nodes() == 4
|
|
assert block.num_edges() == 0
|
|
|
|
block = dgl.create_block(([0, 1, 2], [1, 2, 3]), 4, 5, idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes() == 4
|
|
assert block.num_dst_nodes() == 5
|
|
assert block.num_edges() == 3
|
|
|
|
sx = F.randn((4, 5))
|
|
dx = F.randn((5, 6))
|
|
ex = F.randn((3, 4))
|
|
block.srcdata['x'] = sx
|
|
block.dstdata['x'] = dx
|
|
block.edata['x'] = ex
|
|
|
|
g = dgl.block_to_graph(block)
|
|
assert g.num_src_nodes() == 4
|
|
assert g.num_dst_nodes() == 5
|
|
assert g.num_edges() == 3
|
|
assert g.srcdata['x'] is sx
|
|
assert g.dstdata['x'] is dx
|
|
assert g.edata['x'] is ex
|
|
|
|
block = dgl.create_block({
|
|
('A', 'AB', 'B'): ([1, 2, 3], [2, 1, 0]),
|
|
('B', 'BA', 'A'): ([2, 3], [3, 4])},
|
|
idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes('A') == 4
|
|
assert block.num_src_nodes('B') == 4
|
|
assert block.num_dst_nodes('B') == 3
|
|
assert block.num_dst_nodes('A') == 5
|
|
assert block.num_edges('AB') == 3
|
|
assert block.num_edges('BA') == 2
|
|
|
|
block = dgl.create_block({
|
|
('A', 'AB', 'B'): ([], []),
|
|
('B', 'BA', 'A'): ([], [])},
|
|
idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes('A') == 0
|
|
assert block.num_src_nodes('B') == 0
|
|
assert block.num_dst_nodes('B') == 0
|
|
assert block.num_dst_nodes('A') == 0
|
|
assert block.num_edges('AB') == 0
|
|
assert block.num_edges('BA') == 0
|
|
|
|
block = dgl.create_block({
|
|
('A', 'AB', 'B'): ([], []),
|
|
('B', 'BA', 'A'): ([], [])},
|
|
num_src_nodes={'A': 5, 'B': 5},
|
|
num_dst_nodes={'A': 6, 'B': 4},
|
|
idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes('A') == 5
|
|
assert block.num_src_nodes('B') == 5
|
|
assert block.num_dst_nodes('B') == 4
|
|
assert block.num_dst_nodes('A') == 6
|
|
assert block.num_edges('AB') == 0
|
|
assert block.num_edges('BA') == 0
|
|
|
|
block = dgl.create_block({
|
|
('A', 'AB', 'B'): ([1, 2, 3], [2, 1, 0]),
|
|
('B', 'BA', 'A'): ([2, 3], [3, 4])},
|
|
num_src_nodes={'A': 5, 'B': 5},
|
|
num_dst_nodes={'A': 6, 'B': 4},
|
|
idtype=idtype, device=F.ctx())
|
|
assert block.num_src_nodes('A') == 5
|
|
assert block.num_src_nodes('B') == 5
|
|
assert block.num_dst_nodes('B') == 4
|
|
assert block.num_dst_nodes('A') == 6
|
|
assert block.num_edges(('A', 'AB', 'B')) == 3
|
|
assert block.num_edges(('B', 'BA', 'A')) == 2
|
|
|
|
sax = F.randn((5, 3))
|
|
sbx = F.randn((5, 4))
|
|
dax = F.randn((6, 5))
|
|
dbx = F.randn((4, 6))
|
|
eabx = F.randn((3, 7))
|
|
ebax = F.randn((2, 8))
|
|
block.srcnodes['A'].data['x'] = sax
|
|
block.srcnodes['B'].data['x'] = sbx
|
|
block.dstnodes['A'].data['x'] = dax
|
|
block.dstnodes['B'].data['x'] = dbx
|
|
block.edges['AB'].data['x'] = eabx
|
|
block.edges['BA'].data['x'] = ebax
|
|
|
|
hg = dgl.block_to_graph(block)
|
|
assert hg.num_nodes('A_src') == 5
|
|
assert hg.num_nodes('B_src') == 5
|
|
assert hg.num_nodes('A_dst') == 6
|
|
assert hg.num_nodes('B_dst') == 4
|
|
assert hg.num_edges(('A_src', 'AB', 'B_dst')) == 3
|
|
assert hg.num_edges(('B_src', 'BA', 'A_dst')) == 2
|
|
assert hg.nodes['A_src'].data['x'] is sax
|
|
assert hg.nodes['B_src'].data['x'] is sbx
|
|
assert hg.nodes['A_dst'].data['x'] is dax
|
|
assert hg.nodes['B_dst'].data['x'] is dbx
|
|
assert hg.edges['AB'].data['x'] is eabx
|
|
assert hg.edges['BA'].data['x'] is ebax
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize('fmt', ['coo', 'csr', 'csc'])
|
|
def test_adj_sparse(idtype, fmt):
|
|
if fmt == 'coo':
|
|
A = ssp.random(10, 10, 0.2).tocoo()
|
|
A.data = np.arange(20)
|
|
row = F.tensor(A.row, idtype)
|
|
col = F.tensor(A.col, idtype)
|
|
g = dgl.graph((row, col))
|
|
elif fmt == 'csr':
|
|
A = ssp.random(10, 10, 0.2).tocsr()
|
|
A.data = np.arange(20)
|
|
indptr = F.tensor(A.indptr, idtype)
|
|
indices = F.tensor(A.indices, idtype)
|
|
g = dgl.graph(('csr', (indptr, indices, [])))
|
|
with pytest.raises(DGLError):
|
|
g2 = dgl.graph(('csr', (indptr[:-1], indices, [])), num_nodes=10)
|
|
elif fmt == 'csc':
|
|
A = ssp.random(10, 10, 0.2).tocsc()
|
|
A.data = np.arange(20)
|
|
indptr = F.tensor(A.indptr, idtype)
|
|
indices = F.tensor(A.indices, idtype)
|
|
g = dgl.graph(('csc', (indptr, indices, [])))
|
|
with pytest.raises(DGLError):
|
|
g2 = dgl.graph(('csr', (indptr[:-1], indices, [])), num_nodes=10)
|
|
|
|
A_coo = A.tocoo()
|
|
A_csr = A.tocsr()
|
|
A_csc = A.tocsc()
|
|
row, col = g.adj_sparse('coo')
|
|
assert np.array_equal(F.asnumpy(row), A_coo.row)
|
|
assert np.array_equal(F.asnumpy(col), A_coo.col)
|
|
|
|
indptr, indices, eids = g.adj_sparse('csr')
|
|
assert np.array_equal(F.asnumpy(indptr), A_csr.indptr)
|
|
if fmt == 'csr':
|
|
assert len(eids) == 0
|
|
assert np.array_equal(F.asnumpy(indices), A_csr.indices)
|
|
else:
|
|
indices_sorted = F.zeros(len(indices), idtype)
|
|
indices_sorted = F.scatter_row(indices_sorted, eids, indices)
|
|
indices_sorted_np = np.zeros(len(indices), dtype=A_csr.indices.dtype)
|
|
indices_sorted_np[A_csr.data] = A_csr.indices
|
|
assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)
|
|
|
|
indptr, indices, eids = g.adj_sparse('csc')
|
|
assert np.array_equal(F.asnumpy(indptr), A_csc.indptr)
|
|
if fmt == 'csc':
|
|
assert len(eids) == 0
|
|
assert np.array_equal(F.asnumpy(indices), A_csc.indices)
|
|
else:
|
|
indices_sorted = F.zeros(len(indices), idtype)
|
|
indices_sorted = F.scatter_row(indices_sorted, eids, indices)
|
|
indices_sorted_np = np.zeros(len(indices), dtype=A_csc.indices.dtype)
|
|
indices_sorted_np[A_csc.data] = A_csc.indices
|
|
assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)
|
|
|
|
|
|
def _test_forking_pickler_entry(g, q):
|
|
q.put(g.formats())
|
|
|
|
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support spawning")
|
|
def test_forking_pickler():
|
|
ctx = mp.get_context('spawn')
|
|
g = dgl.graph(([0,1,2],[1,2,3]))
|
|
g.create_formats_()
|
|
q = ctx.Queue(1)
|
|
proc = ctx.Process(target=_test_forking_pickler_entry, args=(g, q))
|
|
proc.start()
|
|
fmt = q.get()['created']
|
|
proc.join()
|
|
assert 'coo' in fmt
|
|
assert 'csr' in fmt
|
|
assert 'csc' in fmt
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# test_create()
|
|
# test_query()
|
|
# test_hypersparse()
|
|
# test_adj("int32")
|
|
# test_inc()
|
|
# test_view("int32")
|
|
# test_view1("int32")
|
|
# test_flatten(F.int32)
|
|
# test_convert_bound()
|
|
# test_convert()
|
|
# test_to_device("int32")
|
|
# test_transform("int32")
|
|
# test_subgraph("int32")
|
|
# test_subgraph_mask("int32")
|
|
# test_apply()
|
|
# test_level1()
|
|
# test_level2()
|
|
# test_updates()
|
|
# test_backward()
|
|
# test_empty_heterograph('int32')
|
|
# test_types_in_function()
|
|
# test_stack_reduce()
|
|
# test_isolated_ntype()
|
|
# test_bipartite()
|
|
# test_dtype_cast()
|
|
# test_float_cast()
|
|
# test_reverse("int32")
|
|
# test_format()
|
|
#test_add_edges(F.int32)
|
|
#test_add_nodes(F.int32)
|
|
#test_remove_edges(F.int32)
|
|
#test_remove_nodes(F.int32)
|
|
#test_clone(F.int32)
|
|
#test_frame(F.int32)
|
|
#test_frame_device(F.int32)
|
|
#test_empty_query(F.int32)
|
|
#test_create_block(F.int32)
|
|
pass
|