| import argparse |
|
|
|
|
| def main(args): |
|
|
| import glob |
| import gzip |
| import json |
| import os |
| import time |
|
|
| import numpy as np |
|
|
| folder_with_pdbs_path = args.input_path |
| save_path = args.output_path |
| ca_only = args.ca_only |
|
|
| 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): |
| |
| 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): |
| |
| 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] |
|
|
| def parse_PDB_biounits(x, atoms=["N", "CA", "C"], chain=None): |
| """ |
| input: x = PDB filename |
| atoms = atoms to extract (optional) |
| output: (length, atoms, coords=(x,y,z)), sequence |
| """ |
| xyz, seq, min_resn, max_resn = {}, {}, 1e6, -1e6 |
| for line in open(x, "rb"): |
| 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: |
| atom = line[12 : 12 + 4].strip() |
| resi = line[17 : 17 + 3] |
| resn = line[22 : 22 + 5].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 |
| else: |
| resa, resn = "", int(resn) - 1 |
| |
| 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]) |
|
|
| |
| seq_, xyz_ = [], [] |
| 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)) |
| else: |
| seq_.append(20) |
| 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]) |
| else: |
| xyz_.append(np.full(3, np.nan)) |
| else: |
| for atom in atoms: |
| xyz_.append(np.full(3, np.nan)) |
| return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)) |
| except TypeError: |
| return "no_chain", "no_chain" |
|
|
| pdb_dict_list = [] |
| c = 0 |
|
|
| if folder_with_pdbs_path[-1] != "/": |
| folder_with_pdbs_path = folder_with_pdbs_path + "/" |
|
|
| 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 = glob.glob(folder_with_pdbs_path + "*.pdb") |
| for biounit in biounit_names: |
| my_dict = {} |
| s = 0 |
| concat_seq = "" |
| concat_N = [] |
| concat_CA = [] |
| concat_C = [] |
| concat_O = [] |
| concat_mask = [] |
| coords_dict = {} |
| for letter in chain_alphabet: |
| if ca_only: |
| sidechain_atoms = ["CA"] |
| else: |
| sidechain_atoms = ["N", "CA", "C", "O"] |
| xyz, seq = parse_PDB_biounits(biounit, atoms=sidechain_atoms, chain=letter) |
| if type(xyz) != str: |
| concat_seq += seq[0] |
| my_dict["seq_chain_" + letter] = seq[0] |
| coords_dict_chain = {} |
| if ca_only: |
| coords_dict_chain["CA_chain_" + letter] = xyz.tolist() |
| else: |
| coords_dict_chain["N_chain_" + letter] = xyz[:, 0, :].tolist() |
| coords_dict_chain["CA_chain_" + letter] = xyz[:, 1, :].tolist() |
| coords_dict_chain["C_chain_" + letter] = xyz[:, 2, :].tolist() |
| coords_dict_chain["O_chain_" + letter] = xyz[:, 3, :].tolist() |
| my_dict["coords_chain_" + letter] = coords_dict_chain |
| s += 1 |
| fi = biounit.rfind("/") |
| my_dict["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 |
|
|
| with open(save_path, "w") as f: |
| for entry in pdb_dict_list: |
| f.write(json.dumps(entry) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| argparser = argparse.ArgumentParser( |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| ) |
|
|
| argparser.add_argument( |
| "--input_path", |
| type=str, |
| help="Path to a folder with pdb files, e.g. /home/my_pdbs/", |
| ) |
| argparser.add_argument( |
| "--output_path", |
| type=str, |
| help="Path where to save .jsonl dictionary of parsed pdbs", |
| ) |
| argparser.add_argument( |
| "--ca_only", |
| action="store_true", |
| default=False, |
| help="parse a backbone-only structure (default: false)", |
| ) |
|
|
| args = argparser.parse_args() |
| main(args) |
|
|
