Node frontiers generator using breadth-first search. Parameters ---------- graph : DGLGraph The graph object. source : list, tensor of nodes Source nodes. reverse : bool, default False If True, traverse following the in-edge direction. Returns --
(graph, source, reverse=False)
| 15 | |
| 16 | |
| 17 | def bfs_nodes_generator(graph, source, reverse=False): |
| 18 | """Node frontiers generator using breadth-first search. |
| 19 | |
| 20 | Parameters |
| 21 | ---------- |
| 22 | graph : DGLGraph |
| 23 | The graph object. |
| 24 | source : list, tensor of nodes |
| 25 | Source nodes. |
| 26 | reverse : bool, default False |
| 27 | If True, traverse following the in-edge direction. |
| 28 | |
| 29 | Returns |
| 30 | ------- |
| 31 | list of node frontiers |
| 32 | Each node frontier is a list or tensor of node ids. |
| 33 | |
| 34 | Examples |
| 35 | -------- |
| 36 | Given a graph (directed, edges from small node id to large): |
| 37 | :: |
| 38 | |
| 39 | 2 - 4 |
| 40 | / \\ |
| 41 | 0 - 1 - 3 - 5 |
| 42 | |
| 43 | >>> g = dgl.graph(([0, 1, 1, 2, 2, 3], [1, 2, 3, 3, 4, 5])) |
| 44 | >>> list(dgl.bfs_nodes_generator(g, 0)) |
| 45 | [tensor([0]), tensor([1]), tensor([2, 3]), tensor([4, 5])] |
| 46 | """ |
| 47 | assert isinstance( |
| 48 | graph, DGLGraph |
| 49 | ), "DGLHeteroGraph is merged with DGLGraph, Please use DGLGraph" |
| 50 | assert ( |
| 51 | len(graph.canonical_etypes) == 1 |
| 52 | ), "bfs_nodes_generator only support homogeneous graph" |
| 53 | # Workaround before support for GPU graph |
| 54 | gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu())) |
| 55 | source = utils.toindex(source, dtype=graph._idtype_str) |
| 56 | ret = _CAPI_DGLBFSNodes_v2(gidx, source.todgltensor(), reverse) |
| 57 | all_nodes = utils.toindex(ret(0), dtype=graph._idtype_str).tousertensor() |
| 58 | # TODO(minjie): how to support directly creating python list |
| 59 | sections = utils.toindex(ret(1)).tonumpy().tolist() |
| 60 | node_frontiers = F.split(all_nodes, sections, dim=0) |
| 61 | return node_frontiers |
| 62 | |
| 63 | |
| 64 | def bfs_edges_generator(graph, source, reverse=False): |
nothing calls this directly
no test coverage detected