112 lines
5.4 KiB
Python
112 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
|
|
|
|
PRIMARY_SUFFIX = ".primary.pep.fa"
|
|
SAFE_ACCESSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
|
|
|
|
|
|
def extract_archive(archive_path: Path, destination: Path) -> None:
|
|
destination_root = destination.resolve()
|
|
with tarfile.open(archive_path, "r:*") as archive:
|
|
for member in archive.getmembers():
|
|
target = (destination / member.name).resolve()
|
|
if target != destination_root and destination_root not in target.parents:
|
|
raise ValueError(f"unsafe archive member: {member.name}")
|
|
archive.extractall(destination)
|
|
|
|
|
|
def discover_proteomes(root: Path) -> dict[str, Path]:
|
|
found: dict[str, Path] = {}
|
|
for path in root.rglob(f"*{PRIMARY_SUFFIX}"):
|
|
accession = path.name[:-len(PRIMARY_SUFFIX)]
|
|
if not SAFE_ACCESSION.fullmatch(accession):
|
|
raise ValueError(f"unsafe accession derived from {path.name}")
|
|
if accession in found:
|
|
raise ValueError(f"duplicate primary protein file for {accession}")
|
|
found[accession] = path
|
|
if not found:
|
|
raise ValueError("primary bundle contains no *.primary.pep.fa files")
|
|
return found
|
|
|
|
|
|
def run_command(command: list[str], accession: str, label: str) -> None:
|
|
completed = subprocess.run(command, text=True, capture_output=True)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(f"{accession} {label} failed:\n{completed.stderr or completed.stdout}")
|
|
|
|
|
|
def run_one(accession: str, proteins: Path, reference: Path, output_dir: Path, threads: int, parser_script: Path) -> dict[str, str | int]:
|
|
db_prefix = output_dir / accession
|
|
raw = output_dir / f"{accession}.raw.tsv"
|
|
hits = output_dir / f"{accession}.hits.tsv"
|
|
run_command(["diamond", "makedb", "--in", str(proteins), "--db", str(db_prefix), "--threads", str(threads)], accession, "makedb")
|
|
run_command([
|
|
"diamond", "blastp", "--query", str(reference), "--db", str(db_prefix), "--out", str(raw),
|
|
"--outfmt", "6", "qseqid", "sseqid", "pident", "length", "mismatch", "gapopen", "qstart", "qend", "sstart", "send", "evalue", "bitscore", "qlen", "slen",
|
|
"--evalue", "1e-5", "--more-sensitive", "--max-target-seqs", "1000", "--threads", str(threads),
|
|
], accession, "blastp")
|
|
run_command([
|
|
sys.executable, str(parser_script), "--input", str(raw), "--output", str(hits),
|
|
"--max-evalue", "1e-5", "--min-query-coverage", "0.50", "--min-subject-coverage", "0.30", "--min-identity", "25",
|
|
], accession, "DIAMOND parsing")
|
|
database = Path(f"{db_prefix}.dmnd")
|
|
return {"accession": accession, "database": database.name, "raw_hits": raw.name, "filtered_hits": hits.name, "hits_bytes": hits.stat().st_size}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run DIAMOND family searches for every proteome in a primary bundle")
|
|
parser.add_argument("--primary-bundle", required=True, type=Path)
|
|
parser.add_argument("--reference", required=True, type=Path)
|
|
parser.add_argument("--jobs", type=int, default=4, help="maximum concurrent samples; does not limit sample count")
|
|
parser.add_argument("--threads-per-job", type=int, default=2)
|
|
parser.add_argument("--output-prefix", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
if args.jobs < 1 or args.threads_per_job < 1:
|
|
parser.error("--jobs and --threads-per-job must be at least 1")
|
|
archive_output = Path(f"{args.output_prefix}.homology.tar.gz")
|
|
manifest_output = Path(f"{args.output_prefix}.homology_manifest.tsv")
|
|
archive_output.parent.mkdir(parents=True, exist_ok=True)
|
|
parser_script = Path(__file__).with_name("parse_diamond.py").resolve()
|
|
with tempfile.TemporaryDirectory(prefix="diamond_batch_") as temp_name:
|
|
temp = Path(temp_name)
|
|
extract_root = temp / "input"
|
|
output_dir = temp / "homology"
|
|
extract_root.mkdir()
|
|
output_dir.mkdir()
|
|
extract_archive(args.primary_bundle.resolve(), extract_root)
|
|
proteomes = discover_proteomes(extract_root)
|
|
results: list[dict[str, str | int]] = []
|
|
with ThreadPoolExecutor(max_workers=min(args.jobs, len(proteomes))) as executor:
|
|
futures = {
|
|
executor.submit(run_one, accession, proteins, args.reference.resolve(), output_dir, args.threads_per_job, parser_script): accession
|
|
for accession, proteins in proteomes.items()
|
|
}
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
results.append(result)
|
|
print(f"completed DIAMOND search: {result['accession']}", flush=True)
|
|
fields = ["accession", "database", "raw_hits", "filtered_hits", "hits_bytes"]
|
|
with manifest_output.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(sorted(results, key=lambda row: str(row["accession"])))
|
|
with tarfile.open(archive_output, "w:gz") as archive:
|
|
for file_path in sorted(output_dir.iterdir()):
|
|
archive.add(file_path, arcname=f"homology/{file_path.name}")
|
|
print(f"processed {len(proteomes)} samples")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|