Update plot_pdb.py

plot_pdb.py

@@ -4,18 +4,20 @@ from pathlib import Path
 
 import Bio.PDB
 
-def display_pdb_by_pdb(pdb):
+def display_pdb_by_pdb(pdb_string,pdb_filename):
     # function to display pdb in py3dmol
     # ref: https://huggingface.co/spaces/AIGE/A_B
-    
+
+    #get pdb column values
     p = Bio.PDB.PDBParser()
-    structure = p.get_structure('myStructureName', pdb)
+    structure = p.get_structure('myStructureName', pdb_filename)
     ids = [a.get_id() for a in structure.get_atoms()]
     print("structure ids:", ids)
-    
+
+    #plot 3D structure
     view = py3Dmol.view(width=500, height=500)
     view.setBackgroundColor('black');
-    view.addModel(pdb, "pdb")
+    view.addModel(pdb_string, "pdb")
     view.setStyle({'cartoon': {'color': 'spectrum'}})
     view.zoomTo()
     output = view._make_html().replace("'", '"')
@@ -38,12 +40,22 @@ def get_pdb(sequence):
             retries += 1
             time.sleep(0.1)
             pdb = str = None
-    return pdb_str
+
+    #save a pdb format file
+    name = sequence[:3] + sequence[-3:]  #combine the firt and last 3 AAs of sequence as a filename.
+    outpath = (
+        Path.cwd() / f"PDB-{name}.pdb")
+    with open(outpath.name, "w") as f:
+        f.write(pdb_string)
+    outpath_str = str(outpath)
+    
+    return pdb_str, outpath_str
 
 def plot_struc(sequence):
 
-    pdb_string = get_pdb(sequence) 
-    
+    pdb_string, pdb_filename = get_pdb(sequence) 
+
+    '''
     name = sequence[:3] + sequence[-3:]  #combine the firt and last 3 AAs of sequence as a filename.
 
     outpath = (
@@ -53,7 +65,8 @@ def plot_struc(sequence):
         f.write(pdb_string)
 
     outpath_str = str(outpath)
-
-    html_view = display_pdb_by_pdb(pdb_string)
+    '''
+    
+    html_view = display_pdb_by_pdb(pdb_string, pdb_filename)
 
     return outpath_str, html_view