Spaces:
Running
Running
| import json | |
| import os | |
| import tempfile | |
| import Bio.PDB | |
| import Bio.SeqUtils | |
| from Bio import pairwise2 | |
| from run_pretrained_model import run_on_folder | |
| def get_seq_based_on_template(seq: str, template_path: str, output_path: str): | |
| # get a list of all residues in template | |
| parser = Bio.PDB.PDBParser() | |
| template_structure = parser.get_structure("template", template_path) | |
| chain = template_structure[0].get_chains().__next__() | |
| template_residues = [i for i in chain.get_residues() if "CA" in i | |
| and Bio.SeqUtils.seq1(i.get_resname()) not in ("X", "", " ")] | |
| template_seq = "".join([Bio.SeqUtils.seq1(i.get_resname()) for i in template_residues]) | |
| # align the sequence to the template | |
| alignment = pairwise2.align.globalxx(seq, template_seq, one_alignment_only=True)[0] | |
| aligned_seq, aligned_template_seq = alignment.seqA, alignment.seqB | |
| # create a new pdb file with the aligned residues | |
| new_structure = Bio.PDB.Structure.Structure("new_structure") | |
| new_model = Bio.PDB.Model.Model(0) | |
| new_structure.add(new_model) | |
| new_chain = Bio.PDB.Chain.Chain("A") # Using chain ID 'A' for the output | |
| new_model.add(new_chain) | |
| template_ind = -1 | |
| seq_ind = 0 | |
| print(aligned_seq, aligned_template_seq, len(template_residues)) | |
| for seq_res, template_res in zip(aligned_seq, aligned_template_seq): | |
| if template_res != "-": | |
| template_ind += 1 | |
| if seq_res != "-": | |
| seq_ind += 1 | |
| if seq_res == "-": | |
| continue | |
| if template_res == "-": | |
| seq_res_3_letter = Bio.SeqUtils.seq3(seq_res).upper() | |
| residue = Bio.PDB.Residue.Residue((' ', seq_ind, ' '), seq_res_3_letter, '') | |
| atom = Bio.PDB.Atom.Atom("C", (0.0, 0.0, 0.0), 1.0, 1.0, ' ', "CA", 0, element="C") | |
| residue.add(atom) | |
| new_chain.add(residue) | |
| else: | |
| residue = template_residues[template_ind].copy() | |
| residue.detach_parent() | |
| residue.id = (' ', seq_ind, ' ') | |
| new_chain.add(residue) | |
| io = Bio.PDB.PDBIO() | |
| io.set_structure(new_structure) | |
| io.save(output_path) | |
| def run_on_sample_seqs(seq1: str, template1_path: str, seq2: str, template_path2: str, output_path: str, | |
| run_config_path: str): | |
| temp_dir = tempfile.TemporaryDirectory() | |
| temp_dir_path = temp_dir.name | |
| get_seq_based_on_template(seq1, template1_path, f"{temp_dir_path}/prot_r.pdb") | |
| get_seq_based_on_template(seq2, template_path2, f"{temp_dir_path}/prot_l.pdb") | |
| json_data = { | |
| "input_r_structure": f"prot_r.pdb", | |
| "input_l_structure": f"prot_l.pdb", | |
| } | |
| tmp_json_folder = f"{temp_dir_path}/jsons" | |
| os.makedirs(tmp_json_folder, exist_ok=True) | |
| json.dump(json_data, open(f"{tmp_json_folder}/input.json", "w")) | |
| tmp_output_folder = f"{temp_dir_path}/output" | |
| run_on_folder(tmp_json_folder, tmp_output_folder, run_config_path, skip_relaxation=True, | |
| long_sequence_inference=False, skip_exists=False) | |
| os.rename(tmp_output_folder + "/predictions/input_predicted_joined.pdb", output_path) | |
| print("moved output to ", output_path) | |
| temp_dir.cleanup() | |