Update app.py

app.py

@@ -18,7 +18,7 @@ from transformers import AutoTokenizer, AutoModel, pipeline as hf_pipeline
 # =================================================================
 # PART 0: THE BRIDGE (Automatic Brain Setup)
 # =================================================================
-st.set_page_config(page_title="PFAS Discovery AI", layout="wide")
+st.set_page_config(page_title="PFAS Discovery AI", layout="wide", initial_sidebar_state="expanded")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_NAME = "JuIm/SMILES_BERT"
 
@@ -85,18 +85,16 @@ def sanity_check_class(smiles, predicted_class):
 def mutate_smart(s):
     """
     Aggressive 'Safe-by-Design' Mutation
-    Prioritizes breaking chains and adding degradation handles.
     """
     try:
         chars = list(sf.split_selfies(sf.encoder(s)))
         # Action 1: Insert Polar Handles (High Probability)
         if random.random() < 0.9: 
             insert_idx = random.randint(0, len(chars))
-            # [O]=Ether, [N]=Amine, [C][=O]=Carbonyl (Degradable)
             atom = random.choice(["[O]", "[N]", "[C][=O]", "[C][O]"])
             chars.insert(insert_idx, atom)
         
-        # Action 2: Cap Ends to increase Solubility
+        # Action 2: Cap Ends
         if random.random() < 0.6:
             chars.append(random.choice(["[O]", "[N]", "[C][=O][O]"])) 
             
@@ -121,11 +119,13 @@ if clf is None:
     st.stop()
 
 # =================================================================
-# PART 2: THE UI (With 4-Graph Dashboard)
+# PART 2: THE UI (Aesthetic Upgrade)
 # =================================================================
 st.title("🧪 End-to-End PFAS Discovery AI")
-st.markdown("Powered by **Evolutionary Optimization** & **Advanced Visualization**")
+st.markdown("### 🤖 Powered by Evolutionary Optimization & Deep Learning")
+st.markdown("---")
 
+# Sidebar
 st.sidebar.header("1. Input Data")
 input_type = st.sidebar.radio("Source:", ["Single Molecule", "Batch CSV"])
 
@@ -143,6 +143,15 @@ else:
 st.sidebar.header("2. Pipeline Mode")
 mode = st.sidebar.selectbox("Mode:", ["Screening (Analyze)", "Discovery (Optimize)"])
 
+# Define Custom Color Palette (Professional & Safe)
+COLOR_MAP = {
+    "Non-PFAS": "#2ecc71",       # Emerald Green
+    "PFCA": "#e74c3c",           # Alizarin Red
+    "PFSA": "#9b59b6",           # Amethyst Purple
+    "General PFAS": "#f39c12",   # Orange
+    "Invalid": "#95a5a6"         # Grey
+}
+
 if st.sidebar.button("🚀 Run Pipeline") and data:
     st.info(f"Running **{mode}** on {len(data)} molecules...")
     
@@ -214,13 +223,12 @@ if st.sidebar.button("🚀 Run Pipeline") and data:
             })
 
     # ------------------------------------------------------------------
-    # VISUALIZATION DASHBOARD
+    # AESTHETIC DASHBOARD
     # ------------------------------------------------------------------
     res_df = pd.DataFrame(results)
     
-    # 1. RESULTS TABLE
-    st.subheader("📊 Data Table")
-    st.dataframe(res_df)
+    st.markdown("### 📊 Analysis Results")
+    st.dataframe(res_df.style.highlight_min(axis=0, subset=['Bioaccumulation'], color='#d4edda'))
     st.download_button("Download CSV", res_df.to_csv(index=False).encode('utf-8'), "results.csv", "text/csv")
 
     st.markdown("---")
@@ -228,53 +236,65 @@ if st.sidebar.button("🚀 Run Pipeline") and data:
     
     col1, col2 = st.columns(2)
     
-    color_map = {"Non-PFAS": "green", "PFCA": "red", "PFSA": "purple", "General PFAS": "orange"}
-    
-    # GRAPH 1: 3D DISCOVERY CUBE
+    # GRAPH 1: 3D DISCOVERY CUBE (Dark Mode)
     with col1:
-        st.subheader("🧊 1. Multi-Dimensional Risk")
+        st.subheader("🧊 1. Multi-Dimensional Risk Space")
         fig_3d = px.scatter_3d(
             res_df, 
             x='Bioaccumulation', y='Mobility', z='Persistence',
-            color='Subclass', symbol='Type' if 'Type' in res_df.columns else 'Subclass',
-            color_discrete_map=color_map, opacity=0.8, size_max=10,
-            title="Bioacc vs Mobility vs Persistence"
+            color='Subclass', 
+            symbol='Type' if 'Type' in res_df.columns else 'Subclass',
+            color_discrete_map=COLOR_MAP, 
+            opacity=0.9, 
+            size_max=12,
+            template="plotly_dark", # <--- PRETTY DARK MODE
+            title="Risk Landscape (Interactive)"
         )
-        fig_3d.update_layout(margin=dict(l=0, r=0, b=0, t=30))
+        fig_3d.update_layout(margin=dict(l=0, r=0, b=0, t=40), height=500)
         st.plotly_chart(fig_3d, use_container_width=True)
 
     # GRAPH 2: CLASS DISTRIBUTION (Bar Chart)
     with col2:
-        st.subheader("📊 2. Class Composition")
+        st.subheader("📊 2. Safety Classification")
+        count_df = res_df['Subclass'].value_counts().reset_index()
+        count_df.columns = ['Subclass', 'Count']
+        
         fig_bar = px.bar(
-            res_df, x="Subclass", color="Subclass", 
-            title="Count of Molecules by Class",
-            color_discrete_map=color_map
+            count_df, x="Subclass", y="Count", color="Subclass", 
+            title="Molecule Counts by Class",
+            color_discrete_map=COLOR_MAP,
+            template="plotly_dark",
+            text_auto=True # <--- Adds numbers on bars
         )
+        fig_bar.update_layout(height=500)
         st.plotly_chart(fig_bar, use_container_width=True)
 
     col3, col4 = st.columns(2)
 
-    # GRAPH 3: PARALLEL COORDINATES (The "Trace" Graph)
+    # GRAPH 3: PARALLEL COORDINATES (Fixed Color Scale)
     with col3:
-        st.subheader("📉 3. Property Tracing")
-        # Normalize Subclass to integer for coloring if needed, or use Bioacc
+        st.subheader("📉 3. Property Trace")
         fig_para = px.parallel_coordinates(
             res_df, 
             dimensions=['Persistence', 'Mobility', 'Bioaccumulation'],
             color="Bioaccumulation",
-            color_continuous_scale=px.colors.diverging.TealRose,
-            title="Trace: Persist -> Mobile -> Bioacc"
+            color_continuous_scale="Spectral_r", # <--- FIXED: Standard built-in scale (Red=High Risk)
+            title="Trace: Persist → Mobile → Bioacc",
+            template="plotly_dark"
         )
+        fig_para.update_layout(height=500)
         st.plotly_chart(fig_para, use_container_width=True)
 
     # GRAPH 4: DISTRIBUTION VIOLIN PLOT
     with col4:
-        st.subheader("🎻 4. Risk Distribution")
+        st.subheader("🎻 4. Bioaccumulation Spread")
         fig_vio = px.violin(
             res_df, y="Bioaccumulation", x="Subclass", 
             color="Subclass", box=True, points="all",
-            color_discrete_map=color_map,
-            title="Bioaccumulation Spread per Class"
+            color_discrete_map=COLOR_MAP,
+            template="plotly_dark",
+            title="Distribution Density"
         )
+        fig_vio.add_hline(y=3.5, line_dash="dash", line_color="orange", annotation_text="Safety Limit")
+        fig_vio.update_layout(height=500)
         st.plotly_chart(fig_vio, use_container_width=True)