diff --git a/CHANGELOG.md b/CHANGELOG.md index 89800ddba..a137ee55e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,31 @@ 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 + 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. + + - 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` + to retain the previous behavior. + ### Fixes - Avoid deprecated AnnData `obsm_keys()` calls when plotting clonotype networks or exporting them to igraph. 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/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: diff --git a/src/scirpy/ir_dist/__init__.py b/src/scirpy/ir_dist/__init__.py index 8104795d1..84729617a 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. @@ -84,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, - 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`. """ @@ -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/_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 71c926062..ad643cb71 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,24 +1287,30 @@ 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. + 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 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 - 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: @@ -1328,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 = 4, + gap_penalty: int | Literal["default"] = "default", # resolves to 12 ntrim: int = 3, ctrim: int = 2, fixed_gappos: bool = True, @@ -1342,6 +1349,29 @@ 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( + "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": + default_changes.append( + "The default value of `gap_penalty` has changed from 4 to 12. " + "Set `gap_penalty=12` explicitly to acknowledge the new default, " + "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, + ) self.dist_weight = dist_weight self.gap_penalty = gap_penalty self.ntrim = ntrim @@ -1394,6 +1424,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, *, @@ -1485,10 +1527,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): @@ -1498,9 +1544,13 @@ 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) + # 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 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: @@ -1552,6 +1602,592 @@ 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. + + 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 + 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 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. 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 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. 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. + 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 = 24, + *, + dist_weight: int = 3, + gap_penalty: int = 12, + 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) + 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""" + 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 long long* __restrict__ weighted_aa_distance_matrix, + const int alphabet_size, + 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; + + 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]; + 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 += weighted_aa_distance_matrix[val1 * alphabet_size + val2]; + if (distance > cutoff + 1) { + break; + } + } + } + else { + int short_len = min(seq1_len, seq2_len); + + // 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++) { + 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 += weighted_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 += weighted_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_weighted, + self.tcr_nb_distance_matrix.shape[0], + 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 + min_lengths = np.maximum(possible_lengths - max_length_diff, self.ntrim + self.ctrim + 1) + length_bounds_blocks = [ + ( + 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 + ] + + 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_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.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..6c3f5d597 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 @@ -13,6 +14,7 @@ DistanceCalculator, FastAlignmentDistanceCalculator, GPUHammingDistanceCalculator, + GPUTCRdistDistanceCalculator, HammingDistanceCalculator, IdentityDistanceCalculator, LevenshteinDistanceCalculator, @@ -633,7 +635,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, @@ -648,7 +650,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 ( @@ -665,9 +667,9 @@ 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 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, @@ -682,7 +684,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 @@ -860,10 +862,88 @@ 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 previous cutoff with gap_penalty=12. + ( + {"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]]), + ), + # 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]]), + ), + # 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): - 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) + 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) + 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) @@ -871,6 +951,133 @@ def test_tcrdist(test_parameters, test_input, expected_result): assert np.array_equal(res.todense(), expected_result) +@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", + [ + # 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, 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 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) + + +@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(cutoff=24, 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", [ @@ -1019,13 +1226,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 +1242,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 +1258,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 +1269,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 +1280,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 +1522,108 @@ 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 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): + # 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 + ) + 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_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. + 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.array([[13, 13, 13, 1], [0, 0, 1, 13]])) + + +@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 +1666,71 @@ 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, gap_penalty=4, **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]) +@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, seqs2) + expected = TCRdistDistanceCalculator(cutoff=4, gap_penalty=gap_penalty, ntrim=0, ctrim=0, n_jobs=1).calc_dist_mat( + seqs, seqs2 + ) + 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 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} + 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()) 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