| 2258 | |
| 2259 | @parametrize_idtype |
| 2260 | def test_reorder_graph(idtype): |
| 2261 | g = dgl.graph( |
| 2262 | ([0, 1, 2, 3, 4], [2, 2, 3, 2, 3]), idtype=idtype, device=F.ctx() |
| 2263 | ) |
| 2264 | g.ndata["h"] = F.copy_to(F.randn((g.num_nodes(), 3)), ctx=F.ctx()) |
| 2265 | g.edata["w"] = F.copy_to(F.randn((g.num_edges(), 2)), ctx=F.ctx()) |
| 2266 | |
| 2267 | # call with default: node_permute_algo=None, edge_permute_algo='src' |
| 2268 | rg = dgl.reorder_graph(g) |
| 2269 | assert dgl.EID in rg.edata.keys() |
| 2270 | src = F.asnumpy(rg.edges()[0]) |
| 2271 | assert np.array_equal(src, np.sort(src)) |
| 2272 | |
| 2273 | # call with 'rcmk' node_permute_algo |
| 2274 | rg = dgl.reorder_graph(g, node_permute_algo="rcmk") |
| 2275 | assert dgl.NID in rg.ndata.keys() |
| 2276 | assert dgl.EID in rg.edata.keys() |
| 2277 | src = F.asnumpy(rg.edges()[0]) |
| 2278 | assert np.array_equal(src, np.sort(src)) |
| 2279 | |
| 2280 | # call with 'dst' edge_permute_algo |
| 2281 | rg = dgl.reorder_graph(g, edge_permute_algo="dst") |
| 2282 | dst = F.asnumpy(rg.edges()[1]) |
| 2283 | assert np.array_equal(dst, np.sort(dst)) |
| 2284 | |
| 2285 | # call with unknown edge_permute_algo |
| 2286 | raise_error = False |
| 2287 | try: |
| 2288 | dgl.reorder_graph(g, edge_permute_algo="none") |
| 2289 | except: |
| 2290 | raise_error = True |
| 2291 | assert raise_error |
| 2292 | |
| 2293 | # reorder back to original according to stored ids |
| 2294 | rg = dgl.reorder_graph(g, node_permute_algo="rcmk") |
| 2295 | rg2 = dgl.reorder_graph( |
| 2296 | rg, |
| 2297 | "custom", |
| 2298 | permute_config={"nodes_perm": np.argsort(F.asnumpy(rg.ndata[dgl.NID]))}, |
| 2299 | ) |
| 2300 | assert F.array_equal(g.ndata["h"], rg2.ndata["h"]) |
| 2301 | assert F.array_equal(g.edata["w"], rg2.edata["w"]) |
| 2302 | |
| 2303 | # do not store ids |
| 2304 | rg = dgl.reorder_graph(g, store_ids=False) |
| 2305 | assert not dgl.NID in rg.ndata.keys() |
| 2306 | assert not dgl.EID in rg.edata.keys() |
| 2307 | |
| 2308 | # metis does not work on windows. |
| 2309 | if os.name == "nt": |
| 2310 | pass |
| 2311 | else: |
| 2312 | # metis_partition may fail for small graph. |
| 2313 | mg = create_large_graph(1000).to(F.ctx()) |
| 2314 | |
| 2315 | # call with metis strategy, but k is not specified |
| 2316 | raise_error = False |
| 2317 | try: |