| 2506 | |
| 2507 | @parametrize_idtype |
| 2508 | def test_module_remove_self_loop(idtype): |
| 2509 | transform = dgl.RemoveSelfLoop() |
| 2510 | |
| 2511 | # Case1: homogeneous graph |
| 2512 | g = dgl.graph(([1, 1], [1, 2]), idtype=idtype, device=F.ctx()) |
| 2513 | g.ndata["h"] = F.randn((g.num_nodes(), 2)) |
| 2514 | g.edata["w"] = F.randn((g.num_edges(), 3)) |
| 2515 | new_g = transform(g) |
| 2516 | assert new_g.device == g.device |
| 2517 | assert new_g.idtype == g.idtype |
| 2518 | assert new_g.num_nodes() == g.num_nodes() |
| 2519 | assert new_g.num_edges() == 1 |
| 2520 | src, dst = new_g.edges() |
| 2521 | eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) |
| 2522 | assert eset == {(1, 2)} |
| 2523 | assert "h" in new_g.ndata |
| 2524 | assert "w" in new_g.edata |
| 2525 | |
| 2526 | # Case2: heterogeneous graph |
| 2527 | g = dgl.heterograph( |
| 2528 | { |
| 2529 | ("user", "plays", "game"): ([0, 1], [1, 1]), |
| 2530 | ("user", "follows", "user"): ([1, 2], [2, 2]), |
| 2531 | }, |
| 2532 | idtype=idtype, |
| 2533 | device=F.ctx(), |
| 2534 | ) |
| 2535 | g.nodes["user"].data["h1"] = F.randn((3, 2)) |
| 2536 | g.edges["plays"].data["w1"] = F.randn((2, 3)) |
| 2537 | g.nodes["game"].data["h2"] = F.randn((2, 4)) |
| 2538 | g.edges["follows"].data["w2"] = F.randn((2, 5)) |
| 2539 | |
| 2540 | new_g = transform(g) |
| 2541 | assert new_g.device == g.device |
| 2542 | assert new_g.idtype == g.idtype |
| 2543 | assert new_g.ntypes == g.ntypes |
| 2544 | assert new_g.canonical_etypes == g.canonical_etypes |
| 2545 | for nty in new_g.ntypes: |
| 2546 | assert new_g.num_nodes(nty) == g.num_nodes(nty) |
| 2547 | assert new_g.num_edges("plays") == 2 |
| 2548 | assert new_g.num_edges("follows") == 1 |
| 2549 | assert "h1" in new_g.nodes["user"].data |
| 2550 | assert "h2" in new_g.nodes["game"].data |
| 2551 | assert "w1" in new_g.edges["plays"].data |
| 2552 | assert "w2" in new_g.edges["follows"].data |
| 2553 | |
| 2554 | |
| 2555 | @parametrize_idtype |