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    import numpy as np25    with open(args.jsonl_input_path, 'r') as json_file:26        json_list = list(json_file)27    28    my_dict = {}29    for json_str in json_list:30        result = json.loads(json_str)31        all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain']32        path_to_PSSM = args.PSSM_input_path+"/"+result['name'] + ".npz"33        print(path_to_PSSM)34        pssm_input = np.load(path_to_PSSM)35        pssm_dict = {}36        for chain in all_chain_list:37            pssm_dict[chain] = {}38            pssm_dict[chain]['pssm_coef'] = pssm_input[chain+'_coef'].tolist() #[L] per position coefficient to trust PSSM; 0.0 - do not use it; 1.0 - just use PSSM only39            pssm_dict[chain]['pssm_bias'] = pssm_input[chain+'_bias'].tolist() #[L,21] probability (sums up to 1.0 over alphabet of size 21) from PSSM40            pssm_dict[chain]['pssm_log_odds'] = pssm_input[chain+'_odds'].tolist() #[L,21] log_odds ratios coming from PSSM; optional/not needed41        my_dict[result['name']] = pssm_dict42    43    #Write output to:    44    with open(args.output_path, 'w') as f:45        f.write(json.dumps(my_dict) + '\n')46    47if __name__ == "__main__":48    argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)49 50    argparser.add_argument("--PSSM_input_path", type=str, help="Path to PSSMs saved as npz files.")51    argparser.add_argument("--jsonl_input_path", type=str, help="Path where to load .jsonl dictionary of parsed pdbs.")52    argparser.add_argument("--output_path", type=str, help="Path where to save .jsonl dictionary with PSSM bias.")53 54    args = argparser.parse_args()55    main(args)56