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 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'])