added visualisation of generated molecules and restart button

app.py

@@ -6,8 +6,20 @@ import pygad
 
 from VQGAE.models import VQGAE, OrderingNetwork
 from CGRtools.containers import QueryContainer
+from CGRtools.utils import grid_depict
 from VQGAE.utils import frag_counts_to_inds, restore_order, decode_molecules
 
+import base64
+from streamlit.components.v1 import html
+
+
+def render_svg(svg_string):
+    """Renders the given svg string."""
+    c = st.container()
+    with c:
+        html(svg_string)
+
+
 # define groups to filter
 allene = QueryContainer()
 allene.add_atom("C")
@@ -121,7 +133,7 @@ def load_data(batch_size):
 st.title('Inverse QSAR of Tubulin with VQGAE')
 
 with st.sidebar:
-    with st.form("my_form"):
+    with st.form("ga_options"):
         num_generations = st.slider(
             'Number of generations for GA',
             min_value=3,
@@ -190,7 +202,7 @@ with st.sidebar:
         )
         # 2/3 of num_parents_mating
         use_ordering_score = st.toggle('Use ordering score', value=True)
-        batch_size = int(st.number_input("Random seed", value=200, placeholder="Type a number..."))
+        batch_size = int(st.number_input("Batch size", value=200, placeholder="Type a number..."))
         random_seed = int(st.number_input("Random seed", value=42, placeholder="Type a number..."))
         submit = st.form_submit_button('Start optimisation')
 
@@ -291,31 +303,53 @@ if submit:
         decoding_bar.progress(decoding_step / total_decoding_steps, text=decoding_progress_text)
     gen_stats = pd.DataFrame(results)
     decoding_bar.empty()
-    full_stats = pd.concat([gen_stats, chosen_gen[["similarity_score", "rf_score"]].reset_index(), ], axis=1, ignore_index=False)
-
-    st.dataframe(full_stats)
-
+    full_stats = pd.concat([gen_stats, chosen_gen.reset_index()[["similarity_score", "rf_score"]]], axis=1)
+    valid_gen_stats = full_stats[full_stats.validity == 1]
+
+    valid_gen_mols = []
+    for i, record in zip(list(valid_gen_stats.index), valid_gen_stats.to_dict("records")):
+        valid_gen_mols.append(gen_molecules[i])
+
+    filtered_gen_mols = []
+    filtered_indices = []
+    for mol_i, mol in enumerate(valid_gen_mols):
+        is_frag = allene < mol or peroxide_charge < mol or peroxide < mol
+        is_ring = False
+        for ring in mol.sssr:
+            if len(ring) > 8 or len(ring) < 4:
+                is_ring = True
+                break
+        if not is_frag and not is_ring:
+            filtered_gen_mols.append(mol)
+            filtered_indices.append(mol_i)
+
+    filtered_gen_stats = valid_gen_stats.iloc[filtered_indices]
+
+    st.subheader('Generation results', divider='rainbow')
+    st.dataframe(filtered_gen_stats)
     st.download_button(
         label="Download results as CSV",
-        data=convert_df(full_stats),
-        file_name='vqgae_tubulin_inhibitors.csv',
+        data=convert_df(filtered_gen_stats),
+        file_name='vqgae_tubulin_inhibitors_valid.csv',
         mime='text/csv',
     )
 
-    # valid_gen_stats = full_stats[full_stats.valid == 1]
-    #
-    # valid_gen_mols = []
-    # for i, record in zip(list(valid_gen_stats.index), valid_gen_stats.to_dict("records")):
-    #     mol = gen_molecules[i]
-    #     valid_gen_mols.append(mol)
-    #
-    # filtered_gen_mols = []
-    # for mol in valid_gen_mols:
-    #     is_frag = allene < mol or peroxide_charge < mol or peroxide < mol
-    #     is_macro = False
-    #     for ring in mol.sssr:
-    #         if len(ring) > 8 or len(ring) < 4:
-    #             is_macro = True
-    #             break
-    #     if not is_frag and not is_macro:
-    #         filtered_gen_mols.append(mol)
+    st.subheader('Examples of generated molecules')
+    examples_inds = sorted(filtered_gen_stats.sort_values(by=["rf_score"], ascending=False).index[:9])
+    examples = [filtered_gen_mols[i] for i in examples_inds]
+    svg = grid_depict(examples, 2)
+    render_svg(svg)
+
+    show_full_stats = st.checkbox('Show full stats')
+    if show_full_stats:
+        st.dataframe(full_stats)
+
+        st.download_button(
+            label="Download full results as CSV",
+            data=convert_df(full_stats),
+            file_name='vqgae_tubulin_inhibitors_full.csv',
+            mime='text/csv',
+        )
+
+    if st.button("Restart"):
+        st.rerun()