Spaces:
Sleeping
Sleeping
File size: 2,782 Bytes
00aa807 |
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 57 58 59 60 |
import argparse
def main(args):
import glob
import random
import numpy as np
import json
import itertools
with open(args.input_path, 'r') as json_file:
json_list = list(json_file)
fixed_list = [[int(item) for item in one.split()] for one in args.position_list.split(",")]
global_designed_chain_list = [str(item) for item in args.chain_list.split()]
my_dict = {}
if not args.specify_non_fixed:
for json_str in json_list:
result = json.loads(json_str)
all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain']
fixed_position_dict = {}
for i, chain in enumerate(global_designed_chain_list):
fixed_position_dict[chain] = fixed_list[i]
for chain in all_chain_list:
if chain not in global_designed_chain_list:
fixed_position_dict[chain] = []
my_dict[result['name']] = fixed_position_dict
else:
for json_str in json_list:
result = json.loads(json_str)
all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain']
fixed_position_dict = {}
for chain in all_chain_list:
seq_length = len(result[f'seq_chain_{chain}'])
all_residue_list = (np.arange(seq_length)+1).tolist()
if chain not in global_designed_chain_list:
fixed_position_dict[chain] = all_residue_list
else:
idx = np.argwhere(np.array(global_designed_chain_list) == chain)[0][0]
fixed_position_dict[chain] = list(set(all_residue_list)-set(fixed_list[idx]))
my_dict[result['name']] = fixed_position_dict
with open(args.output_path, 'w') as f:
f.write(json.dumps(my_dict) + '\n')
#e.g. output
#{"5TTA": {"A": [1, 2, 3, 7, 8, 9, 22, 25, 33], "B": []}, "3LIS": {"A": [], "B": []}}
if __name__ == "__main__":
argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
argparser.add_argument("--input_path", type=str, help="Path to the parsed PDBs")
argparser.add_argument("--output_path", type=str, help="Path to the output dictionary")
argparser.add_argument("--chain_list", type=str, default='', help="List of the chains that need to be fixed")
argparser.add_argument("--position_list", type=str, default='', help="Position lists, e.g. 11 12 14 18, 1 2 3 4 for first chain and the second chain")
argparser.add_argument("--specify_non_fixed", action="store_true", default=False, help="Allows specifying just residues that need to be designed (default: false)")
args = argparser.parse_args()
main(args)
|