Update main.py

main.py

@@ -180,11 +180,18 @@ def predict(pdb_code, pdb_file):
     data = []
 
 
-    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]])
+    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})    
 
-    topN = 100
-    topN_ind = np.argsort(adaptability)[::-1][:topN]    
+        # 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]    
 
     pdb = open(pdb_file.name, "r").read()