OneScience-Group/ProteinMPNN

1import argparse2import os3import sys4 5_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))6while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")):7    _PARENT = os.path.dirname(_PROJECT_ROOT)8    if _PARENT == _PROJECT_ROOT:9        break10    _PROJECT_ROOT = _PARENT11_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model")12_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT")13for _path in (_MODEL_ROOT, _PROJECT_ROOT):14    if os.path.exists(_path) and _path not in sys.path:15        sys.path.insert(0, _path)16if _ONESCIENCE_ROOT:17    _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src")18    for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT):19        if os.path.exists(_path) and _path not in sys.path:20            sys.path.insert(0, _path)21 22def main(args):23    import json24 25    with open(args.input_path, 'r') as json_file:26        json_list = list(json_file)27    28    global_designed_chain_list = []29    if args.chain_list != '':30        global_designed_chain_list = [str(item) for item in args.chain_list.split()]31    my_dict = {}32    for json_str in json_list:33        result = json.loads(json_str)34        all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain'] #['A','B', 'C',...]35        if len(global_designed_chain_list) > 0:36            designed_chain_list = global_designed_chain_list37        else:38            #manually specify, e.g.39            designed_chain_list = ["A"]40        fixed_chain_list = [letter for letter in all_chain_list if letter not in designed_chain_list] #fix/do not redesign these chains 41        my_dict[result['name']]= (designed_chain_list, fixed_chain_list)42    43    with open(args.output_path, 'w') as f:44        f.write(json.dumps(my_dict) + '\n')45 46 47if __name__ == "__main__":48    argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)49    argparser.add_argument("--input_path", type=str, help="Path to the parsed PDBs")50    argparser.add_argument("--output_path", type=str, help="Path to the output dictionary")51    argparser.add_argument("--chain_list", type=str, default='', help="List of the chains that need to be designed")    52 53    args = argparser.parse_args()54    main(args) 55 56# Output looks like this:57# {"5TTA": [["A"], ["B"]], "3LIS": [["A"], ["B"]]}58 59