def level_order(forest, roots): try: edges = bfs_edges_generator(forest, roots) _, leaves = forest.find_edges(edges[-1]) edges_back = bfs_edges_generator(forest, roots, reverse=True) yield from reversed(edges_back) yield from edges except: return None
def level_order(forest, roots): device = forest.device edges = list(bfs_edges_generator(forest, roots)) edges = [e.to(device) for e in edges] _, leaves = forest.find_edges(edges[-1]) edges_back = list(bfs_edges_generator(forest, roots, reverse=True)) edges_back = [e.to(device) for e in edges_back] yield from reversed(edges_back) yield from edges
def level_order(forest, roots): forest = tocpu(forest) edges = bfs_edges_generator(forest, roots) if len(edges) == 0: # no edges in the tree; do not perform loopy BP return _, leaves = forest.find_edges(edges[-1]) edges_back = bfs_edges_generator(forest, roots, reverse=True) yield from reversed(edges_back) yield from edges
def test_bfs(index_dtype, n=100): def _bfs_nx(g_nx, src): edges = nx.bfs_edges(g_nx, src) layers_nx = [set([src])] edges_nx = [] frontier = set() edge_frontier = set() for u, v in edges: if u in layers_nx[-1]: frontier.add(v) edge_frontier.add(g.edge_id(u, v)) else: layers_nx.append(frontier) edges_nx.append(edge_frontier) frontier = set([v]) edge_frontier = set([g.edge_id(u, v)]) # avoids empty successors if len(frontier) > 0 and len(edge_frontier) > 0: layers_nx.append(frontier) edges_nx.append(edge_frontier) return layers_nx, edges_nx g = dgl.DGLGraph() a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n)) g.from_scipy_sparse_matrix(a) if index_dtype == 'int32': g = dgl.graph(g.edges()).int() else: g = dgl.graph(g.edges()).long() g_nx = g.to_networkx() src = random.choice(range(n)) layers_nx, _ = _bfs_nx(g_nx, src) layers_dgl = dgl.bfs_nodes_generator(g, src) assert len(layers_dgl) == len(layers_nx) assert all(toset(x) == y for x, y in zip(layers_dgl, layers_nx)) g_nx = nx.random_tree(n, seed=42) g = dgl.DGLGraph() g.from_networkx(g_nx) if index_dtype == 'int32': g = dgl.graph(g.edges()).int() else: g = dgl.graph(g.edges()).long() src = 0 _, edges_nx = _bfs_nx(g_nx, src) edges_dgl = dgl.bfs_edges_generator(g, src) assert len(edges_dgl) == len(edges_nx) assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))