diff --git a/README.md b/README.md index 1f13ef1..10c700a 100644 --- a/README.md +++ b/README.md @@ -76,11 +76,20 @@ cropability snp -i samples.fa --min-af 0.05 --min-depth 10 # 计算 LD 矩阵 cropability ld --n-samples 500 --n-snps 1000 +# 运行原生 mpileup（真实 BAM/CRAM 输入） +cropability pileup -r ref.fa -b sample1.bam sample2.bam -o cohort.mpileup + +# 运行变异检测流程（默认 hybrid: 原生 mpileup + 原生 FastCall3 逻辑） +cropability call-variants -r ref.fa -b sample1.bam sample2.bam -o cohort.vcf --mode hybrid + # 导出 TorchScript 模型（供 Java 调用） cropability export --model add --output model.pt cropability export --model embedding --output embed.pt ``` +> NGS 流程默认在 CropAbility 内部执行，不依赖外部 `samtools/FastCall3` 命令。 +> 运行时需要 `pysam`（`pip install "cropability[io]"`）。可选 Rust 后端可用于性能加速。 + ### Python API ```python diff --git a/cropability/cli/main.py b/cropability/cli/main.py index ade3ddc..9758ff3 100644 --- a/cropability/cli/main.py +++ b/cropability/cli/main.py @@ -9,6 +9,8 @@ ld — 计算连锁不平衡矩阵 gwas — 全基因组关联分析 align — 批量序列比对 + pileup — 运行原生 mpileup（CropAbility 内置） + call-variants — 运行 NGS 变异检测流程 export — 导出 TorchScript 模型 benchmark — GPU 性能基准测试 """ @@ -199,6 +201,60 @@ def cmd_ld(args: argparse.Namespace) -> int: return 0 +def cmd_pileup(args: argparse.Namespace) -> int: + """运行原生 mpileup 并输出位点汇总文本。""" + from cropability.genomics.pipeline import QCThresholds, VariantPipeline + + qc = QCThresholds( + min_depth=args.min_depth, + min_base_quality=args.min_baseq, + min_mapping_quality=args.min_mapq, + min_alt_freq=args.min_af, + ) + pipeline = VariantPipeline() + report = pipeline.run( + mode="mpileup", + reference=args.reference, + bam_files=args.bam, + output=args.output, + qc=qc, + regions=args.region, + dry_run=args.dry_run, + ) + print("native pileup completed") + print(f" output: {report['mpileup']['output']}") + return 0 + + +def cmd_call_variants(args: argparse.Namespace) -> int: + """运行原生变异检测流程（mpileup / fastcall3 / hybrid）。""" + from cropability.genomics.pipeline import QCThresholds, VariantPipeline + + qc = QCThresholds( + min_depth=args.min_depth, + min_base_quality=args.min_baseq, + min_mapping_quality=args.min_mapq, + min_alt_freq=args.min_af, + ) + pipeline = VariantPipeline() + report = pipeline.run( + mode=args.mode, + reference=args.reference, + bam_files=args.bam, + output=args.output, + qc=qc, + regions=args.region, + mpileup_output=args.mpileup_output, + dry_run=args.dry_run, + ) + print(f"pipeline mode={report['mode']} completed") + if "mpileup" in report: + print(f" mpileup: {report['mpileup']['output']}") + if "fastcall3" in report: + print(f" vcf: {report['fastcall3']['output']}") + return 0 + + # --------------------------------------------------------------------------- # 主解析器 # --------------------------------------------------------------------------- @@ -243,6 +299,37 @@ def build_parser() -> argparse.ArgumentParser: ld.add_argument("--n-samples", type=int, default=200) ld.add_argument("--n-snps", type=int, default=500) + # pileup + pileup = sub.add_parser("pileup", help="运行原生 mpileup") + pileup.add_argument("-r", "--reference", required=True, help="参考基因组 FASTA") + pileup.add_argument("-b", "--bam", required=True, nargs="+", help="输入 BAM/CRAM 文件列表") + pileup.add_argument("-o", "--output", required=True, help="输出位点汇总路径（TSV）") + pileup.add_argument("--region", default=None, help="区域过滤，如 chr1:1-100000") + pileup.add_argument("--min-depth", type=int, default=10, help="最小深度阈值") + pileup.add_argument("--min-baseq", type=int, default=20, help="最小碱基质量") + pileup.add_argument("--min-mapq", type=int, default=20, help="最小比对质量") + pileup.add_argument("--min-af", type=float, default=0.05, help="最小替代等位频率") + pileup.add_argument("--dry-run", action="store_true", help="仅打印命令，不实际执行") + + # call-variants + call = sub.add_parser("call-variants", help="运行变异检测流程") + call.add_argument("-r", "--reference", required=True, help="参考基因组 FASTA") + call.add_argument("-b", "--bam", required=True, nargs="+", help="输入 BAM/CRAM 文件列表") + call.add_argument("-o", "--output", required=True, help="输出 VCF（或 mpileup）路径") + call.add_argument( + "--mode", + choices=["mpileup", "fastcall3", "hybrid"], + default="hybrid", + help="流程模式", + ) + call.add_argument("--mpileup-output", default=None, help="hybrid 模式的中间 mpileup 输出路径") + call.add_argument("--region", default=None, help="区域过滤，如 chr1:1-100000") + call.add_argument("--min-depth", type=int, default=10, help="最小深度阈值") + call.add_argument("--min-baseq", type=int, default=20, help="最小碱基质量") + call.add_argument("--min-mapq", type=int, default=20, help="最小比对质量") + call.add_argument("--min-af", type=float, default=0.05, help="最小替代等位频率") + call.add_argument("--dry-run", action="store_true", help="仅打印命令，不实际执行") + return parser @@ -252,6 +339,8 @@ def build_parser() -> argparse.ArgumentParser: "snp": cmd_snp, "export": cmd_export, "ld": cmd_ld, + "pileup": cmd_pileup, + "call-variants": cmd_call_variants, } diff --git a/cropability/genomics/__init__.py b/cropability/genomics/__init__.py index 7d0d71d..561f20d 100644 --- a/cropability/genomics/__init__.py +++ b/cropability/genomics/__init__.py @@ -9,10 +9,13 @@ - alignment: 序列比对评分 """ -from cropability.genomics.variant import VariantCaller, SNPResult -from cropability.genomics.ld import LDCalculator, LDResult -from cropability.genomics.gwas import GWASEngine, GWASResult from cropability.genomics.alignment import SmithWatermanGPU +from cropability.genomics.fastcall3 import FastCall3Config, FastCall3Runner, FastCall3RunResult +from cropability.genomics.gwas import GWASEngine, GWASResult +from cropability.genomics.ld import LDCalculator, LDResult +from cropability.genomics.pileup import MpileupParser, PileupRecord, PileupSample, PileupSiteSummary +from cropability.genomics.pipeline import PipelineConfig, QCThresholds, VariantPipeline +from cropability.genomics.variant import SNPResult, VariantCaller __all__ = [ "VariantCaller", @@ -22,4 +25,14 @@ "GWASEngine", "GWASResult", "SmithWatermanGPU", + "FastCall3Runner", + "FastCall3Config", + "FastCall3RunResult", + "MpileupParser", + "PileupRecord", + "PileupSample", + "PileupSiteSummary", + "VariantPipeline", + "PipelineConfig", + "QCThresholds", ] diff --git a/cropability/genomics/fastcall3.py b/cropability/genomics/fastcall3.py new file mode 100644 index 0000000..360de5d --- /dev/null +++ b/cropability/genomics/fastcall3.py @@ -0,0 +1,257 @@ +""" +FastCall3 native adapter +======================== +Implements in-process variant calling in CropAbility without invoking external +FastCall3 or samtools binaries, with an optional Rust acceleration hook. +""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +from cropability.genomics.pileup import NativePileupEngine, PileupRecord +from cropability.io.vcf import VCFRecord, VCFWriter +from cropability.utils.logging import get_logger + +logger = get_logger(__name__) + + +@dataclass +class FastCall3Config: + timeout_seconds: int = 3600 + prefer_rust_backend: bool = True + extra_args: list[str] = field(default_factory=list) + + +@dataclass +class FastCall3RunResult: + command: list[str] + returncode: int + stdout: str + stderr: str + output_vcf: Path + backend: str = "python" + n_records: int = 0 + elapsed_seconds: float = 0.0 + + @property + def ok(self) -> bool: + return self.returncode == 0 and self.output_vcf.exists() + + +def _choose_genotype(ref_count: int, alt_count: int, depth: int, min_depth: int) -> str: + if depth < min_depth: + return "./." + if depth == 0: + return "./." + af = alt_count / depth + if af >= 0.8: + return "1/1" + if af >= 0.2: + return "0/1" + return "0/0" + + +class FastCall3Runner: + """Native FastCall3-style variant caller executed fully in CropAbility.""" + + def __init__(self, config: FastCall3Config | None = None) -> None: + self.config = config or FastCall3Config() + self.pileup_engine = NativePileupEngine(prefer_rust_backend=self.config.prefer_rust_backend) + + def _run_rust_backend( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + output_vcf: str | Path, + regions: str | None, + min_base_quality: int, + min_mapping_quality: int, + min_depth: int, + min_alt_freq: float, + ) -> FastCall3RunResult | None: + if not self.config.prefer_rust_backend: + return None + backend = self.pileup_engine.try_get_rust_backend() + if backend is None: + return None + logger.info("Using optional Rust backend for native FastCall3 call") + return backend.run_fastcall3( + reference=str(reference), + bam_files=[str(x) for x in bam_files], + output_vcf=str(output_vcf), + regions=regions, + min_base_quality=min_base_quality, + min_mapping_quality=min_mapping_quality, + min_depth=min_depth, + min_alt_freq=min_alt_freq, + ) + + def run( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + output_vcf: str | Path, + regions: str | None = None, + min_base_quality: int | None = None, + min_mapping_quality: int | None = None, + min_depth: int | None = None, + min_alt_freq: float = 0.05, + extra_args: Sequence[str] | None = None, + dry_run: bool = False, + ) -> FastCall3RunResult: + output_path = Path(output_vcf) + mbq = 20 if min_base_quality is None else min_base_quality + mmq = 20 if min_mapping_quality is None else min_mapping_quality + mdp = 10 if min_depth is None else min_depth + + command = [ + "cropability-fastcall3-native", + "--reference", + str(reference), + "--output", + str(output_path), + "--min-base-quality", + str(mbq), + "--min-mapping-quality", + str(mmq), + "--min-depth", + str(mdp), + "--min-alt-freq", + str(min_alt_freq), + ] + for bam in bam_files: + command.extend(["--bam", str(bam)]) + if regions: + command.extend(["--regions", regions]) + if self.config.extra_args: + command.extend(self.config.extra_args) + if extra_args: + command.extend([str(x) for x in extra_args]) + + if dry_run: + return FastCall3RunResult( + command=command, + returncode=0, + stdout="dry-run", + stderr="", + output_vcf=output_path, + backend="dry-run", + n_records=0, + elapsed_seconds=0.0, + ) + + rust_result = self._run_rust_backend( + reference=reference, + bam_files=bam_files, + output_vcf=output_vcf, + regions=regions, + min_base_quality=mbq, + min_mapping_quality=mmq, + min_depth=mdp, + min_alt_freq=min_alt_freq, + ) + if rust_result is not None: + return rust_result + + start = time.time() + pileup_records = self.pileup_engine.generate_records( + reference=reference, + bam_files=bam_files, + regions=regions, + min_base_quality=mbq, + min_mapping_quality=mmq, + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + n_records = 0 + sample_names = [Path(x).stem for x in bam_files] + meta = [ + '##INFO=', + '##INFO=', + '##INFO=', + '##FORMAT=', + '##FORMAT=', + '##FORMAT=', + ] + + with VCFWriter(output_path, sample_names=sample_names) as writer: + writer.write_header(source="CropAbility.NativeFastCall3", extra_meta=meta) + for rec in pileup_records: + out = self._call_record( + rec=rec, + sample_names=sample_names, + min_depth=mdp, + min_alt_freq=min_alt_freq, + ) + if out is None: + continue + writer.write_record(out) + n_records += 1 + + elapsed = time.time() - start + return FastCall3RunResult( + command=command, + returncode=0, + stdout=f"native fastcall3 completed with {n_records} records", + stderr="", + output_vcf=output_path, + backend="python", + n_records=n_records, + elapsed_seconds=elapsed, + ) + + def _call_record( + self, + rec: PileupRecord, + sample_names: Sequence[str], + min_depth: int, + min_alt_freq: float, + ) -> VCFRecord | None: + total_depth = 0 + merged: dict[str, int] = {k: 0 for k in ("A", "C", "G", "T", "N")} + for sample in rec.samples.values(): + total_depth += sample.depth + for base, count in sample.base_counts.items(): + merged[base] = merged.get(base, 0) + count + + if total_depth < min_depth: + return None + + ref = rec.ref_base.upper() + candidates = {b: c for b, c in merged.items() if b in {"A", "C", "G", "T"} and b != ref} + if not candidates: + return None + alt, alt_count = max(candidates.items(), key=lambda kv: kv[1]) + alt_freq = alt_count / max(1, total_depth) + if alt_freq < min_alt_freq: + return None + + sample_values: list[str] = [] + for sample_name in sample_names: + s = rec.samples[sample_name] + depth = s.depth + ref_count = s.base_counts.get(ref, 0) + sample_alt_count = s.base_counts.get(alt, 0) + gt = _choose_genotype(ref_count, sample_alt_count, depth, min_depth) + sample_values.append(f"{gt}:{depth}:{ref_count},{sample_alt_count}") + + return VCFRecord( + chrom=rec.chrom, + pos=rec.pos, + id=".", + ref=ref, + alt=[alt], + qual=None, + filter=[], + info={ + "DP": str(total_depth), + "AF": f"{alt_freq:.6f}", + "AC": str(alt_count), + }, + format=["GT", "DP", "AD"], + samples=sample_values, + ) + diff --git a/cropability/genomics/pileup.py b/cropability/genomics/pileup.py new file mode 100644 index 0000000..f3d35f7 --- /dev/null +++ b/cropability/genomics/pileup.py @@ -0,0 +1,371 @@ +""" +mpileup parsing and native generation +==================================== +Provides parser helpers for text mpileup and in-process pileup generation from +alignment files using pysam, with an optional Rust acceleration hook. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from pathlib import Path + +from cropability.utils.logging import get_logger + +logger = get_logger(__name__) + +_BASES = ("A", "C", "G", "T", "N") + + +@dataclass +class PileupSample: + """Single-sample per-site pileup summary.""" + + depth: int + base_counts: dict[str, int] + insertions: int = 0 + deletions: int = 0 + + @property + def alt_count(self) -> int: + return sum(v for k, v in self.base_counts.items() if k in {"A", "C", "G", "T"}) + + +@dataclass +class PileupRecord: + """Single genomic site pileup summary across one or more samples.""" + + chrom: str + pos: int + ref_base: str + samples: dict[str, PileupSample] = field(default_factory=dict) + + def total_depth(self) -> int: + return sum(s.depth for s in self.samples.values()) + + +@dataclass +class PileupSiteSummary: + """Merged cross-sample site summary for threshold-based filtering.""" + + chrom: str + pos: int + ref_base: str + depth: int + alt_base: str | None + alt_count: int + alt_freq: float + + +def _parse_pileup_bases(bases: str, ref_base: str) -> tuple[dict[str, int], int, int]: + counts = {b: 0 for b in _BASES} + insertions = 0 + deletions = 0 + i = 0 + ref_u = ref_base.upper() + + while i < len(bases): + c = bases[i] + if c == "^": + i += 2 + continue + if c == "$": + i += 1 + continue + if c in "+-": + sign = c + i += 1 + nbuf = [] + while i < len(bases) and bases[i].isdigit(): + nbuf.append(bases[i]) + i += 1 + length = int("".join(nbuf)) if nbuf else 0 + if sign == "+": + insertions += 1 + else: + deletions += 1 + i += length + continue + if c == "*": + deletions += 1 + i += 1 + continue + if c in ".,": # matches reference + if ref_u in counts: + counts[ref_u] += 1 + else: + counts["N"] += 1 + i += 1 + continue + + b = c.upper() + if b in counts: + counts[b] += 1 + else: + counts["N"] += 1 + i += 1 + + return counts, insertions, deletions + + +class MpileupParser: + """ + Parser for text-based mpileup output lines. + + Input format: + CHROM POS REF [DP BASES QUAL]... + where each sample occupies three columns. + """ + + def __init__(self, sample_names: Sequence[str] | None = None) -> None: + self.sample_names = list(sample_names) if sample_names is not None else None + + def parse_line(self, line: str) -> PileupRecord | None: + line = line.rstrip("\n") + if not line: + return None + cols = line.split("\t") + if len(cols) < 6: + return None + if (len(cols) - 3) % 3 != 0: + return None + + chrom, pos_s, ref_base = cols[0], cols[1], cols[2] + per_sample = cols[3:] + n_samples = len(per_sample) // 3 + if self.sample_names is None: + names = [f"sample{i + 1}" for i in range(n_samples)] + else: + if len(self.sample_names) != n_samples: + raise ValueError( + f"sample_names count ({len(self.sample_names)}) != mpileup samples ({n_samples})" + ) + names = list(self.sample_names) + + samples: dict[str, PileupSample] = {} + for i, name in enumerate(names): + depth = int(per_sample[i * 3]) + bases = per_sample[i * 3 + 1] + counts, ins, dels = _parse_pileup_bases(bases, ref_base) + samples[name] = PileupSample( + depth=depth, + base_counts=counts, + insertions=ins, + deletions=dels, + ) + + return PileupRecord( + chrom=chrom, + pos=int(pos_s), + ref_base=ref_base.upper(), + samples=samples, + ) + + def parse_file(self, path: str | Path) -> Iterator[PileupRecord]: + p = Path(path) + with p.open("r", encoding="utf-8") as f: + for line in f: + rec = self.parse_line(line) + if rec is not None: + yield rec + + def summarize_sites( + self, + records: Iterator[PileupRecord], + min_depth: int = 10, + min_alt_freq: float = 0.05, + ) -> list[PileupSiteSummary]: + summaries: list[PileupSiteSummary] = [] + for rec in records: + merged_counts = {b: 0 for b in _BASES} + depth = 0 + for sample in rec.samples.values(): + depth += sample.depth + for b, c in sample.base_counts.items(): + merged_counts[b] = merged_counts.get(b, 0) + c + + if depth < min_depth: + continue + ref = rec.ref_base + alt_candidates = { + b: c + for b, c in merged_counts.items() + if b in {"A", "C", "G", "T"} and b != ref + } + if not alt_candidates: + continue + alt_base, alt_count = max(alt_candidates.items(), key=lambda kv: kv[1]) + alt_freq = alt_count / max(depth, 1) + if alt_freq < min_alt_freq: + continue + summaries.append( + PileupSiteSummary( + chrom=rec.chrom, + pos=rec.pos, + ref_base=ref, + depth=depth, + alt_base=alt_base, + alt_count=alt_count, + alt_freq=alt_freq, + ) + ) + + logger.info("Generated %d pileup site summaries", len(summaries)) + return summaries + + +class NativePileupEngine: + """In-process pileup engine replacing external `samtools mpileup` calls.""" + + def __init__(self, prefer_rust_backend: bool = True) -> None: + self.prefer_rust_backend = prefer_rust_backend + + def try_get_rust_backend(self): + if not self.prefer_rust_backend: + return None + try: + # Optional extension point. This module is intentionally optional. + from cropability.native import ngs_core # type: ignore + + return ngs_core + except Exception: + return None + + def _require_pysam(self): + try: + import pysam # noqa: PLC0415 + + return pysam + except Exception as e: # pragma: no cover - import path dependent + raise RuntimeError( + "pysam is required for native pileup. Install optional dependency: cropability[io]" + ) from e + + def _parse_region(self, region: str | None) -> tuple[str | None, int | None, int | None]: + if not region: + return None, None, None + if ":" not in region: + return region, None, None + chrom, span = region.split(":", 1) + if "-" not in span: + return chrom, None, None + start_s, end_s = span.split("-", 1) + start = max(0, int(start_s.replace(",", "")) - 1) + end = int(end_s.replace(",", "")) + return chrom, start, end + + def _count_sample_column( + self, + column, + ref_base: str, + min_base_quality: int, + min_mapping_quality: int, + ) -> PileupSample: + counts = {b: 0 for b in _BASES} + insertions = 0 + deletions = 0 + depth = 0 + ref_base_u = ref_base.upper() + for pr in column.pileups: + aln = pr.alignment + if aln.mapping_quality < min_mapping_quality: + continue + if pr.is_refskip: + continue + if pr.is_del: + deletions += 1 + continue + qpos = pr.query_position + if qpos is None: + continue + if qpos >= len(aln.query_qualities): + continue + if aln.query_qualities[qpos] < min_base_quality: + continue + depth += 1 + if pr.indel > 0: + insertions += 1 + elif pr.indel < 0: + deletions += 1 + + base = aln.query_sequence[qpos].upper() + if base not in counts: + base = "N" + if base == ref_base_u: + counts[ref_base_u] += 1 + else: + counts[base] += 1 + return PileupSample(depth=depth, base_counts=counts, insertions=insertions, deletions=deletions) + + def generate_records( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + regions: str | None = None, + min_base_quality: int = 20, + min_mapping_quality: int = 20, + ) -> Iterator[PileupRecord]: + rust_backend = self.try_get_rust_backend() + if rust_backend is not None and hasattr(rust_backend, "generate_pileup_records"): + logger.info("Using optional Rust backend for pileup generation") + yield from rust_backend.generate_pileup_records( + reference=str(reference), + bam_files=[str(x) for x in bam_files], + regions=regions, + min_base_quality=min_base_quality, + min_mapping_quality=min_mapping_quality, + ) + return + + pysam = self._require_pysam() + sample_names = [Path(x).stem for x in bam_files] + chrom_filter, start, end = self._parse_region(regions) + ref = pysam.FastaFile(str(reference)) + bam_handles = [pysam.AlignmentFile(str(p), "rc" if str(p).endswith(".cram") else "rb") for p in bam_files] + try: + site_table: dict[tuple[str, int], dict[str, PileupSample]] = {} + ref_table: dict[tuple[str, int], str] = {} + for sample_name, bam in zip(sample_names, bam_handles): + pileup_iter = bam.pileup( + contig=chrom_filter, + start=start, + stop=end, + truncate=chrom_filter is not None, + min_base_quality=0, # handled by custom thresholds + stepper="all", + ) + for col in pileup_iter: + chrom = col.reference_name + pos1 = col.pos + 1 + key = (chrom, pos1) + if key not in ref_table: + ref_base = ref.fetch(chrom, col.pos, col.pos + 1).upper() + ref_table[key] = ref_base if ref_base else "N" + sample_pileup = self._count_sample_column( + column=col, + ref_base=ref_table[key], + min_base_quality=min_base_quality, + min_mapping_quality=min_mapping_quality, + ) + site_table.setdefault(key, {})[sample_name] = sample_pileup + + for (chrom, pos1), sample_data in sorted(site_table.items()): + for sample_name in sample_names: + if sample_name not in sample_data: + sample_data[sample_name] = PileupSample( + depth=0, + base_counts={b: 0 for b in _BASES}, + insertions=0, + deletions=0, + ) + yield PileupRecord( + chrom=chrom, + pos=pos1, + ref_base=ref_table[(chrom, pos1)], + samples=sample_data, + ) + finally: + ref.close() + for h in bam_handles: + h.close() diff --git a/cropability/genomics/pipeline.py b/cropability/genomics/pipeline.py new file mode 100644 index 0000000..4b431f7 --- /dev/null +++ b/cropability/genomics/pipeline.py @@ -0,0 +1,211 @@ +""" +NGS variant detection pipeline orchestration +=========================================== +Native in-process pipeline for mpileup-like summaries and FastCall3-style VCF +calling without shelling out to external binaries. +""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +from cropability.genomics.fastcall3 import FastCall3Config, FastCall3Runner +from cropability.genomics.pileup import MpileupParser, NativePileupEngine +from cropability.io.bam import AlignmentInputManager +from cropability.utils.logging import get_logger + +logger = get_logger(__name__) + + +@dataclass +class QCThresholds: + min_depth: int = 10 + min_base_quality: int = 20 + min_mapping_quality: int = 20 + min_alt_freq: float = 0.05 + + +@dataclass +class PipelineConfig: + fastcall3: FastCall3Config = field(default_factory=FastCall3Config) + prefer_rust_backend: bool = True + timeout_seconds: int = 3600 + + +class VariantPipeline: + """Native NGS variant pipeline.""" + + def __init__(self, config: PipelineConfig | None = None) -> None: + self.config = config or PipelineConfig() + self.fastcall3_runner = FastCall3Runner(self.config.fastcall3) + self.pileup_engine = NativePileupEngine(prefer_rust_backend=self.config.prefer_rust_backend) + + def _write_pileup_summary( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + output_path: str | Path, + qc: QCThresholds, + regions: str | None = None, + ) -> dict[str, object]: + out = Path(output_path) + records = self.pileup_engine.generate_records( + reference=reference, + bam_files=bam_files, + regions=regions, + min_base_quality=qc.min_base_quality, + min_mapping_quality=qc.min_mapping_quality, + ) + out.parent.mkdir(parents=True, exist_ok=True) + n_sites = 0 + parser = MpileupParser() + with out.open("w", encoding="utf-8") as f: + f.write("#CHROM\tPOS\tREF\tDP\tALT\tAC\tAF\n") + summaries = parser.summarize_sites( + records=records, + min_depth=qc.min_depth, + min_alt_freq=qc.min_alt_freq, + ) + for s in summaries: + f.write( + f"{s.chrom}\t{s.pos}\t{s.ref_base}\t{s.depth}\t" + f"{s.alt_base or '.'}\t{s.alt_count}\t{s.alt_freq:.6f}\n" + ) + n_sites += 1 + return { + "engine": "native", + "output": str(out), + "returncode": 0, + "n_sites": n_sites, + } + + def run_mpileup( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + output_path: str | Path, + qc: QCThresholds, + regions: str | None = None, + dry_run: bool = False, + ) -> dict[str, object]: + if dry_run: + return { + "engine": "native", + "command": [ + "cropability-native-pileup", + "--reference", + str(reference), + "--output", + str(output_path), + ], + "output": str(output_path), + "returncode": 0, + } + return self._write_pileup_summary( + reference=reference, + bam_files=bam_files, + output_path=output_path, + qc=qc, + regions=regions, + ) + + def run_fastcall3( + self, + reference: str | Path, + bam_files: Sequence[str | Path], + output_vcf: str | Path, + qc: QCThresholds, + regions: str | None = None, + dry_run: bool = False, + ) -> dict[str, object]: + result = self.fastcall3_runner.run( + reference=reference, + bam_files=bam_files, + output_vcf=output_vcf, + regions=regions, + min_base_quality=qc.min_base_quality, + min_mapping_quality=qc.min_mapping_quality, + min_depth=qc.min_depth, + min_alt_freq=qc.min_alt_freq, + dry_run=dry_run, + ) + return { + "engine": result.backend, + "command": result.command, + "output": str(result.output_vcf), + "returncode": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + "n_records": result.n_records, + "elapsed_seconds": result.elapsed_seconds, + } + + def run( + self, + mode: str, + reference: str | Path, + bam_files: Sequence[str | Path], + output: str | Path, + qc: QCThresholds | None = None, + regions: str | None = None, + mpileup_output: str | Path | None = None, + dry_run: bool = False, + ) -> dict[str, object]: + mode = mode.lower() + if mode not in {"mpileup", "fastcall3", "hybrid"}: + raise ValueError("mode must be one of: mpileup, fastcall3, hybrid") + + qc = qc or QCThresholds() + input_manager = AlignmentInputManager(bam_files) + input_manager.validate_tools(require_samtools=False, require_pysam=not dry_run) + alignment_files = input_manager.collect(require_index=True) + bam_paths = [f.path for f in alignment_files] + + start = time.time() + report: dict[str, object] = {"mode": mode, "reference": str(reference), "engine": "native"} + if mode == "mpileup": + report["mpileup"] = self.run_mpileup( + reference=reference, + bam_files=bam_paths, + output_path=output, + qc=qc, + regions=regions, + dry_run=dry_run, + ) + elif mode == "fastcall3": + report["fastcall3"] = self.run_fastcall3( + reference=reference, + bam_files=bam_paths, + output_vcf=output, + qc=qc, + regions=regions, + dry_run=dry_run, + ) + else: + mpileup_out = ( + Path(mpileup_output) + if mpileup_output is not None + else Path(f"{output}.pileup.summary.tsv") + ) + report["mpileup"] = self.run_mpileup( + reference=reference, + bam_files=bam_paths, + output_path=mpileup_out, + qc=qc, + regions=regions, + dry_run=dry_run, + ) + report["fastcall3"] = self.run_fastcall3( + reference=reference, + bam_files=bam_paths, + output_vcf=output, + qc=qc, + regions=regions, + dry_run=dry_run, + ) + report["elapsed_seconds"] = time.time() - start + return report + diff --git a/cropability/io/__init__.py b/cropability/io/__init__.py index 384a703..6e939b1 100644 --- a/cropability/io/__init__.py +++ b/cropability/io/__init__.py @@ -1,9 +1,12 @@ """基因组数据 I/O 模块：FASTA/FASTQ/VCF 读写与格式转换。""" +from cropability.io.bam import AlignmentFile, AlignmentInputManager from cropability.io.fasta import FastaReader, FastaWriter -from cropability.io.vcf import VCFReader, VCFWriter, VCFRecord +from cropability.io.vcf import VCFReader, VCFRecord, VCFWriter __all__ = [ + "AlignmentFile", + "AlignmentInputManager", "FastaReader", "FastaWriter", "VCFReader", diff --git a/cropability/io/bam.py b/cropability/io/bam.py new file mode 100644 index 0000000..d3a2095 --- /dev/null +++ b/cropability/io/bam.py @@ -0,0 +1,119 @@ +""" +BAM/CRAM 输入抽象 +================= +为 NGS 变异检测流水线提供 BAM/CRAM 文件检查、样本命名与索引校验功能。 +""" + +from __future__ import annotations + +import shutil +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +from cropability.utils.logging import get_logger + +logger = get_logger(__name__) + + +@dataclass +class AlignmentFile: + """单个对齐输入文件信息。""" + + path: Path + sample_name: str + fmt: str # bam | cram + has_index: bool + index_path: Path | None + + @property + def is_bam(self) -> bool: + return self.fmt == "bam" + + @property + def is_cram(self) -> bool: + return self.fmt == "cram" + + +def _detect_format(path: Path) -> str: + suffix = path.suffix.lower() + if suffix == ".bam": + return "bam" + if suffix == ".cram": + return "cram" + raise ValueError(f"Unsupported alignment format for: {path} (expected .bam/.cram)") + + +def _resolve_index(path: Path, fmt: str) -> Path | None: + if fmt == "bam": + candidates = [Path(f"{path}.bai"), path.with_suffix(".bai")] + else: + candidates = [Path(f"{path}.crai"), path.with_suffix(".crai")] + for idx in candidates: + if idx.exists(): + return idx + return None + + +class AlignmentInputManager: + """ + BAM/CRAM 输入管理器。 + + 支持： + - 基础文件存在性校验 + - BAM/CRAM 格式检测 + - 索引文件存在性校验 + - samtools / optional pysam 依赖检查 + """ + + def __init__(self, paths: Sequence[str | Path]) -> None: + if not paths: + raise ValueError("At least one alignment file is required") + self.paths = [Path(p) for p in paths] + + @staticmethod + def check_samtools() -> bool: + return shutil.which("samtools") is not None + + @staticmethod + def check_pysam() -> bool: + try: + import pysam # noqa: F401 + + return True + except Exception: + return False + + def validate_tools( + self, + require_samtools: bool = True, + require_pysam: bool = False, + ) -> None: + if require_samtools and not self.check_samtools(): + raise RuntimeError("samtools is required but not found in PATH") + if require_pysam and not self.check_pysam(): + raise RuntimeError("pysam is required but not installed") + + def collect(self, require_index: bool = True) -> list[AlignmentFile]: + files: list[AlignmentFile] = [] + for p in self.paths: + if not p.exists(): + raise FileNotFoundError(f"Alignment file not found: {p}") + fmt = _detect_format(p) + idx = _resolve_index(p, fmt) + if require_index and idx is None: + ext = ".bai" if fmt == "bam" else ".crai" + raise FileNotFoundError( + f"Index not found for {p}. Expected {p}{ext} or sibling index file." + ) + files.append( + AlignmentFile( + path=p, + sample_name=p.stem, + fmt=fmt, + has_index=idx is not None, + index_path=idx, + ) + ) + logger.info(f"Collected {len(files)} alignment files") + return files diff --git a/cropability/io/vcf.py b/cropability/io/vcf.py index d5e03c6..ee13cb5 100644 --- a/cropability/io/vcf.py +++ b/cropability/io/vcf.py @@ -182,11 +182,22 @@ def __init__(self, path: Union[str, Path], sample_names: Optional[List[str]] = N self.sample_names = sample_names or [] self._f = self.path.open("w", encoding="utf-8") self._wrote_header = False - - def write_header(self, extra_meta: Optional[List[str]] = None) -> None: - self._f.write("##fileformat=VCFv4.2\n") + self._meta_lines: List[str] = [] + + def write_header( + self, + extra_meta: Optional[List[str]] = None, + source: Optional[str] = None, + fileformat: str = "VCFv4.2", + ) -> None: + self._f.write(f"##fileformat={fileformat}\n") + if source: + meta = f"##source={source}" + self._f.write(meta + "\n") + self._meta_lines.append(meta) for line in (extra_meta or []): self._f.write(line + "\n") + self._meta_lines.append(line) cols = ["#CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO"] if self.sample_names: cols += ["FORMAT"] + self.sample_names @@ -208,6 +219,12 @@ def write_record(self, rec: VCFRecord) -> None: if rec.format: cols.append(":".join(rec.format)) cols.extend(rec.samples) + elif self.sample_names: + cols.append("GT") + if rec.samples: + cols.extend(rec.samples) + else: + cols.extend(["./."] * len(self.sample_names)) self._f.write("\t".join(cols) + "\n") def close(self) -> None: diff --git a/tests/test_ngs_pipeline.py b/tests/test_ngs_pipeline.py new file mode 100644 index 0000000..a84d7b1 --- /dev/null +++ b/tests/test_ngs_pipeline.py @@ -0,0 +1,144 @@ +"""测试 NGS pipeline 新增模块。""" + + +import pytest + +from cropability.cli.main import build_parser +from cropability.genomics.fastcall3 import FastCall3Config, FastCall3Runner +from cropability.genomics.pileup import MpileupParser, PileupRecord, PileupSample +from cropability.genomics.pipeline import QCThresholds, VariantPipeline +from cropability.io.bam import AlignmentInputManager +from cropability.io.vcf import VCFReader + + +class TestAlignmentInputManager: + def test_collect_bam_with_index(self, tmp_path): + bam = tmp_path / "sample1.bam" + bai = tmp_path / "sample1.bam.bai" + bam.write_bytes(b"") + bai.write_bytes(b"") + mgr = AlignmentInputManager([bam]) + files = mgr.collect(require_index=True) + assert len(files) == 1 + assert files[0].is_bam + assert files[0].has_index + + def test_missing_index_raises(self, tmp_path): + bam = tmp_path / "sample1.bam" + bam.write_bytes(b"") + mgr = AlignmentInputManager([bam]) + with pytest.raises(FileNotFoundError): + mgr.collect(require_index=True) + + +class TestMpileupParser: + def test_parse_line_and_summary(self): + parser = MpileupParser(sample_names=["s1"]) + # chr1 10 A depth=5 bases=".,TtA" quals dummy + rec = parser.parse_line("chr1\t10\tA\t5\t.,TtA\tFFFFF") + assert rec is not None + assert rec.chrom == "chr1" + assert rec.pos == 10 + assert rec.samples["s1"].depth == 5 + summaries = parser.summarize_sites(iter([rec]), min_depth=1, min_alt_freq=0.1) + assert len(summaries) == 1 + assert summaries[0].alt_base in {"T", "C", "G", "A"} + + +class TestFastCall3Runner: + def test_build_and_dry_run(self, tmp_path): + out = tmp_path / "out.vcf" + runner = FastCall3Runner(FastCall3Config()) + result = runner.run( + reference="ref.fa", + bam_files=["a.bam", "b.bam"], + output_vcf=out, + regions="chr1:1-100", + min_base_quality=20, + min_mapping_quality=30, + min_depth=10, + dry_run=True, + ) + assert result.returncode == 0 + assert result.backend == "dry-run" + assert "--reference" in result.command + assert "--output" in result.command + + def test_native_run_writes_vcf(self, monkeypatch, tmp_path): + out = tmp_path / "out.vcf" + runner = FastCall3Runner(FastCall3Config(prefer_rust_backend=False)) + + fake_records = [ + PileupRecord( + chrom="chr1", + pos=42, + ref_base="A", + samples={ + "a": PileupSample(depth=10, base_counts={"A": 7, "C": 3, "G": 0, "T": 0, "N": 0}), + "b": PileupSample(depth=8, base_counts={"A": 3, "C": 5, "G": 0, "T": 0, "N": 0}), + }, + ) + ] + monkeypatch.setattr(runner.pileup_engine, "generate_records", lambda **_: iter(fake_records)) + + res = runner.run( + reference="ref.fa", + bam_files=["a.bam", "b.bam"], + output_vcf=out, + min_depth=1, + min_alt_freq=0.1, + dry_run=False, + ) + assert res.ok + assert res.n_records == 1 + + records = list(VCFReader(out)) + assert len(records) == 1 + assert records[0].chrom == "chr1" + assert records[0].pos == 42 + assert records[0].ref == "A" + assert records[0].alt == ["C"] + + +class TestVariantPipeline: + def test_run_hybrid_dry_run(self, tmp_path): + bam = tmp_path / "cohort.bam" + bai = tmp_path / "cohort.bam.bai" + bam.write_bytes(b"") + bai.write_bytes(b"") + + pipeline = VariantPipeline() + report = pipeline.run( + mode="hybrid", + reference="ref.fa", + bam_files=[bam], + output=tmp_path / "cohort.vcf", + qc=QCThresholds(), + dry_run=True, + ) + assert report["mode"] == "hybrid" + assert "mpileup" in report + assert "fastcall3" in report + assert report["engine"] == "native" + assert report["fastcall3"]["engine"] == "dry-run" + + +class TestCliExtensions: + def test_new_subcommands_exist(self): + parser = build_parser() + ns = parser.parse_args( + [ + "call-variants", + "-r", + "ref.fa", + "-b", + "a.bam", + "-o", + "out.vcf", + "--mode", + "hybrid", + "--dry-run", + ] + ) + assert ns.command == "call-variants" + assert ns.mode == "hybrid"