optimized clustering

app.py

@@ -34,6 +34,82 @@ disable_progress_bars("huggingface_hub")
 smiles_parser = SMILESRead.create_parser(ignore=True)
 
 
+def reassign_nums_chunk(route_dict):
+    """Process a chunk of routes for reassigning numbers"""
+    return {k: reassign_nums(v) for k, v in route_dict.items()}
+
+def cluster_molecules_optimized(fingerprints_dict, max_clusters):
+    """Memory-optimized version of cluster_molecules.
+    
+    Args:
+        fingerprints_dict (dict): Dictionary of pre-computed fingerprints
+        max_clusters (int): Maximum number of clusters
+        
+    Returns:
+        dict: Clustering results containing clusters_dict and cluster_labels
+    """
+    try:
+        # Convert dictionary to arrays for efficient processing
+        labels = np.array(list(fingerprints_dict.keys()))
+        fingerprints = np.array(list(fingerprints_dict.values()))
+        
+        # Calculate similarity matrix in chunks to save memory
+        chunk_size = 100
+        n_samples = len(fingerprints)
+        similarity_matrix = np.zeros((n_samples, n_samples))
+        
+        for i in range(0, n_samples, chunk_size):
+            chunk_end = min(i + chunk_size, n_samples)
+            chunk = fingerprints[i:chunk_end]
+            
+            # Calculate similarity for this chunk against all fingerprints
+            similarity_chunk = tanimoto_similarity_continuous(chunk, fingerprints)
+            similarity_matrix[i:chunk_end] = similarity_chunk
+            
+            # Clear memory
+            del similarity_chunk
+            gc.collect()
+        
+        # Convert to distance matrix
+        distance_matrix = 1 - similarity_matrix
+        
+        # Free memory
+        del similarity_matrix
+        gc.collect()
+        
+        # Calculate condensed distance matrix
+        condensed_distance = squareform(distance_matrix)
+        
+        # Free memory
+        del distance_matrix
+        gc.collect()
+        
+        # Calculate linkage
+        Z = fastcluster.linkage(condensed_distance, method='average')
+        
+        # Free memory
+        del condensed_distance
+        gc.collect()
+        
+        # Perform clustering
+        cluster_labels = fcluster(Z, max_clusters, criterion='maxclust')
+        
+        # Create clusters dictionary
+        clusters_dict = {}
+        for cluster in range(1, max_clusters + 1):
+            cluster_indices = np.where(cluster_labels == cluster)[0]
+            clusters_dict[cluster] = list(labels[cluster_indices])
+        
+        return {
+            'clusters_dict': clusters_dict,
+            'cluster_labels': cluster_labels,
+            'linkage_matrix': Z
+        }
+        
+    except Exception as e:
+        print(f"Error in cluster_molecules_optimized: {str(e)}")
+        raise e
+
 def download_button(object_to_download, download_filename, button_text, pickle_it=False):
     """
     Issued from
@@ -341,86 +417,139 @@ if submit_planning:
             #     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}")
 
-        # Add these functions outside the if submit_planning block
         @st.cache_data
         def prepare_clustering_data(tree):
-            """Pre-compute and cache the clustering data"""
-            super_cgrs_dict = reassign_nums(tree)
-            reduced_super_cgrs_dict = process_all_rs_cgrs(super_cgrs_dict)
-            return reduced_super_cgrs_dict
+            """Pre-compute and cache the clustering data in chunks"""
+            try:
+                # Free memory before starting
+                gc.collect()
+                
+                # Process in chunks
+                chunk_size = 10
+                super_cgrs_dict = {}
+                
+                for i in range(0, len(tree.winning_nodes), chunk_size):
+                    chunk = list(tree.winning_nodes)[i:i+chunk_size]
+                    temp_dict = {node: tree.synthesis_route(node) for node in chunk}
+                    chunk_super_cgrs = reassign_nums_chunk(temp_dict)
+                    super_cgrs_dict.update(chunk_super_cgrs)
+                    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}
+                    reduced_chunk = process_all_rs_cgrs(chunk_dict)
+                    reduced_super_cgrs_dict.update(reduced_chunk)
+                    del chunk_dict
+                    gc.collect()
+                
+                del super_cgrs_dict
+                gc.collect()
+                
+                return reduced_super_cgrs_dict
+            except Exception as e:
+                st.error(f"Error in prepare_clustering_data: {str(e)}")
+                return None
 
         @st.cache_data
-        def perform_clustering(_reduced_super_cgrs_dict, num_clusters):
-            """Perform the actual clustering with cached results"""
-            mfp = MorganFingerprint()
-            return cluster_molecules(_reduced_super_cgrs_dict, mfp, max_clusters=num_clusters)
+        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"
 
         cluster_box, z = st.columns(2, gap="medium")
         with cluster_box:
-            # Initialize session state for clustering data
-            if 'clustering_prepared' not in st.session_state:
-                st.session_state.clustering_prepared = False
-                st.session_state.reduced_super_cgrs_dict = None
+            # Initialize session state
+            if 'clustering_state' not in st.session_state:
+                st.session_state.clustering_state = {
+                    'prepared': False,
+                    'data': None,
+                    'last_memory': 0
+                }
             
-            current_memory = psutil.Process().memory_info().rss / 1024 / 1024
-            st.write(f"Current memory usage: {current_memory:.2f} MB")
+            st.write(memory_status())
             st.write(f"Number of winning nodes: {len(tree.winning_nodes)}")
             
-            # Prepare data button
-            if not st.session_state.clustering_prepared:
+            # 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()
+            
+            # Prepare data with progress tracking
+            if not st.session_state.clustering_state['prepared']:
                 if st.button('Prepare clustering data'):
-                    with st.spinner("Preparing clustering data..."):
+                    with st.spinner("Preparing data..."):
                         try:
-                            gc.collect()
-                            st.session_state.reduced_super_cgrs_dict = prepare_clustering_data(tree)
-                            st.session_state.clustering_prepared = True
-                            st.success("Data prepared successfully!")
+                            progress_bar = st.progress(0)
+                            st.session_state.clustering_state['data'] = prepare_clustering_data(tree)
+                            st.session_state.clustering_state['prepared'] = True
+                            progress_bar.progress(100)
+                            st.success("Data prepared!")
                         except Exception as e:
-                            st.error(f"Failed to prepare data: {str(e)}")
-                            st.write(f"Memory at error: {current_memory:.2f} MB")
+                            st.error(f"Preparation failed: {str(e)}")
             
-            # Only show clustering controls if data is prepared
-            if st.session_state.clustering_prepared:
+            # Clustering controls
+            if st.session_state.clustering_state['prepared']:
                 num_clusters = st.slider(
-                    'Number of clusters to display', 
-                    min_value=2, 
-                    max_value=min(10, len(tree.winning_nodes)), 
+                    'Number of clusters',
+                    min_value=2,
+                    max_value=min(10, len(tree.winning_nodes)),
                     value=2
                 )
-
+                
                 if st.button('Generate clusters'):
-                    with st.spinner("Generating clusters..."):
+                    with st.spinner("Clustering..."):
                         try:
                             results = perform_clustering(
-                                st.session_state.reduced_super_cgrs_dict,
+                                st.session_state.clustering_state['data'],
                                 num_clusters
                             )
                             
-                            # Display clusters
-                            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}"
-                                )
-                                
+                            if results:
+                                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)
+                                            st.image(
+                                                get_route_svg(tree, node_id),
+                                                caption=f"Route {node_id}; {num_steps} steps; Score: {route_score}"
+                                            )
                         except Exception as e:
                             st.error(f"Clustering failed: {str(e)}")
-                            st.write(f"Memory at error: {current_memory:.2f} MB")
-
-            # Add clear cache button
-            if st.button('Clear cache and memory'):
-                st.cache_data.clear()
-                st.session_state.clustering_prepared = False
-                st.session_state.reduced_super_cgrs_dict = None
-                gc.collect()
-                st.success("Cache and memory cleared!")
-                st.rerun()
-    
+            
         stat_col, download_col = st.columns(2, gap="medium")
 
         with stat_col: