Function test_subgraph_relabel_nodes

tests/python/common/test_subgraph.py:58–110

(relabel_nodes)

@pytest.mark.parametrize("relabel_nodes", [True, False])
def test_subgraph_relabel_nodes(relabel_nodes):
    g = generate_graph()
    h = g.ndata["h"]
    l = g.edata["l"]
    nid = [0, 2, 3, 6, 7, 9]
    sg = g.subgraph(nid, relabel_nodes=relabel_nodes)
    eid = {2, 3, 4, 5, 10, 11, 12, 13, 16}
    assert set(F.asnumpy(sg.edata[dgl.EID])) == eid
    eid = sg.edata[dgl.EID]
    # the subgraph is empty initially except for EID field
    # the subgraph is empty initially except for NID field if relabel_nodes
    if relabel_nodes:
        assert len(sg.ndata) == 2
    assert len(sg.edata) == 2
    sh = sg.ndata["h"]
    # The node number is not reduced if relabel_node=False.
    # The subgraph keeps the same node information as the original graph.
    if relabel_nodes:
        assert F.allclose(F.gather_row(h, F.tensor(nid)), sh)
    else:
        assert F.allclose(
            F.gather_row(h, F.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])), sh
        )
    # The s,d,eid means the source node, destination node and edge id of the subgraph.
    # The edges labeled 1 are those selected by the subgraph.
    """
    s, d, eid
    0, 1, 0
    1, 9, 1
    0, 2, 2    1
    2, 9, 3    1
    0, 3, 4    1
    3, 9, 5    1
    0, 4, 6
    4, 9, 7
    0, 5, 8
    5, 9, 9       3
    0, 6, 10   1
    6, 9, 11   1  3
    0, 7, 12   1
    7, 9, 13   1  3
    0, 8, 14
    8, 9, 15      3
    9, 0, 16   1
    """
    assert F.allclose(F.gather_row(l, eid), sg.edata["l"])
    # update the node/edge features on the subgraph should NOT
    # reflect to the parent graph.
    if relabel_nodes:
        sg.ndata["h"] = F.zeros((6, D))
    else:
        sg.ndata["h"] = F.zeros((10, D))
    assert F.allclose(h, g.ndata["h"])


def _test_map_to_subgraph():