def test_multi_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + nodes.mailbox['m'].sum(1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + th.zeros(shape, dtype=dtype, device=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2)
g.add_nodes(2)
g.add_edge(0, 1)
# recv both 0deg and non-0deg nodes
old = th.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata['h']
# 0deg check: initialized with the func and got applied
assert U.allclose(new[0], th.full((5,), 4))
# non-0deg check
assert U.allclose(new[1], th.sum(old, 0) * 2)
# recv again on zero degree node
g.recv([0])
assert U.allclose(g.nodes[0].data['h'], th.full((5,), 8))
# recv again on node with no incoming message
g.recv([1])
assert U.allclose(g.nodes[1].data['h'], th.sum(old, 0) * 4)
def _disabled_test_send_twice():
# TODO(minjie): please re-enable this unittest after the send code problem is fixed.
g = DGLGraph()
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(2, 1)
def _message_a(edges):
return {'a': edges.src['a']}
def _message_b(edges):
return {'a': edges.src['a'] * 3}
def _reduce(nodes):
return {'a': nodes.mailbox['a'].max(1)[0]}
old_repr = th.randn(3, 5)
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((0, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert U.allclose(new_repr[1], old_repr[0] * 3)
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((2, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert U.allclose(new_repr[1],
th.stack([old_repr[0], old_repr[2] * 3], 0).max(0)[0])
def test_send_twice_different_msg():
g = DGLGraph()
g.set_n_initializer(dgl.init.zero_initializer)
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(2, 1)
def _message_a(edges):
return {'a': edges.src['a']}
def _message_b(edges):
return {'a': edges.src['a'] * 3}
def _reduce(nodes):
return {'a': nodes.mailbox['a'].max(1)[0]}
old_repr = th.randn(3, 5)
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((0, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert U.allclose(new_repr[1], old_repr[0] * 3)
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((2, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert U.allclose(new_repr[1], th.stack([old_repr[0], old_repr[2] * 3], 0).max(0)[0])
def test_sharing():
data = Frame(create_test_data())
f1 = FrameRef(data, index=toindex([0, 1, 2, 3]))
f2 = FrameRef(data, index=toindex([2, 3, 4, 5, 6]))
# test read
for k, v in f1.items():
assert U.allclose(data[k].data[0:4], v)
for k, v in f2.items():
assert U.allclose(data[k].data[2:7], v)
f2_a1 = f2['a1'].data
# test write
# update own ref should not been seen by the other.
f1[Index(th.tensor([0, 1]))] = {
'a1': th.zeros([2, D]),
'a2': th.zeros([2, D]),
'a3': th.zeros([2, D]),
}
assert U.allclose(f2['a1'], f2_a1)
# update shared space should been seen by the other.
f1[Index(th.tensor([2, 3]))] = {
'a1': th.ones([2, D]),
'a2': th.ones([2, D]),
'a3': th.ones([2, D]),
}
f2_a1[0:2] = th.ones([2, D])
assert U.allclose(f2['a1'], f2_a1)
def _check_nx_feature(nxg, nf, ef):
# check node and edge feature of nxg
# this is used to check to_networkx
num_nodes = len(nxg)
num_edges = nxg.size()
if num_nodes > 0:
node_feat = ddict(list)
for nid, attr in nxg.nodes(data=True):
assert len(attr) == len(nf)
for k in nxg.nodes[nid]:
node_feat[k].append(attr[k].unsqueeze(0))
for k in node_feat:
feat = th.cat(node_feat[k], dim=0)
assert U.allclose(feat, nf[k])
else:
assert len(nf) == 0
if num_edges > 0:
edge_feat = ddict(lambda: [0] * num_edges)
for u, v, attr in nxg.edges(data=True):
assert len(attr) == len(ef) + 1 # extra id
eid = attr['id']
for k in ef:
edge_feat[k][eid] = attr[k].unsqueeze(0)
for k in edge_feat:
feat = th.cat(edge_feat[k], dim=0)
assert U.allclose(feat, ef[k])
else:
assert len(ef) == 0
def test_slicing():
data = Frame(create_test_data(grad=True))
f1 = FrameRef(data, index=toindex(slice(1, 5)))
f2 = FrameRef(data, index=toindex(slice(3, 8)))
# test read
for k, v in f1.items():
assert U.allclose(data[k].data[1:5], v)
f2_a1 = f2['a1'].data
# test write
f1[Index(th.tensor([0, 1]))] = {
'a1': th.zeros([2, D]),
'a2': th.zeros([2, D]),
'a3': th.zeros([2, D]),
}
assert U.allclose(f2['a1'], f2_a1)
f1[Index(th.tensor([2, 3]))] = {
'a1': th.ones([2, D]),
'a2': th.ones([2, D]),
'a3': th.ones([2, D]),
}
f2_a1[toindex(slice(0, 2))] = 1
assert U.allclose(f2['a1'], f2_a1)
f1[toindex(slice(2, 4))] = {
'a1': th.zeros([2, D]),
'a2': th.zeros([2, D]),
'a3': th.zeros([2, D]),
}
f2_a1[toindex(slice(0, 2))] = 0
assert U.allclose(f2['a1'], f2_a1)
def test_apply_nodes():
def _upd(nodes):
return {'h': nodes.data['h'] * 2}
g = generate_graph()
g.register_apply_node_func(_upd)
old = g.ndata['h']
g.apply_nodes()
assert U.allclose(old * 2, g.ndata['h'])
u = th.tensor([0, 3, 4, 6])
g.apply_nodes(lambda nodes: {'h': nodes.data['h'] * 0.}, u)
assert U.allclose(g.ndata['h'][u], th.zeros((4, D)))
def test_spmv_3d_feat():
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * edges.data['h'].unsqueeze(1).unsqueeze(1)}
def sum_udf(nodes):
return {'h': nodes.mailbox['sum'].sum(1)}
n = 100
p = 0.1
a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
g = dgl.DGLGraph(a)
m = g.number_of_edges()
# test#1: v2v with adj data
h = th.randn((n, 5, 5))
e = th.randn((m,))
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=fn.src_mul_edge('h', 'h', 'sum'), reduce_func=fn.sum('sum', 'h')) # 1
ans = g.ndata['h']
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=fn.sum('sum', 'h')) # 2
assert U.allclose(g.ndata['h'], ans)
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf) # 3
assert U.allclose(g.ndata['h'], ans)
# test#2: e2v
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * edges.data['h']}
h = th.randn((n, 5, 5))
e = th.randn((m, 5, 5))
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=fn.src_mul_edge('h', 'h', 'sum'), reduce_func=fn.sum('sum', 'h')) # 1
ans = g.ndata['h']
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=fn.sum('sum', 'h')) # 2
assert U.allclose(g.ndata['h'], ans)
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf) # 3
assert U.allclose(g.ndata['h'], ans)
def test_batch_unbatch2():
# test setting/getting features after batch
a = dgl.DGLGraph()
a.add_nodes(4)
a.add_edges(0, [1, 2, 3])
b = dgl.DGLGraph()
b.add_nodes(3)
b.add_edges(0, [1, 2])
c = dgl.batch([a, b])
c.ndata['h'] = th.ones(7, 1)
c.edata['w'] = th.ones(5, 1)
assert U.allclose(c.ndata['h'], th.ones(7, 1))
assert U.allclose(c.edata['w'], th.ones(5, 1))
def _assert_is_identical(g, g2):
assert g.number_of_nodes() == g2.number_of_nodes()
src, dst = g.all_edges()
src2, dst2 = g2.all_edges()
assert torch.equal(src, src2)
assert torch.equal(dst, dst2)
assert len(g.ndata) == len(g2.ndata)
assert len(g.edata) == len(g2.edata)
for k in g.ndata:
assert U.allclose(g.ndata[k], g2.ndata[k])
for k in g.edata:
assert U.allclose(g.edata[k], g2.edata[k])
def test_send_recv_after_conversion():
# test send and recv after converting from a graph with edges
g = generate_graph()
# nx graph
nxg = g.to_networkx(node_attrs=['h'])
g1 = DGLGraph()
# some random node and edges
g1.add_nodes(4)
g1.add_edges([1, 2], [2, 3])
g1.set_n_initializer(dgl.init.zero_initializer)
g1.from_networkx(nxg, node_attrs=['h'])
# sparse matrix
row, col= g.all_edges()
data = range(len(row))
n = g.number_of_nodes()
a = sp.coo_matrix((data, (row, col)), shape=(n, n))
g2 = DGLGraph()
# some random node and edges
g2.add_nodes(5)
g2.add_edges([1, 2, 4], [2, 3, 0])
g2.set_n_initializer(dgl.init.zero_initializer)
g2.from_scipy_sparse_matrix(a)
g2.ndata['h'] = g.ndata['h']
# on dgl graph
g.send(message_func=message_func)
g.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
# nx
g1.send(message_func=message_func)
g1.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g1.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
# sparse matrix
g2.send(message_func=message_func)
g2.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g2.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
assert U.allclose(g.ndata['h'], g1.ndata['h'])
assert U.allclose(g.ndata['h'], g2.ndata['h'])
def test_apply_edges():
def _upd(edges):
return {'w': edges.data['w'] * 2}
g = generate_graph()
g.register_apply_edge_func(_upd)
old = g.edata['w']
g.apply_edges()
assert U.allclose(old * 2, g.edata['w'])
u = th.tensor([0, 0, 0, 4, 5, 6])
v = th.tensor([1, 2, 3, 9, 9, 9])
g.apply_edges(lambda edges: {'w': edges.data['w'] * 0.}, (u, v))
eid = g.edge_ids(u, v)
assert U.allclose(g.edata['w'][eid], th.zeros((6, D)))
def test_add_rows():
data = Frame()
f1 = FrameRef(data)
f1.add_rows(4)
x = th.randn(1, 4)
f1[Index(th.tensor([0]))] = {'x': x}
ans = th.cat([x, th.zeros(3, 4)])
assert U.allclose(f1['x'], ans)
f1.add_rows(4)
f1[toindex(slice(4, 8))] = {'x': th.ones(4, 4), 'y': th.ones(4, 5)}
ans = th.cat([ans, th.ones(4, 4)])
assert U.allclose(f1['x'], ans)
ans = th.cat([th.zeros(4, 5), th.ones(4, 5)])
assert U.allclose(f1['y'], ans)
def testJacobian(self):
import pybullet as p
clid = p.connect(p.SHARED_MEMORY)
if (clid < 0):
p.connect(p.DIRECT)
time_step = 0.001
gravity_constant = -9.81
urdfs = [
"TwoJointRobot_w_fixedJoints.urdf",
"TwoJointRobot_w_fixedJoints.urdf", "kuka_iiwa/model.urdf",
"kuka_lwr/kuka.urdf"
]
for urdf in urdfs:
p.resetSimulation()
p.setTimeStep(time_step)
p.setGravity(0.0, 0.0, gravity_constant)
robotId = p.loadURDF(urdf, useFixedBase=True)
p.resetBasePositionAndOrientation(robotId, [0, 0, 0], [0, 0, 0, 1])
numJoints = p.getNumJoints(robotId)
endEffectorIndex = numJoints - 1
# Set a joint target for the position control and step the sim.
self.setJointPosition(robotId,
[0.1 * (i % 3) for i in range(numJoints)])
p.stepSimulation()
# Get the joint and link state directly from Bullet.
mpos, mvel, mtorq = self.getMotorJointStates(robotId)
result = p.getLinkState(robotId,
endEffectorIndex,
computeLinkVelocity=1,
computeForwardKinematics=1)
link_trn, link_rot, com_trn, com_rot, frame_pos, frame_rot, link_vt, link_vr = result
# Get the Jacobians for the CoM of the end-effector link.
# Note that in this example com_rot = identity, and we would need to use com_rot.T * com_trn.
# The localPosition is always defined in terms of the link frame coordinates.
zero_vec = [0.0] * len(mpos)
jac_t, jac_r = p.calculateJacobian(robotId, endEffectorIndex,
com_trn, mpos, zero_vec,
zero_vec)
assert (allclose(dot(jac_t, mvel), link_vt))
assert (allclose(dot(jac_r, mvel), link_vr))
p.disconnect()
def _pull_nodes(nodes):
# compute ground truth
g.pull(nodes, _mfunc_hxw1, _rfunc_m1, _afunc)
o1 = g.ndata.pop('o1')
g.pull(nodes, _mfunc_hxw2, _rfunc_m2, _afunc)
o2 = g.ndata.pop('o2')
g.pull(nodes, _mfunc_hxw1, _rfunc_m1max, _afunc)
o3 = g.ndata.pop('o3')
# v2v spmv
g.pull(nodes, fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.sum(msg='m1', out='o1'),
_afunc)
assert U.allclose(o1, g.ndata.pop('o1'))
# v2v fallback to e2v
g.pull(nodes, fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.sum(msg='m2', out='o2'),
_afunc)
assert U.allclose(o2, g.ndata.pop('o2'))
# v2v fallback to degree bucketing
g.pull(nodes, fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.max(msg='m1', out='o3'),
_afunc)
assert U.allclose(o3, g.ndata.pop('o3'))
# multi builtins, both v2v spmv
g.pull(nodes,
[fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
[fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
_afunc)
assert U.allclose(o1, g.ndata.pop('o1'))
assert U.allclose(o1, g.ndata.pop('o2'))
# multi builtins, one v2v spmv, one fallback to e2v
g.pull(nodes,
[fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w2', out='m2')],
[fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
_afunc)
assert U.allclose(o1, g.ndata.pop('o1'))
assert U.allclose(o2, g.ndata.pop('o2'))
# multi builtins, one v2v spmv, one fallback to e2v, one fallback to degree-bucketing
g.pull(nodes,
[fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.src_mul_edge(src='h', edge='w1', out='m3')],
[fn.sum(msg='m1', out='o1'),
fn.sum(msg='m2', out='o2'),
fn.max(msg='m3', out='o3')],
_afunc)
assert U.allclose(o1, g.ndata.pop('o1'))
assert U.allclose(o2, g.ndata.pop('o2'))
assert U.allclose(o3, g.ndata.pop('o3'))
def _test(fld):
def message_func(edges):
return {'m1' : edges.src[fld] + edges.dst[fld],
'm2' : edges.src[fld] * edges.dst[fld]}
def reduce_func(nodes):
return {fld : th.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
def apply_func_2(nodes):
return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
g = generate_graph()
# recv
v1 = g.get_n_repr()[fld]
# no specialization
g.send((u, v), message_func)
g.recv([0,1,2,3,9], reduce_func, apply_func)
v2 = g.get_n_repr()[fld]
# user break reduce func into 2 builtin
g.set_n_repr({fld : v1})
g.send((u, v), message_func)
g.recv([0,1,2,3,9],
[fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
apply_func_2)
v3 = g.get_n_repr()[fld]
assert U.allclose(v2, v3)
def test_pickling_frame():
x = torch.randn(3, 7)
y = torch.randn(3, 5)
c = Column(x)
c2 = _reconstruct_pickle(c)
assert U.allclose(c.data, c2.data)
fr = Frame({'x': x, 'y': y})
fr2 = _reconstruct_pickle(fr)
assert U.allclose(fr2['x'].data, x)
assert U.allclose(fr2['y'].data, y)
fr = Frame()
def test_append3():
# test append on empty frame
f = Frame(num_rows=5)
data = {'h': th.ones((3, 2))}
f.append(data)
assert f.num_rows == 8
ans = th.cat([th.zeros((5, 2)), th.ones((3, 2))], dim=0)
assert U.allclose(f['h'].data, ans)
# test append with new column
data = {'h': 2 * th.ones((3, 2)), 'w': 2 * th.ones((3, 2))}
f.append(data)
assert f.num_rows == 11
ans1 = th.cat([ans, 2 * th.ones((3, 2))], 0)
ans2 = th.cat([th.zeros((8, 2)), 2 * th.ones((3, 2))], 0)
assert U.allclose(f['h'].data, ans1)
assert U.allclose(f['w'].data, ans2)
def test_batch_unbatch():
t1 = tree1()
t2 = tree2()
bg = dgl.batch([t1, t2])
assert bg.number_of_nodes() == 10
assert bg.number_of_edges() == 8
assert bg.batch_size == 2
assert bg.batch_num_nodes == [5, 5]
assert bg.batch_num_edges == [4, 4]
tt1, tt2 = dgl.unbatch(bg)
assert U.allclose(t1.ndata['h'], tt1.ndata['h'])
assert U.allclose(t1.edata['h'], tt1.edata['h'])
assert U.allclose(t2.ndata['h'], tt2.ndata['h'])
assert U.allclose(t2.edata['h'], tt2.edata['h'])
def test_v2v_snr_multi_fn():
u = th.tensor([0, 0, 0, 3, 4, 9])
v = th.tensor([1, 2, 3, 9, 9, 0])
def message_func(edges):
return {'m2': edges.src['f2']}
def message_func_edge(edges):
return {'m2': edges.src['f2'] * edges.data['e2']}
def reduce_func(nodes):
return {'v1' : th.sum(nodes.mailbox['m2'], 1)}
g = generate_graph()
g.set_n_repr({'v1' : th.zeros((10, D)), 'v2' : th.zeros((10, D)),
'v3' : th.zeros((10, D))})
fld = 'f2'
g.send_and_recv((u, v), message_func, reduce_func)
v1 = g.ndata['v1']
# 1 message, 2 reduces
g.send_and_recv((u, v),
fn.copy_src(src=fld, out='m'),
[fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')],
None)
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert U.allclose(v1, v2)
assert U.allclose(v1, v3)
# send and recv with edge weights, 2 message, 3 reduces
g.send_and_recv((u, v),
[fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
[fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
None)
v1 = g.ndata['v1']
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert U.allclose(v1, v2)
assert U.allclose(v1, v3)
# run UDF with single message and reduce
g.send_and_recv((u, v), message_func_edge,
reduce_func, None)
v2 = g.ndata['v2']
assert U.allclose(v1, v2)
def testJacobian(self):
import pybullet as p
clid = p.connect(p.SHARED_MEMORY)
if (clid < 0):
p.connect(p.DIRECT)
time_step = 0.001
gravity_constant = -9.81
urdfs = [
"TwoJointRobot_w_fixedJoints.urdf", "TwoJointRobot_w_fixedJoints.urdf",
"kuka_iiwa/model.urdf", "kuka_lwr/kuka.urdf"
]
for urdf in urdfs:
p.resetSimulation()
p.setTimeStep(time_step)
p.setGravity(0.0, 0.0, gravity_constant)
robotId = p.loadURDF(urdf, useFixedBase=True)
p.resetBasePositionAndOrientation(robotId, [0, 0, 0], [0, 0, 0, 1])
numJoints = p.getNumJoints(robotId)
endEffectorIndex = numJoints - 1
# Set a joint target for the position control and step the sim.
self.setJointPosition(robotId, [0.1 * (i % 3) for i in range(numJoints)])
p.stepSimulation()
# Get the joint and link state directly from Bullet.
mpos, mvel, mtorq = self.getMotorJointStates(robotId)
result = p.getLinkState(robotId,
endEffectorIndex,
computeLinkVelocity=1,
computeForwardKinematics=1)
link_trn, link_rot, com_trn, com_rot, frame_pos, frame_rot, link_vt, link_vr = result
# Get the Jacobians for the CoM of the end-effector link.
# Note that in this example com_rot = identity, and we would need to use com_rot.T * com_trn.
# The localPosition is always defined in terms of the link frame coordinates.
zero_vec = [0.0] * len(mpos)
jac_t, jac_r = p.calculateJacobian(robotId, endEffectorIndex, com_trn, mpos, zero_vec,
zero_vec)
assert (allclose(dot(jac_t, mvel), link_vt))
assert (allclose(dot(jac_r, mvel), link_vr))
p.disconnect()
def test_column2():
# Test frameref column getter/setter
data = Frame(create_test_data())
f = FrameRef(data, toindex([3, 4, 5, 6, 7]))
assert f.num_rows == 5
assert len(f) == 3
assert U.allclose(f['a1'], data['a1'].data[3:8])
# set column should reflect on the referenced data
f['a1'] = th.zeros([5, D])
assert U.allclose(data['a1'].data[3:8], th.zeros([5, D]))
# add new partial column should fail with error initializer
f.set_initializer(lambda shape, dtype: assert_(False))
def failed_add_col():
f['a4'] = th.ones([5, D])
assert check_fail(failed_add_col)
def test_row4():
# test updating row with empty frame but has preset num_rows
f = FrameRef(Frame(num_rows=5))
rowid = Index(th.tensor([0, 2, 4]))
f[rowid] = {'h': th.ones((3, 2))}
ans = th.zeros((5, 2))
ans[th.tensor([0, 2, 4])] = th.ones((3, 2))
assert U.allclose(f['h'], ans)
def test_src_mul_edge():
# src_mul_edge with all fields
g = generate_graph()
g.register_message_func(fn.src_mul_edge(src='h', edge='h', out='m'))
g.register_reduce_func(reducer_both)
g.update_all()
assert U.allclose(g.ndata['h'],
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
def test_copy_edge():
# copy_edge with both fields
g = generate_graph()
g.register_message_func(fn.copy_edge(edge='h', out='m'))
g.register_reduce_func(reducer_both)
g.update_all()
assert U.allclose(g.ndata['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
def test_reverse_shared_frames():
g = dgl.DGLGraph()
g.add_nodes(3)
g.add_edges([0, 1, 2], [1, 2, 1])
g.ndata['h'] = th.tensor([[0.], [1.], [2.]], requires_grad=True)
g.edata['h'] = th.tensor([[3.], [4.], [5.]], requires_grad=True)
rg = g.reverse(share_ndata=True, share_edata=True)
assert U.allclose(g.ndata['h'], rg.ndata['h'])
assert U.allclose(g.edata['h'], rg.edata['h'])
assert U.allclose(g.edges[[0, 2], [1, 1]].data['h'],
rg.edges[[1, 1], [0, 2]].data['h'])
rg.ndata['h'] = rg.ndata['h'] + 1
assert U.allclose(rg.ndata['h'], g.ndata['h'])
g.edata['h'] = g.edata['h'] - 1
assert U.allclose(rg.edata['h'], g.edata['h'])
src_msg = fn.copy_src(src='h', out='m')
sum_reduce = fn.sum(msg='m', out='h')
rg.update_all(src_msg, sum_reduce)
assert U.allclose(g.ndata['h'], rg.ndata['h'])
# Grad check
g.ndata['h'].retain_grad()
rg.ndata['h'].retain_grad()
loss_func = th.nn.MSELoss()
target = th.zeros(3, 1)
loss = loss_func(rg.ndata['h'], target)
loss.backward()
assert U.allclose(g.ndata['h'].grad, rg.ndata['h'].grad)
def test_recv_0deg_newfld():
# test recv with 0deg nodes; the reducer also creates a new field
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h1': nodes.data['h'] + nodes.mailbox['m'].sum(1)}
def _apply(nodes):
return {'h1': nodes.data['h1'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + th.zeros(shape, dtype=dtype, device=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
# test#1: recv both 0deg and non-0deg nodes
old = th.randn((2, 5))
g.set_n_initializer(_init2, 'h1')
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h1')
# 0deg check: initialized with the func and got applied
assert U.allclose(new[0], th.full((5, ), 4))
# non-0deg check
assert U.allclose(new[1], th.sum(old, 0) * 2)
# test#2: recv only 0deg node
old = th.randn((2, 5))
g.ndata['h'] = old
g.ndata['h1'] = th.full((2, 5), -1) # this is necessary
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h1')
# 0deg check: fallback to apply
assert U.allclose(new[0], th.full((5, ), -2))
# non-0deg check: not changed
assert U.allclose(new[1], th.full((5, ), -1))
def test_basics():
g = generate_graph()
h = g.ndata['h']
l = g.edata['l']
nid = [0, 2, 3, 6, 7, 9]
sg = g.subgraph(nid)
eid = {2, 3, 4, 5, 10, 11, 12, 13, 16}
assert set(sg.parent_eid.numpy()) == eid
eid = sg.parent_eid
# the subgraph is empty initially
assert len(sg.ndata) == 0
assert len(sg.edata) == 0
# the data is copied after explict copy from
sg.copy_from_parent()
assert len(sg.ndata) == 1
assert len(sg.edata) == 1
sh = sg.ndata['h']
assert U.allclose(h[nid], sh)
'''
s, d, eid
0, 1, 0
1, 9, 1
0, 2, 2 1
2, 9, 3 1
0, 3, 4 1
3, 9, 5 1
0, 4, 6
4, 9, 7
0, 5, 8
5, 9, 9 3
0, 6, 10 1
6, 9, 11 1 3
0, 7, 12 1
7, 9, 13 1 3
0, 8, 14
8, 9, 15 3
9, 0, 16 1
'''
assert U.allclose(l[eid], sg.edata['l'])
# update the node/edge features on the subgraph should NOT
# reflect to the parent graph.
sg.ndata['h'] = th.zeros((6, D))
assert U.allclose(h, g.ndata['h'])
def test_prop_nodes_bfs():
g = dgl.DGLGraph(nx.path_graph(5))
g.ndata['x'] = th.ones((5, 2))
g.register_message_func(mfunc)
g.register_reduce_func(rfunc)
dgl.prop_nodes_bfs(g, 0)
# pull nodes using bfs order will result in a cumsum[i] + data[i] + data[i+1]
assert U.allclose(g.ndata['x'],
th.tensor([[2., 2.], [4., 4.], [6., 6.], [8., 8.], [9., 9.]]))
def test_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h': nodes.data['h'] + nodes.mailbox['m'].sum(1)}
def _apply(nodes):
return {'h': nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + th.zeros(shape, dtype=dtype, device=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2, 'h')
# test#1: recv both 0deg and non-0deg nodes
old = th.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert U.allclose(new[0], th.full((5, ), 4))
# non-0deg check
assert U.allclose(new[1], th.sum(old, 0) * 2)
# test#2: recv only 0deg node is equal to apply
old = th.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h')
# 0deg check: equal to apply_nodes
assert U.allclose(new[0], 2 * old[0])
# non-0deg check: untouched
assert U.allclose(new[1], old[1])
def test_graph_creation():
g = dgl.DGLGraph()
# test add nodes with data
g.add_nodes(5)
g.add_nodes(5, {'h': th.ones((5, 2))})
ans = th.cat([th.zeros(5, 2), th.ones(5, 2)], 0)
U.allclose(ans, g.ndata['h'])
g.ndata['w'] = 2 * th.ones((10, 2))
assert U.allclose(2 * th.ones((10, 2)), g.ndata['w'])
# test add edges with data
g.add_edges([2, 3], [3, 4])
g.add_edges([0, 1], [1, 2], {'m': th.ones((2, 2))})
ans = th.cat([th.zeros(2, 2), th.ones(2, 2)], 0)
assert U.allclose(ans, g.edata['m'])
# test clear and add again
g.clear()
g.add_nodes(5)
g.ndata['h'] = 3 * th.ones((5, 2))
assert U.allclose(3 * th.ones((5, 2)), g.ndata['h'])
