From 880a8823dd343352339118487b3d865339937531 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 13:27:36 +0200 Subject: [PATCH 01/54] Add GPU hamming improvements --- src/scirpy/ir_dist/metrics.py | 256 +++++++++++++++++------ src/scirpy/tests/test_ir_dist_metrics.py | 10 +- 2 files changed, 201 insertions(+), 65 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 61a64d1b8..e67fce2f8 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -776,15 +776,15 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. For performance reasons, the computation of the final result matrix is split up into several blocks. The parameter - gpu_n_blocks determines the number of those blocks. The parameter gpu_block_width determines how much GPU memory + gpu_col_blocks determines the number of those blocks. The parameter gpu_block_width determines how much GPU memory is reserved for the computed result of each block in SPARSE representation. - E.g. there is a 1000x1000 (dense represenation) not yet computed result matrix with gpu_n_blocks=10 and gpu_block_width=20. + E.g. there is a 1000x1000 (dense represenation) not yet computed result matrix with gpu_col_blocks=10 and gpu_block_width=20. Then the result matrix is computed in 10 blocks of 1000x100 (dense representation). Each of these blocks needs to fit into a 1000x20 block in SPARSE representation once computed and this 1000x20 block needs to fit into GPU memory. So there shouldn't be a resulting row in a block that has more than 20 values <= cutoff. - The parameter gpu_block_width should be chosen based on the available GPU memory. Choosing lower values for gpu_n_blocks increases + The parameter gpu_block_width should be chosen based on the available GPU memory. Choosing lower values for gpu_col_blocks increases the performance but also increases the risk of running out of reserved memory, since the result blocks that need to fit into the reserved GPU memory in sparse representation get bigger. @@ -793,29 +793,43 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): cutoff: Will eleminate distances > cutoff to make efficient use of sparse matrices. - gpu_n_blocks: + gpu_col_blocks: Number of blocks in which the final result matrix should be computed. Each block reserves GPU memory in which the computed result block has to fit in sparse representation. Lower values give better performance but increase the risk of running out of reserved memory. This value should be chosen based on the estimated sparsity of the result matrix and the size of the GPU device memory. + gpu_row_blocks: + Number of row blocks in which the final result matrix should be computed. Higher values can reduce + memory pressure and improve length homogeneity within a row block, but also add block-management overhead. gpu_block_width: Maximum width of blocks in which the final result matrix should be computed. Each block reserves GPU memory in which the computed result block has to fit in sparse representation. Higher values allow for a lower - number of result blocks (gpu_n_blocks) which increases the performance. This value should be chosen based on + number of result blocks (gpu_col_blocks) which increases the performance. This value should be chosen based on the GPU device memory. + benchmark: + If True, print coarse-grained timings for the main phases of the GPU hamming calculation. """ def __init__( self, *, cutoff: int = 2, - gpu_n_blocks: int = 10, + gpu_col_blocks: int = 10, + gpu_row_blocks: int = 1, gpu_block_width: int = 1000, + benchmark: bool = False, ): super().__init__(n_jobs=1, n_blocks=1) + if cutoff > 126: + raise ValueError( + "GPUHammingDistanceCalculator only supports cutoff <= 126 because the intermediate " + "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." + ) self.cutoff = cutoff - self.gpu_n_blocks = gpu_n_blocks + self.gpu_col_blocks = gpu_col_blocks + self.gpu_row_blocks = gpu_row_blocks self.gpu_block_width = gpu_block_width + self.benchmark = benchmark def _gpu_hamming_mat( self, @@ -854,18 +868,49 @@ def _gpu_hamming_mat( not implemented for the GPU hamming calculator yet. """ import cupy as cp + import time from tqdm import tqdm - seqs_lengths = np.vectorize(len)(seqs) - seqs_original_indices = np.argsort(seqs_lengths) - seqs = seqs[seqs_original_indices] - - seqs2_lengths = np.vectorize(len)(seqs2) - seqs2_original_indices = np.argsort(seqs2_lengths) - seqs2 = seqs2[seqs2_original_indices] + benchmark = self.benchmark + benchmark_timers: dict[str, float] = {} + benchmark_wall_start = time.perf_counter() + + def _benchmark_add(label: str, start: float) -> None: + if benchmark: + benchmark_timers[label] = benchmark_timers.get(label, 0.0) + (time.perf_counter() - start) + + n_col_blocks = self.gpu_col_blocks + n_row_blocks = self.gpu_row_blocks + + phase_start = time.perf_counter() + seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) + seqs_sorted_per_block = [] + seqs_original_indices_blocks = [] + + for seqs_block in seqs_blocks: + seqs_block_lengths = np.vectorize(len)(seqs_block) + seqs_block_sort_indices = np.argsort(seqs_block_lengths) + seqs_sorted_per_block.append(seqs_block[seqs_block_sort_indices]) + seqs_original_indices_blocks.append(cp.asarray(seqs_block_sort_indices.astype(np.int32))) + + seqs = np.concatenate(seqs_sorted_per_block) + + seqs2_blocks = np.array_split(np.asarray(seqs2), n_col_blocks) + seqs2_sorted_per_block = [] + seqs2_original_indices_blocks = [] + seqs2_block_start = 0 + + for seqs2_block in seqs2_blocks: + seqs2_block_lengths = np.vectorize(len)(seqs2_block) + seqs2_block_sort_indices = np.argsort(seqs2_block_lengths) + seqs2_sorted_per_block.append(seqs2_block[seqs2_block_sort_indices]) + seqs2_original_indices_blocks.append( + cp.asarray((seqs2_block_sort_indices + seqs2_block_start).astype(np.int32)) + ) + seqs2_block_start += len(seqs2_block) - seqs_original_indices = cp.asarray(seqs_original_indices, dtype=np.int32) - seqs2_original_indices = cp.asarray(seqs2_original_indices, dtype=np.int32) + seqs2 = np.concatenate(seqs2_sorted_per_block) + _benchmark_add("preprocess_and_sort", phase_start) is_symmetric = False @@ -881,17 +926,23 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. L[i] = len(seq) return mat, L + phase_start = time.perf_counter() try: seqs_mat1, seqs_L1 = _seqs2mat_fast(seqs, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat_fast(seqs2, max_len=max_seq_len) + _benchmark_add("seqs2mat_fast", phase_start) except UnicodeError: + _benchmark_add("seqs2mat_fast", phase_start) logging.info( "UnicodeError error occurred while converting sequences, retrying with implementation for non ascii sequences" ) unique_characters = "".join(sorted({char for string in (*seqs, *seqs2) for char in string})) + phase_start = time.perf_counter() seqs_mat1, seqs_L1 = _seqs2mat(seqs, alphabet=unique_characters, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat(seqs2, alphabet=unique_characters, max_len=max_seq_len) + _benchmark_add("seqs2mat_fallback", phase_start) + phase_start = time.perf_counter() hamming_kernel = cp.RawKernel( r""" extern "C" __global__ __launch_bounds__(256) @@ -933,12 +984,15 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. if(val1 != val2) { distance++; + if (distance > cutoff + 1) { + break; + } } } if (distance <= cutoff + 1) { int seqs2_original_index = seqs2_original_indices[col]; - data[seqs_original_index * data_cols + row_end_index] = distance; - indices[seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; + data[(long long)seqs_original_index * data_cols + row_end_index] = distance; + indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; row_end_index++; } } @@ -970,24 +1024,32 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. int data_index = row_start + col; if ((data_index < data_rows) && (col < row_end_index)) { - data[data_index] = data_matrix[row * data_matrix_cols + col]; - indices[data_index] = indices_matrix[row * indices_matrix_cols + col]; + data[data_index] = data_matrix[(long long)row * data_matrix_cols + col]; + indices[data_index] = indices_matrix[(long long)row * indices_matrix_cols + col]; } } } - """, + """, "create_csr_kernel", ) - - def calc_block_gpu( - seqs_mat1, seqs_mat2_block, seqs_L1_block, seqs_L2, seqs2_original_indices_blocks, block_offset + _benchmark_add("kernel_setup", phase_start) + + def calc_col_block_gpu( + seqs_mat1, + seqs_mat2_block, + seqs_L1_block, + seqs_L2, + seqs_original_indices_block, + seqs2_original_indices_block, + block_offset, ): - import cupy as cp - + block_start = time.perf_counter() + phase_start = time.perf_counter() d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8)) d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8)) d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32)) d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32)) + _benchmark_add("gpu_transfer_to_device", phase_start) # Due to performance reasons and since we expect the result matrix to be very sparse, we # set a maximum result width for the current block @@ -1005,9 +1067,12 @@ def calc_block_gpu( d_data_matrix_cols = max_block_width d_indices_matrix_cols = max_block_width + phase_start = time.perf_counter() d_seqs_mat1_transposed = cp.transpose(d_seqs_mat1).copy() d_seqs_mat2_transposed = cp.transpose(d_seqs_mat2).copy() + _benchmark_add("gpu_transpose_copy", phase_start) + phase_start = time.perf_counter() hamming_kernel( (blocks_per_grid,), (threads_per_block,), @@ -1016,8 +1081,8 @@ def calc_block_gpu( d_seqs_mat2_transposed, d_seqs_L1, d_seqs_L2, - seqs_original_indices, - seqs2_original_indices_blocks, + seqs_original_indices_block, + seqs2_original_indices_block, self.cutoff, d_data_matrix, d_indices_matrix, @@ -1032,8 +1097,12 @@ def calc_block_gpu( is_symmetric, ), ) + cp.cuda.Stream.null.synchronize() + _benchmark_add("gpu_hamming_kernel", phase_start) + phase_start = time.perf_counter() row_element_counts = d_row_element_counts.get() + _benchmark_add("gpu_row_counts_get", phase_start) row_max_len = np.max(row_element_counts) row_element_sum = np.sum(row_element_counts, dtype=np.int64) @@ -1064,6 +1133,7 @@ def calc_block_gpu( blocks_per_grid_y = (d_data_matrix.shape[1] + threads_per_block[1] - 1) // threads_per_block[1] blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y) + phase_start = time.perf_counter() create_csr_kernel( (blocks_per_grid_x, blocks_per_grid_y), threads_per_block, @@ -1079,51 +1149,30 @@ def calc_block_gpu( d_indices_matrix.shape[1], ), ) + cp.cuda.Stream.null.synchronize() + _benchmark_add("gpu_create_csr_kernel", phase_start) + phase_start = time.perf_counter() data = d_data.get() indptr = d_indptr.get() indices = d_indices.get() + _benchmark_add("gpu_transfer_to_host", phase_start) res = csr_matrix((data, indices, indptr), shape=(seqs_mat1.shape[0], seqs_mat2.shape[0])) + _benchmark_add("gpu_col_blocks_total", block_start) return res - # Set the number of blocks for the calculation. A higher number can be more memory friendly, whereas - # a lower number can improve the performance. - n_blocks = self.gpu_n_blocks - - seqs_mat2_blocks = np.array_split(seqs_mat2, n_blocks) - seqs_L2_blocks = np.array_split(seqs_L2, n_blocks) - seqs2_original_indices_blocks = np.array_split(seqs2_original_indices, n_blocks) - result_blocks = [None] * n_blocks - - block_offset = start_column + phase_start = time.perf_counter() + seqs_mat1_blocks = np.array_split(seqs_mat1, n_row_blocks) + seqs_L1_blocks = np.array_split(seqs_L1, n_row_blocks) + seqs_mat2_blocks = np.array_split(seqs_mat2, n_col_blocks) + seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) + _benchmark_add("matrix_block_split", phase_start) logging.info( - f"\nStart GPU calculations for {n_blocks} sparse matrix result blocks of max width {self.gpu_block_width}:" + f"\nStart GPU calculations for {n_row_blocks} row blocks x {n_col_blocks} column blocks of max width {self.gpu_block_width}:" ) - for i in tqdm(range(0, n_blocks), desc="Processing", unit="block"): - result_blocks[i] = calc_block_gpu( - seqs_mat1, - seqs_mat2_blocks[i], - seqs_L1, - seqs_L2_blocks[i], - seqs2_original_indices_blocks[i], - block_offset, - ) - block_offset += seqs_mat2_blocks[i].shape[0] - - num_elements = 0 - for i in range(0, len(result_blocks)): - num_elements += result_blocks[i].indptr[-1] - - assert ( - num_elements <= np.iinfo(np.int32).max - ), f"""ERROR: The overall number of result values is too high to construct the final CSR matrix by combining - the already calculated blocks. - Current number: {num_elements}, Maximum number: {np.iinfo(np.int32).max}. - Consider choosing a smaller cutoff to resolve this issue.""" - @nb.njit def csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elements): data = np.empty(num_elements, dtype=block_data[0].dtype) @@ -1147,6 +1196,7 @@ def csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elem return data, indices, indptr def csr_union(blocks): + phase_start = time.perf_counter() num_rows = blocks[0].shape[0] num_elements = sum(b.nnz for b in blocks) @@ -1157,12 +1207,92 @@ def csr_union(blocks): data, indices, indptr = csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elements) shape = blocks[0].shape - return csr_matrix((data, indices, indptr), shape=shape) + result = csr_matrix((data, indices, indptr), shape=shape) + _benchmark_add("csr_union", phase_start) + return result + + def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): + row_block_start = time.perf_counter() + result_blocks = [None] * n_col_blocks + block_offset = start_column + + for i in range(0, n_col_blocks): + result_blocks[i] = calc_col_block_gpu( + seqs_mat1_block, + seqs_mat2_blocks[i], + seqs_L1_block, + seqs_L2_blocks[i], + seqs_original_indices_block, + seqs2_original_indices_blocks[i], + block_offset, + ) + block_offset += seqs_mat2_blocks[i].shape[0] + + num_elements = 0 + for i in range(0, len(result_blocks)): + num_elements += result_blocks[i].indptr[-1] + + assert ( + num_elements <= np.iinfo(np.int32).max + ), f"""ERROR: The overall number of result values is too high to construct the final CSR matrix by combining + the already calculated blocks. + Current number: {num_elements}, Maximum number: {np.iinfo(np.int32).max}. + Consider choosing a smaller cutoff to resolve this issue.""" + + result_sparse = csr_union(result_blocks) + phase_start = time.perf_counter() + result_sparse.sort_indices() + _benchmark_add("row_block_sort_indices", phase_start) + _benchmark_add("row_block_total", row_block_start) + return result_sparse + + row_blocks = [None] * n_row_blocks + phase_start = time.perf_counter() + with tqdm(total=n_row_blocks * n_col_blocks, desc="Processing", unit="block") as progress_bar: + for row_block_idx in range(n_row_blocks): + row_blocks[row_block_idx] = calc_row_block_gpu( + seqs_mat1_blocks[row_block_idx], + seqs_L1_blocks[row_block_idx], + seqs_original_indices_blocks[row_block_idx], + ) + progress_bar.update(n_col_blocks) + _benchmark_add("all_row_blocks", phase_start) - result_sparse = csr_union(result_blocks) + phase_start = time.perf_counter() + result_sparse = scipy.sparse.vstack(row_blocks, format="csr") + _benchmark_add("final_vstack", phase_start) row_element_counts_gpu = np.diff(result_sparse.indptr) + phase_start = time.perf_counter() result_sparse.sort_indices() + _benchmark_add("final_sort_indices", phase_start) + + if benchmark: + total_wall = time.perf_counter() - benchmark_wall_start + print("GPU hamming benchmark") + print(f" total_wall: {total_wall:.3f}s") + for label in ( + "preprocess_and_sort", + "seqs2mat_fast", + "seqs2mat_fallback", + "kernel_setup", + "matrix_block_split", + "gpu_transfer_to_device", + "gpu_transpose_copy", + "gpu_hamming_kernel", + "gpu_row_counts_get", + "gpu_create_csr_kernel", + "gpu_transfer_to_host", + "gpu_col_blocks_total", + "csr_union", + "row_block_sort_indices", + "row_block_total", + "all_row_blocks", + "final_vstack", + "final_sort_indices", + ): + if label in benchmark_timers: + print(f" {label}: {benchmark_timers[label]:.3f}s") # Returns the results in a way that fits the current interface, could be improved later return [result_sparse.data], [result_sparse.indices], row_element_counts_gpu, np.array([None]) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index f970df33c..2a5c3cbf7 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -957,13 +957,19 @@ def test_hamming_long_sequence(): @pytest.mark.gpu def test_gpu_hamming_long_sequence(): """Regression test for #626 and #682""" - hamming_calculator = GPUHammingDistanceCalculator(cutoff=50, gpu_n_blocks=1, gpu_block_width=50) + hamming_calculator = GPUHammingDistanceCalculator(cutoff=50, gpu_col_blocks=1, gpu_block_width=50) seq1 = np.array(["A" * 128, "AAB", "AABB", "ABA"]) seq2 = np.array(["A" * 128, "ABBB", "ABBB"]) res = hamming_calculator.calc_dist_mat(seq1, seq2) assert isinstance(res, scipy.sparse.csr_matrix) +@pytest.mark.gpu +def test_gpu_hamming_cutoff_guard(): + with pytest.raises(ValueError, match="cutoff <= 126"): + GPUHammingDistanceCalculator(cutoff=127) + + def test_hamming_histogram_reference(): from . import TESTDATA @@ -990,7 +996,7 @@ def test_gpu_hamming_reference(): seqs = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_seqs.npy") reference_result = scipy.sparse.load_npz(TESTDATA / "hamming_test_data/hamming_WU3k_csr_result.npz") - gpu_hamming_calculator = GPUHammingDistanceCalculator(cutoff=2, gpu_n_blocks=5, gpu_block_width=500) + gpu_hamming_calculator = GPUHammingDistanceCalculator(cutoff=2, gpu_col_blocks=5, gpu_block_width=500) res = gpu_hamming_calculator.calc_dist_mat(seqs, seqs) assert np.array_equal(res.data, reference_result.data) From 3d403917371bd3cf1caed9ba7761a2e2fe84de12 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 13:58:56 +0200 Subject: [PATCH 02/54] Change cutoff guard to <=125 --- src/scirpy/ir_dist/metrics.py | 4 ++-- src/scirpy/tests/test_ir_dist_metrics.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index e67fce2f8..f54e9ddba 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -820,9 +820,9 @@ def __init__( benchmark: bool = False, ): super().__init__(n_jobs=1, n_blocks=1) - if cutoff > 126: + if cutoff > 125: raise ValueError( - "GPUHammingDistanceCalculator only supports cutoff <= 126 because the intermediate " + "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." ) self.cutoff = cutoff diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 2a5c3cbf7..b659cc8bd 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -966,8 +966,8 @@ def test_gpu_hamming_long_sequence(): @pytest.mark.gpu def test_gpu_hamming_cutoff_guard(): - with pytest.raises(ValueError, match="cutoff <= 126"): - GPUHammingDistanceCalculator(cutoff=127) + with pytest.raises(ValueError, match="cutoff <= 125"): + GPUHammingDistanceCalculator(cutoff=126) def test_hamming_histogram_reference(): From d053bee4c70dd9c8c4e7a0ee66574237efc1bc0b Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:06:21 +0200 Subject: [PATCH 03/54] Add parameter guards for block size parameters --- src/scirpy/ir_dist/metrics.py | 6 ++++++ src/scirpy/tests/test_ir_dist_metrics.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index f54e9ddba..635dfe52d 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -825,6 +825,12 @@ def __init__( "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." ) + if gpu_col_blocks < 1: + raise ValueError("`gpu_col_blocks` must be >= 1.") + if gpu_row_blocks < 1: + raise ValueError("`gpu_row_blocks` must be >= 1.") + if gpu_block_width < 1: + raise ValueError("`gpu_block_width` must be >= 1.") self.cutoff = cutoff self.gpu_col_blocks = gpu_col_blocks self.gpu_row_blocks = gpu_row_blocks diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index b659cc8bd..50104f852 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -970,6 +970,20 @@ def test_gpu_hamming_cutoff_guard(): GPUHammingDistanceCalculator(cutoff=126) +@pytest.mark.gpu +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"gpu_col_blocks": 0}, "`gpu_col_blocks` must be >= 1"), + ({"gpu_row_blocks": 0}, "`gpu_row_blocks` must be >= 1"), + ({"gpu_block_width": 0}, "`gpu_block_width` must be >= 1"), + ], +) +def test_gpu_hamming_block_parameter_guards(kwargs, message): + with pytest.raises(ValueError, match=message): + GPUHammingDistanceCalculator(**kwargs) + + def test_hamming_histogram_reference(): from . import TESTDATA From 5200588f41569b15704fab6fef37fdb992305cf6 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:14:09 +0200 Subject: [PATCH 04/54] Add tests for different block size parameters --- src/scirpy/tests/test_ir_dist_metrics.py | 32 +++++++++++++++--------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 50104f852..193fb2b18 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -954,6 +954,16 @@ def test_hamming_long_sequence(): assert isinstance(res, scipy.sparse.csr_matrix) +def test_hamming_histogram_reference(): + from . import TESTDATA + + seqs = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_seqs.npy") + hamming_calculator = HammingDistanceCalculator(2, 2, 100, normalize=True, histogram=True) + row_mins_ref = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_histogram_result.npy") + _, _, _, row_mins = hamming_calculator._hamming_mat(seqs=seqs, seqs2=seqs) + assert np.array_equal(row_mins_ref, row_mins) + + @pytest.mark.gpu def test_gpu_hamming_long_sequence(): """Regression test for #626 and #682""" @@ -984,16 +994,6 @@ def test_gpu_hamming_block_parameter_guards(kwargs, message): GPUHammingDistanceCalculator(**kwargs) -def test_hamming_histogram_reference(): - from . import TESTDATA - - seqs = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_seqs.npy") - hamming_calculator = HammingDistanceCalculator(2, 2, 100, normalize=True, histogram=True) - row_mins_ref = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_histogram_result.npy") - _, _, _, row_mins = hamming_calculator._hamming_mat(seqs=seqs, seqs2=seqs) - assert np.array_equal(row_mins_ref, row_mins) - - def test_tcrdist_histogram_not_implemented(): # Change once histogram is implemented for tcrdist with pytest.raises(NotImplementedError, match=None): @@ -1003,14 +1003,22 @@ def test_tcrdist_histogram_not_implemented(): @pytest.mark.gpu -def test_gpu_hamming_reference(): +@pytest.mark.parametrize( + "kwargs", + [ + {"gpu_col_blocks": 5, "gpu_block_width": 500}, + {"gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, + {"gpu_col_blocks": 7, "gpu_block_width": 500}, + ], +) +def test_gpu_hamming_reference(kwargs): # test hamming distance against reference implementation from . import TESTDATA seqs = np.load(TESTDATA / "hamming_test_data/hamming_WU3k_seqs.npy") reference_result = scipy.sparse.load_npz(TESTDATA / "hamming_test_data/hamming_WU3k_csr_result.npz") - gpu_hamming_calculator = GPUHammingDistanceCalculator(cutoff=2, gpu_col_blocks=5, gpu_block_width=500) + gpu_hamming_calculator = GPUHammingDistanceCalculator(cutoff=2, **kwargs) res = gpu_hamming_calculator.calc_dist_mat(seqs, seqs) assert np.array_equal(res.data, reference_result.data) From 5b89ab22ac192aaa864d0ef3712f17b10fc67bc9 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:24:30 +0200 Subject: [PATCH 05/54] Adapt GPUHammingDistanceCalculator doc string --- src/scirpy/ir_dist/metrics.py | 37 ++++++++++++++++------------------- 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 635dfe52d..1cb2e1532 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -775,18 +775,19 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): The code of this class is based on `pwseqdist `_. Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. - For performance reasons, the computation of the final result matrix is split up into several blocks. The parameter - gpu_col_blocks determines the number of those blocks. The parameter gpu_block_width determines how much GPU memory - is reserved for the computed result of each block in SPARSE representation. + For performance reasons, the computation of the final result matrix is split into a grid of row and column blocks. + `gpu_row_blocks` splits the query sequences into horizontal chunks, while `gpu_col_blocks` splits the reference + sequences into vertical chunks. Each row/column block pair is computed on the GPU and converted to a sparse CSR + block before the blocks are combined again. - E.g. there is a 1000x1000 (dense represenation) not yet computed result matrix with gpu_col_blocks=10 and gpu_block_width=20. - Then the result matrix is computed in 10 blocks of 1000x100 (dense representation). Each of these blocks needs to fit into - a 1000x20 block in SPARSE representation once computed and this 1000x20 block needs to fit into GPU memory. So there shouldn't - be a resulting row in a block that has more than 20 values <= cutoff. + `gpu_block_width` controls how many sparse result entries are reserved per row for each row/column block pair. + For example, with a 1000 x 1000 dense result matrix, `gpu_row_blocks=2`, `gpu_col_blocks=10`, and + `gpu_block_width=20`, each GPU block covers roughly 500 x 100 dense comparisons and reserves enough GPU memory + for at most 20 retained distances below the cutoff per row in that block. Therefore, no row within a single row/column block pair + may contain more than `gpu_block_width` distances `<= cutoff`. - The parameter gpu_block_width should be chosen based on the available GPU memory. Choosing lower values for gpu_col_blocks increases - the performance but also increases the risk of running out of reserved memory, since the result blocks that need to fit into the - reserved GPU memory in sparse representation get bigger. + Larger numbers of row or column blocks reduce per-block memory pressure but add block-management overhead. Larger + values for `gpu_block_width` can avoid sparse buffer overflows but require more GPU memory. Parameters ---------- @@ -794,18 +795,14 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): Will eleminate distances > cutoff to make efficient use of sparse matrices. gpu_col_blocks: - Number of blocks in which the final result matrix should be computed. Each block reserves GPU memory - in which the computed result block has to fit in sparse representation. Lower values give better performance - but increase the risk of running out of reserved memory. This value should be chosen based on the - estimated sparsity of the result matrix and the size of the GPU device memory. + Number of column blocks used to split the final result matrix. Lower values reduce block-management overhead + but increase the size of each GPU block and therefore memory pressure. gpu_row_blocks: - Number of row blocks in which the final result matrix should be computed. Higher values can reduce - memory pressure and improve length homogeneity within a row block, but also add block-management overhead. + Number of row blocks used to split the final result matrix. Higher values reduce per-block memory pressure + and can improve sequence length homogeneity within a block, but add block-management overhead. gpu_block_width: - Maximum width of blocks in which the final result matrix should be computed. Each block reserves GPU memory - in which the computed result block has to fit in sparse representation. Higher values allow for a lower - number of result blocks (gpu_col_blocks) which increases the performance. This value should be chosen based on - the GPU device memory. + Maximum number of retained sparse entries per row and row/column block pair. Higher values tolerate denser + results within a block but require more GPU memory. benchmark: If True, print coarse-grained timings for the main phases of the GPU hamming calculation. """ From 7718e2cd7b11db1c548c12a3fcbcba37b4d78185 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:39:05 +0200 Subject: [PATCH 06/54] Change sequence length array data type to int32 --- src/scirpy/ir_dist/metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 1cb2e1532..117f30721 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -923,7 +923,7 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. if max_len is None: max_len = np.max([len(s) for s in seqs]) mat = -1 * np.ones((len(seqs), max_len), dtype=np.int8) - L = np.zeros(len(seqs), dtype=np.int8 if max_len <= np.iinfo(np.int8).max else np.int16) + L = np.zeros(len(seqs), dtype=np.int32) for i, seq in enumerate(seqs): mat[i][0 : len(seq)] = np.frombuffer(seq.encode("ascii"), dtype=np.uint8) L[i] = len(seq) From 9d3252d710c7e125e3c94eb3fd737395a8f92c36 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:48:19 +0200 Subject: [PATCH 07/54] Avoid unnecessary dtype copies in GPU hamming transfers --- src/scirpy/ir_dist/metrics.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 117f30721..0416091f2 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1048,10 +1048,10 @@ def calc_col_block_gpu( ): block_start = time.perf_counter() phase_start = time.perf_counter() - d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8)) - d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8)) - d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32)) - d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32)) + d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8, copy=False)) + d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8, copy=False)) + d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) + d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32, copy=False)) _benchmark_add("gpu_transfer_to_device", phase_start) # Due to performance reasons and since we expect the result matrix to be very sparse, we From 80ee4c5acc53f5cf769830e800bc0884aef8d6fe Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 1 Jul 2026 14:50:12 +0200 Subject: [PATCH 08/54] Replace asserts with exceptions --- src/scirpy/ir_dist/metrics.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 0416091f2..84f63d2af 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1109,16 +1109,18 @@ def calc_col_block_gpu( row_max_len = np.max(row_element_counts) row_element_sum = np.sum(row_element_counts, dtype=np.int64) - assert ( - row_max_len <= max_block_width - ), f"""ERROR: The chosen result block width is too small to hold all result values of the current block. - Chosen width: {max_block_width}, Necessary width: {row_max_len}.""" + if row_max_len > max_block_width: + raise ValueError( + "The chosen result block width is too small to hold all result values of the current block. " + f"Chosen width: {max_block_width}, necessary width: {row_max_len}." + ) - assert ( - row_element_sum <= np.iinfo(np.int32).max - ), f"""ERROR: There are too many result values to be held by the resulting CSR matrix of the current block. - Current number: {row_element_sum}, Maximum number: {np.iinfo(np.int32).max}. - Consider choosing a smaller cutoff to resolve this issue.""" + if row_element_sum > np.iinfo(np.int32).max: + raise ValueError( + "There are too many result values to be held by the resulting CSR matrix of the current block. " + f"Current number: {row_element_sum}, maximum number: {np.iinfo(np.int32).max}. " + "Consider choosing a smaller cutoff to resolve this issue." + ) indptr = np.zeros(seqs_mat1.shape[0] + 1, dtype=np.int32) indptr[1:] = np.cumsum(row_element_counts) From dada164e6864d80463afb02ed5fe8de6508d0a9d Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 13:09:42 +0200 Subject: [PATCH 09/54] Transform guarding assert into error --- src/scirpy/ir_dist/metrics.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 84f63d2af..dbb51834f 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1237,12 +1237,13 @@ def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_blo for i in range(0, len(result_blocks)): num_elements += result_blocks[i].indptr[-1] - assert ( - num_elements <= np.iinfo(np.int32).max - ), f"""ERROR: The overall number of result values is too high to construct the final CSR matrix by combining - the already calculated blocks. - Current number: {num_elements}, Maximum number: {np.iinfo(np.int32).max}. - Consider choosing a smaller cutoff to resolve this issue.""" + if num_elements > np.iinfo(np.int32).max: + raise ValueError( + "The overall number of result values is too high to construct the final CSR matrix by combining " + "the already calculated blocks. " + f"Current number: {num_elements}, maximum number: {np.iinfo(np.int32).max}. " + "Consider choosing a smaller cutoff to resolve this issue." + ) result_sparse = csr_union(result_blocks) phase_start = time.perf_counter() From 6887bf3bcbcbe80474527f29d13516a2b10bd28a Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 13:13:49 +0200 Subject: [PATCH 10/54] Remove benchmarking option --- src/scirpy/ir_dist/metrics.py | 79 ----------------------------------- 1 file changed, 79 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index dbb51834f..70bd1a9fa 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -803,8 +803,6 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): gpu_block_width: Maximum number of retained sparse entries per row and row/column block pair. Higher values tolerate denser results within a block but require more GPU memory. - benchmark: - If True, print coarse-grained timings for the main phases of the GPU hamming calculation. """ def __init__( @@ -814,7 +812,6 @@ def __init__( gpu_col_blocks: int = 10, gpu_row_blocks: int = 1, gpu_block_width: int = 1000, - benchmark: bool = False, ): super().__init__(n_jobs=1, n_blocks=1) if cutoff > 125: @@ -832,7 +829,6 @@ def __init__( self.gpu_col_blocks = gpu_col_blocks self.gpu_row_blocks = gpu_row_blocks self.gpu_block_width = gpu_block_width - self.benchmark = benchmark def _gpu_hamming_mat( self, @@ -871,21 +867,11 @@ def _gpu_hamming_mat( not implemented for the GPU hamming calculator yet. """ import cupy as cp - import time from tqdm import tqdm - benchmark = self.benchmark - benchmark_timers: dict[str, float] = {} - benchmark_wall_start = time.perf_counter() - - def _benchmark_add(label: str, start: float) -> None: - if benchmark: - benchmark_timers[label] = benchmark_timers.get(label, 0.0) + (time.perf_counter() - start) - n_col_blocks = self.gpu_col_blocks n_row_blocks = self.gpu_row_blocks - phase_start = time.perf_counter() seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) seqs_sorted_per_block = [] seqs_original_indices_blocks = [] @@ -913,7 +899,6 @@ def _benchmark_add(label: str, start: float) -> None: seqs2_block_start += len(seqs2_block) seqs2 = np.concatenate(seqs2_sorted_per_block) - _benchmark_add("preprocess_and_sort", phase_start) is_symmetric = False @@ -929,23 +914,17 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. L[i] = len(seq) return mat, L - phase_start = time.perf_counter() try: seqs_mat1, seqs_L1 = _seqs2mat_fast(seqs, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat_fast(seqs2, max_len=max_seq_len) - _benchmark_add("seqs2mat_fast", phase_start) except UnicodeError: - _benchmark_add("seqs2mat_fast", phase_start) logging.info( "UnicodeError error occurred while converting sequences, retrying with implementation for non ascii sequences" ) unique_characters = "".join(sorted({char for string in (*seqs, *seqs2) for char in string})) - phase_start = time.perf_counter() seqs_mat1, seqs_L1 = _seqs2mat(seqs, alphabet=unique_characters, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat(seqs2, alphabet=unique_characters, max_len=max_seq_len) - _benchmark_add("seqs2mat_fallback", phase_start) - phase_start = time.perf_counter() hamming_kernel = cp.RawKernel( r""" extern "C" __global__ __launch_bounds__(256) @@ -1035,7 +1014,6 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. """, "create_csr_kernel", ) - _benchmark_add("kernel_setup", phase_start) def calc_col_block_gpu( seqs_mat1, @@ -1046,13 +1024,10 @@ def calc_col_block_gpu( seqs2_original_indices_block, block_offset, ): - block_start = time.perf_counter() - phase_start = time.perf_counter() d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8, copy=False)) d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8, copy=False)) d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32, copy=False)) - _benchmark_add("gpu_transfer_to_device", phase_start) # Due to performance reasons and since we expect the result matrix to be very sparse, we # set a maximum result width for the current block @@ -1070,12 +1045,9 @@ def calc_col_block_gpu( d_data_matrix_cols = max_block_width d_indices_matrix_cols = max_block_width - phase_start = time.perf_counter() d_seqs_mat1_transposed = cp.transpose(d_seqs_mat1).copy() d_seqs_mat2_transposed = cp.transpose(d_seqs_mat2).copy() - _benchmark_add("gpu_transpose_copy", phase_start) - phase_start = time.perf_counter() hamming_kernel( (blocks_per_grid,), (threads_per_block,), @@ -1100,12 +1072,8 @@ def calc_col_block_gpu( is_symmetric, ), ) - cp.cuda.Stream.null.synchronize() - _benchmark_add("gpu_hamming_kernel", phase_start) - phase_start = time.perf_counter() row_element_counts = d_row_element_counts.get() - _benchmark_add("gpu_row_counts_get", phase_start) row_max_len = np.max(row_element_counts) row_element_sum = np.sum(row_element_counts, dtype=np.int64) @@ -1138,7 +1106,6 @@ def calc_col_block_gpu( blocks_per_grid_y = (d_data_matrix.shape[1] + threads_per_block[1] - 1) // threads_per_block[1] blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y) - phase_start = time.perf_counter() create_csr_kernel( (blocks_per_grid_x, blocks_per_grid_y), threads_per_block, @@ -1154,25 +1121,18 @@ def calc_col_block_gpu( d_indices_matrix.shape[1], ), ) - cp.cuda.Stream.null.synchronize() - _benchmark_add("gpu_create_csr_kernel", phase_start) - phase_start = time.perf_counter() data = d_data.get() indptr = d_indptr.get() indices = d_indices.get() - _benchmark_add("gpu_transfer_to_host", phase_start) res = csr_matrix((data, indices, indptr), shape=(seqs_mat1.shape[0], seqs_mat2.shape[0])) - _benchmark_add("gpu_col_blocks_total", block_start) return res - phase_start = time.perf_counter() seqs_mat1_blocks = np.array_split(seqs_mat1, n_row_blocks) seqs_L1_blocks = np.array_split(seqs_L1, n_row_blocks) seqs_mat2_blocks = np.array_split(seqs_mat2, n_col_blocks) seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) - _benchmark_add("matrix_block_split", phase_start) logging.info( f"\nStart GPU calculations for {n_row_blocks} row blocks x {n_col_blocks} column blocks of max width {self.gpu_block_width}:" @@ -1201,7 +1161,6 @@ def csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elem return data, indices, indptr def csr_union(blocks): - phase_start = time.perf_counter() num_rows = blocks[0].shape[0] num_elements = sum(b.nnz for b in blocks) @@ -1213,11 +1172,9 @@ def csr_union(blocks): shape = blocks[0].shape result = csr_matrix((data, indices, indptr), shape=shape) - _benchmark_add("csr_union", phase_start) return result def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): - row_block_start = time.perf_counter() result_blocks = [None] * n_col_blocks block_offset = start_column @@ -1246,14 +1203,10 @@ def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_blo ) result_sparse = csr_union(result_blocks) - phase_start = time.perf_counter() result_sparse.sort_indices() - _benchmark_add("row_block_sort_indices", phase_start) - _benchmark_add("row_block_total", row_block_start) return result_sparse row_blocks = [None] * n_row_blocks - phase_start = time.perf_counter() with tqdm(total=n_row_blocks * n_col_blocks, desc="Processing", unit="block") as progress_bar: for row_block_idx in range(n_row_blocks): row_blocks[row_block_idx] = calc_row_block_gpu( @@ -1262,43 +1215,11 @@ def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_blo seqs_original_indices_blocks[row_block_idx], ) progress_bar.update(n_col_blocks) - _benchmark_add("all_row_blocks", phase_start) - phase_start = time.perf_counter() result_sparse = scipy.sparse.vstack(row_blocks, format="csr") - _benchmark_add("final_vstack", phase_start) row_element_counts_gpu = np.diff(result_sparse.indptr) - phase_start = time.perf_counter() result_sparse.sort_indices() - _benchmark_add("final_sort_indices", phase_start) - - if benchmark: - total_wall = time.perf_counter() - benchmark_wall_start - print("GPU hamming benchmark") - print(f" total_wall: {total_wall:.3f}s") - for label in ( - "preprocess_and_sort", - "seqs2mat_fast", - "seqs2mat_fallback", - "kernel_setup", - "matrix_block_split", - "gpu_transfer_to_device", - "gpu_transpose_copy", - "gpu_hamming_kernel", - "gpu_row_counts_get", - "gpu_create_csr_kernel", - "gpu_transfer_to_host", - "gpu_col_blocks_total", - "csr_union", - "row_block_sort_indices", - "row_block_total", - "all_row_blocks", - "final_vstack", - "final_sort_indices", - ): - if label in benchmark_timers: - print(f" {label}: {benchmark_timers[label]:.3f}s") # Returns the results in a way that fits the current interface, could be improved later return [result_sparse.data], [result_sparse.indices], row_element_counts_gpu, np.array([None]) From 86391e73d3aca0b617d1a72f647ffa3b8def90f8 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 13:50:14 +0200 Subject: [PATCH 11/54] Skip calculation of blocks under diagonal --- src/scirpy/ir_dist/metrics.py | 35 ++++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 70bd1a9fa..f2256a1a7 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -900,8 +900,6 @@ def _gpu_hamming_mat( seqs2 = np.concatenate(seqs2_sorted_per_block) - is_symmetric = False - max_seq_len = max(len(s) for s in (*seqs, *seqs2)) def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np.ndarray, np.ndarray]: @@ -1069,7 +1067,7 @@ def calc_col_block_gpu( seqs_mat2_cols, d_data_matrix_cols, d_indices_matrix_cols, - is_symmetric, + False, ), ) @@ -1174,22 +1172,28 @@ def csr_union(blocks): result = csr_matrix((data, indices, indptr), shape=shape) return result - def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): - result_blocks = [None] * n_col_blocks + def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): + result_blocks = [] block_offset = start_column - - for i in range(0, n_col_blocks): - result_blocks[i] = calc_col_block_gpu( - seqs_mat1_block, - seqs_mat2_blocks[i], - seqs_L1_block, - seqs_L2_blocks[i], - seqs_original_indices_block, - seqs2_original_indices_blocks[i], - block_offset, + start_col_block = row_block_idx if is_symmetric else 0 + + for i in range(start_col_block, n_col_blocks): + result_blocks.append( + calc_col_block_gpu( + seqs_mat1_block, + seqs_mat2_blocks[i], + seqs_L1_block, + seqs_L2_blocks[i], + seqs_original_indices_block, + seqs2_original_indices_blocks[i], + block_offset, + ) ) block_offset += seqs_mat2_blocks[i].shape[0] + if not result_blocks: + return csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])) + num_elements = 0 for i in range(0, len(result_blocks)): num_elements += result_blocks[i].indptr[-1] @@ -1210,6 +1214,7 @@ def calc_row_block_gpu(seqs_mat1_block, seqs_L1_block, seqs_original_indices_blo with tqdm(total=n_row_blocks * n_col_blocks, desc="Processing", unit="block") as progress_bar: for row_block_idx in range(n_row_blocks): row_blocks[row_block_idx] = calc_row_block_gpu( + row_block_idx, seqs_mat1_blocks[row_block_idx], seqs_L1_blocks[row_block_idx], seqs_original_indices_blocks[row_block_idx], From 1ab2ffb2c79c1e66f0d230e3ed2f0278989c19c1 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 14:00:51 +0200 Subject: [PATCH 12/54] Update block skipping --- src/scirpy/ir_dist/metrics.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index f2256a1a7..7c72c6044 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -873,6 +873,7 @@ def _gpu_hamming_mat( n_row_blocks = self.gpu_row_blocks seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) + seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) seqs_sorted_per_block = [] seqs_original_indices_blocks = [] @@ -885,6 +886,7 @@ def _gpu_hamming_mat( seqs = np.concatenate(seqs_sorted_per_block) seqs2_blocks = np.array_split(np.asarray(seqs2), n_col_blocks) + seqs2_block_starts = np.cumsum([0] + [len(block) for block in seqs2_blocks[:-1]]) seqs2_sorted_per_block = [] seqs2_original_indices_blocks = [] seqs2_block_start = 0 @@ -1175,9 +1177,15 @@ def csr_union(blocks): def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): result_blocks = [] block_offset = start_column - start_col_block = row_block_idx if is_symmetric else 0 + row_start = seqs_block_starts[row_block_idx] + + for i in range(0, n_col_blocks): + col_end = seqs2_block_starts[i] + seqs_mat2_blocks[i].shape[0] + # Skip calculation of blocks below the diagonal if the result matrix is symmetric. + if is_symmetric and col_end <= row_start: + block_offset += seqs_mat2_blocks[i].shape[0] + continue - for i in range(start_col_block, n_col_blocks): result_blocks.append( calc_col_block_gpu( seqs_mat1_block, From 701b8606de51422aee055541a8fb097d9c6f0a31 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 14:17:10 +0200 Subject: [PATCH 13/54] Update block skipping to work with tqdm --- src/scirpy/ir_dist/metrics.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 7c72c6044..c1bdee51f 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1174,15 +1174,27 @@ def csr_union(blocks): result = csr_matrix((data, indices, indptr), shape=shape) return result + def skip_col_block(row_block_idx, col_block_idx): + row_start = seqs_block_starts[row_block_idx] + col_end = seqs2_block_starts[col_block_idx] + seqs_mat2_blocks[col_block_idx].shape[0] + return is_symmetric and col_end <= row_start + + def count_blocks_to_compute(): + n_blocks_to_compute = 0 + for row_block_idx in range(n_row_blocks): + for col_block_idx in range(n_col_blocks): + if not skip_col_block(row_block_idx, col_block_idx): + n_blocks_to_compute += 1 + return n_blocks_to_compute + def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): result_blocks = [] + n_calculated_blocks = 0 block_offset = start_column - row_start = seqs_block_starts[row_block_idx] for i in range(0, n_col_blocks): - col_end = seqs2_block_starts[i] + seqs_mat2_blocks[i].shape[0] # Skip calculation of blocks below the diagonal if the result matrix is symmetric. - if is_symmetric and col_end <= row_start: + if skip_col_block(row_block_idx, i): block_offset += seqs_mat2_blocks[i].shape[0] continue @@ -1197,10 +1209,11 @@ def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_origi block_offset, ) ) + n_calculated_blocks += 1 block_offset += seqs_mat2_blocks[i].shape[0] if not result_blocks: - return csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])) + return csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])), n_calculated_blocks num_elements = 0 for i in range(0, len(result_blocks)): @@ -1216,18 +1229,18 @@ def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_origi result_sparse = csr_union(result_blocks) result_sparse.sort_indices() - return result_sparse + return result_sparse, n_calculated_blocks row_blocks = [None] * n_row_blocks - with tqdm(total=n_row_blocks * n_col_blocks, desc="Processing", unit="block") as progress_bar: + with tqdm(total=count_blocks_to_compute(), desc="Processing", unit="block") as progress_bar: for row_block_idx in range(n_row_blocks): - row_blocks[row_block_idx] = calc_row_block_gpu( + row_blocks[row_block_idx], n_calculated_blocks = calc_row_block_gpu( row_block_idx, seqs_mat1_blocks[row_block_idx], seqs_L1_blocks[row_block_idx], seqs_original_indices_blocks[row_block_idx], ) - progress_bar.update(n_col_blocks) + progress_bar.update(n_calculated_blocks) result_sparse = scipy.sparse.vstack(row_blocks, format="csr") From c82add0f974e970a3690d56af9606155c98c604d Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 14:26:19 +0200 Subject: [PATCH 14/54] Remove logic to only compute upper triangular half for symmetric result blocks within Hamming kernel itself --- src/scirpy/ir_dist/metrics.py | 42 ++++++++++++++++------------------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index c1bdee51f..096724837 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -945,8 +945,7 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. const int seqs_mat1_cols, const int seqs_mat2_cols, const int data_cols, - const int indices_cols, - const bool is_symmetric + const int indices_cols ) { int row = blockDim.x * blockIdx.x + threadIdx.x; if (row < seqs_mat1_rows) { @@ -955,28 +954,26 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. int row_end_index = 0; for (int col = 0; col < seqs_mat2_rows; col++) { - if ((! is_symmetric ) || (col + block_offset) >= row) { - int seq2_len = seqs_L2[col]; - char distance = 1; - - if (seq1_len == seq2_len) { - for (int i = 0; i < seq1_len; i++) { - char val1 = seqs_mat1[i*seqs_mat1_rows+row]; - char val2 = seqs_mat2[i*seqs_mat2_rows+col]; - - if(val1 != val2) { - distance++; - if (distance > cutoff + 1) { - break; - } + int seq2_len = seqs_L2[col]; + char distance = 1; + + if (seq1_len == seq2_len) { + for (int i = 0; i < seq1_len; i++) { + char val1 = seqs_mat1[i*seqs_mat1_rows+row]; + char val2 = seqs_mat2[i*seqs_mat2_rows+col]; + + if(val1 != val2) { + distance++; + if (distance > cutoff + 1) { + break; } } - if (distance <= cutoff + 1) { - int seqs2_original_index = seqs2_original_indices[col]; - data[(long long)seqs_original_index * data_cols + row_end_index] = distance; - indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; - row_end_index++; - } + } + if (distance <= cutoff + 1) { + int seqs2_original_index = seqs2_original_indices[col]; + data[(long long)seqs_original_index * data_cols + row_end_index] = distance; + indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; + row_end_index++; } } } @@ -1069,7 +1066,6 @@ def calc_col_block_gpu( seqs_mat2_cols, d_data_matrix_cols, d_indices_matrix_cols, - False, ), ) From 4884d50271ecf38eb73ba53d2ee3fe990a69b319 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 14:43:01 +0200 Subject: [PATCH 15/54] Extend test cases to cover block skipping --- src/scirpy/tests/test_ir_dist_metrics.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 193fb2b18..7c45bbf1b 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -1009,6 +1009,8 @@ def test_tcrdist_histogram_not_implemented(): {"gpu_col_blocks": 5, "gpu_block_width": 500}, {"gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, {"gpu_col_blocks": 7, "gpu_block_width": 500}, + {"gpu_row_blocks": 5, "gpu_col_blocks": 3, "gpu_block_width": 500}, + {"gpu_row_blocks": 11, "gpu_col_blocks": 17, "gpu_block_width": 503}, ], ) def test_gpu_hamming_reference(kwargs): From efa31eb2d96383dd755f6a71c7a8aff943e54aa7 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 15:52:51 +0200 Subject: [PATCH 16/54] Update changelog --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f3c2abfb..7f7e66fbc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,9 @@ and this project adheres to [Semantic Versioning][]. ### Performance improvements - Speed up identity distance metric computation for comparisons between two different sequence arrays ([#701](https://github.com/scverse/scirpy/pull/701)). + - Improve the GPU implementation of the Hamming distance metric with row and column block splitting, block skipping + for symmetric distance matrices, and faster sparse block assembly. The `gpu_n_blocks` parameter has been replaced + by `gpu_col_blocks`, and `gpu_row_blocks` has been added. ### Chore From a2bba485ad8b0688f6347b480f2aad8cb8f4d014 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Tue, 7 Jul 2026 16:01:30 +0200 Subject: [PATCH 17/54] Add GPU hamming integration test --- src/scirpy/tests/test_ir_dist.py | 35 ++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/src/scirpy/tests/test_ir_dist.py b/src/scirpy/tests/test_ir_dist.py index 3f1b4ba96..339f2ee43 100644 --- a/src/scirpy/tests/test_ir_dist.py +++ b/src/scirpy/tests/test_ir_dist.py @@ -163,6 +163,41 @@ def test_ir_dist( npt.assert_array_equal(res["VDJ"]["distances"].toarray(), expected_dist_vdj) +@pytest.mark.gpu +def test_ir_dist_gpu_hamming(adata_cdr3): + expected_seq_vj = np.array(["AAA", "AHA"]) + expected_seq_vdj = np.array(["AAA", "KK", "KKK", "KKY", "LLL"]) + ir.pp.ir_dist( + adata_cdr3, + metric="gpu_hamming", + sequence="aa", + cutoff=2, + gpu_col_blocks=2, + gpu_row_blocks=2, + gpu_block_width=3, + ) + res = ( + adata_cdr3.mod["airr"].uns["ir_dist_aa_gpu_hamming"] + if isinstance(adata_cdr3, MuData) + else adata_cdr3.uns["ir_dist_aa_gpu_hamming"] + ) + npt.assert_array_equal(res["VJ"]["seqs"], expected_seq_vj) + npt.assert_array_equal(res["VDJ"]["seqs"], expected_seq_vdj) + npt.assert_array_equal(res["VJ"]["distances"].toarray(), np.array([[1, 2], [2, 1]])) + npt.assert_array_equal( + res["VDJ"]["distances"].toarray(), + np.array( + [ + [1, 0, 0, 0, 0], + [0, 1, 0, 0, 0], + [0, 0, 1, 2, 0], + [0, 0, 2, 1, 0], + [0, 0, 0, 0, 1], + ] + ), + ) + + @pytest.mark.parametrize("mudata", [False, True], ids=["AnnData", "MuData"]) def test_ir_dist_tcrdist_tcrblosum_chain_routing(mudata): # `ir_dist` should automatically route VJ to TCRBLOSUM alpha and VDJ to beta. From 940d3a01a7f1c6cfae30010c3036c20585739273 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 22 Jul 2026 12:05:27 +0200 Subject: [PATCH 18/54] Prevent out-of-bounds writes in GPU Hamming kernel --- src/scirpy/ir_dist/metrics.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 096724837..cbd0a82ed 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -970,9 +970,11 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. } } if (distance <= cutoff + 1) { - int seqs2_original_index = seqs2_original_indices[col]; - data[(long long)seqs_original_index * data_cols + row_end_index] = distance; - indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; + if (row_end_index < data_cols) { + int seqs2_original_index = seqs2_original_indices[col]; + data[(long long)seqs_original_index * data_cols + row_end_index] = distance; + indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; + } row_end_index++; } } From d5eef6e0a10139a45db813f79b3bffbce30dfd1d Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 22 Jul 2026 12:08:57 +0200 Subject: [PATCH 19/54] Handle excess blocks in GPU Hamming calculator --- src/scirpy/ir_dist/metrics.py | 4 ++-- src/scirpy/tests/test_ir_dist.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index cbd0a82ed..e29cb5334 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -869,8 +869,8 @@ def _gpu_hamming_mat( import cupy as cp from tqdm import tqdm - n_col_blocks = self.gpu_col_blocks - n_row_blocks = self.gpu_row_blocks + n_col_blocks = min(self.gpu_col_blocks, len(seqs2)) + n_row_blocks = min(self.gpu_row_blocks, len(seqs)) seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) diff --git a/src/scirpy/tests/test_ir_dist.py b/src/scirpy/tests/test_ir_dist.py index 339f2ee43..1e515f839 100644 --- a/src/scirpy/tests/test_ir_dist.py +++ b/src/scirpy/tests/test_ir_dist.py @@ -172,8 +172,8 @@ def test_ir_dist_gpu_hamming(adata_cdr3): metric="gpu_hamming", sequence="aa", cutoff=2, - gpu_col_blocks=2, - gpu_row_blocks=2, + gpu_col_blocks=10, + gpu_row_blocks=10, gpu_block_width=3, ) res = ( From 0530c30b8890672810246b63364908d83cfa1c1a Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 22 Jul 2026 15:02:24 +0200 Subject: [PATCH 20/54] Support joblib blocks in GPU Hamming calculator --- src/scirpy/ir_dist/metrics.py | 29 +++++------ src/scirpy/tests/test_ir_dist_metrics.py | 61 +++++++++++++++++++++--- 2 files changed, 66 insertions(+), 24 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index e29cb5334..f46019720 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -794,6 +794,9 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): cutoff: Will eleminate distances > cutoff to make efficient use of sparse matrices. + n_blocks: + Number of outer row blocks submitted through joblib. This can be used with a distributed joblib backend to + distribute the calculation across multiple GPU workers. gpu_col_blocks: Number of column blocks used to split the final result matrix. Lower values reduce block-management overhead but increase the size of each GPU block and therefore memory pressure. @@ -809,11 +812,12 @@ def __init__( self, *, cutoff: int = 2, + n_blocks: int = 1, gpu_col_blocks: int = 10, gpu_row_blocks: int = 1, gpu_block_width: int = 1000, ): - super().__init__(n_jobs=1, n_blocks=1) + super().__init__(n_jobs=1, n_blocks=n_blocks) if cutoff > 125: raise ValueError( "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " @@ -837,7 +841,7 @@ def _gpu_hamming_mat( seqs2: Sequence[str], is_symmetric: bool = False, start_column: int = 0, - ) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: + ) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray]: """Computes the pairwise hamming distances for sequences in seqs and seqs2 with GPU support. Parameters @@ -848,8 +852,8 @@ def _gpu_hamming_mat( Determines whether the final result matrix is symmetric, assuming that this function is only used to compute a block of a bigger result matrix start_column: - Determines at which column the calculation should be started. This is only used if this function is - used to compute a block of a bigger result matrix that is symmetric + Global row offset of an outer row block scheduled by joblib. Used to skip column blocks below the diagonal + when computing a symmetric result matrix. Returns ------- @@ -939,11 +943,8 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. char* __restrict__ data, int* __restrict__ indices, int* __restrict__ row_element_counts, - const int block_offset, const int seqs_mat1_rows, const int seqs_mat2_rows, - const int seqs_mat1_cols, - const int seqs_mat2_cols, const int data_cols, const int indices_cols ) { @@ -1021,7 +1022,6 @@ def calc_col_block_gpu( seqs_L2, seqs_original_indices_block, seqs2_original_indices_block, - block_offset, ): d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8, copy=False)) d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8, copy=False)) @@ -1039,8 +1039,8 @@ def calc_col_block_gpu( threads_per_block = 256 blocks_per_grid = (seqs_mat1.shape[0] + (threads_per_block - 1)) // threads_per_block - seqs_mat1_rows, seqs_mat1_cols = seqs_mat1.shape - seqs_mat2_rows, seqs_mat2_cols = seqs_mat2_block.shape + seqs_mat1_rows = seqs_mat1.shape[0] + seqs_mat2_rows = seqs_mat2_block.shape[0] d_data_matrix_cols = max_block_width d_indices_matrix_cols = max_block_width @@ -1061,11 +1061,8 @@ def calc_col_block_gpu( d_data_matrix, d_indices_matrix, d_row_element_counts, - block_offset, seqs_mat1_rows, seqs_mat2_rows, - seqs_mat1_cols, - seqs_mat2_cols, d_data_matrix_cols, d_indices_matrix_cols, ), @@ -1173,7 +1170,7 @@ def csr_union(blocks): return result def skip_col_block(row_block_idx, col_block_idx): - row_start = seqs_block_starts[row_block_idx] + row_start = start_column + seqs_block_starts[row_block_idx] col_end = seqs2_block_starts[col_block_idx] + seqs_mat2_blocks[col_block_idx].shape[0] return is_symmetric and col_end <= row_start @@ -1188,12 +1185,10 @@ def count_blocks_to_compute(): def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): result_blocks = [] n_calculated_blocks = 0 - block_offset = start_column for i in range(0, n_col_blocks): # Skip calculation of blocks below the diagonal if the result matrix is symmetric. if skip_col_block(row_block_idx, i): - block_offset += seqs_mat2_blocks[i].shape[0] continue result_blocks.append( @@ -1204,11 +1199,9 @@ def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_origi seqs_L2_blocks[i], seqs_original_indices_block, seqs2_original_indices_blocks[i], - block_offset, ) ) n_calculated_blocks += 1 - block_offset += seqs_mat2_blocks[i].shape[0] if not result_blocks: return csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])), n_calculated_blocks diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 7c45bbf1b..7472f58d0 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -965,13 +965,61 @@ def test_hamming_histogram_reference(): @pytest.mark.gpu -def test_gpu_hamming_long_sequence(): - """Regression test for #626 and #682""" - hamming_calculator = GPUHammingDistanceCalculator(cutoff=50, gpu_col_blocks=1, gpu_block_width=50) - seq1 = np.array(["A" * 128, "AAB", "AABB", "ABA"]) - seq2 = np.array(["A" * 128, "ABBB", "ABBB"]) - res = hamming_calculator.calc_dist_mat(seq1, seq2) +@pytest.mark.parametrize( + "test_parameters,test_input,expected_result", + [ + # Regression test for #626 and #682: sequence lengths exceeding the int8 range. + ( + {"cutoff": 50, "gpu_col_blocks": 1, "gpu_block_width": 50}, + (np.array(["A" * 128, "AAB", "AABB", "ABA"]), np.array(["A" * 128, "ABBB", "ABBB"])), + np.array([[1, 0, 0], [0, 0, 0], [0, 2, 2], [0, 0, 0]]), + ), + # Symmetric calculation split into outer joblib blocks and internal GPU blocks. + ( + {"cutoff": 2, "n_blocks": 2, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 3}, + (np.array(["AAAA", "AAAT", "AATT", "TTTT"]), None), + np.array([[1, 2, 3, 0], [2, 1, 2, 0], [3, 2, 1, 3], [0, 0, 3, 1]]), + ), + # Asymmetric calculation split into outer joblib blocks and internal GPU blocks. + ( + {"cutoff": 2, "n_blocks": 2, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 3}, + (np.array(["AAAA", "AATA", "HHHH", "WWWW"]), np.array(["WWWW", "AAAA", "ATAA"])), + np.array([[0, 1, 2], [0, 2, 3], [0, 0, 0], [1, 0, 0]]), + ), + # Distances above the cutoff and comparisons between sequences of unequal length are omitted. + ( + {"cutoff": 1, "gpu_row_blocks": 3, "gpu_col_blocks": 3, "gpu_block_width": 2}, + (np.array(["AAA", "AAT", "AAAA", "TTT"]), None), + np.array([[1, 2, 0, 0], [2, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]), + ), + # Duplicate sequences are retained as separate zero-distance entries with distance + 1 encoding. + ( + {"cutoff": 0, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 2}, + (np.array(["AAA", "AAA", "AAT"]), None), + np.array([[1, 1, 0], [1, 1, 0], [0, 0, 1]]), + ), + # Block counts exceeding the number of sequences are reduced to avoid empty internal GPU blocks. + ( + {"cutoff": 1, "n_blocks": 5, "gpu_row_blocks": 5, "gpu_col_blocks": 5, "gpu_block_width": 2}, + (np.array(["AAA", "AAT"]), None), + np.array([[1, 2], [2, 1]]), + ), + ], +) +def test_gpu_hamming(test_parameters, test_input, expected_result): + hamming_calculator = GPUHammingDistanceCalculator(**test_parameters) + res = hamming_calculator.calc_dist_mat(*test_input) + assert isinstance(res, scipy.sparse.csr_matrix) + npt.assert_array_equal(res.toarray(), expected_result) + + +@pytest.mark.gpu +def test_gpu_hamming_block_width_guard(): + hamming_calculator = GPUHammingDistanceCalculator(cutoff=0, gpu_col_blocks=1, gpu_block_width=1) + + with pytest.raises(ValueError, match="result block width is too small"): + hamming_calculator.calc_dist_mat(np.array(["AAA", "AAA"])) @pytest.mark.gpu @@ -1008,6 +1056,7 @@ def test_tcrdist_histogram_not_implemented(): [ {"gpu_col_blocks": 5, "gpu_block_width": 500}, {"gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, + {"n_blocks": 3, "gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, {"gpu_col_blocks": 7, "gpu_block_width": 500}, {"gpu_row_blocks": 5, "gpu_col_blocks": 3, "gpu_block_width": 500}, {"gpu_row_blocks": 11, "gpu_col_blocks": 17, "gpu_block_width": 503}, From 4c9f30e1d9b806f1d3c5d693bf41391f84b19493 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 11:27:48 +0200 Subject: [PATCH 21/54] Balance symmetric distance blocks and use fixed GPU tile sizes --- src/scirpy/ir_dist/metrics.py | 51 ++++++++++++++++++++++++----------- 1 file changed, 35 insertions(+), 16 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index f46019720..09f555a84 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -545,10 +545,31 @@ def calc_dist_mat(self, seqs: Sequence[str], seqs2: Sequence[str] | None = None) seqs2 = np.array(seqs2) is_symmetric = np.array_equal(seqs, seqs2) - if self.n_blocks > 1: - split_seqs = np.array_split(seqs, self.n_blocks) - start_columns = np.cumsum([0] + [len(seq) for seq in split_seqs[:-1]]) - arguments = [(split_seqs[x], seqs2, is_symmetric, start_columns[x]) for x in range(self.n_blocks)] + if self.n_blocks < 2: + distance_matrix_csr, row_mins = self._calc_dist_mat_block(seqs, seqs2, is_symmetric) + else: + if is_symmetric: + # Computing only the upper triangle of a symmetric result matrix gives earlier row partitions more + # work, so use shorter partitions at the beginning to balance the number of comparisons across + # parallel jobs. Increasing n_blocks for better load balancing instead would add block-processing + # overhead. + partition_fractions = np.arange(self.n_blocks + 1) / self.n_blocks + partition_boundaries = np.rint( + len(seqs) * (1 - np.sqrt(1 - partition_fractions)) + ).astype(int) + else: + partition_boundaries = np.rint(np.linspace(0, len(seqs), self.n_blocks + 1)).astype(int) + + split_seqs = np.split(seqs, partition_boundaries[1:-1]) + arguments = [ + ( + split_seqs[x], + seqs2, + is_symmetric, + partition_boundaries[x], + ) + for x in range(self.n_blocks) + ] delayed_jobs = [joblib.delayed(self._calc_dist_mat_block)(*args) for args in arguments] results = joblib.Parallel(return_as="list")(delayed_jobs) @@ -556,8 +577,6 @@ def calc_dist_mat(self, seqs: Sequence[str], seqs2: Sequence[str] | None = None) block_matrices_csr, block_row_mins = zip(*results, strict=False) distance_matrix_csr = scipy.sparse.vstack(block_matrices_csr) row_mins = np.concatenate(block_row_mins) - else: - distance_matrix_csr, row_mins = self._calc_dist_mat_block(seqs, seqs2, is_symmetric) if is_symmetric: upper_triangular_distance_matrix = distance_matrix_csr @@ -813,8 +832,8 @@ def __init__( *, cutoff: int = 2, n_blocks: int = 1, - gpu_col_blocks: int = 10, - gpu_row_blocks: int = 1, + gpu_col_block_size: int = 100_000, + gpu_row_block_size: int = 100_000, gpu_block_width: int = 1000, ): super().__init__(n_jobs=1, n_blocks=n_blocks) @@ -823,15 +842,15 @@ def __init__( "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." ) - if gpu_col_blocks < 1: - raise ValueError("`gpu_col_blocks` must be >= 1.") - if gpu_row_blocks < 1: - raise ValueError("`gpu_row_blocks` must be >= 1.") + if gpu_col_block_size < 1: + raise ValueError("`gpu_col_block_size` must be >= 1.") + if gpu_row_block_size < 1: + raise ValueError("`gpu_row_block_size` must be >= 1.") if gpu_block_width < 1: raise ValueError("`gpu_block_width` must be >= 1.") self.cutoff = cutoff - self.gpu_col_blocks = gpu_col_blocks - self.gpu_row_blocks = gpu_row_blocks + self.gpu_col_block_size = gpu_col_block_size + self.gpu_row_block_size = gpu_row_block_size self.gpu_block_width = gpu_block_width def _gpu_hamming_mat( @@ -873,8 +892,8 @@ def _gpu_hamming_mat( import cupy as cp from tqdm import tqdm - n_col_blocks = min(self.gpu_col_blocks, len(seqs2)) - n_row_blocks = min(self.gpu_row_blocks, len(seqs)) + n_col_blocks = (len(seqs2) + self.gpu_col_block_size - 1) // self.gpu_col_block_size + n_row_blocks = (len(seqs) + self.gpu_row_block_size - 1) // self.gpu_row_block_size seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) From 6d22b8bb21b9811a19fce59365b9927460967177 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 15:03:08 +0200 Subject: [PATCH 22/54] Improve GPU Hamming block sizing and buffer handling --- src/scirpy/ir_dist/metrics.py | 189 ++++++++++++++++++---------------- 1 file changed, 102 insertions(+), 87 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 09f555a84..74d56f796 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -794,19 +794,17 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): The code of this class is based on `pwseqdist `_. Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. - For performance reasons, the computation of the final result matrix is split into a grid of row and column blocks. - `gpu_row_blocks` splits the query sequences into horizontal chunks, while `gpu_col_blocks` splits the reference - sequences into vertical chunks. Each row/column block pair is computed on the GPU and converted to a sparse CSR - block before the blocks are combined again. + For performance reasons, the rows and columns of the final result matrix are grouped into blocks for GPU + computation. `gpu_block_rows` and `gpu_block_cols` control how many matrix rows and columns are grouped into each + block. Each block is computed on the GPU and converted to a sparse CSR block before the blocks are combined again. - `gpu_block_width` controls how many sparse result entries are reserved per row for each row/column block pair. - For example, with a 1000 x 1000 dense result matrix, `gpu_row_blocks=2`, `gpu_col_blocks=10`, and - `gpu_block_width=20`, each GPU block covers roughly 500 x 100 dense comparisons and reserves enough GPU memory - for at most 20 retained distances below the cutoff per row in that block. Therefore, no row within a single row/column block pair - may contain more than `gpu_block_width` distances `<= cutoff`. + `gpu_buffer_cols` controls how many buffer columns are initially reserved for sparse result entries in each row of + a block. Because only distances at or below the cutoff are retained, the number of entries that need to be stored + is usually considerably smaller than the number of columns in the block. If necessary, the buffer is enlarged and + the calculation is retried. - Larger numbers of row or column blocks reduce per-block memory pressure but add block-management overhead. Larger - values for `gpu_block_width` can avoid sparse buffer overflows but require more GPU memory. + Smaller blocks reduce per-block memory pressure but add block-management overhead. Larger values for + `gpu_buffer_cols` can avoid retries but require more GPU memory. Parameters ---------- @@ -816,15 +814,13 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): n_blocks: Number of outer row blocks submitted through joblib. This can be used with a distributed joblib backend to distribute the calculation across multiple GPU workers. - gpu_col_blocks: - Number of column blocks used to split the final result matrix. Lower values reduce block-management overhead - but increase the size of each GPU block and therefore memory pressure. - gpu_row_blocks: - Number of row blocks used to split the final result matrix. Higher values reduce per-block memory pressure - and can improve sequence length homogeneity within a block, but add block-management overhead. - gpu_block_width: - Maximum number of retained sparse entries per row and row/column block pair. Higher values tolerate denser - results within a block but require more GPU memory. + gpu_block_rows: + Number of result matrix rows per GPU block. + gpu_block_cols: + Number of result matrix columns per GPU block. + gpu_buffer_cols: + Initial number of retained sparse entries reserved per row of each block. Higher values can avoid retries for + denser results but require more GPU memory. """ def __init__( @@ -832,9 +828,9 @@ def __init__( *, cutoff: int = 2, n_blocks: int = 1, - gpu_col_block_size: int = 100_000, - gpu_row_block_size: int = 100_000, - gpu_block_width: int = 1000, + gpu_block_rows: int = 100_000, + gpu_block_cols: int = 100_000, + gpu_buffer_cols: int = 1000, ): super().__init__(n_jobs=1, n_blocks=n_blocks) if cutoff > 125: @@ -842,16 +838,16 @@ def __init__( "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." ) - if gpu_col_block_size < 1: - raise ValueError("`gpu_col_block_size` must be >= 1.") - if gpu_row_block_size < 1: - raise ValueError("`gpu_row_block_size` must be >= 1.") - if gpu_block_width < 1: - raise ValueError("`gpu_block_width` must be >= 1.") + if gpu_block_rows < 1: + raise ValueError("`gpu_block_rows` must be >= 1.") + if gpu_block_cols < 1: + raise ValueError("`gpu_block_cols` must be >= 1.") + if gpu_buffer_cols < 1: + raise ValueError("`gpu_buffer_cols` must be >= 1.") self.cutoff = cutoff - self.gpu_col_block_size = gpu_col_block_size - self.gpu_row_block_size = gpu_row_block_size - self.gpu_block_width = gpu_block_width + self.gpu_block_rows = gpu_block_rows + self.gpu_block_cols = gpu_block_cols + self.gpu_buffer_cols = gpu_buffer_cols def _gpu_hamming_mat( self, @@ -892,8 +888,8 @@ def _gpu_hamming_mat( import cupy as cp from tqdm import tqdm - n_col_blocks = (len(seqs2) + self.gpu_col_block_size - 1) // self.gpu_col_block_size - n_row_blocks = (len(seqs) + self.gpu_row_block_size - 1) // self.gpu_row_block_size + n_col_blocks = (len(seqs2) + self.gpu_block_cols - 1) // self.gpu_block_cols + n_row_blocks = (len(seqs) + self.gpu_block_rows - 1) // self.gpu_block_rows seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) @@ -1041,62 +1037,69 @@ def calc_col_block_gpu( seqs_L2, seqs_original_indices_block, seqs2_original_indices_block, + buffer_width, ): d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8, copy=False)) d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8, copy=False)) d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32, copy=False)) - # Due to performance reasons and since we expect the result matrix to be very sparse, we - # set a maximum result width for the current block - max_block_width = self.gpu_block_width - - d_data_matrix = cp.empty((seqs_mat1.shape[0], max_block_width), dtype=cp.int8) - d_indices_matrix = cp.empty((seqs_mat1.shape[0], max_block_width), dtype=np.int32) - d_row_element_counts = cp.zeros(seqs_mat1.shape[0], dtype=np.int32) - threads_per_block = 256 blocks_per_grid = (seqs_mat1.shape[0] + (threads_per_block - 1)) // threads_per_block seqs_mat1_rows = seqs_mat1.shape[0] seqs_mat2_rows = seqs_mat2_block.shape[0] - d_data_matrix_cols = max_block_width - d_indices_matrix_cols = max_block_width d_seqs_mat1_transposed = cp.transpose(d_seqs_mat1).copy() d_seqs_mat2_transposed = cp.transpose(d_seqs_mat2).copy() - hamming_kernel( - (blocks_per_grid,), - (threads_per_block,), - ( - d_seqs_mat1_transposed, - d_seqs_mat2_transposed, - d_seqs_L1, - d_seqs_L2, - seqs_original_indices_block, - seqs2_original_indices_block, - self.cutoff, - d_data_matrix, - d_indices_matrix, - d_row_element_counts, - seqs_mat1_rows, - seqs_mat2_rows, - d_data_matrix_cols, - d_indices_matrix_cols, - ), - ) + def run_hamming_kernel(buffer_width): + d_data_matrix = cp.empty((seqs_mat1_rows, buffer_width), dtype=cp.int8) + d_indices_matrix = cp.empty((seqs_mat1_rows, buffer_width), dtype=np.int32) + d_row_element_counts = cp.zeros(seqs_mat1_rows, dtype=np.int32) + + hamming_kernel( + (blocks_per_grid,), + (threads_per_block,), + ( + d_seqs_mat1_transposed, + d_seqs_mat2_transposed, + d_seqs_L1, + d_seqs_L2, + seqs_original_indices_block, + seqs2_original_indices_block, + self.cutoff, + d_data_matrix, + d_indices_matrix, + d_row_element_counts, + seqs_mat1_rows, + seqs_mat2_rows, + buffer_width, + buffer_width, + ), + ) + row_element_counts = d_row_element_counts.get() + required_buffer_width = int(np.max(row_element_counts)) + if required_buffer_width > buffer_width: + # Release undersized buffers before allocating larger ones for the retry. + d_data_matrix = None + d_indices_matrix = None + return d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width + + d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width = run_hamming_kernel(buffer_width) + + if required_buffer_width > buffer_width: + # The buffer was too small, so retry with the required buffer size. + print( + f"GPU Hamming buffer retry for a {seqs_mat1_rows} x {seqs_mat2_rows} block: " + f"{buffer_width} -> {required_buffer_width}", + flush=True, + ) + buffer_width = required_buffer_width + d_data_matrix, d_indices_matrix, row_element_counts, _ = run_hamming_kernel(buffer_width) - row_element_counts = d_row_element_counts.get() - row_max_len = np.max(row_element_counts) row_element_sum = np.sum(row_element_counts, dtype=np.int64) - if row_max_len > max_block_width: - raise ValueError( - "The chosen result block width is too small to hold all result values of the current block. " - f"Chosen width: {max_block_width}, necessary width: {row_max_len}." - ) - if row_element_sum > np.iinfo(np.int32).max: raise ValueError( "There are too many result values to be held by the resulting CSR matrix of the current block. " @@ -1141,7 +1144,7 @@ def calc_col_block_gpu( indices = d_indices.get() res = csr_matrix((data, indices, indptr), shape=(seqs_mat1.shape[0], seqs_mat2.shape[0])) - return res + return res, buffer_width seqs_mat1_blocks = np.array_split(seqs_mat1, n_row_blocks) seqs_L1_blocks = np.array_split(seqs_L1, n_row_blocks) @@ -1149,7 +1152,7 @@ def calc_col_block_gpu( seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) logging.info( - f"\nStart GPU calculations for {n_row_blocks} row blocks x {n_col_blocks} column blocks of max width {self.gpu_block_width}:" + f"\nStart GPU calculations for {n_row_blocks} row blocks x {n_col_blocks} column blocks:" ) @nb.njit @@ -1201,7 +1204,13 @@ def count_blocks_to_compute(): n_blocks_to_compute += 1 return n_blocks_to_compute - def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_original_indices_block): + def calc_row_block_gpu( + row_block_idx, + seqs_mat1_block, + seqs_L1_block, + seqs_original_indices_block, + buffer_width, + ): result_blocks = [] n_calculated_blocks = 0 @@ -1210,20 +1219,24 @@ def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_origi if skip_col_block(row_block_idx, i): continue - result_blocks.append( - calc_col_block_gpu( - seqs_mat1_block, - seqs_mat2_blocks[i], - seqs_L1_block, - seqs_L2_blocks[i], - seqs_original_indices_block, - seqs2_original_indices_blocks[i], - ) + result_block, buffer_width = calc_col_block_gpu( + seqs_mat1_block, + seqs_mat2_blocks[i], + seqs_L1_block, + seqs_L2_blocks[i], + seqs_original_indices_block, + seqs2_original_indices_blocks[i], + buffer_width, ) + result_blocks.append(result_block) n_calculated_blocks += 1 if not result_blocks: - return csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])), n_calculated_blocks + return ( + csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])), + n_calculated_blocks, + buffer_width, + ) num_elements = 0 for i in range(0, len(result_blocks)): @@ -1239,16 +1252,18 @@ def calc_row_block_gpu(row_block_idx, seqs_mat1_block, seqs_L1_block, seqs_origi result_sparse = csr_union(result_blocks) result_sparse.sort_indices() - return result_sparse, n_calculated_blocks + return result_sparse, n_calculated_blocks, buffer_width row_blocks = [None] * n_row_blocks + buffer_width = self.gpu_buffer_cols with tqdm(total=count_blocks_to_compute(), desc="Processing", unit="block") as progress_bar: for row_block_idx in range(n_row_blocks): - row_blocks[row_block_idx], n_calculated_blocks = calc_row_block_gpu( + row_blocks[row_block_idx], n_calculated_blocks, buffer_width = calc_row_block_gpu( row_block_idx, seqs_mat1_blocks[row_block_idx], seqs_L1_blocks[row_block_idx], seqs_original_indices_blocks[row_block_idx], + buffer_width, ) progress_bar.update(n_calculated_blocks) From c5938abbd34939a1ff3e8ef7b474bfdb41ef19b3 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 15:14:57 +0200 Subject: [PATCH 23/54] Change parameter names --- CHANGELOG.md | 7 +-- src/scirpy/ir_dist/metrics.py | 66 ++++++++++++------------ src/scirpy/tests/test_ir_dist.py | 6 +-- src/scirpy/tests/test_ir_dist_metrics.py | 46 ++++++++--------- 4 files changed, 63 insertions(+), 62 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f7e66fbc..994020f79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,9 +18,10 @@ and this project adheres to [Semantic Versioning][]. ### Performance improvements - Speed up identity distance metric computation for comparisons between two different sequence arrays ([#701](https://github.com/scverse/scirpy/pull/701)). - - Improve the GPU implementation of the Hamming distance metric with row and column block splitting, block skipping - for symmetric distance matrices, and faster sparse block assembly. The `gpu_n_blocks` parameter has been replaced - by `gpu_col_blocks`, and `gpu_row_blocks` has been added. + - Improve the GPU implementation of the Hamming distance metric with row and column tiling, tile skipping for + symmetric distance matrices, adaptive result buffers, and faster sparse tile assembly. GPU tiles can be configured + with `gpu_tile_rows` and `gpu_tile_cols`; `gpu_n_blocks` has been replaced by `gpu_tile_cols`, and `gpu_block_width` + has been replaced by `gpu_tile_buffer_cols`. ### Chore diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 74d56f796..793bb76d7 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -794,17 +794,17 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): The code of this class is based on `pwseqdist `_. Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. - For performance reasons, the rows and columns of the final result matrix are grouped into blocks for GPU - computation. `gpu_block_rows` and `gpu_block_cols` control how many matrix rows and columns are grouped into each - block. Each block is computed on the GPU and converted to a sparse CSR block before the blocks are combined again. + For performance reasons, the rows and columns of the final result matrix are grouped into tiles for GPU + computation. `gpu_tile_rows` and `gpu_tile_cols` control how many matrix rows and columns are grouped into each + tile. Each tile is computed on the GPU and converted to a sparse CSR matrix before the tiles are combined again. - `gpu_buffer_cols` controls how many buffer columns are initially reserved for sparse result entries in each row of - a block. Because only distances at or below the cutoff are retained, the number of entries that need to be stored - is usually considerably smaller than the number of columns in the block. If necessary, the buffer is enlarged and - the calculation is retried. + `gpu_tile_buffer_cols` controls how many buffer columns are initially reserved for sparse result entries in each + row of a tile. Because only distances at or below the cutoff are retained, the number of entries that need to be + stored is usually considerably smaller than the number of columns in the tile. If necessary, the buffer is + enlarged and the calculation is retried. - Smaller blocks reduce per-block memory pressure but add block-management overhead. Larger values for - `gpu_buffer_cols` can avoid retries but require more GPU memory. + Smaller tiles reduce per-tile memory pressure but add tile-management overhead. Larger values for + `gpu_tile_buffer_cols` can avoid retries but require more GPU memory. Parameters ---------- @@ -812,14 +812,14 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): Will eleminate distances > cutoff to make efficient use of sparse matrices. n_blocks: - Number of outer row blocks submitted through joblib. This can be used with a distributed joblib backend to + Number of outer row partitions submitted through joblib. This can be used with a distributed joblib backend to distribute the calculation across multiple GPU workers. - gpu_block_rows: - Number of result matrix rows per GPU block. - gpu_block_cols: - Number of result matrix columns per GPU block. - gpu_buffer_cols: - Initial number of retained sparse entries reserved per row of each block. Higher values can avoid retries for + gpu_tile_rows: + Number of result matrix rows per GPU tile. + gpu_tile_cols: + Number of result matrix columns per GPU tile. + gpu_tile_buffer_cols: + Initial number of retained sparse entries reserved per row of each tile. Higher values can avoid retries for denser results but require more GPU memory. """ @@ -828,9 +828,9 @@ def __init__( *, cutoff: int = 2, n_blocks: int = 1, - gpu_block_rows: int = 100_000, - gpu_block_cols: int = 100_000, - gpu_buffer_cols: int = 1000, + gpu_tile_rows: int = 100_000, + gpu_tile_cols: int = 100_000, + gpu_tile_buffer_cols: int = 1000, ): super().__init__(n_jobs=1, n_blocks=n_blocks) if cutoff > 125: @@ -838,16 +838,16 @@ def __init__( "GPUHammingDistanceCalculator only supports cutoff <= 125 because the intermediate " "GPU buffer stores distances as signed int8 values and uses distance + 1 encoding." ) - if gpu_block_rows < 1: - raise ValueError("`gpu_block_rows` must be >= 1.") - if gpu_block_cols < 1: - raise ValueError("`gpu_block_cols` must be >= 1.") - if gpu_buffer_cols < 1: - raise ValueError("`gpu_buffer_cols` must be >= 1.") + if gpu_tile_rows < 1: + raise ValueError("`gpu_tile_rows` must be >= 1.") + if gpu_tile_cols < 1: + raise ValueError("`gpu_tile_cols` must be >= 1.") + if gpu_tile_buffer_cols < 1: + raise ValueError("`gpu_tile_buffer_cols` must be >= 1.") self.cutoff = cutoff - self.gpu_block_rows = gpu_block_rows - self.gpu_block_cols = gpu_block_cols - self.gpu_buffer_cols = gpu_buffer_cols + self.gpu_tile_rows = gpu_tile_rows + self.gpu_tile_cols = gpu_tile_cols + self.gpu_tile_buffer_cols = gpu_tile_buffer_cols def _gpu_hamming_mat( self, @@ -888,8 +888,8 @@ def _gpu_hamming_mat( import cupy as cp from tqdm import tqdm - n_col_blocks = (len(seqs2) + self.gpu_block_cols - 1) // self.gpu_block_cols - n_row_blocks = (len(seqs) + self.gpu_block_rows - 1) // self.gpu_block_rows + n_col_blocks = (len(seqs2) + self.gpu_tile_cols - 1) // self.gpu_tile_cols + n_row_blocks = (len(seqs) + self.gpu_tile_rows - 1) // self.gpu_tile_rows seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) @@ -1091,7 +1091,7 @@ def run_hamming_kernel(buffer_width): if required_buffer_width > buffer_width: # The buffer was too small, so retry with the required buffer size. print( - f"GPU Hamming buffer retry for a {seqs_mat1_rows} x {seqs_mat2_rows} block: " + f"GPU Hamming buffer retry for a {seqs_mat1_rows} x {seqs_mat2_rows} tile: " f"{buffer_width} -> {required_buffer_width}", flush=True, ) @@ -1152,7 +1152,7 @@ def run_hamming_kernel(buffer_width): seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) logging.info( - f"\nStart GPU calculations for {n_row_blocks} row blocks x {n_col_blocks} column blocks:" + f"\nStart GPU calculations for {n_row_blocks} row tiles x {n_col_blocks} column tiles:" ) @nb.njit @@ -1255,7 +1255,7 @@ def calc_row_block_gpu( return result_sparse, n_calculated_blocks, buffer_width row_blocks = [None] * n_row_blocks - buffer_width = self.gpu_buffer_cols + buffer_width = self.gpu_tile_buffer_cols with tqdm(total=count_blocks_to_compute(), desc="Processing", unit="block") as progress_bar: for row_block_idx in range(n_row_blocks): row_blocks[row_block_idx], n_calculated_blocks, buffer_width = calc_row_block_gpu( diff --git a/src/scirpy/tests/test_ir_dist.py b/src/scirpy/tests/test_ir_dist.py index 1e515f839..8fdd24580 100644 --- a/src/scirpy/tests/test_ir_dist.py +++ b/src/scirpy/tests/test_ir_dist.py @@ -172,9 +172,9 @@ def test_ir_dist_gpu_hamming(adata_cdr3): metric="gpu_hamming", sequence="aa", cutoff=2, - gpu_col_blocks=10, - gpu_row_blocks=10, - gpu_block_width=3, + gpu_tile_rows=1, + gpu_tile_cols=1, + gpu_tile_buffer_cols=3, ) res = ( adata_cdr3.mod["airr"].uns["ir_dist_aa_gpu_hamming"] diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 7472f58d0..f01b98a95 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -970,37 +970,37 @@ def test_hamming_histogram_reference(): [ # Regression test for #626 and #682: sequence lengths exceeding the int8 range. ( - {"cutoff": 50, "gpu_col_blocks": 1, "gpu_block_width": 50}, + {"cutoff": 50, "gpu_tile_cols": 3, "gpu_tile_buffer_cols": 50}, (np.array(["A" * 128, "AAB", "AABB", "ABA"]), np.array(["A" * 128, "ABBB", "ABBB"])), np.array([[1, 0, 0], [0, 0, 0], [0, 2, 2], [0, 0, 0]]), ), - # Symmetric calculation split into outer joblib blocks and internal GPU blocks. + # Symmetric calculation split into outer joblib blocks and internal GPU tiles. ( - {"cutoff": 2, "n_blocks": 2, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 3}, + {"cutoff": 2, "n_blocks": 2, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 3}, (np.array(["AAAA", "AAAT", "AATT", "TTTT"]), None), np.array([[1, 2, 3, 0], [2, 1, 2, 0], [3, 2, 1, 3], [0, 0, 3, 1]]), ), - # Asymmetric calculation split into outer joblib blocks and internal GPU blocks. + # Asymmetric calculation split into outer joblib blocks and internal GPU tiles. ( - {"cutoff": 2, "n_blocks": 2, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 3}, + {"cutoff": 2, "n_blocks": 2, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 3}, (np.array(["AAAA", "AATA", "HHHH", "WWWW"]), np.array(["WWWW", "AAAA", "ATAA"])), np.array([[0, 1, 2], [0, 2, 3], [0, 0, 0], [1, 0, 0]]), ), # Distances above the cutoff and comparisons between sequences of unequal length are omitted. ( - {"cutoff": 1, "gpu_row_blocks": 3, "gpu_col_blocks": 3, "gpu_block_width": 2}, + {"cutoff": 1, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 2}, (np.array(["AAA", "AAT", "AAAA", "TTT"]), None), np.array([[1, 2, 0, 0], [2, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]), ), # Duplicate sequences are retained as separate zero-distance entries with distance + 1 encoding. ( - {"cutoff": 0, "gpu_row_blocks": 2, "gpu_col_blocks": 2, "gpu_block_width": 2}, + {"cutoff": 0, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 2}, (np.array(["AAA", "AAA", "AAT"]), None), np.array([[1, 1, 0], [1, 1, 0], [0, 0, 1]]), ), - # Block counts exceeding the number of sequences are reduced to avoid empty internal GPU blocks. + # Outer joblib block counts exceeding the number of sequences are supported. ( - {"cutoff": 1, "n_blocks": 5, "gpu_row_blocks": 5, "gpu_col_blocks": 5, "gpu_block_width": 2}, + {"cutoff": 1, "n_blocks": 5, "gpu_tile_rows": 1, "gpu_tile_cols": 1, "gpu_tile_buffer_cols": 2}, (np.array(["AAA", "AAT"]), None), np.array([[1, 2], [2, 1]]), ), @@ -1015,11 +1015,11 @@ def test_gpu_hamming(test_parameters, test_input, expected_result): @pytest.mark.gpu -def test_gpu_hamming_block_width_guard(): - hamming_calculator = GPUHammingDistanceCalculator(cutoff=0, gpu_col_blocks=1, gpu_block_width=1) +def test_gpu_hamming_buffer_retry(): + hamming_calculator = GPUHammingDistanceCalculator(cutoff=0, gpu_tile_cols=2, gpu_tile_buffer_cols=1) + result = hamming_calculator.calc_dist_mat(np.array(["AAA", "AAA"])) - with pytest.raises(ValueError, match="result block width is too small"): - hamming_calculator.calc_dist_mat(np.array(["AAA", "AAA"])) + npt.assert_array_equal(result.toarray(), np.ones((2, 2))) @pytest.mark.gpu @@ -1032,12 +1032,12 @@ def test_gpu_hamming_cutoff_guard(): @pytest.mark.parametrize( "kwargs, message", [ - ({"gpu_col_blocks": 0}, "`gpu_col_blocks` must be >= 1"), - ({"gpu_row_blocks": 0}, "`gpu_row_blocks` must be >= 1"), - ({"gpu_block_width": 0}, "`gpu_block_width` must be >= 1"), + ({"gpu_tile_rows": 0}, "`gpu_tile_rows` must be >= 1"), + ({"gpu_tile_cols": 0}, "`gpu_tile_cols` must be >= 1"), + ({"gpu_tile_buffer_cols": 0}, "`gpu_tile_buffer_cols` must be >= 1"), ], ) -def test_gpu_hamming_block_parameter_guards(kwargs, message): +def test_gpu_hamming_tile_parameter_guards(kwargs, message): with pytest.raises(ValueError, match=message): GPUHammingDistanceCalculator(**kwargs) @@ -1054,12 +1054,12 @@ def test_tcrdist_histogram_not_implemented(): @pytest.mark.parametrize( "kwargs", [ - {"gpu_col_blocks": 5, "gpu_block_width": 500}, - {"gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, - {"n_blocks": 3, "gpu_row_blocks": 3, "gpu_col_blocks": 5, "gpu_block_width": 500}, - {"gpu_col_blocks": 7, "gpu_block_width": 500}, - {"gpu_row_blocks": 5, "gpu_col_blocks": 3, "gpu_block_width": 500}, - {"gpu_row_blocks": 11, "gpu_col_blocks": 17, "gpu_block_width": 503}, + {"gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 517, "gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"n_blocks": 3, "gpu_tile_rows": 517, "gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_cols": 222, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 310, "gpu_tile_cols": 517, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 141, "gpu_tile_cols": 92, "gpu_tile_buffer_cols": 503}, ], ) def test_gpu_hamming_reference(kwargs): From 7f4b593f690eaec60adf4b4d6cfaeb4f6bb5cc9a Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 15:23:04 +0200 Subject: [PATCH 24/54] Log tile retry info --- src/scirpy/ir_dist/metrics.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 793bb76d7..bd2c8eddf 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1090,10 +1090,9 @@ def run_hamming_kernel(buffer_width): if required_buffer_width > buffer_width: # The buffer was too small, so retry with the required buffer size. - print( - f"GPU Hamming buffer retry for a {seqs_mat1_rows} x {seqs_mat2_rows} tile: " - f"{buffer_width} -> {required_buffer_width}", - flush=True, + logging.info( + f"GPU Hamming tile buffer increased from {buffer_width} to {required_buffer_width}; " + f"retrying the {seqs_mat1_rows} x {seqs_mat2_rows} tile." ) buffer_width = required_buffer_width d_data_matrix, d_indices_matrix, row_element_counts, _ = run_hamming_kernel(buffer_width) From 14e0bcf4720d7044be6c00b63ed8292f197e1f00 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 15:35:56 +0200 Subject: [PATCH 25/54] Test maximum GPU Hamming cutoff --- src/scirpy/tests/test_ir_dist_metrics.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index f01b98a95..f3175283c 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -974,6 +974,12 @@ def test_hamming_histogram_reference(): (np.array(["A" * 128, "AAB", "AABB", "ABA"]), np.array(["A" * 128, "ABBB", "ABBB"])), np.array([[1, 0, 0], [0, 0, 0], [0, 2, 2], [0, 0, 0]]), ), + # The maximum supported cutoff retains distance 125 with distance + 1 encoding and omits distance 126. + ( + {"cutoff": 125}, + (np.array(["A" * 127]), np.array(["B" * 125 + "AA", "B" * 126 + "A"])), + np.array([[126, 0]]), + ), # Symmetric calculation split into outer joblib blocks and internal GPU tiles. ( {"cutoff": 2, "n_blocks": 2, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 3}, From 9fea04dd99224333ee7cacf10c4555fa0e312248 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 16:06:45 +0200 Subject: [PATCH 26/54] Extend large dataset tutorial --- docs/tutorials/large-datasets.md | 43 +++++++++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/docs/tutorials/large-datasets.md b/docs/tutorials/large-datasets.md index ebdb8b8d8..995754765 100644 --- a/docs/tutorials/large-datasets.md +++ b/docs/tutorials/large-datasets.md @@ -59,10 +59,47 @@ First, install the optional `cupy` dependency: !pip install scirpy[cupy] ``` -Then simply run +Then run -``` +```python ir.pp.ir_dist(mdata, metric="gpu_hamming") ``` -to take advantage of GPU acceleration. +to use the first available GPU. The calculation is split into tiles. Their size can be adjusted with +`gpu_tile_rows` and `gpu_tile_cols`: smaller tiles use less GPU memory but increase processing overhead. +`gpu_tile_buffer_cols` controls the initially reserved space for retained distances and is enlarged automatically +if necessary. + +### Using multiple GPUs + +Multiple GPUs can be used through the joblib dask backend and +[dask-cuda](https://docs.rapids.ai/api/dask-cuda/stable/): + +``` +!pip install "dask[distributed]" dask-cuda +``` + +The following example creates one dask worker per GPU and divides the distance calculation into two outer +partitions: + +```python +import joblib +from dask.distributed import Client +from dask_cuda import LocalCUDACluster + +with LocalCUDACluster( + CUDA_VISIBLE_DEVICES="0,1", + n_workers=2, + threads_per_worker=1, +) as cluster, Client(cluster): + with joblib.parallel_config(backend="dask", n_jobs=2): + ir.pp.ir_dist( + mdata, + metric="gpu_hamming", + n_blocks=2, + ) +``` + +Set `CUDA_VISIBLE_DEVICES`, `n_workers`, `n_jobs`, and `n_blocks` to the number of GPUs to use. `n_blocks` controls +the outer partitions distributed between workers, whereas `gpu_tile_rows` and `gpu_tile_cols` control the smaller +tiles computed within each worker. From 717974b292e914b70d6a7f237589b75bfc9e7ced Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Thu, 23 Jul 2026 20:26:34 +0200 Subject: [PATCH 27/54] Reformat --- src/scirpy/ir_dist/metrics.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index bd2c8eddf..ab0ee0afd 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -554,9 +554,7 @@ def calc_dist_mat(self, seqs: Sequence[str], seqs2: Sequence[str] | None = None) # parallel jobs. Increasing n_blocks for better load balancing instead would add block-processing # overhead. partition_fractions = np.arange(self.n_blocks + 1) / self.n_blocks - partition_boundaries = np.rint( - len(seqs) * (1 - np.sqrt(1 - partition_fractions)) - ).astype(int) + partition_boundaries = np.rint(len(seqs) * (1 - np.sqrt(1 - partition_fractions))).astype(int) else: partition_boundaries = np.rint(np.linspace(0, len(seqs), self.n_blocks + 1)).astype(int) @@ -1086,7 +1084,9 @@ def run_hamming_kernel(buffer_width): d_indices_matrix = None return d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width - d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width = run_hamming_kernel(buffer_width) + d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width = run_hamming_kernel( + buffer_width + ) if required_buffer_width > buffer_width: # The buffer was too small, so retry with the required buffer size. @@ -1150,9 +1150,7 @@ def run_hamming_kernel(buffer_width): seqs_mat2_blocks = np.array_split(seqs_mat2, n_col_blocks) seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) - logging.info( - f"\nStart GPU calculations for {n_row_blocks} row tiles x {n_col_blocks} column tiles:" - ) + logging.info(f"\nStart GPU calculations for {n_row_blocks} row tiles x {n_col_blocks} column tiles:") @nb.njit def csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elements): From 7feb2b4a50bc5cab4b5576501277a326e0f04da0 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Sun, 26 Jul 2026 13:03:47 +0200 Subject: [PATCH 28/54] Adapt input array memory management --- src/scirpy/ir_dist/metrics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index ab0ee0afd..211088012 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -541,8 +541,8 @@ def calc_dist_mat(self, seqs: Sequence[str], seqs2: Sequence[str] | None = None) if seqs2 is None: seqs2 = seqs - seqs = np.array(seqs) - seqs2 = np.array(seqs2) + seqs = np.asarray(seqs) + seqs2 = np.asarray(seqs2) is_symmetric = np.array_equal(seqs, seqs2) if self.n_blocks < 2: From 3b62993a15b7d61c57529d81d363a6df8ffffb7c Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 09:35:09 +0200 Subject: [PATCH 29/54] Avoid redundant host allocations for GPU Hamming result buffers --- src/scirpy/ir_dist/metrics.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 211088012..040802e2b 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1111,11 +1111,8 @@ def run_hamming_kernel(buffer_width): d_indptr = cp.asarray(indptr) n_elements = indptr[-1] - data = np.zeros(n_elements, dtype=np.int32) - d_data = cp.zeros_like(data) - - indices = np.zeros(n_elements, dtype=np.int32) - d_indices = cp.zeros_like(indices) + d_data = cp.zeros(n_elements, dtype=cp.int32) + d_indices = cp.zeros(n_elements, dtype=cp.int32) threads_per_block = (1, 256) blocks_per_grid_x = (d_data_matrix.shape[0] + threads_per_block[0] - 1) // threads_per_block[0] From ef7b651d713404bb8ec8ced3bccdf048a80395c1 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 09:37:31 +0200 Subject: [PATCH 30/54] Reuse host indptr in GPU Hamming CSR assembly --- src/scirpy/ir_dist/metrics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 040802e2b..3f8dc3b02 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1136,7 +1136,6 @@ def run_hamming_kernel(buffer_width): ) data = d_data.get() - indptr = d_indptr.get() indices = d_indices.get() res = csr_matrix((data, indices, indptr), shape=(seqs_mat1.shape[0], seqs_mat2.shape[0])) From c85a43a9627fd699050c781dbbed5d0f7eb25600 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 09:50:16 +0200 Subject: [PATCH 31/54] Avoid retaining redundant GPU Hamming input buffers --- src/scirpy/ir_dist/metrics.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 3f8dc3b02..005760bda 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1037,8 +1037,6 @@ def calc_col_block_gpu( seqs2_original_indices_block, buffer_width, ): - d_seqs_mat1 = cp.asarray(seqs_mat1.astype(np.int8, copy=False)) - d_seqs_mat2 = cp.asarray(seqs_mat2_block.astype(np.int8, copy=False)) d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32, copy=False)) @@ -1048,8 +1046,8 @@ def calc_col_block_gpu( seqs_mat1_rows = seqs_mat1.shape[0] seqs_mat2_rows = seqs_mat2_block.shape[0] - d_seqs_mat1_transposed = cp.transpose(d_seqs_mat1).copy() - d_seqs_mat2_transposed = cp.transpose(d_seqs_mat2).copy() + d_seqs_mat1_transposed = cp.transpose(cp.asarray(seqs_mat1.astype(np.int8, copy=False))).copy() + d_seqs_mat2_transposed = cp.transpose(cp.asarray(seqs_mat2_block.astype(np.int8, copy=False))).copy() def run_hamming_kernel(buffer_width): d_data_matrix = cp.empty((seqs_mat1_rows, buffer_width), dtype=cp.int8) From 648365001b055f64f773f60c5cf762e33c9f468e Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 10:10:17 +0200 Subject: [PATCH 32/54] Avoid temporary sequence tuples in GPU Hamming --- src/scirpy/ir_dist/metrics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 005760bda..98fe423f2 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -919,7 +919,7 @@ def _gpu_hamming_mat( seqs2 = np.concatenate(seqs2_sorted_per_block) - max_seq_len = max(len(s) for s in (*seqs, *seqs2)) + max_seq_len = max(len(s) for s in itertools.chain(seqs, seqs2)) def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np.ndarray, np.ndarray]: if max_len is None: @@ -938,7 +938,7 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. logging.info( "UnicodeError error occurred while converting sequences, retrying with implementation for non ascii sequences" ) - unique_characters = "".join(sorted({char for string in (*seqs, *seqs2) for char in string})) + unique_characters = "".join(sorted({char for string in itertools.chain(seqs, seqs2) for char in string})) seqs_mat1, seqs_L1 = _seqs2mat(seqs, alphabet=unique_characters, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat(seqs2, alphabet=unique_characters, max_len=max_seq_len) From 32d097c72b5afedce1e0e3d0c338cd4e1860de51 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 10:35:46 +0200 Subject: [PATCH 33/54] Prevent theoretical overflow in GPU Hamming combined entry count --- src/scirpy/ir_dist/metrics.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 98fe423f2..fa885c470 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1229,9 +1229,7 @@ def calc_row_block_gpu( buffer_width, ) - num_elements = 0 - for i in range(0, len(result_blocks)): - num_elements += result_blocks[i].indptr[-1] + num_elements = sum(int(block.indptr[-1]) for block in result_blocks) if num_elements > np.iinfo(np.int32).max: raise ValueError( From feffe20ff14fd3f93d75278de789d3f307fb5d59 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 10:49:57 +0200 Subject: [PATCH 34/54] Use int32 for empty GPU Hamming result tiles --- src/scirpy/ir_dist/metrics.py | 2 +- src/scirpy/tests/test_ir_dist_metrics.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index fa885c470..fe3246f10 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1224,7 +1224,7 @@ def calc_row_block_gpu( if not result_blocks: return ( - csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0])), + csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0]), dtype=np.int32), n_calculated_blocks, buffer_width, ) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index f3175283c..7d85de070 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -1017,6 +1017,8 @@ def test_gpu_hamming(test_parameters, test_input, expected_result): res = hamming_calculator.calc_dist_mat(*test_input) assert isinstance(res, scipy.sparse.csr_matrix) + if test_parameters.get("n_blocks", 1) <= len(test_input[0]): + assert res.dtype == np.dtype("int32") npt.assert_array_equal(res.toarray(), expected_result) From 42468c259f3a42337c7d807df6a2b94c17b2b7c7 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 12:31:48 +0200 Subject: [PATCH 35/54] Precompute sequence length ranges for GPU Hamming tiles --- src/scirpy/ir_dist/metrics.py | 56 +++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 23 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index fe3246f10..581bbc185 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -949,7 +949,8 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. const char* __restrict__ seqs_mat1, const char* __restrict__ seqs_mat2, const int* __restrict__ seqs_L1, - const int* seqs_L2, + const int* __restrict__ length_starts, + const int* __restrict__ length_ends, const int* __restrict__ seqs_original_indices, const int* seqs2_original_indices, const int cutoff, @@ -967,30 +968,29 @@ def _seqs2mat_fast(seqs: Sequence[str], max_len: None | int = None) -> tuple[np. int seq1_len = seqs_L1[row]; int row_end_index = 0; - for (int col = 0; col < seqs_mat2_rows; col++) { - int seq2_len = seqs_L2[col]; + int col_start = length_starts[seq1_len]; + int col_end = length_ends[seq1_len]; + for (int col = col_start; col < col_end; col++) { char distance = 1; - if (seq1_len == seq2_len) { - for (int i = 0; i < seq1_len; i++) { - char val1 = seqs_mat1[i*seqs_mat1_rows+row]; - char val2 = seqs_mat2[i*seqs_mat2_rows+col]; + for (int i = 0; i < seq1_len; i++) { + char val1 = seqs_mat1[i*seqs_mat1_rows+row]; + char val2 = seqs_mat2[i*seqs_mat2_rows+col]; - if(val1 != val2) { - distance++; - if (distance > cutoff + 1) { - break; - } + if(val1 != val2) { + distance++; + if (distance > cutoff + 1) { + break; } } - if (distance <= cutoff + 1) { - if (row_end_index < data_cols) { - int seqs2_original_index = seqs2_original_indices[col]; - data[(long long)seqs_original_index * data_cols + row_end_index] = distance; - indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; - } - row_end_index++; + } + if (distance <= cutoff + 1) { + if (row_end_index < data_cols) { + int seqs2_original_index = seqs2_original_indices[col]; + data[(long long)seqs_original_index * data_cols + row_end_index] = distance; + indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; } + row_end_index++; } } row_element_counts[seqs_original_index] = row_end_index; @@ -1032,13 +1032,13 @@ def calc_col_block_gpu( seqs_mat1, seqs_mat2_block, seqs_L1_block, - seqs_L2, + length_bounds, seqs_original_indices_block, seqs2_original_indices_block, buffer_width, ): d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) - d_seqs_L2 = cp.asarray(seqs_L2.astype(np.int32, copy=False)) + d_length_starts, d_length_ends = length_bounds threads_per_block = 256 blocks_per_grid = (seqs_mat1.shape[0] + (threads_per_block - 1)) // threads_per_block @@ -1061,7 +1061,8 @@ def run_hamming_kernel(buffer_width): d_seqs_mat1_transposed, d_seqs_mat2_transposed, d_seqs_L1, - d_seqs_L2, + d_length_starts, + d_length_ends, seqs_original_indices_block, seqs2_original_indices_block, self.cutoff, @@ -1144,6 +1145,15 @@ def run_hamming_kernel(buffer_width): seqs_mat2_blocks = np.array_split(seqs_mat2, n_col_blocks) seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) + possible_lengths = np.arange(max_seq_len + 1) + length_bounds_blocks = [ + ( + cp.asarray(np.searchsorted(lengths, possible_lengths, side="left").astype(np.int32)), + cp.asarray(np.searchsorted(lengths, possible_lengths, side="right").astype(np.int32)), + ) + for lengths in seqs_L2_blocks + ] + logging.info(f"\nStart GPU calculations for {n_row_blocks} row tiles x {n_col_blocks} column tiles:") @nb.njit @@ -1214,7 +1224,7 @@ def calc_row_block_gpu( seqs_mat1_block, seqs_mat2_blocks[i], seqs_L1_block, - seqs_L2_blocks[i], + length_bounds_blocks[i], seqs_original_indices_block, seqs2_original_indices_blocks[i], buffer_width, From 1ad0221bea9665bcb5f3a5f28a37b483f31e49d1 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 13:28:18 +0200 Subject: [PATCH 36/54] Deprecate legacy GPU Hamming parameters --- CHANGELOG.md | 12 ++++++++---- src/scirpy/ir_dist/metrics.py | 18 ++++++++++++++++++ src/scirpy/tests/test_deprecations.py | 26 ++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 66bba91b8..2c8d731eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,14 @@ and this project adheres to [Semantic Versioning][]. `metric="needleman_wunsch"`. The `alignment` and `fastalignment` metrics are now deprecated. When `gap_open == gap_extend` (which has always been the default), use `needleman_wunsch` instead. +### Performance improvements + + - Improve the GPU implementation of the Hamming distance metric with row and column tiling, adaptive result buffers, + and other performance optimizations. GPU tiles can be configured + with `gpu_tile_rows`, `gpu_tile_cols`, and `gpu_tile_buffer_cols`. The previous `gpu_n_blocks` and `gpu_block_width` + parameters are deprecated and ignored; calls using them emit a warning and use the new parameters' defaults + unless the new parameters are explicitly provided. + ## v0.25.1 ### Fixes @@ -47,10 +55,6 @@ and this project adheres to [Semantic Versioning][]. ### Performance improvements - Speed up identity distance metric computation for comparisons between two different sequence arrays ([#701](https://github.com/scverse/scirpy/pull/701)). - - Improve the GPU implementation of the Hamming distance metric with row and column tiling, tile skipping for - symmetric distance matrices, adaptive result buffers, and faster sparse tile assembly. GPU tiles can be configured - with `gpu_tile_rows` and `gpu_tile_cols`; `gpu_n_blocks` has been replaced by `gpu_tile_cols`, and `gpu_block_width` - has been replaced by `gpu_tile_buffer_cols`. ### Chore diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 1a4dd0507..71c926062 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -804,8 +804,24 @@ class GPUHammingDistanceCalculator(_MetricDistanceCalculator): gpu_tile_buffer_cols: Initial number of retained sparse entries reserved per row of each tile. Higher values can avoid retries for denser results but require more GPU memory. + gpu_n_blocks: + Deprecated since v0.26.0 and ignored. Use `gpu_tile_cols` to set the number of columns per GPU tile. + gpu_block_width: + Deprecated since v0.26.0 and ignored. Use `gpu_tile_buffer_cols` to set the initial result buffer width. """ + @deprecated_arg( + "gpu_n_blocks", + Deprecation( + "0.26.0", "This argument is ignored. Use `gpu_tile_cols` to set the number of columns per GPU tile." + ), + ) + @deprecated_arg( + "gpu_block_width", + Deprecation( + "0.26.0", "This argument is ignored. Use `gpu_tile_buffer_cols` to set the initial result buffer width." + ), + ) def __init__( self, *, @@ -814,6 +830,8 @@ def __init__( gpu_tile_rows: int = 100_000, gpu_tile_cols: int = 100_000, gpu_tile_buffer_cols: int = 1000, + gpu_n_blocks: int | None = None, + gpu_block_width: int | None = None, ): super().__init__(n_jobs=1, n_blocks=n_blocks) if cutoff > 125: diff --git a/src/scirpy/tests/test_deprecations.py b/src/scirpy/tests/test_deprecations.py index 4bea4acc1..abeded67f 100644 --- a/src/scirpy/tests/test_deprecations.py +++ b/src/scirpy/tests/test_deprecations.py @@ -14,6 +14,7 @@ from scirpy.ir_dist.metrics import ( AlignmentDistanceCalculator, FastAlignmentDistanceCalculator, + GPUHammingDistanceCalculator, LevenshteinDistanceCalculator, ) @@ -50,6 +51,31 @@ def test_no_spurious_block_size_warning(calculator): _assert_no_deprecation_warning(calculator, "block_size") +@pytest.mark.parametrize( + "legacy_kwargs", + [{"gpu_n_blocks": 2}, {"gpu_block_width": 3}, {"gpu_n_blocks": 2, "gpu_block_width": 3}], +) +@pytest.mark.parametrize( + "tile_kwargs", + [{}, {"gpu_tile_rows": 11, "gpu_tile_cols": 13, "gpu_tile_buffer_cols": 5}], +) +def test_gpu_hamming_deprecated_parameters(legacy_kwargs, tile_kwargs): + expected = GPUHammingDistanceCalculator(**tile_kwargs) + with pytest.warns(FutureWarning) as record: + actual = GPUHammingDistanceCalculator(**legacy_kwargs, **tile_kwargs) + + assert len(record) == len(legacy_kwargs) + for arg in legacy_kwargs: + assert any(f"argument {arg} is deprecated" in str(w.message) and "ignored" in str(w.message) for w in record) + for attr in ("gpu_tile_rows", "gpu_tile_cols", "gpu_tile_buffer_cols"): + assert getattr(actual, attr) == getattr(expected, attr) + + +def test_no_spurious_gpu_hamming_warning(): + _assert_no_deprecation_warning(GPUHammingDistanceCalculator, "gpu_") + _assert_no_deprecation_warning(GPUHammingDistanceCalculator, "gpu_", gpu_tile_cols=13, gpu_tile_buffer_cols=5) + + def test_alignment_distance_calculator_deprecated(): with pytest.warns(FutureWarning, match="NeedlemanWunschDistanceCalculator"): AlignmentDistanceCalculator() From a211e40029ec62f8d014d75c8a8a18870df57d73 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Fri, 11 Sep 2026 14:08:18 +0200 Subject: [PATCH 37/54] Clarify multi-GPU configuration and link to advanced usage docs --- docs/tutorials/large-datasets.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/tutorials/large-datasets.md b/docs/tutorials/large-datasets.md index 995754765..7fa0c3e77 100644 --- a/docs/tutorials/large-datasets.md +++ b/docs/tutorials/large-datasets.md @@ -100,6 +100,10 @@ with LocalCUDACluster( ) ``` -Set `CUDA_VISIBLE_DEVICES`, `n_workers`, `n_jobs`, and `n_blocks` to the number of GPUs to use. `n_blocks` controls -the outer partitions distributed between workers, whereas `gpu_tile_rows` and `gpu_tile_cols` control the smaller -tiles computed within each worker. +Set `CUDA_VISIBLE_DEVICES` to a comma-separated list of GPU IDs, such as `"0,1"` for GPUs 0 and 1. +For this setup, set `n_workers`, `n_jobs`, and `n_blocks` to the number of selected GPUs. `n_blocks` controls +the number of outer partitions distributed between workers, whereas `gpu_tile_rows` and `gpu_tile_cols` control +the size (rows and columns) of the smaller tiles computed within each worker. + +For advanced Dask-CUDA configuration and GPU memory management, see the +[rapids-singlecell documentation](https://rapids-singlecell.readthedocs.io/en/latest/out_of_core.html). From 772055e9cf9c841a4f9269ec2f4184235c520a7c Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 14 Sep 2026 12:24:09 +0200 Subject: [PATCH 38/54] Add TCRdist GPU implementation --- docs/api.rst | 1 + src/scirpy/ir_dist/__init__.py | 6 + src/scirpy/ir_dist/metrics.py | 578 +++++++++++++++++++++++ src/scirpy/tests/test_ir_dist.py | 5 +- src/scirpy/tests/test_ir_dist_metrics.py | 184 +++++++- 5 files changed, 763 insertions(+), 11 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 8e8ed3172..ce19a0e2c 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -311,4 +311,5 @@ distance metrics ir_dist.metrics.AlignmentDistanceCalculator ir_dist.metrics.FastAlignmentDistanceCalculator ir_dist.metrics.TCRdistDistanceCalculator + ir_dist.metrics.GPUTCRdistDistanceCalculator ir_dist.metrics.NeedlemanWunschDistanceCalculator diff --git a/src/scirpy/ir_dist/__init__.py b/src/scirpy/ir_dist/__init__.py index 8104795d1..0a1126ea3 100644 --- a/src/scirpy/ir_dist/__init__.py +++ b/src/scirpy/ir_dist/__init__.py @@ -39,6 +39,7 @@ def IrNeighbors(*args, **kwargs): "gpu_hamming", "normalized_hamming", "tcrdist", + "gpu_tcrdist", "needleman_wunsch", ] | metrics.DistanceCalculator @@ -56,6 +57,9 @@ def IrNeighbors(*args, **kwargs): Uses the BLOSUM62 substitution matrix by default. TCRBLOSUM alpha/beta substitution matrices (:cite:`TCRBLOSUM`) can be selected with `base_matrix="tcrblosum"`. See :class:`~scirpy.ir_dist.metrics.TCRdistDistanceCalculator`. + * `gpu_tcrdist` -- TCRdist calculated with a GPU, with support for BLOSUM62 and TCRBLOSUM matrices. + Currently only supports `fixed_gappos=True`. + See :class:`~scirpy.ir_dist.metrics.GPUTCRdistDistanceCalculator`. * `needleman_wunsch` -- Distance based on linear-gap Needleman-Wunsch global alignment. Uses the BLOSUM62 substitution matrix. This option is incompatible with nucleotide sequences. @@ -127,6 +131,8 @@ def _get_distance_calculator( dist_calc = metrics.GPUHammingDistanceCalculator(**kwargs) elif metric == "tcrdist": dist_calc = metrics.TCRdistDistanceCalculator(n_jobs=n_jobs, chain_type=chain_type, **kwargs) + elif metric == "gpu_tcrdist": + dist_calc = metrics.GPUTCRdistDistanceCalculator(chain_type=chain_type, **kwargs) elif metric == "needleman_wunsch": dist_calc = metrics.NeedlemanWunschDistanceCalculator(n_jobs=n_jobs, **kwargs) else: diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 71c926062..1cf41b621 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1552,6 +1552,584 @@ def _nb_tcrdist_mat(): _metric_mat = _tcrdist_mat +class GPUTCRdistDistanceCalculator(TCRdistDistanceCalculator): + """Computes pairwise distances between TCR CDR3 sequences based on the "tcrdist" distance metric with GPU support. + + The code of this class is heavily based on `pwseqdist `_. + Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. + + Using default weight, gap penalty, ntrim and ctrim is equivalent to the + original distance published in :cite:`TCRdist`. + + For performance reasons, the rows and columns of the final result matrix are grouped into tiles for GPU + computation. `gpu_tile_rows` and `gpu_tile_cols` control how many matrix rows and columns are grouped into each + tile. Each tile is computed on the GPU and converted to a sparse CSR matrix before the tiles are combined again. + + `gpu_tile_buffer_cols` controls how many buffer columns are initially reserved for sparse result entries in each + row of a tile. Because only distances at or below the cutoff are retained, the number of entries that need to be + stored is usually considerably smaller than the number of columns in the tile. If necessary, the buffer is + enlarged and the calculation is retried. + + Smaller tiles reduce per-tile memory pressure but add tile-management overhead. Larger values for + `gpu_tile_buffer_cols` can avoid retries but require more GPU memory. + + Only `fixed_gappos=True` is currently supported. Distances are accumulated using 64-bit integers; + the GPU result buffers and the final CSR data use 32-bit integers. + + Parameters + ---------- + dist_weight: + Weight applied to the mismatch distances before summing with the gap penalties + gap_penalty: + Distance penalty for the difference in the length of the two sequences + ntrim/ctrim: + Positions trimmed off the N-terminus (0) and C-terminus (L-1) ends of the peptide sequence. These symbols will be ignored + in the distance calculation. + fixed_gappos: + If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6). + Only True is currently supported; False raises NotImplementedError. + cutoff: + Will eliminate distances > cutoff to make efficient use of sparse matrices. Must be between 0 and + 2**31 - 2 to allow distance + 1 encoding in the GPU result buffer. + n_blocks: + Number of outer row partitions submitted through joblib. This can be used with a distributed joblib backend to + distribute the calculation across multiple GPU workers. + gpu_tile_rows: + Number of result matrix rows per GPU tile. + gpu_tile_cols: + Number of result matrix columns per GPU tile. + gpu_tile_buffer_cols: + Initial number of retained sparse entries reserved per row of each tile. Higher values can avoid retries for + denser results but require more GPU memory. + base_matrix: + Amino acid substitution matrix used by TCRdist. `"blosum62"` uses the original + BLOSUM62 substitution matrix, while `"tcrblosum"` uses TCRBLOSUM substitution + matrices (:cite:`TCRBLOSUM`). Depending on `chain_type`, either the TCRBLOSUM + alpha- or beta-chain matrix is used. + distance_cap: + Maximum distance assigned to a mismatch after converting substitution scores to distances. + The default value, `"default"`, keeps the original behavior: BLOSUM62 uses a cap of `4`, + while TCRBLOSUM distances are uncapped. Set to an integer to choose a cap explicitly, or + `None` for uncapped distances. + chain_type: + Required when `base_matrix="tcrblosum"`. `"VJ"` selects the alpha-chain matrix + and `"VDJ"` selects the beta-chain matrix. When called via `ir_dist`, this value + is set automatically and should not be provided. + """ + + def __init__( + self, + cutoff: int = 20, + *, + dist_weight: int = 3, + gap_penalty: int = 4, + ntrim: int = 3, + ctrim: int = 2, + fixed_gappos: bool = True, + n_blocks: int = 1, + gpu_tile_rows: int = 100_000, + gpu_tile_cols: int = 100_000, + gpu_tile_buffer_cols: int = 1000, + base_matrix: Literal["blosum62", "tcrblosum"] = "blosum62", + distance_cap: int | None | Literal["default"] = "default", + chain_type: Literal["VJ", "VDJ"] | None = None, + ): + if not fixed_gappos: + raise NotImplementedError("Only `fixed_gappos=True` is currently supported for GPU TCRdist.") + if not isinstance(cutoff, (int, np.integer)) or not 0 <= cutoff < np.iinfo(np.int32).max: + raise ValueError("`cutoff` must be an integer between 0 and 2**31 - 2.") + for name, value in ( + ("dist_weight", dist_weight), + ("gap_penalty", gap_penalty), + ("ntrim", ntrim), + ("ctrim", ctrim), + ): + if not isinstance(value, (int, np.integer)) or not 0 <= value <= np.iinfo(np.int32).max: + raise ValueError(f"`{name}` must be a non-negative 32-bit integer.") + if gpu_tile_rows < 1: + raise ValueError("`gpu_tile_rows` must be >= 1.") + if gpu_tile_cols < 1: + raise ValueError("`gpu_tile_cols` must be >= 1.") + if gpu_tile_buffer_cols < 1: + raise ValueError("`gpu_tile_buffer_cols` must be >= 1.") + + super().__init__( + cutoff=cutoff, + dist_weight=dist_weight, + gap_penalty=gap_penalty, + ntrim=ntrim, + ctrim=ctrim, + fixed_gappos=fixed_gappos, + n_jobs=1, + n_blocks=n_blocks, + base_matrix=base_matrix, + distance_cap=distance_cap, + chain_type=chain_type, + ) + self.gpu_tile_rows = gpu_tile_rows + self.gpu_tile_cols = gpu_tile_cols + self.gpu_tile_buffer_cols = gpu_tile_buffer_cols + + def _gpu_tcrdist_mat( + self, + *, + seqs: Sequence[str], + seqs2: Sequence[str], + is_symmetric: bool = False, + start_column: int = 0, + ) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray]: + """Computes the pairwise tcrdist distances for sequences in seqs and seqs2 with GPU support. + + Parameters + ---------- + seqs/2: + A python sequence of strings representing gene sequences + is_symmetric: + Determines whether the final result matrix is symmetric, assuming that this function is + only used to compute a block of a bigger result matrix + start_column: + Global row offset of an outer row block scheduled by joblib. Used to skip column blocks below the diagonal + when computing a symmetric result matrix. + + Returns + ------- + data_rows: + List with array containing the non-zero data values of the result matrix, + needed to create the final scipy CSR result matrix later + indices_rows: + List with array containing the non-zero entry column indeces of the result matrix, + needed to create the final scipy CSR result matrix later + row_element_counts: + Array with integers that indicate the amount of non-zero values of the result matrix per row, + needed to create the final scipy CSR result matrix later + row_mins: + Always returns a numpy array containing None because the computation of the minimum distance per row is + not implemented for the GPU tcrdist calculator yet. + """ + import cupy as cp + from tqdm import tqdm + + n_col_blocks = (len(seqs2) + self.gpu_tile_cols - 1) // self.gpu_tile_cols + n_row_blocks = (len(seqs) + self.gpu_tile_rows - 1) // self.gpu_tile_rows + + seqs_blocks = np.array_split(np.asarray(seqs), n_row_blocks) + seqs_block_starts = np.cumsum([0] + [len(block) for block in seqs_blocks[:-1]]) + seqs_sorted_per_block = [] + seqs_original_indices_blocks = [] + + for seqs_block in seqs_blocks: + seqs_block_lengths = np.vectorize(len)(seqs_block) + seqs_block_sort_indices = np.argsort(seqs_block_lengths) + seqs_sorted_per_block.append(seqs_block[seqs_block_sort_indices]) + seqs_original_indices_blocks.append(cp.asarray(seqs_block_sort_indices.astype(np.int32))) + + seqs = np.concatenate(seqs_sorted_per_block) + + seqs2_blocks = np.array_split(np.asarray(seqs2), n_col_blocks) + seqs2_block_starts = np.cumsum([0] + [len(block) for block in seqs2_blocks[:-1]]) + seqs2_sorted_per_block = [] + seqs2_original_indices_blocks = [] + seqs2_block_start = 0 + + for seqs2_block in seqs2_blocks: + seqs2_block_lengths = np.vectorize(len)(seqs2_block) + seqs2_block_sort_indices = np.argsort(seqs2_block_lengths) + seqs2_sorted_per_block.append(seqs2_block[seqs2_block_sort_indices]) + seqs2_original_indices_blocks.append( + cp.asarray((seqs2_block_sort_indices + seqs2_block_start).astype(np.int32)) + ) + seqs2_block_start += len(seqs2_block) + + seqs2 = np.concatenate(seqs2_sorted_per_block) + + max_seq_len = max(len(s) for s in itertools.chain(seqs, seqs2)) + + seqs_mat1, seqs_L1 = _seqs2mat(seqs, max_len=max_seq_len) + seqs_mat2, seqs_L2 = _seqs2mat(seqs2, max_len=max_seq_len) + d_dist_mat = cp.asarray(self.tcr_nb_distance_matrix.astype(np.int32, copy=False)) + + tcrdist_kernel = cp.RawKernel( + r""" + extern "C" __global__ __launch_bounds__(256) + void tcrdist_kernel( + const char* __restrict__ seqs_mat1, + const char* __restrict__ seqs_mat2, + const int* __restrict__ seqs_L1, + const int* __restrict__ seqs_L2, + const int* __restrict__ length_starts, + const int* __restrict__ length_ends, + const int* __restrict__ seqs_original_indices, + const int* seqs2_original_indices, + const int cutoff, + int* __restrict__ data, + int* __restrict__ indices, + int* __restrict__ row_element_counts, + const int seqs_mat1_rows, + const int seqs_mat2_rows, + const int data_cols, + const int indices_cols, + const int* __restrict__ aa_distance_matrix, + const int alphabet_size, + const int dist_weight, + const int gap_penalty, + const int ntrim, + const int ctrim + ) { + int row = blockDim.x * blockIdx.x + threadIdx.x; + if (row < seqs_mat1_rows) { + int seqs_original_index = seqs_original_indices[row]; + int seq1_len = seqs_L1[row]; + int row_end_index = 0; + + int col_start = length_starts[seq1_len]; + int col_end = length_ends[seq1_len]; + for (int col = col_start; col < col_end; col++) { + int seq2_len = seqs_L2[col]; + long long distance = 1 + (long long)gap_penalty * abs(seq1_len - seq2_len); + + if (seq1_len == seq2_len) { + for (int i = ntrim; i < seq1_len - ctrim; i++) { + char val1 = seqs_mat1[(long long)i * seqs_mat1_rows + row]; + char val2 = seqs_mat2[(long long)i * seqs_mat2_rows + col]; + distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + if (distance > cutoff + 1) { + break; + } + } + } + else { + int short_len = min(seq1_len, seq2_len); + + // Equivalent to min(6, 3 + (short_len - 5) // 2) in the CPU implementation. + // Needed since Python rounds integer division down, whereas CUDA C++ truncates toward zero. + int gappos = min(6, (short_len + 1) / 2); + int remainder = short_len - gappos; + + for (int n_i = ntrim; n_i < gappos; n_i++) { + char val1 = seqs_mat1[(long long)n_i * seqs_mat1_rows + row]; + char val2 = seqs_mat2[(long long)n_i * seqs_mat2_rows + col]; + distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + if (distance > cutoff + 1) { + break; + } + } + if (distance <= cutoff + 1) { + for (int c_i = ctrim; c_i < remainder; c_i++) { + int i = seq1_len - 1 - c_i; + int j = seq2_len - 1 - c_i; + char val1 = seqs_mat1[(long long)i * seqs_mat1_rows + row]; + char val2 = seqs_mat2[(long long)j * seqs_mat2_rows + col]; + distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + if (distance > cutoff + 1) { + break; + } + } + } + } + if (distance <= cutoff + 1) { + if (row_end_index < data_cols) { + int seqs2_original_index = seqs2_original_indices[col]; + data[(long long)seqs_original_index * data_cols + row_end_index] = distance; + indices[(long long)seqs_original_index * indices_cols + row_end_index] = seqs2_original_index; + } + row_end_index++; + } + } + row_element_counts[seqs_original_index] = row_end_index; + } + } + """, + "tcrdist_kernel", + options=("--maxrregcount=256",), + ) + + create_csr_kernel = cp.RawKernel( + r""" + extern "C" __global__ + void create_csr_kernel( + int* data, int* indices, + int* data_matrix, int* indices_matrix, + int* indptr, int data_matrix_rows, int data_matrix_cols, int data_rows, int indices_matrix_cols + ) { + int row = blockDim.x * blockIdx.x + threadIdx.x; + int col = blockDim.y * blockIdx.y + threadIdx.y; + + if (row < data_matrix_rows && col < data_matrix_cols) { + int row_start = indptr[row]; + int row_end = indptr[row + 1]; + int row_end_index = row_end - row_start; + int data_index = row_start + col; + + if ((data_index < data_rows) && (col < row_end_index)) { + data[data_index] = data_matrix[(long long)row * data_matrix_cols + col]; + indices[data_index] = indices_matrix[(long long)row * indices_matrix_cols + col]; + } + } + } + """, + "create_csr_kernel", + ) + + def calc_col_block_gpu( + seqs_mat1, + seqs_mat2_block, + seqs_L1_block, + seqs_L2_block, + length_bounds, + seqs_original_indices_block, + seqs2_original_indices_block, + buffer_width, + ): + d_seqs_L1 = cp.asarray(seqs_L1_block.astype(np.int32, copy=False)) + d_seqs_L2 = cp.asarray(seqs_L2_block.astype(np.int32, copy=False)) + d_length_starts, d_length_ends = length_bounds + + threads_per_block = 256 + blocks_per_grid = (seqs_mat1.shape[0] + (threads_per_block - 1)) // threads_per_block + + seqs_mat1_rows = seqs_mat1.shape[0] + seqs_mat2_rows = seqs_mat2_block.shape[0] + + d_seqs_mat1_transposed = cp.transpose(cp.asarray(seqs_mat1.astype(np.int8, copy=False))).copy() + d_seqs_mat2_transposed = cp.transpose(cp.asarray(seqs_mat2_block.astype(np.int8, copy=False))).copy() + + def run_tcrdist_kernel(buffer_width): + d_data_matrix = cp.empty((seqs_mat1_rows, buffer_width), dtype=cp.int32) + d_indices_matrix = cp.empty((seqs_mat1_rows, buffer_width), dtype=np.int32) + d_row_element_counts = cp.zeros(seqs_mat1_rows, dtype=np.int32) + + tcrdist_kernel( + (blocks_per_grid,), + (threads_per_block,), + ( + d_seqs_mat1_transposed, + d_seqs_mat2_transposed, + d_seqs_L1, + d_seqs_L2, + d_length_starts, + d_length_ends, + seqs_original_indices_block, + seqs2_original_indices_block, + self.cutoff, + d_data_matrix, + d_indices_matrix, + d_row_element_counts, + seqs_mat1_rows, + seqs_mat2_rows, + buffer_width, + buffer_width, + d_dist_mat, + self.tcr_nb_distance_matrix.shape[0], + self.dist_weight, + self.gap_penalty, + self.ntrim, + self.ctrim, + ), + ) + row_element_counts = d_row_element_counts.get() + required_buffer_width = int(np.max(row_element_counts)) + if required_buffer_width > buffer_width: + # Release undersized buffers before allocating larger ones for the retry. + d_data_matrix = None + d_indices_matrix = None + return d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width + + d_data_matrix, d_indices_matrix, row_element_counts, required_buffer_width = run_tcrdist_kernel( + buffer_width + ) + + if required_buffer_width > buffer_width: + # The buffer was too small, so retry with the required buffer size. + logging.info( + f"GPU TCRdist tile buffer increased from {buffer_width} to {required_buffer_width}; " + f"retrying the {seqs_mat1_rows} x {seqs_mat2_rows} tile." + ) + buffer_width = required_buffer_width + d_data_matrix, d_indices_matrix, row_element_counts, _ = run_tcrdist_kernel(buffer_width) + + row_element_sum = np.sum(row_element_counts, dtype=np.int64) + + if row_element_sum > np.iinfo(np.int32).max: + raise ValueError( + "There are too many result values to be held by the resulting CSR matrix of the current block. " + f"Current number: {row_element_sum}, maximum number: {np.iinfo(np.int32).max}. " + "Consider choosing a smaller cutoff to resolve this issue." + ) + + indptr = np.zeros(seqs_mat1.shape[0] + 1, dtype=np.int32) + indptr[1:] = np.cumsum(row_element_counts) + d_indptr = cp.asarray(indptr) + + n_elements = indptr[-1] + d_data = cp.zeros(n_elements, dtype=cp.int32) + d_indices = cp.zeros(n_elements, dtype=cp.int32) + + threads_per_block = (1, 256) + blocks_per_grid_x = (d_data_matrix.shape[0] + threads_per_block[0] - 1) // threads_per_block[0] + blocks_per_grid_y = (d_data_matrix.shape[1] + threads_per_block[1] - 1) // threads_per_block[1] + blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y) + + create_csr_kernel( + (blocks_per_grid_x, blocks_per_grid_y), + threads_per_block, + ( + d_data, + d_indices, + d_data_matrix, + d_indices_matrix, + d_indptr, + d_data_matrix.shape[0], + d_data_matrix.shape[1], + d_data.shape[0], + d_indices_matrix.shape[1], + ), + ) + + data = d_data.get() + indices = d_indices.get() + + res = csr_matrix((data, indices, indptr), shape=(seqs_mat1.shape[0], seqs_mat2.shape[0])) + return res, buffer_width + + seqs_mat1_blocks = np.array_split(seqs_mat1, n_row_blocks) + seqs_L1_blocks = np.array_split(seqs_L1, n_row_blocks) + seqs_mat2_blocks = np.array_split(seqs_mat2, n_col_blocks) + seqs_L2_blocks = np.array_split(seqs_L2, n_col_blocks) + + possible_lengths = np.arange(max_seq_len + 1) + max_length_diff = self.cutoff // self.gap_penalty if self.gap_penalty else max_seq_len + length_bounds_blocks = [ + ( + cp.asarray(np.searchsorted(lengths, possible_lengths - max_length_diff, side="left").astype(np.int32)), + cp.asarray(np.searchsorted(lengths, possible_lengths + max_length_diff, side="right").astype(np.int32)), + ) + for lengths in seqs_L2_blocks + ] + + logging.info(f"\nStart GPU calculations for {n_row_blocks} row tiles x {n_col_blocks} column tiles:") + + @nb.njit + def csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elements): + data = np.empty(num_elements, dtype=block_data[0].dtype) + indices = np.empty(num_elements, dtype=block_indices[0].dtype) + indptr = np.zeros(num_rows + 1, dtype=np.int32) + + ptr = 0 + for row in range(num_rows): + for b in range(len(block_indptrs)): + start = block_indptrs[b][row] + end = block_indptrs[b][row + 1] + count = end - start + + for j in range(count): + data[ptr + j] = block_data[b][start + j] + indices[ptr + j] = block_indices[b][start + j] + + ptr += count + indptr[row + 1] = ptr + + return data, indices, indptr + + def csr_union(blocks): + num_rows = blocks[0].shape[0] + num_elements = sum(b.nnz for b in blocks) + + block_data = [b.data for b in blocks] + block_indices = [b.indices for b in blocks] + block_indptrs = [b.indptr for b in blocks] + + data, indices, indptr = csr_union_numba(block_data, block_indices, block_indptrs, num_rows, num_elements) + + shape = blocks[0].shape + result = csr_matrix((data, indices, indptr), shape=shape) + return result + + def skip_col_block(row_block_idx, col_block_idx): + row_start = start_column + seqs_block_starts[row_block_idx] + col_end = seqs2_block_starts[col_block_idx] + seqs_mat2_blocks[col_block_idx].shape[0] + return is_symmetric and col_end <= row_start + + def count_blocks_to_compute(): + n_blocks_to_compute = 0 + for row_block_idx in range(n_row_blocks): + for col_block_idx in range(n_col_blocks): + if not skip_col_block(row_block_idx, col_block_idx): + n_blocks_to_compute += 1 + return n_blocks_to_compute + + def calc_row_block_gpu( + row_block_idx, + seqs_mat1_block, + seqs_L1_block, + seqs_original_indices_block, + buffer_width, + ): + result_blocks = [] + n_calculated_blocks = 0 + + for i in range(0, n_col_blocks): + # Skip calculation of blocks below the diagonal if the result matrix is symmetric. + if skip_col_block(row_block_idx, i): + continue + + result_block, buffer_width = calc_col_block_gpu( + seqs_mat1_block, + seqs_mat2_blocks[i], + seqs_L1_block, + seqs_L2_blocks[i], + length_bounds_blocks[i], + seqs_original_indices_block, + seqs2_original_indices_blocks[i], + buffer_width, + ) + result_blocks.append(result_block) + n_calculated_blocks += 1 + + if not result_blocks: + return ( + csr_matrix((seqs_mat1_block.shape[0], seqs_mat2.shape[0]), dtype=np.int32), + n_calculated_blocks, + buffer_width, + ) + + num_elements = sum(int(block.indptr[-1]) for block in result_blocks) + + if num_elements > np.iinfo(np.int32).max: + raise ValueError( + "The overall number of result values is too high to construct the final CSR matrix by combining " + "the already calculated blocks. " + f"Current number: {num_elements}, maximum number: {np.iinfo(np.int32).max}. " + "Consider choosing a smaller cutoff to resolve this issue." + ) + + result_sparse = csr_union(result_blocks) + result_sparse.sort_indices() + return result_sparse, n_calculated_blocks, buffer_width + + row_blocks = [None] * n_row_blocks + buffer_width = self.gpu_tile_buffer_cols + with tqdm(total=count_blocks_to_compute(), desc="Processing", unit="block") as progress_bar: + for row_block_idx in range(n_row_blocks): + row_blocks[row_block_idx], n_calculated_blocks, buffer_width = calc_row_block_gpu( + row_block_idx, + seqs_mat1_blocks[row_block_idx], + seqs_L1_blocks[row_block_idx], + seqs_original_indices_blocks[row_block_idx], + buffer_width, + ) + progress_bar.update(n_calculated_blocks) + + result_sparse = scipy.sparse.vstack(row_blocks, format="csr") + + row_element_counts_gpu = np.diff(result_sparse.indptr) + result_sparse.sort_indices() + + # Returns the results in a way that fits the current interface, could be improved later + return [result_sparse.data], [result_sparse.indices], row_element_counts_gpu, np.array([None]) + + _metric_mat = _gpu_tcrdist_mat + + class NeedlemanWunschDistanceCalculator(_MetricDistanceCalculator): """Computes pairwise global-alignment distances with linear-gap Needleman-Wunsch. diff --git a/src/scirpy/tests/test_ir_dist.py b/src/scirpy/tests/test_ir_dist.py index 763e941f1..a021e3d24 100644 --- a/src/scirpy/tests/test_ir_dist.py +++ b/src/scirpy/tests/test_ir_dist.py @@ -199,7 +199,8 @@ def test_ir_dist_gpu_hamming(adata_cdr3): @pytest.mark.parametrize("mudata", [False, True], ids=["AnnData", "MuData"]) -def test_ir_dist_tcrdist_tcrblosum_chain_routing(mudata): +@pytest.mark.parametrize("metric", ["tcrdist", pytest.param("gpu_tcrdist", marks=pytest.mark.gpu)]) +def test_ir_dist_tcrdist_tcrblosum_chain_routing(mudata, metric): # `ir_dist` should automatically route VJ to TCRBLOSUM alpha and VDJ to beta. adata = _make_adata( pd.DataFrame( @@ -220,7 +221,7 @@ def test_ir_dist_tcrdist_tcrblosum_chain_routing(mudata): ir.pp.ir_dist( adata, - metric="tcrdist", + metric=metric, sequence="aa", cutoff=20, base_matrix="tcrblosum", diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 1f9b5f95e..edbe87ee9 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -13,6 +13,7 @@ DistanceCalculator, FastAlignmentDistanceCalculator, GPUHammingDistanceCalculator, + GPUTCRdistDistanceCalculator, HammingDistanceCalculator, IdentityDistanceCalculator, LevenshteinDistanceCalculator, @@ -862,8 +863,19 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), ], ) -def test_tcrdist(test_parameters, test_input, expected_result): - tcrdist_calculator = TCRdistDistanceCalculator(**test_parameters) +@pytest.mark.parametrize( + "calculator_class", + [TCRdistDistanceCalculator, pytest.param(GPUTCRdistDistanceCalculator, marks=pytest.mark.gpu)], +) +def test_tcrdist(test_parameters, test_input, expected_result, calculator_class): + test_parameters = test_parameters.copy() + if calculator_class is GPUTCRdistDistanceCalculator: + test_parameters.pop("n_jobs", None) + if not test_parameters.get("fixed_gappos", True): + with pytest.raises(NotImplementedError, match="fixed_gappos=True"): + calculator_class(**test_parameters) + return + tcrdist_calculator = calculator_class(**test_parameters) seq1, seq2 = test_input res = tcrdist_calculator.calc_dist_mat(seq1, seq2) assert isinstance(res, scipy.sparse.csr_matrix) @@ -1019,13 +1031,14 @@ def test_needleman_wunsch(test_parameters, test_input, expected_result): assert np.array_equal(res.todense(), expected_result) -def test_sequence_dist_tcrdist_tcrblosum(): +@pytest.mark.parametrize("metric", ["tcrdist", pytest.param("gpu_tcrdist", marks=pytest.mark.gpu)]) +def test_sequence_dist_tcrdist_tcrblosum(metric): # `sequence_dist` needs an explicit `chain_type` for `tcrdist` with `tcrblosum`; # `ir_dist` handles this automatically. seqs = np.array(["AAACAAAA", "AAARAAAA"]) res = ir.ir_dist.sequence_dist( seqs, - metric="tcrdist", + metric=metric, cutoff=20, n_jobs=1, base_matrix="tcrblosum", @@ -1034,12 +1047,13 @@ def test_sequence_dist_tcrdist_tcrblosum(): npt.assert_array_equal(res.toarray(), np.array([[1, 7], [7, 1]])) -def test_sequence_dist_tcrdist_distance_cap(): +@pytest.mark.parametrize("metric", ["tcrdist", pytest.param("gpu_tcrdist", marks=pytest.mark.gpu)]) +def test_sequence_dist_tcrdist_distance_cap(metric): seqs = np.array(["AAACAAAA", "AAAHAAAA"]) default_tcrblosum = ir.ir_dist.sequence_dist( seqs, - metric="tcrdist", + metric=metric, cutoff=20, n_jobs=1, base_matrix="tcrblosum", @@ -1049,7 +1063,7 @@ def test_sequence_dist_tcrdist_distance_cap(): capped_tcrblosum = ir.ir_dist.sequence_dist( seqs, - metric="tcrdist", + metric=metric, cutoff=20, n_jobs=1, base_matrix="tcrblosum", @@ -1060,7 +1074,7 @@ def test_sequence_dist_tcrdist_distance_cap(): custom_capped_tcrblosum = ir.ir_dist.sequence_dist( seqs, - metric="tcrdist", + metric=metric, cutoff=20, n_jobs=1, base_matrix="tcrblosum", @@ -1071,7 +1085,7 @@ def test_sequence_dist_tcrdist_distance_cap(): uncapped_blosum62 = ir.ir_dist.sequence_dist( np.array(["AAACAAAA", "AAARAAAA"]), - metric="tcrdist", + metric=metric, cutoff=1000, n_jobs=1, base_matrix="blosum62", @@ -1313,6 +1327,95 @@ def test_gpu_hamming_tile_parameter_guards(kwargs, message): GPUHammingDistanceCalculator(**kwargs) +@pytest.mark.gpu +@pytest.mark.parametrize( + "test_parameters,test_input,expected_result", + [ + # Regression test for #626 and #682: sequence lengths exceeding the int8 range. + ( + {"cutoff": 50, "gpu_tile_cols": 3, "gpu_tile_buffer_cols": 50}, + (np.array(["A" * 128, "AAR", "AARR", "ARA"]), np.array(["A" * 128, "ARRR", "ARRR"])), + np.array([[1, 0, 0], [0, 0, 0], [0, 2, 2], [0, 0, 0]]), + ), + # Retain a distance exactly at the cutoff and omit one just above it. + ( + {"cutoff": 125}, + (np.array(["A" * 127]), np.array(["R" * 125 + "AA", "R" * 126 + "A"])), + np.array([[126, 0]]), + ), + # Symmetric calculation split into outer joblib blocks and internal GPU tiles. + ( + {"cutoff": 2, "n_blocks": 2, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 3}, + (np.array(["AAAA", "AAAT", "AATT", "TTTT"]), None), + np.array([[1, 2, 3, 0], [2, 1, 2, 0], [3, 2, 1, 3], [0, 0, 3, 1]]), + ), + # Asymmetric calculation split into outer joblib blocks and internal GPU tiles. + ( + {"cutoff": 2, "n_blocks": 2, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 3}, + (np.array(["AAAA", "AATA", "HHHH", "WWWW"]), np.array(["WWWW", "AAAA", "ATAA"])), + np.array([[0, 1, 2], [0, 2, 3], [0, 0, 0], [1, 0, 0]]), + ), + # Distances above the cutoff and comparisons between sequences of unequal length are omitted. + ( + {"cutoff": 1, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 2}, + (np.array(["AAA", "AAT", "AAAA", "TTT"]), None), + np.array([[1, 2, 0, 0], [2, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]), + ), + # Duplicate sequences are retained as separate zero-distance entries with distance + 1 encoding. + ( + {"cutoff": 0, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 2}, + (np.array(["AAA", "AAA", "AAT"]), None), + np.array([[1, 1, 0], [1, 1, 0], [0, 0, 1]]), + ), + # Outer joblib block counts exceeding the number of sequences are supported. + ( + {"cutoff": 1, "n_blocks": 5, "gpu_tile_rows": 1, "gpu_tile_cols": 1, "gpu_tile_buffer_cols": 2}, + (np.array(["AAA", "AAT"]), None), + np.array([[1, 2], [2, 1]]), + ), + ], +) +def test_gpu_tcrdist(test_parameters, test_input, expected_result): + tcrdist_calculator = GPUTCRdistDistanceCalculator( + dist_weight=1, distance_cap=1, ntrim=0, ctrim=0, gap_penalty=1000, **test_parameters + ) + res = tcrdist_calculator.calc_dist_mat(*test_input) + + assert isinstance(res, scipy.sparse.csr_matrix) + if test_parameters.get("n_blocks", 1) <= len(test_input[0]): + assert res.dtype == np.dtype("int32") + npt.assert_array_equal(res.toarray(), expected_result) + + +@pytest.mark.gpu +def test_gpu_tcrdist_buffer_retry(): + tcrdist_calculator = GPUTCRdistDistanceCalculator(cutoff=0, gpu_tile_cols=2, gpu_tile_buffer_cols=1) + result = tcrdist_calculator.calc_dist_mat(np.array(["AAA", "AAA"])) + + npt.assert_array_equal(result.toarray(), np.ones((2, 2))) + + +@pytest.mark.gpu +@pytest.mark.parametrize("cutoff", [-1, 1.5, np.iinfo(np.int32).max]) +def test_gpu_tcrdist_cutoff_guard(cutoff): + with pytest.raises(ValueError, match="`cutoff` must be an integer"): + GPUTCRdistDistanceCalculator(cutoff=cutoff) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"gpu_tile_rows": 0}, "`gpu_tile_rows` must be >= 1"), + ({"gpu_tile_cols": 0}, "`gpu_tile_cols` must be >= 1"), + ({"gpu_tile_buffer_cols": 0}, "`gpu_tile_buffer_cols` must be >= 1"), + ], +) +def test_gpu_tcrdist_tile_parameter_guards(kwargs, message): + with pytest.raises(ValueError, match=message): + GPUTCRdistDistanceCalculator(**kwargs) + + def test_tcrdist_histogram_not_implemented(): # Change once histogram is implemented for tcrdist with pytest.raises(NotImplementedError, match=None): @@ -1355,3 +1458,66 @@ def test_gpu_hamming_reference(kwargs): assert np.array_equal(res.indices, reference_result.indices) assert np.array_equal(res.indptr, reference_result.indptr) assert np.array_equal(res.todense(), reference_result.todense()) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "kwargs", + [ + {"gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 517, "gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"n_blocks": 3, "gpu_tile_rows": 517, "gpu_tile_cols": 310, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_cols": 222, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 310, "gpu_tile_cols": 517, "gpu_tile_buffer_cols": 500}, + {"gpu_tile_rows": 141, "gpu_tile_cols": 92, "gpu_tile_buffer_cols": 503}, + ], +) +def test_gpu_tcrdist_reference(kwargs): + # test tcrdist distance against reference implementation + from . import TESTDATA + + seqs = np.load(TESTDATA / "tcrdist_test_data/tcrdist_WU3k_seqs.npy") + reference_result = scipy.sparse.load_npz(TESTDATA / "tcrdist_test_data/tcrdist_WU3k_csr_result.npz") + + gpu_tcrdist_calculator = GPUTCRdistDistanceCalculator(cutoff=15, **kwargs) + res = gpu_tcrdist_calculator.calc_dist_mat(seqs, seqs) + + assert np.array_equal(res.data, reference_result.data) + assert np.array_equal(res.indices, reference_result.indices) + assert np.array_equal(res.indptr, reference_result.indptr) + assert np.array_equal(res.todense(), reference_result.todense()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("name", ["dist_weight", "gap_penalty", "ntrim", "ctrim"]) +@pytest.mark.parametrize("value", [-1, 1.5, int(np.iinfo(np.int32).max) + 1]) +def test_gpu_tcrdist_parameter_guards(name, value): + with pytest.raises(ValueError, match=f"`{name}` must be a non-negative 32-bit integer"): + GPUTCRdistDistanceCalculator(**{name: value}) + + +@pytest.mark.gpu +@pytest.mark.parametrize("gap_penalty", [0, 4, 5]) +def test_gpu_tcrdist_length_bounds(gap_penalty): + # Cover both ends of the permitted length range, and the unrestricted case with no gap penalty. + seqs = np.array(["AAAAA", "AAA", "AAAA"]) + calculator = GPUTCRdistDistanceCalculator( + cutoff=4, gap_penalty=gap_penalty, ntrim=0, ctrim=0, gpu_tile_cols=2, gpu_tile_buffer_cols=1 + ) + result = calculator.calc_dist_mat(seqs) + expected = TCRdistDistanceCalculator(cutoff=4, gap_penalty=gap_penalty, ntrim=0, ctrim=0, n_jobs=1).calc_dist_mat( + seqs + ) + npt.assert_array_equal(result.toarray(), expected.toarray()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("ntrim,ctrim", [(0, 0), (3, 2), (10, 10)]) +def test_gpu_tcrdist_short_sequences(ntrim, ctrim): + # Short, unequal sequences exercise the CPU's floor division when locating the fixed gap. + seqs = np.array(["AR", "ARN", "ARND", "ARNDC"]) + seqs2 = np.array(["ACR", "ARNC", "RARND", "ARN"]) + kwargs = {"cutoff": 100, "ntrim": ntrim, "ctrim": ctrim} + result = GPUTCRdistDistanceCalculator(gpu_tile_rows=2, gpu_tile_cols=2, **kwargs).calc_dist_mat(seqs, seqs2) + expected = TCRdistDistanceCalculator(n_jobs=1, **kwargs).calc_dist_mat(seqs, seqs2) + npt.assert_array_equal(result.toarray(), expected.toarray()) From 3285c8f84335669d4c4ffc05842918bd88e0f690 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 14 Sep 2026 13:42:06 +0200 Subject: [PATCH 39/54] Strengthen GPU TCRdist buffer retry and length boundary tests --- src/scirpy/tests/test_ir_dist_metrics.py | 26 ++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index edbe87ee9..59a3824e8 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -1355,7 +1355,7 @@ def test_gpu_hamming_tile_parameter_guards(kwargs, message): (np.array(["AAAA", "AATA", "HHHH", "WWWW"]), np.array(["WWWW", "AAAA", "ATAA"])), np.array([[0, 1, 2], [0, 2, 3], [0, 0, 0], [1, 0, 0]]), ), - # Distances above the cutoff and comparisons between sequences of unequal length are omitted. + # Distances above the cutoff are omitted. ( {"cutoff": 1, "gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 2}, (np.array(["AAA", "AAT", "AAAA", "TTT"]), None), @@ -1376,6 +1376,8 @@ def test_gpu_hamming_tile_parameter_guards(kwargs, message): ], ) def test_gpu_tcrdist(test_parameters, test_input, expected_result): + # Use unit substitution costs, no trimming, and a high gap penalty to keep expected distances simple + # while testing tile splitting and result assembly. tcrdist_calculator = GPUTCRdistDistanceCalculator( dist_weight=1, distance_cap=1, ntrim=0, ctrim=0, gap_penalty=1000, **test_parameters ) @@ -1389,10 +1391,13 @@ def test_gpu_tcrdist(test_parameters, test_input, expected_result): @pytest.mark.gpu def test_gpu_tcrdist_buffer_retry(): - tcrdist_calculator = GPUTCRdistDistanceCalculator(cutoff=0, gpu_tile_cols=2, gpu_tile_buffer_cols=1) - result = tcrdist_calculator.calc_dist_mat(np.array(["AAA", "AAA"])) + # Multiple retained distances per row force a retry; column lengths are deliberately unsorted. + tcrdist_calculator = GPUTCRdistDistanceCalculator( + cutoff=12, ntrim=0, ctrim=0, gpu_tile_cols=4, gpu_tile_buffer_cols=1 + ) + result = tcrdist_calculator.calc_dist_mat(np.array(["AAAA", "AAA"]), np.array(["AAAR", "AAAAA", "AAA", "AAAA"])) - npt.assert_array_equal(result.toarray(), np.ones((2, 2))) + npt.assert_array_equal(result.toarray(), np.array([[13, 5, 5, 1], [0, 9, 1, 5]])) @pytest.mark.gpu @@ -1498,15 +1503,20 @@ def test_gpu_tcrdist_parameter_guards(name, value): @pytest.mark.gpu @pytest.mark.parametrize("gap_penalty", [0, 4, 5]) -def test_gpu_tcrdist_length_bounds(gap_penalty): - # Cover both ends of the permitted length range, and the unrestricted case with no gap penalty. +@pytest.mark.parametrize( + "seqs2", + [None, np.array(["AAAAAA", "AA", "AAAA", "AAA", "AAAAA"])], + ids=["symmetric", "rectangular"], +) +def test_gpu_tcrdist_length_bounds(gap_penalty, seqs2): + # Cover both length bounds without relying on mirroring, and the unrestricted case with no gap penalty. seqs = np.array(["AAAAA", "AAA", "AAAA"]) calculator = GPUTCRdistDistanceCalculator( cutoff=4, gap_penalty=gap_penalty, ntrim=0, ctrim=0, gpu_tile_cols=2, gpu_tile_buffer_cols=1 ) - result = calculator.calc_dist_mat(seqs) + result = calculator.calc_dist_mat(seqs, seqs2) expected = TCRdistDistanceCalculator(cutoff=4, gap_penalty=gap_penalty, ntrim=0, ctrim=0, n_jobs=1).calc_dist_mat( - seqs + seqs, seqs2 ) npt.assert_array_equal(result.toarray(), expected.toarray()) From 329408c4aac32f11a60e5c687d0833ec44bfb2b9 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 14 Sep 2026 13:43:25 +0200 Subject: [PATCH 40/54] Add changelog entry for GPU TCRdist --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c8d731eb..b886d63bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning][]. ### Features + - Add `metric="gpu_tcrdist"` for GPU-accelerated TCRdist with row and column tiling, adaptive result buffers, + and support for BLOSUM62 and TCRBLOSUM matrices. + - Add a linear-gap Needleman-Wunsch distance metric for CDR3 amino acid sequences via `metric="needleman_wunsch"`. The `alignment` and `fastalignment` metrics are now deprecated. When `gap_open == gap_extend` (which has always been the default), use `needleman_wunsch` instead. From 6a595f9e9405171012651fb4595171b12387d237 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 14 Sep 2026 13:45:56 +0200 Subject: [PATCH 41/54] Document GPU TCRdist support in the large datasets tutorial --- docs/tutorials/large-datasets.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/tutorials/large-datasets.md b/docs/tutorials/large-datasets.md index 7fa0c3e77..8e98fb07b 100644 --- a/docs/tutorials/large-datasets.md +++ b/docs/tutorials/large-datasets.md @@ -49,9 +49,10 @@ with joblib.parallel_config(backend="dask", n_jobs=200, verbose=10): ) ``` -## Using GPU acceleration for hamming distance +## Using GPU acceleration for Hamming and TCRdist -The Hamming distance metric supports GPU acceleration via [cupy](https://cupy.dev/). +The Hamming and TCRdist distance metrics support GPU acceleration via [cupy](https://cupy.dev/). +The examples below use `gpu_hamming`; `gpu_tcrdist` supports the same GPU tile parameters. First, install the optional `cupy` dependency: From f64582df152917d276159a7e04c19a3a89ae72f2 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 11:28:54 +0200 Subject: [PATCH 42/54] Increase default TCRdist gap penalty to 12 and warn about the change --- CHANGELOG.md | 7 +++ src/scirpy/ir_dist/metrics.py | 21 +++++++-- src/scirpy/tests/test_ir_dist_metrics.py | 59 ++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89800ddba..e3f0f1e01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,13 @@ and this project adheres to [Semantic Versioning][]. ## Unreleased +### Changed + + - Increase the default TCRdist `gap_penalty` from 4 to 12 to match the authors' later implementation. + This changes distances between sequences of different lengths. Omitting `gap_penalty` now emits a + `UserWarning`; explicitly set `gap_penalty=12` to acknowledge the new default, or `gap_penalty=4` + to retain the previous behavior. + ### Fixes - Avoid deprecated AnnData `obsm_keys()` calls when plotting clonotype networks or exporting them to igraph. diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 71c926062..d70e8aa59 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1,5 +1,6 @@ import abc import itertools +import warnings from collections.abc import Sequence from typing import Literal @@ -1286,15 +1287,17 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): The code of this class is heavily based on `pwseqdist `_. Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. - Using default weight, gap penalty, ntrim and ctrim is equivalent to the - original distance published in :cite:`TCRdist`. + This implements the CDR3 component of :cite:`TCRdist`. The default gap penalty + is 12, following the authors' later implementation; the original paper used 8. Parameters ---------- dist_weight: Weight applied to the mismatch distances before summing with the gap penalties gap_penalty: - Distance penalty for the difference in the length of the two sequences + Distance penalty per amino acid of length difference. Defaults to 12. + If omitted, emits a warning about the change from the previous default of 4. + Set explicitly to 12 to acknowledge the new default, or to 4 to retain the previous behavior. ntrim/ctrim: Positions trimmed off the N-terminus (0) and C-terminus (L-1) ends of the peptide sequence. These symbols will be ignored in the distance calculation. @@ -1331,7 +1334,7 @@ def __init__( cutoff: int = 20, *, dist_weight: int = 3, - gap_penalty: int = 4, + gap_penalty: int | Literal["default"] = "default", ntrim: int = 3, ctrim: int = 2, fixed_gappos: bool = True, @@ -1342,6 +1345,16 @@ def __init__( distance_cap: int | None | Literal["default"] = "default", chain_type: Literal["VJ", "VDJ"] | None = None, ): + if gap_penalty == "default": + warnings.warn( + "The default value of `gap_penalty` has changed from 4 to 12. " + "This changes TCRdist distances between sequences of different lengths. " + "Set `gap_penalty=12` explicitly to acknowledge the new default, " + "or `gap_penalty=4` to retain the previous behavior.", + UserWarning, + stacklevel=2, + ) + gap_penalty = 12 self.dist_weight = dist_weight self.gap_penalty = gap_penalty self.ntrim = ntrim diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 1f9b5f95e..14196eead 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -1,3 +1,4 @@ +import warnings from functools import partial from unittest.mock import patch @@ -860,6 +861,45 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), np.array([[1, 34, 28], [34, 1, 19], [28, 19, 1]]), ), + # Test combined mismatch and gap costs, including distances above the cutoff. + ( + {"gap_penalty": 12, "n_jobs": 1, "cutoff": 24}, + ( + np.array(["AAAAAAAA", "AAARAAAA", "AAAAAAAAAA"]), + np.array(["AAARAAAAA", "AAAAAAAA", "AAAAAAAAAAA", "AAAAAAAAA"]), + ), + np.array([[25, 1, 0, 13], [13, 13, 0, 25], [25, 25, 13, 13]]), + ), + # Test that two extra residues exceed the default cutoff with gap_penalty=12. + ( + {"gap_penalty": 12, "n_jobs": 1}, + (np.array(["AAAAAAAAAA", "AAAAAAAA", "AAAAAAAAA", "AAAAAAAAAAA"]), None), + np.array([[1, 0, 13, 13], [0, 1, 13, 0], [13, 13, 1, 0], [13, 0, 0, 1]]), + ), + # Test that two extra residues are retained exactly at cutoff=24. + ( + {"gap_penalty": 12, "n_jobs": 1, "cutoff": 24}, + (np.array(["AAAAAAAAAA", "AAAAAAAA", "AAAAAAAAA", "AAAAAAAAAAA"]), None), + np.array([[1, 25, 13, 13], [25, 1, 13, 0], [13, 13, 1, 25], [13, 0, 25, 1]]), + ), + # Test that a single gap or mismatch is retained exactly at cutoff=12. + ( + {"gap_penalty": 12, "n_jobs": 1, "cutoff": 12}, + ( + np.array(["AAAAAAAA", "AAARAAAA", "AAAAAAAAA"]), + np.array(["AAAAAAAAA", "AAARAAAAA", "AAAAAAAA", "AAARAAAA"]), + ), + np.array([[13, 0, 1, 13], [0, 13, 13, 1], [1, 13, 13, 0]]), + ), + # Test that cutoff=11 excludes a single gap or mismatch but retains identical sequences. + ( + {"gap_penalty": 12, "n_jobs": 1, "cutoff": 11}, + ( + np.array(["AAAAAAAA", "AAARAAAA", "AAAAAAAAA"]), + np.array(["AAAAAAAAA", "AAARAAAAA", "AAAAAAAA", "AAARAAAA"]), + ), + np.array([[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]]), + ), ], ) def test_tcrdist(test_parameters, test_input, expected_result): @@ -871,6 +911,25 @@ def test_tcrdist(test_parameters, test_input, expected_result): assert np.array_equal(res.todense(), expected_result) +@pytest.mark.parametrize("gap_penalty", [None, 4, 12]) +def test_tcrdist_gap_penalty_default(gap_penalty): + kwargs = {} if gap_penalty is None else {"gap_penalty": gap_penalty} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + calculator = TCRdistDistanceCalculator(n_jobs=1, **kwargs) + if gap_penalty is None: + assert len(caught) == 1 + assert caught[0].category is UserWarning + assert "has changed from 4 to 12" in str(caught[0].message) + else: + assert not caught + expected_penalty = 12 if gap_penalty is None else gap_penalty + assert calculator.gap_penalty == expected_penalty + # One extra residue contributes only the gap penalty; stored distances are offset by 1. + result = calculator.calc_dist_mat(np.array(["AAAAAAAA"]), np.array(["AAAAAAAAA"])) + assert result[0, 0] == expected_penalty + 1 + + @pytest.mark.parametrize( "test_parameters,test_input,expected_result", [ From e1c5a4754efd58bc8a89fa521a3f4281cd8259bb Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 13:20:51 +0200 Subject: [PATCH 43/54] Exclude fully trimmed sequences from TCRdist comparisons --- CHANGELOG.md | 4 ++ src/scirpy/ir_dist/_util.py | 2 +- src/scirpy/ir_dist/metrics.py | 20 ++++++- src/scirpy/tests/test_clonotypes.py | 31 ++++++++++ src/scirpy/tests/test_ir_dist_metrics.py | 74 ++++++++++++++++++++++-- src/scirpy/tests/test_ir_dist_util.py | 10 ++++ 6 files changed, 135 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3f0f1e01..61656f3a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning][]. ### Changed + - Exclude TCRdist comparisons involving sequences with no positions remaining after trimming, + including self-comparisons. A warning reports the number of affected input sequences; matrix + shape and sequence order are preserved. + - Increase the default TCRdist `gap_penalty` from 4 to 12 to match the authors' later implementation. This changes distances between sequences of different lengths. Omitting `gap_penalty` now emits a `UserWarning`; explicitly set `gap_penalty=12` to acknowledge the new default, or `gap_penalty=4` diff --git a/src/scirpy/ir_dist/_util.py b/src/scirpy/ir_dist/_util.py index d82a9ffd7..d9b373cda 100644 --- a/src/scirpy/ir_dist/_util.py +++ b/src/scirpy/ir_dist/_util.py @@ -297,7 +297,7 @@ def lookup( distance_matrix = self.distance_matrices[distance_matrix_name] - if np.max(distance_matrix.data) > np.iinfo(np.uint8).max: + if np.max(distance_matrix.data, initial=0) > np.iinfo(np.uint8).max: raise OverflowError( "The data values in the distance scipy.sparse.csr_matrix exceed the maximum value for uint8 (255)" ) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index d70e8aa59..ea6b9d1fd 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1300,7 +1300,9 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): Set explicitly to 12 to acknowledge the new default, or to 4 to retain the previous behavior. ntrim/ctrim: Positions trimmed off the N-terminus (0) and C-terminus (L-1) ends of the peptide sequence. These symbols will be ignored - in the distance calculation. + in the distance calculation. Sequences with length <= ntrim + ctrim have no remaining + positions and trigger a warning. All their comparisons, including self-comparisons, are + excluded; their rows and columns remain empty without changing the matrix shape. fixed_gappos: If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6). If False, find the "optimal" position for inserting the gaps to make up the difference in length @@ -1407,6 +1409,18 @@ def __init__( super().__init__(n_jobs=n_jobs, n_blocks=n_blocks, histogram=histogram) + def _validate_seqs(self, seqs: Sequence[str], seqs2: Sequence[str]) -> None: + sequences = seqs if seqs2 is seqs else itertools.chain(seqs, seqs2) + n_excluded = sum(len(seq) <= self.ntrim + self.ctrim for seq in sequences) + if n_excluded: + warnings.warn( + f"{n_excluded} input sequences have no positions remaining after trimming " + f"(ntrim={self.ntrim}, ctrim={self.ctrim}). All comparisons involving these sequences " + "are excluded, including self-comparisons. Their matrix rows and columns remain empty.", + UserWarning, + stacklevel=2, + ) + def _tcrdist_mat( self, *, @@ -1498,10 +1512,14 @@ def _nb_tcrdist_mat(): thread_id = nb.get_thread_id() row_end_index = 0 seq1_len = seqs_L1[row_index] + if seq1_len <= ntrim + ctrim: + continue for col_index in range(start_column + row_index * is_symmetric, num_cols): distance = 1 seq2_len = seqs_L2[col_index] + if seq2_len <= ntrim + ctrim: + continue if seq1_len == seq2_len: for i in range(ntrim, seq1_len - ctrim): diff --git a/src/scirpy/tests/test_clonotypes.py b/src/scirpy/tests/test_clonotypes.py index e7c490672..6fa401ec5 100644 --- a/src/scirpy/tests/test_clonotypes.py +++ b/src/scirpy/tests/test_clonotypes.py @@ -68,6 +68,37 @@ def test_define_clonotypes_diagonal_connectivities(adata_define_clonotype_cluste npt.assert_equal(clonotype_size, np.array([1, 1, 1, 1])) +@pytest.mark.parametrize("receptor_arms", ["VJ", "VDJ", "all", "any"]) +@pytest.mark.parametrize("dual_ir", ["primary_only", "all", "any"]) +@pytest.mark.parametrize("partitions", ["connected", "leiden", "fastgreedy"]) +@pytest.mark.parametrize("all_trimmed", [False, True]) +def test_clonotype_clusters_fully_trimmed(receptor_arms, dual_ir, partitions, all_trimmed): + from .util import _make_adata + + sequences = ["AAA", "AAAAA", "AAAA", "AA"] if all_trimmed else ["AAA", "AAAAA", "AAAAAA", "AAAAAAA"] + adata = _make_adata(pd.DataFrame({"IR_VJ_1_junction_aa": sequences, "IR_VDJ_1_junction_aa": sequences})) + with pytest.warns(UserWarning, match="no positions remaining"): + ir.pp.ir_dist(adata, metric="tcrdist", sequence="aa", gap_penalty=12, n_jobs=1) + ir.tl.define_clonotype_clusters( + adata, + metric="tcrdist", + sequence="aa", + receptor_arms=receptor_arms, + dual_ir=dual_ir, + partitions=partitions, + within_group=None, + ) + labels = adata.obs["cc_aa_tcrdist"] + assert labels.notna().all() + assert labels.iloc[0] != labels.iloc[1] + if all_trimmed: + assert labels.nunique() == 4 + assert adata.uns["cc_aa_tcrdist"]["distances"].nnz == 0 + else: + assert labels.iloc[2] == labels.iloc[3] + assert labels.nunique() == 3 + + def test_clonotypes_end_to_end1(adata_define_clonotypes): """Test that default parameters of define_clonotypes yields clonotypes based on nt-identity. diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 14196eead..74cab284d 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -634,7 +634,7 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), np.array([[1, 25, 33], [25, 1, 33], [33, 33, 1]]), ), - # test with high ntrim - trimmimg with ntrim is only possible until the beginning of the gap for sequences of unequal length + # Test that excessive N-terminal trimming excludes all comparisons. ( { "dist_weight": 3, @@ -649,7 +649,7 @@ def test_sequence_dist_all_metrics(metric, n_jobs): np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), ), - np.array([[1, 1, 9], [1, 1, 9], [9, 9, 1]]), + np.zeros((3, 3)), ), # test with ctrim = 0 and cutoff set high to neglect it ( @@ -668,7 +668,7 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), np.array([[1, 25, 21], [25, 1, 21], [21, 21, 1]]), ), - # test with high ctrim - trimmimg with ctrim is only possible until the end of the gap for sequences of unequal length + # Test that excessive C-terminal trimming excludes all comparisons. ( { "dist_weight": 3, @@ -683,7 +683,7 @@ def test_sequence_dist_all_metrics(metric, n_jobs): np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), ), - np.array([[1, 1, 21], [1, 1, 21], [21, 21, 1]]), + np.zeros((3, 3)), ), # test with fixed_gappos = False and a high cutoff to neglect it # AAAAA added at the beginning of the usual sequences to make to difference of min_gappos and max_gappos more significant @@ -911,6 +911,72 @@ def test_tcrdist(test_parameters, test_input, expected_result): assert np.array_equal(res.todense(), expected_result) +@pytest.mark.parametrize("fixed_gappos", [True, False]) +@pytest.mark.parametrize( + "seqs,seqs2,ntrim,ctrim,n_excluded,expected", + [ + # Invalid lengths below and at the trimming boundary keep empty rows and columns. + ( + ["AAA", "AAAAAA", "", "AAAAA", "AAAAAAA"], + None, + 3, + 2, + 3, + [[0, 0, 0, 0, 0], [0, 1, 0, 0, 13], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 13, 0, 0, 1]], + ), + # Both input lists contribute to the warning count in rectangular comparisons. + (["AAAAAA", "AAA"], ["AAAAA", "AAAAAAA", ""], 5, 0, 3, [[0, 13, 0], [0, 0, 0]]), + # Fully trimmed inputs have no self-matches either. + (["AAA", "AAAAA"], None, 0, 5, 2, [[0, 0], [0, 0]]), + # Without trimming, only empty sequences are excluded. + (["", "AA"], None, 0, 0, 1, [[0, 0], [0, 1]]), + # Mixed valid and fully trimmed sequences retain only valid matches within the cutoff. + ( + ["AAA", "AAAAAAAA", "AAARAAAA", "AAAAAAAAA", "AAAAA", "AAAAAAAAAA"], + None, + 3, + 2, + 2, + [ + [0, 0, 0, 0, 0, 0], + [0, 1, 13, 13, 0, 0], + [0, 13, 1, 0, 0, 0], + [0, 13, 0, 1, 0, 13], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 13, 0, 1], + ], + ), + # Rectangular inputs combine custom trimming, mismatches, gaps, and excluded rows/columns. + ( + ["AA", "AAAAAA", "AARAAA", "AAAAAAA"], + ["AAAAAAAA", "", "AAAAAA", "AAA", "AAAAAAA"], + 2, + 1, + 3, + [[0, 0, 0, 0, 0], [0, 0, 1, 0, 13], [0, 0, 13, 0, 0], [13, 0, 13, 0, 1]], + ), + # No matches remain when all row sequences are fully trimmed, even with valid columns. + ( + ["", "A", "AAAA", "AAAAAA"], + ["AAAAAAA", "AAAAAA", "AAAAAAAA", "AA", "AAAAAAA"], + 4, + 2, + 6, + [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]], + ), + ], +) +def test_tcrdist_fully_trimmed(seqs, seqs2, ntrim, ctrim, n_excluded, expected, fixed_gappos): + calculator = TCRdistDistanceCalculator( + gap_penalty=12, ntrim=ntrim, ctrim=ctrim, fixed_gappos=fixed_gappos, n_jobs=1, n_blocks=2 + ) + with pytest.warns(UserWarning, match=f"{n_excluded} input sequences have no positions remaining") as caught: + result = calculator.calc_dist_mat(seqs, seqs2) + assert len(caught) == 1 + npt.assert_array_equal(result.toarray(), expected) + assert result.nnz == np.count_nonzero(expected) + + @pytest.mark.parametrize("gap_penalty", [None, 4, 12]) def test_tcrdist_gap_penalty_default(gap_penalty): kwargs = {} if gap_penalty is None else {"gap_penalty": gap_penalty} diff --git a/src/scirpy/tests/test_ir_dist_util.py b/src/scirpy/tests/test_ir_dist_util.py index 3c8f605da..a0f4063c9 100644 --- a/src/scirpy/tests/test_ir_dist_util.py +++ b/src/scirpy/tests/test_ir_dist_util.py @@ -337,3 +337,13 @@ def test_dlnf_lookup_with_different_forward_and_reverse_tables_rect(dlnf_with_lo == list(dlnf_with_lookup.lookup(6, "VDJ_test", "VJ_test").todense().A1) == [0, 0, 0, 0, 0] ) + + +@pytest.mark.parametrize("dlnf_with_lookup", ["dlnf_square", "dlnf_rectangle"], indirect=True) +def test_dlnf_lookup_empty_distance_matrix(dlnf_with_lookup): + matrix = dlnf_with_lookup.distance_matrices["test"] + matrix.data[:] = 0 + matrix.eliminate_zeros() + result = dlnf_with_lookup.lookup(np.array([0, 1]), "VJ_test") + assert result.shape == (2, dlnf_with_lookup.lookups["VJ_test"][2].size) + assert result.nnz == 0 From 24efc9f980fca26d6f6c3d5944cda90d15b48b77 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 14:38:29 +0200 Subject: [PATCH 44/54] Respect custom trimming when positioning fixed TCRdist gaps --- CHANGELOG.md | 4 ++++ src/scirpy/ir_dist/metrics.py | 7 ++++-- src/scirpy/tests/test_ir_dist_metrics.py | 29 +++++++++++++++++++++++- 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 61656f3a2..2ad4ec202 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning][]. ### Changed + - Make the fixed TCRdist gap position depend on `ntrim` and `ctrim` so it stays within the + remaining region. Default trimming retains the previous behavior; custom trimming can change + distances. The variable gap-position search (`fixed_gappos=False`) is unchanged. + - Exclude TCRdist comparisons involving sequences with no positions remaining after trimming, including self-comparisons. A warning reports the number of affected input sequences; matrix shape and sequence order are preserved. diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index ea6b9d1fd..4331362cb 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1304,7 +1304,8 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): positions and trigger a warning. All their comparisons, including self-comparisons, are excluded; their rows and columns remain empty without changing the matrix shape. fixed_gappos: - If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6). + If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6) - + the gap position falls within the untrimmed region. If False, find the "optimal" position for inserting the gaps to make up the difference in length cutoff: Will eliminate distances > cutoff to make efficient @@ -1529,9 +1530,11 @@ def _nb_tcrdist_mat(): short_len = min(seq1_len, seq2_len) len_diff = abs(seq1_len - seq2_len) if fixed_gappos: - min_gappos = min(6, 3 + (short_len - 5) // 2) + gappos_limit = max(6, ntrim) + min_gappos = min(gappos_limit, ntrim + (short_len - ntrim - ctrim) // 2) max_gappos = min_gappos else: + # TODO: Adapt the dynamic gap-position search bounds to custom ntrim and ctrim values. min_gappos = 5 max_gappos = short_len - 1 - 4 while min_gappos > max_gappos: diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 74cab284d..a44ee698b 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -666,7 +666,7 @@ def test_sequence_dist_all_metrics(metric, n_jobs): np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), np.array(["AAAAAAAAAA", "AAAARRAAAA", "AANDAAAA"]), ), - np.array([[1, 25, 21], [25, 1, 21], [21, 21, 1]]), + np.array([[1, 25, 21], [25, 1, 33], [21, 33, 1]]), ), # Test that excessive C-terminal trimming excludes all comparisons. ( @@ -900,6 +900,33 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), np.array([[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]]), ), + # Fixed gaps respect N-terminal trimming beyond the original maximum gap position. + ( + {"gap_penalty": 12, "ntrim": 8, "ctrim": 2, "cutoff": 30, "n_jobs": 1}, + ( + np.array(["AAAAAAAAAAAA", "AAAAAAARAAAA", "AAAAAAAARAAA"]), + np.array(["AAAAAAAAAAAAA", "AAAAAAAARAAAA", "AAAAAAAAARAAA"]), + ), + np.array([[13, 13, 25], [13, 13, 25], [25, 25, 13]]), + ), + # Fixed gaps respect C-terminal trimming that leaves just one position to compare. + ( + {"gap_penalty": 12, "ntrim": 3, "ctrim": 8, "cutoff": 30, "n_jobs": 1}, + ( + np.array(["AAAAAAAAAAAA", "AAAARAAAAAAA", "AAARAAAAAAAA"]), + np.array(["AAAAAAAAAAAAA", "AAAAARAAAAAAA", "AAAARAAAAAAAA"]), + ), + np.array([[13, 13, 25], [13, 13, 25], [25, 25, 13]]), + ), + # Custom trimming moves the fixed gap within the remaining region of short sequences. + ( + {"gap_penalty": 12, "ntrim": 3, "ctrim": 0, "cutoff": 30, "n_jobs": 1}, + ( + np.array(["AAAAAAAA", "AAAARAAA"]), + np.array(["AAAAAAAAA", "AAAARAAAA", "AAAAARAAA"]), + ), + np.array([[13, 25, 13], [25, 13, 25]]), + ), ], ) def test_tcrdist(test_parameters, test_input, expected_result): From b24fab95bbe6c2accb8fcf0e6d3aa2f7f8ff81c2 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 15:00:21 +0200 Subject: [PATCH 45/54] Increase default TCRdist cutoff to 24 and consolidate default-change warnings --- CHANGELOG.md | 6 ++++ src/scirpy/ir_dist/__init__.py | 2 +- src/scirpy/ir_dist/metrics.py | 28 ++++++++++++----- src/scirpy/tests/test_ir_dist_metrics.py | 40 +++++++++++++++++++++--- 4 files changed, 63 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ad4ec202..3325ac8c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,12 @@ and this project adheres to [Semantic Versioning][]. ### Changed + - Increase the default TCRdist `cutoff` from 20 to 24, allowing two gap positions with the new + default gap penalty of 12. Omitting `cutoff` emits a `UserWarning`, combined with the gap-penalty + warning when both defaults are used. Set `cutoff=24` explicitly to acknowledge the new default, + or `cutoff=20` to retain the previous cutoff. To retain both previous parameter values, also set + `gap_penalty=4`. The increased cutoff can add sequence matches and change clonotype clusters. + - Make the fixed TCRdist gap position depend on `ntrim` and `ctrim` so it stays within the remaining region. Default trimming retains the previous behavior; custom trimming can change distances. The variable gap-position search (`fixed_gappos=False`) is unchanged. diff --git a/src/scirpy/ir_dist/__init__.py b/src/scirpy/ir_dist/__init__.py index 8104795d1..823232492 100644 --- a/src/scirpy/ir_dist/__init__.py +++ b/src/scirpy/ir_dist/__init__.py @@ -84,7 +84,7 @@ def IrNeighbors(*args, **kwargs): matrix. A sensible cutoff depends on the distance metric, you can find information in the corresponding docs. If set to `None`, the cutoff will be `10` for the `alignment`, `fastalignment`, and `needleman_wunsch` metric, - and `2` for `levenshtein` and `hamming`. + `24` for `tcrdist`, and `2` for `levenshtein` and `hamming`. For the identity metric, the cutoff is ignored and always set to `0`. """ diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 4331362cb..da0ba0c10 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1308,8 +1308,9 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): the gap position falls within the untrimmed region. If False, find the "optimal" position for inserting the gaps to make up the difference in length cutoff: - Will eliminate distances > cutoff to make efficient - use of sparse matrices. + Will eliminate distances > cutoff to make efficient use of sparse matrices. Defaults to 24. + If omitted, emits a warning about the change from the previous default of 20. + Set explicitly to 24 to acknowledge the new default, or to 20 to retain the previous cutoff. n_jobs: Number of numba parallel threads to use for the pairwise distance calculation n_blocks: @@ -1334,10 +1335,10 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): def __init__( self, - cutoff: int = 20, + cutoff: int | Literal["default"] = "default", # resolves to 24 *, dist_weight: int = 3, - gap_penalty: int | Literal["default"] = "default", + gap_penalty: int | Literal["default"] = "default", # resolves to 12 ntrim: int = 3, ctrim: int = 2, fixed_gappos: bool = True, @@ -1348,16 +1349,27 @@ def __init__( distance_cap: int | None | Literal["default"] = "default", chain_type: Literal["VJ", "VDJ"] | None = None, ): + default_changes = [] + if cutoff == "default": + default_changes.append( + "The default value of `cutoff` has changed from 20 to 24. " + "Set `cutoff=24` explicitly to acknowledge the new default, " + "or `cutoff=20` to retain the previous cutoff." + ) + cutoff = 24 if gap_penalty == "default": - warnings.warn( + default_changes.append( "The default value of `gap_penalty` has changed from 4 to 12. " - "This changes TCRdist distances between sequences of different lengths. " "Set `gap_penalty=12` explicitly to acknowledge the new default, " - "or `gap_penalty=4` to retain the previous behavior.", + "or `gap_penalty=4` to retain the previous gap penalty." + ) + gap_penalty = 12 + if default_changes: + warnings.warn( + " ".join(default_changes) + " These changes may affect distance results and clonotype clusters.", UserWarning, stacklevel=2, ) - gap_penalty = 12 self.dist_weight = dist_weight self.gap_penalty = gap_penalty self.ntrim = ntrim diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index a44ee698b..abd90b407 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -870,9 +870,9 @@ def test_sequence_dist_all_metrics(metric, n_jobs): ), np.array([[25, 1, 0, 13], [13, 13, 0, 25], [25, 25, 13, 13]]), ), - # Test that two extra residues exceed the default cutoff with gap_penalty=12. + # Test that two extra residues exceed the previous cutoff with gap_penalty=12. ( - {"gap_penalty": 12, "n_jobs": 1}, + {"gap_penalty": 12, "n_jobs": 1, "cutoff": 20}, (np.array(["AAAAAAAAAA", "AAAAAAAA", "AAAAAAAAA", "AAAAAAAAAAA"]), None), np.array([[1, 0, 13, 13], [0, 1, 13, 0], [13, 13, 1, 0], [13, 0, 0, 1]]), ), @@ -995,7 +995,7 @@ def test_tcrdist(test_parameters, test_input, expected_result): ) def test_tcrdist_fully_trimmed(seqs, seqs2, ntrim, ctrim, n_excluded, expected, fixed_gappos): calculator = TCRdistDistanceCalculator( - gap_penalty=12, ntrim=ntrim, ctrim=ctrim, fixed_gappos=fixed_gappos, n_jobs=1, n_blocks=2 + cutoff=20, gap_penalty=12, ntrim=ntrim, ctrim=ctrim, fixed_gappos=fixed_gappos, n_jobs=1, n_blocks=2 ) with pytest.warns(UserWarning, match=f"{n_excluded} input sequences have no positions remaining") as caught: result = calculator.calc_dist_mat(seqs, seqs2) @@ -1004,12 +1004,44 @@ def test_tcrdist_fully_trimmed(seqs, seqs2, ntrim, ctrim, n_excluded, expected, assert result.nnz == np.count_nonzero(expected) +@pytest.mark.parametrize("cutoff", [None, 20, 24]) +@pytest.mark.parametrize("gap_penalty", [None, 4, 12]) +def test_tcrdist_default_warning(cutoff, gap_penalty): + kwargs = {name: value for name, value in (("cutoff", cutoff), ("gap_penalty", gap_penalty)) if value is not None} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + calculator = TCRdistDistanceCalculator(n_jobs=1, **kwargs) + assert calculator.cutoff == (24 if cutoff is None else cutoff) + assert calculator.gap_penalty == (12 if gap_penalty is None else gap_penalty) + if cutoff is None or gap_penalty is None: + assert len(caught) == 1 + assert caught[0].category is UserWarning + message = str(caught[0].message) + assert ("`cutoff` has changed from 20 to 24" in message) == (cutoff is None) + assert ("`gap_penalty` has changed from 4 to 12" in message) == (gap_penalty is None) + else: + assert not caught + + +@pytest.mark.parametrize("cutoff,expected", [(None, 25), (20, 0), (24, 25)]) +def test_sequence_dist_tcrdist_default_cutoff(cutoff, expected): + kwargs = {} if cutoff is None else {"cutoff": cutoff} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", UserWarning) + result = ir.ir_dist.sequence_dist( + ["AAAAAAAA"], ["AAAAAAAAAA"], metric="tcrdist", gap_penalty=12, n_jobs=1, **kwargs + ) + migration_warnings = [w for w in caught if "default value of `cutoff`" in str(w.message)] + assert len(migration_warnings) == (cutoff is None) + assert result[0, 0] == expected + + @pytest.mark.parametrize("gap_penalty", [None, 4, 12]) def test_tcrdist_gap_penalty_default(gap_penalty): kwargs = {} if gap_penalty is None else {"gap_penalty": gap_penalty} with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - calculator = TCRdistDistanceCalculator(n_jobs=1, **kwargs) + calculator = TCRdistDistanceCalculator(cutoff=24, n_jobs=1, **kwargs) if gap_penalty is None: assert len(caught) == 1 assert caught[0].category is UserWarning From 802f5ac0826e9f7bf490fe2078c28867c68f4069 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 15:08:26 +0200 Subject: [PATCH 46/54] Add TODO for removing temporary TCRdist default-change warnings --- src/scirpy/ir_dist/metrics.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index da0ba0c10..6356e3eb1 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1349,6 +1349,8 @@ def __init__( distance_cap: int | None | Literal["default"] = "default", chain_type: Literal["VJ", "VDJ"] | None = None, ): + # TODO: Remove the migration warning after a transition period and set + # cutoff=24 and gap_penalty=12 directly in the signature. default_changes = [] if cutoff == "default": default_changes.append( From e20a3c6c3d823e3bac1b8e95001f2fff15b189b8 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Mon, 28 Sep 2026 15:54:13 +0200 Subject: [PATCH 47/54] Document original fixed gap-position formula and custom trimming adaptation --- src/scirpy/ir_dist/metrics.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 6356e3eb1..c794a99b0 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1544,6 +1544,8 @@ def _nb_tcrdist_mat(): short_len = min(seq1_len, seq2_len) len_diff = abs(seq1_len - seq2_len) if fixed_gappos: + # Adapt the original min_gappos = min(6, 3 + (short_len - 5) // 2) formula to custom + # ntrim and ctrim values, keeping the gap within the remaining region. gappos_limit = max(6, ntrim) min_gappos = min(gappos_limit, ntrim + (short_len - ntrim - ctrim) // 2) max_gappos = min_gappos From e61f789e9c9e689d60eb87a50fb5b6d57d94877a Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 10:20:24 +0200 Subject: [PATCH 48/54] Update default cutoff and gap_penalty --- src/scirpy/ir_dist/metrics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 56fac5281..4030871fb 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1669,10 +1669,10 @@ class GPUTCRdistDistanceCalculator(TCRdistDistanceCalculator): def __init__( self, - cutoff: int = 20, + cutoff: int = 24, *, dist_weight: int = 3, - gap_penalty: int = 4, + gap_penalty: int = 12, ntrim: int = 3, ctrim: int = 2, fixed_gappos: bool = True, From 4b54cb14ac0800801ff50110205052de1ca1f24d Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 10:26:43 +0200 Subject: [PATCH 49/54] Avoid comparison of fully trimmed sequences --- src/scirpy/ir_dist/metrics.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 4030871fb..c2e226038 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1831,10 +1831,18 @@ def _gpu_tcrdist_mat( int seq1_len = seqs_L1[row]; int row_end_index = 0; + if (seq1_len <= ntrim + ctrim) { + row_element_counts[seqs_original_index] = 0; + return; + } + int col_start = length_starts[seq1_len]; int col_end = length_ends[seq1_len]; for (int col = col_start; col < col_end; col++) { int seq2_len = seqs_L2[col]; + if (seq2_len <= ntrim + ctrim) { + continue; + } long long distance = 1 + (long long)gap_penalty * abs(seq1_len - seq2_len); if (seq1_len == seq2_len) { From 7ea6d994b0efd9ad63e235cd912be42021bfc832 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 10:30:23 +0200 Subject: [PATCH 50/54] Respect custom trimming when calculating fixed GPU TCRdist gap position --- src/scirpy/ir_dist/metrics.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index c2e226038..20a4a9a78 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1858,9 +1858,10 @@ def _gpu_tcrdist_mat( else { int short_len = min(seq1_len, seq2_len); - // Equivalent to min(6, 3 + (short_len - 5) // 2) in the CPU implementation. - // Needed since Python rounds integer division down, whereas CUDA C++ truncates toward zero. - int gappos = min(6, (short_len + 1) / 2); + // Adapt the original min_gappos = min(6, 3 + (short_len - 5) // 2) formula to custom + // ntrim and ctrim values, keeping the gap within the remaining region. + int gappos_limit = max(6, ntrim); + int gappos = min(gappos_limit, ntrim + (short_len - ntrim - ctrim) / 2); int remainder = short_len - gappos; for (int n_i = ntrim; n_i < gappos; n_i++) { From 93d5a576d4800a947cb3d1441b3c7c9bf80d208f Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 10:38:24 +0200 Subject: [PATCH 51/54] Update GPU TCRdist documentation for new defaults and trimming behavior --- CHANGELOG.md | 7 ++++--- src/scirpy/ir_dist/__init__.py | 2 +- src/scirpy/ir_dist/metrics.py | 16 +++++++++------- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f77b63913..a137ee55e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning][]. ## Unreleased +### Features + + - Add `metric="gpu_tcrdist"` for GPU-accelerated TCRdist. Currently only `fixed_gappos=True` is supported. + ### Changed - Increase the default TCRdist `cutoff` from 20 to 24, allowing two gap positions with the new @@ -39,9 +43,6 @@ and this project adheres to [Semantic Versioning][]. ### Features - - Add `metric="gpu_tcrdist"` for GPU-accelerated TCRdist with row and column tiling, adaptive result buffers, - and support for BLOSUM62 and TCRBLOSUM matrices. - - Add a linear-gap Needleman-Wunsch distance metric for CDR3 amino acid sequences via `metric="needleman_wunsch"`. The `alignment` and `fastalignment` metrics are now deprecated. When `gap_open == gap_extend` (which has always been the default), use `needleman_wunsch` instead. diff --git a/src/scirpy/ir_dist/__init__.py b/src/scirpy/ir_dist/__init__.py index d6bb327a5..84729617a 100644 --- a/src/scirpy/ir_dist/__init__.py +++ b/src/scirpy/ir_dist/__init__.py @@ -88,7 +88,7 @@ def IrNeighbors(*args, **kwargs): matrix. A sensible cutoff depends on the distance metric, you can find information in the corresponding docs. If set to `None`, the cutoff will be `10` for the `alignment`, `fastalignment`, and `needleman_wunsch` metric, - `24` for `tcrdist`, and `2` for `levenshtein` and `hamming`. + `24` for `tcrdist` and `gpu_tcrdist`, and `2` for `levenshtein` and `hamming`. For the identity metric, the cutoff is ignored and always set to `0`. """ diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 20a4a9a78..f432c4145 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1608,8 +1608,8 @@ class GPUTCRdistDistanceCalculator(TCRdistDistanceCalculator): The code of this class is heavily based on `pwseqdist `_. Reused under MIT license, Copyright (c) 2020 Andrew Fiore-Gartland. - Using default weight, gap penalty, ntrim and ctrim is equivalent to the - original distance published in :cite:`TCRdist`. + This implements the CDR3 component of :cite:`TCRdist`. The default gap penalty + is 12, following the authors' later implementation; the original paper used 8. For performance reasons, the rows and columns of the final result matrix are grouped into tiles for GPU computation. `gpu_tile_rows` and `gpu_tile_cols` control how many matrix rows and columns are grouped into each @@ -1631,16 +1631,18 @@ class GPUTCRdistDistanceCalculator(TCRdistDistanceCalculator): dist_weight: Weight applied to the mismatch distances before summing with the gap penalties gap_penalty: - Distance penalty for the difference in the length of the two sequences + Distance penalty per amino acid of length difference. Defaults to 12. ntrim/ctrim: Positions trimmed off the N-terminus (0) and C-terminus (L-1) ends of the peptide sequence. These symbols will be ignored - in the distance calculation. + in the distance calculation. Sequences with length <= ntrim + ctrim have no remaining + positions and trigger a warning. All their comparisons, including self-comparisons, are + excluded; their rows and columns remain empty without changing the matrix shape. fixed_gappos: - If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6). + If True, insert gaps at a fixed position after the cysteine residue starting the CDR3 (typically position 6) - + the gap position accounts for custom ntrim and ctrim values and falls within the untrimmed region. Only True is currently supported; False raises NotImplementedError. cutoff: - Will eliminate distances > cutoff to make efficient use of sparse matrices. Must be between 0 and - 2**31 - 2 to allow distance + 1 encoding in the GPU result buffer. + Will eliminate distances > cutoff to make efficient use of sparse matrices. Defaults to 24. n_blocks: Number of outer row partitions submitted through joblib. This can be used with a distributed joblib backend to distribute the calculation across multiple GPU workers. From 00e62abdf6e82bd0d62069cb7575ddf2af875b72 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 10:45:16 +0200 Subject: [PATCH 52/54] Update GPU TCRdist tests for new defaults and trimming behavior --- src/scirpy/tests/test_ir_dist_metrics.py | 35 ++++++++++++++++++------ 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index 12953ccc2..6c3f5d597 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -938,6 +938,7 @@ def test_tcrdist(test_parameters, test_input, expected_result, calculator_class) test_parameters = test_parameters.copy() if calculator_class is GPUTCRdistDistanceCalculator: test_parameters.pop("n_jobs", None) + test_parameters.update(gpu_tile_rows=2, gpu_tile_cols=2, gpu_tile_buffer_cols=1) if not test_parameters.get("fixed_gappos", True): with pytest.raises(NotImplementedError, match="fixed_gappos=True"): calculator_class(**test_parameters) @@ -950,7 +951,14 @@ def test_tcrdist(test_parameters, test_input, expected_result, calculator_class) assert np.array_equal(res.todense(), expected_result) -@pytest.mark.parametrize("fixed_gappos", [True, False]) +@pytest.mark.parametrize( + "calculator_class,fixed_gappos", + [ + (TCRdistDistanceCalculator, True), + (TCRdistDistanceCalculator, False), + pytest.param(GPUTCRdistDistanceCalculator, True, marks=pytest.mark.gpu), + ], +) @pytest.mark.parametrize( "seqs,seqs2,ntrim,ctrim,n_excluded,expected", [ @@ -1005,13 +1013,16 @@ def test_tcrdist(test_parameters, test_input, expected_result, calculator_class) ), ], ) -def test_tcrdist_fully_trimmed(seqs, seqs2, ntrim, ctrim, n_excluded, expected, fixed_gappos): - calculator = TCRdistDistanceCalculator( - cutoff=20, gap_penalty=12, ntrim=ntrim, ctrim=ctrim, fixed_gappos=fixed_gappos, n_jobs=1, n_blocks=2 +def test_tcrdist_fully_trimmed(seqs, seqs2, ntrim, ctrim, n_excluded, expected, calculator_class, fixed_gappos): + kwargs = {"n_jobs": 1} + if calculator_class is GPUTCRdistDistanceCalculator: + kwargs = {"gpu_tile_rows": 2, "gpu_tile_cols": 2, "gpu_tile_buffer_cols": 1} + calculator = calculator_class( + cutoff=20, gap_penalty=12, ntrim=ntrim, ctrim=ctrim, fixed_gappos=fixed_gappos, n_blocks=2, **kwargs ) with pytest.warns(UserWarning, match=f"{n_excluded} input sequences have no positions remaining") as caught: result = calculator.calc_dist_mat(seqs, seqs2) - assert len(caught) == 1 + assert sum("input sequences have no positions remaining" in str(w.message) for w in caught) == 1 npt.assert_array_equal(result.toarray(), expected) assert result.nnz == np.count_nonzero(expected) @@ -1573,6 +1584,14 @@ def test_gpu_tcrdist(test_parameters, test_input, expected_result): npt.assert_array_equal(res.toarray(), expected_result) +@pytest.mark.gpu +def test_gpu_tcrdist_defaults(): + # Retain up to two extra residues at the default cutoff; three exceed it. + calculator = GPUTCRdistDistanceCalculator(gpu_tile_rows=2, gpu_tile_cols=2, gpu_tile_buffer_cols=1) + result = calculator.calc_dist_mat(np.array(["A" * 8, "A" * 9, "A" * 10, "A" * 11])) + npt.assert_array_equal(result.toarray(), [[1, 13, 25, 0], [13, 1, 13, 25], [25, 13, 1, 13], [0, 25, 13, 1]]) + + @pytest.mark.gpu def test_gpu_tcrdist_buffer_retry(): # Multiple retained distances per row force a retry; column lengths are deliberately unsorted. @@ -1581,7 +1600,7 @@ def test_gpu_tcrdist_buffer_retry(): ) result = tcrdist_calculator.calc_dist_mat(np.array(["AAAA", "AAA"]), np.array(["AAAR", "AAAAA", "AAA", "AAAA"])) - npt.assert_array_equal(result.toarray(), np.array([[13, 5, 5, 1], [0, 9, 1, 5]])) + npt.assert_array_equal(result.toarray(), np.array([[13, 13, 13, 1], [0, 0, 1, 13]])) @pytest.mark.gpu @@ -1668,7 +1687,7 @@ def test_gpu_tcrdist_reference(kwargs): seqs = np.load(TESTDATA / "tcrdist_test_data/tcrdist_WU3k_seqs.npy") reference_result = scipy.sparse.load_npz(TESTDATA / "tcrdist_test_data/tcrdist_WU3k_csr_result.npz") - gpu_tcrdist_calculator = GPUTCRdistDistanceCalculator(cutoff=15, **kwargs) + gpu_tcrdist_calculator = GPUTCRdistDistanceCalculator(cutoff=15, gap_penalty=4, **kwargs) res = gpu_tcrdist_calculator.calc_dist_mat(seqs, seqs) assert np.array_equal(res.data, reference_result.data) @@ -1708,7 +1727,7 @@ def test_gpu_tcrdist_length_bounds(gap_penalty, seqs2): @pytest.mark.gpu @pytest.mark.parametrize("ntrim,ctrim", [(0, 0), (3, 2), (10, 10)]) def test_gpu_tcrdist_short_sequences(ntrim, ctrim): - # Short, unequal sequences exercise the CPU's floor division when locating the fixed gap. + # Short sequences cover custom gap placement and exclusion after trimming. seqs = np.array(["AR", "ARN", "ARND", "ARNDC"]) seqs2 = np.array(["ACR", "ARNC", "RARND", "ARN"]) kwargs = {"cutoff": 100, "ntrim": ntrim, "ctrim": ctrim} From 2e7983559fe182513728920a2b695f8ef8dfa8b8 Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 11:12:14 +0200 Subject: [PATCH 53/54] Precompute weighted amino acid distance matrix --- src/scirpy/ir_dist/metrics.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index f432c4145..82d0825cd 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1798,7 +1798,8 @@ def _gpu_tcrdist_mat( seqs_mat1, seqs_L1 = _seqs2mat(seqs, max_len=max_seq_len) seqs_mat2, seqs_L2 = _seqs2mat(seqs2, max_len=max_seq_len) - d_dist_mat = cp.asarray(self.tcr_nb_distance_matrix.astype(np.int32, copy=False)) + dist_mat_weighted = self.tcr_nb_distance_matrix.astype(np.int64, copy=False) * self.dist_weight + d_dist_mat_weighted = cp.asarray(dist_mat_weighted) tcrdist_kernel = cp.RawKernel( r""" @@ -1820,9 +1821,8 @@ def _gpu_tcrdist_mat( const int seqs_mat2_rows, const int data_cols, const int indices_cols, - const int* __restrict__ aa_distance_matrix, + const long long* __restrict__ weighted_aa_distance_matrix, const int alphabet_size, - const int dist_weight, const int gap_penalty, const int ntrim, const int ctrim @@ -1851,7 +1851,7 @@ def _gpu_tcrdist_mat( for (int i = ntrim; i < seq1_len - ctrim; i++) { char val1 = seqs_mat1[(long long)i * seqs_mat1_rows + row]; char val2 = seqs_mat2[(long long)i * seqs_mat2_rows + col]; - distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + distance += weighted_aa_distance_matrix[val1 * alphabet_size + val2]; if (distance > cutoff + 1) { break; } @@ -1869,7 +1869,7 @@ def _gpu_tcrdist_mat( for (int n_i = ntrim; n_i < gappos; n_i++) { char val1 = seqs_mat1[(long long)n_i * seqs_mat1_rows + row]; char val2 = seqs_mat2[(long long)n_i * seqs_mat2_rows + col]; - distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + distance += weighted_aa_distance_matrix[val1 * alphabet_size + val2]; if (distance > cutoff + 1) { break; } @@ -1880,7 +1880,7 @@ def _gpu_tcrdist_mat( int j = seq2_len - 1 - c_i; char val1 = seqs_mat1[(long long)i * seqs_mat1_rows + row]; char val2 = seqs_mat2[(long long)j * seqs_mat2_rows + col]; - distance += (long long)dist_weight * aa_distance_matrix[val1 * alphabet_size + val2]; + distance += weighted_aa_distance_matrix[val1 * alphabet_size + val2]; if (distance > cutoff + 1) { break; } @@ -1979,9 +1979,8 @@ def run_tcrdist_kernel(buffer_width): seqs_mat2_rows, buffer_width, buffer_width, - d_dist_mat, + d_dist_mat_weighted, self.tcr_nb_distance_matrix.shape[0], - self.dist_weight, self.gap_penalty, self.ntrim, self.ctrim, From b17b75b967aa8f1dfe73f8993737d31f6eddc96c Mon Sep 17 00:00:00 2001 From: felixpetschko Date: Wed, 30 Sep 2026 11:23:40 +0200 Subject: [PATCH 54/54] Exclude fully trimmed columns from GPU TCRdist length bounds --- src/scirpy/ir_dist/metrics.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index 82d0825cd..ad643cb71 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -1842,9 +1842,6 @@ def _gpu_tcrdist_mat( int col_end = length_ends[seq1_len]; for (int col = col_start; col < col_end; col++) { int seq2_len = seqs_L2[col]; - if (seq2_len <= ntrim + ctrim) { - continue; - } long long distance = 1 + (long long)gap_penalty * abs(seq1_len - seq2_len); if (seq1_len == seq2_len) { @@ -2058,9 +2055,10 @@ def run_tcrdist_kernel(buffer_width): possible_lengths = np.arange(max_seq_len + 1) max_length_diff = self.cutoff // self.gap_penalty if self.gap_penalty else max_seq_len + min_lengths = np.maximum(possible_lengths - max_length_diff, self.ntrim + self.ctrim + 1) length_bounds_blocks = [ ( - cp.asarray(np.searchsorted(lengths, possible_lengths - max_length_diff, side="left").astype(np.int32)), + cp.asarray(np.searchsorted(lengths, min_lengths, side="left").astype(np.int32)), cp.asarray(np.searchsorted(lengths, possible_lengths + max_length_diff, side="right").astype(np.int32)), ) for lengths in seqs_L2_blocks