model/comp_surface/prepare_target/computeTargetMesh_test_samples.py

import os
import sys
import numpy as np
import shutil
import glob
import pymesh
import Bio.PDB
from Bio.PDB import * 
from rdkit import Chem
import warnings
warnings.filterwarnings("ignore")
from IPython.utils import io
from sklearn.neighbors import KDTree
from scipy.spatial import distance

from default_config.masif_opts import masif_opts
from compute_normal import compute_normal
from computeAPBS import computeAPBS
from computeCharges import computeCharges, assignChargesToNewMesh
from computeHydrophobicity import computeHydrophobicity
from computeMSMS import computeMSMS
from fixmesh import fix_mesh
from save_ply import save_ply
from mol2graph import *


def compute_inp_surface(target_filename, ligand_filename,out_dir = None, dist_threshold=10):
    try:

        sufix = '_'+str(dist_threshold)+'A.pdb'
        # out_filename = os.path.splitext(target_filename)[0]
        if out_dir is not None: 
            out_filename = os.path.join(out_dir,ligand_filename.split('/')[-2]) 
            os.makedirs(out_filename,exist_ok=True)
            sufix = '/' + os.path.splitext(target_filename)[0].split('/')[-1] + '_'+str(dist_threshold)+'A.pdb'
        else:
            out_filename = os.path.splitext(ligand_filename)[0]
        if os.path.exists(out_filename+f"/{sufix.split('.pdb')[0]}.ply"):
            print('have done skip!')
            return 0
        input_filename = os.path.splitext(target_filename)[0]
        # Get atom coordinates
        # try:
        #     mol = Chem.MolFromMol2File(ligand_filename, sanitize=False, cleanupSubstructures=False)
        # except:
        #     print('mol2 faild try sdf')
        if ligand_filename.endswith('.sdf'):
            mol = Chem.SDMolSupplier(ligand_filename, sanitize=False)[0]
        elif ligand_filename.endswith('.pdb'):
            mol = Chem.MolFromPDBFile(ligand_filename, sanitize=False)
        g = mol_to_nx(mol)
        atomCoords = np.array([g.nodes[i]['pos'].tolist() for i in g.nodes])
        
        # Read protein and select aminino acids in the binding pocket
        parser = Bio.PDB.PDBParser(QUIET=True) # QUIET=True avoids comments on errors in the pdb.

        structures = parser.get_structure('target', input_filename+'.pdb')
        structure = structures[0] # 'structures' may contain several proteins in this case only one.

        atoms  = Bio.PDB.Selection.unfold_entities(structure, 'A')
        ns = Bio.PDB.NeighborSearch(atoms)
        
        close_residues= []
        for a in atomCoords:  
            close_residues.extend(ns.search(a, dist_threshold, level='R'))
        close_residues = Bio.PDB.Selection.uniqueify(close_residues)

        class SelectNeighbors(Select):
            def accept_residue(self, residue):
                if residue in close_residues:
                    if all(a in [i.get_name() for i in residue.get_unpacked_list()] for a in ['N', 'CA', 'C', 'O']) or residue.resname=='HOH':
                        return True
                    else:
                        return False
                else:
                    return False

        pdbio = PDBIO()
        pdbio.set_structure(structure)
        pdbio.save(out_filename+sufix, SelectNeighbors())
        
        # Identify closes atom to the ligand
        structures = parser.get_structure('target', out_filename+sufix)
        structure = structures[0] # 'structures' may contain several proteins in this case only one.
        atoms  = Bio.PDB.Selection.unfold_entities(structure, 'A')

        #dist = [distance.euclidean(atomCoords.mean(axis=0), a.get_coord()) for a in atoms]
        #atom_idx = np.argmin(dist)
        #dist = [[distance.euclidean(ac, a.get_coord()) for ac in atomCoords] for a in atoms]
        #atom_idx = np.argsort(np.min(dist, axis=1))[0]
        
        # Compute MSMS of surface w/hydrogens, 
        try:
            dist = [distance.euclidean(atomCoords.mean(axis=0), a.get_coord()) for a in atoms]
            atom_idx = np.argmin(dist)
            vertices1, faces1, normals1, names1, areas1 = computeMSMS(out_filename+sufix,\
                                                                    protonate=True, one_cavity=atom_idx)
            
            # Find the distance between every vertex in binding site surface and each atom in the ligand.
            kdt = KDTree(atomCoords)
            d, r = kdt.query(vertices1)
            assert(len(d) == len(vertices1))
            iface_v = np.where(d <= dist_threshold-5)[0]
            faces_to_keep = [idx for idx, face in enumerate(faces1) if all(v in iface_v  for v in face)] 
        
            # Compute "charged" vertices
            if masif_opts['use_hbond']:
                vertex_hbond = computeCharges(input_filename, vertices1, names1)    
        
            # For each surface residue, assign the hydrophobicity of its amino acid. 
            if masif_opts['use_hphob']:
                vertex_hphobicity = computeHydrophobicity(names1)    
            
            # If protonate = false, recompute MSMS of surface, but without hydrogens (set radius of hydrogens to 0).
            vertices2 = vertices1
            faces2 = faces1
        
            # Fix the mesh.
            mesh = pymesh.form_mesh(vertices2, faces2)
            mesh = pymesh.submesh(mesh, faces_to_keep, 0)
            with io.capture_output() as captured:
                regular_mesh = fix_mesh(mesh, masif_opts['mesh_res'])
            
        except:
            try:
                dist = [[distance.euclidean(ac, a.get_coord()) for ac in atomCoords] for a in atoms]
                atom_idx = np.argsort(np.min(dist, axis=1))[0]
                vertices1, faces1, normals1, names1, areas1 = computeMSMS(out_filename+sufix,\
                                                                        protonate=True, one_cavity=atom_idx)

                # Find the distance between every vertex in binding site surface and each atom in the ligand.
                kdt = KDTree(atomCoords)
                d, r = kdt.query(vertices1)
                assert(len(d) == len(vertices1))
                iface_v = np.where(d <= dist_threshold-5)[0]
                faces_to_keep = [idx for idx, face in enumerate(faces1) if all(v in iface_v  for v in face)] 
        
                # Compute "charged" vertices
                if masif_opts['use_hbond']:
                    vertex_hbond = computeCharges(input_filename, vertices1, names1)    
        
                # For each surface residue, assign the hydrophobicity of its amino acid. 
                if masif_opts['use_hphob']:
                    vertex_hphobicity = computeHydrophobicity(names1)    
            
                # If protonate = false, recompute MSMS of surface, but without hydrogens (set radius of hydrogens to 0).
                vertices2 = vertices1
                faces2 = faces1
        
                # Fix the mesh.
                mesh = pymesh.form_mesh(vertices2, faces2)
                mesh = pymesh.submesh(mesh, faces_to_keep, 0)
                with io.capture_output() as captured:
                    regular_mesh = fix_mesh(mesh, masif_opts['mesh_res'])
                    
            except:
                vertices1, faces1, normals1, names1, areas1 = computeMSMS(out_filename+sufix,\
                                                                        protonate=True, one_cavity=None)

                # Find the distance between every vertex in binding site surface and each atom in the ligand.
                kdt = KDTree(atomCoords)
                d, r = kdt.query(vertices1)
                assert(len(d) == len(vertices1))
                iface_v = np.where(d <= dist_threshold-5)[0]
                faces_to_keep = [idx for idx, face in enumerate(faces1) if all(v in iface_v  for v in face)] 
        
                # Compute "charged" vertices
                if masif_opts['use_hbond']:
                    vertex_hbond = computeCharges(input_filename, vertices1, names1)    
        
                # For each surface residue, assign the hydrophobicity of its amino acid. 
                if masif_opts['use_hphob']:
                    vertex_hphobicity = computeHydrophobicity(names1)    
            
                # If protonate = false, recompute MSMS of surface, but without hydrogens (set radius of hydrogens to 0).
                vertices2 = vertices1
                faces2 = faces1
        
                # Fix the mesh.
                mesh = pymesh.form_mesh(vertices2, faces2)
                mesh = pymesh.submesh(mesh, faces_to_keep, 0)
                with io.capture_output() as captured:
                    regular_mesh = fix_mesh(mesh, masif_opts['mesh_res'])
            
        # Compute the normals
        vertex_normal = compute_normal(regular_mesh.vertices, regular_mesh.faces)
        # Assign charges on new vertices based on charges of old vertices (nearest
        # neighbor)
        
        if masif_opts['use_hbond']:
            vertex_hbond = assignChargesToNewMesh(regular_mesh.vertices, vertices1,\
                vertex_hbond, masif_opts)
        
        if masif_opts['use_hphob']:
            vertex_hphobicity = assignChargesToNewMesh(regular_mesh.vertices, vertices1,\
                vertex_hphobicity, masif_opts)
        
        if masif_opts['use_apbs']:
            vertex_charges = computeAPBS(regular_mesh.vertices, out_filename+sufix, out_filename+"_temp")
            
        # Compute the principal curvature components for the shape index. 
        regular_mesh.add_attribute("vertex_mean_curvature")
        H = regular_mesh.get_attribute("vertex_mean_curvature")
        regular_mesh.add_attribute("vertex_gaussian_curvature")
        K = regular_mesh.get_attribute("vertex_gaussian_curvature")
        elem = np.square(H) - K
        # In some cases this equation is less than zero, likely due to the method that computes the mean and gaussian curvature.
        # set to an epsilon.
        elem[elem<0] = 1e-8
        k1 = H + np.sqrt(elem)
        k2 = H - np.sqrt(elem)
        # Compute the shape index 
        si = (k1+k2)/(k1-k2)
        si = np.arctan(si)*(2/np.pi)
        
        # Convert to ply and save.
        save_ply(out_filename+f"/{sufix.split('.pdb')[0]}.ply", regular_mesh.vertices,\
                regular_mesh.faces, normals=vertex_normal, charges=vertex_charges,\
                normalize_charges=True, hbond=vertex_hbond, hphob=vertex_hphobicity,\
                si=si)
        
        return 0
    except:
        return target_filename

    
if __name__ == "__main__":
    from joblib import delayed,Parallel
    # arguments
    from argparse import ArgumentParser, Namespace, FileType
    parser = ArgumentParser()
    parser.add_argument('--data_dir', type=str, default='~/SurfDock/model/data/test_samples', help='')
    parser.add_argument('--out_dir', type=str, default='~/SurfDock/model/data/test_samples_8A_surface', help='')
    parser.add_argument('--n_jobs',type=int, default=1, help='Number of parallel jobs (-1 for all CPUs)')
    args = parser.parse_args()
    os.makedirs(args.out_dir,exist_ok=True)

    sys.path.append(args.out_dir)
    from tqdm import tqdm
    args_list = []
    for protein in tqdm(os.listdir(args.data_dir)):
            if os.path.exists(os.path.join(args.out_dir,protein,f'{protein}_protein_processed_obabel_reduce_obabel.pdb')):
                target_filename = os.path.join(args.out_dir,protein,f'{protein}_protein_processed_obabel_reduce_obabel.pdb')
            elif os.path.exists(os.path.join(args.data_dir,protein,f'{protein}_protein_processed.pdb')):
                target_filename = os.path.join(args.data_dir,protein,f'{protein}_protein_processed.pdb')
                print(f'{protein} use {target_filename}; Please check this protein file was processed by openbabel reduce! in protein_process')
            else:
                print(f'{protein} not exists , Please check file name or path')
                continue
            ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.sdf')
            if not os.path.exists(ligand_filename):
                ligand_filename = os.path.join(args.data_dir,protein,f'{protein}_ligand.mol2')
            args_list.append((target_filename,ligand_filename))
    print(f'number {len(args_list)} need to processed.....')
    results = Parallel(n_jobs = args.n_jobs,backend = 'multiprocessing')(delayed(compute_inp_surface)(target_filename, ligand_filename,args.out_dir, dist_threshold=8) for (target_filename, ligand_filename) in tqdm(args_list))
    # print(results)
    # Find all files in args.out_dir that end with _temp
    
    files = glob.glob(os.path.join(args.out_dir, '*_temp*')) + glob.glob(os.path.join(args.out_dir, '*msms*'))
    # Delete all found files
    for f in files:
        os.remove(f)