smftools 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

smftools/preprocessing/filter_reads_on_modification_thresholds.py

@@ -1,9 +1,16 @@
-import math
+from __future__ import annotations
+
 import gc
+from typing import List, Optional, Sequence
+
+import anndata as ad
 import numpy as np
 import pandas as pd
-import anndata as ad
-from typing import Optional, Sequence, List
+
+from smftools.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
 
 def filter_reads_on_modification_thresholds(
     adata: ad.AnnData,
@@ -15,32 +22,40 @@ def filter_reads_on_modification_thresholds(
     a_thresholds: Optional[Sequence[float]] = None,
     use_other_c_as_background: bool = False,
     min_valid_fraction_positions_in_read_vs_ref: Optional[float] = None,
-    uns_flag: str = 'filter_reads_on_modification_thresholds_performed',
+    uns_flag: str = "filter_reads_on_modification_thresholds_performed",
     bypass: bool = False,
     force_redo: bool = False,
-    reference_column: str = 'Reference_strand',
+    reference_column: str = "Reference_strand",
     # memory-control options:
     batch_size: int = 200,
     compute_obs_if_missing: bool = True,
-    treat_zero_as_invalid: bool = False
+    treat_zero_as_invalid: bool = False,
 ) -> ad.AnnData:
-    """
-    Memory-efficient filtering by per-read modification thresholds.
-
-    - If required obs columns exist, uses them directly (fast).
-    - Otherwise, computes the relevant per-read metrics per-reference in batches
-      and writes them into adata.obs before filtering.
-
-    Parameters of interest :
-      - gpc_thresholds, cpg_thresholds, any_c_thresholds, a_thresholds:
-          each should be [min, max] (floats 0..1) or None. Thresholds are inclusive.
-      - use_other_c_as_background: require GpC/CpG > other_C background (if present).
-      - min_valid_fraction_positions_in_read_vs_ref: minimum fraction of valid sites
-          in the read vs reference (0..1). If None, this check is skipped.
-      - compute_obs_if_missing: if True, compute Fraction_* and Valid_* obs columns
-          if they are not already present, using a low-memory per-ref strategy.
-      - treat_zero_as_invalid: if True, a zero in X counts as invalid (non-site).
-          If False, zeros are considered valid positions (adjust to your data semantics).
+    """Filter reads based on per-read modification thresholds.
+
+    If required obs columns exist, they are used directly. Otherwise, the function
+    computes the relevant per-read metrics in batches and stores them in ``adata.obs``.
+
+    Args:
+        adata: AnnData object to filter.
+        smf_modality: SMF modality identifier.
+        mod_target_bases: List of target bases to evaluate.
+        gpc_thresholds: ``[min, max]`` thresholds for GpC (0..1) or ``None``.
+        cpg_thresholds: ``[min, max]`` thresholds for CpG (0..1) or ``None``.
+        any_c_thresholds: ``[min, max]`` thresholds for any C (0..1) or ``None``.
+        a_thresholds: ``[min, max]`` thresholds for A (0..1) or ``None``.
+        use_other_c_as_background: Require GpC/CpG > other_C background if present.
+        min_valid_fraction_positions_in_read_vs_ref: Minimum valid-site fraction per read.
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+        bypass: Whether to skip processing.
+        force_redo: Whether to rerun even if ``uns_flag`` is set.
+        reference_column: Obs column containing reference identifiers.
+        batch_size: Batch size for low-memory computation.
+        compute_obs_if_missing: Whether to compute missing obs summaries.
+        treat_zero_as_invalid: Whether zeros should be treated as invalid positions.
+
+    Returns:
+        anndata.AnnData: Filtered AnnData object.
     """
 
     # quick exit flags:
@@ -50,17 +65,23 @@ def filter_reads_on_modification_thresholds(
 
     # helper: check whether obs columns exist for a particular mod type
     def obs_has_columns_for(mod_type):
+        """Return True if per-read summary columns exist for a mod type."""
         col_pref = {
-            "GpC": ("Fraction_GpC_site_modified", f"Valid_GpC_site_in_read_vs_reference"),
-            "CpG": ("Fraction_CpG_site_modified", f"Valid_CpG_site_in_read_vs_reference"),
-            "C": ("Fraction_C_site_modified", f"Valid_C_site_in_read_vs_reference"),
-            "A": ("Fraction_A_site_modified", f"Valid_A_site_in_read_vs_reference"),
+            "GpC": ("Fraction_GpC_site_modified", "Valid_GpC_site_in_read_vs_reference"),
+            "CpG": ("Fraction_CpG_site_modified", "Valid_CpG_site_in_read_vs_reference"),
+            "C": ("Fraction_C_site_modified", "Valid_C_site_in_read_vs_reference"),
+            "A": ("Fraction_A_site_modified", "Valid_A_site_in_read_vs_reference"),
         }.get(mod_type, (None, None))
         return (col_pref[0] in adata.obs.columns) and (col_pref[1] in adata.obs.columns)
 
     # if all required obs columns are present, use them directly (fast path)
     required_present = True
-    for mt, thr in (("GpC", gpc_thresholds), ("CpG", cpg_thresholds), ("C", any_c_thresholds), ("A", a_thresholds)):
+    for mt, thr in (
+        ("GpC", gpc_thresholds),
+        ("CpG", cpg_thresholds),
+        ("C", any_c_thresholds),
+        ("A", a_thresholds),
+    ):
         if thr is not None and mt in mod_target_bases:
             if not obs_has_columns_for(mt):
                 required_present = False
@@ -75,9 +96,10 @@ def filter_reads_on_modification_thresholds(
     # Build mapping from reference -> var column names (expected pattern)
     # e.g. var column names: "{ref}_GpC_site", "{ref}_CpG_site", "{ref}_any_C_site", "{ref}_other_C_site", "{ref}_A_site"
     # If your var column naming differs, adjust these suffixes.
-    refs = list(adata.obs[reference_column].astype('category').cat.categories)
+    refs = list(adata.obs[reference_column].astype("category").cat.categories)
 
     def _find_var_col_for(ref, suffix):
+        """Resolve a var column name for a reference/suffix pair."""
         name = f"{ref}_{suffix}"
         if name in adata.var.columns:
             return name
@@ -121,7 +143,9 @@ def filter_reads_on_modification_thresholds(
                 var_mask_bool = np.asarray(adata.var[var_colname].values).astype(bool)
             except Exception:
                 # if var has values not boolean, attempt coercion
-                var_mask_bool = np.asarray(pd.to_numeric(adata.var[var_colname], errors='coerce').fillna(0).astype(bool))
+                var_mask_bool = np.asarray(
+                    pd.to_numeric(adata.var[var_colname], errors="coerce").fillna(0).astype(bool)
+                )
 
             if not var_mask_bool.any():
                 return
@@ -154,16 +178,20 @@ def filter_reads_on_modification_thresholds(
                         #   valid_count = (non-nan if float data else non-zero) per row
                         # For sparse, .data are only stored nonzeros, so (X_block > 0).sum is fine
                         modified_count = np.asarray((X_block > 0).sum(axis=1)).ravel()
-                        if np.isnan(X_block.data).any() if hasattr(X_block, 'data') else False:
+                        if np.isnan(X_block.data).any() if hasattr(X_block, "data") else False:
                             # if sparse with stored NaNs (!) handle differently - unlikely
-                            valid_count = np.asarray(~np.isnan(X_block.toarray()).sum(axis=1)).ravel()
+                            valid_count = np.asarray(
+                                ~np.isnan(X_block.toarray()).sum(axis=1)
+                            ).ravel()
                         else:
                             if treat_zero_as_invalid:
                                 # valid = number of non-zero entries
                                 valid_count = np.asarray((X_block != 0).sum(axis=1)).ravel()
                             else:
                                 # treat all positions as valid positions (they exist in reference) -> denominator = n_cols_for_ref
-                                valid_count = np.full_like(modified_count, n_cols_for_ref, dtype=float)
+                                valid_count = np.full_like(
+                                    modified_count, n_cols_for_ref, dtype=float
+                                )
                     else:
                         # dense numpy
                         Xb = np.asarray(X_block)
@@ -193,14 +221,18 @@ def filter_reads_on_modification_thresholds(
 
                 # fraction modified = modified_count / valid_count (guard divide-by-zero)
                 frac = np.zeros_like(modified_count, dtype=float)
-                mask_valid_nonzero = (valid_count > 0)
-                frac[mask_valid_nonzero] = modified_count[mask_valid_nonzero] / valid_count[mask_valid_nonzero]
+                mask_valid_nonzero = valid_count > 0
+                frac[mask_valid_nonzero] = (
+                    modified_count[mask_valid_nonzero] / valid_count[mask_valid_nonzero]
+                )
 
                 # write to out arrays
                 out_frac_arr[block_rows_idx] = frac
                 # valid fraction relative to reference = valid_count / n_cols_for_ref
                 out_valid_arr[block_rows_idx] = np.zeros_like(valid_count, dtype=float)
-                out_valid_arr[block_rows_idx][mask_valid_nonzero] = (valid_count[mask_valid_nonzero] / float(n_cols_for_ref))
+                out_valid_arr[block_rows_idx][mask_valid_nonzero] = valid_count[
+                    mask_valid_nonzero
+                ] / float(n_cols_for_ref)
 
                 # free block memory ASAP
                 del X_block, modified_count, valid_count, frac
@@ -210,29 +242,51 @@ def filter_reads_on_modification_thresholds(
         # GpC
         if "GpC" in mod_target_bases:
             for ref in refs:
-                _compute_for_ref_and_suffix(ref, "GpC_site", create_cols["Fraction_GpC_site_modified"], create_cols["Valid_GpC_site_in_read_vs_reference"])
+                _compute_for_ref_and_suffix(
+                    ref,
+                    "GpC_site",
+                    create_cols["Fraction_GpC_site_modified"],
+                    create_cols["Valid_GpC_site_in_read_vs_reference"],
+                )
         # other_C (for background)
         # We'll also compute 'other_C' per reference if it exists
         other_c_per_ref = {}
         for ref in refs:
             other_col = _find_var_col_for(ref, "other_C_site")
             if other_col:
-                other_c_per_ref[ref] = np.where(np.asarray(adata.var[other_col].values).astype(bool))[0]
+                other_c_per_ref[ref] = np.where(
+                    np.asarray(adata.var[other_col].values).astype(bool)
+                )[0]
 
         # CpG
         if "CpG" in mod_target_bases:
             for ref in refs:
-                _compute_for_ref_and_suffix(ref, "CpG_site", create_cols["Fraction_CpG_site_modified"], create_cols["Valid_CpG_site_in_read_vs_reference"])
+                _compute_for_ref_and_suffix(
+                    ref,
+                    "CpG_site",
+                    create_cols["Fraction_CpG_site_modified"],
+                    create_cols["Valid_CpG_site_in_read_vs_reference"],
+                )
 
         # any C
         if "C" in mod_target_bases:
             for ref in refs:
-                _compute_for_ref_and_suffix(ref, "C_site", create_cols["Fraction_C_site_modified"], create_cols["Valid_C_site_in_read_vs_reference"])
+                _compute_for_ref_and_suffix(
+                    ref,
+                    "C_site",
+                    create_cols["Fraction_C_site_modified"],
+                    create_cols["Valid_C_site_in_read_vs_reference"],
+                )
 
         # A
         if "A" in mod_target_bases:
             for ref in refs:
-                _compute_for_ref_and_suffix(ref, "A_site", create_cols["Fraction_A_site_modified"], create_cols["Valid_A_site_in_read_vs_reference"])
+                _compute_for_ref_and_suffix(
+                    ref,
+                    "A_site",
+                    create_cols["Fraction_A_site_modified"],
+                    create_cols["Valid_A_site_in_read_vs_reference"],
+                )
 
         # write created arrays into adata.obs
         for cname, arr in create_cols.items():
@@ -243,16 +297,20 @@ def filter_reads_on_modification_thresholds(
             # compute per-ref background ratio if both exist
             # Simplest approach: if 'Fraction_GpC_site_modified' and 'Fraction_other_C_site_modified' exist, compute ratio
             if "Fraction_other_C_site_modified" in adata.obs.columns:
-                with np.errstate(divide='ignore', invalid='ignore'):
-                    ratio = adata.obs["Fraction_GpC_site_modified"].astype(float) / adata.obs["Fraction_other_C_site_modified"].astype(float)
+                with np.errstate(divide="ignore", invalid="ignore"):
+                    ratio = adata.obs["Fraction_GpC_site_modified"].astype(float) / adata.obs[
+                        "Fraction_other_C_site_modified"
+                    ].astype(float)
                 adata.obs["GpC_to_other_C_mod_ratio"] = ratio.fillna(0.0)
             else:
                 adata.obs["GpC_to_other_C_mod_ratio"] = np.nan
 
         if "CpG" in mod_target_bases and use_other_c_as_background:
             if "Fraction_other_C_site_modified" in adata.obs.columns:
-                with np.errstate(divide='ignore', invalid='ignore'):
-                    ratio = adata.obs["Fraction_CpG_site_modified"].astype(float) / adata.obs["Fraction_other_C_site_modified"].astype(float)
+                with np.errstate(divide="ignore", invalid="ignore"):
+                    ratio = adata.obs["Fraction_CpG_site_modified"].astype(float) / adata.obs[
+                        "Fraction_other_C_site_modified"
+                    ].astype(float)
                 adata.obs["CpG_to_other_C_mod_ratio"] = ratio.fillna(0.0)
             else:
                 adata.obs["CpG_to_other_C_mod_ratio"] = np.nan
@@ -266,10 +324,14 @@ def filter_reads_on_modification_thresholds(
 
     # helper to get min/max from param like [min, max] or tuple(None,..)
     def _unpack_minmax(thr):
+        """Normalize a threshold pair to ordered (min, max) floats."""
         if thr is None:
             return None, None
         try:
-            lo, hi = float(thr[0]) if thr[0] is not None else None, float(thr[1]) if thr[1] is not None else None
+            lo, hi = (
+                float(thr[0]) if thr[0] is not None else None,
+                float(thr[1]) if thr[1] is not None else None,
+            )
             if lo is not None and hi is not None and lo > hi:
                 lo, hi = hi, lo
             return lo, hi
@@ -277,76 +339,124 @@ def filter_reads_on_modification_thresholds(
             return None, None
 
     # GpC thresholds
-    if gpc_thresholds and 'GpC' in mod_target_bases:
+    if gpc_thresholds and "GpC" in mod_target_bases:
         lo, hi = _unpack_minmax(gpc_thresholds)
-        if use_other_c_as_background and smf_modality != 'deaminase' and "GpC_to_other_C_mod_ratio" in filtered.obs.columns:
+        if (
+            use_other_c_as_background
+            and smf_modality != "deaminase"
+            and "GpC_to_other_C_mod_ratio" in filtered.obs.columns
+        ):
             filtered = filtered[filtered.obs["GpC_to_other_C_mod_ratio"].astype(float) > 1]
         if lo is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_GpC_site_modified"].astype(float) >= lo]
-            print(f"Removed {s0 - filtered.n_obs} reads below min GpC fraction {lo}")
+            logger.info("Removed %s reads below min GpC fraction %s", s0 - filtered.n_obs, lo)
         if hi is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_GpC_site_modified"].astype(float) <= hi]
-            print(f"Removed {s0 - filtered.n_obs} reads above max GpC fraction {hi}")
-        if (min_valid_fraction_positions_in_read_vs_ref is not None) and ("Valid_GpC_site_in_read_vs_reference" in filtered.obs.columns):
+            logger.info("Removed %s reads above max GpC fraction %s", s0 - filtered.n_obs, hi)
+        if (min_valid_fraction_positions_in_read_vs_ref is not None) and (
+            "Valid_GpC_site_in_read_vs_reference" in filtered.obs.columns
+        ):
             s0 = filtered.n_obs
-            filtered = filtered[filtered.obs["Valid_GpC_site_in_read_vs_reference"].astype(float) >= float(min_valid_fraction_positions_in_read_vs_ref)]
-            print(f"Removed {s0 - filtered.n_obs} reads with insufficient valid GpC site fraction vs ref")
+            filtered = filtered[
+                filtered.obs["Valid_GpC_site_in_read_vs_reference"].astype(float)
+                >= float(min_valid_fraction_positions_in_read_vs_ref)
+            ]
+            logger.info(
+                "Removed %s reads with insufficient valid GpC site fraction vs ref",
+                s0 - filtered.n_obs,
+            )
 
     # CpG thresholds
-    if cpg_thresholds and 'CpG' in mod_target_bases:
+    if cpg_thresholds and "CpG" in mod_target_bases:
         lo, hi = _unpack_minmax(cpg_thresholds)
-        if use_other_c_as_background and smf_modality != 'deaminase' and "CpG_to_other_C_mod_ratio" in filtered.obs.columns:
+        if (
+            use_other_c_as_background
+            and smf_modality != "deaminase"
+            and "CpG_to_other_C_mod_ratio" in filtered.obs.columns
+        ):
             filtered = filtered[filtered.obs["CpG_to_other_C_mod_ratio"].astype(float) > 1]
         if lo is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_CpG_site_modified"].astype(float) >= lo]
-            print(f"Removed {s0 - filtered.n_obs} reads below min CpG fraction {lo}")
+            logger.info("Removed %s reads below min CpG fraction %s", s0 - filtered.n_obs, lo)
         if hi is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_CpG_site_modified"].astype(float) <= hi]
-            print(f"Removed {s0 - filtered.n_obs} reads above max CpG fraction {hi}")
-        if (min_valid_fraction_positions_in_read_vs_ref is not None) and ("Valid_CpG_site_in_read_vs_reference" in filtered.obs.columns):
+            logger.info("Removed %s reads above max CpG fraction %s", s0 - filtered.n_obs, hi)
+        if (min_valid_fraction_positions_in_read_vs_ref is not None) and (
+            "Valid_CpG_site_in_read_vs_reference" in filtered.obs.columns
+        ):
             s0 = filtered.n_obs
-            filtered = filtered[filtered.obs["Valid_CpG_site_in_read_vs_reference"].astype(float) >= float(min_valid_fraction_positions_in_read_vs_ref)]
-            print(f"Removed {s0 - filtered.n_obs} reads with insufficient valid CpG site fraction vs ref")
+            filtered = filtered[
+                filtered.obs["Valid_CpG_site_in_read_vs_reference"].astype(float)
+                >= float(min_valid_fraction_positions_in_read_vs_ref)
+            ]
+            logger.info(
+                "Removed %s reads with insufficient valid CpG site fraction vs ref",
+                s0 - filtered.n_obs,
+            )
 
     # any C thresholds
-    if any_c_thresholds and 'C' in mod_target_bases:
+    if any_c_thresholds and "C" in mod_target_bases:
         lo, hi = _unpack_minmax(any_c_thresholds)
         if lo is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_C_site_modified"].astype(float) >= lo]
-            print(f"Removed {s0 - filtered.n_obs} reads below min any-C fraction {lo}")
+            logger.info(
+                "Removed %s reads below min any-C fraction %s",
+                s0 - filtered.n_obs,
+                lo,
+            )
         if hi is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_C_site_modified"].astype(float) <= hi]
-            print(f"Removed {s0 - filtered.n_obs} reads above max any-C fraction {hi}")
-        if (min_valid_fraction_positions_in_read_vs_ref is not None) and ("Valid_C_site_in_read_vs_reference" in filtered.obs.columns):
+            logger.info(
+                "Removed %s reads above max any-C fraction %s",
+                s0 - filtered.n_obs,
+                hi,
+            )
+        if (min_valid_fraction_positions_in_read_vs_ref is not None) and (
+            "Valid_C_site_in_read_vs_reference" in filtered.obs.columns
+        ):
             s0 = filtered.n_obs
-            filtered = filtered[filtered.obs["Valid_C_site_in_read_vs_reference"].astype(float) >= float(min_valid_fraction_positions_in_read_vs_ref)]
-            print(f"Removed {s0 - filtered.n_obs} reads with insufficient valid any-C site fraction vs ref")
+            filtered = filtered[
+                filtered.obs["Valid_C_site_in_read_vs_reference"].astype(float)
+                >= float(min_valid_fraction_positions_in_read_vs_ref)
+            ]
+            logger.info(
+                "Removed %s reads with insufficient valid any-C site fraction vs ref",
+                s0 - filtered.n_obs,
+            )
 
     # A thresholds
-    if a_thresholds and 'A' in mod_target_bases:
+    if a_thresholds and "A" in mod_target_bases:
         lo, hi = _unpack_minmax(a_thresholds)
         if lo is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_A_site_modified"].astype(float) >= lo]
-            print(f"Removed {s0 - filtered.n_obs} reads below min A fraction {lo}")
+            logger.info("Removed %s reads below min A fraction %s", s0 - filtered.n_obs, lo)
         if hi is not None:
             s0 = filtered.n_obs
             filtered = filtered[filtered.obs["Fraction_A_site_modified"].astype(float) <= hi]
-            print(f"Removed {s0 - filtered.n_obs} reads above max A fraction {hi}")
-        if (min_valid_fraction_positions_in_read_vs_ref is not None) and ("Valid_A_site_in_read_vs_reference" in filtered.obs.columns):
+            logger.info("Removed %s reads above max A fraction %s", s0 - filtered.n_obs, hi)
+        if (min_valid_fraction_positions_in_read_vs_ref is not None) and (
+            "Valid_A_site_in_read_vs_reference" in filtered.obs.columns
+        ):
             s0 = filtered.n_obs
-            filtered = filtered[filtered.obs["Valid_A_site_in_read_vs_reference"].astype(float) >= float(min_valid_fraction_positions_in_read_vs_ref)]
-            print(f"Removed {s0 - filtered.n_obs} reads with insufficient valid A site fraction vs ref")
+            filtered = filtered[
+                filtered.obs["Valid_A_site_in_read_vs_reference"].astype(float)
+                >= float(min_valid_fraction_positions_in_read_vs_ref)
+            ]
+            logger.info(
+                "Removed %s reads with insufficient valid A site fraction vs ref",
+                s0 - filtered.n_obs,
+            )
 
     filtered = filtered.copy()
 
     # mark as done
     filtered.uns[uns_flag] = True
 
-    return filtered
+    return filtered