DMcloudComplex.py

import os
import sys

from ops.argparser import argparser
from ops.os_operation import mkdir
import time
import shutil
import pickle


# Merge CIF files
# -*- coding: utf-8 -*-
from Bio.PDB import PDBParser, MMCIFParser, PDBIO, MMCIFIO
from Bio.PDB.Model import Model
from Bio.PDB.Structure import Structure
import os
import string

# ---- chain ID generator ----

_PDB_CHAIN_POOL = list(string.ascii_uppercase + string.digits + string.ascii_lowercase)

def next_chain_id(used_ids):
    for cid in _PDB_CHAIN_POOL:
        if cid not in used_ids:
            return cid
    return None  # 62 chains?

def _parse_structure(path):
    ext = os.path.splitext(path)[1].lower()
    sid = os.path.basename(path)
    if ext in [".cif", ".mmcif"]:
        parser = MMCIFParser(QUIET=True)
    elif ext in [".pdb", ".ent"]:
        parser = PDBParser(QUIET=True, PERMISSIVE=True)
    else:
        raise ValueError(f"Unsupported file type: {path}")
    return parser.get_structure(sid, path)

def merge_structures_to_cif_and_pdb(input_files, out_cif, out_pdb):
    """
    """
    if not input_files:
        raise ValueError("input_files is empty.")

    # 単一ファイルなら読み込んでそのまま両形式に保存（可能なら）
    if len(input_files) == 1:
        s = _parse_structure(input_files[0])
        # mmCIF
        mmio = MMCIFIO(); mmio.set_structure(s); mmio.save(out_cif)
        # PDB
        n_chains = sum(1 for _ in s.get_chains())
        wrote_pdb = False
        if n_chains <= len(_PDB_CHAIN_POOL):
            pio = PDBIO(); pio.set_structure(s); pio.save(out_pdb)
            wrote_pdb = True
        return dict(status="ok(single)", n_chains=n_chains, wrote_pdb=wrote_pdb, wrote_cif=True)

    # Merge
    merged = Structure("merged")
    model0 = Model(0)
    merged.add(model0)

    used_chain_ids = set()
    total_chains = 0

    for f in input_files:
        s = _parse_structure(f)
        for m in s:
            for ch in m:
                total_chains += 1
                cid = next_chain_id(used_chain_ids)
                if cid is None:

                    cid = f"chain_{total_chains}"
                ch.id = cid
                used_chain_ids.add(cid)
                model0.add(ch)

    mmio = MMCIFIO(); mmio.set_structure(merged); mmio.save(out_cif)

    wrote_pdb = False
    if all((len(cid) == 1 and cid in _PDB_CHAIN_POOL) for cid in used_chain_ids):
        pio = PDBIO(); pio.set_structure(merged); pio.save(out_pdb)
        wrote_pdb = True
    else:
        pass

    return dict(status="ok(merged)", n_chains=total_chains, wrote_pdb=wrote_pdb, wrote_cif=True)

##=============================Main function=========================================

def init_save_path(origin_map_path):
    save_path = os.path.join(os.getcwd(), 'Predict_Result')
    mkdir(save_path)
    map_name = os.path.split(origin_map_path)[1].replace(".mrc", "")
    map_name = map_name.replace(".map", "")
    map_name = map_name.replace("(","").replace(")","")
    save_path = os.path.join(save_path, map_name)
    mkdir(save_path)
    return save_path,map_name
def set_up_envrionment(params):
    if params['resolution']>20:
        print("maps with %.2f resolution is not supported! We only support maps with resolution 0-20A!"%params['resolution'])
        exit()
    gpu_id = params["gpu"]
    if gpu_id is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
    cur_map_path = os.path.abspath(params['map'])
    if cur_map_path.endswith(".gz"):
        from ops.os_operation import unzip_gz
        cur_map_path = unzip_gz(cur_map_path)

    if params['output'] is None:
        save_path,map_name = init_save_path(cur_map_path)
    else:
        save_path=params['output']
        map_name="input_complex"
        mkdir(save_path)
    save_path = os.path.abspath(save_path)
    from CryoREAD.data_processing.Unify_Map import Unify_Map
    cur_map_path = Unify_Map(cur_map_path,os.path.join(save_path,map_name+"_unified.mrc"))
    from CryoREAD.data_processing.Resize_Map import Resize_Map
    cur_map_path = Resize_Map(cur_map_path,os.path.join(save_path,map_name+".mrc"))
    if params['contour_nuc']<0:
        #change contour level to 0 and increase all the density
        from CryoREAD.data_processing.map_utils import increase_map_density
        cur_map_path = increase_map_density(cur_map_path,os.path.join(save_path,map_name+"_increase.mrc"),params['contour_nuc'])
        params['contour_nuc']=0
    from CryoREAD.data_processing.map_utils import segment_map
    cur_new_map_path = os.path.join(save_path,map_name+"_segment.mrc")
    cur_map_path = segment_map(cur_map_path,cur_new_map_path,contour=0) #save the final prediction prob array space
    return save_path,cur_map_path

if __name__ == "__main__":
    params = argparser()
    save_path,cur_map_path = set_up_envrionment(params)
    #parse fasta files to determine if run fasta CryoREAD version or not
    from ops.fasta_utils import read_fasta,write_drna_fasta,write_fasta
    fasta_path = os.path.abspath(params['map'])
    gpu_id = params["gpu"]

    #first run DNA/RNA predictions
    running_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(running_dir)
    cryoread_script = os.path.join(os.getcwd(),"CryoREAD")
    best_model_path = os.path.join(cryoread_script,"best_model")
    cryoread_script = os.path.join(cryoread_script,"main.py")
    contour_nuc = params['contour_nuc']
    cryoread_out_path = os.path.join(save_path,"CryoREAD_run")
    command_line=f"python3 %s --mode=0 -F={cur_map_path} -M {best_model_path} " \
                 f"--contour={contour_nuc} --gpu {gpu_id} " \
                 f"--batch_size=4 --no_seqinfo --output {cryoread_out_path} "%(cryoread_script)
    print("Run command: %s"%command_line)
    protein_map_path = os.path.join(cryoread_out_path,"mask_protein.mrc")
    pho_ldp_path = os.path.join(cryoread_out_path,"graph_atomic_modeling/pho_LDP/pho_LDPdens.cif")
    sugar_ldp_path = os.path.join(cryoread_out_path,"graph_atomic_modeling/sugar_LDP/sugar_LDPdens.cif")
    if not os.path.exists(protein_map_path):
        os.system(command_line)
    
    if not os.path.exists(protein_map_path):
        print("protein detection failed")
        sys.exit(1)

    if not os.path.exists(pho_ldp_path) or not os.path.exists(sugar_ldp_path):
        print("RNA/DNA detection failed")
        sys.exit(1)
    
    print("#RNA/DNA detection finished!")
    print("#Start DMcloudRNA")
    protein_models = params["protein"]
    nucleic_models = params["nucleic"]

    dmcloud_script = os.path.join(os.getcwd(),"DMcloudRNA/DMcloudRNA.py")

    #Three parameter combinations
    rna_params = [
        {"LocalRadius":8,"AveLDDT_cut":0.4,"icp_dcut":4.0,"n_ransac_iter":2000000,"MinCkstSize":50},
        {"LocalRadius":14,"AveLDDT_cut":0.4,"icp_dcut":4.0,"n_ransac_iter":2000000,"MinCkstSize":50},
        {"LocalRadius":20,"AveLDDT_cut":0.4,"icp_dcut":4.0,"n_ransac_iter":2000000,"MinCkstSize":50}
    ]
    for i,param in enumerate(rna_params):
        rna_out_path = os.path.join(save_path,"DMcloudRNA_run%d"%i)
        command_line=f"python3 %s --OutPath={rna_out_path} " \
        f"--Source {nucleic_models} --SugLDPs {sugar_ldp_path} --PhoLDPs {pho_ldp_path} " \
        f"--MinClstSize {param['MinCkstSize']} --n_ransac_iter {param['n_ransac_iter']} --icp_dcut {param['icp_dcut']} " \
        f"--LocalRadius {param['LocalRadius']} --AveLDDT_cut {param['AveLDDT_cut']} "%(dmcloud_script)
        print("Run command: %s"%command_line)
        if not os.path.exists(rna_out_path):
            os.system(command_line)

    print("#RNA/DNA structure modeling finished!")
    #Detect best RNA Model
    RNA_SCOREs = []
    for i,param in enumerate(rna_params):
        rna_out_path = os.path.join(save_path,"DMcloudRNA_run%d"%i)
        pkl_file=f"{rna_out_path}/selected_alignment.pkl"
        try:
            with open(pkl_file,'rb') as f:
                TMP_DATA = pickle.load(f)
                total_score = sum(d['score'] for d in TMP_DATA)
                print(f"DMcloudRNA_run{i} {total_score}")
                RNA_SCOREs.append({"total_score": total_score, "rna_out_path": rna_out_path})
        except:
            print(f"No Hits: Cannot find {pkl_file}")
    if len(RNA_SCOREs)==0:
        print("No RNA/DNA models are built!")
        sys.exit(1)
    RNA_SCOREs = sorted(RNA_SCOREs, key=lambda x: x['total_score'], reverse=True)
    best_rna_out_path = RNA_SCOREs[0]['rna_out_path']
    best_merged_cif = os.path.join(save_path,"best_rna_models.cif")
    best_merged_pdb = os.path.join(save_path,"best_rna_models.pdb")
    cif_files = []
    for file in os.listdir(f"{best_rna_out_path}/Selected/"):
        if file.endswith(".cif") and file.startswith("Filled"):
            cif_files.append(os.path.join(best_rna_out_path,f"Selected/{file}"))

    info = merge_structures_to_cif_and_pdb(cif_files,best_merged_cif,best_merged_pdb)
    print(info)
    merge_structures_to_cif_and_pdb(cif_files,best_merged_cif,best_merged_pdb)
    print(f"Best RNA/DNA model from {best_rna_out_path} with score {RNA_SCOREs[0]['total_score']}")
    print(f"Merged CIF file saved to {best_merged_cif}")

    print("#Start DMcloudProtein")

    dmcloud_script = os.path.join(os.getcwd(),"DMCloud","DMcloud.py")

    #Three parameter combinations
    protein_params = [
        {"LocalRadius":8.0,"AveLDDT_cut":0.4,"icp_dcut":3.0,"n_ransac_iter":2000000,"MinCkstSize":100},
    ]
    if len(protein_params)==0:
        print("No protein parameters are set!")
        sys.exit(1)
    if protein_models is None:
        print("No protein models are provided!")
        sys.exit(1)

    diffusion_config = os.path.join(os.getcwd(),"DMCloud","config","diffusion.json")
    dmcloud_dir = os.path.join(os.getcwd(),"DMCloud")
    contour_prot = params['contour_prot']
    for i,param in enumerate(protein_params):
        protein_out_path = os.path.join(save_path,"DMcloudProtein_run%d"%i)
        command_line=f"python3 %s --OutPath={protein_out_path} " \
        f"--Source {protein_models} --Target {protein_map_path} --contour {contour_prot} " \
        f"--MinClstSize {param['MinCkstSize']} --n_ransac_iter {param['n_ransac_iter']} --icp_dcut {param['icp_dcut']} " \
        f"--diffusion_gpu {gpu_id} --diffusion_config {diffusion_config} " \
        f"--LocalRadius {param['LocalRadius']} --AveLDDT_cut {param['AveLDDT_cut']} "%(dmcloud_script)
        #print("Run command: %s"%command_line)
        if not os.path.exists(protein_out_path):
            print("Run command: %s"%command_line)
            os.system(command_line)

    Protein_SCOREs = []
    for i,param in enumerate(protein_params):
        protein_out_path = os.path.join(save_path,"DMcloudProtein_run%d"%i)
        pkl_file=f"{protein_out_path}/selected_alignment.pkl"
        try:
            with open(pkl_file,'rb') as f:
                TMP_DATA = pickle.load(f)
                total_score = sum(d['score'] for d in TMP_DATA)
                print(f"DMcloudProtein_run{i} {total_score}")
                Protein_SCOREs.append({"total_score": total_score, "protein_out_path": protein_out_path})
        except:
            print(f"No Hits: Cannot find {pkl_file}")
    if len(Protein_SCOREs)==0:
        print("No Protein models are built!")
        sys.exit(1)
    
    Protein_SCOREs = sorted(Protein_SCOREs, key=lambda x: x['total_score'], reverse=True)
    best_protein_out_path = Protein_SCOREs[0]['protein_out_path']

    cif_files = []
    for file in os.listdir(f"{best_protein_out_path}/Selected/"):
        if file.startswith("Filled"):
            cif_files.append(os.path.join(best_protein_out_path,f"Selected/{file}"))

    best_merged_cif = os.path.join(save_path,"best_protein_models.cif")
    best_merged_pdb = os.path.join(save_path,"best_protein_models.pdb")
    print(f"Best Protein model from {best_protein_out_path} with score {Protein_SCOREs[0]['total_score']}")
    info = merge_structures_to_cif_and_pdb(cif_files,best_merged_cif,best_merged_pdb)
    print(info)
    print(f"Merged CIF file saved to {best_merged_cif}")