bug fixes

app.py

@@ -127,7 +127,6 @@ if 'clustering_started' not in st.session_state:
 if 'clusters_downloaded' not in st.session_state:
     st.session_state.clusters_downloaded = False
 
-# st.write("Initial session state:", dict(st.session_state))
 
 intro_text = '''
 This is a demo of the graphical user interface of 
@@ -256,7 +255,7 @@ if submit_planning and not st.session_state.planning_done:
     st.session_state.planning_done = True
 
     # Display results if planning has been completed
-if st.session_state.planning_done and st.session_state.res is not None and st.session_state.clustering_started:
+if st.session_state.planning_done and st.session_state.res is not None:
     res = st.session_state.res
     tree = st.session_state.tree
 
@@ -275,235 +274,9 @@ if st.session_state.planning_done and st.session_state.res is not None and st.se
                 image_counter += 1
                 num_steps = len(tree.synthesis_route(node_id))
                 route_score = round(tree.route_score(node_id), 3)
-                st.image(get_route_svg(tree, node_id), caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
-
+                st.image(get_route_svg(tree, node_id), 
+                        caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
 
-        ### Modified part
-        # cluster_box, z = st.columns(2, gap="medium")
-        # with cluster_box:
-        #     num_clusters = st.slider('Number of clusters to display', min_value=2, max_value=10, value=2)
-
-        # submit_clustering = st.button('Start clustering')
-
-        # if submit_clustering:
-        #     st.subheader("Examples of clusters")
-        #     super_cgrs_dict = reassign_nums(tree)
-
-        #     reduced_super_cgrs_dict = process_all_rs_cgrs(super_cgrs_dict)
-
-        #     mfp = MorganFingerprint()
-
-        #     results = cluster_molecules(reduced_super_cgrs_dict, mfp, max_clusters=num_clusters)
-        # cluster_box, z = st.columns(2, gap="medium")
-        # with cluster_box:
-        #     # Initialize session state if not exists
-        #     if 'memory_warning_shown' not in st.session_state:
-        #         st.session_state.memory_warning_shown = False
-            
-        #     current_memory = psutil.Process().memory_info().rss / 1024 / 1024
-        #     st.write(f"Current memory usage: {current_memory:.2f} MB")
-        #     st.write(f"Number of winning nodes: {len(tree.winning_nodes)}")
-            
-        #     # Memory warning
-        #     if current_memory > 1000 and not st.session_state.memory_warning_shown:
-        #         st.warning("Memory usage is high. Consider reducing the number of routes or clearing cache.")
-        #         st.session_state.memory_warning_shown = True
-            
-        #     # Store the previous value in session state
-        #     if 'prev_num_clusters' not in st.session_state:
-        #         st.session_state.prev_num_clusters = 2
-            
-        #     num_clusters = st.slider(
-        #         'Number of clusters to display', 
-        #         min_value=2, 
-        #         max_value=min(10, len(tree.winning_nodes)), 
-        #         value=st.session_state.prev_num_clusters
-        #     )
-            
-        #     # Update the stored value only if it changed
-        #     if num_clusters != st.session_state.prev_num_clusters:
-        #         st.session_state.prev_num_clusters = num_clusters
-
-        # submit_clustering = st.button('Start clustering')
-
-        # if submit_clustering:
-        #     try:
-        #         with st.spinner("Processing clusters..."):
-        #             # Clear memory before starting
-        #             gc.collect()
-                    
-        #             st.write("Starting clustering process...")
-        #             memory_before = psutil.Process().memory_info().rss / 1024 / 1024
-        #             st.write(f"Memory before clustering: {memory_before:.2f} MB")
-                    
-        #             super_cgrs_dict = reassign_nums(tree)
-        #             del tree  # Free up memory from the tree object since we don't need it anymore
-        #             gc.collect()
-                    
-        #             reduced_super_cgrs_dict = process_all_rs_cgrs(super_cgrs_dict)
-        #             del super_cgrs_dict  # Free up memory
-        #             gc.collect()
-                    
-        #             memory_after = psutil.Process().memory_info().rss / 1024 / 1024
-        #             st.write(f"Memory after CGR processing: {memory_after:.2f} MB")
-                    
-        #             mfp = MorganFingerprint()
-        #             results = cluster_molecules(reduced_super_cgrs_dict, mfp, max_clusters=num_clusters)
-        #             del reduced_super_cgrs_dict  # Free up memory
-        #             gc.collect()
-                    
-        #             st.write("Clustering completed")
-                    
-        #     except Exception as e:
-        #         st.error(f"Clustering failed with error: {str(e)}")
-        #         st.write(f"Memory at error: {psutil.Process().memory_info().rss / 1024 / 1024:.2f} MB")
-        #         raise e
-
-            
-            # Access results
-            # clusters = results['clusters_dict']
-
-            # for cluster_num, node_id_list in clusters.items():
-            #     st.markdown(f"Cluster's number: ``{cluster_num}``")
-            #     node_id = node_id_list[0]
-            #     num_steps = len(tree.synthesis_route(node_id))
-            #     route_score = round(tree.route_score(node_id), 3)
-            #     st.image(get_route_svg(tree, node_id), caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
-
-        @st.cache_data(hash_funcs={Tree: lambda _: None})
-        def prepare_clustering_data(tree):
-            try:
-                # Log the start and basic info from the Tree
-                print("Starting clustering data preparation.")
-                total_nodes = len(tree.winning_nodes)
-                print(f"Total winning nodes: {total_nodes}")
-                print(f"Tree id: {id(tree)}")
-                
-                chunk_size = 10
-                super_cgrs_dict = {}
-                
-                # Process winning nodes in chunks
-                for i in range(0, total_nodes, chunk_size):
-                    current_chunk = list(tree.winning_nodes)[i:i+chunk_size]
-                    print(f"Processing chunk {i // chunk_size + 1}: Nodes {current_chunk}")
-                    
-                    temp_dict = {}
-                    for node in current_chunk:
-                        try:
-                            # Log before processing each node
-                            print(f"Processing node {node}")
-                            route = tree.synthesis_route(node)
-                            temp_dict[node] = route
-                            print(f"Node {node} processed successfully (route length: {len(route)}).")
-                        except Exception as e:
-                            print(f"Error processing node {node}: {e}")
-                    
-                    # Log before calling reassign_nums_chunk
-                    print(f"Calling reassign_nums_chunk for nodes: {list(temp_dict.keys())}")
-                    chunk_super_cgrs = reassign_nums_chunk(temp_dict)
-                    super_cgrs_dict.update(chunk_super_cgrs)
-                    print(f"Chunk {i // chunk_size + 1} processed. Keys: {list(chunk_super_cgrs.keys())}")
-                    
-                    del temp_dict
-                    gc.collect()
-                
-                # Process reduced CGRs in chunks
-                reduced_super_cgrs_dict = {}
-                for i in range(0, len(super_cgrs_dict), chunk_size):
-                    keys = list(super_cgrs_dict.keys())[i:i+chunk_size]
-                    chunk_dict = {k: super_cgrs_dict[k] for k in keys}
-                    print(f"Reducing chunk for keys: {keys}")
-                    reduced_chunk = process_all_rs_cgrs(chunk_dict)
-                    reduced_super_cgrs_dict.update(reduced_chunk)
-                    print(f"Reduced chunk processed for keys: {list(reduced_chunk.keys())}")
-                    
-                    del chunk_dict
-                    gc.collect()
-                
-                print("Clustering data preparation complete.")
-                return reduced_super_cgrs_dict
-            except Exception as e:
-                print(f"Error in prepare_clustering_data: {str(e)}")
-                st.error(f"Error in prepare_clustering_data: {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:
-            st.write(memory_status())
-            st.write(f"Number of winning nodes: {len(st.session_state.tree_data.winning_nodes)}")
-            
-            # Step 1: Prepare Data Button
-            if not st.session_state.clustering_state['prepared']:
-                if st.button('Step 1: Prepare clustering data'):
-                    with st.spinner("Preparing data..."):
-                        try:
-                            st.session_state.clustering_state['data'] = prepare_clustering_data(st.session_state.tree_data)
-                            st.session_state.clustering_state['prepared'] = True
-                            st.success("Data prepared! Now you can proceed to Step 2.")
-                        except Exception as e:
-                            st.error(f"Preparation failed: {str(e)}")
-            
-            # Step 2: Only show clustering controls if data is prepared
-            if st.session_state.clustering_state['prepared']:
-                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(st.session_state.tree_data.winning_nodes)),
-                    value=st.session_state.clustering_state['num_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'],
-                                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(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(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:
             st.subheader("Statistics")
@@ -519,7 +292,6 @@ if st.session_state.planning_done and st.session_state.res is not None and st.se
 
         st.header("Clustering the retrosynthetic routes")
 
-        # Initialize slider state if not already set
         if 'num_clusters' not in st.session_state:
             st.session_state['num_clusters'] = 10
 
@@ -532,13 +304,10 @@ if st.session_state.planning_done and st.session_state.res is not None and st.se
                 value=st.session_state['num_clusters'],
                 key='cluster_slider'
             )
-            # Save the current slider value to session_state
             st.session_state['num_clusters'] = num_clusters
 
         if st.button('Start clustering', key='submit_clustering'):
             st.session_state.clustering_started = True
-            # st.write("Clustering started; session state now:", dict(st.session_state))
-        #     st.write("Clustering started!")
             st.subheader("Examples of clusters")
             super_cgrs_dict = reassign_nums(tree)
 
@@ -562,16 +331,13 @@ if st.session_state.planning_done and st.session_state.res is not None and st.se
             cluster_stat_col, cluster_download_col = st.columns(2, gap="medium")
             with cluster_stat_col:
                 st.subheader("Statistics")
-                # st.write(cluster_sizes)
                 cluster_df = pd.DataFrame({'Cluster': range(len(cluster_sizes)), 'Routes': cluster_sizes})
-                # cluster_df = pd.DataFrame(cluster_sizes, index=[0])
                 st.write(cluster_df)
 
             def on_download_click():
                 st.session_state.clusters_downloaded = True
                 st.write("Download clusters button pressed via on_click. Updated session state:", dict(st.session_state))
                 save_route_images(tree, reactions_dict, cluster_dict=clusters_converted)
-                # Here you can call save_route_images(...) if desired.
                 
             with cluster_download_col:
                 st.subheader("Downloads: Don't work. Resets evey time")
@@ -581,24 +347,8 @@ if st.session_state.planning_done and st.session_state.res is not None and st.se
                 # Use on_click to capture the click event reliably.
                 st.button('Download clusters', key='download_clusters_button', on_click=on_download_click)
                 
-                # Log whether the flag has been set after the button definition.
                 st.write("Clusters downloaded flag (from session_state):", st.session_state.get("clusters_downloaded"))
 
-            #             # save_route_images(tree, reactions_dict, cluster_dict=clusters_converted)
-            # with cluster_download_col:
-            #     st.subheader("Downloads")
-            #     reactions_dict = extract_reactions(tree)
-            #     clusters_converted = {int(key): value for key, value in clusters.items()} if clusters else clusters
-
-            #     if st.session_state.clustering_started:
-            #         st.write("Rendering download clusters button. Session state:", dict(st.session_state))
-            #         # Use a more unique key for the download button.
-            #         download_clusters = st.button('Download clusters', key='download_clusters_button')
-            #         st.write("download_clusters value:", download_clusters)
-            #         if download_clusters:
-            #             st.session_state.clusters_downloaded = True
-            #             st.write("Download clusters button pressed. Updated session state:", dict(st.session_state))
-
             col1, _ = st.columns([.2, .8])
             with col1:
                 fig = pie_chart(cluster_sizes)
@@ -631,4 +381,8 @@ st.divider()
 st.header('Restart from the beginning?')
 if st.button("Restart"):
     st.session_state.planning_done = False
+    st.session_state.tree = None
+    st.session_state.res = None
+    st.session_state.clustering_started = False
+    st.session_state.clusters_downloaded = False
     st.rerun()