Add GPU-accelerated TCRdist distance calculation

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.

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

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:
 

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:

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)"
             )