masster 0.5.28__py3-none-any.whl → 0.6.2__py3-none-any.whl

masster/study/importers.py

@@ -2,12 +2,13 @@
 import.py
 
 Module providing import functionality for Study class, specifically for importing
-oracle identification data into consensus features.
+oracle and TIMA identification data into consensus features.
 """
 
 from __future__ import annotations
 
 import os
+import glob
 import pandas as pd
 import polars as pl
 
@@ -320,3 +321,385 @@ def import_oracle(
             "id_matches": len(self.id_df),
         },
     )
+
+
+def import_tima(self, folder, file="results_annotation"):
+    """
+    Import TIMA identification data and map it to consensus features.
+
+    This method reads TIMA identification results from folder/results_annotation_*.tsv
+    and creates lib_df and id_df DataFrames with detailed library and identification information.
+    It also updates consensus_df with top identification results.
+
+    Parameters:
+        folder (str): Path to TIMA folder containing results_annotation_*.tsv files
+        file (str, optional): Base name of TIMA results file (default: "results_annotation")
+
+    Returns:
+        None: Updates consensus_df, creates lib_df and id_df in-place with TIMA identification data
+
+    Raises:
+        FileNotFoundError: If the TIMA results file doesn't exist
+        ValueError: If consensus_df is empty or doesn't have required columns
+
+    Example:
+        >>> study.import_tima(folder="path/to/tima_results")
+    """
+
+    self.logger.info(f"Starting TIMA import from folder: {folder}")
+
+    # Validate inputs
+    if self.consensus_df is None or self.consensus_df.is_empty():
+        raise ValueError("consensus_df is empty or not available. Run merge() first.")
+
+    if "consensus_id" not in self.consensus_df.columns:
+        raise ValueError("consensus_df must contain 'consensus_id' column")
+
+    # Find TIMA file
+    tima_pattern = os.path.join(folder, f"*{file}*.tsv")
+    tima_files = glob.glob(tima_pattern)
+
+    if not tima_files:
+        raise FileNotFoundError(f"TIMA results file not found with pattern: {tima_pattern}")
+
+    tima_file_path = tima_files[0]
+    self.logger.debug(f"Loading TIMA data from: {tima_file_path}")
+
+    try:
+        # Read TIMA data using polars
+        tima_data = pl.read_csv(
+            tima_file_path,
+            separator="\t",
+            schema_overrides={
+                "feature_id": pl.Utf8,  # Read as Utf8 string
+            },
+            infer_schema_length=10000,
+        )
+        self.logger.info(f"TIMA data loaded successfully with {len(tima_data)} rows")
+    except Exception as e:
+        self.logger.error(f"Could not read {tima_file_path}: {e}")
+        raise
+
+    # Check if TIMA feature_ids match consensus_df consensus_id column
+    if "consensus_id" not in self.consensus_df.columns:
+        raise ValueError("consensus_df must contain 'consensus_id' column")
+
+    # Compare TIMA feature_ids with consensus_df consensus_ids
+    consensus_ids = set(self.consensus_df["consensus_id"].to_list())
+    tima_ids = set(tima_data["feature_id"].to_list())
+
+    matching_ids = consensus_ids.intersection(tima_ids)
+    non_matching_ids = tima_ids - consensus_ids
+
+    if non_matching_ids:
+        self.logger.warning(
+            f"Found {len(non_matching_ids)} feature_ids in TIMA data that do not match any consensus_id in consensus_df. "
+            f"These will be filtered out. Matching features: {len(matching_ids)}/{len(tima_ids)}"
+        )
+        # Filter to only matching feature_ids
+        tima_data = tima_data.filter(pl.col("feature_id").is_in(list(consensus_ids)))
+
+    if len(tima_data) == 0:
+        self.logger.error("No TIMA feature_ids match consensus_df consensus_id values")
+        raise ValueError("No matching features found between TIMA data and consensus_df")
+
+    self.logger.debug(f"Matched {len(tima_data)} TIMA entries to consensus_df consensus_id values")
+
+    # Filter to only rows with identification data (non-empty label_compound)
+    initial_count = len(tima_data)
+    tima_data = tima_data.filter(
+        pl.col("label_compound").is_not_null() & (pl.col("label_compound").cast(pl.Utf8).str.strip_chars() != "")
+    )
+
+    self.logger.debug(f"Filtered to {len(tima_data)}/{initial_count} TIMA entries with identifications")
+
+    if len(tima_data) == 0:
+        self.logger.warning("No TIMA entries with identifications found")
+        return
+
+    # === CREATE LIB_DF ===
+    self.logger.debug("Creating lib_df from TIMA annotation data")
+    self.logger.debug(f"TIMA data shape before lib_df creation: {tima_data.shape}")
+
+    # Create unique lib_uid for each library entry
+    tima_data = tima_data.with_columns(pl.arange(0, len(tima_data)).alias("lib_uid"))
+
+    # Map TIMA columns to lib_df schema
+    lib_data = []
+    for row in tima_data.iter_rows(named=True):
+        # Extract z (charge) from adduct
+        z = None
+        adduct_str = str(row.get("adduct", ""))
+        if "+" in adduct_str:
+            z = 1
+        elif "-" in adduct_str:
+            z = -1
+
+        # Get SMILES
+        smiles = row.get("smiles_no_stereo", None)
+        if smiles is None or (isinstance(smiles, str) and smiles.strip() == ""):
+            smiles = None
+
+        # Calculate InChI from SMILES if available
+        inchi = None
+        if smiles:
+            try:
+                from rdkit import Chem
+
+                mol_rdkit = Chem.MolFromSmiles(smiles)
+                if mol_rdkit:
+                    inchi = Chem.MolToInchi(mol_rdkit)
+            except ImportError:
+                pass  # RDKit not available
+            except Exception:
+                pass
+
+        # Calculate formula from SMILES if available
+        formula = None
+        if smiles:
+            try:
+                from rdkit import Chem
+
+                mol_rdkit = Chem.MolFromSmiles(smiles)
+                if mol_rdkit:
+                    formula = Chem.rdMolDescriptors.CalcMolFormula(mol_rdkit)
+            except ImportError:
+                pass  # RDKit not available
+            except Exception:
+                pass
+
+        # Calculate mass from m/z and charge
+        m = None
+        mz_value = row.get("mz", None)
+        if mz_value is not None and z is not None:
+            try:
+                m = float(mz_value) * abs(z)
+            except (ValueError, TypeError):
+                pass
+
+        # Get class and clean NaN values
+        class_value = row.get("label_classyfire", None)
+        if class_value is None or (isinstance(class_value, str) and class_value.upper() == "NAN"):
+            class_value = None
+
+        lib_entry = {
+            "lib_uid": row["lib_uid"],
+            "cmpd_uid": row["lib_uid"],  # Use lib_uid as compound identifier
+            "source_id": None,  # Leave empty as requested
+            "name": row.get("label_compound", None),
+            "shortname": None,  # Not available in TIMA data
+            "class": class_value,
+            "smiles": smiles,
+            "inchi": inchi,
+            "inchikey": row.get("inchikey_connectivity_layer", None),
+            "formula": formula,
+            "iso": 0,  # Fixed isotope value
+            "adduct": row.get("adduct", None),
+            "probability": row.get("score", None),
+            "m": m,
+            "z": z,
+            "mz": row.get("mz", None),
+            "rt": None,  # Set to null as requested
+            "quant_group": None,
+            "db_id": None,  # Not available in TIMA data
+            "db": row.get("library", None),
+        }
+        lib_data.append(lib_entry)
+
+    self.logger.debug(f"Created {len(lib_data)} lib_data entries")
+
+    # Create lib_df as Polars DataFrame with error handling for mixed types
+    try:
+        lib_df_temp = pl.DataFrame(lib_data)
+    except Exception as e:
+        self.logger.warning(f"Error creating lib_df with polars: {e}")
+        # Fallback: convert to pandas first, then to polars
+        lib_df_pandas = pd.DataFrame(lib_data)
+        lib_df_temp = pl.from_pandas(lib_df_pandas)
+
+    # Ensure uniqueness by name and adduct combination
+    # Sort by lib_uid and keep first occurrence (earliest in processing order)
+    self.lib_df = lib_df_temp.sort("lib_uid").unique(subset=["name", "adduct"], keep="first")
+
+    self.logger.info(
+        f"Created lib_df with {len(self.lib_df)} library entries ({len(lib_data) - len(self.lib_df)} duplicates removed)"
+    )
+
+    # === CREATE ID_DF ===
+    self.logger.debug("Creating id_df from TIMA identification matches")
+
+    # Create a mapping from consensus_id to consensus_uid
+    # TIMA data has feature_id which matches consensus_id, map to consensus_uid for id_df
+    consensus_id_to_uid_map = dict(
+        zip(self.consensus_df["consensus_id"].to_list(), self.consensus_df["consensus_uid"].to_list())
+    )
+
+    # Create identification matches
+    id_data = []
+    for row in tima_data.iter_rows(named=True):
+        # Map TIMA feature_id to consensus_df consensus_uid
+        tima_feature_id = row["feature_id"]
+        consensus_uid = consensus_id_to_uid_map.get(tima_feature_id)
+
+        if consensus_uid is None:
+            # Skip if we can't find the mapping (shouldn't happen after filtering)
+            continue
+
+        # Use error_mz for mz_delta
+        mz_delta = None
+        error_mz = row.get("error_mz", None)
+        if error_mz is not None:
+            try:
+                mz_delta = float(error_mz)
+            except (ValueError, TypeError):
+                pass
+
+        # Use error_rt for rt_delta
+        rt_delta = None
+        rt_err_value = row.get("error_rt", None)
+        if rt_err_value is not None:
+            try:
+                rt_delta = float(rt_err_value)
+            except (ValueError, TypeError):
+                pass
+
+        # Create matcher as "tima-" + library
+        matcher = "tima"  # default fallback
+        library_value = row.get("library", None)
+        if library_value is not None:
+            try:
+                library = str(library_value)
+                matcher = f"tima-{library}"
+            except (ValueError, TypeError):
+                pass
+
+        id_entry = {
+            "consensus_uid": consensus_uid,  # Use mapped consensus_uid from consensus_df
+            "lib_uid": row["lib_uid"],
+            "mz_delta": mz_delta,
+            "rt_delta": rt_delta,
+            "matcher": matcher,
+            "score": row.get("score", None),
+        }
+        id_data.append(id_entry)
+
+    # Create id_df as Polars DataFrame with explicit schema to avoid inference issues
+    # Match consensus_uid type to consensus_df
+    consensus_uid_dtype = self.consensus_df["consensus_uid"].dtype
+    id_schema = {
+        "consensus_uid": consensus_uid_dtype,  # Match the type from consensus_df
+        "lib_uid": pl.Int64,
+        "mz_delta": pl.Float64,
+        "rt_delta": pl.Float64,
+        "matcher": pl.Utf8,
+        "score": pl.Float64,
+    }
+    id_df_temp = pl.DataFrame(id_data, schema=id_schema)
+
+    # Filter id_df to only include lib_uids that exist in the final unique lib_df
+    unique_lib_uids = self.lib_df.select("lib_uid").to_series()
+    self.id_df = id_df_temp.filter(pl.col("lib_uid").is_in(unique_lib_uids))
+
+    self.logger.info(f"Created id_df with {len(self.id_df)} identification matches")
+
+    # === UPDATE CONSENSUS_DF ===
+    self.logger.debug("Updating consensus_df with top identification results")
+
+    # tima_data is already a polars DataFrame
+    tima_pl = tima_data
+
+    # Group by feature_id and select the best identification (highest score)
+    # In case of ties, take the first one
+    best_ids = (
+        tima_pl.group_by("feature_id")
+        .agg([pl.col("score").max().alias("max_score")])
+        .join(tima_pl, on="feature_id")
+        .filter(pl.col("score") == pl.col("max_score"))
+        .group_by("feature_id")
+        .first()  # In case of ties, take the first
+    )
+
+    # Join with consensus_df to map consensus_id to consensus_uid
+    best_ids = best_ids.join(
+        self.consensus_df.select(["consensus_id", "consensus_uid"]), left_on="feature_id", right_on="consensus_id", how="left"
+    )
+
+    self.logger.debug(f"Selected best identifications for {len(best_ids)} consensus features")
+
+    # Prepare the identification columns
+    id_columns = {
+        "id_top_name": best_ids.select("consensus_uid", "label_compound"),
+        "id_top_adduct": best_ids.select("consensus_uid", "adduct"),
+        "id_top_class": best_ids.select("consensus_uid", "label_classyfire"),
+        "id_top_score": best_ids.select("consensus_uid", pl.col("score").round(3).alias("score")),
+    }
+
+    # Initialize identification columns in consensus_df if they don't exist
+    for col_name in id_columns.keys():
+        if col_name not in self.consensus_df.columns:
+            if col_name == "id_top_score":
+                self.consensus_df = self.consensus_df.with_columns(pl.lit(None, dtype=pl.Float64).alias(col_name))
+            else:
+                self.consensus_df = self.consensus_df.with_columns(pl.lit(None, dtype=pl.String).alias(col_name))
+
+    # Update consensus_df with TIMA identifications
+    for col_name, id_data_col in id_columns.items():
+        tima_column = id_data_col.columns[1]  # second column (after consensus_uid)
+
+        # Create update dataframe
+        update_data = id_data_col.rename({tima_column: col_name})
+
+        # Join and update
+        self.consensus_df = (
+            self.consensus_df.join(update_data, on="consensus_uid", how="left", suffix="_tima")
+            .with_columns(pl.coalesce([f"{col_name}_tima", col_name]).alias(col_name))
+            .drop(f"{col_name}_tima")
+        )
+
+    # Replace NaN values with None in identification columns
+    id_col_names = ["id_top_name", "id_top_adduct", "id_top_class", "id_top_score"]
+    for col_name in id_col_names:
+        if col_name in self.consensus_df.columns:
+            # For string columns, replace empty strings and "nan" with None
+            if col_name != "id_top_score":
+                self.consensus_df = self.consensus_df.with_columns(
+                    pl.when(
+                        pl.col(col_name).is_null()
+                        | (pl.col(col_name) == "")
+                        | (pl.col(col_name) == "nan")
+                        | (pl.col(col_name) == "NaN")
+                    )
+                    .then(None)
+                    .otherwise(pl.col(col_name))
+                    .alias(col_name)
+                )
+            # For numeric columns, replace NaN with None
+            else:
+                self.consensus_df = self.consensus_df.with_columns(
+                    pl.when(pl.col(col_name).is_null() | pl.col(col_name).is_nan())
+                    .then(None)
+                    .otherwise(pl.col(col_name))
+                    .alias(col_name)
+                )
+
+    # Count how many consensus features were updated
+    updated_count = self.consensus_df.filter(pl.col("id_top_name").is_not_null()).height
+    total_consensus = len(self.consensus_df)
+
+    self.logger.success(
+        f"TIMA import completed. {updated_count}/{total_consensus} "
+        f"consensus features now have identifications ({updated_count / total_consensus * 100:.1f}%)"
+    )
+
+    # Update history
+    self.update_history(
+        ["import_tima"],
+        {
+            "folder": folder,
+            "file": file,
+            "updated_features": updated_count,
+            "total_features": total_consensus,
+            "lib_entries": len(self.lib_df),
+            "id_matches": len(self.id_df),
+        },
+    )

masster/study/load.py

@@ -191,17 +191,6 @@ def load(self, filename=None):
 
     _load_study5(self, filename)
 
-    # After loading the study, check if we have consensus features before loading consensus XML
-    # if (self.consensus_df is not None and not self.consensus_df.is_empty()):
-    #    consensus_xml_path = filename.replace(".study5", ".consensusXML")
-    #    if os.path.exists(consensus_xml_path):
-    #        self._load_consensusXML(filename=consensus_xml_path)
-    # self.logger.info(f"Automatically loaded consensus from {consensus_xml_path}")
-    #    else:
-    #        self.logger.warning(f"No consensus XML file found at {consensus_xml_path}")
-    # else:
-    #    self.logger.debug("No consensus features found, skipping consensusXML loading")
-
     self.filename = filename
 
 

masster/study/merge.py

@@ -441,9 +441,15 @@ def merge(study, **kwargs) -> None:
     cached_valid_adducts = None
     try:
         cached_adducts_df = study._get_adducts()
+        # Remove all adducts with wrong polarity
+        if study.polarity == "positive":
+            cached_adducts_df = cached_adducts_df.filter(pl.col("charge") >= 0)
+        else:
+            cached_adducts_df = cached_adducts_df.filter(pl.col("charge") <= 0)
         if not cached_adducts_df.is_empty():
             cached_valid_adducts = set(cached_adducts_df["name"].to_list())
         else:
+            study.logger.warning(f"No valid adducts found for polarity '{study.polarity}'")
             cached_valid_adducts = set()
     except Exception as e:
         study.logger.warning(f"Could not retrieve study adducts: {e}")
@@ -452,6 +458,13 @@ def merge(study, **kwargs) -> None:
     # Always allow '?' adducts
     cached_valid_adducts.add("?")
 
+    # Bypass for single sample case
+    if len(study.samples_df) == 1:
+        study.logger.info("Single sample detected - bypassing merge algorithm and using direct feature mapping")
+        _handle_single_sample_merge(study, cached_adducts_df, cached_valid_adducts)
+        # Skip all post-processing for single sample case
+        return
+
     # Route to algorithm implementation
     if params.method == "kd":
         consensus_map = _merge_kd(study, params)
@@ -1719,6 +1732,10 @@ def _calculate_consensus_statistics(
         mz_values: m/z values from chunk consensus features
         intensity_values: Intensity values from chunk consensus features
         quality_values: Quality values from chunk consensus features
+        number_features: Number of unique features contributing
+        number_samples: Number of unique samples contributing
+        cached_adducts_df: Cached DataFrame of valid adducts for the study
+        cached_valid_adducts: Cached set of valid adduct names for the study
 
     Returns:
         Dictionary with consensus feature metadata
@@ -3612,6 +3629,142 @@ def __merge_adduct_grouping(study, consensus_data, rt_tol, mz_tol):
     return adduct_group_list, adduct_of_list
 
 
+def _handle_single_sample_merge(study, cached_adducts_df=None, cached_valid_adducts=None):
+    """
+    Handle merge for the special case of a single sample.
+    Directly populate consensus_df from the sample's features_df without any filtering.
+    
+    Args:
+        study: Study object with single sample
+        cached_adducts_df: Pre-computed adducts DataFrame (optional)
+        cached_valid_adducts: Set of valid adduct names (optional)
+    """
+    import polars as pl
+    import uuid
+    
+    if len(study.samples_df) != 1:
+        raise ValueError("_handle_single_sample_merge should only be called with exactly one sample")
+    
+    # Get the single sample's features
+    sample_row = study.samples_df.row(0, named=True)
+    sample_uid = sample_row["sample_uid"]
+    
+    # Filter features for this sample  
+    sample_features = study.features_df.filter(pl.col("sample_uid") == sample_uid)
+    
+    if len(sample_features) == 0:
+        study.logger.warning("No features found for single sample")
+        study.consensus_df = pl.DataFrame()
+        study.consensus_mapping_df = pl.DataFrame()
+        return
+    
+    study.logger.info(f"Creating consensus from {len(sample_features)} features in single sample")
+    
+    # Create consensus features directly from sample features
+    consensus_list = []
+    mapping_list = []
+    
+    # Cache valid adducts
+    valid_adducts = cached_valid_adducts if cached_valid_adducts is not None else set()
+    valid_adducts.add("?")  # Always allow '?' adducts
+    
+    for i, feature_row in enumerate(sample_features.iter_rows(named=True)):
+        # Generate unique consensus ID
+        consensus_id_str = str(uuid.uuid4()).replace("-", "")[:16]
+        
+        # Handle adduct information
+        adduct = feature_row.get("adduct")
+        if adduct is None or adduct not in valid_adducts:
+            # Set default adduct based on study polarity
+            study_polarity = getattr(study, "polarity", "positive")
+            if study_polarity in ["negative", "neg"]:
+                adduct = "[M-?]1-"
+                adduct_charge = -1
+                adduct_mass_shift = -1.007825
+            else:
+                adduct = "[M+?]1+"
+                adduct_charge = 1
+                adduct_mass_shift = 1.007825
+        else:
+            # Try to get charge and mass shift from cached adducts
+            adduct_charge = 1
+            adduct_mass_shift = 1.007825
+            if cached_adducts_df is not None and not cached_adducts_df.is_empty():
+                matching_adduct = cached_adducts_df.filter(pl.col("name") == adduct)
+                if not matching_adduct.is_empty():
+                    adduct_row = matching_adduct.row(0, named=True)
+                    adduct_charge = adduct_row["charge"]
+                    adduct_mass_shift = adduct_row["mass_shift"]
+        
+        # Calculate neutral mass
+        mz = feature_row.get("mz", 0.0)
+        if adduct_charge and adduct_mass_shift is not None:
+            adduct_mass_neutral = mz * abs(adduct_charge) - adduct_mass_shift
+        else:
+            adduct_mass_neutral = None
+        
+        # Count MS2 scans
+        ms2_scans = feature_row.get("ms2_scans", [])
+        ms2_count = len(ms2_scans) if ms2_scans else 0
+        
+        # Create consensus feature metadata
+        consensus_feature = {
+            "consensus_uid": i,
+            "consensus_id": consensus_id_str,
+            "quality": feature_row.get("quality", 1.0),
+            "number_samples": 1,  # Always 1 for single sample
+            "rt": feature_row.get("rt", 0.0),
+            "mz": mz,
+            "rt_min": feature_row.get("rt", 0.0),
+            "rt_max": feature_row.get("rt", 0.0),
+            "rt_mean": feature_row.get("rt", 0.0),
+            "rt_start_mean": feature_row.get("rt_start", 0.0),
+            "rt_end_mean": feature_row.get("rt_end", 0.0),
+            "rt_delta_mean": feature_row.get("rt_delta", 0.0),
+            "mz_min": mz,
+            "mz_max": mz,
+            "mz_mean": mz,
+            "mz_start_mean": feature_row.get("mz_start", 0.0),
+            "mz_end_mean": feature_row.get("mz_end", 0.0),
+            "inty_mean": feature_row.get("inty", 0.0),
+            "bl": -1.0,
+            "chrom_coherence_mean": feature_row.get("chrom_coherence", 0.0),
+            "chrom_prominence_mean": feature_row.get("chrom_prominence", 0.0),
+            "chrom_prominence_scaled_mean": feature_row.get("chrom_prominence_scaled", 0.0),
+            "chrom_height_scaled_mean": feature_row.get("chrom_height_scaled", 0.0),
+            "iso": None,  # Will be filled by find_iso() function
+            "iso_mean": feature_row.get("iso", 0.0),
+            "charge_mean": feature_row.get("charge", 0.0),
+            "number_ms2": ms2_count,
+            "adducts": [[adduct, 1, 100.0]],  # Single adduct with 100% frequency
+            "adduct_top": adduct,
+            "adduct_charge_top": adduct_charge,
+            "adduct_mass_neutral_top": adduct_mass_neutral,
+            "adduct_mass_shift_top": adduct_mass_shift,
+            "id_top_name": None,
+            "id_top_class": None,
+            "id_top_adduct": None,
+            "id_top_score": None,
+            "id_source": None,
+        }
+        
+        consensus_list.append(consensus_feature)
+        
+        # Create mapping entry
+        mapping_entry = {
+            "consensus_uid": i,
+            "sample_uid": sample_uid,
+            "feature_uid": feature_row.get("feature_uid"),
+        }
+        mapping_list.append(mapping_entry)
+    
+    # Create DataFrames
+    study.consensus_df = pl.DataFrame(consensus_list, strict=False)
+    study.consensus_mapping_df = pl.DataFrame(mapping_list, strict=False)
+    
+    study.logger.info(f"Created {len(consensus_list)} consensus features from single sample")
+
+
 def _fast_correlation(x, y):
     """
     Fast correlation coefficient calculation for consensus matrix data.