Test ir_dist TCRBLOSUM chain routing

Author: felixpetschko

src/scirpy/tests/test_ir_dist.py

@@ -12,7 +12,7 @@
 from scirpy.ir_dist.metrics import DistanceCalculator
 from scirpy.util import DataHandler, _is_symmetric
 
-from .util import _squarify
+from .util import _make_adata, _squarify
 
 
 def _assert_frame_equal(left, right):
@@ -163,6 +163,44 @@ def test_ir_dist(
     npt.assert_array_equal(res["VDJ"]["distances"].toarray(), expected_dist_vdj)
 
 
+@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.
+    adata = _make_adata(
+        pd.DataFrame(
+            [
+                ["cell1", "AAACAAAA", "AAACAAAA", "TRA", "TRB"],
+                ["cell2", "AAAHAAAA", "AAAHAAAA", "TRA", "TRB"],
+            ],
+            columns=[
+                "cell_id",
+                "IR_VJ_1_junction_aa",
+                "IR_VDJ_1_junction_aa",
+                "IR_VJ_1_locus",
+                "IR_VDJ_1_locus",
+            ],
+        ).set_index("cell_id"),
+        mudata,
+    )
+
+    ir.pp.ir_dist(
+        adata,
+        metric="tcrdist",
+        sequence="aa",
+        cutoff=20,
+        base_matrix="tcrblosum",
+        key_added="ir_dist_tcrblosum",
+        n_jobs=1,
+    )
+
+    res = adata.mod["airr"].uns["ir_dist_tcrblosum"] if isinstance(adata, MuData) else adata.uns["ir_dist_tcrblosum"]
+    expected_seqs = np.array(["AAACAAAA", "AAAHAAAA"])
+    npt.assert_array_equal(res["VJ"]["seqs"], expected_seqs)
+    npt.assert_array_equal(res["VDJ"]["seqs"], expected_seqs)
+    npt.assert_array_equal(res["VJ"]["distances"].toarray(), np.array([[1, 16], [16, 1]]))
+    npt.assert_array_equal(res["VDJ"]["distances"].toarray(), np.array([[1, 19], [19, 1]]))
+
+
 @pytest.mark.parametrize("with_adata2", [False, True])
 @pytest.mark.parametrize("joblib_backend", ["loky", "multiprocessing", "threading"])
 @pytest.mark.parametrize("n_jobs", [1, 2])