106 lines
5.0 KiB
Python
106 lines
5.0 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_one(accession: str, proteins: Path, family_hmm: Path, output_dir: Path, cpu: int, parser_script: Path) -> dict[str, str | int]:
|
|
domtbl = output_dir / f"{accession}.domtblout"
|
|
log = output_dir / f"{accession}.log"
|
|
hits = output_dir / f"{accession}.hits.tsv"
|
|
command = ["hmmsearch", "--cpu", str(cpu), "-E", "1e-5", "--domE", "1e-5", "--domtblout", str(domtbl), str(family_hmm), str(proteins)]
|
|
with log.open("w", encoding="utf-8") as log_handle:
|
|
completed = subprocess.run(command, stdout=log_handle, stderr=subprocess.PIPE, text=True)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(f"{accession} hmmsearch failed:\n{completed.stderr}")
|
|
parsed = subprocess.run([
|
|
sys.executable, str(parser_script), "--input", str(domtbl), "--mode", "hmmsearch",
|
|
"--max-ievalue", "1e-5", "--min-hmm-coverage", "0.50", "--output", str(hits),
|
|
], text=True, capture_output=True)
|
|
if parsed.returncode != 0:
|
|
raise RuntimeError(f"{accession} HMM parsing failed:\n{parsed.stderr or parsed.stdout}")
|
|
return {"accession": accession, "domtblout": domtbl.name, "hits": hits.name, "log": log.name, "hits_bytes": hits.stat().st_size}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run HMM family searches for every proteome in a primary bundle")
|
|
parser.add_argument("--primary-bundle", required=True, type=Path)
|
|
parser.add_argument("--family-hmm", required=True, type=Path)
|
|
parser.add_argument("--jobs", type=int, default=4, help="maximum concurrent samples; does not limit sample count")
|
|
parser.add_argument("--cpu-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.cpu_per_job < 1:
|
|
parser.error("--jobs and --cpu-per-job must be at least 1")
|
|
archive_output = Path(f"{args.output_prefix}.hmm.tar.gz")
|
|
manifest_output = Path(f"{args.output_prefix}.hmm_manifest.tsv")
|
|
archive_output.parent.mkdir(parents=True, exist_ok=True)
|
|
parser_script = Path(__file__).with_name("parse_hmmer.py").resolve()
|
|
with tempfile.TemporaryDirectory(prefix="hmm_batch_") as temp_name:
|
|
temp = Path(temp_name)
|
|
extract_root = temp / "input"
|
|
output_dir = temp / "hmm"
|
|
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.family_hmm.resolve(), output_dir, args.cpu_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 HMM search: {result['accession']}", flush=True)
|
|
fields = ["accession", "domtblout", "hits", "log", "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"hmm/{file_path.name}")
|
|
print(f"processed {len(proteomes)} samples")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|