Function test_remove_edges

tests/python/common/transforms/test_transform.py:1130–1203

(idtype)

@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
@parametrize_idtype
def test_remove_edges(idtype):
    def check(g1, etype, g, edges_removed):
        src, dst, eid = g.edges(etype=etype, form="all")
        src1, dst1 = g1.edges(etype=etype, order="eid")
        if etype is not None:
            eid1 = g1.edges[etype].data[dgl.EID]
        else:
            eid1 = g1.edata[dgl.EID]
        src1 = F.asnumpy(src1)
        dst1 = F.asnumpy(dst1)
        eid1 = F.asnumpy(eid1)
        src = F.asnumpy(src)
        dst = F.asnumpy(dst)
        eid = F.asnumpy(eid)
        sde_set = set(zip(src, dst, eid))

        for s, d, e in zip(src1, dst1, eid1):
            assert (s, d, e) in sde_set
        assert not np.isin(edges_removed, eid1).any()
        assert g1.idtype == g.idtype

    for fmt in ["coo", "csr", "csc"]:
        for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]:
            g = dgl.graph(([0, 2, 1, 3], [1, 3, 2, 4]), idtype=idtype).formats(
                fmt
            )
            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
            check(g1, None, g, edges_to_remove)

            g = dgl.from_scipy(
                spsp.csr_matrix(
                    ([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])), shape=(5, 5)
                ),
                idtype=idtype,
            ).formats(fmt)
            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
            check(g1, None, g, edges_to_remove)

    g = dgl.heterograph(
        {
            ("A", "AA", "A"): ([0, 2, 1, 3], [1, 3, 2, 4]),
            ("A", "AB", "B"): ([0, 1, 3, 1], [1, 3, 5, 6]),
            ("B", "BA", "A"): ([2, 3], [3, 2]),
        },
        idtype=idtype,
    )
    g2 = dgl.remove_edges(
        g,
        {
            "AA": F.tensor([2], idtype),
            "AB": F.tensor([3], idtype),
            "BA": F.tensor([1], idtype),
        },
    )
    check(g2, "AA", g, [2])
    check(g2, "AB", g, [3])
    check(g2, "BA", g, [1])