Upload main.py

main.py

@@ -1,3 +1,4 @@
+
 import gradio as gr
 import py3Dmol
 from Bio.PDB import *
@@ -13,7 +14,6 @@ from transformMD import GNNTransformMD
 import sys
 import pytraj as pt
 import pickle 
-import nglview as nv
 
 # JavaScript functions
 resid_hover = """function(atom,viewer) {{
@@ -143,17 +143,6 @@ def get_pdbid_from_filename(filename: str):
     # Assuming the filename would be of the standard form 11GS.pdb
     return filename.split(".")[0]
 
-
-    for i in range(10):
-        data.append([i, atom_mapping[atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1], atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")],
-        atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],
-        atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")],
-        adaptability[topN_ind[i]]
-    ])
-   
-    pdb = open(pdb_file.name, "r").read()
-
-
 def predict(pdb_code, pdb_file):
     #path_to_pdb = get_pdb(pdb_code=pdb_code, filepath=pdb_file)
 
@@ -190,18 +179,12 @@ def predict(pdb_code, pdb_file):
     
     data = []
 
-    for i in range(10):
-        data.append([i, atom_mapping[atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1], atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")],
-        atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],
-        atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")],
-        adaptability[topN_ind[i]]
-    ])
-   
+
     for i in range(adaptability.shape[0]):
         data.append([i, atom_mapping[atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1], atoms_protein.iloc[i, atoms_protein.columns.get_loc("x")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("y")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("z")],adaptability[i]])
 
-    #topN = 100
-    #topN_ind = np.argsort(adaptability)[::-1][:topN]    
+    topN = 100
+    topN_ind = np.argsort(adaptability)[::-1][:topN]    
 
     pdb = open(pdb_file.name, "r").read()
 
@@ -227,7 +210,6 @@ def predict(pdb_code, pdb_file):
 
 callback = gr.CSVLogger()
 
-
 def run():
     with gr.Blocks() as demo:
         gr.Markdown("# Protein Adaptability Prediction")
@@ -254,4 +236,4 @@ def run():
 
 
 if __name__ == "__main__":
-    run()
+    run()