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

smftools/preprocessing/__init__.py

@@ -5,6 +5,10 @@ from importlib import import_module
 _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:
@@ -133,23 +133,23 @@ def append_base_context(
                 adata.var[f"{ref}_{site_type}_valid_coverage"] = (
                     (adata.var[f"{ref}_{site_type}"]) & (adata.var[f"position_in_{ref}"])
                 )
-                if native:
-                    adata.obsm[f"{ref}_{site_type}_valid_coverage"] = adata[
-                        :, adata.var[f"{ref}_{site_type}_valid_coverage"]
-                    ].layers["binarized_methylation"]
-                else:
-                    adata.obsm[f"{ref}_{site_type}_valid_coverage"] = adata[
-                        :, adata.var[f"{ref}_{site_type}_valid_coverage"]
-                    ].X
-            else:
-                pass
-
-            if native:
-                adata.obsm[f"{ref}_{site_type}"] = adata[:, adata.var[f"{ref}_{site_type}"]].layers[
-                    "binarized_methylation"
-                ]
-            else:
-                adata.obsm[f"{ref}_{site_type}"] = adata[:, adata.var[f"{ref}_{site_type}"]].X
+            #     if native:
+            #         adata.obsm[f"{ref}_{site_type}_valid_coverage"] = adata[
+            #             :, adata.var[f"{ref}_{site_type}_valid_coverage"]
+            #         ].layers["binarized_methylation"]
+            #     else:
+            #         adata.obsm[f"{ref}_{site_type}_valid_coverage"] = adata[
+            #             :, adata.var[f"{ref}_{site_type}_valid_coverage"]
+            #         ].X
+            # else:
+            #     pass
+
+            # if native:
+            #     adata.obsm[f"{ref}_{site_type}"] = adata[:, adata.var[f"{ref}_{site_type}"]].layers[
+            #         "binarized_methylation"
+            #     ]
+            # else:
+            #     adata.obsm[f"{ref}_{site_type}"] = adata[:, adata.var[f"{ref}_{site_type}"]].X
 
     # mark as done
     adata.uns[uns_flag] = True

smftools/preprocessing/append_mismatch_frequency_sites.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Iterable, Sequence
+
+import numpy as np
+import pandas as pd
+
+from smftools.constants import BASE_QUALITY_SCORES, MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def append_mismatch_frequency_sites(
+    adata: "ad.AnnData",
+    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,
+    bypass: bool = False,
+) -> None:
+    """Append mismatch frequency metadata and variable-site flags per reference.
+
+    Args:
+        adata: AnnData object.
+        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.
+        bypass: Whether to skip running this step.
+    """
+    if bypass:
+        return
+
+    already = bool(adata.uns.get(uns_flag, False))
+    if already and not force_redo:
+        return
+
+    if mismatch_layer not in adata.layers:
+        logger.debug(
+            "Mismatch layer '%s' not found; skipping mismatch frequency step.", mismatch_layer
+        )
+        return
+
+    mismatch_map = adata.uns.get("mismatch_integer_encoding_map", {})
+    if not mismatch_map:
+        logger.debug("Mismatch encoding map not found; skipping mismatch frequency step.")
+        return
+
+    n_value = mismatch_map.get("N", MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT["N"])
+    pad_value = mismatch_map.get("PAD", MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT["PAD"])
+
+    base_int_to_label = {
+        int(value): str(base)
+        for base, value in mismatch_map.items()
+        if base not in {"N", "PAD"} and isinstance(value, (int, np.integer))
+    }
+    if not base_int_to_label:
+        logger.debug("Mismatch encoding map missing base labels; skipping mismatch frequency step.")
+        return
+
+    has_span_mask = read_span_layer in adata.layers
+    if not has_span_mask:
+        logger.debug(
+            "Read span mask '%s' not found; mismatch frequencies will be computed over all reads.",
+            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]
+
+    if mismatch_frequency_range is None:
+        mismatch_frequency_range = (0.0, 1.0)
+
+    lower_bound, upper_bound = mismatch_frequency_range
+
+    for ref in references:
+        ref_mask = adata.obs[ref_column] == 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)
+
+        frequency_values = np.full(n_vars, np.nan, dtype=float)
+        variable_flags = np.zeros(n_vars, dtype=bool)
+        mismatch_base_frequencies: list[list[tuple[str, float]]] = [[] for _ in range(n_vars)]
+
+        if ref_mask.sum() == 0:
+            adata.var[f"{ref}_mismatch_frequency"] = pd.Series(
+                frequency_values, index=adata.var.index
+            )
+            adata.var[f"{ref}_variable_sequence_site"] = pd.Series(
+                variable_flags, index=adata.var.index
+            )
+            adata.var[f"{ref}_mismatch_base_frequencies"] = pd.Series(
+                mismatch_base_frequencies, index=adata.var.index
+            )
+            continue
+
+        mismatch_matrix = np.asarray(adata.layers[mismatch_layer][ref_mask])
+        if has_span_mask:
+            span_matrix = np.asarray(adata.layers[read_span_layer][ref_mask])
+            coverage_mask = span_matrix > 0
+            coverage_counts = coverage_mask.sum(axis=0).astype(float)
+        else:
+            coverage_mask = np.ones_like(mismatch_matrix, dtype=bool)
+            coverage_counts = np.full(n_vars, ref_mask.sum(), dtype=float)
+
+        mismatch_mask = (~np.isin(mismatch_matrix, [n_value, pad_value])) & coverage_mask
+        mismatch_counts = mismatch_mask.sum(axis=0)
+
+        frequency_values = np.divide(
+            mismatch_counts,
+            coverage_counts,
+            out=np.full(n_vars, np.nan, dtype=float),
+            where=coverage_counts > 0,
+        )
+        frequency_values = np.where(ref_position_mask.values, frequency_values, np.nan)
+
+        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:
+            base_counts_by_int[base_int] = ((mismatch_matrix == base_int) & coverage_mask).sum(
+                axis=0
+            )
+
+        for idx in range(n_vars):
+            if not ref_position_mask.iloc[idx] or coverage_counts[idx] == 0:
+                continue
+            base_freqs: list[tuple[str, float]] = []
+            for base_int, base_label in base_int_to_label.items():
+                count = base_counts_by_int[base_int][idx]
+                if count > 0:
+                    base_freqs.append((base_label, float(count / coverage_counts[idx])))
+            mismatch_base_frequencies[idx] = base_freqs
+
+        adata.var[f"{ref}_mismatch_frequency"] = pd.Series(frequency_values, index=adata.var.index)
+        adata.var[f"{ref}_variable_sequence_site"] = pd.Series(
+            variable_flags, index=adata.var.index
+        )
+        adata.var[f"{ref}_mismatch_base_frequencies"] = pd.Series(
+            mismatch_base_frequencies, index=adata.var.index
+        )
+
+    adata.uns[uns_flag] = True

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,
+    )