Function test_remove_selfloop

tests/python/common/transforms/test_transform.py:2198–2256

(idtype)

@parametrize_idtype
def test_remove_selfloop(idtype):
    # homogeneous graph
    g = dgl.graph(([0, 0, 0, 1], [1, 0, 0, 2]), idtype=idtype, device=F.ctx())
    g.edata["he"] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_self_loop(g)
    assert g.num_nodes() == 3
    assert g.num_edges() == 2
    assert F.array_equal(g.edata["he"], F.tensor([1, 4], dtype=idtype))

    # bipartite graph
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1, 2], [1, 2, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    # nothing will happend
    raise_error = False
    try:
        g = dgl.remove_self_loop(g, etype="plays")
    except:
        raise_error = True
    assert raise_error

    g = create_test_heterograph4(idtype)
    g = dgl.remove_self_loop(g, etype="follows")
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 2
    assert g.num_edges("follows") == 2
    assert g.num_edges("plays") == 2
    u, v = g.edges(form="uv", order="eid", etype="follows")
    assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))
    assert F.array_equal(
        g.edges["follows"].data["h"], F.tensor([2, 4], dtype=idtype)
    )
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([1, 2], dtype=idtype)
    )

    raise_error = False
    try:
        g = dgl.remove_self_loop(g, etype="plays")
    except:
        raise_error = True
    assert raise_error

    # batch information
    g = dgl.graph(
        ([0, 0, 0, 1, 3, 3, 4], [1, 0, 0, 2, 3, 4, 4]),
        idtype=idtype,
        device=F.ctx(),
    )
    g.set_batch_num_nodes([3, 2])
    g.set_batch_num_edges([4, 3])
    g = dgl.remove_self_loop(g)
    assert g.num_nodes() == 5
    assert g.num_edges() == 3
    assert F.array_equal(g.batch_num_nodes(), F.tensor([3, 2], dtype=idtype))