SBVS/process_data.py

import os, esm, torch, pickle, molvs, requests, argparse
import pandas as pd
import numpy as np
from tqdm import tqdm
from rdkit import Chem
from typing import List
from Bio.PDB import PDBParser
from Bio.SeqUtils import seq1
from sklearn.cluster import SpectralClustering
from sklearn.model_selection import train_test_split
from graphein.protein.graphs import construct_graph
from graphein.protein.config import ProteinGraphConfig
from graphein.ml import GraphFormatConvertor
from graphein.protein.utils import download_alphafold_structure
from graphein.protein.edges.distance import (add_peptide_bonds,
                                             add_hydrogen_bond_interactions,
                                             add_disulfide_interactions,
                                             add_ionic_interactions,
                                             add_aromatic_interactions,
                                             add_aromatic_sulphur_interactions,
                                             add_cation_pi_interactions
                                            )
from data_prep import split_data
from utils import set_seed, get_protein_sequence, check_molecule
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

def extract_sequence_from_pdb(pdb_file):
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure('protein', pdb_file)

    sequences = []
    for model in structure:
        for chain in model:
            sequence = ''
            for residue in chain:
                sequence += seq1(residue.get_resname())
            sequences.append(sequence)
    return sequences[0]

def generate_protein_graph(df):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
    model = model.to(device)
    batch_converter = alphabet.get_batch_converter()
    model.eval()
    # path = 'refined-set'
    # target = [i.split('.')[0] for i in df['Uniprot_id'].unique()]
    target = df['Uniprot_id'].unique()
    # exist_pdb_file = [i.split('.')[0] for i in os.listdir('data/PDB')]
    exist_pdb_file = os.listdir('data/PDB')
    print(exist_pdb_file)
    exist_file = os.listdir(args.prot_dir)
    error_list = []
    for pro in tqdm(target):
        # protein graph residue topology
        if f'{pro}.pkl' in exist_file:
            continue
        else:
            print(pro)
            sequence = df[df['Uniprot_id'] == pro]['Sequence'].values[0]
            new_edge_funcs = {"edge_construction_functions": [add_peptide_bonds,
                                                            # add_aromatic_interactions,
                                                            add_hydrogen_bond_interactions,
                                                            add_disulfide_interactions,
                                                            add_ionic_interactions,
                                                            add_aromatic_sulphur_interactions,
                                                            add_cation_pi_interactions]
                                }
            convertor = GraphFormatConvertor(src_format="nx", dst_format="pyg")
            if pro in exist_pdb_file:
                config = ProteinGraphConfig(**new_edge_funcs, verbose=0, 
                                        pdb_path=f'data/PDB/{pro}')
                g = convertor(construct_graph(config=config, path=f'data/PDB/{pro}',
                                verbose=False))
            else:
                try:
                    config = ProteinGraphConfig(**new_edge_funcs)
                    g = convertor(construct_graph(config=config, uniprot_id=pro, verbose=False))
                except:
                    error_list.append(pro)
                    pass
            # protein graph node feature
            try:
                prot_data = [(pro, sequence)]
                _, _, batch_tokens = batch_converter(prot_data)
                batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
                with torch.no_grad():
                    batch_tokens = batch_tokens.to(device)
                    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
                esm_emb = results["representations"][33][0, 1: batch_lens-1].cpu().numpy()

                with open(f'{args.prot_dir}/{pro}.pkl', 'wb') as f:
                    pickle.dump({pro: [sequence, esm_emb, g]}, f)
                    print(f'{pro} save to {args.prot_dir}')
            except:
                if pro not in error_list:
                    error_list.append(pro)
                pass
    return error_list

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--dataset', type=str,
                        help='Dataset file name')
    parser.add_argument('--task', type=str, default='CPI',
                        choices=['CPI', 'QSAR'],
                        help='Task type for your data (CPI or target-specific).'
                             'The data processing could be a little different')
    parser.add_argument('--split', type=str, default='random',
                        choices=['random', 'ac'],
                        help='Data splitting method')
    parser.add_argument('--train_ratio', type=float, default=0.8,
                        help='Train ratio to split data into train/test sets')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--prot_dir', type=str, 
                        default='data/Protein_pretrained_feat',
                        help='Directory to store protein graph data')
    args = parser.parse_args()
    args.dataset += '.csv'

    df = pd.read_csv(f'data/{args.dataset}')
    set_seed(args.seed)

    # if df['y'].dtype == int or len(df['y'].unique()) == 2:
    #     df['y'] = df['y'].astype(int)
    #     task_type = 'classification'
    # else:
    #     df['y'] = df['y'].astype(float)
    #     task_type = 'regression'
    df['cliff_mol'] = 0
    # print(f'Task type: {task_type}')

    if 'exp_mean [nM]' not in df.columns:
        df['exp_mean [nM]'] = 0

    # check ligand structure
    smi = df['smiles'].unique()
    smi_modify = [check_molecule(i) for i in tqdm(smi)]
    smi_dict = dict(zip(smi, smi_modify))
    df['smiles'] = df['smiles'].apply(lambda x: smi_dict[x])
    len_before = len(df)
    df = df.dropna(subset=['smiles'])
    print(f'{len_before - len(df)} invalid ligands are removed.')
    error_idx = set(range(len(df))) - set(df.dropna(subset=['smiles']).index)
    print(f'Error index: {error_idx}')

    # check protein structure
    assert 'Uniprot_id' in df.columns, \
        '"Uniprot_id" column not found in the dataset, please provide the Uniprot_id.\n' \
        'If you dont have UniProt IDs but have PDB files, ' \
        'you can fill the "Uniprot_id" column with PDB file name which should be stored in "data/PDB" folder.'
    
    # get protein sequence
    if 'Sequence' not in df.columns or df['Sequence'].isnull().sum() > 0:
        non_seq = df[df['Sequence'].isnull()]['Uniprot_id'].unique() if 'Sequence' in df.columns else df['Uniprot_id'].unique()
        uni_seq = [extract_sequence_from_pdb(f'data/PDB/{i}')  if '.pdb' in i \
                    else get_protein_sequence(i) if i in non_seq \
                    else df[df['Uniprot_id'] == i]['Sequence'].values[0]
                    for i in tqdm(df['Uniprot_id'].unique())]
        uni_seq = dict(zip(df['Uniprot_id'].unique(), uni_seq))
        df['Sequence'] = df['Uniprot_id'].map(uni_seq)
        df['Uniprot_id'] = df['Uniprot_id'].map(lambda x: x.split('.')[0])

    # get protein graph
    error_list = generate_protein_graph(df)
    print(f'Error list: {error_list}')
    len_before = len(df)
    df = df[~df['Uniprot_id'].isin(error_list)]
    len_after = len(df)
    print(f'{len_before - len_after} data are removed due to '\
            'unable to find AF2-predicted structures.')

    # data splitting
    df_all = []
    if 'split' not in df.columns or (args.task == 'QSAR' and not os.path.exist(f'data/{args.dataset}')):
        for target in df['Uniprot_id'].unique():
            subset = df[df['Uniprot_id'] == target]
            subset = subset.reset_index(drop=True)
            if args.split == 'random':
                subset['split'] = np.random.choice(['train', 'test'], len(subset), p=[args.train_ratio, 1-args.train_ratio])
            elif args.split == 'ac':
                subset = split_data(subset['smiles'].values.tolist(),
                                    bioactivity=subset['y'].values.tolist(),
                                    in_log10=True, similarity=0.9, test_size=1-args.train_ratio, random_state=args.seed)
            else:
                raise ValueError('Cannot use activity cliff-based splitting for classification tasks.')

            if args.task == 'QSAR':
                if not os.path.exist(f'data/{args.dataset}'):
                    subset.to_csv(f'data/{args.dataset}/{target}', index=False)
                subset.to_csv(f'data/{target}', index=False)
            elif args.task == 'CPI':
                df_all.append(subset)
        if args.task == 'CPI':
            df_all = pd.concat(df_all)
            df_all.to_csv(f'data/{args.dataset}', index=False)
    else:
        if args.task == 'QSAR':
            for target in df['Uniprot_id'].unique():
                subset = df[df['Uniprot_id'] == target]
                subset = subset.reset_index(drop=True)
                if not os.path.exist(f'data/{args.dataset}'):
                    subset.to_csv(f'data/{args.dataset}/{target}', index=False)
                subset.to_csv(f'data/{target}.csv', index=False)
        elif args.task == 'CPI':
            df.to_csv(f'data/{args.dataset}', index=False)
    print('Data processing finished. You can run model training/testing now.')