Repository: BiocPy/scranpy (language: python)

Issue:
Added bindings for nearest neighbor detection.
Little tricky as we need to move some files in Annoy to get the header files in the right place. Most easily done by running CMake but we can also do it manually. Note that this only affects our builds, as we can just package the already-moved headers in the source.


Discussion / hints:

Fix patch:
diff --git a/extern/libscran b/extern/libscran
index 8659e9c..96a292c 160000
--- a/extern/libscran
+++ b/extern/libscran
@@ -1,1 +1,1 @@
-Subproject commit 8659e9c76bd7656b1d75051f26eb5aa08c935d9c
+Subproject commit 96a292cb4b7daee0d10af6ddd6fd9ac659d1823f
diff --git a/setup.cfg b/setup.cfg
index 16f199c..2ccab7a 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -52,6 +52,7 @@ install_requires =
     mattress
     numpy>=1.22.4
     BiocFrame
+    igraph
 
 [options.packages.find]
 where = src
diff --git a/setup.py b/setup.py
index acc8bae..bec1ec3 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@
                         "src/scranpy/lib/model_gene_variances.cpp",
                         "src/scranpy/lib/run_pca.cpp",
                         "src/scranpy/lib/find_nearest_neighbors.cpp",
+                        "src/scranpy/lib/build_snn_graph.cpp",
                         "src/scranpy/lib/run_tsne.cpp",
                         "src/scranpy/lib/run_umap.cpp",
                     ],
diff --git a/src/scranpy/clustering/__init__.py b/src/scranpy/clustering/__init__.py
new file mode 100644
index 0000000..210374f
--- /dev/null
+++ b/src/scranpy/clustering/__init__.py
@@ -0,0 +1,1 @@
+from .build_snn_graph import build_snn_graph
diff --git a/src/scranpy/clustering/build_snn_graph.py b/src/scranpy/clustering/build_snn_graph.py
new file mode 100644
index 0000000..ac65f67
--- /dev/null
+++ b/src/scranpy/clustering/build_snn_graph.py
@@ -0,0 +1,98 @@
+import ctypes as ct
+from copy import deepcopy
+from typing import Literal
+
+import igraph as ig
+import numpy as np
+
+from ..cpphelpers import lib
+from ..nearest_neighbors import (
+    NeighborIndex,
+    NeighborResults,
+    build_neighbor_index,
+    find_nearest_neighbors,
+)
+from ..types import NeighborIndexOrResults, is_neighbor_class
+
+__author__ = "ltla, jkanche"
+__copyright__ = "ltla, jkanche"
+__license__ = "MIT"
+
+
+def build_snn_graph(
+    input: NeighborIndexOrResults,
+    num_neighbors: int = 10,
+    approximate: bool = True,
+    weight_scheme: Literal["ranked", "jaccard", "number"] = "ranked",
+    num_threads: int = 1,
+) -> ig.Graph:
+    """Build Shared nearest neighbor graph.
+
+    `input` is either a pre-built neighbor search index for the dataset, or a
+    pre-computed set of neighbor search results for all cells. If `input` is a matrix,
+    we compute the nearest neighbors for each cell, assuming it represents the
+    coordinates for each cell, usually the result of PCA step.
+    rows are variables, columns are cells.
+
+    Args:
+        input (NeighborIndexOrResults): Input matrix or pre-computed neighbors.
+        num_neighbors (int, optional): Number of neighbors to use.
+            Ignored if `input` is a `NeighborResults` object. Defaults to 15.
+        approximate (bool, optional): Whether to build an index for an approximate
+            neighbor search. Defaults to True.
+        weight_scheme (Literal["ranked", "jaccard", "number"], optional): Weighting
+            scheme for the edges between cells. This can be based on the top ranks
+            of the shared neighbors ("rank"), the number of shared neighbors ("number")
+            or the Jaccard index of the neighbor sets between cells ("jaccard").
+            Defaults to "ranked".
+        num_threads (int, optional): Number of threads to use. Defaults to 1.
+
+    Raises:
+        TypeError, ValueError: if inputs do not match expectations.
+
+    Returns:
+        ig.Graph: an igraph object.
+    """
+    if not is_neighbor_class(input):
+        raise TypeError(
+            "`input` must be either the nearest neighbor search index, search results "
+            "or a matrix."
+        )
+
+    if weight_scheme not in ["ranked", "jaccard", "number"]:
+        raise ValueError(
+            '\'weight_scheme\' must be one of "ranked", "jaccard", "number"'
+            f"provided {weight_scheme}"
+        )
+
+    graph = None
+    scheme = weight_scheme.encode("UTF-8")
+
+    if not isinstance(input, NeighborResults):
+        if not isinstance(input, NeighborIndex):
+            input = build_neighbor_index(input, approximate=approximate)
+        built = lib.build_snn_graph_from_nn_index(
+            input.ptr, num_neighbors, scheme, num_threads
+        )
+    else:
+        built = lib.build_snn_graph_from_nn_results(input.ptr, scheme, num_threads)
+
+    try:
+        nedges = lib.fetch_snn_graph_edges(built)
+        idx_pointer = ct.cast(lib.fetch_snn_graph_indices(built), ct.POINTER(ct.c_int))
+        idx_array = np.ctypeslib.as_array(idx_pointer, shape=(nedges * 2,))
+        w_pointer = ct.cast(lib.fetch_snn_graph_weights(built), ct.POINTER(ct.c_double))
+        w_array = np.ctypeslib.as_array(w_pointer, shape=(nedges,))
+
+        edge_list = []
+        for i in range(nedges):
+            edge_list.append((idx_array[2 * i], idx_array[2 * i + 1]))
+
+        nc = input.num_cells()
+        graph = ig.Graph(n=nc, edges=edge_list)
+        graph.es["weight"] = deepcopy(w_array)
+
+    finally:
+        lib.free_snn_graph(built)
+
+    return graph
diff --git a/src/scranpy/cpphelpers.py b/src/scranpy/cpphelpers.py
index 7a1a109..f009146 100644
--- a/src/scranpy/cpphelpers.py
+++ b/src/scranpy/cpphelpers.py
@@ -178,6 +178,23 @@ def load_dll() -> ct.CDLL:
 ]
 lib.run_multibatch_pca.restype = ct.c_void_p
 
+lib.build_snn_graph_from_nn_results.restype = ct.c_void_p
+lib.build_snn_graph_from_nn_results.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_int]
+lib.build_snn_graph_from_nn_index.restype = ct.c_void_p
+lib.build_snn_graph_from_nn_index.argtypes = [
+    ct.c_void_p,
+    ct.c_int,
+    ct.c_char_p,
+    ct.c_int,
+]
+lib.fetch_snn_graph_edges.restype = ct.c_int
+lib.fetch_snn_graph_edges.argtypes = [ct.c_void_p]
+lib.fetch_snn_graph_indices.restype = ct.c_void_p
+lib.fetch_snn_graph_indices.argtypes = [ct.c_void_p]
+lib.fetch_snn_graph_weights.restype = ct.c_void_p
+lib.fetch_snn_graph_weights.argtypes = [ct.c_void_p]
+lib.free_snn_graph.argtypes = [ct.c_void_p]
+
 lib.initialize_tsne.restype = ct.c_void_p
 lib.initialize_tsne.argtypes = [ct.c_void_p, ct.c_double, ct.c_int]
 lib.randomize_tsne_start.argtypes = [ct.c_int, ct.c_void_p, ct.c_int]
diff --git a/src/scranpy/lib/build_snn_graph.cpp b/src/scranpy/lib/build_snn_graph.cpp
new file mode 100644
index 0000000..9eab6ab
--- /dev/null
+++ b/src/scranpy/lib/build_snn_graph.cpp
@@ -0,0 +1,55 @@
+#include "parallel.h" // must be first, to set all macros.
+
+#include "knncolle/knncolle.hpp"
+#include "scran/clustering/BuildSnnGraph.hpp"
+#include <cstring>
+
+scran::BuildSnnGraph::Scheme resolve_weighting_scheme(const char* weight_scheme) {
+    if (std::strcmp(weight_scheme, "ranked") == 0) {
+        return scran::BuildSnnGraph::RANKED;
+    } else if (std::strcmp(weight_scheme, "number") == 0) {
+        return scran::BuildSnnGraph::NUMBER;
+    } else if (std::strcmp(weight_scheme, "jaccard") == 0) {
+        return scran::BuildSnnGraph::JACCARD;
+    } 
+
+    throw std::runtime_error("unknown weighting scheme '" + std::string(weight_scheme) + "'");
+    return scran::BuildSnnGraph::RANKED; // this has no purpose but to avoid compiler warnings.
+}
+
+extern "C" {
+
+scran::BuildSnnGraph::Results* build_snn_graph_from_nn_results(const knncolle::NeighborList<>* x, const char* weight_scheme, int num_threads) {
+    scran::BuildSnnGraph runner;
+    runner.set_num_threads(num_threads);
+    runner.set_weighting_scheme(resolve_weighting_scheme(weight_scheme));
+    auto res = runner.run(*x);
+    return new scran::BuildSnnGraph::Results(std::move(res));
+}
+
+scran::BuildSnnGraph::Results* build_snn_graph_from_nn_index(const knncolle::Base<>* x, int num_neighbors, const char* weight_scheme, int num_threads) {
+    scran::BuildSnnGraph runner;
+    runner.set_neighbors(num_neighbors);
+    runner.set_num_threads(num_threads);
+    runner.set_weighting_scheme(resolve_weighting_scheme(weight_scheme));
+    auto res = runner.run(x);
+    return new scran::BuildSnnGraph::Results(std::move(res));
+}
+
+int fetch_snn_graph_edges(const scran::BuildSnnGraph::Results* ptr) {
+    return ptr->weights.size();
+}
+
+const int* fetch_snn_graph_indices(const scran::BuildSnnGraph::Results* ptr) {
+    return ptr->edges.data();
+}
+
+const double* fetch_snn_graph_weights(const scran::BuildSnnGraph::Results* ptr) {
+    return ptr->weights.data();
+}
+
+void free_snn_graph(scran::BuildSnnGraph::Results* ptr) {
+    delete ptr;
+}
+
+}


Test patch:
diff --git a/tests/test_snn_graph.py b/tests/test_snn_graph.py
new file mode 100644
index 0000000..610a9f4
--- /dev/null
+++ b/tests/test_snn_graph.py
@@ -0,0 +1,18 @@
+import igraph as ig
+from scranpy.clustering import build_snn_graph
+
+__author__ = "ltla, jkanche"
+__copyright__ = "ltla, jkanche"
+__license__ = "MIT"
+
+
+def test_build_snn_graph(mock_data):
+    y = mock_data.pcs
+    out = build_snn_graph(y)
+
+    assert isinstance(out, ig.Graph)
+
+    clustering = out.community_multilevel()
+
+    assert clustering is not None
+    assert clustering.membership is not None