samsampleX 0.1.0__py3-none-any.whl

samsamplex/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

samsamplex/bed.py

@@ -0,0 +1,146 @@
+"""BED file I/O and depth-array combination utilities."""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import TextIO
+
+import numpy as np
+
+from .depth import DepthArray
+
+
+def contig_names_match(name1: str, name2: str) -> bool:
+    """Check whether two contig names match, ignoring a leading 'chr' prefix."""
+    if name1 == name2:
+        return True
+    strip = lambda s: s[3:] if s.startswith("chr") else s
+    return strip(name1) == strip(name2)
+
+
+# ── Writing ──────────────────────────────────────────────────────────────────
+
+
+def write_bed_entry(fp: TextIO, chrom: str, start: int, end: int, depth: int) -> None:
+    """Write a single BED4 line."""
+    fp.write(f"{chrom}\t{start}\t{end}\t{depth}\n")
+
+
+def write_bed_output(fp: TextIO, arr: DepthArray, collapse: int = 0) -> None:
+    """Write a DepthArray to BED format with optional collapsing.
+
+    When *collapse* is 0 every position gets its own line.  When > 0,
+    consecutive positions whose depth differs by <= *collapse* are merged
+    into a single interval (using the depth of the first position in the
+    interval.
+    """
+    if arr.length == 0:
+        return
+
+    if collapse == 0:
+        for i in range(arr.length):
+            pos = arr.start + i
+            write_bed_entry(fp, arr.contig, pos, pos + 1, int(arr.depths[i]))
+    else:
+        interval_start = arr.start
+        interval_depth = int(arr.depths[0])
+
+        for i in range(1, arr.length):
+            current_depth = int(arr.depths[i])
+            if abs(current_depth - interval_depth) > collapse:
+                write_bed_entry(fp, arr.contig, interval_start, arr.start + i, interval_depth)
+                interval_start = arr.start + i
+                interval_depth = current_depth
+
+        write_bed_entry(fp, arr.contig, interval_start, arr.end, interval_depth)
+
+
+# ── Reading ──────────────────────────────────────────────────────────────────
+
+
+def bed_read_depths(
+    bed_path: str,
+    contig: str,
+    region_start: int,
+    region_end: int,
+) -> DepthArray:
+    """Read depth values from a BED4 file into a DepthArray for a region."""
+    depths = np.zeros(region_end - region_start, dtype=np.int32)
+
+    with open(bed_path) as fp:
+        for line in fp:
+            line = line.rstrip("\n")
+            if not line or line.startswith("#"):
+                continue
+
+            parts = line.split("\t")
+            if len(parts) < 4:
+                parts = line.split()
+            if len(parts) < 4:
+                continue
+
+            chrom = parts[0]
+            bed_start = int(parts[1])
+            bed_end = int(parts[2])
+            depth = int(parts[3])
+
+            if not contig_names_match(chrom, contig):
+                continue
+
+            ov_start = max(bed_start, region_start)
+            ov_end = min(bed_end, region_end)
+            if ov_start >= ov_end:
+                continue
+
+            depths[ov_start - region_start : ov_end - region_start] = depth
+
+    return DepthArray(contig=contig, start=region_start, end=region_end, depths=depths)
+
+
+# ── Combining multiple templates ─────────────────────────────────────────────
+
+
+def bed_combine_depths(
+    arrays: list[DepthArray],
+    mode: str = "min",
+    seed: int = 42,
+) -> DepthArray:
+    """Combine multiple DepthArrays position-by-position.
+
+    Supported *mode* values: ``"min"``, ``"mean"``, ``"median"``, ``"max"``, ``"random"``.
+    """
+    if len(arrays) == 1:
+        return DepthArray(
+            contig=arrays[0].contig,
+            start=arrays[0].start,
+            end=arrays[0].end,
+            depths=arrays[0].depths.copy(),
+        )
+
+    stacked = np.stack([a.depths for a in arrays], axis=0)  # (n_arrays, length)
+
+    if mode == "min":
+        combined = stacked.min(axis=0)
+    elif mode == "max":
+        combined = stacked.max(axis=0)
+    elif mode == "mean":
+        combined = (stacked.sum(axis=0) // len(arrays)).astype(np.int32)
+    elif mode == "median":
+        combined = np.rint(np.median(stacked, axis=0)).astype(np.int32)
+    elif mode == "random":
+        rng = random.Random(seed)
+        mins = stacked.min(axis=0)
+        maxs = stacked.max(axis=0)
+        combined = np.array(
+            [
+                v if mn == mx else rng.randint(mn, mx)
+                for mn, mx, v in zip(mins, maxs, mins)
+            ],
+            dtype=np.int32,
+        )
+    else:
+        raise ValueError(f"Unknown combine mode: {mode}")
+
+    ref = arrays[0]
+    return DepthArray(contig=ref.contig, start=ref.start, end=ref.end, depths=combined)

samsamplex/cli.py

@@ -0,0 +1,342 @@
+"""Command-line interface for samsamplex."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import __version__
+from .modes import DEPTH_MODES
+
+
+def _add_map_parser(subparsers: argparse._SubParsersAction) -> None:
+    p = subparsers.add_parser(
+        "map",
+        help="Extract depth of coverage from BAM to BED template",
+    )
+    p.add_argument(
+        "--template-bam",
+        nargs="+",
+        required=True,
+        help="Input BAM file(s)",
+    )
+    p.add_argument("--region", required=True, help="Target region (samtools-style)")
+    p.add_argument("--out-bed", required=True, help="Output BED file")
+    p.add_argument(
+        "--collapse",
+        type=int,
+        default=0,
+        help="Merge consecutive positions with depth diff <= INT [default: 0]",
+    )
+    p.add_argument(
+        "--mode",
+        default="mean",
+        choices=DEPTH_MODES,
+        help="How to combine depths when multiple BAMs given [default: mean]",
+    )
+    p.add_argument("--seed", type=int, default=42, help="Random seed for --mode random [default: 42]")
+
+
+def _add_sample_parser(subparsers: argparse._SubParsersAction) -> None:
+    p = subparsers.add_parser(
+        "sample",
+        help="Sample reads from BAM to match template depth distribution",
+    )
+    p.add_argument("--source-bam", required=True, help="Input BAM file to sample from")
+    p.add_argument(
+        "--template-bed",
+        nargs="*",
+        default=[],
+        help="Template BED file(s) with depth values (required unless --uniform)",
+    )
+    p.add_argument(
+        "--uniform",
+        type=float,
+        default=None,
+        metavar="FRACTION",
+        help="Uniform sampling: retain fraction of reads by hash (0-1). Bypasses template/depth logic.",
+    )
+    p.add_argument("--region", required=True, help="Target region (samtools-style)")
+    p.add_argument("--out-bam", required=True, help="Output BAM file")
+    p.add_argument(
+        "--mode",
+        default="random",
+        choices=DEPTH_MODES,
+        help="How to combine multiple template BEDs [default: random]",
+    )
+    p.add_argument(
+        "--stat",
+        default="mean",
+        choices=DEPTH_MODES,
+        help=(
+            "Statistic for summarising ratio over read span; "
+            "'random' picks one ratio from covered bases deterministically from span+seed "
+            "[default: mean]"
+        ),
+    )
+    p.add_argument("--seed", type=int, default=42, help="Random seed [default: 42]")
+    # p.add_argument("--no-sort", action="store_true", help="Skip sorting/indexing output BAM")
+    p.add_argument("--threads", type=int, default=2, help="Number of writer threads to use [default: 2]")
+
+
+def _add_plot_parser(subparsers: argparse._SubParsersAction) -> None:
+    p = subparsers.add_parser(
+        "plot",
+        help="Compare depth of coverage and output as PNG plot or TSV data",
+    )
+    p.add_argument("--source-bam", required=True, help="Source BAM file")
+    p.add_argument("--out-bam", required=True, help="Output BAM file (from sampling)")
+    p.add_argument("--region", required=True, help="Target region (samtools-style)")
+
+    tpl = p.add_mutually_exclusive_group(required=True)
+    tpl.add_argument("--template-bam", help="Template BAM file")
+    tpl.add_argument("--template-bed", help="Template BED file")
+
+    out = p.add_mutually_exclusive_group(required=True)
+    out.add_argument("--out-png", help="Output PNG plot file")
+    out.add_argument("--out-tsv", help="Output TSV data file")
+
+
+def _add_mapback_parser(subparsers: argparse._SubParsersAction) -> None:
+    p = subparsers.add_parser(
+        "mapback",
+        help="Remap HLA*LA PRG-mapped reads back to chr6 coordinates",
+    )
+    p.add_argument("--source-bam", required=True, help="HLA*LA-remapped BAM file")
+    p.add_argument("--region", required=True, help="Target region on chr6 (samtools-style)")
+    p.add_argument("--out-bam", required=True, help="Output BAM file")
+    p.add_argument(
+        "--genome-build", required=True, choices=("GRCh38", "GRCh37"),
+        help="Reference genome build",
+    )
+    p.add_argument(
+        "--prg-seq",
+        default="HLA-LA/graphs/PRG_MHC_GRCh38_withIMGT/sequences.txt",
+        help="Path to HLA*LA sequences.txt [default: HLA-LA/graphs/PRG_MHC_GRCh38_withIMGT/sequences.txt]",
+    )
+    p.add_argument("--no-sort", action="store_true", help="Skip sorting/indexing output BAM")
+
+
+def _add_stats_parser(subparsers: argparse._SubParsersAction) -> None:
+    p = subparsers.add_parser(
+        "stats",
+        help="Compare depth distributions between two BAM or BED files",
+    )
+    p.add_argument("--a", required=True, help="First input file — BAM or BED (reference)")
+    p.add_argument("--b", required=True, help="Second input file — BAM or BED (comparison)")
+    p.add_argument("--region", required=True, help="Region to compare (samtools-style)")
+
+
+# ── Subcommand handlers ─────────────────────────────────────────────────────
+
+
+def _run_map(args: argparse.Namespace) -> int:
+    from .bed import bed_combine_depths, write_bed_output
+    from .depth import depth_from_bam, region_parse, resolve_contig_name
+
+    import pysam
+
+    log = lambda msg: print(msg, file=sys.stderr)
+
+    region = region_parse(args.region)
+    template_bams = args.template_bam  # list of one or more paths
+
+    log(f"[map] Template BAM(s): {template_bams}")
+    log(f"[map] Region: {args.region}")
+    log(f"[map] Collapse: {args.collapse}")
+    log(f"[map] Output BED: {args.out_bed}")
+    if len(template_bams) > 1:
+        log(f"[map] Mode: {args.mode}")
+
+    with pysam.AlignmentFile(template_bams[0], "rb") as bam:
+        resolved = resolve_contig_name(bam.header, region.contig)
+        if resolved is None:
+            log(f"Error: Contig '{region.contig}' not found in BAM")
+            return 1
+        region.contig = resolved
+        contig_len = bam.get_reference_length(resolved)
+
+    if region.start < 0:
+        region.start = 0
+    if region.end < 0:
+        region.end = contig_len
+
+    log(f"[map] Parsed region: {region.contig}:{region.start}-{region.end}")
+    log("[map] Computing depth array(s) (this may take a while)...")
+
+    depth_arrays = []
+    for bam_path in template_bams:
+        depth_arrays.append(
+            depth_from_bam(bam_path, region.contig, region.start, region.end)
+        )
+
+    if len(depth_arrays) == 1:
+        depth = depth_arrays[0]
+    else:
+        depth = bed_combine_depths(
+            depth_arrays, mode=args.mode, seed=args.seed
+        )
+
+    log(f"[map] Computed depth for {depth.length} positions")
+
+    suffix = " (collapsed)" if args.collapse > 0 else ""
+    log(f"[map] Writing BED file{suffix}...")
+
+    with open(args.out_bed, "w") as fp:
+        write_bed_output(fp, depth, collapse=args.collapse)
+
+    log(f"[map] Done. Output written to: {args.out_bed}")
+    return 0
+
+
+def _run_sample(args: argparse.Namespace) -> int:
+    from .sample import sample_run
+
+    log = lambda msg: print(msg, file=sys.stderr)
+
+    if args.uniform is None and not args.template_bed:
+        log("Error: Either --template-bed or --uniform is required")
+        return 1
+
+    if args.uniform is not None:
+        if args.uniform <= 0 or args.uniform > 1:
+            log(f"Error: --uniform must be in (0, 1], got {args.uniform}")
+            return 1
+
+    return sample_run(
+        source_bam=args.source_bam,
+        template_beds=args.template_bed if args.template_bed else [],
+        region_str=args.region,
+        out_bam=args.out_bam,
+        mode=args.mode,
+        stat=args.stat,
+        seed=args.seed,
+        # no_sort=args.no_sort,
+        uniform_fraction=args.uniform,
+        threads=args.threads,
+    )
+
+
+def _run_plot(args: argparse.Namespace) -> int:
+    from .plot import plot_run
+
+    return plot_run(
+        source_bam=args.source_bam,
+        out_bam=args.out_bam,
+        region_str=args.region,
+        template_bam=args.template_bam,
+        template_bed=args.template_bed,
+        out_png=args.out_png,
+        out_tsv=args.out_tsv,
+    )
+
+
+def _run_mapback(args: argparse.Namespace) -> int:
+    from .mapback import mapback_run
+
+    return mapback_run(
+        source_bam=args.source_bam,
+        region_str=args.region,
+        out_bam=args.out_bam,
+        genome_build=args.genome_build,
+        prg_seq=args.prg_seq,
+        no_sort=args.no_sort,
+    )
+
+
+def _run_stats(args: argparse.Namespace) -> int:
+    from .depth import region_parse
+    from .metrics import depth_from_path, metrics_calculate, metrics_print
+
+    log = lambda msg: print(msg, file=sys.stderr)
+
+    region = region_parse(args.region)
+
+    # Resolve region bounds from a BAM input (BED files don't carry contig lengths).
+    bam_input = None
+    for path in (args.a, args.b):
+        if not path.endswith(".bed"):
+            bam_input = path
+            break
+
+    if bam_input is not None:
+        import pysam
+        from .depth import resolve_contig_name
+
+        with pysam.AlignmentFile(bam_input, "rb") as bam:
+            resolved = resolve_contig_name(bam.header, region.contig)
+            if resolved is None:
+                log(f"Error: Contig '{region.contig}' not found in {bam_input}")
+                return 1
+            region.contig = resolved
+            if region.start < 0:
+                region.start = 0
+            if region.end < 0:
+                region.end = bam.get_reference_length(resolved)
+    else:
+        if region.start < 0 or region.end < 0:
+            log("Error: --region must specify explicit start and end when both inputs are BED files")
+            return 1
+
+    def _fmt_type(path: str) -> str:
+        return "BED" if path.endswith(".bed") else "BAM"
+
+    log(f"Input A ({_fmt_type(args.a)}): {args.a}")
+    log(f"Input B ({_fmt_type(args.b)}): {args.b}")
+    log(f"Region: {region.contig}:{region.start + 1}-{region.end}")
+
+    depth_a = depth_from_path(args.a, region.contig, region.start, region.end)
+    depth_b = depth_from_path(args.b, region.contig, region.start, region.end)
+
+    result = metrics_calculate(depth_a, depth_b)
+
+    log("")
+    log("========== Comparison Metrics ==========")
+    log(f"A (reference):           {args.a}")
+    log(f"B (comparison):          {args.b}")
+    log(f"Region:                  {region.contig}:{region.start + 1}-{region.end}")
+    log(f"Positions compared:      {depth_a.length}")
+    log("-" * 41)
+    metrics_print(result, label_a="A", label_b="B")
+
+    return 0
+
+
+# ── Entry point ──────────────────────────────────────────────────────────────
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="samsamplex",
+        description=f"samsamplex v{__version__} — Depth-aware BAM file sampling",
+    )
+    parser.add_argument(
+        "-v", "--version", action="version", version=f"samsamplex v{__version__}"
+    )
+
+    subparsers = parser.add_subparsers(dest="command")
+    _add_map_parser(subparsers)
+    _add_sample_parser(subparsers)
+    _add_plot_parser(subparsers)
+    _add_mapback_parser(subparsers)
+    _add_stats_parser(subparsers)
+
+    args = parser.parse_args(argv)
+
+    if args.command is None:
+        parser.print_help()
+        sys.exit(1)
+
+    dispatch = {
+        "map": _run_map,
+        "sample": _run_sample,
+        "plot": _run_plot,
+        "mapback": _run_mapback,
+        "stats": _run_stats,
+    }
+
+    sys.exit(dispatch[args.command](args))
+
+
+if __name__ == "__main__":
+    main()

samsamplex/depth.py

@@ -0,0 +1,126 @@
+"""Depth array computation from BAM files and region parsing utilities."""
+
+from __future__ import annotations
+
+import re
+import sys
+from dataclasses import dataclass, field
+
+import numpy as np
+import pysam
+
+
+@dataclass
+class Region:
+    """Genomic region parsed from a samtools-style string."""
+
+    contig: str
+    start: int = -1  # 0-based inclusive; -1 means whole contig
+    end: int = -1  # 0-based exclusive; -1 means whole contig
+
+
+@dataclass
+class DepthArray:
+    """Per-position depth values over a contiguous genomic region."""
+
+    contig: str
+    start: int
+    end: int
+    depths: np.ndarray = field(repr=False)
+
+    @property
+    def length(self) -> int:
+        return self.end - self.start
+
+
+def region_parse(region_str: str) -> Region:
+    """Parse a samtools-style region string (e.g. chr1, chr1:1000-2000).
+
+    Input coordinates are 1-based; the returned Region uses 0-based half-open
+    coordinates.
+    """
+    m = re.match(r"^([^:]+)(?::(\d+)(?:-(\d+))?)?$", region_str)
+    if not m:
+        raise ValueError(f"Invalid region format: {region_str}")
+
+    contig = m.group(1)
+    start = int(m.group(2)) - 1 if m.group(2) else -1  # 1-based → 0-based
+    end = int(m.group(3)) if m.group(3) else -1  # exclusive, stays as-is
+
+    return Region(contig=contig, start=start, end=end)
+
+
+def resolve_contig_name(header: pysam.AlignmentHeader, contig: str) -> str | None:
+    """Resolve *contig* against a BAM header, trying chr-prefix variants.
+
+    Returns the matching name from the header, or None.
+    """
+    references = set(header.references)
+
+    if contig in references:
+        return contig
+
+    if contig.startswith("chr"):
+        alt = contig[3:]
+        if alt in references:
+            print(f"Note: Using contig '{alt}' (matched from '{contig}')", file=sys.stderr)
+            return alt
+    else:
+        alt = f"chr{contig}"
+        if alt in references:
+            print(f"Note: Using contig '{alt}' (matched from '{contig}')", file=sys.stderr)
+            return alt
+
+    return None
+
+
+def get_contig_length(header: pysam.AlignmentHeader, contig: str) -> int | None:
+    """Return contig length from a BAM header, handling chr-prefix mismatch."""
+    resolved = resolve_contig_name(header, contig)
+    if resolved is None:
+        return None
+    return header.get_reference_length(resolved)
+
+
+def depth_from_bam(
+    bam_path: str,
+    contig: str,
+    start: int,
+    end: int,
+) -> DepthArray:
+    """Compute per-position depth for a region from an indexed BAM file.
+
+    Uses a simplified read-iteration approach:
+    for each read overlapping the region, increment depth at every covered
+    position (CIGAR-unaware).
+    """
+    with pysam.AlignmentFile(bam_path, "rb") as bam:
+        resolved = resolve_contig_name(bam.header, contig)
+        if resolved is None:
+            raise ValueError(f"Contig '{contig}' not found in BAM header")
+
+        tid = bam.get_tid(resolved)
+        contig_len = bam.get_reference_length(resolved)
+
+        if start < 0:
+            start = 0
+        if end < 0 or end > contig_len:
+            end = contig_len
+
+        depths = np.zeros(end - start, dtype=np.int32)
+
+        for read in bam.fetch(resolved, start, end):
+
+            r_start = read.reference_start
+            r_end = read.reference_end
+            if r_end is None:
+                continue
+
+            ov_start = max(r_start, start)
+            ov_end = min(r_end, end)
+            if ov_start >= ov_end:
+                continue
+
+            depths[ov_start - start : ov_end - start] += 1
+
+    return DepthArray(contig=resolved, start=start, end=end, depths=depths)