brisc 0.2.1__tar.gz → 0.2.2__tar.gz

brisc-0.2.2/.git_archival.txt

@@ -0,0 +1,2 @@
+node: $Format:%H$
+describe-name: $Format:%(describe:tags=true,match=*[0-9]*)$

brisc-0.2.2/.gitattributes

@@ -0,0 +1 @@
+.git_archival.txt export-subst

brisc-0.2.2/.gitignore

@@ -0,0 +1,21 @@
+# Python and build artifacts
+brisc/__pycache__/
+build/
+dist/
+*.egg-info/
+wheelhouse/
+conda-bld/
+brisc/_v2/
+brisc/_v3/
+brisc/_v4/
+brisc/_sve/
+
+# Cython artifacts (generated C/C++ files and compiled libraries)
+*.c
+*.cpp
+*.so
+*.pyd
+
+# Miscellaneous
+*.txt
+!.git_archival.txt

{brisc-0.2.1 → brisc-0.2.2}/MANIFEST.in

@@ -1,4 +1,6 @@
 recursive-include brisc *.pyx *.pxd
+prune docs
+prune .github
 prune brisc/x86_64_v2
 prune brisc/x86_64_v3
 prune brisc/x86_64_v4

{brisc-0.2.1/brisc.egg-info → brisc-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: brisc
-Version: 0.2.1
+Version: 0.2.2
 Summary: Fast single-cell toolkit
 Author-email: Michael Wainberg <m.wainberg@utoronto.ca>
 License-Expression: MIT

{brisc-0.2.1 → brisc-0.2.2}/brisc/hvg.pyx

@@ -12,20 +12,24 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
                                const long long[::1] cell_indices,
                                const unsigned long long num_cells,
                                const unsigned num_dataset_genes,
+                               unsigned long long[:, :, ::1] thread_sums,
+                               unsigned long long[::1] counts_sum,
+                               unsigned long long[::1] squared_counts_sum,
                                float[::1] mean,
                                float[::1] var,
                                unsigned[::1] nonzero_count,
                                unsigned num_threads):
+    # `thread_sums[thread_index, gene]` holds the sum, squared sum and number
+    # of non-zero counts of each gene across the elements processed by each
+    # thread, interleaved so that each element updates a single cache line.
+    # It is allocated once by the caller and reused across batches.
 
     cdef unsigned gene, cell, thread_index
     cdef unsigned long long num_elements, i, j, value, chunk_size, start, \
         end, total_sum, total_sum_of_squares, total_nonzero_count
     cdef float inv_num_cells = <float> 1 / num_cells
     cdef double inv_num_pairs_of_cells = 1.0 / (num_cells * (num_cells - 1))
-    cdef vector[unsigned long long] sum_buffer, sum_of_squares_buffer
-    cdef vector[vector[unsigned long long]] thread_sum, thread_sum_of_squares
-    cdef vector[vector[unsigned]] thread_nonzero_counts
-    cdef unsigned long long[::1] sum, sum_of_squares
+    cdef unsigned long long[:, ::1] sums
 
     if cell_indices.shape[0] == 0:
         num_elements = indices.shape[0]
@@ -34,21 +38,20 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
         num_threads = min(num_threads, num_cells)
 
     if num_threads <= 1:
-        sum_buffer.resize(num_dataset_genes)
-        sum = <unsigned long long[:num_dataset_genes]> sum_buffer.data()
-        sum_of_squares_buffer.resize(num_dataset_genes)
-        sum_of_squares = <unsigned long long[:num_dataset_genes]> \
-            sum_of_squares_buffer.data()
-        nonzero_count[:] = 0
+        sums = thread_sums[0]
+        for gene in range(num_dataset_genes):
+            sums[gene, 0] = 0
+            sums[gene, 1] = 0
+            sums[gene, 2] = 0
         if cell_indices.shape[0] == 0:
             # Iterate over all elements of the count matrix, ignoring which
             # cell they're from
             for i in range(num_elements):
                 gene = indices[i]
                 value = <unsigned long long> data[i]
-                sum[gene] += value
-                sum_of_squares[gene] += value * value
-                nonzero_count[gene] += 1
+                sums[gene, 0] += value
+                sums[gene, 1] += value * value
+                sums[gene, 2] += 1
         else:
             # Only iterate over the elements from cells in `cell_indices` (i.e.
             # cells in this batch, and/or passing QC)
@@ -58,20 +61,20 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
                                <unsigned long long> indptr[cell + 1]):
                     gene = indices[i]
                     value = <unsigned long long> data[i]
-                    sum[gene] += value
-                    sum_of_squares[gene] += value * value
-                    nonzero_count[gene] += 1
+                    sums[gene, 0] += value
+                    sums[gene, 1] += value * value
+                    sums[gene, 2] += 1
 
         # Calculate means and variances from the sums and squared sums
         for gene in range(num_dataset_genes):
-            mean[gene] = sum[gene] * inv_num_cells
+            counts_sum[gene] = sums[gene, 0]
+            squared_counts_sum[gene] = sums[gene, 1]
+            nonzero_count[gene] = <unsigned> sums[gene, 2]
+            mean[gene] = counts_sum[gene] * inv_num_cells
             var[gene] = <float> (inv_num_pairs_of_cells * (
-                <double> num_cells * <double> sum_of_squares[gene] -
-                <double> sum[gene] * <double> sum[gene]))
+                <double> num_cells * <double> squared_counts_sum[gene] -
+                <double> counts_sum[gene] * <double> counts_sum[gene]))
     else:
-        thread_sum.resize(num_threads)
-        thread_sum_of_squares.resize(num_threads)
-        thread_nonzero_counts.resize(num_threads)
         with nogil:
             if cell_indices.shape[0] == 0:
                 # Partition the work by elements, not cells, for better
@@ -80,31 +83,28 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
                 chunk_size = (num_elements + num_threads - 1) / num_threads
                 with parallel(num_threads=num_threads):
                     thread_index = threadid()
-                    thread_sum[thread_index].resize(num_dataset_genes)
-                    thread_sum_of_squares[thread_index].resize(
-                        num_dataset_genes)
-                    thread_nonzero_counts[thread_index].resize(
-                        num_dataset_genes)
+                    for gene in range(num_dataset_genes):
+                        thread_sums[thread_index, gene, 0] = 0
+                        thread_sums[thread_index, gene, 1] = 0
+                        thread_sums[thread_index, gene, 2] = 0
 
                     start = thread_index * chunk_size
                     end = min(start + chunk_size, num_elements)
                     for i in range(start, end):
                         gene = indices[i]
                         value = <unsigned long long> data[i]
-                        thread_sum[thread_index][gene] += value
-                        thread_sum_of_squares[thread_index][gene] += \
-                            value * value
-                        thread_nonzero_counts[thread_index][gene] += 1
+                        thread_sums[thread_index, gene, 0] += value
+                        thread_sums[thread_index, gene, 1] += value * value
+                        thread_sums[thread_index, gene, 2] += 1
             else:
                 # Partition the work by cells
                 chunk_size = (num_cells + num_threads - 1) / num_threads
                 with parallel(num_threads=num_threads):
                     thread_index = threadid()
-                    thread_sum[thread_index].resize(num_dataset_genes)
-                    thread_sum_of_squares[thread_index].resize(
-                        num_dataset_genes)
-                    thread_nonzero_counts[thread_index].resize(
-                        num_dataset_genes)
+                    for gene in range(num_dataset_genes):
+                        thread_sums[thread_index, gene, 0] = 0
+                        thread_sums[thread_index, gene, 1] = 0
+                        thread_sums[thread_index, gene, 2] = 0
 
                     start = thread_index * chunk_size
                     end = min(start + chunk_size, num_cells)
@@ -114,10 +114,9 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
                                        <unsigned long long> indptr[cell + 1]):
                             gene = indices[i]
                             value = <unsigned long long> data[i]
-                            thread_sum[thread_index][gene] += value
-                            thread_sum_of_squares[thread_index][gene] += \
-                                value * value
-                            thread_nonzero_counts[thread_index][gene] += 1
+                            thread_sums[thread_index, gene, 0] += value
+                            thread_sums[thread_index, gene, 1] += value * value
+                            thread_sums[thread_index, gene, 2] += 1
 
             # Calculate means and variances by aggregating the sums and squared
             # sums across threads
@@ -126,17 +125,18 @@ def gene_mean_and_variance_csr(const numeric[::1] data,
                 total_sum_of_squares = 0
                 total_nonzero_count = 0
                 for thread_index in range(num_threads):
-                    total_sum = total_sum + \
-                        thread_sum[thread_index][gene]
+                    total_sum = total_sum + thread_sums[thread_index, gene, 0]
                     total_sum_of_squares = total_sum_of_squares + \
-                        thread_sum_of_squares[thread_index][gene]
+                        thread_sums[thread_index, gene, 1]
                     total_nonzero_count = total_nonzero_count + \
-                        thread_nonzero_counts[thread_index][gene]
+                        thread_sums[thread_index, gene, 2]
+                counts_sum[gene] = total_sum
+                squared_counts_sum[gene] = total_sum_of_squares
+                nonzero_count[gene] = <unsigned> total_nonzero_count
                 mean[gene] = total_sum * inv_num_cells
                 var[gene] = <float> (inv_num_pairs_of_cells * (
                     <double> num_cells * <double> total_sum_of_squares -
                     <double> total_sum * <double> total_sum))
-                nonzero_count[gene] = total_nonzero_count
 
 
 def gene_mean_and_variance_csc(const numeric[::1] data,
@@ -258,38 +258,33 @@ def clipped_sum_csr(const numeric[::1] data,
                     const unsigned num_dataset_genes,
                     const long long[::1] cell_indices,
                     const float[::1] clip_val,
+                    unsigned long long[:, :, ::1] thread_sums,
+                    const unsigned long long[::1] counts_sum,
+                    const unsigned long long[::1] squared_counts_sum,
                     float[::1] batch_counts_sum,
                     float[::1] squared_batch_counts_sum,
                     unsigned num_threads):
-    cdef unsigned cell, gene, thread_index, chunk_size, start, end
-    cdef unsigned long long i, j, value, total_batch_counts_sum, \
-        total_squared_batch_counts_sum, total_num_out_of_range
-    cdef vector[unsigned long long] batch_counts_sum_int_buffer, \
-        squared_batch_counts_sum_int_buffer, num_out_of_range_buffer
-    cdef vector[vector[unsigned long long]] thread_batch_counts_sum_int, \
-        thread_squared_batch_counts_sum_int, thread_num_out_of_range
+    cdef unsigned cell, gene, thread_index
+    cdef unsigned long long i, j, value, chunk_size, start, end, \
+        out_of_range_sum, squared_out_of_range_sum, num_out_of_range
     cdef pair[unsigned, unsigned] row_range
-    cdef unsigned long long[::1] batch_counts_sum_int, \
-        squared_batch_counts_sum_int, num_out_of_range
+    cdef unsigned long long[:, ::1] sums
 
     # Key insight: the things we're summing are integers except when
     # `value > clip_val[gene]`, in which case we are adding a (floating-point)
     # constant, so keep track of this case separately and add it at the end, to
-    # minimize floating-point error. Also, use an integer version of `clip_val`
-    # for these comparisons, to avoid any potential inconsistency in
-    # floating-point roundoff between code paths.
+    # minimize floating-point error. This case is rare, so only update
+    # `thread_sums` for these elements, and subtract them from the unclipped
+    # sums calculated by `gene_mean_and_variance_csr()`. Also, use an integer
+    # version of `clip_val` for these comparisons, to avoid any potential
+    # inconsistency in floating-point roundoff between code paths.
     num_threads = min(num_threads, num_cells)
     if num_threads <= 1:
-        batch_counts_sum_int_buffer.resize(num_dataset_genes)
-        batch_counts_sum_int = <unsigned long long[:num_dataset_genes]> \
-            batch_counts_sum_int_buffer.data()
-        squared_batch_counts_sum_int_buffer.resize(num_dataset_genes)
-        squared_batch_counts_sum_int = \
-            <unsigned long long[:num_dataset_genes]> \
-            squared_batch_counts_sum_int_buffer.data()
-        num_out_of_range_buffer.resize(num_dataset_genes)
-        num_out_of_range = <unsigned long long[:num_dataset_genes]> \
-            num_out_of_range_buffer.data()
+        sums = thread_sums[0]
+        for gene in range(num_dataset_genes):
+            sums[gene, 0] = 0
+            sums[gene, 1] = 0
+            sums[gene, 2] = 0
         if cell_indices.shape[0] == 0:
             for cell in range(num_cells):
                 for i in range(<unsigned long long> indptr[cell],
@@ -297,10 +292,9 @@ def clipped_sum_csr(const numeric[::1] data,
                     gene = indices[i]
                     value = <unsigned long long> data[i]
                     if value > clip_val[gene]:
-                        num_out_of_range[gene] += 1
-                    else:
-                        batch_counts_sum_int[gene] += value
-                        squared_batch_counts_sum_int[gene] += value * value
+                        sums[gene, 0] += value
+                        sums[gene, 1] += value * value
+                        sums[gene, 2] += 1
         else:
             for j in range(<unsigned long long> cell_indices.shape[0]):
                 cell = cell_indices[j]
@@ -309,31 +303,25 @@ def clipped_sum_csr(const numeric[::1] data,
                     gene = indices[i]
                     value = <unsigned long long> data[i]
                     if value > clip_val[gene]:
-                        num_out_of_range[gene] += 1
-                    else:
-                        batch_counts_sum_int[gene] += value
-                        squared_batch_counts_sum_int[gene] += value * value
+                        sums[gene, 0] += value
+                        sums[gene, 1] += value * value
+                        sums[gene, 2] += 1
         for gene in range(num_dataset_genes):
             batch_counts_sum[gene] = \
-                batch_counts_sum_int[gene] + \
-                num_out_of_range[gene] * clip_val[gene]
+                (counts_sum[gene] - sums[gene, 0]) + \
+                sums[gene, 2] * clip_val[gene]
             squared_batch_counts_sum[gene] = \
-                squared_batch_counts_sum_int[gene] + \
-                num_out_of_range[gene] * clip_val[gene] * clip_val[gene]
+                (squared_counts_sum[gene] - sums[gene, 1]) + \
+                sums[gene, 2] * clip_val[gene] * clip_val[gene]
     else:
-        thread_batch_counts_sum_int.resize(num_threads)
-        thread_squared_batch_counts_sum_int.resize(num_threads)
-        thread_num_out_of_range.resize(num_threads)
         with nogil:
             if cell_indices.shape[0] == 0:
                 with parallel(num_threads=num_threads):
                     thread_index = threadid()
-                    thread_batch_counts_sum_int[thread_index].resize(
-                        num_dataset_genes)
-                    thread_squared_batch_counts_sum_int[thread_index].resize(
-                        num_dataset_genes)
-                    thread_num_out_of_range[thread_index].resize(
-                        num_dataset_genes)
+                    for gene in range(num_dataset_genes):
+                        thread_sums[thread_index, gene, 0] = 0
+                        thread_sums[thread_index, gene, 1] = 0
+                        thread_sums[thread_index, gene, 2] = 0
 
                     row_range = get_thread_offset(indptr, thread_index,
                                                   num_threads)
@@ -343,23 +331,18 @@ def clipped_sum_csr(const numeric[::1] data,
                             gene = indices[i]
                             value = <unsigned long long> data[i]
                             if value > clip_val[gene]:
-                                thread_num_out_of_range[
-                                    thread_index][gene] += 1
-                            else:
-                                thread_batch_counts_sum_int[
-                                    thread_index][gene] += value
-                                thread_squared_batch_counts_sum_int[
-                                    thread_index][gene] += value * value
+                                thread_sums[thread_index, gene, 0] += value
+                                thread_sums[thread_index, gene, 1] += \
+                                    value * value
+                                thread_sums[thread_index, gene, 2] += 1
             else:
                 chunk_size = (num_cells + num_threads - 1) / num_threads
                 with parallel(num_threads=num_threads):
                     thread_index = threadid()
-                    thread_batch_counts_sum_int[thread_index].resize(
-                        num_dataset_genes)
-                    thread_squared_batch_counts_sum_int[thread_index].resize(
-                        num_dataset_genes)
-                    thread_num_out_of_range[thread_index].resize(
-                        num_dataset_genes)
+                    for gene in range(num_dataset_genes):
+                        thread_sums[thread_index, gene, 0] = 0
+                        thread_sums[thread_index, gene, 1] = 0
+                        thread_sums[thread_index, gene, 2] = 0
 
                     start = thread_index * chunk_size
                     end = min(start + chunk_size, num_cells)
@@ -370,36 +353,29 @@ def clipped_sum_csr(const numeric[::1] data,
                             gene = indices[i]
                             value = <unsigned long long> data[i]
                             if value > clip_val[gene]:
-                                thread_num_out_of_range[
-                                    thread_index][gene] += 1
-                            else:
-                                thread_batch_counts_sum_int[
-                                    thread_index][gene] += value
-                                thread_squared_batch_counts_sum_int[
-                                    thread_index][gene] += value * value
+                                thread_sums[thread_index, gene, 0] += value
+                                thread_sums[thread_index, gene, 1] += \
+                                    value * value
+                                thread_sums[thread_index, gene, 2] += 1
 
             for gene in prange(num_dataset_genes,
                                num_threads=num_threads):
-                total_batch_counts_sum = 0
-                total_squared_batch_counts_sum = 0
-                total_num_out_of_range = 0
+                out_of_range_sum = 0
+                squared_out_of_range_sum = 0
+                num_out_of_range = 0
                 for thread_index in range(num_threads):
-                    total_batch_counts_sum = \
-                        total_batch_counts_sum + \
-                        thread_batch_counts_sum_int[thread_index][gene]
-                    total_squared_batch_counts_sum = \
-                        total_squared_batch_counts_sum + \
-                        thread_squared_batch_counts_sum_int[thread_index][gene]
-                    total_num_out_of_range = \
-                        total_num_out_of_range + \
-                        thread_num_out_of_range[thread_index][gene]
+                    out_of_range_sum = out_of_range_sum + \
+                        thread_sums[thread_index, gene, 0]
+                    squared_out_of_range_sum = squared_out_of_range_sum + \
+                        thread_sums[thread_index, gene, 1]
+                    num_out_of_range = num_out_of_range + \
+                        thread_sums[thread_index, gene, 2]
                 batch_counts_sum[gene] = \
-                    total_batch_counts_sum + \
-                    total_num_out_of_range * clip_val[gene]
+                    (counts_sum[gene] - out_of_range_sum) + \
+                    num_out_of_range * clip_val[gene]
                 squared_batch_counts_sum[gene] = \
-                    total_squared_batch_counts_sum + \
-                    total_num_out_of_range * \
-                    clip_val[gene] * clip_val[gene]
+                    (squared_counts_sum[gene] - squared_out_of_range_sum) + \
+                    num_out_of_range * clip_val[gene] * clip_val[gene]
 
 
 def clipped_sum_csc(const numeric[::1] data,

{brisc-0.2.1 → brisc-0.2.2}/brisc/knn.pyx

@@ -37,12 +37,14 @@ def knn_self(const float[:, ::1] X,
         num_cells_in_block, query_index_in_block, query_sorted, num_pairs, \
         pair_index, next_pair_index, query_in_cluster_index, \
         num_queries_in_cluster, max_cluster_size, cluster_start, \
-        cluster_size, neighbor_offset, num_dimensions = X.shape[1], \
+        cluster_size, neighbor_offset, block_offset, block_end, \
+        num_dimensions = X.shape[1], \
         num_clusters = centroids.shape[0], \
         num_chunks = (num_X + chunk_size_kmeans - 1) / chunk_size_kmeans, \
         num_blocks = (num_X + chunk_size_search - 1) / chunk_size_search
     cdef float norm, worst_distance, distance, query_norm, true_distance, \
-        partial_distance
+        block_min_distance
+    cdef const float* row
     cdef uninitialized_vector[unsigned] cluster_offsets_buffer
     cdef uninitialized_vector[float] centroid_norms_buffer, \
         temp_distances_buffer, query_X_buffer, block_distances_buffer
@@ -242,34 +244,25 @@ def knn_self(const float[:, ::1] X,
                     query_sorted = block_start + query_index_in_block
                     query_norm = cell_norms[query_sorted]
                     worst_distance = distances[query_sorted, 0]
-                    if not cluster_start <= query_sorted < \
-                            cluster_start + cluster_size:
-                        # Different cluster from the one the query cell is in;
-                        # leave out the self-neighbors check
-                        for neighbor_offset in range(cluster_size):
-                            neighbor = cluster_start + neighbor_offset
-                            partial_distance = block_distances[
-                                query_in_cluster_index * cluster_size +
-                                neighbor_offset]
-                            true_distance = query_norm + partial_distance
-                            if true_distance < worst_distance:
-                                max_heap_replace_top(
-                                    &neighbors[query_sorted, 0],
-                                    &distances[query_sorted, 0],
-                                    sorted_order[neighbor], true_distance,
-                                    num_neighbors)
-                                worst_distance = distances[query_sorted, 0]
-                    else:
-                        # Same cluster the query cell is in; check for
-                        # self-neighbors
-                        for neighbor_offset in range(cluster_size):
+                    row = &block_distances[
+                        query_in_cluster_index * cluster_size]
+                    # Skip blocks of 64 cells where even the nearest one
+                    # does not beat the current worst neighbor
+                    for block_offset in range(0, cluster_size, 64):
+                        block_end = min(block_offset + 64, cluster_size)
+                        block_min_distance = row[block_offset]
+                        for neighbor_offset in range(block_offset + 1,
+                                                     block_end):
+                            block_min_distance = min(block_min_distance,
+                                                     row[neighbor_offset])
+                        if query_norm + block_min_distance >= worst_distance:
+                            continue
+                        for neighbor_offset in range(block_offset,
+                                                     block_end):
                             neighbor = cluster_start + neighbor_offset
                             if query_sorted == neighbor:
                                 continue  # skip self-neighbors
-                            partial_distance = block_distances[
-                                query_in_cluster_index * cluster_size +
-                                neighbor_offset]
-                            true_distance = query_norm + partial_distance
+                            true_distance = query_norm + row[neighbor_offset]
                             if true_distance < worst_distance:
                                 max_heap_replace_top(
                                     &neighbors[query_sorted, 0],
@@ -458,37 +451,28 @@ def knn_self(const float[:, ::1] X,
                             query_sorted = block_start + query_index_in_block
                             query_norm = cell_norms[query_sorted]
                             worst_distance = distances[query_sorted, 0]
-                            if not cluster_start <= query_sorted < \
-                                    cluster_start + cluster_size:
-                                # Different cluster from the one the query cell
-                                # is in; leave out the self-neighbors check
-                                for neighbor_offset in range(cluster_size):
-                                    neighbor = cluster_start + neighbor_offset
-                                    partial_distance = thread_block_distances[
-                                        thread_index][query_in_cluster_index *
-                                            cluster_size + neighbor_offset]
-                                    true_distance = \
-                                        query_norm + partial_distance
-                                    if true_distance < worst_distance:
-                                        max_heap_replace_top(
-                                            &neighbors[query_sorted, 0],
-                                            &distances[query_sorted, 0],
-                                            sorted_order[neighbor],
-                                            true_distance, num_neighbors)
-                                        worst_distance = \
-                                            distances[query_sorted, 0]
-                            else:
-                                # Same cluster the query cell is in; check for
-                                # self-neighbors
-                                for neighbor_offset in range(cluster_size):
+                            row = \
+                                thread_block_distances[thread_index].data() + \
+                                query_in_cluster_index * cluster_size
+                            for block_offset in range(0, cluster_size, 64):
+                                block_end = \
+                                    min(block_offset + 64, cluster_size)
+                                block_min_distance = row[block_offset]
+                                for neighbor_offset in range(
+                                        block_offset + 1, block_end):
+                                    block_min_distance = min(
+                                        block_min_distance,
+                                        row[neighbor_offset])
+                                if query_norm + block_min_distance >= \
+                                        worst_distance:
+                                    continue
+                                for neighbor_offset in range(block_offset,
+                                                             block_end):
                                     neighbor = cluster_start + neighbor_offset
                                     if query_sorted == neighbor:
                                         continue  # skip self-neighbors
-                                    partial_distance = thread_block_distances[
-                                        thread_index][query_in_cluster_index *
-                                            cluster_size + neighbor_offset]
                                     true_distance = \
-                                        query_norm + partial_distance
+                                        query_norm + row[neighbor_offset]
                                     if true_distance < worst_distance:
                                         max_heap_replace_top(
                                             &neighbors[query_sorted, 0],
@@ -544,14 +528,16 @@ def knn_cross(const float[:, ::1] Y,
         num_cells_in_block, query_index_in_block, query_sorted, num_pairs, \
         pair_index, next_pair_index, query_in_cluster_index, \
         num_queries_in_cluster, max_cluster_size, cluster_start, \
-        cluster_size, neighbor_offset, num_dimensions = X.shape[1], \
+        cluster_size, neighbor_offset, block_offset, block_end, \
+        num_dimensions = X.shape[1], \
         num_clusters = centroids.shape[0], \
         num_chunks = (num_Y + chunk_size_kmeans - 1) / \
             chunk_size_kmeans, \
         num_blocks = (num_Y + chunk_size_search - 1) / \
             chunk_size_search
     cdef float norm, worst_distance, distance, query_norm, true_distance, \
-        partial_distance
+        block_min_distance
+    cdef const float* row
     cdef uninitialized_vector[unsigned] cluster_offsets_buffer
     cdef uninitialized_vector[float] centroid_norms_buffer, \
         temp_distances_buffer, query_Y_buffer, block_distances_buffer
@@ -743,18 +729,29 @@ def knn_cross(const float[:, ::1] Y,
                     query_sorted = block_start + query_index_in_block
                     query_norm = query_norms[query_sorted]
                     worst_distance = distances[query_sorted, 0]
-                    for neighbor_offset in range(cluster_size):
-                        neighbor = cluster_start + neighbor_offset
-                        partial_distance = block_distances[
-                            query_in_cluster_index * cluster_size +
-                            neighbor_offset]
-                        true_distance = query_norm + partial_distance
-                        if true_distance < worst_distance:
-                            max_heap_replace_top(&neighbors[query_sorted, 0],
-                                                 &distances[query_sorted, 0],
-                                                 neighbor, true_distance,
-                                                 num_neighbors)
-                            worst_distance = distances[query_sorted, 0]
+                    row = &block_distances[
+                        query_in_cluster_index * cluster_size]
+                    # Skip blocks of 64 cells where even the nearest one
+                    # does not beat the current worst neighbor
+                    for block_offset in range(0, cluster_size, 64):
+                        block_end = min(block_offset + 64, cluster_size)
+                        block_min_distance = row[block_offset]
+                        for neighbor_offset in range(block_offset + 1,
+                                                     block_end):
+                            block_min_distance = min(block_min_distance,
+                                                     row[neighbor_offset])
+                        if query_norm + block_min_distance >= worst_distance:
+                            continue
+                        for neighbor_offset in range(block_offset,
+                                                     block_end):
+                            neighbor = cluster_start + neighbor_offset
+                            true_distance = query_norm + row[neighbor_offset]
+                            if true_distance < worst_distance:
+                                max_heap_replace_top(
+                                    &neighbors[query_sorted, 0],
+                                    &distances[query_sorted, 0], neighbor,
+                                    true_distance, num_neighbors)
+                                worst_distance = distances[query_sorted, 0]
 
                 pair_index = next_pair_index
 
@@ -944,18 +941,34 @@ def knn_cross(const float[:, ::1] Y,
                             query_sorted = block_start + query_index_in_block
                             query_norm = query_norms[query_sorted]
                             worst_distance = distances[query_sorted, 0]
-                            for neighbor_offset in range(cluster_size):
-                                neighbor = cluster_start + neighbor_offset
-                                partial_distance = thread_block_distances[
-                                    thread_index][query_in_cluster_index *
-                                        cluster_size + neighbor_offset]
-                                true_distance = query_norm + partial_distance
-                                if true_distance < worst_distance:
-                                    max_heap_replace_top(
-                                        &neighbors[query_sorted, 0],
-                                        &distances[query_sorted, 0],
-                                        neighbor, true_distance, num_neighbors)
-                                    worst_distance = distances[query_sorted, 0]
+                            row = \
+                                thread_block_distances[thread_index].data() + \
+                                query_in_cluster_index * cluster_size
+                            for block_offset in range(0, cluster_size, 64):
+                                block_end = \
+                                    min(block_offset + 64, cluster_size)
+                                block_min_distance = row[block_offset]
+                                for neighbor_offset in range(
+                                        block_offset + 1, block_end):
+                                    block_min_distance = min(
+                                        block_min_distance,
+                                        row[neighbor_offset])
+                                if query_norm + block_min_distance >= \
+                                        worst_distance:
+                                    continue
+                                for neighbor_offset in range(block_offset,
+                                                             block_end):
+                                    neighbor = cluster_start + neighbor_offset
+                                    true_distance = \
+                                        query_norm + row[neighbor_offset]
+                                    if true_distance < worst_distance:
+                                        max_heap_replace_top(
+                                            &neighbors[query_sorted, 0],
+                                            &distances[query_sorted, 0],
+                                            neighbor, true_distance,
+                                            num_neighbors)
+                                        worst_distance = \
+                                            distances[query_sorted, 0]
 
                         pair_index = next_pair_index
 

{brisc-0.2.1 → brisc-0.2.2}/brisc/single_cell.py

@@ -11385,6 +11385,12 @@ class SingleCell:
         norm_gene_vars = []
         total_nonzero_count = np.zeros(len(genes_in_all_datasets),
                                        dtype=np.uint64)
+        # For CSR, allocate the per-thread sums once and reuse them across
+        # batches, rather than reallocating them for every batch
+        if any(isinstance(dataset._X, csr_array) for dataset in datasets):
+            thread_sums = np.empty((num_threads, max(
+                dataset._X.shape[1] for dataset in datasets), 3),
+                dtype=np.uint64)
         for X, cell_mask, gene_indices in batches:
             num_dataset_genes = X.shape[1]
             mean = np.empty(num_dataset_genes, dtype=np.float32)
@@ -11400,10 +11406,15 @@ class SingleCell:
                 else:
                     cell_indices = cell_mask
                     num_cells = len(cell_indices)
+                counts_sum = np.empty(num_dataset_genes, dtype=np.uint64)
+                squared_counts_sum = np.empty(num_dataset_genes,
+                                              dtype=np.uint64)
                 gene_mean_and_variance_csr(
                     data=X.data, indices=X.indices, indptr=X.indptr,
                     cell_indices=cell_indices, num_cells=num_cells,
-                    num_dataset_genes=num_dataset_genes, mean=mean, var=var,
+                    num_dataset_genes=num_dataset_genes,
+                    thread_sums=thread_sums, counts_sum=counts_sum,
+                    squared_counts_sum=squared_counts_sum, mean=mean, var=var,
                     nonzero_count=nonzero_count, num_threads=num_threads)
             else:
                 if cell_mask is None:
@@ -11452,6 +11463,8 @@ class SingleCell:
                     data=X.data, indices=X.indices, indptr=X.indptr,
                     num_cells=num_cells, num_dataset_genes=num_dataset_genes,
                     cell_indices=cell_indices, clip_val=clip_val,
+                    thread_sums=thread_sums, counts_sum=counts_sum,
+                    squared_counts_sum=squared_counts_sum,
                     batch_counts_sum=batch_counts_sum,
                     squared_batch_counts_sum=squared_batch_counts_sum,
                     num_threads=num_threads)

{brisc-0.2.1 → brisc-0.2.2/brisc.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: brisc
-Version: 0.2.1
+Version: 0.2.2
 Summary: Fast single-cell toolkit
 Author-email: Michael Wainberg <m.wainberg@utoronto.ca>
 License-Expression: MIT

{brisc-0.2.1 → brisc-0.2.2}/brisc.egg-info/SOURCES.txt

@@ -1,3 +1,6 @@
+.git_archival.txt
+.gitattributes
+.gitignore
 LICENSE
 MANIFEST.in
 README.md
@@ -37,7 +40,11 @@ brisc.egg-info/PKG-INFO
 brisc.egg-info/SOURCES.txt
 brisc.egg-info/dependency_links.txt
 brisc.egg-info/requires.txt
+brisc.egg-info/scm_file_list.json
+brisc.egg-info/scm_version.json
 brisc.egg-info/top_level.txt
+tests/conftest.py
+tests/make_test_data.py
 tests/test_cluster.py
 tests/test_de.py
 tests/test_embedding.py

brisc-0.2.2/brisc.egg-info/scm_file_list.json

@@ -0,0 +1,110 @@
+{
+  "files": [
+    ".git_archival.txt",
+    ".gitattributes",
+    ".github/workflows/build-and-publish.yml",
+    ".github/workflows/docs.yml",
+    ".gitignore",
+    "LICENSE",
+    "MANIFEST.in",
+    "README.md",
+    "brisc/__init__.py",
+    "brisc/colorspacious.py",
+    "brisc/concatenate.py",
+    "brisc/cysparse.pyx",
+    "brisc/cyutils.pxd",
+    "brisc/cyutils.pyx",
+    "brisc/de.py",
+    "brisc/doublets.pyx",
+    "brisc/embed.pyx",
+    "brisc/harmonize.pyx",
+    "brisc/hdf5.pyx",
+    "brisc/hvg.pyx",
+    "brisc/kmeans.pyx",
+    "brisc/knn.pyx",
+    "brisc/labels.pyx",
+    "brisc/leiden.pyx",
+    "brisc/loess.pyx",
+    "brisc/norm_factors.pyx",
+    "brisc/normalize.pyx",
+    "brisc/pca.pyx",
+    "brisc/pseudobulk.py",
+    "brisc/pseudobulk_and_markers.pyx",
+    "brisc/qc.pyx",
+    "brisc/single_cell.py",
+    "brisc/snn.pyx",
+    "brisc/sparse.py",
+    "brisc/threadpoolctl.py",
+    "brisc/type_aliases.py",
+    "brisc/utils.py",
+    "brisc/validated_dict.py",
+    "docs/.gitignore",
+    "docs/Makefile",
+    "docs/deploy.sh",
+    "docs/deploy/_worker.js",
+    "docs/deploy/robots.txt",
+    "docs/make_tutorial_data.py",
+    "docs/source/_pygments_ristretto.py",
+    "docs/source/_static/css/custom.css",
+    "docs/source/_static/fonts/DejaVuSansMono-subset.woff2",
+    "docs/source/_static/images/favicon.ico",
+    "docs/source/_static/images/runner_logo_dark.svg",
+    "docs/source/_static/images/runner_logo_light.svg",
+    "docs/source/_static/images/runner_title_wide_dark.svg",
+    "docs/source/_static/images/runner_title_wide_light.svg",
+    "docs/source/_static/js/api-scrollspy.js",
+    "docs/source/_static/js/benchmark-data.js",
+    "docs/source/_static/js/benchmark.js",
+    "docs/source/_static/js/details-anchor.js",
+    "docs/source/_static/js/header-reveal.js",
+    "docs/source/_static/js/theme-fix.js",
+    "docs/source/_templates/guide-nav.html",
+    "docs/source/api/index.md",
+    "docs/source/conf.py",
+    "docs/source/de/analysis.rst",
+    "docs/source/de/index.rst",
+    "docs/source/de/io.rst",
+    "docs/source/de/properties.rst",
+    "docs/source/index.md",
+    "docs/source/installation.md",
+    "docs/source/pseudobulk/analysis.rst",
+    "docs/source/pseudobulk/data_access.rst",
+    "docs/source/pseudobulk/dictionary_interface.rst",
+    "docs/source/pseudobulk/index.rst",
+    "docs/source/pseudobulk/io.rst",
+    "docs/source/pseudobulk/manipulation.rst",
+    "docs/source/pseudobulk/properties.rst",
+    "docs/source/singlecell/analysis.rst",
+    "docs/source/singlecell/data_access.rst",
+    "docs/source/singlecell/index.rst",
+    "docs/source/singlecell/io.rst",
+    "docs/source/singlecell/manipulation.rst",
+    "docs/source/singlecell/properties.rst",
+    "docs/source/singlecell/typedefs.rst",
+    "docs/source/tutorials/basic_workflow.md",
+    "docs/source/tutorials/data_manipulation.md",
+    "docs/source/tutorials/differential_expression.md",
+    "docs/source/tutorials/images/label_transfer.png",
+    "docs/source/tutorials/images/markers.png",
+    "docs/source/tutorials/images/umap.png",
+    "docs/source/tutorials/images/volcano.png",
+    "docs/source/tutorials/index.md",
+    "docs/source/tutorials/integration_and_label_transfer.md",
+    "docs/source/tutorials/interoperability.md",
+    "pyproject.toml",
+    "setup.py",
+    "tests/conftest.py",
+    "tests/make_test_data.py",
+    "tests/test_cluster.py",
+    "tests/test_de.py",
+    "tests/test_embedding.py",
+    "tests/test_hvg.py",
+    "tests/test_integration.py",
+    "tests/test_markers.py",
+    "tests/test_neighbors.py",
+    "tests/test_normalize.py",
+    "tests/test_pca.py",
+    "tests/test_pseudobulk.py",
+    "tests/test_qc.py"
+  ]
+}

brisc-0.2.2/brisc.egg-info/scm_version.json

@@ -0,0 +1,8 @@
+{
+  "tag": "0.2.2",
+  "distance": 0,
+  "node": "g0a16f206c3772da54539139bb13d522c030a32ff",
+  "dirty": false,
+  "branch": "HEAD",
+  "node_date": "2026-10-23"
+}

{brisc-0.2.1 → brisc-0.2.2}/pyproject.toml

@@ -1,10 +1,10 @@
 [build-system]
-requires = ["setuptools>=61", "wheel", "Cython>=3.0", "numpy>=1.17.0", "scipy>=0.16.0"]
+requires = ["setuptools>=64", "setuptools-scm>=8", "wheel", "Cython>=3.0", "numpy>=1.17.0", "scipy>=0.16.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "brisc"
-version = "0.2.1"
+dynamic = ["version"]
 description = "Fast single-cell toolkit"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -39,6 +39,8 @@ include = ["brisc", "brisc.*"]
 [tool.setuptools.exclude-package-data]
 "*" = ["*.cpp", "*.c"]
 
+[tool.setuptools_scm]
+
 [tool.cibuildwheel]
 build = "cp39-* cp310-* cp311-* cp312-* cp313-* cp314-* cp315-*"
 skip = "*musllinux* *i686* *ppc64le* *s390x* cp314t-* cp315t-*"

brisc-0.2.2/tests/conftest.py

@@ -0,0 +1,30 @@
+"""Shared fixtures for the brisc test suite."""
+import shutil
+import pytest
+
+_FIXTURE = "Parse_10K_PBMC_cytokines.h5ad"
+_HF_REPO = "dashingcell/Parse_10K_PBMC_cytokines"
+
+
+@pytest.fixture(scope="session")
+def data_path():
+    # Pull the ~0.15 GB fixture from the public HF dataset (cached under
+    # ~/.cache/huggingface).
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(_HF_REPO, _FIXTURE, repo_type="dataset")
+
+
+@pytest.fixture(scope="session")
+def sc_orig(data_path):
+    from brisc import SingleCell
+    return SingleCell(data_path)
+
+
+@pytest.fixture(scope="session")
+def require_r():
+    # R is installed in CI, so a broken ryp/limma should fail loudly, not skip.
+    if shutil.which("R") is None and shutil.which("Rscript") is None:
+        pytest.skip("R not installed")
+    import ryp
+    ryp.r("suppressMessages(library(limma))")
+    return True

brisc-0.2.2/tests/make_test_data.py

@@ -0,0 +1,105 @@
+"""
+Build the brisc test fixture and publish it to the HF dataset CI pulls from.
+Kept for provenance -- CI never runs this; conftest fetches the finished file.
+
+PBS vs IFN-gamma, balanced 50/50, subsampled to 10,000 cells x full genes:
+24 pseudobulk samples (2 x 12 donors) instead of 1,092, so DE by
+(sample, cell_type) stays well-populated.
+
+Run where the 10M source (~227 GB) lives; skips the download if it's present.
+Needs huggingface_hub and a write token (HF_TOKEN or `huggingface-cli login`).
+"""
+import subprocess
+from pathlib import Path
+import polars as pl
+from brisc import SingleCell
+from huggingface_hub import HfApi
+
+REPO_ID = 'dashingcell/Parse_10K_PBMC_cytokines'
+N_CELLS = 10_000
+CONDITIONS = ['PBS', 'IFN-gamma']
+OBS_COLUMNS = ['sample', 'donor', 'cell_type', 'treatment', 'cytokine']
+
+data_dir = Path.home() / 'single-cell' / 'Parse'
+src = data_dir / 'Parse_10M_PBMC_cytokines.h5ad'
+out = data_dir / f'Parse_{N_CELLS // 1000}K_PBMC_cytokines.h5ad'
+
+CARD = """\
+---
+license: other
+license_name: parse-biosciences-data-terms
+license_link: https://www.parsebiosciences.com/datasets/
+pretty_name: Parse 10K PBMC Cytokines (IFN-γ vs PBS)
+tags:
+- single-cell
+- scRNA-seq
+- PBMC
+- immunology
+- cytokines
+- brisc
+size_categories:
+- 1K<n<10K
+---
+
+# Parse 10K PBMC Cytokines (IFN-γ vs PBS)
+
+Balanced 10,000-cell subsample of the ~10 million cell cytokine stimulation
+dataset from Parse Biosciences.
+
+In the original experiment, peripheral blood mononuclear cells (PBMCs) from
+twelve healthy donors were treated with either one of 90 different cytokines or
+a phosphate-buffered saline (PBS) control for 24 hours, yielding
+(90 + 1) × 12 = 1,092 experimental conditions.
+
+For the brisc test suite, the data is restricted to the interferon-gamma
+(IFN-γ) versus PBS contrast and subsampled to 5,000 cells per condition,
+stratified by cell type. The result is 10,000 cells × 40,352 genes of raw
+counts across 24 samples (12 donors × 2 conditions) and 18 annotated cell
+types, with `obs` columns `sample`, `donor`, `cell_type`, `treatment`, and
+`cytokine`.
+
+## Source
+
+- https://www.parsebiosciences.com/datasets/10-million-human-pbmcs-in-a-single-experiment/
+- https://www.biorxiv.org/content/10.64898/2025.12.12.693897v1
+
+All credit to Parse Biosciences.
+"""
+
+def remove_unused_categories(s: pl.Series) -> pl.Expr:
+    """polars-Enum version of pandas' remove_unused_categories: rebuild each
+    Enum from the values present."""
+    present = set(s.unique().to_list())
+    kept = [c for c in s.dtype.categories.to_list() if c in present]
+    return pl.col(s.name).cast(pl.Enum(kept))
+
+# Fetch the 10M source once
+data_dir.mkdir(parents=True, exist_ok=True)
+if not src.exists():
+    url = 'https://parse-wget.s3.us-west-2.amazonaws.com/10m/' + src.name
+    subprocess.run(['wget', '-O', str(src), url], check=True)
+
+# Build the balanced 10K fixture
+sc = SingleCell(src, obs_columns=OBS_COLUMNS)\
+    .filter_obs(pl.col('cytokine').is_in(CONDITIONS))
+half = N_CELLS // 2
+pbs = sc.filter_obs(pl.col('cytokine') == 'PBS')\
+        .subsample_obs(n=half, by_column='cell_type', QC_column=None, seed=0)
+ifn = sc.filter_obs(pl.col('cytokine') == 'IFN-gamma')\
+        .subsample_obs(n=half, by_column='cell_type', QC_column=None, seed=0)
+fixture = pbs.concat_obs(ifn)
+
+# Drop unused Enum categories (sample 1,092 -> 24, cytokine 91 -> 2)
+enum_cols = [name for name, dt in fixture.obs.schema.items() if dt == pl.Enum]
+fixture = fixture.with_columns_obs(
+    [remove_unused_categories(fixture.obs[c]) for c in enum_cols])
+fixture.save(out, overwrite=True)
+
+# Publish data + card to the HF dataset repo (no-op if unchanged)
+api = HfApi()
+api.create_repo(REPO_ID, repo_type='dataset', exist_ok=True)
+api.upload_file(path_or_fileobj=str(out), path_in_repo=out.name,
+                repo_id=REPO_ID, repo_type='dataset')
+api.upload_file(path_or_fileobj=CARD.encode(), path_in_repo='README.md',
+                repo_id=REPO_ID, repo_type='dataset')
+print(f'Published to https://huggingface.co/datasets/{REPO_ID}')