masster 0.4.19__py3-none-any.whl → 0.4.20__py3-none-any.whl

masster/study/merge.py

@@ -10,7 +10,8 @@ from datetime import datetime
 from tqdm import tqdm
 import pyopenms as oms
 import polars as pl
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
+from concurrent.futures.process import BrokenProcessPool
 from masster.study.defaults import merge_defaults
 
 
@@ -367,6 +368,17 @@ def merge(self, **kwargs) -> None:
     if params.method not in ['sensitivity', 'qt', 'nowarp', 'kd_chunked', 'qt_chunked', 'quality']:
         raise ValueError(f"Invalid method '{params.method}'. Must be one of: ['sensitivity', 'qt', 'nowarp', 'kd_chunked', 'qt_chunked', 'quality']")
     
+    # Check if chunked method is advisable for large datasets
+    num_samples = len(self.samples_df) if hasattr(self, 'samples_df') and self.samples_df is not None else 0
+    if num_samples > 500:
+        chunked_methods = {'kd_chunked', 'qt_chunked'}
+        if params.method not in chunked_methods:
+            self.logger.warning(
+                f"Large dataset detected ({num_samples} samples > 500). "
+                f"For better performance and memory efficiency, consider using a chunked method: "
+                f"'kd_chunked' or 'qt_chunked' instead of '{params.method}'"
+            )
+    
     # Persist last used params for diagnostics
     try:
         self._merge_params_last = params.to_dict()
@@ -385,7 +397,7 @@ def merge(self, **kwargs) -> None:
     # Ensure feature maps are available for merging (regenerate if needed)
     if len(self.features_maps) < len(self.samples_df):
         self.features_maps = []
-        self.load_features()
+        # Feature maps will be generated on-demand within each merge method
     
     self.logger.info(
         f"Merge: {params.method}, samples={params.min_samples}, rt_tol={params.rt_tol}s, mz_tol={params.mz_tol}Da, min_rel_cc_size={params.min_rel_cc_size}, max_pairwise_log_fc={params.max_pairwise_log_fc}, max_nr_conflicts={params.max_nr_conflicts}"
@@ -433,9 +445,16 @@ def merge(self, **kwargs) -> None:
         consensus_map = _merge_qt_chunked(self, params, cached_adducts_df, cached_valid_adducts)
         # Note: _merge_qt_chunked populates consensus_df directly, no need to extract
     
+    # Enhanced post-clustering to merge over-segmented features (for qt and kd methods)
+    if params.method in ['qt', 'sensitivity', 'qt_chunked', 'kd_chunked']:
+        self._consensus_cleanup(params.rt_tol, params.mz_tol)
+    
     # Perform adduct grouping
     self._perform_adduct_grouping(params.rt_tol, params.mz_tol)
     
+    # Identify coeluting consensus features by mass shifts and update adduct information
+    self._identify_adduct_by_mass_shift(params.rt_tol, cached_adducts_df)
+    
     # Link MS2 if requested
     if params.link_ms2:
         self._finalize_merge(params.link_ms2, params.min_samples)
@@ -448,10 +467,13 @@ def merge(self, **kwargs) -> None:
 def _merge_kd(self, params: merge_defaults) -> oms.ConsensusMap:
     """KD-tree based merge (fast, recommended)"""
     
+    # Generate temporary feature maps on-demand from features_df
+    temp_feature_maps = _generate_feature_maps_on_demand(self)
+    
     consensus_map = oms.ConsensusMap()
     file_descriptions = consensus_map.getColumnHeaders()
     
-    for i, feature_map in enumerate(self.features_maps):
+    for i, feature_map in enumerate(temp_feature_maps):
         file_description = file_descriptions.get(i, oms.ColumnHeader())
         file_description.filename = self.samples_df.row(i, named=True)["sample_name"]
         file_description.size = feature_map.size()
@@ -477,22 +499,145 @@ def _merge_kd(self, params: merge_defaults) -> oms.ConsensusMap:
     #params_oms.setValue("link:charge_merging", "With_charge_zero") THIS LEADS TO A CRASH
     
     grouper.setParameters(params_oms)
-    grouper.group(self.features_maps, consensus_map)
+    grouper.group(temp_feature_maps, consensus_map)
     
     return consensus_map
 
 
+def _generate_feature_maps_on_demand(study):
+    """
+    Generate feature maps on-demand from study.features_df for merge operations.
+    Returns temporary feature maps that are not cached in the study.
+    
+    Args:
+        study: Study object containing features_df and samples_df
+    
+    Returns:
+        list: List of temporary FeatureMap objects
+    """
+    import polars as pl
+    import pyopenms as oms
+    import numpy as np
+    
+    if study.features_df is None or len(study.features_df) == 0:
+        study.logger.error("No features_df available for generating feature maps")
+        return []
+    
+    temp_feature_maps = []
+    n_samples = len(study.samples_df)
+    n_features = len(study.features_df)
+    
+    # Performance optimization: use efficient polars groupby for large datasets
+    use_groupby_optimization = n_features > 5000
+    if use_groupby_optimization:
+        study.logger.debug(f"Using polars groupby optimization for {n_features} features across {n_samples} samples")
+        
+        # Pre-group features by sample_uid - this is much more efficient than repeated filtering
+        features_by_sample = study.features_df.group_by("sample_uid").agg([
+            pl.col("feature_id"),
+            pl.col("mz"), 
+            pl.col("rt"),
+            pl.col("inty"),
+            pl.col("quality").fill_null(1.0),
+            pl.col("charge").fill_null(0)
+        ])
+        
+        # Convert to dictionary for fast lookups
+        sample_feature_dict = {}
+        for row in features_by_sample.iter_rows(named=True):
+            sample_uid = row["sample_uid"]
+            # Convert lists to numpy arrays for vectorized operations
+            sample_feature_dict[sample_uid] = {
+                "feature_id": np.array(row["feature_id"]),
+                "mz": np.array(row["mz"]),
+                "rt": np.array(row["rt"]),
+                "inty": np.array(row["inty"]),
+                "quality": np.array(row["quality"]),
+                "charge": np.array(row["charge"])
+            }
+    
+    # Process each sample in order
+    for sample_index, row_dict in enumerate(study.samples_df.iter_rows(named=True)):
+        sample_uid = row_dict["sample_uid"]
+        
+        if use_groupby_optimization:
+            # Use pre-grouped data with vectorized operations
+            if sample_uid not in sample_feature_dict:
+                feature_map = oms.FeatureMap()
+                temp_feature_maps.append(feature_map)
+                continue
+                
+            sample_data = sample_feature_dict[sample_uid]
+            n_sample_features = len(sample_data["feature_id"])
+            
+            if n_sample_features == 0:
+                feature_map = oms.FeatureMap()
+                temp_feature_maps.append(feature_map)
+                continue
+            
+            # Create new FeatureMap
+            feature_map = oms.FeatureMap()
+            
+            # Use vectorized data directly (no conversion needed)
+            for i in range(n_sample_features):
+                try:
+                    feature = oms.Feature()
+                    feature.setUniqueId(int(sample_data["feature_id"][i]))
+                    feature.setMZ(float(sample_data["mz"][i]))
+                    feature.setRT(float(sample_data["rt"][i]))
+                    feature.setIntensity(float(sample_data["inty"][i]))
+                    feature.setOverallQuality(float(sample_data["quality"][i]))
+                    feature.setCharge(int(sample_data["charge"][i]))
+                    feature_map.push_back(feature)
+                except (ValueError, TypeError) as e:
+                    study.logger.warning(f"Skipping feature due to conversion error: {e}")
+                    continue
+        else:
+            # Use original polars-based approach for smaller datasets
+            sample_features = study.features_df.filter(pl.col("sample_uid") == sample_uid)
+            
+            # Create new FeatureMap
+            feature_map = oms.FeatureMap()
+            
+            # Convert DataFrame features to OpenMS Features
+            for feature_row in sample_features.iter_rows(named=True):
+                feature = oms.Feature()
+                
+                # Set properties from DataFrame (handle missing values gracefully)
+                try:
+                    feature.setUniqueId(int(feature_row["feature_id"]))
+                    feature.setMZ(float(feature_row["mz"]))
+                    feature.setRT(float(feature_row["rt"]))
+                    feature.setIntensity(float(feature_row["inty"]))
+                    feature.setOverallQuality(float(feature_row["quality"]))
+                    feature.setCharge(int(feature_row["charge"]))
+                    
+                    # Add to feature map
+                    feature_map.push_back(feature)
+                except (ValueError, TypeError) as e:
+                    study.logger.warning(f"Skipping feature due to conversion error: {e}")
+                    continue
+        
+        temp_feature_maps.append(feature_map)
+    
+    study.logger.debug(f"Generated {len(temp_feature_maps)} temporary feature maps from features_df")
+    return temp_feature_maps
+
+
 def _merge_qt(self, params: merge_defaults) -> oms.ConsensusMap:
     """QT (Quality Threshold) based merge"""
     
-    n_samples = len(self.features_maps)
+    # Generate temporary feature maps on-demand from features_df
+    temp_feature_maps = _generate_feature_maps_on_demand(self)
+    
+    n_samples = len(temp_feature_maps)
     if n_samples > 1000:
         self.logger.warning(f"QT with {n_samples} samples may be slow [O(n²)]. Consider KD [O(n log n)]")
     
     consensus_map = oms.ConsensusMap()
     file_descriptions = consensus_map.getColumnHeaders()
     
-    for i, feature_map in enumerate(self.features_maps):
+    for i, feature_map in enumerate(temp_feature_maps):
         file_description = file_descriptions.get(i, oms.ColumnHeader())
         file_description.filename = self.samples_df.row(i, named=True)["sample_name"]
         file_description.size = feature_map.size()
@@ -515,7 +660,7 @@ def _merge_qt(self, params: merge_defaults) -> oms.ConsensusMap:
     params_oms.setValue("nr_partitions", params.nr_partitions)
 
     grouper.setParameters(params_oms)
-    grouper.group(self.features_maps, consensus_map)
+    grouper.group(temp_feature_maps, consensus_map)
     
     return consensus_map
 
@@ -1013,10 +1158,13 @@ def _filter_coherence(self, features: list, min_coherence: float) -> list:
 def _merge_kd_nowarp(self, params: merge_defaults) -> oms.ConsensusMap:
     """KD-tree based merge without RT warping"""
     
+    # Generate temporary feature maps on-demand from features_df
+    temp_feature_maps = _generate_feature_maps_on_demand(self)
+    
     consensus_map = oms.ConsensusMap()
     file_descriptions = consensus_map.getColumnHeaders()
     
-    for i, feature_map in enumerate(self.features_maps):
+    for i, feature_map in enumerate(temp_feature_maps):
         file_description = file_descriptions.get(i, oms.ColumnHeader())
         file_description.filename = self.samples_df.row(i, named=True)["sample_name"]
         file_description.size = feature_map.size()
@@ -1040,7 +1188,7 @@ def _merge_kd_nowarp(self, params: merge_defaults) -> oms.ConsensusMap:
     #params_oms.setValue("link:charge_merging", "Any")
     
     grouper.setParameters(params_oms)
-    grouper.group(self.features_maps, consensus_map)
+    grouper.group(temp_feature_maps, consensus_map)
     
     return consensus_map
 
@@ -1048,7 +1196,10 @@ def _merge_kd_nowarp(self, params: merge_defaults) -> oms.ConsensusMap:
 def _merge_kd_chunked(self, params: merge_defaults, cached_adducts_df=None, cached_valid_adducts=None) -> oms.ConsensusMap:
     """KD-based chunked merge with proper cross-chunk consensus building and optional parallel processing"""
     
-    n_samples = len(self.features_maps)
+    # Generate temporary feature maps on-demand from features_df
+    temp_feature_maps = _generate_feature_maps_on_demand(self)
+    
+    n_samples = len(temp_feature_maps)
     if n_samples <= params.chunk_size:
         self.logger.info(f"Dataset size ({n_samples}) ≤ chunk_size, using KD merge")
         consensus_map = _merge_kd(self, params)
@@ -1060,7 +1211,7 @@ def _merge_kd_chunked(self, params: merge_defaults, cached_adducts_df=None, cach
     chunks = []
     for i in range(0, n_samples, params.chunk_size):
         chunk_end = min(i + params.chunk_size, n_samples)
-        chunks.append((i, self.features_maps[i:chunk_end]))
+        chunks.append((i, temp_feature_maps[i:chunk_end]))
     
     self.logger.debug(f"Processing {len(chunks)} chunks of max {params.chunk_size} samples using {params.threads or 'sequential'} thread(s)")
     
@@ -1146,28 +1297,66 @@ def _merge_kd_chunked(self, params: merge_defaults, cached_adducts_df=None, cach
             }
             chunk_data_list.append(chunk_data)
         
-        # Process chunks in parallel
-        with ProcessPoolExecutor(max_workers=params.threads) as executor:
-            # Submit all chunk processing tasks
-            future_to_chunk = {executor.submit(_process_kd_chunk_parallel, chunk_data): i 
-                             for i, chunk_data in enumerate(chunk_data_list)}
-            
-            # Collect results with progress tracking
-            completed_chunks = 0
-            total_chunks = len(chunk_data_list)
-            serialized_chunk_results = []
-            
-            for future in as_completed(future_to_chunk):
-                chunk_idx = future_to_chunk[future]
-                try:
-                    chunk_start_idx, consensus_features = future.result()
-                    serialized_chunk_results.append((chunk_start_idx, consensus_features))
-                    completed_chunks += 1
-                    n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
-                    self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
-                except Exception as exc:
-                    self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
-                    raise exc
+        # Process chunks in parallel - try ProcessPoolExecutor first, fallback to ThreadPoolExecutor on Windows
+        try:
+            with ProcessPoolExecutor(max_workers=params.threads) as executor:
+                # Submit all chunk processing tasks
+                future_to_chunk = {executor.submit(_process_kd_chunk_parallel, chunk_data): i 
+                                 for i, chunk_data in enumerate(chunk_data_list)}
+                
+                # Collect results with progress tracking
+                completed_chunks = 0
+                total_chunks = len(chunk_data_list)
+                serialized_chunk_results = []
+                
+                for future in as_completed(future_to_chunk):
+                    chunk_idx = future_to_chunk[future]
+                    try:
+                        chunk_start_idx, consensus_features = future.result()
+                        serialized_chunk_results.append((chunk_start_idx, consensus_features))
+                        completed_chunks += 1
+                        n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
+                        self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
+                    except Exception as exc:
+                        # Check if this is a BrokenProcessPool exception from Windows multiprocessing issues
+                        if isinstance(exc, BrokenProcessPool) or "process pool" in str(exc).lower():
+                            # Convert to RuntimeError so outer except block can catch it for fallback
+                            raise RuntimeError(f"Windows multiprocessing failure: {exc}")
+                        else:
+                            self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
+                            raise exc
+                        
+        except (RuntimeError, OSError, BrokenProcessPool) as e:
+            # Handle Windows multiprocessing issues - fallback to ThreadPoolExecutor
+            if ("freeze_support" in str(e) or "spawn" in str(e) or "bootstrapping" in str(e) or 
+                "process pool" in str(e).lower() or "Windows multiprocessing failure" in str(e)):
+                self.logger.warning(f"ProcessPoolExecutor failed (likely Windows multiprocessing issue): {e}")
+                self.logger.info(f"Falling back to ThreadPoolExecutor with {params.threads} threads")
+                
+                with ThreadPoolExecutor(max_workers=params.threads) as executor:
+                    # Submit all chunk processing tasks
+                    future_to_chunk = {executor.submit(_process_kd_chunk_parallel, chunk_data): i 
+                                     for i, chunk_data in enumerate(chunk_data_list)}
+                    
+                    # Collect results with progress tracking
+                    completed_chunks = 0
+                    total_chunks = len(chunk_data_list)
+                    serialized_chunk_results = []
+                    
+                    for future in as_completed(future_to_chunk):
+                        chunk_idx = future_to_chunk[future]
+                        try:
+                            chunk_start_idx, consensus_features = future.result()
+                            serialized_chunk_results.append((chunk_start_idx, consensus_features))
+                            completed_chunks += 1
+                            n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
+                            self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
+                        except Exception as exc:
+                            self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
+                            raise exc
+            else:
+                # Re-raise other exceptions
+                raise
         
         # Store serialized results for _merge_chunk_results to handle directly  
         chunk_consensus_maps = []
@@ -1187,7 +1376,10 @@ def _merge_kd_chunked(self, params: merge_defaults, cached_adducts_df=None, cach
 def _merge_qt_chunked(self, params: merge_defaults, cached_adducts_df=None, cached_valid_adducts=None) -> oms.ConsensusMap:
     """QT-based chunked merge with proper cross-chunk consensus building and optional parallel processing"""
     
-    n_samples = len(self.features_maps)
+    # Generate temporary feature maps on-demand from features_df
+    temp_feature_maps = _generate_feature_maps_on_demand(self)
+    
+    n_samples = len(temp_feature_maps)
     if n_samples <= params.chunk_size:
         self.logger.info(f"Dataset size ({n_samples}) ≤ chunk_size, using QT merge")
         consensus_map = _merge_qt(self, params)
@@ -1199,7 +1391,7 @@ def _merge_qt_chunked(self, params: merge_defaults, cached_adducts_df=None, cach
     chunks = []
     for i in range(0, n_samples, params.chunk_size):
         chunk_end = min(i + params.chunk_size, n_samples)
-        chunks.append((i, self.features_maps[i:chunk_end]))
+        chunks.append((i, temp_feature_maps[i:chunk_end]))
     
     self.logger.debug(f"Processing {len(chunks)} chunks of max {params.chunk_size} samples using {params.threads or 'sequential'} thread(s)")
     
@@ -1277,28 +1469,69 @@ def _merge_qt_chunked(self, params: merge_defaults, cached_adducts_df=None, cach
             }
             chunk_data_list.append(chunk_data)
         
-        # Process chunks in parallel
-        with ProcessPoolExecutor(max_workers=params.threads) as executor:
-            # Submit all chunk processing tasks
-            future_to_chunk = {executor.submit(_process_qt_chunk_parallel, chunk_data): i 
-                             for i, chunk_data in enumerate(chunk_data_list)}
-            
-            # Collect results with progress tracking
-            completed_chunks = 0
-            total_chunks = len(chunk_data_list)
-            serialized_chunk_results = []
-            
-            for future in as_completed(future_to_chunk):
-                chunk_idx = future_to_chunk[future]
-                try:
-                    chunk_start_idx, consensus_features = future.result()
-                    serialized_chunk_results.append((chunk_start_idx, consensus_features))
-                    completed_chunks += 1
-                    n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
-                    self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
-                except Exception as exc:
-                    self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
-                    raise exc
+        # Process chunks in parallel - try ProcessPoolExecutor first, fallback to ThreadPoolExecutor on Windows
+        executor_class = ProcessPoolExecutor
+        executor_name = "processes"
+        
+        try:
+            with ProcessPoolExecutor(max_workers=params.threads) as executor:
+                # Submit all chunk processing tasks
+                future_to_chunk = {executor.submit(_process_qt_chunk_parallel, chunk_data): i 
+                                 for i, chunk_data in enumerate(chunk_data_list)}
+                
+                # Collect results with progress tracking
+                completed_chunks = 0
+                total_chunks = len(chunk_data_list)
+                serialized_chunk_results = []
+                
+                for future in as_completed(future_to_chunk):
+                    chunk_idx = future_to_chunk[future]
+                    try:
+                        chunk_start_idx, consensus_features = future.result()
+                        serialized_chunk_results.append((chunk_start_idx, consensus_features))
+                        completed_chunks += 1
+                        n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
+                        self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
+                    except Exception as exc:
+                        # Check if this is a BrokenProcessPool exception from Windows multiprocessing issues
+                        if isinstance(exc, BrokenProcessPool) or "process pool" in str(exc).lower():
+                            # Convert to RuntimeError so outer except block can catch it for fallback
+                            raise RuntimeError(f"Windows multiprocessing failure: {exc}")
+                        else:
+                            self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
+                            raise exc
+                        
+        except (RuntimeError, OSError, BrokenProcessPool) as e:
+            # Handle Windows multiprocessing issues - fallback to ThreadPoolExecutor
+            if ("freeze_support" in str(e) or "spawn" in str(e) or "bootstrapping" in str(e) or 
+                "process pool" in str(e).lower() or "Windows multiprocessing failure" in str(e)):
+                self.logger.warning(f"ProcessPoolExecutor failed (likely Windows multiprocessing issue): {e}")
+                self.logger.info(f"Falling back to ThreadPoolExecutor with {params.threads} threads")
+                
+                with ThreadPoolExecutor(max_workers=params.threads) as executor:
+                    # Submit all chunk processing tasks
+                    future_to_chunk = {executor.submit(_process_qt_chunk_parallel, chunk_data): i 
+                                     for i, chunk_data in enumerate(chunk_data_list)}
+                    
+                    # Collect results with progress tracking
+                    completed_chunks = 0
+                    total_chunks = len(chunk_data_list)
+                    serialized_chunk_results = []
+                    
+                    for future in as_completed(future_to_chunk):
+                        chunk_idx = future_to_chunk[future]
+                        try:
+                            chunk_start_idx, consensus_features = future.result()
+                            serialized_chunk_results.append((chunk_start_idx, consensus_features))
+                            completed_chunks += 1
+                            n_samples_in_chunk = len(chunk_data_list[chunk_idx]['chunk_samples_data'])
+                            self.logger.info(f"Completed chunk {completed_chunks}/{total_chunks} (samples {chunk_start_idx + 1}-{chunk_start_idx + n_samples_in_chunk})")
+                        except Exception as exc:
+                            self.logger.error(f"Chunk {chunk_idx} generated an exception: {exc}")
+                            raise exc
+            else:
+                # Re-raise other exceptions
+                raise
         
         # Store serialized results for _merge_chunk_results to handle directly  
         chunk_consensus_maps = []
@@ -2445,6 +2678,595 @@ def _perform_adduct_grouping(self, rt_tol, mz_tol):
         )
 
 
+def _count_tight_clusters(self, mz_tol: float = 0.04, rt_tol: float = 0.3) -> int:
+    """
+    Count consensus features grouped in tight clusters.
+    
+    Args:
+        mz_tol: m/z tolerance in Daltons for cluster detection
+        rt_tol: RT tolerance in seconds for cluster detection
+        
+    Returns:
+        Number of tight clusters found
+    """
+    if len(self.consensus_df) < 2:
+        return 0
+    
+    # Extract consensus feature data
+    consensus_data = []
+    for row in self.consensus_df.iter_rows(named=True):
+        consensus_data.append({
+            'consensus_uid': row['consensus_uid'],
+            'mz': row['mz'], 
+            'rt': row['rt']
+        })
+    
+    # Build spatial index using bins
+    rt_bin_size = rt_tol / 2
+    mz_bin_size = mz_tol / 2
+    
+    bins = defaultdict(list)
+    for feature in consensus_data:
+        rt_bin = int(feature['rt'] / rt_bin_size)
+        mz_bin = int(feature['mz'] / mz_bin_size)
+        bins[(rt_bin, mz_bin)].append(feature)
+    
+    processed_features = set()
+    tight_clusters_count = 0
+    
+    for bin_key, bin_features in bins.items():
+        if len(bin_features) < 2:
+            continue
+            
+        # Check neighboring bins for additional features
+        rt_bin, mz_bin = bin_key
+        all_nearby_features = list(bin_features)
+        
+        # Check 8 neighboring bins
+        for drt in [-1, 0, 1]:
+            for dmz in [-1, 0, 1]:
+                if drt == 0 and dmz == 0:
+                    continue
+                neighbor_key = (rt_bin + drt, mz_bin + dmz)
+                if neighbor_key in bins:
+                    all_nearby_features.extend(bins[neighbor_key])
+        
+        # Filter to features within actual tolerances and not yet processed
+        valid_cluster_features = []
+        for feature in all_nearby_features:
+            if feature['consensus_uid'] in processed_features:
+                continue
+                
+            # Check if this feature is within tolerances of any bin feature
+            for bin_feature in bin_features:
+                rt_diff = abs(feature['rt'] - bin_feature['rt'])
+                mz_diff = abs(feature['mz'] - bin_feature['mz'])
+                
+                if rt_diff <= rt_tol and mz_diff <= mz_tol:
+                    valid_cluster_features.append(feature)
+                    break
+        
+        # Count as tight cluster if we have multiple features
+        if len(valid_cluster_features) >= 2:
+            tight_clusters_count += 1
+            for feature in valid_cluster_features:
+                processed_features.add(feature['consensus_uid'])
+    
+    return tight_clusters_count
+
+
+def _consensus_cleanup(self, rt_tol, mz_tol):
+    """
+    Consensus cleanup to merge over-segmented consensus features and remove isotopic features.
+    
+    This function:
+    1. Identifies and merges consensus features that are likely over-segmented 
+       (too many features in very tight m/z and RT windows)
+    2. Performs deisotoping to remove +1 and +2 isotopic features
+    """
+    if len(self.consensus_df) == 0:
+        return
+    
+    initial_count = len(self.consensus_df)
+    
+    # Only perform enhanced post-clustering if there are many features
+    if initial_count < 50:
+        return
+    
+    self.logger.debug(f"Enhanced post-clustering: processing {initial_count} consensus features")
+    
+    # Find tight clusters using spatial binning
+    consensus_data = []
+    for row in self.consensus_df.iter_rows(named=True):
+        consensus_data.append({
+            'consensus_uid': row['consensus_uid'],
+            'mz': row['mz'], 
+            'rt': row['rt'],
+            'inty_mean': row.get('inty_mean', 0),
+            'number_samples': row.get('number_samples', 0)
+        })
+    
+    # Parameters for tight clustering detection - more lenient for effective merging
+    tight_rt_tol = min(0.5, rt_tol * 0.5)  # More lenient RT tolerance (max 0.5s)
+    tight_mz_tol = min(0.05, max(0.03, mz_tol * 2.0))  # More lenient m/z tolerance (min 30 mDa, max 50 mDa)
+    
+    # Build spatial index using smaller RT and m/z bins for better coverage
+    rt_bin_size = tight_rt_tol / 4  # Smaller bins to ensure nearby features are captured
+    mz_bin_size = tight_mz_tol / 4  # Smaller bins to ensure nearby features are captured
+    
+    bins = defaultdict(list)
+    for feature in consensus_data:
+        rt_bin = int(feature['rt'] / rt_bin_size)
+        mz_bin = int(feature['mz'] / mz_bin_size)
+        bins[(rt_bin, mz_bin)].append(feature)
+    
+    # Find clusters that need merging
+    merge_groups = []
+    processed_uids = set()
+    
+    for bin_key, bin_features in bins.items():
+        # Check current bin and extended neighboring bins for complete cluster
+        rt_bin, mz_bin = bin_key
+        cluster_features = list(bin_features)
+        
+        # Check a larger neighborhood (±2 bins) to ensure we capture all nearby features
+        for dr in [-2, -1, 0, 1, 2]:
+            for dm in [-2, -1, 0, 1, 2]:
+                if dr == 0 and dm == 0:
+                    continue
+                neighbor_key = (rt_bin + dr, mz_bin + dm)
+                if neighbor_key in bins:
+                    cluster_features.extend(bins[neighbor_key])
+        
+        # Remove duplicates
+        seen_uids = set()
+        unique_features = []
+        for f in cluster_features:
+            if f['consensus_uid'] not in seen_uids:
+                unique_features.append(f)
+                seen_uids.add(f['consensus_uid'])
+        
+        # Only proceed if we have at least 2 features after including neighbors
+        if len(unique_features) < 2:
+            continue
+            
+        # Calculate cluster bounds
+        mzs = [f['mz'] for f in unique_features]
+        rts = [f['rt'] for f in unique_features]
+        
+        mz_spread = max(mzs) - min(mzs)
+        rt_spread = max(rts) - min(rts)
+        
+        # Only merge if features are tightly clustered
+        if mz_spread <= tight_mz_tol and rt_spread <= tight_rt_tol:
+            # Filter out features that were already processed
+            uids_in_cluster = {f['consensus_uid'] for f in unique_features}
+            unprocessed_features = [f for f in unique_features if f['consensus_uid'] not in processed_uids]
+            
+            # Only proceed if we have at least 2 unprocessed features that still form a tight cluster
+            if len(unprocessed_features) >= 2:
+                # Recalculate bounds for unprocessed features only
+                unprocessed_mzs = [f['mz'] for f in unprocessed_features]
+                unprocessed_rts = [f['rt'] for f in unprocessed_features]
+                
+                unprocessed_mz_spread = max(unprocessed_mzs) - min(unprocessed_mzs)
+                unprocessed_rt_spread = max(unprocessed_rts) - min(unprocessed_rts)
+                
+                # Check if unprocessed features still meet tight clustering criteria
+                if unprocessed_mz_spread <= tight_mz_tol and unprocessed_rt_spread <= tight_rt_tol:
+                    merge_groups.append(unprocessed_features)
+                    processed_uids.update({f['consensus_uid'] for f in unprocessed_features})
+    
+    if not merge_groups:
+        return
+    
+    self.logger.debug(f"Found {len(merge_groups)} over-segmented clusters to merge")
+    
+    # Merge clusters by keeping the most representative feature
+    uids_to_remove = set()
+    
+    for group in merge_groups:
+        if len(group) < 2:
+            continue
+            
+        # Find the most representative feature (highest intensity and sample count)
+        best_feature = max(group, key=lambda x: (x['number_samples'], x['inty_mean']))
+        
+        # Mark other features for removal
+        for f in group:
+            if f['consensus_uid'] != best_feature['consensus_uid']:
+                uids_to_remove.add(f['consensus_uid'])
+    
+    if uids_to_remove:
+        # Remove merged features from consensus_df
+        self.consensus_df = self.consensus_df.filter(
+            ~pl.col('consensus_uid').is_in(list(uids_to_remove))
+        )
+        
+        # Also update consensus_mapping_df if it exists
+        if hasattr(self, 'consensus_mapping_df') and not self.consensus_mapping_df.is_empty():
+            self.consensus_mapping_df = self.consensus_mapping_df.filter(
+                ~pl.col('consensus_uid').is_in(list(uids_to_remove))
+            )
+        
+        final_count = len(self.consensus_df)
+        reduction = initial_count - final_count
+        reduction_pct = (reduction / initial_count) * 100
+        
+        if reduction > 0:
+            self.logger.debug(f"Enhanced post-clustering: {initial_count} → {final_count} features ({reduction_pct:.1f}% reduction)")
+    
+    # Step 2: Deisotoping - Remove +1 and +2 isotopic consensus features
+    pre_deisotoping_count = len(self.consensus_df)
+    isotope_uids_to_remove = set()
+    
+    # Use strict tolerances for deisotoping (same as declustering)
+    deisotope_rt_tol = min(0.3, rt_tol * 0.3)  # Strict RT tolerance for isotope detection
+    deisotope_mz_tol = min(0.01, mz_tol * 0.5)  # Strict m/z tolerance for isotope detection
+    
+    # Get current consensus data for isotope detection
+    current_consensus_data = []
+    for row in self.consensus_df.iter_rows(named=True):
+        current_consensus_data.append({
+            'consensus_uid': row['consensus_uid'],
+            'mz': row['mz'], 
+            'rt': row['rt'],
+            'number_samples': row.get('number_samples', 0)
+        })
+    
+    # Sort by m/z for efficient searching
+    current_consensus_data.sort(key=lambda x: x['mz'])
+    n_current = len(current_consensus_data)
+    
+    for i in range(n_current):
+        feature_i = current_consensus_data[i]
+        
+        # Skip if already marked for removal
+        if feature_i['consensus_uid'] in isotope_uids_to_remove:
+            continue
+            
+        # Look for potential +1 and +2 isotopes (higher m/z)
+        for j in range(i + 1, n_current):
+            feature_j = current_consensus_data[j]
+            
+            # Skip if already marked for removal
+            if feature_j['consensus_uid'] in isotope_uids_to_remove:
+                continue
+                
+            mz_diff = feature_j['mz'] - feature_i['mz']
+            
+            # Break if m/z difference is too large (features are sorted by m/z)
+            if mz_diff > 2.1:  # Beyond +2 isotope range
+                break
+                
+            rt_diff = abs(feature_j['rt'] - feature_i['rt'])
+            
+            # Check for +1 isotope (C13 mass difference ≈ 1.003354 Da)
+            if (0.995 <= mz_diff <= 1.011) and rt_diff <= deisotope_rt_tol:
+                # Potential +1 isotope - should have fewer samples than main feature
+                if feature_j['number_samples'] < feature_i['number_samples']:
+                    isotope_uids_to_remove.add(feature_j['consensus_uid'])
+                    continue
+            
+            # Check for +2 isotope (2 * C13 mass difference ≈ 2.006708 Da) 
+            if (1.995 <= mz_diff <= 2.018) and rt_diff <= deisotope_rt_tol:
+                # Potential +2 isotope - should have fewer samples than main feature
+                if feature_j['number_samples'] < feature_i['number_samples']:
+                    isotope_uids_to_remove.add(feature_j['consensus_uid'])
+                    continue
+    
+    # Remove isotopic features
+    if isotope_uids_to_remove:
+        self.consensus_df = self.consensus_df.filter(
+            ~pl.col('consensus_uid').is_in(list(isotope_uids_to_remove))
+        )
+        
+        # Also update consensus_mapping_df if it exists
+        if hasattr(self, 'consensus_mapping_df') and not self.consensus_mapping_df.is_empty():
+            self.consensus_mapping_df = self.consensus_mapping_df.filter(
+                ~pl.col('consensus_uid').is_in(list(isotope_uids_to_remove))
+            )
+        
+        post_deisotoping_count = len(self.consensus_df)
+        isotope_reduction = pre_deisotoping_count - post_deisotoping_count
+        
+        if isotope_reduction > 0:
+            self.logger.debug(f"Deisotoping: {pre_deisotoping_count} → {post_deisotoping_count} features ({isotope_reduction} isotopic features removed)")
+    
+    # Final summary
+    final_count = len(self.consensus_df)
+    total_reduction = initial_count - final_count
+    if total_reduction > 0:
+        total_reduction_pct = (total_reduction / initial_count) * 100
+        self.logger.debug(f"Consensus cleanup complete: {initial_count} → {final_count} features ({total_reduction_pct:.1f}% total reduction)")
+
+
+def _identify_adduct_by_mass_shift(self, rt_tol, cached_adducts_df=None):
+    """
+    Identify coeluting consensus features by characteristic mass shifts between adducts
+    and update their adduct information accordingly.
+    
+    This function:
+    1. Generates a catalogue of mass shifts between adducts using _get_adducts()
+    2. Searches for pairs of consensus features with same RT (within strict RT tolerance)
+       and matching m/z shifts (±0.005 Da)
+    3. Updates adduct_* columns based on identified relationships
+    
+    Args:
+        rt_tol: RT tolerance in seconds (strict tolerance for coelution detection)
+        cached_adducts_df: Pre-computed adducts DataFrame for performance
+    """
+    import polars as pl
+    import numpy as np
+    from collections import defaultdict
+    
+    # Check if consensus_df exists and has features
+    if len(self.consensus_df) == 0:
+        self.logger.debug("No consensus features for adduct identification by mass shift")
+        return
+    
+    self.logger.info(f"Identifying coeluting adducts by mass shifts in {len(self.consensus_df)} consensus features...")
+    
+    # Get adducts DataFrame if not provided
+    if cached_adducts_df is None or cached_adducts_df.is_empty():
+        try:
+            # Use lower min_probability for better adduct coverage in mass shift identification
+            cached_adducts_df = self._get_adducts(min_probability=0.01)
+        except Exception as e:
+            self.logger.warning(f"Could not retrieve adducts for mass shift identification: {e}")
+            return
+    
+    if cached_adducts_df.is_empty():
+        self.logger.debug("No adducts available for mass shift identification")
+        return
+    
+    # Build catalogue of mass shifts between adducts
+    mass_shift_catalog = {}
+    adduct_info = {}
+    
+    # Extract adduct information
+    adducts_data = cached_adducts_df.select(["name", "charge", "mass_shift"]).to_dicts()
+    
+    for adduct in adducts_data:
+        name = adduct["name"]
+        charge = adduct["charge"] 
+        mass_shift = adduct["mass_shift"]
+        
+        adduct_info[name] = {
+            "charge": charge,
+            "mass_shift": mass_shift
+        }
+    
+    # Generate pairwise mass differences for catalog
+    for adduct1 in adducts_data:
+        for adduct2 in adducts_data:
+            if adduct1["name"] == adduct2["name"]:
+                continue
+                
+            name1, charge1, ms1 = adduct1["name"], adduct1["charge"], adduct1["mass_shift"]
+            name2, charge2, ms2 = adduct2["name"], adduct2["charge"], adduct2["mass_shift"]
+            
+            # Only consider shifts between adducts that have the same charge (same ionization state)
+            if charge1 != charge2:
+                continue
+            
+            # Calculate expected m/z difference
+            if charge1 != 0 and charge2 != 0:
+                mz_diff = (ms1 - ms2) / abs(charge1)
+            else:
+                continue  # Skip neutral adducts for this analysis
+            
+            # Store the mass shift relationship
+            shift_key = round(mz_diff, 4)  # Round to 4 decimal places for matching
+            if shift_key not in mass_shift_catalog:
+                mass_shift_catalog[shift_key] = []
+            mass_shift_catalog[shift_key].append({
+                "from_adduct": name1,
+                "to_adduct": name2, 
+                "mz_shift": mz_diff,
+                "from_charge": charge1,
+                "to_charge": charge2
+            })
+    
+    self.logger.debug(f"Generated mass shift catalog with {len(mass_shift_catalog)} unique shifts")
+    
+    # Get consensus features data
+    consensus_data = []
+    for i, row in enumerate(self.consensus_df.iter_rows(named=True)):
+        consensus_data.append({
+            "index": i,
+            "consensus_uid": row["consensus_uid"],
+            "rt": row["rt"], 
+            "mz": row["mz"],
+            "adduct_top": row.get("adduct_top", "[M+?]1+"),
+            "adduct_charge_top": row.get("adduct_charge_top", 1),
+            "adduct_mass_neutral_top": row.get("adduct_mass_neutral_top"),
+            "adduct_mass_shift_top": row.get("adduct_mass_shift_top"),
+            "inty_mean": row.get("inty_mean", 0)
+        })
+    
+    # Sort by RT for efficient searching
+    consensus_data.sort(key=lambda x: x["rt"])
+    n_features = len(consensus_data)
+    
+    # Track updates to make
+    adduct_updates = {}  # consensus_uid -> new_adduct_info
+    
+    # Strict RT tolerance for coelution (convert to minutes)
+    rt_tol_strict = rt_tol * 0.5  # Use half the merge tolerance for strict coelution
+    mz_tol_shift = 0.005  # ±5 mDa tolerance for mass shift matching
+    
+    # Search for coeluting pairs with characteristic mass shifts
+    updated_count = 0
+    
+    for i in range(n_features):
+        feature1 = consensus_data[i]
+        rt1 = feature1["rt"]
+        mz1 = feature1["mz"]
+        adduct1 = feature1["adduct_top"]
+        
+        # Skip if already has identified adduct (not [M+?]) - DISABLED to allow re-evaluation
+        # if adduct1 and "?" not in adduct1:
+        #     continue
+            
+        # Search for coeluting features within strict RT tolerance
+        for j in range(i + 1, n_features):
+            feature2 = consensus_data[j]
+            rt2 = feature2["rt"]
+            
+            # Break if RT difference exceeds tolerance (sorted by RT)
+            if abs(rt2 - rt1) > rt_tol_strict:
+                break
+                
+            mz2 = feature2["mz"]
+            adduct2 = feature2["adduct_top"]
+            
+            # Skip if already has identified adduct (not [M+?]) - DISABLED to allow re-evaluation
+            # if adduct2 and "?" not in adduct2:
+            #     continue
+            
+            # Calculate observed m/z difference
+            mz_diff = mz2 - mz1
+            shift_key = round(mz_diff, 4)
+            
+            # Check if this mass shift matches any known adduct relationships
+            for catalog_shift, relationships in mass_shift_catalog.items():
+                if abs(shift_key - catalog_shift) <= mz_tol_shift:
+                    # Found a matching mass shift!
+                    
+                    # Choose the best relationship based on common adducts
+                    best_rel = None
+                    best_score = 0
+                    
+                    for rel in relationships:
+                        # Prioritize common adducts ([M+H]+, [M+Na]+, [M+NH4]+)
+                        score = 0
+                        if "H]" in rel["from_adduct"]: score += 3
+                        if "Na]" in rel["from_adduct"]: score += 2  
+                        if "NH4]" in rel["from_adduct"]: score += 2
+                        if "H]" in rel["to_adduct"]: score += 3
+                        if "Na]" in rel["to_adduct"]: score += 2
+                        if "NH4]" in rel["to_adduct"]: score += 2
+                        
+                        if score > best_score:
+                            best_score = score
+                            best_rel = rel
+                    
+                    if best_rel:
+                        # Determine which feature gets which adduct based on intensity
+                        inty1 = feature1["inty_mean"]
+                        inty2 = feature2["inty_mean"] 
+                        
+                        # Assign higher intensity to [M+H]+ if possible
+                        if "H]" in best_rel["from_adduct"] and inty1 >= inty2:
+                            # Feature 1 = from_adduct, Feature 2 = to_adduct
+                            from_feature = feature1
+                            to_feature = feature2
+                            from_adduct_name = best_rel["from_adduct"]
+                            to_adduct_name = best_rel["to_adduct"]
+                        elif "H]" in best_rel["to_adduct"] and inty2 >= inty1:
+                            # Feature 2 = to_adduct (reverse), Feature 1 = from_adduct
+                            from_feature = feature2
+                            to_feature = feature1
+                            from_adduct_name = best_rel["to_adduct"]
+                            to_adduct_name = best_rel["from_adduct"]
+                        else:
+                            # Assignment based on mass shift direction
+                            # catalog_shift = (ms1 - ms2) / abs(charge1) where ms1 = from_adduct mass shift, ms2 = to_adduct mass shift
+                            # If catalog_shift > 0: from_adduct has higher m/z than to_adduct
+                            # If catalog_shift < 0: from_adduct has lower m/z than to_adduct
+                            # observed mz_diff = mz2 - mz1
+                            # If mz_diff matches catalog_shift: feature2 should get to_adduct, feature1 should get from_adduct
+                            # If mz_diff matches -catalog_shift: assignments are swapped
+                            
+                            if abs(mz_diff - catalog_shift) <= abs(mz_diff - (-catalog_shift)):
+                                # mz_diff matches catalog_shift direction
+                                from_feature = feature1
+                                to_feature = feature2
+                                from_adduct_name = best_rel["from_adduct"]
+                                to_adduct_name = best_rel["to_adduct"]
+                            else:
+                                # mz_diff matches reverse direction of catalog_shift
+                                from_feature = feature2
+                                to_feature = feature1
+                                from_adduct_name = best_rel["to_adduct"] 
+                                to_adduct_name = best_rel["from_adduct"]
+                        
+                        # Get adduct details from catalog
+                        from_adduct_info = adduct_info.get(from_adduct_name, {})
+                        to_adduct_info = adduct_info.get(to_adduct_name, {})
+                        
+                        # Calculate neutral masses
+                        from_charge = from_adduct_info.get("charge", 1)
+                        to_charge = to_adduct_info.get("charge", 1)
+                        from_mass_shift = from_adduct_info.get("mass_shift", 1.007825)
+                        to_mass_shift = to_adduct_info.get("mass_shift", 1.007825)
+                        
+                        from_neutral_mass = from_feature["mz"] * abs(from_charge) - from_mass_shift
+                        to_neutral_mass = to_feature["mz"] * abs(to_charge) - to_mass_shift
+                        
+                        # Store updates
+                        adduct_updates[from_feature["consensus_uid"]] = {
+                            "adduct_top": from_adduct_name,
+                            "adduct_charge_top": from_charge,
+                            "adduct_mass_neutral_top": from_neutral_mass,
+                            "adduct_mass_shift_top": from_mass_shift
+                        }
+                        
+                        adduct_updates[to_feature["consensus_uid"]] = {
+                            "adduct_top": to_adduct_name,
+                            "adduct_charge_top": to_charge,
+                            "adduct_mass_neutral_top": to_neutral_mass,
+                            "adduct_mass_shift_top": to_mass_shift
+                        }
+                        
+                        updated_count += 2
+                        self.logger.debug(
+                            f"Identified adduct pair: {from_adduct_name} (m/z {from_feature['mz']:.4f}) "
+                            f"<-> {to_adduct_name} (m/z {to_feature['mz']:.4f}), "
+                            f"RT {rt1:.2f}s, Δm/z {mz_diff:.4f}"
+                        )
+                        break  # Found match, no need to check other relationships
+    
+    # Apply updates to consensus_df
+    if adduct_updates:
+        # Prepare update data
+        consensus_uids = self.consensus_df["consensus_uid"].to_list()
+        
+        new_adduct_top = []
+        new_adduct_charge_top = []
+        new_adduct_mass_neutral_top = []
+        new_adduct_mass_shift_top = []
+        
+        for uid in consensus_uids:
+            if uid in adduct_updates:
+                update = adduct_updates[uid]
+                new_adduct_top.append(update["adduct_top"])
+                new_adduct_charge_top.append(update["adduct_charge_top"])
+                new_adduct_mass_neutral_top.append(update["adduct_mass_neutral_top"])
+                new_adduct_mass_shift_top.append(update["adduct_mass_shift_top"])
+            else:
+                # Keep existing values
+                row_idx = consensus_uids.index(uid)
+                row = self.consensus_df.row(row_idx, named=True)
+                new_adduct_top.append(row.get("adduct_top"))
+                new_adduct_charge_top.append(row.get("adduct_charge_top"))
+                new_adduct_mass_neutral_top.append(row.get("adduct_mass_neutral_top"))
+                new_adduct_mass_shift_top.append(row.get("adduct_mass_shift_top"))
+        
+        # Update the DataFrame
+        self.consensus_df = self.consensus_df.with_columns([
+            pl.Series("adduct_top", new_adduct_top),
+            pl.Series("adduct_charge_top", new_adduct_charge_top), 
+            pl.Series("adduct_mass_neutral_top", new_adduct_mass_neutral_top),
+            pl.Series("adduct_mass_shift_top", new_adduct_mass_shift_top)
+        ])
+        
+        self.logger.info(f"Updated adduct assignments for {updated_count} consensus features based on mass shifts")
+    else:
+        self.logger.debug("No consensus features updated based on mass shift analysis")
+
+
 def _finalize_merge(self, link_ms2, min_samples):
     """Complete the merge process with final calculations and cleanup."""
     import polars as pl
@@ -2483,14 +3305,20 @@ def _finalize_merge(self, link_ms2, min_samples):
     )
 
     # Calculate the completeness of the consensus map
+    # Log completion with tight cluster metrics
     if len(self.consensus_df) > 0 and len(self.samples_df) > 0:
         c = (
             len(self.consensus_mapping_df)
             / len(self.consensus_df)
             / len(self.samples_df)
         )
+        
+        # Count tight clusters with specified thresholds
+        tight_clusters = _count_tight_clusters(self,mz_tol=0.04, rt_tol=0.3)
+        
         self.logger.info(
-            f"Merging completed. Consensus features: {len(self.consensus_df)}. Completeness: {c:.2f}.",
+            f"Merging completed. Consensus features: {len(self.consensus_df)}. "
+            f"Completeness: {c:.2f}. Tight clusters left: {tight_clusters}.",
         )
     else:
         self.logger.warning(