(g, idtype, descending)
| 125 | @pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"])) |
| 126 | @pytest.mark.parametrize("descending", [True, False]) |
| 127 | def test_topk(g, idtype, descending): |
| 128 | g = g.astype(idtype).to(F.ctx()) |
| 129 | g.ndata["x"] = F.randn((g.num_nodes(), 3)) |
| 130 | |
| 131 | # Test.1: to test the case where k > number of nodes. |
| 132 | dgl.topk_nodes(g, "x", 100, sortby=-1) |
| 133 | |
| 134 | # Test.2: test correctness |
| 135 | min_nnodes = F.asnumpy(g.batch_num_nodes()).min() |
| 136 | if min_nnodes <= 1: |
| 137 | return |
| 138 | k = min_nnodes - 1 |
| 139 | val, indices = dgl.topk_nodes(g, "x", k, descending=descending, sortby=-1) |
| 140 | print(k) |
| 141 | print(g.ndata["x"]) |
| 142 | print("val", val) |
| 143 | print("indices", indices) |
| 144 | subg = dgl.unbatch(g) |
| 145 | subval, subidx = [], [] |
| 146 | for sg in subg: |
| 147 | subx = F.asnumpy(sg.ndata["x"]) |
| 148 | ai = np.argsort(subx[:, -1:].flatten()) |
| 149 | if descending: |
| 150 | ai = np.ascontiguousarray(ai[::-1]) |
| 151 | subx = np.expand_dims(subx[ai[:k]], 0) |
| 152 | subval.append(F.tensor(subx)) |
| 153 | subidx.append(F.tensor(np.expand_dims(ai[:k], 0))) |
| 154 | print(F.cat(subval, dim=0)) |
| 155 | assert F.allclose(val, F.cat(subval, dim=0)) |
| 156 | assert F.allclose(indices, F.cat(subidx, dim=0)) |
| 157 | |
| 158 | # Test.3: sorby=None |
| 159 | dgl.topk_nodes(g, "x", k, sortby=None) |
| 160 | |
| 161 | g.edata["x"] = F.randn((g.num_edges(), 3)) |
| 162 | |
| 163 | # Test.4: topk edges where k > number of edges. |
| 164 | dgl.topk_edges(g, "x", 100, sortby=-1) |
| 165 | |
| 166 | # Test.5: topk edges test correctness |
| 167 | min_nedges = F.asnumpy(g.batch_num_edges()).min() |
| 168 | if min_nedges <= 1: |
| 169 | return |
| 170 | k = min_nedges - 1 |
| 171 | val, indices = dgl.topk_edges(g, "x", k, descending=descending, sortby=-1) |
| 172 | print(k) |
| 173 | print(g.edata["x"]) |
| 174 | print("val", val) |
| 175 | print("indices", indices) |
| 176 | subg = dgl.unbatch(g) |
| 177 | subval, subidx = [], [] |
| 178 | for sg in subg: |
| 179 | subx = F.asnumpy(sg.edata["x"]) |
| 180 | ai = np.argsort(subx[:, -1:].flatten()) |
| 181 | if descending: |
| 182 | ai = np.ascontiguousarray(ai[::-1]) |
| 183 | subx = np.expand_dims(subx[ai[:k]], 0) |
| 184 | subval.append(F.tensor(subx)) |
nothing calls this directly
no test coverage detected