Merge branch 'main' into add-niche-detection

Author: Intron7

.pre-commit-config.yaml

@@ -37,7 +37,7 @@ repos:
               docs/references.bib|
           )
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.4
+    rev: v22.1.5
     hooks:
     -   id: clang-format
         args: [--style=file, -i]

docs/installation.md

@@ -202,7 +202,17 @@ Then run the following command to execute the container:
 ```bash
 apptainer run --nv rsc.sif
 ```
+### Running on HPC systems with SLURM
 
+When running on HPC systems via SLURM, conda must be explicitly activated before running Python scripts. Use `apptainer exec` instead of `apptainer run`:
+
+```bash
+apptainer exec --nv \
+    --bind /path/to/your/data:/path/to/your/data \
+    rsc.sif \
+    bash -c "source /opt/conda/etc/profile.d/conda.sh && conda activate base && python"
+```
+Without sourcing conda first, `CONDA_PREFIX` will be unset and CuPy will fail to locate the CUDA libraries inside the container, resulting in a `TypeError: expected str, bytes or os.PathLike object, not NoneType` error.
 
 # System requirements
 

docs/release-notes/0.15.1.md

@@ -4,6 +4,17 @@
 ```
 * Add `rsc.gr.calculate_niche` with flavors `neighborhood`, `utag` , and `cellcharter`. Mirrors `squidpy.gr.calculate_niche` {pr}`644` {smaller}`S Dicks`
 * Add a minimal full-covariance GMM (`squidpy_gpu._gmm.gmm_fit_predict`) used by the `cellcharter` {pr}`644` {smaller}`S Dicks`
+* ``tl.leiden`` and ``tl.louvain`` now record the final modularity value in ``adata.uns[key_added]["modularity"]``
+(scalar for a single resolution, list for multiple resolutions) {pr}`648` {smaller}`J Pintar`
+
 ```{rubric} Bug fixes
 ```
 * Fixes `tl.rank_genes_groups` returning NaN/zero `logfoldchanges`/`pvals` with `groups=[subset]` and `reference='rest'` {pr}`651` {smaller}`S Dicks`
+* Fixes `pp.bbknn` connectivities diverging from upstream `bbknn`: per-batch neighbours are now sorted by distance before `fuzzy_simplicial_set` (so weights no longer collapse near 1.0), and the default `trim` matches upstream (`10 * neighbors_within_batch * n_batches`). Trimming kernel no longer crashes for large `trim`, and a new block-cooperative sort kernel is auto-dispatched for large `trim` for substantial speedups {pr}`659` {smaller}`S Dicks`
+* Fixes float64 precision loss in `pp.normalize_pearson_residuals` on CSR/CSC input {pr}`658` {smaller}`A Mikaeili & S Dicks`
+
+```{rubric} Misc
+```
+* ``adata.uns[key_added]["params"]["resolution"]`` is now stored as a scalar ``float`` when a single resolution
+is passed to ``tl.leiden`` and ``tl.louvain`` to match behaviour in Scanpy, and as a ``list`` when multiple
+resolutions are passed. Previously it was always stored as a list. {pr}`648`. {smaller}`J Pintar`

docs/release-notes/blank.md

@@ -6,11 +6,9 @@
 ```{rubric} Performance
 ```
 
-
 ```{rubric} Bug fixes
 ```
 
-
 ```{rubric} Misc
 ```
 

src/rapids_singlecell/_cuda/bbknn/bbknn.cu

@@ -6,20 +6,54 @@
 using namespace nb::literals;
 
 constexpr int BLOCK_SIZE = 64;
+// Block-cooperative sort kernel: BLOCK_THREADS * ITEMS_PER_THREAD = 2048.
+// Rows larger than this must use the per-thread kernel (kernel 1).
+constexpr int SORT_BLOCK_THREADS = 128;
+constexpr int SORT_ITEMS_PER_THREAD = 16;
+constexpr int SORT_TILE_SIZE = SORT_BLOCK_THREADS * SORT_ITEMS_PER_THREAD;
 
 static inline void launch_find_top_k_per_row(const float* data,
                                              const int* indptr, int n_rows,
                                              int trim, float* vals,
                                              cudaStream_t stream) {
-    dim3 block(BLOCK_SIZE);
-    dim3 grid((n_rows + BLOCK_SIZE - 1) / BLOCK_SIZE);
-    size_t shared_mem_size = static_cast<size_t>(BLOCK_SIZE) *
-                             static_cast<size_t>(trim) * sizeof(float);
+    // Each thread keeps its row's top-`trim` values in shared memory, so the
+    // per-block shared-mem request is BLOCK_SIZE * trim * sizeof(float).
+    // The default per-block shared-mem cap is ~48 KB; halve the block size
+    // until the request fits so the launch succeeds for any reasonable trim.
+    constexpr size_t SHARED_MEM_BUDGET = 48 * 1024;
+    const size_t per_thread_bytes = static_cast<size_t>(trim) * sizeof(float);
+    int block_size = BLOCK_SIZE;
+    while (block_size > 1 &&
+           static_cast<size_t>(block_size) * per_thread_bytes >
+               SHARED_MEM_BUDGET) {
+        block_size /= 2;
+    }
+    if (static_cast<size_t>(block_size) * per_thread_bytes >
+        SHARED_MEM_BUDGET) {
+        throw std::runtime_error(
+            "find_top_k_per_row: trim too large for shared-memory budget; "
+            "use find_top_k_per_row_sorted instead");
+    }
+    dim3 block(block_size);
+    dim3 grid((n_rows + block_size - 1) / block_size);
+    size_t shared_mem_size = static_cast<size_t>(block_size) * per_thread_bytes;
     find_top_k_per_row_kernel<<<grid, block, shared_mem_size, stream>>>(
         data, indptr, n_rows, trim, vals);
     CUDA_CHECK_LAST_ERROR(find_top_k_per_row_kernel);
 }
 
+static inline void launch_find_top_k_per_row_sorted(const float* data,
+                                                    const int* indptr,
+                                                    int n_rows, int trim,
+                                                    float* vals,
+                                                    cudaStream_t stream) {
+    dim3 block(SORT_BLOCK_THREADS);
+    dim3 grid(n_rows);
+    find_top_k_per_row_sorted_kernel<SORT_BLOCK_THREADS, SORT_ITEMS_PER_THREAD>
+        <<<grid, block, 0, stream>>>(data, indptr, n_rows, trim, vals);
+    CUDA_CHECK_LAST_ERROR(find_top_k_per_row_sorted_kernel);
+}
+
 static inline void launch_cut_smaller(int* indptr, int* index, float* data,
                                       float* vals, int n_rows,
                                       cudaStream_t stream) {
@@ -43,6 +77,20 @@ void register_bindings(nb::module_& m) {
         "data"_a, "indptr"_a, nb::kw_only(), "n_rows"_a, "trim"_a, "vals"_a,
         "stream"_a = 0);
 
+    m.def(
+        "find_top_k_per_row_sorted",
+        [](gpu_array_c<const float, Device> data,
+           gpu_array_c<const int, Device> indptr, int n_rows, int trim,
+           gpu_array_c<float, Device> vals, std::uintptr_t stream) {
+            launch_find_top_k_per_row_sorted(data.data(), indptr.data(), n_rows,
+                                             trim, vals.data(),
+                                             (cudaStream_t)stream);
+        },
+        "data"_a, "indptr"_a, nb::kw_only(), "n_rows"_a, "trim"_a, "vals"_a,
+        "stream"_a = 0);
+
+    m.def("sort_tile_size", []() { return SORT_TILE_SIZE; });
+
     m.def(
         "cut_smaller",
         [](gpu_array_c<int, Device> indptr, gpu_array_c<int, Device> index,

src/rapids_singlecell/_cuda/bbknn/kernels_bbknn.cuh

@@ -1,6 +1,8 @@
 #pragma once
 
+#include <cub/block/block_radix_sort.cuh>
 #include <cuda_runtime.h>
+#include <math_constants.h>
 
 __global__ void find_top_k_per_row_kernel(const float* __restrict__ data,
                                           const int* __restrict__ indptr,
@@ -49,6 +51,59 @@ __global__ void find_top_k_per_row_kernel(const float* __restrict__ data,
     vals[row] = top_k[min_index];
 }
 
+// Block-cooperative variant: one CUDA block per row, sorts the row with
+// BlockRadixSort, returns the `trim`-th largest as the cut value. Shared
+// memory is the CUB sort temp storage only, independent of `trim`, so it
+// scales to large `trim` values where the per-thread top-k kernel runs out
+// of shared memory. Requires every row to fit in BLOCK_THREADS *
+// ITEMS_PER_THREAD.
+template <int BLOCK_THREADS, int ITEMS_PER_THREAD>
+__global__ void find_top_k_per_row_sorted_kernel(const float* __restrict__ data,
+                                                 const int* __restrict__ indptr,
+                                                 const int n_rows,
+                                                 const int trim,
+                                                 float* __restrict__ vals) {
+    int row = blockIdx.x;
+    if (row >= n_rows) {
+        return;
+    }
+
+    int start = indptr[row];
+    int end = indptr[row + 1];
+    int length = end - start;
+
+    if (length <= trim) {
+        if (threadIdx.x == 0) {
+            vals[row] = 0.0f;  // insufficient elements
+        }
+        return;
+    }
+
+    using BlockRadixSort =
+        cub::BlockRadixSort<float, BLOCK_THREADS, ITEMS_PER_THREAD>;
+    __shared__ typename BlockRadixSort::TempStorage temp_storage;
+
+    float thread_keys[ITEMS_PER_THREAD];
+#pragma unroll
+    for (int i = 0; i < ITEMS_PER_THREAD; ++i) {
+        int idx = threadIdx.x * ITEMS_PER_THREAD + i;
+        // Pad out-of-range with -inf so they sort to the bottom of a
+        // descending sort and never appear among the trim largest.
+        thread_keys[i] = (idx < length) ? data[start + idx] : -CUDART_INF_F;
+    }
+
+    BlockRadixSort(temp_storage).SortDescending(thread_keys);
+
+    // After SortDescending with blocked arrangement, sorted index i lives at
+    // thread (i / ITEMS_PER_THREAD), local slot (i % ITEMS_PER_THREAD).
+    int target_idx = trim - 1;
+    int target_thread = target_idx / ITEMS_PER_THREAD;
+    int target_item = target_idx % ITEMS_PER_THREAD;
+    if (threadIdx.x == target_thread) {
+        vals[row] = thread_keys[target_item];
+    }
+}
+
 __global__ void cut_smaller_kernel(const int* __restrict__ indptr,
                                    const int* __restrict__ index,
                                    float* __restrict__ data,

src/rapids_singlecell/_cuda/pr/kernels_pr.cuh

@@ -24,7 +24,7 @@ __global__ void sparse_norm_res_csc_kernel(
             ++sparse_idx;
         }
         residuals[res_index] -= mu;
-        residuals[res_index] /= sqrtf(mu + mu * mu * inv_theta);
+        residuals[res_index] *= rsqrt(mu + mu * mu * inv_theta);
         // clamp to [-clip, clip]
         if (residuals[res_index] < -clip) residuals[res_index] = -clip;
         if (residuals[res_index] > clip) residuals[res_index] = clip;
@@ -53,7 +53,7 @@ __global__ void sparse_norm_res_csr_kernel(
             ++sparse_idx;
         }
         residuals[res_index] -= mu;
-        residuals[res_index] /= sqrtf(mu + mu * mu * inv_theta);
+        residuals[res_index] *= rsqrt(mu + mu * mu * inv_theta);
 
         if (residuals[res_index] < -clip) residuals[res_index] = -clip;
         if (residuals[res_index] > clip) residuals[res_index] = clip;

src/rapids_singlecell/preprocessing/_neighbors/__init__.py

@@ -407,8 +407,16 @@ def bbknn(
         knn_indices[:, col_range] = ind_to[sub_ind]
         knn_dist[:, col_range] = sub_dist
 
+    # Sort each row so neighbors are ordered closest-first across all batches.
+    # fuzzy_simplicial_set uses the first non-zero distance per row as the
+    # local-connectivity rho; unsorted input collapses sigma and weights.
+    order = cp.argsort(knn_dist, axis=1)
+    row_idx = cp.arange(n_obs)[:, None]
+    knn_dist = knn_dist[row_idx, order]
+    knn_indices = knn_indices[row_idx, order]
+
     if trim is None:
-        trim = 10 * neighbors_within_batch
+        trim = 10 * total_neighbors
 
     params = dict(
         n_neighbors=total_neighbors,

src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py

@@ -130,30 +130,63 @@ def _fix_self_distances(knn_dist: cp.ndarray, metric: _Metrics) -> cp.ndarray:
     return knn_dist
 
 
-def _trimming(cnts: cp_sparse.csr_matrix, trim: int) -> cp_sparse.csr_matrix:
+# Empirically, the block-cooperative CUB sort kernel is faster for trim >= 100;
+# below this threshold the per-thread top-k kernel has less launch overhead.
+_TRIM_SORT_THRESHOLD = 100
+
+
+def _trimming(
+    cnts: cp_sparse.csr_matrix,
+    trim: int,
+    *,
+    kernel: str = "auto",
+) -> cp_sparse.csr_matrix:
     from rapids_singlecell._cuda._bbknn_cuda import (
         cut_smaller,
         find_top_k_per_row,
+        find_top_k_per_row_sorted,
+        sort_tile_size,
     )
 
     n_rows = cnts.shape[0]
     vals_gpu = cp.zeros(n_rows, dtype=cp.float32)
+    stream = cp.cuda.get_current_stream().ptr
+
+    if kernel == "auto":
+        if trim >= _TRIM_SORT_THRESHOLD:
+            max_row_nnz = int(cp.diff(cnts.indptr).max().get())
+            kernel = "sorted" if max_row_nnz <= sort_tile_size() else "thread"
+        else:
+            kernel = "thread"
+
+    if kernel == "sorted":
+        find_top_k_per_row_sorted(
+            cnts.data,
+            cnts.indptr,
+            n_rows=n_rows,
+            trim=trim,
+            vals=vals_gpu,
+            stream=stream,
+        )
+    elif kernel == "thread":
+        find_top_k_per_row(
+            cnts.data,
+            cnts.indptr,
+            n_rows=n_rows,
+            trim=trim,
+            vals=vals_gpu,
+            stream=stream,
+        )
+    else:
+        raise ValueError(f"Unknown trim kernel: {kernel!r}")
 
-    find_top_k_per_row(
-        cnts.data,
-        cnts.indptr,
-        n_rows=cnts.shape[0],
-        trim=trim,
-        vals=vals_gpu,
-        stream=cp.cuda.get_current_stream().ptr,
-    )
     cut_smaller(
         cnts.indptr,
         cnts.indices,
         cnts.data,
         vals=vals_gpu,
-        n_rows=cnts.shape[0],
-        stream=cp.cuda.get_current_stream().ptr,
+        n_rows=n_rows,
+        stream=stream,
     )
     cnts.eliminate_zeros()
     return cnts

src/rapids_singlecell/tools/_clustering.py

@@ -229,8 +229,9 @@ def leiden(
         resolutions = [resolution]
     else:
         resolutions = resolution
+    modularities = []
     for resolution in resolutions:
-        leiden_parts, _ = culeiden(
+        leiden_parts, modularity = culeiden(
             g,
             resolution=resolution,
             random_state=random_state,
@@ -241,6 +242,7 @@ def leiden(
             leiden_parts = leiden_parts.to_backend("pandas").compute()
         else:
             leiden_parts = leiden_parts.to_pandas()
+        modularities.append(modularity)
 
         # Format output
         groups = leiden_parts.sort_values("vertex")[["partition"]].to_numpy().ravel()
@@ -270,10 +272,13 @@ def leiden(
     # store information on the clustering parameters
     adata.uns[key_added] = {}
     adata.uns[key_added]["params"] = {
-        "resolution": resolutions,
+        "resolution": resolutions if len(resolutions) > 1 else resolutions[0],
         "random_state": random_state,
         "n_iterations": n_iterations,
     }
+    adata.uns[key_added]["modularity"] = (
+        modularities if len(modularities) > 1 else modularities[0]
+    )
     return adata if copy else None
 
 
@@ -383,8 +388,9 @@ def louvain(
         resolutions = [resolution]
     else:
         resolutions = resolution
+    modularities = []
     for resolution in resolutions:
-        louvain_parts, _ = culouvain(
+        louvain_parts, modularity = culouvain(
             g,
             resolution=resolution,
             max_level=n_iterations,
@@ -394,6 +400,7 @@ def louvain(
             louvain_parts = louvain_parts.to_backend("pandas").compute()
         else:
             louvain_parts = louvain_parts.to_pandas()
+        modularities.append(modularity)
 
         # Format output
         groups = louvain_parts.sort_values("vertex")[["partition"]].to_numpy().ravel()
@@ -422,10 +429,13 @@ def louvain(
         Comms.destroy()
     adata.uns[key_added] = {}
     adata.uns[key_added]["params"] = {
-        "resolution": resolutions,
+        "resolution": resolutions if len(resolutions) > 1 else resolutions[0],
         "n_iterations": n_iterations,
         "threshold": threshold,
     }
+    adata.uns[key_added]["modularity"] = (
+        modularities if len(modularities) > 1 else modularities[0]
+    )
     return adata if copy else None