analysis_pdb/modelbuilder.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :modelbuilder.py
@Description:       :
@Date               :2024/01/11 15:29:28
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''
import argparse  
from dataclasses import dataclass, field
from pathlib import Path
import logging
from analysis_pdb import PDBAnalyzer
from build_modeller import PDBModeler
from modeller import ModellerError
import pymol
from typing import Dict
from Bio.SeqRecord import SeqRecord

@dataclass
class PDBAlign:
    template_file: Path
    target_file: Path
    pymol_instance: object = field(init=False)
    out_file: Path = field(default=None)  # 输出文件路径

    def __post_init__(self):
        self.initialize_pymol()

    def initialize_pymol(self):
        self.pymol_instance = pymol.cmd
        self.pymol_instance.reinitialize()

    def align(self) -> str:
        self.pymol_instance.reinitialize()
        # 首先，加载模板结构
        self.pymol_instance.load(self.template_file.as_posix(), "template")

        # 加载并对齐所有目标结构
        self.pymol_instance.load(self.target_file.as_posix(), "target")
        self.pymol_instance.align("target", "template")

        return self.pymol_instance.get_pdbstr('target')
    
    def save(self, out_file: Path):
        self.pymol_instance.save(out_file.as_posix(), "target")

@dataclass
class LoopModelBuilder:
    pdb_file: Path
    output_dir: Path = field(default=None)
    logger: logging.Logger = field(init=False)
    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
        self.output_dir = self.pdb_file.parent if self.output_dir is None else self.output_dir
        self.analyzer_instance = PDBAnalyzer(self.pdb_file)
        self.setup_logging()
        self.initialize_pymol()

    def setup_logging(self):
        log_file = self.output_dir / "loop_modeling.log"
        if log_file.exists():
            log_file.unlink()

        # 创建 logger 实例
        self.logger = logging.getLogger('LoopModelBuilder')
        self.logger.setLevel(logging.INFO)

        # 创建 file handler
        file_handler = logging.FileHandler(log_file, mode='a')
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

        # 创建 console handler
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

        # 将 handlers 添加到 logger
        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)


    def initialize_pymol(self):
        self.pymol_instance = pymol.cmd
        self.pymol_instance.reinitialize()

    def split_chains(self) -> dict:
        split_dict = {}
        for chain_id in self.analyzer_instance.chain_id_list:
            self.logger.info(f"Chain {chain_id} info(from file not fasta): \n{self.sequences[chain_id]}")
            chain_file = self.output_dir.joinpath(f"{self.analyzer_instance.pid}_{chain_id}.pdb")
            self.analyzer_instance.split_chain(chain_id).to_pdb(chain_file.as_posix())
            split_dict[chain_id] = chain_file.read_text()
        return split_dict

    def merge_and_save_pdb(self, pdb_strings: dict, pdb_id: str):
        merged_file = Path(f"{pdb_id}_merged.pdb")
        self.import_and_merge_pdb_strings(pdb_strings, "merged_object", merged_file)

    @property
    def sequences(self) -> Dict:
        return self.analyzer_instance.extract_sequences(missing_char='-')  # 或者 'X', 或者 ''

    @property
    def missing_info(self) -> Dict:
        return self.analyzer_instance.extract_sequences_info()

    def split_all_chains(self):
        self.logger.info(f'Residues info for {self.pdb_file}: \n',self.sequences)
        self.logger.info(f'Missing residues info for {self.pdb_file}:\n {self.missing_info}')
        split_dict = {} # split all chains
        for j in self.analyzer_instance.chain_id_list:
            self.logger.info(f"Chain {j} info(from file not fasta): \n{self.sequences[j]}")
            fn = self.output_dir.joinpath(f"{self.analyzer_instance.pid}_{j}.pdb")
            self.analyzer_instance.split_chain(j).to_pdb(fn.as_posix())
            split_dict[j]=fn.read_text()
        return split_dict

    def model_missing_loops(self, typestr:str = 'refine.very_fast', buildnumber: str = 1) -> dict:
        mc_dict = {}
        if not self.missing_info:
            self.logger.info(f'No missing residues found in {self.pdb_file}. Skipping model_missing_loops.')
            return mc_dict
        self.logger.info(f'Missing residues info for {self.pdb_file}:\n {self.missing_info}')
        # create workdir
        for mc in self.missing_info.keys():
            self.logger.info(f'Building model for chain {mc}')
            workdir = self.pdb_file.parent.joinpath(f'./{self.pdb_id}modellerfix_{mc}')
            workdir.mkdir(exist_ok=True, parents=True)
            mc_fasta = self.analyzer_instance.find_most_similar(self.sequences[mc]) # get misschain fasta file
            # write fasta file
            mc_fasta_record = SeqRecord(mc_fasta, id=f'{mc}', description=f'{self.pdb_id}_{mc}')
            out_fasta_file = workdir.joinpath(f'{self.analyzer_instance.pid}_{mc}.fasta')
            self.analyzer_instance.write_seq_to_fasta_single_line(mc_fasta_record, out_fasta_file)
            self.logger.info(f'>{self.pdb_id}|{mc}|missing site:{self.missing_info[mc]}|length:{len(mc_fasta)}')
            self.logger.info(mc_fasta)
            # build model
            modeller = PDBModeler(self.pdb_file, out_fasta_file, workdir, mc, buildnumber, typestr)
            try:
                modeller_results = modeller.make_model()
            except ModellerError as mod_err:
                self.logger.info(f'Failed to build model for chain {mc}')
                self.logger.info(f'No loops detected in {out_fasta_file.name}')
                self.logger.info(f'may pdb file sequence is not correct')
                self.logger.error(f'Modeller error for chain {mc}: {mod_err}')
                continue
            except Exception as e:
                self.logger.error(f'Unexpected error in model_missing_loops: {e}')
            self.logger.info(f'Model files: {[file.name for file in modeller_results]}')
            # change id to original
            for i in modeller_results:
                self.change_chain_identifier(i, 'A', mc, split=False)
            # 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
    def change_chain_identifier(pdb_file: Path, chain_id:str, new_chain_id:str, split:bool = True) -> Path:
        o = PDBAnalyzer(pdb_file)
        o.change_chain_identifier(chain_id, new_chain_id, split=split).to_pdb(pdb_file)
        return pdb_file

    def run(self, typestr:str = 'refine.very_fast', buildnumber: str = 1) -> Path:
        split_dict = self.split_all_chains()
        mc_dict = self.model_missing_loops(typestr=typestr, buildnumber=buildnumber)
        split_dict.update(mc_dict)  # 更新 split_dict
        out_file = self.output_dir.joinpath(f'{self.analyzer_instance.pid}.modellerfix.pdb')
        self.import_and_merge_pdb_strings(split_dict, "merged_object", out_file.as_posix())

    def import_and_merge_pdb_strings(self, pdb_strings, merged_object_name, output_file):
        # 使用 PyMOL 实例导入和合并 PDB
        for chain_id, pdb_str in pdb_strings.items():
            object_name = f"chain_{chain_id}"
            self.pymol_instance.read_pdbstr(pdb_str, object_name)

        object_names = [f"chain_{chain_id}" for chain_id in pdb_strings.keys()]
        self.pymol_instance.create(merged_object_name, ' or '.join(object_names))
        self.pymol_instance.save(output_file, merged_object_name)

def main():
    # 创建ArgumentParser对象
    parser = argparse.ArgumentParser(description="Loop Model Builder for PDB files")
    # 添加参数
    parser.add_argument("pdb_file", type=str, help="Path to the PDB file")
    parser.add_argument("--typestr", type=str, default="refine.very_fast", 
                        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()

    # 调用LoopModelBuilder
    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, 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}")

if __name__ == "__main__":
    main()