| import argparse |
| import os |
| import sys |
|
|
| _PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) |
| while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): |
| _PARENT = os.path.dirname(_PROJECT_ROOT) |
| if _PARENT == _PROJECT_ROOT: |
| break |
| _PROJECT_ROOT = _PARENT |
| _MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") |
| _ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") |
| for _path in (_MODEL_ROOT, _PROJECT_ROOT): |
| if os.path.exists(_path) and _path not in sys.path: |
| sys.path.insert(0, _path) |
| if _ONESCIENCE_ROOT: |
| _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") |
| for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): |
| if os.path.exists(_path) and _path not in sys.path: |
| sys.path.insert(0, _path) |
|
|
| 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') |
| |
| |
| |
|
|
| 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) |
|
|
|
|
