masster 0.3.14__py3-none-any.whl → 0.3.16__py3-none-any.whl

masster/study/helpers.py

@@ -211,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.
 
@@ -225,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")
 
@@ -1298,7 +1305,7 @@ def compress_chrom(self):
 # =====================================================================================
 
 
-def name_replace(self, replace_dict):
+def sample_name_replace(self, replace_dict):
     """
     Replace sample names in samples_df based on a dictionary mapping.
     
@@ -1364,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.
     
@@ -1446,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 = []
@@ -1924,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)
@@ -1943,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
@@ -1957,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)}",
         )
@@ -1969,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)}",
         )
@@ -2170,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
 
 
@@ -2276,6 +2356,357 @@ def consensus_delete(self, consensus):
     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
 # =====================================================================================
@@ -2712,3 +3143,43 @@ def _ensure_features_df_schema_order(self):
         
     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}")