Update main.py

main.py

@@ -1,4 +1,5 @@
 
+
 import gradio as gr
 import py3Dmol
 from Bio.PDB import *
@@ -180,18 +181,11 @@ def predict(pdb_code, pdb_file):
     data = []
 
 
-    for i in range(10):
-        view.addSphere({'center':{'x':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")], 'y':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],'z':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")]},'radius':adaptability[topN_ind[i]]/1.5,'color':'orange','alpha':0.75})    
+    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]])
 
-        # Add lighting and shading options to the view object:
-        view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
-        view.setStyle({'sphere': {}})
-        view.addLight([0, 0, 10], [1, 1, 1], 1)  # Add directional light from the z-axis
-        view.setSpecular(0.5)  # Adjust the specular lighting effect
-        view.setAmbient(0.5)  # Adjust the ambient lighting effect
-    
-    #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()
 
@@ -200,9 +194,26 @@ def predict(pdb_code, pdb_file):
     view.addModel(pdb, "pdb")
     view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
    
-    for i in range(topN):
-        view.addSphere({'center':{'x':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")], 'y':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],'z':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")]},'radius':adaptability[topN_ind[i]]/1.5,'color':'orange','alpha':0.75})    
-
+    for i in range(10):
+        adaptability_value = adaptability[topN_ind[i]]
+        color = adaptability_color_map(adaptability_value)  # Define your own color mapping function
+        view.addSphere({
+            'center': {
+                'x': atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")],
+                'y': atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],
+                'z': atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")]
+            },
+            'radius': adaptability_value / 1.5,
+            'color': color,
+            'alpha': 0.75
+        })
+        # Add lighting and shading options to the view object:
+        view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
+        view.setStyle({'sphere': {}})
+        view.addLight([0, 0, 10], [1, 1, 1], 1)  # Add directional light from the z-axis
+        view.setSpecular(0.5)  # Adjust the specular lighting effect
+        view.setAmbient(0.5)  # Adjust the ambient lighting effect
+        
     view.zoomTo()
 
     output = view._make_html().replace("'", '"')