add filter for umis per target

FrickTobias · Nov 19, 2021 · d1334e3 · d1334e3
1 parent 5e1d750
commit d1334e3
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/dbspro/notebook.py b/src/dbspro/notebook.py
@@ -87,6 +87,16 @@ def filter_ratio(df: pd.DataFrame, threshold: int, opr=operator.gt) -> pd.DataFr
 pd.DataFrame.filter_ratio = filter_ratio
 
 
+@_filter_wrapper
+def filter_target_count(df: pd.DataFrame, threshold: int, opr=operator.gt) -> pd.DataFrame:
+    """Filter UMI count per target"""
+    print(f"Filtering for targets with UMI count {opr.__name__} {threshold}")
+    return df[opr(df.groupby(["Barcode", "Target"])["UMI"].transform('count'), threshold)]
+
+
+pd.DataFrame.filter_target_count = filter_target_count
+
+
 @_filter_wrapper
 def filter_dups(df: pd.DataFrame, threshold: Union[float, int] = 2, min_len: int = 3) -> pd.DataFrame:
     """Remove barcodes that share a portion of their UMI-Targets combos based on the given threshold. If float