|
|
| from alphafold.common import protein, residue_constants
|
| from alphafold.relax import relax
|
| from pathlib import Path
|
| import sys
|
| import os
|
| import time
|
| import Bio.PDB
|
|
|
|
|
| def relax_pdb_path(pdb_path, output_path, use_gpu=False):
|
| pdb_lines = Path(pdb_path).read_text()
|
| pdb_obj = protein.from_pdb_string(pdb_lines)
|
|
|
| amber_relaxer = relax.AmberRelaxation(
|
| max_iterations=0,
|
| tolerance=2.39,
|
| stiffness=10.0,
|
| exclude_residues=[],
|
| max_outer_iterations=3,
|
| use_gpu=use_gpu)
|
|
|
| start_time = time.time()
|
| print("Relaxing", pdb_path, flush=True)
|
| relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=pdb_obj)
|
| open(output_path, "w").write(relaxed_pdb_lines)
|
| print("Relax took", time.time() - start_time, flush=True)
|
|
|
| return relaxed_pdb_lines
|
|
|
|
|
| def relax_folder(input_folder: str, output_folder: str, use_gpu=False):
|
| os.makedirs(output_folder, exist_ok=True)
|
| input_names = [i for i in os.listdir(input_folder) if i.endswith(".pdb")]
|
|
|
| input_names.sort(key=lambda x: os.path.getsize(os.path.join(input_folder, x)))
|
| print("Relaxing", len(input_names), "pdbs", input_names, flush=True)
|
|
|
| for pdb_name in input_names:
|
| pdb_path = os.path.join(input_folder, pdb_name)
|
| output_path = os.path.join(output_folder, pdb_name[:-4] + "_relaxed.pdb")
|
|
|
|
|
|
|
|
|
| if os.path.exists(output_path):
|
| print(f"--fixing chain ids {pdb_name}", flush=True)
|
| fix_chain_ids(pdb_path, output_path, output_path)
|
|
|
|
|
| def fix_chain_ids(original_pdb_path: str, relaxed_pdb_path: str, output_path: str):
|
|
|
| pdb_parser = Bio.PDB.PDBParser(QUIET=True)
|
| original_pdb_model = next(iter(pdb_parser.get_structure("o_pdb", original_pdb_path)))
|
| original_chain_ids = [c.id for c in original_pdb_model.get_chains()]
|
|
|
| relaxed_pdb_struct = pdb_parser.get_structure("r_pdb", relaxed_pdb_path)
|
| relaxed_pdb_model = next(iter(relaxed_pdb_struct))
|
| relaxed_chain_ids = [c.id for c in relaxed_pdb_model.get_chains()]
|
|
|
| PDB_CHAIN_IDS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
|
| available_chain_ids = [c for c in PDB_CHAIN_IDS if (c not in relaxed_chain_ids and c not in original_chain_ids)]
|
| assert len(original_chain_ids) == len(relaxed_chain_ids)
|
| chain_id_to_chain = {c.id: c for c in relaxed_pdb_model.get_chains()}
|
|
|
| for original_name, new_name in zip(original_chain_ids, relaxed_chain_ids):
|
| if original_name in relaxed_pdb_model:
|
| tmp_name = available_chain_ids.pop()
|
| relaxed_pdb_model[original_name].id = tmp_name
|
| print("TempRenaming", original_name, "to", tmp_name, flush=True)
|
| chain_id_to_chain[new_name].id = original_name
|
| print("Renaming", new_name, "to", original_name, flush=True)
|
|
|
| io = Bio.PDB.PDBIO()
|
| io.set_structure(relaxed_pdb_struct)
|
| io.save(output_path)
|
|
|
|
|
| if __name__ == "__main__":
|
| assert len(sys.argv) == 3, "Usage: python af_relax.py <input_pdb_path> <output_pdb_path>"
|
|
|
| path1, path2 = os.path.abspath(sys.argv[1]), os.path.abspath(sys.argv[2])
|
| if os.path.isdir(path1):
|
| relax_folder(path1, path2, True)
|
| else:
|
| relax_pdb_path(path1, path2, True)
|
|
|
