MCPcopy
hub / github.com/dmlc/dgl / test_node_removal

Function test_node_removal

tests/python/common/test_heterograph-remove.py:16–41  ·  view source on GitHub ↗
(idtype)

Source from the content-addressed store, hash-verified

14
15@parametrize_idtype
16def test_node_removal(idtype):
17 g = create_graph(idtype, 10)
18 g.add_edges(0, 0)
19 assert g.num_nodes() == 10
20 g.ndata["id"] = F.arange(0, 10)
21
22 # remove nodes
23 g.remove_nodes(range(4, 7))
24 assert g.num_nodes() == 7
25 assert F.array_equal(g.ndata["id"], F.tensor([0, 1, 2, 3, 7, 8, 9]))
26 assert dgl.NID not in g.ndata
27 assert dgl.EID not in g.edata
28
29 # add nodes
30 g.add_nodes(3)
31 assert g.num_nodes() == 10
32 assert F.array_equal(
33 g.ndata["id"], F.tensor([0, 1, 2, 3, 7, 8, 9, 0, 0, 0])
34 )
35
36 # remove nodes
37 g.remove_nodes(range(1, 4), store_ids=True)
38 assert g.num_nodes() == 7
39 assert F.array_equal(g.ndata["id"], F.tensor([0, 7, 8, 9, 0, 0, 0]))
40 assert dgl.NID in g.ndata
41 assert dgl.EID in g.edata
42
43
44@parametrize_idtype

Callers 1

Calls 5

remove_nodesMethod · 0.80
create_graphFunction · 0.70
add_edgesMethod · 0.45
num_nodesMethod · 0.45
add_nodesMethod · 0.45

Tested by

no test coverage detected