214 lines
6.6 KiB
Python
Executable File
214 lines
6.6 KiB
Python
Executable File
# pipelines/blast_pipeline.py
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
from dataclasses import dataclass, asdict
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any, List, Tuple
|
|
|
|
|
|
@dataclass
|
|
class BlastOutputs:
|
|
out_dir: str
|
|
db_fasta: str
|
|
blastout: str
|
|
filtered_tsv: str
|
|
id_list: str
|
|
hit_count: int
|
|
unique_id_count: int
|
|
params: Dict[str, Any]
|
|
|
|
|
|
def _ensure_exe(name: str) -> str:
|
|
exe = shutil.which(name)
|
|
if not exe:
|
|
raise RuntimeError(f"找不到可执行程序: {name} (请确认已安装并在 PATH 中)")
|
|
return exe
|
|
|
|
|
|
def _parse_blast_line(line: str) -> Optional[Dict[str, Any]]:
|
|
parts = line.rstrip("\n").split("\t")
|
|
if len(parts) < 14:
|
|
return None
|
|
try:
|
|
qseqid = parts[0]
|
|
sseqid = parts[1]
|
|
pident = float(parts[2])
|
|
length = int(parts[3])
|
|
evalue = float(parts[10])
|
|
bitscore = float(parts[11])
|
|
qlen = int(parts[12])
|
|
slen = int(parts[13])
|
|
qcov = (length / qlen * 100.0) if qlen else 0.0
|
|
scov = (length / slen * 100.0) if slen else 0.0
|
|
return {
|
|
"qseqid": qseqid,
|
|
"sseqid": sseqid,
|
|
"pident": pident,
|
|
"length": length,
|
|
"evalue": evalue,
|
|
"bitscore": bitscore,
|
|
"qlen": qlen,
|
|
"slen": slen,
|
|
"qcov": qcov,
|
|
"scov": scov,
|
|
"raw": parts,
|
|
}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def run_blastp_pipeline(
|
|
reference_fasta: str,
|
|
query_fasta: str,
|
|
out_dir: str,
|
|
prefix: str,
|
|
evalue_max: float = 1e-5,
|
|
pident_min: float = 30.0,
|
|
qcov_min: float = 40.0,
|
|
scov_min: float = 40.0,
|
|
threads: int = 8,
|
|
max_target_seqs: int = 50,
|
|
) -> BlastOutputs:
|
|
"""
|
|
makeblastdb + blastp + 过滤 + 输出 id_list
|
|
输出文件：
|
|
{prefix}.blastout
|
|
{prefix}.blastout.filter.tsv (带 qcov/scov)
|
|
{prefix}.blastout.filter.id
|
|
{prefix}.summary.json
|
|
"""
|
|
_ensure_exe("makeblastdb")
|
|
_ensure_exe("blastp")
|
|
|
|
outp = Path(out_dir).resolve()
|
|
outp.mkdir(parents=True, exist_ok=True)
|
|
|
|
ref = Path(reference_fasta).resolve()
|
|
qry = Path(query_fasta).resolve()
|
|
|
|
blastout = (outp / f"{prefix}.blastout").resolve()
|
|
filtered_tsv = (outp / f"{prefix}.blastout.filter.tsv").resolve()
|
|
id_list = (outp / f"{prefix}.blastout.filter.id").resolve()
|
|
summary_json = (outp / f"{prefix}.summary.json").resolve()
|
|
|
|
# 1) 建库：强制把数据库生成在 out_dir，并使用固定前缀 db_prefix
|
|
db_prefix = (outp / Path(ref).stem).resolve() # 例如 .../blast/SPL_Ath
|
|
pin = Path(str(db_prefix) + ".pin")
|
|
psq = Path(str(db_prefix) + ".psq")
|
|
phr = Path(str(db_prefix) + ".phr")
|
|
|
|
if not (pin.exists() and psq.exists() and phr.exists()):
|
|
cmd_db = [
|
|
"makeblastdb",
|
|
"-in", str(ref),
|
|
"-dbtype", "prot",
|
|
"-out", str(db_prefix),
|
|
]
|
|
proc_db = subprocess.run(cmd_db, cwd=str(outp), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
if proc_db.returncode != 0:
|
|
raise RuntimeError(f"makeblastdb 失败\nCMD: {' '.join(cmd_db)}\nSTDERR:\n{proc_db.stderr[:4000]}")
|
|
# blastp 的 -db 使用 reference fasta 名称即可（在 out_dir 里）
|
|
db_for_blast = str(db_prefix)
|
|
|
|
# 2) blastp
|
|
outfmt = "6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore qlen slen"
|
|
cmd_blast = [
|
|
"blastp",
|
|
"-query", str(qry),
|
|
"-db", db_for_blast,
|
|
"-evalue", str(evalue_max),
|
|
"-outfmt", outfmt,
|
|
"-max_target_seqs", str(max_target_seqs),
|
|
"-num_threads", str(threads),
|
|
"-out", str(blastout),
|
|
]
|
|
proc = subprocess.run(cmd_blast, cwd=str(outp), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"blastp 失败\nCMD: {' '.join(cmd_blast)}\nSTDERR:\n{proc.stderr[:4000]}")
|
|
|
|
# 3) 过滤 + 每个 qseqid 只保留 best hit
|
|
best: Dict[str, Dict[str, Any]] = {}
|
|
kept_rows: List[Dict[str, Any]] = []
|
|
|
|
with blastout.open("r", encoding="utf-8", errors="ignore") as f:
|
|
for line in f:
|
|
if not line.strip():
|
|
continue
|
|
rec = _parse_blast_line(line)
|
|
if not rec:
|
|
continue
|
|
|
|
if rec["pident"] < pident_min:
|
|
continue
|
|
if rec["qcov"] < qcov_min:
|
|
continue
|
|
if rec["scov"] < scov_min:
|
|
continue
|
|
|
|
q = rec["qseqid"]
|
|
cur = best.get(q)
|
|
if cur is None:
|
|
best[q] = rec
|
|
else:
|
|
# bitscore 越大越好；若相同，evalue 越小越好
|
|
if (rec["bitscore"] > cur["bitscore"]) or (
|
|
rec["bitscore"] == cur["bitscore"] and rec["evalue"] < cur["evalue"]
|
|
):
|
|
best[q] = rec
|
|
|
|
kept_rows = list(best.values())
|
|
kept_rows.sort(key=lambda x: (-x["bitscore"], x["evalue"]))
|
|
|
|
# 写 filter.tsv
|
|
with filtered_tsv.open("w", encoding="utf-8") as w:
|
|
w.write("\t".join(["qseqid", "sseqid", "pident", "evalue", "bitscore", "qcov", "scov", "aln_len", "qlen", "slen"]) + "\n")
|
|
for r in kept_rows:
|
|
w.write(
|
|
"\t".join([
|
|
r["qseqid"], r["sseqid"],
|
|
f"{r['pident']:.3f}",
|
|
f"{r['evalue']:.3g}",
|
|
f"{r['bitscore']:.1f}",
|
|
f"{r['qcov']:.1f}",
|
|
f"{r['scov']:.1f}",
|
|
str(r["length"]),
|
|
str(r["qlen"]),
|
|
str(r["slen"]),
|
|
]) + "\n"
|
|
)
|
|
|
|
# 写 id list
|
|
ids = sorted({r["qseqid"] for r in kept_rows})
|
|
with id_list.open("w", encoding="utf-8") as w:
|
|
for i in ids:
|
|
w.write(i + "\n")
|
|
|
|
outputs = BlastOutputs(
|
|
out_dir=str(outp),
|
|
db_fasta=str(ref),
|
|
blastout=str(blastout),
|
|
filtered_tsv=str(filtered_tsv),
|
|
id_list=str(id_list),
|
|
hit_count=len(kept_rows),
|
|
unique_id_count=len(ids),
|
|
params={
|
|
"reference_fasta": str(ref),
|
|
"query_fasta": str(qry),
|
|
"prefix": prefix,
|
|
"evalue_max": evalue_max,
|
|
"pident_min": pident_min,
|
|
"qcov_min": qcov_min,
|
|
"scov_min": scov_min,
|
|
"threads": threads,
|
|
"max_target_seqs": max_target_seqs,
|
|
},
|
|
)
|
|
|
|
with summary_json.open("w", encoding="utf-8") as w:
|
|
json.dump(asdict(outputs), w, ensure_ascii=False, indent=2)
|
|
|
|
return outputs
|