| import os
|
| import shutil
|
| import sys
|
|
|
| import numpy as np
|
| from typing import List, Optional
|
| from math import cos, sin
|
|
|
| import Bio
|
| import Bio.PDB, Bio.SeqIO, Bio.SeqUtils
|
| import Bio.PDB.Residue
|
|
|
|
|
| def read_model_path(pdb_path: str):
|
| if pdb_path.endswith(".cif"):
|
| return Bio.PDB.MMCIFParser().get_structure("s_cif", pdb_path)
|
| return Bio.PDB.PDBParser(QUIET=True).get_structure("s_pdb", pdb_path)
|
|
|
|
|
| def _merge_models(model_path1: str, model_path2: str, output_path: str, output_cif: bool = False):
|
|
|
| model_struct1 = read_model_path(model_path1)
|
| model_struct2 = read_model_path(model_path2)
|
|
|
| model_1 = next(iter(model_struct1))
|
| model_2 = next(iter(model_struct2))
|
|
|
| chains_1 = {i.get_id(): i for i in model_1.get_chains()}
|
| chains_2 = {i.get_id(): i for i in model_2.get_chains()}
|
|
|
| for chain in chains_2.values():
|
| if chain.id in chains_1:
|
| res_to_move = [i for i in chain.get_residues()]
|
| chains_1[chain.id].child_list += res_to_move
|
| chains_1[chain.id].child_list.sort(key=lambda x: x.id[1])
|
| for res in res_to_move:
|
| chain.detach_child(res.id)
|
| res.detach_parent()
|
| res.parent = chains_1[chain.id]
|
|
|
| chains_1[chain.id].child_list += chain.child_list
|
| else:
|
| model_1.child_list.append(chain)
|
| chain.parent = model_1
|
|
|
| if output_cif:
|
| io = Bio.PDB.MMCIFIO()
|
| else:
|
| io = Bio.PDB.PDBIO()
|
| io.set_structure(model_struct1)
|
| io.save(output_path)
|
|
|
|
|
| def _rotate_atom(coord, euler_rotation_tuple):
|
|
|
| x, y, z = euler_rotation_tuple
|
| cx, cy, cz = cos(x), cos(y), cos(z)
|
| sx, sy, sz = sin(x), sin(y), sin(z)
|
| v0 = np.array([cz*cy, -sy*sx*cz - sz*cx, -sy*cx*cz + sz*sx])
|
| v1 = np.array([sz*cy, -sy*sx*sz + cx*cz, -sy*cx*sz - sx*cz])
|
| v2 = np.array([sy, cy*sx, cy*cx])
|
|
|
| return np.array([
|
| v0[0] * coord[0] + v0[1] * coord[1] + v0[2] * coord[2],
|
| v1[0] * coord[0] + v1[1] * coord[1] + v1[2] * coord[2],
|
| v2[0] * coord[0] + v2[1] * coord[1] + v2[2] * coord[2],
|
| ])
|
|
|
|
|
| def apply_transform(pdb_path: str, output_path: str, transform_numbers: List[float]):
|
| assert len(transform_numbers) == 6
|
|
|
| pdb_parser = Bio.PDB.PDBParser(QUIET=True)
|
| pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
|
| assert len(list(pdb_struct)) == 1, f"Too many models! {pdb_path}"
|
|
|
| rotation_tuple, translation_tuple = transform_numbers[:3], transform_numbers[3:]
|
| for atom in next(iter(pdb_struct)).get_atoms():
|
| coord = atom.get_coord()
|
| new_coord = _rotate_atom(coord, rotation_tuple)
|
| new_coord += np.array(translation_tuple)
|
| atom.set_coord(new_coord)
|
|
|
| io = Bio.PDB.PDBIO()
|
| io.set_structure(pdb_struct)
|
| io.save(output_path)
|
|
|
|
|
| def create_transformation_pdb(assembly_path: str, transforms_str: str, output_path: str, output_cif: bool = False):
|
|
|
|
|
| tmp_pdb_path = os.path.join(assembly_path, "tmp_comb_trans.pdb")
|
| chains_path = os.path.join(assembly_path, "chain.list")
|
| if not os.path.exists(chains_path):
|
| chains_path = os.path.join(assembly_path, "chains.txt")
|
| subunit_filenames = [i.split(" ")[0] for i in open(chains_path, "r").read().split("\n") if len(i) > 0]
|
|
|
| if os.path.exists(output_path):
|
| os.remove(output_path)
|
| for transform_str in transforms_str.split(","):
|
| chain_ind = int(transform_str.split("(")[0])
|
| subunit_path = os.path.join(assembly_path, subunit_filenames[chain_ind])
|
| transform_numbers_as_str = transform_str.split("(")[1][:-1]
|
| transform_numbers = list(map(float, transform_numbers_as_str.split(" ")))
|
|
|
| apply_transform(subunit_path, tmp_pdb_path, transform_numbers)
|
| if os.path.exists(output_path):
|
| _merge_models(output_path, tmp_pdb_path, output_path, output_cif=output_cif)
|
| else:
|
| if output_cif:
|
| io = Bio.PDB.MMCIFIO()
|
| io.set_structure(read_model_path(tmp_pdb_path))
|
| io.save(output_path)
|
| else:
|
| shutil.copy(tmp_pdb_path, output_path)
|
| if output_cif:
|
|
|
| cif_lines = open(output_path, "r").read().split("\n")
|
| cif_lines = [i + " " if i.startswith("_atom_site") else i for i in cif_lines]
|
| open(output_path, "w").write("\n".join(cif_lines) + "\n")
|
| os.remove(tmp_pdb_path)
|
|
|
|
|
| def create_complexes(result_path: str, first_result: Optional[int] = None, last_result: Optional[int] = None,
|
| output_folder: Optional[str] = None, output_cif: bool = False) \
|
| -> List[str]:
|
| output_folder = output_folder or os.path.dirname(result_path)
|
| os.makedirs(output_folder, exist_ok=True)
|
| assembly_path = os.path.dirname(result_path)
|
| transforms_strs = [i[i.index("[") + 1:i.index("]")] for i in open(result_path, "r").read().split("\n") if i]
|
|
|
| if first_result is None:
|
| first_result = 0
|
| if last_result is None or last_result > len(transforms_strs):
|
| last_result = len(transforms_strs)
|
|
|
| output_files = []
|
| for i in range(first_result, last_result):
|
| file_ext = ".cif" if output_cif else ".pdb"
|
| output_path = os.path.join(output_folder,
|
| os.path.basename(result_path).split(".")[0] + "_" + str(i) + file_ext)
|
| create_transformation_pdb(assembly_path, transforms_strs[i], output_path=output_path, output_cif=output_cif)
|
| output_files.append(output_path)
|
| return output_files
|
|
|
|
|
| if __name__ == '__main__':
|
| if len(sys.argv) != 4:
|
| raise Exception("Usage: <script> result_path first_result_num last_result_num")
|
| _result_path = os.path.abspath(sys.argv[1])
|
| _first_result = int(sys.argv[2]) - 1
|
| _last_result = int(sys.argv[3])
|
|
|
| create_complexes(_result_path, _first_result, _last_result)
|
|
|
