From ef30cc99088bddc5abc10bb9af332740887cd301 Mon Sep 17 00:00:00 2001
From: lingyuzeng <pylyzeng@gmail.com>
Date: Wed, 6 Mar 2024 18:05:57 +0800
Subject: [PATCH] use_pymol_align_merge args add

---
 modelbuilder.py | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/modelbuilder.py b/modelbuilder.py
index 109ff9f..9469689 100644
--- a/modelbuilder.py
+++ b/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,13 +18,11 @@ 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:
     template_file: Path
-    target_file: Path  
+    target_file: Path
     pymol_instance: object = field(init=False)
     out_file: Path = field(default=None)  # 输出文件路径
 
@@ -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,14 +156,17 @@ 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
-            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})
-            else:
-                self.logger.warning('more than one model file, please set num_loop to 1')
+            # 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:
+                    aligner = PDBAlign(self.pdb_file, modeller_results[0])
+                    pdbstr = aligner.align()
+                    mc_dict.update({mc: pdbstr})
+                else:
+                    self.logger.warning('more than one model file, please set num_loop to 1')
         return mc_dict
     
     @staticmethod
@@ -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}")