in hugegraph-ml/src/hugegraph_ml/data/hugegraph2dgl.py [0:0]
def _convert_graph_from_ogb(vertices, edges, feat_key, year_key, weight_key):
if len(vertices) == 0:
warnings.warn("This graph has no vertices", Warning)
return dgl.graph(())
vertex_ids = [v["id"] for v in vertices]
vertex_id_to_idx = {vertex_id: idx for idx, vertex_id in enumerate(vertex_ids)}
src_idx = [vertex_id_to_idx[e["outV"]] for e in edges]
dst_idx = [vertex_id_to_idx[e["inV"]] for e in edges]
graph_dgl = dgl.graph((src_idx, dst_idx))
if feat_key and feat_key in vertices[0]["properties"]:
node_feats = [
v["properties"][feat_key]
for v in vertices[0 : graph_dgl.number_of_nodes()]
]
graph_dgl.ndata["feat"] = torch.tensor(node_feats, dtype=torch.float32)
if year_key and year_key in edges[0]["properties"]:
year = [e["properties"][year_key] for e in edges]
graph_dgl.edata["year"] = torch.tensor(year, dtype=torch.int64)
if weight_key and weight_key in edges[0]["properties"]:
weight = [e["properties"][weight_key] for e in edges]
graph_dgl.edata["weight"] = torch.tensor(weight, dtype=torch.int64)
return graph_dgl, vertex_id_to_idx