add contral args

mainfix.py

@@ -2,22 +2,22 @@
 # -*- encoding: utf-8 -*-
 '''
 @file               :mainfix.py
-@Description:       :批量调用modelbuilder.py进行结构修复
+@Description        :批量调用modelbuilder.py进行结构修复
 @Date               :2024/01/12 13:38:53
 @Author             :lyzeng
 @Email              :pylyzeng@gmail.com
 @version            :1.0
 '''
 
-
 import argparse
 from pathlib import Path
 import subprocess
 import shutil
 import concurrent.futures
 import os
+import sys
 
-def process_pdb_file(pdb_file: Path, script_path: Path, working_dir: Path):
+def process_pdb_file(pdb_file: Path, script_path: Path, working_dir: Path, python_interpreter: str, typestr="refine.very_fast", use_pymol_align_merge=False):
     pdb_id = pdb_file.stem
     runner_dir = working_dir / f"runner_{pdb_id}"
     runner_dir.mkdir(exist_ok=True)
@@ -26,11 +26,27 @@ def process_pdb_file(pdb_file: Path, script_path: Path, working_dir: Path):
     shutil.copy(pdb_file, copied_pdb_file)
 
     # 构建并运行modelbuilder.py脚本的命令
-    command = ["/home/era/.conda/envs/pyfastx/bin/python", script_path.absolute().as_posix(), copied_pdb_file.absolute().as_posix()]
+    command = [python_interpreter, script_path.absolute().as_posix(), copied_pdb_file.absolute().as_posix()]
+
+    # 添加额外的参数
+    command += ["--typestr", typestr]
+    if use_pymol_align_merge:
+        command += ["--use_pymol_align_merge"]
+
+    # 运行modelbuilder.py脚本
     subprocess.run(command, cwd=runner_dir)
 
+def main():
+    parser = argparse.ArgumentParser(description="Process multiple PDB files using modelbuilder.py")
+    parser.add_argument("pdb_dir_path", help="Directory containing PDB files")
+    parser.add_argument("script_path", help="Path to the modelbuilder.py script")
+    parser.add_argument("--python_interpreter", default=sys.executable, help="Path to the Python interpreter")
+    args = parser.parse_args()
+
+    pdb_dir_path = args.pdb_dir_path
+    script_path = args.script_path
+    python_interpreter = args.python_interpreter
 
-def main(pdb_dir_path, script_path):
     pdb_dir = Path(pdb_dir_path)
     pdb_files = list(pdb_dir.glob("*.pdb"))
 
@@ -39,19 +55,15 @@ def main(pdb_dir_path, script_path):
     max_workers = max(1, os.cpu_count() - 1)  # 保留一个核心给系统
 
     with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-        futures = [executor.submit(process_pdb_file, pdb_file, script_path, pdb_dir) for pdb_file in pdb_files]
+        futures = [executor.submit(process_pdb_file, pdb_file, script_path, pdb_dir, python_interpreter) for pdb_file in pdb_files]
         for future in concurrent.futures.as_completed(futures):
             future.result()
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Process multiple PDB files using modelbuilder.py")
+    main()
-    parser.add_argument("pdb_dir_path", help="Directory containing PDB files")
-    parser.add_argument("script_path", help="Path to the modelbuilder.py script")
-    args = parser.parse_args()
 
-    main(args.pdb_dir_path, args.script_path)
 
-# python mainfix.py pdb_test1 modelbuilder.py
+# python mainfix.py pdb_test1 modelbuilder.py --python_interpreter /root/micromamba/envs/pyfastx/bin/python
 # 复制修复完成文件到指定文件夹
 # import os
 # import shutil