Refactor app.py to improve user experience by updating SMILES input display, enhancing error handling for invalid inputs, and refining the layout for displaying DCM layers. Adjusted slider label for clarity and ensured consistent rendering dimensions for 3D models.

app.py

@@ -15,14 +15,14 @@ import streamlit.components.v1 as components
 
 import streamlit as st
 from streamlit_ketcher import st_ketcher
-
+WIDTH = 500
 st.title("DCM-Net")
 smiles = "OC(=O)[C@@H]1CCCN1"
 # Create a molecule editor
 smiles = st_ketcher(smiles)
 
 # Output the SMILES string of the drawn molecule
-st.write(f"SMILES: {smiles}")
+st.write(f"SMILES input: {smiles}")
 
 
 # color charges by magnitude
@@ -83,58 +83,60 @@ def render_3d(atoms, dcm_charges=None, dcmcolors=None):
     return view
 
 # create a slider for n_dcm
-n_dcm = st.slider("Number of DCM layers", min_value=1, max_value=4, value=1)
+n_dcm = st.slider("Number of Distributed Charges per Atom", min_value=1, max_value=4, value=4)
 
 # create a button to run the DCM-Net
 if st.button("Run DCM-Net"):
-    results = run_dcm(smiles, n_dcm)
-    print(results)
-
-    st.subheader("Model Output")
-    # ase object
-    ase_atoms = results["atoms"]
-    
-    # Display the 3D model
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        showmol(render_3d(ase_atoms), height=300, width=1000)
-    with col2:
-        st.image(results["smiles_image"])
-
-    st.subheader("DCM-1")
-    dcmcolors1 = color_charges(results["mono_dc1"], 1, ase_atoms)
-    # two columns showing combined and just the dcmol
-    col1, col2 = st.columns(2)
-    with col1:
-        showmol(render_3d(ase_atoms, results["dcmol"]), height=300, width=1000)
-    with col2:
-        showmol(render_3d(None, results["dcmol"], dcmcolors1), height=300, width=1000)
-
-    if n_dcm >= 2:
-        st.subheader("DCM-2")
-        dcmcolors2 = color_charges(results["mono_dc2"], 2, ase_atoms)
+    # bail if smiles is not valid
+    if smiles is None or smiles == "":
+        st.error("Please enter a valid SMILES string")
+        st.stop()
+
+    # start a spinner
+    with st.spinner("Running DCM-Net..."):
+        results = run_dcm(smiles, n_dcm)
+        print(results)
+        ase_atoms = results["atoms"]
         col1, col2 = st.columns(2)
         with col1:
-            showmol(render_3d(ase_atoms, results["dcmol2"]), height=300, width=1000)
+            showmol(render_3d(ase_atoms), height=300, width=WIDTH)
         with col2:
-            showmol(render_3d(None, results["dcmol2"], dcmcolors2), height=300, width=1000)
-
-    if n_dcm >= 3:
-        st.subheader("DCM-3")
-        dcmcolors3 = color_charges(results["mono_dc3"], 3, ase_atoms)
-        col1, col2 = st.columns(2)
-        with col1:
-            showmol(render_3d(ase_atoms, results["dcmol3"]), height=300, width=1000)
-        with col2:
-            showmol(render_3d(None, results["dcmol3"], dcmcolors3), height=300, width=1000)
-
-    if n_dcm >= 4:
-        st.subheader("DCM-4")
-        dcmcolors4 = color_charges(results["mono_dc4"], 4, ase_atoms)
+            st.image(results["smiles_image"])
+    st.subheader("Model Output")
+    if n_dcm >= 1:
+        st.markdown("### DCM-1")
+        dcmcolors1 = color_charges(results["mono_dc1"], 1, ase_atoms)
+        # two columns showing combined and just the dcmol
         col1, col2 = st.columns(2)
         with col1:
-            showmol(render_3d(ase_atoms, results["dcmol4"]), height=300, width=1000)
+            showmol(render_3d(ase_atoms, results["dcmol"]), height=300, width=WIDTH)
         with col2:
-            showmol(render_3d(None, results["dcmol4"], dcmcolors4), height=300, width=1000)
+            showmol(render_3d(None, results["dcmol"], dcmcolors1), height=300, width=WIDTH)
+
+        if n_dcm >= 2:
+            st.markdown("### DCM-2")
+            dcmcolors2 = color_charges(results["mono_dc2"], 2, ase_atoms)
+            col1, col2 = st.columns(2)
+            with col1:
+                showmol(render_3d(ase_atoms, results["dcmol2"]), height=300, width=WIDTH)
+            with col2:
+                showmol(render_3d(None, results["dcmol2"], dcmcolors2), height=300, width=WIDTH)
+
+        if n_dcm >= 3:
+            st.markdown("### DCM-3")
+            dcmcolors3 = color_charges(results["mono_dc3"], 3, ase_atoms)
+            col1, col2 = st.columns(2)
+            with col1:
+                showmol(render_3d(ase_atoms, results["dcmol3"]), height=300, width=WIDTH)
+            with col2:
+                showmol(render_3d(None, results["dcmol3"], dcmcolors3), height=300, width=WIDTH)
+
+        if n_dcm >= 4:
+            st.markdown("### DCM-4")
+            dcmcolors4 = color_charges(results["mono_dc4"], 4, ase_atoms)
+            col1, col2 = st.columns(2)
+            with col1:
+                showmol(render_3d(ase_atoms, results["dcmol4"]), height=300, width=WIDTH)
+            with col2:
+                showmol(render_3d(None, results["dcmol4"], dcmcolors4), height=300, width=WIDTH)