From 6cb91a58983ecb051fc5beb762c384e9fab7c79d Mon Sep 17 00:00:00 2001
From: Zewen Kelvin Tuong <z.tuong@uq.edu.au>
Date: Wed, 6 Dec 2023 12:24:56 +1000
Subject: [PATCH] add a new option to toggle whether to keep the
 id_cp_interaction value when plotting

---
 ktplotspy/plot/plot_cpdb.py |  8 ++++++++
 tests/test_plot_cpdb.py     | 16 ++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/ktplotspy/plot/plot_cpdb.py b/ktplotspy/plot/plot_cpdb.py
index f622240..a97deb1 100644
--- a/ktplotspy/plot/plot_cpdb.py
+++ b/ktplotspy/plot/plot_cpdb.py
@@ -80,6 +80,7 @@ def plot_cpdb(
     scale_alpha_by_interaction_scores: bool = False,
     scale_alpha_by_cellsign: bool = False,
     filter_by_cellsign: bool = False,
+    keep_id_cp_interaction: bool = False,
 ) -> Union[ggplot, pd.DataFrame]:
     """Plotting CellPhoneDB results as a dot plot.
 
@@ -157,6 +158,8 @@ def plot_cpdb(
         Whether or not to filter the transparency of interactions by the cellsign.
     filter_by_cellsign: bool, optional
         Filter out interactions with a 0 value cellsign.
+    keep_id_cp_interaction: bool, optional
+        Whether to keep the original `id_cp_interaction` value when plotting.
     Returns
     -------
     Union[ggplot, pd.DataFrame]
@@ -384,6 +387,11 @@ def plot_cpdb(
     if return_table:
         return df
     else:
+        # change the labelling of interaction_group
+        if keep_id_cp_interaction:
+            df.interaction_group = [re.sub(DEFAULT_SEP * 3, "_", c) for c in df.interaction_group]
+        else:
+            df.interaction_group = [c.split(DEFAULT_SEP * 3)[1] for c in df.interaction_group]
         # set global figure size
         options.figure_size = figsize
 
diff --git a/tests/test_plot_cpdb.py b/tests/test_plot_cpdb.py
index 857b4bd..5292865 100644
--- a/tests/test_plot_cpdb.py
+++ b/tests/test_plot_cpdb.py
@@ -23,6 +23,22 @@ def test_plot_cpdb(mock_show, adata, means, pvals):
     g
 
 
+@patch("matplotlib.pyplot.show")
+@pytest.mark.usefixtures("adata", "means", "pvals")
+def test_plot_cpdb_keep_id(mock_show, adata, means, pvals):
+    g = plot_cpdb(
+        adata=adata,
+        cell_type1="B cell",
+        cell_type2="CD4T cell",
+        means=means,
+        pvals=pvals,
+        celltype_key="celltype",
+        genes=["CXCL13", "CD274", "CXCR5"],
+        keep_id_cp_interaction=True,
+    )
+    g
+
+
 @patch("matplotlib.pyplot.show")
 @pytest.mark.usefixtures("adata", "means", "pvals")
 def test_plot_cpdb_title(mock_show, adata, means, pvals):