corems.mass_spectra.calc.lc_calc

   1import numpy as np
   2import pandas as pd
   3import warnings, scipy, multiprocessing
   4
   5if not hasattr(np, 'trapezoid'):  # numpy < 2.0
   6    np.trapezoid = np.trapz
   7from ripser import ripser
   8from scipy import sparse
   9from scipy.spatial import KDTree
  10from sklearn.svm import SVR
  11from sklearn.cluster import AgglomerativeClustering
  12import matplotlib.pyplot as plt
  13from tqdm import tqdm
  14
  15from corems.chroma_peak.factory.chroma_peak_classes import LCMSMassFeature
  16from corems.mass_spectra.calc import SignalProcessing as sp
  17from corems.mass_spectra.factory.chromat_data import EIC_Data
  18from corems.mass_spectrum.input.numpyArray import ms_from_array_profile
  19
  20warnings.filterwarnings("ignore", category=RuntimeWarning)
  21
  22def find_closest(A, target):
  23    """Find the index of closest value in A to each value in target.
  24
  25    Parameters
  26    ----------
  27    A : :obj:`~numpy.array`
  28        The array to search (blueprint). A must be sorted.
  29    target : :obj:`~numpy.array`
  30        The array of values to search for. target must be sorted.
  31
  32    Returns
  33    -------
  34    :obj:`~numpy.array`
  35        The indices of the closest values in A to each value in target.
  36    """
  37    idx = A.searchsorted(target)
  38    idx = np.clip(idx, 1, len(A) - 1)
  39    left = A[idx - 1]
  40    right = A[idx]
  41    idx -= target - left < right - target
  42    return idx
  43
  44
  45class LCCalculations:
  46    """Methods for performing LC calculations on mass spectra data.
  47
  48    Notes
  49    -----
  50    This class is intended to be used as a mixin for the LCMSBase class.
  51
  52    Methods
  53    -------
  54    * get_max_eic(eic_data).
  55        Returns the maximum EIC value from the given EIC data. A static method.
  56    * smooth_tic(tic).
  57        Smooths the TIC data using the specified smoothing method and settings.
  58    * eic_centroid_detector(rt, eic, max_eic).
  59        Performs EIC centroid detection on the given EIC data.
  60    * find_nearest_scan(rt).
  61        Finds the nearest scan to the given retention time.
  62    * get_average_mass_spectrum(scan_list, apex_scan, spectrum_mode="profile", ms_level=1, auto_process=True, use_parser=False, perform_checks=True, polarity=None).
  63        Returns an averaged mass spectrum object.
  64    * find_mass_features(ms_level=1).
  65        Find regions of interest for a given MS level (default is MS1).
  66    * integrate_mass_features(drop_if_fail=False, ms_level=1).
  67        Integrate mass features of interest and extracts EICs.
  68    * find_c13_mass_features().
  69        Evaluate mass features and mark likely C13 isotopes.
  70    * deconvolute_ms1_mass_features().
  71        Deconvolute mass features' ms1 mass spectra.
  72    """
  73
  74    @staticmethod
  75    def get_max_eic(eic_data: dict):
  76        """Returns the maximum EIC value from the given EIC data.
  77
  78        Notes
  79        -----
  80        This is a static method.
  81
  82        Parameters
  83        ----------
  84        eic_data : dict
  85            A dictionary containing EIC data.
  86
  87        Returns
  88        -------
  89        float
  90            The maximum EIC value.
  91        """
  92        max_eic = 0
  93        for eic_data in eic_data.values():
  94            ind_max_eic = max(eic_data.get("EIC"))
  95            max_eic = ind_max_eic if ind_max_eic > max_eic else max_eic
  96
  97        return max_eic
  98
  99    def smooth_tic(self, tic):
 100        """Smooths the TIC or EIC data using the specified smoothing method and settings.
 101
 102        Parameters
 103        ----------
 104        tic : numpy.ndarray
 105            The TIC (or EIC) data to be smoothed.
 106
 107        Returns
 108        -------
 109        numpy.ndarray
 110            The smoothed TIC data.
 111        """
 112        implemented_smooth_method = self.parameters.lc_ms.implemented_smooth_method
 113
 114        pol_order = self.parameters.lc_ms.savgol_pol_order
 115
 116        window_len = self.parameters.lc_ms.smooth_window
 117
 118        window = self.parameters.lc_ms.smooth_method
 119
 120        return sp.smooth_signal(
 121            tic, window_len, window, pol_order, implemented_smooth_method
 122        )
 123
 124    def eic_centroid_detector(self, rt, eic, max_eic, apex_indexes=[]):
 125        """Performs EIC centroid detection on the given EIC data.
 126
 127        Parameters
 128        ----------
 129        rt : numpy.ndarray
 130            The retention time data.
 131        eic : numpy.ndarray
 132            The EIC data.
 133        max_eic : float
 134            The maximum EIC value.
 135        apex_indexes : list, optional
 136            The apexes of the EIC peaks. Defaults to [], which means that the apexes will be calculated by the function.
 137
 138        Returns
 139        -------
 140        numpy.ndarray
 141            The indexes of left, apex, and right limits as a generator.
 142        """
 143
 144        max_prominence = self.parameters.lc_ms.peak_max_prominence_percent
 145
 146        max_height = self.parameters.lc_ms.peak_height_max_percent
 147
 148        signal_threshold = self.parameters.lc_ms.eic_signal_threshold
 149
 150        min_peak_datapoints = self.parameters.lc_ms.min_peak_datapoints
 151
 152        peak_derivative_threshold = self.parameters.lc_ms.peak_derivative_threshold
 153
 154        include_indexes = sp.peak_picking_first_derivative(
 155            domain=rt,
 156            signal=eic,
 157            max_height=max_height,
 158            max_prominence=max_prominence,
 159            max_signal=max_eic,
 160            min_peak_datapoints=min_peak_datapoints,
 161            peak_derivative_threshold=peak_derivative_threshold,
 162            signal_threshold=signal_threshold,
 163            correct_baseline=False,
 164            plot_res=False,
 165            apex_indexes=apex_indexes,
 166        )
 167        #include_indexes is a generator of tuples (left_index, apex_index, right_index)
 168        include_indexes = list(include_indexes)
 169        # Add check to make sure that there are at least 1/2 of min_peak_datapoints on either side of the apex
 170        indicies = [x for x in include_indexes]
 171        for idx in indicies:
 172            if (idx[1] - idx[0] < min_peak_datapoints / 2) or (
 173                idx[2] - idx[1] < min_peak_datapoints / 2
 174            ):
 175                include_indexes.remove(idx)
 176        return include_indexes
 177
 178    def find_nearest_scan(self, rt):
 179        """Finds the nearest scan to the given retention time.
 180
 181        Parameters
 182        ----------
 183        rt : float
 184            The retention time (in minutes) to find the nearest scan for.
 185
 186        Returns
 187        -------
 188        int
 189            The scan number of the nearest scan.
 190        """
 191        array_rt = np.array(self.retention_time)
 192
 193        scan_index = (np.abs(array_rt - rt)).argmin()
 194
 195        real_scan = self.scans_number[scan_index]
 196
 197        return real_scan
 198
 199    def add_peak_metrics(self, remove_by_metrics=True, induced_features=False):
 200        """Add peak metrics to the mass features.
 201
 202        This function calculates the peak metrics for each mass feature and adds them to the mass feature objects.
 203
 204        Parameters
 205        ----------
 206        remove_by_metrics : bool, optional
 207            If True, remove mass features based on their peak metrics such as S/N, Gaussian similarity,
 208            dispersity index, and noise score. Default is True, which checks the setting in the processing parameters.
 209            If False, peak metrics are calculated but no mass features are removed, regardless of the setting in the processing parameters.
 210        induced_features : bool, optional
 211            Whether the mass features to be integrated were induced. Default is False.
 212        """
 213        # Check that at least some mass features have eic data
 214        if induced_features:
 215            mf_dict_values = self.induced_mass_features.values()
 216        else:
 217            mf_dict_values = self.mass_features.values()
 218
 219        if not any([mf._eic_data is not None for mf in mf_dict_values]):
 220            raise ValueError(
 221                "No mass features have EIC data. Run integrate_mass_features first."
 222            )
 223
 224        for mass_feature in mf_dict_values:
 225            # Check if the mass feature has been integrated
 226            if mass_feature._eic_data is not None and mass_feature.area is not None:
 227                # Calculate peak metrics
 228                mass_feature.calc_half_height_width()
 229                mass_feature.calc_tailing_factor()
 230                mass_feature.calc_dispersity_index()
 231                mass_feature.calc_gaussian_similarity()
 232                mass_feature.calc_noise_score()
 233        
 234        # Remove mass features by peak metrics if designated in parameters
 235        if self.parameters.lc_ms.remove_mass_features_by_peak_metrics and remove_by_metrics:
 236            self._remove_mass_features_by_peak_metrics(induced_features=induced_features)
 237
 238    def get_average_mass_spectrum(
 239        self,
 240        scan_list,
 241        apex_scan,
 242        spectrum_mode="profile",
 243        ms_level=1,
 244        auto_process=True,
 245        use_parser=False,
 246        perform_checks=True,
 247        polarity=None,
 248        ms_params=None,
 249    ):
 250        """Returns an averaged mass spectrum object
 251
 252        Parameters
 253        ----------
 254        scan_list : list
 255            List of scan numbers to average.
 256        apex_scan : int
 257            Number of the apex scan
 258        spectrum_mode : str, optional
 259            The spectrum mode to use. Defaults to "profile". Not that only "profile" mode is supported for averaging.
 260        ms_level : int, optional
 261            The MS level to use. Defaults to 1.
 262        auto_process : bool, optional
 263            If True, the averaged mass spectrum will be auto-processed. Defaults to True.
 264        use_parser : bool, optional
 265            If True, the mass spectra will be obtained from the parser. Defaults to False.
 266        perform_checks : bool, optional
 267            If True, the function will check if the data are within the ms_unprocessed dictionary and are the correct mode. Defaults to True. Only set to False if you are sure the data are profile, and (if not using the parser) are in the ms_unprocessed dictionary!  ms_unprocessed dictionary also must be indexed on scan
 268        polarity : int, optional
 269            The polarity of the mass spectra (1 or -1). If not set, the polarity will be determined from the dataset. Defaults to None. (fastest if set to -1 or 1)
 270        ms_params : MSParameters, optional
 271            The mass spectrum parameters to use. If not set (None), the globally set parameters will be used. Defaults to None.
 272
 273        Returns
 274        -------
 275        MassSpectrumProfile
 276            The averaged mass spectrum object.
 277
 278        Raises
 279        ------
 280        ValueError
 281            If the spectrum mode is not "profile".
 282            If the MS level is not found in the unprocessed mass spectra dictionary.
 283            If not all scan numbers are found in the unprocessed mass spectra dictionary.
 284        """
 285        if perform_checks:
 286            if spectrum_mode != "profile":
 287                raise ValueError("Averaging only supported for profile mode")
 288
 289        if polarity is None:
 290            # set polarity to -1 if negative mode, 1 if positive mode (for mass spectrum creation)
 291            if self.polarity == "negative":
 292                polarity = -1
 293            elif self.polarity == "positive":
 294                polarity = 1
 295            else:
 296                raise ValueError(
 297                    "Polarity not set for dataset, must be a set containing either 'positive' or 'negative'"
 298                )
 299
 300        # if not using_parser, check that scan numbers are in _ms_unprocessed
 301        if not use_parser:
 302            if perform_checks:
 303                # Set index to scan for faster lookup
 304                ms_df = (
 305                    self._ms_unprocessed[ms_level]
 306                    .copy()
 307                    .set_index("scan", drop=False)
 308                    .sort_index()
 309                )
 310                my_ms_df = ms_df.loc[scan_list]
 311                # Check that all scan numbers are in the ms_df
 312                if not all(np.isin(scan_list, ms_df.index)):
 313                    raise ValueError(
 314                        "Not all scan numbers found in the unprocessed mass spectra dictionary"
 315                    )
 316            else:
 317                my_ms_df = (
 318                    pd.DataFrame({"scan": scan_list})
 319                    .set_index("scan")
 320                    .join(self._ms_unprocessed[ms_level], how="left")
 321                )
 322
 323        if use_parser:
 324            ms_list = [
 325                self.spectra_parser.get_mass_spectrum_from_scan(
 326                    x, spectrum_mode=spectrum_mode, auto_process=False
 327                )
 328                for x in scan_list
 329            ]
 330            ms_mz = [x._mz_exp for x in ms_list]
 331            ms_int = [x._abundance for x in ms_list]
 332            my_ms_df = []
 333            for i in np.arange(len(ms_mz)):
 334                my_ms_df.append(
 335                    pd.DataFrame(
 336                        {"mz": ms_mz[i], "intensity": ms_int[i], "scan": scan_list[i]}
 337                    )
 338                )
 339            my_ms_df = pd.concat(my_ms_df)
 340
 341        if not self.check_if_grid(my_ms_df):
 342            my_ms_df = self.grid_data(my_ms_df)
 343
 344        my_ms_ave = my_ms_df.groupby("mz")["intensity"].sum().reset_index()
 345
 346        ms = ms_from_array_profile(
 347            my_ms_ave.mz,
 348            my_ms_ave.intensity,
 349            self.file_location,
 350            polarity=polarity,
 351            auto_process=False,
 352        )
 353
 354        # Set the mass spectrum parameters, auto-process if auto_process is True, and add to the dataset
 355        if ms is not None:
 356            if ms_params is not None:
 357                ms.parameters = ms_params
 358            ms.scan_number = apex_scan
 359            if auto_process:
 360                ms.process_mass_spec()
 361        return ms
 362
 363    def find_mass_features(self, ms_level=1, grid=True, assign_ms2_scans=False, ms2_scan_filter=None, 
 364                          targeted_search=False, target_search_dict=None, accumulate_features=False):
 365        """Find mass features within an LCMSBase object
 366
 367        Note that this is a wrapper function that calls the find_mass_features_ph function, but can be extended to support other peak picking methods in the future.
 368
 369        Parameters
 370        ----------
 371        ms_level : int, optional
 372            The MS level to use for peak picking Default is 1.
 373        grid : bool, optional
 374            If True, will regrid the data before running the persistent homology calculations (after checking if the data is gridded),
 375            used for persistent homology peak picking for profile data only. Default is True.
 376        assign_ms2_scans : bool, optional
 377            If True, assign MS2 scan numbers to mass features after peak picking.
 378            This populates the ms2_scan_numbers attribute on each mass feature, which enables
 379            choosing representative features based on MS2 availability. Default is False.
 380        ms2_scan_filter : str or None, optional
 381            Filter string for MS2 scans when assign_ms2_scans is True (e.g., 'hcd').
 382            If None, all MS2 scans are considered. Default is None.
 383        targeted_search : bool, optional
 384            If True, perform targeted mass feature search using the target_search_dict.
 385            This mode filters data to only m/z and RT windows of interest and bypasses
 386            intensity and persistence thresholds. Default is False.
 387        target_search_dict : dict or None, optional
 388            Dictionary containing target search parameters. Required if targeted_search is True.
 389            Must contain:
 390                - 'target_mz_list': list of target m/z values
 391                - 'target_rt_list': list of target retention times (in minutes)
 392                - 'mz_tolerance_ppm': m/z tolerance in ppm
 393                - 'rt_tolerance': retention time tolerance (in minutes)
 394            Optionally can contain:
 395                - 'type': type label for mass features (e.g., "internal standard")
 396                  If not provided, defaults to "targeted"
 397            Default is None.
 398        accumulate_features : bool, optional
 399            If True, new mass features will be added to existing features rather than replacing them.
 400            This allows multiple sequential calls to find_mass_features to build up a combined set.
 401            Default is False (replace existing features for backwards compatibility).
 402
 403        Raises
 404        ------
 405        ValueError
 406            If no MS level data is found on the object.
 407            If persistent homology peak picking is attempted on non-profile mode data.
 408            If data is not gridded and grid is False.
 409            If peak picking method is not implemented.
 410            If targeted_search is True but target_search_dict is None or invalid.
 411
 412        Returns
 413        -------
 414        None, but assigns the mass_features and eics attributes to the object.
 415
 416        """
 417        # Validate targeted search parameters
 418        if targeted_search:
 419            if target_search_dict is None:
 420                raise ValueError("target_search_dict must be provided when targeted_search is True")
 421            required_keys = ['target_mz_list', 'target_rt_list', 'mz_tolerance_ppm', 'rt_tolerance']
 422            for key in required_keys:
 423                if key not in target_search_dict:
 424                    raise ValueError(f"target_search_dict must contain '{key}'")
 425            if len(target_search_dict['target_mz_list']) != len(target_search_dict['target_rt_list']):
 426                raise ValueError("target_mz_list and target_rt_list must have the same length")
 427        
 428        pp_method = self.parameters.lc_ms.peak_picking_method
 429
 430        if pp_method == "persistent homology":
 431            msx_scan_df = self.scan_df[self.scan_df["ms_level"] == ms_level]
 432            if all(msx_scan_df["ms_format"] == "profile"):
 433                # Determine mass feature type
 434                if targeted_search:
 435                    mf_type = target_search_dict.get('type', 'targeted')
 436                else:
 437                    mf_type = 'untargeted'
 438                self.find_mass_features_ph(ms_level=ms_level, grid=grid, 
 439                                          targeted_search=targeted_search, 
 440                                          target_search_dict=target_search_dict,
 441                                          mf_type=mf_type,
 442                                          accumulate_features=accumulate_features)
 443            else:
 444                raise ValueError(
 445                    "MS{} scans are not profile mode, which is required for persistent homology peak picking.".format(
 446                        ms_level
 447                    )
 448                )
 449        elif pp_method == "centroided_persistent_homology":
 450            msx_scan_df = self.scan_df[self.scan_df["ms_level"] == ms_level]
 451            if all(msx_scan_df["ms_format"] == "centroid"):
 452                # Determine mass feature type
 453                if targeted_search:
 454                    mf_type = target_search_dict.get('type', 'targeted')
 455                else:
 456                    mf_type = 'untargeted'
 457                self.find_mass_features_ph_centroid(ms_level=ms_level, 
 458                                                    targeted_search=targeted_search, 
 459                                                    target_search_dict=target_search_dict,
 460                                                    mf_type=mf_type,
 461                                                    accumulate_features=accumulate_features)
 462            else:
 463                raise ValueError(
 464                    "MS{} scans are not centroid mode, which is required for persistent homology centroided peak picking.".format(
 465                        ms_level
 466                    )
 467                )
 468        else:
 469            raise ValueError("Peak picking method not implemented")
 470        
 471        # Cluster mass features to remove redundant features
 472        self.cluster_mass_features(drop_children=True)
 473        
 474        # Optionally assign MS2 scan numbers to mass features during peak picking
 475        # This helps with choosing representative features that have MS2 data
 476        if assign_ms2_scans:
 477            try:
 478                self._find_ms2_scans_for_mass_features(
 479                    mf_ids=None,  # Process all mass features
 480                    scan_filter=ms2_scan_filter
 481                )
 482            except ValueError:
 483                # No MS2 scans found - this is okay, just skip
 484                pass
 485        
 486        # Remove noisey mass features if designated in parameters
 487        if self.parameters.lc_ms.remove_redundant_mass_features and not targeted_search:
 488            self._remove_redundant_mass_features()
 489
 490    def integrate_mass_features(
 491        self, drop_if_fail=True, drop_duplicates=True, ms_level=1, induced_features=False
 492    ):
 493        """Integrate mass features and extract EICs.
 494
 495        Populates the _eics attribute on the LCMSBase object for each unique mz in the mass_features dataframe and adds data (start_scan, final_scan, area) to the mass_features attribute.
 496
 497        Parameters
 498        ----------
 499        drop_if_fail : bool, optional
 500            Whether to drop mass features if the EIC limit calculations fail.
 501            Default is True.
 502        drop_duplicates : bool, optional
 503            Whether to mass features that appear to be duplicates
 504            (i.e., mz is similar to another mass feature and limits of the EIC are similar or encapsulating).
 505            Default is True.
 506        ms_level : int, optional
 507            The MS level to use. Default is 1.
 508        induced_features : bool, optional
 509            Whether the mass features to be intergrated were induced. Default is False.
 510
 511        Raises
 512        ------
 513        ValueError
 514            If no mass features are found.
 515            If no MS level data is found for the given MS level (either in data or in the scan data)
 516
 517        Returns
 518        -------
 519        None, but populates the eics attribute on the LCMSBase object and adds data (start_scan, final_scan, area) to the mass_features attribute.
 520
 521        Notes
 522        -----
 523        drop_if_fail is useful for discarding mass features that do not have good shapes, usually due to a detection on a shoulder of a peak or a noisy region (especially if minimal smoothing is used during mass feature detection).
 524        """
 525        
 526        # Check if there is data
 527        if ms_level in self._ms_unprocessed.keys():
 528            raw_data = self._ms_unprocessed[ms_level].copy()
 529        else:
 530            raise ValueError("No MS level " + str(ms_level) + " data found")
 531
 532        # Check if mass_spectrum exists on each mass feature
 533        if induced_features:
 534            mf_dict = self.induced_mass_features
 535            if len(mf_dict) == 0:
 536                raise ValueError(
 537                    "No induced mass features found, did you run fill_missing_cluster_features() first?"
 538                )
 539
 540            ## remove not found induced mass features by mz <= 0 (-99 indicator)
 541            # also remove any where mz is nan
 542            mf_dict = {k:v for k, v in mf_dict.items() if v.mz > 0 and not np.isnan(v.mz)}
 543
 544        else:
 545            mf_dict = self.mass_features
 546            if len(mf_dict) == 0:
 547                raise ValueError(
 548                    "No mass features found, did you run find_mass_features() first?"
 549                )
 550
 551        # Subset scan data to only include correct ms_level
 552        scan_df_sub = self.scan_df[
 553            self.scan_df["ms_level"] == int(ms_level)
 554        ].reset_index(drop=True)
 555        if scan_df_sub.empty:
 556            raise ValueError("No MS level " + ms_level + " data found in scan data")
 557        scan_df_sub = scan_df_sub[["scan", "scan_time"]].copy()
 558
 559        mzs_to_extract = np.unique([mf.mz for mf in mf_dict.values()])
 560        mzs_to_extract.sort()
 561
 562        # Pre-sort raw_data by mz for faster filtering
 563        raw_data_sorted = raw_data.sort_values(["mz", "scan"]).reset_index(drop=True)
 564        raw_data_mz = raw_data_sorted["mz"].values
 565
 566        # Get EICs for each unique mz in mass features list
 567        for mz in mzs_to_extract:
 568            mz_max = mz + self.parameters.lc_ms.eic_tolerance_ppm * mz / 1e6
 569            mz_min = mz - self.parameters.lc_ms.eic_tolerance_ppm * mz / 1e6
 570
 571            # Use binary search for faster mz range filtering
 572            left_idx = np.searchsorted(raw_data_mz, mz_min, side="left")
 573            right_idx = np.searchsorted(raw_data_mz, mz_max, side="right")
 574            raw_data_sub = raw_data_sorted.iloc[left_idx:right_idx].copy()
 575
 576            raw_data_sub = (
 577                raw_data_sub.groupby(["scan"])["intensity"].sum().reset_index()
 578            )
 579            raw_data_sub = scan_df_sub.merge(raw_data_sub, on="scan", how="left")
 580            raw_data_sub["intensity"] = raw_data_sub["intensity"].fillna(0)
 581            myEIC = EIC_Data(
 582                scans=raw_data_sub["scan"].values,
 583                time=raw_data_sub["scan_time"].values,
 584                eic=raw_data_sub["intensity"].values,
 585            )
 586            # Smooth EIC
 587            smoothed_eic = self.smooth_tic(myEIC.eic)
 588            smoothed_eic[smoothed_eic < 0] = 0
 589            myEIC.eic_smoothed = smoothed_eic
 590            self.eics[mz] = myEIC
 591
 592        # Get limits of mass features using EIC centroid detector and integrate
 593        for idx, mass_feature in list(mf_dict.items()):
 594            mz = mass_feature.mz
 595            apex_scan = mass_feature.apex_scan
 596
 597            # Pull EIC data and find apex scan index
 598            myEIC = self.eics[mz]
 599            mf_dict[idx]._eic_data = myEIC
 600            mf_dict[idx]._eic_mz = mz
 601            apex_index = np.searchsorted(myEIC.scans, apex_scan)
 602
 603            # Find left and right limits of peak using EIC centroid detector, add to EICData
 604            centroid_eics = self.eic_centroid_detector(
 605                myEIC.time,
 606                myEIC.eic_smoothed,
 607                mass_feature.intensity * 1.1,
 608                apex_indexes=[int(apex_index)],
 609            )
 610            l_a_r_scan_idx = [i for i in centroid_eics]
 611            if len(l_a_r_scan_idx) > 0:
 612                # Calculate number of consecutive scans with intensity > 0 and check if it is above the minimum consecutive scans
 613                # Find the number of consecutive non-zero values in the EIC segment
 614                mask = myEIC.eic[l_a_r_scan_idx[0][0] : l_a_r_scan_idx[0][2] + 1] > 0
 615                # Find the longest run of consecutive True values
 616                if np.any(mask):
 617                    # Find indices where mask changes value
 618                    diff = np.diff(np.concatenate(([0], mask.astype(int), [0])))
 619                    starts = np.where(diff == 1)[0]
 620                    ends = np.where(diff == -1)[0]
 621                    consecutive_scans = (ends - starts).max()
 622                else:
 623                    consecutive_scans = 0
 624                if consecutive_scans < self.parameters.lc_ms.consecutive_scan_min:
 625                    mf_dict.pop(idx)
 626                    continue
 627                # Add start and final scan to mass_features and EICData
 628                left_scan, right_scan = (
 629                    myEIC.scans[l_a_r_scan_idx[0][0]],
 630                    myEIC.scans[l_a_r_scan_idx[0][2]],
 631                )
 632                mf_scan_apex = [(left_scan, int(apex_scan), right_scan)]
 633                myEIC.apexes = myEIC.apexes + mf_scan_apex
 634                mf_dict[idx].start_scan = left_scan
 635                mf_dict[idx].final_scan = right_scan
 636
 637                # Find area under peak using limits from EIC centroid detector, add to mass_features and EICData
 638                area = np.trapezoid(
 639                    myEIC.eic_smoothed[l_a_r_scan_idx[0][0] : l_a_r_scan_idx[0][2] + 1],
 640                    myEIC.time[l_a_r_scan_idx[0][0] : l_a_r_scan_idx[0][2] + 1],
 641                )
 642                myEIC.areas = myEIC.areas + [area]
 643                self.eics[mz] = myEIC
 644                mf_dict[idx]._area = area
 645            else:
 646                if drop_if_fail is True:
 647                    mf_dict.pop(idx)
 648
 649        if drop_duplicates:
 650            # Prepare mass feature dataframe
 651            if induced_features:
 652                mf_df = self.mass_features_to_df(induced_features = True).copy()
 653                mf_df = mf_df[mf_df.start_scan.notna()]        
 654            else:
 655                mf_df = self.mass_features_to_df(induced_features = False).copy()
 656
 657            # For each mass feature, find all mass features within the clustering tolerance ppm and drop if their start and end times are within another mass feature
 658            # Keep the first mass feature (highest persistence)
 659            for idx, mass_feature in mf_df.iterrows():
 660                mz = mass_feature.mz
 661                apex_scan = mass_feature.apex_scan
 662
 663                mf_df["mz_diff_ppm"] = np.abs(mf_df["mz"] - mz) / mz * 10**6
 664                mf_df_sub = mf_df[
 665                    mf_df["mz_diff_ppm"]
 666                    < self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel
 667                    * 10**6
 668                ].copy()
 669
 670                # For all mass features within the clustering tolerance, check if the start and end times are within the start and end times of the mass feature
 671                for idx2, mass_feature2 in mf_df_sub.iterrows():
 672                    if idx2 != idx:
 673                        if (
 674                            mass_feature2.start_scan >= mass_feature.start_scan
 675                            and mass_feature2.final_scan <= mass_feature.final_scan
 676                        ):
 677                            if idx2 in self.mass_features.keys():
 678                                self.mass_features.pop(idx2)
 679        
 680        # Filter MS2 scans to only include those within integration bounds
 681        # This ensures MS2 scans outside start_scan to final_scan are removed
 682        if induced_features:
 683            self._filter_ms2_scans_by_integration_bounds(mf_dict=self.induced_mass_features)
 684        else:
 685            self._filter_ms2_scans_by_integration_bounds(mf_dict=self.mass_features)
 686
 687    def find_c13_mass_features(self):
 688        """Mark likely C13 isotopes and connect to monoisoitopic mass features.
 689
 690        Returns
 691        -------
 692        None, but populates the monoisotopic_mf_id and isotopologue_type attributes to the indivual LCMSMassFeatures within the mass_features attribute of the LCMSBase object.
 693
 694        Raises
 695        ------
 696        ValueError
 697            If no mass features are found.
 698        """
 699        verbose = self.parameters.lc_ms.verbose_processing
 700        if verbose:
 701            print("evaluating mass features for C13 isotopes")
 702        if self.mass_features is None:
 703            raise ValueError("No mass features found, run find_mass_features() first")
 704
 705        # Data prep fo sparse distance matrix
 706        dims = ["mz", "scan_time"]
 707        mf_df = self.mass_features_to_df().copy()
 708        # Drop mass features that have no area (these are likely to be noise)
 709        mf_df = mf_df[mf_df["area"].notnull()]
 710        mf_df["mf_id"] = mf_df.index.values
 711        dims = ["mz", "scan_time"]
 712
 713        # Sort my ascending mz so we always get the monoisotopic mass first, regardless of the order/intensity of the mass features
 714        mf_df = mf_df.sort_values(by=["mz"]).reset_index(drop=True).copy()
 715
 716        mz_diff = 1.003355  # C13-C12 mass difference
 717        tol = [
 718            mf_df["mz"].median()
 719            * self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
 720            self.parameters.lc_ms.mass_feature_cluster_rt_tolerance * 0.5,
 721        ]  # mz, in relative; scan_time in minutes
 722
 723        # Compute inter-feature distances
 724        distances = None
 725        for i in range(len(dims)):
 726            # Construct k-d tree
 727            values = mf_df[dims[i]].values
 728            tree = KDTree(values.reshape(-1, 1))
 729
 730            max_tol = tol[i]
 731            if dims[i] == "mz":
 732                # Maximum absolute tolerance
 733                max_tol = mz_diff + tol[i]
 734
 735            # Compute sparse distance matrix
 736            # the larger the max_tol, the slower this operation is
 737            sdm = tree.sparse_distance_matrix(tree, max_tol, output_type="coo_matrix")
 738
 739            # Only consider forward case, exclude diagonal
 740            sdm = sparse.triu(sdm, k=1)
 741
 742            if dims[i] == "mz":
 743                min_tol = mz_diff - tol[i]
 744                # Get only the ones that are above the min tol
 745                idx = sdm.data > min_tol
 746
 747                # Reconstruct sparse distance matrix
 748                sdm = sparse.coo_matrix(
 749                    (sdm.data[idx], (sdm.row[idx], sdm.col[idx])),
 750                    shape=(len(values), len(values)),
 751                )
 752
 753            # Cast as binary matrix
 754            sdm.data = np.ones_like(sdm.data)
 755
 756            # Stack distances
 757            if distances is None:
 758                distances = sdm
 759            else:
 760                distances = distances.multiply(sdm)
 761
 762        # Extract indices of within-tolerance points
 763        distances = distances.tocoo()
 764        pairs = np.stack((distances.row, distances.col), axis=1)  # C12 to C13 pairs
 765
 766        # Turn pairs (which are index of mf_df) into mf_id and then into two dataframes to join to mf_df
 767        pairs_mf = pairs.copy()
 768        pairs_mf[:, 0] = mf_df.iloc[pairs[:, 0]].mf_id.values
 769        pairs_mf[:, 1] = mf_df.iloc[pairs[:, 1]].mf_id.values
 770
 771        # Connect monoisotopic masses with isotopologes within mass_features
 772        monos = np.setdiff1d(np.unique(pairs_mf[:, 0]), np.unique(pairs_mf[:, 1]))
 773        for mono in monos:
 774            self.mass_features[mono].monoisotopic_mf_id = mono
 775        pairs_iso_df = pd.DataFrame(pairs_mf, columns=["parent", "child"])
 776        while not pairs_iso_df.empty:
 777            pairs_iso_df = pairs_iso_df.set_index("parent", drop=False)
 778            m1_isos = pairs_iso_df.loc[monos, "child"].unique()
 779            for iso in m1_isos:
 780                # Set monoisotopic_mf_id and isotopologue_type for isotopologues
 781                parent = pairs_mf[pairs_mf[:, 1] == iso, 0]
 782                if len(parent) > 1:
 783                    # Choose the parent that is closest in time to the isotopologue
 784                    parent_time = [self.mass_features[p].retention_time for p in parent]
 785                    time_diff = [
 786                        np.abs(self.mass_features[iso].retention_time - x)
 787                        for x in parent_time
 788                    ]
 789                    parent = parent[np.argmin(time_diff)]
 790                else:
 791                    parent = parent[0]
 792                self.mass_features[iso].monoisotopic_mf_id = self.mass_features[
 793                    parent
 794                ].monoisotopic_mf_id
 795                if self.mass_features[iso].monoisotopic_mf_id is not None:
 796                    mass_diff = (
 797                        self.mass_features[iso].mz
 798                        - self.mass_features[
 799                            self.mass_features[iso].monoisotopic_mf_id
 800                        ].mz
 801                    )
 802                    self.mass_features[iso].isotopologue_type = "13C" + str(
 803                        int(round(mass_diff, 0))
 804                    )
 805
 806            # Drop the mono and iso from the pairs_iso_df
 807            pairs_iso_df = pairs_iso_df.drop(
 808                index=monos, errors="ignore"
 809            )  # Drop pairs where the parent is a child that is a child of a root
 810            pairs_iso_df = pairs_iso_df.set_index("child", drop=False)
 811            pairs_iso_df = pairs_iso_df.drop(index=m1_isos, errors="ignore")
 812
 813            if not pairs_iso_df.empty:
 814                # Get new monos, recognizing that these are just 13C isotopologues that are connected to other 13C isotopologues to repeat the process
 815                monos = np.setdiff1d(
 816                    np.unique(pairs_iso_df.parent), np.unique(pairs_iso_df.child)
 817                )
 818        if verbose:
 819            # Report fraction of compounds annotated with isotopes
 820            mf_df["c13_flag"] = np.where(
 821                np.logical_or(
 822                    np.isin(mf_df["mf_id"], pairs_mf[:, 0]),
 823                    np.isin(mf_df["mf_id"], pairs_mf[:, 1]),
 824                ),
 825                1,
 826                0,
 827            )
 828            print(
 829                str(round(len(mf_df[mf_df["c13_flag"] == 1]) / len(mf_df), ndigits=3))
 830                + " of mass features have or are C13 isotopes"
 831            )
 832
 833    def deconvolute_ms1_mass_features(self):
 834        """Deconvolute MS1 mass features
 835
 836        Deconvolute mass features ms1 spectrum based on the correlation of all masses within a spectrum over the EIC of the mass features
 837
 838        Parameters
 839        ----------
 840        None
 841
 842        Returns
 843        -------
 844        None, but assigns the _ms_deconvoluted_idx, mass_spectrum_deconvoluted_parent,
 845        and associated_mass_features_deconvoluted attributes to the mass features in the
 846        mass_features attribute of the LCMSBase object.
 847
 848        Raises
 849        ------
 850        ValueError
 851            If no mass features are found, must run find_mass_features() first.
 852            If no EICs are found, did you run integrate_mass_features() first?
 853
 854        """
 855        # Checks for set mass_features and eics
 856        if self.mass_features is None:
 857            raise ValueError(
 858                "No mass features found, did you run find_mass_features() first?"
 859            )
 860
 861        if self.eics == {}:
 862            raise ValueError(
 863                "No EICs found, did you run integrate_mass_features() first?"
 864            )
 865
 866        if 1 not in self._ms_unprocessed.keys():
 867            raise ValueError("No unprocessed MS1 spectra found.")
 868
 869        # Prep ms1 data
 870        ms1_data = self._ms_unprocessed[1].copy()
 871        ms1_data = ms1_data.set_index("scan")
 872
 873        # Prep mass feature summary
 874        mass_feature_df = self.mass_features_to_df()
 875
 876        # Loop through each mass feature
 877        for mf_id, mass_feature in self.mass_features.items():
 878            # Check that the mass_feature.mz attribute == the mz of the mass feature in the mass_feature_df
 879            if mass_feature.mz != mass_feature.ms1_peak.mz_exp:
 880                continue
 881
 882            # Get the left and right limits of the EIC of the mass feature
 883            l_scan, _, r_scan = mass_feature._eic_data.apexes[0]
 884
 885            # Pull from the _ms1_unprocessed data the scan range of interest and sort by mz
 886            ms1_data_sub = ms1_data.loc[l_scan:r_scan].copy()
 887            ms1_data_sub = ms1_data_sub.sort_values(by=["mz"]).reset_index(drop=False)
 888
 889            # Get the centroided masses of the mass feature
 890            mf_mspeak_mzs = mass_feature.mass_spectrum.mz_exp
 891
 892            # Find the closest mz in the ms1 data to the centroided masses of the mass feature
 893            ms1_data_sub["mass_feature_mz"] = mf_mspeak_mzs[
 894                find_closest(mf_mspeak_mzs, ms1_data_sub.mz.values)
 895            ]
 896
 897            # Drop rows with mz_diff > 0.01 between the mass feature mz and the ms1 data mz
 898            ms1_data_sub["mz_diff_rel"] = (
 899                np.abs(ms1_data_sub["mass_feature_mz"] - ms1_data_sub["mz"])
 900                / ms1_data_sub["mz"]
 901            )
 902            ms1_data_sub = ms1_data_sub[
 903                ms1_data_sub["mz_diff_rel"]
 904                < self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel
 905            ].reset_index(drop=True)
 906
 907            # Group by mass_feature_mz and scan and sum intensity
 908            ms1_data_sub_group = (
 909                ms1_data_sub.groupby(["mass_feature_mz", "scan"])["intensity"]
 910                .sum()
 911                .reset_index()
 912            )
 913
 914            # Calculate the correlation of the intensities of the mass feature and the ms1 data (set to 0 if no intensity)
 915            corr = (
 916                ms1_data_sub_group.pivot(
 917                    index="scan", columns="mass_feature_mz", values="intensity"
 918                )
 919                .fillna(0)
 920                .corr()
 921            )
 922
 923            # Subset the correlation matrix to only include the masses of the mass feature and those with a correlation > 0.8
 924            decon_corr_min = self.parameters.lc_ms.ms1_deconvolution_corr_min
 925
 926            # Try catch for KeyError in case the mass feature mz is not in the correlation matrix
 927            try:
 928                corr_subset = corr.loc[mass_feature.mz]
 929            except KeyError:
 930                # If the mass feature mz is not in the correlation matrix, skip to the next mass feature
 931                continue
 932
 933            corr_subset = corr_subset[corr_subset > decon_corr_min]
 934
 935            # Get the masses from the mass spectrum that are the result of the deconvolution
 936            mzs_decon = corr_subset.index.values
 937
 938            # Get the indices of the mzs_decon in mass_feature.mass_spectrum.mz_exp and assign to the mass feature
 939            mzs_decon_idx = [
 940                id
 941                for id, mz in enumerate(mass_feature.mass_spectrum.mz_exp)
 942                if mz in mzs_decon
 943            ]
 944            mass_feature._ms_deconvoluted_idx = mzs_decon_idx
 945
 946            # Check if the mass feature's ms1 peak is the largest in the deconvoluted mass spectrum
 947            if (
 948                mass_feature.ms1_peak.abundance
 949                == mass_feature.mass_spectrum.abundance[mzs_decon_idx].max()
 950            ):
 951                mass_feature.mass_spectrum_deconvoluted_parent = True
 952            else:
 953                mass_feature.mass_spectrum_deconvoluted_parent = False
 954
 955            # Check for other mass features that are in the deconvoluted mass spectrum and add the deconvoluted mass spectrum to the mass feature
 956            # Subset mass_feature_df to only include mass features that are within the clustering tolerance
 957            mass_feature_df_sub = mass_feature_df[
 958                abs(mass_feature.retention_time - mass_feature_df["scan_time"])
 959                < self.parameters.lc_ms.mass_feature_cluster_rt_tolerance
 960            ].copy()
 961            # Calculate the mz difference in ppm between the mass feature and the peaks in the deconvoluted mass spectrum
 962            mass_feature_df_sub["mz_diff_ppm"] = [
 963                np.abs(mzs_decon - mz).min() / mz * 10**6
 964                for mz in mass_feature_df_sub["mz"]
 965            ]
 966            # Subset mass_feature_df to only include mass features that are within 1 ppm of the deconvoluted masses
 967            mfs_associated_decon = mass_feature_df_sub[
 968                mass_feature_df_sub["mz_diff_ppm"]
 969                < self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel * 10**6
 970            ].index.values
 971
 972            mass_feature.associated_mass_features_deconvoluted = mfs_associated_decon
 973
 974    def _remove_redundant_mass_features(
 975        self,
 976        ) -> None:
 977        """
 978        Identify and remove redundant mass features that are likely contaminants based on their m/z values and scan frequency. 
 979        Especially useful for HILIC data where signals do not return to baseline between peaks or for data with significant background noise.
 980        
 981        Contaminants are characterized by:
 982        1. Similar m/z values (within ppm_tolerance)
 983        2. High frequency across scan numbers (ubiquitous presence)
 984    
 985        Notes
 986        -----
 987        Depends on self.mass_features being populated, uses the parameters in self.parameters.lc_ms for tolerances (mass_feature_cluster_mz_tolerance_rel)
 988        """
 989        ppm_tolerance = self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel*1e6
 990        min_scan_frequency = self.parameters.lc_ms.redundant_scan_frequency_min
 991        n_retain = self.parameters.lc_ms.redundant_feature_retain_n
 992
 993        df = self.mass_features_to_df()
 994
 995        if df.empty:
 996            return pd.DataFrame()
 997        # df index should be mf_id
 998        if 'mf_id' not in df.columns:
 999            if 'mf_id' in df.index.names:
1000                df = df.reset_index()
1001            else:
1002                raise ValueError("DataFrame must contain 'mf_id' column or index.")
1003        
1004        # Sort by m/z for efficient grouping
1005        df_sorted = df.sort_values('mz').reset_index(drop=True)
1006        
1007        # Calculate total number of unique scans for frequency calculation
1008        # Calculate total possible scans (check the cluster rt tolerance and the min rt and max rt of the data)
1009        total_time = self.scan_df['scan_time'].max() - self.scan_df['scan_time'].min()
1010        cluster_rt_tolerance = self.parameters.lc_ms.mass_feature_cluster_rt_tolerance
1011        # If the feature was detected in every possible scan (and then rolled up), it would be in this many scans
1012        total_scans =  int(total_time / cluster_rt_tolerance) + 1
1013
1014        # Group similar m/z values using ppm tolerance
1015        mz_groups = []
1016        current_group = []
1017        
1018        for i, row in df_sorted.iterrows():
1019            current_mz = row['mz']
1020            
1021            if not current_group:
1022                # Start first group
1023                current_group = [i]
1024            else:
1025                # Check if current m/z is within tolerance of group representative
1026                group_representative_mz = df_sorted.iloc[current_group[0]]['mz']
1027                ppm_diff = abs(current_mz - group_representative_mz) / group_representative_mz * 1e6
1028                
1029                if ppm_diff <= ppm_tolerance:
1030                    # Add to current group
1031                    current_group.append(i)
1032                else:
1033                    # Start new group, but first process current group
1034                    if len(current_group) > 0:
1035                        mz_groups.append(current_group)
1036                    current_group = [i]
1037        
1038        # Don't forget the last group
1039        if current_group:
1040            mz_groups.append(current_group)
1041        
1042        # Analyze each m/z group for contaminant characteristics
1043        
1044        for group_indices in mz_groups:
1045            group_data = df_sorted.iloc[group_indices]
1046            
1047            # Calculate group statistics
1048            unique_scans = group_data['apex_scan'].nunique()
1049            scan_frequency = unique_scans / total_scans
1050            
1051            # Check if this group meets contaminant criteria
1052            if scan_frequency >= min_scan_frequency:
1053                group_data = group_data.sort_values('intensity', ascending=False)
1054                non_representative_mf_id = group_data.iloc[n_retain:]['mf_id'].tolist()  # These will be removed
1055
1056                self.mass_features = {
1057                    k: v for k, v in self.mass_features.items() if k not in non_representative_mf_id
1058                }
1059
1060    def _remove_mass_features_by_peak_metrics(self, induced_features=False) -> None:
1061        """Remove mass features based on peak metrics defined in mass_feature_attribute_filter_dict.
1062        
1063        This method filters mass features based on various peak shape metrics and quality indicators
1064        such as noise scores, Gaussian similarity, tailing factors, dispersity index, etc.
1065        
1066        The filtering criteria are defined in the mass_feature_attribute_filter_dict parameter,
1067        which should contain attribute names as keys and filter specifications as values.
1068        
1069        Filter specification format:
1070        {attribute_name: {'value': threshold, 'operator': comparison}}
1071        
1072        Available operators:
1073        - '>' or 'greater': Keep features where attribute > threshold
1074        - '<' or 'less': Keep features where attribute < threshold  
1075        - '>=' or 'greater_equal': Keep features where attribute >= threshold
1076        - '<=' or 'less_equal': Keep features where attribute <= threshold
1077        
1078        Examples:
1079        - {'noise_score_max': {'value': 0.5, 'operator': '>='}} - Keep features with noise_score_max >= 0.5
1080        - {'dispersity_index': {'value': 0.1, 'operator': '<'}} - Keep features with dispersity_index < 0.1
1081        - {'gaussian_similarity': {'value': 0.7, 'operator': '>='}} - Keep features with gaussian_similarity >= 0.7
1082        
1083        Parameters
1084        ----------
1085        induced_features : bool, optional
1086            If True, filter induced_mass_features instead of regular mass_features. Default is False.
1087        
1088        Returns
1089        -------
1090        None
1091            Modifies self.mass_features or self.induced_mass_features in place by removing filtered features.
1092            
1093        Raises
1094        ------
1095        ValueError
1096            If no mass features are found, if an invalid attribute is specified, or if filter specification is malformed.
1097        """
1098        # Select the appropriate mass features dictionary
1099        if induced_features:
1100            mf_dict = self.induced_mass_features
1101            mf_type = "induced mass features"
1102        else:
1103            mf_dict = self.mass_features
1104            mf_type = "mass features"
1105            
1106        if mf_dict is None or len(mf_dict) == 0:
1107            raise ValueError(f"No {mf_type} found, run {'gap filling' if induced_features else 'find_mass_features()'} first")
1108            
1109        filter_dict = self.parameters.lc_ms.mass_feature_attribute_filter_dict
1110        
1111        if not filter_dict:
1112            # No filtering criteria specified, return early
1113            return
1114            
1115        verbose = self.parameters.lc_ms.verbose_processing
1116        initial_count = len(mf_dict)
1117        
1118        if verbose:
1119            print(f"Filtering {mf_type} using peak metrics. Initial count: {initial_count}")
1120            
1121        # List to collect IDs of mass features to remove
1122        features_to_remove = []
1123        
1124        for mf_id, mass_feature in mf_dict.items():
1125            should_remove = False
1126            
1127            for attribute_name, filter_spec in filter_dict.items():
1128                # Validate filter specification structure
1129                if not isinstance(filter_spec, dict):
1130                    raise ValueError(f"Filter specification for '{attribute_name}' must be a dictionary with 'value' and 'operator' keys")
1131                
1132                if 'value' not in filter_spec or 'operator' not in filter_spec:
1133                    raise ValueError(f"Filter specification for '{attribute_name}' must contain both 'value' and 'operator' keys")
1134                
1135                threshold_value = filter_spec['value']
1136                operator = filter_spec['operator'].lower().strip()
1137                
1138                # Validate operator
1139                valid_operators = {'>', '<', '>=', '<=', 'greater', 'less', 'greater_equal', 'less_equal'}
1140                if operator not in valid_operators:
1141                    raise ValueError(f"Invalid operator '{operator}' for attribute '{attribute_name}'. Valid operators: {valid_operators}")
1142                
1143                # Normalize operator names
1144                operator_map = {
1145                    'greater': '>',
1146                    'less': '<', 
1147                    'greater_equal': '>=',
1148                    'less_equal': '<='
1149                }
1150                operator = operator_map.get(operator, operator)
1151                
1152                # Get the attribute value from the mass feature
1153                try:
1154                    if hasattr(mass_feature, attribute_name):
1155                        attribute_value = getattr(mass_feature, attribute_name)
1156                    else:
1157                        raise ValueError(f"Mass feature does not have attribute '{attribute_name}'")
1158                        
1159                    # Handle None values or attributes that haven't been calculated
1160                    if attribute_value is None:
1161                        if verbose:
1162                            print(f"Warning: Mass feature {mf_id} has None value for '{attribute_name}'. Removing feature.")
1163                        should_remove = True
1164                        break
1165                        
1166                    # Handle numpy arrays (like half_height_width which returns mean)
1167                    if hasattr(attribute_value, '__len__') and not isinstance(attribute_value, str):
1168                        # For arrays, we use the mean or appropriate summary statistic
1169                        if attribute_name == 'half_height_width':
1170                            # half_height_width property already returns the mean
1171                            pass
1172                        else:
1173                            attribute_value = float(np.mean(attribute_value))
1174                    
1175                    # Handle NaN values
1176                    if np.isnan(float(attribute_value)):
1177                        if verbose:
1178                            print(f"Warning: Mass feature {mf_id} has NaN value for '{attribute_name}'. Removing feature.")
1179                        should_remove = True
1180                        break
1181                    
1182                    # Apply the threshold comparison based on operator
1183                    attribute_value = float(attribute_value)
1184                    threshold_value = float(threshold_value)
1185                    
1186                    if operator == '>' and not (attribute_value > threshold_value):
1187                        should_remove = True
1188                        break
1189                    elif operator == '<' and not (attribute_value < threshold_value):
1190                        should_remove = True
1191                        break
1192                    elif operator == '>=' and not (attribute_value >= threshold_value):
1193                        should_remove = True
1194                        break  
1195                    elif operator == '<=' and not (attribute_value <= threshold_value):
1196                        should_remove = True
1197                        break
1198                        
1199                except (AttributeError, ValueError, TypeError) as e:
1200                    if verbose:
1201                        print(f"Error evaluating filter '{attribute_name}' for mass feature {mf_id}: {e}")
1202                    should_remove = True
1203                    break
1204            
1205            if should_remove:
1206                features_to_remove.append(mf_id)
1207        
1208        # Remove filtered mass features
1209        for mf_id in features_to_remove:
1210            del mf_dict[mf_id]
1211        
1212        if verbose and len(features_to_remove) > 0:
1213            print(f"Removed {len(features_to_remove)} {mf_type} based on peak metrics. Remaining: {len(mf_dict)}")
1214        
1215        # Update the appropriate dictionary
1216        if induced_features:
1217            self.induced_mass_features = mf_dict
1218        else:
1219            self.mass_features = mf_dict
1220        
1221        # Clean up unassociated EICs and ms1 data (only for regular features)
1222        self._remove_unassociated_eics()
1223        self._remove_unassociated_ms1_spectra()
1224            
1225    def _remove_unassociated_eics(self) -> None:
1226        """Remove EICs that are not associated with any mass features.
1227
1228        This method cleans up the eics attribute by removing any EICs that do not correspond to
1229        any mass features currently stored in the mass_features attribute. This is useful for
1230        freeing up memory and ensuring that only relevant EICs are retained.
1231
1232        Returns
1233        -------
1234        None
1235            Modifies self.eics in place by removing unassociated EICs.
1236        """
1237        if self.mass_features is None or len(self.mass_features) == 0:
1238            self.eics = {}
1239            return
1240
1241        # Get the set of m/z values associated with current mass features
1242        associated_mzs = {mf.mz for mf in self.mass_features.values()}
1243
1244        # Remove EICs that are not associated with any mass features
1245        self.eics = {mz: eic for mz, eic in self.eics.items() if mz in associated_mzs}
1246    
1247    def _remove_unassociated_ms1_spectra(self) -> None:
1248        """Remove MS1 spectra that are not associated with any mass features.
1249        This method cleans up the _ms_unprocessed attribute by removing any MS1 spectra that do not correspond to
1250        any mass features currently stored in the mass_features attribute. This is useful for freeing up memory
1251        and ensuring that only relevant MS1 spectra are retained.
1252
1253        Returns
1254        -------
1255        None
1256        """
1257        if self.mass_features is None or len(self.mass_features) == 0:
1258            self._ms_unprocessed = {}
1259            return
1260
1261        # Get the set of m/z values associated with current mass features
1262        associated_ms1_scans = {mf.apex_scan for mf in self.mass_features.values()}
1263        associated_ms1_scans = [int(scan) for scan in associated_ms1_scans]
1264        
1265        # Get keys within the _ms attribute (these are individual MassSpectrum objects)
1266        current_stored_spectra = list(set(self._ms.keys()))
1267        if len(current_stored_spectra) == 0:
1268            return
1269        current_stored_spectra = [int(scan) for scan in current_stored_spectra]
1270
1271        # Filter the current_stored_spectra to only ms1 scans
1272        current_stored_spectra_ms1 = [ scan for scan in current_stored_spectra if scan in self.ms1_scans ]
1273
1274        # Remove MS1 spectra that are not associated with any mass features
1275        scans_to_drop = [scan for scan in current_stored_spectra_ms1 if scan not in associated_ms1_scans]
1276        for scan in scans_to_drop:
1277            if scan in self._ms:
1278                del self._ms[scan]
1279
1280class PHCalculations:
1281    """Methods for performing calculations related to 2D peak picking via persistent homology on LCMS data.
1282
1283    Notes
1284    -----
1285    This class is intended to be used as a mixin for the LCMSBase class.
1286
1287    Methods
1288    -------
1289    * sparse_mean_filter(idx, V, radius=[0, 1, 1]).
1290        Sparse implementation of a mean filter.
1291    * embed_unique_indices(a).
1292        Creates an array of indices, sorted by unique element.
1293    * sparse_upper_star(idx, V).
1294        Sparse implementation of an upper star filtration.
1295    * check_if_grid(data).
1296        Check if the data is gridded in mz space.
1297    * grid_data(data).
1298        Grid the data in the mz dimension.
1299    * find_mass_features_ph(ms_level=1, grid=True).
1300        Find mass features within an LCMSBase object using persistent homology.
1301    * cluster_mass_features(drop_children=True).
1302        Cluster regions of interest.
1303    """
1304
1305    @staticmethod
1306    def sparse_mean_filter(idx, V, radius=[0, 1, 1]):
1307        """Sparse implementation of a mean filter.
1308
1309        Parameters
1310        ----------
1311        idx : :obj:`~numpy.array`
1312            Edge indices for each dimension (MxN).
1313        V : :obj:`~numpy.array`
1314            Array of intensity data (Mx1).
1315        radius : float or list
1316            Radius of the sparse filter in each dimension. Values less than
1317            zero indicate no connectivity in that dimension.
1318
1319        Returns
1320        -------
1321        :obj:`~numpy.array`
1322            Filtered intensities (Mx1).
1323
1324        Notes
1325        -----
1326        This function has been adapted from the original implementation in the Deimos package: https://github.com/pnnl/deimos.
1327        This is a static method.
1328        """
1329
1330        # Copy indices
1331        idx = idx.copy().astype(V.dtype)
1332
1333        # Scale
1334        for i, r in enumerate(radius):
1335            # Increase inter-index distance
1336            if r < 1:
1337                idx[:, i] *= 2
1338
1339            # Do nothing
1340            elif r == 1:
1341                pass
1342
1343            # Decrease inter-index distance
1344            else:
1345                idx[:, i] /= r
1346
1347        # Connectivity matrix
1348        cmat = KDTree(idx)
1349        cmat = cmat.sparse_distance_matrix(cmat, 1, p=np.inf, output_type="coo_matrix")
1350        cmat.setdiag(1)
1351
1352        # Pair indices
1353        I, J = cmat.nonzero()
1354
1355        # Delete cmat
1356        cmat_shape = cmat.shape
1357        del cmat
1358
1359        # Sum over columns
1360        V_sum = sparse.bsr_matrix(
1361            (V[J], (I, I)), shape=cmat_shape, dtype=V.dtype
1362        ).diagonal(0)
1363
1364        # Count over columns
1365        V_count = sparse.bsr_matrix(
1366            (np.ones_like(J), (I, I)), shape=cmat_shape, dtype=V.dtype
1367        ).diagonal(0)
1368
1369        return V_sum / V_count
1370
1371    @staticmethod
1372    def embed_unique_indices(a):
1373        """Creates an array of indices, sorted by unique element.
1374
1375        Parameters
1376        ----------
1377        a : :obj:`~numpy.array`
1378            Array of unique elements (Mx1).
1379
1380        Returns
1381        -------
1382        :obj:`~numpy.array`
1383            Array of indices (Mx1).
1384
1385        Notes
1386        -----
1387        This function has been adapted from the original implementation in the Deimos package: https://github.com/pnnl/deimos
1388        This is a static method.
1389        """
1390
1391        def count_tens(n):
1392            # Count tens
1393            ntens = (n - 1) // 10
1394
1395            while True:
1396                ntens_test = (ntens + n - 1) // 10
1397
1398                if ntens_test == ntens:
1399                    return ntens
1400                else:
1401                    ntens = ntens_test
1402
1403        def arange_exclude_10s(n):
1404            # How many 10s will there be?
1405            ntens = count_tens(n)
1406
1407            # Base array
1408            arr = np.arange(0, n + ntens)
1409
1410            # Exclude 10s
1411            arr = arr[(arr == 0) | (arr % 10 != 0)][:n]
1412
1413            return arr
1414
1415        # Creates an array of indices, sorted by unique element
1416        idx_sort = np.argsort(a)
1417        idx_unsort = np.argsort(idx_sort)
1418
1419        # Sorts records array so all unique elements are together
1420        sorted_a = a[idx_sort]
1421
1422        # Returns the unique values, the index of the first occurrence,
1423        # and the count for each element
1424        vals, idx_start, count = np.unique(
1425            sorted_a, return_index=True, return_counts=True
1426        )
1427
1428        # Splits the indices into separate arrays
1429        splits = np.split(idx_sort, idx_start[1:])
1430
1431        # Creates unique indices for each split
1432        idx_unq = np.concatenate([arange_exclude_10s(len(x)) for x in splits])
1433
1434        # Reorders according to input array
1435        idx_unq = idx_unq[idx_unsort]
1436
1437        # Magnitude of each index
1438        exp = np.log10(
1439            idx_unq, where=idx_unq > 0, out=np.zeros_like(idx_unq, dtype=np.float64)
1440        )
1441        idx_unq_mag = np.power(10, np.floor(exp) + 1)
1442
1443        # Result
1444        return a + idx_unq / idx_unq_mag
1445
1446    @staticmethod
1447    def roll_up_dataframe(
1448        df: pd.DataFrame,
1449        sort_by: str,
1450        tol: list,
1451        relative: list,
1452        dims: list,
1453        memory_opt_threshold: int = 10000,
1454    ):
1455        """Subset data by rolling up into apex in appropriate dimensions.
1456
1457        Parameters
1458        ----------
1459        data : pd.DataFrame
1460            The input data containing "dims" columns and the "sort_by" column.
1461        sort_by : str
1462            The column to sort the data by, this will determine which mass features get rolled up into a parent mass feature
1463            (i.e., the mass feature with the highest value in the sort_by column).
1464        dims : list
1465            A list of dimension names (column names in the data DataFrame) to roll up the mass features by.
1466        tol : list
1467            A list of tolerances for each dimension. The length of the list must match the number of dimensions.
1468            The tolerances can be relative (as a fraction of the maximum value in that dimension) or absolute (in the units of that dimension).
1469            If relative is True, the tolerance will be multiplied by the maximum value in that dimension.
1470            If relative is False, the tolerance will be used as is.
1471        relative : list
1472            A list of booleans indicating whether the tolerance for each dimension is relative (True) or absolute (False).
1473        memory_opt_threshold : int, optional
1474            Minimum number of rows to trigger memory-optimized processing. Default is 10000.
1475
1476        Returns
1477        -------
1478        pd.DataFrame
1479            A DataFrame with only the rolled up mass features, with the original index and columns.
1480
1481
1482        Raises
1483        ------
1484        ValueError
1485            If the input data is not a pandas DataFrame.
1486            If the input data does not have columns for each of the dimensions in "dims".
1487            If the length of "dims", "tol", and "relative" do not match.
1488        """
1489        og_columns = df.columns.copy()
1490
1491        # Unindex the data, but keep the original index
1492        if df.index.name is not None:
1493            og_index = df.index.name
1494        else:
1495            og_index = "index"
1496        df = df.reset_index(drop=False)
1497
1498        # Sort data by sort_by column, and reindex
1499        df = df.sort_values(by=sort_by, ascending=False).reset_index(drop=True)
1500
1501        # Check that data is a DataFrame and has columns for each of the dims
1502        if not isinstance(df, pd.DataFrame):
1503            raise ValueError("Data must be a pandas DataFrame")
1504        for dim in dims:
1505            if dim not in df.columns:
1506                raise ValueError(f"Data must have a column for {dim}")
1507        if len(dims) != len(tol) or len(dims) != len(relative):
1508            raise ValueError(
1509                "Dimensions, tolerances, and relative flags must be the same length"
1510            )
1511
1512        # Pre-compute all values arrays
1513        all_values = [df[dim].values for dim in dims]
1514
1515        # Choose processing method based on dataframe size
1516        if len(df) >= memory_opt_threshold:
1517            # Memory-optimized approach for large dataframes
1518            distances = PHCalculations._compute_distances_memory_optimized(
1519                all_values, tol, relative
1520            )
1521        else:
1522            # Faster approach for smaller dataframes
1523            distances = PHCalculations._compute_distances_original(
1524                all_values, tol, relative
1525            )
1526
1527        # Process pairs with original logic but memory optimizations
1528        distances = distances.tocoo()
1529        pairs = np.stack((distances.row, distances.col), axis=1)
1530        pairs_df = pd.DataFrame(pairs, columns=["parent", "child"]).set_index("parent")
1531        del distances, pairs  # Free memory immediately
1532
1533        to_drop = []
1534        while not pairs_df.empty:
1535            # Find root_parents and their children (original logic preserved)
1536            root_parents = np.setdiff1d(
1537                np.unique(pairs_df.index.values), np.unique(pairs_df.child.values)
1538            )
1539            children_of_roots = pairs_df.loc[root_parents, "child"].unique()
1540            to_drop.extend(children_of_roots)  # Use extend instead of append
1541
1542            # Remove root_children as possible parents from pairs_df for next iteration
1543            pairs_df = pairs_df.drop(index=children_of_roots, errors="ignore")
1544            pairs_df = pairs_df.reset_index().set_index("child")
1545            # Remove root_children as possible children from pairs_df for next iteration
1546            pairs_df = pairs_df.drop(index=children_of_roots)
1547
1548            # Prepare for next iteration
1549            pairs_df = pairs_df.reset_index().set_index("parent")
1550
1551        # Convert to numpy array for efficient dropping
1552        to_drop = np.array(to_drop)
1553
1554        # Drop mass features that are not cluster parents
1555        df_sub = df.drop(index=to_drop)
1556
1557        # Set index back to og_index and only keep original columns
1558        df_sub = df_sub.set_index(og_index).sort_index()[og_columns]
1559
1560        return df_sub
1561
1562    @staticmethod
1563    def _compute_distances_original(all_values, tol, relative):
1564        """Original distance computation method for smaller datasets.
1565
1566        This method computes the pairwise distances between features in the dataset
1567        using a straightforward approach. It is suitable for smaller datasets where
1568        memory usage is not a primary concern.
1569
1570        Parameters
1571        ----------
1572        all_values : list of :obj:`~numpy.array`
1573            List of arrays containing the values for each dimension.
1574        tol : list of float
1575            List of tolerances for each dimension.
1576        relative : list of bool
1577            List of booleans indicating whether the tolerance for each dimension is relative (True) or absolute (False).
1578
1579        Returns
1580        -------
1581        :obj:`~scipy.sparse.coo_matrix`
1582            Sparse matrix indicating pairwise distances within tolerances.
1583        """
1584        # Compute inter-feature distances with memory optimization
1585        distances = None
1586        for i in range(len(all_values)):
1587            values = all_values[i]
1588            # Use single precision if possible to reduce memory
1589            tree = KDTree(values.reshape(-1, 1).astype(np.float32))
1590
1591            max_tol = tol[i]
1592            if relative[i] is True:
1593                max_tol = tol[i] * values.max()
1594
1595            # Compute sparse distance matrix with smaller chunks if memory is an issue
1596            sdm = tree.sparse_distance_matrix(tree, max_tol, output_type="coo_matrix")
1597
1598            # Only consider forward case, exclude diagonal
1599            sdm = sparse.triu(sdm, k=1)
1600
1601            # Process relative distances more efficiently
1602            if relative[i] is True:
1603                # Vectorized computation without creating intermediate arrays
1604                row_values = values[sdm.row]
1605                valid_idx = sdm.data <= tol[i] * row_values
1606
1607                # Reconstruct sparse matrix more efficiently
1608                sdm = sparse.coo_matrix(
1609                    (
1610                        np.ones(valid_idx.sum(), dtype=np.uint8),
1611                        (sdm.row[valid_idx], sdm.col[valid_idx]),
1612                    ),
1613                    shape=(len(values), len(values)),
1614                )
1615            else:
1616                # Cast as binary matrix with smaller data type
1617                sdm.data = np.ones(len(sdm.data), dtype=np.uint8)
1618
1619            # Stack distances with memory-efficient multiplication
1620            if distances is None:
1621                distances = sdm
1622            else:
1623                # Use in-place operations where possible
1624                distances = distances.multiply(sdm)
1625                del sdm  # Free memory immediately
1626
1627        return distances
1628
1629    @staticmethod
1630    def _compute_distances_memory_optimized(all_values, tol, relative):
1631        """Memory-optimized distance computation for large datasets.
1632
1633        This method computes the pairwise distances between features in the dataset
1634        using a more memory-efficient approach. It is suitable for larger datasets
1635        where memory usage is a primary concern.
1636
1637        Parameters
1638        ----------
1639        all_values : list of :obj:`~numpy.array`
1640            List of arrays containing the values for each dimension.
1641        tol : list of float
1642            List of tolerances for each dimension.
1643        relative : list of bool
1644            List of booleans indicating whether the tolerance for each dimension is relative (True) or absolute (False).
1645
1646        Returns
1647        -------
1648        :obj:`~scipy.sparse.coo_matrix`
1649            Sparse matrix indicating pairwise distances within tolerances.
1650        """
1651        # Compute distance matrix for first dimension (full matrix as before)
1652        values_0 = all_values[0].astype(np.float32)
1653        tree_0 = KDTree(values_0.reshape(-1, 1))
1654
1655        max_tol_0 = tol[0]
1656        if relative[0] is True:
1657            max_tol_0 = tol[0] * values_0.max()
1658
1659        # Compute sparse distance matrix for first dimension
1660        distances = tree_0.sparse_distance_matrix(
1661            tree_0, max_tol_0, output_type="coo_matrix"
1662        )
1663        distances = sparse.triu(distances, k=1)
1664
1665        # Process relative distances for first dimension
1666        if relative[0] is True:
1667            row_values = values_0[distances.row]
1668            valid_idx = distances.data <= tol[0] * row_values
1669            distances = sparse.coo_matrix(
1670                (
1671                    np.ones(valid_idx.sum(), dtype=np.uint8),
1672                    (distances.row[valid_idx], distances.col[valid_idx]),
1673                ),
1674                shape=(len(values_0), len(values_0)),
1675            )
1676        else:
1677            distances.data = np.ones(len(distances.data), dtype=np.uint8)
1678
1679        # For remaining dimensions, work only on chunks defined by first dimension pairs
1680        if len(all_values) > 1:
1681            distances_coo = distances.tocoo()
1682            valid_pairs = []
1683
1684            # Process each pair from first dimension
1685            for idx in range(len(distances_coo.data)):
1686                i, j = distances_coo.row[idx], distances_coo.col[idx]
1687                is_valid_pair = True
1688
1689                # Check remaining dimensions for this specific pair
1690                for dim_idx in range(1, len(all_values)):
1691                    values = all_values[dim_idx]
1692                    val_i, val_j = values[i], values[j]
1693
1694                    max_tol = tol[dim_idx]
1695                    if relative[dim_idx] is True:
1696                        max_tol = tol[dim_idx] * values.max()
1697
1698                    distance_ij = abs(val_i - val_j)
1699
1700                    # Check if this pair satisfies the tolerance for this dimension
1701                    if relative[dim_idx] is True:
1702                        if distance_ij > tol[dim_idx] * val_i:
1703                            is_valid_pair = False
1704                            break
1705                    else:
1706                        if distance_ij > max_tol:
1707                            is_valid_pair = False
1708                            break
1709
1710                if is_valid_pair:
1711                    valid_pairs.append((i, j))
1712
1713            # Rebuild distances matrix with only valid pairs
1714            if valid_pairs:
1715                valid_pairs = np.array(valid_pairs)
1716                distances = sparse.coo_matrix(
1717                    (
1718                        np.ones(len(valid_pairs), dtype=np.uint8),
1719                        (valid_pairs[:, 0], valid_pairs[:, 1]),
1720                    ),
1721                    shape=(len(values_0), len(values_0)),
1722                )
1723            else:
1724                # No valid pairs found
1725                distances = sparse.coo_matrix(
1726                    (len(values_0), len(values_0)), dtype=np.uint8
1727                )
1728
1729        return distances
1730
1731    def sparse_upper_star(self, idx, V):
1732        """Sparse implementation of an upper star filtration.
1733
1734        Parameters
1735        ----------
1736        idx : :obj:`~numpy.array`
1737            Edge indices for each dimension (MxN).
1738        V : :obj:`~numpy.array`
1739            Array of intensity data (Mx1).
1740        Returns
1741        -------
1742        idx : :obj:`~numpy.array`
1743            Index of filtered points (Mx1).
1744        persistence : :obj:`~numpy.array`
1745            Persistence of each filtered point (Mx1).
1746
1747        Notes
1748        -----
1749        This function has been adapted from the original implementation in the Deimos package: https://github.com/pnnl/deimos
1750        """
1751
1752        # Invert
1753        V = -1 * V.copy().astype(int)
1754
1755        # Embed indices
1756        V = self.embed_unique_indices(V)
1757
1758        # Connectivity matrix
1759        cmat = KDTree(idx)
1760        cmat = cmat.sparse_distance_matrix(cmat, 1, p=np.inf, output_type="coo_matrix")
1761        cmat.setdiag(1)
1762        cmat = sparse.triu(cmat)
1763
1764        # Pairwise minimums
1765        I, J = cmat.nonzero()
1766        d = np.maximum(V[I], V[J])
1767
1768        # Delete connectiity matrix
1769        cmat_shape = cmat.shape
1770        del cmat
1771
1772        # Sparse distance matrix
1773        sdm = sparse.coo_matrix((d, (I, J)), shape=cmat_shape)
1774
1775        # Delete pairwise mins
1776        del d, I, J
1777
1778        # Persistence homology
1779        ph = ripser(sdm, distance_matrix=True, maxdim=0)["dgms"][0]
1780
1781        # Bound death values
1782        ph[ph[:, 1] == np.inf, 1] = np.max(V)
1783
1784        # Construct tree to query against
1785        tree = KDTree(V.reshape((-1, 1)))
1786
1787        # Get the indexes of the first nearest neighbor by birth
1788        _, nn = tree.query(ph[:, 0].reshape((-1, 1)), k=1, workers=-1)
1789
1790        return nn, -(ph[:, 0] // 1 - ph[:, 1] // 1)
1791
1792    def check_if_grid(self, data):
1793        """Check if the data are gridded in mz space.
1794
1795        Parameters
1796        ----------
1797        data : DataFrame
1798            DataFrame containing the mass spectrometry data.  Needs to have mz and scan columns.
1799
1800        Returns
1801        -------
1802        bool
1803            True if the data is gridded in the mz direction, False otherwise.
1804
1805        Notes
1806        -----
1807        This function is used within the grid_data function and the find_mass_features function and is not intended to be called directly.
1808        """
1809        # Calculate the difference between consecutive mz values in a single scan
1810        dat_check = data.copy().reset_index(drop=True)
1811        dat_check["mz_diff"] = np.abs(dat_check["mz"].diff())
1812        mz_diff_min = (
1813            dat_check.groupby("scan")["mz_diff"].min().min()
1814        )  # within each scan, what is the smallest mz difference between consecutive mz values
1815
1816        # Find the mininum mz difference between mz values in the data; regardless of scan
1817        dat_check_mz = dat_check[["mz"]].drop_duplicates().copy()
1818        dat_check_mz = dat_check_mz.sort_values(by=["mz"]).reset_index(drop=True)
1819        dat_check_mz["mz_diff"] = np.abs(dat_check_mz["mz"].diff())
1820
1821        # Get minimum mz_diff between mz values in the data
1822        mz_diff_min_raw = dat_check_mz["mz_diff"].min()
1823
1824        # If the minimum mz difference between mz values in the data is less than the minimum mz difference between mz values within a single scan, then the data is not gridded
1825        if mz_diff_min_raw < mz_diff_min:
1826            return False
1827        else:
1828            return True
1829
1830    def grid_data(self, data, attempts=5):
1831        """Grid the data in the mz dimension.
1832
1833        Data must be gridded prior to persistent homology calculations and computing average mass spectrum
1834
1835        Parameters
1836        ----------
1837        data : DataFrame
1838            The input data containing mz, scan, scan_time, and intensity columns.
1839        attempts : int, optional
1840            The number of attempts to grid the data. Default is 5.
1841
1842        Returns
1843        -------
1844        DataFrame
1845            The gridded data with mz, scan, scan_time, and intensity columns.
1846
1847        Raises
1848        ------
1849        ValueError
1850            If gridding fails after the specified number of attempts.
1851        """
1852        attempt_i = 0
1853        while attempt_i < attempts:
1854            attempt_i += 1
1855            data = self._grid_data(data)
1856
1857            if self.check_if_grid(data):
1858                return data
1859
1860        if not self.check_if_grid(data):
1861            raise ValueError(
1862                "Gridding failed after "
1863                + str(attempt_i)
1864                + " attempts. Please check the data."
1865            )
1866        else:
1867            return data
1868
1869    def _grid_data(self, data):
1870        """Internal method to grid the data in the mz dimension.
1871
1872        Notes
1873        -----
1874        This method is called by the grid_data method and should not be called directly.
1875        It will attempt to grid the data in the mz dimension by creating a grid of mz values based on the minimum mz difference within each scan,
1876        but it does not check if the data is already gridded or if the gridding is successful.
1877
1878        Parameters
1879        ----------
1880        data : pd.DataFrame or pl.DataFrame
1881            The input data to grid.
1882
1883        Returns
1884        -------
1885        pd.DataFrame or pl.DataFrame
1886            The data after attempting to grid it in the mz dimension.
1887        """
1888        # Calculate the difference between consecutive mz values in a single scan for grid spacing
1889        data_w = data.copy().reset_index(drop=True)
1890        data_w["mz_diff"] = np.abs(data_w["mz"].diff())
1891        mz_diff_min = data_w.groupby("scan")["mz_diff"].min().min() * 0.99999
1892
1893        # Need high intensity mz values first so they are parents in the output pairs stack
1894        dat_mz = data_w[["mz", "intensity"]].sort_values(
1895            by=["intensity"], ascending=False
1896        )
1897        dat_mz = dat_mz[["mz"]].drop_duplicates().reset_index(drop=True).copy()
1898
1899        # Construct KD tree
1900        tree = KDTree(dat_mz.mz.values.reshape(-1, 1))
1901        sdm = tree.sparse_distance_matrix(tree, mz_diff_min, output_type="coo_matrix")
1902        sdm = sparse.triu(sdm, k=1)
1903        sdm.data = np.ones_like(sdm.data)
1904        distances = sdm.tocoo()
1905        pairs = np.stack((distances.row, distances.col), axis=1)
1906
1907        # Cull pairs to just get root
1908        to_drop = []
1909        while len(pairs) > 0:
1910            root_parents = np.setdiff1d(np.unique(pairs[:, 0]), np.unique(pairs[:, 1]))
1911            id_root_parents = np.isin(pairs[:, 0], root_parents)
1912            children_of_roots = np.unique(pairs[id_root_parents, 1])
1913            to_drop = np.append(to_drop, children_of_roots)
1914
1915            # Set up pairs array for next iteration by removing pairs with children or parents already dropped
1916            pairs = pairs[~np.isin(pairs[:, 1], to_drop), :]
1917            pairs = pairs[~np.isin(pairs[:, 0], to_drop), :]
1918        dat_mz = dat_mz.reset_index(drop=True).drop(index=np.array(to_drop))
1919        mz_dat_np = (
1920            dat_mz[["mz"]]
1921            .sort_values(by=["mz"])
1922            .reset_index(drop=True)
1923            .values.flatten()
1924        )
1925
1926        # Sort data by mz and recast mz to nearest value in mz_dat_np
1927        data_w = data_w.sort_values(by=["mz"]).reset_index(drop=True).copy()
1928        data_w["mz_new"] = mz_dat_np[find_closest(mz_dat_np, data_w["mz"].values)]
1929        data_w["mz_diff"] = np.abs(data_w["mz"] - data_w["mz_new"])
1930
1931        # Rename mz_new as mz; drop mz_diff; groupby scan and mz and sum intensity
1932        new_data_w = data_w.rename(columns={"mz": "mz_orig", "mz_new": "mz"}).copy()
1933        new_data_w = (
1934            new_data_w.drop(columns=["mz_diff", "mz_orig"])
1935            .groupby(["scan", "mz"])["intensity"]
1936            .sum()
1937            .reset_index()
1938        )
1939        new_data_w = (
1940            new_data_w.sort_values(by=["scan", "mz"], ascending=[True, True])
1941            .reset_index(drop=True)
1942            .copy()
1943        )
1944
1945        return new_data_w
1946
1947    def _filter_data_by_targets(self, data, target_search_dict):
1948        """Filter MS data to only include m/z and RT windows around target values.
1949        
1950        Parameters
1951        ----------
1952        data : pd.DataFrame
1953            MS data with 'mz' and 'scan_time' columns
1954        target_search_dict : dict
1955            Dictionary with target_mz_list, target_rt_list, mz_tolerance_ppm, rt_tolerance
1956            
1957        Returns
1958        -------
1959        pd.DataFrame
1960            Filtered data containing only points within target windows
1961        """
1962        target_mz_list = target_search_dict['target_mz_list']
1963        target_rt_list = target_search_dict['target_rt_list']
1964        mz_tolerance_ppm = target_search_dict['mz_tolerance_ppm']
1965        rt_tolerance = target_search_dict['rt_tolerance']
1966        
1967        # Create a mask for data points that fall within any target window
1968        mask = np.zeros(len(data), dtype=bool)
1969        
1970        for target_mz, target_rt in zip(target_mz_list, target_rt_list):
1971            # Calculate m/z window
1972            mz_tol = target_mz * mz_tolerance_ppm / 1e6
1973            mz_min = target_mz - mz_tol
1974            mz_max = target_mz + mz_tol
1975            
1976            # Calculate RT window
1977            rt_min = target_rt - rt_tolerance
1978            rt_max = target_rt + rt_tolerance
1979            
1980            # Create mask for this target
1981            target_mask = (
1982                (data['mz'] >= mz_min) & (data['mz'] <= mz_max) &
1983                (data['scan_time'] >= rt_min) & (data['scan_time'] <= rt_max)
1984            )
1985            
1986            # Combine with overall mask
1987            mask |= target_mask
1988        
1989        return data[mask].reset_index(drop=True)
1990    
1991    def find_mass_features_ph(self, ms_level=1, grid=True, targeted_search=False, target_search_dict=None, mf_type="untargeted", accumulate_features=False):
1992        """Find mass features within an LCMSBase object using persistent homology.
1993
1994        Assigns the mass_features attribute to the object (a dictionary of LCMSMassFeature objects, keyed by mass feature id)
1995
1996        Parameters
1997        ----------
1998        ms_level : int, optional
1999            The MS level to use. Default is 1.
2000        grid : bool, optional
2001            If True, will regrid the data before running the persistent homology calculations (after checking if the data is gridded). Default is True.
2002        targeted_search : bool, optional
2003            If True, perform targeted search mode. Default is False.
2004        target_search_dict : dict or None, optional
2005            Dictionary with target parameters for targeted search. Default is None.
2006        mf_type : str, optional
2007            Type label for the mass features. Default is "untargeted".
2008        accumulate_features : bool, optional
2009            If True, add to existing features rather than replacing them. Default is False.
2010
2011        Raises
2012        ------
2013        ValueError
2014            If no MS level data is found on the object.
2015            If data is not gridded and grid is False.
2016
2017        Returns
2018        -------
2019        None, but assigns the mass_features attribute to the object.
2020
2021        Notes
2022        -----
2023        This function has been adapted from the original implementation in the Deimos package: https://github.com/pnnl/deimos
2024        """
2025        # Check that ms_level is a key in self._ms_uprocessed
2026        if ms_level not in self._ms_unprocessed.keys():
2027            raise ValueError(
2028                "No MS level "
2029                + str(ms_level)
2030                + " data found, did you instantiate with parser specific to MS level?"
2031            )
2032
2033        # Get ms data
2034        data = self._ms_unprocessed[ms_level].copy()
2035
2036        # Drop rows with missing intensity values and reset index
2037        data = data.dropna(subset=["intensity"]).reset_index(drop=True)
2038        
2039        # Add scan_time for filtering if in targeted mode
2040        if targeted_search:
2041            data = data.merge(self.scan_df[["scan", "scan_time"]], on="scan", how="left")
2042
2043        # Threshold data (bypass thresholds in targeted mode)
2044        dims = ["mz", "scan_time"]
2045        if targeted_search:
2046            # In targeted mode, bypass intensity and persistence thresholds
2047            threshold = 0
2048            persistence_threshold = 0
2049            # Filter data to only target windows
2050            data_thres = self._filter_data_by_targets(data, target_search_dict)
2051            if len(data_thres) == 0:
2052                if self.parameters.lc_ms.verbose_processing:
2053                    print("No data found in target windows")
2054                self.mass_features = {}
2055                return
2056        else:
2057            threshold = self.parameters.lc_ms.ph_inten_min_rel * data.intensity.max()
2058            persistence_threshold = (
2059                self.parameters.lc_ms.ph_persis_min_rel * data.intensity.max()
2060            )
2061            data_thres = data[data["intensity"] > threshold].reset_index(drop=True).copy()
2062
2063        # Check if gridded, if not, grid
2064        gridded_mz = self.check_if_grid(data_thres)
2065        if gridded_mz is False:
2066            if grid is False:
2067                raise ValueError(
2068                    "Data are not gridded in mz dimension, try reprocessing with a different params or grid data before running this function"
2069                )
2070            else:
2071                data_thres = self.grid_data(data_thres)
2072
2073        # Add scan_time (skip if already present from targeted mode)
2074        if 'scan_time' not in data_thres.columns:
2075            data_thres = data_thres.merge(self.scan_df[["scan", "scan_time"]], on="scan")
2076        # Process in chunks if required
2077        if len(data_thres) > 10000:
2078            return self._find_mass_features_ph_partition(
2079                data_thres, dims, persistence_threshold, mf_type, accumulate_features
2080            )
2081        else:
2082            # Process all at once
2083            return self._find_mass_features_ph_single(
2084                data_thres, dims, persistence_threshold, mf_type, accumulate_features
2085            )
2086            return self._find_mass_features_ph_single(
2087                data_thres, dims, persistence_threshold, mf_type
2088            )
2089
2090    def _find_mass_features_ph_single(self, data_thres, dims, persistence_threshold, mf_type="untargeted", accumulate_features=False):
2091        """Process all data at once (original logic)."""
2092        # Build factors
2093        factors = {
2094            dim: pd.factorize(data_thres[dim], sort=True)[1].astype(np.float32)
2095            for dim in dims
2096        }
2097
2098        # Build indexes
2099        index = {
2100            dim: np.searchsorted(factors[dim], data_thres[dim]).astype(np.float32)
2101            for dim in factors
2102        }
2103
2104        # Smooth and process
2105        mass_features_df = self._process_partition_ph(
2106            data_thres, index, dims, persistence_threshold
2107        )
2108
2109        # Roll up within chunk to remove duplicates
2110        mass_features_df = self.roll_up_dataframe(
2111            df=mass_features_df,
2112            sort_by="persistence",
2113            dims=["mz", "scan_time"],
2114            tol=[
2115                self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2116                self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2117            ],
2118            relative=[True, False],
2119        )
2120        mass_features_df = mass_features_df.reset_index(drop=True)
2121
2122        # Populate mass_features attribute
2123        self._populate_mass_features(mass_features_df, mf_type, accumulate_features)
2124
2125    def _find_mass_features_ph_partition(self, data_thres, dims, persistence_threshold, mf_type="untargeted", accumulate_features=False):
2126        """Partition the persistent homology mass feature detection for large datasets.
2127
2128        This method splits the input data into overlapping scan partitions, processes each partition to detect mass features
2129        using persistent homology, rolls up duplicates within and across partitions, and populates the mass_features attribute.
2130
2131        Parameters
2132        ----------
2133        data_thres : pd.DataFrame
2134            The thresholded input data containing mass spectrometry information.
2135        dims : list
2136            List of dimension names (e.g., ["mz", "scan_time"]) used for feature detection.
2137        persistence_threshold : float
2138            Minimum persistence value required for a detected mass feature to be retained.
2139        mf_type : str, optional
2140            Type label for the mass features. Default is "untargeted".
2141        accumulate_features : bool, optional
2142            If True, add to existing features rather than replacing them. Default is False.
2143
2144        Returns
2145        -------
2146        None
2147            Populates the mass_features attribute of the object with detected mass features.
2148        """
2149        all_mass_features = []
2150
2151        # Split scans into partitions
2152        unique_scans = sorted(data_thres["scan"].unique())
2153        unique_scans_n = len(unique_scans)
2154
2155        # Calculate partition size in scans based on goal
2156        partition_size_goal = 5000
2157        scans_per_partition = max(
2158            1, partition_size_goal // (len(data_thres) // unique_scans_n)
2159        )
2160        if scans_per_partition == 0:
2161            scans_per_partition = 1
2162
2163        # Make partitions based on scans, with overlapping in partitioned scans
2164        scan_overlap = 4
2165        partition_scans = []
2166        for i in range(0, unique_scans_n, scans_per_partition):
2167            start_idx = max(0, i - scan_overlap)
2168            end_idx = min(
2169                unique_scans_n - 1, i + scans_per_partition - 1 + scan_overlap
2170            )
2171            scans_group = [int(s) for s in unique_scans[start_idx : end_idx + 1]]
2172            partition_scans.append(scans_group)
2173
2174        # Set index to scan for faster filtering
2175        data_thres = data_thres.set_index("scan")
2176        for scans in partition_scans:
2177            # Determine start and end scan for partition, with 5 scans overlap
2178            partition_data = data_thres.loc[scans].reset_index(drop=False).copy()
2179
2180            if len(partition_data) == 0:
2181                continue
2182
2183            # Build factors for this partition
2184            factors = {
2185                dim: pd.factorize(partition_data[dim], sort=True)[1].astype(np.float32)
2186                for dim in dims
2187            }
2188
2189            # Build indexes
2190            index = {
2191                dim: np.searchsorted(factors[dim], partition_data[dim]).astype(
2192                    np.float32
2193                )
2194                for dim in factors
2195            }
2196
2197            # Process partition
2198            partition_features = self._process_partition_ph(
2199                partition_data, index, dims, persistence_threshold
2200            )
2201
2202            if len(partition_features) == 0:
2203                continue
2204
2205            # Roll up within partition
2206            partition_features = self.roll_up_dataframe(
2207                df=partition_features,
2208                sort_by="persistence",
2209                dims=["mz", "scan_time"],
2210                tol=[
2211                    self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2212                    self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2213                ],
2214                relative=[True, False],
2215            )
2216            partition_features = partition_features.reset_index(drop=True)
2217
2218            if len(partition_features) > 0:
2219                all_mass_features.append(partition_features)
2220
2221        # Combine results from all partitions
2222        if all_mass_features:
2223            combined_features = pd.concat(all_mass_features, ignore_index=True)
2224
2225            # Sort by persistence
2226            combined_features = combined_features.sort_values(
2227                by="persistence", ascending=False
2228            ).reset_index(drop=True)
2229
2230            # Remove duplicates from overlapping regions
2231            combined_features = self.roll_up_dataframe(
2232                df=combined_features,
2233                sort_by="persistence",
2234                dims=["mz", "scan_time"],
2235                tol=[
2236                    self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2237                    self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2238                ],
2239                relative=[True, False],
2240            )
2241
2242            # resort by persistence and reset index
2243            combined_features = combined_features.reset_index(drop=True)
2244
2245            # Populate mass_features attribute
2246            self._populate_mass_features(combined_features, mf_type, accumulate_features)
2247        else:
2248            self.mass_features = {}
2249
2250    def _process_partition_ph(self, partition_data, index, dims, persistence_threshold):
2251        """Process a single partition with persistent homology."""
2252        # Smooth data
2253        iterations = self.parameters.lc_ms.ph_smooth_it
2254        smooth_radius = [
2255            self.parameters.lc_ms.ph_smooth_radius_mz,
2256            self.parameters.lc_ms.ph_smooth_radius_scan,
2257        ]
2258
2259        index_array = np.vstack([index[dim] for dim in dims]).T
2260        V = partition_data["intensity"].values
2261        resid = np.inf
2262
2263        for i in range(iterations):
2264            # Previous iteration
2265            V_prev = V.copy()
2266            resid_prev = resid
2267            V = self.sparse_mean_filter(index_array, V, radius=smooth_radius)
2268
2269            # Calculate residual with previous iteration
2270            resid = np.sqrt(np.mean(np.square(V - V_prev)))
2271
2272            # Evaluate convergence
2273            if i > 0:
2274                # Percent change in residual
2275                test = np.abs(resid - resid_prev) / resid_prev
2276
2277                # Exit criteria
2278                if test <= 0:
2279                    break
2280
2281        # Overwrite values
2282        partition_data = partition_data.copy()
2283        partition_data["intensity"] = V
2284
2285        # Use persistent homology to find regions of interest
2286        pidx, pers = self.sparse_upper_star(index_array, V)
2287        pidx = pidx[pers > 1]
2288        pers = pers[pers > 1]
2289
2290        if len(pidx) == 0:
2291            return pd.DataFrame()
2292
2293        # Get peaks
2294        peaks = partition_data.iloc[pidx, :].reset_index(drop=True)
2295
2296        # Add persistence column
2297        peaks["persistence"] = pers
2298        mass_features = peaks.sort_values(
2299            by="persistence", ascending=False
2300        ).reset_index(drop=True)
2301
2302        # Filter by persistence threshold
2303        mass_features = mass_features.loc[
2304            mass_features["persistence"] > persistence_threshold, :
2305        ].reset_index(drop=True)
2306
2307        return mass_features
2308
2309    def _populate_mass_features(self, mass_features_df, mf_type="untargeted", accumulate_features=False):
2310        """Populate the mass_features attribute from a DataFrame.
2311
2312        Parameters
2313        ----------
2314        mass_features_df : pd.DataFrame
2315            DataFrame containing mass feature information.
2316            Note that the order of this DataFrame will determine the order of mass features in the mass_features attribute.
2317        mf_type : str, optional
2318            Type label for the mass features. Default is "untargeted".
2319        accumulate_features : bool, optional
2320            If True, new features will be added to existing features rather than replacing them.
2321            Mass feature IDs will be offset to avoid conflicts. Default is False.
2322
2323        Returns
2324        -------
2325        None, but assigns or updates the mass_features attribute to the object.
2326        """
2327        # Rename scan column to apex_scan
2328        mass_features_df = mass_features_df.rename(
2329            columns={"scan": "apex_scan", "scan_time": "retention_time"}
2330        )
2331
2332        # Initialize or preserve existing mass_features attribute
2333        if accumulate_features and self.mass_features is not None and len(self.mass_features) > 0:
2334            # Find the maximum existing ID to offset new IDs and avoid conflicts
2335            id_offset = max(self.mass_features.keys()) + 1
2336            initial_count = len(self.mass_features)
2337        else:
2338            # Replace mode (default/backwards compatible)
2339            self.mass_features = {}
2340            id_offset = 0
2341            initial_count = 0
2342        
2343        # Add new mass features
2344        for idx, row in enumerate(mass_features_df.itertuples()):
2345            row_dict = mass_features_df.iloc[row.Index].to_dict()
2346            lcms_feature = LCMSMassFeature(self, **row_dict)
2347            lcms_feature.type = mf_type
2348            # Use sequential ID starting from id_offset to avoid conflicts with existing features
2349            new_id = idx + id_offset
2350            lcms_feature._id = new_id  # Update the internal ID
2351            self.mass_features[new_id] = lcms_feature
2352
2353        if self.parameters.lc_ms.verbose_processing:
2354            if accumulate_features and initial_count > 0:
2355                print(f"Found {len(mass_features_df)} new mass features (total: {len(self.mass_features)})")
2356            else:
2357                print("Found " + str(len(mass_features_df)) + " initial mass features")
2358
2359    def find_mass_features_ph_centroid(self, ms_level=1, targeted_search=False, target_search_dict=None, mf_type="untargeted", accumulate_features=False):
2360        """Find mass features within an LCMSBase object using persistent homology-type approach but with centroided data.
2361
2362        Parameters
2363        ----------
2364        ms_level : int, optional
2365            The MS level to use. Default is 1.
2366        targeted_search : bool, optional
2367            If True, perform targeted search mode. Default is False.
2368        target_search_dict : dict or None, optional
2369            Dictionary with target parameters for targeted search. Default is None.
2370        mf_type : str, optional
2371            Type label for the mass features. Default is "untargeted".
2372        accumulate_features : bool, optional
2373            If True, add to existing features rather than replacing them. Default is False.
2374
2375        Raises
2376        ------
2377        ValueError
2378            If no MS level data is found on the object.
2379
2380        Returns
2381        -------
2382        None, but assigns the mass_features attribute to the object.
2383        """
2384        # Check that ms_level is a key in self._ms_uprocessed
2385        if ms_level not in self._ms_unprocessed.keys():
2386            raise ValueError(
2387                "No MS level "
2388                + str(ms_level)
2389                + " data found, did you instantiate with parser specific to MS level?"
2390            )
2391
2392        # Work with reference instead of copy
2393        data = self._ms_unprocessed[ms_level]
2394
2395        # Merge with scan data first (needed for filtering in targeted mode)
2396        scan_subset = self.scan_df[["scan", "scan_time"]]
2397        data_with_time = data.merge(scan_subset, on="scan", how="inner")
2398        
2399        # Calculate threshold and filter (bypass in targeted mode)
2400        if targeted_search:
2401            # In targeted mode, bypass intensity threshold
2402            threshold = 0
2403            valid_mask = data_with_time["intensity"].notna()
2404            required_cols = ["mz", "intensity", "scan", "scan_time"]
2405            data_thres = data_with_time.loc[valid_mask, required_cols].copy()
2406            
2407            # Filter to target windows
2408            data_thres = self._filter_data_by_targets(data_thres, target_search_dict)
2409            
2410            if len(data_thres) == 0:
2411                if self.parameters.lc_ms.verbose_processing:
2412                    print("No data found in target windows")
2413                self.mass_features = {}
2414                return
2415        else:
2416            # Normal mode with threshold
2417            max_intensity = data_with_time["intensity"].max()
2418            threshold = self.parameters.lc_ms.ph_inten_min_rel * max_intensity
2419            valid_mask = data_with_time["intensity"].notna() & (data_with_time["intensity"] > threshold)
2420            required_cols = ["mz", "intensity", "scan", "scan_time"]
2421            data_thres = data_with_time.loc[valid_mask, required_cols].copy()
2422        
2423        data_thres["persistence"] = data_thres["intensity"]
2424        mf_df = data_thres
2425        del data_thres, scan_subset, data_with_time
2426
2427        # Order by scan_time and then mz to ensure features near in rt are processed together
2428        # It's ok that different scans are in different partitions; we will roll up later
2429        mf_df = mf_df.sort_values(
2430            by=["scan_time", "mz"], ascending=[True, True]
2431        ).reset_index(drop=True)
2432        partition_size = 10000
2433        partitions = [
2434            mf_df.iloc[i : i + partition_size].reset_index(drop=True)
2435            for i in range(0, len(mf_df), partition_size)
2436        ]
2437        del mf_df
2438
2439        # Run roll_up_dataframe on each partition
2440        rolled_partitions = []
2441        for part in partitions:
2442            rolled = self.roll_up_dataframe(
2443                df=part,
2444                sort_by="persistence",
2445                dims=["mz", "scan_time"],
2446                tol=[
2447                    self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2448                    self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2449                ],
2450                relative=[True, False],
2451            )
2452            rolled_partitions.append(rolled)
2453        del partitions
2454
2455        # Run roll_up_dataframe on the rolled_up partitions to merge features near partition boundaries
2456
2457        # Combine results and run a final roll-up to merge features near partition boundaries
2458        mf_df_final = pd.concat(rolled_partitions, ignore_index=True)
2459        del rolled_partitions
2460
2461        # Reorder by persistence before final roll-up
2462        mf_df_final = mf_df_final.sort_values(
2463            by="persistence", ascending=False
2464        ).reset_index(drop=True)
2465
2466        mf_df_final = self.roll_up_dataframe(
2467            df=mf_df_final,
2468            sort_by="persistence",
2469            dims=["mz", "scan_time"],
2470            tol=[
2471                self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2472                self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2473            ],
2474            relative=[True, False],
2475        )
2476        # reset index
2477        mf_df_final = mf_df_final.reset_index(drop=True)
2478
2479        # Combine rename and sort operations
2480        mass_features = (
2481            mf_df_final.rename(
2482                columns={"scan": "apex_scan", "scan_time": "retention_time"}
2483            )
2484            .sort_values(by="persistence", ascending=False)
2485            .reset_index(drop=True)
2486        )
2487        del mf_df_final  # Free memory
2488
2489        # Order by persistence and reset index
2490        mass_features = mass_features.sort_values(
2491            by="persistence", ascending=False
2492        ).reset_index(drop=True)
2493
2494        self.mass_features = {}
2495        for idx, row in mass_features.iterrows():
2496            row_dict = row.to_dict()
2497            lcms_feature = LCMSMassFeature(self, **row_dict)
2498            lcms_feature.type = mf_type
2499            self.mass_features[lcms_feature.id] = lcms_feature
2500
2501        if self.parameters.lc_ms.verbose_processing:
2502            print("Found " + str(len(mass_features)) + " initial mass features")
2503    
2504    def cluster_mass_features(self, drop_children=True, sort_by="persistence"):
2505        """Cluster mass features
2506
2507        Based on their proximity in the mz and scan_time dimensions, priorizies the mass features with the highest persistence.
2508
2509        Parameters
2510        ----------
2511        drop_children : bool, optional
2512            Whether to drop the mass features that are not cluster parents. Default is True.
2513        sort_by : str, optional
2514            The column to sort the mass features by, this will determine which mass features get rolled up into a parent mass feature. Default is "persistence".
2515
2516        Raises
2517        ------
2518        ValueError
2519            If no mass features are found.
2520            If too many mass features are found.
2521
2522        Returns
2523        -------
2524        None if drop_children is True, otherwise returns a list of mass feature ids that are not cluster parents.
2525        """
2526        if self.mass_features is None:
2527            raise ValueError("No mass features found, run find_mass_features() first")
2528        if len(self.mass_features) > 400000:
2529            raise ValueError(
2530                "Too many mass features of interest found, run find_mass_features() with a higher intensity threshold"
2531            )
2532        dims = ["mz", "scan_time"]
2533        mf_df_og = self.mass_features_to_df()
2534        mf_df = mf_df_og.copy()
2535
2536        tol = [
2537            self.parameters.lc_ms.mass_feature_cluster_mz_tolerance_rel,
2538            self.parameters.lc_ms.mass_feature_cluster_rt_tolerance,
2539        ]  # mz, in relative; scan_time in minutes
2540        relative = [True, False]
2541
2542        # Roll up mass features based on their proximity in the declared dimensions
2543        mf_df_new = self.roll_up_dataframe(
2544            df=mf_df, sort_by=sort_by, dims=dims, tol=tol, relative=relative
2545        )
2546
2547        mf_df["cluster_parent"] = np.where(
2548            np.isin(mf_df.index, mf_df_new.index), True, False
2549        )
2550
2551        # get mass feature ids of features that are not cluster parents
2552        cluster_daughters = mf_df[~mf_df["cluster_parent"]].index.values
2553        if drop_children is True:
2554            # Drop mass features that are not cluster parents from self
2555            self.mass_features = {
2556                k: v
2557                for k, v in self.mass_features.items()
2558                if k not in cluster_daughters
2559            }
2560        else:
2561            return cluster_daughters
2562
2563
2564class LCMSCollectionCalculations:
2565    """Methods for performing calculations related to LCMSCollection objects.
2566
2567    Notes
2568    -----
2569    This class is intended as a mixin for the LCMSCollection class.
2570    """
2571
2572    @staticmethod
2573    def _plot_multiple_eics(ax, cluster_mfs, induced_cluster_mfs, rep_sample_id, rep_mf_id,
2574                           median_rt, eic_buffer_time, plot_smoothed=False, 
2575                           plot_datapoints=False, label_samples=False, lcms_collection=None):
2576        """Internal method to plot multiple EICs from different samples on a given axis.
2577        
2578        Parameters
2579        ----------
2580        ax : matplotlib.axes.Axes
2581            The axis to plot on.
2582        cluster_mfs : pd.DataFrame
2583            DataFrame containing cluster mass features (non-induced).
2584        induced_cluster_mfs : pd.DataFrame or None
2585            DataFrame containing induced (gap-filled) mass features.
2586        rep_sample_id : int
2587            Sample ID of the representative mass feature.
2588        rep_mf_id : int
2589            Mass feature ID of the representative mass feature.
2590        median_rt : float
2591            Median retention time for the cluster.
2592        eic_buffer_time : float
2593            Time buffer around the peak (minutes).
2594        plot_smoothed : bool, optional
2595            If True, plot smoothed EICs. Default is False.
2596        plot_datapoints : bool, optional
2597            If True, plot EIC datapoints. Default is False.
2598        label_samples : bool, optional
2599            If True, label each sample individually. Default is False.
2600        lcms_collection : LCMSCollection, optional
2601            The parent collection object for accessing samples. Required.
2602        """
2603        ax.set_title("EICs from all samples", loc="left")
2604        
2605        # Track if we've added labels for legend (to avoid duplicates)
2606        rep_labeled = False
2607        regular_labeled = False
2608        induced_labeled = False
2609        
2610        # Plot regular (non-induced) mass features
2611        for _, row in cluster_mfs.iterrows():
2612            sample_id = int(row['sample_id'])
2613            mf_id = row['mf_id']
2614            sample = lcms_collection[sample_id]
2615            sample_name = row['sample_name']
2616            
2617            # Get EIC using eic_mz column from dataframe
2618            eic_mz = row.get('_eic_mz')
2619            if eic_mz is not None and not pd.isna(eic_mz) and hasattr(sample, 'eics') and sample.eics:
2620                eic_data = sample.eics.get(eic_mz)
2621            else:
2622                eic_data = None
2623            
2624            if eic_data:
2625                # Determine line style and width
2626                if sample_id == rep_sample_id and mf_id == rep_mf_id:
2627                    # Representative feature - bold line
2628                    linewidth = 2.5
2629                    alpha = 1.0
2630                    color = 'tab:blue'
2631                    if label_samples:
2632                        label = f"{sample_name} (representative)"
2633                    else:
2634                        label = "Representative" if not rep_labeled else None
2635                        rep_labeled = True
2636                else:
2637                    # Other features - thinner line
2638                    linewidth = 1.0
2639                    alpha = 0.5
2640                    color = 'tab:blue'
2641                    if label_samples:
2642                        label = sample_name
2643                    else:
2644                        label = "Regular features" if not regular_labeled else None
2645                        regular_labeled = True
2646                
2647                ax.plot(
2648                    eic_data.time,
2649                    eic_data.eic,
2650                    c=color,
2651                    linewidth=linewidth,
2652                    alpha=alpha,
2653                    linestyle='-',
2654                    label=label
2655                )
2656                
2657                if plot_datapoints:
2658                    ax.scatter(
2659                        eic_data.time,
2660                        eic_data.eic,
2661                        c=color,
2662                        alpha=alpha,
2663                        s=10
2664                    )
2665                
2666                if plot_smoothed and hasattr(eic_data, 'eic_smoothed'):
2667                    ax.plot(
2668                        eic_data.time,
2669                        eic_data.eic_smoothed,
2670                        c=color,
2671                        linestyle='--',
2672                        alpha=alpha * 0.8,
2673                        linewidth=linewidth * 0.8
2674                    )
2675        
2676        # Plot induced (gap-filled) mass features if available
2677        if induced_cluster_mfs is not None and not induced_cluster_mfs.empty:
2678            for _, row in induced_cluster_mfs.iterrows():
2679                sample_id = int(row['sample_id'])
2680                mf_id = row['mf_id']
2681                sample = lcms_collection[sample_id]
2682                sample_name = row['sample_name']
2683                
2684                # Get EIC using eic_mz column from dataframe
2685                eic_mz = row.get('_eic_mz')
2686                if eic_mz is not None and not pd.isna(eic_mz) and hasattr(sample, 'eics') and sample.eics:
2687                    eic_data = sample.eics.get(eic_mz)
2688                else:
2689                    eic_data = None
2690                
2691                if eic_data:
2692                    # Induced features - even thinner line
2693                    linewidth = 0.5
2694                    alpha = 0.4
2695                    color = 'tab:orange'
2696                    
2697                    if label_samples:
2698                        label = f"{sample_name} (induced)"
2699                    else:
2700                        label = "Gap-filled features" if not induced_labeled else None
2701                        induced_labeled = True
2702                    
2703                    ax.plot(
2704                        eic_data.time,
2705                        eic_data.eic,
2706                        c=color,
2707                        linewidth=linewidth,
2708                        alpha=alpha,
2709                        linestyle='-',
2710                        label=label
2711                    )
2712                    
2713                    if plot_datapoints:
2714                        ax.scatter(
2715                            eic_data.time,
2716                            eic_data.eic,
2717                            c=color,
2718                            alpha=alpha,
2719                            s=5
2720                        )
2721                    
2722                    if plot_smoothed and hasattr(eic_data, 'eic_smoothed'):
2723                        ax.plot(
2724                            eic_data.time,
2725                            eic_data.eic_smoothed,
2726                            c=color,
2727                            linestyle='--',
2728                            alpha=alpha * 0.8,
2729                            linewidth=linewidth * 0.8
2730                        )
2731        
2732        # Add vertical line at median RT
2733        ax.axvline(
2734            x=median_rt,
2735            color='k',
2736            linestyle='--',
2737            alpha=0.7,
2738            label='Median RT'
2739        )
2740        
2741        ax.set_ylabel("Intensity")
2742        ax.set_xlabel("Time (minutes)")
2743        ax.set_xlim(
2744            median_rt - eic_buffer_time,
2745            median_rt + eic_buffer_time,
2746        )
2747        ax.legend(loc='upper left', fontsize=8)
2748        ax.yaxis.get_major_formatter().set_useOffset(False)
2749
2750    def clean_sparse_matrix(self, sparse_matrix):
2751        """Clean a sparse matrix by removing duplicates and sorting.
2752
2753        Parameters
2754        ----------
2755        sparse_matrix : :obj:`~numpy.array`
2756            A sparse matrix to clean.
2757
2758        Returns
2759        -------
2760        :obj:`~numpy.array`
2761            A cleaned sparse matrix.
2762        """
2763        for match in sparse_matrix:
2764            match.sort()
2765        sparse_matrix.sort()
2766        dereplicated_sparse_matrix = np.unique(sparse_matrix, axis=0)
2767        return dereplicated_sparse_matrix
2768
2769    def match_mfs(self, mf_c, mf_i):
2770        """Match mass features between two LCMS objects.
2771
2772        Parameters
2773        ----------
2774        mf_c : :obj:`~pandas.DataFrame`
2775            The mass features to match against.
2776        mf_i : :obj:`~pandas.DataFrame`
2777            The mass features to match.
2778
2779        Returns
2780        -------
2781        :obj:`~pandas.DataFrame`
2782            The matched mass features from mf_c.
2783        :obj:`~pandas.DataFrame`
2784            The matched mass features from mf_i.
2785
2786        Notes
2787        -----
2788        This function has been adapted from the original implementation in the Deimos package:
2789        https://github.com/pnnl/deimos
2790        """
2791        if mf_c is None or mf_i is None or len(mf_c.index) < 1 or len(mf_i.index) < 1:
2792            return None, None
2793
2794        # Prepare dataframes
2795        mf_c = mf_c.copy()
2796        mf_c["id_i"] = 0
2797        mf_i = mf_i.copy()
2798        mf_i["id_i"] = 1
2799
2800        # Set dimensions for matching
2801        dims = ["mz", "scan_time"]
2802        relative = [True, False]
2803        mz_tol = self.parameters.lcms_collection.alignment_mz_tol_ppm * 1e-6
2804        rt_tol = self.parameters.lcms_collection.alignment_rt_tol
2805        tol = [mz_tol, rt_tol]
2806
2807        # Compute inter-feature distances
2808        idx = []
2809        for i, f in enumerate(dims):
2810            # vectors
2811            v1 = mf_c[f].values.reshape(-1, 1)
2812            v2 = mf_i[f].values.reshape(-1, 1)
2813
2814            # Distances
2815            d = scipy.spatial.distance.cdist(v1, v2)
2816
2817            if relative[i] is True:
2818                # Divisor
2819                basis = np.repeat(v1, v2.shape[0], axis=1)
2820                fix = np.repeat(v2, v1.shape[0], axis=1).T
2821                basis = np.where(basis == 0, fix, basis)
2822
2823                # Divide
2824                d = np.divide(d, basis, out=np.zeros_like(basis), where=basis != 0)
2825
2826            # Check tol
2827            idx.append(d <= tol[i])
2828
2829        # Stack truth arrays
2830        idx = np.prod(np.dstack(idx), axis=-1, dtype=bool)
2831
2832        # Compute normalized 3d distance
2833        v1 = mf_c[dims].values / tol
2834        v2 = mf_i[dims].values / tol
2835        dist3d = scipy.spatial.distance.cdist(v1, v2, "cityblock")
2836        
2837        # Separate features within tolerance from those outside
2838        # Features outside tolerance should be inf, features within tolerance keep their distance
2839        # Use idx mask: True for within tolerance, False for outside
2840        dist3d_within_tol = np.where(idx, dist3d, np.inf)
2841
2842        # Normalize to 0-1 (only affects within-tolerance distances)
2843        mx = np.max(dist3d_within_tol[idx]) if np.sum(idx) > 0 else 0
2844        if mx > 0:
2845            # Lower distance is better - normalize only the within-tolerance values
2846            dist3d_within_tol = np.where(idx, dist3d_within_tol / mx, np.inf)
2847        else:
2848            # All matches are perfect (distance=0), assign tiny value to within-tolerance pairs
2849            dist3d_within_tol = np.where(idx, 1e-10, np.inf)
2850        
2851        # Use the masked distance matrix
2852        dist3d = dist3d_within_tol
2853
2854        # Min over dims
2855        mincols = np.min(dist3d, axis=0, keepdims=True)
2856
2857        # Zero out mincols over dims
2858        dist3d[dist3d != mincols] = np.inf
2859
2860        # Min over clusters
2861        minrows = np.min(dist3d, axis=1, keepdims=True)
2862
2863        # Where max and nonzero
2864        ii, jj = np.where((dist3d == minrows) & (dist3d < np.inf))
2865
2866        # Reorder
2867        mf_c = mf_c.iloc[ii]
2868        mf_i = mf_i.iloc[jj]
2869
2870        if len(mf_c.index) < 1 or len(mf_i.index) < 1:
2871            return None, None
2872
2873        return mf_c, mf_i
2874
2875    def fit_rts(self, a, b, align="scan_time", **kwargs):
2876        """
2877        Fit a support vector regressor to matched features.
2878
2879        Parameters
2880        ----------
2881        a : :obj:`~pandas.DataFrame`
2882            First set of input feature coordinates and intensities; the center object and the object to align to.
2883        b : :obj:`~pandas.DataFrame`
2884            Second set of input feature coordinates and intensities; the object to align to the center object.
2885        align : str
2886            Dimension to align.
2887        kwargs
2888            Keyword arguments for support vector regressor
2889            (:class:`sklearn.svm.SVR`).
2890
2891        Returns
2892        -------
2893        :obj:`~function`
2894            An interpolation function where one can input a retention time and get the predicted retention time.
2895
2896        Notes
2897        -----
2898        This function has been adapted from the original implementation in the Deimos package:
2899        https://github.com/pnnl/deimos
2900
2901        """
2902
2903        # Uniqueify
2904        x = a[align].values
2905        y = b[align].values
2906        arr = np.vstack((x, y)).T
2907        arr = np.unique(arr, axis=0)
2908
2909        # Safety check: ensure we have data to work with
2910        if len(arr) == 0:
2911            warnings.warn("No data points available for retention time fitting. Returning identity function.")
2912            return lambda x: x
2913
2914        # Check kwargs
2915        if "kernel" in kwargs:
2916            kernel = kwargs.get("kernel")
2917        else:
2918            kernel = "linear"
2919
2920        # Construct interpolation axis
2921        newx = np.linspace(arr[:, 0].min(), arr[:, 0].max(), 1000)
2922
2923        # Linear kernel
2924        if kernel == "linear":
2925            reg = scipy.stats.linregress(x, y)
2926            newy = reg.slope * newx + reg.intercept
2927
2928        # Other kernels
2929        else:
2930            # Fit
2931            svr = SVR(**kwargs)
2932            svr.fit(arr[:, 1].reshape(-1, 1), arr[:, 0])
2933
2934            # Predict
2935            newy = svr.predict(newx.reshape(-1, 1))
2936
2937        # Pad x and y_pred with zeros to force interpolation to start at 0
2938        newx = np.concatenate(([0], newx))
2939        newy = np.concatenate(([0], newy))
2940
2941        # Pad x and y_pred with max time to force interpolation to end at max time to force interpolation to match at end max time
2942        max_time = self[0].scan_df["scan_time"].max()
2943        newx = np.concatenate((newx, [max_time]))
2944        newy = np.concatenate((newy, [max_time]))
2945
2946        # Return an interpolation function for the x and y_pred
2947        def interp(x):
2948            pred_y = np.interp(x, newx, newy)
2949            return pred_y
2950
2951        return interp
2952
2953    def get_anchor_mass_features(self, mf_df):
2954        """
2955        Get the anchor mass features from a DataFrame of mass features.
2956
2957        Parameters
2958        ----------
2959        mf_df : :obj:`~pandas.DataFrame`
2960            The mass features to filter to just the anchor mass features.
2961
2962        Returns
2963        -------
2964        :obj:`~pandas.DataFrame`
2965            The anchor mass features dataframe.
2966        """
2967        mf_df = mf_df.copy()
2968
2969        if (
2970            "deconvoluted_mass_spectra"
2971            in self.parameters.lcms_collection.mass_feature_anchor_technique
2972        ):
2973            # Drop features that are not mass_spectrum_deconvoluted_parent or are NA as mass_spectrum_deconvoluted_parent
2974            mf_df = mf_df.dropna(subset=["mass_spectrum_deconvoluted_parent"])
2975            mf_df = mf_df[mf_df["mass_spectrum_deconvoluted_parent"]]
2976
2977        if (
2978            "absolute_intensity"
2979            in self.parameters.lcms_collection.mass_feature_anchor_technique
2980        ):
2981            # Drop features that have an intensity lower than the threshold
2982            threshold = self.parameters.lcms_collection.mass_feature_anchor_absolute_intensity_threshold
2983            mf_df = mf_df[mf_df["intensity"] > threshold]
2984
2985        if (
2986            "relative_intensity"
2987            in self.parameters.lcms_collection.mass_feature_anchor_technique
2988        ):
2989            # Drop features in the lower fraction of intensities
2990            threshold_quantile = self.parameters.lcms_collection.mass_feature_anchor_relative_intensity_threshold
2991            intensity_threshold = mf_df["intensity"].quantile(threshold_quantile)
2992            mf_df = mf_df[mf_df["intensity"] >= intensity_threshold]
2993
2994        return mf_df
2995
2996    def attempt_alignment(self, matches_c, matches_i):
2997        """
2998        Check if alignment is needed for the LCMS objects in the collection.
2999        """
3000
3001        # Hold out a subset of matches_c and matches_i for spline fitting
3002        matches_c.reset_index(drop=False, inplace=True)
3003        matches_i.reset_index(drop=False, inplace=True)
3004
3005        # Check if there are enough matches to attempt alignment
3006        minimum_matches = self.parameters.lcms_collection.alignment_minimum_matches
3007        if len(matches_c) < minimum_matches:
3008            # Return False (no alignment) and identity function (returns original time) 
3009            # which isn't used but is a placeholder to avoid errors in downstream code since 
3010            # the function expects a callable to be returned
3011            return False, lambda x: x
3012
3013        # Rearrange matches_c and matches_i to be in the order of the scan_time of matches_c
3014        matches_c = matches_c.sort_values(by="scan_time")
3015        matches_i = matches_i.iloc[matches_c.index.values]
3016
3017        hold_out_fraction = self.parameters.lcms_collection.alignment_hold_out_fraction
3018        # starting with an array of length len(matches_c), select equally spaced indices to hold out
3019        idx_holdout = matches_c.index.values[
3020            np.arange(0, len(matches_c), int(1 / hold_out_fraction))
3021        ]
3022
3023        matches_c_holdout = matches_c.loc[idx_holdout].copy()
3024        matches_i_holdout = matches_i.loc[idx_holdout].copy()
3025
3026        # Remove the holdout matches from the matches_c and matches_i DataFrames and reset the index
3027        matches_c = matches_c.drop(index=idx_holdout).set_index("sample_name")
3028        matches_i = matches_i.drop(index=idx_holdout).set_index("sample_name")
3029
3030        # Reset the scan_time to the original scan_time
3031        matches_i = matches_i.copy()
3032        matches_i["scan_time"] = matches_i["scan_time_og"]
3033
3034        # Fit the retention times of the LCMS object to the center LCMS object using the matched mass features
3035        spl = self.fit_rts(matches_c, matches_i, kernel="rbf", C=1000)
3036
3037        # Check if the spline fitting improved the alignment for the holdout matches
3038        matches_i_holdout["scan_time_fit"] = spl(matches_i_holdout["scan_time"])
3039        og_diff = np.abs(
3040            matches_i_holdout["scan_time"] - matches_c_holdout["scan_time"]
3041        )
3042        fit_diff = np.abs(
3043            matches_i_holdout["scan_time_fit"] - matches_c_holdout["scan_time"]
3044        )
3045
3046        if (
3047            "fraction_improved"
3048            in self.parameters.lcms_collection.alignment_acceptance_technique
3049        ):
3050            fraction_improved = np.sum(fit_diff < og_diff) / len(og_diff)
3051            use_spline_alignment = (
3052                fraction_improved
3053                > self.parameters.lcms_collection.alignment_acceptance_fraction_improved_threshold
3054            )
3055        if (
3056            "mean_squared_error_improved"
3057            in self.parameters.lcms_collection.alignment_acceptance_technique
3058        ):
3059            mse_og = np.mean(og_diff**2)
3060            mse = np.mean(fit_diff**2)
3061            use_spline_alignment = mse < mse_og
3062            # Convert to boolean
3063            use_spline_alignment = bool(use_spline_alignment)
3064
3065        return use_spline_alignment, spl
3066
3067    def align_lcms_objects(self, overwrite=False):
3068        """
3069        Align LCMS objects in the collection.
3070
3071        Aligns the LCMS objects in the collection by aligning the retention times of the mass features in the LCMS objects.
3072        First, the mass features in the center LCMS object are matched to the mass features in the other LCMS objects,
3073        starting with the LCMS object immediately following the center LCMS object. The retention times of the LCMS objects
3074        are then fit to the center LCMS object using the matched mass features.
3075
3076        Returns
3077        -------
3078        None, but aligns the LCMS objects in the collection and sets the scan_time_aligned column in the scan_df attribute of each LCMS object.
3079
3080        Notes
3081        -----
3082        This function has been adapted from the original implementation in the Deimos package:
3083        https://github.com/pnnl/deimos
3084        """
3085       
3086        # Prepare the center LCMS object
3087        center_obj_ids = self.manifest_dataframe[
3088            self.manifest_dataframe["center"]
3089        ].collection_id.values
3090
3091        full_mf_df = self.mass_features_dataframe
3092        # re-index to sample_name for faster lookups
3093        full_mf_df = full_mf_df.reset_index().set_index("sample_name")
3094        samples_with_features = set(full_mf_df.index.get_level_values("sample_name"))
3095
3096        if "scan_time_aligned" in full_mf_df.columns and not overwrite:
3097            raise ValueError("Mass features have already been aligned")
3098
3099        def _set_scan_time_alignment_for_sample(sample_idx, use_alignment, spline):
3100            """Set scan_time_aligned for one sample using spline or identity mapping."""
3101            if use_alignment and spline is not None:
3102                self[sample_idx]._scan_info["scan_time_aligned"] = {
3103                    k: spline(v) for k, v in self[sample_idx]._scan_info["scan_time"].items()
3104                }
3105                return True
3106
3107            self[sample_idx]._scan_info["scan_time_aligned"] = self[sample_idx]._scan_info[
3108                "scan_time"
3109            ].copy()
3110            return False
3111
3112        def _get_feature_df_at_or_after(start_idx, index_step, use_alignment, spline):
3113            """Return next sample index/dataframe with features, aligning empty samples on the way."""
3114            i = start_idx
3115            while 0 <= i < len(self):
3116                sample_name = self.samples[i]
3117                if sample_name in samples_with_features:
3118                    mf_df_i = full_mf_df.loc[sample_name].copy()
3119                    mf_df_i["scan_time_og"] = mf_df_i["scan_time"]
3120                    mf_df_i = mf_df_i.reset_index(drop=False)
3121                    if use_alignment and spline is not None:
3122                        # Use previous step transform as a better matching starting point.
3123                        mf_df_i["scan_time"] = spline(mf_df_i["scan_time"])
3124                    return i, mf_df_i
3125
3126                _set_scan_time_alignment_for_sample(i, use_alignment, spline)
3127                self.rt_alignment_attempted = True
3128                i += index_step
3129
3130            return i, None
3131
3132        anchor_mf_dfs = []
3133        for center_obj_id in center_obj_ids:
3134            # Get the anchor mass features from the center LCMS object
3135            mf_df_c = full_mf_df.loc[self.samples[center_obj_id]]
3136            mf_df_c = self.get_anchor_mass_features(mf_df_c)
3137            anchor_mf_dfs.append(mf_df_c)
3138
3139            # Set scan_time_aligned to scan_time for the center LCMS object
3140            center_scan_df = self[center_obj_id].scan_df.copy()
3141            center_scan_df["scan_time_aligned"] = center_scan_df["scan_time"]
3142            self[center_obj_id].scan_df = center_scan_df
3143            
3144            # Store alignment data for center object (identity mapping)
3145            center_sample_name = self.samples[center_obj_id]
3146
3147            index_steps = (1, -1)
3148            # Run this twice, once going forward (+1 indexing) and once going backward (-1 indexing)
3149            for index_step in index_steps:
3150                # Initialize spline for propagation to samples without features
3151                spl = None
3152                use_spline_alignment = False
3153
3154                # Loop through the other LCMS objects in this direction.
3155                i, mf_df_i = _get_feature_df_at_or_after(
3156                    center_obj_id + index_step,
3157                    index_step,
3158                    use_spline_alignment,
3159                    spl,
3160                )
3161
3162                while mf_df_i is not None:
3163                    mf_df_i = self.get_anchor_mass_features(mf_df_i)
3164
3165                    # Match the mass features in the LCMS object to the anchor mass features in the center LCMS object.
3166                    matches_c, matches_i = self.match_mfs(mf_df_c, mf_df_i)
3167
3168                    if matches_c is not None:
3169                        use_spline_alignment, spl = self.attempt_alignment(
3170                            matches_c, matches_i
3171                        )
3172
3173                        # Record if we used alignment for this sample
3174                        sample_name = self.samples[i]
3175                        self._manifest_dict[sample_name]["use_rt_alignment"] = (
3176                            use_spline_alignment
3177                        )
3178
3179                        if use_spline_alignment:
3180                            # Set new retention times on scan_df for lc_obj using the spline fitting
3181                            matches_i["scan_time_fit"] = spl(matches_i["scan_time"])
3182
3183                        self.rt_aligned = _set_scan_time_alignment_for_sample(
3184                            i, use_spline_alignment, spl
3185                        )
3186                        self.rt_alignment_attempted = True
3187
3188                        i, mf_df_i = _get_feature_df_at_or_after(
3189                            i + index_step,
3190                            index_step,
3191                            use_spline_alignment,
3192                            spl,
3193                        )
3194                    else:
3195                        # If no matches are found, propagate prior alignment from this index step.
3196                        sample_name = self.samples[i]
3197                        used_previous_alignment = use_spline_alignment and spl is not None
3198                        self._manifest_dict[sample_name]["use_rt_alignment"] = (
3199                            used_previous_alignment
3200                        )
3201
3202                        self.rt_aligned = _set_scan_time_alignment_for_sample(
3203                            i, used_previous_alignment, spl
3204                        )
3205                        self.rt_alignment_attempted = True
3206
3207                        i, mf_df_i = _get_feature_df_at_or_after(
3208                            i + index_step,
3209                            index_step,
3210                            used_previous_alignment,
3211                            spl,
3212                        )
3213
3214        # Now align each batch using the center objects as anchors with the other batches
3215        mf_df_c = anchor_mf_dfs[0]
3216        for i in center_obj_ids[1:]:
3217            mf_df_i = full_mf_df.loc[self.samples[i]].copy()
3218            mf_df_i["scan_time_og"] = mf_df_i["scan_time"]
3219            mf_df_i = self.get_anchor_mass_features(mf_df_i)
3220
3221            matches_c, matches_i = self.match_mfs(mf_df_c, mf_df_i)
3222            if matches_c is not None:
3223                use_spline_alignment, spl = self.attempt_alignment(matches_c, matches_i)
3224
3225                # Record if we used alignment for this sample
3226                sample_name = self.samples[i]
3227                self._manifest_dict[sample_name]["use_rt_alignment"] = (
3228                    use_spline_alignment
3229                )
3230
3231                if use_spline_alignment:
3232                    # Set new retention times on all this object's
3233                    new_times = spl(self[i].scan_df["scan_time"])
3234                    new_scan_info = self[i].scan_df.copy()
3235                    new_scan_info["scan_time_aligned"] = new_times
3236                    self[i].scan_df = new_scan_info
3237                    
3238
3239                    # Get the batch that this object belongs to
3240                    batch = self.manifest[self.samples[i]]["batch"]
3241
3242                    for j in range(len(self)):
3243                        if self.manifest[self.samples[j]]["batch"] == batch:
3244                            if j != i:
3245                                sample_name_j = self.samples[j]
3246                                self._manifest_dict[sample_name_j]["use_rt_alignment"] = (
3247                                    use_spline_alignment
3248                                )
3249                                new_scan_info = self[j].scan_df.copy()
3250                                aligned_times = spl(self[j].scan_df["scan_time_aligned"])
3251                                new_scan_info["scan_time_aligned"] = aligned_times
3252                                self[j].scan_df = new_scan_info
3253                                
3254        # Set final mass_features_dataframe with the aligned scan_time
3255        center_sample_name = self.samples[center_obj_ids[0]]
3256        self._manifest_dict[center_sample_name]["use_rt_alignment"] = False
3257        new_scan_info = self[center_obj_ids[0]].scan_df.copy()
3258        new_scan_info["scan_time_aligned"] = new_scan_info["scan_time"]
3259
3260    def add_consensus_mass_features(self):
3261        """
3262        Create consensus mass features by clustering aligned features across samples.
3263        
3264        This method clusters mass features from all samples in the collection based on
3265        their m/z and aligned retention time proximity. Features that cluster together
3266        across samples are assigned a common cluster ID, creating consensus features
3267        that represent the same compound detected across multiple samples.
3268        
3269        The clustering process:
3270        1. Partitions features by m/z to avoid large sparse matrices and enable parallelization
3271        2. Clusters features within each partition using hierarchical clustering
3272        3. Merges partition-boundary clusters that represent the same feature
3273        4. Filters out clusters not present in minimum fraction of samples
3274        
3275        Must be run after align_lcms_objects(). Results are stored in the 
3276        mass_features_dataframe with a 'cluster' column added.
3277        
3278        Parameters
3279        ----------
3280        None
3281            Uses parameters from self.parameters.lcms_collection:
3282            - consensus_mz_tol_ppm: m/z tolerance for clustering (ppm)
3283            - consensus_rt_tol: retention time tolerance for clustering (minutes)
3284            - consensus_partition_size: target partition size for managing memory and parallelization
3285            - consensus_min_sample_fraction: minimum fraction of samples a cluster
3286              must appear in to be retained (0-1)
3287            - cores: number of CPU cores to use for parallel partition processing
3288            
3289        Returns
3290        -------
3291        None
3292            Updates self.mass_features_dataframe in place by adding 'cluster' column
3293            and filtering to retain only clusters meeting minimum sample presence.
3294            
3295        Raises
3296        ------
3297        ValueError
3298            If mass features have not been aligned (run align_lcms_objects() first).
3299            
3300        Notes
3301        -----
3302        - Partitioning prevents memory issues with large sparse distance matrices
3303        - Each partition is processed in parallel (up to cores limit)
3304        - Clusters not meeting consensus_min_sample_fraction are automatically removed
3305        - Access cluster_summary_dataframe property for summary statistics
3306        - Use fill_missing_cluster_features() for gap-filling after clustering
3307        
3308        See Also
3309        --------
3310        align_lcms_objects : Aligns retention times before consensus clustering
3311        cluster_summary_dataframe : Property that generates summary statistics for clusters
3312        fill_missing_cluster_features : Gap-fill missing features in clusters
3313        """
3314        # Get the combined mass features from all LCMS objects, keep the original index as a separate column
3315        combined_mfs = self.mass_features_dataframe.copy()
3316        combined_mfs["coll_mf_id"] = combined_mfs.index
3317
3318        # Check if the mass features have been aligned
3319        if "scan_time_aligned" not in combined_mfs.columns:
3320            raise ValueError(
3321                "Mass features have not been aligned, run align_lcms_objects() first"
3322            )
3323
3324        # Partition the mass features by mz so we can parallelize the matching before clustering
3325        from corems.chroma_peak.calc import subset as corems_subset
3326
3327        # get max mz from combined_mfs and calculate tolerance from ppm
3328        mz_tol = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
3329        n_partition_size = self.parameters.lcms_collection.consensus_partition_size
3330        lazy_partitions = corems_subset.multi_sample_partition(
3331            combined_mfs,
3332            split_on="mz",
3333            size=n_partition_size,
3334            tol=mz_tol,
3335            relative=True,
3336        )
3337
3338        # If any of lazy_partitions._counts is 2xn_partition_size, issue a warning
3339        if np.array(lazy_partitions._counts).max() > 2 * n_partition_size:
3340            warnings.warn(
3341                "Some partitions are larger than 2x the goal partition size. Consider increasing the partition or decreasing the mz_tol."
3342            )
3343
3344        # Cluster the mass features within each partition
3345        if self.parameters.lcms_collection.cores > lazy_partitions.n_partitions:
3346            cores_to_use = lazy_partitions.n_partitions
3347        else:
3348            cores_to_use = self.parameters.lcms_collection.cores
3349        # mfs_with_clusters = lazy_partitions.map(self.cluster_mass_features, processes=cores_to_use)
3350        mfs_with_clusters = lazy_partitions.map(
3351            self.cluster_mass_features_agg_cluster, processes=cores_to_use
3352        )
3353
3354        # Clean up cluster id names after partitioning
3355        new_cluster_ids = (
3356            mfs_with_clusters[["cluster", "partition_idx"]]
3357            .drop_duplicates()
3358            .reset_index(drop=True)
3359        )
3360        new_cluster_ids["cluster_unqiue"] = new_cluster_ids.index
3361        mfs_with_clusters = mfs_with_clusters.merge(
3362            new_cluster_ids, on=["cluster", "partition_idx"]
3363        )
3364        mfs_with_clusters["cluster"] = mfs_with_clusters["cluster_unqiue"]
3365        mfs_with_clusters = mfs_with_clusters.drop(columns=["cluster_unqiue"])
3366
3367        # Embed a new cluster id into the mass features dataframe and set as index
3368        mfs_with_clusters["idx"] = mfs_with_clusters.index
3369
3370        try:
3371            # Check if any clusters can be merged into a single cluster
3372            eval_dict = self.evaluate_clusters_for_repeats(mfs_with_clusters)
3373
3374            # Merge clusters identified in eval_dict
3375            while len(eval_dict["merge_these_clusters"]) > 0:
3376                list_of_clusters_to_merge = [
3377                    [x[0], x[1]] for x in eval_dict["merge_these_clusters"]
3378                ]
3379                # Convert to a dataframe with columns "new_cluster" and "cluster"
3380                df = pd.DataFrame(
3381                    np.array(list_of_clusters_to_merge), columns=["new_cluster", "cluster"]
3382                )
3383                # Drop duplicates of "child" clusters
3384                df = df.drop_duplicates("cluster", keep="first")
3385                df = df.drop_duplicates("new_cluster", keep="first")
3386                mfs_with_clusters = mfs_with_clusters.merge(df, on="cluster", how="left")
3387                mfs_with_clusters["cluster"] = mfs_with_clusters["new_cluster"].fillna(
3388                    mfs_with_clusters["cluster"]
3389                )
3390                mfs_with_clusters = mfs_with_clusters.drop(columns=["new_cluster"])
3391
3392                # Re-evaluate clusters for repeats
3393                eval_dict = self.evaluate_clusters_for_repeats(mfs_with_clusters)
3394                self.mass_features_dataframe = mfs_with_clusters
3395
3396        except:
3397            mfs_with_clusters.set_index('coll_mf_id', inplace = True)
3398            self.mass_features_dataframe = mfs_with_clusters
3399            
3400        # Filter out clusters that don't meet minimum sample fraction
3401        self._filter_clusters_by_sample_presence()
3402            
3403        # TODO KRH: Deal with isomers better? Pool them together and then split them out using samples with 2 as the template?
3404    
3405    def _filter_clusters_by_sample_presence(self):
3406        """
3407        Filter out clusters that don't meet the minimum sample fraction threshold.
3408        
3409        Removes clusters (and their associated mass features) from the mass_features_dataframe
3410        if they don't appear in at least consensus_min_sample_fraction of samples.
3411        
3412        This is called automatically at the end of add_consensus_mass_features().
3413        
3414        Returns
3415        -------
3416        None
3417            Updates self.mass_features_dataframe in place by removing clusters that don't
3418            meet the minimum sample presence threshold.
3419        """
3420        if self.mass_features_dataframe is None or len(self.mass_features_dataframe) == 0:
3421            return
3422        
3423        min_sample_fraction = self.parameters.lcms_collection.consensus_min_sample_fraction
3424        
3425        # Validate parameter
3426        if not 0 <= min_sample_fraction <= 1:
3427            raise ValueError("consensus_min_sample_fraction must be between 0 and 1")
3428        
3429        # Calculate minimum number of samples required
3430        total_samples = len(self.samples)
3431        min_samples_required = min_sample_fraction * total_samples
3432        
3433        # Count unique samples per cluster
3434        cluster_sample_counts = (
3435            self.mass_features_dataframe.groupby('cluster')['sample_id']
3436            .nunique()
3437            .reset_index(name='sample_count')
3438        )
3439        
3440        # Identify clusters to keep
3441        clusters_to_keep = cluster_sample_counts[
3442            cluster_sample_counts['sample_count'] > min_samples_required
3443        ]['cluster'].values
3444        
3445        # Filter mass features dataframe
3446        self.mass_features_dataframe = self.mass_features_dataframe[
3447            self.mass_features_dataframe['cluster'].isin(clusters_to_keep)
3448        ]
3449        
3450    def summarize_clusters(self):
3451        """
3452        Generate summary statistics for consensus mass feature clusters.
3453        
3454        Computes aggregate statistics (median, mean, std, min, max) for each cluster
3455        across all samples. Combines both regular mass features and induced mass features
3456        (from gap-filling) when available to provide complete cluster statistics.
3457        
3458        Must be run after add_consensus_mass_features() which creates the cluster assignments.
3459        Results are stored in cluster_summary_dataframe property and used by plotting methods.
3460        
3461        Parameters
3462        ----------
3463        None
3464            Operates on self.mass_features_dataframe and self.induced_mass_features_dataframe.
3465            Both must contain 'cluster' column.
3466            
3467        Returns
3468        -------
3469        :obj:`~pandas.DataFrame` or None
3470            DataFrame with one row per cluster containing summary statistics:
3471            - cluster: cluster ID
3472            - mz_{median,mean,std,max,min}: m/z statistics
3473            - scan_time_aligned_{median,mean,std,max,min}: aligned RT statistics
3474            - half_height_width_{median,mean,std,max,min}: peak width statistics
3475            - tailing_factor_{median,mean,std,max,min}: peak shape statistics
3476            - dispersity_index_{median,mean,std,max,min}: peak quality statistics
3477            - sample_id_nunique: number of unique samples containing the cluster
3478            - intensity_{max,median,mean,std,min}: intensity statistics
3479            - persistence_{max,median,mean,std,min}: persistence statistics
3480            
3481            Returns None if mass_features_dataframe is empty.
3482            
3483        Notes
3484        -----
3485        - Summary DataFrame is automatically stored in cluster_summary_dataframe property
3486        - Includes both regular and induced (gap-filled) mass features when available
3487        - Used by plotting methods: plot_consensus_mz_features, plot_mz_features_per_cluster
3488        - Sample count (sample_id_nunique) indicates cluster prevalence across samples
3489        - Filters applied by consensus_min_sample_fraction affect which clusters appear
3490        
3491        See Also
3492        --------
3493        add_consensus_mass_features : Creates clusters before summarization
3494        fill_missing_cluster_features : Creates induced mass features via gap-filling
3495        plot_consensus_mz_features : Visualizes cluster summaries
3496        plot_mz_features_per_cluster : Shows cluster size distribution
3497        """
3498        # First check if there are minimum columns in the features dataframe
3499        if len(self.mass_features_dataframe.columns) < 1:
3500            return None
3501
3502        # Combine regular and induced mass features
3503        mf_df = self.mass_features_dataframe.copy()
3504        mf_df = mf_df.reset_index(drop=False)
3505        
3506        # Check if induced mass features are available and combine them
3507        if self.induced_mass_features_dataframe is not None and len(self.induced_mass_features_dataframe) > 0:
3508            imf_df = self.induced_mass_features_dataframe.copy()
3509            imf_df = imf_df.reset_index(drop=False)
3510            # Cluster column extracted from mf_id in _prepare_lcms_mass_features_for_combination
3511            # Combine regular and induced features
3512            mf_df = pd.concat([mf_df, imf_df], axis=0)
3513            mf_df = mf_df.reset_index(drop=True)
3514        
3515        # Filter out any rows with NaN cluster values before converting to int
3516        if 'cluster' in mf_df.columns:
3517            mf_df = mf_df.dropna(subset=['cluster'])
3518            mf_df['cluster'] = mf_df['cluster'].astype(int)
3519
3520        # Build aggregation dictionary based on available columns
3521        agg_dict = {
3522            "mz": ["median", "mean", "std", "max", "min"],
3523            "scan_time_aligned": ["median", "mean", "std", "max", "min"],
3524            "sample_id": ["nunique"],
3525            "intensity": ["max", "median", "mean", "std", "min"],
3526        }
3527        
3528        # Add optional columns if they exist
3529        optional_columns = {
3530            "half_height_width": ["median", "mean", "std", "max", "min"],
3531            "tailing_factor": ["median", "mean", "std", "max", "min"],
3532            "dispersity_index": ["median", "mean", "std", "max", "min"],
3533            "persistence": ["max", "median", "mean", "std", "min"],
3534        }
3535        
3536        for col, funcs in optional_columns.items():
3537            if col in mf_df.columns:
3538                agg_dict[col] = funcs
3539
3540        summary_df = (
3541            mf_df.groupby("cluster")
3542            .agg(agg_dict)
3543            .reset_index()
3544        )
3545
3546        # Fix the column names
3547        summary_df.columns = [
3548            "_".join(col).strip()
3549            for col in summary_df.columns.values
3550            if col != "cluster"
3551        ]
3552        summary_df = summary_df.rename(columns={"cluster_": "cluster"})
3553        # Set cluster as the index for easy lookup
3554        summary_df = summary_df.set_index('cluster')
3555        return summary_df
3556
3557    def plot_mz_features_per_cluster(self, return_fig = False):
3558        """
3559        Plot the number of mass features in a cluster against how many clusters
3560        contain that number of mass features
3561
3562        Parameters
3563        -----------
3564        return_fig : boolean
3565            Indicates whether to plot composite feature map (False) or return figure object (True). Defaults to False.
3566
3567        Returns
3568        --------
3569        matplotlib.pyplot.Figure
3570            A figure displaying the frequency with which clusters contain the given number of m/z features
3571
3572        Raises
3573        ------
3574        Warning
3575            If consensus features haven't been added to the object yet
3576        """
3577
3578        if not hasattr(self, 'cluster_summary_dataframe'):
3579            raise ValueError(
3580                'cluster_summary_dataframe is not set, must run add_consensus_mass_features() first'
3581            )
3582        else:
3583            sum_data = self.cluster_summary_dataframe
3584            fig, ax = plt.subplots()
3585            sum_data.sample_id_nunique.value_counts().sort_index().plot(ax = ax, kind = 'bar')
3586            plt.xlabel('Number of mass features in a cluster')
3587            plt.ylabel('Number of clusters with this many mass features')
3588            if return_fig:
3589                plt.close(fig)
3590                return fig
3591            else:
3592                plt.show()
3593        
3594    def plot_mz_features_across_samples(self, alpha = 0.75, s = 0.005, return_fig = False):
3595        """
3596        Generate Scan Time vs m/z plot of all the mass features across all 
3597        samples in collection where intensity of color on the plot indicates
3598        density of mass features, NOT INTENSITY
3599
3600        Parameters
3601        -----------
3602        alpha :  float
3603            Desired transparency for plotted m/z features.  Defaults to 0.75.
3604        s : float
3605            Desired size of plotted m/z features. Defaults to 0.005.
3606        return_fig : boolean
3607            Indicates whether to plot composite feature map (False) or return figure object (True). Defaults to False.
3608
3609        Returns
3610        --------
3611        matplotlib.pyplot.Figure
3612            A figure displaying a scan time vs m/z scatterplot of all the m/z features identified in the collection.
3613            Parameters alpha (transparency) and s (marker size) allow the user to emphasize the density of features.
3614            Intensity of features is not represented.
3615        """
3616        df = self.mass_features_dataframe.copy()
3617        fig = plt.figure()
3618        plt.scatter(
3619            df.scan_time_aligned,
3620            df.mz,
3621            c = 'tab:gray',
3622            alpha = alpha,
3623            s = s
3624        )
3625
3626        plt.xlabel('Scan time')
3627        plt.ylabel('m/z')
3628        plt.ylim(0, np.ceil(np.max(df.mz)))
3629        plt.xlim(0, np.ceil(np.max(df.scan_time)))
3630        plt.title('All mass features, all samples')
3631        
3632        if return_fig:
3633            plt.close(fig)
3634            return fig
3635        else:
3636            plt.show()
3637
3638    def plot_consensus_mz_features(self, xb = 'xb', xt = 'xt', yb = 'yb', yt = 'yt', show_all = True, return_fig = False):
3639        """
3640        Generate Scan Time vs m/z plot of the consensus features scaled by size
3641        with option ('show_all') of leaving the individual m/z features in the figure.
3642
3643        Parameters
3644        -----------
3645        xb :  float
3646            Desired starting scan time value for the x-axis. Defaults to 0.
3647        xt : float
3648            Desired ending scan time for the x-axis. Defaults to the maximum scan time value in the provided data.
3649        yb :  float
3650            Desired starting m/z value for the y-axis. Defaults to 0.
3651        yt : float
3652            Desired ending m/z for the y-axis. Defaults to the maximum m/z value in the provided data.
3653        show_all : boolean
3654            Indicates whether to display all identified m/z features (True) or just the consensus features (False). Defaults to True.
3655        return_fig : boolean
3656            Indicates whether to plot composite feature map (False) or return figure object (True). Defaults to False.
3657
3658        Returns
3659        --------
3660        matplotlib.pyplot.Figure
3661            A scalable figure that overlays the consensus features over all the m/z features identified in the collection.
3662            Consensus features are scaled by how many m/z features are represented in the consensus. Figure can be scaled by
3663            inputting desired boundaries on the scan time (xb, xt) and m/z values (yb, yt).
3664        """
3665        df = self.cluster_summary_dataframe.copy()
3666        mfdf = self.mass_features_dataframe.copy()
3667
3668        fig = plt.figure()
3669        if show_all:
3670            plt.scatter(
3671                mfdf.scan_time_aligned,
3672                mfdf.mz,
3673                c = 'tab:gray',
3674                s = 1
3675            )
3676
3677        m = plt.scatter(
3678            df.scan_time_aligned_median,
3679            df.mz_median, 
3680            c = 'tab:orange',
3681            alpha = 0.7, 
3682            s = (df.sample_id_nunique**2)/5
3683        )
3684
3685        plt.xlabel('Scan time')
3686        plt.ylabel('m/z')
3687        
3688        if xt == 'xt':
3689            xt = np.ceil(np.max(mfdf.mz))
3690        if yt == 'yt':
3691            yt = np.ceil(np.max(mfdf.scan_time))
3692        if xb == 'xb':
3693            xb = 0
3694        if yb == 'yb':
3695            yb = 0
3696        plt.ylim(xb, xt)
3697        plt.xlim(yb, yt)
3698
3699        kw = dict(
3700            prop = 'sizes',
3701            num = max(1, int(len(df.sample_id_nunique.unique())/3)),
3702            color = 'tab:orange',
3703            alpha = 0.7,
3704            func = lambda s: np.sqrt(s*5)
3705        )
3706
3707        plt.legend(
3708            *m.legend_elements(**kw), 
3709            title = 'Features\nper cluster',
3710            bbox_to_anchor = (1.01, 0.4, 0.225, 0.5)
3711        )
3712        plt.tight_layout()
3713        plt.title('Consensus Features')
3714
3715        if return_fig:
3716            plt.close(fig)
3717            return fig
3718        else:
3719            plt.show()
3720    
3721    def plot_cluster(
3722        self,
3723        cluster_id,
3724        to_plot=["EIC", "MS1", "MS2"],
3725        return_fig=False,
3726        plot_smoothed_eic=False,
3727        plot_eic_datapoints=False,
3728        eic_buffer_time=None,
3729        label_samples=False,
3730        molecular_metadata=None,
3731        spectral_library=None,
3732    ):
3733        """
3734        Plot a consensus mass feature cluster across all samples.
3735        
3736        Similar to LCMSMassFeature.plot() but shows EICs from all samples in the cluster,
3737        highlighting the representative mass feature.
3738        
3739        Parameters
3740        ----------
3741        cluster_id : int
3742            The cluster ID to plot
3743        to_plot : list, optional
3744            List of strings specifying what to plot: "EIC", "MS1", "MS2", "MS2_mirror".
3745            Default is ["EIC", "MS1", "MS2"].
3746        return_fig : bool, optional
3747            If True, returns the figure object. Default is False.
3748        plot_smoothed_eic : bool, optional
3749            If True, plots smoothed EICs. Default is False.
3750        plot_eic_datapoints : bool, optional
3751            If True, plots EIC data points. Default is False.
3752        eic_buffer_time : float, optional
3753            Time buffer around the peak for EIC plotting (minutes).
3754            If None, uses parameter setting. Default is None.
3755        label_samples : bool, optional
3756            If True, labels each sample in the legend. Default is False.
3757        molecular_metadata : dict, optional
3758            Dictionary mapping molecular IDs to MetaboliteMetadata objects.
3759            Required for MS2_mirror plots. Default is None.
3760        spectral_library : FlashEntropySearch, optional
3761            FlashEntropy spectral library containing MS2 spectra.
3762            Required for MS2_mirror plots to retrieve library spectra. Default is None.
3763            
3764        Returns
3765        -------
3766        matplotlib.figure.Figure or None
3767            The figure object if return_fig=True, otherwise None
3768            
3769        Raises
3770        ------
3771        ValueError
3772            If cluster_id is not found or if required data is not loaded
3773        """
3774        import matplotlib.pyplot as plt
3775        
3776        # Get cluster summary for median values
3777        if cluster_id not in self.cluster_summary_dataframe.index:
3778            raise ValueError(
3779                f"Cluster {cluster_id} not found in cluster_summary_dataframe. "
3780                f"Run add_consensus_mass_features() first."
3781            )
3782        
3783        cluster_summary = self.cluster_summary_dataframe.loc[cluster_id]
3784        
3785        # Get representative mass feature info
3786        rep_info = self.get_most_representative_sample_for_cluster(cluster_id)
3787        rep_sample_id = rep_info['sample_id']
3788        rep_mf_id = rep_info['mf_id']
3789        rep_sample = self[rep_sample_id]
3790        
3791        # Check if representative mass feature is loaded
3792        if rep_mf_id not in rep_sample.mass_features:
3793            raise ValueError(
3794                f"Representative mass feature {rep_mf_id} not loaded in sample {rep_sample.sample_name}. "
3795                f"Run reload_representative_mass_features() or process_consensus_features() first."
3796            )
3797        
3798        rep_mf = rep_sample.mass_features[rep_mf_id]
3799        
3800        # Get eic buffer time
3801        if eic_buffer_time is None:
3802            eic_buffer_time = self[0].parameters.lc_ms.eic_buffer_time
3803        
3804        # Adjust to_plot based on available data
3805        if rep_mf.mass_spectrum is None:
3806            to_plot = [x for x in to_plot if x != "MS1"]
3807        if len(rep_mf.ms2_mass_spectra) == 0:
3808            to_plot = [x for x in to_plot if x not in ["MS2", "MS2_mirror"]]
3809        
3810        # Check if EICs are available
3811        cluster_mfs = self.mass_features_dataframe[
3812            self.mass_features_dataframe['cluster'] == cluster_id
3813        ]
3814        
3815        has_eics = False
3816        # Check regular features
3817        for _, row in cluster_mfs.iterrows():
3818            sample_id = int(row['sample_id'])
3819            sample = self[sample_id]
3820            if hasattr(sample, 'eics') and sample.eics:
3821                if len(sample.eics) > 0:
3822                    has_eics = True
3823                    break
3824        
3825        # Also check induced features if available
3826        induced_cluster_mfs = None
3827        if not has_eics and self.induced_mass_features_dataframe is not None:
3828            induced_cluster_mfs = self.induced_mass_features_dataframe[
3829                self.induced_mass_features_dataframe['cluster'] == cluster_id
3830            ]
3831            for _, row in induced_cluster_mfs.iterrows():
3832                sample_id = int(row['sample_id'])
3833                sample = self[sample_id]
3834                if hasattr(sample, 'eics') and sample.eics:
3835                    if len(sample.eics) > 0:
3836                        has_eics = True
3837                        break
3838        
3839        if not has_eics:
3840            to_plot = [x for x in to_plot if x != "EIC"]
3841            if len(to_plot) == 0:
3842                raise ValueError(
3843                    f"No plottable data available for cluster {cluster_id}. "
3844                    f"Run process_consensus_features(gather_eics=True, add_ms1=True, add_ms2=True) first."
3845                )
3846        
3847        # Get induced features if not already retrieved
3848        if induced_cluster_mfs is None and self.induced_mass_features_dataframe is not None:
3849            induced_cluster_mfs = self.induced_mass_features_dataframe[
3850                self.induced_mass_features_dataframe['cluster'] == cluster_id
3851            ]
3852        
3853        # Check if MS1 is deconvoluted
3854        deconvoluted = rep_mf._ms_deconvoluted_idx is not None
3855        
3856        # Create figure
3857        fig, axs = plt.subplots(
3858            len(to_plot), 1, figsize=(10, len(to_plot) * 4), squeeze=False
3859        )
3860        
3861        fig.suptitle(
3862            f"Consensus Cluster {cluster_id}: "
3863            f"m/z = {cluster_summary['mz_median']:.4f} "
3864            f"(±{cluster_summary['mz_std']:.4f}); "
3865            f"RT = {cluster_summary['scan_time_aligned_median']:.2f} min "
3866            f"(±{cluster_summary['scan_time_aligned_std']:.2f}); "
3867            f"{int(cluster_summary['sample_id_nunique'])} samples"
3868        )
3869        
3870        i = 0
3871        
3872        # EIC plot - show all samples using helper method
3873        if "EIC" in to_plot:
3874            self._plot_multiple_eics(
3875                axs[i][0],
3876                cluster_mfs,
3877                induced_cluster_mfs,
3878                rep_sample_id,
3879                rep_mf_id,
3880                cluster_summary['scan_time_aligned_median'],
3881                eic_buffer_time,
3882                plot_smoothed=plot_smoothed_eic,
3883                plot_datapoints=plot_eic_datapoints,
3884                label_samples=label_samples,
3885                lcms_collection=self
3886            )
3887            i += 1
3888        
3889        # MS1 plot - from representative using helper method
3890        if "MS1" in to_plot:
3891            rep_mf._plot_ms1_spectrum(
3892                axs[i][0], 
3893                deconvoluted=deconvoluted, 
3894                sample_name=rep_sample.sample_name
3895            )
3896            i += 1
3897        
3898        # MS2 plot - from representative using helper method
3899        if "MS2" in to_plot:
3900            rep_mf._plot_ms2_spectrum(axs[i][0], sample_name=rep_sample.sample_name)
3901            i += 1
3902        
3903        # MS2 mirror plot - from representative using helper method
3904        if "MS2_mirror" in to_plot:
3905            rep_mf._plot_ms2_mirror(axs[i][0], molecular_metadata=molecular_metadata, spectral_library=spectral_library)
3906            i += 1
3907        
3908        plt.tight_layout()
3909        
3910        if return_fig:
3911            plt.close(fig)
3912            return fig
3913        else:
3914            plt.show()
3915            return None
3916    
3917    def get_representative_mass_features_for_all_clusters(self, representative_metric=None):
3918        """
3919        Get the most representative mass feature for all clusters in bulk.
3920        
3921        This is much more efficient than calling get_most_representative_sample_for_cluster
3922        in a loop, as it processes all clusters in a single pass over the dataframe.
3923        
3924        Parameters
3925        ----------
3926        representative_metric : str, optional
3927            The metric to use to determine the most representative sample.
3928            If None, uses the value from self.parameters.lcms_collection.consensus_representative_metric.
3929            Options:
3930            - 'intensity': Selects the mass feature with the highest intensity
3931            - 'intensity_prefer_ms2': Selects the highest intensity feature that has MS2 scans,
3932              or the highest intensity overall if none have MS2
3933            Default is None (uses parameter setting).
3934            
3935        Returns
3936        -------
3937        :obj:`~pandas.DataFrame`
3938            DataFrame with one row per cluster containing:
3939            - cluster: cluster ID
3940            - sample_id: The sample ID of the most representative sample
3941            - mf_id: The mass feature ID in the sample
3942            - coll_mf_id: The collection-level mass feature ID (index)
3943            - has_ms2: Whether this mass feature has MS2 scan numbers
3944            - intensity: The intensity value of the representative mass feature
3945        """
3946        # Use default from parameters if not specified
3947        if representative_metric is None:
3948            representative_metric = self.parameters.lcms_collection.consensus_representative_metric
3949        
3950        mf_df = self.mass_features_dataframe.copy()
3951        # Reset index to make coll_mf_id a column we can work with
3952        mf_df = mf_df.reset_index(drop=False)
3953        
3954        # Handle special metric 'intensity_prefer_ms2'
3955        if representative_metric == 'intensity_prefer_ms2':
3956            if 'intensity' not in mf_df.columns:
3957                raise ValueError(
3958                    f"'intensity' column not found in mass_features_dataframe. "
3959                    f"Available columns: {mf_df.columns.tolist()}"
3960                )
3961            
3962            # Add has_ms2 flag if ms2_scan_numbers column exists
3963            if 'ms2_scan_numbers' in mf_df.columns:
3964                def has_ms2_scans(val):
3965                    if val is None:
3966                        return False
3967                    try:
3968                        return len(val) > 0
3969                    except (TypeError, ValueError):
3970                        return False
3971                
3972                mf_df['has_ms2'] = mf_df['ms2_scan_numbers'].apply(has_ms2_scans)
3973                
3974                # Sort by has_ms2 (descending) then intensity (descending)
3975                # This ensures features with MS2 are preferred when intensities are equal
3976                mf_df = mf_df.sort_values(['has_ms2', 'intensity'], ascending=[False, False])
3977            else:
3978                mf_df['has_ms2'] = False
3979                mf_df = mf_df.sort_values('intensity', ascending=False)
3980            
3981            # Group by cluster and take the first (highest intensity, preferring MS2)
3982            representatives = mf_df.groupby('cluster').first().reset_index()
3983            
3984        else:
3985            # Standard metric - check if it exists
3986            if representative_metric not in mf_df.columns:
3987                raise ValueError(
3988                    f"Metric '{representative_metric}' not found. Available columns: {mf_df.columns.tolist()}"
3989                )
3990            
3991            # Add has_ms2 flag for consistency
3992            if 'ms2_scan_numbers' in mf_df.columns:
3993                def has_ms2_scans(val):
3994                    if val is None:
3995                        return False
3996                    try:
3997                        return len(val) > 0
3998                    except (TypeError, ValueError):
3999                        return False
4000                mf_df['has_ms2'] = mf_df['ms2_scan_numbers'].apply(has_ms2_scans)
4001            else:
4002                mf_df['has_ms2'] = False
4003            
4004            # Get the index of max value for each cluster
4005            idx = mf_df.groupby('cluster')[representative_metric].idxmax()
4006            representatives = mf_df.loc[idx].copy()
4007        
4008        # Select only the columns we need
4009        result_cols = ['cluster', 'sample_id', 'mf_id', 'coll_mf_id', 'has_ms2', 'intensity']
4010        representatives = representatives[result_cols]
4011        
4012        return representatives
4013    
4014    def get_sample_mf_map_for_representatives(self, representative_metric=None, include_cluster_id=True):
4015        """
4016        Build a mapping of sample_id -> list of representative mass feature IDs to load.
4017        
4018        This is a DRY helper method used by both process_consensus_features() and
4019        ReadSavedLCMSCollection to determine which mass features should be loaded
4020        for each sample when loading representatives.
4021        
4022        Parameters
4023        ----------
4024        representative_metric : str, optional
4025            The metric to use to determine the most representative sample.
4026            If None, uses the value from self.parameters.lcms_collection.consensus_representative_metric.
4027            Default is None.
4028        include_cluster_id : bool, optional
4029            If True, returns tuples of (mf_id, cluster_id). If False, returns just mf_id.
4030            Default is True.
4031        
4032        Returns
4033        -------
4034        dict
4035            Dictionary mapping sample_id (int) to list of mass feature identifiers.
4036            If include_cluster_id=True: list of tuples (mf_id, cluster_id)
4037            If include_cluster_id=False: list of mf_id integers
4038        
4039        Examples
4040        --------
4041        >>> # Get map with cluster IDs for loading
4042        >>> sample_mf_map = collection.get_sample_mf_map_for_representatives()
4043        >>> # sample_mf_map = {0: [(123, 0), (456, 1)], 1: [(789, 2)], ...}
4044        >>> 
4045        >>> # Get map without cluster IDs for pipeline
4046        >>> sample_mf_map = collection.get_sample_mf_map_for_representatives(include_cluster_id=False)
4047        >>> # sample_mf_map = {0: [123, 456], 1: [789], ...}
4048        """
4049        # Get all representative mass features in bulk (much faster than looping)
4050        representatives = self.get_representative_mass_features_for_all_clusters(
4051            representative_metric=representative_metric
4052        )
4053        
4054        # Build sample_mf_map
4055        sample_mf_map = {}
4056        for _, row in representatives.iterrows():
4057            sample_id = row['sample_id']
4058            mf_id = row['mf_id']
4059            cluster_id = row['cluster']
4060            
4061            if sample_id not in sample_mf_map:
4062                sample_mf_map[sample_id] = []
4063            
4064            if include_cluster_id:
4065                sample_mf_map[sample_id].append((mf_id, cluster_id))
4066            else:
4067                sample_mf_map[sample_id].append(mf_id)
4068        
4069        return sample_mf_map
4070    
4071    def get_most_representative_sample_for_cluster(self, cluster_id, representative_metric=None):
4072        """
4073        Get the most representative sample for a given cluster based on a metric.
4074        
4075        Parameters
4076        ----------
4077        cluster_id : int
4078            The cluster ID to find the representative sample for.
4079        representative_metric : str, optional
4080            The metric to use to determine the most representative sample.
4081            If None, uses the value from self.parameters.lcms_collection.consensus_representative_metric.
4082            Options:
4083            - 'intensity': Selects the mass feature with the highest intensity
4084            - 'intensity_prefer_ms2': Selects the highest intensity feature that has MS2 scans,
4085              or the highest intensity overall if none have MS2
4086            Default is None (uses parameter setting).
4087            
4088        Returns
4089        -------
4090        dict
4091            Dictionary containing:
4092            - 'sample_id': The sample ID of the most representative sample
4093            - 'sample_name': The sample name of the most representative sample
4094            - 'mf_id': The mass feature ID in the sample
4095            - 'coll_mf_id': The collection-level mass feature ID (index)
4096            - 'has_ms2': Whether this mass feature has MS2 scan numbers
4097            - 'intensity': The intensity value of the representative mass feature
4098        
4099        Raises
4100        ------
4101        ValueError
4102            If cluster_id is not found or if representative_metric is not a valid column.
4103        """
4104        # Use the bulk method to get all representatives, then filter to this cluster
4105        # This follows DRY principle and ensures consistency
4106        all_representatives = self.get_representative_mass_features_for_all_clusters(
4107            representative_metric=representative_metric
4108        )
4109        
4110        # Filter to the requested cluster
4111        cluster_rep = all_representatives[all_representatives['cluster'] == cluster_id]
4112        
4113        if len(cluster_rep) == 0:
4114            # Try to provide helpful error message
4115            available_clusters = self.mass_features_dataframe['cluster'].unique()
4116            raise ValueError(
4117                f"Cluster {cluster_id} not found in mass_features_dataframe. "
4118                f"Available clusters: {sorted(available_clusters[:10].tolist())}... "
4119                f"(showing first 10 of {len(available_clusters)} total clusters)"
4120            )
4121        
4122        # Get the representative row (should only be one)
4123        rep_row = cluster_rep.iloc[0]
4124        
4125        # Get sample name from sample_id (convert to int for list indexing)
4126        sample_id = int(rep_row['sample_id'])
4127        sample_name = self.samples[sample_id]
4128        
4129        return {
4130            'sample_id': sample_id,
4131            'sample_name': sample_name,
4132            'mf_id': rep_row['mf_id'],
4133            'coll_mf_id': rep_row['coll_mf_id'],
4134            'has_ms2': rep_row['has_ms2'],
4135            'intensity': rep_row['intensity']
4136        }
4137    
4138    def reload_representative_mass_features(self, add_ms2=False, auto_process_ms2=True, ms2_spectrum_mode=None, ms2_scan_filter=None):
4139        """
4140        Reload mass features for all representative samples in the cluster summary.
4141        
4142        This method is useful when the collection was loaded with load_light=True,
4143        which stores mass features only in the collection dataframe. This reloads
4144        the specific mass features that are representatives for each cluster,
4145        allowing them to be accessed as LCMSMassFeature objects.
4146        
4147        Parameters
4148        ----------
4149        add_ms2 : bool, optional
4150            If True, also loads and associates MS2 spectra with mass features. Default is False.
4151        auto_process_ms2 : bool, optional
4152            If True and add_ms2=True, auto-processes MS2 spectra. Default is True.
4153        ms2_spectrum_mode : str or None, optional
4154            Spectrum mode for MS2 spectra. If None, determines from parser. Default is None.
4155        ms2_scan_filter : str or None, optional
4156            Filter string for MS2 scans (e.g., 'hcd'). Default is None.
4157        
4158        Returns
4159        -------
4160        dict
4161            Dictionary mapping sample_id to list of reloaded mf_ids.
4162            
4163        Raises
4164        ------
4165        ValueError
4166            If cluster_summary_dataframe is not set (run add_consensus_mass_features first).
4167            
4168        Notes
4169        -----
4170        - Only reloads mass features that are cluster representatives
4171        - Uses get_most_representative_sample_for_cluster() to determine which to reload
4172        - More memory-efficient than reloading all mass features
4173        - Parallelized based on lcms_collection.cores parameter
4174        - MS2 association uses same logic as add_associated_ms2_dda()
4175        
4176        See Also
4177        --------
4178        _reload_sample_mass_features : Low-level method to reload specific mass features
4179        get_most_representative_sample_for_cluster : Gets representative sample for cluster
4180        """
4181        # Validate prerequisites
4182        if not hasattr(self, 'cluster_summary_dataframe') or self.cluster_summary_dataframe is None:
4183            raise ValueError(
4184                "cluster_summary_dataframe not found. Must run add_consensus_mass_features() first."
4185            )
4186        
4187        # Get all representative mass features in bulk (much faster than looping)
4188        representatives = self.get_representative_mass_features_for_all_clusters()
4189        
4190        # Build a dictionary of sample_id -> list of mf_ids that are representatives
4191        sample_mf_map = {}
4192        for _, row in representatives.iterrows():
4193            sample_id = row['sample_id']
4194            mf_id = row['mf_id']
4195            
4196            if sample_id not in sample_mf_map:
4197                sample_mf_map[sample_id] = []
4198            sample_mf_map[sample_id].append(mf_id)
4199        
4200        # Reload mass features for each sample (parallelized)
4201        if self.parameters.lcms_collection.cores == 1:
4202            # Serial processing
4203            from tqdm import tqdm
4204            for sample_id in tqdm(sample_mf_map.keys(), desc="Reloading representative mass features", unit="sample"):
4205                mf_ids = sample_mf_map[sample_id]
4206                self._reload_sample_mass_features(sample_id, mf_ids_to_load=mf_ids, add_ms2=add_ms2, 
4207                                                  auto_process_ms2=auto_process_ms2, ms2_spectrum_mode=ms2_spectrum_mode,
4208                                                  ms2_scan_filter=ms2_scan_filter)
4209        else:
4210            # Parallel processing
4211            import multiprocessing
4212            from tqdm import tqdm
4213            
4214            if self.parameters.lcms_collection.cores > len(sample_mf_map):
4215                ncores = len(sample_mf_map)
4216            else:
4217                ncores = self.parameters.lcms_collection.cores
4218            
4219            pool = multiprocessing.Pool(ncores)
4220            
4221            # Build arguments list for starmap
4222            args_list = [
4223                (sample_id, sample_mf_map[sample_id], add_ms2, auto_process_ms2, 
4224                 ms2_spectrum_mode, ms2_scan_filter, False)
4225                for sample_id in sample_mf_map.keys()
4226            ]
4227            
4228            # Execute in parallel
4229            mp_result = pool.starmap(self._reload_sample_mass_features, args_list)
4230            pool.close()
4231            pool.join()
4232            
4233            # Collect results back into samples
4234            for i, sample_id in enumerate(tqdm(sample_mf_map.keys(), desc="Collecting reloaded mass features", unit="sample")):
4235                self[sample_id].mass_features = mp_result[i]
4236        
4237        return sample_mf_map
4238    
4239    def _associate_ms2_with_mass_features(self, sample, local_mf_ids, auto_process=True, 
4240                                          spectrum_mode=None, scan_filter=None):
4241        """
4242        Associate MS2 spectra with specific mass features in a sample.
4243        
4244        Uses the LCMSBase helper method to find and load MS2 scans for the specified mass features.
4245        
4246        Parameters
4247        ----------
4248        sample : LCMSBase
4249            The sample object containing mass features and scan data.
4250        local_mf_ids : list of int
4251            List of local (sample-level) mass feature IDs to find MS2 for.
4252        auto_process : bool, optional
4253            If True, auto-processes the MS2 spectra. Default is True.
4254        spectrum_mode : str or None, optional
4255            Spectrum mode for MS2 spectra. If None, determines from parser. Default is None.
4256        scan_filter : str or None, optional
4257            Filter string for MS2 scans (e.g., 'hcd'). Default is None.
4258            
4259        Returns
4260        -------
4261        dict
4262            Dictionary of scan_number -> MassSpectrum objects for the loaded MS2 spectra.
4263        """
4264        # Check if we have scan data
4265        if not hasattr(sample, 'scan_df') or sample.scan_df is None:
4266            return {}
4267        
4268        # Separate mass features into those that need scan finding vs those that already have scans
4269        mfs_needing_scan_finding = []
4270        unique_dda_scans = set()
4271        
4272        for mf_id in local_mf_ids:
4273            if mf_id not in sample.mass_features:
4274                continue
4275            mf = sample.mass_features[mf_id]
4276            # If this mass feature already has MS2 scans, add them to our set
4277            if mf.ms2_scan_numbers is not None and len(mf.ms2_scan_numbers) > 0:
4278                # Convert to integers in case they come from HDF5 as numpy types
4279                unique_dda_scans.update([int(scan) for scan in mf.ms2_scan_numbers])
4280            else:
4281                # Otherwise, we need to find scans for this mass feature
4282                mfs_needing_scan_finding.append(mf_id)
4283        
4284        # Only run the scan finding for mass features that need it
4285        if mfs_needing_scan_finding:
4286            found_scans = sample._find_ms2_scans_for_mass_features(
4287                mf_ids=mfs_needing_scan_finding,
4288                scan_filter=scan_filter
4289            )
4290            unique_dda_scans.update(found_scans)
4291
4292        if len(unique_dda_scans) == 0:
4293            return {}
4294        
4295        # Get ms2 parameters from sample
4296        #TODO KRH: deal with different ms2 scan types here (CID vs HCD), may need to add scan translator to the initializeion
4297        ms_params = sample.parameters.mass_spectrum['ms2']
4298
4299        # Load MS2 spectra (convert set to list)
4300        sample.add_mass_spectra(
4301            scan_list=list(unique_dda_scans),
4302            auto_process=auto_process,
4303            spectrum_mode=spectrum_mode,
4304            ms_level=2,
4305            use_parser=True,
4306            ms_params=ms_params,
4307        )
4308        
4309        # Associate MS2 spectra with mass features
4310        for mf_id in local_mf_ids:
4311            if mf_id not in sample.mass_features:
4312                continue
4313            if sample.mass_features[mf_id].ms2_scan_numbers is not None and len(sample.mass_features[mf_id].ms2_scan_numbers) > 0:
4314                for dda_scan in sample.mass_features[mf_id].ms2_scan_numbers:
4315                    if dda_scan in sample._ms:
4316                        sample.mass_features[mf_id].ms2_mass_spectra[dda_scan] = sample._ms[dda_scan]
4317        
4318        # Return only the MS2 spectra we loaded (for parallel processing)
4319        return {scan: sample._ms[scan] for scan in unique_dda_scans if scan in sample._ms}
4320    
4321    def _reload_sample_mass_features(self, sample_id, mf_ids_to_load=None, add_ms2=False, 
4322                                     auto_process_ms2=True, ms2_spectrum_mode=None, ms2_scan_filter=None,
4323                                     inplace=True):
4324        """
4325        Reload specific mass features for a sample from HDF5.
4326        
4327        This is useful when the collection was loaded with load_light=True,
4328        which stores mass features only in the collection dataframe and not
4329        as LCMSMassFeature objects in individual samples.
4330        
4331        Parameters
4332        ----------
4333        sample_id : int
4334            The sample ID to reload mass features for.
4335        mf_ids_to_load : list of str, optional
4336            List of collection-level mf_ids (format: '{sample_id}_{local_mf_id}') to load.
4337            If None, loads all mass features for the sample.
4338        add_ms2 : bool, optional
4339            If True, also loads and associates MS2 spectra. Default is False.
4340        auto_process_ms2 : bool, optional
4341            If True, auto-processes MS2 spectra. Default is True.
4342        ms2_spectrum_mode : str or None, optional
4343            Spectrum mode for MS2 spectra. Default is None.
4344        ms2_scan_filter : str or None, optional
4345            Filter string for MS2 scans. Default is None.
4346        inplace : bool, optional
4347            If True, updates the sample's mass_features in place. If False, returns the
4348            mass_features dictionary (for multiprocessing). Default is True.
4349            
4350        Returns
4351        -------
4352        dict or None
4353            If inplace=False, returns dictionary of mass features.
4354            Otherwise returns None and updates object in place.
4355        """
4356        sample = self[sample_id]
4357        sample_name = self.samples[sample_id]
4358        
4359        # Check if we have a collection parser that can reload
4360        if not hasattr(self, 'collection_parser') or self.collection_parser is None:
4361            print("Warning: Cannot reload mass features - no collection_parser available")
4362            if not inplace:
4363                return {}
4364            return
4365        
4366        # Get the HDF5 file for this sample
4367        hdf5_file = self.collection_parser.folder_location / f"{sample_name}.corems/{sample_name}.hdf5"
4368        
4369        if not hdf5_file.exists():
4370            print(f"Warning: HDF5 file not found for sample {sample_name}: {hdf5_file}")
4371            if not inplace:
4372                return {}
4373            return
4374        
4375        # Import here to avoid circular imports
4376        from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
4377        
4378        # If specific mf_ids requested, use them directly
4379        local_mf_ids_to_load = None
4380        if mf_ids_to_load is not None:
4381            # mf_ids_to_load is already a list of sample-level mf_ids (integers)
4382            # No parsing needed - they come from the mf_id column in the dataframe
4383            local_mf_ids_to_load = set(mf_ids_to_load)
4384        
4385        # Reload mass features from HDF5
4386        with ReadCoreMSHDFMassSpectra(hdf5_file) as parser:
4387            # Load mass features - if specific IDs requested, only load those
4388            parser.import_mass_features(sample, mf_ids=local_mf_ids_to_load)
4389        
4390        # If add_ms2, associate MS2 spectra with the loaded mass features
4391        if add_ms2 and local_mf_ids_to_load is not None:
4392            self._associate_ms2_with_mass_features(
4393                sample, 
4394                list(local_mf_ids_to_load),
4395                auto_process=auto_process_ms2,
4396                spectrum_mode=ms2_spectrum_mode,
4397                scan_filter=ms2_scan_filter
4398            )
4399        
4400        # Return mass features if not inplace (for multiprocessing)
4401        if not inplace:
4402            return sample.mass_features
4403        
4404    def add_sparse_distance_matrix(self, features):
4405        if features is None:
4406            return None
4407        else:
4408            features = features.copy()
4409
4410        # Parameters for calculating distance between features
4411        dims = ["mz", "scan_time_aligned"]
4412        relative = [True, False]
4413        mz_tol_relative = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
4414        tol = [mz_tol_relative, self.parameters.lcms_collection.consensus_rt_tol]
4415        dist_weight = [1, 1]
4416
4417        # Check that the dimensions and tolerances are the same length
4418        if (
4419            len(dims) != len(tol)
4420            or len(dims) != len(relative)
4421            or len(dims) != len(dist_weight)
4422        ):
4423            raise ValueError(
4424                "The dimensions, tolerances, relative, dist_weight, and na_allow lists must be the same length"
4425            )
4426
4427        # Make connectivity matrix for masking within sample mass features
4428        ## Masking matrix cmat will mark all features from the same sample as 0
4429        ## To mask, a matrix can be multiplied by cmat and features from same
4430        ## samples are multiplied by 0, while features from different samples 
4431        ## are multiplied by 1
4432
4433        if "sample_id" not in features.columns:
4434            cmat = None
4435        else:
4436            vals = features["sample_id"].values.reshape(-1, 1)
4437            cmat = scipy.spatial.distance.cdist(vals, vals)
4438            # Convert to binary (0 if same sample, 1 if different)
4439            cmat = np.where(cmat == 0, 0, 1)
4440            # Convert to coordinate matrix for sparse operations later
4441            cmat = sparse.coo_matrix(cmat)
4442
4443        # Compute inter-feature distances using sparse matrix approach
4444        distances = None # clear the distances object before starting
4445        for i in range(len(dims)): # iterate through all dimensions to be considered
4446            # Construct k-d tree
4447            values = features[dims[i]].values
4448
4449            tree = KDTree(values.reshape(-1, 1))
4450
4451            max_tol = tol[i]
4452            if relative[i] is True:
4453                # Maximum absolute tolerance
4454                max_tol = tol[i] * values.max()
4455
4456            # Compute sparse distance matrix
4457            # the larger the max_tol, the slower this operation is
4458            sdm = tree.sparse_distance_matrix(tree, max_tol, output_type="coo_matrix")
4459
4460            # Only consider forward case, exclude diagonal
4461            sdm = sparse.triu(sdm, k=1)
4462
4463            # Filter relative distances
4464            if relative[i] is True:
4465                # Compute relative distances
4466                rel_dists = sdm.data / values[sdm.row]
4467
4468                # Indices of relative distances less than tolerance
4469                idx = rel_dists <= tol[i]
4470
4471                # Reconstruct sparse distance matrix
4472                sdm = sparse.coo_matrix(
4473                    (rel_dists[idx], (sdm.row[idx], sdm.col[idx])),
4474                    shape=(len(values), len(values)),
4475                )
4476
4477            # Scaled distances wrt the maximum tolerance for the dimension
4478            sdm.data = sdm.data / tol[i]
4479
4480            # Stack distances for dimensions where na_allow is False
4481            if distances is None:
4482                sdm.data = sdm.data * dist_weight[i]
4483                # Replace zeros with epsilon to handle perfect matches
4484                sdm.data[sdm.data == 0] = 1e-10
4485                distances = sdm
4486            else:
4487                # Prepare sdm to match shape of existing distances
4488                distances_truth = distances.copy()
4489                # make new sparse matrix with same positions as previous 
4490                # distance matrix but all ones for values
4491                distances_truth.data = np.ones_like(distances_truth.data)
4492                
4493                # Replace zeros with epsilon BEFORE multiply to prevent sparse matrix from dropping them
4494                sdm.data[sdm.data == 0] = 1e-10
4495                
4496                # multiply the new sparse matrix (sdm) by this mask to remove 
4497                # data that doesn't exist in original sparse matrix
4498                sdm = distances_truth.multiply(sdm)
4499                
4500                sdm.data = sdm.data * dist_weight[i]
4501                # Replace zeros with epsilon to handle perfect matches
4502                sdm.data[sdm.data == 0] = 1e-10
4503
4504                # use same process as before to remove data from previous
4505                # distances matrix that isn't in new distances matrix
4506                sdm_truth = sdm.copy()
4507                sdm_truth.data = np.ones_like(sdm_truth.data)
4508
4509                # remove the distances that are not sdm
4510                distances = distances.multiply(sdm_truth)
4511
4512                # Sum the new distances
4513                distances = distances + sdm
4514
4515        # Multiply by connectivity matrix for more masking
4516        distances = distances.multiply(cmat)
4517
4518        # Set attribute holding distance matrix
4519        self._sparse_distance_matrix = distances
4520
4521    def evaluate_clusters_for_repeats(self, features):
4522        raise NotImplementedError('evaluate_clusters_for_repeats not implemented yet')
4523        summary_df = self.cluster_summary_dataframe.copy()
4524
4525        # Arrange by decreasing median intensity
4526        summary_df = summary_df.sort_values(
4527            by="intensity_median", ascending=False
4528        ).reset_index(drop=True)
4529
4530        # Find clusters that are within the mz_tol and rt_tol of each other (on the medians)
4531        # Create a distance matrix
4532        # Define how to calculate the distance between features
4533        dims = ["mz_median", "scan_time_aligned_median"]
4534        relative = [True, False]
4535        mz_tol_relative = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
4536        tol = [mz_tol_relative, self.parameters.lcms_collection.consensus_rt_tol]
4537
4538        # Compute inter-feature distances
4539        distances = None
4540        for i in range(len(dims)):
4541            # Construct k-d tree
4542            values = summary_df[dims[i]].values
4543            tree = KDTree(values.reshape(-1, 1))
4544
4545            max_tol = tol[i]
4546            if relative[i] is True:
4547                # Maximum absolute tolerance
4548                max_tol = tol[i] * values.max()
4549
4550            # Compute sparse distance matrix
4551            # the larger the max_tol, the slower this operation is
4552            sdm = tree.sparse_distance_matrix(tree, max_tol, output_type="coo_matrix")
4553
4554            # Only consider forward case, exclude diagonal
4555            sdm = sparse.triu(sdm, k=1)
4556
4557            # Filter relative distances
4558            if relative[i] is True:
4559                # Compute relative distances
4560                rel_dists = sdm.data / values[sdm.row]  # or col?
4561
4562                # Indices of relative distances less than tolerance
4563                idx = rel_dists <= tol[i]
4564
4565                # Reconstruct sparse distance matrix
4566                sdm = sparse.coo_matrix(
4567                    (rel_dists[idx], (sdm.row[idx], sdm.col[idx])),
4568                    shape=(len(values), len(values)),
4569                )
4570
4571            # Cast as binary matrix
4572            sdm.data = np.ones_like(sdm.data)
4573
4574            # Stack distances
4575            if distances is None:
4576                distances = sdm
4577            else:
4578                distances = distances.multiply(sdm)
4579
4580        # Roll up features
4581        # Extract indices of within-tolerance points
4582        distances = distances.tocoo()
4583        pairs = np.stack(
4584            (distances.row, distances.col), axis=1
4585        )  # These are the index values of the clusters, not the cluster ids
4586        # Conver to cluster ids
4587        pairs_df = pd.DataFrame(pairs, columns=["parent", "child"])
4588        pairs_df["parent"] = summary_df.loc[pairs[:, 0]]["cluster"].values
4589        pairs_df["child"] = summary_df.loc[pairs[:, 1]]["cluster"].values
4590        pairs_df = pairs_df.set_index("parent")
4591
4592        merge_these_clusters = []
4593        possible_overlaps = []
4594        root_parents = np.setdiff1d(
4595            np.unique(pairs_df.index.values), np.unique(pairs_df.child.values)
4596        )
4597        for parent in root_parents:
4598            parent_features = features[features["cluster"] == parent]
4599            children = pairs_df.loc[[parent], "child"].tolist()
4600            for child in children:
4601                overlap = self.check_merge(parent_features, child, features)
4602                if len(overlap) == 0:
4603                    merge_these_clusters.append((parent, child, len(overlap)))
4604                else:
4605                    possible_overlaps.append((parent, child, len(overlap)))
4606
4607        result_dict = {}
4608        result_dict["merge_these_clusters"] = merge_these_clusters
4609        result_dict["possible_overlaps"] = possible_overlaps
4610
4611        return result_dict
4612
4613    def check_merge(self, parent_features, child, features):
4614        # Grab the features of the parent and children
4615        child_features = features[features["cluster"] == child]
4616
4617        # Check if there is an overlap between mf_coll_id in the parent and child clusters
4618        overlap = np.intersect1d(
4619            parent_features["sample_id"].values, child_features["sample_id"].values
4620        )
4621
4622        return overlap
4623
4624    def cluster_mass_features_agg_cluster(self, features):
4625        if features is None:
4626            return None
4627
4628        features = features.copy()
4629
4630        self.add_sparse_distance_matrix(features)
4631
4632        distances = self._sparse_distance_matrix
4633
4634        # Convert to full matrix
4635        distances = distances.todense()
4636        
4637        # Cast all 0s to 1s for a distance matrix
4638        distances[distances == 0] = 1
4639        distances = np.asarray(distances)
4640
4641        # Perform clustering
4642        try:
4643            clustering = AgglomerativeClustering(
4644                n_clusters=None,
4645                linkage="complete",
4646                # using complete linkage will prevent one sample from being assigned to multiple clusters
4647                metric="precomputed",
4648                distance_threshold=1,
4649            ).fit(distances)
4650            features["cluster"] = clustering.labels_
4651
4652        # All data points are singleton clusters
4653        except:
4654            features["cluster"] = np.arange(len(features.index))
4655
4656        return features
4657
4658    def cluster_inspection_plot(self, clu, return_fig = False):        
4659        """
4660        Generate Scan Time vs m/z plot for a narrow range around a given 
4661        cluster. This tool is meant to support the user in fine tuning the
4662        tolerances used for the clustering algorithm. The user-provided cluster
4663        ID is highlighted in larger, magenta marker and the ten largest of the
4664        remaining clusters are idenfitied with different colors while the
4665        smallest clusters are light gray.
4666
4667        Parameters
4668        -----------
4669        clu :  integer
4670            A cluster ID that exists in self.mass_features_dataframe
4671        return_fig : boolean
4672            Indicates whether to plot cluster inspection figure (False) or 
4673            return figure object (True). Defaults to False.
4674
4675        Returns
4676        --------
4677        matplotlib.pyplot.Figure
4678            A figure displaying a scan time vs m/z scatterplot of small region
4679            around a given cluster with the ten largest clusters in the region
4680            distinctly identified
4681
4682        Raises
4683        ------
4684        Warning
4685            If cluster data haven't been added to the object yet
4686        """
4687
4688        if 'cluster' not in self.mass_features_dataframe.columns:
4689            raise ValueError(
4690            'Cluster information is not yet added to mass_features_dataframe, must run add_consensus_mass_features() first'
4691            )
4692        
4693        else:
4694            mztol = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
4695            rttol = self.parameters.lcms_collection.consensus_rt_tol
4696            clu_features = self.mass_features_dataframe.copy()
4697
4698            inclu = clu_features[clu_features.cluster == clu]
4699            exclu = clu_features[clu_features.cluster != clu]
4700
4701            dt_ymin = np.floor(min(inclu.mz)) - 1
4702            dt_ymax = np.ceil(max(inclu.mz)) + 1
4703            dt_xmin = np.floor(min(inclu.scan_time_aligned)) - 1
4704            dt_xmax = np.ceil(max(inclu.scan_time_aligned)) + 1
4705
4706            exclu = exclu[
4707                (
4708                    exclu.mz.between(dt_ymin, dt_ymax, inclusive = 'both')
4709                ) & (
4710                    exclu.scan_time_aligned.between(dt_xmin, dt_xmax, inclusive = 'both')
4711                )
4712            ]
4713
4714            bigclulist = list(exclu.cluster.value_counts()[:10].index)
4715            bigclu = exclu[exclu.cluster.isin(bigclulist)]
4716            smclu = exclu[~exclu.cluster.isin(bigclulist)]
4717
4718            colors = np.arange(0, 10)
4719            colordict = dict(zip(bigclulist, colors))
4720            bigclu['color'] = bigclu.cluster.apply(lambda x: colordict[x])
4721
4722            fig = plt.figure(figsize = (7.5, 5))
4723
4724            plt.scatter(
4725                inclu.scan_time_aligned,
4726                inclu.mz,
4727                c = 'm',
4728                s = 3,
4729                label = 'Cluster ' + str(clu)
4730            )
4731
4732            plt.scatter(
4733                bigclu.scan_time_aligned,
4734                bigclu.mz,
4735                c = bigclu.color,
4736                cmap = 'tab10',
4737                s = 1.5
4738            )
4739
4740            plt.scatter(
4741                smclu.scan_time_aligned,
4742                smclu.mz,
4743                c = 'silver',
4744                s = 2,
4745                label = 'Small clusters'
4746            )
4747
4748            plt.ylim(dt_ymin, dt_ymax)
4749            plt.xlim(dt_xmin, dt_xmax)
4750            plt.legend(ncol = 2, bbox_to_anchor = (0.8, -0.1))
4751            plt.xlabel('Scan time')
4752            plt.ylabel('m/z')
4753            title_str = 'Cluster ' + str(clu)
4754            title_str += ': representing ' + str(len(inclu.sample_id.unique())) 
4755            title_str += ' of ' + str(len(clu_features.sample_id.unique())) 
4756            title_str += ' samples\n'
4757            title_str += 'M/Z tolerance: ' + str(mztol) + '\n'
4758            title_str += 'Scan Time tolerance: ' + str(rttol)
4759            plt.title(title_str, fontsize = 10)
4760
4761            if return_fig:
4762                plt.close(fig)
4763                return fig
4764            else:
4765                plt.show()
4766
4767    def plot_cluster_outlier_frequency(self, dim_list = ['mz', 'scan_time_aligned'], clu_size_thresh = 0.5, return_fig = False):
4768        """
4769        Generate histogram showing the frequency of outlier occurrences by
4770        clustering dimension across all clusters
4771
4772        Parameters
4773        -----------
4774        dim_list :  list
4775            List of strings describing dimensions that can be used in 
4776            clustering. Available list items:
4777                - 'mz'
4778                - 'scan_time_aligned'
4779                - 'half_height_width'
4780                - 'tailing_factor'
4781                - 'dispersity_index'
4782                - 'intensity'
4783                - 'persistence'
4784        clu_size_thresh : float
4785            Value between 0 and 1 that indicates what percentage of samples 
4786            need to be present in a cluster before it's evaluated for outliers.
4787            Defaults to 0.5.
4788        return_fig : boolean
4789            Indicates whether to plot cluster inspection figure (False) or 
4790            return figure object (True). Defaults to False.
4791
4792        Returns
4793        --------
4794        matplotlib.pyplot.Figure
4795            A figure displaying the frequency of outlier occurrences across all
4796            clusters in the provided measurement dimensions
4797
4798        Raises
4799        ------
4800        Warning
4801            If cluster data haven't been added to the object yet
4802        """
4803
4804        if not hasattr(self, 'cluster_summary_dataframe'):
4805            raise ValueError(
4806                'cluster_summary_dataframe is not yet added, must run add_consensus_mass_features() first'
4807            )
4808
4809        mfdf = self.mass_features_dataframe.copy()
4810        summarydf = self.cluster_summary_dataframe
4811
4812        numsamples = len(self)
4813        sumdf = summarydf[summarydf.sample_id_nunique > numsamples * clu_size_thresh].reset_index(drop = True).copy()
4814
4815        ## find the ranges for non-outlier values and add them to sumdf
4816        mergelist = ['cluster']
4817        for dim in dim_list:
4818            maxtag = dim + '_outmax'
4819            mintag = dim + '_outmin'
4820            mergelist.append(maxtag)
4821            mergelist.append(mintag)
4822            # Calculate outlier thresholds using vectorized operations
4823            sumdf[mintag] = sumdf[dim + '_mean'] - 3*sumdf[dim + '_std']
4824            sumdf[maxtag] = sumdf[dim + '_mean'] + 3*sumdf[dim + '_std']
4825            ## If NaN shows up anywhere in dim_min, dim_max calculations, value is set to NaN and it's 
4826            ## not flagged. This happens when there's not enough values to compute median/std for that 
4827            ## dimension therefore can't have outliers
4828
4829        ## add ranges to mfdf and identify mass features that fall outside the ranges
4830        # Merge without dropping NaN - we'll handle it per-dimension
4831        outdf = pd.merge(mfdf, sumdf[mergelist], on = 'cluster')
4832
4833        outtags = ['cluster']
4834        for dim in dim_list:
4835            dimtag = dim + '_outlier'
4836            outtags.append(dimtag)
4837            maxtag = dim + '_outmax'
4838            mintag = dim + '_outmin'
4839            # Only flag as outlier if thresholds are valid (not NaN)
4840            outdf[dimtag] = np.where(
4841                (outdf[maxtag].notna() & outdf[mintag].notna()) &
4842                (((outdf[dim] > outdf[maxtag])) | ((outdf[dim] < outdf[mintag]))), 
4843                True, 
4844                False
4845            )
4846
4847        ## identify number of outliers in each cluster
4848        outliers = outdf[outtags]
4849        outliers = outliers.groupby(['cluster']).sum()
4850
4851        ## plot number of clusters that contain any outliers
4852        fig = plt.figure()
4853        plt.bar(dim_list, outliers.sum().values, width = 0.5)
4854        plt.xticks(rotation = 90)
4855        plt.title('Frequency of outliers across all clusters by category')
4856        
4857        if return_fig:
4858            plt.close(fig)
4859            return fig
4860        else:
4861            plt.show()
4862            
4863    def _search_for_targeted_mass_features_in_sample(self, obj_idx, missingdf, cluster_dict, expand_on_miss=False, inplace=True):
4864        """
4865        Helper method to search for missing mass features in a single sample.
4866        
4867        Internal method called by fill_missing_cluster_features() to perform
4868        gap-filling for one sample in the collection.
4869        
4870        Parameters
4871        ----------
4872        obj_idx : int
4873            Index of the sample being processed
4874        missingdf : pd.DataFrame
4875            DataFrame containing cluster information with columns:
4876            'cluster', 'sample_id_nunique', 'mz_min', 'mz_max', 
4877            'scan_time_aligned_min', 'scan_time_aligned_max', 'mz_min_allowed', 
4878            'mz_max_allowed', 'scan_time_aligned_min_allowed', 
4879            'scan_time_aligned_max_allowed', 'missing_samples'
4880        cluster_dict : dict
4881            Pre-computed cluster feature dictionary to avoid recomputation
4882        expand_on_miss : bool
4883            If True, expands search window when no peak found initially
4884        inplace : bool
4885            If True, assigns induced_mass_features in place. If False, returns the
4886            induced features dictionary (for multiprocessing)
4887            
4888        Returns
4889        -------
4890        dict or None
4891            If inplace=False, returns dictionary of induced mass features.
4892            Otherwise returns None and updates object in place.
4893        """
4894        ## Use the pre-computed cluster dictionary passed as parameter
4895        
4896        ## to get clusters missing data based on sample index:
4897        sampledf = missingdf[
4898            missingdf.missing_samples.apply(lambda x: obj_idx in x)
4899        ].reset_index(drop = True).copy()
4900
4901        # Skip if no missing features for this sample
4902        if len(sampledf) == 0:
4903            if not inplace:
4904                return {}
4905            return
4906
4907        self.load_raw_data(obj_idx, 1)
4908               
4909        ## this is the line that bugs due to _ms_unprocessed not having key 1
4910        ms1df = self[obj_idx]._ms_unprocessed[1].copy()
4911        scan_df = self[obj_idx].scan_df[['scan', 'scan_time_aligned']]
4912        ms1df = pd.merge(ms1df, scan_df, on = 'scan')
4913
4914        # Pre-extract all values from sampledf to avoid repeated .iloc calls
4915        clusters = sampledf.cluster.values
4916        mz_mins = sampledf.mz_min.values
4917        mz_maxs = sampledf.mz_max.values
4918        st_mins = sampledf.scan_time_aligned_min.values
4919        st_maxs = sampledf.scan_time_aligned_max.values
4920        
4921        if expand_on_miss:
4922            mz_mins_allowed = sampledf.mz_min_allowed.values
4923            mz_maxs_allowed = sampledf.mz_max_allowed.values
4924            st_mins_allowed = sampledf.sta_min_allowed.values
4925            st_maxs_allowed = sampledf.sta_max_allowed.values
4926
4927        # Pre-filter ms1df to reduce search space
4928        mz_global_min = mz_mins.min()
4929        mz_global_max = mz_maxs.max()
4930        st_global_min = st_mins.min()
4931        st_global_max = st_maxs.max()
4932        
4933        if expand_on_miss:
4934            mz_global_min = min(mz_global_min, mz_mins_allowed.min())
4935            mz_global_max = max(mz_global_max, mz_maxs_allowed.max())
4936            st_global_min = min(st_global_min, st_mins_allowed.min())
4937            st_global_max = max(st_global_max, st_maxs_allowed.max())
4938        
4939        ms1df_filtered = ms1df[
4940            (ms1df.mz >= mz_global_min) & 
4941            (ms1df.mz <= mz_global_max) &
4942            (ms1df.scan_time_aligned >= st_global_min) &
4943            (ms1df.scan_time_aligned <= st_global_max)
4944        ].copy()
4945
4946        # Generate set_ids for all features
4947        set_ids = [f'c{clusters[i]}_{i}_i' for i in range(len(sampledf))]
4948        
4949        # Use batch method to process all features at once
4950        if expand_on_miss:
4951            # First try with normal bounds
4952            peaks_dict = self[obj_idx].search_for_targeted_mass_features_batch(
4953                ms1df_filtered,
4954                mz_mins,
4955                mz_maxs,
4956                st_mins,
4957                st_maxs,
4958                set_ids,
4959                obj_idx=obj_idx,
4960                st_aligned=True
4961            )
4962            
4963            # Retry failed features with expanded bounds
4964            failed_indices = [i for i, sid in enumerate(set_ids) if peaks_dict[sid].apex_scan == -99]
4965            if failed_indices:
4966                failed_ids = [set_ids[i] for i in failed_indices]
4967                retry_peaks = self[obj_idx].search_for_targeted_mass_features_batch(
4968                    ms1df_filtered,
4969                    mz_mins_allowed[failed_indices],
4970                    mz_maxs_allowed[failed_indices],
4971                    st_mins_allowed[failed_indices],
4972                    st_maxs_allowed[failed_indices],
4973                    failed_ids,
4974                    obj_idx=obj_idx,
4975                    st_aligned=True
4976                )
4977                peaks_dict.update(retry_peaks)
4978        else:
4979            peaks_dict = self[obj_idx].search_for_targeted_mass_features_batch(
4980                ms1df_filtered,
4981                mz_mins,
4982                mz_maxs,
4983                st_mins,
4984                st_maxs,
4985                set_ids,
4986                obj_idx=obj_idx,
4987                st_aligned=True
4988            )
4989        
4990        # Assign peaks to induced_mass_features and cluster_dict
4991        for i in range(len(sampledf)):
4992            peak = peaks_dict[set_ids[i]]
4993            self[obj_idx].induced_mass_features[peak.id] = peak
4994            cluster_dict[clusters[i]] += [set_ids[i]]
4995
4996        # TODO KRH: Let's try to avoid these steps unless asked for by parameters to pick up speed
4997        if False:
4998            self[obj_idx].add_associated_ms1(induced_features = True)
4999            # need to set drop_if_fail to false for induced features as they will fail
5000            self[obj_idx].add_peak_metrics(induced_features = True)
5001            
5002        self[obj_idx].integrate_mass_features(drop_if_fail = False, induced_features = True)
5003
5004        if not inplace:
5005            return self[obj_idx].induced_mass_features
5006    
5007    def fill_missing_cluster_features(self):
5008        """
5009        Gap-filling for consensus mass features across collection samples.
5010        
5011        For clusters present in multiple samples but missing from others, searches
5012        raw MS1 data to find peaks in expected m/z and retention time windows. This
5013        creates "induced" mass features for peaks that exist in the data but weren't
5014        detected in the initial peak detection.
5015        
5016        Must be run after add_consensus_mass_features(). Results are accessible via
5017        induced_mass_features_dataframe property and included in collection_pivot_table
5018        and collection_consensus_report outputs.
5019
5020        Parameters
5021        ----------
5022        None
5023            Uses parameters from self.parameters.lcms_collection:
5024            - consensus_min_sample_fraction: Minimum fraction of samples (0-1) that must contain
5025              a cluster before gap-filling is attempted
5026            - gap_fill_expand_on_miss: If True, expands search window when no peak is found
5027            
5028        Returns
5029        -------
5030        None
5031            Updates induced_mass_features attribute for each LCMSBase object and
5032            combines them into induced_mass_features_dataframe.
5033            
5034        Raises
5035        ------
5036        ValueError
5037            If cluster_summary_dataframe is not set (must run add_consensus_mass_features first).
5038            
5039        Notes
5040        -----
5041        - Loads raw MS1 data for each sample, which may be memory intensive
5042        - Induced features are integrated and metrics calculated automatically
5043        - Processing can be parallelized using parameters.lcms_collection.cores
5044        
5045        See Also
5046        --------
5047        add_consensus_mass_features : Creates consensus features before gap-filling
5048        collection_pivot_table : Includes both regular and induced features
5049        collection_consensus_report : Reports on complete feature matrix
5050        """
5051        
5052        # Validate prerequisites
5053        if not hasattr(self, 'cluster_summary_dataframe') or self.cluster_summary_dataframe is None:
5054            raise ValueError(
5055                "cluster_summary_dataframe not found. Must run add_consensus_mass_features() first."
5056            )
5057        
5058        # Get parameters from settings
5059        min_cluster_presence = self.parameters.lcms_collection.consensus_min_sample_fraction
5060        expand_on_miss = self.parameters.lcms_collection.gap_fill_expand_on_miss
5061        
5062        # Validate parameters
5063        if not 0 <= min_cluster_presence <= 1:
5064            raise ValueError("consensus_min_sample_fraction must be between 0 and 1")
5065        
5066        summarydf = self.cluster_summary_dataframe
5067        mfdf = self.mass_features_dataframe
5068        
5069        sample_ct = len(self.samples)
5070        
5071        # Identify clusters present in sufficient samples but not all samples
5072        missingdf = summarydf[[
5073            'cluster', 
5074            'sample_id_nunique', 
5075            'mz_min', 
5076            'mz_max', 
5077            'scan_time_aligned_min', 
5078            'scan_time_aligned_max'
5079        ]]
5080        missingdf = missingdf[missingdf.sample_id_nunique > min_cluster_presence * sample_ct]
5081        missingdf = missingdf[missingdf.sample_id_nunique != sample_ct]
5082        
5083        # Check if there are any clusters to gap-fill
5084        if len(missingdf) == 0:
5085            return
5086
5087        # Find which samples are missing for each cluster
5088        # Use range(sample_ct) to include all samples, even those with no mass features
5089        all_sample_ids = list(range(sample_ct))
5090        missing_samples_list = []
5091        for c in missingdf.cluster.to_numpy():
5092            cludf = mfdf[mfdf.cluster == c]
5093            missing = [x for x in all_sample_ids if x not in cludf.sample_id.unique()]
5094            missing_samples_list.append(missing)
5095        missingdf['missing_samples'] = missing_samples_list
5096        
5097        # Calculate expanded search windows for expand_on_miss option
5098        mz_clu_tol = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
5099        rt_clu_tol = self.parameters.lcms_collection.consensus_rt_tol
5100        missingdf['mz_max_allowed'] = missingdf.mz_max + mz_clu_tol * missingdf.mz_max
5101        missingdf['mz_min_allowed'] = missingdf.mz_min - mz_clu_tol * missingdf.mz_min
5102        missingdf['sta_max_allowed'] = missingdf.scan_time_aligned_max + rt_clu_tol * missingdf.scan_time_aligned_max
5103        missingdf['sta_min_allowed'] = missingdf.scan_time_aligned_min - rt_clu_tol * missingdf.scan_time_aligned_min
5104
5105        # Compute cluster dictionary once to avoid recomputing for each sample
5106        cluster_dict = self.cluster_feature_dictionary
5107        
5108        # Process each sample to search for missing features
5109        if self.parameters.lcms_collection.cores == 1:
5110            for i in tqdm(range(sample_ct), desc="Gap-filling samples", unit="sample"):
5111                self._search_for_targeted_mass_features_in_sample(i, missingdf, cluster_dict, expand_on_miss)
5112
5113        if self.parameters.lcms_collection.cores > 1:
5114            if self.parameters.lcms_collection.cores > len(self):
5115                ncores = len(self)
5116            else:
5117                ncores = self.parameters.lcms_collection.cores
5118            pool = multiprocessing.Pool(ncores)
5119            mp_result = pool.starmap(
5120                self._search_for_targeted_mass_features_in_sample, 
5121                [(x, missingdf, cluster_dict, expand_on_miss, False) for x in range(sample_ct)]
5122            )
5123
5124            for i in tqdm(range(sample_ct), desc="Collecting gap-filled features", unit="sample"):
5125                self[i].induced_mass_features = mp_result[i]
5126                
5127        self._combine_mass_features(induced_features = True)
5128        
5129        # Mark that gap-filling has been performed
5130        self.missing_mass_features_searched = True
5131        
5132        for sample_name in self.samples:
5133            self._lcms[sample_name].mass_features = {}
5134    
5135    def process_samples_pipeline(self, operations, description=None, keep_raw_data=False, show_progress=True):
5136        """
5137        Execute a pipeline of operations on all samples in parallel.
5138        
5139        This method provides a flexible framework for performing multiple
5140        sample-level operations in a single parallelized pass, which is more
5141        efficient than calling separate methods sequentially.
5142        
5143        Parameters
5144        ----------
5145        operations : list of SampleOperation
5146            List of operations to perform on each sample, in order.
5147            Each operation should be an instance of a class derived from
5148            SampleOperation (see lc_calc_operations module).
5149        description : str or None, optional
5150            Progress bar description. If None, automatically generates description
5151            from operation descriptions (e.g., "gap-filling, reloading features").
5152            Default is None.
5153        keep_raw_data : bool, optional
5154            If True, keeps raw MS data loaded in memory after pipeline completes.
5155            If False, cleans up raw data to free memory. Default is False.
5156        show_progress : bool, optional
5157            If True, displays progress bars during processing. If False, runs silently.
5158            Default is True.
5159            
5160        Returns
5161        -------
5162        dict
5163            Dictionary with results from pipeline execution, keyed by operation name.
5164            Structure: {operation_name: {sample_id: result, ...}, ...}
5165            
5166        Raises
5167        ------
5168        ValueError
5169            If operations list is empty or contains invalid operations.
5170            
5171        Notes
5172        -----
5173        - Operations are executed sequentially within each sample
5174        - Samples are processed in parallel based on parameters.lcms_collection.cores
5175        - Each operation can have conditional execution via can_execute()
5176        - Results are collected back via collect_results() method of each operation
5177        - Failed operations for a sample are logged but don't halt processing
5178        - Raw MS data loaded by operations is automatically cleaned up unless keep_raw_data=True
5179        
5180        Examples
5181        --------
5182        >>> from corems.mass_spectra.calc.lc_calc_operations import (
5183        ...     GapFillOperation, ReloadFeaturesOperation
5184        ... )
5185        >>> ops = [
5186        ...     GapFillOperation('gap_fill', expand_on_miss=True),
5187        ...     ReloadFeaturesOperation('reload', add_ms2=True)
5188        ... ]
5189        >>> results = lcms_collection.process_samples_pipeline(ops)
5190        
5191        See Also
5192        --------
5193        lc_calc_operations : Module containing built-in operation classes
5194        fill_and_process_features : Convenience method combining common operations
5195        """
5196        from corems.mass_spectra.calc.lc_calc_operations import SampleOperation
5197        
5198        # Validate operations
5199        if not operations or len(operations) == 0:
5200            raise ValueError("operations list cannot be empty")
5201        
5202        for op in operations:
5203            if not isinstance(op, SampleOperation):
5204                raise ValueError(f"All operations must be SampleOperation instances, got {type(op)}")
5205        
5206        # Generate description from operations if not provided
5207        if description is None:
5208            operation_descriptions = [op.description for op in operations]
5209            description = ", ".join(operation_descriptions).capitalize()
5210        
5211        # Prepare runtime parameters for each operation
5212        # This is where we gather collection-level data that operations need
5213        runtime_params = self._prepare_pipeline_runtime_params(operations)
5214        runtime_params['keep_raw_data'] = keep_raw_data
5215        
5216        # Execute pipeline
5217        sample_ct = len(self.samples)
5218        
5219        if self.parameters.lcms_collection.cores == 1:
5220            # Serial processing
5221            results_by_operation = {op.name: {} for op in operations}
5222            
5223            if show_progress:
5224                from tqdm import tqdm
5225                # Print description on its own line before progress bar
5226                print(f"\n{description.capitalize()}:")
5227                iterator = tqdm(range(sample_ct), unit="sample", ncols=80)
5228            else:
5229                iterator = range(sample_ct)
5230            
5231            for sample_id in iterator:
5232                sample_results = self._execute_sample_pipeline(
5233                    sample_id, operations, runtime_params, inplace=True
5234                )
5235                # Collect results (collect_results already called in _execute_sample_pipeline when inplace=True)
5236                # Skip 'sample_id' key which is added for tracking
5237                for op_name, result in sample_results.items():
5238                    if op_name != 'sample_id':
5239                        results_by_operation[op_name][sample_id] = result
5240        else:
5241            # Parallel processing
5242            import multiprocessing
5243            
5244            if self.parameters.lcms_collection.cores > sample_ct:
5245                ncores = sample_ct
5246            else:
5247                ncores = self.parameters.lcms_collection.cores
5248            
5249            pool = multiprocessing.Pool(ncores)
5250            
5251            # Build arguments for each sample
5252            args_list = [
5253                (sample_id, operations, runtime_params, False)
5254                for sample_id in range(sample_ct)
5255            ]
5256            
5257            # Execute in parallel with progress tracking
5258            results_by_operation = {op.name: {} for op in operations}
5259            
5260            if show_progress:
5261                from tqdm import tqdm
5262                import time
5263                
5264                # Use starmap_async for parallel execution with progress tracking
5265                async_result = pool.starmap_async(self._execute_sample_pipeline, args_list)
5266                
5267                # Poll for completion and update progress bar
5268                print(description)
5269                pbar = tqdm(
5270                    total=sample_ct, 
5271                    desc="",
5272                    unit="sample", 
5273                    position=0,
5274                    leave=True,
5275                    dynamic_ncols=True
5276                )
5277                prev_completed = 0
5278                while not async_result.ready():
5279                    # Get number of completed tasks by checking remaining
5280                    completed = sample_ct - async_result._number_left
5281                    if completed > prev_completed:
5282                        pbar.update(completed - prev_completed)
5283                        prev_completed = completed
5284                    time.sleep(0.5)  # Poll every 500ms to avoid spam
5285                
5286                # Final update to 100%
5287                if prev_completed < sample_ct:
5288                    pbar.update(sample_ct - prev_completed)
5289                pbar.close()
5290                
5291                # Get all results
5292                mp_results = async_result.get()
5293            else:
5294                # Execute without progress
5295                mp_results = pool.starmap(self._execute_sample_pipeline, args_list)
5296            
5297            pool.close()
5298            pool.join()
5299            
5300            # Collect results back into collection
5301            for result in mp_results:
5302                sample_id = result.get('sample_id')
5303                for op in operations:
5304                    op_result = result.get(op.name)
5305                    if op_result is not None:
5306                        op.collect_results(sample_id, op_result, self)
5307                        results_by_operation[op.name][sample_id] = op_result
5308        
5309        return results_by_operation
5310    
5311    def _prepare_pipeline_runtime_params(self, operations):
5312        """
5313        Prepare runtime parameters needed by operations in the pipeline.
5314        
5315        This method gathers collection-level data that operations need,
5316        such as cluster information for gap-filling or mf_ids for reloading.
5317        
5318        Parameters
5319        ----------
5320        operations : list of SampleOperation
5321            List of operations that will be executed
5322            
5323        Returns
5324        -------
5325        dict
5326            Dictionary of runtime parameters for operations
5327        """
5328        from corems.mass_spectra.calc.lc_calc_operations import (
5329            GapFillOperation, ReloadFeaturesOperation, MS2SpectralSearchOperation,
5330            LoadEICsOperation
5331        )
5332        
5333        runtime_params = {}
5334        
5335        # Check if any operation needs gap-fill parameters
5336        needs_gap_fill = any(isinstance(op, GapFillOperation) for op in operations)
5337        if needs_gap_fill:
5338            # Prepare gap-fill parameters (same as fill_missing_cluster_features)
5339            min_cluster_presence = self.parameters.lcms_collection.consensus_min_sample_fraction
5340            expand_on_miss = self.parameters.lcms_collection.gap_fill_expand_on_miss
5341            
5342            summarydf = self.cluster_summary_dataframe
5343            mfdf = self.mass_features_dataframe
5344            sample_ct = len(self.samples)
5345            
5346            # Identify clusters needing gap-filling
5347            # Note: cluster_summary_dataframe has 'cluster' as index, need to reset it
5348            missingdf = summarydf.reset_index()[[
5349                'cluster', 
5350                'sample_id_nunique', 
5351                'mz_min', 
5352                'mz_max', 
5353                'scan_time_aligned_min', 
5354                'scan_time_aligned_max'
5355            ]].copy()
5356            missingdf = missingdf[missingdf.sample_id_nunique > min_cluster_presence * sample_ct]
5357            missingdf = missingdf[missingdf.sample_id_nunique != sample_ct]
5358            
5359            if len(missingdf) > 0:
5360                # Find which samples are missing for each cluster
5361                # Use range(sample_ct) to include all samples, even those with no mass features
5362                all_sample_ids = list(range(sample_ct))
5363                missing_samples_list = []
5364                for c in missingdf.cluster.to_numpy():
5365                    cludf = mfdf[mfdf.cluster == c]
5366                    missing = [x for x in all_sample_ids if x not in cludf.sample_id.unique()]
5367                    missing_samples_list.append(missing)
5368                missingdf['missing_samples'] = missing_samples_list
5369                
5370                # Calculate expanded search windows
5371                mz_clu_tol = self.parameters.lcms_collection.consensus_mz_tol_ppm * 1e-6
5372                rt_clu_tol = self.parameters.lcms_collection.consensus_rt_tol
5373                missingdf['mz_max_allowed'] = missingdf.mz_max + mz_clu_tol * missingdf.mz_max
5374                missingdf['mz_min_allowed'] = missingdf.mz_min - mz_clu_tol * missingdf.mz_min
5375                missingdf['sta_max_allowed'] = missingdf.scan_time_aligned_max + rt_clu_tol * missingdf.scan_time_aligned_max
5376                missingdf['sta_min_allowed'] = missingdf.scan_time_aligned_min - rt_clu_tol * missingdf.scan_time_aligned_min
5377                
5378                runtime_params['missingdf'] = missingdf
5379                runtime_params['cluster_dict'] = self.cluster_feature_dictionary
5380                runtime_params['expand_on_miss'] = expand_on_miss
5381        
5382        # Check if any operation needs reload parameters
5383        needs_reload = any(isinstance(op, ReloadFeaturesOperation) for op in operations)
5384        if needs_reload:
5385            # Use DRY helper method to build sample_mf_map
5386            sample_mf_map = self.get_sample_mf_map_for_representatives(include_cluster_id=False)
5387            runtime_params['sample_mf_map'] = sample_mf_map
5388        
5389        # Check if any operation needs MS2 spectral search parameters
5390        needs_ms2_search = any(isinstance(op, MS2SpectralSearchOperation) for op in operations)
5391        if needs_ms2_search:
5392            # Pass through pre-prepared spectral library
5393            if hasattr(self, '_spectral_lib') and self._spectral_lib is not None:
5394                runtime_params['fe_lib'] = self._spectral_lib
5395            if hasattr(self, '_spectral_search_molecular_metadata'):
5396                runtime_params['molecular_metadata'] = self._spectral_search_molecular_metadata
5397        
5398        # Check if any operation needs EIC loading parameters
5399        needs_eic_loading = any(isinstance(op, LoadEICsOperation) for op in operations)
5400        if needs_eic_loading:
5401            # Build cluster_mz_dict: map of sample_id -> list of m/z values in clusters
5402            mfdf = self.mass_features_dataframe
5403            cluster_mz_dict = {}
5404            
5405            # Get all mass features that belong to clusters (cluster is not NaN)
5406            clustered_mf = mfdf[mfdf['cluster'].notna()]
5407            
5408            # Group by sample_id and collect all m/z values associated with eics
5409            for sample_id in clustered_mf['sample_id'].unique():
5410                sample_df = clustered_mf[clustered_mf['sample_id'] == sample_id]
5411                sample = self[sample_id]  # Get the LCMS object for this sample
5412                
5413                # Extract _eic_mz from actual mass feature objects, not from dataframe
5414                eic_mz_list = []
5415                for mf_id in sample_df['mf_id'].values:
5416                    if mf_id in sample.mass_features:
5417                        mf = sample.mass_features[mf_id]
5418                        if hasattr(mf, '_eic_mz') and mf._eic_mz is not None:
5419                            eic_mz_list.append(mf._eic_mz)
5420                
5421                # Use the collected m/z values, or fallback to empty list if none found
5422                cluster_mz_dict[sample_id] = list(set(eic_mz_list)) if eic_mz_list else []
5423            
5424            runtime_params['cluster_mz_dict'] = cluster_mz_dict
5425        
5426        return runtime_params
5427    
5428    def _execute_sample_pipeline(self, sample_id, operations, runtime_params, inplace=True):
5429        """
5430        Execute a pipeline of operations on a single sample.
5431        
5432        This is the worker function called (potentially in parallel) for each sample.
5433        
5434        Parameters
5435        ----------
5436        sample_id : int
5437            Sample ID to process
5438        operations : list of SampleOperation
5439            Operations to execute in order
5440        runtime_params : dict
5441            Runtime parameters prepared by _prepare_pipeline_runtime_params
5442        inplace : bool, optional
5443            If True, updates sample in place. If False, returns results for
5444            multiprocessing. Default is True.
5445            
5446        Returns
5447        -------
5448        dict
5449            Dictionary with results from each operation, keyed by operation name.
5450            If inplace=True, returns results that need to be collected.
5451            If inplace=False, returns all results for multiprocessing collection.
5452        """
5453        results = {}
5454        
5455        # Check if any operations need raw MS data
5456        needs_raw_data = {}  # {ms_level: True/False}
5457        for op in operations:
5458            needs_raw, ms_level = op.needs_raw_ms_data()
5459            if needs_raw and ms_level:
5460                needs_raw_data[ms_level] = True
5461        
5462        # Load raw data once if any operations need it
5463        # Note: For gap-filling, it loads data internally, so we just track it here
5464        for ms_level in needs_raw_data.keys():
5465            # Gap-filling loads its own data, but we want to keep track that it's loaded
5466            # Other operations can then use the loaded data
5467            pass
5468        
5469        for op in operations:
5470            # Check if operation can execute on this sample
5471            sample = self[sample_id]
5472            if not op.can_execute(sample, self):
5473                # Skip this operation for this sample if prerequisites aren't met
5474                # This allows processing to continue for samples that don't have
5475                # all required data (e.g., MS2 spectra)
5476                results[op.name] = None
5477                continue
5478            
5479            # Prepare operation-specific runtime params
5480            op_runtime_params = {}
5481            
5482            # Add gap-fill params if this is a gap-fill operation
5483            from corems.mass_spectra.calc.lc_calc_operations import (
5484                GapFillOperation, ReloadFeaturesOperation, MS2SpectralSearchOperation, LoadEICsOperation
5485            )
5486            
5487            if isinstance(op, GapFillOperation):
5488                if 'missingdf' in runtime_params:
5489                    op_runtime_params['missingdf'] = runtime_params['missingdf']
5490                    op_runtime_params['cluster_dict'] = runtime_params['cluster_dict']
5491                    op_runtime_params['expand_on_miss'] = runtime_params['expand_on_miss']
5492            
5493            elif isinstance(op, ReloadFeaturesOperation):
5494                if 'sample_mf_map' in runtime_params:
5495                    sample_mf_map = runtime_params['sample_mf_map']
5496                    # Always pass mf_ids_to_load to ensure we only load what's needed
5497                    # If sample not in map, it has no representatives - pass empty list
5498                    op_runtime_params['mf_ids_to_load'] = sample_mf_map.get(sample_id, [])
5499            
5500            elif isinstance(op, MS2SpectralSearchOperation):
5501                # Add MS2 spectral search parameters
5502                if 'fe_lib' in runtime_params:
5503                    op_runtime_params['fe_lib'] = runtime_params['fe_lib']
5504                if 'molecular_metadata' in runtime_params:
5505                    op_runtime_params['molecular_metadata'] = runtime_params['molecular_metadata']
5506            
5507            elif isinstance(op, LoadEICsOperation):
5508                # Add EIC loading parameters
5509                if 'cluster_mz_dict' in runtime_params:
5510                    op_runtime_params['cluster_mz_dict'] = runtime_params['cluster_mz_dict']
5511            
5512            # Execute the operation
5513            result = op.execute(sample_id, self, **op_runtime_params)
5514            results[op.name] = result
5515            
5516            # If inplace, collect immediately
5517            if inplace and result is not None:
5518                op.collect_results(sample_id, result, self)
5519        
5520        # Clean up raw data if requested
5521        keep_raw_data = runtime_params.get('keep_raw_data', False)
5522        if not keep_raw_data:
5523            for ms_level in needs_raw_data.keys():
5524                if ms_level in self[sample_id]._ms_unprocessed:
5525                    del self[sample_id]._ms_unprocessed[ms_level]
5526        
5527        # Include sample_id in results for tracking (especially important for imap_unordered)
5528        results['sample_id'] = sample_id
5529        return results
5530    
5531    def process_consensus_features(self, load_representatives=True, perform_gap_filling=True,
5532                                   add_ms1=False, add_ms2=False,
5533                                   ms2_scan_filter=None, molecular_formula_search=False,
5534                                   ms2_spectral_search=False, spectral_lib=None,
5535                                   molecular_metadata=None,
5536                                   gather_eics=False,
5537                                   keep_raw_data=False,
5538                                   show_progress=True):
5539        """
5540        Process consensus mass features across the collection in a single parallelized pass.
5541        
5542        This method provides a convenient interface to the sample processing pipeline,
5543        allowing multiple operations (gap-filling, feature reloading, MS1/MS2 association,
5544        molecular formula search, and MS2 spectral search) to be performed efficiently in 
5545        a single pass through all samples.
5546        
5547        Parameters
5548        ----------
5549        load_representatives : bool, optional
5550            If True, loads representative mass features from HDF5. Default is True.
5551        perform_gap_filling : bool, optional
5552            If True, performs gap-filling for missing cluster features. Default is True.
5553            This operation loads raw MS1 data which can be reused by subsequent operations.
5554        add_ms1 : bool, optional
5555            If True and load_representatives=True, associates MS1 spectra with
5556            loaded features. Automatically uses raw data from gap-filling if available,
5557            otherwise uses parser. Spectrum mode is auto-detected. Default is False.
5558        add_ms2 : bool, optional
5559            If True and load_representatives=True, associates MS2 spectra with
5560            loaded features and automatically processes them. Spectrum mode is auto-detected. Default is False.
5561        ms2_scan_filter : str or None, optional
5562            Filter string for MS2 scans (e.g., 'hcd'). Default is None.
5563        molecular_formula_search : bool, optional
5564            If True, performs molecular formula search on mass features using
5565            associated MS1 spectra. Requires add_ms1=True or that MS1 spectra
5566            are already associated. Uses parameters from 
5567            parameters.mass_spectrum["ms1"].molecular_search. Default is False.
5568        ms2_spectral_search : bool, optional
5569            If True, performs MS2 spectral library search using FlashEntropy.
5570            Requires add_ms2=True and spectral_lib to be provided. Default is False.
5571        spectral_lib : FlashEntropy library, optional
5572            Pre-prepared FlashEntropy spectral library for MS2 search.
5573            Create using MSPInterface.get_metabolomics_spectra_library().
5574            Required if ms2_spectral_search=True. Default is None.
5575        molecular_metadata : pd.DataFrame, optional
5576            Molecular metadata corresponding to spectral_lib.
5577            Returned from MSPInterface.get_metabolomics_spectra_library().
5578            Stored as self.spectral_search_molecular_metadata for later export.
5579            Default is None.
5580        gather_eics : bool, optional
5581            If True, loads extracted ion chromatograms (EICs) from HDF5 for all
5582            mass features with assigned cluster_index (including gap-filled features).
5583            Enables access to EICs via get_eics_for_cluster(cluster_id) method.
5584            Requires that EICs were previously exported with export_eics=True.
5585            Default is False.
5586        keep_raw_data : bool, optional
5587            If True, keeps raw MS data loaded in memory after pipeline completes.
5588            If False, cleans up raw data to free memory. Default is False.
5589        show_progress : bool, optional
5590            If True, displays progress bars during processing. If False, runs silently.
5591            Default is True.
5592            
5593        Returns
5594        -------
5595        dict
5596            Dictionary with pipeline results. Keys include:
5597            - 'gap_fill': dict mapping sample_id to induced mass features (if gap-filling)
5598            - 'reload': dict mapping sample_id to reloaded mass features (if reloading)
5599            - 'mf_search': dict mapping sample_id to number of features searched (if molecular formula search)
5600            - 'ms2_search': dict mapping sample_id to number of spectra searched (if MS2 spectral search)
5601            
5602        Raises
5603        ------
5604        ValueError
5605            If neither operation is enabled, or if required parameters are missing.
5606            
5607        Notes
5608        -----
5609        - Must run add_consensus_mass_features() before calling this method
5610        - Processes samples in parallel based on parameters.lcms_collection.cores
5611        - Raw MS1 data loaded by gap-filling is automatically reused by MS1 association
5612        - MS2 spectral search requires add_ms2=True and msp_file_path
5613        - FlashEntropy library is created once and reused across all samples
5614        - More efficient than calling individual methods separately
5615        - After gap-filling, sets missing_mass_features_searched = True
5616        - Mass features remain loaded in memory for downstream processing
5617        - For more advanced workflows, use process_samples_pipeline() directly
5618        
5619        Examples
5620        --------
5621        >>> # Prepare spectral library for MS2 search
5622        >>> from corems.molecular_id.search.database_interfaces import MSPInterface
5623        >>> my_msp = MSPInterface(file_path='path/to/library.msp')
5624        >>> spectral_lib, molecular_metadata = my_msp.get_metabolomics_spectra_library(
5625        ...     polarity='negative',
5626        ...     format='flashentropy',
5627        ...     normalize=True,
5628        ...     fe_kwargs={
5629        ...         'normalize_intensity': True,
5630        ...         'min_ms2_difference_in_da': 0.02,
5631        ...         'max_ms2_tolerance_in_da': 0.01,
5632        ...         'max_indexed_mz': 3000,
5633        ...         'precursor_ions_removal_da': None,
5634        ...         'noise_threshold': 0,
5635        ...     }
5636        ... )
5637        >>> 
5638        >>> # Gap-fill, reload with MS1/MS2, perform molecular formula and spectral search
5639        >>> results = lcms_collection.process_consensus_features(
5640        ...     load_representatives=True,
5641        ...     perform_gap_filling=True,
5642        ...     add_ms1=True,
5643        ...     add_ms2=True,
5644        ...     molecular_formula_search=True,
5645        ...     ms2_spectral_search=True,
5646        ...     spectral_lib=spectral_lib,
5647        ...     molecular_metadata=molecular_metadata
5648        ... )
5649        
5650        See Also
5651        --------
5652        process_samples_pipeline : Generic pipeline executor for custom workflows
5653        fill_missing_cluster_features : Original gap-filling method
5654        reload_representative_mass_features : Original reload method
5655        """
5656        from corems.mass_spectra.calc.lc_calc_operations import (
5657            GapFillOperation, ReloadFeaturesOperation, MolecularFormulaSearchOperation,
5658            MS2SpectralSearchOperation, LoadEICsOperation
5659        )
5660        
5661        # Validate that at least one meaningful operation is enabled
5662        has_operations = (
5663            perform_gap_filling or 
5664            load_representatives or 
5665            molecular_formula_search or 
5666            ms2_spectral_search or 
5667            gather_eics or
5668            add_ms1 or
5669            add_ms2
5670        )
5671        
5672        if not has_operations:
5673            raise ValueError(
5674                "At least one operation must be enabled: perform_gap_filling, load_representatives, "
5675                "molecular_formula_search, ms2_spectral_search, gather_eics, add_ms1, or add_ms2"
5676            )
5677        
5678        # Validate prerequisites for gap-filling
5679        if perform_gap_filling:
5680            if not hasattr(self, 'cluster_summary_dataframe') or self.cluster_summary_dataframe is None:
5681                raise ValueError(
5682                    "Cannot perform gap-filling: cluster_summary_dataframe not set. "
5683                    "You must run add_consensus_mass_features() before calling process_consensus_features()."
5684                )
5685        
5686        # Validate prerequisites for MS2 spectral search
5687        if ms2_spectral_search:
5688            if spectral_lib is None:
5689                raise ValueError(
5690                    "MS2 spectral search requires spectral_lib to be provided. "
5691                    "Create it using MSPInterface.get_metabolomics_spectra_library() before calling this method."
5692                )
5693            # Check if mass features will be loaded OR are already loaded
5694            # (The operation's can_execute will check if MS2 spectra are actually present)
5695            if not load_representatives and not perform_gap_filling:
5696                # Check if at least one sample has mass features loaded
5697                # This allows MS2 search on already-loaded features
5698                has_loaded_features = any(
5699                    len(self[i].mass_features) > 0 if hasattr(self[i], 'mass_features') and self[i].mass_features is not None else False
5700                    for i in range(len(self.samples))
5701                )
5702                if not has_loaded_features:
5703                    raise ValueError(
5704                        "MS2 spectral search requires mass features to be loaded. "
5705                        "Either set load_representatives=True or perform_gap_filling=True to load them, "
5706                        "or load them in a previous call to process_consensus_features() before calling "
5707                        "with ms2_spectral_search=True."
5708                    )
5709        
5710        # Build pipeline
5711        operations = []
5712        
5713        if perform_gap_filling:
5714            expand_on_miss = self.parameters.lcms_collection.gap_fill_expand_on_miss
5715            operations.append(GapFillOperation('gap_fill', expand_on_miss=expand_on_miss))
5716        
5717        if load_representatives:
5718            operations.append(ReloadFeaturesOperation(
5719                'reload',
5720                add_ms1=add_ms1,
5721                add_ms2=add_ms2,
5722                auto_process_ms2=add_ms2,  # Auto-process MS2 if add_ms2 is enabled
5723                ms2_scan_filter=ms2_scan_filter
5724            ))
5725        
5726        if molecular_formula_search:
5727            operations.append(MolecularFormulaSearchOperation('mf_search'))
5728        
5729        if ms2_spectral_search:
5730            operations.append(MS2SpectralSearchOperation(
5731                'ms2_search',
5732                ms2_scan_filter=ms2_scan_filter
5733            ))
5734            # Store spectral library and metadata for runtime preparation
5735            self._spectral_lib = spectral_lib
5736            self._spectral_search_molecular_metadata = molecular_metadata
5737        
5738        if gather_eics:
5739            operations.append(LoadEICsOperation('load_eics'))
5740        
5741        # Execute pipeline (description auto-generated from operations)
5742        results = self.process_samples_pipeline(
5743            operations,
5744            keep_raw_data=keep_raw_data,
5745            show_progress=show_progress
5746        )
5747        
5748        # Store molecular metadata if spectral search was performed
5749        if ms2_spectral_search and hasattr(self, '_spectral_search_molecular_metadata'):
5750            # This allows users to access the metadata for reporting
5751            self.spectral_search_molecular_metadata = self._spectral_search_molecular_metadata
5752        # Post-processing
5753        if perform_gap_filling:
5754            # Combine induced mass features into dataframe
5755            self._combine_mass_features(induced_features=True)
5756            # Mark that gap-filling has been performed
5757            self.missing_mass_features_searched = True
5758
5759            # Add ._eic_mz to induced_mass_features_dataframe if it exists
5760            if self.induced_mass_features_dataframe is not None and len(self.induced_mass_features_dataframe) > 0:
5761                eics_mz = []
5762                for i, row in self.induced_mass_features_dataframe.iterrows():
5763                    sample_id = row['sample_id']
5764                    sample = self[sample_id]
5765                    if row['mf_id'] in sample.induced_mass_features.keys():
5766                        eic_mz = sample.induced_mass_features[row['mf_id']]._eic_mz
5767                        eics_mz.append(eic_mz)
5768                    else:
5769                        eics_mz.append(None)
5770                self.induced_mass_features_dataframe['_eic_mz'] = eics_mz
5771
5772            # Clear mass features from samples to free memory
5773            for sample_name in self.samples:
5774                self._lcms[sample_name].induced_mass_features = {}
5775        
5776        # Associate EICs with mass features if they were loaded
5777        # This must happen after all operations complete to work on the actual sample objects
5778        if gather_eics:
5779            print("\nAssociating EICs with mass features:")
5780            from tqdm import tqdm
5781            
5782            for sample_id in tqdm(range(len(self.samples)), unit="sample", ncols=80):
5783                sample = self[sample_id]
5784                if sample.eics:  # Only if EICs were loaded
5785                    # Associate EICs with regular mass features
5786                    sample.associate_eics_with_mass_features(induced=False)
5787                    # Associate EICs with induced mass features
5788                    sample.associate_eics_with_mass_features(induced=True)
5789                
5790        return results