Verify the node/edge features of the partitioned graph with the original graph Parameters: ----------- g : DGL Graph Object of the original graph gpb : global partition book created for the partitioned graph object node_feats : dictionary with key, va
(
g, gpb, part, node_feats, edge_feats, orig_nids, orig_eids
)
| 77 | |
| 78 | |
| 79 | def verify_graph_feats( |
| 80 | g, gpb, part, node_feats, edge_feats, orig_nids, orig_eids |
| 81 | ): |
| 82 | """Verify the node/edge features of the partitioned graph with |
| 83 | the original graph |
| 84 | |
| 85 | Parameters: |
| 86 | ----------- |
| 87 | g : DGL Graph Object |
| 88 | of the original graph |
| 89 | gpb : global partition book |
| 90 | created for the partitioned graph object |
| 91 | node_feats : dictionary |
| 92 | with key, value pairs as node-types and features as numpy arrays |
| 93 | edge_feats : dictionary |
| 94 | with key, value pairs as edge-types and features as numpy arrays |
| 95 | orig_nids : dictionary |
| 96 | with key, value pairs as node-types and (global) nids from the |
| 97 | original graph |
| 98 | orig_eids : dictionary |
| 99 | with key, value pairs as edge-types and (global) eids from the |
| 100 | original graph |
| 101 | """ |
| 102 | for ntype in g.ntypes: |
| 103 | ntype_id = g.get_ntype_id(ntype) |
| 104 | inner_node_mask = _get_inner_node_mask(part, ntype_id) |
| 105 | inner_nids = part.ndata[dgl.NID][inner_node_mask] |
| 106 | ntype_ids, inner_type_nids = gpb.map_to_per_ntype(inner_nids) |
| 107 | partid = gpb.nid2partid(inner_type_nids, ntype) |
| 108 | assert np.all(ntype_ids.numpy() == ntype_id) |
| 109 | assert np.all(partid.numpy() == gpb.partid) |
| 110 | |
| 111 | orig_id = orig_nids[ntype][inner_type_nids] |
| 112 | local_nids = gpb.nid2localnid(inner_type_nids, gpb.partid, ntype) |
| 113 | |
| 114 | for name in g.nodes[ntype].data: |
| 115 | if name in [dgl.NID, "inner_node"]: |
| 116 | continue |
| 117 | true_feats = g.nodes[ntype].data[name][orig_id] |
| 118 | ndata = node_feats[ntype + "/" + name][local_nids] |
| 119 | assert np.array_equal(ndata.numpy(), true_feats.numpy()) |
| 120 | |
| 121 | for etype in g.canonical_etypes: |
| 122 | etype_id = g.get_etype_id(etype) |
| 123 | inner_edge_mask = _get_inner_edge_mask(part, etype_id) |
| 124 | inner_eids = part.edata[dgl.EID][inner_edge_mask] |
| 125 | etype_ids, inner_type_eids = gpb.map_to_per_etype(inner_eids) |
| 126 | partid = gpb.eid2partid(inner_type_eids, etype) |
| 127 | assert np.all(etype_ids.numpy() == etype_id) |
| 128 | assert np.all(partid.numpy() == gpb.partid) |
| 129 | |
| 130 | orig_id = orig_eids[_etype_tuple_to_str(etype)][inner_type_eids] |
| 131 | local_eids = gpb.eid2localeid(inner_type_eids, gpb.partid, etype) |
| 132 | |
| 133 | for name in g.edges[etype].data: |
| 134 | if name in [dgl.EID, "inner_edge"]: |
| 135 | continue |
| 136 | true_feats = g.edges[etype].data[name][orig_id] |