WIP: integrating

python/dgl/cgraph.py

@@ -4,94 +4,294 @@ from ._ffi.function import _init_api
 from . import backend as F
 from . import utils
 
-class DGLGraph(object):
-    def __init__(self):
+class GraphIndex(object):
+    """Graph index object.
+
+    Parameters
+    ----------
+    graph_data : graph data
+        Data to initialize graph. Same as networkx's semantics.
+    """
+    def __init__(self, graph_data=None):
+        # TODO: convert from graph data
         self._handle = _CAPI_DGLGraphCreate()
 
     def __del__(self):
+        """Free this graph index object."""
         _CAPI_DGLGraphFree(self._handle)
 
     def add_nodes(self, num):
+        """Add nodes.
+        
+        Parameters
+        ----------
+        num : int
+            Number of nodes to be added.
+        """
         _CAPI_DGLGraphAddVertices(self._handle, num);
 
     def add_edge(self, u, v):
+        """Add one edge.
+        
+        Parameters
+        ----------
+        u : int
+            The src node.
+        v : int
+            The dst node.
+        """
         _CAPI_DGLGraphAddEdge(self._handle, u, v);
 
     def add_edges(self, u, v):
-        u = utils.Index(u)
-        v = utils.Index(v)
-        u_array = F.asdglarray(u.totensor())
-        v_array = F.asdglarray(v.totensor())
+        """Add many edges.
+        
+        Parameters
+        ----------
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+        """
+        #u = utils.Index(u)
+        #v = utils.Index(v)
+        #u_array = F.asdglarray(u.totensor())
+        #v_array = F.asdglarray(v.totensor())
+        u_array = u.todgltensor()
+        v_array = v.todgltensor()
         _CAPI_DGLGraphAddEdges(self._handle, u_array, v_array)
 
     def clear(self):
+        """Clear the graph."""
         _CAPI_DGLGraphClear(self._handle)
 
     def number_of_nodes(self):
+        """Return the number of nodes.
+
+        Returns
+        -------
+        int
+            The number of nodes
+        """
         return _CAPI_DGLGraphNumVertices(self._handle)
 
     def number_of_edges(self):
+        """Return the number of edges.
+
+        Returns
+        -------
+        int
+            The number of edges
+        """
         return _CAPI_DGLGraphNumEdges(self._handle)
 
-    def has_vertex(self, vid):
+    def has_node(self, vid):
+        """Return true if the node exists.
+
+        Parameters
+        ----------
+        vid : int
+            The nodes
+
+        Returns
+        -------
+        bool
+            True if the node exists
+        """
         return _CAPI_DGLGraphHasVertex(self._handle, vid)
 
-    def has_vertices(self, vids):
-        vids = utils.Index(vids)
-        vid_array = F.asdglarray(vids.totensor())
-        return _CAPI_DGLGraphHasVertices(self._handle, vid_array)
+    def has_nodes(self, vids):
+        """Return true if the nodes exist.
+
+        Parameters
+        ----------
+        vid : utils.Index
+            The nodes
+
+        Returns
+        -------
+        utils.Index
+            0-1 array indicating existence
+        """
+        vid_array = vids.todgltensor()
+        return utils.Index(_CAPI_DGLGraphHasVertices(self._handle, vid_array))
 
     def has_edge(self, u, v):
+        """Return true if the edge exists.
+
+        Parameters
+        ----------
+        u : int
+            The src node.
+        v : int
+            The dst node.
+
+        Returns
+        -------
+        bool
+            True if the edge exists
+        """
         return _CAPI_DGLGraphHasEdge(self._handle, u, v)
 
     def has_edges(self, u, v):
-        u = utils.Index(u)
-        v = utils.Index(v)
-        u_array = F.asdglarray(u.totensor())
-        v_array = F.asdglarray(v.totensor())
-        return _CAPI_DGLGraphHasEdges(self._handle, u_array, v_array)
+        """Return true if the edge exists.
+
+        Parameters
+        ----------
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+
+        Returns
+        -------
+        utils.Index
+            0-1 array indicating existence
+        """
+        u_array = u.todgltensor()
+        v_array = v.todgltensor()
+        return utils.Index(_CAPI_DGLGraphHasEdges(self._handle, u_array, v_array))
 
     def predecessors(self, v, radius=1):
-        return _CAPI_DGLGraphPredecessors(self._handle, v, radius)
+        """Return the predecessors of the node.
+
+        Parameters
+        ----------
+        v : int
+            The node.
+        radius : int, optional
+            The radius of the neighborhood.
+
+        Returns
+        -------
+        utils.Index
+            Array of predecessors
+        """
+        return utils.Index(_CAPI_DGLGraphPredecessors(self._handle, v, radius))
 
     def successors(self, v, radius=1):
-        return _CAPI_DGLGraphSuccessors(self._handle, v, radius)
+        """Return the successors of the node.
+
+        Parameters
+        ----------
+        v : int
+            The node.
+        radius : int, optional
+            The radius of the neighborhood.
+
+        Returns
+        -------
+        utils.Index
+            Array of successors
+        """
+        return utils.Index(_CAPI_DGLGraphSuccessors(self._handle, v, radius))
 
     def edge_id(self, u, v):
+        """Return the id of the edge.
+
+        Parameters
+        ----------
+        u : int
+            The src node.
+        v : int
+            The dst node.
+
+        Returns
+        -------
+        int
+            The edge id.
+        """
         return _CAPI_DGLGraphEdgeId(self._handle, u, v)
 
     def edge_ids(self, u, v):
-        u = utils.Index(u)
-        v = utils.Index(v)
-        u_array = F.asdglarray(u.totensor())
-        v_array = F.asdglarray(v.totensor())
-        return _CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array)
+        """Return the edge ids.
+
+        Parameters
+        ----------
+        u : utils.Index
+            The src nodes.
+        v : utils.Index
+            The dst nodes.
+
+        Returns
+        -------
+        utils.Index
+            Teh edge id array.
+        """
+        u_array = u.todgltensor()
+        v_array = v.todgltensor()
+        return utils.Index(_CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array))
 
     def in_edges(self, v):
-        if isinstance(v, int):
-            edge_array = _CAPI_DGLGraphInEdges_1(self._handle, v)
+        """Return the in edges of the node(s).
+
+        Parameters
+        ----------
+        v : utils.Index
+            The node(s).
+        
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        if len(v) == 1:
+            edge_array = _CAPI_DGLGraphInEdges_1(self._handle, v[0])
         else:
-            v = utils.Index(v)
-            v_array = F.asdglarray(v.totensor())
+            v_array = v.todgltensor()
             edge_array = _CAPI_DGLGraphInEdges_2(self._handle, v_array)
-        src = edge_array(0)
-        dst = edge_array(1)
-        eid = edge_array(2)
+        src = utils.Index(edge_array(0))
+        dst = utils.Index(edge_array(1))
+        eid = utils.Index(edge_array(2))
         return src, dst, eid
 
     def out_edges(self, v):
-        if isinstance(v, int):
-            edge_array = _CAPI_DGLGraphOutEdges_1(self._handle, v)
+        """Return the out edges of the node(s).
+
+        Parameters
+        ----------
+        v : utils.Index
+            The node(s).
+        
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        if len(v) == 1:
+            edge_array = _CAPI_DGLGraphOutEdges_1(self._handle, v[0])
         else:
-            v = utils.Index(v)
-            v_array = F.asdglarray(v.totensor())
+            v_array = v.todgltensor()
             edge_array = _CAPI_DGLGraphOutEdges_2(self._handle, v_array)
-        src = edge_array(0)
-        dst = edge_array(1)
-        eid = edge_array(2)
+        src = utils.Index(edge_array(0))
+        dst = utils.Index(edge_array(1))
+        eid = utils.Index(edge_array(2))
         return src, dst, eid
 
     def edges(self, sorted=False):
+        """Return all the edges
+
+        Parameters
+        ----------
+        sorted : bool
+            True if the returned edges are sorted by their ids.
+        
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
         edge_array = _CAPI_DGLGraphEdges(self._handle, sorted)
         src = edge_array(0)
         dst = edge_array(1)
@@ -99,19 +299,76 @@ class DGLGraph(object):
         return src, dst, eid
 
     def in_degree(self, v):
+        """Return the in degree of the node.
+
+        Parameters
+        ----------
+        v : int
+            The node.
+
+        Returns
+        -------
+        int
+            The in degree.
+        """
         return _CAPI_DGLGraphInDegree(self._handle, v)
 
     def in_degrees(self, v):
-        v = utils.Index(v)
-        v_array = F.asdglarray(v.totensor())
-        return _CAPI_DGLGraphInDegrees(self._handle, v_array)
+        """Return the in degrees of the nodes.
+
+        Parameters
+        ----------
+        v : utils.Index
+            The nodes.
+
+        Returns
+        -------
+        int
+            The in degree array.
+        """
+        v_array = v.todgltensor()
+        return utils.Index(_CAPI_DGLGraphInDegrees(self._handle, v_array))
 
     def out_degree(self, v):
+        """Return the out degree of the node.
+
+        Parameters
+        ----------
+        v : int
+            The node.
+
+        Returns
+        -------
+        int
+            The out degree.
+        """
         return _CAPI_DGLGraphOutDegree(self._handle, v)
 
     def out_degrees(self, v):
-        v = utils.Index(v)
-        v_array = F.asdglarray(v.totensor())
-        return _CAPI_DGLGraphOutDegrees(self._handle, v_array)
+        """Return the out degrees of the nodes.
+
+        Parameters
+        ----------
+        v : utils.Index
+            The nodes.
+
+        Returns
+        -------
+        int
+            The out degree array.
+        """
+        v_array = v.todgltensor()
+        return utils.Index(_CAPI_DGLGraphOutDegrees(self._handle, v_array))
+
+    def asnetworkx(self):
+        """Convert to networkx graph.
+
+        Returns
+        -------
+        networkx.DiGraph
+            The nx graph
+        """
+        # TODO
+        return None
 
-_init_api("dgl.cgraph")
+_init_api("dgl.graph")

src/graph/graph_apis.cc

@@ -31,20 +31,20 @@ PackedFunc ConvertEdgeArrayToPackedFunc(const Graph::EdgeArray& ea) {
 // Graph handler type
 typedef void* GraphHandle;
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphCreate")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphCreate")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = new Graph();
     *rv = ghandle;
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphFree")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphFree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
     delete gptr;
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddVertices")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
@@ -52,7 +52,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddVertices")
     gptr->AddVertices(num_vertices);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddEdge")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddEdge")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
@@ -61,7 +61,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddEdge")
     gptr->AddEdge(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddEdges")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
@@ -70,28 +70,28 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphAddEdges")
     gptr->AddEdges(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphClear")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphClear")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
     gptr->Clear();
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphNumVertices")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphNumVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
     *rv = static_cast<int64_t>(gptr->NumVertices());
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphNumEdges")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphNumEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
     *rv = static_cast<int64_t>(gptr->NumEdges());
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasVertex")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasVertex")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -99,7 +99,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasVertex")
     *rv = gptr->HasVertex(vid);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasVertices")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -107,7 +107,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasVertices")
     *rv = gptr->HasVertices(vids);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasEdge")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasEdge")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -116,7 +116,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasEdge")
     *rv = gptr->HasEdge(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasEdges")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -125,7 +125,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphHasEdges")
     *rv = gptr->HasEdges(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphPredecessors")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphPredecessors")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -134,7 +134,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphPredecessors")
     *rv = gptr->Predecessors(vid, radius);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphSuccessors")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphSuccessors")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -143,7 +143,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphSuccessors")
     *rv = gptr->Successors(vid, radius);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdgeId")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdgeId")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -152,7 +152,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdgeId")
     *rv = static_cast<int64_t>(gptr->EdgeId(src, dst));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdgeIds")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdgeIds")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -161,7 +161,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdgeIds")
     *rv = gptr->EdgeIds(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInEdges_1")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInEdges_1")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -169,7 +169,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInEdges_1")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInEdges_2")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInEdges_2")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -177,7 +177,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInEdges_2")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutEdges_1")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutEdges_1")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -185,7 +185,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutEdges_1")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutEdges_2")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutEdges_2")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -193,7 +193,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutEdges_2")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdges")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -201,7 +201,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphEdges")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(sorted));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInDegree")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInDegree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -209,7 +209,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInDegree")
     *rv = static_cast<int64_t>(gptr->InDegree(vid));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInDegrees")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInDegrees")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -217,7 +217,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphInDegrees")
     *rv = gptr->InDegrees(vids);
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutDegree")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutDegree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -225,7 +225,7 @@ TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutDegree")
     *rv = static_cast<int64_t>(gptr->OutDegree(vid));
   });
 
-TVM_REGISTER_GLOBAL("cgraph._CAPI_DGLGraphOutDegrees")
+TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutDegrees")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);