分子对接与虚拟筛选

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
vs_batch.py —— 批量虚拟筛选模板（AutoDock Vina 命令行 + 多进程 + 结果汇总）

用途
    给一个已经准备好的受体（PDBQT）和一个装满配体（PDBQT）的文件夹，
    用多进程并行跑 Vina，把每个配体的最优打分汇总成一张 CSV。

依赖
    - AutoDock Vina 可执行文件（vina 或 vina.exe），并已加入 PATH
      安装：conda install -c conda-forge vina  或  pip install -U vina
    - 仅需 Python 标准库（subprocess / multiprocessing / csv / argparse / re）

典型用法
    # 1) 先准备受体（含极性氢）与配体（PDBQT）
    python vs_batch.py \
        --receptor 1iep_receptor.pdbqt \
        --ligands  ligands/ \
        --config   1iep_receptor.box.txt \
        --out-dir  poses \
        --csv      results.csv \
        --nproc    8 \
        --exhaustiveness 16 \
        --top-n    5

    # 2) 只想重新汇总已有结果，不重跑对接
    python vs_batch.py --ligands ligands/ --out-dir poses --csv results.csv --summarize-only

config 文件（--config 指向的文本文件）内容示例：
    center_x = 15.190
    center_y = 53.903
    center_z = 16.917
    size_x   = 20.0
    size_y   = 20.0
    size_z   = 20.0

若没有 config 文件，也可以直接用命令行给盒子：
    --center 15.190 53.903 16.917 --size 20 20 20
"""

import argparse
import csv
import os
import re
import shutil
import subprocess
import sys
from collections import Counter
from concurrent.futures import ProcessPoolExecutor, as_completed
from glob import glob

# 匹配输出 PDBQT 里的结果行，例如：
# REMARK VINA RESULT:    -13.2      0.000      0.000
_RESULT_RE = re.compile(
    r"^REMARK\s+VINA RESULT:\s+(-?\d+\.?\d*)\s+(-?\d+\.?\d*)\s+(-?\d+\.?\d*)"
)


def parse_scores(pdbqt_path):
    """读取一个 Vina 输出 PDBQT，返回 [(affinity, rmsd_lb, rmsd_ub), ...]（按出现顺序=按打分排序）。"""
    scores = []
    try:
        with open(pdbqt_path, "r", errors="ignore") as fh:
            for line in fh:
                m = _RESULT_RE.match(line.strip())
                if m:
                    scores.append((float(m.group(1)), float(m.group(2)), float(m.group(3))))
    except OSError:
        return []
    return scores


def find_ligands(paths):
    """把 --ligands 传进来的目录 / glob / 文件统一展开成 PDBQT 文件绝对路径列表。"""
    files = []
    for p in paths:
        if os.path.isdir(p):
            files.extend(glob(os.path.join(p, "*.pdbqt")))
        elif any(ch in p for ch in "*?["):
            files.extend(glob(p))
        elif os.path.isfile(p):
            files.append(p)
    # 去重并排序，保证可复现
    files = sorted(set(os.path.abspath(f) for f in files))
    # 排除已经带 _out 的结果文件，避免重复对接
    return [f for f in files if not os.path.basename(f).endswith("_out.pdbqt")]


def assign_stems(ligands):
    """给每个配体分配唯一的输出名（stem），避免不同目录下的同名配体互相覆盖。

    基线是文件名去掉扩展名；同一个 basename 出现多次时，统一追加 __1、__2 … 后缀。
    返回 (stems, dups)：
        stems —— {配体绝对路径: 输出名}
        dups  —— 出现冲突的 basename 列表（排序后），供提示用
    """
    bases = [os.path.splitext(os.path.basename(f))[0] for f in ligands]
    total = Counter(bases)
    seen = Counter()
    stems = {}
    for f, b in zip(ligands, bases):
        if total[b] == 1:
            stems[f] = b
        else:
            seen[b] += 1
            stems[f] = "%s__%d" % (b, seen[b])
    dups = sorted(b for b, c in total.items() if c > 1)
    return stems, dups


def build_box_config(box_path, center, size):
    """没有现成 config 时，按 center/size 生成一个临时 config 文件。"""
    with open(box_path, "w") as fh:
        fh.write("center_x = %.3f\n" % center[0])
        fh.write("center_y = %.3f\n" % center[1])
        fh.write("center_z = %.3f\n" % center[2])
        fh.write("size_x = %.1f\n" % size[0])
        fh.write("size_y = %.1f\n" % size[1])
        fh.write("size_z = %.1f\n" % size[2])
    return box_path


def run_one(job):
    """单个配体对接。job 是一个 dict，方便进程池序列化。"""
    ligand = job["ligand"]
    name = job.get("stem") or os.path.splitext(os.path.basename(ligand))[0]
    out_file = os.path.join(job["out_dir"], name + "_out.pdbqt")
    log_file = os.path.join(job["out_dir"], name + ".log")

    cmd = [
        job["vina"],
        "--receptor", job["receptor"],
        "--ligand", ligand,
        "--config", job["config"],
        "--out", out_file,
        "--exhaustiveness", str(job["exhaustiveness"]),
        "--cpu", str(job["cpu_per_job"]),
    ]
    if job["scoring"]:
        cmd += ["--scoring", job["scoring"]]
    if job["seed"] is not None:
        cmd += ["--seed", str(job["seed"])]

    try:
        with open(log_file, "w") as log:
            proc = subprocess.run(
                cmd, stdout=log, stderr=subprocess.STDOUT,
                timeout=job["timeout"], check=False,
            )
    except subprocess.TimeoutExpired:
        return {"ligand": name, "status": "timeout", "out_file": out_file}
    except FileNotFoundError:
        return {"ligand": name, "status": "vina-not-found", "out_file": out_file}

    if proc.returncode != 0 or not os.path.isfile(out_file):
        return {"ligand": name, "status": "failed(rc=%d)" % proc.returncode, "out_file": out_file}

    scores = parse_scores(out_file)
    if not scores:
        return {"ligand": name, "status": "no-result", "out_file": out_file}

    return {
        "ligand": name,
        "status": "ok",
        "out_file": out_file,
        "best_affinity": scores[0][0],
        "best_rmsd_lb": scores[0][1],
        "best_rmsd_ub": scores[0][2],
        "n_poses": len(scores),
    }


def summarize(out_dir, csv_path, top_n):
    """扫描 out_dir 下所有 *_out.pdbqt，汇总最优打分到 CSV。"""
    rows = []
    for out_file in sorted(glob(os.path.join(out_dir, "*_out.pdbqt"))):
        name = os.path.basename(out_file)[:-len("_out.pdbqt")]
        scores = parse_scores(out_file)
        if scores:
            best = scores[0]
            rows.append({
                "ligand": name,
                "best_affinity": best[0],
                "best_rmsd_lb": best[1],
                "best_rmsd_ub": best[2],
                "n_poses": len(scores),
                "out_file": out_file,
            })
    rows.sort(key=lambda r: r["best_affinity"])  # 越负越靠前

    with open(csv_path, "w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(["rank", "ligand", "best_affinity_kcal_mol",
                         "rmsd_lb", "rmsd_ub", "n_poses", "out_file"])
        for i, r in enumerate(rows, 1):
            writer.writerow([i, r["ligand"], r["best_affinity"],
                             r["best_rmsd_lb"], r["best_rmsd_ub"],
                             r["n_poses"], r["out_file"]])
    return rows


def main():
    ap = argparse.ArgumentParser(description="AutoDock Vina 批量虚拟筛选模板")
    ap.add_argument("--receptor", help="受体 PDBQT 文件")
    ap.add_argument("--ligands", nargs="+", required=True,
                    help="配体 PDBQT：目录 / glob / 文件，可给多个")
    ap.add_argument("--config", help="Vina config 文本（含 center/size）")
    ap.add_argument("--center", nargs=3, type=float, metavar=("X", "Y", "Z"),
                    help="盒子中心；与 --size 搭配，替代 --config")
    ap.add_argument("--size", nargs=3, type=float, metavar=("X", "Y", "Z"),
                    help="盒子边长（埃）；与 --center 搭配")
    ap.add_argument("--out-dir", default="poses", help="输出目录（默认 poses）")
    ap.add_argument("--csv", default="vs_results.csv", help="汇总 CSV 路径")
    ap.add_argument("--nproc", type=int, default=os.cpu_count() or 1,
                    help="并行进程数（同时跑多少个配体）")
    ap.add_argument("--cpu-per-job", type=int, default=1,
                    help="每个 vina 进程内部使用几个 CPU 核（默认 1）")
    ap.add_argument("--exhaustiveness", type=int, default=16,
                    help="搜索强度；越大越慢越稳（默认 16，正式筛选可用 32）")
    ap.add_argument("--scoring", default=None, choices=["vina", "ad4", "vinardo"],
                    help="打分函数；默认用 vina。用 ad4 时需先算 affinity maps 并改用 --maps")
    ap.add_argument("--seed", type=int, default=None, help="随机种子，固定可复现")
    ap.add_argument("--timeout", type=int, default=3600, help="单个配体超时秒数")
    ap.add_argument("--top-n", type=int, default=0, help="打印前 N 名（0=不打印）")
    ap.add_argument("--summarize-only", action="store_true",
                    help="跳过对接，只重新汇总 out-dir 里已有结果")
    ap.add_argument("--vina", default=None, help="vina 可执行文件路径（默认从 PATH 找）")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    vina_bin = args.vina or shutil.which("vina") or "vina"

    if args.summarize_only:
        rows = summarize(args.out_dir, args.csv, args.top_n)
        print("汇总完成：%d 条结果 -> %s" % (len(rows), args.csv))
        _print_top(rows, args.top_n)
        return

    if not args.receptor:
        ap.error("需要 --receptor（受体 PDBQT）")
    if not args.config:
        if not (args.center and args.size):
            ap.error("需要 --config，或同时给出 --center 与 --size")
        args.config = build_box_config(os.path.join(args.out_dir, "_box.txt"),
                                       args.center, args.size)

    ligands = find_ligands(args.ligands)
    if not ligands:
        print("没有找到任何配体 PDBQT，检查 --ligands 路径。", file=sys.stderr)
        sys.exit(1)

    # 提前校验输入：这些问题逐个配体报一遍既费时间又难定位，直接拦在前面
    if not os.path.isfile(args.receptor):
        ap.error("找不到受体文件：%s" % args.receptor)
    if not os.path.isfile(args.config):
        ap.error("找不到 config 文件：%s" % args.config)
    if not (os.path.isfile(vina_bin) or shutil.which(vina_bin)):
        ap.error("找不到 vina 可执行文件（%s）。请先 conda install -c conda-forge vina，"
                 "或用 --vina 指定完整路径。" % vina_bin)

    # 不同目录下的同名配体输出会互相覆盖，这里统一改名规避
    stems, dups = assign_stems(ligands)
    if dups:
        print("提示：以下配体重名，输出名已自动加 __N 后缀避免覆盖：%s"
              % ", ".join(dups))

    print("vina        : %s" % vina_bin)
    print("受体        : %s" % os.path.abspath(args.receptor))
    print("config      : %s" % os.path.abspath(args.config))
    print("配体数量    : %d" % len(ligands))
    print("并行进程    : %d（每进程 %d 核）" % (args.nproc, args.cpu_per_job))
    print("搜索强度    : %d" % args.exhaustiveness)

    jobs = [{
        "ligand": lig,
        "stem": stems[lig],
        "receptor": os.path.abspath(args.receptor),
        "config": os.path.abspath(args.config),
        "out_dir": os.path.abspath(args.out_dir),
        "vina": vina_bin,
        "exhaustiveness": args.exhaustiveness,
        "cpu_per_job": args.cpu_per_job,
        "scoring": args.scoring,
        "seed": args.seed,
        "timeout": args.timeout,
    } for lig in ligands]

    done = 0
    results = []
    with ProcessPoolExecutor(max_workers=args.nproc) as pool:
        futures = {pool.submit(run_one, j): j for j in jobs}
        for fut in as_completed(futures):
            r = fut.result()
            results.append(r)
            done += 1
            tag = "%s(%s)" % (r["ligand"], r["status"])
            print("[%d/%d] %s" % (done, len(jobs), tag))

    rows = summarize(args.out_dir, args.csv, args.top_n)
    ok = sum(1 for r in results if r["status"] == "ok")
    print("\n对接结束：成功 %d / 共 %d，汇总 -> %s" % (ok, len(jobs), args.csv))
    _print_top(rows, args.top_n)


def _print_top(rows, top_n):
    if top_n and rows:
        print("\n打分最优的前 %d 个配体（kcal/mol，越负越可能结合）：" % top_n)
        for i, r in enumerate(rows[:top_n], 1):
            print("  %2d. %-30s %8.2f" % (i, r["ligand"], r["best_affinity"]))


if __name__ == "__main__":
    main()

# ---------------------------------------------------------------
# 环境依赖（批量虚拟筛选）
# 建议用 conda / mamba 建独立环境，避免与系统 Python 冲突
# ---------------------------------------------------------------

# 对接引擎（同时提供命令行 vina 和 Python API）
# conda 更稳；也可 pip install -U vina
vina>=1.2.5

# 分子读写与构象
numpy
scipy
rdkit

# 受体 / 配体 PDBQT 准备（提供 mk_prepare_ligand.py、mk_prepare_receptor.py、mk_export.py）
meeko>=0.6

# 给配体加氢 / 生成 3D 构象 / 枚举质子化与互变异构（提供 scrub.py）
# 注意：scrub.py 内部 import joblib，但 molscrub 没把它写进依赖，必须单独装，
#       否则跑 scrub.py 会报 ModuleNotFoundError: No module named 'joblib'
molscrub
joblib

# 蛋白结构处理（官方 Colab 工作流里用到）
gemmi
prody

# 可视化（可选，用于看结果）
py3Dmol

# 说明：
# - AutoDock-GPU（GPU 加速大规模对接）：https://github.com/ccsb-scripps/AutoDock-GPU
# - Ringtail（虚拟筛选结果入库与富集分析）：https://github.com/forlilab/ringtail
# - ad4 打分需要 autogrid4 预计算 affinity maps：conda install -c conda-forge autogrid