model/datasets/get_pocket_embedding.py

# by caoduanhua email : caodh@zju.edu.cn
# cycle chain and residues
import pickle
import os
from Bio.PDB import *
import warnings
warnings.filterwarnings('ignore')
import tqdm
from Bio.PDB import PDBParser
biopython_parser = PDBParser()
import torch
import pandas as pd
from argparse import ArgumentParser
parser = ArgumentParser()
# parser.add_argument('--pocket_dir', type=str, default='~/PDBBind_pocket_8A', help='pocket dir locations')
# parser.add_argument('--full_protein_dir', type=str, default='~/PDBBind_processed', help='full protein dir locations')
parser.add_argument('--protein_pocket_csv', type=str, default='~/processsed/PDBBIND.csv', help='save pocket and full protein csv locations')
parser.add_argument('--embeddings_dir', type=str, default='~/esm_embedding/esm_embedding_output', help='full protein embedding dir locations')
parser.add_argument('--pocket_emb_save_dir', type=str, default='~/esm_embedding/esm_embedding_output_pocket_new', help='')
args = parser.parse_args()

df = pd.read_csv(args.protein_pocket_csv)
full_protein_paths = list(df['protein_path'].tolist())
pocket_paths = list(df['pocket_path'].tolist())
# pocket_dir = args.pocket_dir
# full_protein_dir = args.full_protein_dir
protein_pocket_csv = args.protein_pocket_csv
embeddings_dir = args.embeddings_dir
pocket_emb_save_dir = args.pocket_emb_save_dir

three_to_one = {'ALA':	'A',
'ARG':	'R',
'ASN':	'N',
'ASP':	'D',
'CYS':	'C',
'GLN':	'Q',
'GLU':	'E',
'GLY':	'G',
'HIS':	'H',
'ILE':	'I',
'LEU':	'L',
'LYS':	'K',
'MET':	'M',
'MSE':  'M', # this is almost the same AA as MET. The sulfur is just replaced by Selen
'PHE':	'F',
'PRO':	'P',
'PYL':	'O',
'SER':	'S',
'SEC':	'U',
'THR':	'T',
'TRP':	'W',
'TYR':	'Y',
'VAL':	'V',
'ASX':	'B',
'GLX':	'Z',
'XAA':	'X',
'XLE':	'J'}
# if os.path.exists(
Assertion_list = []
os.makedirs(pocket_emb_save_dir,exist_ok = True)
# pbar = tqdm.tqdm(os.listdir(pocket_dir),total=len(os.listdir(pocket_dir)))
pbar = tqdm.tqdm(zip(full_protein_paths,pocket_paths),total=len(full_protein_paths))

for pbar_idx,(full_protein_path,pocket_path) in enumerate(pbar):
    protein_name = os.path.splitext(os.path.basename(pocket_path))[0]
    # raise  AssertionError(protein_name)# if False else None
    if os.path.exists(os.path.join(pocket_emb_save_dir,f'{protein_name}.pt')):
        pbar.set_description(f'have done ,just skip!')
        continue
    try:
        pocket = pocket_path
        full_protein = full_protein_path
        # full_protein = f'{full_protein_dir}/{pdb_id}.pdb'
        pocket_structure = biopython_parser.get_structure(f"{protein_name}", pocket)[0]
        full_structure = biopython_parser.get_structure(f"{protein_name}", full_protein)[0]
        pocket_embeddings =[]
        pocket_infos_all = []
        for i,chain in enumerate(full_structure.get_chains()):
            chain_id = chain.get_id()
            try:
                pocket_chain = pocket_structure[chain_id]
            except KeyError:
                pbar.set_description(f'{chain_id} not in {protein_name} pocket skip this chain')
                continue
            try:

                embeddings_path_chain = os.path.join(embeddings_dir,f'{os.path.basename(full_protein_path)}_chain_{i}.pt')
                # embeddings_path_chain = os.path.join(embeddings_dir,f'{pdb_id}.pdb_chain_{i}.pt')
                embeddings = torch.load(embeddings_path_chain)['representations'][33]
                assert len(list(chain.get_residues())) == len(embeddings),'embedding must equal to res nums!'
            except AssertionError:
                # pbar.set_description(f'{pdb_id} has error!,{len(list(chain.get_residues()))},{len(embeddings)}')
                # Assertion_list.append(pdb_id)
                residue_list = list(chain.get_residues())
                for res_idx, residue in enumerate(residue_list):
                # for res_idx, residue in enumerate(chain):
                    if residue.get_resname() == 'HOH':
                        chain.detach_child(residue.get_id())
                        continue
                    c_alpha, n, c = None, None, None
                    for atom in residue:
                        if atom.name == 'CA':
                            c_alpha = list(atom.get_vector())
                        if atom.name == 'N':
                            n = list(atom.get_vector())
                        if atom.name == 'C':
                            c = list(atom.get_vector())
                    if c_alpha != None and n != None and c != None:  
                        continue
                    else:
                        chain.detach_child(residue.get_id())
                        continue

                assert len(list(chain.get_residues())) == len(embeddings),f'embedding must equal to res nums! {len(list(chain.get_residues()))},{len(embeddings)}'
            pocket_infos = []
            pocket_residue_list = list(pocket_chain.get_residues())
            for res_idx, residue in enumerate(pocket_residue_list):

                if residue.get_resname() == 'HOH':

                    continue
                c_alpha, n, c = None, None, None
                for atom in residue:
                    if atom.name == 'CA':
                        c_alpha = list(atom.get_vector())
                    if atom.name == 'N':
                        n = list(atom.get_vector())
                    if atom.name == 'C':
                        c = list(atom.get_vector())
                if c_alpha != None and n != None and c != None:  
                    pocket_infos += [residue.get_id()]
            #         continue
                else:

                    print(residue.get_resname())
                    continue

            pocket_infos_all += pocket_infos
            # check the res in pocket 
            pocket_idx_list = []
            for res_idx,res in enumerate(chain.get_residues()):
                if res.get_id() in pocket_infos:
                    pocket_idx_list.append(res_idx)
                # else:

            pocket_embeddings.append(embeddings[pocket_idx_list])
        pocket_embeddings = torch.cat(pocket_embeddings,dim = 0)
        assert len(pocket_embeddings) == len(pocket_infos_all),f'pocket embedding must equal to res nums! {len(pocket_embeddings)},{len(pocket_infos_all)}'
        torch.save(pocket_embeddings,os.path.join(pocket_emb_save_dir,f'{protein_name}.pt'))
    except AssertionError as e:
        print(e,protein_name)

        Assertion_list.append(protein_name)

        continue
    except FileNotFoundError as e:
        Assertion_list.append(protein_name)

        continue
    except Exception as e:
        Assertion_list.append(protein_name)

        continue
    pbar.set_description(f'{pbar_idx}/{len(full_protein_paths)} done!')
print('Assertion_list:',Assertion_list)