def clean_protein(prot: protein.Protein, checks: bool = True):
"""Adds missing atoms to Protein instance.
Args:
prot: A `protein.Protein` instance.
checks: A `bool` specifying whether to add additional checks to the cleaning
process.
Returns:
pdb_string: A string of the cleaned protein.
"""
_check_atom_mask_is_ideal(prot)
# Clean pdb.
prot_pdb_string = protein.to_pdb(prot)
pdb_file = io.StringIO(prot_pdb_string)
alterations_info = {}
fixed_pdb = cleanup.fix_pdb(pdb_file, alterations_info)
fixed_pdb_file = io.StringIO(fixed_pdb)
pdb_structure = PdbStructure(fixed_pdb_file)
cleanup.clean_structure(pdb_structure, alterations_info)
logger.info("alterations info: %s", alterations_info)
# Write pdb file of cleaned structure.
as_file = openmm_app.PDBFile(pdb_structure)
pdb_string = _get_pdb_string(as_file.getTopology(), as_file.getPositions())
if checks:
_check_cleaned_atoms(pdb_string, prot_pdb_string)
return pdb_string
def _result_to_pdb_file(self, res_dict, pdb_file):
feat = {
'aatype': res_dict['aatype'],
'residue_index': res_dict['residue_index'],
}
result = {
"structure_module": {
"final_atom_mask": res_dict["final_atom_mask"],
"final_atom_positions": res_dict["final_atom_positions"],
}
}
pdb_str = protein.to_pdb(protein.from_prediction(feat, result))
open(pdb_file, 'w').write(pdb_str)
def test_to_pdb(self):
with open(
os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
'2rbg.pdb')) as f:
pdb_string = f.read()
prot = protein.from_pdb_string(pdb_string, chain_id='A')
pdb_string_reconstr = protein.to_pdb(prot)
prot_reconstr = protein.from_pdb_string(pdb_string_reconstr)
np.testing.assert_array_equal(prot_reconstr.aatype, prot.aatype)
np.testing.assert_array_almost_equal(prot_reconstr.atom_positions,
prot.atom_positions)
np.testing.assert_array_almost_equal(prot_reconstr.atom_mask,
prot.atom_mask)
np.testing.assert_array_equal(prot_reconstr.residue_index,
prot.residue_index)
np.testing.assert_array_almost_equal(prot_reconstr.b_factors,
prot.b_factors)
def generate_unrelaxed_pdb(aatype,
residue_index,
model_output,
pdb_path,
b_factors=None):
fold_output = model_output['structure_module']
if b_factors is None:
b_factors = np.zeros_like(fold_output['final_atom_mask'])
prot = protein.Protein(aatype=aatype,
atom_positions=fold_output['final_atom_positions'],
atom_mask=fold_output['final_atom_mask'],
residue_index=residue_index + 1,
b_factors=b_factors)
with open(pdb_path, 'w') as f:
f.write(protein.to_pdb(prot))
return prot