use_pymol_align_merge args add

modelbuilder.py

@@ -11,7 +11,6 @@
 import argparse  
 from dataclasses import dataclass, field
 from pathlib import Path
-import os
 import logging
 from analysis_pdb import PDBAnalyzer
 from build_modeller import PDBModeler
@@ -19,8 +18,6 @@ from modeller import ModellerError
 import pymol
 from typing import Dict
 from Bio.SeqRecord import SeqRecord
-from concurrent.futures import ProcessPoolExecutor
-from functools import partial
 
 @dataclass
 class PDBAlign:
@@ -58,6 +55,7 @@ class LoopModelBuilder:
     pymol_instance: object = field(init=False)
     analyzer_instance: PDBAnalyzer = field(init=False)
     pdb_id: str = field(init=False)
+    use_pymol_align_merge: bool = False  # 控制是否使用 PyMOL 进行对齐和合并的标志
 
     def __post_init__(self):
         self.pdb_id = self.pdb_file.stem
@@ -158,9 +156,12 @@ class LoopModelBuilder:
             # change id to original
             for i in modeller_results:
                 self.change_chain_identifier(i, 'A', mc, split=False)
-            # use pymol to align and merge
+            # Use PyMOL to align and merge, controlled by use_pymol_align_merge flag
+            if self.use_pymol_align_merge: 
+                self.logger.warning("PyMOL alignment and merging is disabled or multiple model files exist. "
+                    "This may result in clashes between multiple protein chains, leading to failures in MD simulations. "
+                    "Please ensure the structures are compatible for simulations.")
                 if len(modeller_results) == 1:
-                # use pymol to align
                     aligner = PDBAlign(self.pdb_file, modeller_results[0])
                     pdbstr = aligner.align()
                     mc_dict.update({mc: pdbstr})
@@ -200,6 +201,7 @@ def main():
                         help="Model refinement type (default: refine.very_fast)")
     parser.add_argument("--output_dir", type=str, default=None, 
                         help="Output directory (default: same as PDB file directory)")
+    parser.add_argument("--use_pymol_align_merge", action='store_true', help="Use PyMOL for alignment and merging. Note: This may result in clashes between multiple protein chains, leading to failures in MD simulations.")
 
     # 解析参数
     args = parser.parse_args()
@@ -208,7 +210,7 @@ def main():
     pdb_path = Path(args.pdb_file).absolute()
     output_dir = Path(args.output_dir).absolute() if args.output_dir else None
 
-    model_builder = LoopModelBuilder(pdb_file=pdb_path, output_dir=output_dir)
+    model_builder = LoopModelBuilder(pdb_file=pdb_path, output_dir=output_dir, use_pymol_align_merge=args.use_pymol_align_merge)
     result_path = model_builder.run(typestr=args.typestr)
 
     print(f"Modeling completed. Output file: {result_path}")