File size: 2,030 Bytes
feb41b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | import os
import argparse
import pandas as pd
import biotite.structure.io as bsio
from tqdm import tqdm
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--pdb_dir", type=str, default=None)
parser.add_argument("--pdb_file", type=str, default=None)
parser.add_argument("--out_file", type=str, default="plddt.csv")
parser.add_argument("--type", type=str, choices=["residue", "protein"], default="protein")
args = parser.parse_args()
if args.pdb_dir is not None:
if args.type == "protein":
out_info = {"pdb": [], "plddt": []}
elif args.type == "residue":
out_info = {"pdb": []}
pdbs = sorted(os.listdir(args.pdb_dir))
for pdb in tqdm(pdbs):
pdb_file = os.path.join(args.pdb_dir, pdb)
struct = bsio.load_structure(pdb_file, extra_fields=["b_factor"])
if args.type == "protein":
plddt = struct.b_factor.mean()
out_info["pdb"].append(pdb)
out_info["plddt"].append(plddt)
elif args.type == "residue":
out_info["pdb"].append(pdb)
res_list = []
for res, plddt in zip(struct.res_id, struct.b_factor):
if res not in res_list:
if not out_info.get(res):
out_info[res] = []
res_list.append(res)
out_info[res].append(plddt)
pd.DataFrame(out_info).to_csv(args.out_file, index=False)
else:
struct = bsio.load_structure(args.pdb_file, extra_fields=["b_factor"])
if args.type == "protein":
plddt = struct.b_factor.mean()
elif args.type == "residue":
res_list = []
plddt = []
for res, b_factor in zip(struct.res_id, struct.b_factor):
if res not in res_list:
res_list.append(res)
plddt.append(b_factor)
print(plddt) |