1from pathlib import Path
   2
   3import numpy as np
   4import pandas as pd
   5import warnings
   6
   7import matplotlib.pyplot as plt
   8
   9from corems.encapsulation.factory.parameters import (
  10    LCMSParameters,
  11    LCMSCollectionParameters,
  12)
  13from corems.encapsulation.plot_utils import _finalize_plot
  14from corems.mass_spectra.calc.lc_calc import (
  15    LCCalculations,
  16    PHCalculations,
  17    LCMSCollectionCalculations,
  18    find_closest,
  19)
  20from corems.molecular_id.search.lcms_spectral_search import LCMSSpectralSearch
  21from corems.mass_spectrum.input.numpyArray import (
  22    ms_from_array_profile,
  23    ms_from_array_centroid,
  24)
  25from corems.chroma_peak.factory.chroma_peak_classes import LCMSMassFeature
  26
  27
  28def release_parser(parser):
  29    """Close a spectra parser.
  30
  31    Thermo readers expose ``close_file``. mzML has neither. HDF5 exposes
  32    ``close``. Call only the method that exists.
  33    """
  34    if parser is None:
  35        return
  36    closer = getattr(parser, "close_file", None)
  37    if not callable(closer):
  38        closer = getattr(parser, "close", None)
  39    if callable(closer):
  40        closer()
  41
  42
  43def _paths_equal(left, right):
  44    if left == right:
  45        return True
  46    try:
  47        return Path(str(left)) == Path(str(right))
  48    except (TypeError, ValueError):
  49        return False
  50
  51
  52def _drop_all_na_float_columns(frames):
  53    """Drop all-NA columns that another frame stores as float.
  54
  55    An all-NA object column next to a float makes ``concat`` warn, and a
  56    future pandas can change that float's dtype. Leaving the column out keeps
  57    the float dtype and fills the missing rows with NA. An all-NA object
  58    column next to an int is kept. Dropping that column turns the id into
  59    float64 (None, 2 -> nan, 2.0). Empty frames are returned unchanged.
  60    """
  61    float_value_cols = set()
  62    for frame in frames:
  63        if len(frame) == 0:
  64            continue
  65        for col in frame.columns:
  66            values = frame[col]
  67            if values.notna().any() and pd.api.types.is_float_dtype(values):
  68                float_value_cols.add(col)
  69    cleaned = []
  70    for frame in frames:
  71        if len(frame) == 0:
  72            cleaned.append(frame)
  73            continue
  74        empty_float_cols = [
  75            col
  76            for col in frame.columns
  77            if col in float_value_cols and frame[col].isna().all()
  78        ]
  79        if empty_float_cols:
  80            frame = frame.drop(columns=empty_float_cols)
  81        cleaned.append(frame)
  82    return cleaned
  83
  84
  85class MassSpectraBase:
  86    """Base class for mass spectra objects.
  87
  88    Parameters
  89    -----------
  90    file_location : str or Path
  91        The location of the file containing the mass spectra data.
  92    analyzer : str, optional
  93        The type of analyzer used to generate the mass spectra data. Defaults to 'Unknown'.
  94    instrument_label : str, optional
  95        The type of instrument used to generate the mass spectra data. Defaults to 'Unknown'.
  96    sample_name : str, optional
  97        The name of the sample; defaults to the file name if not provided to the parser. Defaults to None.
  98    spectra_parser : object, optional
  99        The spectra parser object used to create the mass spectra object. Defaults to None.
 100
 101    Attributes
 102    -----------
 103    spectra_parser_class : class
 104        The class of the spectra parser used to create the mass spectra object.
 105    file_location : str or Path
 106        The location of the file containing the mass spectra data.
 107    sample_name : str
 108        The name of the sample; defaults to the file name if not provided to the parser.
 109    analyzer : str
 110        The type of analyzer used to generate the mass spectra data. Derived from the spectra parser.
 111    instrument_label : str
 112        The type of instrument used to generate the mass spectra data. Derived from the spectra parser.
 113    _scan_info : dict
 114        A dictionary containing the scan data with columns for scan number, scan time, ms level, precursor m/z,
 115        scan text, and scan window (lower and upper).
 116        Associated with the property scan_df, which returns a pandas DataFrame or can set this attribute from a pandas DataFrame.
 117    _ms : dict
 118        A dictionary containing mass spectra for the dataset, keys of dictionary are scan numbers. Initialized as an empty dictionary.
 119    _ms_unprocessed: dictionary of pandas.DataFrames or None
 120        A dictionary of unprocssed mass spectra data, as an (optional) intermediate data product for peak picking.
 121        Key is ms_level, and value is dataframe with columns for scan number, m/z, and intensity. Default is None.
 122
 123    Methods
 124    --------
 125    * add_mass_spectra(scan_list, spectrum_mode: str = 'profile', use_parser = True, auto_process=True).
 126        Add mass spectra (or singlel mass spectrum) to _ms slot, from a list of scans
 127    * get_time_of_scan_id(scan).
 128        Returns the scan time for the specified scan number.
 129    """
 130
 131    def __init__(
 132        self,
 133        file_location,
 134        analyzer="Unknown",
 135        instrument_label="Unknown",
 136        sample_name=None,
 137        spectra_parser=None,
 138    ):
 139        if isinstance(file_location, str):
 140            file_location = Path(file_location)
 141        else:
 142            file_location = file_location
 143        if not file_location.exists():
 144            raise FileExistsError("File does not exist: " + str(file_location))
 145
 146        if sample_name:
 147            self.sample_name = sample_name
 148        else:
 149            self.sample_name = file_location.stem
 150
 151        self.file_location = file_location
 152        self.analyzer = analyzer
 153        self.instrument_label = instrument_label
 154        self._raw_file_location = None
 155        # Live reader for this process. Omitted from pickle; see __getstate__.
 156        self._spectra_parser = None
 157        # False when the cache is the caller's parser. close() must not Dispose it.
 158        self._spectra_parser_owned = False
 159
 160        # Add the spectra parser class to the object if it is not None
 161        if spectra_parser is not None:
 162            self.spectra_parser_class = spectra_parser.__class__
 163            if self.spectra_parser_class.__name__ == "ReadCoreMSHDFMassSpectra":
 164                self.raw_file_location = spectra_parser.get_raw_file_location()
 165
 166            # Compare the instance that was passed in. The property opens a reader.
 167            if (
 168                self.sample_name is not None
 169                and self.sample_name != spectra_parser.sample_name
 170                and self.spectra_parser_class.__name__ != "ReadCoreMSHDFMassSpectra"
 171            ):
 172                warnings.warn(
 173                    "sample_name provided to MassSpectraBase object does not match sample_name provided to spectra parser object",
 174                    UserWarning,
 175                )
 176            if self.analyzer != spectra_parser.analyzer:
 177                warnings.warn(
 178                    "analyzer provided to MassSpectraBase object does not match analyzer provided to spectra parser object",
 179                    UserWarning,
 180                )
 181            if self.instrument_label != spectra_parser.instrument_label:
 182                warnings.warn(
 183                    "instrument provided to MassSpectraBase object does not match instrument provided to spectra parser object",
 184                    UserWarning,
 185                )
 186            if self.file_location != spectra_parser.file_location:
 187                warnings.warn(
 188                    "file_location provided to MassSpectraBase object does not match file_location provided to spectra parser object",
 189                    UserWarning,
 190                )
 191            if self._parser_matches(spectra_parser):
 192                self._spectra_parser = spectra_parser
 193
 194        # Instantiate empty dictionaries for scan information and mass spectra
 195        self._scan_info = {}
 196        self._ms = {}
 197        self._ms_unprocessed = {}
 198
 199    def _parser_matches(self, parser):
 200        if parser is None:
 201            return False
 202        parser_path = getattr(parser, "file_location", None)
 203        if parser_path is None:
 204            return False
 205        return _paths_equal(parser_path, self.raw_file_location)
 206
 207    def _release_cached_parser(self, force=False):
 208        parser = getattr(self, "_spectra_parser", None)
 209        owned = getattr(self, "_spectra_parser_owned", False)
 210        if parser is None or not (owned or force):
 211            return
 212        self._spectra_parser = None
 213        self._spectra_parser_owned = False
 214        release_parser(parser)
 215
 216    @property
 217    def spectra_parser(self):
 218        """Return one parser for ``raw_file_location``.
 219
 220        The first access opens it. Later accesses in this process reuse it.
 221        A worker that received this object through pickle opens its own reader
 222        on the first access there. The caller's parser, when one was passed
 223        in, is reused and is not closed by this object.
 224        """
 225        cached = self._spectra_parser
 226        if cached is not None and self._parser_matches(cached):
 227            return cached
 228        if cached is not None:
 229            self._release_cached_parser(force=True)
 230        if not hasattr(self, "spectra_parser_class"):
 231            raise AttributeError("spectra_parser_class is not set")
 232        location = self.raw_file_location
 233        if not Path(location).exists():
 234            raise FileNotFoundError(
 235                f"Raw file not found at location: {location}, update raw_file_location property to point to correct location."
 236            )
 237        parser = self.spectra_parser_class(location)
 238        self._spectra_parser = parser
 239        self._spectra_parser_owned = True
 240        return parser
 241
 242    def close(self):
 243        """Release a reader this object opened.
 244
 245        A parser passed into the constructor stays open. Its owner closes it.
 246        """
 247        self._release_cached_parser(force=False)
 248
 249    def __del__(self):
 250        try:
 251            self.close()
 252        except Exception:
 253            return
 254
 255    def __getstate__(self):
 256        # The Thermo reader is a pythonnet handle. It cannot cross a process
 257        # boundary, and Mac multiprocessing uses spawn. The child reopens it.
 258        state = self.__dict__.copy()
 259        state["_spectra_parser"] = None
 260        state["_spectra_parser_owned"] = False
 261        return state
 262
 263    def __setstate__(self, state):
 264        self.__dict__.update(state)
 265        self._spectra_parser = None
 266        self._spectra_parser_owned = False
 267
 268    @property
 269    def raw_file_location(self):
 270        """Returns the file_location unless the _raw_file_location is not None."""
 271        return (
 272            self._raw_file_location
 273            if self._raw_file_location is not None
 274            else self.file_location
 275        )
 276
 277    @raw_file_location.setter
 278    def raw_file_location(self, value):
 279        self._raw_file_location = value
 280        cached = getattr(self, "_spectra_parser", None)
 281        if cached is not None and not self._parser_matches(cached):
 282            self._release_cached_parser(force=True)
 283
 284    def add_mass_spectrum(self, mass_spec):
 285        """Adds a mass spectrum to the dataset.
 286
 287        Parameters
 288        -----------
 289        mass_spec : MassSpectrum
 290            The corems MassSpectrum object to be added to the dataset.
 291
 292        Notes
 293        -----
 294        This is a helper function for the add_mass_spectra() method, and is not intended to be called directly.
 295        """
 296        # check if mass_spec has a scan_number attribute
 297        if not hasattr(mass_spec, "scan_number"):
 298            raise ValueError(
 299                "Mass spectrum must have a scan_number attribute to be added to the dataset correctly"
 300            )
 301        self._ms[mass_spec.scan_number] = mass_spec
 302
 303    def add_mass_spectra(
 304        self,
 305        scan_list,
 306        spectrum_mode=None,
 307        ms_level=1,
 308        use_parser=True,
 309        auto_process=True,
 310        ms_params=None,
 311    ):
 312        """Add mass spectra to _ms dictionary, from a list of scans or single scan
 313
 314        Notes
 315        -----
 316        The mass spectra will inherit the mass_spectrum, ms_peak, and molecular_search parameters from the LCMSBase object.
 317
 318
 319        Parameters
 320        -----------
 321        scan_list : list of ints
 322            List of scans to use to populate _ms slot
 323        spectrum_mode : str or None
 324            The spectrum mode to use for the mass spectra.
 325            If None, method will use the spectrum mode from the spectra parser to ascertain the spectrum mode (this allows for mixed types).
 326            Defaults to None.
 327        ms_level : int, optional
 328            The MS level to use for the mass spectra.
 329            This is used to pass the molecular_search parameters from the LCMS object to the individual MassSpectrum objects.
 330            Defaults to 1.
 331        using_parser : bool
 332            Whether to use the mass spectra parser to get the mass spectra.  Defaults to True.
 333        auto_process : bool
 334            Whether to auto-process the mass spectra.  Defaults to True.
 335        ms_params : MSParameters or None
 336            The mass spectrum parameters to use for the mass spectra.  If None, uses the globally set MSParameters.
 337
 338        Raises
 339        ------
 340        TypeError
 341            If scan_list is not a list of ints
 342        ValueError
 343            If polarity is not 'positive' or 'negative'
 344            If ms_level is not 1 or 2
 345        """
 346
 347        # check if scan_list is a list or a single int; if single int, convert to list
 348        if isinstance(scan_list, int):
 349            scan_list = [scan_list]
 350        if not isinstance(scan_list, list):
 351            raise TypeError("scan_list must be a list of integers")
 352        for scan in scan_list:
 353            if not isinstance(scan, int):
 354                raise TypeError("scan_list must be a list of integers")
 355
 356        # set polarity to -1 if negative mode, 1 if positive mode (for mass spectrum creation)
 357        if self.polarity == "negative":
 358            polarity = -1
 359        elif self.polarity == "positive":
 360            polarity = 1
 361        else:
 362            raise ValueError(
 363                "Polarity not set for dataset, must be a either 'positive' or 'negative'"
 364            )
 365
 366        # is not using_parser, check that ms1 and ms2 are not None
 367        if not use_parser:
 368            if ms_level not in self._ms_unprocessed.keys():
 369                raise ValueError(
 370                    "ms_level {} not found in _ms_unprocessed dictionary".format(
 371                        ms_level
 372                    )
 373                )
 374
 375        scan_list = list(set(scan_list))
 376        scan_list.sort()
 377
 378        # Skip scans that have already been added to _ms to avoid redundant reprocessing
 379        already_added = [s for s in scan_list if s in self._ms]
 380        if already_added:
 381            scan_list = [s for s in scan_list if s not in self._ms]
 382        if not scan_list:
 383            return
 384
 385        if not use_parser:
 386            if self._ms_unprocessed[ms_level] is None:
 387                raise ValueError(
 388                    "No unprocessed data found for ms_level {}".format(ms_level)
 389                )
 390            if (
 391                len(
 392                    np.setdiff1d(
 393                        scan_list, self._ms_unprocessed[ms_level].scan.tolist()
 394                    )
 395                )
 396                > 0
 397            ):
 398                raise ValueError(
 399                    "Not all scans in scan_list are present in the unprocessed data"
 400                )
 401            # Prepare the ms_df for parsing
 402            ms_df = self._ms_unprocessed[ms_level].copy().set_index("scan", drop=False)
 403
 404        if use_parser:
 405            # Use batch function to get all mass spectra at once
 406            if spectrum_mode is None:
 407                # get spectrum mode from _scan_info for each scan
 408                spectrum_modes = [
 409                    self.scan_df.loc[scan, "ms_format"] for scan in scan_list
 410                ]
 411                spectrum_mode_batch = (
 412                    spectrum_modes[0] if len(set(spectrum_modes)) == 1 else None
 413                )
 414            else:
 415                spectrum_mode_batch = spectrum_mode
 416
 417            ms_list = self.spectra_parser.get_mass_spectra_from_scan_list(
 418                scan_list=scan_list,
 419                spectrum_mode=spectrum_mode_batch,
 420                auto_process=False,
 421            )
 422
 423            # Process each mass spectrum
 424            for i, scan in enumerate(scan_list):
 425                ms = ms_list[i] if i < len(ms_list) else None
 426                if ms is not None:
 427                    if ms_params is not None:
 428                        ms.parameters = ms_params
 429                    ms.scan_number = scan
 430                    if auto_process:
 431                        ms.process_mass_spec()
 432                    self.add_mass_spectrum(ms)
 433        else:
 434            # Original non-parser logic remains unchanged
 435            for scan in scan_list:
 436                ms = None
 437                if spectrum_mode is None:
 438                    # get spectrum mode from _scan_info
 439                    spectrum_mode_scan = self.scan_df.loc[scan, "ms_format"]
 440                else:
 441                    spectrum_mode_scan = spectrum_mode
 442
 443                my_ms_df = ms_df.loc[scan]
 444                if spectrum_mode_scan == "profile":
 445                    # Check this - it might be better to use the MassSpectrumProfile class to instantiate the mass spectrum
 446                    ms = ms_from_array_profile(
 447                        my_ms_df.mz,
 448                        my_ms_df.intensity,
 449                        self.file_location,
 450                        polarity=polarity,
 451                        auto_process=False,
 452                    )
 453                else:
 454                    ms = ms_from_array_centroid(
 455                        mz=my_ms_df.mz,
 456                        abundance=my_ms_df.intensity,
 457                        rp=[np.nan] * len(my_ms_df.mz),
 458                        s2n=[np.nan] * len(my_ms_df.mz),
 459                        dataname=self.file_location,
 460                        polarity=polarity,
 461                        auto_process=False,
 462                    )
 463
 464                # Set the mass spectrum parameters, auto-process if auto_process is True, and add to the dataset
 465                if ms is not None:
 466                    if ms_params is not None:
 467                        ms.parameters = ms_params
 468                    ms.scan_number = scan
 469                    if auto_process:
 470                        ms.process_mass_spec()
 471                    self.add_mass_spectrum(ms)
 472
 473    def get_time_of_scan_id(self, scan):
 474        """Returns the scan time for the specified scan number.
 475
 476        Parameters
 477        -----------
 478        scan : int
 479            The scan number of the desired scan time.
 480
 481        Returns
 482        --------
 483        float
 484            The scan time for the specified scan number (in minutes).
 485
 486        Raises
 487        ------
 488        ValueError
 489            If no scan time is found for the specified scan number.
 490        """
 491        # Check if _retenion_time_list is empty and raise error if so
 492        if len(self._retention_time_list) == 0:
 493            raise ValueError("No retention times found in dataset")
 494        rt = self._retention_time_list[self._scans_number_list.index(scan)]
 495        return rt
 496
 497    @property
 498    def scan_df(self):
 499        """
 500        pandas.DataFrame : A pandas DataFrame containing the scan info data with columns for scan number, scan time, ms level, precursor m/z, scan text, and scan window (lower and upper).
 501        """
 502        scan_df = pd.DataFrame.from_dict(self._scan_info)
 503        return scan_df
 504
 505    @property
 506    def ms(self):
 507        """
 508        dictionary : contains the key associated with mass spectra and values are the associated MassSpecProfiles
 509        """
 510        return self._ms
 511
 512    @scan_df.setter
 513    def scan_df(self, df):
 514        """
 515        Sets the scan data for the dataset.
 516
 517        Parameters
 518        -----------
 519        df : pandas.DataFrame
 520            A pandas DataFrame containing the scan data with columns for scan number, scan time, ms level,
 521            precursor m/z, scan text, and scan window (lower and upper).
 522        """
 523        self._scan_info = df.to_dict()
 524
 525    def __getitem__(self, scan_number):
 526        return self._ms.get(scan_number)
 527
 528    def __iter__(self):
 529        """Returns an iterator over the mass spectra in the dataset.
 530
 531        Returns
 532        --------
 533        iterator
 534            An iterator over the mass spectra in the dataset.
 535        """
 536        return iter(self._ms.values())
 537
 538
 539class LCMSBase(MassSpectraBase, LCCalculations, PHCalculations, LCMSSpectralSearch):
 540    """A class representing a liquid chromatography-mass spectrometry (LC-MS) data object.
 541
 542    This class is not intended to be instantiated directly, but rather to be instantiated by an appropriate mass spectra parser using the get_lcms_obj() method.
 543
 544    Parameters
 545    -----------
 546    file_location : str or Path
 547        The location of the file containing the mass spectra data.
 548    analyzer : str, optional
 549        The type of analyzer used to generate the mass spectra data. Defaults to 'Unknown'.
 550    instrument_label : str, optional
 551        The type of instrument used to generate the mass spectra data. Defaults to 'Unknown'.
 552    sample_name : str, optional
 553        The name of the sample; defaults to the file name if not provided to the parser. Defaults to None.
 554    spectra_parser : object, optional
 555        The spectra parser object used to create the mass spectra object. Defaults to None.
 556
 557    Attributes
 558    -----------
 559    polarity : str
 560        The polarity of the ionization mode used for the dataset.
 561    _parameters : LCMSParameters
 562        The parameters used for all methods called on the LCMSBase object. Set upon instantiation from LCMSParameters.
 563    _retention_time_list : numpy.ndarray
 564        An array of retention times for the dataset.
 565    _scans_number_list : list
 566        A list of scan numbers for the dataset.
 567    _tic_list : numpy.ndarray
 568        An array of total ion current (TIC) values for the dataset.
 569    eics : dict
 570        A dictionary containing extracted ion chromatograms (EICs) for the dataset.
 571        Key is the mz of the EIC. Initialized as an empty dictionary.
 572    mass_features : dictionary of LCMSMassFeature objects
 573        A dictionary containing mass features for the dataset.
 574        Key is mass feature ID. Initialized as an empty dictionary.
 575    induced_mass_features: dictionary of LCMSMassFeature objects
 576        A dictionary containing mass features from a collection that don't
 577        satisfy criteria for initial mass features. Key is mass feature ID.
 578        Initialized as an empty dictionary.
 579    missing_mass_features: pandas.DataFrame
 580        A table of clusters in a given sample for which a mass feature was
 581        sought and not found
 582    spectral_search_results : dictionary of MS2SearchResults objects
 583        A dictionary containing spectral search results for the dataset.
 584        Key is scan number : precursor mz. Initialized as an empty dictionary.
 585
 586    Methods
 587    --------
 588    * get_parameters_json().
 589        Returns the parameters used for the LC-MS analysis in JSON format.
 590    * add_associated_ms2_dda(add_to_lcmsobj=True, auto_process=True, use_parser=True)
 591        Adds which MS2 scans are associated with each mass feature to the
 592        mass_features dictionary and optionally adds the MS2 spectra to the _ms dictionary.
 593    * add_associated_ms1(add_to_lcmsobj=True, auto_process=True, use_parser=True)
 594        Adds the MS1 spectra associated with each mass feature to the
 595        mass_features dictionary and adds the MS1 spectra to the _ms dictionary.
 596    * mass_features_to_df()
 597        Returns a pandas dataframe summarizing the mass features in the dataset.
 598    * to_mgf(out_file_path, feature_ids=None, ...)
 599        Export mass features to a SIRIUS-compatible MGF file.
 600    * set_tic_list_from_data(overwrite=False)
 601        Sets the TIC list from the mass spectrum objects within the _ms dictionary.
 602    * set_retention_time_from_data(overwrite=False)
 603        Sets the retention time list from the data in the _ms dictionary.
 604    * set_scans_number_from_data(overwrite=False)
 605        Sets the scan number list from the data in the _ms dictionary.
 606    * plot_composite_mz_features(binsize = 1e-4, ph_int_min_thresh = 0.001, mf_plot = True, ms2_plot = True, return_fig = False, path = None)
 607        Generates plot of M/Z features comparing scan time vs M/Z value
 608    * search_for_targeted_mass_feature(ms1df: pd.DataFrame, sample: pd.Series, tol_flag = 0)
 609        Searches for mass features in specific M/Z and scan time windows that
 610        were missed by the persistent homology search
 611    """
 612
 613    def __init__(
 614        self,
 615        file_location,
 616        analyzer="Unknown",
 617        instrument_label="Unknown",
 618        sample_name=None,
 619        spectra_parser=None,
 620    ):
 621        super().__init__(
 622            file_location, analyzer, instrument_label, sample_name, spectra_parser
 623        )
 624        self.polarity = ""
 625        self._parameters = LCMSParameters()
 626        self._retention_time_list = []
 627        self._scans_number_list = []
 628        self._tic_list = []
 629        self.eics = {}
 630        self.mass_features = {}
 631        self.induced_mass_features = {}
 632        self.spectral_search_results = {}
 633
 634    def get_eic_mz_for_mass_feature(self, mf_mz, tolerance=0.0001):
 635        """Get the EIC dictionary key (m/z) that best matches a mass feature's m/z.
 636
 637        Finds the closest EIC m/z key within the specified tolerance.
 638
 639        Parameters
 640        ----------
 641        mf_mz : float
 642            The m/z value of the mass feature to match.
 643        tolerance : float, optional
 644            Maximum m/z difference for matching. Default is 0.0001 Da.
 645
 646        Returns
 647        -------
 648        float or None
 649            The EIC dictionary key (m/z) of the closest matching EIC,
 650            or None if no EIC is within tolerance.
 651        """
 652        if not hasattr(self, "eics") or not self.eics:
 653            return None
 654
 655        best_eic_mz = None
 656        best_diff = tolerance
 657        for eic_mz in self.eics.keys():
 658            diff = abs(mf_mz - eic_mz)
 659            if diff < best_diff:
 660                best_diff = diff
 661                best_eic_mz = eic_mz
 662        return best_eic_mz
 663
 664    def associate_eics_with_mass_features(self, tolerance=0.0001, induced=False):
 665        """Associate EICs with mass features using tolerance-based m/z matching.
 666
 667        Associates EIC_Data objects from self.eics with mass features by finding
 668        the closest EIC within the specified m/z tolerance. This is more robust
 669        than exact matching which can fail due to floating point precision issues.
 670
 671        Parameters
 672        ----------
 673        tolerance : float, optional
 674            Maximum m/z difference for matching EICs to mass features. Default is 0.0001 Da.
 675        induced : bool, optional
 676            If True, associates EICs with induced_mass_features instead of mass_features.
 677            Default is False.
 678
 679        Notes
 680        -----
 681        For each mass feature, this method finds the EIC with the closest m/z value
 682        within the tolerance window and assigns it to the mass feature's _eic_data attribute.
 683        If multiple EICs are within tolerance, the one with the smallest m/z difference is chosen.
 684        """
 685        # Select which mass features dictionary to use
 686        mf_dict = self.induced_mass_features if induced else self.mass_features
 687
 688        # Use the _eic_mz attribute on each mass_feature to find the closest matching EIC
 689        for idx in mf_dict.keys():
 690            mf_mz = mf_dict[idx]._eic_mz
 691            # Find closest EIC within tolerance
 692            best_match = None
 693            best_diff = tolerance
 694            for eic_mz, eic_data in self.eics.items():
 695                diff = abs(mf_mz - eic_mz)
 696                if diff < best_diff:
 697                    best_diff = diff
 698                    best_match = eic_data
 699            if best_match is not None:
 700                mf_dict[idx]._eic_data = best_match
 701
 702    def get_parameters_json(self):
 703        """Returns the parameters stored for the LC-MS object in JSON format.
 704
 705        Returns
 706        --------
 707        str
 708            The parameters used for the LC-MS analysis in JSON format.
 709        """
 710        return self.parameters.to_json()
 711
 712    def remove_unprocessed_data(self, ms_level=None):
 713        """Removes the unprocessed data from the LCMSBase object.
 714
 715        Parameters
 716        -----------
 717        ms_level : int, optional
 718            The MS level to remove the unprocessed data for. If None, removes unprocessed data for all MS levels.
 719
 720        Raises
 721        ------
 722        ValueError
 723            If ms_level is not 1 or 2.
 724
 725        Notes
 726        -----
 727        This method is useful for freeing up memory after the data has been processed.
 728        """
 729        if ms_level is None:
 730            for ms_level in self._ms_unprocessed.keys():
 731                self._ms_unprocessed[ms_level] = None
 732        if ms_level not in [1, 2]:
 733            raise ValueError("ms_level must be 1 or 2")
 734        self._ms_unprocessed[ms_level] = None
 735
 736    def _filter_ms2_scans_by_integration_bounds(self, mf_dict=None):
 737        """Filter MS2 scans to only include those within integration bounds.
 738
 739        Removes MS2 scan numbers that fall outside the start_scan to final_scan range
 740        for each mass feature. This should be called after integration sets the bounds.
 741
 742        Parameters
 743        ----------
 744        mf_dict : dict, optional
 745            Dictionary of mass features to filter. If None, uses self.mass_features.
 746
 747        Returns
 748        -------
 749        int
 750            Number of MS2 scans removed across all mass features.
 751        """
 752        if mf_dict is None:
 753            mf_dict = self.mass_features
 754
 755        total_removed = 0
 756
 757        for mf_id, mf in mf_dict.items():
 758            # Only filter if integration bounds are set and MS2 scans exist
 759            if (
 760                hasattr(mf, "start_scan")
 761                and hasattr(mf, "final_scan")
 762                and mf.start_scan is not None
 763                and mf.final_scan is not None
 764                and mf.ms2_scan_numbers is not None
 765                and len(mf.ms2_scan_numbers) > 0
 766            ):
 767                # Filter scan numbers to only those within bounds
 768                original_count = len(mf.ms2_scan_numbers)
 769                mf.ms2_scan_numbers = [
 770                    scan
 771                    for scan in mf.ms2_scan_numbers
 772                    if mf.start_scan <= scan <= mf.final_scan
 773                ]
 774                removed = original_count - len(mf.ms2_scan_numbers)
 775                total_removed += removed
 776
 777        return total_removed
 778
 779    def _find_ms2_scans_for_mass_features(self, mf_ids=None, scan_filter=None):
 780        """Find MS2 scans associated with mass features.
 781
 782        This helper method finds MS2 scans that match mass features based on RT and m/z tolerances.
 783        It updates the ms2_scan_numbers attribute on each mass feature.
 784
 785        Parameters
 786        ----------
 787        mf_ids : list of int, optional
 788            List of mass feature IDs to find MS2 for. If None, finds for all mass features.
 789        scan_filter : str, optional
 790            Filter string for MS2 scans (e.g., 'hcd'). Default is None.
 791
 792        Returns
 793        -------
 794        list
 795            List of unique MS2 scan numbers found across all mass features.
 796
 797        Raises
 798        ------
 799        ValueError
 800            If no MS2 scans are found in the dataset.
 801        """
 802        # Get mass features to process
 803        if mf_ids is None:
 804            mf_ids = list(self.mass_features.keys())
 805
 806        # Get mass features dataframe
 807        mf_df = self.mass_features_to_df()
 808        mf_df = mf_df.loc[mf_ids].copy()
 809
 810        # Find ms2 scans that have a precursor m/z value
 811        ms2_scans = self.scan_df[self.scan_df.ms_level == 2]
 812        ms2_scans = ms2_scans[~ms2_scans.precursor_mz.isna()]
 813        ms2_scans = ms2_scans[ms2_scans.tic > 0]
 814
 815        if len(ms2_scans) == 0:
 816            raise ValueError("No DDA scans found in dataset")
 817
 818        if scan_filter is not None:
 819            ms2_scans = ms2_scans[ms2_scans.scan_text.str.contains(scan_filter)]
 820
 821        # Get tolerances from parameters
 822        time_tol = self.parameters.lc_ms.ms2_dda_rt_tolerance
 823        mz_tol = self.parameters.lc_ms.ms2_dda_mz_tolerance
 824
 825        # For each mass feature, find the ms2 scans that are within the roi scan time and mz range
 826        dda_scans = []
 827        for i, row in mf_df.iterrows():
 828            ms2_scans_filtered = ms2_scans[
 829                ms2_scans.scan_time.between(
 830                    row.scan_time - time_tol, row.scan_time + time_tol
 831                )
 832            ]
 833            ms2_scans_filtered = ms2_scans_filtered[
 834                ms2_scans_filtered.precursor_mz.between(
 835                    row.mz - mz_tol, row.mz + mz_tol
 836                )
 837            ]
 838            scan_list = ms2_scans_filtered.scan.tolist()
 839            if scan_list:
 840                # Filter scans by integration bounds if they exist
 841                mf = self.mass_features[i]
 842                if (
 843                    hasattr(mf, "start_scan")
 844                    and hasattr(mf, "final_scan")
 845                    and mf.start_scan is not None
 846                    and mf.final_scan is not None
 847                ):
 848                    # Only keep scans within integration bounds
 849                    scan_list = [
 850                        s for s in scan_list if mf.start_scan <= s <= mf.final_scan
 851                    ]
 852
 853                if scan_list:  # Only add if there are still scans after filtering
 854                    self.mass_features[i].ms2_scan_numbers = scan_list + list(
 855                        self.mass_features[i].ms2_scan_numbers
 856                    )
 857                    dda_scans.extend(scan_list)
 858
 859        return list(set(dda_scans))
 860
 861    def add_associated_ms2_dda(
 862        self,
 863        auto_process=True,
 864        use_parser=True,
 865        spectrum_mode=None,
 866        ms_params_key="ms2",
 867        scan_filter=None,
 868    ):
 869        """Add MS2 spectra associated with mass features to the dataset.
 870
 871        Populates the mass_features ms2_scan_numbers attribute (on mass_features dictionary on LCMSObject)
 872
 873        Parameters
 874        -----------
 875        auto_process : bool, optional
 876            If True, auto-processes the MS2 spectra before adding it to the object's _ms dictionary. Default is True.
 877        use_parser : bool, optional
 878            If True, envoke the spectra parser to get the MS2 spectra. Default is True.
 879        spectrum_mode : str or None, optional
 880            The spectrum mode to use for the mass spectra.  If None, method will use the spectrum mode
 881            from the spectra parser to ascertain the spectrum mode (this allows for mixed types).
 882            Defaults to None. (faster if defined, otherwise will check each scan)
 883        ms_params_key : string, optional
 884            The key of the mass spectrum parameters to use for the mass spectra, accessed from the LCMSObject.parameters.mass_spectrum attribute.
 885            Defaults to 'ms2'.
 886        scan_filter : str
 887            A string to filter the scans to add to the _ms dictionary.  If None, all scans are added.  Defaults to None.
 888            "hcd" will pull out only HCD scans.
 889
 890        Raises
 891        ------
 892        ValueError
 893            If mass_features is not set, must run find_mass_features() first.
 894            If no MS2 scans are found in the dataset.
 895            If no precursor m/z values are found in MS2 scans, not a DDA dataset.
 896        """
 897        # Check if mass_features is set, raise error if not
 898        if self.mass_features is None:
 899            raise ValueError(
 900                "mass_features not set, must run find_mass_features() first"
 901            )
 902
 903        # reconfigure ms_params to get the correct mass spectrum parameters from the key
 904        ms_params = self.parameters.mass_spectrum[ms_params_key]
 905
 906        # Find MS2 scans for all mass features
 907        dda_scans = self._find_ms2_scans_for_mass_features(scan_filter=scan_filter)
 908
 909        # Load MS2 spectra
 910        self.add_mass_spectra(
 911            scan_list=dda_scans,
 912            auto_process=auto_process,
 913            spectrum_mode=spectrum_mode,
 914            use_parser=use_parser,
 915            ms_params=ms_params,
 916        )
 917
 918        # Associate appropriate _ms attribute to appropriate mass feature's ms2_mass_spectra attribute
 919        for mf_id in self.mass_features:
 920            if self.mass_features[mf_id].ms2_scan_numbers is not None:
 921                for dda_scan in self.mass_features[mf_id].ms2_scan_numbers:
 922                    if dda_scan in self._ms.keys():
 923                        self.mass_features[mf_id].ms2_mass_spectra[dda_scan] = self._ms[
 924                            dda_scan
 925                        ]
 926
 927    def add_associated_ms1(
 928        self,
 929        auto_process=True,
 930        use_parser=True,
 931        spectrum_mode=None,
 932        induced_features=False,
 933    ):
 934        """Add MS1 spectra associated with mass features to the dataset.
 935
 936        Parameters
 937        -----------
 938        auto_process : bool, optional
 939            If True, auto-processes the MS1 spectra before adding it to the object's _ms dictionary. Default is True.
 940        use_parser : bool, optional
 941            If True, envoke the spectra parser to get the MS1 spectra. Default is True.
 942        spectrum_mode : str or None, optional
 943            The spectrum mode to use for the mass spectra.  If None, method will use the spectrum mode
 944            from the spectra parser to ascertain the spectrum mode (this allows for mixed types).
 945            Defaults to None. (faster if defined, otherwise will check each scan)
 946        induced_features : bool, optional
 947            If True, add associated MS1 of the induced mass features instead of the primary mass features
 948
 949        Raises
 950        ------
 951        ValueError
 952            If mass_features is not set, must run find_mass_features() first.
 953            If apex scans are not profile mode, all apex scans must be profile mode for averaging.
 954            If number of scans to average is not  1 or an integer with an integer median (i.e. 3, 5, 7, 9).
 955            If deconvolute is True and no EICs are found, did you run integrate_mass_features() first?
 956        """
 957        # Check if mass_features is set, raise error if not
 958        if self.mass_features is None:
 959            raise ValueError(
 960                "mass_features not set, must run find_mass_features() first"
 961            )
 962
 963        if induced_features:
 964            mf_dict = self.induced_mass_features
 965        else:
 966            mf_dict = self.mass_features
 967
 968        scans_to_average = self.parameters.lc_ms.ms1_scans_to_average
 969
 970        ## sketchy work around for induced mass features
 971        scan_list = [
 972            int(mf_dict[x].apex_scan)
 973            for x in mf_dict
 974            if int(mf_dict[x].apex_scan) != -99
 975        ]
 976
 977        if scans_to_average == 1:
 978            # Add to LCMSobj
 979            self.add_mass_spectra(
 980                scan_list=scan_list,
 981                auto_process=auto_process,
 982                use_parser=use_parser,
 983                spectrum_mode=spectrum_mode,
 984                ms_params=self.parameters.mass_spectrum["ms1"],
 985            )
 986
 987        elif (
 988            (scans_to_average - 1) % 2
 989        ) == 0:  # scans_to_average = 3, 5, 7 etc, mirror l/r around apex
 990            apex_scans = list(set(scan_list))
 991            # Check if all apex scans are profile mode, raise error if not
 992            if not all(self.scan_df.loc[apex_scans, "ms_format"] == "profile"):
 993                raise ValueError("All apex scans must be profile mode for averaging")
 994
 995            # First get sets of scans to average
 996            def get_scans_from_apex(ms1_scans, apex_scan, scans_to_average):
 997                ms1_idx_start = ms1_scans.index(apex_scan) - int(
 998                    (scans_to_average - 1) / 2
 999                )
1000                if ms1_idx_start < 0:
1001                    ms1_idx_start = 0
1002                ms1_idx_end = (
1003                    ms1_scans.index(apex_scan) + int((scans_to_average - 1) / 2) + 1
1004                )
1005                if ms1_idx_end > (len(ms1_scans) - 1):
1006                    ms1_idx_end = len(ms1_scans) - 1
1007                scan_list = ms1_scans[ms1_idx_start:ms1_idx_end]
1008                return scan_list
1009
1010            ms1_scans = self.ms1_scans
1011            scans_lists = [
1012                get_scans_from_apex(ms1_scans, apex_scan, scans_to_average)
1013                for apex_scan in apex_scans
1014            ]
1015
1016            # set polarity to -1 if negative mode, 1 if positive mode (for mass spectrum creation)
1017            if self.polarity == "negative":
1018                polarity = -1
1019            elif self.polarity == "positive":
1020                polarity = 1
1021
1022            if not use_parser:
1023                # Perform checks and prepare _ms_unprocessed dictionary if use_parser is False (saves time to do this once)
1024                ms1_unprocessed = self._ms_unprocessed[1].copy()
1025                # Set the index on _ms_unprocessed[1] to scan number
1026                ms1_unprocessed = ms1_unprocessed.set_index("scan", drop=False)
1027                self._ms_unprocessed[1] = ms1_unprocessed
1028
1029                # Check that all the scans in scan_lists are indexs in self._ms_unprocessed[1]
1030                scans_lists_flat = list(
1031                    set([scan for sublist in scans_lists for scan in sublist])
1032                )
1033                if (
1034                    len(
1035                        np.setdiff1d(
1036                            np.sort(scans_lists_flat),
1037                            np.sort(ms1_unprocessed.index.values),
1038                        )
1039                    )
1040                    > 0
1041                ):
1042                    raise ValueError(
1043                        "Not all scans to average are present in the unprocessed data"
1044                    )
1045
1046            for scan_list_average, apex_scan in zip(scans_lists, apex_scans):
1047                # Get unprocessed mass spectrum from scans
1048                ms = self.get_average_mass_spectrum(
1049                    scan_list=scan_list_average,
1050                    apex_scan=apex_scan,
1051                    spectrum_mode="profile",
1052                    ms_level=1,
1053                    auto_process=auto_process,
1054                    use_parser=use_parser,
1055                    perform_checks=False,
1056                    polarity=polarity,
1057                    ms_params=self.parameters.mass_spectrum["ms1"],
1058                )
1059                # Add mass spectrum to LCMS object and associated with mass feature
1060                self.add_mass_spectrum(ms)
1061
1062            if not use_parser:
1063                # Reset the index on _ms_unprocessed[1] to not be scan number
1064                ms1_unprocessed = ms1_unprocessed.reset_index(drop=True)
1065                self._ms_unprocessed[1] = ms1_unprocessed
1066        else:
1067            raise ValueError(
1068                "Number of scans to average must be 1 or an integer with an integer median (i.e. 3, 5, 7, 9)"
1069            )
1070
1071        # Associate the ms1 spectra with the mass features
1072        for k in mf_dict.keys():
1073            ## another induced feature work around
1074            if mf_dict[k].apex_scan != -99:
1075                mf_dict[k].mass_spectrum = self._ms[mf_dict[k].apex_scan]
1076                mf_dict[k].update_mz()
1077
1078    def mass_features_to_df(
1079        self, induced_features=False, drop_na_cols=False, include_cols=None
1080    ):
1081        """Returns a pandas dataframe summarizing the mass features.
1082
1083        The dataframe contains the following columns: mf_id, mz, _eic_mz, apex_scan,
1084        scan_time, intensity, persistence, area, monoisotopic_mf_id, and
1085        isotopologue_type. The index is set to mf_id (mass feature ID).
1086
1087        ``_eic_mz`` is the m/z used for EIC extraction (set on integrate). If
1088        not set on the feature object, it falls back to the feature ``mz`` so
1089        collection dataframes always carry a usable EIC key for regular and
1090        induced features.
1091
1092        Parameters
1093        -----------
1094        induced_features : bool, optional
1095            If True, calls the induced_mass_features dictionary. Defaults to False.
1096        drop_na_cols : bool, optional
1097            If True, drops columns that are entirely NA. Defaults to False.
1098        include_cols : list of str, optional
1099            If provided, only includes the specified columns in the output (in addition to 'mf_id' which is always included as the index).
1100            If None, includes all available columns. Defaults to None.
1101
1102        Raises
1103        --------
1104        ValueError
1105            If the sample provided doesn't contain the mass feature data.
1106
1107        Returns
1108        --------
1109        pandas.DataFrame
1110            A pandas dataframe of mass features with the following columns:
1111            mf_id, mz, _eic_mz, apex_scan, scan_time, intensity, persistence, area.
1112        """
1113        import pandas as pd
1114
1115        def mass_spectrum_to_string(
1116            mass_spec, normalize=True, min_normalized_abun=0.01
1117        ):
1118            """Converts a mass spectrum to a string of m/z:abundance pairs.
1119
1120            Parameters
1121            -----------
1122            mass_spec : MassSpectrum
1123                A MassSpectrum object to be converted to a string.
1124            normalize : bool, optional
1125                If True, normalizes the abundance values to a maximum of 1. Defaults to True.
1126            min_normalized_abun : float, optional
1127                The minimum normalized abundance value to include in the string, only used if normalize is True. Defaults to 0.01.
1128
1129            Returns
1130            --------
1131            str
1132                A string of m/z:abundance pairs from the mass spectrum, separated by a semicolon.
1133            """
1134            mz_np = mass_spec.to_dataframe()["m/z"].values
1135            abun_np = mass_spec.to_dataframe()["Peak Height"].values
1136            if normalize:
1137                abun_np = abun_np / abun_np.max()
1138            mz_abun = np.column_stack((mz_np, abun_np))
1139            if normalize:
1140                mz_abun = mz_abun[mz_abun[:, 1] > min_normalized_abun]
1141            mz_abun_str = [
1142                str(round(mz, ndigits=4)) + ":" + str(round(abun, ndigits=2))
1143                for mz, abun in mz_abun
1144            ]
1145            return "; ".join(mz_abun_str)
1146
1147        if induced_features:
1148            mf_dict = self.induced_mass_features
1149        else:
1150            mf_dict = self.mass_features
1151
1152        if len(mf_dict) == 0:
1153            # Return an empty dataframe with the expected structure
1154            # This allows collection processing to continue even if some samples have no features
1155            return pd.DataFrame()
1156
1157        cols_in_df = [
1158            "id",
1159            "apex_scan",
1160            "start_scan",
1161            "final_scan",
1162            "retention_time",
1163            "intensity",
1164            "persistence",
1165            "area",
1166            "dispersity_index",
1167            "normalized_dispersity_index",
1168            "tailing_factor",
1169            "gaussian_similarity",
1170            "noise_score",
1171            "noise_score_min",
1172            "noise_score_max",
1173            "monoisotopic_mf_id",
1174            "isotopologue_type",
1175            "mass_spectrum_deconvoluted_parent",
1176            "ms2_scan_numbers",
1177            "type",
1178        ]
1179
1180        df_mf_list = []
1181        for mf_id in mf_dict.keys():
1182            # Find cols_in_df that are in single_mf
1183            df_keys = list(set(cols_in_df).intersection(mf_dict[mf_id].__dir__()))
1184            dict_mf = {}
1185            # Get the values for each key in df_keys from the mass feature object
1186            for key in df_keys:
1187                value = getattr(mf_dict[mf_id], key)
1188                # Wrap list/array values in a list so pandas treats them as single cell values
1189                if key == "ms2_scan_numbers" and isinstance(value, (list, np.ndarray)):
1190                    dict_mf[key] = [value]
1191                else:
1192                    dict_mf[key] = value
1193            if len(mf_dict[mf_id].ms2_scan_numbers) > 0:
1194                # Add MS2 spectra info
1195                best_ms2_spectrum = mf_dict[mf_id].best_ms2
1196                if best_ms2_spectrum is not None:
1197                    dict_mf["ms2_spectrum"] = mass_spectrum_to_string(best_ms2_spectrum)
1198            if len(mf_dict[mf_id].associated_mass_features_deconvoluted) > 0:
1199                dict_mf["associated_mass_features"] = ", ".join(
1200                    map(
1201                        str,
1202                        mf_dict[mf_id].associated_mass_features_deconvoluted,
1203                    )
1204                )
1205            if mf_dict[mf_id]._half_height_width is not None:
1206                dict_mf["half_height_width"] = mf_dict[mf_id].half_height_width
1207            # Check if EIC for mass feature is set
1208            df_mf_single = pd.DataFrame(dict_mf, index=[mf_id])
1209            df_mf_single["mz"] = mf_dict[mf_id].mz
1210            # Always expose _eic_mz: prefer value set on integrate, else feature mz
1211            eic_mz = getattr(mf_dict[mf_id], "_eic_mz", None)
1212            if eic_mz is None or (isinstance(eic_mz, float) and np.isnan(eic_mz)):
1213                eic_mz = mf_dict[mf_id].mz
1214            df_mf_single["_eic_mz"] = eic_mz
1215            df_mf_list.append(df_mf_single)
1216        # Same rule as collection concat: drop all-NA float columns only.
1217        cleaned_frames = _drop_all_na_float_columns(
1218            [frame for frame in df_mf_list if len(frame) > 0]
1219        )
1220        df_mf = pd.concat(cleaned_frames)
1221
1222        # rename _area to area and id to mf_id
1223        df_mf = df_mf.rename(
1224            columns={
1225                "id": "mf_id",
1226                "retention_time": "scan_time",
1227            }
1228        )
1229
1230        # reorder columns
1231        col_order = [
1232            "mf_id",
1233            "type",
1234            "scan_time",
1235            "mz",
1236            "_eic_mz",
1237            "apex_scan",
1238            "start_scan",
1239            "final_scan",
1240            "intensity",
1241            "persistence",
1242            "area",
1243            "half_height_width",
1244            "tailing_factor",
1245            "dispersity_index",
1246            "normalized_dispersity_index",
1247            "gaussian_similarity",
1248            "noise_score",
1249            "noise_score_min",
1250            "noise_score_max",
1251            "monoisotopic_mf_id",
1252            "isotopologue_type",
1253            "mass_spectrum_deconvoluted_parent",
1254            "associated_mass_features",
1255            "ms2_scan_numbers",
1256            "ms2_spectrum",
1257        ]
1258        # drop columns that are not in col_order
1259        cols_to_order = [col for col in col_order if col in df_mf.columns]
1260        df_mf = df_mf[cols_to_order]
1261
1262        # reset index to mf_id
1263        df_mf = df_mf.set_index("mf_id")
1264        df_mf.index.name = "mf_id"
1265
1266        if "half_height_width" in df_mf.columns:
1267            df_mf["half_height_width"] = df_mf["half_height_width"].astype("float64")
1268        if "tailing_factor" in df_mf.columns:
1269            df_mf["tailing_factor"] = df_mf["tailing_factor"].astype("float64")
1270        if "dispersity_index" in df_mf.columns:
1271            df_mf["dispersity_index"] = df_mf["dispersity_index"].astype("float64")
1272        if "normalized_dispersity_index" in df_mf.columns:
1273            df_mf["normalized_dispersity_index"] = df_mf[
1274                "normalized_dispersity_index"
1275            ].astype("float64")
1276
1277        # Filter columns if include_cols is specified
1278        if include_cols is not None:
1279            # Ensure include_cols is a list
1280            if not isinstance(include_cols, list):
1281                raise ValueError("include_cols must be a list of column names")
1282            # Keep only requested columns that exist in the dataframe
1283            available_cols = [col for col in include_cols if col in df_mf.columns]
1284            df_mf = df_mf[available_cols]
1285
1286        # Drop columns that are entirely NA if requested
1287        if drop_na_cols:
1288            df_mf = df_mf.dropna(axis=1, how="all")
1289
1290        return df_mf
1291
1292    def to_mgf(
1293        self,
1294        out_file_path,
1295        feature_ids=None,
1296        *,
1297        ms2_mode="best",
1298        overwrite=False,
1299    ):
1300        """Export mass features to a SIRIUS-compatible MGF file.
1301
1302        Each complete feature is written as a one-peak MS1 block (feature
1303        precursor m/z and intensity) and one or more MS2 blocks sharing
1304        FEATURE_ID. ``ms2_mode='best'`` writes ``LCMSMassFeature.best_ms2``:
1305        the highest library-similarity scan after a spectral search,
1306        otherwise the non-chimeric scan closest to the apex. Features
1307        missing a usable MS2 spectrum are skipped. A feature with no
1308        precursor m/z aborts the export before any bytes are written.
1309
1310        Parameters
1311        ----------
1312        out_file_path : str or Path
1313            Output path. ``.mgf`` is appended when no suffix is given.
1314        feature_ids : iterable, optional
1315            Mass-feature ids to export. Default is all keys in ``mass_features``.
1316        ms2_mode : {'best', 'all'}, optional
1317            Which associated MS2 spectra to write. ``'best'`` (default)
1318            writes ``LCMSMassFeature.best_ms2``. After a spectral search
1319            that is the scan with the highest library entropy similarity;
1320            with no search results it is the non-chimeric scan closest to
1321            the apex. ``'all'`` writes every MS2 scan that has peaks.
1322        overwrite : bool, optional
1323            Replace an existing file. Default is False.
1324
1325        Returns
1326        -------
1327        pathlib.Path
1328            Path of the written MGF file.
1329
1330        Warns
1331        -----
1332        UserWarning
1333            If one or more features are skipped because they have no usable
1334            MS2 spectrum.
1335
1336        Raises
1337        ------
1338        FileExistsError
1339            If the path exists and ``overwrite`` is False.
1340        ValueError
1341            If a requested mass-feature id is unknown, polarity is not
1342            ``'positive'`` or ``'negative'``, a feature has no precursor
1343            m/z, or no complete MS1/MS2 pairs remain. A missing precursor
1344            m/z stops the export before the file is written.
1345        """
1346        from corems.mass_spectra.output.mgf import write_feature_records_to_mgf
1347
1348        if feature_ids is None:
1349            ids = list(self.mass_features.keys())
1350        else:
1351            ids = list(feature_ids)
1352            missing = [i for i in ids if i not in self.mass_features]
1353            if missing:
1354                raise ValueError(f"Mass feature id(s) not found: {missing}")
1355
1356        records = [
1357            (mf_id, self.mass_features[mf_id], getattr(self, "sample_name", None))
1358            for mf_id in ids
1359        ]
1360        return write_feature_records_to_mgf(
1361            records,
1362            polarity=self.polarity,
1363            out_file_path=out_file_path,
1364            ms2_mode=ms2_mode,
1365            overwrite=overwrite,
1366        )
1367
1368    def mass_features_ms1_annot_to_df(self, suppress_warnings=False):
1369        """Returns a pandas dataframe summarizing the MS1 annotations for the mass features in the dataset.
1370
1371        Parameters
1372        -----------
1373        suppress_warnings : bool, optional
1374            If True, suppresses the warning when no MS1 annotations are found.
1375            Useful when calling from collection-level methods. Default is False.
1376
1377        Returns
1378        --------
1379        pandas.DataFrame
1380            A pandas dataframe of MS1 annotations for the mass features in the dataset.
1381            The index is set to mf_id (mass feature ID)
1382
1383        Raises
1384        ------
1385        Warning
1386            If no MS1 annotations were found for the mass features in the dataset
1387            (unless suppress_warnings=True).
1388        """
1389        annot_df_list_ms1 = []
1390        for mf_id in self.mass_features.keys():
1391            if self.mass_features[mf_id].mass_spectrum is None:
1392                pass
1393            else:
1394                # Add ms1 annotations to ms1 annotation list
1395                if (
1396                    np.abs(
1397                        (
1398                            self.mass_features[mf_id].ms1_peak.mz_exp
1399                            - self.mass_features[mf_id].mz
1400                        )
1401                    )
1402                    < 0.01
1403                ):
1404                    # Get the molecular formula from the mass spectrum
1405                    annot_df = self.mass_features[mf_id].mass_spectrum.to_dataframe()
1406                    # Subset to pull out only the peak associated with the mass feature
1407                    annot_df = annot_df[
1408                        annot_df["Index"] == self.mass_features[mf_id].ms1_peak.index
1409                    ].copy()
1410
1411                    # If there are more than 1 row, remove any rows without a molecular formula
1412                    if len(annot_df) > 1:
1413                        annot_df = annot_df[~annot_df["Molecular Formula"].isna()]
1414
1415                    # Remove the index column and add column for mf_id
1416                    annot_df = annot_df.drop(columns=["Index"])
1417                    annot_df["mf_id"] = mf_id
1418                    annot_df_list_ms1.append(annot_df)
1419
1420        if len(annot_df_list_ms1) > 0:
1421            annot_ms1_df_full = pd.concat(annot_df_list_ms1)
1422            annot_ms1_df_full = annot_ms1_df_full.set_index("mf_id")
1423            annot_ms1_df_full.index.name = "mf_id"
1424
1425        else:
1426            annot_ms1_df_full = None
1427            # Warn that no ms1 annotations were found (unless suppressed)
1428            if not suppress_warnings:
1429                warnings.warn(
1430                    "No MS1 annotations found for mass features in dataset, were MS1 spectra added and processed within the dataset?",
1431                    UserWarning,
1432                )
1433
1434        return annot_ms1_df_full
1435
1436    def mass_features_ms2_annot_to_df(
1437        self, molecular_metadata=None, suppress_warnings=False
1438    ):
1439        """Returns a pandas dataframe summarizing the MS2 annotations for the mass features in the dataset.
1440
1441        Parameters
1442        -----------
1443        molecular_metadata :  dict of MolecularMetadata objects
1444            A dictionary of MolecularMetadata objects, keyed by ref_mol_id. Defaults to None.
1445        suppress_warnings : bool, optional
1446            If True, suppresses the warning when no MS2 annotations are found.
1447            Useful when calling from collection-level methods. Default is False.
1448
1449        Returns
1450        --------
1451        pandas.DataFrame
1452            A pandas dataframe of MS2 annotations for the mass features in the dataset,
1453            and optionally molecular metadata. The index is set to mf_id (mass feature ID)
1454
1455        Raises
1456        ------
1457        Warning
1458            If no MS2 annotations were found for the mass features in the dataset
1459            (unless suppress_warnings=True).
1460        """
1461        annot_df_list_ms2 = []
1462        for mf_id in self.mass_features.keys():
1463            if len(self.mass_features[mf_id].ms2_similarity_results) > 0:
1464                # Add ms2 annotations to ms2 annotation list
1465                for result in self.mass_features[mf_id].ms2_similarity_results:
1466                    annot_df_ms2 = result.to_dataframe()
1467                    annot_df_ms2["mf_id"] = mf_id
1468                    annot_df_list_ms2.append(annot_df_ms2)
1469
1470        if len(annot_df_list_ms2) > 0:
1471            annot_ms2_df_full = pd.concat(annot_df_list_ms2)
1472            if molecular_metadata is not None:
1473                molecular_metadata_df = pd.concat(
1474                    [
1475                        pd.DataFrame.from_dict(v.__dict__, orient="index").transpose()
1476                        for k, v in molecular_metadata.items()
1477                    ],
1478                    ignore_index=True,
1479                )
1480                molecular_metadata_df = molecular_metadata_df.rename(
1481                    columns={"id": "ref_mol_id"}
1482                )
1483                annot_ms2_df_full = annot_ms2_df_full.merge(
1484                    molecular_metadata_df, on="ref_mol_id", how="left"
1485                )
1486            annot_ms2_df_full = annot_ms2_df_full.drop_duplicates(
1487                subset=["mf_id", "query_spectrum_id", "ref_ms_id"]
1488            ).copy()
1489            annot_ms2_df_full = annot_ms2_df_full.set_index("mf_id")
1490            annot_ms2_df_full.index.name = "mf_id"
1491        else:
1492            annot_ms2_df_full = None
1493            # Warn that no ms2 annotations were found (unless suppressed)
1494            if not suppress_warnings:
1495                warnings.warn(
1496                    "No MS2 annotations found for mass features in dataset, were MS2 spectra added and searched against a database?",
1497                    UserWarning,
1498                )
1499
1500        return annot_ms2_df_full
1501
1502    def plot_composite_mz_features(
1503        self,
1504        binsize=1e-4,
1505        ph_int_min_thresh=0.001,
1506        mf_plot=True,
1507        ms2_plot=True,
1508        return_fig=False,
1509        path=None,
1510    ):
1511        """Returns a figure displaying
1512            (1) thresholded, unprocessed data
1513            (2) the m/z features
1514            (3) which m/z features are associated with MS2 spectra
1515
1516        Parameters
1517        -----------
1518        binsize :  float
1519            Desired binsize for the m/z axis of the composite feature map.  Defaults to 1e-4.
1520        mf_plot : boolean
1521            Indicates whether to plot the m/z features. Defaults to True.
1522        ms2_plot : boolean
1523            Indicates whether to identify m/z features with associated MS2 spectra. Defaults to True.
1524        return_fig : bool, optional
1525            If True, return the open figure (caller owns lifecycle).
1526            Default is False.
1527        path : str or path-like, optional
1528            If set, save the figure to this path. When ``return_fig`` is False,
1529            the figure is closed after saving and ``plt.show()`` is not called.
1530
1531        Returns
1532        --------
1533        matplotlib.figure.Figure or None
1534            A figure with the thresholded, unprocessed data on an axis of m/z
1535            value with respect to scan time if ``return_fig`` is True; otherwise
1536            None. Unprocessed data is displayed in gray scale with darker colors
1537            indicating higher intensities. If m/z features are plotted, they are
1538            displayed in cyan. If m/z features associated with MS2 spectra are
1539            plotted, they are displayed in red.
1540
1541        Raises
1542        ------
1543        ValueError
1544            If m/z features are set to be plot but aren't in the dataset.
1545            If m/z features with associated MS2 data are set to be plot but no
1546            MS2 annotations were found for the m/z features in the dataset.
1547        """
1548        if mf_plot:
1549            # Check if mass_features is set, raise error if not
1550            if self.mass_features is None:
1551                raise ValueError(
1552                    "mass_features not set, must run find_mass_features() first"
1553                )
1554            ## call mass feature data
1555            mf_df = self.mass_features_to_df()
1556
1557        if ms2_plot:
1558            if not mf_plot:
1559                # Check if mass_features is set, raise error if not
1560                if self.mass_features is None:
1561                    raise ValueError(
1562                        "mass_features not set, must run find_mass_features() first"
1563                    )
1564
1565            ## call m/z feature data
1566            mf_df = self.mass_features_to_df()
1567
1568            # Check if ms2_spectrum is set, raise error if not
1569            if "ms2_spectrum" not in mf_df.columns:
1570                raise ValueError(
1571                    "ms2_spectrum not set, must run add_associated_ms2_dda() first"
1572                )
1573
1574        ## threshold and grid unprocessed data
1575        df = self._ms_unprocessed[1].copy()
1576        df = df.dropna(subset=["intensity"]).reset_index(drop=True)
1577        threshold = ph_int_min_thresh * df.intensity.max()
1578        df_thres = df[df["intensity"] > threshold].reset_index(drop=True).copy()
1579        df = self.grid_data(df_thres)
1580
1581        ## format unprocessed data for plotting
1582        df = df.merge(self.scan_df[["scan", "scan_time"]], on="scan")
1583        mz_grid = np.arange(0, np.max(df.mz), binsize)
1584        mz_data = np.array(df.mz)
1585        df["mz_bin"] = find_closest(mz_grid, mz_data)
1586        df["ab_bin"] = df.groupby(["mz_bin", "scan_time"]).intensity.transform(sum)
1587        unproc_df = df[["scan_time", "mz_bin", "ab_bin"]].drop_duplicates(
1588            ignore_index=True
1589        )
1590
1591        ## generate figure
1592        fig = plt.figure()
1593        plt.scatter(
1594            unproc_df.scan_time,
1595            unproc_df.mz_bin * binsize,
1596            c=unproc_df.ab_bin / np.max(unproc_df.ab_bin),
1597            alpha=unproc_df.ab_bin / np.max(unproc_df.ab_bin),
1598            cmap="Greys_r",
1599            s=1,
1600        )
1601
1602        if mf_plot:
1603            if ms2_plot:
1604                plt.scatter(
1605                    mf_df[mf_df.ms2_spectrum.isna()].scan_time,
1606                    mf_df[mf_df.ms2_spectrum.isna()].mz,
1607                    c="c",
1608                    s=4,
1609                    label="M/Z features without MS2",
1610                )
1611            else:
1612                plt.scatter(mf_df.scan_time, mf_df.mz, c="c", s=4, label="M/Z features")
1613
1614        if ms2_plot:
1615            plt.scatter(
1616                mf_df[~mf_df.ms2_spectrum.isna()].scan_time,
1617                mf_df[~mf_df.ms2_spectrum.isna()].mz,
1618                c="r",
1619                s=2,
1620                label="M/Z features with MS2",
1621            )
1622
1623        if mf_plot == True or ms2_plot == True:
1624            plt.legend(loc="lower center", bbox_to_anchor=(0.5, -0.25), ncol=2)
1625        plt.xlabel("Scan time")
1626        plt.ylabel("m/z")
1627        plt.ylim(0, np.ceil(np.max(df.mz)))
1628        plt.xlim(0, np.ceil(np.max(df.scan_time)))
1629        plt.title("Composite Feature Map")
1630
1631        return _finalize_plot(fig, return_fig=return_fig, path=path)
1632
1633    def search_for_targeted_mass_features_batch(
1634        self,
1635        ms1df,
1636        mz_mins,
1637        mz_maxs,
1638        st_mins,
1639        st_maxs,
1640        set_ids,
1641        obj_idx=0,
1642        st_aligned=False,
1643    ):
1644        """
1645        Returns multiple LCMSMassFeatures from a specific sample within specific mass and time ranges.
1646        Sorts m/z once and binary-searches each window. The selected row is the
1647        highest intensity in the window, with the same tie break as a full scan.
1648
1649        Parameters
1650        -----------
1651        ms1df : pd.DataFrame
1652            Dataframe containing all the possible MS1 values to consider, collected by calling _ms_unprocessed[1] on the sample.
1653        mz_mins : np.ndarray
1654            Array of lower bounds of m/z values to use to find peaks.
1655        mz_maxs : np.ndarray
1656            Array of upper bounds of m/z values to use to find peaks.
1657        st_mins : np.ndarray
1658            Array of lower bounds of scan times to use to find peaks.
1659        st_maxs : np.ndarray
1660            Array of upper bounds of scan times to use to find peaks.
1661        set_ids : np.ndarray or list
1662            Array of strings used as IDs in LCMSMassFeatures.
1663        obj_idx : int
1664            Identifies index of sample in a collection. Defaults to 0.
1665        st_aligned : bool
1666            Whether to use scan_time_aligned or scan_time. Defaults to False.
1667
1668        Returns
1669        --------
1670        dict
1671            Dictionary mapping set_id to LCMSMassFeature objects.
1672
1673        Raises
1674        ------
1675        ValueError
1676            If appropriate scan time data is not contained in ms1df or if array lengths don't match.
1677        """
1678        # Validate inputs
1679        n_features = len(mz_mins)
1680        if not all(
1681            len(arr) == n_features for arr in [mz_maxs, st_mins, st_maxs, set_ids]
1682        ):
1683            raise ValueError("All input arrays must have the same length")
1684
1685        # Validate scan time column
1686        time_col = "scan_time_aligned" if st_aligned else "scan_time"
1687        if time_col not in ms1df.columns:
1688            raise ValueError(f"{time_col} not contained in ms1df")
1689
1690        # Pre-extract columns for faster access
1691        mz_vals = ms1df.mz.values
1692        st_vals = ms1df[time_col].values
1693        scan_vals = ms1df.scan.values
1694        intensity_vals = ms1df.intensity.values
1695
1696        # Index finite m/z once. Comparisons promote to float64, and a NaN m/z
1697        # never falls inside a window, so those rows stay out of the index.
1698        mz64 = np.asarray(mz_vals, dtype=np.float64)
1699        finite_idx = np.flatnonzero(np.isfinite(mz64))
1700        mz_order = finite_idx[np.argsort(mz64[finite_idx])]
1701        mz_sorted = mz64[mz_order]
1702        mz_mins = np.asarray(mz_mins, dtype=np.float64)
1703        mz_maxs = np.asarray(mz_maxs, dtype=np.float64)
1704        st_mins = np.asarray(st_mins, dtype=np.float64)
1705        st_maxs = np.asarray(st_maxs, dtype=np.float64)
1706        st64 = np.asarray(st_vals, dtype=np.float64)
1707        lo = np.searchsorted(mz_sorted, mz_mins, side="left")
1708        hi = np.searchsorted(mz_sorted, mz_maxs, side="right")
1709        intensity_is_float = np.issubdtype(intensity_vals.dtype, np.floating)
1710
1711        def _empty_feature(set_id):
1712            return LCMSMassFeature(
1713                self,
1714                apex_scan=-99,
1715                mz=np.nan,
1716                intensity=np.nan,
1717                retention_time=np.nan,
1718                persistence=np.nan,
1719                id=set_id,
1720            )
1721
1722        results = {}
1723        for i in range(n_features):
1724            set_id = set_ids[i]
1725            bounds_ok = (
1726                np.isfinite(mz_mins[i])
1727                and np.isfinite(mz_maxs[i])
1728                and mz_mins[i] <= mz_maxs[i]
1729            )
1730            if not bounds_ok or lo[i] >= hi[i]:
1731                results[set_id] = _empty_feature(set_id)
1732                continue
1733
1734            idx = mz_order[lo[i] : hi[i]]
1735            in_time = (st64[idx] >= st_mins[i]) & (st64[idx] <= st_maxs[i])
1736            idx = idx[in_time]
1737            if idx.size == 0:
1738                results[set_id] = _empty_feature(set_id)
1739                continue
1740
1741            # Same row as np.where(mask) + np.argmax: highest intensity, and
1742            # the earliest original row when several share that intensity.
1743            # np.argmax also prefers a NaN intensity over any finite value.
1744            intensities = intensity_vals[idx]
1745            if intensity_is_float and np.isnan(intensities).any():
1746                idx = np.sort(idx)
1747                peak_idx = idx[np.argmax(intensity_vals[idx])]
1748            else:
1749                max_intensity = intensities.max()
1750                peak_idx = idx[intensities == max_intensity].min()
1751
1752            results[set_id] = LCMSMassFeature(
1753                self,
1754                apex_scan=scan_vals[peak_idx],
1755                mz=mz_vals[peak_idx],
1756                intensity=intensity_vals[peak_idx],
1757                retention_time=st_vals[peak_idx],
1758                persistence=np.nan,
1759                id=set_id,
1760            )
1761
1762        return results
1763
1764    def search_for_targeted_mass_feature(
1765        self, ms1df, mz_min, mz_max, st_min, st_max, set_id, obj_idx=0, st_aligned=False
1766    ):
1767        """
1768        Returns an LCMSMassFeature from a specific sample within a specific mass and time range. Returns an empty
1769        LCMSMassFeature if no satisfactory peak is found in the given window.
1770
1771        Parameters
1772        -----------
1773        ms1df :  Pandas DataFrame
1774            Dataframe containing all the possible MS1 values to consider, collected by calling _ms_unprocessed[1] on the sample.
1775        mz_min : float
1776            Identifies lower bound of the weights to use to find a peak.
1777        mz_max : float
1778            Identifies upper bound of the weights to use to find a peak.
1779        st_min : float
1780            Identifies lower bound of the scan times to use to find a peak.
1781        st_max : float
1782            Identifies upper bound of the scan times to use to find a peak.
1783        set_id : str
1784            Indicates string used as ID in LCMSMassFeature.
1785        obj_idx : int
1786            Identifies index of sample in a collection that LCMSMassFeature should be assigned to. Defaults to 0 and is not used
1787            if data provided is an LCMSBase instead of an LCMSCollection.
1788        st_aligned : boolean
1789            Indicates whether to call scan time from scan_time or from scan_time_aligned if using a collection. Defaults to False.
1790
1791        Returns
1792        --------
1793        LCMSMassFeature
1794            Object from ChromaPeak that contains data on selected MS1 peak. If no peak is found, will contain missing
1795            information and list the apex scan value as -99.
1796
1797        Raises
1798        ------
1799        Warning
1800            If appropriate scan time data is not contained in ms1df.
1801        """
1802        # Convert single feature to arrays and call batch method
1803        results = self.search_for_targeted_mass_features_batch(
1804            ms1df,
1805            np.array([mz_min]),
1806            np.array([mz_max]),
1807            np.array([st_min]),
1808            np.array([st_max]),
1809            [set_id],
1810            obj_idx=obj_idx,
1811            st_aligned=st_aligned,
1812        )
1813        return results[set_id]
1814
1815    def __len__(self):
1816        """
1817        Returns the number of mass spectra in the dataset.
1818
1819        Returns
1820        --------
1821        int
1822            The number of mass spectra in the dataset.
1823        """
1824        return len(self._ms)
1825
1826    def __getitem__(self, scan_number):
1827        """
1828        Returns the mass spectrum corresponding to the specified scan number.
1829
1830        Parameters
1831        -----------
1832        scan_number : int
1833            The scan number of the desired mass spectrum.
1834
1835        Returns
1836        --------
1837        MassSpectrum
1838            The mass spectrum corresponding to the specified scan number.
1839        """
1840        return self._ms.get(scan_number)
1841
1842    def __iter__(self):
1843        """Returns an iterator over the mass spectra in the dataset.
1844
1845        Returns
1846        --------
1847        iterator
1848            An iterator over the mass spectra in the dataset.
1849        """
1850        return iter(self._ms.values())
1851
1852    def set_tic_list_from_data(self, overwrite=False):
1853        """Sets the TIC list from the mass spectrum objects within the _ms dictionary.
1854
1855        Parameters
1856        -----------
1857        overwrite : bool, optional
1858            If True, overwrites the TIC list if it is already set. Defaults to False.
1859
1860        Notes
1861        -----
1862        If the _ms dictionary is incomplete, sets the TIC list to an empty list.
1863
1864        Raises
1865        ------
1866        ValueError
1867            If no mass spectra are found in the dataset.
1868            If the TIC list is already set and overwrite is False.
1869        """
1870        # Check if _ms is empty and raise error if so
1871        if len(self._ms) == 0:
1872            raise ValueError("No mass spectra found in dataset")
1873
1874        # Check if tic_list is already set and raise error if so
1875        if len(self.tic) > 0 and not overwrite:
1876            raise ValueError("TIC list already set, use overwrite=True to overwrite")
1877
1878        self.tic = [self._ms.get(i).tic for i in self.scans_number]
1879
1880    def set_retention_time_from_data(self, overwrite=False):
1881        """Sets the retention time list from the data in the _ms dictionary.
1882
1883        Parameters
1884        -----------
1885        overwrite : bool, optional
1886            If True, overwrites the retention time list if it is already set. Defaults to False.
1887
1888        Notes
1889        -----
1890        If the _ms dictionary is empty or incomplete, sets the retention time list to an empty list.
1891
1892        Raises
1893        ------
1894        ValueError
1895            If no mass spectra are found in the dataset.
1896            If the retention time list is already set and overwrite is False.
1897        """
1898        # Check if _ms is empty and raise error if so
1899        if len(self._ms) == 0:
1900            raise ValueError("No mass spectra found in dataset")
1901
1902        # Check if retention_time_list is already set and raise error if so
1903        if len(self.retention_time) > 0 and not overwrite:
1904            raise ValueError(
1905                "Retention time list already set, use overwrite=True to overwrite"
1906            )
1907
1908        retention_time_list = []
1909        for key_ms in sorted(self._ms.keys()):
1910            retention_time_list.append(self._ms.get(key_ms).retention_time)
1911        self.retention_time = retention_time_list
1912
1913    def set_scans_number_from_data(self, overwrite=False):
1914        """Sets the scan number list from the data in the _ms dictionary.
1915
1916        Notes
1917        -----
1918        If the _ms dictionary is empty or incomplete, sets the scan number list to an empty list.
1919
1920        Raises
1921        ------
1922        ValueError
1923            If no mass spectra are found in the dataset.
1924            If the scan number list is already set and overwrite is False.
1925        """
1926        # Check if _ms is empty and raise error if so
1927        if len(self._ms) == 0:
1928            raise ValueError("No mass spectra found in dataset")
1929
1930        # Check if scans_number_list is already set and raise error if so
1931        if len(self.scans_number) > 0 and not overwrite:
1932            raise ValueError(
1933                "Scan number list already set, use overwrite=True to overwrite"
1934            )
1935
1936        self.scans_number = sorted(self._ms.keys())
1937
1938    @property
1939    def ms1_scans(self):
1940        """
1941        list : A list of MS1 scan numbers for the dataset.
1942        """
1943        return self.scan_df[self.scan_df.ms_level == 1].index.tolist()
1944
1945    @property
1946    def parameters(self):
1947        """
1948        LCMSParameters : The parameters used for the LC-MS analysis.
1949        """
1950        return self._parameters
1951
1952    @parameters.setter
1953    def parameters(self, paramsinstance):
1954        """
1955        Sets the parameters used for the LC-MS analysis.
1956
1957        Parameters
1958        -----------
1959        paramsinstance : LCMSParameters
1960            The parameters used for the LC-MS analysis.
1961        """
1962        self._parameters = paramsinstance
1963
1964    @property
1965    def scans_number(self):
1966        """
1967        list : A list of scan numbers for the dataset.
1968        """
1969        return self._scans_number_list
1970
1971    @scans_number.setter
1972    def scans_number(self, scan_numbers_list):
1973        """
1974        Sets the scan numbers for the dataset.
1975
1976        Parameters
1977        -----------
1978        scan_numbers_list : list
1979            A list of scan numbers for the dataset.
1980        """
1981        self._scans_number_list = scan_numbers_list
1982
1983    @property
1984    def retention_time(self):
1985        """
1986        numpy.ndarray : An array of retention times for the dataset.
1987        """
1988        return self._retention_time_list
1989
1990    @retention_time.setter
1991    def retention_time(self, rt_list):
1992        """
1993        Sets the retention times for the dataset.
1994
1995        Parameters
1996        -----------
1997        rt_list : list
1998            A list of retention times for the dataset.
1999        """
2000        self._retention_time_list = np.array(rt_list)
2001
2002    @property
2003    def tic(self):
2004        """
2005        numpy.ndarray : An array of TIC values for the dataset.
2006        """
2007        return self._tic_list
2008
2009    @tic.setter
2010    def tic(self, tic_list):
2011        """
2012        Sets the TIC values for the dataset.
2013
2014        Parameters
2015        -----------
2016        tic_list : list
2017            A list of TIC values for the dataset.
2018        """
2019        self._tic_list = np.array(tic_list)
2020
2021
2022class LCMSCollection(LCMSCollectionCalculations):
2023    """A class representing a collection of liquid chromatography-mass spectrometry (LC-MS) runs.
2024    These runs can be from the same or different samples, but must be from the same instrument and have the same parameters
2025    for the initial processing steps.  The LCMS objects are stored in an ordered dictionary with the sample name as the key.
2026
2027    Parameters
2028    -----------
2029
2030    Attributes
2031    -----------
2032
2033    Methods
2034    --------
2035
2036    Notes
2037    ------
2038    This class is not intended to be instantiated directly, but rather instantiated using a parser object and then interacted with.
2039    #TODO KRH: add docstrings
2040    """
2041
2042    def __init__(self, collection_location, manifest, collection_parser=None):
2043        self.collection_location = collection_location
2044        self._manifest_dict = manifest
2045        self.collection_parser = collection_parser
2046        self.raw_files_relocated = False
2047
2048        # These attributes are generally set by the parser during instantiation of this class
2049        self._lcms = {}
2050        self._combined_mass_features = None
2051        self._combined_induced_mass_features = None
2052        self.consensus_mass_features = {}
2053        self._parameters = LCMSCollectionParameters()
2054        self.isotopes_dropped = False
2055        self._mass_features_locked = (
2056            False  # Prevents rebuilding mass_features_dataframe from samples
2057        )
2058        # Consensus feature grouping labels (isotopes + z=1 adducts);
2059        # index = cluster. Not tracer/enriched labeling.
2060        self.feature_group_dataframe = None
2061
2062        # These attributes are set during processing
2063        self.rt_aligned = False
2064        self.rt_alignment_attempted = False
2065        self.missing_mass_features_searched = False
2066
2067    def _reorder_lcms_objects(self):
2068        """
2069        Reorders the LCMS objects in the collection based on the order in the manifest.
2070        """
2071        ordered_samples = self.samples
2072        self._lcms = {k: self._lcms[k] for k in ordered_samples}
2073
2074    def __getitem__(self, index):
2075        if isinstance(index, (float, np.floating, np.ndarray)):
2076            index = int(index)
2077        elif isinstance(index, np.integer):
2078            index = int(index)
2079        samp_name = self.samples[index]
2080        self._lcms[samp_name]
2081        return self._lcms[samp_name]
2082
2083    def __len__(self):
2084        return len(self.samples)
2085
2086    def _prepare_lcms_mass_features_for_combination(
2087        self, lcms_obj, induced_features=False
2088    ):
2089        """
2090        Prepares the mass features in the LCMS objects in the collection for combination.
2091        """
2092        if induced_features:
2093            mf_df = lcms_obj.mass_features_to_df(induced_features=True)
2094        # Check if lcms_obj has attribute light_mf_df
2095        elif hasattr(lcms_obj, "light_mf_df"):
2096            mf_df = lcms_obj.light_mf_df
2097        else:
2098            mf_df = lcms_obj.mass_features_to_df()
2099
2100        # If dataframe is empty, add minimal required columns and return
2101        if len(mf_df) == 0:
2102
2103            mf_df["sample_name"] = []
2104            mf_df["sample_id"] = []
2105            mf_df["coll_mf_id"] = []
2106            mf_df["mf_id"] = []
2107            mf_df[
2108                "_eic_mz"
2109            ] = []  # Include _eic_mz for consistency with non-empty dataframes
2110            if induced_features:
2111                mf_df["cluster"] = []
2112            return mf_df
2113
2114        # Remove index
2115        mf_df = mf_df.reset_index(drop=False)
2116        # Add sample name and sample id to the dataframe
2117        mf_df["sample_name"] = lcms_obj.sample_name
2118        # Ensure sample_id is stored as an integer to avoid float indices later
2119        try:
2120            mf_df["sample_id"] = int(
2121                self.manifest[lcms_obj.sample_name]["collection_id"]
2122            )
2123        except Exception:
2124            mf_df["sample_id"] = self.manifest[lcms_obj.sample_name]["collection_id"]
2125        mf_df["coll_mf_id"] = (
2126            mf_df["sample_id"].astype(str) + "_" + mf_df["mf_id"].astype(str)
2127        )
2128
2129        # For induced features, extract cluster from mf_id (format: c{cluster}_{index}_i)
2130        # and add as a column since cluster_index attribute may not be set on the object
2131        if induced_features:
2132
2133            def extract_cluster(mf_id):
2134                # mf_id format: c{cluster}_{index}_i
2135                # Example: c123_5_i -> cluster 123
2136                if isinstance(mf_id, str) and mf_id.startswith("c") and "_i" in mf_id:
2137                    parts = mf_id.split("_")
2138                    if len(parts) >= 2:
2139                        cluster_str = parts[0][1:]  # Remove 'c' prefix
2140                        try:
2141                            return int(cluster_str)
2142                        except ValueError:
2143                            return None
2144                return None
2145
2146            mf_df["cluster"] = mf_df["mf_id"].apply(extract_cluster)
2147
2148        # Check if scan_df has scan_time_aligned and add to mf_df if so
2149        if "scan_time_aligned" in lcms_obj.scan_df.columns:
2150            scan_df = lcms_obj.scan_df[["scan", "scan_time_aligned"]].copy()
2151            scan_df = scan_df.rename(columns={"scan": "apex_scan"})
2152            mf_df = mf_df.merge(scan_df, on="apex_scan")
2153
2154        return mf_df
2155
2156    def _combine_mass_features(self, induced_features=False):
2157        """
2158        Concatenates the mass features from all the LCMS objects in the collection.
2159
2160        Returns
2161        --------
2162        None, sets the _combined_mass_features or _combined_induced_mass_feature attribute.
2163
2164        Notes
2165        -----
2166        If _mass_features_locked is True (e.g., when only representative features are loaded),
2167        this method will skip rebuilding the regular mass features dataframe to preserve
2168        the full collection-level dataframe. Induced features are always rebuilt since they
2169        are created during processing.
2170        """
2171
2172        # Skip rebuilding regular mass features if locked (preserves full dataframe)
2173        if not induced_features and self._mass_features_locked:
2174            return
2175
2176        ## TODO: See why this function runs slower on multiprocessing,
2177        ## especially for induced features
2178        ## has only been considered so far on ~20 samples
2179        #        if self.parameters.lcms_collection.cores == 1:
2180        #            # Prepare mass features for combination sequentially
2181        #            mf_df_list = []
2182        #            for lcms_obj in self:
2183        #                mf_df = self._prepare_lcms_mass_features_for_combination(lcms_obj, induced_features)
2184        #                mf_df_list.append(mf_df)
2185
2186        #        if self.parameters.lcms_collection.cores > 1:
2187        #            # Parallelize the mass feature preparation
2188        #            if self.parameters.lcms_collection.cores > len(self):
2189        #                ncores = len(self)
2190        #            else:
2191        #                ncores = self.parameters.lcms_collection.cores
2192        #            pool = multiprocessing.Pool(ncores)
2193        #            mf_df_list = pool.starmap(
2194        #                self._prepare_lcms_mass_features_for_combination,
2195        #                [(lcms_obj, induced_features) for lcms_obj in self]
2196        #            )
2197
2198        # Prepare mass features for combination sequentially
2199        mf_df_list = []
2200        for lcms_obj in self:
2201            # Skip samples with no induced mass features if processing induced features
2202            if induced_features:
2203                if (
2204                    not hasattr(lcms_obj, "induced_mass_features")
2205                    or len(lcms_obj.induced_mass_features) == 0
2206                ):
2207                    continue
2208            mf_df = self._prepare_lcms_mass_features_for_combination(
2209                lcms_obj, induced_features
2210            )
2211            mf_df_list.append(mf_df)
2212
2213        # If no mass features were collected (e.g., no induced features exist), return early
2214        if len(mf_df_list) == 0:
2215            # Add a warning here, not sure how one might reach this state, clearly saying if they are induced features or not
2216            warnings.warn(
2217                "No mass features found to combine in the collection.", UserWarning
2218            )
2219            if induced_features:
2220                self._combined_induced_mass_features = None
2221            else:
2222                self._combined_mass_features = None
2223            return
2224
2225        # A sample of gap-fill misses stores area as all-NA object. A sample
2226        # with a measured area stores it as float. Concat warns on pandas 2
2227        # and can change that float's dtype. Drop the empty column first.
2228        combined_mass_features = pd.concat(_drop_all_na_float_columns(mf_df_list))
2229        # Ensure sample_id and cluster columns have integer dtypes where possible
2230        if "sample_id" in combined_mass_features.columns:
2231            try:
2232                combined_mass_features["sample_id"] = combined_mass_features[
2233                    "sample_id"
2234                ].astype(int)
2235            except Exception:
2236                combined_mass_features["sample_id"] = pd.to_numeric(
2237                    combined_mass_features["sample_id"], errors="coerce"
2238                ).astype("Int64")
2239        if "cluster" in combined_mass_features.columns:
2240            try:
2241                combined_mass_features["cluster"] = combined_mass_features[
2242                    "cluster"
2243                ].astype(int)
2244            except Exception:
2245                combined_mass_features["cluster"] = pd.to_numeric(
2246                    combined_mass_features["cluster"], errors="coerce"
2247                ).astype("Int64")
2248        # Move coll_mf_id, sample_name, sample_id, and mf_id to front
2249        cols = combined_mass_features.columns.tolist()
2250        top_cols = [
2251            "coll_mf_id",
2252            "sample_name",
2253            "sample_id",
2254            "mf_id",
2255            "mz",
2256            "scan_time_aligned",
2257            "cluster",
2258        ]
2259        cols = [
2260            x
2261            for x in top_cols + [col for col in cols if col not in top_cols]
2262            if x in cols
2263        ]
2264        combined_mass_features = combined_mass_features[cols]
2265        # Make coll_mf_id the index
2266        combined_mass_features = combined_mass_features.set_index("coll_mf_id")
2267        if induced_features == True:
2268            self._combined_induced_mass_features = combined_mass_features
2269        else:
2270            self._combined_mass_features = combined_mass_features
2271
2272    def _check_mass_features_df(self, induced_features=False):
2273        """Checks if the mass features dataframe has expected columns.  If not, adds them.
2274
2275        Returns
2276        --------
2277        pandas.DataFrame
2278            A pandas dataframe of mass features in the collection.
2279
2280        Notes
2281        ------
2282        If scan_time_aligned is not in the _combined_mass_features or
2283        _combined_induced_mass_features, tries to add it.
2284
2285        """
2286
2287        if induced_features:
2288            cmf_df = self._combined_induced_mass_features
2289        else:
2290            cmf_df = self._combined_mass_features
2291        # Check if parameters are set to drop isotopologues and drop if so
2292        if self.parameters.lcms_collection.drop_isotopologues:
2293            if not self.isotopes_dropped:
2294                self._drop_isotopologues()
2295        # Check if scan_time_aligned is in combined_mass_features, try to add if not
2296        if cmf_df is not None and "scan_time_aligned" not in cmf_df.columns:
2297            cmb_mf = cmf_df.copy()
2298            cmb_mf = cmb_mf.reset_index(drop=False)
2299            lcms_aligned = [
2300                True for x in self if "scan_time_aligned" in x.scan_df.columns
2301            ]
2302            if len(lcms_aligned) == len(self):
2303                # Add scan_time_aligned to combined_mass_features dataframe
2304                scan_time_aligned_list = []
2305                for lcms_obj in self:
2306                    scan_time_df_i = lcms_obj.scan_df[["scan", "scan_time_aligned"]]
2307                    scan_time_df_i["sample_name"] = lcms_obj.sample_name
2308                    scan_time_aligned_list.append(scan_time_df_i)
2309                scan_time_aligned_df = pd.concat(scan_time_aligned_list)
2310                # Rename scan to apex_scan
2311                scan_time_aligned_df = scan_time_aligned_df.rename(
2312                    columns={"scan": "apex_scan"}
2313                )
2314                cmb_mf_merged = cmb_mf.merge(
2315                    scan_time_aligned_df, on=["apex_scan", "sample_name"]
2316                )
2317                cmb_mf_merged = cmb_mf_merged.set_index("coll_mf_id")
2318                # Merge scan_time_aligned_df with combined_mass_features on apex_scan and sample_name
2319                if induced_features:
2320                    self._combined_induced_mass_features = cmb_mf_merged
2321                else:
2322                    self._combined_mass_features = cmb_mf_merged
2323
2324    def plot_tics(
2325        self, ms_level=1, type="raw", plot_legend=False, return_fig=False, path=None
2326    ):
2327        """Plots the TICs for all the LCMS objects in the collection.
2328
2329        Parameters
2330        -----------
2331        ms_level : int, optional
2332            The MS level to plot the TICs for. Defaults to 1.
2333        type : str, optional
2334            The type of TIC to plot, either "raw" or "corrected" or "both". Defaults to "raw".
2335        plot_legend : bool, optional
2336            If True, plots a legend on the TIC plot that labels each sample. Defaults to False.
2337        return_fig : bool, optional
2338            If True, return the open figure (caller owns lifecycle).
2339            Default is False.
2340        path : str or path-like, optional
2341            If set, save the figure to this path. When ``return_fig`` is False,
2342            the figure is closed after saving and ``plt.show()`` is not called.
2343
2344        Returns
2345        --------
2346        matplotlib.figure.Figure or None
2347            The figure if ``return_fig`` is True; otherwise None.
2348        """
2349        if type == "both":
2350            to_plot = ["raw", "corrected"]
2351        else:
2352            to_plot = [type]
2353
2354        fig, axs = plt.subplots(
2355            len(to_plot), 1, figsize=(10, 5 * len(to_plot)), sharex=True, squeeze=False
2356        )
2357
2358        for i, plot_type in enumerate(to_plot):
2359            ax = axs[i, 0]
2360            colors = iter(plt.cm.rainbow(np.linspace(0, 1, len(self))))
2361            for lcms_obj in self:
2362                c = next(colors)
2363                scan_df = lcms_obj.scan_df
2364                scan_df = scan_df[scan_df.ms_level == ms_level]
2365                if plot_type == "corrected":
2366                    # Check that scan_time_aligned is key in scan_df
2367                    if "scan_time_aligned" not in scan_df.columns:
2368                        raise ValueError(
2369                            f"scan_time_aligned not found in scan_df for {lcms_obj.sample_name}"
2370                        )
2371                    else:
2372                        ax.plot(
2373                            scan_df.scan_time_aligned,
2374                            scan_df.tic,
2375                            label=lcms_obj.sample_name,
2376                            c=c,
2377                            linewidth=0.3,
2378                        )
2379                elif plot_type == "raw":
2380                    ax.plot(
2381                        scan_df.scan_time,
2382                        scan_df.tic,
2383                        label=lcms_obj.sample_name,
2384                        c=c,
2385                        linewidth=0.3,
2386                    )
2387            ax.set_xlabel("Retention Time (min," + f" {plot_type})")
2388            ax.set_ylabel("TIC")
2389            if plot_legend:
2390                ax.legend()
2391        return _finalize_plot(fig, return_fig=return_fig, path=path)
2392
2393    def plot_alignments(self, plot_legend=False, return_fig=False, path=None):
2394        """Plots the alignment of the LCMS objects in the collection.
2395
2396        Parameters
2397        -----------
2398        plot_legend : bool, optional
2399            If True, plots a legend on the alignment plot that labels each sample. Defaults to False.
2400        return_fig : bool, optional
2401            If True, return the open figure (caller owns lifecycle).
2402            Default is False.
2403        path : str or path-like, optional
2404            If set, save the figure to this path. When ``return_fig`` is False,
2405            the figure is closed after saving and ``plt.show()`` is not called.
2406
2407        Returns
2408        --------
2409        matplotlib.figure.Figure or None
2410            The figure if ``return_fig`` is True; otherwise None.
2411        """
2412        fig, ax = plt.subplots(figsize=(10, 5))
2413        colors = iter(plt.cm.rainbow(np.linspace(0, 1, len(self))))
2414
2415        for lcms_obj in self:
2416            c = next(colors)
2417            scan_df = lcms_obj.scan_df
2418            if "scan_time_aligned" not in scan_df.columns:
2419                raise ValueError(
2420                    f"scan_time_aligned not found in scan_df for {lcms_obj.sample_name}"
2421                )
2422            scan_df["time_diff"] = scan_df.scan_time - scan_df.scan_time_aligned
2423            ax.plot(
2424                scan_df.scan_time_aligned,
2425                scan_df.time_diff,
2426                label=lcms_obj.sample_name,
2427                c=c,
2428                linewidth=0.3,
2429            )
2430
2431        ax.set_xlabel("Aligned Retention Time (min)")
2432        ax.set_ylabel("Time Difference (min)")
2433        if plot_legend:
2434            ax.legend()
2435        return _finalize_plot(fig, return_fig=return_fig, path=path)
2436
2437    def _drop_isotopologues(self):
2438        """Drops isotopologues from the mass features in combined_mass_features dataframe."""
2439        cmb_mf_df = self._combined_mass_features
2440
2441        # Keep monos or if no monos
2442        cmb_monos = cmb_mf_df[cmb_mf_df.monoisotopic_mf_id == cmb_mf_df.mf_id]
2443        cmb_nomonos = cmb_mf_df[cmb_mf_df.monoisotopic_mf_id.isnull()]
2444        # Keep deconvoluted parent or if no deconvoluted parent
2445        cmb_decon_parent = cmb_mf_df[
2446            cmb_mf_df.mass_spectrum_deconvoluted_parent
2447            | cmb_mf_df.monoisotopic_mf_id.isnull()
2448        ]
2449
2450        cmb_mf_df2 = pd.concat([cmb_monos, cmb_nomonos, cmb_decon_parent])
2451        cmb_mf_df2 = cmb_mf_df2[~cmb_mf_df2.index.duplicated(keep="first")]
2452        self.isotopes_dropped = True
2453        self._combined_mass_features = cmb_mf_df2
2454
2455    def load_raw_data(self, sample_idx: int, ms_level=1, time_range=None) -> None:
2456        """Load raw data for a specific sample index in the collection.
2457
2458        Parameters
2459        -----------
2460        sample_idx : int
2461            The index of the sample in the collection.
2462        ms_level : int, optional
2463            The MS level to load raw data for. Defaults to 1.
2464        time_range : tuple or list of tuples, optional
2465            Retention time range(s) to load. Can be a single tuple (min, max) or
2466            a list of tuples for multiple ranges. If None, loads all data. Defaults to None.
2467
2468        Raises
2469        -------
2470        IndexError
2471            If the sample index is out of range.
2472        ValueError
2473            If raw data for the specified MS level is already loaded for the sample index.
2474        ValueError
2475            If the spectra parser is not set for the LCMS object or if the parser type does not support loading raw data.
2476
2477        Returns
2478        --------
2479        None, but updates the LCMS object with the raw data for the specified MS level.
2480        """
2481        if sample_idx < 0 or sample_idx >= len(self.samples):
2482            raise IndexError("Sample index out of range.")
2483
2484        # Check that the sample does not already have raw data loaded
2485        if ms_level in self[sample_idx]._ms_unprocessed:
2486            raise ValueError(
2487                f"Raw data for MS{ms_level} already loaded for sample index {sample_idx}. Drop data first if you want to reload it."
2488            )
2489
2490        # Check the parser class. The spectra_parser property opens a reader.
2491        if getattr(self[sample_idx], "spectra_parser_class", None) is None:
2492            raise ValueError("Spectra parser is not set for this LCMS object.")
2493
2494        # Instantiate the parser and load the raw data using the correct method
2495        parser = self[sample_idx].spectra_parser
2496        parser_class_name = self[sample_idx].spectra_parser_class.__name__
2497        scan_df = self[sample_idx].scan_df
2498
2499        # Get raw data for the specified MS level using the appropriate method
2500        if parser_class_name == "ImportMassSpectraThermoMSFileReader":
2501            self[sample_idx]._ms_unprocessed[ms_level] = parser.get_ms_raw(
2502                spectra=f"ms{ms_level}", scan_df=scan_df, time_range=time_range
2503            )[f"ms{ms_level}"]
2504
2505        elif parser_class_name == "MZMLSpectraParser":
2506            data = parser.load()
2507            self[sample_idx]._ms_unprocessed[ms_level] = parser.get_ms_raw(
2508                spectra=f"ms{ms_level}",
2509                scan_df=scan_df,
2510                data=data,
2511                time_range=time_range,
2512            )[f"ms{ms_level}"]
2513
2514        elif parser_class_name == "ReadCoreMSHDFMassSpectra":
2515            raise ValueError(
2516                "ReadCoreMSHDFMassSpectra does not have a method to load raw data. Need to instantiate the original parser to access the raw data."
2517            )
2518
2519    def drop_raw_data(self, sample_idx: int, ms_level=1) -> None:
2520        """Drop raw data for a specific sample index in the collection.
2521
2522        Parameters
2523        -----------
2524        sample_idx : int
2525            The index of the sample in the collection.
2526        ms_level : int, optional
2527            The MS level to drop raw data for. Defaults to 1.
2528
2529        Raises
2530        -------
2531        IndexError
2532            If the sample index is out of range.
2533        ValueError
2534            If raw data for the specified MS level is not loaded for the sample index.
2535
2536        Returns
2537        --------
2538        None
2539        """
2540        if sample_idx < 0 or sample_idx >= len(self.samples):
2541            raise IndexError("Sample index out of range.")
2542
2543        # Check that the sample has raw data loaded
2544        if ms_level not in self[sample_idx]._ms_unprocessed:
2545            raise ValueError(
2546                f"No raw data for MS{ms_level} found for sample index {sample_idx}. Load data first if you want to drop it."
2547            )
2548
2549        # Drop the raw data
2550        del self[sample_idx]._ms_unprocessed[ms_level]
2551
2552    def update_raw_file_locations(self, new_raw_folder):
2553        """Update the raw file locations for all LCMS objects in the collection.
2554
2555        This method updates the path to the original raw data files (.raw, .mzML, etc.)
2556        that were used to create the processed HDF5 files stored in .corems folders.
2557
2558        Parameters
2559        -----------
2560        new_raw_folder : str or Path
2561            The new folder location containing the raw data files (.raw, .mzML, etc.).
2562            The method will look for raw files with the same base name as each sample.
2563
2564        Raises
2565        -------
2566        FileNotFoundError
2567            If the new raw folder does not exist.
2568        FileNotFoundError
2569            If a raw file for a sample is not found in the new folder.
2570
2571        Returns
2572        --------
2573        None, but updates the raw_file_location for each LCMS object in the collection.
2574
2575        Examples
2576        --------
2577        If raw files were moved from /old/path/ to /new/path/:
2578        >>> lcms_collection.update_raw_file_locations("/new/path/")
2579        """
2580        from pathlib import Path
2581
2582        if isinstance(new_raw_folder, str):
2583            new_raw_folder = Path(new_raw_folder)
2584
2585        if not new_raw_folder.exists():
2586            raise FileNotFoundError(f"Raw data folder does not exist: {new_raw_folder}")
2587
2588        # Common raw file extensions
2589        raw_extensions = [".raw", ".mzML", ".mzml"]
2590
2591        for sample_name in self.samples:
2592            lcms_obj = self._lcms[sample_name]
2593
2594            # Try to find the raw file with common extensions
2595            new_raw_file = None
2596            for ext in raw_extensions:
2597                candidate = new_raw_folder / f"{sample_name}{ext}"
2598                if candidate.exists():
2599                    new_raw_file = candidate
2600                    break
2601
2602            if new_raw_file is None:
2603                raise FileNotFoundError(
2604                    f"Raw file for sample '{sample_name}' not found in {new_raw_folder}. "
2605                    f"Tried extensions: {', '.join(raw_extensions)}"
2606                )
2607
2608            # Update the raw file location and set flag that raw files have been relocated
2609            lcms_obj.raw_file_location = new_raw_file
2610        self.raw_files_relocated = True
2611
2612    def collection_pivot_table(self, attribute="coll_mf_id", verbose=True):
2613        """Generate a pivot table of all regular and induced mass features in
2614        a collection. Default attribute presented is the mass feature ID, also
2615        prints a list of other available attributes.
2616
2617        Parameters
2618        -----------
2619        attribute : str
2620            The desired attribute to be presented in the pivot table. Defaults
2621            to mass feature ID
2622        verbose : boolean
2623            Print out all the possible values the fill the pivot table and list
2624            attributes that are not collected for induced mass features
2625
2626        Returns
2627        --------
2628        pd.DataFrame
2629            A DataFrame that displays one given attribute across all clusters
2630            and samples in a collection
2631
2632        """
2633
2634        mf_pivot = self.mass_features_dataframe.copy()
2635        mf_pivot.reset_index(inplace=True)
2636
2637        # Only include induced mass features if gap-filling has been performed
2638        if self.induced_mass_features_dataframe is not None:
2639            imf_pivot = self.induced_mass_features_dataframe.copy()
2640            imf_pivot.reset_index(inplace=True)
2641            # Cluster column extracted from mf_id in _prepare_lcms_mass_features_for_combination
2642            mf_pivot = pd.concat([mf_pivot, imf_pivot], axis=0)
2643            mf_pivot.reset_index(drop=True, inplace=True)
2644        else:
2645            imf_pivot = None
2646
2647        mf_pivot["cluster"] = mf_pivot["cluster"].astype(int)
2648
2649        if verbose:
2650            print(
2651                "Attributes available for pivot table:\n",
2652                [
2653                    x
2654                    for x in mf_pivot.columns
2655                    if x
2656                    not in ["cluster", "sample_name", "mf_id", "partition_idx", "idx"]
2657                ],
2658            )
2659            if imf_pivot is not None:
2660                print(
2661                    "\nAttributes that have no value for induced mass features:\n",
2662                    imf_pivot.columns[imf_pivot.isna().all()].tolist(),
2663                )
2664
2665        # Create pivot table and reindex to include all samples (even those with no features)
2666        pivot = mf_pivot.pivot(index="cluster", columns="sample_name", values=attribute)
2667
2668        # Reindex columns to include all samples in the collection
2669        all_samples = self.samples
2670        pivot = pivot.reindex(columns=all_samples)
2671
2672        return pivot
2673
2674    def cluster_representatives_table(self):
2675        """Generate a table of representative mass features from each consensus cluster.
2676
2677        This method returns a DataFrame containing all attributes for the
2678        representative mass feature from each consensus cluster. The representative
2679        is selected using the same logic as process_consensus_features().
2680
2681        Returns
2682        --------
2683        pd.DataFrame
2684            A DataFrame with one row per cluster containing all attributes for
2685            each cluster's representative mass feature. Includes:
2686            - cluster: cluster ID (as a column for easy joining)
2687            - polarity: ionization polarity from the collection
2688            - n_samples_detected: number of samples where the cluster was detected
2689            - All other mass feature attributes from the representative
2690
2691        Notes
2692        -----
2693        The representative metric used is determined by
2694        self.parameters.lcms_collection.consensus_representative_metric and
2695        is the same metric used by process_consensus_features() for consistency.
2696        Common options include 'intensity' (highest intensity) or
2697        'intensity_prefer_ms2' (highest intensity with preference for MS2 data).
2698        """
2699
2700        mf_df = self.mass_features_dataframe.copy()
2701        mf_df.reset_index(inplace=True)
2702
2703        # Include induced mass features if they exist (from gap-filling)
2704        if self.induced_mass_features_dataframe is not None:
2705            imf_df = self.induced_mass_features_dataframe.copy()
2706            imf_df.reset_index(inplace=True)
2707            # Cluster column extracted from mf_id in _prepare_lcms_mass_features_for_combination
2708            mf_df = pd.concat([mf_df, imf_df], axis=0)
2709            mf_df.reset_index(drop=True, inplace=True)
2710        mf_df["cluster"] = mf_df["cluster"].astype(int)
2711
2712        # Calculate number of samples per cluster
2713        cluster_sample_counts = (
2714            mf_df.groupby("cluster")["sample_id"].nunique().to_dict()
2715        )
2716
2717        # Use the same representative selection logic as process_consensus_features
2718        # This uses the configured representative_metric from parameters
2719        representatives = self.get_representative_mass_features_for_all_clusters()
2720
2721        # Get the coll_mf_ids of the representatives
2722        representative_ids = representatives["coll_mf_id"].tolist()
2723
2724        # Filter mf_df to only include representative features
2725        consensus_report = mf_df[mf_df.coll_mf_id.isin(representative_ids)].copy()
2726
2727        # Add polarity (get from first sample in collection)
2728        if len(self) > 0:
2729            polarity = self[0].polarity
2730        else:
2731            polarity = "unknown"
2732        consensus_report["polarity"] = polarity
2733
2734        # Add number of samples detected
2735        consensus_report["n_samples_detected"] = consensus_report["cluster"].map(
2736            cluster_sample_counts
2737        )
2738
2739        # Reorder columns to put cluster at the front
2740        cols = consensus_report.columns.tolist()
2741        if "cluster" in cols:
2742            cols.remove("cluster")
2743            cols = ["cluster"] + cols
2744            consensus_report = consensus_report[cols]
2745
2746        # Do not present a single preferred ion_type; possible_ion_types holds
2747        # all residual-tied candidates (same policy as cluster_summary).
2748        if "ion_type" in consensus_report.columns:
2749            if "possible_ion_types" in consensus_report.columns:
2750                from corems.mass_spectra.calc.feature_grouping import (
2751                    merge_possible_ion_types,
2752                )
2753
2754                for i in consensus_report.index:
2755                    filled = merge_possible_ion_types(
2756                        consensus_report.at[i, "ion_type"],
2757                        consensus_report.at[i, "possible_ion_types"],
2758                    )
2759                    if filled is not None:
2760                        consensus_report.at[i, "possible_ion_types"] = filled
2761            consensus_report = consensus_report.drop(columns=["ion_type"])
2762
2763        # Sort by cluster and return with cluster as a regular column
2764        return consensus_report.sort_values(by="cluster")
2765
2766    def to_mgf(
2767        self,
2768        out_file_path,
2769        cluster_ids=None,
2770        *,
2771        ms2_mode="best",
2772        overwrite=False,
2773    ):
2774        """Export consensus representative features to a SIRIUS-compatible MGF file.
2775
2776        FEATURE_ID is the consensus cluster id. ``ms2_mode='best'`` writes
2777        each representative's ``LCMSMassFeature.best_ms2``: the highest
2778        library-similarity scan after a spectral search, otherwise the
2779        non-chimeric scan closest to the apex. Representatives missing a
2780        usable MS2 spectrum are skipped, the same as a single-sample export.
2781        A representative with no precursor m/z aborts the export before any
2782        bytes are written. Requires representative mass features loaded on
2783        samples (typically via
2784        ``process_consensus_features(load_representatives=True, add_ms2=True)``).
2785
2786        Parameters
2787        ----------
2788        out_file_path : str or Path
2789            Output path. ``.mgf`` is appended when no suffix is given.
2790        cluster_ids : iterable, optional
2791            Consensus cluster ids to export. Default is all representatives.
2792        ms2_mode : {'best', 'all'}, optional
2793            Which associated MS2 spectra to write. ``'best'`` (default)
2794            writes ``LCMSMassFeature.best_ms2``. After a spectral search
2795            that is the scan with the highest library entropy similarity;
2796            with no search results it is the non-chimeric scan closest to
2797            the apex. ``'all'`` writes every MS2 scan that has peaks.
2798        overwrite : bool, optional
2799            Replace an existing file. Default is False.
2800
2801        Returns
2802        -------
2803        pathlib.Path
2804            Path of the written MGF file.
2805
2806        Warns
2807        -----
2808        UserWarning
2809            If one or more representatives are skipped because they have no
2810            usable MS2 spectrum.
2811
2812        Raises
2813        ------
2814        FileExistsError
2815            If the path exists and ``overwrite`` is False.
2816        ValueError
2817            If representatives are missing or not loaded, a requested
2818            cluster id is unknown, polarity is not ``'positive'`` or
2819            ``'negative'``, a representative has no precursor m/z, or no
2820            complete MS1/MS2 pairs remain. A missing precursor m/z stops
2821            the export before the file is written.
2822        """
2823        from corems.mass_spectra.output.mgf import (
2824            iter_collection_mgf_records,
2825            write_feature_records_to_mgf,
2826        )
2827
2828        records = iter_collection_mgf_records(self, cluster_ids=cluster_ids)
2829        if not records:
2830            raise ValueError(
2831                "No consensus representatives found. Run process_consensus_features() "
2832                "with load_representatives=True, add_ms2=True first."
2833            )
2834        sample_name = records[0][2]
2835        polarity = self._lcms[sample_name].polarity
2836        return write_feature_records_to_mgf(
2837            records,
2838            polarity=polarity,
2839            out_file_path=out_file_path,
2840            ms2_mode=ms2_mode,
2841            overwrite=overwrite,
2842        )
2843
2844    def feature_annotations_table(
2845        self, molecular_metadata=None, drop_unannotated=False, report_best_only=False
2846    ):
2847        """Generate a comprehensive annotation table for all loaded mass features across samples.
2848
2849        This method consolidates MS1 molecular formula assignments and MS2 spectral
2850        search results for all mass features across all samples in the collection.
2851        Only includes representative mass features (one per cluster per sample).
2852
2853        Parameters
2854        ----------
2855        molecular_metadata : dict, optional
2856            Dictionary of MolecularMetadata objects, keyed by metabref_mol_id.
2857            Required for including molecular metadata in MS2 annotations.
2858            Default is None.
2859        drop_unannotated : bool, optional
2860            If True, drops rows where all annotation columns (everything except
2861            cluster, MS2 Spectrum, and representative_sample) are NaN.
2862            Default is False.
2863        report_best_only : bool, optional
2864            If True, only includes the best MS2 annotation per mass feature based on confidence score.
2865            Default is False, which includes all MS2 annotations for each mass feature.
2866
2867        Returns
2868        -------
2869        pd.DataFrame
2870            Consolidated annotation report with columns including:
2871            - cluster: cluster ID
2872            - sample_name: sample name
2873            - sample_id: sample ID
2874            - Mass Feature ID: mass feature ID within the sample
2875            - Mass feature attributes (mz, scan_time, intensity, etc.)
2876            - MS1 annotations (if molecular_formula_search was run):
2877              ``Molecular Formula``, ``Ion Formula``, ``Calculated m/z``, etc.
2878            - MS2 annotations (if ms2_spectral_search was run):
2879              ``Library Molecular Formula``, ``Library Ion Formula``,
2880              ``Entropy Similarity``, ``name``, etc.
2881
2882        Notes
2883        -----
2884        This method uses the standard LCMSMetabolomicsExport.to_report() workflow
2885        for each sample, then consolidates all results and adds cluster information.
2886
2887        Only mass features that are loaded in each sample's mass_features dict
2888        are included (typically the representative features if load_representatives
2889        was used in process_consensus_features).
2890
2891        ``Molecular Formula`` / ``Ion Formula`` are filled only from MS1 molecular
2892        formula search. Spectral-library formulas appear under
2893        ``Library Molecular Formula`` / ``Library Ion Formula`` so MS2-only hits
2894        are not mistaken for MS1 formula assignments. When MS1 and MS2 match
2895        (same ion formula), the library formula columns are cleared on that
2896        row to avoid repeating the MS1 values.
2897
2898        Raises
2899        ------
2900        ValueError
2901            If no representative features have been loaded. Call process_consensus_features
2902            with load_representatives=True first.
2903        ValueError
2904            If no samples with loaded mass features are found in the collection.
2905        """
2906        from corems.mass_spectra.output.export import LCMSMetabolomicsExport
2907        import warnings
2908
2909        # Check if representative features have been loaded
2910        # Count samples with mass features loaded
2911        samples_with_features = sum(
2912            1
2913            for lcms_obj in self
2914            if hasattr(lcms_obj, "mass_features") and len(lcms_obj.mass_features) > 0
2915        )
2916
2917        if samples_with_features == 0:
2918            raise ValueError(
2919                "No representative mass features have been loaded into individual samples. "
2920                "Call process_consensus_features() with load_representatives=True before "
2921                "calling feature_annotations_table()."
2922            )
2923
2924        # Collect reports from all samples
2925        all_sample_reports = []
2926        has_any_ms2_annotations = False
2927
2928        for sample_id, lcms_obj in enumerate(self):
2929            # Skip samples with no loaded mass features
2930            if (
2931                not hasattr(lcms_obj, "mass_features")
2932                or len(lcms_obj.mass_features) == 0
2933            ):
2934                continue
2935
2936            sample_name = self.samples[sample_id]
2937
2938            # Create exporter and generate report using standard workflow
2939            # Suppress individual warnings - we'll warn at collection level if needed
2940            exporter = LCMSMetabolomicsExport("temp", lcms_obj)
2941            sample_report = exporter.to_report(
2942                molecular_metadata=molecular_metadata, suppress_warnings=True
2943            )
2944
2945            # Check if this sample has any MS2 annotations
2946            ms2_cols = [
2947                col
2948                for col in sample_report.columns
2949                if "Entropy Similarity" in col or "spectral_similarity" in col.lower()
2950            ]
2951            if ms2_cols and sample_report[ms2_cols].notna().any().any():
2952                has_any_ms2_annotations = True
2953
2954            # Add sample information
2955            sample_report["representative_sample"] = sample_name
2956            sample_report["sample_id"] = sample_id
2957
2958            # Get cluster information from the mass_features_dataframe
2959            # Build coll_mf_id for each row to look up cluster
2960            sample_report["coll_mf_id"] = (
2961                sample_report["sample_id"].astype(str)
2962                + "_"
2963                + sample_report["Mass Feature ID"].astype(str)
2964            )
2965
2966            # Get cluster from mass_features_dataframe
2967            if (
2968                self.mass_features_dataframe is not None
2969                and "cluster" in self.mass_features_dataframe.columns
2970            ):
2971                mf_df = self.mass_features_dataframe.reset_index()
2972                cluster_lookup = mf_df.set_index("coll_mf_id")["cluster"].to_dict()
2973                sample_report["cluster"] = sample_report["coll_mf_id"].map(
2974                    cluster_lookup
2975                )
2976            else:
2977                sample_report["cluster"] = None
2978
2979            # Drop temporary coll_mf_id column
2980            sample_report = sample_report.drop(columns=["coll_mf_id"])
2981
2982            all_sample_reports.append(sample_report)
2983
2984        # Combine all sample reports
2985        if len(all_sample_reports) == 0:
2986            raise ValueError("No samples with loaded mass features found in collection")
2987
2988        collection_report = pd.concat(all_sample_reports, ignore_index=True)
2989
2990        # Feature-group constrain: drop consensus isotopes; filter library ion types
2991        from corems.mass_spectra.calc.feature_grouping import (
2992            allowed_ion_types_for_row,
2993            constrain_annotation_active,
2994            ion_type_allowed,
2995            should_skip_isotope_for_annotation,
2996        )
2997
2998        labels = getattr(self, "feature_group_dataframe", None)
2999        constrain = constrain_annotation_active(
3000            labels,
3001            bool(
3002                getattr(
3003                    self.parameters.lcms_collection,
3004                    "feature_group_constrain_annotation",
3005                    True,
3006                )
3007            ),
3008        )
3009        if constrain and labels is not None and "cluster" in collection_report.columns:
3010            keep_rows = []
3011            for _, row in collection_report.iterrows():
3012                cid = row.get("cluster")
3013                if cid is None or (isinstance(cid, float) and pd.isna(cid)):
3014                    keep_rows.append(True)
3015                    continue
3016                try:
3017                    cid_i = int(cid)
3018                except (TypeError, ValueError):
3019                    keep_rows.append(True)
3020                    continue
3021                if cid_i not in labels.index:
3022                    keep_rows.append(True)
3023                    continue
3024                lab = labels.loc[cid_i]
3025                if should_skip_isotope_for_annotation(lab.get("ion_role"), True):
3026                    keep_rows.append(False)
3027                    continue
3028                allowed = allowed_ion_types_for_row(
3029                    lab.get("ion_type"), lab.get("possible_ion_types")
3030                )
3031                if allowed is None:
3032                    keep_rows.append(True)
3033                    continue
3034                drop = False
3035                for col in ("Ion Type", "Library Ion Type"):
3036                    it = row.get(col)
3037                    if it is None or (isinstance(it, float) and pd.isna(it)):
3038                        continue
3039                    if not ion_type_allowed(it, allowed):
3040                        drop = True
3041                        break
3042                keep_rows.append(not drop)
3043            collection_report = collection_report.loc[keep_rows].reset_index(drop=True)
3044
3045        # Join consensus feature-group labels (drop ion_type; possibles only)
3046        from corems.mass_spectra.calc.feature_grouping import (
3047            GROUP_COLUMNS,
3048            merge_possible_ion_types,
3049        )
3050
3051        if (
3052            labels is not None
3053            and len(labels)
3054            and "cluster" in collection_report.columns
3055        ):
3056            gcols = [c for c in GROUP_COLUMNS if c in labels.columns]
3057            grouping = labels[gcols].copy().reset_index()
3058            first = grouping.columns[0]
3059            if first != "cluster":
3060                grouping = grouping.rename(columns={first: "cluster"})
3061            drop_existing = [c for c in gcols if c in collection_report.columns]
3062            if drop_existing:
3063                collection_report = collection_report.drop(columns=drop_existing)
3064            collection_report = collection_report.merge(
3065                grouping, on="cluster", how="left"
3066            )
3067            if (
3068                "ion_type" in collection_report.columns
3069                and "possible_ion_types" in collection_report.columns
3070            ):
3071                for i in collection_report.index:
3072                    filled = merge_possible_ion_types(
3073                        collection_report.at[i, "ion_type"],
3074                        collection_report.at[i, "possible_ion_types"],
3075                    )
3076                    if filled is not None:
3077                        collection_report.at[i, "possible_ion_types"] = filled
3078            if "ion_type" in collection_report.columns:
3079                collection_report = collection_report.drop(columns=["ion_type"])
3080
3081        # Warn only if NO samples in the collection have MS2 annotations
3082        if not has_any_ms2_annotations:
3083            warnings.warn(
3084                "No MS2 annotations found across any samples in collection, were MS2 spectra added and searched against a database?",
3085                UserWarning,
3086            )
3087
3088        # Reorder columns to match specified order
3089        group_front = [
3090            c
3091            for c in (
3092                "feature_group_id",
3093                "ion_role",
3094                "possible_ion_types",
3095                "isotope_state",
3096                "mono_cluster_id",
3097            )
3098            if c in collection_report.columns
3099        ]
3100        desired_cols = [
3101            "cluster",
3102            *group_front,
3103            "Isotopologue Type",
3104            "Is Largest Ion after Deconvolution",
3105            "MS2 Spectrum",
3106            "Calculated m/z",
3107            "m/z Error (ppm)",
3108            "m/z Error Score",
3109            "Isotopologue Similarity",
3110            "Confidence Score",
3111            "Ion Formula",
3112            "Molecular Formula",
3113            "Library Ion Formula",
3114            "Library Ion Type",
3115            "Library Molecular Formula",
3116            "inchikey",
3117            "name",
3118            "ref_ms_id",
3119            "Entropy Similarity",
3120            "Library mzs in Query (fraction)",
3121            "Spectra with Annotation (n)",
3122            "representative_sample",
3123        ]
3124
3125        leading = [c for c in desired_cols if c in collection_report.columns]
3126        rest = [c for c in collection_report.columns if c not in leading]
3127        collection_report = collection_report[leading + rest]
3128        nonempty = [
3129            c
3130            for c in collection_report.columns
3131            if c == "cluster" or collection_report[c].notna().any()
3132        ]
3133        collection_report = collection_report[nonempty]
3134
3135        # Optionally drop rows without any annotations
3136        if drop_unannotated:
3137            # Columns to exclude from the "all NA" check
3138            exclude_cols = ["cluster", "MS2 Spectrum", "representative_sample"]
3139            # Get annotation columns (everything except the excluded ones)
3140            annot_cols = [
3141                col for col in collection_report.columns if col not in exclude_cols
3142            ]
3143            # Keep rows where at least one annotation column is not NA
3144            if len(annot_cols) > 0:
3145                collection_report = collection_report[
3146                    collection_report[annot_cols].notna().any(axis=1)
3147                ]
3148
3149        # Sort by cluster, then by annotation quality (only columns that remain)
3150        sort_cols = ["cluster"]
3151        sort_asc = [True]
3152        if "Entropy Similarity" in collection_report.columns:
3153            sort_cols.append("Entropy Similarity")
3154            sort_asc.append(False)
3155        if "Confidence Score" in collection_report.columns:
3156            sort_cols.append("Confidence Score")
3157            sort_asc.append(False)
3158        collection_report = collection_report.sort_values(
3159            by=sort_cols, ascending=sort_asc
3160        )
3161
3162        if report_best_only:
3163            # Keep only the best annotation per cluster based on the first annotation column available
3164            if "Entropy Similarity" in collection_report.columns:
3165                best_annot_col = "Entropy Similarity"
3166            elif "Confidence Score" in collection_report.columns:
3167                best_annot_col = "Confidence Score"
3168            else:
3169                best_annot_col = None
3170
3171            if best_annot_col is not None:
3172                collection_report = collection_report.sort_values(
3173                    by=["cluster", best_annot_col], ascending=[True, False]
3174                )
3175                collection_report = collection_report.drop_duplicates(
3176                    subset=["cluster"], keep="first"
3177                )
3178
3179        return collection_report
3180
3181    @property
3182    def parameters(self):
3183        """
3184        LCMSCollectionParameters : The parameters used for the LCMS collection.
3185        """
3186        return self._parameters
3187
3188    @parameters.setter
3189    def parameters(self, paramsinstance):
3190        """
3191        Sets the parameters used for the LCMS analysis collection.
3192
3193        Parameters
3194        -----------
3195        paramsinstance : LCMSCollectionParameters
3196            The parameters used for the LC-MS analysis.
3197        """
3198        self._parameters = paramsinstance
3199
3200    @property
3201    def mass_features_dataframe(self):
3202        self._check_mass_features_df()
3203        return self._combined_mass_features
3204
3205    @mass_features_dataframe.setter
3206    def mass_features_dataframe(self, df):
3207        # Check that the dataframe has the expected columns
3208        expected_cols = ["sample_name", "sample_id", "mz", "scan_time"]
3209        if not all([col in df.columns for col in expected_cols]):
3210            raise ValueError(
3211                f"Expected columns not found in dataframe: {expected_cols}"
3212            )
3213
3214        # Check that coll_mf_id is the index and it is unique
3215        if df.index.name != "coll_mf_id":
3216            raise ValueError("coll_mf_id must be the index of the dataframe")
3217        if not df.index.is_unique:
3218            raise ValueError("coll_mf_id must be unique")
3219        self._combined_mass_features = df
3220
3221    @property
3222    def induced_mass_features_dataframe(self):
3223        self._check_mass_features_df(induced_features=True)
3224        if (
3225            self._combined_induced_mass_features is not None
3226            and len(self._combined_induced_mass_features) > 0
3227        ):
3228            # The cluster column is extracted from mf_id in _prepare_lcms_mass_features_for_combination
3229            # mf_id format for induced features: c{cluster}_{index}_i
3230            pass
3231        return self._combined_induced_mass_features
3232
3233    @induced_mass_features_dataframe.setter
3234    def induced_mass_features_dataframe(self, df):
3235        # Check that the dataframe has the expected columns
3236        expected_cols = ["sample_name", "sample_id", "mz", "scan_time"]
3237        if not all([col in df.columns for col in expected_cols]):
3238            raise ValueError(
3239                f"Expected columns not found in dataframe: {expected_cols}"
3240            )
3241
3242        # Check that coll_mf_id is the index and it is unique
3243        if df.index.name != "coll_mf_id":
3244            raise ValueError("coll_mf_id must be the index of the dataframe")
3245        if not df.index.is_unique:
3246            raise ValueError("coll_mf_id must be unique")
3247        self._combined_induced_mass_features = df
3248
3249    @property
3250    def cluster_summary_dataframe(self):
3251        return self.summarize_clusters()
3252
3253    @property
3254    def samples(self):
3255        manifest_df = self.manifest_dataframe
3256        # order by batch, then by order
3257        manifest_df = manifest_df.sort_values(by=["batch", "order"])
3258        return manifest_df.index.tolist()
3259
3260    @property
3261    def manifest(self):
3262        return self._manifest_dict
3263
3264    @property
3265    def manifest_dataframe(self):
3266        return pd.DataFrame(self._manifest_dict).T
3267
3268    @property
3269    def raw_files(self):
3270        """Returns a list of raw files in the collection."""
3271        return [x.raw_file_location for x in self]
3272
3273    @property
3274    def rt_alignments(self):
3275        """Returns a dictionary of retention time alignments for the collection."""
3276        if self.rt_aligned:
3277            _rt_alignments = {}
3278            # Construct a dictionary of aligned retention times (stored on each LCMS object within the collection, not the collection itself)
3279            for i, lcms_obj in enumerate(self):
3280                aligned_times = [
3281                    x
3282                    for k, x in sorted(lcms_obj._scan_info["scan_time_aligned"].items())
3283                ]
3284                _rt_alignments[i] = aligned_times
3285            return _rt_alignments
3286        else:
3287            return None
3288
3289    @property
3290    def cluster_feature_dictionary(self):
3291        """Generates a dictionary with clusters for keys and mass feature IDs as entries"""
3292        df = self.mass_features_dataframe
3293        cluster_dict = {
3294            cluster: list(ids) for cluster, ids in df.groupby("cluster").groups.items()
3295        }
3296        return cluster_dict
3297
3298    def get_eics_for_cluster(self, cluster_id):
3299        """
3300        Retrieve all EICs for mass features in a specific cluster across all samples.
3301
3302        Returns a dictionary mapping sample names to EIC_Data objects for the given cluster.
3303        Useful for visualizing and comparing chromatographic peaks across samples.
3304
3305        Parameters
3306        ----------
3307        cluster_id : int
3308            The cluster ID to retrieve EICs for
3309
3310        Returns
3311        -------
3312        dict
3313            Dictionary with structure: {sample_name: EIC_Data object}
3314            Only includes samples where the EIC was loaded.
3315
3316        Examples
3317        --------
3318        >>> # Load EICs first
3319        >>> collection.process_consensus_features(gather_eics=True, ...)
3320        >>>
3321        >>> # Get all EICs for cluster 5
3322        >>> eics = collection.get_eics_for_cluster(5)
3323        >>> for sample_name, eic_data in eics.items():
3324        ...     print(f"{sample_name}: {len(eic_data.scans)} scans")
3325
3326        Notes
3327        -----
3328        Requires that EICs have been loaded using gather_eics=True in
3329        process_consensus_features() or manually loaded via LoadEICsOperation.
3330        """
3331        eics_by_sample = {}
3332
3333        # Iterate through all samples
3334        for sample_id, sample in enumerate(self):
3335            sample_name = self.samples[sample_id]
3336
3337            # Check if sample has EICs loaded
3338            if not hasattr(sample, "eics") or not sample.eics:
3339                continue
3340
3341            # Find mass features in this cluster for this sample
3342            # Check both regular and induced mass features
3343            for mf in list(sample.mass_features.values()) + list(
3344                sample.induced_mass_features.values()
3345            ):
3346                if hasattr(mf, "cluster_index") and mf.cluster_index == cluster_id:
3347                    # Get the EIC for this mass feature's m/z
3348                    if mf.mz in sample.eics:
3349                        eics_by_sample[sample_name] = sample.eics[mf.mz]
3350                        break  # Found the EIC for this sample, move to next sample
3351
3352        return eics_by_sample