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

smftools/informatics/h5ad_functions.py

@@ -10,6 +10,7 @@ import numpy as np
 import pandas as pd
 import scipy.sparse as sp
 
+from smftools.constants import BASE_QUALITY_SCORES, READ_SPAN_MASK, REFERENCE_STRAND
 from smftools.logging_utils import get_logger
 from smftools.optional_imports import require
 
@@ -84,6 +85,194 @@ def add_demux_type_annotation(
     return adata
 
 
+def append_reference_strand_quality_stats(
+    adata,
+    ref_column: str = REFERENCE_STRAND,
+    quality_layer: str = BASE_QUALITY_SCORES,
+    read_span_layer: str = READ_SPAN_MASK,
+    uns_flag: str = "append_reference_strand_quality_stats_performed",
+    force_redo: bool = False,
+    bypass: bool = False,
+) -> None:
+    """Append per-position quality and error rate stats for each reference strand.
+
+    Args:
+        adata: AnnData object to annotate in-place.
+        ref_column: Obs column defining reference strand groups.
+        quality_layer: Layer containing base quality scores.
+        read_span_layer: Optional layer marking covered positions (1=covered, 0=not covered).
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+        force_redo: Whether to rerun even if ``uns_flag`` is set.
+        bypass: Whether to skip this step.
+    """
+    if bypass:
+        return
+
+    already = bool(adata.uns.get(uns_flag, False))
+    if already and not force_redo:
+        return
+
+    if ref_column not in adata.obs:
+        logger.debug("Reference column '%s' not found; skipping quality stats.", ref_column)
+        return
+
+    if quality_layer not in adata.layers:
+        logger.debug("Quality layer '%s' not found; skipping quality stats.", quality_layer)
+        return
+
+    ref_values = adata.obs[ref_column]
+    references = (
+        ref_values.cat.categories if hasattr(ref_values, "cat") else pd.Index(pd.unique(ref_values))
+    )
+    n_vars = adata.shape[1]
+    has_span_mask = read_span_layer in adata.layers
+
+    for ref in references:
+        ref_mask = ref_values == ref
+        ref_position_mask = adata.var.get(f"position_in_{ref}")
+        if ref_position_mask is None:
+            ref_position_mask = pd.Series(np.ones(n_vars, dtype=bool), index=adata.var.index)
+        else:
+            ref_position_mask = ref_position_mask.astype(bool)
+
+        mean_quality = np.full(n_vars, np.nan, dtype=float)
+        std_quality = np.full(n_vars, np.nan, dtype=float)
+        mean_error = np.full(n_vars, np.nan, dtype=float)
+        std_error = np.full(n_vars, np.nan, dtype=float)
+
+        if ref_mask.sum() > 0:
+            quality_matrix = np.asarray(adata.layers[quality_layer][ref_mask]).astype(float)
+            quality_matrix[quality_matrix < 0] = np.nan
+            if has_span_mask:
+                coverage_mask = np.asarray(adata.layers[read_span_layer][ref_mask]) > 0
+                quality_matrix = np.where(coverage_mask, quality_matrix, np.nan)
+
+            mean_quality = np.nanmean(quality_matrix, axis=0)
+            std_quality = np.nanstd(quality_matrix, axis=0)
+
+            error_matrix = np.power(10.0, -quality_matrix / 10.0)
+            mean_error = np.nanmean(error_matrix, axis=0)
+            std_error = np.nanstd(error_matrix, axis=0)
+
+        mean_quality = np.where(ref_position_mask.values, mean_quality, np.nan)
+        std_quality = np.where(ref_position_mask.values, std_quality, np.nan)
+        mean_error = np.where(ref_position_mask.values, mean_error, np.nan)
+        std_error = np.where(ref_position_mask.values, std_error, np.nan)
+
+        adata.var[f"{ref}_mean_base_quality"] = pd.Series(mean_quality, index=adata.var.index)
+        adata.var[f"{ref}_std_base_quality"] = pd.Series(std_quality, index=adata.var.index)
+        adata.var[f"{ref}_mean_error_rate"] = pd.Series(mean_error, index=adata.var.index)
+        adata.var[f"{ref}_std_error_rate"] = pd.Series(std_error, index=adata.var.index)
+
+    adata.uns[uns_flag] = True
+
+
+def add_read_tag_annotations(
+    adata,
+    bam_files: Optional[List[str]] = None,
+    read_tags: Optional[Dict[str, Dict[str, object]]] = None,
+    tag_names: Optional[List[str]] = None,
+    include_flags: bool = True,
+    include_cigar: bool = True,
+    extract_read_tags_from_bam_callable=None,
+    samtools_backend: str | None = "auto",
+):
+    """Populate adata.obs with read tag metadata.
+
+    Args:
+        adata: AnnData to annotate (modified in-place).
+        bam_files: Optional list of BAM files to extract tags from.
+        read_tags: Optional mapping of read name to tag dict.
+        tag_names: Optional list of BAM tag names to extract (e.g. ["NM", "MD", "MM", "ML"]).
+        include_flags: Whether to add a FLAGS list column.
+        include_cigar: Whether to add the CIGAR string column.
+        extract_read_tags_from_bam_callable: Optional callable to extract tags from a BAM.
+        samtools_backend: Backend selection for samtools-compatible operations (auto|python|cli).
+
+    Returns:
+        None (mutates adata in-place).
+    """
+    if read_tags is None:
+        read_tags = {}
+        if bam_files:
+            extractor = extract_read_tags_from_bam_callable or globals().get(
+                "extract_read_tags_from_bam"
+            )
+            if extractor is None:
+                raise ValueError(
+                    "No `read_tags` provided and `extract_read_tags_from_bam` not found."
+                )
+            for bam in bam_files:
+                bam_read_tags = extractor(
+                    bam,
+                    tag_names=tag_names,
+                    include_flags=include_flags,
+                    include_cigar=include_cigar,
+                    samtools_backend=samtools_backend,
+                )
+                if not isinstance(bam_read_tags, dict):
+                    raise ValueError(f"extract_read_tags_from_bam returned non-dict for {bam}")
+                read_tags.update(bam_read_tags)
+
+    if not read_tags:
+        return
+
+    df = pd.DataFrame.from_dict(read_tags, orient="index")
+    df_reindexed = df.reindex(adata.obs_names)
+    for column in df_reindexed.columns:
+        adata.obs[column] = df_reindexed[column].values
+
+
+def add_secondary_supplementary_alignment_flags(
+    adata,
+    bam_path: str | Path,
+    *,
+    uns_flag: str = "add_secondary_supplementary_flags_performed",
+    bypass: bool = False,
+    force_redo: bool = False,
+    samtools_backend: str | None = "auto",
+) -> None:
+    """Annotate whether reads have secondary/supplementary alignments.
+
+    Args:
+        adata: AnnData to annotate (modified in-place).
+        bam_path: Path to the aligned/sorted BAM to scan.
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+        bypass: Whether to skip annotation.
+        force_redo: Whether to recompute even if ``uns_flag`` is set.
+        samtools_backend: Backend selection for samtools-compatible operations (auto|python|cli).
+    """
+    already = bool(adata.uns.get(uns_flag, False))
+    if (already and not force_redo) or bypass:
+        return
+
+    from .bam_functions import (
+        extract_secondary_supplementary_alignment_spans,
+        find_secondary_supplementary_read_names,
+    )
+
+    secondary_reads, supplementary_reads = find_secondary_supplementary_read_names(
+        bam_path,
+        adata.obs_names,
+        samtools_backend=samtools_backend,
+    )
+    secondary_spans, supplementary_spans = extract_secondary_supplementary_alignment_spans(
+        bam_path,
+        adata.obs_names,
+        samtools_backend=samtools_backend,
+    )
+
+    adata.obs["has_secondary_alignment"] = adata.obs_names.isin(secondary_reads)
+    adata.obs["has_supplementary_alignment"] = adata.obs_names.isin(supplementary_reads)
+    adata.obs["secondary_alignment_spans"] = [
+        secondary_spans.get(read_name) for read_name in adata.obs_names
+    ]
+    adata.obs["supplementary_alignment_spans"] = [
+        supplementary_spans.get(read_name) for read_name in adata.obs_names
+    ]
+    adata.uns[uns_flag] = True
+
+
 def add_read_length_and_mapping_qc(
     adata,
     bam_files: Optional[List[str]] = None,
@@ -104,7 +293,8 @@ def add_read_length_and_mapping_qc(
     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]
+        Optional dict mapping obs_name -> [read_length, read_quality, reference_length, mapped_length,
+        mapping_quality, reference_start, reference_end]
         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).
@@ -154,10 +344,12 @@ def add_read_length_and_mapping_qc(
         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)
+        adata.obs["reference_start"] = np.full(n, np.nan)
+        adata.obs["reference_end"] = np.full(n, np.nan)
     else:
         # Build DF robustly
         # Convert values to lists where possible, else to [val, val, val...]
-        max_cols = 5
+        max_cols = 7
         rows = {}
         for k, v in read_metrics.items():
             if isinstance(v, (list, tuple, np.ndarray)):
@@ -179,6 +371,8 @@ def add_read_length_and_mapping_qc(
                 "reference_length",
                 "mapped_length",
                 "mapping_quality",
+                "reference_start",
+                "reference_end",
             ],
         )
 
@@ -191,6 +385,8 @@ def add_read_length_and_mapping_qc(
         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
+        adata.obs["reference_start"] = df_reindexed["reference_start"].values
+        adata.obs["reference_end"] = df_reindexed["reference_end"].values
 
     # Compute ratio columns safely (avoid divide-by-zero and preserve NaN)
     # read_length_to_reference_length_ratio