add graphsaint

pgl/graph_kernel.pyx

@@ -321,3 +321,43 @@ def alias_sample_build_table(np.ndarray[np.float64_t, ndim=1] probs):
             if alias[l_i] < 1:
                 smaller_num.push_back(l_i)
     return alias, events
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def adj_extract(
+    np.ndarray[np.int64_t, ndim=1] adj_indptr,
+    np.ndarray[np.int64_t, ndim=1] sorted_v,
+    vector[long long] sampled_nodes,
+):
+    """
+    Extract all eids of given sampled_nodes for the origin graph.
+    ret_edge_index: edge ids between sampled_nodes.
+
+    Refers: https://github.com/GraphSAINT/GraphSAINT
+    """
+    cdef long long i, v, j
+    cdef long long num_v_orig, num_v_sub
+    cdef long long start_neigh, end_neigh
+    cdef vector[int] _arr_bit
+    cdef vector[long long] ret_edge_index
+    num_v_orig = adj_indptr.size-1
+    _arr_bit = vector[int](num_v_orig,-1)
+    num_v_sub = sampled_nodes.size()
+    i = 0
+    with nogil:
+        while i < num_v_sub:
+            _arr_bit[sampled_nodes[i]] = i
+            i = i + 1
+        i = 0
+        while i < num_v_sub:
+            v = sampled_nodes[i]
+            start_neigh = adj_indptr[v]
+            end_neigh = adj_indptr[v+1]
+            j = start_neigh
+            while j < end_neigh:
+                if _arr_bit[sorted_v[j]] > -1:
+                    ret_edge_index.push_back(j)
+                j = j + 1
+            i = i + 1
+    return ret_edge_index

pgl/sample.py

@@ -476,3 +476,38 @@ def pinsage_sample(graph,
             layer_nodes[0], dtype="int64")
 
     return subgraphs
+
+
+def graph_saint_random_walk_sample(graph,
+                                   nodes,
+                                   max_depth,
+                                   alias_name=None,
+                                   events_name=None):
+    """Implement of graph saint random walk sample.
+
+    First, this function will get random walks path for given nodes and depth.
+    Then, it will create subgraph from all sampled nodes.
+
+    Reference Paper: https://arxiv.org/abs/1907.04931
+
+    Args:
+        graph: A pgl graph instance
+        nodes: Walk starting from nodes
+        max_depth: Max walking depth
+
+    Return:
+        a subgraph of sampled nodes.
+    """
+    graph.indegree()
+    walks = deepwalk_sample(graph, nodes, max_depth, alias_name, events_name)
+    sample_nodes = []
+    for walk in walks:
+        sample_nodes.extend(walk)
+    sample_nodes = np.unique(sample_nodes)
+    eids = graph_kernel.adj_extract(graph._adj_dst_index._indptr,
+                                    graph._adj_dst_index._sorted_v,
+                                    sample_nodes)
+    subgraph = graph.subgraph(
+        nodes=sample_nodes, eid=eids, with_node_feat=True, with_edge_feat=True)
+    subgraph.node_feat["index"] = np.array(sample_nodes, dtype="int64")
+    return subgraph