dmlc--dgl
nv-dlasalle
50 行
1.8 KiB
Python
50 行
1.8 KiB
Python
import backend as F
|
|
from test_utils import parametrize_idtype
|
|
import dgl
|
|
|
|
@parametrize_idtype
|
|
def test_heterograph_merge(idtype):
|
|
g1 = dgl.heterograph({("a", "to", "b"): ([0,1], [1,0])}).astype(idtype).to(F.ctx())
|
|
g1_n_edges = g1.num_edges(etype="to")
|
|
g1.nodes["a"].data["nh"] = F.randn((2,3))
|
|
g1.nodes["b"].data["nh"] = F.randn((2,3))
|
|
g1.edges["to"].data["eh"] = F.randn((2,3))
|
|
|
|
g2 = dgl.heterograph({("a", "to", "b"): ([1,2,3], [2,3,5])}).astype(idtype).to(F.ctx())
|
|
g2.nodes["a"].data["nh"] = F.randn((4,3))
|
|
g2.nodes["b"].data["nh"] = F.randn((6,3))
|
|
g2.edges["to"].data["eh"] = F.randn((3,3))
|
|
g2.add_nodes(3, ntype="a")
|
|
g2.add_nodes(3, ntype="b")
|
|
|
|
m = dgl.merge([g1, g2])
|
|
|
|
# Check g2's edges and nodes were added to g1's in m.
|
|
m_us = F.asnumpy(m.edges()[0][g1_n_edges:])
|
|
g2_us = F.asnumpy(g2.edges()[0])
|
|
assert all(m_us == g2_us)
|
|
m_vs = F.asnumpy(m.edges()[1][g1_n_edges:])
|
|
g2_vs = F.asnumpy(g2.edges()[1])
|
|
assert all(m_vs == g2_vs)
|
|
for ntype in m.ntypes:
|
|
assert m.num_nodes(ntype=ntype) == max(
|
|
g1.num_nodes(ntype=ntype), g2.num_nodes(ntype=ntype)
|
|
)
|
|
|
|
# Check g1's node data was updated with g2's in m.
|
|
for key in m.nodes[ntype].data:
|
|
g2_n_nodes = g2.num_nodes(ntype=ntype)
|
|
updated_g1_ndata = F.asnumpy(m.nodes[ntype].data[key][:g2_n_nodes])
|
|
g2_ndata = F.asnumpy(g2.nodes[ntype].data[key])
|
|
assert all(
|
|
(updated_g1_ndata == g2_ndata).flatten()
|
|
)
|
|
|
|
# Check g1's edge data was updated with g2's in m.
|
|
for key in m.edges["to"].data:
|
|
updated_g1_edata = F.asnumpy(m.edges["to"].data[key][g1_n_edges:])
|
|
g2_edata = F.asnumpy(g2.edges["to"].data[key])
|
|
assert all(
|
|
(updated_g1_edata == g2_edata).flatten()
|
|
)
|