OneScience-Group/ProteinMPNN

1import argparse2import os3import sys4 5_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))6while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")):7    _PARENT = os.path.dirname(_PROJECT_ROOT)8    if _PARENT == _PROJECT_ROOT:9        break10    _PROJECT_ROOT = _PARENT11_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model")12_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT")13for _path in (_MODEL_ROOT, _PROJECT_ROOT):14    if os.path.exists(_path) and _path not in sys.path:15        sys.path.insert(0, _path)16if _ONESCIENCE_ROOT:17    _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src")18    for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT):19        if os.path.exists(_path) and _path not in sys.path:20            sys.path.insert(0, _path)21 22def main(args):23 24    import numpy as np25    import os, time, gzip, json26    import glob 27    28    folder_with_pdbs_path = args.input_path29    save_path = args.output_path30    ca_only = args.ca_only31    32    alpha_1 = list("ARNDCQEGHILKMFPSTWYV-")33    states = len(alpha_1)34    alpha_3 = ['ALA','ARG','ASN','ASP','CYS','GLN','GLU','GLY','HIS','ILE',35               'LEU','LYS','MET','PHE','PRO','SER','THR','TRP','TYR','VAL','GAP']36    37    aa_1_N = {a:n for n,a in enumerate(alpha_1)}38    aa_3_N = {a:n for n,a in enumerate(alpha_3)}39    aa_N_1 = {n:a for n,a in enumerate(alpha_1)}40    aa_1_3 = {a:b for a,b in zip(alpha_1,alpha_3)}41    aa_3_1 = {b:a for a,b in zip(alpha_1,alpha_3)}42    43    def AA_to_N(x):44      # ["ARND"] -> [[0,1,2,3]]45      x = np.array(x);46      if x.ndim == 0: x = x[None]47      return [[aa_1_N.get(a, states-1) for a in y] for y in x]48    49    def N_to_AA(x):50      # [[0,1,2,3]] -> ["ARND"]51      x = np.array(x);52      if x.ndim == 1: x = x[None]53      return ["".join([aa_N_1.get(a,"-") for a in y]) for y in x]54    55    56    def parse_PDB_biounits(x, atoms=['N','CA','C'], chain=None):57      '''58      input:  x = PDB filename59              atoms = atoms to extract (optional)60      output: (length, atoms, coords=(x,y,z)), sequence61      '''62      xyz,seq,min_resn,max_resn = {},{},1e6,-1e663      for line in open(x,"rb"):64        line = line.decode("utf-8","ignore").rstrip()65    66        if line[:6] == "HETATM" and line[17:17+3] == "MSE":67          line = line.replace("HETATM","ATOM  ")68          line = line.replace("MSE","MET")69    70        if line[:4] == "ATOM":71          ch = line[21:22]72          if ch == chain or chain is None:73            atom = line[12:12+4].strip()74            resi = line[17:17+3]75            resn = line[22:22+5].strip()76            x,y,z = [float(line[i:(i+8)]) for i in [30,38,46]]77    78            if resn[-1].isalpha(): 79                resa,resn = resn[-1],int(resn[:-1])-180            else: 81                resa,resn = "",int(resn)-182    #         resn = int(resn)83            if resn < min_resn: 84                min_resn = resn85            if resn > max_resn: 86                max_resn = resn87            if resn not in xyz: 88                xyz[resn] = {}89            if resa not in xyz[resn]: 90                xyz[resn][resa] = {}91            if resn not in seq: 92                seq[resn] = {}93            if resa not in seq[resn]: 94                seq[resn][resa] = resi95    96            if atom not in xyz[resn][resa]:97              xyz[resn][resa][atom] = np.array([x,y,z])98    99      # convert to numpy arrays, fill in missing values100      seq_,xyz_ = [],[]101      try:102          for resn in range(min_resn,max_resn+1):103            if resn in seq:104              for k in sorted(seq[resn]): seq_.append(aa_3_N.get(seq[resn][k],20))105            else: seq_.append(20)106            if resn in xyz:107              for k in sorted(xyz[resn]):108                for atom in atoms:109                  if atom in xyz[resn][k]: xyz_.append(xyz[resn][k][atom])110                  else: xyz_.append(np.full(3,np.nan))111            else:112              for atom in atoms: xyz_.append(np.full(3,np.nan))113          return np.array(xyz_).reshape(-1,len(atoms),3), N_to_AA(np.array(seq_))114      except TypeError:115          return 'no_chain', 'no_chain'116    117    118    119    pdb_dict_list = []120    c = 0121    122    if folder_with_pdbs_path[-1]!='/':123        folder_with_pdbs_path = folder_with_pdbs_path+'/'124    125    126    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']127    extra_alphabet = [str(item) for item in list(np.arange(300))]128    chain_alphabet = init_alphabet + extra_alphabet129    130    biounit_names = glob.glob(folder_with_pdbs_path+'*.pdb')131    for biounit in biounit_names:132        my_dict = {}133        s = 0134        concat_seq = ''135        concat_N = []136        concat_CA = []137        concat_C = []138        concat_O = []139        concat_mask = []140        coords_dict = {}141        for letter in chain_alphabet:142            if ca_only:143                sidechain_atoms = ['CA']144            else:145                sidechain_atoms = ['N', 'CA', 'C', 'O']146            xyz, seq = parse_PDB_biounits(biounit, atoms=sidechain_atoms, chain=letter)147            if type(xyz) != str:148                concat_seq += seq[0]149                my_dict['seq_chain_'+letter]=seq[0]150                coords_dict_chain = {}151                if ca_only:152                    coords_dict_chain['CA_chain_'+letter]=xyz.tolist()153                else:154                    coords_dict_chain['N_chain_' + letter] = xyz[:, 0, :].tolist()155                    coords_dict_chain['CA_chain_' + letter] = xyz[:, 1, :].tolist()156                    coords_dict_chain['C_chain_' + letter] = xyz[:, 2, :].tolist()157                    coords_dict_chain['O_chain_' + letter] = xyz[:, 3, :].tolist()158                my_dict['coords_chain_'+letter]=coords_dict_chain159                s += 1160        fi = biounit.rfind("/")161        my_dict['name']=biounit[(fi+1):-4]162        my_dict['num_of_chains'] = s163        my_dict['seq'] = concat_seq164        if s < len(chain_alphabet):165            pdb_dict_list.append(my_dict)166            c+=1167            168            169    with open(save_path, 'w') as f:170        for entry in pdb_dict_list:171            f.write(json.dumps(entry) + '\n')172           173 174if __name__ == "__main__":175    argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)176 177    argparser.add_argument("--input_path", type=str, help="Path to a folder with pdb files, e.g. /home/my_pdbs/")178    argparser.add_argument("--output_path", type=str, help="Path where to save .jsonl dictionary of parsed pdbs")179    argparser.add_argument("--ca_only", action="store_true", default=False, help="parse a backbone-only structure (default: false)")180 181    args = argparser.parse_args()182    main(args)183