masster 0.5.1__py3-none-any.whl → 0.5.4__py3-none-any.whl

masster/study/helpers.py

@@ -362,10 +362,6 @@ def get_chrom(self, uids=None, samples=None):
     ids = self._get_consensus_uids(uids)
     sample_uids = self._get_samples_uids(samples)
 
-    if self.consensus_map is None:
-        self.logger.error("No consensus map found.")
-        return None
-
     # Pre-filter all DataFrames to reduce join sizes
     filtered_consensus_mapping = self.consensus_mapping_df.filter(
         pl.col("consensus_uid").is_in(ids),
@@ -529,30 +525,132 @@ def get_consensus(self, quant="chrom_area"):
     return df
 
 
-# TODO I don't get this param
-def get_consensus_matrix(self, quant="chrom_area"):
+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.
-    Optimized implementation that avoids expensive join operations.
+    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  
+                - 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.",
-        )
+        self.logger.error(f"Quantification method {quant} not found in features_df.")
         return None
 
-    # Create a lookup dictionary from features_df for O(1) value access
-    feature_values = {}
+    # 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()
+
+    # Filter datasets upfront to reduce processing load
+    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))
+    )
+    
+    # 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")
+    )
+    
+    # 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"))
+    )
+    
+    # 4. Pivot to create the matrix format
+    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]
+    ])
+    
+    return matrix_df
+
+
+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 map 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
+
+    # 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()
+
+    # Create a lookup dictionary from features_df for gap detection (exclude filled features)
+    # Key: (feature_uid, sample_uid) -> Value: 1 (detected)
+    feature_detection = {}
     for row in self.features_df.iter_rows(named=True):
-        feature_uid = row["feature_uid"]
         sample_uid = row["sample_uid"]
-        value = row[quant] if row[quant] is not None else 0
-        feature_values[(feature_uid, sample_uid)] = value
+        if sample_uid in sample_uids:  # Only include specified samples
+            # 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
 
-    # Build consensus matrix directly using the consensus_mapping_df
+    # Build gaps matrix directly using the consensus_mapping_df
     matrix_dict = {}
     sample_mapping = dict(
-        self.samples_df.select(["sample_uid", "sample_name"]).iter_rows(),
+        self.samples_df.filter(pl.col("sample_uid").is_in(sample_uids))
+        .select(["sample_uid", "sample_name"]).iter_rows(),
     )
 
     for row in self.consensus_mapping_df.iter_rows(named=True):
@@ -560,65 +658,53 @@ def get_consensus_matrix(self, quant="chrom_area"):
         sample_uid = row["sample_uid"]
         feature_uid = row["feature_uid"]
 
-        # Look up the quantification value
+        # Only process samples and consensus features in our filtered lists
+        if sample_uid not in sample_uids or consensus_uid not in uids:
+            continue
+
+        # Check if feature was detected (not filled)
         key = (feature_uid, sample_uid)
-        value = feature_values.get(key, 0)
+        detected = feature_detection.get(key, 0)  # 0 if not found (gap), 1 if detected
 
         if consensus_uid not in matrix_dict:
             matrix_dict[consensus_uid] = {}
 
         sample_name = sample_mapping.get(sample_uid, f"sample_{sample_uid}")
 
-        # Take max if multiple features map to same consensus/sample combination
+        # For gaps matrix, we want to know if ANY feature was detected for this consensus/sample
+        # So we take max (if any feature is detected, the consensus feature is detected)
         if sample_name in matrix_dict[consensus_uid]:
             matrix_dict[consensus_uid][sample_name] = max(
                 matrix_dict[consensus_uid][sample_name],
-                value,
+                detected,
             )
         else:
-            matrix_dict[consensus_uid][sample_name] = value
-
-    # Convert to Polars DataFrame with proper formatting
-    import polars as pl
+            matrix_dict[consensus_uid][sample_name] = detected
 
-    # Convert matrix_dict to list of records for Polars
+    # Convert to Polars DataFrame
     records = []
     for consensus_uid, sample_values in matrix_dict.items():
         record = {"consensus_uid": consensus_uid}
         record.update(sample_values)
         records.append(record)
 
+    if not records:
+        self.logger.warning("No gaps data found for specified consensus features and samples")
+        return pl.DataFrame()
+
     # Create Polars DataFrame and set proper data types
-    df2 = pl.DataFrame(records)
+    df_gaps = pl.DataFrame(records)
 
-    # Fill null values with 0 and round numeric columns
-    numeric_cols = [col for col in df2.columns if col != "consensus_uid"]
-    df2 = df2.with_columns(
+    # Fill null values with 0 (gaps) and ensure integer type for gap indicators
+    numeric_cols = [col for col in df_gaps.columns if col != "consensus_uid"]
+    df_gaps = df_gaps.with_columns(
         [
             pl.col("consensus_uid").cast(pl.UInt64),
-            *[pl.col(col).fill_null(0).round(0) for col in numeric_cols],
+            *[pl.col(col).fill_null(0).cast(pl.Int8) for col in numeric_cols],
         ],
     )
 
-    return df2
-
-
-def get_gaps_matrix(self, uids=None):
-    """
-    Get a matrix of gaps between consensus features with samples as columns and consensus features as rows.
-    """
-    if self.consensus_df is None:
-        self.logger.error("No consensus map found.")
-        return None
-    uids = self._get_consensus_uids(uids)
-
-    df1 = self.get_consensus_matrix(quant="filled")
-    if df1 is None or df1.empty:
-        self.logger.warning("No gap data found.")
-        return None
-    # keep only rows where consensus_id is in ids - use pandas indexing since df1 is already pandas
-    df1 = df1[df1.index.isin(uids)]
-    return df1
+    return df_gaps
 
 
 def get_gaps_stats(self, uids=None):
@@ -705,6 +791,53 @@ 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_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
+    if hasattr(self, "history") and self.history:
+        for key in keys_to_remove:
+            if key in self.history:
+                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}",
+    )
+
+
 def fill_reset(self):
     # remove all features with filled=True
     if self.features_df is None:
@@ -719,7 +852,7 @@ def fill_reset(self):
         pl.col("feature_uid").is_in(feature_uids_to_keep),
     )
     self.logger.info(
-        f"Reset filled chromatograms. Chroms removed: {l1 - len(self.features_df)}",
+        f"Removed {l1 - len(self.features_df)} gap-filled features",
     )
 
 
@@ -1044,6 +1177,250 @@ def get_sample_stats(self):
     )
 
 
+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
+            - rt: Retention time
+            - rt_delta_mean: Mean retention time delta
+            - mz: Mass-to-charge ratio
+            - mz_range: Mass range (mz_max - mz_min)
+            - log10_inty_mean: Log10 of mean intensity
+            - number_samples: Number of samples
+            - number_ms2: Number of MS2 spectra
+            - charge_mean: Mean charge
+            - quality: Feature quality
+            - chrom_coherence_mean: Mean chromatographic coherence
+            - chrom_height_scaled_mean: Mean scaled chromatographic height
+            - chrom_prominence_scaled_mean: Mean scaled chromatographic prominence
+            - qc_ratio: Ratio of QC samples where feature was detected
+            - qc_cv: RSD (relative standard deviation) of intensity for QC samples
+            - qc_to_blank: Ratio of average QC intensity to average blank intensity
+    """
+    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.")
+        return pl.DataFrame()
+
+    # Get all columns and their data types - work with original dataframe
+    data_df = self.consensus_df.clone()
+
+    # Define specific columns to include in the exact order requested
+    desired_columns = [
+        "consensus_uid",  # Include consensus_uid for identification
+        "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"
+    ]
+    
+    # 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)}")
+        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_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")
+                        )
+                        
+                        # Join with result DataFrame
+                        result_df = result_df.join(
+                            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():
+                    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...")
+                    
+                    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"])
+                            
+                            # 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"])
+                            
+                            # Join QC and blank averages and calculate ratio
+                            qc_blank_ratios = qc_averages.join(
+                                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)
+                                .then(pl.col("qc_avg") / pl.col("blank_avg"))
+                                .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_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)
+                            )
+                        
+                        # 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")
+                        )
+                        
+                        # Join with result DataFrame
+                        result_df = result_df.join(
+                            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:
+            result_df = result_df.with_columns(pl.lit(None, dtype=pl.Float64).alias("qc_ratio"))
+        if qc_cv_values is not None:
+            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")
+        ])
+    
+    return result_df
+
+
 # =====================================================================================
 # DATA COMPRESSION AND RESTORATION FUNCTIONS
 # =====================================================================================
@@ -2131,7 +2508,7 @@ def _apply_chunked_select(self, filter_expr, chunk_size: int):
     else:
         return pl.DataFrame()
 
-
+'''
 def features_select_benchmarked(
     self,
     mz=None,
@@ -2224,7 +2601,7 @@ def monkey_patch_study():
     Study.features_select_benchmarked = features_select_benchmarked
 
     print("Patched Study.features_select with consolidated optimized implementation")
-
+'''
 
 def features_filter(
     self,