Add tcrblosum support to TCRdist

src/scirpy/ir_dist/__init__.py

@@ -85,7 +85,9 @@ def _get_metric_key(metric: MetricType) -> str:
     return "custom" if isinstance(metric, metrics.DistanceCalculator) else metric  # type: ignore
 
 
-def _get_distance_calculator(metric: MetricType, cutoff: int | None, *, n_jobs=-1, **kwargs):
+def _get_distance_calculator(
+    metric: MetricType, cutoff: int | None, *, n_jobs=-1, chain_type: Literal["VJ", "VDJ"] | None = None, **kwargs
+):
     """Returns an instance of :class:`~scirpy.ir_dist.metrics.DistanceCalculator`
     given a metric.
 
@@ -116,7 +118,7 @@ def _get_distance_calculator(metric: MetricType, cutoff: int | None, *, n_jobs=-
     elif metric == "gpu_hamming":
         dist_calc = metrics.GPUHammingDistanceCalculator(**kwargs)
     elif metric == "tcrdist":
-        dist_calc = metrics.TCRdistDistanceCalculator(n_jobs=n_jobs, **kwargs)
+        dist_calc = metrics.TCRdistDistanceCalculator(n_jobs=n_jobs, chain_type=chain_type, **kwargs)
     else:
         raise ValueError("Invalid distance metric.")
 
@@ -252,9 +254,9 @@ def _get_unique_seqs(tmp_adata, chain_type):
                 result[chain_type][tmp_key] = unique_seqs
 
     # compute distance matrices
-    dist_calc = _get_distance_calculator(metric, cutoff, n_jobs=n_jobs, **kwargs)
     for chain_type in ["VJ", "VDJ"]:
         logging.info(f"Computing sequence x sequence distance matrix for {chain_type} sequences.")  # type: ignore
+        dist_calc = _get_distance_calculator(metric, cutoff, n_jobs=n_jobs, chain_type=chain_type, **kwargs)
         result[chain_type]["distances"] = dist_calc.calc_dist_mat(
             result[chain_type]["seqs"], result[chain_type].get("seqs2", None)
         ).tocsr()