smftools 0.1.6__py3-none-any.whl → 0.2.1__py3-none-any.whl

smftools/plotting/qc_plotting.py

@@ -0,0 +1,270 @@
+import os
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+import os
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+
+def plot_read_qc_histograms(
+    adata,
+    outdir,
+    obs_keys,
+    sample_key,
+    bins=60,
+    clip_quantiles=(0.0, 0.995),
+    min_non_nan=10,
+    rows_per_fig=6,
+    topn_categories=15,
+    figsize_cell=(3.6, 2.6),
+    dpi=150,
+):
+    """
+    Plot a grid of QC histograms: rows = samples (from `sample_key`), columns = `obs_keys`.
+
+    Numeric columns -> histogram per sample.
+    Categorical columns -> bar chart of top categories per sample.
+
+    Saves paginated PNGs to `outdir`.
+
+    Parameters
+    ----------
+    adata : AnnData
+    outdir : str
+    obs_keys : list[str]
+    sample_key : str
+        Column in adata.obs defining rows (samples/barcodes).
+    bins : int
+        Histogram bins for numeric metrics.
+    clip_quantiles : tuple or None
+        Clip numeric data globally per metric for consistent axes, e.g. (0.0, 0.995).
+    min_non_nan : int
+        Minimum finite values to plot a panel.
+    rows_per_fig : int
+        Number of samples per page.
+    topn_categories : int
+        For categorical metrics, show top-N categories (per sample).
+    figsize_cell : (float, float)
+        Size of each subplot cell (width, height).
+    dpi : int
+        Figure resolution.
+    """
+    os.makedirs(outdir, exist_ok=True)
+
+    if sample_key not in adata.obs.columns:
+        raise KeyError(f"'{sample_key}' not found in adata.obs")
+
+    # Ensure sample_key is categorical for stable ordering
+    samples = adata.obs[sample_key]
+    if not pd.api.types.is_categorical_dtype(samples):
+        samples = samples.astype("category")
+    sample_levels = list(samples.cat.categories)
+
+    # Validate keys, and classify numeric vs categorical
+    valid_keys = []
+    is_numeric = {}
+    for key in obs_keys:
+        if key not in adata.obs.columns:
+            print(f"[WARN] '{key}' not found in obs; skipping.")
+            continue
+        s = adata.obs[key]
+        num = pd.api.types.is_numeric_dtype(s)
+        valid_keys.append(key)
+        is_numeric[key] = num
+    if not valid_keys:
+        print("[plot_read_qc_grid] No valid obs_keys to plot.")
+        return
+
+    # Precompute global numeric ranges (after clipping) so rows share x-axis per column
+    global_ranges = {}
+    for key in valid_keys:
+        if not is_numeric[key]:
+            continue
+        s = pd.to_numeric(adata.obs[key], errors="coerce").replace([np.inf, -np.inf], np.nan).dropna()
+        if s.size < min_non_nan:
+            # still set something to avoid errors; just use min/max or (0,1)
+            lo, hi = (0.0, 1.0) if s.size == 0 else (float(s.min()), float(s.max()))
+        else:
+            if clip_quantiles:
+                qlo = s.quantile(clip_quantiles[0]) if clip_quantiles[0] is not None else s.min()
+                qhi = s.quantile(clip_quantiles[1]) if clip_quantiles[1] is not None else s.max()
+                lo, hi = float(qlo), float(qhi)
+                if not (np.isfinite(lo) and np.isfinite(hi) and hi > lo):
+                    lo, hi = float(s.min()), float(s.max())
+            else:
+                lo, hi = float(s.min()), float(s.max())
+        global_ranges[key] = (lo, hi)
+
+    def _sanitize(name: str) -> str:
+        return "".join(c if c.isalnum() or c in "-._" else "_" for c in str(name))
+
+    ncols = len(valid_keys)
+    fig_w = figsize_cell[0] * ncols
+    # rows per page is rows_per_fig; figure height scales accordingly
+    fig_h_unit = figsize_cell[1]
+
+    for start in range(0, len(sample_levels), rows_per_fig):
+        chunk = sample_levels[start:start + rows_per_fig]
+        nrows = len(chunk)
+        fig, axes = plt.subplots(
+            nrows=nrows, ncols=ncols,
+            figsize=(fig_w, fig_h_unit * nrows),
+            dpi=dpi,
+            squeeze=False,
+        )
+
+        for r, sample_val in enumerate(chunk):
+            row_mask = (adata.obs[sample_key].values == sample_val)
+            n_in_row = int(row_mask.sum())
+
+            for c, key in enumerate(valid_keys):
+                ax = axes[r, c]
+                series = adata.obs.loc[row_mask, key]
+
+                if is_numeric[key]:
+                    x = pd.to_numeric(series, errors="coerce").replace([np.inf, -np.inf], np.nan).dropna()
+                    if x.size < min_non_nan:
+                        ax.text(0.5, 0.5, f"n={x.size} (<{min_non_nan})", ha="center", va="center")
+                    else:
+                        # clip to global range for consistent axes
+                        lo, hi = global_ranges[key]
+                        x = x.clip(lo, hi)
+                        ax.hist(x.values, bins=bins, range=(lo, hi), edgecolor="black", alpha=0.7)
+                        ax.set_xlim(lo, hi)
+                    if r == 0:
+                        ax.set_title(key)
+                    if c == 0:
+                        ax.set_ylabel(f"{sample_val}\n(n={n_in_row})")
+                    ax.grid(alpha=0.25)
+                    ax.set_xlabel("")  # keep uncluttered; x-limit conveys scale
+                else:
+                    vc = series.astype("category").value_counts(dropna=False)
+                    if vc.sum() < min_non_nan:
+                        ax.text(0.5, 0.5, f"n={vc.sum()} (<{min_non_nan})", ha="center", va="center")
+                    else:
+                        vc_top = vc.iloc[:topn_categories][::-1]  # show top-N, reversed for barh
+                        ax.barh(vc_top.index.astype(str), vc_top.values)
+                        ax.invert_yaxis()
+                    if r == 0:
+                        ax.set_title(f"{key} (cat)")
+                    if c == 0:
+                        ax.set_ylabel(f"{sample_val}\n(n={n_in_row})")
+                    ax.grid(alpha=0.25)
+                    # trim labels to reduce clutter
+                    if vc.sum() >= min_non_nan:
+                        ax.tick_params(axis="y", labelsize=8)
+
+        plt.tight_layout()
+        page = start // rows_per_fig + 1
+        out_png = os.path.join(outdir, f"qc_grid_{_sanitize(sample_key)}_page{page}.png")
+        plt.savefig(out_png, bbox_inches="tight")
+        plt.close(fig)
+
+
+# def plot_read_qc_histograms(
+#     adata,
+#     outdir,
+#     obs_keys,
+#     sample_key=None,
+#     *,
+#     bins=100,
+#     clip_quantiles=(0.0, 0.995),
+#     min_non_nan=10,
+#     figsize=(6, 4),
+#     dpi=150
+# ):
+#     """
+#     Plots histograms for given obs_keys, optionally grouped by sample_key.
+
+#     Parameters
+#     ----------
+#     adata : AnnData
+#         AnnData object.
+#     outdir : str
+#         Output directory for PNG files.
+#     obs_keys : list[str]
+#         List of obs columns to plot.
+#     sample_key : str or None
+#         Column in adata.obs to group by (e.g., 'Barcode').
+#         If None, plots are for the full dataset only.
+#     bins : int
+#         Number of histogram bins for numeric data.
+#     clip_quantiles : tuple or None
+#         (low_q, high_q) to clip extreme values for plotting.
+#     min_non_nan : int
+#         Minimum number of finite values to plot.
+#     figsize : tuple
+#         Figure size.
+#     dpi : int
+#         Figure resolution.
+#     """
+#     os.makedirs(outdir, exist_ok=True)
+
+#     # Define grouping
+#     if sample_key and sample_key in adata.obs.columns:
+#         groups = adata.obs.groupby(sample_key)
+#     else:
+#         groups = [(None, adata.obs)]  # single group
+
+#     for group_name, group_df in groups:
+#         # For each metric
+#         for key in obs_keys:
+#             if key not in group_df.columns:
+#                 print(f"[WARN] '{key}' not found in obs; skipping.")
+#                 continue
+
+#             series = group_df[key]
+
+#             # Numeric columns
+#             if pd.api.types.is_numeric_dtype(series):
+#                 x = pd.to_numeric(series, errors="coerce").replace([np.inf, -np.inf], np.nan).dropna()
+#                 if len(x) < min_non_nan:
+#                     continue
+
+#                 # Clip for better visualization
+#                 if clip_quantiles:
+#                     lo = x.quantile(clip_quantiles[0]) if clip_quantiles[0] is not None else x.min()
+#                     hi = x.quantile(clip_quantiles[1]) if clip_quantiles[1] is not None else x.max()
+#                     if np.isfinite(lo) and np.isfinite(hi) and hi > lo:
+#                         x = x.clip(lo, hi)
+
+#                 fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
+#                 ax.hist(x, bins=bins, edgecolor="black", alpha=0.7)
+#                 ax.set_xlabel(key)
+#                 ax.set_ylabel("Count")
+
+#                 title = f"{key}" if group_name is None else f"{key} — {sample_key}={group_name}"
+#                 ax.set_title(title)
+
+#                 plt.tight_layout()
+
+#                 # Save PNG
+#                 safe_group = "all" if group_name is None else str(group_name)
+#                 fname = f"{key}_{sample_key}_{safe_group}.png" if sample_key else f"{key}.png"
+#                 fname = fname.replace("/", "_")
+#                 fig.savefig(os.path.join(outdir, fname))
+#                 plt.close(fig)
+
+#             else:
+#                 # Categorical columns
+#                 vc = series.astype("category").value_counts(dropna=False)
+#                 if vc.sum() < min_non_nan:
+#                     continue
+
+#                 fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
+#                 vc.plot(kind="barh", ax=ax)
+#                 ax.set_xlabel("Count")
+
+#                 title = f"{key} (categorical)" if group_name is None else f"{key} — {sample_key}={group_name}"
+#                 ax.set_title(title)
+
+#                 plt.tight_layout()
+
+#                 safe_group = "all" if group_name is None else str(group_name)
+#                 fname = f"{key}_{sample_key}_{safe_group}.png" if sample_key else f"{key}.png"
+#                 fname = fname.replace("/", "_")
+#                 fig.savefig(os.path.join(outdir, fname))
+#                 plt.close(fig)

smftools/preprocessing/__init__.py

@@ -0,0 +1,38 @@
+from .add_read_length_and_mapping_qc import add_read_length_and_mapping_qc
+from .append_base_context import append_base_context
+from .append_binary_layer_by_base_context import append_binary_layer_by_base_context
+from .binarize_on_Youden import binarize_on_Youden
+from .calculate_complexity import calculate_complexity
+from .calculate_complexity_II import calculate_complexity_II
+from .calculate_read_modification_stats import calculate_read_modification_stats
+from .calculate_coverage import calculate_coverage
+from .calculate_position_Youden import calculate_position_Youden
+from .calculate_read_length_stats import calculate_read_length_stats
+from .clean_NaN import clean_NaN
+from .filter_adata_by_nan_proportion import filter_adata_by_nan_proportion
+from .filter_reads_on_modification_thresholds import filter_reads_on_modification_thresholds
+from .filter_reads_on_length_quality_mapping import filter_reads_on_length_quality_mapping
+from .invert_adata import invert_adata
+from .load_sample_sheet import load_sample_sheet
+from .flag_duplicate_reads import flag_duplicate_reads
+from .subsample_adata import subsample_adata
+
+__all__ = [
+    "add_read_length_and_mapping_qc",
+    "append_base_context",
+    "append_binary_layer_by_base_context",
+    "binarize_on_Youden",
+    "calculate_complexity",
+    "calculate_read_modification_stats",
+    "calculate_coverage",    
+    "calculate_position_Youden",
+    "calculate_read_length_stats",
+    "clean_NaN",   
+    "filter_adata_by_nan_proportion",
+    "filter_reads_on_modification_thresholds",
+    "filter_reads_on_length_quality_mapping",
+    "invert_adata",
+    "load_sample_sheet",
+    "flag_duplicate_reads",
+    "subsample_adata"
+]

smftools/preprocessing/add_read_length_and_mapping_qc.py

@@ -0,0 +1,129 @@
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from typing import Optional, List, Dict, Union
+
+def add_read_length_and_mapping_qc(
+    adata,
+    bam_files: Optional[List[str]] = None,
+    read_metrics: Optional[Dict[str, Union[list, tuple]]] = None,
+    uns_flag: str = "read_lenth_and_mapping_qc_performed",
+    extract_read_features_from_bam_callable = None,
+    bypass: bool = False,
+    force_redo: bool = True
+):
+    """
+    Populate adata.obs with read/mapping QC columns.
+
+    Parameters
+    ----------
+    adata
+        AnnData to annotate (modified in-place).
+    bam_files
+        Optional list of BAM files to extract metrics from. Ignored if read_metrics supplied.
+    read_metrics
+        Optional dict mapping obs_name -> [read_length, read_quality, reference_length, mapped_length, mapping_quality]
+        If provided, this will be used directly and bam_files will be ignored.
+    uns_flag
+        key in final_adata.uns used to record that QC was performed (kept the name with original misspelling).
+    extract_read_features_from_bam_callable
+        Optional callable(bam_path) -> dict mapping read_name -> list/tuple of metrics.
+        If not provided and bam_files is given, function will attempt to call `extract_read_features_from_bam`
+        from the global namespace (your existing helper).
+    Returns
+    -------
+    None (mutates final_adata in-place)
+    """
+
+    # Only run if not already performed
+    already = bool(adata.uns.get(uns_flag, False))
+    if (already and not force_redo) or bypass:
+        # QC already performed; nothing to do
+        return
+
+    # Build read_metrics dict either from provided arg or by extracting from bam files
+    if read_metrics is None:
+        read_metrics = {}
+        if bam_files:
+            extractor = extract_read_features_from_bam_callable or globals().get("extract_read_features_from_bam")
+            if extractor is None:
+                raise ValueError("No `read_metrics` provided and `extract_read_features_from_bam` not found.")
+            for bam in bam_files:
+                bam_read_metrics = extractor(bam)
+                if not isinstance(bam_read_metrics, dict):
+                    raise ValueError(f"extract_read_features_from_bam returned non-dict for {bam}")
+                read_metrics.update(bam_read_metrics)
+        else:
+            # nothing to do
+            read_metrics = {}
+
+    # Convert read_metrics dict -> DataFrame (rows = read id)
+    # Values may be lists/tuples or scalars; prefer lists/tuples with 5 entries.
+    if len(read_metrics) == 0:
+        # fill with NaNs
+        n = adata.n_obs
+        adata.obs['read_length'] = np.full(n, np.nan)
+        adata.obs['mapped_length'] = np.full(n, np.nan)
+        adata.obs['reference_length'] = np.full(n, np.nan)
+        adata.obs['read_quality'] = np.full(n, np.nan)
+        adata.obs['mapping_quality'] = np.full(n, np.nan)
+    else:
+        # Build DF robustly
+        # Convert values to lists where possible, else to [val, val, val...]
+        max_cols = 5
+        rows = {}
+        for k, v in read_metrics.items():
+            if isinstance(v, (list, tuple, np.ndarray)):
+                vals = list(v)
+            else:
+                # scalar -> replicate into 5 columns to preserve original behavior
+                vals = [v] * max_cols
+            # Ensure length >= 5
+            if len(vals) < max_cols:
+                vals = vals + [np.nan] * (max_cols - len(vals))
+            rows[k] = vals[:max_cols]
+
+        df = pd.DataFrame.from_dict(rows, orient='index', columns=[
+            'read_length', 'read_quality', 'reference_length', 'mapped_length', 'mapping_quality'
+        ])
+
+        # Reindex to final_adata.obs_names so order matches adata
+        # If obs_names are not present as keys in df, the results will be NaN
+        df_reindexed = df.reindex(adata.obs_names).astype(float)
+
+        adata.obs['read_length'] = df_reindexed['read_length'].values
+        adata.obs['mapped_length'] = df_reindexed['mapped_length'].values
+        adata.obs['reference_length'] = df_reindexed['reference_length'].values
+        adata.obs['read_quality'] = df_reindexed['read_quality'].values
+        adata.obs['mapping_quality'] = df_reindexed['mapping_quality'].values
+
+    # Compute ratio columns safely (avoid divide-by-zero and preserve NaN)
+    # read_length_to_reference_length_ratio
+    rl = pd.to_numeric(adata.obs['read_length'], errors='coerce').to_numpy(dtype=float)
+    ref_len = pd.to_numeric(adata.obs['reference_length'], errors='coerce').to_numpy(dtype=float)
+    mapped_len = pd.to_numeric(adata.obs['mapped_length'], errors='coerce').to_numpy(dtype=float)
+
+    # safe divisions: use np.where to avoid warnings and replace inf with nan
+    with np.errstate(divide='ignore', invalid='ignore'):
+        rl_to_ref = np.where((ref_len != 0) & np.isfinite(ref_len), rl / ref_len, np.nan)
+        mapped_to_ref = np.where((ref_len != 0) & np.isfinite(ref_len), mapped_len / ref_len, np.nan)
+        mapped_to_read = np.where((rl != 0) & np.isfinite(rl), mapped_len / rl, np.nan)
+
+    adata.obs['read_length_to_reference_length_ratio'] = rl_to_ref
+    adata.obs['mapped_length_to_reference_length_ratio'] = mapped_to_ref
+    adata.obs['mapped_length_to_read_length_ratio'] = mapped_to_read
+
+    # Add read level raw modification signal: sum over X rows
+    X = adata.X
+    if sp.issparse(X):
+        # sum returns (n_obs, 1) sparse matrix; convert to 1d array
+        raw_sig = np.asarray(X.sum(axis=1)).ravel()
+    else:
+        raw_sig = np.asarray(X.sum(axis=1)).ravel()
+
+    adata.obs['Raw_modification_signal'] = raw_sig
+
+    # mark as done
+    adata.uns[uns_flag] = True
+
+    return None

smftools/preprocessing/append_base_context.py

@@ -0,0 +1,122 @@
+def append_base_context(adata, 
+                        obs_column='Reference_strand', 
+                        use_consensus=False,
+                        native=False, 
+                        mod_target_bases=['GpC', 'CpG'],
+                        bypass=False,
+                        force_redo=False,
+                        uns_flag='base_context_added'
+):
+    """
+    Adds nucleobase context to the position within the given category. When use_consensus is True, it uses the consensus sequence, otherwise it defaults to the FASTA sequence.
+
+    Parameters:
+        adata (AnnData): The input adata object.
+        obs_column (str): The observation column in which to stratify on. Default is 'Reference_strand', which should not be changed for most purposes.
+        use_consensus (bool): A truth statement indicating whether to use the consensus sequence from the reads mapped to the reference. If False, the reference FASTA is used instead.
+        native (bool): If False, perform conversion SMF assumptions. If True, perform native SMF assumptions
+        mod_target_bases (list): Base contexts that may be modified.
+    
+    Returns:
+        None
+    """
+    import numpy as np
+    import anndata as ad
+
+    # Only run if not already performed
+    already = bool(adata.uns.get(uns_flag, False))
+    if (already and not force_redo) or bypass:
+        # QC already performed; nothing to do
+        return
+
+    print('Adding base context based on reference FASTA sequence for sample')
+    categories = adata.obs[obs_column].cat.categories
+    site_types = []
+    
+    if any(base in mod_target_bases for base in ['GpC', 'CpG', 'C']):
+        site_types += ['GpC_site', 'CpG_site', 'ambiguous_GpC_CpG_site', 'other_C_site', 'any_C_site']
+    
+    if 'A' in mod_target_bases:
+        site_types += ['A_site']
+
+    for cat in categories:
+        # Assess if the strand is the top or bottom strand converted
+        if 'top' in cat:
+            strand = 'top'
+        elif 'bottom' in cat:
+            strand = 'bottom'
+
+        if native:
+            basename = cat.split(f"_{strand}")[0]
+            if use_consensus:
+                sequence = adata.uns[f'{basename}_consensus_sequence']
+            else:
+                # This sequence is the unconverted FASTA sequence of the original input FASTA for the locus
+                sequence = adata.uns[f'{basename}_FASTA_sequence']
+        else:
+            basename = cat.split(f"_{strand}")[0]
+            if use_consensus:
+                sequence = adata.uns[f'{basename}_consensus_sequence']
+            else:
+                # This sequence is the unconverted FASTA sequence of the original input FASTA for the locus
+                sequence = adata.uns[f'{basename}_FASTA_sequence']
+        # Init a dict keyed by reference site type that points to a bool of whether the position is that site type.    
+        boolean_dict = {}
+        for site_type in site_types:
+            boolean_dict[f'{cat}_{site_type}'] = np.full(len(sequence), False, dtype=bool)
+
+        if any(base in mod_target_bases for base in ['GpC', 'CpG', 'C']):
+            if strand == 'top':
+                # Iterate through the sequence and apply the criteria
+                for i in range(1, len(sequence) - 1):
+                    if sequence[i] == 'C':
+                        boolean_dict[f'{cat}_any_C_site'][i] = True
+                        if sequence[i - 1] == 'G' and sequence[i + 1] != 'G':
+                            boolean_dict[f'{cat}_GpC_site'][i] = True
+                        elif sequence[i - 1] == 'G' and sequence[i + 1] == 'G':
+                            boolean_dict[f'{cat}_ambiguous_GpC_CpG_site'][i] = True
+                        elif sequence[i - 1] != 'G' and sequence[i + 1] == 'G':
+                            boolean_dict[f'{cat}_CpG_site'][i] = True
+                        elif sequence[i - 1] != 'G' and sequence[i + 1] != 'G':
+                            boolean_dict[f'{cat}_other_C_site'][i] = True
+            elif strand == 'bottom':
+                # Iterate through the sequence and apply the criteria
+                for i in range(1, len(sequence) - 1):
+                    if sequence[i] == 'G':
+                        boolean_dict[f'{cat}_any_C_site'][i] = True
+                        if sequence[i + 1] == 'C' and sequence[i - 1] != 'C':
+                            boolean_dict[f'{cat}_GpC_site'][i] = True
+                        elif sequence[i - 1] == 'C' and sequence[i + 1] == 'C':
+                            boolean_dict[f'{cat}_ambiguous_GpC_CpG_site'][i] = True
+                        elif sequence[i - 1] == 'C' and sequence[i + 1] != 'C':
+                            boolean_dict[f'{cat}_CpG_site'][i] = True
+                        elif sequence[i - 1] != 'C' and sequence[i + 1] != 'C':
+                            boolean_dict[f'{cat}_other_C_site'][i] = True
+            else:
+                print('Error: top or bottom strand of conversion could not be determined. Ensure this value is in the Reference name.')
+
+        if 'A' in mod_target_bases:
+            if strand == 'top':
+                # Iterate through the sequence and apply the criteria
+                for i in range(1, len(sequence) - 1):
+                    if sequence[i] == 'A':
+                        boolean_dict[f'{cat}_A_site'][i] = True
+            elif strand == 'bottom':
+                # Iterate through the sequence and apply the criteria
+                for i in range(1, len(sequence) - 1):
+                    if sequence[i] == 'T':
+                        boolean_dict[f'{cat}_A_site'][i] = True
+            else:
+                print('Error: top or bottom strand of conversion could not be determined. Ensure this value is in the Reference name.')
+
+        for site_type in site_types:
+            adata.var[f'{cat}_{site_type}'] = boolean_dict[f'{cat}_{site_type}'].astype(bool)
+            if native:
+                adata.obsm[f'{cat}_{site_type}'] = adata[:, adata.var[f'{cat}_{site_type}'] == True].layers['binarized_methylation']
+            else:
+                adata.obsm[f'{cat}_{site_type}'] = adata[:, adata.var[f'{cat}_{site_type}'] == True].X
+
+    # mark as done
+    adata.uns[uns_flag] = True
+
+    return None