masster 0.5.22__py3-none-any.whl → 0.5.24__py3-none-any.whl

masster/study/helpers.py

@@ -71,12 +71,7 @@ def get_bpc(owner, sample=None, rt_unit="s", label=None, original=False):
         # fallback to pandas
         try:
             bpc_pd = s.ms1_df.to_pandas()[["rt", "inty"]]
-            bpc_pd = (
-                bpc_pd.groupby("rt")
-                .agg({"inty": "max"})
-                .reset_index()
-                .sort_values("rt")
-            )
+            bpc_pd = bpc_pd.groupby("rt").agg({"inty": "max"}).reset_index().sort_values("rt")
         except Exception:
             raise
 
@@ -375,8 +370,7 @@ def get_chrom(self, uids=None, samples=None):
     )
     # Pre-filter features_df to only relevant features and samples
     filtered_features = self.features_df.filter(
-        pl.col("feature_uid").is_in(relevant_feature_uids)
-        & pl.col("sample_uid").is_in(sample_uids),
+        pl.col("feature_uid").is_in(relevant_feature_uids) & pl.col("sample_uid").is_in(sample_uids),
     ).select(
         [
             "feature_uid",
@@ -489,6 +483,7 @@ def align_reset(self):
 
     # Ensure column order is maintained after with_columns operation
     from masster.study.helpers import _ensure_features_df_schema_order
+
     _ensure_features_df_schema_order(self)
     self.logger.info("Alignment reset: all feature RTs set to original_RT.")
 
@@ -530,24 +525,24 @@ def get_consensus_matrix(self, quant="chrom_area", samples=None):
     """
     Get a matrix of consensus features with samples as columns and consensus features as rows.
     Highly optimized implementation using vectorized Polars operations.
-    
+
     Parameters:
         quant (str): Quantification method column name (default: "chrom_area")
         samples: Sample identifier(s) to include. Can be:
                 - None: include all samples (default)
-                - int: single sample_uid  
+                - int: single sample_uid
                 - str: single sample_name
                 - list: multiple sample_uids or sample_names
     """
     import polars as pl
-    
+
     if quant not in self.features_df.columns:
         self.logger.error(f"Quantification method {quant} not found in features_df.")
         return None
 
     # Get sample_uids to include in the matrix
     sample_uids = self._get_samples_uids(samples) if samples is not None else self.samples_df["sample_uid"].to_list()
-    
+
     if not sample_uids:
         self.logger.warning("No valid samples found for consensus matrix")
         return pl.DataFrame()
@@ -556,44 +551,31 @@ def get_consensus_matrix(self, quant="chrom_area", samples=None):
     features_filtered = self.features_df.filter(pl.col("sample_uid").is_in(sample_uids))
     samples_filtered = self.samples_df.filter(pl.col("sample_uid").is_in(sample_uids))
     consensus_mapping_filtered = self.consensus_mapping_df.filter(pl.col("sample_uid").is_in(sample_uids))
-    
+
     # Join operations to combine data efficiently
     # 1. Join consensus mapping with features to get quantification values
-    consensus_with_values = (
-        consensus_mapping_filtered
-        .join(features_filtered.select(["feature_uid", "sample_uid", quant]), 
-              on=["feature_uid", "sample_uid"], how="left")
-        .with_columns(pl.col(quant).fill_null(0))
-    )
-    
+    consensus_with_values = consensus_mapping_filtered.join(
+        features_filtered.select(["feature_uid", "sample_uid", quant]), on=["feature_uid", "sample_uid"], how="left"
+    ).with_columns(pl.col(quant).fill_null(0))
+
     # 2. Join with samples to get sample names
-    consensus_with_names = (
-        consensus_with_values
-        .join(samples_filtered.select(["sample_uid", "sample_name"]), 
-              on="sample_uid", how="left")
+    consensus_with_names = consensus_with_values.join(
+        samples_filtered.select(["sample_uid", "sample_name"]), on="sample_uid", how="left"
     )
-    
+
     # 3. Group by consensus_uid and sample_name, taking max value per group
-    aggregated = (
-        consensus_with_names
-        .group_by(["consensus_uid", "sample_name"])
-        .agg(pl.col(quant).max().alias("value"))
-    )
-    
+    aggregated = consensus_with_names.group_by(["consensus_uid", "sample_name"]).agg(pl.col(quant).max().alias("value"))
+
     # 4. Pivot to create the matrix format
-    matrix_df = (
-        aggregated
-        .pivot(on="sample_name", index="consensus_uid", values="value")
-        .fill_null(0)
-    )
-    
+    matrix_df = aggregated.pivot(on="sample_name", index="consensus_uid", values="value").fill_null(0)
+
     # 5. Round numeric columns and ensure proper types
     numeric_cols = [col for col in matrix_df.columns if col != "consensus_uid"]
     matrix_df = matrix_df.with_columns([
         pl.col("consensus_uid").cast(pl.UInt64),
-        *[pl.col(col).round(0) for col in numeric_cols]
+        *[pl.col(col).round(0) for col in numeric_cols],
     ])
-    
+
     return matrix_df
 
 
@@ -601,26 +583,26 @@ def get_gaps_matrix(self, uids=None, samples=None):
     """
     Get a matrix of gaps between consensus features with samples as columns and consensus features as rows.
     Optimized implementation that builds the gaps matrix directly without calling get_consensus_matrix().
-    
+
     Parameters:
         uids: Consensus UID(s) to include. If None, includes all consensus features.
         samples: Sample identifier(s) to include. If None, includes all samples.
                 Can be int (sample_uid), str (sample_name), or list of either.
-    
+
     Returns:
         pl.DataFrame: Gaps matrix with consensus_uid as first column and samples as other columns.
                      Values are 1 (detected) or 0 (missing/gap).
     """
     import polars as pl
-    
+
     if self.consensus_df is None or self.consensus_df.is_empty():
         self.logger.error("No consensus found.")
         return None
-    
+
     if self.consensus_mapping_df is None or self.consensus_mapping_df.is_empty():
         self.logger.error("No consensus mapping found.")
         return None
-        
+
     if self.features_df is None or self.features_df.is_empty():
         self.logger.error("No features found.")
         return None
@@ -628,7 +610,7 @@ def get_gaps_matrix(self, uids=None, samples=None):
     # Get consensus UIDs and sample UIDs to include
     uids = self._get_consensus_uids(uids)
     sample_uids = self._get_samples_uids(samples) if samples is not None else self.samples_df["sample_uid"].to_list()
-    
+
     if not uids or not sample_uids:
         self.logger.warning("No valid consensus features or samples found for gaps matrix")
         return pl.DataFrame()
@@ -642,7 +624,7 @@ def get_gaps_matrix(self, uids=None, samples=None):
             # Skip filled features (gaps should only show original detections)
             if row.get("filled", False):
                 continue
-                
+
             feature_uid = row["feature_uid"]
             # If feature exists and is not filled, it's detected (1)
             feature_detection[(feature_uid, sample_uid)] = 1
@@ -651,7 +633,8 @@ def get_gaps_matrix(self, uids=None, samples=None):
     matrix_dict = {}
     sample_mapping = dict(
         self.samples_df.filter(pl.col("sample_uid").is_in(sample_uids))
-        .select(["sample_uid", "sample_name"]).iter_rows(),
+        .select(["sample_uid", "sample_name"])
+        .iter_rows(),
     )
 
     for row in self.consensus_mapping_df.iter_rows(named=True):
@@ -732,7 +715,7 @@ def get_gaps_stats(self, uids=None):
 def get_consensus_matches(self, uids=None, filled=True):
     """
     Get feature matches for consensus UIDs with optimized join operation.
-    
+
     Parameters:
         uids: Consensus UID(s) to get matches for. Can be:
               - None: get matches for all consensus features
@@ -740,50 +723,47 @@ def get_consensus_matches(self, uids=None, filled=True):
               - list: multiple consensus UIDs
         filled (bool): Whether to include filled rows (True) or exclude them (False).
                       Default is True to maintain backward compatibility.
-    
+
     Returns:
         pl.DataFrame: Feature matches for the specified consensus UIDs
     """
     # Handle single int by converting to list
     if isinstance(uids, int):
         uids = [uids]
-    
+
     uids = self._get_consensus_uids(uids)
-    
+
     if not uids:
         return pl.DataFrame()
-    
+
     # Early validation checks
     if self.consensus_mapping_df is None or self.consensus_mapping_df.is_empty():
         self.logger.warning("No consensus mapping data available")
         return pl.DataFrame()
-    
+
     if self.features_df is None or self.features_df.is_empty():
         self.logger.warning("No feature data available")
         return pl.DataFrame()
-    
+
     # Build the query with optional filled filter
     features_query = self.features_df.lazy()
-    
+
     # Apply filled filter if specified
     if not filled and "filled" in self.features_df.columns:
         features_query = features_query.filter(~pl.col("filled"))
-    
+
     # Optimized single-pass operation using join instead of two separate filters
     # This avoids creating intermediate Python lists and leverages Polars' optimized joins
     matches = (
-        features_query
-        .join(
-            self.consensus_mapping_df
-            .lazy()
+        features_query.join(
+            self.consensus_mapping_df.lazy()
             .filter(pl.col("consensus_uid").is_in(uids))
             .select("feature_uid"),  # Only select what we need for the join
             on="feature_uid",
-            how="inner"
-        )
-        .collect(streaming=True)  # Use streaming for memory efficiency with large datasets
+            how="inner",
+        ).collect(streaming=True)  # Use streaming for memory efficiency with large datasets
     )
-    
+
     return matches
 
 
@@ -795,34 +775,34 @@ def get_consensus_matches(self, uids=None, filled=True):
 def consensus_reset(self):
     """
     Reset consensus data by clearing consensus DataFrames and removing filled features.
-    
+
     This function:
     1. Sets consensus_df, consensus_ms2, consensus_mapping_df, id_df to empty pl.DataFrame()
     2. Removes all filled features from features_df
     3. Removes relevant operations from history (merge, integrate, find_ms2, fill, identify)
     4. Logs the number of features removed
-    
+
     This effectively undoes the merge() operation and any gap-filling.
     """
     self.logger.debug("Resetting consensus data.")
-    
+
     # Reset consensus DataFrames to empty
     self.consensus_df = pl.DataFrame()
-    self.consensus_ms2 = pl.DataFrame()  
+    self.consensus_ms2 = pl.DataFrame()
     self.consensus_mapping_df = pl.DataFrame()
     self.id_df = pl.DataFrame()
-    
+
     # Remove filled features from features_df
     if self.features_df is None:
         self.logger.warning("No features found.")
         return
-        
+
     l1 = len(self.features_df)
-    
+
     # Filter out filled features (keep only non-filled features)
     if "filled" in self.features_df.columns:
         self.features_df = self.features_df.filter(~pl.col("filled") | pl.col("filled").is_null())
-    
+
     # Remove consensus-related operations from history
     keys_to_remove = ["merge", "integrate", "integrate_chrom", "find_ms2", "fill", "fill_single", "identify"]
     history_removed_count = 0
@@ -832,7 +812,7 @@ def consensus_reset(self):
                 del self.history[key]
                 history_removed_count += 1
                 self.logger.debug(f"Removed '{key}' from history")
-    
+
     removed_count = l1 - len(self.features_df)
     self.logger.info(
         f"Reset consensus data. Consensus DataFrames cleared. Features removed: {removed_count}. History entries removed: {history_removed_count}",
@@ -1049,13 +1029,13 @@ def get_orphans(self):
 def get_sample_stats(self):
     """
     Get statistics for all samples in the study.
-    
+
     Returns:
         pl.DataFrame: DataFrame with the following columns:
             - sample_uid: Sample unique identifier
             - num_features: Total number of features per sample
             - num_ms1: Number of MS1 features per sample
-            - num_ms2: Number of MS2 features per sample  
+            - num_ms2: Number of MS2 features per sample
             - num_linked_ms1: Number of non-filled features present in consensus_mapping_df
             - num_orphans: Number of non-filled features not present in consensus_mapping_df
             - max_rt_correction: Maximum RT correction applied
@@ -1065,19 +1045,19 @@ def get_sample_stats(self):
     if self.samples_df is None or self.samples_df.is_empty():
         self.logger.warning("No samples found in study.")
         return pl.DataFrame()
-    
+
     if self.features_df is None or self.features_df.is_empty():
         self.logger.warning("No features found in study.")
         return pl.DataFrame()
-    
+
     # Get base sample information
     sample_uids = self.samples_df["sample_uid"].to_list()
     stats_data = []
-    
+
     for sample_uid in sample_uids:
         # Filter features for this sample
         sample_features = self.features_df.filter(pl.col("sample_uid") == sample_uid)
-        
+
         if sample_features.is_empty():
             # Sample has no features
             stats_data.append({
@@ -1089,66 +1069,60 @@ def get_sample_stats(self):
                 "num_orphans": 0,
                 "max_rt_correction": None,
                 "average_rt_correction": None,
-                "num_linked_ms2": 0
+                "num_linked_ms2": 0,
             })
             continue
-        
+
         # Basic feature counts
         num_features = len(sample_features)
-        
+
         # Count MS1 and MS2 features
         # Assume features with ms_level=1 or missing ms_level are MS1
-        num_ms1 = sample_features.filter(
-            pl.col("ms_level").is_null() | (pl.col("ms_level") == 1)
-        ).height if "ms_level" in sample_features.columns else num_features
-        
-        num_ms2 = sample_features.filter(
-            pl.col("ms_level") == 2
-        ).height if "ms_level" in sample_features.columns else 0
-        
+        num_ms1 = (
+            sample_features.filter(pl.col("ms_level").is_null() | (pl.col("ms_level") == 1)).height
+            if "ms_level" in sample_features.columns
+            else num_features
+        )
+
+        num_ms2 = sample_features.filter(pl.col("ms_level") == 2).height if "ms_level" in sample_features.columns else 0
+
         # Get non-filled features for this sample
         if "filled" in sample_features.columns:
             non_filled_features = sample_features.filter(~pl.col("filled") | pl.col("filled").is_null())
         else:
             non_filled_features = sample_features
-        
+
         # Count linked MS1 features (non-filled and present in consensus_mapping_df)
         num_linked_ms1 = 0
         if not self.consensus_mapping_df.is_empty() and not non_filled_features.is_empty():
-            linked_feature_uids = self.consensus_mapping_df.filter(
-                pl.col("sample_uid") == sample_uid
-            )["feature_uid"].to_list()
-            
-            num_linked_ms1 = non_filled_features.filter(
-                pl.col("feature_uid").is_in(linked_feature_uids)
-            ).height
-        
+            linked_feature_uids = self.consensus_mapping_df.filter(pl.col("sample_uid") == sample_uid)[
+                "feature_uid"
+            ].to_list()
+
+            num_linked_ms1 = non_filled_features.filter(pl.col("feature_uid").is_in(linked_feature_uids)).height
+
         # Count orphan features (non-filled and NOT present in consensus_mapping_df)
         num_orphans = len(non_filled_features) - num_linked_ms1
-        
+
         # Calculate RT correction statistics
         max_rt_correction = None
         average_rt_correction = None
-        
+
         if "rt" in sample_features.columns and "rt_original" in sample_features.columns:
             rt_corrections = sample_features.with_columns(
                 (pl.col("rt") - pl.col("rt_original")).alias("rt_correction")
-            ).filter(
-                pl.col("rt_correction").is_not_null()
-            )["rt_correction"]
-            
+            ).filter(pl.col("rt_correction").is_not_null())["rt_correction"]
+
             if not rt_corrections.is_empty():
                 max_rt_correction = rt_corrections.abs().max()
                 average_rt_correction = rt_corrections.abs().mean()
-        
+
         # Count linked MS2 spectra from consensus_ms2_df
         num_linked_ms2 = 0
-        if hasattr(self, 'consensus_ms2') and self.consensus_ms2 is not None and not self.consensus_ms2.is_empty():
+        if hasattr(self, "consensus_ms2") and self.consensus_ms2 is not None and not self.consensus_ms2.is_empty():
             if "sample_uid" in self.consensus_ms2.columns:
-                num_linked_ms2 = self.consensus_ms2.filter(
-                    pl.col("sample_uid") == sample_uid
-                ).height
-        
+                num_linked_ms2 = self.consensus_ms2.filter(pl.col("sample_uid") == sample_uid).height
+
         stats_data.append({
             "sample_uid": sample_uid,
             "num_features": num_features,
@@ -1158,9 +1132,9 @@ def get_sample_stats(self):
             "num_orphans": num_orphans,
             "max_rt_correction": max_rt_correction,
             "average_rt_correction": average_rt_correction,
-            "num_linked_ms2": num_linked_ms2
+            "num_linked_ms2": num_linked_ms2,
         })
-    
+
     # Create DataFrame with proper schema
     return pl.DataFrame(
         stats_data,
@@ -1173,15 +1147,15 @@ def get_sample_stats(self):
             "num_orphans": pl.UInt32,
             "max_rt_correction": pl.Float64,
             "average_rt_correction": pl.Float64,
-            "num_linked_ms2": pl.UInt32
-        }
+            "num_linked_ms2": pl.UInt32,
+        },
     )
 
 
 def get_consensus_stats(self):
     """
     Get key performance indicators for each consensus feature.
-    
+
     Returns:
         pl.DataFrame: DataFrame with the following columns:
             - consensus_uid: Consensus unique identifier
@@ -1203,7 +1177,7 @@ def get_consensus_stats(self):
     """
     import polars as pl
     import numpy as np
-    
+
     # Check if consensus_df exists and has data
     if self.consensus_df is None or self.consensus_df.is_empty():
         self.logger.error("No consensus data available. Run merge/find_consensus first.")
@@ -1215,134 +1189,146 @@ def get_consensus_stats(self):
     # Define specific columns to include in the exact order requested
     desired_columns = [
         "consensus_uid",  # Include consensus_uid for identification
-        "rt", 
-        "rt_delta_mean", 
-        "mz", 
+        "rt",
+        "rt_delta_mean",
+        "mz",
         "mz_range",  # mz_max-mz_min (will be calculated)
         "log10_inty_mean",  # log10(inty_mean) (will be calculated)
-        "number_samples", 
-        "number_ms2", 
-        "charge_mean", 
-        "quality", 
-        "chrom_coherence_mean", 
-        "chrom_height_scaled_mean", 
-        "chrom_prominence_scaled_mean"
+        "number_samples",
+        "number_ms2",
+        "charge_mean",
+        "quality",
+        "chrom_coherence_mean",
+        "chrom_height_scaled_mean",
+        "chrom_prominence_scaled_mean",
     ]
-    
+
     # Calculate derived columns if they don't exist
     if "mz_range" not in data_df.columns and "mz_max" in data_df.columns and "mz_min" in data_df.columns:
         data_df = data_df.with_columns((pl.col("mz_max") - pl.col("mz_min")).alias("mz_range"))
-    
+
     if "log10_inty_mean" not in data_df.columns and "inty_mean" in data_df.columns:
         data_df = data_df.with_columns(pl.col("inty_mean").log10().alias("log10_inty_mean"))
-    
+
     # Filter to only include columns that exist in the dataframe, preserving order
     available_columns = [col for col in desired_columns if col in data_df.columns]
-    
+
     if len(available_columns) <= 1:  # Only consensus_uid would be 1
-        self.logger.error(f"None of the requested consensus statistics columns were found. Available columns: {list(data_df.columns)}")
+        self.logger.error(
+            f"None of the requested consensus statistics columns were found. Available columns: {list(data_df.columns)}"
+        )
         return pl.DataFrame()
 
     self.logger.debug(f"Creating consensus stats DataFrame with {len(available_columns)} columns: {available_columns}")
 
     # Get base result DataFrame with selected columns
     result_df = data_df.select(available_columns)
-    
+
     # Add QC-related columns
     try:
         # Identify QC and blank samples based on naming patterns
         all_sample_names = self.samples_df["sample_name"].to_list()
-        
+
         # Define patterns for QC and blank identification
         qc_patterns = ["qc", "QC", "quality", "Quality", "control", "Control"]
         blank_patterns = ["blank", "Blank", "BLANK", "blk", "BLK"]
-        
+
         # Get QC and blank sample names
         qc_sample_names = [name for name in all_sample_names if any(pattern in name for pattern in qc_patterns)]
         blank_sample_names = [name for name in all_sample_names if any(pattern in name for pattern in blank_patterns)]
-        
+
         self.logger.debug(f"Found {len(qc_sample_names)} QC samples and {len(blank_sample_names)} blank samples")
-        
+
         # Initialize QC columns with null values
         qc_ratio_values = [None] * len(result_df)
-        qc_cv_values = [None] * len(result_df) 
+        qc_cv_values = [None] * len(result_df)
         qc_to_blank_values = [None] * len(result_df)
-        
+
         if len(qc_sample_names) > 0:
             # Calculate QC metrics using optimized approach - get only QC+blank data
             self.logger.debug("Fetching optimized consensus matrices for QC calculations...")
-            
+
             # Get QC consensus matrix (only QC samples)
             qc_consensus_matrix = self.get_consensus_matrix(samples=qc_sample_names)
-            
+
             # Get blank consensus matrix (only blank samples) if blanks exist
             blank_consensus_matrix = None
             if len(blank_sample_names) > 0:
                 blank_consensus_matrix = self.get_consensus_matrix(samples=blank_sample_names)
-            
+
             if qc_consensus_matrix is not None and not qc_consensus_matrix.is_empty():
                 available_qc_cols = [col for col in qc_consensus_matrix.columns if col != "consensus_uid"]
                 self.logger.debug(f"Found {len(available_qc_cols)} QC columns in optimized QC matrix")
-                
+
                 # 2. QC CV: Calculate CV for QC samples
                 if len(available_qc_cols) > 0:
                     self.logger.debug("Calculating QC CV...")
                     try:
                         # Calculate CV (coefficient of variation) for QC samples
                         qc_data = qc_consensus_matrix.select(["consensus_uid"] + available_qc_cols)
-                        
+
                         # Calculate mean and std for each row across QC columns
-                        qc_stats = qc_data.with_columns([
-                            pl.concat_list([pl.col(col) for col in available_qc_cols]).alias("qc_values")
-                        ]).with_columns([
-                            pl.col("qc_values").list.mean().alias("qc_mean"),
-                            pl.col("qc_values").list.std().alias("qc_std")
-                        ]).with_columns(
-                            # CV = std / mean (NOT multiplied by 100 to keep between 0-1)
-                            pl.when(pl.col("qc_mean") > 0)
-                            .then(pl.col("qc_std") / pl.col("qc_mean"))
-                            .otherwise(None)
-                            .alias("qc_cv")
+                        qc_stats = (
+                            qc_data.with_columns([
+                                pl.concat_list([pl.col(col) for col in available_qc_cols]).alias("qc_values")
+                            ])
+                            .with_columns([
+                                pl.col("qc_values").list.mean().alias("qc_mean"),
+                                pl.col("qc_values").list.std().alias("qc_std"),
+                            ])
+                            .with_columns(
+                                # CV = std / mean (NOT multiplied by 100 to keep between 0-1)
+                                pl.when(pl.col("qc_mean") > 0)
+                                .then(pl.col("qc_std") / pl.col("qc_mean"))
+                                .otherwise(None)
+                                .alias("qc_cv")
+                            )
                         )
-                        
+
                         # Join with result DataFrame
                         result_df = result_df.join(
-                            qc_stats.select(["consensus_uid", "qc_cv"]), 
-                            on="consensus_uid", 
-                            how="left"
+                            qc_stats.select(["consensus_uid", "qc_cv"]), on="consensus_uid", how="left"
                         )
                         qc_cv_values = None  # Indicate we successfully added the column
-                        
+
                     except Exception as e:
                         self.logger.debug(f"Could not calculate QC CV: {e}")
-                
+
                 # 3. QC to blank ratio: Compare average QC to average blank intensity
-                if len(available_qc_cols) > 0 and blank_consensus_matrix is not None and not blank_consensus_matrix.is_empty():
+                if (
+                    len(available_qc_cols) > 0
+                    and blank_consensus_matrix is not None
+                    and not blank_consensus_matrix.is_empty()
+                ):
                     available_blank_cols = [col for col in blank_consensus_matrix.columns if col != "consensus_uid"]
-                    self.logger.debug(f"Calculating QC to blank ratio with {len(available_blank_cols)} blank columns...")
-                    
+                    self.logger.debug(
+                        f"Calculating QC to blank ratio with {len(available_blank_cols)} blank columns..."
+                    )
+
                     if len(available_blank_cols) > 0:
                         try:
                             # Calculate average intensity for QC samples
-                            qc_averages = qc_data.with_columns([
-                                pl.concat_list([pl.col(col) for col in available_qc_cols]).alias("qc_values")
-                            ]).with_columns(
-                                pl.col("qc_values").list.mean().alias("qc_avg")
-                            ).select(["consensus_uid", "qc_avg"])
-                            
+                            qc_averages = (
+                                qc_data.with_columns([
+                                    pl.concat_list([pl.col(col) for col in available_qc_cols]).alias("qc_values")
+                                ])
+                                .with_columns(pl.col("qc_values").list.mean().alias("qc_avg"))
+                                .select(["consensus_uid", "qc_avg"])
+                            )
+
                             # Calculate average intensity for blank samples
                             blank_data = blank_consensus_matrix.select(["consensus_uid"] + available_blank_cols)
-                            blank_averages = blank_data.with_columns([
-                                pl.concat_list([pl.col(col) for col in available_blank_cols]).alias("blank_values")
-                            ]).with_columns(
-                                pl.col("blank_values").list.mean().alias("blank_avg")
-                            ).select(["consensus_uid", "blank_avg"])
-                            
+                            blank_averages = (
+                                blank_data.with_columns([
+                                    pl.concat_list([pl.col(col) for col in available_blank_cols]).alias("blank_values")
+                                ])
+                                .with_columns(pl.col("blank_values").list.mean().alias("blank_avg"))
+                                .select(["consensus_uid", "blank_avg"])
+                            )
+
                             # Join QC and blank averages and calculate ratio
                             qc_blank_ratios = qc_averages.join(
-                                blank_averages, 
-                                on="consensus_uid", 
-                                how="left"
+                                blank_averages, on="consensus_uid", how="left"
                             ).with_columns(
                                 # Ratio = qc_avg / blank_avg, but only where blank_avg > 0
                                 pl.when(pl.col("blank_avg") > 0)
@@ -1350,57 +1336,53 @@ def get_consensus_stats(self):
                                 .otherwise(None)
                                 .alias("qc_to_blank")
                             )
-                            
+
                             # Join with result DataFrame
                             result_df = result_df.join(
-                                qc_blank_ratios.select(["consensus_uid", "qc_to_blank"]), 
-                                on="consensus_uid", 
-                                how="left"
+                                qc_blank_ratios.select(["consensus_uid", "qc_to_blank"]), on="consensus_uid", how="left"
                             )
                             qc_to_blank_values = None  # Indicate we successfully added the column
-                            
+
                         except Exception as e:
                             self.logger.debug(f"Could not calculate QC to blank ratio: {e}")
-            
+
             # 1. QC ratio: Get optimized gaps matrix for QC samples only
             self.logger.debug("Calculating QC detection ratio with optimized gaps matrix...")
             try:
                 # Use optimized get_gaps_matrix with QC samples filtering for faster performance
                 qc_gaps_matrix = self.get_gaps_matrix(samples=qc_sample_names)
-                
+
                 if qc_gaps_matrix is not None and not qc_gaps_matrix.is_empty():
                     # Get QC columns (should be all columns except consensus_uid since we filtered)
                     available_qc_cols_gaps = [col for col in qc_gaps_matrix.columns if col != "consensus_uid"]
                     self.logger.debug(f"Found {len(available_qc_cols_gaps)} QC columns in optimized gaps matrix")
-                    
+
                     if len(available_qc_cols_gaps) > 0:
                         # Calculate QC detection ratio for each consensus feature
                         qc_detection = qc_gaps_matrix.select(["consensus_uid"] + available_qc_cols_gaps)
-                        
+
                         # Data should already be properly typed from get_gaps_matrix, but ensure consistency
                         for col in available_qc_cols_gaps:
-                            qc_detection = qc_detection.with_columns(
-                                pl.col(col).fill_null(0).cast(pl.Int8).alias(col)
-                            )
-                        
+                            qc_detection = qc_detection.with_columns(pl.col(col).fill_null(0).cast(pl.Int8).alias(col))
+
                         # Calculate ratio (sum of detections / number of QC samples)
                         qc_ratios = qc_detection.with_columns(
                             pl.concat_list([pl.col(col) for col in available_qc_cols_gaps]).alias("qc_detections")
                         ).with_columns(
-                            (pl.col("qc_detections").list.sum().cast(pl.Float64) / len(available_qc_cols_gaps)).alias("qc_ratio")
+                            (pl.col("qc_detections").list.sum().cast(pl.Float64) / len(available_qc_cols_gaps)).alias(
+                                "qc_ratio"
+                            )
                         )
-                        
+
                         # Join with result DataFrame
                         result_df = result_df.join(
-                            qc_ratios.select(["consensus_uid", "qc_ratio"]), 
-                            on="consensus_uid", 
-                            how="left"
+                            qc_ratios.select(["consensus_uid", "qc_ratio"]), on="consensus_uid", how="left"
                         )
                         qc_ratio_values = None  # Indicate we successfully added the column
-                        
+
             except Exception as e:
                 self.logger.debug(f"Could not calculate QC ratio: {e}")
-        
+
         # Add null columns for any QC metrics that couldn't be calculated
         # Add null columns for any QC metrics that couldn't be calculated
         if qc_ratio_values is not None:
@@ -1409,16 +1391,16 @@ def get_consensus_stats(self):
             result_df = result_df.with_columns(pl.lit(None, dtype=pl.Float64).alias("qc_cv"))
         if qc_to_blank_values is not None:
             result_df = result_df.with_columns(pl.lit(None, dtype=pl.Float64).alias("qc_to_blank"))
-            
+
     except Exception as e:
         self.logger.warning(f"Error calculating QC metrics: {e}")
         # Add null columns if QC calculation fails
         result_df = result_df.with_columns([
             pl.lit(None, dtype=pl.Float64).alias("qc_ratio"),
             pl.lit(None, dtype=pl.Float64).alias("qc_cv"),
-            pl.lit(None, dtype=pl.Float64).alias("qc_to_blank")
+            pl.lit(None, dtype=pl.Float64).alias("qc_to_blank"),
         ])
-    
+
     return result_df
 
 
@@ -1565,9 +1547,7 @@ def restore_features(self, samples=None, maps=False):
                 continue
 
             # Check which columns are actually available in the sample
-            available_columns = [
-                col for col in columns_to_update if col in sample.features_df.columns
-            ]
+            available_columns = [col for col in columns_to_update if col in sample.features_df.columns]
             if not available_columns:
                 self.logger.debug(f"No target columns found in sample {sample_name}")
                 continue
@@ -1590,9 +1570,7 @@ def restore_features(self, samples=None, maps=False):
                             original_dtype = self.features_df[col].dtype
 
                             # Update the specific row and column, preserving dtype
-                            mask = (pl.col("feature_uid") == feature_uid) & (
-                                pl.col("sample_uid") == sample_uid
-                            )
+                            mask = (pl.col("feature_uid") == feature_uid) & (pl.col("sample_uid") == sample_uid)
 
                             # Handle object columns (like Chromatogram) differently
                             if original_dtype == pl.Object:
@@ -1730,9 +1708,7 @@ def restore_chrom(self, samples=None, mz_tol=0.010, rt_tol=10.0):
                     feature_uid = study_feature_mapping[key]
 
                     # Update only the chrom column
-                    mask = (pl.col("feature_uid") == feature_uid) & (
-                        pl.col("sample_uid") == sample_uid
-                    )
+                    mask = (pl.col("feature_uid") == feature_uid) & (pl.col("sample_uid") == sample_uid)
                     self.features_df = self.features_df.with_columns(
                         pl.when(mask)
                         .then(pl.lit(chrom, dtype=pl.Object, allow_object=True))
@@ -1807,11 +1783,7 @@ def restore_chrom(self, samples=None, mz_tol=0.010, rt_tol=10.0):
             sample = Sample(log_level="ERROR")
             sample._load_sample5(sample_path, map=False)
 
-            if (
-                not hasattr(sample, "ms1_df")
-                or sample.ms1_df is None
-                or sample.ms1_df.is_empty()
-            ):
+            if not hasattr(sample, "ms1_df") or sample.ms1_df is None or sample.ms1_df.is_empty():
                 continue
 
             # Process each missing feature
@@ -1920,9 +1892,7 @@ def compress_ms2(self, max_replicates=5):
     # Handle None values by treating them as 0
     self.consensus_ms2 = self.consensus_ms2.with_columns(
         [
-            (
-                pl.col("number_frags").fill_null(0) * pl.col("prec_inty").fill_null(0)
-            ).alias("ranking_score"),
+            (pl.col("number_frags").fill_null(0) * pl.col("prec_inty").fill_null(0)).alias("ranking_score"),
         ],
     )
 
@@ -2259,57 +2229,86 @@ def features_select(
         return pl.DataFrame()
 
     # Early return optimization
-    filter_params = [mz, rt, inty, sample_uid, sample_name, consensus_uid, 
-                    feature_uid, filled, quality, chrom_coherence, 
-                    chrom_prominence, chrom_prominence_scaled, chrom_height_scaled]
-    
+    filter_params = [
+        mz,
+        rt,
+        inty,
+        sample_uid,
+        sample_name,
+        consensus_uid,
+        feature_uid,
+        filled,
+        quality,
+        chrom_coherence,
+        chrom_prominence,
+        chrom_prominence_scaled,
+        chrom_height_scaled,
+    ]
+
     if all(param is None for param in filter_params):
         return self.features_df.clone()
 
     import time
+
     start_time = time.perf_counter()
     initial_count = len(self.features_df)
 
     # Build optimized filter expression
     filter_expr = _build_optimized_filter_expression(
-        self, mz, rt, inty, sample_uid, sample_name, consensus_uid,
-        feature_uid, filled, quality, chrom_coherence,
-        chrom_prominence, chrom_prominence_scaled, chrom_height_scaled
+        self,
+        mz,
+        rt,
+        inty,
+        sample_uid,
+        sample_name,
+        consensus_uid,
+        feature_uid,
+        filled,
+        quality,
+        chrom_coherence,
+        chrom_prominence,
+        chrom_prominence_scaled,
+        chrom_height_scaled,
     )
-    
+
     if filter_expr is None:
         return pl.DataFrame()
-    
+
     # Apply filter with optimized execution strategy
     if use_lazy_streaming and initial_count > chunk_size:
         result = _apply_chunked_select(self, filter_expr, chunk_size)
     else:
-        result = (
-            self.features_df
-            .lazy()
-            .filter(filter_expr)
-            .collect(streaming=use_lazy_streaming)
-        )
-    
+        result = self.features_df.lazy().filter(filter_expr).collect(streaming=use_lazy_streaming)
+
     # Log performance
     elapsed_time = time.perf_counter() - start_time
     final_count = len(result)
     removed_count = initial_count - final_count
-    
+
     if final_count == 0:
         self.logger.warning("No features remaining after applying selection criteria.")
     else:
-        self.logger.debug(
-            f"Selected features: {final_count:,} (removed: {removed_count:,}) in {elapsed_time:.4f}s"
-        )
-    
+        self.logger.debug(f"Selected features: {final_count:,} (removed: {removed_count:,}) in {elapsed_time:.4f}s")
+
     return result
 
 
-def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_name, 
-                                     consensus_uid, feature_uid, filled, quality,
-                                     chrom_coherence, chrom_prominence, 
-                                     chrom_prominence_scaled, chrom_height_scaled):
+def _build_optimized_filter_expression(
+    self,
+    mz,
+    rt,
+    inty,
+    sample_uid,
+    sample_name,
+    consensus_uid,
+    feature_uid,
+    filled,
+    quality,
+    chrom_coherence,
+    chrom_prominence,
+    chrom_prominence_scaled,
+    chrom_height_scaled,
+):
     """
     Build optimized filter expression with efficient column checking and expression combining.
     """
@@ -2317,7 +2316,7 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
     available_columns = set(self.features_df.columns)
     filter_conditions = []
     warnings = []
-    
+
     # Build filter conditions with optimized expressions
     if mz is not None:
         if isinstance(mz, tuple) and len(mz) == 2:
@@ -2422,7 +2421,9 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
         if "chrom_coherence" in available_columns:
             if isinstance(chrom_coherence, tuple) and len(chrom_coherence) == 2:
                 min_coherence, max_coherence = chrom_coherence
-                filter_conditions.append(pl.col("chrom_coherence").is_between(min_coherence, max_coherence, closed="both"))
+                filter_conditions.append(
+                    pl.col("chrom_coherence").is_between(min_coherence, max_coherence, closed="both")
+                )
             else:
                 filter_conditions.append(pl.col("chrom_coherence") >= chrom_coherence)
         else:
@@ -2433,7 +2434,9 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
         if "chrom_prominence" in available_columns:
             if isinstance(chrom_prominence, tuple) and len(chrom_prominence) == 2:
                 min_prominence, max_prominence = chrom_prominence
-                filter_conditions.append(pl.col("chrom_prominence").is_between(min_prominence, max_prominence, closed="both"))
+                filter_conditions.append(
+                    pl.col("chrom_prominence").is_between(min_prominence, max_prominence, closed="both")
+                )
             else:
                 filter_conditions.append(pl.col("chrom_prominence") >= chrom_prominence)
         else:
@@ -2445,7 +2448,10 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
             if isinstance(chrom_prominence_scaled, tuple) and len(chrom_prominence_scaled) == 2:
                 min_prominence_scaled, max_prominence_scaled = chrom_prominence_scaled
                 filter_conditions.append(
-                    pl.col("chrom_prominence_scaled").is_between(min_prominence_scaled, max_prominence_scaled, closed="both"))
+                    pl.col("chrom_prominence_scaled").is_between(
+                        min_prominence_scaled, max_prominence_scaled, closed="both"
+                    )
+                )
             else:
                 filter_conditions.append(pl.col("chrom_prominence_scaled") >= chrom_prominence_scaled)
         else:
@@ -2457,7 +2463,8 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
             if isinstance(chrom_height_scaled, tuple) and len(chrom_height_scaled) == 2:
                 min_height_scaled, max_height_scaled = chrom_height_scaled
                 filter_conditions.append(
-                    pl.col("chrom_height_scaled").is_between(min_height_scaled, max_height_scaled, closed="both"))
+                    pl.col("chrom_height_scaled").is_between(min_height_scaled, max_height_scaled, closed="both")
+                )
             else:
                 filter_conditions.append(pl.col("chrom_height_scaled") >= chrom_height_scaled)
         else:
@@ -2470,12 +2477,13 @@ def _build_optimized_filter_expression(self, mz, rt, inty, sample_uid, sample_na
     # Combine all conditions efficiently
     if not filter_conditions:
         return None
-    
+
     # Use reduce for efficient expression combination
     from functools import reduce
     import operator
+
     combined_expr = reduce(operator.and_, filter_conditions)
-    
+
     return combined_expr
 
 
@@ -2485,30 +2493,27 @@ def _apply_chunked_select(self, filter_expr, chunk_size: int):
     """
     total_features = len(self.features_df)
     num_chunks = (total_features + chunk_size - 1) // chunk_size
-    
+
     self.logger.debug(f"Using chunked select with {num_chunks} chunks")
-    
+
     filtered_chunks = []
     for i in range(num_chunks):
         start_idx = i * chunk_size
         end_idx = min((i + 1) * chunk_size, total_features)
-        
+
         chunk_result = (
-            self.features_df
-            .lazy()
-            .slice(start_idx, end_idx - start_idx)
-            .filter(filter_expr)
-            .collect(streaming=True)
+            self.features_df.lazy().slice(start_idx, end_idx - start_idx).filter(filter_expr).collect(streaming=True)
         )
-        
+
         if not chunk_result.is_empty():
             filtered_chunks.append(chunk_result)
-    
+
     if filtered_chunks:
         return pl.concat(filtered_chunks, how="vertical")
     else:
         return pl.DataFrame()
 
+
 '''
 def features_select_benchmarked(
     self,
@@ -2604,19 +2609,14 @@ def monkey_patch_study():
     print("Patched Study.features_select with consolidated optimized implementation")
 '''
 
-def features_filter(
-    self, 
-    features,
-    chunk_size: int = 50000,
-    use_index_based: bool = True,
-    parallel: bool = True
-):
+
+def features_filter(self, features, chunk_size: int = 50000, use_index_based: bool = True, parallel: bool = True):
     """
     Filter features_df by keeping only features that match the given criteria.
     This keeps only the specified features and removes all others.
 
     FULLY OPTIMIZED VERSION: Index-based filtering, chunked processing, and lazy evaluation.
-    
+
     Performance improvements:
     - Index-based filtering using sorted arrays (O(n log n) instead of O(n²))
     - Chunked processing to handle large datasets without memory issues
@@ -2646,26 +2646,24 @@ def features_filter(
         return
 
     initial_count = len(self.features_df)
-    
+
     # Extract feature UIDs efficiently
     feature_uids_to_keep = _extract_feature_uids_optimized(self, features)
     if not feature_uids_to_keep:
         self.logger.warning("No feature UIDs provided for filtering.")
         return
-    
+
     # Choose optimal filtering strategy based on data size and characteristics
     if use_index_based and len(self.features_df) > 10000:
         _apply_index_based_filter(self, feature_uids_to_keep, chunk_size, parallel)
     else:
         _apply_standard_filter(self, feature_uids_to_keep)
-    
+
     # Calculate results and log performance
     final_count = len(self.features_df)
     removed_count = initial_count - final_count
-    
-    self.logger.info(
-        f"Filtered features. Kept: {final_count:,}. Removed: {removed_count:,}."
-    )
+
+    self.logger.info(f"Filtered features. Kept: {final_count:,}. Removed: {removed_count:,}.")
 
 
 def _extract_feature_uids_optimized(self, features):
@@ -2679,13 +2677,13 @@ def _extract_feature_uids_optimized(self, features):
             return set()
         # Use polars native operations for efficiency
         return set(features.select("feature_uid").to_series().to_list())
-    
+
     elif isinstance(features, (list, tuple)):
         return set(features)  # Convert to set immediately for O(1) lookups
-    
+
     elif isinstance(features, int):
         return {features}
-    
+
     else:
         self.logger.error("features parameter must be a DataFrame, list, tuple, or int")
         return set()
@@ -2694,7 +2692,7 @@ def _extract_feature_uids_optimized(self, features):
 def _apply_index_based_filter(self, feature_uids_to_keep, chunk_size: int, parallel: bool):
     """
     Apply index-based filtering with chunked processing and lazy evaluation.
-    
+
     This method uses:
     1. Sorted arrays and binary search for O(log n) lookups
     2. Chunked processing to manage memory usage
@@ -2702,9 +2700,9 @@ def _apply_index_based_filter(self, feature_uids_to_keep, chunk_size: int, paral
     4. Hash-based set operations for optimal performance
     """
     self.logger.debug(f"Using index-based filtering with chunks of {chunk_size:,}")
-    
+
     total_features = len(self.features_df)
-    
+
     if total_features <= chunk_size:
         # Small dataset - process in single chunk with optimized operations
         _filter_single_chunk_optimized(self, feature_uids_to_keep)
@@ -2720,30 +2718,21 @@ def _filter_single_chunk_optimized(self, feature_uids_to_keep):
     """
     # Create boolean mask using hash-based set lookup (O(1) per element)
     filter_expr = pl.col("feature_uid").is_in(list(feature_uids_to_keep))
-    
+
     # Apply filter using lazy evaluation with optimized execution
     self.features_df = (
-        self.features_df
-        .lazy()
-        .filter(filter_expr)
-        .collect(streaming=True)  # Use streaming for memory efficiency
+        self.features_df.lazy().filter(filter_expr).collect(streaming=True)  # Use streaming for memory efficiency
     )
-    
+
     # Apply same filter to consensus_mapping_df if it exists
-    if (self.consensus_mapping_df is not None and 
-        not self.consensus_mapping_df.is_empty()):
-        self.consensus_mapping_df = (
-            self.consensus_mapping_df
-            .lazy()
-            .filter(filter_expr)
-            .collect(streaming=True)
-        )
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
+        self.consensus_mapping_df = self.consensus_mapping_df.lazy().filter(filter_expr).collect(streaming=True)
 
 
 def _filter_chunked_lazy(self, feature_uids_to_keep, chunk_size: int, parallel: bool):
     """
     Chunked processing with lazy evaluation for large datasets.
-    
+
     This approach:
     1. Processes data in manageable chunks to control memory usage
     2. Uses lazy evaluation to optimize query execution
@@ -2752,35 +2741,34 @@ def _filter_chunked_lazy(self, feature_uids_to_keep, chunk_size: int, parallel:
     """
     total_features = len(self.features_df)
     num_chunks = (total_features + chunk_size - 1) // chunk_size
-    
+
     self.logger.debug(f"Processing {total_features:,} features in {num_chunks} chunks")
-    
+
     # Process features_df in chunks using lazy evaluation
     filtered_chunks = []
-    
+
     for i in range(num_chunks):
         start_idx = i * chunk_size
         end_idx = min((i + 1) * chunk_size, total_features)
-        
+
         # Create lazy query for this chunk
         chunk_query = (
-            self.features_df
-            .lazy()
+            self.features_df.lazy()
             .slice(start_idx, end_idx - start_idx)
             .filter(pl.col("feature_uid").is_in(list(feature_uids_to_keep)))
         )
-        
+
         # Collect chunk with streaming for memory efficiency
         chunk_result = chunk_query.collect(streaming=True)
         if not chunk_result.is_empty():
             filtered_chunks.append(chunk_result)
-    
+
     # Combine all filtered chunks efficiently
     if filtered_chunks:
         self.features_df = pl.concat(filtered_chunks, how="vertical")
     else:
         self.features_df = pl.DataFrame()  # No features remain
-    
+
     # Apply same chunked processing to consensus_mapping_df
     _filter_consensus_mapping_chunked(self, feature_uids_to_keep, chunk_size)
 
@@ -2789,17 +2777,15 @@ def _filter_consensus_mapping_chunked(self, feature_uids_to_keep, chunk_size: in
     """
     Apply chunked filtering to consensus_mapping_df with same optimization strategy.
     """
-    if (self.consensus_mapping_df is None or 
-        self.consensus_mapping_df.is_empty()):
+    if self.consensus_mapping_df is None or self.consensus_mapping_df.is_empty():
         return
-    
+
     total_mappings = len(self.consensus_mapping_df)
-    
+
     if total_mappings <= chunk_size:
         # Single chunk processing
         self.consensus_mapping_df = (
-            self.consensus_mapping_df
-            .lazy()
+            self.consensus_mapping_df.lazy()
             .filter(pl.col("feature_uid").is_in(list(feature_uids_to_keep)))
             .collect(streaming=True)
         )
@@ -2807,22 +2793,21 @@ def _filter_consensus_mapping_chunked(self, feature_uids_to_keep, chunk_size: in
         # Multi-chunk processing
         num_chunks = (total_mappings + chunk_size - 1) // chunk_size
         filtered_chunks = []
-        
+
         for i in range(num_chunks):
             start_idx = i * chunk_size
             end_idx = min((i + 1) * chunk_size, total_mappings)
-            
+
             chunk_query = (
-                self.consensus_mapping_df
-                .lazy()
+                self.consensus_mapping_df.lazy()
                 .slice(start_idx, end_idx - start_idx)
                 .filter(pl.col("feature_uid").is_in(list(feature_uids_to_keep)))
             )
-            
+
             chunk_result = chunk_query.collect(streaming=True)
             if not chunk_result.is_empty():
                 filtered_chunks.append(chunk_result)
-        
+
         if filtered_chunks:
             self.consensus_mapping_df = pl.concat(filtered_chunks, how="vertical")
         else:
@@ -2835,24 +2820,13 @@ def _apply_standard_filter(self, feature_uids_to_keep):
     Still uses optimized set operations and lazy evaluation.
     """
     filter_expr = pl.col("feature_uid").is_in(list(feature_uids_to_keep))
-    
+
     # Apply filter with lazy evaluation
-    self.features_df = (
-        self.features_df
-        .lazy()
-        .filter(filter_expr)
-        .collect(streaming=True)
-    )
-    
+    self.features_df = self.features_df.lazy().filter(filter_expr).collect(streaming=True)
+
     # Apply to consensus_mapping_df
-    if (self.consensus_mapping_df is not None and 
-        not self.consensus_mapping_df.is_empty()):
-        self.consensus_mapping_df = (
-            self.consensus_mapping_df
-            .lazy()
-            .filter(filter_expr)
-            .collect(streaming=True)
-        )
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
+        self.consensus_mapping_df = self.consensus_mapping_df.lazy().filter(filter_expr).collect(streaming=True)
 
 
 def features_delete(self, features):
@@ -2914,14 +2888,9 @@ def features_delete(self, features):
 
     # Apply filter to consensus_mapping_df if it exists - batch operation
     mapping_removed_count = 0
-    if (
-        self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         initial_mapping_count = len(self.consensus_mapping_df)
-        self.consensus_mapping_df = (
-            self.consensus_mapping_df.lazy().filter(filter_condition).collect()
-        )
+        self.consensus_mapping_df = self.consensus_mapping_df.lazy().filter(filter_condition).collect()
         mapping_removed_count = initial_mapping_count - len(self.consensus_mapping_df)
 
     # Calculate results once and log efficiently
@@ -3028,18 +2997,41 @@ def consensus_select(
         return pl.DataFrame()
 
     # Early return optimization - check if any filters are provided
-    filter_params = [uid, mz, rt, inty_mean, consensus_uid, consensus_id, number_samples, 
-                    number_ms2, quality, bl, chrom_coherence_mean, chrom_prominence_mean,
-                    chrom_prominence_scaled_mean, chrom_height_scaled_mean, 
-                    rt_delta_mean, id_top_score, identified,
-                    # New adduct and identification parameters
-                    adduct_top, adduct_charge_top, adduct_mass_neutral_top, adduct_mass_shift_top,
-                    adduct_group, adduct_of, id_top_name, id_top_class, id_top_adduct]
-    
+    filter_params = [
+        uid,
+        mz,
+        rt,
+        inty_mean,
+        consensus_uid,
+        consensus_id,
+        number_samples,
+        number_ms2,
+        quality,
+        bl,
+        chrom_coherence_mean,
+        chrom_prominence_mean,
+        chrom_prominence_scaled_mean,
+        chrom_height_scaled_mean,
+        rt_delta_mean,
+        id_top_score,
+        identified,
+        # New adduct and identification parameters
+        adduct_top,
+        adduct_charge_top,
+        adduct_mass_neutral_top,
+        adduct_mass_shift_top,
+        adduct_group,
+        adduct_of,
+        id_top_name,
+        id_top_class,
+        id_top_adduct,
+    ]
+
     if all(param is None for param in filter_params) and sortby is None:
         return self.consensus_df.clone()
 
     import time
+
     start_time = time.perf_counter()
     initial_count = len(self.consensus_df)
 
@@ -3082,8 +3074,9 @@ def consensus_select(
                 default_mz_tol = default_mz_tol.eic_mz_tol
             else:
                 from masster.study.defaults.align_def import align_defaults
+
                 default_mz_tol = align_defaults().mz_max_diff
-            
+
             min_mz = mz - default_mz_tol
             max_mz = mz + default_mz_tol
             filter_conditions.append((pl.col("mz") >= min_mz) & (pl.col("mz") <= max_mz))
@@ -3106,8 +3099,9 @@ def consensus_select(
                 default_rt_tol = default_rt_tol.eic_rt_tol
             else:
                 from masster.study.defaults.align_def import align_defaults
+
                 default_rt_tol = align_defaults().rt_tol
-            
+
             min_rt = rt - default_rt_tol
             max_rt = rt + default_rt_tol
             filter_conditions.append((pl.col("rt") >= min_rt) & (pl.col("rt") <= max_rt))
@@ -3192,8 +3186,8 @@ def consensus_select(
         if "adduct_charge_top" in available_columns:
             if isinstance(adduct_charge_top, tuple) and len(adduct_charge_top) == 2:
                 filter_conditions.append(
-                    (pl.col("adduct_charge_top") >= adduct_charge_top[0]) &
-                    (pl.col("adduct_charge_top") <= adduct_charge_top[1])
+                    (pl.col("adduct_charge_top") >= adduct_charge_top[0])
+                    & (pl.col("adduct_charge_top") <= adduct_charge_top[1])
                 )
             elif isinstance(adduct_charge_top, list):
                 filter_conditions.append(pl.col("adduct_charge_top").is_in(adduct_charge_top))
@@ -3207,8 +3201,8 @@ def consensus_select(
         if "adduct_mass_neutral_top" in available_columns:
             if isinstance(adduct_mass_neutral_top, tuple) and len(adduct_mass_neutral_top) == 2:
                 filter_conditions.append(
-                    (pl.col("adduct_mass_neutral_top") >= adduct_mass_neutral_top[0]) &
-                    (pl.col("adduct_mass_neutral_top") <= adduct_mass_neutral_top[1])
+                    (pl.col("adduct_mass_neutral_top") >= adduct_mass_neutral_top[0])
+                    & (pl.col("adduct_mass_neutral_top") <= adduct_mass_neutral_top[1])
                 )
             elif isinstance(adduct_mass_neutral_top, list):
                 filter_conditions.append(pl.col("adduct_mass_neutral_top").is_in(adduct_mass_neutral_top))
@@ -3222,8 +3216,8 @@ def consensus_select(
         if "adduct_mass_shift_top" in available_columns:
             if isinstance(adduct_mass_shift_top, tuple) and len(adduct_mass_shift_top) == 2:
                 filter_conditions.append(
-                    (pl.col("adduct_mass_shift_top") >= adduct_mass_shift_top[0]) &
-                    (pl.col("adduct_mass_shift_top") <= adduct_mass_shift_top[1])
+                    (pl.col("adduct_mass_shift_top") >= adduct_mass_shift_top[0])
+                    & (pl.col("adduct_mass_shift_top") <= adduct_mass_shift_top[1])
                 )
             elif isinstance(adduct_mass_shift_top, list):
                 filter_conditions.append(pl.col("adduct_mass_shift_top").is_in(adduct_mass_shift_top))
@@ -3287,8 +3281,7 @@ def consensus_select(
         if "id_top_score" in available_columns:
             if isinstance(id_top_score, tuple) and len(id_top_score) == 2:
                 filter_conditions.append(
-                    (pl.col("id_top_score") >= id_top_score[0]) &
-                    (pl.col("id_top_score") <= id_top_score[1])
+                    (pl.col("id_top_score") >= id_top_score[0]) & (pl.col("id_top_score") <= id_top_score[1])
                 )
             elif isinstance(id_top_score, list):
                 filter_conditions.append(pl.col("id_top_score").is_in(id_top_score))
@@ -3306,14 +3299,10 @@ def consensus_select(
         # Combine all conditions efficiently using reduce
         from functools import reduce
         import operator
+
         combined_filter = reduce(operator.and_, filter_conditions)
-        
-        consensus = (
-            self.consensus_df
-            .lazy()
-            .filter(combined_filter)
-            .collect(streaming=True)
-        )
+
+        consensus = self.consensus_df.lazy().filter(combined_filter).collect(streaming=True)
     else:
         consensus = self.consensus_df.clone()
 
@@ -3334,10 +3323,10 @@ def consensus_select(
         elif isinstance(sortby, (list, tuple)):
             valid_columns = [col for col in sortby if col in consensus.columns]
             invalid_columns = [col for col in sortby if col not in consensus.columns]
-            
+
             if invalid_columns:
                 self.logger.warning(f"Sort columns not found in consensus DataFrame: {invalid_columns}")
-            
+
             if valid_columns:
                 consensus = consensus.sort(valid_columns, descending=descending)
         else:
@@ -3346,8 +3335,10 @@ def consensus_select(
     # Log performance metrics
     elapsed_time = time.perf_counter() - start_time
     removed_count = initial_count - final_count
-    
-    self.logger.info(f"Selected consensus features: {final_count:,} (removed: {removed_count:,}) in {elapsed_time:.4f}s")
+
+    self.logger.info(
+        f"Selected consensus features: {final_count:,} (removed: {removed_count:,}) in {elapsed_time:.4f}s"
+    )
 
     return consensus
 
@@ -3393,10 +3384,7 @@ def consensus_filter(self, consensus):
 
     # Get feature_uids that need to be kept in features_df
     feature_uids_to_keep = []
-    if (
-        self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         feature_uids_to_keep = self.consensus_mapping_df.filter(
             pl.col("consensus_uid").is_in(consensus_uids_to_keep),
         )["feature_uid"].to_list()
@@ -3407,10 +3395,7 @@ def consensus_filter(self, consensus):
     )
 
     # Keep only relevant entries in consensus_mapping_df
-    if (
-        self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         initial_mapping_count = len(self.consensus_mapping_df)
         self.consensus_mapping_df = self.consensus_mapping_df.filter(
             pl.col("consensus_uid").is_in(consensus_uids_to_keep),
@@ -3423,11 +3408,7 @@ def consensus_filter(self, consensus):
             )
 
     # Keep only corresponding features in features_df
-    if (
-        feature_uids_to_keep
-        and self.features_df is not None
-        and not self.features_df.is_empty()
-    ):
+    if feature_uids_to_keep and self.features_df is not None and not self.features_df.is_empty():
         initial_features_count = len(self.features_df)
         self.features_df = self.features_df.filter(
             pl.col("feature_uid").is_in(feature_uids_to_keep),
@@ -3440,11 +3421,7 @@ def consensus_filter(self, consensus):
             )
 
     # Keep only relevant entries in consensus_ms2 if it exists
-    if (
-        hasattr(self, "consensus_ms2")
-        and self.consensus_ms2 is not None
-        and not self.consensus_ms2.is_empty()
-    ):
+    if hasattr(self, "consensus_ms2") and self.consensus_ms2 is not None and not self.consensus_ms2.is_empty():
         initial_ms2_count = len(self.consensus_ms2)
         self.consensus_ms2 = self.consensus_ms2.filter(
             pl.col("consensus_uid").is_in(consensus_uids_to_keep),
@@ -3514,10 +3491,7 @@ def consensus_delete(self, consensus):
 
     # Get feature_uids that need to be removed from features_df
     feature_uids_to_remove = []
-    if (
-        self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         feature_uids_to_remove = self.consensus_mapping_df.filter(
             pl.col("consensus_uid").is_in(consensus_uids_to_remove),
         )["feature_uid"].to_list()
@@ -3529,10 +3503,7 @@ def consensus_delete(self, consensus):
 
     # Remove from consensus_mapping_df
     mapping_removed_count = 0
-    if (
-        self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         initial_mapping_count = len(self.consensus_mapping_df)
         self.consensus_mapping_df = self.consensus_mapping_df.filter(
             ~pl.col("consensus_uid").is_in(consensus_uids_to_remove),
@@ -3541,11 +3512,7 @@ def consensus_delete(self, consensus):
 
     # Remove corresponding features from features_df
     features_removed_count = 0
-    if (
-        feature_uids_to_remove
-        and self.features_df is not None
-        and not self.features_df.is_empty()
-    ):
+    if feature_uids_to_remove and self.features_df is not None and not self.features_df.is_empty():
         initial_features_count = len(self.features_df)
         self.features_df = self.features_df.filter(
             ~pl.col("feature_uid").is_in(feature_uids_to_remove),
@@ -3554,11 +3521,7 @@ def consensus_delete(self, consensus):
 
     # Remove from consensus_ms2 if it exists
     ms2_removed_count = 0
-    if (
-        hasattr(self, "consensus_ms2")
-        and self.consensus_ms2 is not None
-        and not self.consensus_ms2.is_empty()
-    ):
+    if hasattr(self, "consensus_ms2") and self.consensus_ms2 is not None and not self.consensus_ms2.is_empty():
         initial_ms2_count = len(self.consensus_ms2)
         self.consensus_ms2 = self.consensus_ms2.filter(
             ~pl.col("consensus_uid").is_in(consensus_uids_to_remove),
@@ -3577,7 +3540,7 @@ def consensus_delete(self, consensus):
         log_parts.append(f"{features_removed_count} features")
     if ms2_removed_count > 0:
         log_parts.append(f"{ms2_removed_count} MS2 spectra")
-    
+
     log_message = ". ".join(log_parts) + f". Remaining consensus: {final_consensus_count}"
     self.logger.info(log_message)
 
@@ -3651,8 +3614,7 @@ def samples_select(
                 # Treat as range
                 min_uid, max_uid = sample_uid
                 filter_conditions.append(
-                    (pl.col("sample_uid") >= min_uid)
-                    & (pl.col("sample_uid") <= max_uid),
+                    (pl.col("sample_uid") >= min_uid) & (pl.col("sample_uid") <= max_uid),
                 )
             else:
                 # Treat as list
@@ -3695,8 +3657,7 @@ def samples_select(
                     # Treat as range
                     min_batch, max_batch = sample_batch
                     filter_conditions.append(
-                        (pl.col("sample_batch") >= min_batch)
-                        & (pl.col("sample_batch") <= max_batch),
+                        (pl.col("sample_batch") >= min_batch) & (pl.col("sample_batch") <= max_batch),
                     )
                 else:
                     # Treat as list
@@ -3714,8 +3675,7 @@ def samples_select(
                     # Treat as range
                     min_seq, max_seq = sample_sequence
                     filter_conditions.append(
-                        (pl.col("sample_sequence") >= min_seq)
-                        & (pl.col("sample_sequence") <= max_seq),
+                        (pl.col("sample_sequence") >= min_seq) & (pl.col("sample_sequence") <= max_seq),
                     )
                 else:
                     # Treat as list
@@ -3733,8 +3693,7 @@ def samples_select(
             if isinstance(num_features, tuple) and len(num_features) == 2:
                 min_features, max_features = num_features
                 filter_conditions.append(
-                    (pl.col("num_features") >= min_features)
-                    & (pl.col("num_features") <= max_features),
+                    (pl.col("num_features") >= min_features) & (pl.col("num_features") <= max_features),
                 )
             else:
                 filter_conditions.append(pl.col("num_features") >= num_features)
@@ -3883,11 +3842,7 @@ def samples_delete(self, samples):
 
     # 2. Remove corresponding features from features_df
     removed_features_count = 0
-    if (
-        feature_uids_to_remove
-        and self.features_df is not None
-        and not self.features_df.is_empty()
-    ):
+    if feature_uids_to_remove and self.features_df is not None and not self.features_df.is_empty():
         self.features_df = self.features_df.filter(
             ~pl.col("sample_uid").is_in(sample_uids_to_remove),
         )
@@ -3895,11 +3850,7 @@ def samples_delete(self, samples):
 
     # 3. Remove from consensus_mapping_df
     removed_mapping_count = 0
-    if (
-        feature_uids_to_remove
-        and self.consensus_mapping_df is not None
-        and not self.consensus_mapping_df.is_empty()
-    ):
+    if feature_uids_to_remove and self.consensus_mapping_df is not None and not self.consensus_mapping_df.is_empty():
         initial_mapping_count = len(self.consensus_mapping_df)
         self.consensus_mapping_df = self.consensus_mapping_df.filter(
             ~pl.col("feature_uid").is_in(feature_uids_to_remove),
@@ -3908,11 +3859,7 @@ def samples_delete(self, samples):
 
     # 4. Remove from consensus_ms2 if it exists
     removed_ms2_count = 0
-    if (
-        hasattr(self, "consensus_ms2")
-        and self.consensus_ms2 is not None
-        and not self.consensus_ms2.is_empty()
-    ):
+    if hasattr(self, "consensus_ms2") and self.consensus_ms2 is not None and not self.consensus_ms2.is_empty():
         initial_ms2_count = len(self.consensus_ms2)
         self.consensus_ms2 = self.consensus_ms2.filter(
             ~pl.col("sample_uid").is_in(sample_uids_to_remove),
@@ -3921,11 +3868,7 @@ def samples_delete(self, samples):
 
     # 5. Remove from feature_maps and update map_id
     removed_maps_count = 0
-    if (
-        hasattr(self, "feature_maps")
-        and self.feature_maps is not None
-        and map_ids_to_remove
-    ):
+    if hasattr(self, "feature_maps") and self.feature_maps is not None and map_ids_to_remove:
         # Remove feature maps in reverse order to maintain indices
         for map_id in sorted(map_ids_to_remove, reverse=True):
             if 0 <= map_id < len(self.feature_maps):
@@ -4148,9 +4091,6 @@ def set_samples_color(self, by=None, palette="Turbo256"):
         self.logger.debug(f"Set sample colors based on {by} using {palette} palette")
 
 
-
-
-
 def _get_color_palette(palette_name):
     """
     Get color palette as a list of hex color codes using the cmap library.
@@ -4304,9 +4244,7 @@ def _sample_colors_from_colormap(palette_name, n_colors):
         # Distribute samples evenly across the full colormap range (same approach as set_samples_color(by=None))
         for i in range(n_colors):
             # Evenly distribute samples across colormap (avoiding endpoints to prevent white/black)
-            normalized_value = (
-                i + 0.5
-            ) / n_colors  # +0.5 to center samples in their bins
+            normalized_value = (i + 0.5) / n_colors  # +0.5 to center samples in their bins
             # Map to a subset of colormap to avoid extreme colors (use 10% to 90% range)
             normalized_value = 0.1 + (normalized_value * 0.8)
 
@@ -4441,18 +4379,14 @@ def restore_ms2(self, samples=None, **kwargs):
     self.logger.info(f"Restoring MS2 data from {len(sample_uids)} samples...")
 
     # Clear existing consensus_ms2 to rebuild from scratch
-    initial_ms2_count = (
-        len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
-    )
+    initial_ms2_count = len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
     self.consensus_ms2 = pl.DataFrame()
 
     # Re-run find_ms2 which will rebuild consensus_ms2
     try:
         self.find_ms2(**kwargs)
 
-        final_ms2_count = (
-            len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
-        )
+        final_ms2_count = len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
 
         self.logger.info(
             f"MS2 restoration completed: {initial_ms2_count} -> {final_ms2_count} MS2 spectra",
@@ -4551,12 +4485,8 @@ def decompress(self, features=True, ms2=True, chrom=True, samples=None, **kwargs
     # Check if MS2 data might need restoration (compare expected vs actual)
     ms2_need_restoration = False
     if ms2:
-        current_ms2_count = (
-            len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
-        )
-        consensus_count = (
-            len(self.consensus_df) if not self.consensus_df.is_empty() else 0
-        )
+        current_ms2_count = len(self.consensus_ms2) if not self.consensus_ms2.is_empty() else 0
+        consensus_count = len(self.consensus_df) if not self.consensus_df.is_empty() else 0
 
         if consensus_count > 0:
             # Calculate expected MS2 count based on consensus features with MS2 potential
@@ -4607,15 +4537,11 @@ def decompress(self, features=True, ms2=True, chrom=True, samples=None, **kwargs
             # Then do additional chrom gap-filling if needed
             self.restore_chrom(samples=samples, **restore_kwargs)
 
-        elif (
-            "features" in operations_needed and "chromatograms" not in operations_needed
-        ):
+        elif "features" in operations_needed and "chromatograms" not in operations_needed:
             self.logger.info("Phase 1: Restoring features data...")
             self.restore_features(samples=samples)
 
-        elif (
-            "chromatograms" in operations_needed and "features" not in operations_needed
-        ):
+        elif "chromatograms" in operations_needed and "features" not in operations_needed:
             self.logger.info("Phase 1: Restoring chromatograms...")
             restore_kwargs = {}
             if "mz_tol" in kwargs: