wangjin2000/ESM2PPI
0
1import py3Dmol2import requests3from pathlib import Path4 5import Bio.PDB6 7def display_pdb_by_pdb(pdb_string,pdb_filename,color_map="pLDDT"):8 # function to display pdb in py3dmol9 # ref: https://huggingface.co/spaces/AIGE/A_B10 11 #get pdb column values12 p = Bio.PDB.PDBParser()13 structure = p.get_structure('myStructureName', pdb_filename)14 ids = [a.get_id() for a in structure.get_atoms()]15 pLDDTs = [a.get_bfactor() for a in structure.get_atoms()]16 17 plddt_bands = ['#FF7D45','#FFDB13','#65CBF3','#0053D6']18 style = "cartoon"19 20 #plot 3D structure21 view = py3Dmol.view(width=500, height=500)22 view.setBackgroundColor('black');23 view.addModel(pdb_string, "pdb")24 25 if color_map == "spectrum":26 view.setStyle({'cartoon': {'color': 'spectrum'}})27 else:28 for i,plddt in enumerate(pLDDTs):29 if plddt >0.90:30 color = plddt_bands[3]31 elif plddt >0.70:32 color = plddt_bands[2]33 elif plddt >0.50:34 color = plddt_bands[1]35 else:36 color = plddt_bands[0]37 ''' 38 if style == "cartoon" and show_sidechains == True:39 view.setStyle({'model': -1, 'serial': i+1}, {style: {'color': color},"stick":{}})40 else:41 view.setStyle({'model': -1, 'serial': i+1}, {style: {'color': color}})42 '''43 view.setStyle({'model': -1, 'serial': i+1}, {style:{'color': color}})44 45 view.zoomTo()46 output = view._make_html().replace("'", '"')47 x = f"""<!DOCTYPE html><html></center> {output} </center></html>""" # do not use ' in this input48 49 return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;50 display-capture; encrypted-media;" sandbox="allow-modals allow-forms51 allow-scripts allow-same-origin allow-popups52 allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""53 allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""54 55def get_pdb(sequence):56 retries = 057 pdb_str = None58 url = "https://api.esmatlas.com/foldSequence/v1/pdb/"59 while retries < 3 and pdb_str is None:60 response = requests.post(url, data=sequence, verify=False)61 pdb_str = response.text62 if pdb_str == "INTERNAL SERVER ERROR":63 retries += 164 time.sleep(0.1)65 pdb = None #pdb = str = None66 67 #save a pdb format file68 name = sequence[:3] + sequence[-3:] #combine the firt and last 3 AAs of sequence as a filename.69 outpath = (70 Path.cwd() / f"PDB-{name}.pdb")71 with open(outpath.name, "w") as f:72 f.write(pdb_str)73 outpath_str = str(outpath)74 75 return pdb_str, outpath_str76 77def plot_struc(sequence):78 79 pdb_string, pdb_filename = get_pdb(sequence) 80 81 html_view = display_pdb_by_pdb(pdb_string, pdb_filename, color_map = "spectrum")82 83 return pdb_filename, html_view