model/datasets/pdbbind.py

import binascii
import glob
import hashlib
import os
# import pickle
import _pickle as pickle # use cPickle to speed up
import MDAnalysis as mda
from plyfile import PlyData
from torch_geometric.data import Data
from torch_geometric.transforms import FaceToEdge, Cartesian
from collections import defaultdict
from multiprocessing import Pool
import random
import copy
from joblib import Parallel, delayed
import numpy as np
import torch
from rdkit.Chem import MolToSmiles, MolFromSmiles, AddHs
from torch_geometric.data import Dataset, HeteroData
from torch_geometric.loader import DataLoader, DataListLoader
from torch_geometric.transforms import BaseTransform
from tqdm import tqdm
from loguru import logger
from datasets.process_mols import read_molecule, get_rec_graph, generate_conformer, \
    get_lig_graph_with_matching, extract_receptor_structure, parse_receptor, parse_pdb_from_path
from utils.diffusion_utils import modify_conformer, set_time
from utils.utils import read_strings_from_txt
from utils import so3, torus
import MDAnalysis as mda
from prefetch_generator import BackgroundGenerator
class DataLoaderX(DataLoader):
    def __iter__(self):
        return BackgroundGenerator(super().__iter__())   
class NoiseTransformBERT(BaseTransform):
    def __init__(self, t_to_sigma, no_torsion, all_atom):
        self.t_to_sigma = t_to_sigma
        self.no_torsion = no_torsion
        self.all_atom = all_atom
    def __call__(self, data):
        t = np.random.uniform()
        # t_rot = np.random.uniform()
        # t_tor = np.random.uniform()
        t_tr, t_rot, t_tor = t, t, t
        return self.apply_noise(data, t_tr, t_rot, t_tor)
    def apply_noise(self, data, t_tr, t_rot, t_tor, tr_update = None, rot_update=None, torsion_updates=None):
        # mdn mode make no update
        if not torch.is_tensor(data['ligand'].pos):
            data['ligand'].pos = random.choice(data['ligand'].pos)
        set_time(data, t_tr, t_rot, t_tor, 1, self.all_atom, device=None)
        # set noise scale and time
        # in the first steps ,tor and rot set to zero ,because the ligand is not in the binding pocket,\
        # modify the ligand in those freendom degree is useless
        eps_tor_sigma = 0.0314
        eps_tr_sigma = 0.1
        eps_rot_sigma = 0.1
        tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma(t_tr, t_rot, t_tor)
        # random like BERT style but probility is 0.8 and 0.2
        prob = random.random()
        # eps_sigma = 1e-10
        # 15% randomly change a freedom degree to noise and not noise any freedom degree
        if prob < 0.05:
            prob /= 0.05
            # 85% randomly change a freedom degree to noise
            if prob < 0.95:
                """
                0:tr
                1:rot
                2:tor
                """
                freedom_to_noise = random.choice([0,1,2])
                if freedom_to_noise == 0:
                    tr_sigma, rot_sigma, tor_sigma = tr_sigma,eps_rot_sigma,eps_tor_sigma
                elif freedom_to_noise == 1:
                    tr_sigma, rot_sigma, tor_sigma = eps_tr_sigma,rot_sigma,eps_tor_sigma
                else:
                    tr_sigma, rot_sigma, tor_sigma = eps_tr_sigma,eps_rot_sigma,tor_sigma
                    torsion_updates = np.random.normal(loc=0.0, scale=tor_sigma, size=data['ligand'].edge_mask.sum()) if torsion_updates is None else torsion_updates
                    # random selected opne edge to change
                    if data['ligand'].edge_mask.sum() >= 1:
                        tmp = np.zeros_like(torsion_updates)
                        # selected = np.random.randint(0, 10, 1)
                        selected = np.random.randint(0,data['ligand'].edge_mask.sum(),1)[0]
                        tmp[selected] = 1
                        torsion_updates = tmp*torsion_updates + eps_tor_sigma*torsion_updates
                    # torsion_updates = None if self.no_torsion else torsion_updates
            # 15% not noise any freedom degree
            else:
                # return data
                tr_sigma, rot_sigma, tor_sigma = eps_tr_sigma,eps_rot_sigma,eps_tor_sigma

        tr_update = torch.normal(mean=0, std=tr_sigma, size=(1, 3)) if tr_update is None else tr_update
        rot_update = so3.sample_vec(eps=rot_sigma) if rot_update is None else rot_update
        torsion_updates = np.random.normal(loc=0.0, scale=tor_sigma, size=data['ligand'].edge_mask.sum()) if torsion_updates is None else torsion_updates
        torsion_updates = None if self.no_torsion else torsion_updates

        modify_conformer(data, tr_update, torch.from_numpy(rot_update).float(), torsion_updates)
        data.tr_score = -tr_update / tr_sigma ** 2
        data.rot_score = torch.from_numpy(so3.score_vec(vec=rot_update, eps=rot_sigma)).float().unsqueeze(0)
        data.tor_score = None if self.no_torsion else torch.from_numpy(torus.score(torsion_updates, tor_sigma)).float()
        data.tor_sigma_edge = None if self.no_torsion else np.ones(data['ligand'].edge_mask.sum()) * tor_sigma
        return data
class NoiseTransform(BaseTransform):
    def __init__(self, t_to_sigma, no_torsion, all_atom):
        self.t_to_sigma = t_to_sigma
        self.no_torsion = no_torsion
        self.all_atom = all_atom

    def __call__(self, data):
        t = np.random.uniform()
        t_tr, t_rot, t_tor = t, t, t
        return self.apply_noise(data, t_tr, t_rot, t_tor)

    def apply_noise(self, data, t_tr, t_rot, t_tor, tr_update = None, rot_update=None, torsion_updates=None):
        # mdn mode make no update
        if not torch.is_tensor(data['ligand'].pos):
            data['ligand'].pos = random.choice(data['ligand'].pos)

        tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma(t_tr, t_rot, t_tor)
        set_time(data, t_tr, t_rot, t_tor, 1, self.all_atom, device=None)

        tr_update = torch.normal(mean=0, std=tr_sigma, size=(1, 3)) if tr_update is None else tr_update
        rot_update = so3.sample_vec(eps=rot_sigma) if rot_update is None else rot_update
        torsion_updates = np.random.normal(loc=0.0, scale=tor_sigma, size=data['ligand'].edge_mask.sum()) if torsion_updates is None else torsion_updates
        torsion_updates = None if self.no_torsion else torsion_updates
        modify_conformer(data, tr_update, torch.from_numpy(rot_update).float(), torsion_updates)

        data.tr_score = -tr_update / tr_sigma ** 2
        data.rot_score = torch.from_numpy(so3.score_vec(vec=rot_update, eps=rot_sigma)).float().unsqueeze(0)
        data.tor_score = None if self.no_torsion else torch.from_numpy(torus.score(torsion_updates, tor_sigma)).float()
        data.tor_sigma_edge = None if self.no_torsion else np.ones(data['ligand'].edge_mask.sum()) * tor_sigma
        return data

class PDBBind(Dataset):
    def __init__(self, root, transform=None, cache_path='data/cache', split_path='data/', limit_complexes=0,
                 receptor_radius=30, num_workers=1, c_alpha_max_neighbors=None, popsize=15, maxiter=15,
                 matching=True, keep_original=False, max_lig_size=None, remove_hs=False, num_conformers=1, all_atoms=False,
                 atom_radius=5, atom_max_neighbors=None, esm_embeddings_path=None, require_ligand=False,
                 ligands_list=None, protein_path_list=None, ligand_descriptions=None, keep_local_structures=False,surface_path = None):

        super(PDBBind, self).__init__(root, transform)
        self.surface_path = surface_path
        self.transform = transform
        self.pdbbind_dir = root
        self.max_lig_size = max_lig_size
        self.split_path = split_path
        self.limit_complexes = limit_complexes
        self.receptor_radius = receptor_radius
        self.num_workers = num_workers
        self.c_alpha_max_neighbors = c_alpha_max_neighbors
        self.remove_hs = remove_hs
        self.esm_embeddings_path = esm_embeddings_path
        self.require_ligand = require_ligand
        self.protein_path_list = protein_path_list
        self.ligand_descriptions = ligand_descriptions
        self.keep_local_structures = keep_local_structures
        if matching or protein_path_list is not None and ligand_descriptions is not None:
            cache_path += '_torsion'
        if all_atoms:
            cache_path += '_allatoms'
        self.full_cache_path = os.path.join(cache_path, f'limit{self.limit_complexes}'
                                                        f'_INDEX{os.path.splitext(os.path.basename(self.split_path))[0]}'
                                                        f'_maxLigSize{self.max_lig_size}_H{int(not self.remove_hs)}'
                                                        f'_recRad{self.receptor_radius}_recMax{self.c_alpha_max_neighbors}'
                                            + ('' if not all_atoms else f'_atomRad{atom_radius}_atomMax{atom_max_neighbors}')
                                            + ('' if not matching or num_conformers == 1 else f'_confs{num_conformers}')
                                            + ('' if self.esm_embeddings_path is None else f'_esmEmbeddings')
                                            + ('' if not keep_local_structures else f'_keptLocalStruct')
                                            + ('' if protein_path_list is None or ligand_descriptions is None else str(binascii.crc32(''.join(ligand_descriptions + protein_path_list).encode()))))
        
        self.popsize, self.maxiter = popsize, maxiter
        self.matching, self.keep_original = matching, keep_original
        self.num_conformers = num_conformers
        self.all_atoms = all_atoms
        self.atom_radius, self.atom_max_neighbors = atom_radius, atom_max_neighbors
        if not os.path.exists(os.path.join(self.full_cache_path, "heterographs_0.pkl"))\
                or (require_ligand and not os.path.exists(os.path.join(self.full_cache_path, "rdkit_ligands_0.pkl"))):
            os.makedirs(self.full_cache_path, exist_ok=True)
            if protein_path_list is None or ligand_descriptions is None:
                self.preprocessing()
            else:
                self.inference_preprocessing()
        logger.info('Training dataset size: {}'.format(len(glob.glob(os.path.join(self.full_cache_path,'heterographs_*.pkl')))))
        # logger.info('loading data from memory: ', os.path.join(self.full_cache_path, "heterographs.pkl"))

        # with open(os.path.join(self.full_cache_path, "heterographs.pkl"), 'rb') as f:
        #     self.complex_graphs = pickle.load(f)
        # print_statistics(self.complex_graphs)
        # logger.info('loaded data from memory: ', len(self.complex_graphs))
        # # filter out complexes with ligand not meet required in tarinset!
        # if 'timesplit' not in os.path.basename(self.split_path):

        #     filter_names = [i.strip() for i in open('/home/caoduanhua/DeepLearningForDock/DiffDockForScreen/diffScreen/data/pdbbind_pdbscreen/splits/data_filter_ligpre').readlines()]
        #     logger.info('only ligpre success data for train ')
        #     self.complex_graphs = [data for data in self.complex_graphs if data.name in filter_names]
        #     logger.info('only ligpre success data for train: nums: ',len(self.complex_graphs))

        # logger.info('loaded data from memory: ', len(self.complex_graphs))
        # if require_ligand:
        #     logger.info('loading ligand data from memory: ', os.path.join(self.full_cache_path, "rdkit_ligands.pkl"))
        #     with open(os.path.join(self.full_cache_path, "rdkit_ligands.pkl"), 'rb') as f:
        #         self.rdkit_ligands = pickle.load(f)
        #     logger.info('loaded ligand data from memory!')

    def len(self):
        return len(glob.glob(os.path.join(self.full_cache_path,'heterographs_*.pkl')))
        # return 80
    def get_complexs_list(self,num):
        graphs_list = []
        for idx in range(min(num,len(glob.glob(os.path.join(self.full_cache_path,'heterographs_*.pkl'))))):
            try:
                with open(os.path.join(self.full_cache_path,f'heterographs_{idx}.pkl'),'rb') as f:
                    complex_graph = pickle.load(f)
                complex_graph['ligand'].orig_pos -= complex_graph.original_center.numpy()
                complex_graph['receptor'].center_pos -= complex_graph.original_center.numpy()
                complex_graph['receptor'].atoms_pos -= complex_graph.original_center.numpy()
                graphs_list.append(complex_graph)
            except:
                continue
        return graphs_list

    def get(self, idx):
        if self.require_ligand:
            with open(os.path.join(self.full_cache_path,f'heterographs_{idx}.pkl'),'rb') as f:
                complex_graph = pickle.load(f)
            with open(os.path.join(self.full_cache_path,f'rdkit_ligands_{idx}.pkl'),'rb') as f:
                complex_graph.mol = pickle.load(f)
            # complex_graph = copy.deepcopy(self.complex_graphs[idx])
            # complex_graph.mol = copy.deepcopy(self.rdkit_ligands[idx])
            complex_graph['ligand'].orig_pos -= complex_graph.original_center.numpy()
            complex_graph['receptor'].center_pos -= complex_graph.original_center.numpy()
            complex_graph['receptor'].atoms_pos -= complex_graph.original_center.numpy()
            # for mdn traing
            if self.transform is None and not self.require_ligand:
                logger.info('for mdn traing, use original ligand pos')
                complex_graph['ligand'].pos = torch.from_numpy(complex_graph['ligand'].orig_pos).float()

            return complex_graph
        else:
            with open(os.path.join(self.full_cache_path,f'heterographs_{idx}.pkl'),'rb') as f:
                complex_graph = pickle.load(f)
            # complex_graph = copy.deepcopy(self.complex_graphs[idx])
            complex_graph['ligand'].orig_pos -= complex_graph.original_center.numpy()
            complex_graph['receptor'].center_pos -= complex_graph.original_center.numpy()
            complex_graph['receptor'].atoms_pos -= complex_graph.original_center.numpy()
            if self.transform is None and not self.require_ligand:
                # when use mdn traing ,use original ligand pos,test use rdkit pos
                logger.info('for mdn traing, use original ligand pos')
                complex_graph['ligand'].pos = torch.from_numpy(complex_graph['ligand'].orig_pos).float()
            return complex_graph
        
    def preprocessing(self):
        assert self.surface_path is not None,'surface_path is None please set this param if you want to use surface feature'
        logger.info(f'Processing complexes from [{self.split_path}] and saving it to [{self.full_cache_path}]')
        complex_names_all = read_strings_from_txt(self.split_path)
        logger.info('complex_names_all: ',len(complex_names_all))
        if self.limit_complexes is not None and self.limit_complexes != 0:
            complex_names_all = complex_names_all[:self.limit_complexes]
        logger.info(f'Loading {len(complex_names_all)} complexes.')
        if self.esm_embeddings_path is not None:
            # map protein name to embeddings , such as 5y80_protein_processed -> 1028 embeddings vectors
            id_to_embeddings = torch.load(self.esm_embeddings_path)
            # chain_embeddings_dictlist = defaultdict(list)
            # for key, embedding in id_to_embeddings.items():
            #     key_name = key # key_name is the protein name like 5y80
            #     if key_name in complex_names_all:
            #         chain_embeddings_dictlist[key_name].append(embedding)
            lm_embeddings_chains_all = []
            embedding_names = list(id_to_embeddings.keys())
            complex_names_all = [name for name in embedding_names if name.split('_')[0] in set(complex_names_all)]
            # complex_names_all = list(set(complex_names_all).intersection(set(embedding_names)))
            # logger.info('complex_names_all: ',len(complex_names_all))
            # logger.info(complex_names_all)
            complex_names_all_new = []
            for name in complex_names_all:
                try:
                    
                    lm_embeddings_chains_all.append(id_to_embeddings[name])
                    complex_names_all_new.append(name.split('_')[0])
                except:
                    # complex_names_all.remove(name.split('_')[0])
                    continue
            assert len(complex_names_all) == len(lm_embeddings_chains_all),'len(complex_names_all) {}!= {}len(lm_embeddings_chains_all)'.format(len(complex_names_all),len(lm_embeddings_chains_all))
                    # lm_embeddings_chains_all.append(None)
            complex_names_all = complex_names_all_new
        else:
            lm_embeddings_chains_all = [None] * len(complex_names_all)

        if self.num_workers > 1:
            # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
            for i in range(len(complex_names_all)//1000+1):
                if os.path.exists(os.path.join(self.full_cache_path, f"heterographs{i}.pkl")):
                    continue
                complex_names = complex_names_all[1000*i:1000*(i+1)]
                lm_embeddings_chains = lm_embeddings_chains_all[1000*i:1000*(i+1)]
                complex_graphs, rdkit_ligands = [], []
                # if self.num_workers > 1:
                #     p = Pool(self.num_workers, maxtasksperchild=1)
                #     p.__enter__()
                with tqdm(total=len(complex_names), desc=f'loading complexes {i}/{len(complex_names_all)//1000+1}') as pbar:
                    # map_fn = p.imap_unordered if self.num_workers > 1 else map
                    t_list = Parallel(n_jobs=self.num_workers, backend="multiprocessing")(delayed(self.get_complex)(x) for x in tqdm(zip(complex_names, lm_embeddings_chains, [None] * len(complex_names), [None] * len(complex_names)),total=len(complex_names)))
                    for t in t_list:
                        complex_graphs.extend(t[0])
                        rdkit_ligands.extend(t[1])
                        pbar.update()

                    # for t in map_fn(self.get_complex, zip(complex_names, lm_embeddings_chains, [None] * len(complex_names), [None] * len(complex_names))):
                    #     complex_graphs.extend(t[0])
                    #     rdkit_ligands.extend(t[1])
                    # pbar.update()
                # if self.num_workers > 1: p.__exit__(None, None, None)

                with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), 'wb') as f:
                    pickle.dump((complex_graphs), f,protocol=-1)
                with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), 'wb') as f:
                    pickle.dump((rdkit_ligands), f,protocol=-1)

            complex_graphs_all = []
            for i in range(len(complex_names_all)//1000+1):
                with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), 'rb') as f:
                    l = pickle.load(f)
                    complex_graphs_all.extend(l)
            with open(os.path.join(self.full_cache_path, f"heterographs.pkl"), 'wb') as f:
                pickle.dump((complex_graphs_all), f,protocol=-1)

            rdkit_ligands_all = []
            for i in range(len(complex_names_all) // 1000 + 1):
                with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), 'rb') as f:
                    l = pickle.load(f)
                    rdkit_ligands_all.extend(l)
            with open(os.path.join(self.full_cache_path, f"rdkit_ligands.pkl"), 'wb') as f:
                pickle.dump((rdkit_ligands_all), f,protocol=-1)
        else:
            complex_graphs, rdkit_ligands = [], []
            with tqdm(total=len(complex_names_all), desc='loading complexes') as pbar:
                for t in map(self.get_complex, zip(complex_names_all, lm_embeddings_chains_all, [None] * len(complex_names_all), [None] * len(complex_names_all))):

                    complex_graphs.extend(t[0])
                    rdkit_ligands.extend(t[1])
                    pbar.update()
            with open(os.path.join(self.full_cache_path, "heterographs.pkl"), 'wb') as f:
                pickle.dump((complex_graphs), f,protocol=-1)
            with open(os.path.join(self.full_cache_path, "rdkit_ligands.pkl"), 'wb') as f:
                pickle.dump((rdkit_ligands), f,protocol=-1)
    def inference_preprocessing(self):
        ligands_list = []
        logger.info('Reading molecules and generating local structures with RDKit (unless --keep_local_structures is turned on).')
        failed_ligand_indices = []
        for idx, ligand_description in tqdm(enumerate(self.ligand_descriptions)):
            try:
                mol = MolFromSmiles(ligand_description)  # check if it is a smiles or a path
                if mol is not None:
                    mol = AddHs(mol)
                    generate_conformer(mol)
                    ligands_list.append(mol)
                else:
                    mol = read_molecule(ligand_description, remove_hs=False, sanitize=True)
                    if mol is None:
                        raise Exception('RDKit could not read the molecule ', ligand_description)
                    if not self.keep_local_structures:
                        mol.RemoveAllConformers()
                        mol = AddHs(mol)
                        generate_conformer(mol)
                    ligands_list.append(mol)
            except Exception as e:

                logger.info('Failed to read molecule ', ligand_description, ' We are skipping it. The reason is the exception: ', e)
                failed_ligand_indices.append(idx)
        for index in sorted(failed_ligand_indices, reverse=True):
            del self.protein_path_list[index]
            del self.ligand_descriptions[index]

        if self.esm_embeddings_path is not None:
            logger.info('Reading language model embeddings.')
            lm_embeddings_chains_all = []
            if not os.path.exists(self.esm_embeddings_path): raise Exception('ESM embeddings path does not exist: ',self.esm_embeddings_path)
            for protein_path in self.protein_path_list:
                embeddings_paths = sorted(glob.glob(os.path.join(self.esm_embeddings_path, os.path.basename(protein_path)) + '*'))
                lm_embeddings_chains = []
                for embeddings_path in embeddings_paths:
                    lm_embeddings_chains.append(torch.load(embeddings_path)['representations'][33])
                lm_embeddings_chains_all.append(lm_embeddings_chains)
        else:
            lm_embeddings_chains_all = [None] * len(self.protein_path_list)

        logger.info('Generating graphs for ligands and proteins')
        if self.num_workers > 1:
            # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
            for i in range(len(self.protein_path_list)//1000+1):
                if os.path.exists(os.path.join(self.full_cache_path, f"heterographs{i}.pkl")):
                    continue
                protein_paths_chunk = self.protein_path_list[1000*i:1000*(i+1)]
                ligand_description_chunk = self.ligand_descriptions[1000*i:1000*(i+1)]
                ligands_chunk = ligands_list[1000 * i:1000 * (i + 1)]
                lm_embeddings_chains = lm_embeddings_chains_all[1000*i:1000*(i+1)]
                complex_graphs, rdkit_ligands = [], []
                if self.num_workers > 1:
                    p = Pool(self.num_workers, maxtasksperchild=1)
                    p.__enter__()
                with tqdm(total=len(protein_paths_chunk), desc=f'loading complexes {i}/{len(protein_paths_chunk)//1000+1}') as pbar:
                    map_fn = p.imap_unordered if self.num_workers > 1 else map
                    for t in map_fn(self.get_complex, zip(protein_paths_chunk, lm_embeddings_chains, ligands_chunk,ligand_description_chunk)):
                        complex_graphs.extend(t[0])
                        rdkit_ligands.extend(t[1])
                        pbar.update()
                if self.num_workers > 1: p.__exit__(None, None, None)

                with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), 'wb') as f:
                    pickle.dump((complex_graphs), f,protocol=-1)
                with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), 'wb') as f:
                    pickle.dump((rdkit_ligands), f,protocol=-1)

            complex_graphs_all = []
            for i in range(len(self.protein_path_list)//1000+1):
                with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), 'rb') as f:
                    l = pickle.load(f)
                    complex_graphs_all.extend(l)
            with open(os.path.join(self.full_cache_path, f"heterographs.pkl"), 'wb') as f:
                pickle.dump((complex_graphs_all), f,protocol=-1)

            rdkit_ligands_all = []
            for i in range(len(self.protein_path_list) // 1000 + 1):
                with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), 'rb') as f:
                    l = pickle.load(f)
                    rdkit_ligands_all.extend(l)
            with open(os.path.join(self.full_cache_path, f"rdkit_ligands.pkl"), 'wb') as f:
                pickle.dump((rdkit_ligands_all), f,protocol=-1)
        else:
            complex_graphs, rdkit_ligands = [], []
            with tqdm(total=len(self.protein_path_list), desc='loading complexes') as pbar:
                for t in map(self.get_complex, zip(self.protein_path_list, lm_embeddings_chains_all, ligands_list, self.ligand_descriptions)):
                    complex_graphs.extend(t[0])
                    rdkit_ligands.extend(t[1])
                    pbar.update()
            if complex_graphs == []: raise Exception('Preprocessing did not succeed for any complex')
            with open(os.path.join(self.full_cache_path, "heterographs.pkl"), 'wb') as f:
                pickle.dump((complex_graphs), f,protocol=-1)
            with open(os.path.join(self.full_cache_path, "rdkit_ligands.pkl"), 'wb') as f:
                pickle.dump((rdkit_ligands), f,protocol=-1)
    def get_complex(self, par):
        name, lm_embedding_chains, ligand, ligand_description = par
        if not os.path.exists(os.path.join(self.pdbbind_dir, name)) and ligand is None:
            logger.info(os.path.join(self.pdbbind_dir, name))
            logger.info("Folder not found", name)
            logger.info("Skipping", name)
            return [], []
        if ligand is not None:
            rec_model = parse_pdb_from_path(name)
            pure_pocket_path = os.path.join(os.path.splitext(name)[0],'_pure.pdb')
            # mda_rec_model = mda.Universe(name)
            name = f'{name}_{ligand_description}'
            ligs = [ligand]
        else:
            try:
                rec_path = glob.glob(f'{self.surface_path}/{name}/*.pdb')[0]
                rec_model = parse_pdb_from_path(rec_path)
                pure_pocket_path = rec_path.replace('.pdb','_pure.pdb')
                # mda_rec_model = mda.Universe(os.path.join(self.pdbbind_dir, name, f'{name}_pocket.pdb'))
            except Exception as e:
                logger.info(f'Skipping {name} because of the error:')
                logger.info(e)
                return [], []
            ligs = [read_abs_file_mol(os.path.join(self.pdbbind_dir, name,f'{name}_ligand.sdf'), remove_hs=False, sanitize=True)]
            # ligs = read_mols(self.pdbbind_dir, name, remove_hs=False)
        complex_graphs = []
        failed_indices = []
        if len(ligs)==0:
            logger.info(f'No ligands found for {name}')
            return [],[]
        # assert len(ligs) > 0, f'No ligands found for {name}'
        for i, lig in enumerate(ligs):
            if self.max_lig_size is not None and lig.GetNumHeavyAtoms() > self.max_lig_size:
                logger.info(f'Ligand with {lig.GetNumHeavyAtoms()} heavy atoms is larger than max_lig_size {self.max_lig_size}. Not including {name} in preprocessed data.')
                continue
            complex_graph = HeteroData()
            complex_graph['name'] = name
            try:
                get_lig_graph_with_matching(lig, complex_graph, self.popsize, self.maxiter, self.matching, self.keep_original,
                                            self.num_conformers, remove_hs=self.remove_hs)

                rec, rec_coords, c_alpha_coords, n_coords, c_coords, lm_embeddings = extract_receptor_structure(copy.deepcopy(rec_model), lig, save_file=pure_pocket_path,lm_embedding_chains=lm_embedding_chains)
                if lm_embeddings is not None and c_alpha_coords is not None and len(c_alpha_coords) != len(lm_embeddings):
                    assert lm_embeddings is not None and c_alpha_coords is not None and len(c_alpha_coords) == len(lm_embeddings),'length error'
                    logger.info(f'LM embeddings for complex {name} did not have the right length for the protein. Skipping {name}.')
                    failed_indices.append(i)
                    continue
                mda_rec_model = mda.Universe(pure_pocket_path)
                # raise 'pure_pocket_path : {}'.format(pure_pocket_path)
                get_rec_graph(mda_rec_model, rec_coords, c_alpha_coords, n_coords, c_coords, complex_graph, rec_radius=self.receptor_radius,
                                c_alpha_max_neighbors=self.c_alpha_max_neighbors, all_atoms=self.all_atoms,
                                atom_radius=self.atom_radius, atom_max_neighbors=self.atom_max_neighbors, remove_hs=self.remove_hs, lm_embeddings=lm_embeddings)
            except Exception as e:
                logger.info(f'Skipping {name} because of the rec_model parser error:')
                logger.info(e)
                failed_indices.append(i)
                continue
            protein_center = torch.mean(complex_graph['receptor'].pos, dim=0, keepdim=True)
            complex_graph['receptor'].pos -= protein_center
            if self.all_atoms:
                complex_graph['atom'].pos -= protein_center

            if (not self.matching) or self.num_conformers == 1:
                complex_graph['ligand'].pos -= protein_center
            else:
                for p in complex_graph['ligand'].pos:
                    p -= protein_center

            ligand_center = torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True)
            complex_graph.original_center = protein_center
            complex_graph.original_ligand_center = ligand_center + protein_center
            # add surface
            if self.surface_path is not None:
                try:
                    if len(glob.glob(f'{self.surface_path}/{name}/*.ply'))==0:
                        logger.info('no surface file for ',name)
                        failed_indices.append(i)
                        continue
                    with open(glob.glob(f'{self.surface_path}/{name}/*.ply')[0], 'rb') as f:
                        data = PlyData.read(f)
                    features = ([torch.tensor(data['vertex'][axis.name]) for axis in data['vertex'].properties if axis.name not in ['nx', 'ny', 'nz'] ])
                    pos = torch.stack(features[:3], dim=-1)
                    # pos 需要减去center_protein_pos
                    pos -= complex_graph.original_center
                    features = torch.stack(features[3:], dim=-1)
                    face = None
                    if 'face' in data:
                        faces = data['face']['vertex_indices']
                        faces = [torch.tensor(fa, dtype=torch.long) for fa in faces]
                        face = torch.stack(faces, dim=-1)
                    data = Data(x=features, pos=pos, face=face)
                    data = FaceToEdge()(data)
                    data = Cartesian(cat=False)(data)
                    complex_graph['surface'].pos = data.pos
                    complex_graph['surface'].x = data.x
                    complex_graph['surface','surface_edge','surface'].edge_index = data.edge_index
                    complex_graph['surface','surface_edge','surface'].edge_attr = data.edge_attr
                except Exception as e:
                    logger.info(f'Skipping {name} because of the surface error:')
                    logger.info(e)
                    failed_indices.append(i)
                    continue
            # surface end
            complex_graphs.append(complex_graph)
        for idx_to_delete in sorted(failed_indices, reverse=True):
            del ligs[idx_to_delete]

        return complex_graphs, ligs
def print_statistics(complex_graphs):
    statistics = ([], [], [], [])

    for complex_graph in complex_graphs:
        lig_pos = complex_graph['ligand'].pos if torch.is_tensor(complex_graph['ligand'].pos) else complex_graph['ligand'].pos[0]
        radius_protein = torch.max(torch.linalg.vector_norm(complex_graph['receptor'].pos, dim=1))
        molecule_center = torch.mean(lig_pos, dim=0)
        radius_molecule = torch.max(
            torch.linalg.vector_norm(lig_pos - molecule_center.unsqueeze(0), dim=1))
        distance_center = torch.linalg.vector_norm(molecule_center)
        statistics[0].append(radius_protein)
        statistics[1].append(radius_molecule)
        statistics[2].append(distance_center)
        if "rmsd_matching" in complex_graph:
            statistics[3].append(complex_graph.rmsd_matching)
        else:
            statistics[3].append(0)

    name = ['radius protein', 'radius molecule', 'distance protein-mol', 'rmsd matching']
    logger.info('Number of complexes: ', len(complex_graphs))
    for i in range(4):
        array = np.asarray(statistics[i])
        logger.info(f"{name[i]}: mean {np.mean(array)}, std {np.std(array)}, max {np.max(array)}")

def construct_loader(args, t_to_sigma):
    if args.transformStyle=='BERT':
        transform = NoiseTransformBERT(t_to_sigma=t_to_sigma, no_torsion=args.no_torsion,
                               all_atom=args.all_atoms) if not args.model_type == 'mdn_model' else None
    if args.transformStyle=='diffdock':
        transform = NoiseTransform(t_to_sigma=t_to_sigma, no_torsion=args.no_torsion,
                                all_atom=args.all_atoms) if not args.model_type == 'mdn_model' else None


    common_args = {'transform': transform, 'root': args.data_dir, 'limit_complexes': args.limit_complexes,
                   'receptor_radius': args.receptor_radius,
                   'c_alpha_max_neighbors': args.c_alpha_max_neighbors,
                   'remove_hs': args.remove_hs, 'max_lig_size': args.max_lig_size,
                   'matching': args.matching, 'popsize': args.matching_popsize, 'maxiter': args.matching_maxiter,
                   'num_workers': args.num_workers, 'all_atoms': args.all_atoms,
                   'atom_radius': args.atom_radius, 'atom_max_neighbors': args.atom_max_neighbors,
                   'esm_embeddings_path': args.esm_embeddings_path,'surface_path':args.surface_path}

    train_dataset = PDBBind(cache_path=args.cache_path, split_path=args.split_train, keep_original=True,
                            num_conformers=args.num_conformers, **common_args)
    val_dataset = PDBBind(cache_path=args.cache_path, split_path=args.split_val, keep_original=True, **common_args)

    # loader_class = DataListLoader if torch.cuda.is_available() else DataLoader
    loader_class = DataLoaderX
    # prefetch_factor = 0
    train_loader = loader_class(dataset=train_dataset, batch_size=args.batch_size, num_workers=args.num_dataloader_workers,shuffle=True, pin_memory=args.pin_memory,prefetch_factor = 2,drop_last = True)
    val_loader = loader_class(dataset=val_dataset, batch_size=args.batch_size, num_workers=args.num_dataloader_workers,shuffle=True, pin_memory=args.pin_memory,prefetch_factor = 2)

    return train_loader, val_loader

def read_mol(pdbbind_dir, name, remove_hs=False):
    lig = read_molecule(os.path.join(pdbbind_dir, name, f'{name}_ligand.sdf'), remove_hs=remove_hs, sanitize=True)
    if lig is None:  # read mol2 file if sdf file cannot be sanitized
        logger.info('Using the .sdf file failed. We found a .mol2 file instead and are trying to use that.')
        lig = read_molecule(os.path.join(pdbbind_dir, name, f'{name}_ligand.mol2'), remove_hs=remove_hs, sanitize=True)
    return lig
def read_abs_file_mol(file, remove_hs=False, sanitize=True):
    mol = read_molecule(file, remove_hs=remove_hs, sanitize=True)
    
    if file.endswith(".sdf") and mol is None:
        # mol = read_molecule(file, remove_hs=remove_hs, sanitize=True)
        if os.path.exists(file[:-4] + ".mol2"):
            logger.info('Using the .sdf file failed. We found a .mol2 file instead and are trying to use that.')
            mol = read_molecule(file[:-4] + ".mol2", remove_hs=remove_hs, sanitize=True)
    elif file.endswith(".mol2") and mol is None:
        if os.path.exists(file[:-4] + ".sdf"):
            logger.info('Using the .mol2 file failed. We found a .sdf file instead and are trying to use that.')
            mol = read_molecule(file[:-4] + ".sdf", remove_hs=remove_hs, sanitize=True)

    return mol
from rdkit.Chem import AllChem


def read_mols(pdbbind_dir, name, remove_hs=False):
    ligs = []
    for file in os.listdir(os.path.join(pdbbind_dir, name)):
        if  'rdkit' not in file:
            if file.endswith(".sdf"):
                lig = read_molecule(os.path.join(pdbbind_dir, name, file), remove_hs=remove_hs, sanitize=True)
                if lig is not None:
                    try:
                        mol_rdkit = copy.deepcopy(lig)
                        mol_rdkit.RemoveAllConformers()
                        mol_rdkit = AllChem.AddHs(mol_rdkit)
                        generate_conformer(mol_rdkit)
                        ligs.append(lig)
                        break
                    except:
                        continue
                else:
                    continue
            elif file.endswith(".mol2"):
                lig = read_molecule(os.path.join(pdbbind_dir, name, file), remove_hs=remove_hs, sanitize=True)
                if lig is not None:
                    try:
                        mol_rdkit = copy.deepcopy(lig)
                        mol_rdkit.RemoveAllConformers()
                        mol_rdkit = AllChem.AddHs(mol_rdkit)
                        generate_conformer(mol_rdkit)
                        ligs.append(lig)
                        break
                    except:
                        continue
                else:
                    continue
    return ligs