masster 0.3.13__py3-none-any.whl → 0.3.15__py3-none-any.whl

masster/study/helpers.py

@@ -1,3 +1,18 @@
+"""
+helpers.py
+
+This module contains helper functions for the Study class that handle various operations
+like data retrieval, filtering, compression, and utility functions.
+
+The functions are organized into the following sections:
+1. Chromatogram extraction functions (BPC, TIC, EIC, chrom matrix)
+2. Data retrieval helper functions (get_sample, get_consensus, etc.)  
+3. UID helper functions (_get_*_uids)
+4. Data filtering and selection functions
+5. Data compression and restoration functions
+6. Utility functions (reset, naming, colors, schema ordering)
+"""
+
 from __future__ import annotations
 
 import os
@@ -10,6 +25,11 @@ from tqdm import tqdm
 from masster.chromatogram import Chromatogram
 
 
+# =====================================================================================
+# CHROMATOGRAM EXTRACTION FUNCTIONS
+# =====================================================================================
+
+
 def get_bpc(owner, sample=None, rt_unit="s", label=None, original=False):
     """
     Return a Chromatogram object containing the Base Peak Chromatogram (BPC).
@@ -96,7 +116,6 @@ def get_bpc(owner, sample=None, rt_unit="s", label=None, original=False):
                 if (mapping_rows is None or mapping_rows.is_empty()) and hasattr(s, "sample_path"):
                     # attempt to match by sample_path or file name
                     try:
-                        sample_paths = feats.select(["sample_uid", "sample_name", "sample_path"])  # type: ignore[arg-type]
                         # find row where sample_path matches
                         mapping_rows = feats.filter(pl.col("sample_path") == getattr(s, "file", None))
                     except Exception:
@@ -192,7 +211,7 @@ def get_tic(owner, sample=None, label=None):
     return chrom
 
     
-def get_eic(owner, sample=None, mz=None, mz_tol=0.01, rt_unit="s", label=None):
+def get_eic(owner, sample=None, mz=None, mz_tol=None, rt_unit="s", label=None):
     """
     Return a Chromatogram object containing the Extracted Ion Chromatogram (EIC) for a target m/z.
 
@@ -206,13 +225,20 @@ def get_eic(owner, sample=None, mz=None, mz_tol=0.01, rt_unit="s", label=None):
         owner: Study or Sample instance
         sample: Sample identifier (required if owner is Study)  
         mz (float): Target m/z value
-        mz_tol (float): m/z tolerance (default 0.01)
+        mz_tol (float): m/z tolerance. If None, uses owner.parameters.eic_mz_tol (for Study) or defaults to 0.01
         rt_unit (str): Retention time unit for the chromatogram
         label (str): Optional label for the chromatogram
 
     Returns:
         Chromatogram
     """
+    # Use default mz_tol from study parameters if not provided
+    if mz_tol is None:
+        if hasattr(owner, 'parameters') and hasattr(owner.parameters, 'eic_mz_tol'):
+            mz_tol = owner.parameters.eic_mz_tol
+        else:
+            mz_tol = 0.01  # fallback default
+
     if mz is None:
         raise ValueError("mz must be provided for EIC computation")
 
@@ -290,6 +316,9 @@ def get_eic(owner, sample=None, mz=None, mz_tol=0.01, rt_unit="s", label=None):
     return chrom
 
 
+# =====================================================================================
+# DATA RETRIEVAL AND MATRIX FUNCTIONS  
+# =====================================================================================
 
 
 def get_chrom(self, uids=None, samples=None):
@@ -393,10 +422,14 @@ def get_chrom(self, uids=None, samples=None):
     # Create Polars DataFrame with complex objects
     df2_pivoted = pl.DataFrame(pivot_data)
 
-    # Return as Polars DataFrame (can handle complex objects like Chromatogram)
     return df2_pivoted
 
 
+# =====================================================================================
+# UTILITY AND CONFIGURATION FUNCTIONS
+# =====================================================================================
+
+
 def set_folder(self, folder):
     """
     Set the folder for saving and loading files.
@@ -424,6 +457,12 @@ def align_reset(self):
     # Ensure column order is maintained after with_columns operation
     self._ensure_features_df_schema_order()
 
+
+# =====================================================================================
+# DATA RETRIEVAL HELPER FUNCTIONS
+# =====================================================================================
+
+
 # TODO I don't get this param
 def get_consensus(self, quant="chrom_area"):
     if self.consensus_df is None:
@@ -555,6 +594,11 @@ def get_consensus_matches(self, uids=None):
     return matches
 
 
+# =====================================================================================  
+# UID HELPER FUNCTIONS
+# =====================================================================================
+
+
 def fill_reset(self):
     # remove all features with filled=True
     if self.features_df is None:
@@ -757,6 +801,11 @@ def get_orphans(self):
     return not_in_consensus
 
 
+# =====================================================================================
+# DATA COMPRESSION AND RESTORATION FUNCTIONS  
+# =====================================================================================
+
+
 def compress(self, features=True, ms2=True, chrom=False, ms2_max=5):
     """
     Perform compress_features, compress_ms2, and compress_chrom operations.
@@ -1251,7 +1300,12 @@ def compress_chrom(self):
     self.logger.info(f"Compressed chromatograms: cleared {non_null_count} chromatogram objects from features_df")
 
 
-def name_replace(self, replace_dict):
+# =====================================================================================
+# SAMPLE MANAGEMENT AND NAMING FUNCTIONS
+# =====================================================================================
+
+
+def sample_name_replace(self, replace_dict):
     """
     Replace sample names in samples_df based on a dictionary mapping.
     
@@ -1317,7 +1371,7 @@ def name_replace(self, replace_dict):
     self.logger.info(f"Successfully replaced {replaced_count} sample names")
 
 
-def name_reset(self):
+def sample_name_reset(self):
     """
     Reset sample names to the basename of sample_path without extensions.
     
@@ -1399,7 +1453,7 @@ def set_source(self, filename):
     failed_count = 0
 
     # Get all current file_source values
-    current_sources = self.samples_df.get_column("file_source").to_list()
+    current_sources = self.samples_df.get_column("sample_source").to_list()
     sample_names = self.samples_df.get_column("sample_name").to_list()
 
     new_sources = []
@@ -1447,6 +1501,11 @@ def set_source(self, filename):
         self.logger.warning(f"Failed to update file_source for {failed_count} samples")
 
 
+# =====================================================================================
+# DATA FILTERING AND SELECTION FUNCTIONS
+# =====================================================================================
+
+
 def features_select(
     self,
     mz=None,
@@ -1872,13 +1931,21 @@ def consensus_select(
     chrom_prominence_scaled_mean=None,
     chrom_height_scaled_mean=None,
     rt_delta_mean=None,
+    sortby=None,
+    descending=True,
 ):
     """
     Select consensus features from consensus_df based on specified criteria and return the filtered DataFrame.
 
     Parameters:
-        mz: m/z range filter (tuple for range, single value for minimum)
-        rt: retention time range filter (tuple for range, single value for minimum)
+        mz: m/z filter with flexible formats:
+            - float: m/z value ± default tolerance (uses study.parameters.eic_mz_tol)
+            - tuple (mz_min, mz_max): range where mz_max > mz_min
+            - tuple (mz_center, mz_tol): range where mz_tol < mz_center (interpreted as mz_center ± mz_tol)
+        rt: retention time filter with flexible formats:
+            - float: RT value ± default tolerance (uses study.parameters.eic_rt_tol)
+            - tuple (rt_min, rt_max): range where rt_max > rt_min
+            - tuple (rt_center, rt_tol): range where rt_tol < rt_center (interpreted as rt_center ± rt_tol)
         inty_mean: mean intensity filter (tuple for range, single value for minimum)
         consensus_uid: consensus UID filter (list, single value, or tuple for range)
         consensus_id: consensus ID filter (list or single value)
@@ -1891,6 +1958,8 @@ def consensus_select(
         chrom_prominence_scaled_mean: mean scaled chromatogram prominence filter (tuple for range, single value for minimum)
         chrom_height_scaled_mean: mean scaled chromatogram height filter (tuple for range, single value for minimum)
         rt_delta_mean: mean RT delta filter (tuple for range, single value for minimum)
+        sortby: column name(s) to sort by (string, list of strings, or None for no sorting)
+        descending: sort direction (True for descending, False for ascending, default is True)
 
     Returns:
         polars.DataFrame: Filtered consensus DataFrame
@@ -1905,11 +1974,32 @@ def consensus_select(
     # Filter by m/z
     if mz is not None:
         consensus_len_before_filter = len(consensus)
+        
         if isinstance(mz, tuple) and len(mz) == 2:
-            min_mz, max_mz = mz
+            # Check if second value is smaller than first (indicating mz, mz_tol format)
+            if mz[1] < mz[0]:
+                # First is mz, second is mz_tol
+                mz_center, mz_tol = mz
+                min_mz = mz_center - mz_tol
+                max_mz = mz_center + mz_tol
+            else:
+                # Standard (min_mz, max_mz) format
+                min_mz, max_mz = mz
             consensus = consensus.filter((pl.col("mz") >= min_mz) & (pl.col("mz") <= max_mz))
         else:
-            consensus = consensus.filter(pl.col("mz") >= mz)
+            # Single float value - use default mz tolerance from study parameters
+            default_mz_tol = getattr(self, 'parameters', None)
+            if default_mz_tol and hasattr(default_mz_tol, 'eic_mz_tol'):
+                default_mz_tol = default_mz_tol.eic_mz_tol
+            else:
+                # Fallback to align_defaults if study parameters not available
+                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
+            consensus = consensus.filter((pl.col("mz") >= min_mz) & (pl.col("mz") <= max_mz))
+        
         self.logger.debug(
             f"Selected consensus by mz. Consensus removed: {consensus_len_before_filter - len(consensus)}",
         )
@@ -1917,11 +2007,32 @@ def consensus_select(
     # Filter by retention time
     if rt is not None:
         consensus_len_before_filter = len(consensus)
+        
         if isinstance(rt, tuple) and len(rt) == 2:
-            min_rt, max_rt = rt
+            # Check if second value is smaller than first (indicating rt, rt_tol format)
+            if rt[1] < rt[0]:
+                # First is rt, second is rt_tol
+                rt_center, rt_tol = rt
+                min_rt = rt_center - rt_tol
+                max_rt = rt_center + rt_tol
+            else:
+                # Standard (min_rt, max_rt) format
+                min_rt, max_rt = rt
             consensus = consensus.filter((pl.col("rt") >= min_rt) & (pl.col("rt") <= max_rt))
         else:
-            consensus = consensus.filter(pl.col("rt") >= rt)
+            # Single float value - use default rt tolerance from study parameters
+            default_rt_tol = getattr(self, 'parameters', None)
+            if default_rt_tol and hasattr(default_rt_tol, 'eic_rt_tol'):
+                default_rt_tol = default_rt_tol.eic_rt_tol
+            else:
+                # Fallback to align_defaults if study parameters not available
+                from masster.study.defaults.align_def import align_defaults
+                default_rt_tol = align_defaults().rt_max_diff
+            
+            min_rt = rt - default_rt_tol
+            max_rt = rt + default_rt_tol
+            consensus = consensus.filter((pl.col("rt") >= min_rt) & (pl.col("rt") <= max_rt))
+        
         self.logger.debug(
             f"Selected consensus by rt. Consensus removed: {consensus_len_before_filter - len(consensus)}",
         )
@@ -2118,6 +2229,27 @@ def consensus_select(
     else:
         self.logger.info(f"Selected consensus features. Features remaining: {len(consensus)} (from {initial_count})")
 
+    # Sort the results if sortby is specified
+    if sortby is not None:
+        if isinstance(sortby, str):
+            # Single column
+            if sortby in consensus.columns:
+                consensus = consensus.sort(sortby, descending=descending)
+            else:
+                self.logger.warning(f"Sort column '{sortby}' not found in consensus DataFrame")
+        elif isinstance(sortby, (list, tuple)):
+            # Multiple columns
+            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:
+            self.logger.warning(f"Invalid sortby parameter type: {type(sortby)}. Expected str, list, or tuple.")
+
     return consensus
 
 
@@ -2222,3 +2354,832 @@ def consensus_delete(self, consensus):
         None (modifies self.consensus_df and related DataFrames in place)
     """
     self.consensus_filter(consensus)
+
+
+# =====================================================================================
+# SAMPLE MANAGEMENT AND DELETION FUNCTIONS
+# =====================================================================================
+
+
+def samples_select(
+    self,
+    sample_uid=None,
+    sample_name=None,
+    sample_type=None,
+    sample_group=None,
+    sample_batch=None,
+    sample_sequence=None,
+    num_features=None,
+    num_ms1=None,
+    num_ms2=None,
+):
+    """
+    Select samples from samples_df based on specified criteria and return the filtered DataFrame.
+
+    Parameters:
+        sample_uid: sample UID filter (list, single value, or tuple for range)
+        sample_name: sample name filter (list or single value)
+        sample_type: sample type filter (list or single value)
+        sample_group: sample group filter (list or single value)
+        sample_batch: sample batch filter (list, single value, or tuple for range)
+        sample_sequence: sample sequence filter (list, single value, or tuple for range)
+        num_features: number of features filter (tuple for range, single value for minimum)
+        num_ms1: number of MS1 spectra filter (tuple for range, single value for minimum)
+        num_ms2: number of MS2 spectra filter (tuple for range, single value for minimum)
+
+    Returns:
+        polars.DataFrame: Filtered samples DataFrame
+    """
+    if self.samples_df is None or self.samples_df.is_empty():
+        self.logger.warning("No samples found in study.")
+        return pl.DataFrame()
+
+    # Early return if no filters provided
+    filter_params = [
+        sample_uid,
+        sample_name,
+        sample_type,
+        sample_group,
+        sample_batch,
+        sample_sequence,
+        num_features,
+        num_ms1,
+        num_ms2,
+    ]
+    if all(param is None for param in filter_params):
+        return self.samples_df.clone()
+
+    initial_count = len(self.samples_df)
+
+    # Pre-check available columns once for efficiency
+    available_columns = set(self.samples_df.columns)
+
+    # Build all filter conditions first, then apply them all at once
+    filter_conditions = []
+    warnings = []
+
+    # Filter by sample_uid
+    if sample_uid is not None:
+        if isinstance(sample_uid, (list, tuple)):
+            if len(sample_uid) == 2 and not isinstance(sample_uid, list):
+                # Treat as range
+                min_uid, max_uid = sample_uid
+                filter_conditions.append((pl.col("sample_uid") >= min_uid) & (pl.col("sample_uid") <= max_uid))
+            else:
+                # Treat as list
+                filter_conditions.append(pl.col("sample_uid").is_in(sample_uid))
+        else:
+            filter_conditions.append(pl.col("sample_uid") == sample_uid)
+
+    # Filter by sample_name
+    if sample_name is not None:
+        if isinstance(sample_name, list):
+            filter_conditions.append(pl.col("sample_name").is_in(sample_name))
+        else:
+            filter_conditions.append(pl.col("sample_name") == sample_name)
+
+    # Filter by sample_type
+    if sample_type is not None:
+        if "sample_type" in available_columns:
+            if isinstance(sample_type, list):
+                filter_conditions.append(pl.col("sample_type").is_in(sample_type))
+            else:
+                filter_conditions.append(pl.col("sample_type") == sample_type)
+        else:
+            warnings.append("'sample_type' column not found in samples_df")
+
+    # Filter by sample_group
+    if sample_group is not None:
+        if "sample_group" in available_columns:
+            if isinstance(sample_group, list):
+                filter_conditions.append(pl.col("sample_group").is_in(sample_group))
+            else:
+                filter_conditions.append(pl.col("sample_group") == sample_group)
+        else:
+            warnings.append("'sample_group' column not found in samples_df")
+
+    # Filter by sample_batch
+    if sample_batch is not None:
+        if "sample_batch" in available_columns:
+            if isinstance(sample_batch, (list, tuple)):
+                if len(sample_batch) == 2 and not isinstance(sample_batch, list):
+                    # Treat as range
+                    min_batch, max_batch = sample_batch
+                    filter_conditions.append((pl.col("sample_batch") >= min_batch) & (pl.col("sample_batch") <= max_batch))
+                else:
+                    # Treat as list
+                    filter_conditions.append(pl.col("sample_batch").is_in(sample_batch))
+            else:
+                filter_conditions.append(pl.col("sample_batch") == sample_batch)
+        else:
+            warnings.append("'sample_batch' column not found in samples_df")
+
+    # Filter by sample_sequence
+    if sample_sequence is not None:
+        if "sample_sequence" in available_columns:
+            if isinstance(sample_sequence, (list, tuple)):
+                if len(sample_sequence) == 2 and not isinstance(sample_sequence, list):
+                    # Treat as range
+                    min_seq, max_seq = sample_sequence
+                    filter_conditions.append((pl.col("sample_sequence") >= min_seq) & (pl.col("sample_sequence") <= max_seq))
+                else:
+                    # Treat as list
+                    filter_conditions.append(pl.col("sample_sequence").is_in(sample_sequence))
+            else:
+                filter_conditions.append(pl.col("sample_sequence") == sample_sequence)
+        else:
+            warnings.append("'sample_sequence' column not found in samples_df")
+
+    # Filter by num_features
+    if num_features is not None:
+        if "num_features" in available_columns:
+            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))
+            else:
+                filter_conditions.append(pl.col("num_features") >= num_features)
+        else:
+            warnings.append("'num_features' column not found in samples_df")
+
+    # Filter by num_ms1
+    if num_ms1 is not None:
+        if "num_ms1" in available_columns:
+            if isinstance(num_ms1, tuple) and len(num_ms1) == 2:
+                min_ms1, max_ms1 = num_ms1
+                filter_conditions.append((pl.col("num_ms1") >= min_ms1) & (pl.col("num_ms1") <= max_ms1))
+            else:
+                filter_conditions.append(pl.col("num_ms1") >= num_ms1)
+        else:
+            warnings.append("'num_ms1' column not found in samples_df")
+
+    # Filter by num_ms2
+    if num_ms2 is not None:
+        if "num_ms2" in available_columns:
+            if isinstance(num_ms2, tuple) and len(num_ms2) == 2:
+                min_ms2, max_ms2 = num_ms2
+                filter_conditions.append((pl.col("num_ms2") >= min_ms2) & (pl.col("num_ms2") <= max_ms2))
+            else:
+                filter_conditions.append(pl.col("num_ms2") >= num_ms2)
+        else:
+            warnings.append("'num_ms2' column not found in samples_df")
+
+    # Log all warnings once at the end for efficiency
+    for warning in warnings:
+        self.logger.warning(warning)
+
+    # Apply all filters at once using lazy evaluation for optimal performance
+    if filter_conditions:
+        # Combine all conditions with AND
+        combined_filter = filter_conditions[0]
+        for condition in filter_conditions[1:]:
+            combined_filter = combined_filter & condition
+
+        # Apply the combined filter using lazy evaluation
+        samples = self.samples_df.lazy().filter(combined_filter).collect()
+    else:
+        samples = self.samples_df.clone()
+
+    final_count = len(samples)
+
+    if final_count == 0:
+        self.logger.warning("No samples remaining after applying selection criteria.")
+    else:
+        self.logger.info(f"Samples selected: {final_count} (out of {initial_count})")
+
+    return samples
+
+
+def samples_delete(self, samples):
+    """
+    Delete samples and all related data from the study based on sample identifiers.
+    
+    This function eliminates all data related to the specified samples (and their sample_uids) 
+    from all dataframes including:
+    - samples_df: Removes the sample rows
+    - features_df: Removes all features belonging to these samples
+    - consensus_mapping_df: Removes mappings for features from these samples
+    - consensus_ms2: Removes MS2 spectra for features from these samples
+    - feature_maps: Removes the corresponding feature maps
+    
+    Also updates map_id values to maintain sequential indices after deletion.
+
+    Parameters:
+        samples: Samples to delete. Can be:
+                - list of int: List of sample_uids to delete
+                - polars.DataFrame: DataFrame obtained from samples_select (will use sample_uid column)
+                - int: Single sample_uid to delete
+
+    Returns:
+        None (modifies study DataFrames and feature_maps in place)
+    """
+    if self.samples_df is None or self.samples_df.is_empty():
+        self.logger.warning("No samples found in study.")
+        return
+
+    # Early return if no samples provided
+    if samples is None:
+        self.logger.warning("No samples provided for deletion.")
+        return
+
+    initial_sample_count = len(self.samples_df)
+
+    # Determine sample_uids to remove
+    if isinstance(samples, pl.DataFrame):
+        if "sample_uid" not in samples.columns:
+            self.logger.error("samples DataFrame must contain 'sample_uid' column")
+            return
+        sample_uids_to_remove = samples["sample_uid"].to_list()
+    elif isinstance(samples, (list, tuple)):
+        sample_uids_to_remove = list(samples)  # Convert tuple to list if needed
+    elif isinstance(samples, int):
+        sample_uids_to_remove = [samples]
+    else:
+        self.logger.error("samples parameter must be a DataFrame, list, tuple, or int")
+        return
+
+    # Early return if no UIDs to remove
+    if not sample_uids_to_remove:
+        self.logger.warning("No sample UIDs provided for deletion.")
+        return
+
+    # Convert to set for faster lookup if list is large
+    if len(sample_uids_to_remove) > 100:
+        sample_uids_set = set(sample_uids_to_remove)
+        # Use the set for filtering if it's significantly smaller
+        if len(sample_uids_set) < len(sample_uids_to_remove) * 0.8:
+            sample_uids_to_remove = list(sample_uids_set)
+
+    self.logger.info(f"Deleting {len(sample_uids_to_remove)} samples and all related data...")
+
+    # Get feature_uids that need to be removed from features_df
+    feature_uids_to_remove = []
+    initial_features_count = 0
+    if self.features_df is not None and not self.features_df.is_empty():
+        initial_features_count = len(self.features_df)
+        feature_uids_to_remove = self.features_df.filter(
+            pl.col("sample_uid").is_in(sample_uids_to_remove),
+        )["feature_uid"].to_list()
+
+    # Get map_ids to remove from feature_maps (needed before samples_df deletion)
+    map_ids_to_remove = []
+    if hasattr(self, 'feature_maps') and self.feature_maps is not None:
+        # Get map_ids for samples to be deleted
+        map_ids_df = self.samples_df.filter(
+            pl.col("sample_uid").is_in(sample_uids_to_remove)
+        ).select("map_id")
+        if not map_ids_df.is_empty():
+            map_ids_to_remove = map_ids_df["map_id"].to_list()
+
+    # 1. Remove samples from samples_df
+    self.samples_df = self.samples_df.filter(
+        ~pl.col("sample_uid").is_in(sample_uids_to_remove),
+    )
+
+    # 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():
+        self.features_df = self.features_df.filter(
+            ~pl.col("sample_uid").is_in(sample_uids_to_remove),
+        )
+        removed_features_count = initial_features_count - len(self.features_df)
+
+    # 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():
+        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),
+        )
+        removed_mapping_count = initial_mapping_count - len(self.consensus_mapping_df)
+
+    # 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():
+        initial_ms2_count = len(self.consensus_ms2)
+        self.consensus_ms2 = self.consensus_ms2.filter(
+            ~pl.col("sample_uid").is_in(sample_uids_to_remove),
+        )
+        removed_ms2_count = initial_ms2_count - len(self.consensus_ms2)
+
+    # 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:
+        # 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):
+                self.feature_maps.pop(map_id)
+                removed_maps_count += 1
+
+        # Update map_id values in samples_df to maintain sequential indices
+        if len(self.samples_df) > 0:
+            new_map_ids = list(range(len(self.samples_df)))
+            self.samples_df = self.samples_df.with_columns(
+                pl.lit(new_map_ids).alias("map_id")
+            )
+
+    # Calculate and log results
+    removed_sample_count = initial_sample_count - len(self.samples_df)
+    final_sample_count = len(self.samples_df)
+
+    # Create comprehensive summary message
+    summary_parts = [
+        f"Deleted {removed_sample_count} samples",
+    ]
+    
+    if removed_features_count > 0:
+        summary_parts.append(f"{removed_features_count} features")
+    
+    if removed_mapping_count > 0:
+        summary_parts.append(f"{removed_mapping_count} consensus mappings")
+    
+    if removed_ms2_count > 0:
+        summary_parts.append(f"{removed_ms2_count} MS2 spectra")
+    
+    if removed_maps_count > 0:
+        summary_parts.append(f"{removed_maps_count} feature maps")
+
+    summary_parts.append(f"Remaining samples: {final_sample_count}")
+
+    self.logger.info(". ".join(summary_parts))
+
+    # Update map_id indices if needed
+    if removed_maps_count > 0 and final_sample_count > 0:
+        self.logger.debug(f"Updated map_id values to range from 0 to {final_sample_count - 1}")
+
+
+# =====================================================================================
+# COLOR PALETTE AND VISUALIZATION FUNCTIONS
+# =====================================================================================
+
+
+def sample_color(self, by=None, palette="Turbo256"):
+    """
+    Set sample colors in the sample_color column of samples_df.
+    
+    When a new sample is added, this function resets all colors picking from the specified palette.
+    The default palette is Turbo256.
+
+    Parameters:
+        by (str or list, optional): Property to base colors on. Options:
+                                     - 'sample_uid': Use sample_uid values to assign colors
+                                     - 'sample_index': Use sample index (position) to assign colors  
+                                     - 'sample_type': Use sample_type values to assign colors
+                                     - 'sample_name': Use sample_name values to assign colors
+                                     - list of colors: Use provided list of hex color codes
+                                     - None: Use sequential colors from palette (default)
+        palette (str): Color palette to use. Options:
+                      - 'Turbo256': Turbo colormap (256 colors, perceptually uniform)
+                      - 'Viridis256': Viridis colormap (256 colors, perceptually uniform)
+                      - 'Plasma256': Plasma colormap (256 colors, perceptually uniform)
+                      - 'Inferno256': Inferno colormap (256 colors, perceptually uniform)
+                      - 'Magma256': Magma colormap (256 colors, perceptually uniform)
+                      - 'Cividis256': Cividis colormap (256 colors, colorblind-friendly)
+                      - 'Set1': Qualitative palette (9 distinct colors)
+                      - 'Set2': Qualitative palette (8 distinct colors)  
+                      - 'Set3': Qualitative palette (12 distinct colors)
+                      - 'Tab10': Tableau 10 palette (10 distinct colors)
+                      - 'Tab20': Tableau 20 palette (20 distinct colors)
+                      - 'Dark2': Dark qualitative palette (8 colors)
+                      - 'Paired': Paired qualitative palette (12 colors)
+                      - 'Spectral': Spectral diverging colormap
+                      - 'Rainbow': Rainbow colormap
+                      - 'Coolwarm': Cool-warm diverging colormap
+                      - 'Seismic': Seismic diverging colormap
+                      - Any other colormap name supported by the cmap library
+                      
+                      For a complete catalog of available colormaps, see:
+                      https://cmap-docs.readthedocs.io/en/latest/catalog/
+
+    Returns:
+        None (modifies self.samples_df in place)
+
+    Example:
+        # Set colors based on sample type
+        study.sample_color(by='sample_type', palette='Set1')
+        
+        # Set colors using a custom color list
+        study.sample_color(by=['#FF0000', '#00FF00', '#0000FF'])
+        
+        # Reset to default Turbo256 sequential colors
+        study.sample_color()
+    """
+    if self.samples_df is None or len(self.samples_df) == 0:
+        self.logger.warning("No samples found in study.")
+        return
+
+    sample_count = len(self.samples_df)
+    
+    # Handle custom color list
+    if isinstance(by, list):
+        if len(by) < sample_count:
+            self.logger.warning(f"Provided color list has {len(by)} colors but {sample_count} samples. Repeating colors.")
+            # Cycle through the provided colors if there aren't enough
+            colors = []
+            for i in range(sample_count):
+                colors.append(by[i % len(by)])
+        else:
+            colors = by[:sample_count]
+    else:
+        # Use the new approach: sample colors evenly from the whole colormap
+        if by is None:
+            # Sequential colors evenly sampled from the colormap
+            try:
+                colors = _sample_colors_from_colormap(palette, sample_count)
+            except ValueError as e:
+                self.logger.error(f"Error sampling colors from colormap: {e}")
+                return
+                
+        elif by == 'sample_uid':
+            # Use sample_uid to determine position in evenly sampled colormap
+            sample_uids = self.samples_df['sample_uid'].to_list()
+            try:
+                # Sample colors evenly for the number of samples
+                palette_colors = _sample_colors_from_colormap(palette, sample_count)
+                colors = []
+                for uid in sample_uids:
+                    # Use modulo to cycle through evenly sampled colors
+                    color_index = uid % len(palette_colors)
+                    colors.append(palette_colors[color_index])
+            except ValueError as e:
+                self.logger.error(f"Error sampling colors from colormap: {e}")
+                return
+                
+        elif by == 'sample_index':
+            # Use sample index (position in DataFrame) with evenly sampled colors
+            try:
+                colors = _sample_colors_from_colormap(palette, sample_count)
+            except ValueError as e:
+                self.logger.error(f"Error sampling colors from colormap: {e}")
+                return
+                
+        elif by == 'sample_type':
+            # Use sample_type to assign colors - same type gets same color
+            # Sample colors evenly across colormap for unique types
+            sample_types = self.samples_df['sample_type'].to_list()
+            unique_types = list(set([t for t in sample_types if t is not None]))
+            
+            try:
+                # Sample colors evenly for unique types
+                type_colors = _sample_colors_from_colormap(palette, len(unique_types))
+                type_to_color = {}
+                
+                for i, sample_type in enumerate(unique_types):
+                    type_to_color[sample_type] = type_colors[i]
+                
+                colors = []
+                for sample_type in sample_types:
+                    if sample_type is None:
+                        # Default to first color for None
+                        colors.append(type_colors[0] if type_colors else "#000000")
+                    else:
+                        colors.append(type_to_color[sample_type])
+            except ValueError as e:
+                self.logger.error(f"Error sampling colors from colormap: {e}")
+                return
+                    
+        elif by == 'sample_name':
+            # Use sample_name to assign colors - same name gets same color (unlikely but possible)
+            # Sample colors evenly across colormap for unique names
+            sample_names = self.samples_df['sample_name'].to_list()
+            unique_names = list(set([n for n in sample_names if n is not None]))
+            
+            try:
+                # Sample colors evenly for unique names
+                name_colors = _sample_colors_from_colormap(palette, len(unique_names))
+                name_to_color = {}
+                
+                for i, sample_name in enumerate(unique_names):
+                    name_to_color[sample_name] = name_colors[i]
+                
+                colors = []
+                for sample_name in sample_names:
+                    if sample_name is None:
+                        # Default to first color for None
+                        colors.append(name_colors[0] if name_colors else "#000000")
+                    else:
+                        colors.append(name_to_color[sample_name])
+            except ValueError as e:
+                self.logger.error(f"Error sampling colors from colormap: {e}")
+                return
+        else:
+            self.logger.error(f"Invalid by value: {by}. Must be 'sample_uid', 'sample_index', 'sample_type', 'sample_name', a list of colors, or None.")
+            return
+
+    # Update the sample_color column
+    self.samples_df = self.samples_df.with_columns(
+        pl.Series("sample_color", colors).alias("sample_color")
+    )
+    
+    if isinstance(by, list):
+        self.logger.debug(f"Set sample colors using provided color list ({len(by)} colors)")
+    elif by is None:
+        self.logger.debug(f"Set sequential sample colors using {palette} palette")
+    else:
+        self.logger.debug(f"Set sample colors based on {by} using {palette} palette")
+
+
+def sample_color_reset(self):
+    """
+    Reset sample colors to default coloring using the 'turbo' colormap.
+    
+    This function assigns colors by distributing samples evenly across the full
+    turbo colormap range, ensuring maximum color diversity and visual distinction
+    between samples.
+    
+    Returns:
+        None (modifies self.samples_df in place)
+    """
+    if self.samples_df is None or len(self.samples_df) == 0:
+        self.logger.warning("No samples found in study.")
+        return
+    
+    try:
+        from cmap import Colormap
+        
+        # Use turbo colormap
+        cm = Colormap('turbo')
+        
+        # Get sample count and assign colors evenly distributed across colormap
+        n_samples = len(self.samples_df)
+        colors = []
+        
+        # Distribute samples evenly across the full colormap range
+        for i in range(n_samples):
+            # Evenly distribute samples across colormap (avoiding endpoints to prevent white/black)
+            normalized_value = (i + 0.5) / n_samples  # +0.5 to center samples in their bins
+            # Optionally, map to a subset of colormap to avoid extreme colors
+            # Use 10% to 90% of colormap range for better color diversity
+            normalized_value = 0.1 + (normalized_value * 0.8)
+            
+            color_rgba = cm(normalized_value)
+            
+            # Convert RGBA to hex
+            if len(color_rgba) >= 3:
+                r, g, b = color_rgba[:3]
+                # Convert to 0-255 range if needed
+                if max(color_rgba[:3]) <= 1.0:
+                    r, g, b = int(r * 255), int(g * 255), int(b * 255)
+                hex_color = f"#{r:02x}{g:02x}{b:02x}"
+                colors.append(hex_color)
+        
+        # Update the sample_color column
+        self.samples_df = self.samples_df.with_columns(
+            pl.Series("sample_color", colors).alias("sample_color")
+        )
+        
+        self.logger.debug(f"Reset sample colors using turbo colormap with even distribution ({n_samples} samples)")
+        
+    except ImportError:
+        self.logger.error("cmap library is required for sample color reset. Install with: pip install cmap")
+    except Exception as e:
+        self.logger.error(f"Failed to reset sample colors: {e}")
+
+
+def _get_color_palette(palette_name):
+    """
+    Get color palette as a list of hex color codes using the cmap library.
+    
+    Parameters:
+        palette_name (str): Name of the palette
+        
+    Returns:
+        list: List of hex color codes
+        
+    Raises:
+        ValueError: If palette_name is not supported
+    """
+    try:
+        from cmap import Colormap
+    except ImportError:
+        raise ValueError("cmap library is required for color palettes. Install with: pip install cmap")
+    
+    # Map common palette names to cmap names
+    palette_mapping = {
+        # Scientific colormaps
+        "Turbo256": "turbo",
+        "Viridis256": "viridis", 
+        "Plasma256": "plasma",
+        "Inferno256": "inferno",
+        "Magma256": "magma",
+        "Cividis256": "cividis",
+        
+        # Qualitative palettes
+        "Set1": "Set1",
+        "Set2": "Set2", 
+        "Set3": "Set3",
+        "Tab10": "tab10",
+        "Tab20": "tab20",
+        "Dark2": "Dark2",
+        "Paired": "Paired",
+        
+        # Additional useful palettes
+        "Spectral": "Spectral",
+        "Rainbow": "rainbow",
+        "Coolwarm": "coolwarm",
+        "Seismic": "seismic",
+    }
+    
+    # Get the cmap name
+    cmap_name = palette_mapping.get(palette_name, palette_name.lower())
+    
+    try:
+        # Create colormap
+        cm = Colormap(cmap_name)
+        
+        # Determine number of colors to generate
+        if "256" in palette_name:
+            n_colors = 256
+        elif palette_name in ["Set1"]:
+            n_colors = 9
+        elif palette_name in ["Set2", "Dark2"]:
+            n_colors = 8
+        elif palette_name in ["Set3", "Paired"]:
+            n_colors = 12
+        elif palette_name in ["Tab10"]:
+            n_colors = 10
+        elif palette_name in ["Tab20"]:
+            n_colors = 20
+        else:
+            n_colors = 256  # Default for continuous colormaps
+        
+        # Generate colors
+        if n_colors <= 20:
+            # For discrete palettes, use evenly spaced indices
+            indices = [i / (n_colors - 1) for i in range(n_colors)]
+        else:
+            # For continuous palettes, use full range
+            indices = [i / (n_colors - 1) for i in range(n_colors)]
+        
+        # Get colors as RGBA and convert to hex
+        colors = cm(indices)
+        hex_colors = []
+        
+        for color in colors:
+            if len(color) >= 3:  # RGBA or RGB
+                r, g, b = color[:3]
+                # Convert to 0-255 range if needed
+                if max(color[:3]) <= 1.0:
+                    r, g, b = int(r * 255), int(g * 255), int(b * 255)
+                hex_color = f"#{r:02x}{g:02x}{b:02x}"
+                hex_colors.append(hex_color)
+        
+        return hex_colors
+        
+    except Exception as e:
+        raise ValueError(f"Failed to create colormap '{cmap_name}': {e}. "
+                        f"Available palettes: {list(palette_mapping.keys())}")
+
+
+def _sample_colors_from_colormap(palette_name, n_colors):
+    """
+    Sample colors evenly from the whole colormap range, similar to sample_color_reset.
+    
+    Parameters:
+        palette_name (str): Name of the palette/colormap
+        n_colors (int): Number of colors to sample
+        
+    Returns:
+        list: List of hex color codes sampled evenly from the colormap
+        
+    Raises:
+        ValueError: If palette_name is not supported
+    """
+    try:
+        from cmap import Colormap
+    except ImportError:
+        raise ValueError("cmap library is required for color palettes. Install with: pip install cmap")
+    
+    # Map common palette names to cmap names (same as _get_color_palette)
+    palette_mapping = {
+        # Scientific colormaps
+        "Turbo256": "turbo",
+        "Viridis256": "viridis", 
+        "Plasma256": "plasma",
+        "Inferno256": "inferno",
+        "Magma256": "magma",
+        "Cividis256": "cividis",
+        
+        # Qualitative palettes
+        "Set1": "Set1",
+        "Set2": "Set2", 
+        "Set3": "Set3",
+        "Tab10": "tab10",
+        "Tab20": "tab20",
+        "Dark2": "Dark2",
+        "Paired": "Paired",
+        
+        # Additional useful palettes
+        "Spectral": "Spectral",
+        "Rainbow": "rainbow",
+        "Coolwarm": "coolwarm",
+        "Seismic": "seismic",
+    }
+    
+    # Get the cmap name
+    cmap_name = palette_mapping.get(palette_name, palette_name.lower())
+    
+    try:
+        # Create colormap
+        cm = Colormap(cmap_name)
+        
+        colors = []
+        
+        # Distribute samples evenly across the full colormap range (same approach as sample_color_reset)
+        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
+            # Map to a subset of colormap to avoid extreme colors (use 10% to 90% range)
+            normalized_value = 0.1 + (normalized_value * 0.8)
+            
+            color_rgba = cm(normalized_value)
+            
+            # Convert RGBA to hex
+            if len(color_rgba) >= 3:
+                r, g, b = color_rgba[:3]
+                # Convert to 0-255 range if needed
+                if max(color_rgba[:3]) <= 1.0:
+                    r, g, b = int(r * 255), int(g * 255), int(b * 255)
+                hex_color = f"#{r:02x}{g:02x}{b:02x}"
+                colors.append(hex_color)
+        
+        return colors
+        
+    except Exception as e:
+        raise ValueError(f"Failed to create colormap '{cmap_name}': {e}. "
+                        f"Available palettes: {list(palette_mapping.keys())}")
+
+
+def _matplotlib_to_hex(color_dict):
+    """Convert matplotlib color dictionary to list of hex colors."""
+    return list(color_dict.values())
+
+
+# =====================================================================================
+# SCHEMA AND DATA STRUCTURE FUNCTIONS  
+# =====================================================================================
+
+
+def _ensure_features_df_schema_order(self):
+    """
+    Ensure features_df columns are ordered according to study5_schema.json.
+    
+    This method should be called after operations that might scramble the column order.
+    """
+    if self.features_df is None or self.features_df.is_empty():
+        return
+    
+    try:
+        import os
+        import json
+        from masster.study.h5 import _reorder_columns_by_schema
+        
+        # Load schema
+        schema_path = os.path.join(os.path.dirname(__file__), "study5_schema.json")
+        with open(schema_path, 'r') as f:
+            schema = json.load(f)
+        
+        # Reorder columns to match schema
+        self.features_df = _reorder_columns_by_schema(self.features_df, schema, 'features_df')
+        
+    except Exception as e:
+        self.logger.warning(f"Failed to reorder features_df columns: {e}")
+
+
+def migrate_map_id_to_index(self):
+    """
+    Migrate map_id from string-based OpenMS unique IDs to integer indices.
+    
+    This function converts the map_id column from string type (with OpenMS unique IDs)
+    to integer type where each map_id corresponds to the index of the feature map
+    in self.features_maps.
+    
+    This migration is needed for studies that were created before the map_id format
+    change from OpenMS unique IDs to feature map indices.
+    """
+    if self.samples_df is None or self.samples_df.is_empty():
+        self.logger.warning("No samples to migrate")
+        return
+    
+    # Check if migration is needed
+    current_dtype = self.samples_df['map_id'].dtype
+    if current_dtype == pl.Int64:
+        self.logger.info("map_id column is already Int64 type - no migration needed")
+        return
+    
+    self.logger.info("Migrating map_id from string-based OpenMS IDs to integer indices")
+    
+    # Create new map_id values based on sample order
+    # Each sample gets a map_id that corresponds to its position in features_maps
+    sample_count = len(self.samples_df)
+    new_map_ids = list(range(sample_count))
+    
+    # Update the map_id column
+    self.samples_df = self.samples_df.with_columns(
+        pl.lit(new_map_ids).alias("map_id")
+    )
+    
+    # Ensure the column is Int64 type
+    self.samples_df = self.samples_df.cast({"map_id": pl.Int64})
+    
+    self.logger.info(f"Successfully migrated {sample_count} samples to indexed map_id format")
+    self.logger.info(f"map_id now ranges from 0 to {sample_count - 1}")