Integrate the plddt-uncertainty from AlphaFold into the generated PDBs, "abusing" the b-factor field

alphafold_integrate_uncertainty.py

#!/usr/bin/env python
# coding: utf-8
'''
Simply run this script in the output folder of an AlphaFold-run. For more infos on the pLDDT, refer to the original AlphaFold publication (https://www.nature.com/articles/s41586-021-03819-2)

https://github.com/deepmind/alphafold/issues/8
'''


import json
import pickle

from Bio.PDB import PDBIO, PDBParser

with open('ranking_debug.json') as f:
    order = json.load(f)['order']

for i in range(5):
    with open(f'result_{order[i]}.pkl', 'rb') as f:
        plddt = pickle.load(f)['plddt']
    for model_fn in [f'ranked_{i}.pdb', f'relaxed_{order[i]}.pdb', f'unrelaxed_{order[i]}.pdb']:
        structure = PDBParser().get_structure('model', model_fn)
        assert len(list(structure.get_residues())) == len(plddt)
        for res_i, res in enumerate(structure.get_residues()):
            for atom in res.get_atoms():
                atom.set_bfactor(plddt[res_i])
        # now save model
        io = PDBIO()
        io.set_structure(structure)
        io.save(model_fn.replace('.pdb', '_plddt.pdb'))