MCPcopy
hub / github.com/dmlc/dgl / test_flatten

Function test_flatten

tests/python/common/test_heterograph.py:1012–1128  ·  view source on GitHub ↗
(idtype)

Source from the content-addressed store, hash-verified

1010
1011@parametrize_idtype
1012def test_flatten(idtype):
1013 def check_mapping(g, fg):
1014 if len(fg.ntypes) == 1:
1015 SRC = DST = fg.ntypes[0]
1016 else:
1017 SRC = fg.ntypes[0]
1018 DST = fg.ntypes[1]
1019
1020 etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
1021 eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
1022
1023 for i, (etype, eid) in enumerate(zip(etypes, eids)):
1024 src_g, dst_g = g.find_edges([eid], g.canonical_etypes[etype])
1025 src_fg, dst_fg = fg.find_edges([i])
1026 # TODO(gq): I feel this code is quite redundant; can we just add new members (like
1027 # "induced_srcid") to returned heterograph object and not store them as features?
1028 assert F.asnumpy(src_g) == F.asnumpy(
1029 F.gather_row(fg.nodes[SRC].data[dgl.NID], src_fg)[0]
1030 )
1031 tid = F.asnumpy(
1032 F.gather_row(fg.nodes[SRC].data[dgl.NTYPE], src_fg)
1033 ).item()
1034 assert g.canonical_etypes[etype][0] == g.ntypes[tid]
1035 assert F.asnumpy(dst_g) == F.asnumpy(
1036 F.gather_row(fg.nodes[DST].data[dgl.NID], dst_fg)[0]
1037 )
1038 tid = F.asnumpy(
1039 F.gather_row(fg.nodes[DST].data[dgl.NTYPE], dst_fg)
1040 ).item()
1041 assert g.canonical_etypes[etype][2] == g.ntypes[tid]
1042
1043 # check for wildcard slices
1044 g = create_test_heterograph(idtype)
1045 g.nodes["user"].data["h"] = F.ones((3, 5))
1046 g.nodes["game"].data["i"] = F.ones((2, 5))
1047 g.edges["plays"].data["e"] = F.ones((4, 4))
1048 g.edges["wishes"].data["e"] = F.ones((2, 4))
1049 g.edges["wishes"].data["f"] = F.ones((2, 4))
1050
1051 fg = g["user", :, "game"] # user--plays->game and user--wishes->game
1052 assert len(fg.ntypes) == 2
1053 assert fg.ntypes == ["user", "game"]
1054 assert fg.etypes == ["plays+wishes"]
1055 assert fg.idtype == g.idtype
1056 assert fg.device == g.device
1057 etype = fg.etypes[0]
1058 assert fg[etype] is not None # Issue #2166
1059
1060 assert F.array_equal(fg.nodes["user"].data["h"], F.ones((3, 5)))
1061 assert F.array_equal(fg.nodes["game"].data["i"], F.ones((2, 5)))
1062 assert F.array_equal(fg.edata["e"], F.ones((6, 4)))
1063 assert "f" not in fg.edata
1064
1065 etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
1066 eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
1067 assert set(zip(etypes, eids)) == set(
1068 [(3, 0), (3, 1), (2, 1), (2, 0), (2, 3), (2, 2)]
1069 )

Callers

nothing calls this directly

Calls 5

check_mappingFunction · 0.85
asnumpyMethod · 0.80
create_test_heterographFunction · 0.70
edgesMethod · 0.45
ctxMethod · 0.45

Tested by

no test coverage detected