Skip to content

GitLab

Unverified Commit 9b4d6079 authored by Minjie Wang's avatar Minjie Wang Committed by GitHub
Browse files

[Hetero] New syntax (#824) 

* WIP. remove graph arg in NodeBatch and EdgeBatch

* refactor: use graph adapter for scheduler

* WIP: recv

* draft impl

* stuck at bipartite

* bipartite->unitgraph; support dsttype == srctype

* pass test_query

* pass test_query

* pass test_view

* test apply

* pass udf message passing tests

* pass quan's test using builtins

* WIP: wildcard slicing

* new construct methods

* broken

* good

* add stack cross reducer

* fix bug; fix mx

* fix bug in csrmm2 when the CSR is not square

* lint

* removed FlattenedHeteroGraph class

* WIP

* prop nodes, prop edges, filter nodes/edges

* add DGLGraph tests to heterograph. Fix several bugs

* finish nx<->hetero graph conversion

* create bipartite from nx

* more spec on hetero/homo conversion

* silly fixes

* check node and edge types

* repr

* to api

* adj APIs

* inc

* fix some lints and bugs

* fix some lints

* hetero/homo conversion

* fix flatten test

* more spec in hetero_from_homo and test

* flatten using concat names

* WIP: creators

* rewrite hetero_from_homo in a more efficient way

* remove useless variables

* fix lint

* subgraphs and typed subgraphs

* lint & removed heterosubgraph class

* lint x2

* disable heterograph mutation test

* docstring update

* add edge id for nx graph test

* fix mx unittests

* fix bug

* try fix

* fix unittest when cross_reducer is stack

* fix ci

* fix nx bipartite bug; docstring

* fix scipy creation bug

* lint

* fix bug when converting heterograph from homograph

* fix bug in hetero_from_homo about ntype order

* trailing white

* docstring fixes for add_foo and data views

* docstring for relation slice

* to_hetero and to_homo with feature support

* lint

* lint

* DGLGraph compatibility

* incidence matrix & docstring fixes

* example string fixes

* feature in hetero_from_relations

* deduplication of edge types in to_hetero

* fix lint

* fix
parent ddb5d804
Hide whitespace changes
Inline Side-by-side
Showing with 2084 additions and 322 deletions
src/graph/bipartite.cc → src/graph/unit_graph.cc
View file @ 9b4d6079
/*!
* Copyright (c) 2019 by Contributors
* \file graph/bipartite.cc
* \brief Bipartite graph implementation
* \file graph/unit_graph.cc
* \brief UnitGraph graph implementation
*/
#include <dgl/array.h>
#include <dgl/lazy.h>
#include <dgl/immutable_graph.h>
#include <dgl/base_heterograph.h>
#include "./bipartite.h"
#include "./unit_graph.h"
#include "../c_api_common.h"
namespace dgl {
namespace {
// create metagraph of one node type
inline GraphPtr CreateUnitGraphMetaGraph1() {
// a self-loop edge 0->0
std::vector<int64_t> row_vec(1, 0);
std::vector<int64_t> col_vec(1, 0);
IdArray row = aten::VecToIdArray(row_vec);
IdArray col = aten::VecToIdArray(col_vec);
GraphPtr g = ImmutableGraph::CreateFromCOO(1, row, col);
return g;
}
inline GraphPtr CreateBipartiteMetaGraph() {
std::vector<int64_t> row_vec(1, Bipartite::kSrcVType);
std::vector<int64_t> col_vec(1, Bipartite::kDstVType);
// create metagraph of two node types
inline GraphPtr CreateUnitGraphMetaGraph2() {
// an edge 0->1
std::vector<int64_t> row_vec(1, 0);
std::vector<int64_t> col_vec(1, 1);
IdArray row = aten::VecToIdArray(row_vec);
IdArray col = aten::VecToIdArray(col_vec);
GraphPtr g = ImmutableGraph::CreateFromCOO(2, row, col);
return g;
}
const GraphPtr kBipartiteMetaGraph = CreateBipartiteMetaGraph();
inline GraphPtr CreateUnitGraphMetaGraph(int num_vtypes) {
static GraphPtr mg1 = CreateUnitGraphMetaGraph1();
static GraphPtr mg2 = CreateUnitGraphMetaGraph2();
if (num_vtypes == 1)
return mg1;
else if (num_vtypes == 2)
return mg2;
else
LOG(FATAL) << "Invalid number of vertex types. Must be 1 or 2.";
return {};
}
}; // namespace
......@@ -33,48 +56,55 @@ const GraphPtr kBipartiteMetaGraph = CreateBipartiteMetaGraph();
//
//////////////////////////////////////////////////////////
class Bipartite::COO : public BaseHeteroGraph {
class UnitGraph::COO : public BaseHeteroGraph {
public:
COO(int64_t num_src, int64_t num_dst, IdArray src, IdArray dst)
: BaseHeteroGraph(kBipartiteMetaGraph) {
COO(GraphPtr metagraph, int64_t num_src, int64_t num_dst, IdArray src, IdArray dst)
: BaseHeteroGraph(metagraph) {
adj_ = aten::COOMatrix{num_src, num_dst, src, dst};
}
COO(int64_t num_src, int64_t num_dst, IdArray src, IdArray dst, bool is_multigraph)
: BaseHeteroGraph(kBipartiteMetaGraph),
COO(GraphPtr metagraph, int64_t num_src, int64_t num_dst,
IdArray src, IdArray dst, bool is_multigraph)
: BaseHeteroGraph(metagraph),
is_multigraph_(is_multigraph) {
adj_ = aten::COOMatrix{num_src, num_dst, src, dst};
}
explicit COO(const aten::COOMatrix& coo) : BaseHeteroGraph(kBipartiteMetaGraph), adj_(coo) {}
explicit COO(GraphPtr metagraph, const aten::COOMatrix& coo)
: BaseHeteroGraph(metagraph), adj_(coo) {}
uint64_t NumVertexTypes() const override {
return 2;
inline dgl_type_t SrcType() const {
return 0;
}
uint64_t NumEdgeTypes() const override {
return 1;
inline dgl_type_t DstType() const {
return NumVertexTypes() == 1? 0 : 1;
}
inline dgl_type_t EdgeType() const {
return 0;
}
HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
LOG(FATAL) << "The method shouldn't be called for Bipartite graph. "
LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
<< "The relation graph is simply this graph itself.";
return {};
}
void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void Clear() override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
DLContext Context() const override {
......@@ -96,9 +126,9 @@ class Bipartite::COO : public BaseHeteroGraph {
}
uint64_t NumVertices(dgl_type_t vtype) const override {
if (vtype == Bipartite::kSrcVType) {
if (vtype == SrcType()) {
return adj_.num_rows;
} else if (vtype == Bipartite::kDstVType) {
} else if (vtype == DstType()) {
return adj_.num_cols;
} else {
LOG(FATAL) << "Invalid vertex type: " << vtype;
......@@ -258,7 +288,7 @@ class Bipartite::COO : public BaseHeteroGraph {
const auto new_nsrc = subg.induced_vertices[0]->shape[0];
const auto new_ndst = subg.induced_vertices[1]->shape[0];
subg.graph = std::make_shared<COO>(
new_nsrc, new_ndst, new_src, new_dst);
meta_graph(), new_nsrc, new_ndst, new_src, new_dst);
subg.induced_edges = eids;
} else {
IdArray new_src = aten::IndexSelect(adj_.row, eids[0]);
......@@ -266,7 +296,7 @@ class Bipartite::COO : public BaseHeteroGraph {
subg.induced_vertices.emplace_back(aten::Range(0, NumVertices(0), NumBits(), Context()));
subg.induced_vertices.emplace_back(aten::Range(0, NumVertices(1), NumBits(), Context()));
subg.graph = std::make_shared<COO>(
NumVertices(0), NumVertices(1), new_src, new_dst);
meta_graph(), NumVertices(0), NumVertices(1), new_src, new_dst);
subg.induced_edges = eids;
}
return subg;
......@@ -291,50 +321,55 @@ class Bipartite::COO : public BaseHeteroGraph {
//////////////////////////////////////////////////////////
/*! \brief CSR graph */
class Bipartite::CSR : public BaseHeteroGraph {
class UnitGraph::CSR : public BaseHeteroGraph {
public:
CSR(int64_t num_src, int64_t num_dst,
CSR(GraphPtr metagraph, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids)
: BaseHeteroGraph(kBipartiteMetaGraph) {
: BaseHeteroGraph(metagraph) {
adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
}
CSR(int64_t num_src, int64_t num_dst,
CSR(GraphPtr metagraph, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids, bool is_multigraph)
: BaseHeteroGraph(kBipartiteMetaGraph),
is_multigraph_(is_multigraph) {
: BaseHeteroGraph(metagraph), is_multigraph_(is_multigraph) {
adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
}
explicit CSR(const aten::CSRMatrix& csr) : BaseHeteroGraph(kBipartiteMetaGraph), adj_(csr) {}
explicit CSR(GraphPtr metagraph, const aten::CSRMatrix& csr)
: BaseHeteroGraph(metagraph), adj_(csr) {}
uint64_t NumVertexTypes() const override {
return 2;
inline dgl_type_t SrcType() const {
return 0;
}
uint64_t NumEdgeTypes() const override {
return 1;
inline dgl_type_t DstType() const {
return NumVertexTypes() == 1? 0 : 1;
}
inline dgl_type_t EdgeType() const {
return 0;
}
HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
LOG(FATAL) << "The method shouldn't be called for Bipartite graph. "
LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
<< "The relation graph is simply this graph itself.";
return {};
}
void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void Clear() override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
DLContext Context() const override {
......@@ -350,6 +385,7 @@ class Bipartite::CSR : public BaseHeteroGraph {
return *this;
} else {
CSR ret(
meta_graph_,
adj_.num_rows, adj_.num_cols,
aten::AsNumBits(adj_.indptr, bits),
aten::AsNumBits(adj_.indices, bits),
......@@ -364,6 +400,7 @@ class Bipartite::CSR : public BaseHeteroGraph {
return *this;
} else {
CSR ret(
meta_graph_,
adj_.num_rows, adj_.num_cols,
adj_.indptr.CopyTo(ctx),
adj_.indices.CopyTo(ctx),
......@@ -384,9 +421,9 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
uint64_t NumVertices(dgl_type_t vtype) const override {
if (vtype == Bipartite::kSrcVType) {
if (vtype == SrcType()) {
return adj_.num_rows;
} else if (vtype == Bipartite::kDstVType) {
} else if (vtype == DstType()) {
return adj_.num_cols;
} else {
LOG(FATAL) << "Invalid vertex type: " << vtype;
......@@ -408,8 +445,8 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(1, dst)) << "Invalid dst vertex id: " << dst;
CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
return aten::CSRIsNonZero(adj_, src, dst);
}
......@@ -425,13 +462,13 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
return aten::CSRGetRowColumnIndices(adj_, src);
}
IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
CHECK(HasVertex(0, src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(1, dst)) << "Invalid dst vertex id: " << dst;
CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
return aten::CSRGetData(adj_, src, dst);
}
......@@ -463,7 +500,7 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
CHECK(HasVertex(0, vid)) << "Invalid src vertex id: " << vid;
CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
IdArray ret_dst = aten::CSRGetRowColumnIndices(adj_, vid);
IdArray ret_eid = aten::CSRGetRowData(adj_, vid);
IdArray ret_src = aten::Full(vid, ret_dst->shape[0], NumBits(), ret_dst->ctx);
......@@ -499,7 +536,7 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
CHECK(HasVertex(0, vid)) << "Invalid src vertex id: " << vid;
CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
return aten::CSRGetRowNNZ(adj_, vid);
}
......@@ -545,13 +582,14 @@ class Bipartite::CSR : public BaseHeteroGraph {
}
HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
CHECK_EQ(vids.size(), 2) << "Number of vertex types mismatch";
CHECK(aten::IsValidIdArray(vids[0])) << "Invalid vertex id array.";
CHECK(aten::IsValidIdArray(vids[1])) << "Invalid vertex id array.";
CHECK_EQ(vids.size(), NumVertexTypes()) << "Number of vertex types mismatch";
auto srcvids = vids[SrcType()], dstvids = vids[DstType()];
CHECK(aten::IsValidIdArray(srcvids)) << "Invalid vertex id array.";
CHECK(aten::IsValidIdArray(dstvids)) << "Invalid vertex id array.";
HeteroSubgraph subg;
const auto& submat = aten::CSRSliceMatrix(adj_, vids[0], vids[1]);
const auto& submat = aten::CSRSliceMatrix(adj_, srcvids, dstvids);
IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
subg.graph = std::make_shared<CSR>(submat.num_rows, submat.num_cols,
subg.graph = std::make_shared<CSR>(meta_graph(), submat.num_rows, submat.num_cols,
submat.indptr, submat.indices, sub_eids);
subg.induced_vertices = vids;
subg.induced_edges.emplace_back(submat.data);
......@@ -578,40 +616,54 @@ class Bipartite::CSR : public BaseHeteroGraph {
//////////////////////////////////////////////////////////
//
// bipartite graph implementation
// unit graph implementation
//
//////////////////////////////////////////////////////////
DLContext Bipartite::Context() const {
DLContext UnitGraph::Context() const {
return GetAny()->Context();
}
uint8_t Bipartite::NumBits() const {
uint8_t UnitGraph::NumBits() const {
return GetAny()->NumBits();
}
bool Bipartite::IsMultigraph() const {
bool UnitGraph::IsMultigraph() const {
return GetAny()->IsMultigraph();
}
uint64_t Bipartite::NumVertices(dgl_type_t vtype) const {
return GetAny()->NumVertices(vtype);
uint64_t UnitGraph::NumVertices(dgl_type_t vtype) const {
if (in_csr_) {
vtype = (vtype == SrcType()) ? DstType() : SrcType();
return in_csr_->NumVertices(vtype);
} else if (out_csr_) {
return out_csr_->NumVertices(vtype);
} else {
return GetCOO()->NumVertices(vtype);
}
}
uint64_t Bipartite::NumEdges(dgl_type_t etype) const {
uint64_t UnitGraph::NumEdges(dgl_type_t etype) const {
return GetAny()->NumEdges(etype);
}
bool Bipartite::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
return GetAny()->HasVertex(vtype, vid);
bool UnitGraph::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
if (in_csr_) {
vtype = (vtype == SrcType()) ? DstType() : SrcType();
return in_csr_->HasVertex(vtype, vid);
} else if (out_csr_) {
return out_csr_->HasVertex(vtype, vid);
} else {
return GetCOO()->HasVertex(vtype, vid);
}
}
BoolArray Bipartite::HasVertices(dgl_type_t vtype, IdArray vids) const {
BoolArray UnitGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
CHECK(aten::IsValidIdArray(vids)) << "Invalid id array input";
return aten::LT(vids, NumVertices(vtype));
}
bool Bipartite::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
bool UnitGraph::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
if (in_csr_) {
return in_csr_->HasEdgeBetween(etype, dst, src);
} else {
......@@ -619,7 +671,7 @@ bool Bipartite::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) con
}
}
BoolArray Bipartite::HasEdgesBetween(
BoolArray UnitGraph::HasEdgesBetween(
dgl_type_t etype, IdArray src, IdArray dst) const {
if (in_csr_) {
return in_csr_->HasEdgesBetween(etype, dst, src);
......@@ -628,15 +680,15 @@ BoolArray Bipartite::HasEdgesBetween(
}
}
IdArray Bipartite::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
IdArray UnitGraph::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
return GetInCSR()->Successors(etype, dst);
}
IdArray Bipartite::Successors(dgl_type_t etype, dgl_id_t src) const {
IdArray UnitGraph::Successors(dgl_type_t etype, dgl_id_t src) const {
return GetOutCSR()->Successors(etype, src);
}
IdArray Bipartite::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
IdArray UnitGraph::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
if (in_csr_) {
return in_csr_->EdgeId(etype, dst, src);
} else {
......@@ -644,7 +696,7 @@ IdArray Bipartite::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
}
}
EdgeArray Bipartite::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
EdgeArray UnitGraph::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
if (in_csr_) {
EdgeArray edges = in_csr_->EdgeIds(etype, dst, src);
return EdgeArray{edges.dst, edges.src, edges.id};
......@@ -653,33 +705,33 @@ EdgeArray Bipartite::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
}
}
std::pair<dgl_id_t, dgl_id_t> Bipartite::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
std::pair<dgl_id_t, dgl_id_t> UnitGraph::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
return GetCOO()->FindEdge(etype, eid);
}
EdgeArray Bipartite::FindEdges(dgl_type_t etype, IdArray eids) const {
EdgeArray UnitGraph::FindEdges(dgl_type_t etype, IdArray eids) const {
return GetCOO()->FindEdges(etype, eids);
}
EdgeArray Bipartite::InEdges(dgl_type_t etype, dgl_id_t vid) const {
EdgeArray UnitGraph::InEdges(dgl_type_t etype, dgl_id_t vid) const {
const EdgeArray& ret = GetInCSR()->OutEdges(etype, vid);
return {ret.dst, ret.src, ret.id};
}
EdgeArray Bipartite::InEdges(dgl_type_t etype, IdArray vids) const {
EdgeArray UnitGraph::InEdges(dgl_type_t etype, IdArray vids) const {
const EdgeArray& ret = GetInCSR()->OutEdges(etype, vids);
return {ret.dst, ret.src, ret.id};
}
EdgeArray Bipartite::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
return GetOutCSR()->OutEdges(etype, vid);
}
EdgeArray Bipartite::OutEdges(dgl_type_t etype, IdArray vids) const {
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, IdArray vids) const {
return GetOutCSR()->OutEdges(etype, vids);
}
EdgeArray Bipartite::Edges(dgl_type_t etype, const std::string &order) const {
EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
if (order.empty()) {
// arbitrary order
if (in_csr_) {
......@@ -701,39 +753,39 @@ EdgeArray Bipartite::Edges(dgl_type_t etype, const std::string &order) const {
return {};
}
uint64_t Bipartite::InDegree(dgl_type_t etype, dgl_id_t vid) const {
uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
return GetInCSR()->OutDegree(etype, vid);
}
DegreeArray Bipartite::InDegrees(dgl_type_t etype, IdArray vids) const {
DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
return GetInCSR()->OutDegrees(etype, vids);
}
uint64_t Bipartite::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
uint64_t UnitGraph::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
return GetOutCSR()->OutDegree(etype, vid);
}
DegreeArray Bipartite::OutDegrees(dgl_type_t etype, IdArray vids) const {
DegreeArray UnitGraph::OutDegrees(dgl_type_t etype, IdArray vids) const {
return GetOutCSR()->OutDegrees(etype, vids);
}
DGLIdIters Bipartite::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
DGLIdIters UnitGraph::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
return GetOutCSR()->SuccVec(etype, vid);
}
DGLIdIters Bipartite::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
DGLIdIters UnitGraph::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
return GetOutCSR()->OutEdgeVec(etype, vid);
}
DGLIdIters Bipartite::PredVec(dgl_type_t etype, dgl_id_t vid) const {
DGLIdIters UnitGraph::PredVec(dgl_type_t etype, dgl_id_t vid) const {
return GetInCSR()->SuccVec(etype, vid);
}
DGLIdIters Bipartite::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
DGLIdIters UnitGraph::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
return GetInCSR()->OutEdgeVec(etype, vid);
}
std::vector<IdArray> Bipartite::GetAdj(
std::vector<IdArray> UnitGraph::GetAdj(
dgl_type_t etype, bool transpose, const std::string &fmt) const {
// TODO(minjie): Our current semantics of adjacency matrix is row for dst nodes and col for
// src nodes. Therefore, we need to flip the transpose flag. For example, transpose=False
......@@ -753,43 +805,50 @@ std::vector<IdArray> Bipartite::GetAdj(
}
}
HeteroSubgraph Bipartite::VertexSubgraph(const std::vector<IdArray>& vids) const {
HeteroSubgraph UnitGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
// We prefer to generate a subgraph from out-csr.
auto sg = GetOutCSR()->VertexSubgraph(vids);
CSRPtr subcsr = std::dynamic_pointer_cast<CSR>(sg.graph);
HeteroSubgraph ret;
ret.graph = HeteroGraphPtr(new Bipartite(nullptr, subcsr, nullptr));
ret.graph = HeteroGraphPtr(new UnitGraph(meta_graph(), nullptr, subcsr, nullptr));
ret.induced_vertices = std::move(sg.induced_vertices);
ret.induced_edges = std::move(sg.induced_edges);
return ret;
}
HeteroSubgraph Bipartite::EdgeSubgraph(
HeteroSubgraph UnitGraph::EdgeSubgraph(
const std::vector<IdArray>& eids, bool preserve_nodes) const {
auto sg = GetCOO()->EdgeSubgraph(eids, preserve_nodes);
COOPtr subcoo = std::dynamic_pointer_cast<COO>(sg.graph);
HeteroSubgraph ret;
ret.graph = HeteroGraphPtr(new Bipartite(nullptr, nullptr, subcoo));
ret.graph = HeteroGraphPtr(new UnitGraph(meta_graph(), nullptr, nullptr, subcoo));
ret.induced_vertices = std::move(sg.induced_vertices);
ret.induced_edges = std::move(sg.induced_edges);
return ret;
}
HeteroGraphPtr Bipartite::CreateFromCOO(
int64_t num_src, int64_t num_dst,
IdArray row, IdArray col) {
COOPtr coo(new COO(num_src, num_dst, row, col));
return HeteroGraphPtr(new Bipartite(nullptr, nullptr, coo));
HeteroGraphPtr UnitGraph::CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray row, IdArray col) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(num_src, num_dst);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
COOPtr coo(new COO(mg, num_src, num_dst, row, col));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, nullptr, coo));
}
HeteroGraphPtr Bipartite::CreateFromCSR(
int64_t num_src, int64_t num_dst,
HeteroGraphPtr UnitGraph::CreateFromCSR(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids) {
CSRPtr csr(new CSR(num_src, num_dst, indptr, indices, edge_ids));
return HeteroGraphPtr(new Bipartite(nullptr, csr, nullptr));
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(num_src, num_dst);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
CSRPtr csr(new CSR(mg, num_src, num_dst, indptr, indices, edge_ids));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr));
}
HeteroGraphPtr Bipartite::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
HeteroGraphPtr UnitGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
if (g->NumBits() == bits) {
return g;
} else {
......@@ -797,15 +856,15 @@ HeteroGraphPtr Bipartite::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
// we make sure that this graph (on CPU) has materialized CSR,
// and then copy them to other context (usually GPU). This should
// be fixed later.
auto bg = std::dynamic_pointer_cast<Bipartite>(g);
auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
CHECK_NOTNULL(bg);
CSRPtr new_incsr = CSRPtr(new CSR(bg->GetInCSR()->AsNumBits(bits)));
CSRPtr new_outcsr = CSRPtr(new CSR(bg->GetOutCSR()->AsNumBits(bits)));
return HeteroGraphPtr(new Bipartite(new_incsr, new_outcsr, nullptr));
return HeteroGraphPtr(new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, nullptr));
}
}
HeteroGraphPtr Bipartite::CopyTo(HeteroGraphPtr g, const DLContext& ctx) {
HeteroGraphPtr UnitGraph::CopyTo(HeteroGraphPtr g, const DLContext& ctx) {
if (ctx == g->Context()) {
return g;
}
......@@ -813,78 +872,79 @@ HeteroGraphPtr Bipartite::CopyTo(HeteroGraphPtr g, const DLContext& ctx) {
// we make sure that this graph (on CPU) has materialized CSR,
// and then copy them to other context (usually GPU). This should
// be fixed later.
auto bg = std::dynamic_pointer_cast<Bipartite>(g);
auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
CHECK_NOTNULL(bg);
CSRPtr new_incsr = CSRPtr(new CSR(bg->GetInCSR()->CopyTo(ctx)));
CSRPtr new_outcsr = CSRPtr(new CSR(bg->GetOutCSR()->CopyTo(ctx)));
return HeteroGraphPtr(new Bipartite(new_incsr, new_outcsr, nullptr));
return HeteroGraphPtr(new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, nullptr));
}
Bipartite::Bipartite(CSRPtr in_csr, CSRPtr out_csr, COOPtr coo)
: BaseHeteroGraph(kBipartiteMetaGraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
UnitGraph::UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo)
: BaseHeteroGraph(metagraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
CHECK(GetAny()) << "At least one graph structure should exist.";
}
Bipartite::CSRPtr Bipartite::GetInCSR() const {
UnitGraph::CSRPtr UnitGraph::GetInCSR() const {
if (!in_csr_) {
if (out_csr_) {
const auto& newadj = aten::CSRTranspose(out_csr_->adj());
const_cast<Bipartite*>(this)->in_csr_ = std::make_shared<CSR>(newadj);
const_cast<UnitGraph*>(this)->in_csr_ = std::make_shared<CSR>(meta_graph(), newadj);
} else {
CHECK(coo_) << "None of CSR, COO exist";
const auto& adj = coo_->adj();
const auto& newadj = aten::COOToCSR(
aten::COOMatrix{adj.num_cols, adj.num_rows, adj.col, adj.row});
const_cast<Bipartite*>(this)->in_csr_ = std::make_shared<CSR>(newadj);
const_cast<UnitGraph*>(this)->in_csr_ = std::make_shared<CSR>(meta_graph(), newadj);
}
}
return in_csr_;
}
/* !\brief Return out csr. If not exist, transpose the other one.*/
Bipartite::CSRPtr Bipartite::GetOutCSR() const {
UnitGraph::CSRPtr UnitGraph::GetOutCSR() const {
if (!out_csr_) {
if (in_csr_) {
const auto& newadj = aten::CSRTranspose(in_csr_->adj());
const_cast<Bipartite*>(this)->out_csr_ = std::make_shared<CSR>(newadj);
const_cast<UnitGraph*>(this)->out_csr_ = std::make_shared<CSR>(meta_graph(), newadj);
} else {
CHECK(coo_) << "None of CSR, COO exist";
const auto& newadj = aten::COOToCSR(coo_->adj());
const_cast<Bipartite*>(this)->out_csr_ = std::make_shared<CSR>(newadj);
const_cast<UnitGraph*>(this)->out_csr_ = std::make_shared<CSR>(meta_graph(), newadj);
}
}
return out_csr_;
}
/* !\brief Return coo. If not exist, create from csr.*/
Bipartite::COOPtr Bipartite::GetCOO() const {
UnitGraph::COOPtr UnitGraph::GetCOO() const {
if (!coo_) {
if (in_csr_) {
const auto& newadj = aten::CSRToCOO(in_csr_->adj(), true);
const_cast<Bipartite*>(this)->coo_ = std::make_shared<COO>(
const_cast<UnitGraph*>(this)->coo_ = std::make_shared<COO>(
meta_graph(),
aten::COOMatrix{newadj.num_cols, newadj.num_rows, newadj.col, newadj.row});
} else {
CHECK(out_csr_) << "Both CSR are missing.";
const auto& newadj = aten::CSRToCOO(out_csr_->adj(), true);
const_cast<Bipartite*>(this)->coo_ = std::make_shared<COO>(newadj);
const_cast<UnitGraph*>(this)->coo_ = std::make_shared<COO>(meta_graph(), newadj);
}
}
return coo_;
}
aten::CSRMatrix Bipartite::GetInCSRMatrix() const {
aten::CSRMatrix UnitGraph::GetInCSRMatrix() const {
return GetInCSR()->adj();
}
aten::CSRMatrix Bipartite::GetOutCSRMatrix() const {
aten::CSRMatrix UnitGraph::GetOutCSRMatrix() const {
return GetOutCSR()->adj();
}
aten::COOMatrix Bipartite::GetCOOMatrix() const {
aten::COOMatrix UnitGraph::GetCOOMatrix() const {
return GetCOO()->adj();
}
HeteroGraphPtr Bipartite::GetAny() const {
HeteroGraphPtr UnitGraph::GetAny() const {
if (in_csr_) {
return in_csr_;
} else if (out_csr_) {
......
src/graph/bipartite.h → src/graph/unit_graph.h
View file @ 9b4d6079
/*!
* Copyright (c) 2019 by Contributors
* \file graph/bipartite.h
* \brief Bipartite graph
* \file graph/unit_graph.h
* \brief UnitGraph graph
*/
#ifndef DGL_GRAPH_BIPARTITE_H_
#define DGL_GRAPH_BIPARTITE_H_
#ifndef DGL_GRAPH_UNIT_GRAPH_H_
#define DGL_GRAPH_UNIT_GRAPH_H_
#include <dgl/base_heterograph.h>
#include <dgl/lazy.h>
......@@ -19,55 +19,56 @@
namespace dgl {
/*!
* \brief Bipartite graph
* \brief UnitGraph graph
*
* Bipartite graph is a special type of heterograph which has two types
* of nodes: "Src" and "Dst". All the edges are from "Src" type nodes to
* "Dst" type nodes, so there is no edge among nodes of the same type.
* UnitGraph graph is a special type of heterograph which
* (1) Have two types of nodes: "Src" and "Dst". All the edges are
* from "Src" type nodes to "Dst" type nodes, so there is no edge among
* nodes of the same type. Thus, its metagraph has two nodes and one edge
* between them.
* (2) Have only one type of nodes and edges. Thus, its metagraph has one node
* and one self-loop edge.
*/
class Bipartite : public BaseHeteroGraph {
class UnitGraph : public BaseHeteroGraph {
public:
/*! \brief source node group type */
static constexpr dgl_type_t kSrcVType = 0;
/*! \brief destination node group type */
static constexpr dgl_type_t kDstVType = 1;
/*! \brief edge group type */
static constexpr dgl_type_t kEType = 0;
// internal data structure
class COO;
class CSR;
typedef std::shared_ptr<COO> COOPtr;
typedef std::shared_ptr<CSR> CSRPtr;
uint64_t NumVertexTypes() const override {
return 2;
inline dgl_type_t SrcType() const {
return 0;
}
inline dgl_type_t DstType() const {
return NumVertexTypes() == 1? 0 : 1;
}
uint64_t NumEdgeTypes() const override {
return 1;
inline dgl_type_t EdgeType() const {
return 0;
}
HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
LOG(FATAL) << "The method shouldn't be called for Bipartite graph. "
LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
<< "The relation graph is simply this graph itself.";
return {};
}
void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
void Clear() override {
LOG(FATAL) << "Bipartite graph is not mutable.";
LOG(FATAL) << "UnitGraph graph is not mutable.";
}
DLContext Context() const override;
......@@ -139,13 +140,14 @@ class Bipartite : public BaseHeteroGraph {
const std::vector<IdArray>& eids, bool preserve_nodes = false) const override;
// creators
/*! \brief Create a bipartite graph from COO arrays */
static HeteroGraphPtr CreateFromCOO(int64_t num_src, int64_t num_dst,
/*! \brief Create a graph from COO arrays */
static HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray row, IdArray col);
/*! \brief Create a bipartite graph from (out) CSR arrays */
/*! \brief Create a graph from (out) CSR arrays */
static HeteroGraphPtr CreateFromCSR(
int64_t num_src, int64_t num_dst,
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids);
/*! \brief Convert the graph to use the given number of bits for storage */
......@@ -173,7 +175,14 @@ class Bipartite : public BaseHeteroGraph {
aten::COOMatrix GetCOOMatrix() const;
private:
Bipartite(CSRPtr in_csr, CSRPtr out_csr, COOPtr coo);
/*!
* \brief constructor
* \param metagraph metagraph
* \param in_csr in edge csr
* \param out_csr out edge csr
* \param coo coo
*/
UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo);
/*! \return Return any existing format. */
HeteroGraphPtr GetAny() const;
......@@ -190,4 +199,4 @@ class Bipartite : public BaseHeteroGraph {
}; // namespace dgl
#endif // DGL_GRAPH_BIPARTITE_H_
#endif // DGL_GRAPH_UNIT_GRAPH_H_
src/kernel/binary_reduce.cc
View file @ 9b4d6079
......@@ -10,7 +10,7 @@
#include "./binary_reduce_impl_decl.h"
#include "./utils.h"
#include "../c_api_common.h"
#include "../graph/bipartite.h"
#include "../graph/unit_graph.h"
#include "./csr_interface.h"
using namespace dgl::runtime;
......@@ -243,9 +243,9 @@ class ImmutableGraphCSRWrapper : public CSRWrapper {
const ImmutableGraph* gptr_;
};
class BipartiteCSRWrapper : public CSRWrapper {
class UnitGraphCSRWrapper : public CSRWrapper {
public:
explicit BipartiteCSRWrapper(const Bipartite* graph) :
explicit UnitGraphCSRWrapper(const UnitGraph* graph) :
gptr_(graph) { }
aten::CSRMatrix GetInCSRMatrix() const override {
......@@ -265,7 +265,7 @@ class BipartiteCSRWrapper : public CSRWrapper {
}
private:
const Bipartite* gptr_;
const UnitGraph* gptr_;
};
} // namespace
......@@ -350,9 +350,9 @@ void BinaryOpReduce(
{__VA_ARGS__} \
} else if (ObjectTypeChecker<HeteroGraphRef>::Check(sptr.get())) { \
HeteroGraphRef g = argval; \
auto bgptr = std::dynamic_pointer_cast<Bipartite>(g.sptr()); \
auto bgptr = std::dynamic_pointer_cast<UnitGraph>(g.sptr()); \
CHECK_NOTNULL(bgptr); \
BipartiteCSRWrapper wrapper(bgptr.get()); \
UnitGraphCSRWrapper wrapper(bgptr.get()); \
{__VA_ARGS__} \
} \
} while (0)
......
src/kernel/binary_reduce_impl.h
View file @ 9b4d6079
......@@ -36,7 +36,6 @@ GData<Idx, DType> AllocGData(const std::string& op,
// GData
GData<Idx, DType> gdata;
gdata.x_length = x_len;
gdata.out_size = out_data->shape[0];
gdata.lhs_data = static_cast<DType*>(lhs_data->data);
gdata.rhs_data = static_cast<DType*>(rhs_data->data);
gdata.out_data = static_cast<DType*>(out_data->data);
......@@ -127,7 +126,6 @@ BackwardGData<Idx, DType> AllocBackwardGData(
// GData
BackwardGData<Idx, DType> gdata;
gdata.x_length = x_len;
gdata.out_size = out_data->shape[0];
gdata.lhs_data = static_cast<DType*>(lhs_data->data);
gdata.rhs_data = static_cast<DType*>(rhs_data->data);
gdata.out_data = static_cast<DType*>(out_data->data);
......
src/kernel/binary_reduce_impl_decl.h
View file @ 9b4d6079
......@@ -37,9 +37,7 @@ struct GData {
// length along x(feature) dimension
int64_t x_length{0};
// size of data, can be single value or a vector
int64_t data_len;
// number of rows of the output tensor
int64_t out_size{0};
int64_t data_len{0};
// input data
DType *lhs_data{nullptr}, *rhs_data{nullptr};
// output data
......@@ -122,9 +120,7 @@ struct BackwardGData {
// length along x(feature) dimension
int64_t x_length{0};
// size of data, can be single value or a vector
int64_t data_len;
// number of rows of the output tensor
int64_t out_size{0};
int64_t data_len{0};
// input data
DType *lhs_data{nullptr}, *rhs_data{nullptr}, *out_data{nullptr};
DType *grad_out_data{nullptr};
......@@ -227,7 +223,7 @@ struct BcastGData {
int64_t lhs_shape[NDim]{0}, lhs_stride[NDim]{0};
int64_t rhs_shape[NDim]{0}, rhs_stride[NDim]{0};
// size of data, can be single value or a vector
int64_t data_len;
int64_t data_len{0};
// input data
DType *lhs_data{nullptr}, *rhs_data{nullptr};
// input id mappings
......@@ -333,7 +329,7 @@ struct BackwardBcastGData {
int64_t rhs_shape[NDim]{0}, rhs_stride[NDim]{0};
int64_t out_shape[NDim]{0}, out_stride[NDim]{0};
// size of data, can be single value or a vector
int64_t data_len;
int64_t data_len{0};
// input id mappings
Idx *lhs_mapping{nullptr}, *rhs_mapping{nullptr}, *out_mapping{nullptr};
// input data
......
src/kernel/cuda/binary_reduce_sum.cu
View file @ 9b4d6079
......@@ -12,7 +12,6 @@
#include "../csr_interface.h"
using minigun::advance::RuntimeConfig;
using Csr = minigun::Csr<int32_t>;
namespace dgl {
namespace kernel {
......@@ -84,9 +83,9 @@ cublasStatus_t Xgeam<double>(cublasHandle_t handle, cublasOperation_t transa,
template <typename DType>
void CusparseCsrmm2(
const RuntimeConfig& rtcfg,
const Csr& csr,
const aten::CSRMatrix& csr,
const DType* B_data, DType* C_data,
int out_size, int x_length) {
int x_length) {
// We use csrmm2 to perform following operation:
// C = A x B, where A is a sparse matrix in csr format, B is the dense matrix for node
// feature tensor. However, since cusparse only supports column-major, while our tensor
......@@ -94,18 +93,10 @@ void CusparseCsrmm2(
// C = trans(A x trans(B)).
// Currently, we use cublasXgeam to implement transposition and allocate intermediate
// workspace memory for this.
// TODO(minjie): The given CSR could potentially represent a bipartite graph (e.g. in the
// case of nodeflow). Currently, we don't have bipartite graph support. Here is a small
// hack. In the python side, we create a CSR that includes both the source and destination
// nodes in the bipartite graph (so it is still square matrix). Here, when multiplying
// this sparse matrix, we specify the number of rows (the `m` here) to be equal to the
// number of rows of the output tensor (i.e, the `out_size`).
// In the future, we should make sure the number of rows of the given csr is equal
// to out_size (a.k.a the given csr is a rectangle matrix).
const int m = out_size;
const int k = csr.row_offsets.length - 1;
const int m = csr.num_rows;
const int n = x_length;
const int nnz = csr.column_indices.length;
const int k = csr.num_cols;
const int nnz = csr.indices->shape[0];
const DType alpha = 1.0;
const DType beta = 0.0;
// device
......@@ -130,7 +121,9 @@ void CusparseCsrmm2(
CUSPARSE_OPERATION_NON_TRANSPOSE,
CUSPARSE_OPERATION_TRANSPOSE,
m, n, k, nnz, &alpha,
descr, valptr, csr.row_offsets.data, csr.column_indices.data,
descr, valptr,
static_cast<int32_t*>(csr.indptr->data),
static_cast<int32_t*>(csr.indices->data),
B_data, n, &beta, trans_out, m));
device->FreeWorkspace(rtcfg.ctx, valptr);
// transpose the output matrix
......@@ -242,10 +235,9 @@ void CallBinaryReduce<kDLGPU, int32_t, float, SelectSrc, SelectNone,
cuda::FallbackCallBinaryReduce<float>(rtcfg, graph, gdata);
} else {
// cusparse use rev csr for csrmm
auto incsr = graph.GetInCSRMatrix();
Csr csr = utils::CreateCsr<int32_t>(incsr.indptr, incsr.indices);
auto csr = graph.GetInCSRMatrix();
cuda::CusparseCsrmm2(rtcfg, csr, gdata->lhs_data, gdata->out_data,
gdata->out_size, gdata->x_length);
gdata->x_length);
}
}
......@@ -259,10 +251,9 @@ void CallBinaryReduce<kDLGPU, int32_t, double, SelectSrc, SelectNone,
cuda::FallbackCallBinaryReduce<double>(rtcfg, graph, gdata);
} else {
// cusparse use rev csr for csrmm
auto incsr = graph.GetInCSRMatrix();
Csr csr = utils::CreateCsr<int32_t>(incsr.indptr, incsr.indices);
auto csr = graph.GetInCSRMatrix();
cuda::CusparseCsrmm2(rtcfg, csr, gdata->lhs_data, gdata->out_data,
gdata->out_size, gdata->x_length);
gdata->x_length);
}
}
......@@ -278,10 +269,9 @@ void CallBackwardBinaryReduce<kDLGPU, binary_op::kGradLhs, int32_t, float,
if (gdata->lhs_mapping || gdata->rhs_mapping || gdata->out_mapping) {
cuda::FallbackCallBackwardBinaryReduce<float>(rtcfg, graph, gdata);
} else {
auto outcsr = graph.GetOutCSRMatrix();
Csr csr = utils::CreateCsr<int32_t>(outcsr.indptr, outcsr.indices);
auto csr = graph.GetOutCSRMatrix();
cuda::CusparseCsrmm2(rtcfg, csr, gdata->grad_out_data, gdata->grad_lhs_data,
gdata->out_size, gdata->x_length);
gdata->x_length);
}
}
......@@ -295,10 +285,9 @@ void CallBackwardBinaryReduce<kDLGPU, binary_op::kGradLhs, int32_t, double,
if (gdata->lhs_mapping || gdata->rhs_mapping || gdata->out_mapping) {
cuda::FallbackCallBackwardBinaryReduce<double>(rtcfg, graph, gdata);
} else {
auto outcsr = graph.GetOutCSRMatrix();
Csr csr = utils::CreateCsr<int32_t>(outcsr.indptr, outcsr.indices);
auto csr = graph.GetOutCSRMatrix();
cuda::CusparseCsrmm2(rtcfg, csr, gdata->grad_out_data, gdata->grad_lhs_data,
gdata->out_size, gdata->x_length);
gdata->x_length);
}
}
......
src/runtime/c_object_api.cc
View file @ 9b4d6079
......@@ -8,6 +8,7 @@
#include <dmlc/thread_local.h>
#include <dgl/runtime/c_object_api.h>
#include <dgl/runtime/object.h>
#include <dgl/runtime/ndarray.h>
#include <dgl/packed_func_ext.h>
#include <vector>
#include <string>
......@@ -62,6 +63,9 @@ struct APIAttrGetter : public AttrVisitor {
found_object_ref = true;
}
}
void Visit(const char* key, NDArray* value) final {
if (skey == key) *ret = value[0];
}
};
struct APIAttrDir : public AttrVisitor {
......@@ -88,6 +92,9 @@ struct APIAttrDir : public AttrVisitor {
void Visit(const char* key, ObjectRef* value) final {
names->push_back(key);
}
void Visit(const char* key, NDArray* value) final {
names->push_back(key);
}
};
int DGLObjectFree(ObjectHandle handle) {
......
tests/compute/test_hetero_basics.py 0 → 100644
View file @ 9b4d6079
"""Test from `test_basics.py` but for heterograph. Merge this
with `test_basics.py` once DGLHeteroGraph is compatible with DGLGraph.
"""
import backend as F
import dgl
import networkx as nx
from collections import defaultdict as ddict
D = 5
reduce_msg_shapes = set()
def message_func(edges):
assert F.ndim(edges.src['h']) == 2
assert F.shape(edges.src['h'])[1] == D
return {'m' : edges.src['h']}
def reduce_func(nodes):
msgs = nodes.mailbox['m']
reduce_msg_shapes.add(tuple(msgs.shape))
assert F.ndim(msgs) == 3
assert F.shape(msgs)[2] == D
return {'accum' : F.sum(msgs, 1)}
def apply_node_func(nodes):
return {'h' : nodes.data['h'] + nodes.data['accum']}
def generate_graph(grad=False):
'''
s, d, eid
0, 1, 0
1, 9, 1
0, 2, 2
2, 9, 3
0, 3, 4
3, 9, 5
0, 4, 6
4, 9, 7
0, 5, 8
5, 9, 9
0, 6, 10
6, 9, 11
0, 7, 12
7, 9, 13
0, 8, 14
8, 9, 15
9, 0, 16
'''
g = dgl.graph([(0,1), (1,9), (0,2), (2,9), (0,3), (3,9), (0,4), (4,9),
(0,5), (5,9), (0,6), (6,9), (0,7), (7,9), (0,8), (8,9), (9,0)])
ncol = F.randn((10, D))
ecol = F.randn((17, D))
if grad:
ncol = F.attach_grad(ncol)
ecol = F.attach_grad(ecol)
g.ndata['h'] = ncol
g.edata['w'] = ecol
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
return g
def test_batch_setter_getter():
def _pfc(x):
return list(F.zerocopy_to_numpy(x)[:,0])
g = generate_graph()
# set all nodes
g.ndata['h'] = F.zeros((10, D))
assert F.allclose(g.ndata['h'], F.zeros((10, D)))
# pop nodes
old_len = len(g.ndata)
assert _pfc(g.ndata.pop('h')) == [0.] * 10
assert len(g.ndata) == old_len - 1
g.ndata['h'] = F.zeros((10, D))
# set partial nodes
u = F.tensor([1, 3, 5])
g.nodes[u].data['h'] = F.ones((3, D))
assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
# get partial nodes
u = F.tensor([1, 2, 3])
assert _pfc(g.nodes[u].data['h']) == [1., 0., 1.]
'''
s, d, eid
0, 1, 0
1, 9, 1
0, 2, 2
2, 9, 3
0, 3, 4
3, 9, 5
0, 4, 6
4, 9, 7
0, 5, 8
5, 9, 9
0, 6, 10
6, 9, 11
0, 7, 12
7, 9, 13
0, 8, 14
8, 9, 15
9, 0, 16
'''
# set all edges
g.edata['l'] = F.zeros((17, D))
assert _pfc(g.edata['l']) == [0.] * 17
# pop edges
old_len = len(g.edata)
assert _pfc(g.edata.pop('l')) == [0.] * 17
assert len(g.edata) == old_len - 1
g.edata['l'] = F.zeros((17, D))
# set partial edges (many-many)
u = F.tensor([0, 0, 2, 5, 9])
v = F.tensor([1, 3, 9, 9, 0])
g.edges[u, v].data['l'] = F.ones((5, D))
truth = [0.] * 17
truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (many-one)
u = F.tensor([3, 4, 6])
v = F.tensor([9])
g.edges[u, v].data['l'] = F.ones((3, D))
truth[5] = truth[7] = truth[11] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (one-many)
u = F.tensor([0])
v = F.tensor([4, 5, 6])
g.edges[u, v].data['l'] = F.ones((3, D))
truth[6] = truth[8] = truth[10] = 1.
assert _pfc(g.edata['l']) == truth
# get partial edges (many-many)
u = F.tensor([0, 6, 0])
v = F.tensor([6, 9, 7])
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (many-one)
u = F.tensor([5, 6, 7])
v = F.tensor([9])
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (one-many)
u = F.tensor([0])
v = F.tensor([3, 4, 5])
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
def test_batch_setter_autograd():
g = generate_graph(grad=True)
h1 = g.ndata['h']
# partial set
v = F.tensor([1, 2, 8])
hh = F.attach_grad(F.zeros((len(v), D)))
with F.record_grad():
g.nodes[v].data['h'] = hh
h2 = g.ndata['h']
F.backward(h2, F.ones((10, D)) * 2)
assert F.array_equal(F.grad(h1)[:,0], F.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
assert F.array_equal(F.grad(hh)[:,0], F.tensor([2., 2., 2.]))
def test_nx_conversion():
# check conversion between networkx and DGLGraph
def _check_nx_feature(nxg, nf, ef):
# check node and edge feature of nxg
# this is used to check to_networkx
num_nodes = len(nxg)
num_edges = nxg.size()
if num_nodes > 0:
node_feat = ddict(list)
for nid, attr in nxg.nodes(data=True):
assert len(attr) == len(nf)
for k in nxg.nodes[nid]:
node_feat[k].append(F.unsqueeze(attr[k], 0))
for k in node_feat:
feat = F.cat(node_feat[k], 0)
assert F.allclose(feat, nf[k])
else:
assert len(nf) == 0
if num_edges > 0:
edge_feat = ddict(lambda: [0] * num_edges)
for u, v, attr in nxg.edges(data=True):
assert len(attr) == len(ef) + 1 # extra id
eid = attr['id']
for k in ef:
edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
for k in edge_feat:
feat = F.cat(edge_feat[k], 0)
assert F.allclose(feat, ef[k])
else:
assert len(ef) == 0
n1 = F.randn((5, 3))
n2 = F.randn((5, 10))
n3 = F.randn((5, 4))
e1 = F.randn((4, 5))
e2 = F.randn((4, 7))
g = dgl.graph([(0,2),(1,4),(3,0),(4,3)])
g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
g.edata.update({'e1': e1, 'e2': e2})
# convert to networkx
nxg = dgl.to_networkx(g, node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
assert len(nxg) == 5
assert nxg.size() == 4
_check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})
# convert to DGLGraph, nx graph has id in edge feature
# use id feature to test non-tensor copy
g = dgl.graph(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
# check number of features
# test with existing dglgraph (so existing features should be cleared)
assert len(g.ndata) == 1
assert len(g.edata) == 2
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.edata['id'], F.copy_to(F.arange(0, 4), F.cpu()))
# test conversion after modifying DGLGraph
# TODO(minjie): enable after mutation is supported
#g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
#new_n = F.randn((2, 3))
#new_e = F.randn((3, 5))
#g.add_nodes(2, data={'n1': new_n})
## add three edges, one is a multi-edge
#g.add_edges([3, 6, 0], [4, 5, 2], data={'e1': new_e})
#n1 = F.cat((n1, new_n), 0)
#e1 = F.cat((e1, new_e), 0)
## convert to networkx again
#nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
#assert len(nxg) == 7
#assert nxg.size() == 7
#_check_nx_feature(nxg, {'n1': n1}, {'e1': e1})
# now test convert from networkx without id in edge feature
# first pop id in edge feature
for _, _, attr in nxg.edges(data=True):
attr.pop('id')
# test with a new graph
g = dgl.graph(nxg , node_attrs=['n1'], edge_attrs=['e1'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
# check number of features
assert len(g.ndata) == 1
assert len(g.edata) == 1
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# edge feature order follows nxg.edges()
edge_feat = []
for _, _, attr in nxg.edges(data=True):
edge_feat.append(F.unsqueeze(attr['e1'], 0))
edge_feat = F.cat(edge_feat, 0)
assert F.allclose(g.edata['e1'], edge_feat)
# Test converting from a networkx graph whose nodes are
# not labeled with consecutive-integers.
nxg = nx.cycle_graph(5)
nxg.remove_nodes_from([0, 4])
for u in nxg.nodes():
nxg.node[u]['h'] = F.tensor([u])
for u, v, d in nxg.edges(data=True):
d['h'] = F.tensor([u, v])
g = dgl.DGLGraph()
g.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
assert g.has_edge_between(0, 1)
assert g.has_edge_between(1, 2)
assert F.allclose(g.ndata['h'], F.tensor([[1.], [2.], [3.]]))
assert F.allclose(g.edata['h'], F.tensor([[1., 2.], [1., 2.],
[2., 3.], [2., 3.]]))
def test_batch_send():
g = generate_graph()
def _fmsg(edges):
assert tuple(F.shape(edges.src['h'])) == (5, D)
return {'m' : edges.src['h']}
# many-many send
u = F.tensor([0, 0, 0, 0, 0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v), _fmsg)
# one-many send
u = F.tensor([0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v), _fmsg)
# many-one send
u = F.tensor([1, 2, 3, 4, 5])
v = F.tensor([9])
g.send((u, v), _fmsg)
def test_batch_recv():
# basic recv test
g = generate_graph()
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
reduce_msg_shapes.clear()
g.send((u, v), message_func)
g.recv(F.unique(v), reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
def test_apply_nodes():
def _upd(nodes):
return {'h' : nodes.data['h'] * 2}
g = generate_graph()
old = g.ndata['h']
g.apply_nodes(_upd)
assert F.allclose(old * 2, g.ndata['h'])
u = F.tensor([0, 3, 4, 6])
g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
def test_apply_edges():
def _upd(edges):
return {'w' : edges.data['w'] * 2}
g = generate_graph()
old = g.edata['w']
g.apply_edges(_upd)
assert F.allclose(old * 2, g.edata['w'])
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
eid = F.tensor(g.edge_ids(u, v))
assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
def test_update_routines():
g = generate_graph()
# send_and_recv
reduce_msg_shapes.clear()
u = [0, 0, 0, 4, 5, 6]
v = [1, 2, 3, 9, 9, 9]
g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
try:
g.send_and_recv([u, v], message_func, reduce_func, apply_node_func)
assert False
except dgl.DGLError:
pass
# pull
v = F.tensor([1, 2, 3, 9])
reduce_msg_shapes.clear()
g.pull(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
reduce_msg_shapes.clear()
# push
v = F.tensor([0, 1, 2, 3])
reduce_msg_shapes.clear()
g.push(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
reduce_msg_shapes.clear()
# update_all
reduce_msg_shapes.clear()
g.update_all(message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
reduce_msg_shapes.clear()
def test_recv_0deg():
# test recv with 0deg nodes;
g = dgl.graph([(0,1)])
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.set_n_initializer(_init2, 'h')
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv([0, 1], _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node is equal to apply
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv(0, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: equal to apply_nodes
assert F.allclose(new[0], 2 * old[0])
# non-0deg check: untouched
assert F.allclose(new[1], old[1])
def test_recv_0deg_newfld():
# test recv with 0deg nodes; the reducer also creates a new field
g = dgl.graph([(0,1)])
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h1' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h1' : nodes.data['h1'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.set_n_initializer(_init2, 'h1')
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv([0, 1], _reduce, _apply)
new = g.ndata.pop('h1')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node
old = F.randn((2, 5))
g.ndata['h'] = old
g.ndata['h1'] = F.full((2, 5), -1, F.int64) # this is necessary
g.send((0, 1), _message)
g.recv(0, _reduce, _apply)
new = g.ndata.pop('h1')
# 0deg check: fallback to apply
assert F.allclose(new[0], F.full_1d(5, -2, F.int64))
# non-0deg check: not changed
assert F.allclose(new[1], F.full_1d(5, -1, F.int64))
def test_update_all_0deg():
# test#1
g = dgl.graph([(1,0), (2,0), (3,0), (4,0)])
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.set_n_initializer(_init2, 'h')
old_repr = F.randn((5, 5))
g.ndata['h'] = old_repr
g.update_all(_message, _reduce, _apply)
new_repr = g.ndata['h']
# the first row of the new_repr should be the sum of all the node
# features; while the 0-deg nodes should be initialized by the
# initializer and applied with UDF.
assert F.allclose(new_repr[1:], 2*(2+F.zeros((4,5))))
assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
# test#2: graph with no edge
g = dgl.graph([], card=5)
g.set_n_initializer(_init2, 'h')
g.ndata['h'] = old_repr
g.update_all(_message, _reduce, _apply)
new_repr = g.ndata['h']
# should fallback to apply
assert F.allclose(new_repr, 2*old_repr)
def test_pull_0deg():
g = dgl.graph([(0,1)])
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.set_n_initializer(_init2, 'h')
# test#1: pull both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.pull([0, 1], _message, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: pull only 0deg node
old = F.randn((2, 5))
g.ndata['h'] = old
g.pull(0, _message, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: fallback to apply
assert F.allclose(new[0], 2*old[0])
# non-0deg check: not touched
assert F.allclose(new[1], old[1])
def test_send_multigraph():
g = dgl.graph([(0,1), (0,1), (0,1), (2,1)])
def _message_a(edges):
return {'a': edges.data['a']}
def _message_b(edges):
return {'a': edges.data['a'] * 3}
def _reduce(nodes):
return {'a': F.max(nodes.mailbox['a'], 1)}
def answer(*args):
return F.max(F.stack(args, 0), 0)
assert g.is_multigraph
# send by eid
old_repr = F.randn((4, 5))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2, 3], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
# send on multigraph
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(([0, 2], [1, 1]), _message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], F.max(old_repr, 0))
# consecutive send and send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send((2, 1), _message_a)
g.send([0, 1], message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
# consecutive send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(0, message_func=_message_a)
g.send(1, message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
# send_and_recv_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
# Disabled - Heterograph doesn't support mutation
def _test_dynamic_addition():
N = 3
D = 1
g = dgl.DGLGraph()
# Test node addition
g.add_nodes(N)
g.ndata.update({'h1': F.randn((N, D)),
'h2': F.randn((N, D))})
g.add_nodes(3)
assert g.ndata['h1'].shape[0] == g.ndata['h2'].shape[0] == N + 3
# Test edge addition
g.add_edge(0, 1)
g.add_edge(1, 0)
g.edata.update({'h1': F.randn((2, D)),
'h2': F.randn((2, D))})
assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 2
g.add_edges([0, 2], [2, 0])
g.edata['h1'] = F.randn((4, D))
assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 4
g.add_edge(1, 2)
g.edges[4].data['h1'] = F.randn((1, D))
assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 5
# test add edge with part of the features
g.add_edge(2, 1, {'h1': F.randn((1, D))})
assert len(g.edata['h1']) == len(g.edata['h2'])
def test_repr():
G = dgl.graph([(0,1), (0,2), (1,2)], card=10)
repr_string = G.__repr__()
print(repr_string)
G.ndata['x'] = F.zeros((10, 5))
G.edata['y'] = F.zeros((3, 4))
repr_string = G.__repr__()
print(repr_string)
def test_group_apply_edges():
def edge_udf(edges):
h = F.sum(edges.data['feat'] * (edges.src['h'] + edges.dst['h']), dim=2)
normalized_feat = F.softmax(h, dim=1)
return {"norm_feat": normalized_feat}
elist = []
for v in [1, 2, 3, 4, 5, 6, 7, 8]:
elist.append((0, v))
for v in [2, 3, 4, 6, 7, 8]:
elist.append((1, v))
for v in [2, 3, 4, 5, 6, 7, 8]:
elist.append((2, v))
g = dgl.graph(elist)
g.ndata['h'] = F.randn((g.number_of_nodes(), D))
g.edata['feat'] = F.randn((g.number_of_edges(), D))
def _test(group_by):
g.group_apply_edges(group_by=group_by, func=edge_udf)
if group_by == 'src':
u, v, eid = g.out_edges(1, form='all')
else:
u, v, eid = g.in_edges(5, form='all')
out_feat = g.edges[eid].data['norm_feat']
result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
result = F.softmax(F.sum(result, dim=1), dim=0)
assert F.allclose(out_feat, result)
# test group by source nodes
_test('src')
# test group by destination nodes
_test('dst')
def test_local_var():
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)])
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
def foo(g):
g = g.local_var()
g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
g.edata['w'] = F.ones((g.number_of_edges(), 4))
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test out-place update
def foo(g):
g = g.local_var()
g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
g.edges[[2, 3]].data['w'] = F.ones((2, 4))
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test out-place update 2
def foo(g):
g = g.local_var()
g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test auto-pop
def foo(g):
g = g.local_var()
g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
g.edata['ww'] = F.ones((g.number_of_edges(), 4))
foo(g)
assert 'hh' not in g.ndata
assert 'ww' not in g.edata
# test initializer1
g = dgl.graph([(0,1), (1,1)])
g.set_n_initializer(dgl.init.zero_initializer)
def foo(g):
g = g.local_var()
g.nodes[0].data['h'] = F.ones((1, 1))
assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
foo(g)
# test initializer2
def foo_e_initializer(shape, dtype, ctx, id_range):
return F.ones(shape)
g.set_e_initializer(foo_e_initializer, field='h')
def foo(g):
g = g.local_var()
g.edges[0, 1].data['h'] = F.ones((1, 1))
assert F.allclose(g.edata['h'], F.ones((2, 1)))
g.edges[0, 1].data['w'] = F.ones((1, 1))
assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
foo(g)
def test_local_scope():
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)])
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
def foo(g):
with g.local_scope():
g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
g.edata['w'] = F.ones((g.number_of_edges(), 4))
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test out-place update
def foo(g):
with g.local_scope():
g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
g.edges[[2, 3]].data['w'] = F.ones((2, 4))
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test out-place update 2
def foo(g):
with g.local_scope():
g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test auto-pop
def foo(g):
with g.local_scope():
g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
g.edata['ww'] = F.ones((g.number_of_edges(), 4))
foo(g)
assert 'hh' not in g.ndata
assert 'ww' not in g.edata
# test nested scope
def foo(g):
with g.local_scope():
g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
g.edata['ww'] = F.ones((g.number_of_edges(), 4))
with g.local_scope():
g.ndata['hhh'] = F.ones((g.number_of_nodes(), 3))
g.edata['www'] = F.ones((g.number_of_edges(), 4))
assert 'hhh' not in g.ndata
assert 'www' not in g.edata
foo(g)
assert 'hh' not in g.ndata
assert 'ww' not in g.edata
# test initializer1
g = dgl.graph([(0,1), (1,1)])
g.set_n_initializer(dgl.init.zero_initializer)
def foo(g):
with g.local_scope():
g.nodes[0].data['h'] = F.ones((1, 1))
assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
foo(g)
# test initializer2
def foo_e_initializer(shape, dtype, ctx, id_range):
return F.ones(shape)
g.set_e_initializer(foo_e_initializer, field='h')
def foo(g):
with g.local_scope():
g.edges[0, 1].data['h'] = F.ones((1, 1))
assert F.allclose(g.edata['h'], F.ones((2, 1)))
g.edges[0, 1].data['w'] = F.ones((1, 1))
assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
foo(g)
if __name__ == '__main__':
test_nx_conversion()
test_batch_setter_getter()
test_batch_setter_autograd()
test_batch_send()
test_batch_recv()
test_apply_nodes()
test_apply_edges()
test_update_routines()
test_recv_0deg()
test_recv_0deg_newfld()
test_update_all_0deg()
test_pull_0deg()
test_send_multigraph()
#test_dynamic_addition()
test_repr()
test_group_apply_edges()
test_local_var()
test_local_scope()
tests/compute/test_heterograph.py
View file @ 9b4d6079
import dgl
import dgl.function as fn
from collections import Counter
......@@ -10,58 +9,89 @@ import networkx as nx
def create_test_heterograph():
# test heterograph from the docstring, plus a user -- wishes -- game relation
mg = nx.MultiDiGraph([
('user', 'user', 'follows'),
('user', 'game', 'plays'),
('user', 'game', 'wishes'),
('developer', 'game', 'develops')])
# 3 users, 2 games, 2 developers
# metagraph:
# ('user', 'follows', 'user'),
# ('user', 'plays', 'game'),
# ('user', 'wishes', 'game'),
# ('developer', 'develops', 'game')])
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
g = dgl.DGLHeteroGraph((mg, {
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): plays_spmat,
('user', 'wishes', 'game'): [(0, 1), (2, 0)],
('developer', 'develops', 'game'): [(0, 0), (1, 1)],
}))
wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game')
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game')
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
return g
def create_test_heterograph1():
edges = []
edges.extend([(0,1), (1,2)]) # follows
edges.extend([(0,3), (1,3), (2,4), (1,4)]) # plays
edges.extend([(0,4), (2,3)]) # wishes
edges.extend([(5,3), (6,4)]) # develops
ntypes = F.tensor([0, 0, 0, 1, 1, 2, 2])
etypes = F.tensor([0, 0, 1, 1, 1, 1, 2, 2, 3, 3])
g0 = dgl.graph(edges)
g0.ndata[dgl.NTYPE] = ntypes
g0.edata[dgl.ETYPE] = etypes
return dgl.to_hetero(g0, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'])
def get_redfn(name):
return getattr(F, name)
def test_create():
g0 = create_test_heterograph()
g1 = create_test_heterograph1()
assert set(g0.ntypes) == set(g1.ntypes)
assert set(g0.canonical_etypes) == set(g1.canonical_etypes)
# create from nx complete bipartite graph
nxg = nx.complete_bipartite_graph(3, 4)
g = dgl.bipartite(nxg, 'user', 'plays', 'game')
assert g.ntypes == ['user', 'game']
assert g.etypes == ['plays']
assert g.number_of_edges() == 12
# create from scipy
spmat = ssp.coo_matrix(([1,1,1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
g = dgl.graph(spmat)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 3
def test_query():
g = create_test_heterograph()
ntypes = ['user', 'game', 'developer']
etypes = [
canonical_etypes = [
('user', 'follows', 'user'),
('user', 'plays', 'game'),
('user', 'wishes', 'game'),
('developer', 'develops', 'game')]
edges = {
('user', 'follows', 'user'): ([0, 1], [1, 2]),
('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1]),
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
('developer', 'develops', 'game'): ([0, 1], [0, 1]),
}
# edges that does not exist in the graph
negative_edges = {
('user', 'follows', 'user'): ([0, 1], [0, 1]),
('user', 'plays', 'game'): ([0, 2], [1, 0]),
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
}
etypes = ['follows', 'plays', 'wishes', 'develops']
# node & edge types
assert set(ntypes) == set(g.all_node_types)
assert set(etypes) == set(g.all_edge_types)
assert set(ntypes) == set(g.ntypes)
assert set(etypes) == set(g.etypes)
assert set(canonical_etypes) == set(g.canonical_etypes)
# metagraph
mg = g.metagraph
assert set(g.all_node_types) == set(mg.nodes)
assert set(g.ntypes) == set(mg.nodes)
etype_triplets = [(u, v, e) for u, v, e in mg.edges(keys=True)]
assert set([
('user', 'user', 'follows'),
('user', 'game', 'plays'),
('user', 'game', 'wishes'),
('developer', 'game', 'develops')]) == set(etype_triplets)
for i in range(len(etypes)):
assert g.to_canonical_etype(etypes[i]) == canonical_etypes[i]
# number of nodes
assert [g.number_of_nodes(ntype) for ntype in ntypes] == [3, 2, 2]
......@@ -69,92 +99,610 @@ def test_query():
# number of edges
assert [g.number_of_edges(etype) for etype in etypes] == [2, 4, 2, 2]
assert not g.is_multigraph
assert g.is_readonly
# has_node & has_nodes
for ntype in ntypes:
n = g.number_of_nodes(ntype)
for i in range(n):
assert g.has_node(ntype, i)
assert not g.has_node(ntype, n)
assert g.has_node(i, ntype)
assert not g.has_node(n, ntype)
assert np.array_equal(
F.asnumpy(g.has_nodes(ntype, [0, n])).astype('int32'), [1, 0])
for etype in etypes:
srcs, dsts = edges[etype]
for src, dst in zip(srcs, dsts):
assert g.has_edge_between(etype, src, dst)
assert F.asnumpy(g.has_edges_between(etype, srcs, dsts)).all()
srcs, dsts = negative_edges[etype]
for src, dst in zip(srcs, dsts):
assert not g.has_edge_between(etype, src, dst)
assert not F.asnumpy(g.has_edges_between(etype, srcs, dsts)).any()
srcs, dsts = edges[etype]
n_edges = len(srcs)
# predecessors & in_edges & in_degree
pred = [s for s, d in zip(srcs, dsts) if d == 0]
assert set(F.asnumpy(g.predecessors(etype, 0)).tolist()) == set(pred)
u, v = g.in_edges(etype, [0])
assert F.asnumpy(v).tolist() == [0] * len(pred)
assert set(F.asnumpy(u).tolist()) == set(pred)
assert g.in_degree(etype, 0) == len(pred)
# successors & out_edges & out_degree
succ = [d for s, d in zip(srcs, dsts) if s == 0]
assert set(F.asnumpy(g.successors(etype, 0)).tolist()) == set(succ)
u, v = g.out_edges(etype, [0])
assert F.asnumpy(u).tolist() == [0] * len(succ)
assert set(F.asnumpy(v).tolist()) == set(succ)
assert g.out_degree(etype, 0) == len(succ)
# edge_id & edge_ids
for i, (src, dst) in enumerate(zip(srcs, dsts)):
assert g.edge_id(etype, src, dst) == i
assert F.asnumpy(g.edge_id(etype, src, dst, force_multi=True)).tolist() == [i]
assert F.asnumpy(g.edge_ids(etype, srcs, dsts)).tolist() == list(range(n_edges))
u, v, e = g.edge_ids(etype, srcs, dsts, force_multi=True)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
# find_edges
u, v = g.find_edges(etype, list(range(n_edges)))
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
# all_edges. edges are already in srcdst order
for order in ['srcdst', 'eid']:
u, v, e = g.all_edges(etype, 'all', order)
F.asnumpy(g.has_nodes([0, n], ntype)).astype('int32'), [1, 0])
def _test(g):
for etype in etypes:
srcs, dsts = edges[etype]
for src, dst in zip(srcs, dsts):
assert g.has_edge_between(src, dst, etype)
assert F.asnumpy(g.has_edges_between(srcs, dsts, etype)).all()
srcs, dsts = negative_edges[etype]
for src, dst in zip(srcs, dsts):
assert not g.has_edge_between(src, dst, etype)
assert not F.asnumpy(g.has_edges_between(srcs, dsts, etype)).any()
srcs, dsts = edges[etype]
n_edges = len(srcs)
# predecessors & in_edges & in_degree
pred = [s for s, d in zip(srcs, dsts) if d == 0]
assert set(F.asnumpy(g.predecessors(0, etype)).tolist()) == set(pred)
u, v = g.in_edges([0], etype=etype)
assert F.asnumpy(v).tolist() == [0] * len(pred)
assert set(F.asnumpy(u).tolist()) == set(pred)
assert g.in_degree(0, etype) == len(pred)
# successors & out_edges & out_degree
succ = [d for s, d in zip(srcs, dsts) if s == 0]
assert set(F.asnumpy(g.successors(0, etype)).tolist()) == set(succ)
u, v = g.out_edges([0], etype=etype)
assert F.asnumpy(u).tolist() == [0] * len(succ)
assert set(F.asnumpy(v).tolist()) == set(succ)
assert g.out_degree(0, etype) == len(succ)
# edge_id & edge_ids
for i, (src, dst) in enumerate(zip(srcs, dsts)):
assert g.edge_id(src, dst, etype=etype) == i
assert F.asnumpy(g.edge_id(src, dst, etype=etype, force_multi=True)).tolist() == [i]
assert F.asnumpy(g.edge_ids(srcs, dsts, etype=etype)).tolist() == list(range(n_edges))
u, v, e = g.edge_ids(srcs, dsts, etype=etype, force_multi=True)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
# in_degrees & out_degrees
in_degrees = F.asnumpy(g.in_degrees(etype))
out_degrees = F.asnumpy(g.out_degrees(etype))
src_count = Counter(srcs)
dst_count = Counter(dsts)
utype, _, vtype = etype
for i in range(g.number_of_nodes(utype)):
assert out_degrees[i] == src_count[i]
for i in range(g.number_of_nodes(vtype)):
assert in_degrees[i] == dst_count[i]
def test_frame():
# find_edges
u, v = g.find_edges(list(range(n_edges)), etype)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
# all_edges.
for order in ['eid']:
u, v, e = g.all_edges('all', order, etype)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
# in_degrees & out_degrees
in_degrees = F.asnumpy(g.in_degrees(etype=etype))
out_degrees = F.asnumpy(g.out_degrees(etype=etype))
src_count = Counter(srcs)
dst_count = Counter(dsts)
utype, _, vtype = g.to_canonical_etype(etype)
for i in range(g.number_of_nodes(utype)):
assert out_degrees[i] == src_count[i]
for i in range(g.number_of_nodes(vtype)):
assert in_degrees[i] == dst_count[i]
edges = {
'follows': ([0, 1], [1, 2]),
'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
'wishes': ([0, 2], [1, 0]),
'develops': ([0, 1], [0, 1]),
}
# edges that does not exist in the graph
negative_edges = {
'follows': ([0, 1], [0, 1]),
'plays': ([0, 2], [1, 0]),
'wishes': ([0, 1], [0, 1]),
'develops': ([0, 1], [1, 0]),
}
g = create_test_heterograph()
_test(g)
g = create_test_heterograph1()
_test(g)
etypes = canonical_etypes
edges = {
('user', 'follows', 'user'): ([0, 1], [1, 2]),
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
('developer', 'develops', 'game'): ([0, 1], [0, 1]),
}
# edges that does not exist in the graph
negative_edges = {
('user', 'follows', 'user'): ([0, 1], [0, 1]),
('user', 'plays', 'game'): ([0, 2], [1, 0]),
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
}
g = create_test_heterograph()
_test(g)
g = create_test_heterograph1()
_test(g)
# test repr
print(g)
def test_adj():
g = create_test_heterograph()
adj = F.sparse_to_numpy(g.adj(etype='follows'))
assert np.allclose(
adj,
np.array([[0., 0., 0.],
[1., 0., 0.],
[0., 1., 0.]]))
adj = F.sparse_to_numpy(g.adj(transpose=True, etype='follows'))
assert np.allclose(
adj,
np.array([[0., 1., 0.],
[0., 0., 1.],
[0., 0., 0.]]))
adj = F.sparse_to_numpy(g.adj(etype='plays'))
assert np.allclose(
adj,
np.array([[1., 1., 0.],
[0., 1., 1.]]))
adj = F.sparse_to_numpy(g.adj(transpose=True, etype='plays'))
assert np.allclose(
adj,
np.array([[1., 0.],
[1., 1.],
[0., 1.]]))
adj = g.adj(scipy_fmt='csr', etype='follows')
assert np.allclose(
adj.todense(),
np.array([[0., 0., 0.],
[1., 0., 0.],
[0., 1., 0.]]))
adj = g.adj(scipy_fmt='coo', etype='follows')
assert np.allclose(
adj.todense(),
np.array([[0., 0., 0.],
[1., 0., 0.],
[0., 1., 0.]]))
adj = g.adj(scipy_fmt='csr', etype='plays')
assert np.allclose(
adj.todense(),
np.array([[1., 1., 0.],
[0., 1., 1.]]))
adj = g.adj(scipy_fmt='coo', etype='plays')
assert np.allclose(
adj.todense(),
np.array([[1., 1., 0.],
[0., 1., 1.]]))
adj = F.sparse_to_numpy(g['follows'].adj())
assert np.allclose(
adj,
np.array([[0., 0., 0.],
[1., 0., 0.],
[0., 1., 0.]]))
def test_inc():
g = create_test_heterograph()
#follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
adj = F.sparse_to_numpy(g['follows'].inc('in'))
assert np.allclose(
adj,
np.array([[0., 0.],
[1., 0.],
[0., 1.]]))
adj = F.sparse_to_numpy(g['follows'].inc('out'))
assert np.allclose(
adj,
np.array([[1., 0.],
[0., 1.],
[0., 0.]]))
adj = F.sparse_to_numpy(g['follows'].inc('both'))
assert np.allclose(
adj,
np.array([[-1., 0.],
[1., -1.],
[0., 1.]]))
adj = F.sparse_to_numpy(g.inc('in', etype='plays'))
assert np.allclose(
adj,
np.array([[1., 1., 0., 0.],
[0., 0., 1., 1.]]))
adj = F.sparse_to_numpy(g.inc('out', etype='plays'))
assert np.allclose(
adj,
np.array([[1., 0., 0., 0.],
[0., 1., 0., 1.],
[0., 0., 1., 0.]]))
adj = F.sparse_to_numpy(g.inc('both', etype='follows'))
assert np.allclose(
adj,
np.array([[-1., 0.],
[1., -1.],
[0., 1.]]))
def test_view():
# test data view
g = create_test_heterograph()
f1 = F.randn((3, 6))
g.ndata['user']['h'] = f1 # ok
f2 = g.ndata['user']['h']
g.nodes['user'].data['h'] = f1 # ok
f2 = g.nodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(g.nodes['user'][0].data['h'], f1[0:1])
assert F.array_equal(F.tensor(g.nodes('user')), F.arange(0, 3))
f3 = F.randn((2, 4))
g.edata['user', 'follows', 'user']['h'] = f3
f4 = g.edata['user', 'follows', 'user']['h']
g.edges['user', 'follows', 'user'].data['h'] = f3
f4 = g.edges['user', 'follows', 'user'].data['h']
f5 = g.edges['follows'].data['h']
assert F.array_equal(f3, f4)
assert F.array_equal(g.edges['user', 'follows', 'user'][0].data['h'], f3[0:1])
assert F.array_equal(f3, f5)
assert F.array_equal(F.tensor(g.edges(etype='follows', form='eid')), F.arange(0, 2))
def test_view1():
# test relation view
HG = create_test_heterograph()
ntypes = ['user', 'game', 'developer']
canonical_etypes = [
('user', 'follows', 'user'),
('user', 'plays', 'game'),
('user', 'wishes', 'game'),
('developer', 'develops', 'game')]
etypes = ['follows', 'plays', 'wishes', 'develops']
def _test_query():
for etype in etypes:
utype, _, vtype = HG.to_canonical_etype(etype)
g = HG[etype]
srcs, dsts = edges[etype]
for src, dst in zip(srcs, dsts):
assert g.has_edge_between(src, dst)
assert F.asnumpy(g.has_edges_between(srcs, dsts)).all()
srcs, dsts = negative_edges[etype]
for src, dst in zip(srcs, dsts):
assert not g.has_edge_between(src, dst)
assert not F.asnumpy(g.has_edges_between(srcs, dsts)).any()
srcs, dsts = edges[etype]
n_edges = len(srcs)
# predecessors & in_edges & in_degree
pred = [s for s, d in zip(srcs, dsts) if d == 0]
assert set(F.asnumpy(g.predecessors(0)).tolist()) == set(pred)
u, v = g.in_edges([0])
assert F.asnumpy(v).tolist() == [0] * len(pred)
assert set(F.asnumpy(u).tolist()) == set(pred)
assert g.in_degree(0) == len(pred)
# successors & out_edges & out_degree
succ = [d for s, d in zip(srcs, dsts) if s == 0]
assert set(F.asnumpy(g.successors(0)).tolist()) == set(succ)
u, v = g.out_edges([0])
assert F.asnumpy(u).tolist() == [0] * len(succ)
assert set(F.asnumpy(v).tolist()) == set(succ)
assert g.out_degree(0) == len(succ)
# edge_id & edge_ids
for i, (src, dst) in enumerate(zip(srcs, dsts)):
assert g.edge_id(src, dst) == i
assert F.asnumpy(g.edge_id(src, dst, force_multi=True)).tolist() == [i]
assert F.asnumpy(g.edge_ids(srcs, dsts)).tolist() == list(range(n_edges))
u, v, e = g.edge_ids(srcs, dsts, force_multi=True)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
# find_edges
u, v = g.find_edges(list(range(n_edges)))
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
# all_edges.
for order in ['eid']:
u, v, e = g.all_edges(form='all', order=order)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
# in_degrees & out_degrees
in_degrees = F.asnumpy(g.in_degrees())
out_degrees = F.asnumpy(g.out_degrees())
src_count = Counter(srcs)
dst_count = Counter(dsts)
for i in range(g.number_of_nodes(utype)):
assert out_degrees[i] == src_count[i]
for i in range(g.number_of_nodes(vtype)):
assert in_degrees[i] == dst_count[i]
edges = {
'follows': ([0, 1], [1, 2]),
'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
'wishes': ([0, 2], [1, 0]),
'develops': ([0, 1], [0, 1]),
}
# edges that does not exist in the graph
negative_edges = {
'follows': ([0, 1], [0, 1]),
'plays': ([0, 2], [1, 0]),
'wishes': ([0, 1], [0, 1]),
'develops': ([0, 1], [1, 0]),
}
_test_query()
etypes = canonical_etypes
edges = {
('user', 'follows', 'user'): ([0, 1], [1, 2]),
('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]),
('user', 'wishes', 'game'): ([0, 2], [1, 0]),
('developer', 'develops', 'game'): ([0, 1], [0, 1]),
}
# edges that does not exist in the graph
negative_edges = {
('user', 'follows', 'user'): ([0, 1], [0, 1]),
('user', 'plays', 'game'): ([0, 2], [1, 0]),
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
}
_test_query()
# test features
HG.nodes['user'].data['h'] = F.ones((HG.number_of_nodes('user'), 5))
HG.nodes['game'].data['m'] = F.ones((HG.number_of_nodes('game'), 3)) * 2
# test only one node type
g = HG['follows']
assert g.number_of_nodes() == 3
# test ndata and edata
f1 = F.randn((3, 6))
g.ndata['h'] = f1 # ok
f2 = HG.nodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(F.tensor(g.nodes()), F.arange(0, 3))
f3 = F.randn((2, 4))
g.edata['h'] = f3
f4 = HG.edges['follows'].data['h']
assert F.array_equal(f3, f4)
assert F.array_equal(F.tensor(g.edges(form='eid')), F.arange(0, 2))
# test fail case
# fail due to multiple types
fail = False
try:
HG.ndata['h']
except dgl.DGLError:
fail = True
assert fail
fail = False
try:
HG.edata['h']
except dgl.DGLError:
fail = True
assert fail
def test_flatten():
def check_mapping(g, fg):
if len(fg.ntypes) == 1:
SRC = DST = fg.ntypes[0]
else:
SRC = fg.ntypes[0]
DST = fg.ntypes[1]
etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
for i, (etype, eid) in enumerate(zip(etypes, eids)):
src_g, dst_g = g.find_edges([eid], g.canonical_etypes[etype])
src_fg, dst_fg = fg.find_edges([i])
# TODO(gq): I feel this code is quite redundant; can we just add new members (like
# "induced_srcid") to returned heterograph object and not store them as features?
assert src_g == fg.nodes[SRC].data[dgl.NID][src_fg]
tid = F.asnumpy(fg.nodes[SRC].data[dgl.NTYPE][src_fg])[0]
assert g.canonical_etypes[etype][0] == g.ntypes[tid]
assert dst_g == fg.nodes[DST].data[dgl.NID][dst_fg]
tid = F.asnumpy(fg.nodes[DST].data[dgl.NTYPE][dst_fg])[0]
assert g.canonical_etypes[etype][2] == g.ntypes[tid]
# check for wildcard slices
g = create_test_heterograph()
g.nodes['user'].data['h'] = F.ones((3, 5))
g.nodes['game'].data['i'] = F.ones((2, 5))
g.edges['plays'].data['e'] = F.ones((4, 4))
g.edges['wishes'].data['e'] = F.ones((2, 4))
g.edges['wishes'].data['f'] = F.ones((2, 4))
fg = g['user', :, 'game'] # user--plays->game and user--wishes->game
assert len(fg.ntypes) == 2
assert fg.ntypes == ['user', 'game']
assert fg.etypes == ['plays+wishes']
assert F.array_equal(fg.nodes['user'].data['h'], F.ones((3, 5)))
assert F.array_equal(fg.nodes['game'].data['i'], F.ones((2, 5)))
assert F.array_equal(fg.edata['e'], F.ones((6, 4)))
assert 'f' not in fg.edata
etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
assert set(zip(etypes, eids)) == set([(1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1)])
check_mapping(g, fg)
fg = g['user', :, 'user']
# NOTE(gq): The node/edge types from the parent graph is returned if there is only one
# node/edge type. This differs from the behavior above.
assert fg.ntypes == ['user']
assert fg.etypes == ['follows']
u1, v1 = g.edges(etype='follows', order='eid')
u2, v2 = fg.edges(etype='follows', order='eid')
assert F.array_equal(u1, u2)
assert F.array_equal(v1, v2)
fg = g['developer', :, 'game']
assert fg.ntypes == ['developer', 'game']
assert fg.etypes == ['develops']
u1, v1 = g.edges(etype='develops', order='eid')
u2, v2 = fg.edges(etype='develops', order='eid')
assert F.array_equal(u1, u2)
assert F.array_equal(v1, v2)
fg = g[:, :, :]
assert fg.ntypes == ['developer+user', 'game+user']
assert fg.etypes == ['develops+follows+plays+wishes']
check_mapping(g, fg)
# Test another heterograph
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows')
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows')
g_x.nodes['user'].data['h'] = F.randn((4, 3))
g_x.edges['follows'].data['w'] = F.randn((3, 2))
g_y.nodes['user'].data['hh'] = F.randn((4, 5))
g_y.edges['knows'].data['ww'] = F.randn((2, 10))
g = dgl.hetero_from_relations([g_x, g_y])
assert F.array_equal(g.ndata['h'], g_x.ndata['h'])
assert F.array_equal(g.ndata['hh'], g_y.ndata['hh'])
assert F.array_equal(g.edges['follows'].data['w'], g_x.edata['w'])
assert F.array_equal(g.edges['knows'].data['ww'], g_y.edata['ww'])
fg = g['user', :, 'user']
assert fg.ntypes == ['user']
assert fg.etypes == ['follows+knows']
check_mapping(g, fg)
fg = g['user', :, :]
assert fg.ntypes == ['user']
assert fg.etypes == ['follows+knows']
check_mapping(g, fg)
def test_convert():
hg = create_test_heterograph()
hs = []
for ntype in hg.ntypes:
h = F.randn((hg.number_of_nodes(ntype), 5))
hg.nodes[ntype].data['h'] = h
hs.append(h)
hg.nodes['user'].data['x'] = F.randn((3, 3))
ws = []
for etype in hg.canonical_etypes:
w = F.randn((hg.number_of_edges(etype), 5))
hg.edges[etype].data['w'] = w
ws.append(w)
hg.edges['plays'].data['x'] = F.randn((4, 3))
g = dgl.to_homo(hg)
assert F.array_equal(F.cat(hs, dim=0), g.ndata['h'])
assert 'x' not in g.ndata
assert F.array_equal(F.cat(ws, dim=0), g.edata['w'])
assert 'x' not in g.edata
src, dst = g.all_edges(order='eid')
src = F.asnumpy(src)
dst = F.asnumpy(dst)
etype_id, eid = F.asnumpy(g.edata[dgl.ETYPE]), F.asnumpy(g.edata[dgl.EID])
ntype_id, nid = F.asnumpy(g.ndata[dgl.NTYPE]), F.asnumpy(g.ndata[dgl.NID])
for i in range(g.number_of_edges()):
srctype = hg.ntypes[ntype_id[src[i]]]
dsttype = hg.ntypes[ntype_id[dst[i]]]
etype = hg.etypes[etype_id[i]]
src_i, dst_i = hg.find_edges([eid[i]], (srctype, etype, dsttype))
assert np.asscalar(F.asnumpy(src_i)) == nid[src[i]]
assert np.asscalar(F.asnumpy(dst_i)) == nid[dst[i]]
mg = nx.MultiDiGraph([
('user', 'user', 'follows'),
('user', 'game', 'plays'),
('user', 'game', 'wishes'),
('developer', 'game', 'develops')])
for _mg in [None, mg]:
hg2 = dgl.to_hetero(
g, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'],
ntype_field=dgl.NTYPE, etype_field=dgl.ETYPE, metagraph=_mg)
assert set(hg.ntypes) == set(hg2.ntypes)
assert set(hg.canonical_etypes) == set(hg2.canonical_etypes)
for ntype in hg.ntypes:
assert hg.number_of_nodes(ntype) == hg2.number_of_nodes(ntype)
assert F.array_equal(hg.nodes[ntype].data['h'], hg2.nodes[ntype].data['h'])
for canonical_etype in hg.canonical_etypes:
src, dst = hg.all_edges(etype=canonical_etype, order='eid')
src2, dst2 = hg2.all_edges(etype=canonical_etype, order='eid')
assert F.array_equal(src, src2)
assert F.array_equal(dst, dst2)
assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w'])
# hetero_from_homo test case 2
g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)])
g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2'])
assert set(hg.canonical_etypes) == set(
[('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')])
assert hg.number_of_nodes('l0') == 2
assert hg.number_of_nodes('l1') == 1
assert hg.number_of_nodes('l2') == 1
assert hg.number_of_edges('e0') == 2
assert hg.number_of_edges('e1') == 1
assert hg.number_of_edges('e2') == 1
# hetero_from_homo test case 3
mg = nx.MultiDiGraph([
('user', 'movie', 'watches'),
('user', 'TV', 'watches')])
g = dgl.graph([(0, 1), (0, 2)])
g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0])
for _mg in [None, mg]:
hg = dgl.to_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg)
assert set(hg.canonical_etypes) == set(
[('user', 'watches', 'movie'), ('user', 'watches', 'TV')])
assert hg.number_of_nodes('user') == 1
assert hg.number_of_nodes('TV') == 1
assert hg.number_of_nodes('movie') == 1
assert hg.number_of_edges(('user', 'watches', 'TV')) == 1
assert hg.number_of_edges(('user', 'watches', 'movie')) == 1
assert len(hg.etypes) == 2
def test_subgraph():
g = create_test_heterograph()
x = F.randn((3, 5))
y = F.randn((2, 4))
g.nodes['user'].data['h'] = x
g.edges['follows'].data['h'] = y
def _check_subgraph(g, sg):
assert sg.ntypes == ['user', 'game', 'developer']
assert sg.etypes == ['follows', 'plays', 'wishes', 'develops']
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
F.tensor([1, 2], F.int64))
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
F.tensor([0], F.int64))
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
F.tensor([1], F.int64))
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
F.tensor([1], F.int64))
assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
F.tensor([1], F.int64))
assert sg.number_of_nodes('developer') == 0
assert sg.number_of_edges('develops') == 0
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'][1:2])
sg1 = g.subgraph({'user': [1, 2], 'game': [0]})
_check_subgraph(g, sg1)
sg2 = g.edge_subgraph({'follows': [1], 'plays': [1], 'wishes': [1]})
_check_subgraph(g, sg2)
def _check_typed_subgraph(g, sg):
assert set(sg.ntypes) == {'user', 'game'}
assert set(sg.etypes) == {'follows', 'plays', 'wishes'}
for ntype in sg.ntypes:
assert sg.number_of_nodes(ntype) == g.number_of_nodes(ntype)
for etype in sg.etypes:
src_sg, dst_sg = sg.all_edges(etype=etype, order='eid')
src_g, dst_g = g.all_edges(etype=etype, order='eid')
assert F.array_equal(src_sg, src_g)
assert F.array_equal(dst_sg, dst_g)
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'])
g.nodes['user'].data['h'][2] = F.randn((5,))
g.edges['follows'].data['h'][1] = F.randn((4,))
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'])
sg3 = g.node_type_subgraph(['user', 'game'])
_check_typed_subgraph(g, sg3)
sg4 = g.edge_type_subgraph(['follows', 'plays', 'wishes'])
_check_typed_subgraph(g, sg4)
def test_apply():
def node_udf(nodes):
......@@ -163,13 +711,357 @@ def test_apply():
return {'h': edges.data['h'] * 2 + edges.src['h']}
g = create_test_heterograph()
g.ndata['user']['h'] = F.ones((3, 5))
g.apply_nodes({'user': node_udf})
assert F.array_equal(g.ndata['user']['h'], F.ones((3, 5)) * 2)
g.nodes['user'].data['h'] = F.ones((3, 5))
g.apply_nodes(node_udf, ntype='user')
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 2)
g['plays'].edata['h'] = F.ones((4, 5))
g.apply_edges(edge_udf, etype=('user', 'plays', 'game'))
assert F.array_equal(g['plays'].edata['h'], F.ones((4, 5)) * 4)
# test apply on graph with only one type
g['follows'].apply_nodes(node_udf)
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 4)
g.edata['user', 'plays', 'game']['h'] = F.ones((4, 5))
g.apply_edges({('user', 'plays', 'game'): edge_udf})
assert F.array_equal(g.edata['user', 'plays', 'game']['h'], F.ones((4, 5)) * 4)
g['plays'].apply_edges(edge_udf)
assert F.array_equal(g['plays'].edata['h'], F.ones((4, 5)) * 12)
# test fail case
# fail due to multiple types
fail = False
try:
g.apply_nodes(node_udf)
except dgl.DGLError:
fail = True
assert fail
fail = False
try:
g.apply_edges(edge_udf)
except dgl.DGLError:
fail = True
assert fail
def test_level1():
#edges = {
# 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]),
#}
g = create_test_heterograph()
def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes):
return {'y': F.max(nodes.mailbox['m'], 1)}
def mfunc(edges):
return {'m': edges.src['h']}
def afunc(nodes):
return {'y' : nodes.data['y'] + 1}
g.nodes['user'].data['h'] = F.ones((3, 2))
g.send([2, 3], mfunc, etype='plays')
g.recv([0, 1], rfunc, etype='plays')
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
g.nodes['game'].data.pop('y')
# only one type
play_g = g['plays']
play_g.send([2, 3], mfunc)
play_g.recv([0, 1], rfunc)
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# TODO(minjie): following codes will fail because messages are
# not shared with the base graph. However, since send and recv
# are rarely used, no fix at the moment.
# g['plays'].send([2, 3], mfunc)
# g['plays'].recv([0, 1], mfunc)
# test fail case
# fail due to multiple types
fail = False
try:
g.send([2, 3], mfunc)
except dgl.DGLError:
fail = True
assert fail
fail = False
try:
g.recv([0, 1], rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi recv
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : rfunc, ('user', 'wishes', 'game'): rfunc2}, 'sum')
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))
# test multi recv with apply function
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : (rfunc, afunc), ('user', 'wishes', 'game'): rfunc2}, 'sum', afunc)
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))
# test cross reducer
g.nodes['user'].data['h'] = F.randn((3, 2))
for cred in ['sum', 'max', 'min', 'mean']:
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : (rfunc, afunc), 'wishes': rfunc2}, cred, afunc)
y = g.nodes['game'].data['y']
g1 = g['plays']
g2 = g['wishes']
g1.send(g1.edges(), mfunc)
g1.recv(g1.nodes('game'), rfunc, afunc)
y1 = g.nodes['game'].data['y']
g2.send(g2.edges(), mfunc)
g2.recv(g2.nodes('game'), rfunc2)
y2 = g.nodes['game'].data['y']
yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
yy = yy + 1 # final afunc
assert F.array_equal(y, yy)
# test fail case
# fail because cannot infer ntype
fail = False
try:
g.multi_recv([0, 1], {'plays' : rfunc, 'follows': rfunc2}, 'sum')
except dgl.DGLError:
fail = True
assert fail
def test_level2():
#edges = {
# 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]),
#}
g = create_test_heterograph()
def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes):
return {'y': F.max(nodes.mailbox['m'], 1)}
def mfunc(edges):
return {'m': edges.src['h']}
def afunc(nodes):
return {'y' : nodes.data['y'] + 1}
#############################################################
# send_and_recv
#############################################################
g.nodes['user'].data['h'] = F.ones((3, 2))
g.send_and_recv([2, 3], mfunc, rfunc, etype='plays')
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# only one type
g['plays'].send_and_recv([2, 3], mfunc, rfunc)
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# test fail case
# fail due to multiple types
fail = False
try:
g.send_and_recv([2, 3], mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_send_and_recv(
{'plays' : (g.edges(etype='plays'), mfunc, rfunc),
('user', 'wishes', 'game'): (g.edges(etype='wishes'), mfunc, rfunc2)},
'sum')
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))
# test multi
g.multi_send_and_recv(
{'plays' : (g.edges(etype='plays'), mfunc, rfunc, afunc),
('user', 'wishes', 'game'): (g.edges(etype='wishes'), mfunc, rfunc2)},
'sum', afunc)
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))
# test cross reducer
g.nodes['user'].data['h'] = F.randn((3, 2))
for cred in ['sum', 'max', 'min', 'mean']:
g.multi_send_and_recv(
{'plays' : (g.edges(etype='plays'), mfunc, rfunc, afunc),
'wishes': (g.edges(etype='wishes'), mfunc, rfunc2)},
cred, afunc)
y = g.nodes['game'].data['y']
g['plays'].send_and_recv(g.edges(etype='plays'), mfunc, rfunc, afunc)
y1 = g.nodes['game'].data['y']
g['wishes'].send_and_recv(g.edges(etype='wishes'), mfunc, rfunc2)
y2 = g.nodes['game'].data['y']
yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
yy = yy + 1 # final afunc
assert F.array_equal(y, yy)
# test fail case
# fail because cannot infer ntype
fail = False
try:
g.multi_send_and_recv(
{'plays' : (g.edges(etype='plays'), mfunc, rfunc),
'follows': (g.edges(etype='follows'), mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
#############################################################
# pull
#############################################################
g.nodes['user'].data['h'] = F.ones((3, 2))
g.pull(1, mfunc, rfunc, etype='plays')
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# only one type
g['plays'].pull(1, mfunc, rfunc)
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# test fail case
fail = False
try:
g.pull(1, mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_pull(
1,
{'plays' : (mfunc, rfunc),
('user', 'wishes', 'game'): (mfunc, rfunc2)},
'sum')
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[0., 0.], [3., 3.]]))
# test multi
g.multi_pull(
1,
{'plays' : (mfunc, rfunc, afunc),
('user', 'wishes', 'game'): (mfunc, rfunc2)},
'sum', afunc)
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[0., 0.], [5., 5.]]))
# test cross reducer
g.nodes['user'].data['h'] = F.randn((3, 2))
for cred in ['sum', 'max', 'min', 'mean']:
g.multi_pull(
1,
{'plays' : (mfunc, rfunc, afunc),
'wishes': (mfunc, rfunc2)},
cred, afunc)
y = g.nodes['game'].data['y']
g['plays'].pull(1, mfunc, rfunc, afunc)
y1 = g.nodes['game'].data['y']
g['wishes'].pull(1, mfunc, rfunc2)
y2 = g.nodes['game'].data['y']
g.nodes['game'].data['y'] = get_redfn(cred)(F.stack([y1, y2], 0), 0)
g.apply_nodes(afunc, 1, ntype='game')
yy = g.nodes['game'].data['y']
assert F.array_equal(y, yy)
# test fail case
# fail because cannot infer ntype
fail = False
try:
g.multi_pull(
1,
{'plays' : (mfunc, rfunc),
'follows': (mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
#############################################################
# update_all
#############################################################
g.nodes['user'].data['h'] = F.ones((3, 2))
g.update_all(mfunc, rfunc, etype='plays')
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[2., 2.], [2., 2.]]))
# only one type
g['plays'].update_all(mfunc, rfunc)
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[2., 2.], [2., 2.]]))
# test fail case
# fail due to multiple types
fail = False
try:
g.update_all(mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_update_all(
{'plays' : (mfunc, rfunc),
('user', 'wishes', 'game'): (mfunc, rfunc2)},
'sum')
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))
# test multi
g.multi_update_all(
{'plays' : (mfunc, rfunc, afunc),
('user', 'wishes', 'game'): (mfunc, rfunc2)},
'sum', afunc)
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))
# test cross reducer
g.nodes['user'].data['h'] = F.randn((3, 2))
for cred in ['sum', 'max', 'min', 'mean', 'stack']:
g.multi_update_all(
{'plays' : (mfunc, rfunc, afunc),
'wishes': (mfunc, rfunc2)},
cred, afunc)
y = g.nodes['game'].data['y']
g['plays'].update_all(mfunc, rfunc, afunc)
y1 = g.nodes['game'].data['y']
g['wishes'].update_all(mfunc, rfunc2)
y2 = g.nodes['game'].data['y']
if cred == 'stack':
# stack has two both correct outcomes
yy1 = F.stack([F.unsqueeze(y1, 1), F.unsqueeze(y2, 1)], 1)
yy1 = yy1 + 1 # final afunc
yy2 = F.stack([F.unsqueeze(y2, 1), F.unsqueeze(y1, 1)], 1)
yy2 = yy2 + 1 # final afunc
assert F.array_equal(y, yy1) or F.array_equal(y, yy2)
else:
yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
yy = yy + 1 # final afunc
assert F.array_equal(y, yy)
# test fail case
# fail because cannot infer ntype
fail = False
try:
g.update_all(
{'plays' : (mfunc, rfunc),
'follows': (mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
def test_updates():
def msg_func(edges):
......@@ -180,7 +1072,7 @@ def test_updates():
return {'y': nodes.data['y'] * 2}
g = create_test_heterograph()
x = F.randn((3, 5))
g.ndata['user']['h'] = x
g.nodes['user'].data['h'] = x
for msg, red, apply in itertools.product(
[fn.copy_u('h', 'm'), msg_func], [fn.sum('m', 'y'), reduce_func],
......@@ -188,38 +1080,66 @@ def test_updates():
multiplier = 1 if apply is None else 2
g['user', 'plays', 'game'].update_all(msg, red, apply)
y = g.ndata['game']['y']
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
del g.ndata['game']['y']
del g.nodes['game'].data['y']
g['user', 'plays', 'game'].send_and_recv(([0, 1, 2], [0, 1, 1]), msg, red, apply)
y = g.ndata['game']['y']
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], x[0] * multiplier)
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
del g.ndata['game']['y']
del g.nodes['game'].data['y']
g['user', 'plays', 'game'].send(([0, 1, 2], [0, 1, 1]), msg)
g['user', 'plays', 'game'].recv([0, 1], red, apply)
y = g.ndata['game']['y']
plays_g = g['user', 'plays', 'game']
plays_g.send(([0, 1, 2], [0, 1, 1]), msg)
plays_g.recv([0, 1], red, apply)
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], x[0] * multiplier)
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
del g.ndata['game']['y']
del g.nodes['game'].data['y']
# pulls from destination (game) node 0
g['user', 'plays', 'game'].pull(0, msg, red, apply)
y = g.ndata['game']['y']
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
del g.ndata['game']['y']
del g.nodes['game'].data['y']
# pushes from source (user) node 0
g['user', 'plays', 'game'].push(0, msg, red, apply)
y = g.ndata['game']['y']
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], x[0] * multiplier)
del g.ndata['game']['y']
del g.nodes['game'].data['y']
def test_backward():
g = create_test_heterograph()
x = F.randn((3, 5))
F.attach_grad(x)
g.nodes['user'].data['h'] = x
with F.record_grad():
g.multi_update_all(
{'plays' : (fn.copy_u('h', 'm'), fn.sum('m', 'y')),
'wishes': (fn.copy_u('h', 'm'), fn.sum('m', 'y'))},
'sum')
y = g.nodes['game'].data['y']
F.backward(y, F.ones(y.shape))
print(F.grad(x))
assert F.array_equal(F.grad(x), F.tensor([[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.]]))
if __name__ == '__main__':
test_create()
test_query()
test_frame()
test_adj()
test_inc()
test_view()
test_view1()
test_flatten()
test_convert()
test_subgraph()
test_apply()
test_level1()
test_level2()
test_updates()
test_backward()
tests/compute/test_udf.py
View file @ 9b4d6079
......@@ -11,9 +11,9 @@ def test_node_batch():
g.ndata['x'] = feat
# test all
v = ALL
v = utils.toindex(slice(0, g.number_of_nodes()))
n_repr = g.get_n_repr(v)
nbatch = NodeBatch(g, v, n_repr)
nbatch = NodeBatch(v, n_repr)
assert F.allclose(nbatch.data['x'], feat)
assert nbatch.mailbox is None
assert F.allclose(nbatch.nodes(), g.nodes())
......@@ -23,7 +23,7 @@ def test_node_batch():
# test partial
v = utils.toindex(F.tensor([0, 3, 5, 7, 9]))
n_repr = g.get_n_repr(v)
nbatch = NodeBatch(g, v, n_repr)
nbatch = NodeBatch(v, n_repr)
assert F.allclose(nbatch.data['x'], F.gather_row(feat, F.tensor([0, 3, 5, 7, 9])))
assert nbatch.mailbox is None
assert F.allclose(nbatch.nodes(), F.tensor([0, 3, 5, 7, 9]))
......@@ -39,13 +39,13 @@ def test_edge_batch():
g.edata['x'] = efeat
# test all
eid = ALL
eid = utils.toindex(slice(0, g.number_of_edges()))
u, v, _ = g._graph.edges('eid')
src_data = g.get_n_repr(u)
edge_data = g.get_e_repr(eid)
dst_data = g.get_n_repr(v)
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
ebatch = EdgeBatch((u, v, eid), src_data, edge_data, dst_data)
assert F.shape(ebatch.src['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == g.number_of_edges() and\
......@@ -64,7 +64,7 @@ def test_edge_batch():
src_data = g.get_n_repr(u)
edge_data = g.get_e_repr(eid)
dst_data = g.get_n_repr(v)
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
ebatch = EdgeBatch((u, v, eid), src_data, edge_data, dst_data)
assert F.shape(ebatch.src['x'])[0] == 8 and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == 8 and\
......
tests/lint/pylintrc
View file @ 9b4d6079
......@@ -192,7 +192,7 @@ function-naming-style=snake_case
#function-rgx=
# Good variable names which should always be accepted, separated by a comma.
good-names=i,j,k,u,v,e,n,m,w,x,y,g,fn,ex,Run,_
good-names=i,j,k,u,v,e,n,m,w,x,y,g,G,hg,fn,ex,Run,_
# Include a hint for the correct naming format with invalid-name.
include-naming-hint=no
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment