Merge branch 'feat/clonotype_definition' into 'master'

Support alternative clonotype definition.

See merge request icbi-lab/pipelines/singlecell_tcr!13

Former-commit-id: 3121a0452abc9f9d251fb9c36dce3bd0f0920609

sctcrpy/_tools.py

@@ -4,10 +4,15 @@
 import pandas as pd
 from typing import Union
 from ._util import _is_na, _add_to_uns
+from ._compat import Literal
 
 
 def define_clonotypes(
-    adata: AnnData, *, inplace: bool = True, key_added: str = "clonotype"
+    adata: AnnData,
+    *,
+    flavor: Literal["all_chains", "primary_only"] = "all_chains",
+    inplace: bool = True,
+    key_added: str = "clonotype"
 ) -> Union[None, np.ndarray]:
     """Define clonotypes based on CDR3 region.
 
@@ -20,6 +25,10 @@ def define_clonotypes(
     ----------
     adata
         Annotated data matrix
+    flavor
+        Biological model to define clonotypes. 
+        `all_chains`: All four chains of a cell in a clonotype need to be the same. 
+        `primary_only`: Only primary alpha and beta chain need to be the same. 
     inplace
         If True, adds a column to adata.obs
     key_added
@@ -32,12 +41,14 @@ def define_clonotypes(
     or adds a `clonotype` column to `adata`. 
     
     """
+    groupby_cols = {
+        "all_chains": ["TRA_1_cdr3", "TRB_1_cdr3", "TRA_2_cdr3", "TRA_2_cdr3"],
+        "primary_only": ["TRA_1_cdr3", "TRB_1_cdr3"],
+    }
     clonotype_col = np.array(
         [
             "clonotype_{}".format(x)
-            for x in adata.obs.groupby(
-                ["TRA_1_cdr3", "TRB_1_cdr3", "TRA_2_cdr3", "TRA_2_cdr3"]
-            ).ngroup()
+            for x in adata.obs.groupby(groupby_cols[flavor]).ngroup()
         ]
     )
     clonotype_col[

tests/test_tools.py

@@ -28,6 +28,16 @@ def test_define_clonotypes():
         ["clonotype_1", np.nan, "clonotype_1", "clonotype_0", "clonotype_2"],
     )
 
+    res_primary_only = st.tl.define_clonotypes(
+        adata, flavor="primary_only", inplace=False
+    )
+    npt.assert_equal(
+        # order is by alphabet: BBBB < nan
+        # we don't care about the order of numbers, so this is ok.
+        res_primary_only,
+        ["clonotype_0", np.nan, "clonotype_0", "clonotype_0", "clonotype_1"],
+    )
+
     # test inplace
     st.tl.define_clonotypes(adata, key_added="clonotype_")
     npt.assert_equal(res, adata.obs["clonotype_"].values)