Spaces:
Running on Zero
Running on Zero
| import os | |
| from setup.configProcessor import get_result_dir, write_ri, get_cutoff | |
| from misc.enums import LabelType | |
| def write_predictions_to_file(averaged_probabilities): | |
| is_write_ri = write_ri() | |
| result_dir = get_result_dir() | |
| prediction_folder = os.path.join(result_dir, "predictions") | |
| if not os.path.exists(prediction_folder): | |
| os.mkdir(prediction_folder) | |
| for protein_id, prediction in averaged_probabilities.items(): | |
| file_name = os.path.join(prediction_folder, f"{protein_id}.bindNode_out") | |
| with open(file_name, 'w') as fh: | |
| if is_write_ri: | |
| fh.write("Position\tMetal.RI\tMetal.Class\tNuc.RI\tNuc.Class\tSmall.RI\tSmall.Class\tAny.Class\n") | |
| else: | |
| fh.write("Position\tMetal.Prob\tMetal.Class\tNuc.Prob\tNuc.Class\tSmall.Prob\tSmall.Class\tAny.Class\n") | |
| for position, residue_prediction in enumerate(prediction): | |
| metal_probability = residue_prediction[LabelType.METAL.value] | |
| small_probability = residue_prediction[LabelType.SMALL.value] | |
| nuclear_probability = residue_prediction[LabelType.NUCLEAR.value] | |
| metal_prediction = metal_probability >= get_cutoff() | |
| small_prediction = small_probability >= get_cutoff() | |
| nuclear_prediction = nuclear_probability >= get_cutoff() | |
| any_prediction = any([metal_prediction, nuclear_prediction, small_prediction]) | |
| fh.write(f"{position + 1}\t") | |
| fh.write(f"{compute_ri(metal_probability):.9f}\t" if is_write_ri else f"{metal_probability:.9f}\t") | |
| fh.write("b\t" if metal_prediction else "nb\t") | |
| fh.write(f"{compute_ri(nuclear_probability):.9f}\t" if is_write_ri else f"{nuclear_probability:.9f}\t") | |
| fh.write("b\t" if nuclear_prediction else "nb\t") | |
| fh.write(f"{compute_ri(small_probability):.9f}\t" if is_write_ri else f"{small_probability:.9f}\t") | |
| fh.write("b\t" if small_prediction else "nb\t") | |
| fh.write("b\n" if any_prediction else "nb\n") | |
| def compute_ri(prediction_probability): | |
| if prediction_probability < 0.5: | |
| reliability_index = round((0.5 - prediction_probability) * 9 / 0.5) | |
| else: | |
| reliability_index = round((prediction_probability - 0.5) * 9 / 0.5) | |
| return reliability_index | |