File size: 1,626 Bytes
9ae74ae | 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 51 52 53 54 55 56 | #!/usr/bin/env python
import numpy as np
import os
import argparse
# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument("--pdbdir", type=str, required=True)
parser.add_argument("--trbdir", type=str, required=True)
parser.add_argument("--verbose", action="store_true", default=False)
args = parser.parse_args()
pdb_list = os.listdir(args.pdbdir)
for pdb in pdb_list:
if not pdb.endswith(".pdb"):
if args.verbose:
print(f"Skipping {pdb} as it is not a PDB file")
continue
pdb_path = os.path.join(args.pdbdir, pdb)
trb_path = os.path.join(args.trbdir, f'{pdb.split(".")[0]}.trb')
if not os.path.exists(trb_path):
print(f"Error: TRB file not found for PDB {pdb}")
continue
data = np.load(trb_path, allow_pickle=True)
if 'receptor_con_hal_pdb_idx' in data:
#Identify the last residue number in A chain
last_res_id = int(data['receptor_con_hal_pdb_idx'][0][1]) - 1
#Identify where inpaint seq is True, ie, kept fixed
indices = np.where(data['inpaint_seq'][:last_res_id])[0]
else:
#Identify where inpaint seq is True in the only chain
indices = np.where(data['inpaint_seq'])[0]
if args.verbose:
print(f"Adding FIXED labels to {pdb} at positions {indices}")
remarks = []
for position in indices:
remark = f"REMARK PDBinfo-LABEL:{position+1: >5} FIXED"
remarks.append(remark)
# Uncomment below to add hotspots
remarks_str = '\n'.join(remarks)
with open(pdb_path, 'a') as f:
f.write('\n')
f.write(remarks_str)
|