ProteinMPNN-main/helper_scripts/make_fixed_positions_dict.py

import argparse

def m_f_p_d(input_path, output_path, chain_list, position_list, specify_non_fixed):
    import glob
    import random
    import numpy as np
    import json
    import itertools
    
    with open(input_path.name, 'r') as json_file:
        json_list = list(json_file)
    
    fixed_list = [[int(item) for item in one.split()] for one in position_list.split(",")]
    global_designed_chain_list = [str(item) for item in chain_list.split()]
    my_dict = {}
    
    if not 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(output_path, 'w') as f:
        f.write(json.dumps(my_dict) + '\n')
    return output_path
    
    #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)