566 lines
19 KiB
Python
566 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Stage 05 WGS BAM quality metrics from authoritative Manifests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime, timezone
|
|
import fcntl
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from gfa_core.wgs_scientific_identity import build_wgs_analysis_id, sha256_file
|
|
from gfa_core.wgs_upstream_manifest import resolve_wgs_artifacts
|
|
|
|
|
|
STAGE = "wgs_bam_qc"
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Measure coverage and alignment counts for authoritative WGS BAMs"
|
|
)
|
|
parser.add_argument("--workspace", required=True)
|
|
parser.add_argument("--job-id", required=True)
|
|
parser.add_argument("--prefix", required=True)
|
|
parser.add_argument("--markdup-manifest", required=True)
|
|
parser.add_argument("--index-manifest", required=True)
|
|
parser.add_argument("--samtools", required=True)
|
|
parser.add_argument("--coverage-threshold", type=int, default=1, choices=(1,))
|
|
parser.add_argument("--threads", required=True, type=int)
|
|
parser.add_argument("--max-jobs", required=True, type=int)
|
|
return parser
|
|
|
|
|
|
def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
|
|
temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
|
temporary.write_text(
|
|
json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def _safe_component(value: str, label: str) -> None:
|
|
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", value or ""):
|
|
raise ValueError(f"{label} must be a non-empty safe path component")
|
|
|
|
|
|
def _authority_header(
|
|
path: str | Path,
|
|
*,
|
|
stage: str,
|
|
workspace: Path,
|
|
prefix: str,
|
|
job_id: str,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
manifest = json.loads(Path(path).read_text(encoding="utf-8"))
|
|
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
|
raise ValueError(f"{stage} authority Manifest is unreadable") from exc
|
|
if (
|
|
not isinstance(manifest, dict)
|
|
or manifest.get("stage") != stage
|
|
or manifest.get("status") != "success"
|
|
):
|
|
raise ValueError(f"{stage} authority Manifest is not a successful producer")
|
|
if (
|
|
manifest.get("prefix") != prefix
|
|
or manifest.get("job_id") != job_id
|
|
or manifest.get("workspace") != str(workspace)
|
|
):
|
|
raise ValueError(f"{stage} authority Manifest identity mismatch")
|
|
return manifest
|
|
|
|
|
|
def _checked(command: list[str], name: str) -> subprocess.CompletedProcess[str]:
|
|
result = subprocess.run(command, text=True, capture_output=True, check=False)
|
|
if result.returncode:
|
|
raise RuntimeError(
|
|
f"{name} failed with exit code {result.returncode}: {result.stderr.strip()}"
|
|
)
|
|
return result
|
|
|
|
|
|
def _tool_version(samtools: str) -> str:
|
|
result = _checked([samtools, "--version"], "samtools version check")
|
|
lines = result.stdout.strip().splitlines()
|
|
if not lines:
|
|
raise RuntimeError("samtools version check produced no version")
|
|
return lines[0]
|
|
|
|
|
|
def _reference_lengths(path: Path) -> dict[str, int]:
|
|
lengths: dict[str, int] = {}
|
|
try:
|
|
lines = path.read_text(encoding="utf-8").splitlines()
|
|
except (OSError, UnicodeError) as exc:
|
|
raise RuntimeError("reference FAI is unreadable") from exc
|
|
for line_number, line in enumerate(lines, 1):
|
|
fields = line.split("\t")
|
|
if len(fields) < 2 or not fields[0]:
|
|
raise RuntimeError(f"reference FAI line {line_number} is invalid")
|
|
try:
|
|
length = int(fields[1])
|
|
except ValueError as exc:
|
|
raise RuntimeError(
|
|
f"reference FAI line {line_number} has invalid length"
|
|
) from exc
|
|
if length <= 0 or fields[0] in lengths:
|
|
raise RuntimeError(
|
|
f"reference FAI line {line_number} has duplicate/invalid contig"
|
|
)
|
|
lengths[fields[0]] = length
|
|
if not lengths:
|
|
raise RuntimeError("reference FAI contains no contigs")
|
|
return lengths
|
|
|
|
|
|
def _parse_depth(
|
|
path: Path, sample_id: str, lengths: dict[str, int], threshold: int
|
|
) -> dict[str, int | float | str]:
|
|
depth_sum = 0
|
|
bases_ge_threshold = 0
|
|
observed_positions: set[tuple[str, int]] = set()
|
|
try:
|
|
lines = path.read_text(encoding="utf-8").splitlines()
|
|
except (OSError, UnicodeError) as exc:
|
|
raise RuntimeError(f"sample {sample_id}: samtools depth output is unreadable") from exc
|
|
for line_number, line in enumerate(lines, 1):
|
|
fields = line.split("\t")
|
|
if len(fields) != 3 or fields[0] not in lengths:
|
|
raise RuntimeError(
|
|
f"sample {sample_id}: invalid depth row {line_number}"
|
|
)
|
|
try:
|
|
position, depth = int(fields[1]), int(fields[2])
|
|
except ValueError as exc:
|
|
raise RuntimeError(
|
|
f"sample {sample_id}: invalid depth row {line_number}"
|
|
) from exc
|
|
identity = (fields[0], position)
|
|
if (
|
|
position < 1
|
|
or position > lengths[fields[0]]
|
|
or depth < 0
|
|
or identity in observed_positions
|
|
):
|
|
raise RuntimeError(
|
|
f"sample {sample_id}: invalid depth row {line_number}"
|
|
)
|
|
observed_positions.add(identity)
|
|
depth_sum += depth
|
|
if depth >= threshold:
|
|
bases_ge_threshold += 1
|
|
total_reference_bases = sum(lengths.values())
|
|
return {
|
|
"sample_id": sample_id,
|
|
"total_reference_bases": total_reference_bases,
|
|
"depth_sum": depth_sum,
|
|
"mean_depth": depth_sum / total_reference_bases,
|
|
"bases_ge_1x": bases_ge_threshold,
|
|
"coverage_ge_1x": bases_ge_threshold / total_reference_bases,
|
|
}
|
|
|
|
|
|
def _run_sample(
|
|
sample_id: str,
|
|
input_bam: Path,
|
|
bundle: Path,
|
|
args: argparse.Namespace,
|
|
lengths: dict[str, int],
|
|
) -> tuple[dict[str, Path], dict[str, int | float | str]]:
|
|
base = f"{args.prefix}.{STAGE}.{sample_id}"
|
|
flagstat_path = bundle / f"{base}.flagstat.txt"
|
|
depth_summary_path = bundle / f"{base}.depth.tsv"
|
|
depth_raw = bundle / f".{base}.depth.raw.tsv"
|
|
|
|
flagstat = _checked(
|
|
[args.samtools, "flagstat", "-@", str(args.threads), str(input_bam)],
|
|
f"sample {sample_id} flagstat",
|
|
)
|
|
if " in total " not in f" {flagstat.stdout} ":
|
|
raise RuntimeError(f"sample {sample_id}: flagstat content is invalid")
|
|
flagstat_path.write_text(flagstat.stdout, encoding="utf-8")
|
|
|
|
with depth_raw.open("w", encoding="utf-8") as output:
|
|
depth = subprocess.run(
|
|
[
|
|
args.samtools,
|
|
"depth",
|
|
"-d",
|
|
"0",
|
|
"-@",
|
|
str(args.threads),
|
|
str(input_bam),
|
|
],
|
|
text=True,
|
|
stdout=output,
|
|
stderr=subprocess.PIPE,
|
|
check=False,
|
|
)
|
|
if depth.returncode:
|
|
raise RuntimeError(
|
|
f"sample {sample_id} depth failed with exit code "
|
|
f"{depth.returncode}: {depth.stderr.strip()}"
|
|
)
|
|
metrics = _parse_depth(
|
|
depth_raw, sample_id, lengths, args.coverage_threshold
|
|
)
|
|
depth_raw.unlink()
|
|
columns = (
|
|
"sample_id",
|
|
"total_reference_bases",
|
|
"depth_sum",
|
|
"mean_depth",
|
|
"bases_ge_1x",
|
|
"coverage_ge_1x",
|
|
)
|
|
depth_summary_path.write_text(
|
|
"\t".join(columns)
|
|
+ "\n"
|
|
+ "\t".join(
|
|
(
|
|
str(metrics["sample_id"]),
|
|
str(metrics["total_reference_bases"]),
|
|
str(metrics["depth_sum"]),
|
|
f"{float(metrics['mean_depth']):.12g}",
|
|
str(metrics["bases_ge_1x"]),
|
|
f"{float(metrics['coverage_ge_1x']):.12g}",
|
|
)
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
return (
|
|
{
|
|
"bam_flagstat": flagstat_path,
|
|
"bam_depth_summary": depth_summary_path,
|
|
},
|
|
metrics,
|
|
)
|
|
|
|
|
|
def _artifact(
|
|
workspace: Path,
|
|
source: Path,
|
|
destination: Path,
|
|
role: str,
|
|
sample_id: str = "",
|
|
) -> dict[str, Any]:
|
|
size = source.stat().st_size
|
|
record = {
|
|
"artifact": role,
|
|
"role": role,
|
|
"path": str(destination.relative_to(workspace)),
|
|
"full_path": str(destination.resolve()),
|
|
"type": source.suffix.lstrip(".") or "file",
|
|
"size": size,
|
|
"size_bytes": size,
|
|
"sha256": sha256_file(source),
|
|
}
|
|
if sample_id:
|
|
record["sample_id"] = sample_id
|
|
return record
|
|
|
|
|
|
def _publish(
|
|
bundle: Path,
|
|
destination: Path,
|
|
manifest_work: Path,
|
|
authority: Path,
|
|
started_at: str,
|
|
) -> bool:
|
|
lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with lock_path.open("a+", encoding="utf-8") as lock:
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
|
if authority.is_file():
|
|
try:
|
|
current = json.loads(authority.read_text(encoding="utf-8"))
|
|
if (
|
|
current.get("status") == "success"
|
|
and str(current.get("start_time") or "") > started_at
|
|
):
|
|
return False
|
|
except Exception:
|
|
pass
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
os.replace(bundle, destination)
|
|
os.replace(manifest_work, authority)
|
|
return True
|
|
|
|
|
|
def _failure(
|
|
args: argparse.Namespace, run_id: str, started_at: str, error: str
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"manifest_version": "1.0",
|
|
"stage": STAGE,
|
|
"status": "failed",
|
|
"prefix": args.prefix,
|
|
"job_id": args.job_id,
|
|
"run_id": run_id,
|
|
"start_time": started_at,
|
|
"artifacts": {},
|
|
"errors": [error],
|
|
}
|
|
|
|
|
|
def _publish_failure(
|
|
authority: Path,
|
|
args: argparse.Namespace,
|
|
run_id: str,
|
|
started_at: str,
|
|
error: str,
|
|
) -> None:
|
|
lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with lock_path.open("a+", encoding="utf-8") as lock:
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
|
failure_path = (
|
|
authority
|
|
if not authority.exists()
|
|
else authority.parent / f"{args.prefix}.{STAGE}.failed.{run_id}.manifest.json"
|
|
)
|
|
_atomic_json(failure_path, _failure(args, run_id, started_at, error))
|
|
|
|
|
|
def run(args: argparse.Namespace) -> int:
|
|
workspace = Path(args.workspace).resolve()
|
|
run_id = uuid.uuid4().hex
|
|
started_at = datetime.now(timezone.utc).isoformat()
|
|
final = workspace / "final"
|
|
final.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
_safe_component(args.prefix, "prefix")
|
|
except ValueError as exc:
|
|
_atomic_json(
|
|
final / f"{STAGE}.failed.{run_id}.manifest.json",
|
|
_failure(args, run_id, started_at, str(exc)),
|
|
)
|
|
print(f"{STAGE} failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
authority = final / f"{args.prefix}.{STAGE}.manifest.json"
|
|
try:
|
|
if args.threads < 1 or args.max_jobs < 1:
|
|
raise ValueError("threads and max-jobs must be positive")
|
|
markdup_header = _authority_header(
|
|
args.markdup_manifest,
|
|
stage="wgs_markdup",
|
|
workspace=workspace,
|
|
prefix=args.prefix,
|
|
job_id=args.job_id,
|
|
)
|
|
index_header = _authority_header(
|
|
args.index_manifest,
|
|
stage="wgs_index",
|
|
workspace=workspace,
|
|
prefix=args.prefix,
|
|
job_id=args.job_id,
|
|
)
|
|
sample_set_id = str(markdup_header.get("wgs_sample_set_id") or "")
|
|
reference_id = str(markdup_header.get("wgs_reference_id") or "")
|
|
markdup_analysis_id = str(markdup_header.get("wgs_analysis_id") or "")
|
|
index_analysis_id = str(index_header.get("wgs_analysis_id") or "")
|
|
if not all(
|
|
(sample_set_id, reference_id, markdup_analysis_id, index_analysis_id)
|
|
):
|
|
raise ValueError("upstream scientific identities are missing")
|
|
if index_header.get("wgs_reference_id") != reference_id:
|
|
raise ValueError("upstream reference identities do not match")
|
|
markdup = resolve_wgs_artifacts(
|
|
args.markdup_manifest,
|
|
"wgs_markdup",
|
|
("markdup_bam", "markdup_bam_index"),
|
|
sample_set_id,
|
|
reference_id,
|
|
markdup_analysis_id,
|
|
)
|
|
index = resolve_wgs_artifacts(
|
|
args.index_manifest,
|
|
"wgs_index",
|
|
("reference_fai",),
|
|
expected_reference_id=reference_id,
|
|
expected_analysis_id=index_analysis_id,
|
|
)
|
|
if len(index["reference_fai"]) != 1:
|
|
raise RuntimeError("authoritative index must contain one reference FAI")
|
|
bam_by_sample = {item.sample_id: item for item in markdup["markdup_bam"]}
|
|
bai_by_sample = {
|
|
item.sample_id: item for item in markdup["markdup_bam_index"]
|
|
}
|
|
expected_samples = set(bam_by_sample)
|
|
if not expected_samples or expected_samples != set(bai_by_sample):
|
|
raise RuntimeError("authoritative markdup input sample sets do not match")
|
|
lengths = _reference_lengths(index["reference_fai"][0].path)
|
|
version = _tool_version(args.samtools)
|
|
parameters = {
|
|
"samtools": args.samtools,
|
|
"coverage_threshold": args.coverage_threshold,
|
|
"metric_name": "coverage_ge_1x",
|
|
"threads": args.threads,
|
|
"max_jobs": args.max_jobs,
|
|
}
|
|
consumed = [
|
|
item.sha256
|
|
for group in (markdup, index)
|
|
for values in group.values()
|
|
for item in values
|
|
]
|
|
analysis_id = build_wgs_analysis_id(
|
|
STAGE,
|
|
consumed,
|
|
parameters,
|
|
{"samtools": version},
|
|
method="samtools_flagstat_depth",
|
|
)
|
|
attempt = workspace / ".attempts" / run_id
|
|
bundle = attempt / "bundle"
|
|
bundle.mkdir(parents=True, exist_ok=False)
|
|
completed: dict[str, dict[str, Path]] = {}
|
|
sample_metrics: dict[str, dict[str, int | float | str]] = {}
|
|
failures: list[str] = []
|
|
with ThreadPoolExecutor(max_workers=args.max_jobs) as pool:
|
|
futures = {
|
|
pool.submit(
|
|
_run_sample,
|
|
sample,
|
|
bam_by_sample[sample].path,
|
|
bundle,
|
|
args,
|
|
lengths,
|
|
): sample
|
|
for sample in sorted(expected_samples)
|
|
}
|
|
for future in as_completed(futures):
|
|
sample = futures[future]
|
|
try:
|
|
completed[sample], sample_metrics[sample] = future.result()
|
|
except Exception as exc:
|
|
failures.append(str(exc))
|
|
if failures or set(completed) != expected_samples:
|
|
raise RuntimeError(
|
|
"; ".join(failures)
|
|
or "BAM QC input/output sample-set conservation failed"
|
|
)
|
|
|
|
summary = bundle / f"{args.prefix}.{STAGE}.cohort.summary.tsv"
|
|
summary_columns = (
|
|
"sample_id",
|
|
"total_reference_bases",
|
|
"depth_sum",
|
|
"mean_depth",
|
|
"bases_ge_1x",
|
|
"coverage_ge_1x",
|
|
)
|
|
summary.write_text(
|
|
"\t".join(summary_columns)
|
|
+ "\n"
|
|
+ "".join(
|
|
"\t".join(
|
|
(
|
|
sample,
|
|
str(sample_metrics[sample]["total_reference_bases"]),
|
|
str(sample_metrics[sample]["depth_sum"]),
|
|
f"{float(sample_metrics[sample]['mean_depth']):.12g}",
|
|
str(sample_metrics[sample]["bases_ge_1x"]),
|
|
f"{float(sample_metrics[sample]['coverage_ge_1x']):.12g}",
|
|
)
|
|
)
|
|
+ "\n"
|
|
for sample in sorted(sample_metrics)
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
cohort_coverage = sum(
|
|
float(value["coverage_ge_1x"]) for value in sample_metrics.values()
|
|
) / len(sample_metrics)
|
|
report = bundle / f"{args.prefix}.{STAGE}.cohort.report.md"
|
|
report.write_text(
|
|
"# WGS BAM QC\n\n"
|
|
f"- Samples: {len(sample_metrics)}\n"
|
|
f"- Reference bases per sample: {sum(lengths.values())}\n"
|
|
f"- coverage_ge_1x: {cohort_coverage:.12g}\n"
|
|
"- Coverage denominator includes every reference base, including "
|
|
"zero-coverage positions.\n"
|
|
"- No sample was excluded based on coverage or depth.\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
destination = final / ".wgs_bam_qc_runs" / args.prefix / run_id
|
|
artifacts: dict[str, dict[str, Any]] = {}
|
|
for sample in sorted(completed):
|
|
for role, source in completed[sample].items():
|
|
artifacts[f"{role}.{sample}"] = _artifact(
|
|
workspace,
|
|
source,
|
|
destination / source.name,
|
|
role,
|
|
sample,
|
|
)
|
|
for role, source in (
|
|
("bam_qc_summary", summary),
|
|
("bam_qc_report", report),
|
|
):
|
|
artifacts[role] = _artifact(
|
|
workspace, source, destination / source.name, role
|
|
)
|
|
manifest = {
|
|
"manifest_version": "1.0",
|
|
"stage": STAGE,
|
|
"status": "success",
|
|
"prefix": args.prefix,
|
|
"job_id": args.job_id,
|
|
"run_id": run_id,
|
|
"start_time": started_at,
|
|
"workspace": str(workspace),
|
|
"wgs_sample_set_id": sample_set_id,
|
|
"wgs_reference_id": reference_id,
|
|
"upstream_wgs_analysis_ids": {
|
|
"wgs_markdup": markdup_analysis_id,
|
|
"wgs_index": index_analysis_id,
|
|
},
|
|
"wgs_analysis_id": analysis_id,
|
|
"parameters": parameters,
|
|
"software": {
|
|
"samtools": {"path": args.samtools, "version": version}
|
|
},
|
|
"metrics": {
|
|
"input_sample_ids": sorted(expected_samples),
|
|
"output_sample_ids": sorted(completed),
|
|
"total_samples": len(expected_samples),
|
|
"total_reference_bases": sum(lengths.values()),
|
|
"coverage_ge_1x": cohort_coverage,
|
|
"per_sample": [
|
|
sample_metrics[sample] for sample in sorted(sample_metrics)
|
|
],
|
|
},
|
|
"artifacts": artifacts,
|
|
}
|
|
manifest_work = attempt / "manifest.json"
|
|
_atomic_json(manifest_work, manifest)
|
|
if not _publish(bundle, destination, manifest_work, authority, started_at):
|
|
raise RuntimeError(
|
|
"stale WGS BAM QC attempt cannot replace a newer successful run"
|
|
)
|
|
return 0
|
|
except Exception as exc:
|
|
_publish_failure(authority, args, run_id, started_at, str(exc))
|
|
print(f"{STAGE} failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
return run(build_parser().parse_args(argv))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|