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

Function test_edge_removal

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

Source from the content-addressed store, hash-verified

96
97@parametrize_idtype
98def test_edge_removal(idtype):
99 g = create_graph(idtype, 5)
100 for i in range(5):
101 for j in range(5):
102 g.add_edges(i, j)
103 g.edata["id"] = F.arange(0, 25)
104
105 # remove edges
106 g.remove_edges(range(13, 20))
107 assert g.num_nodes() == 5
108 assert g.num_edges() == 18
109 assert F.array_equal(
110 g.edata["id"], F.tensor(list(range(13)) + list(range(20, 25)))
111 )
112 assert dgl.NID not in g.ndata
113 assert dgl.EID not in g.edata
114
115 # add edges
116 g.add_edges(3, 3)
117 assert g.num_nodes() == 5
118 assert g.num_edges() == 19
119 assert F.array_equal(
120 g.edata["id"], F.tensor(list(range(13)) + list(range(20, 25)) + [0])
121 )
122
123 # remove edges
124 g.remove_edges(range(2, 10), store_ids=True)
125 assert g.num_nodes() == 5
126 assert g.num_edges() == 11
127 assert F.array_equal(
128 g.edata["id"], F.tensor([0, 1, 10, 11, 12, 20, 21, 22, 23, 24, 0])
129 )
130 assert dgl.EID in g.edata
131
132
133@parametrize_idtype

Callers 1

Calls 5

remove_edgesMethod · 0.80
create_graphFunction · 0.70
add_edgesMethod · 0.45
num_nodesMethod · 0.45
num_edgesMethod · 0.45

Tested by

no test coverage detected