add DrugBank database

app.py

@@ -124,12 +124,13 @@ def retrieval():
     
     with col3:
         selected_database = st.selectbox(
-                'Select database',('Lenselink', 'Davis', 'DUD-E')
+                'Select database',('Lenselink', 'Davis', 'DUD-E', 'DrugBank')
             )
         l = {
             'Lenselink': 314707,
             'Davis': 30056,
             'DUDE': 1434019,
+            'DrugBank': 10681,
         }
         if selected_database == 'DUD-E':
             selected_database = 'DUDE'
@@ -205,8 +206,15 @@ def retrieval():
         st.markdown('### Retrieval')
         
         selected_k = st.slider(f'Top-k most active drug compounds {selected_database} predicted by HyperPCM are, for k = ', 5, 20, 5, 5)
-        
-        results = pd.DataFrame({'SMILES': np.concatenate(smiles), 'Prediction': np.concatenate(preds)})
+
+        if selected_database != 'DrugBank':
+            results = pd.DataFrame({'SMILES': np.concatenate(smiles), 'Prediction': np.concatenate(preds)})
+        else: 
+            with open(os.path.join(data_path, f'{selected_database}/processed/drugbank.pickle'), 'rb') as handle:
+                lookup = pickle.load(handle)
+            drug_id = np.concatenate(smiles)
+            structure = [lookup[i] for i in drug_id]
+            results = pd.DataFrame({'SMILES': np.concatenate(structure), 'DrugBank ID': drug_id, 'Prediction': np.concatenate(preds)})
         results = results.sort_values(by='Prediction', ascending=False)
         results = results.reset_index()
         
@@ -216,7 +224,11 @@ def retrieval():
                 for i in range(int(selected_k/5)):
                     mol = Chem.MolFromSmiles(results.loc[j + 5*i, 'SMILES'])
                     mol_img = Chem.Draw.MolToImage(mol)
-                    st.image(mol_img, caption=f"{results.loc[j + 5*i, 'Prediction']:.2f}")
+                    if selected_database != 'DrugBank': 
+                        caption = f"{results.loc[j + 5*i, 'Prediction']:.2f}"
+                    else: 
+                        caption = f"{results.loc[j + 5*i, 'DrugBank ID']}:\n{results.loc[j + 5*i, 'Prediction']:.2f}"
+                    st.image(mol_img, caption=caption)
 
         st.download_button(f'Download retrieved drug compounds from the {selected_database} database.', results.head(selected_k).to_csv(index=False).encode('utf-8'), file_name='retrieved_drugs.csv')