Flexpert-Design/data_utils.py

#From https://github.com/JoreyYan/zetadesign/blob/master/data/data.py
import glob
import json
import numpy as np
import gzip
import re
import multiprocessing
import tqdm
import shutil
SENTINEL = 1
import biotite.structure as struc
import biotite.application.dssp as dssp
import biotite.structure.io.pdb.file as file


def parse_PDB_biounits(x, sse,ssedssp,atoms=['N', 'CA', 'C'], chain=None):
    '''
    input:  x = PDB filename
            atoms = atoms to extract (optional)
    output: (length, atoms, coords=(x,y,z)), sequence
    '''

    alpha_1 = list("ARNDCQEGHILKMFPSTWYV-")
    states = len(alpha_1)
    alpha_3 = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE',
               'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL', 'GAP']

    aa_1_N = {a: n for n, a in enumerate(alpha_1)}
    aa_3_N = {a: n for n, a in enumerate(alpha_3)}
    aa_N_1 = {n: a for n, a in enumerate(alpha_1)}
    aa_1_3 = {a: b for a, b in zip(alpha_1, alpha_3)}
    aa_3_1 = {b: a for a, b in zip(alpha_1, alpha_3)}

    def AA_to_N(x):
        # ["ARND"] -> [[0,1,2,3]]
        x = np.array(x);
        if x.ndim == 0: x = x[None]
        return [[aa_1_N.get(a, states - 1) for a in y] for y in x]

    def N_to_AA(x):
        # [[0,1,2,3]] -> ["ARND"]
        x = np.array(x);
        if x.ndim == 1: x = x[None]
        return ["".join([aa_N_1.get(a, "-") for a in y]) for y in x]

    xyz, seq, plddts, min_resn, max_resn = {}, {}, [],  1e6, -1e6

    pdbcontents = x.split('\n')[0]
    with open(pdbcontents) as f:
        pdbcontents = f.readlines()
    for line in pdbcontents:
        #line = line.decode("utf-8", "ignore").rstrip()

        if line[:6] == "HETATM" and line[17:17 + 3] == "MSE":
            line = line.replace("HETATM", "ATOM  ")
            line = line.replace("MSE", "MET")

        if line[:4] == "ATOM":
            ch = line[21:22]
            if ch == chain or chain is None or ch==' ':
                atom = line[12:12 + 4].strip()
                resi = line[17:17 + 3]
                resn = line[22:22 + 5].strip()
                plddt=line[60:60 + 6].strip()



                x, y, z = [float(line[i:(i + 8)]) for i in [30, 38, 46]]

                if resn[-1].isalpha():
                    resa, resn = resn[-1], int(resn[:-1]) - 1 # in same pos ,use last atoms
                else:
                    resa, resn = "_", int(resn) - 1
                #         resn = int(resn)
                if resn < min_resn:
                    min_resn = resn
                if resn > max_resn:
                    max_resn = resn



                if resn not in xyz:
                    xyz[resn] = {}
                if resa not in xyz[resn]:
                    xyz[resn][resa] = {}
                if resn not in seq:
                    seq[resn] = {}

                if resa not in seq[resn]:
                    seq[resn][resa] = resi

                if atom not in xyz[resn][resa]:
                    xyz[resn][resa][atom] = np.array([x, y, z])



    # convert to numpy arrays, fill in missing values
    seq_, xyz_ ,sse_,ssedssp_= [], [], [], []
    dsspidx=0
    sseidx=0
    # try:
    #     for resn in range(min_resn, max_resn + 1):
    #         if resn in seq:
    #             for k in sorted(seq[resn]):
    #                 seq_.append(aa_3_N.get(seq[resn][k], 20))
    #                 try:
    #                     if 'CA' in xyz[resn][k]:
    #                         sse_.append(sse[sseidx])
    #                         sseidx = sseidx + 1
    #                     else:
    #                         sse_.append('-')
    #                 except:
    #                     print('error sse')


    #         else:
    #             seq_.append(20)
    #             sse_.append('-')

    #         misschianatom = False
    #         if resn in xyz:


    #             for k in sorted(xyz[resn]):
    #                 for atom in atoms:
    #                     if atom in xyz[resn][k]:
    #                         xyz_.append(xyz[resn][k][atom])  #some will miss C and O ,but sse is normal,because sse just depend on CA
    #                     else:
    #                         xyz_.append(np.full(3, np.nan))
    #                         misschianatom=True
    #                 if misschianatom:
    #                     ssedssp_.append('-')
    #                     misschianatom = False
    #                 else:
    #                     try:
    #                         ssedssp_.append(ssedssp[dsspidx])         # if miss chain atom,xyz ,seq think is ok , but dssp miss this
    #                         dsspidx = dsspidx + 1
    #                     except:
    #                         print(dsspidx)


    #         else:
    #             for atom in atoms:
    #                 xyz_.append(np.full(3, np.nan))
    #             ssedssp_.append('-')


    #     return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)),np.array(sse_),np.array(ssedssp_)
    # except TypeError as e:
    #     print(f"TypeError: {e}")
    #     return 'no_chain', 'no_chain','no_chain'

    for resn in range(int(min_resn), int(max_resn + 1)):
        if resn in seq:
            for k in sorted(seq[resn]):
                seq_.append(aa_3_N.get(seq[resn][k], 20))
                try:
                    if 'CA' in xyz[resn][k]:
                        sse_.append(sse[sseidx])
                        sseidx = sseidx + 1
                    else:
                        sse_.append('-')
                except:
                    print('error sse')


        else:
            seq_.append(20)
            sse_.append('-')

        misschianatom = False
        if resn in xyz:


            for k in sorted(xyz[resn]):
                for atom in atoms:
                    if atom in xyz[resn][k]:
                        xyz_.append(xyz[resn][k][atom])  #some will miss C and O ,but sse is normal,because sse just depend on CA
                    else:
                        xyz_.append(np.full(3, np.nan))
                        misschianatom=True
                if misschianatom:
                    ssedssp_.append('-')
                    misschianatom = False
                else:
                    try:
                        ssedssp_.append(ssedssp[dsspidx])         # if miss chain atom,xyz ,seq think is ok , but dssp miss this
                        dsspidx = dsspidx + 1
                    except:
                        print(dsspidx)


        else:
            for atom in atoms:
                xyz_.append(np.full(3, np.nan))
            ssedssp_.append('-')


    return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)),np.array(sse_),np.array(ssedssp_)


def parse_PDB(path_to_pdb,name, input_chain_list=None):
    """
    make sure every time just input 1 line
    """
    c = 0
    pdb_dict_list = []


    if input_chain_list:
        chain_alphabet = input_chain_list
    else:
        init_alphabet = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
                         'T',
                         'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
                         'n',
                         'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
        extra_alphabet = [str(item) for item in list(np.arange(300))]
        chain_alphabet = init_alphabet + extra_alphabet

    biounit_names = [path_to_pdb]
    for biounit in biounit_names:
        my_dict = {}
        s = 0
        concat_seq = ''


        for letter in chain_alphabet:

            PDBFile = file.PDBFile.read(biounit)
            array_stack = PDBFile.get_structure(altloc="all")


            sse1 = struc.annotate_sse(array_stack[0], chain_id=letter).tolist()
            if len(sse1)==0:
                sse1 = struc.annotate_sse(array_stack[0], chain_id='').tolist()
            #ssedssp1 = dssp.DsspApp.annotate_sse(array_stack).tolist()
            ssedssp1 = [] #not annotating dssp for now


            xyz, seq, _, _= parse_PDB_biounits(biounit,sse1,ssedssp1,atoms=['N', 'CA', 'C','O'], chain=letter) #TODO: fix the float error
            #ssedssp = sse  #faking it for now
            # if len(sse)!=len(seq[0]):
            #     xxxx=len(seq[0])
            #     print(name)
            #assert len(sse)==len(seq[0])
            #assert len(ssedssp) == len(seq[0])

            if type(xyz) != str:
                concat_seq += seq[0]
                my_dict['seq_chain_' + letter] = seq[0]

                coords_dict_chain = {}
                coords_dict_chain['N'] = xyz[:, 0, :].tolist()
                coords_dict_chain['CA'] = xyz[:, 1, :].tolist()
                coords_dict_chain['C'] = xyz[:, 2, :].tolist()
                coords_dict_chain['O'] = xyz[:, 3, :].tolist()
                my_dict['coords_chain_' + letter] = coords_dict_chain

                #sse=''.join(sse)
                #ssedssp=''.join(ssedssp)
                #my_dict['sse3' ] = sse
                #my_dict['sse8'] = ssedssp
                s += 1
        #fi = biounit.rfind("/")
        my_dict['name'] = name#biounit[(fi + 1):-4]
        my_dict['num_of_chains'] = s
        my_dict['seq'] = concat_seq
        if s <= len(chain_alphabet):
            pdb_dict_list.append(my_dict)
            c += 1
    return pdb_dict_list





def parse_pdb_split_chain(pdbgzFile):

    with open(pdbgzFile) as f:


        lines = f.readlines()
        # pdbcontent = f.decode()


        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
        match = list(set(list(pattern.findall(lines[0]))))


    name=pdbgzFile.split('/')[-1]
        #for chain in match:
        # parse_PDB
    # match=[name[4]]
    # match=['A']
    pdb_data=parse_PDB(pdbgzFile,name,match)

    return pdb_data
def parse_pdb_split_chain_af(pdbgzFile):
    with gzip.open(pdbgzFile, 'rb') as pdbF:
        try:
            pdbcontent = pdbF.read()
        except:
            print(pdbgzFile)

        pdbcontent = pdbcontent.decode()


        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
        match = list(set(list(pattern.findall(pdbcontent))))


    name=pdbgzFile.split('/')[-1].split('.')[0]
        #for chain in match:
        # parse_PDB
    # match=[name[4]]
    # match=[1]
    pdb_data=parse_PDB('/media/junyu/data/perotin/aftest080_1000/'+pdbgzFile.split('/')[-1].split('.')[0]+'.pdb',name,match)

    return pdb_data

def parse_pdb_split_chain_af_3dcnn(pdbgzFile):
    with gzip.open(pdbgzFile, 'rb') as pdbF:
        try:
            pdbcontent = pdbF.read()
        except:
            print(pdbgzFile)

        pdbcontent = pdbcontent.decode()


        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
        match = list(set(list(pattern.findall(pdbcontent))))


    name=pdbgzFile.split('/')[-1].split('.')[0]
    namelist=[]
    for chain in match:
        namelist.append(name+'__'+chain)
    # match=[name[4]]
    # match=[1]



    return namelist
def run_net(files_path,output_path):
    """
    input is pdbgz's dir
    from pdb to jsonl
    """
    list=glob.glob(files_path+'*.pdb')#[:3110]
    data=[]
    for i in tqdm.tqdm(list):
        data_chains=parse_pdb_split_chain(i)
        #for chian in data_chains:
        data.append(data_chains[0])

    print('we want to write now')
    with open(output_path, 'w') as f:
        for entry in data:
            f.write(json.dumps(entry) + '\n')

    f.close()
    print('finished')

def run_netbyondif(filelist,output_path):
    with open(filelist) as f:

        lines = f.readlines()
    data=[]
    data_1=[]
    # data_2 = []
    # data_3 = []
    # data_4 = []
    # data_5 = []
    # data_6 = []
    # data_7 = []
    # data_8 = []
    # data_9 = []
    # data_10 = []
    nums_dict={1:0,2:0,3:0,4:0,5:0,6:0,7:0,8:0,9:0,10:0,}

    for i in tqdm.tqdm(lines):
        data_chains,match=parse_pdb_split_chain(i.split('"')[1])

        for chian in data_chains:
            for i in match:
                meanplddt = round(float(np.asarray(chian['plddts_chain_' + i]).mean()),2)
                data.append({'name':chian['name'],'lens':len(chian['seq']),'meanplddt':meanplddt})
                if int(meanplddt/10)==1:
                    #data_1.append(chian)
                    nums_dict[1]=nums_dict[1]+1
                elif int(meanplddt/10)==2:
                    #data_2.append(chian)
                    nums_dict[2] = nums_dict[2] + 1
                elif int(meanplddt / 10) == 3:
                    #data_3.append(chian)
                    nums_dict[3] = nums_dict[3] + 1
                elif int(meanplddt / 10) == 4:
                    #data_4.append(chian)
                    nums_dict[4] = nums_dict[4] + 1
                elif int(meanplddt / 10) == 5:
                    #data_5.append(chian)
                    nums_dict[5] = nums_dict[5] + 1
                elif int(meanplddt / 10) == 6:
                    #data_6.append(chian)
                    nums_dict[6] = nums_dict[6] + 1
                elif int(meanplddt / 10) == 7:
                    #data_7.append(chian)
                    nums_dict[7] = nums_dict[7] + 1
                elif int(meanplddt / 10) == 8:
                    #data_8.append(chian)
                    nums_dict[8] = nums_dict[8] + 1
                elif int(meanplddt / 10) == 9:
                    #data_9.append(chian)
                    nums_dict[9] = nums_dict[9] + 1
                elif int(meanplddt / 10) == 10:
                    #data_10.append(chian)
                    nums_dict[10] = nums_dict[10] + 1
                else:
                    print(chian['name'])


    #         data.append(chian)
    #
    f.close()
    output_pathindex=output_path+filelist.split('/')[-1].split('.')[0]+'_detail.jsonl'
    print('we want to write now')
    with open(output_pathindex, 'w') as f:
        for entry in data:
            f.write(json.dumps(entry) + '\n')

    f.close()
    #print(nums_dict)
    # count(output_pathindex)
    print('finished')
def list_of_groups(list_info, per_list_len):
    '''
    :param list_info:   列表
    :param per_list_len:  每个小列表的长度
    :return:
    '''
    list_of_group = zip(*(iter(list_info),) *per_list_len)
    end_list = [list(i) for i in list_of_group] # i is a tuple
    count = len(list_info) % per_list_len
    end_list.append(list_info[-count:]) if count !=0 else end_list
    return end_list

def count(filelist):
    with open(filelist) as f:

        lines = f.readlines()
    plddts=[]

    for i in tqdm.tqdm(lines):
        pl=json.loads(i)['meanplddt']
        plddts.append(int(pl/10))

    for i in range(10):
        print('counts '+str(i),plddts.count(i))

def run_net_aftest(files_path,output_path):
    """
    input is pdbgz's dir
    """
    with open(files_path) as f:
        lines = f.readlines()
    data=[]
    for i in tqdm.tqdm(lines):

        data_chains=parse_pdb_split_chain_af('/media/junyu/data/point_cloud/'+i.split('"')[1])
        for chian in data_chains:
            data.append(chian)

    # print('we want to write now')
    # with open(output_path, 'w') as f:
    #     for entry in data:
    #         f.write(json.dumps(entry) + '\n')
    #
    # f.close()
    # print('finished')

    output_pathindex = output_path + str(80) + 'bigthanclass_1000.text'
    print('we want to write now')
    with open(output_pathindex, 'w') as f:
        for entry in data:
            f.write(entry + '\n')


    f.close()

# if __name__ == "__main__":
#     files_path='/media/junyu/data/perotin/chain_set/AFDATA/details/80bigthanclass_1000.jsonl' #'/home/junyu/下载/splits/'#
#     output_path='/media/junyu/data/perotin/chain_set/'



#     # run_net_aftest(files_path,output_path)

#     fakedata='//home/oem/pdb-tools/pdbtools/fixed/'
#     run_net(fakedata,output_path+'tim184.jsonl')



        #
        # f.close()
        # # print(nums_dict)
        # print('finished ' +str(i))


    # alllist=list_of_groups(lists,10000)

    # for i in range(len(alllist)):
    #     thislist=alllist[i]
    #     with open(output_path+'_'+str(i)+'.jsonl', 'w') as f:
    #         for entry in thislist:
    #             f.write(json.dumps(entry) + '\n')
    #
    #     f.close()
    #     # print(nums_dict)
    #     print('finished ' +str(i))

    # _processes = []



    # q = multiprocessing.Queue()
    #
    # proc.start()
    # for eachlist in alllist:
    #     _process = multiprocessing.Process(target=run_netbyondif, args=(eachlist,))
    #     _process.start()


    # run_netbyondif(lists,output_path)