Skip to content
Permalink
Branch: master
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
76 lines (69 sloc) 2.89 KB
/*!
* Copyright (c) 2019 by Contributors
* \file graph/nodeflow.cc
* \brief DGL NodeFlow related functions.
*/
#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>
#include <string.h>
#include "../c_api_common.h"
namespace dgl {
std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::string &fmt,
size_t layer0_size, size_t layer1_start,
size_t layer1_end, bool remap) {
CHECK_GE(layer1_start, layer0_size);
if (fmt == "csr") {
dgl_id_t first_vid = layer1_start - layer0_size;
ImmutableGraph::CSRArray arrs = graph.GetInCSRArray(layer1_start, layer1_end);
if (remap) {
dgl_id_t *indices_data = static_cast<dgl_id_t*>(arrs.indices->data);
dgl_id_t *eid_data = static_cast<dgl_id_t*>(arrs.id->data);
const size_t len = arrs.indices->shape[0];
dgl_id_t first_eid = eid_data[0];
for (size_t i = 0; i < len; i++) {
CHECK_GE(indices_data[i], first_vid);
indices_data[i] -= first_vid;
CHECK_GE(eid_data[i], first_eid);
eid_data[i] -= first_eid;
}
}
return std::vector<IdArray>{arrs.indptr, arrs.indices, arrs.id};
} else if (fmt == "coo") {
ImmutableGraph::CSR::Ptr csr = graph.GetInCSR();
int64_t nnz = csr->indptr[layer1_end] - csr->indptr[layer1_start];
IdArray idx = IdArray::Empty({2 * nnz}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
IdArray eid = IdArray::Empty({nnz}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
int64_t *idx_data = static_cast<int64_t*>(idx->data);
dgl_id_t *eid_data = static_cast<dgl_id_t*>(eid->data);
size_t num_edges = 0;
for (size_t i = layer1_start; i < layer1_end; i++) {
for (int64_t j = csr->indptr[i]; j < csr->indptr[i + 1]; j++) {
// These nodes are all in a layer. We need to remap them to the node id
// local to the layer.
idx_data[num_edges] = remap ? i - layer1_start : i;
num_edges++;
}
}
CHECK_EQ(num_edges, nnz);
if (remap) {
size_t edge_start = csr->indptr[layer1_start];
dgl_id_t first_eid = csr->edge_ids[edge_start];
dgl_id_t first_vid = layer1_start - layer0_size;
for (int64_t i = 0; i < nnz; i++) {
CHECK_GE(csr->indices[edge_start + i], first_vid);
idx_data[nnz + i] = csr->indices[edge_start + i] - first_vid;
eid_data[i] = csr->edge_ids[edge_start + i] - first_eid;
}
} else {
std::copy(csr->indices.begin() + csr->indptr[layer1_start],
csr->indices.begin() + csr->indptr[layer1_end], idx_data + nnz);
std::copy(csr->edge_ids.begin() + csr->indptr[layer1_start],
csr->edge_ids.begin() + csr->indptr[layer1_end], eid_data);
}
return std::vector<IdArray>{idx, eid};
} else {
LOG(FATAL) << "unsupported adjacency matrix format";
return std::vector<IdArray>();
}
}
} // namespace dgl
You can’t perform that action at this time.
You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session.