Update UI

app.py

@@ -345,31 +345,46 @@ def Optimize_a_molecule():
     # oab_sel_container = st.empty()
     if 'checked_single' not in st.session_state:
         st.session_state.checked_single = 'NO'
-    if 'mode' not in st.session_state:
-        st.session_state.mode = 0
+    # if 'mode' not in st.session_state:
+    #     st.session_state.mode = 0
     if 'single_optimized' not in st.session_state:
         st.session_state.single_optimized = False
     if 'smiles_checked' not in st.session_state:
         st.session_state.smiles_checked = False
     # with oab_sel_container.container():
+    ls_opt = ['-','Sorafenib','Pazopanib','Sunitinib']
+    sample_mode = {
+        '-':'',
+        'Sorafenib':'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',
+        'Pazopanib':'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',
+        'Sunitinib':'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
+        }
     oam_sel_col = st.columns([3,7])
-    mode = oam_sel_col[0].selectbox("Select an example",options=['-','Sorafenib','Pazopanib','Sunitinib'],on_change=reset_oam_state,index=st.session_state.mode)
-    if mode == '-':
-        smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',max_chars=200)
-        st.session_state.mode = 0
-    elif mode == 'Sorafenib':
-        smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
-        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',max_chars=200,disabled=True)
-        st.session_state.mode = 1
-    elif mode == 'Pazopanib':
-        smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
-        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',max_chars=200,disabled=True)
-        st.session_state.mode = 2
-    elif mode == 'Sunitinib':
-        smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
-        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C',max_chars=200,disabled=True)
-        st.session_state.mode = 3
-    st.columns([4,1,4])[1].button('Check SMILES',on_click=check_single,args=[smiles])
+    with st.form('sel_smiles'):
+        mode = oam_sel_col[0].selectbox("Select an example",options=ls_opt,on_change=reset_oam_state)
+        smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',sample_mode[mode],max_chars=200,
+                                  disabled=(mode != '-'))
+        # if mode == '-':
+        #     st.session_state.smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',max_chars=200,key='opt_0')
+        #     # st.session_state.mode = 0
+        # elif mode == 'Sorafenib':
+        #     st.session_state.smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
+        #     oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',max_chars=200,disabled=True,key='opt_1')
+        #     # st.session_state.mode = 1
+        # elif mode == 'Pazopanib':
+        #     st.session_state.smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
+        #     oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',max_chars=200,disabled=True,key='opt_2')
+        #     # st.session_state.mode = 2
+        # elif mode == 'Sunitinib':
+        #     st.session_state.smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
+        #     oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C',max_chars=200,disabled=True,key='opt_3')
+        #     # st.session_state.mode = 3
+        check_single_butt = st.form_submit_button('Check SMILES',use_container_width=True)
+    # st.session_state.smiles = st.session_state['opt_'+str(ls_opt.index(st.session_state.mode))]
+    if check_single_butt:
+        st.session_state.mode = mode
+        st.session_state.smiles = smiles
+        check_single(st.session_state.smiles)
     
     if 'optim_single_butt' not in locals():
         optim_single_butt = False
@@ -392,10 +407,16 @@ def Optimize_a_molecule():
                                 set_step(0)
                                 oam_sidebar(0)
                     elif st.session_state.checked_single == 'YES':
+                        if st.session_state.mode != '-':
+                            st.markdown(f"<h4 style='color:mediumseagreen;'>Using example: <b>{st.session_state.mode}</b></h4>",unsafe_allow_html=True)
+                        else:
+                            st.markdown(f"<h4>Selected SMILES</h4>",unsafe_allow_html=True)
+                            st.code(st.session_state.smiles)
                         st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
-                        st.code(Chem.CanonSmiles(st.session_state.smiles))
+                        st.session_state.canon_smiles = Chem.CanonSmiles(st.session_state.smiles)
+                        st.code(st.session_state.canon_smiles)
                         st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
-                        mol = Chem.MolFromSmiles(st.session_state.smiles)
+                        mol = Chem.MolFromSmiles(st.session_state.canon_smiles)
                         imgByteArr = io.BytesIO()
                         MolToImage(mol,size=(400,400)).save(imgByteArr,format='PNG')
                         st.markdown("<p style='text-align: center;'>"+
@@ -422,7 +443,7 @@ def Optimize_a_molecule():
                                 optim_single_butt = st.form_submit_button("Optimize")
                     else:
                         st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
-                        st.code(st.session_state.smiles)
+                        st.code(st.session_state.canon_smiles)
                         if st.session_state.checked_single == 'NoVocab':
                             st.markdown("<p style='text-align: center; color: red;'><b>The molecule contains unavailable vocab(s). Please use another molecule.</b></p>",unsafe_allow_html=True)
                         else:
@@ -446,7 +467,7 @@ def Optimize_a_molecule():
                 st.markdown('Operation in progress. Please wait...')
                 render_animation()
                 model = load_model()
-                st.session_state.new_smiles,st.session_state.sim = optim_single(st.session_state.smiles,model,st.session_state.lr_s,st.session_state.sim_cutoff_s,st.session_state.n_iter_s)
+                st.session_state.new_smiles,st.session_state.sim = optim_single(st.session_state.canon_smiles,model,st.session_state.lr_s,st.session_state.sim_cutoff_s,st.session_state.n_iter_s)
                 st.session_state.single_optimized = True
             ani_con.empty()
                 # sidebar_con.empty()
@@ -498,10 +519,11 @@ def Optimize_a_molecule():
         else: optim_single_con.empty()
     else: optim_single_con.empty()
 
+
 def check_single(smiles):
     # render_view()
     st.session_state.smiles_selected = True
-    st.session_state.smiles = smiles
+    # st.session_state.smiles = smiles
     # check_single_con = st.empty()
     
     # optim = False
@@ -513,7 +535,7 @@ def check_single(smiles):
         if mol is None:
             st.session_state.checked_single = 'MolError'
         else:
-            st.session_state.smiles = Chem.MolToSmiles(mol)
+            st.session_state.canon_smiles = Chem.MolToSmiles(mol)
             st.session_state.logp,st.session_state.sa,st.session_state.cycle,st.session_state.pen_p = penalized_logp_standard(mol)
             moses_passed = mol_passes_filters_custom(mol)
             st.session_state.checked_single = moses_passed
@@ -610,7 +632,7 @@ def Optimize_a_batch():
         with check_batch_con.container():
             passed_num = st.session_state['df'][st.session_state['df'].checked != 'invalid'].shape[0]
             st.markdown('Number of passed SMILES: '+str(passed_num))
-            st.dataframe(st.session_state['df'],use_container_width=True)
+            st.dataframe(st.session_state['df'].style.applymap(highlight_result, subset=pd.IndexSlice[:, ['checked']]),use_container_width=True)
             if passed_num == 0:
                 st.markdown("<p style='text-align: center; color: red;'><b>The uploaded file contains no suitable SMILES string.</b></p>",unsafe_allow_html=True)
                 st.session_state.batch_left_checked = False
@@ -697,9 +719,17 @@ def Optimize_a_batch():
         optim_batch_con.empty()
         
 
-
+def process_check_single(smiles):
+    mol = Chem.MolFromSmiles(smiles)
+    if (mol is not None) and (mol_passes_filters_custom(mol) == 'YES'):
+        return Chem.MolToSmiles(mol)
+    else:
+        return 'invalid'
 def check_batch(smiles_list):
     check = []
+    # check = Parallel(n_jobs=-1,backend='loky')(
+    #     delayed(process_check_single)(smi) for smi in stqdm(smiles_list)
+    # )
     for smi in stqdm(smiles_list):
             mol = Chem.MolFromSmiles(smi)
             if (mol is not None) and (mol_passes_filters_custom(mol) == 'YES'):
@@ -713,6 +743,9 @@ def check_batch(smiles_list):
 def calc_scores(smiles_list):
     score_df = pd.concat([st.session_state.df[st.session_state.df.checked != 'invalid'].SMILES,pd.DataFrame({'Checked_SMILES':smiles_list})],axis=1)
     scores =[]
+    # scores = Parallel(n_jobs=-1,backend='loky')(
+    #     delayed(penalized_logp_standard)(Chem.MolFromSmiles(smi)) for smi in stqdm(smiles_list)
+    # )
     for smi in stqdm(smiles_list):
         logp,sa,cycle,pen_logp = penalized_logp_standard(Chem.MolFromSmiles(smi))
         scores+=[(logp,sa,cycle,pen_logp)]
@@ -723,8 +756,22 @@ def calc_scores(smiles_list):
     st.session_state.scores_calculated = True
     return score_df
 
+def process_calc_new_score(new_smiles,sim):
+    if new_smiles is None:
+        return ('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)
+    else:
+        new_mol = Chem.MolFromSmiles(new_smiles)
+        if new_mol is None:
+            return ('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)
+        else:
+            logp,sa,cycle,pen_logp = penalized_logp_standard(new_mol)
+            return (new_smiles,sim,logp,sa,cycle,pen_logp)
+
 def calc_scores_new(result):
     new_scores =[]
+    # new_scores = Parallel(n_jobs=-1,backend='loky')(
+    #     delayed(process_calc_new_score)(new_smiles,sim) for new_smiles,sim in stqdm(result)
+    # )
     for new_smiles,sim in stqdm(result):
         if new_smiles is None:
             new_scores+=[('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)]