59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
import os
|
||
import shutil
|
||
from pathlib import Path
|
||
import sys
|
||
|
||
# 读取参数
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="合并 vina 拆分输出并检查转换结果")
|
||
parser.add_argument('--n_splits', type=int, default=12, help='拆分数量')
|
||
parser.add_argument('--out_dir', type=str, default='./result', help='meeko 转化出 sdf 文件所在目录')
|
||
parser.add_argument('--output_prefix', type=str, default='poses', help='每个子任务输出目录前缀')
|
||
parser.add_argument('--poses_dir', type=str, default='./result/poses_all', help='合并后SDF输出目录')
|
||
args = parser.parse_args()
|
||
|
||
n_splits = args.n_splits
|
||
out_dir = Path(args.out_dir)
|
||
poses_dir = Path(args.poses_dir)
|
||
poses_dir.mkdir(exist_ok=True)
|
||
|
||
|
||
merged_count = 0
|
||
# 合并所有 output{1..n_splits}/*.sdf 文件并复制
|
||
for i in range(1, n_splits + 1):
|
||
part_dir = out_dir / f'{args.output_prefix}{i}'
|
||
if part_dir.exists():
|
||
sdfs = list(part_dir.glob('*_converted.sdf'))
|
||
for f in sdfs:
|
||
shutil.copy2(f, poses_dir / f.name)
|
||
merged_count += 1
|
||
else:
|
||
print(f"Warning: {part_dir} does not exist.")
|
||
|
||
# 检查数量
|
||
sdf_stems = set(f.stem.replace('_converted', '') for f in poses_dir.glob('*_converted.sdf'))
|
||
pdbqt_stems = set(f.stem for f in out_dir.glob('*_out.pdbqt'))
|
||
|
||
print(f"pdbqt数量: {len(pdbqt_stems)}, 合并后sdf数量: {len(sdf_stems)}")
|
||
|
||
if len(sdf_stems) == len(pdbqt_stems) and sdf_stems == pdbqt_stems:
|
||
print("所有配体都已成功转换！")
|
||
else:
|
||
missing = tuple(sorted(pdbqt_stems - sdf_stems))
|
||
if missing:
|
||
print(f"缺少{len(missing)}个配体转换结果，缺失文件名如下：")
|
||
for name in missing:
|
||
print(name)
|
||
extra = tuple(sorted(sdf_stems - pdbqt_stems))
|
||
if extra:
|
||
print(f"有{len(extra)}个多余的SDF文件（没有对应的pdbqt），多余文件名如下：")
|
||
for name in extra:
|
||
print(name)
|
||
|
||
# 合并12
|
||
# python vina_merge_and_check.py
|
||
# 合并24
|
||
# python vina_merge_and_check.py --n_splits 128 --out_dir ./result --output_prefix poses --poses_dir ./result/poses_all
|
||
|