Reshuffle node ids and edge IDs of a graph. This function reshuffles nodes and edges in a graph so that all nodes/edges of the same type have contiguous IDs. If a graph is partitioned and nodes are assigned to different partitions, all nodes/edges in a partition should get contiguou
(g, node_part=None)
| 58 | |
| 59 | |
| 60 | def reshuffle_graph(g, node_part=None): |
| 61 | """Reshuffle node ids and edge IDs of a graph. |
| 62 | |
| 63 | This function reshuffles nodes and edges in a graph so that all nodes/edges of the same type |
| 64 | have contiguous IDs. If a graph is partitioned and nodes are assigned to different partitions, |
| 65 | all nodes/edges in a partition should |
| 66 | get contiguous IDs; within a partition, all nodes/edges of the same type have contigous IDs. |
| 67 | |
| 68 | Parameters |
| 69 | ---------- |
| 70 | g : DGLGraph |
| 71 | The input graph. |
| 72 | node_part : Tensor |
| 73 | This is a vector whose length is the same as the number of nodes in the input graph. |
| 74 | Each element indicates the partition ID the corresponding node is assigned to. |
| 75 | |
| 76 | Returns |
| 77 | ------- |
| 78 | (DGLGraph, Tensor) |
| 79 | The graph whose nodes and edges are reshuffled. |
| 80 | The 1D tensor that indicates the partition IDs of the nodes in the reshuffled graph. |
| 81 | """ |
| 82 | # In this case, we don't need to reshuffle node IDs and edge IDs. |
| 83 | if node_part is None: |
| 84 | g.ndata["orig_id"] = F.arange(0, g.num_nodes()) |
| 85 | g.edata["orig_id"] = F.arange(0, g.num_edges()) |
| 86 | return g, None |
| 87 | |
| 88 | start = time.time() |
| 89 | if node_part is not None: |
| 90 | node_part = utils.toindex(node_part) |
| 91 | node_part = node_part.tousertensor() |
| 92 | if NTYPE in g.ndata: |
| 93 | is_hetero = len(F.unique(g.ndata[NTYPE])) > 1 |
| 94 | else: |
| 95 | is_hetero = False |
| 96 | if is_hetero: |
| 97 | num_node_types = F.max(g.ndata[NTYPE], 0) + 1 |
| 98 | if node_part is not None: |
| 99 | sorted_part, new2old_map = F.sort_1d( |
| 100 | node_part * num_node_types + g.ndata[NTYPE] |
| 101 | ) |
| 102 | else: |
| 103 | sorted_part, new2old_map = F.sort_1d(g.ndata[NTYPE]) |
| 104 | sorted_part = F.floor_div(sorted_part, num_node_types) |
| 105 | elif node_part is not None: |
| 106 | sorted_part, new2old_map = F.sort_1d(node_part) |
| 107 | else: |
| 108 | g.ndata["orig_id"] = g.ndata[NID] |
| 109 | g.edata["orig_id"] = g.edata[EID] |
| 110 | return g, None |
| 111 | |
| 112 | new_node_ids = np.zeros((g.num_nodes(),), dtype=np.int64) |
| 113 | new_node_ids[F.asnumpy(new2old_map)] = np.arange(0, g.num_nodes()) |
| 114 | # If the input graph is homogneous, we only need to create an empty array, so that |
| 115 | # _CAPI_DGLReassignEdges_Hetero knows how to handle it. |
| 116 | etype = ( |
| 117 | g.edata[ETYPE] |
no test coverage detected