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

smftools/preprocessing/__init__.py

@@ -6,6 +6,9 @@ _LAZY_ATTRS = {
     "append_base_context": "smftools.preprocessing.append_base_context",
     "append_binary_layer_by_base_context": "smftools.preprocessing.append_binary_layer_by_base_context",
     "append_mismatch_frequency_sites": "smftools.preprocessing.append_mismatch_frequency_sites",
+    "append_variant_call_layer": "smftools.preprocessing.append_variant_call_layer",
+    "append_variant_segment_layer": "smftools.preprocessing.append_variant_call_layer",
+    "append_sequence_mismatch_annotations": "smftools.preprocessing.append_sequence_mismatch_annotations",
     "binarize_adata": "smftools.preprocessing.binarize",
     "binarize_on_Youden": "smftools.preprocessing.binarize_on_Youden",
     "calculate_complexity_II": "smftools.preprocessing.calculate_complexity_II",

smftools/preprocessing/append_base_context.py

@@ -51,7 +51,7 @@ def append_base_context(
         site_types += ["A_site"]
 
     for ref in references:
-        # Assess if the strand is the top or bottom strand converted
+        # Assess if the modified strand is the top or bottom strand.
         if "top" in ref:
             strand = "top"
         elif "bottom" in ref:

smftools/preprocessing/append_mismatch_frequency_sites.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Iterable, Sequence
 import numpy as np
 import pandas as pd
 
-from smftools.constants import MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT
+from smftools.constants import BASE_QUALITY_SCORES, MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT
 from smftools.logging_utils import get_logger
 
 if TYPE_CHECKING:
@@ -19,6 +19,7 @@ def append_mismatch_frequency_sites(
     ref_column: str = "Reference_strand",
     mismatch_layer: str = "mismatch_integer_encoding",
     read_span_layer: str = "read_span_mask",
+    quality_layer: str | None = None,
     mismatch_frequency_range: Sequence[float] | None = (0.05, 0.95),
     uns_flag: str = "append_mismatch_frequency_sites_performed",
     force_redo: bool = False,
@@ -31,6 +32,7 @@ def append_mismatch_frequency_sites(
         ref_column: Obs column defining reference categories.
         mismatch_layer: Layer containing mismatch integer encodings.
         read_span_layer: Layer containing read span masks (1=covered, 0=not covered).
+        quality_layer: Layer containing base quality scores for Q-value based error rates.
         mismatch_frequency_range: Lower/upper bounds (inclusive) for variable site flagging.
         uns_flag: Flag in ``adata.uns`` indicating prior completion.
         force_redo: Whether to rerun even if ``uns_flag`` is set.
@@ -73,6 +75,16 @@ def append_mismatch_frequency_sites(
             read_span_layer,
         )
 
+    if quality_layer is None:
+        if BASE_QUALITY_SCORES in adata.layers:
+            quality_layer = BASE_QUALITY_SCORES
+        elif "base_qualities" in adata.layers:
+            quality_layer = "base_qualities"
+
+    if quality_layer is not None and quality_layer not in adata.layers:
+        logger.debug("Quality layer '%s' not found; falling back to range flagging.", quality_layer)
+        quality_layer = None
+
     references = adata.obs[ref_column].cat.categories
     n_vars = adata.shape[1]
 
@@ -125,11 +137,28 @@ def append_mismatch_frequency_sites(
         )
         frequency_values = np.where(ref_position_mask.values, frequency_values, np.nan)
 
-        variable_flags = (
-            (frequency_values >= lower_bound)
-            & (frequency_values <= upper_bound)
-            & ref_position_mask.values
-        )
+        mean_error_rate: np.ndarray | None = None
+        if quality_layer is not None:
+            quality_matrix = np.asarray(adata.layers[quality_layer][ref_mask]).astype(float)
+            quality_matrix[quality_matrix < 0] = np.nan
+            if has_span_mask:
+                quality_matrix = np.where(coverage_mask, quality_matrix, np.nan)
+            error_matrix = np.power(10.0, -quality_matrix / 10.0)
+            mean_error_rate = np.nanmean(error_matrix, axis=0)
+            mean_error_rate = np.where(ref_position_mask.values, mean_error_rate, np.nan)
+
+        if mean_error_rate is None:
+            variable_flags = (
+                (frequency_values >= lower_bound)
+                & (frequency_values <= upper_bound)
+                & ref_position_mask.values
+            )
+        else:
+            variable_flags = (
+                (frequency_values > mean_error_rate)
+                & ref_position_mask.values
+                & np.isfinite(mean_error_rate)
+            )
 
         base_counts_by_int: dict[int, np.ndarray] = {}
         for base_int in base_int_to_label:

smftools/preprocessing/append_sequence_mismatch_annotations.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+from smftools.logging_utils import get_logger
+from smftools.tools.sequence_alignment import align_sequences_with_mismatches
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def _format_mismatch_identity(event: str, seq1_base: str | None, seq2_base: str | None) -> str:
+    if event == "substitution":
+        return f"{seq1_base}->{seq2_base}"
+    if event == "insertion":
+        return f"ins:{seq2_base}"
+    return f"del:{seq1_base}"
+
+
+def append_sequence_mismatch_annotations(
+    adata: "ad.AnnData",
+    seq1_column: str,
+    seq2_column: str,
+    output_prefix: str | None = None,
+    match_score: int = 1,
+    mismatch_score: int = -1,
+    gap_score: int = -2,
+    ignore_n: bool = True,
+    bypass: bool = False,
+    force_redo: bool = False,
+    uns_flag: str = "append_sequence_mismatch_annotations_performed",
+) -> None:
+    """Append mismatch annotations by aligning full reference sequences.
+
+    Extracts the full reference sequences from per-position base columns in
+    ``adata.var``, performs a single global alignment, and maps mismatches
+    (substitutions, insertions, deletions) back to ``adata.var`` indices.
+
+    Results stored in ``adata.var``:
+      - ``{prefix}_mismatch_type``: Per-position str — ``"substitution"``,
+        ``"insertion"``, ``"deletion"``, or ``""`` (no mismatch).
+      - ``{prefix}_mismatch_identity``: Per-position str — e.g. ``"A->G"``,
+        ``"ins:T"``, ``"del:C"``, or ``""``).
+      - ``{prefix}_is_mismatch``: Per-position bool flag.
+
+    Args:
+        adata: AnnData object.
+        seq1_column: Column in ``adata.var`` with per-position bases for reference 1.
+        seq2_column: Column in ``adata.var`` with per-position bases for reference 2.
+        output_prefix: Prefix for output columns. Defaults to ``{seq1_column}__{seq2_column}``.
+        match_score: Alignment match score.
+        mismatch_score: Alignment mismatch score.
+        gap_score: Alignment gap score.
+        ignore_n: Whether to ignore mismatches involving ``N`` bases.
+        bypass: Whether to skip processing.
+        force_redo: Whether to rerun even if ``uns_flag`` is set.
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+    """
+    already = bool(adata.uns.get(uns_flag, False))
+    if (already and not force_redo) or bypass:
+        return
+
+    if seq1_column not in adata.var:
+        raise KeyError(f"Sequence column '{seq1_column}' not found in adata.var")
+    if seq2_column not in adata.var:
+        raise KeyError(f"Sequence column '{seq2_column}' not found in adata.var")
+
+    output_prefix = output_prefix or f"{seq1_column}__{seq2_column}"
+
+    seq1_series = adata.var[seq1_column]
+    seq2_series = adata.var[seq2_column]
+    n_vars = adata.shape[1]
+
+    # ---- Build full sequences from positions where each ref has a valid (non-N) base ----
+    valid1_mask = seq1_series.notna() & (seq1_series != "N")
+    valid2_mask = seq2_series.notna() & (seq2_series != "N")
+
+    # var indices (integers) for each valid base
+    var_indices_1 = np.where(valid1_mask.values)[0]
+    var_indices_2 = np.where(valid2_mask.values)[0]
+
+    full_seq1 = "".join(str(seq1_series.iloc[i]) for i in var_indices_1)
+    full_seq2 = "".join(str(seq2_series.iloc[i]) for i in var_indices_2)
+
+    logger.info(
+        "Aligning full sequences: '%s' (%d bases) vs '%s' (%d bases).",
+        seq1_column,
+        len(full_seq1),
+        seq2_column,
+        len(full_seq2),
+    )
+
+    # ---- Global alignment ----
+    aligned_seq1, aligned_seq2, mismatches = align_sequences_with_mismatches(
+        full_seq1,
+        full_seq2,
+        match_score=match_score,
+        mismatch_score=mismatch_score,
+        gap_score=gap_score,
+        ignore_n=ignore_n,
+    )
+
+    logger.info(
+        "Alignment complete. Aligned length: %d, mismatches: %d.",
+        len(aligned_seq1),
+        len(mismatches),
+    )
+
+    # ---- Map alignment mismatches back to var indices ----
+    mismatch_type_arr = [""] * n_vars
+    mismatch_identity_arr = [""] * n_vars
+    is_mismatch_arr = np.zeros(n_vars, dtype=bool)
+
+    # For substitutions, store the paired var indices from both references.
+    # This is needed because indels shift the coordinate systems so that the
+    # same alignment column maps to different var indices in each reference.
+    substitution_map: list[dict] = []
+
+    for mm in mismatches:
+        # Determine which var index this mismatch maps to.
+        # For substitutions and deletions, seq1_pos is defined.
+        # For insertions, only seq2_pos is defined (gap in seq1).
+        if mm.seq1_pos is not None:
+            var_idx = int(var_indices_1[mm.seq1_pos])
+        elif mm.seq2_pos is not None:
+            var_idx = int(var_indices_2[mm.seq2_pos])
+        else:
+            continue
+
+        mismatch_type_arr[var_idx] = mm.event
+        mismatch_identity_arr[var_idx] = _format_mismatch_identity(
+            mm.event, mm.seq1_base, mm.seq2_base
+        )
+        is_mismatch_arr[var_idx] = True
+
+        if mm.event == "substitution" and mm.seq1_pos is not None and mm.seq2_pos is not None:
+            substitution_map.append(
+                {
+                    "seq1_var_idx": int(var_indices_1[mm.seq1_pos]),
+                    "seq2_var_idx": int(var_indices_2[mm.seq2_pos]),
+                    "seq1_base": mm.seq1_base,
+                    "seq2_base": mm.seq2_base,
+                }
+            )
+
+    adata.var[f"{output_prefix}_mismatch_type"] = pd.Series(
+        mismatch_type_arr, index=adata.var.index
+    )
+    adata.var[f"{output_prefix}_mismatch_identity"] = pd.Series(
+        mismatch_identity_arr, index=adata.var.index
+    )
+    adata.var[f"{output_prefix}_is_mismatch"] = pd.Series(is_mismatch_arr, index=adata.var.index)
+    # Store substitution map as a DataFrame in adata.uns (h5ad-serializable)
+    if substitution_map:
+        adata.uns[f"{output_prefix}_substitution_map"] = pd.DataFrame(substitution_map)
+    adata.uns[uns_flag] = True
+
+    n_sub = sum(1 for t in mismatch_type_arr if t == "substitution")
+    n_ins = sum(1 for t in mismatch_type_arr if t == "insertion")
+    n_del = sum(1 for t in mismatch_type_arr if t == "deletion")
+    logger.info(
+        "Mismatch annotations: %d substitutions, %d insertions, %d deletions.",
+        n_sub,
+        n_ins,
+        n_del,
+    )