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

smftools/informatics/fasta_functions.py

@@ -1,23 +1,93 @@
 from __future__ import annotations
 
 import gzip
+import shutil
+import subprocess
 from concurrent.futures import ProcessPoolExecutor
+from importlib.util import find_spec
 from pathlib import Path
-from typing import Dict, Iterable, Tuple
+from typing import TYPE_CHECKING, Dict, Iterable, Tuple
 
 import numpy as np
-import pysam
 from Bio import SeqIO
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
-from pyfaidx import Fasta
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 from ..readwrite import time_string
 
 logger = get_logger(__name__)
 
+if TYPE_CHECKING:
+    import pysam as pysam_module
+
+
+def _require_pysam() -> "pysam_module":
+    if pysam_types is not None:
+        return pysam_types
+    return require("pysam", extra="pysam", purpose="FASTA access")
+
+
+pysam_types = None
+if find_spec("pysam") is not None:
+    pysam_types = require("pysam", extra="pysam", purpose="FASTA access")
+
+
+def _resolve_fasta_backend() -> str:
+    """Resolve the backend to use for FASTA access."""
+    if pysam_types is not None:
+        return "python"
+    if shutil is not None and shutil.which("samtools"):
+        return "cli"
+    raise RuntimeError("FASTA access requires pysam or samtools in PATH.")
+
+
+def _ensure_fasta_index(fasta: Path) -> None:
+    fai = fasta.with_suffix(fasta.suffix + ".fai")
+    if fai.exists():
+        return
+    if subprocess is None or shutil is None or not shutil.which("samtools"):
+        pysam_mod = _require_pysam()
+        pysam_mod.faidx(str(fasta))
+        return
+    cp = subprocess.run(
+        ["samtools", "faidx", str(fasta)],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    if cp.returncode != 0:
+        raise RuntimeError(f"samtools faidx failed (exit {cp.returncode}):\n{cp.stderr}")
+
+
+def _bed_to_faidx_region(chrom: str, start: int, end: int) -> str:
+    """Convert 0-based half-open BED coords to samtools faidx region."""
+    start1 = start + 1
+    end1 = end
+    if start1 > end1:
+        start1, end1 = end1, start1
+    return f"{chrom}:{start1}-{end1}"
+
+
+def _fetch_sequence_with_samtools(fasta: Path, chrom: str, start: int, end: int) -> str:
+    if subprocess is None or shutil is None:
+        raise RuntimeError("samtools backend is unavailable.")
+    if not shutil.which("samtools"):
+        raise RuntimeError("samtools is required but not available in PATH.")
+    region = _bed_to_faidx_region(chrom, start, end)
+    cp = subprocess.run(
+        ["samtools", "faidx", str(fasta), region],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    if cp.returncode != 0:
+        raise RuntimeError(f"samtools faidx failed (exit {cp.returncode}):\n{cp.stderr}")
+    lines = [line.strip() for line in cp.stdout.splitlines() if line and not line.startswith(">")]
+    return "".join(lines)
+
 
 def _convert_FASTA_record(
     record: SeqRecord,
@@ -160,7 +230,7 @@ def index_fasta(fasta: str | Path, write_chrom_sizes: bool = True) -> Path:
         Path: Path to the index file or chromosome sizes file.
     """
     fasta = Path(fasta)
-    pysam.faidx(str(fasta))  # creates <fasta>.fai
+    _require_pysam().faidx(str(fasta))  # creates <fasta>.fai
 
     fai = fasta.with_suffix(fasta.suffix + ".fai")
     if write_chrom_sizes:
@@ -307,8 +377,13 @@ def subsample_fasta_from_bed(
     # Ensure output directory exists
     output_directory.mkdir(parents=True, exist_ok=True)
 
-    # Load the FASTA file using pyfaidx
-    fasta = Fasta(str(input_FASTA))  # pyfaidx requires string paths
+    backend = _resolve_fasta_backend()
+    _ensure_fasta_index(input_FASTA)
+
+    fasta_handle = None
+    if backend == "python":
+        pysam_mod = _require_pysam()
+        fasta_handle = pysam_mod.FastaFile(str(input_FASTA))
 
     # Open BED + output FASTA
     with input_bed.open("r") as bed, output_FASTA.open("w") as out_fasta:
@@ -319,15 +394,24 @@ def subsample_fasta_from_bed(
             end = int(fields[2])  # BED is 0-based and end is exclusive
             desc = " ".join(fields[3:]) if len(fields) > 3 else ""
 
-            if chrom not in fasta:
+            if backend == "python":
+                assert fasta_handle is not None
+                if chrom not in fasta_handle.references:
+                    logger.warning(f"{chrom} not found in FASTA")
+                    continue
+                sequence = fasta_handle.fetch(chrom, start, end)
+            else:
+                sequence = _fetch_sequence_with_samtools(input_FASTA, chrom, start, end)
+
+            if not sequence:
                 logger.warning(f"{chrom} not found in FASTA")
                 continue
 
-            # pyfaidx is 1-based indexing internally, but [start:end] works with BED coords
-            sequence = fasta[chrom][start:end].seq
-
             header = f">{chrom}:{start}-{end}"
             if desc:
                 header += f"    {desc}"
 
             out_fasta.write(f"{header}\n{sequence}\n")
+
+    if fasta_handle is not None:
+        fasta_handle.close()

smftools/informatics/h5ad_functions.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import glob
 import os
 from concurrent.futures import ProcessPoolExecutor, as_completed
@@ -7,9 +9,9 @@ from typing import Dict, List, Optional, Union
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp
-from pod5 import Reader
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -82,6 +84,112 @@ def add_demux_type_annotation(
     return adata
 
 
+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,
@@ -90,6 +198,7 @@ def add_read_length_and_mapping_qc(
     extract_read_features_from_bam_callable=None,
     bypass: bool = False,
     force_redo: bool = True,
+    samtools_backend: str | None = "auto",
 ):
     """
     Populate adata.obs with read/mapping QC columns.
@@ -101,7 +210,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).
@@ -133,7 +243,7 @@ def add_read_length_and_mapping_qc(
                     "No `read_metrics` provided and `extract_read_features_from_bam` not found."
                 )
             for bam in bam_files:
-                bam_read_metrics = extractor(bam)
+                bam_read_metrics = extractor(bam, samtools_backend)
                 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)
@@ -151,10 +261,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)):
@@ -176,6 +288,8 @@ def add_read_length_and_mapping_qc(
                 "reference_length",
                 "mapped_length",
                 "mapping_quality",
+                "reference_start",
+                "reference_end",
             ],
         )
 
@@ -188,6 +302,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
@@ -228,6 +344,9 @@ def _collect_read_origins_from_pod5(pod5_path: str, target_ids: set[str]) -> dic
     Worker function: scan one POD5 file and return a mapping
     {read_id: pod5_basename} only for read_ids in `target_ids`.
     """
+    p5 = require("pod5", extra="ont", purpose="POD5 metadata")
+    Reader = p5.Reader
+
     basename = os.path.basename(pod5_path)
     mapping: dict[str, str] = {}