add min_times_seen to PolyclonalCollection.mut_escape_corr

jbloomlab · Aug 13, 2023 · ec06a32 · ec06a32
1 parent 3c2dc6e
commit ec06a32
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -9,6 +9,7 @@ The format is based on `Keep a Changelog <https://keepachangelog.com>`_.
 6.6
 ---
 - Fix bug where ``scale_stat_col`` to per-model values in tooltips when using ``PolyclonalCollection.mut_escape_plot`` and ``PolyclonalCollections.mut_escape_icXX``.
+- Add ``min_times_seen`` filter to ``PolyclonalCollection.mut_escape_corr`` and ``PolyclonalCollection.mut_escape_corr_heatmap``.
 
 6.5
 ---

diff --git a/polyclonal/polyclonal_collection.py b/polyclonal/polyclonal_collection.py
@@ -776,13 +776,15 @@ def mut_icXX_df_w_model_values(self, **kwargs):
             validate="one_to_one",
         )
 
-    def mut_escape_corr(self, method="pearson"):
+    def mut_escape_corr(self, method="pearson", min_times_seen=1):
         """Correlation of mutation escape values across models for each epitope.
 
         Parameters
         ----------
         method : str
             A correlation method passable to `pandas.DataFrame.corr`.
+        min_times_seen : int
+            Only include mutations with a *times_seen* >= this value.
 
         Returns
         -------
@@ -805,7 +807,9 @@ def mut_escape_corr(self, method="pearson"):
 
         corr = (
             polyclonal.utils.tidy_to_corr(
-                self.mut_escape_df_replicates.merge(
+                self.mut_escape_df_replicates.query(
+                    "times_seen >= @min_times_seen"
+                ).merge(
                     ids,
                     on=self.descriptor_names,
                     validate="many_to_one",
@@ -826,13 +830,21 @@ def mut_escape_corr(self, method="pearson"):
 
         return corr
 
-    def mut_escape_corr_heatmap(self, method="pearson", plot_corr2=True, **kwargs):
+    def mut_escape_corr_heatmap(
+        self,
+        method="pearson",
+        min_times_seen=1,
+        plot_corr2=True,
+        **kwargs,
+    ):
         """Heatmap of mutation-escape correlation among models at each epitope.
 
         Parameters
         ----------
         method : str
             A correlation method passable to `pandas.DataFrame.corr`.
+        min_times_seen : int
+            Only include mutations with a *times_seen* >= this value.
         plot_corr2 : bool
             Plot squared correlation (eg, :math:`R^2` rather :math:`R`).
         **kwargs
@@ -841,7 +853,7 @@ def mut_escape_corr_heatmap(self, method="pearson", plot_corr2=True, **kwargs):
         corr_label = {"pearson": "r", "kendall": "tau", "spearman": "rho"}[method]
         corr2_label = f"{corr_label}2"
         corr_df = (
-            self.mut_escape_corr(method)
+            self.mut_escape_corr(method, min_times_seen)
             .assign(correlation2=lambda x: x["correlation"] ** 2)
             .rename(columns={"correlation": corr_label, "correlation2": corr2_label})
         )