rerun bugs fix try

app.py

@@ -454,102 +454,82 @@ if submit_planning:
                 st.error(f"Error in prepare_clustering_data: {str(e)}")
                 return None
 
-        @st.cache_data
-        def perform_clustering(_reduced_super_cgrs_dict, num_clusters, chunk_size=10):
-            """Perform clustering with memory-efficient processing"""
-            try:
-                mfp = MorganFingerprint()
-                
-                # Process fingerprints in chunks
-                all_fingerprints = {}
-                for i in range(0, len(_reduced_super_cgrs_dict), chunk_size):
-                    keys = list(_reduced_super_cgrs_dict.keys())[i:i+chunk_size]
-                    chunk_dict = {k: _reduced_super_cgrs_dict[k] for k in keys}
-                    chunk_fingerprints = {k: mfp.calculate(v) for k, v in chunk_dict.items()}
-                    all_fingerprints.update(chunk_fingerprints)
-                    del chunk_dict
-                    gc.collect()
-                
-                return cluster_molecules_optimized(all_fingerprints, max_clusters=num_clusters)
-            except Exception as e:
-                st.error(f"Error in perform_clustering: {str(e)}")
-                return None
-
         def memory_status():
             """Get current memory status"""
             process = psutil.Process()
             memory = process.memory_info().rss / 1024 / 1024
             return f"Memory usage: {memory:.2f} MB"
 
+        # Initialize session state for tree and clustering data
+        if 'tree_data' not in st.session_state:
+            st.session_state.tree_data = tree
+        if 'clustering_state' not in st.session_state:
+            st.session_state.clustering_state = {
+                'prepared': False,
+                'data': None,
+                'num_clusters': 2
+            }
+
         cluster_box, z = st.columns(2, gap="medium")
         with cluster_box:
-            # Initialize session state
-            if 'clustering_state' not in st.session_state:
-                st.session_state.clustering_state = {
-                    'prepared': False,
-                    'data': None,
-                    'last_memory': 0
-                }
-            
             st.write(memory_status())
-            st.write(f"Number of winning nodes: {len(tree.winning_nodes)}")
-            
-            # Memory management controls
-            if st.button('Clear memory'):
-                st.cache_data.clear()
-                st.session_state.clustering_state = {
-                    'prepared': False,
-                    'data': None,
-                    'last_memory': 0
-                }
-                gc.collect()
-                st.success("Memory cleared!")
-                st.rerun()
+            st.write(f"Number of winning nodes: {len(st.session_state.tree_data.winning_nodes)}")
             
-            # Prepare data with progress tracking
+            # Step 1: Prepare Data Button
             if not st.session_state.clustering_state['prepared']:
-                if st.button('Prepare clustering data'):
+                if st.button('Step 1: Prepare clustering data'):
                     with st.spinner("Preparing data..."):
                         try:
-                            progress_bar = st.progress(0)
-                            st.session_state.clustering_state['data'] = prepare_clustering_data(tree)
+                            st.session_state.clustering_state['data'] = prepare_clustering_data(st.session_state.tree_data)
                             st.session_state.clustering_state['prepared'] = True
-                            progress_bar.progress(100)
-                            st.success("Data prepared!")
+                            st.success("Data prepared! Now you can proceed to Step 2.")
                         except Exception as e:
                             st.error(f"Preparation failed: {str(e)}")
             
-            # Clustering controls
+            # Step 2: Only show clustering controls if data is prepared
             if st.session_state.clustering_state['prepared']:
-                num_clusters = st.slider(
+                st.markdown("### Step 2: Select number of clusters")
+                # Store slider value in session state
+                st.session_state.clustering_state['num_clusters'] = st.slider(
                     'Number of clusters',
                     min_value=2,
-                    max_value=min(10, len(tree.winning_nodes)),
-                    value=2
+                    max_value=min(10, len(st.session_state.tree_data.winning_nodes)),
+                    value=st.session_state.clustering_state['num_clusters']
                 )
                 
-                if st.button('Generate clusters'):
+                # Step 3: Generate Clusters Button
+                if st.button('Step 3: Generate clusters'):
                     with st.spinner("Clustering..."):
                         try:
                             results = perform_clustering(
                                 st.session_state.clustering_state['data'],
-                                num_clusters
+                                st.session_state.clustering_state['num_clusters']
                             )
                             
                             if results:
+                                st.success("Clustering complete!")
                                 for cluster_num, node_ids in results['clusters_dict'].items():
                                     with st.expander(f"Cluster {cluster_num}"):
                                         if node_ids:
                                             node_id = node_ids[0]
-                                            num_steps = len(tree.synthesis_route(node_id))
-                                            route_score = round(tree.route_score(node_id), 3)
+                                            num_steps = len(st.session_state.tree_data.synthesis_route(node_id))
+                                            route_score = round(st.session_state.tree_data.route_score(node_id), 3)
                                             st.image(
-                                                get_route_svg(tree, node_id),
+                                                get_route_svg(st.session_state.tree_data, node_id),
                                                 caption=f"Route {node_id}; {num_steps} steps; Score: {route_score}"
                                             )
                         except Exception as e:
                             st.error(f"Clustering failed: {str(e)}")
             
+            # Clear memory button
+            if st.button('Clear memory and start over'):
+                st.cache_data.clear()
+                del st.session_state.clustering_state
+                del st.session_state.tree_data
+                gc.collect()
+                st.success("Memory cleared! Please refresh the page to start over.")
+                st.rerun()
+            
         stat_col, download_col = st.columns(2, gap="medium")
 
         with stat_col: