corems.mass_spectra.calc.lc_calc_operations

   1"""
   2Sample-level operations for LCMS collection processing pipelines.
   3
   4This module provides a framework for defining reusable, composable operations
   5that can be executed on individual samples in a parallelized manner.
   6
   7Classes
   8-------
   9SampleOperation
  10    Base class for all sample-level operations
  11GapFillOperation
  12    Gap-fill missing cluster features for a sample
  13ReloadFeaturesOperation
  14    Reload mass features from HDF5 for a sample
  15
  16"""
  17
  18from abc import ABC, abstractmethod
  19import pandas as pd
  20
  21
  22class SampleOperation(ABC):
  23    """
  24    Base class for operations that can be performed on a sample.
  25
  26    All sample operations must inherit from this class and implement all
  27    abstract methods. This ensures proper integration with the pipeline framework.
  28
  29    Parameters
  30    ----------
  31    name : str
  32        Name of the operation (for logging and identification)
  33    **kwargs
  34        Additional keyword arguments stored as operation parameters
  35
  36    Attributes
  37    ----------
  38    name : str
  39        Operation name
  40    params : dict
  41        Dictionary of operation parameters
  42    description : str
  43        Human-readable description for progress messages (must override in subclasses)
  44    """
  45
  46    def __init__(self, name, **kwargs):
  47        self.name = name
  48        self.params = kwargs
  49
  50    @property
  51    @abstractmethod
  52    def description(self):
  53        """
  54        Human-readable description for progress messages.
  55
  56        This property must be overridden in subclasses to provide a meaningful
  57        description that will be shown in progress bars (e.g., "gap-filling",
  58        "reloading features", etc.).
  59
  60        Returns
  61        -------
  62        str
  63            Brief description of what this operation does
  64        """
  65        pass
  66
  67    @abstractmethod
  68    def needs_raw_ms_data(self):
  69        """
  70        Declare whether this operation needs raw MS data loaded.
  71
  72        Subclasses must implement this method to specify raw data requirements.
  73        The pipeline executor will ensure raw data is loaded before executing
  74        operations that need it, and can clean it up afterwards.
  75
  76        Returns
  77        -------
  78        tuple of (bool, int or None)
  79            (needs_raw_data, ms_level)
  80            - needs_raw_data: True if operation needs raw MS data
  81            - ms_level: MS level needed (1 for MS1, 2 for MS2, etc.) or None
  82
  83        Examples
  84        --------
  85        >>> def needs_raw_ms_data(self):
  86        ...     return True, 1  # Needs MS1 data
  87        >>> def needs_raw_ms_data(self):
  88        ...     return False, None  # No raw data needed
  89        """
  90        pass
  91
  92    @abstractmethod
  93    def can_execute(self, sample, collection):
  94        """
  95        Check if this operation can be executed on the sample.
  96
  97        Subclasses must implement this method to define prerequisites.
  98        Return True if the operation can execute, False otherwise.
  99
 100        Parameters
 101        ----------
 102        sample : LCMSBase
 103            The sample to check
 104        collection : LCMSBaseCollection
 105            The collection containing the sample
 106
 107        Returns
 108        -------
 109        bool
 110            True if operation can execute, False otherwise
 111
 112        Examples
 113        --------
 114        >>> def can_execute(self, sample, collection):
 115        ...     return True  # Can always execute
 116        >>> def can_execute(self, sample, collection):
 117        ...     return hasattr(sample, 'mass_features') and sample.mass_features
 118        """
 119        pass
 120
 121    @abstractmethod
 122    def execute(self, sample_id, collection, **runtime_params):
 123        """
 124        Execute the operation on a sample.
 125
 126        This method must be implemented by subclasses.
 127
 128        Parameters
 129        ----------
 130        sample_id : int
 131            Sample ID to process
 132        collection : LCMSBaseCollection
 133            The collection containing the sample
 134        **runtime_params
 135            Runtime parameters passed from the pipeline
 136
 137        Returns
 138        -------
 139        result
 140            Operation result (can be None if operation modifies sample in place)
 141        """
 142        pass
 143
 144    @abstractmethod
 145    def collect_results(self, sample_id, result, collection):
 146        """
 147        Collect results back into collection after parallel execution.
 148
 149        Subclasses must implement this method to handle result collection.
 150        If the operation modifies samples in place and doesn't need to collect
 151        results, simply implement as `pass`.
 152
 153        Parameters
 154        ----------
 155        sample_id : int
 156            Sample ID that was processed
 157        result
 158            Result returned from execute()
 159        collection : LCMSBaseCollection
 160            The collection to update
 161
 162        Examples
 163        --------
 164        >>> def collect_results(self, sample_id, result, collection):
 165        ...     pass  # Operation modifies sample in place
 166        >>> def collect_results(self, sample_id, result, collection):
 167        ...     collection[sample_id].induced_mass_features = result
 168        """
 169        pass
 170
 171    def __repr__(self):
 172        return f"{self.__class__.__name__}(name='{self.name}')"
 173
 174
 175class GapFillOperation(SampleOperation):
 176    """
 177    Gap-fill missing cluster features for a sample.
 178
 179    Searches raw MS1 data to find peaks in expected m/z and retention time
 180    windows for clusters that are present in other samples but missing from
 181    this sample.
 182
 183    The operation merges each gap's retention-time bounds, including the
 184    expanded bounds when expand_on_miss is set and eic_buffer_time, and
 185    passes that range to load_raw_data. load_raw_data forwards the range
 186    into get_ms_raw with the sample's full scan table. The Thermo and mzML
 187    readers ignore the range and read every scan in the table. Scattered
 188    gaps merge into one window that covers the run, so the range does not
 189    shrink a typical read. After the load, the search keeps MS1 points
 190    inside the gap m/z and retention-time bounds.
 191
 192    Parameters
 193    ----------
 194    name : str
 195        Operation name
 196    expand_on_miss : bool, optional
 197        If True, a miss is searched again with the m/z window widened by
 198        consensus_mz_tol_ppm and the retention-time window widened by
 199        consensus_rt_tol minutes. Taken from gap_fill_expand_on_miss,
 200        which defaults to True.
 201
 202    Notes
 203    -----
 204    Requires that add_consensus_mass_features() has been run on the collection.
 205    This operation loads raw MS1 data which will be available for subsequent operations.
 206    """
 207
 208    @property
 209    def description(self):
 210        """Human-readable description for progress messages."""
 211        return "gap-filling"
 212
 213    def needs_raw_ms_data(self):
 214        """This operation needs raw MS1 data."""
 215        return True, 1
 216
 217    def can_execute(self, sample, collection):
 218        """Check if cluster summary exists."""
 219        return (
 220            hasattr(collection, "cluster_summary_dataframe")
 221            and collection.cluster_summary_dataframe is not None
 222        )
 223
 224    def execute(self, sample_id, collection, **runtime_params):
 225        """
 226        Execute gap-filling for a single sample.
 227
 228        Parameters
 229        ----------
 230        sample_id : int
 231            Sample index to process
 232        collection : LCMSBaseCollection
 233            The collection
 234        **runtime_params
 235            Runtime parameters including:
 236            - missingdf : pd.DataFrame - Cluster information and missing samples (optional)
 237            - cluster_dict : dict - Cluster feature dictionary (optional)
 238            - expand_on_miss : bool - Whether to expand search window on miss (optional)
 239            If these are not provided, returns empty dict (no gaps to fill).
 240
 241        Returns
 242        -------
 243        dict
 244            Dictionary of induced mass features (empty if no gaps to fill)
 245        """
 246        # Extract gap-fill parameters from runtime_params
 247        # If not present, there are no gaps to fill, so return early
 248        if "missingdf" not in runtime_params:
 249            return {}
 250
 251        missingdf = runtime_params["missingdf"]
 252        cluster_dict = runtime_params["cluster_dict"]
 253        expand_on_miss = runtime_params["expand_on_miss"]
 254
 255        # This is essentially the same logic as _search_for_targeted_mass_features_in_sample
 256        # but extracted into an operation
 257
 258        # Get clusters missing data for this sample
 259        sampledf = (
 260            missingdf[missingdf.missing_samples.apply(lambda x: sample_id in x)]
 261            .reset_index(drop=True)
 262            .copy()
 263        )
 264
 265        # Skip if no missing features for this sample
 266        if len(sampledf) == 0:
 267            return {}
 268
 269        # Get buffer time from LCMS parameters for EIC extraction
 270        # This ensures we capture the full chromatographic peak beyond cluster bounds
 271        buffer_rt = collection[sample_id].parameters.lc_ms.eic_buffer_time
 272
 273        # Windows passed to the loader. Thermo and mzML still read the full scan table.
 274        time_ranges = []
 275
 276        for _, row in sampledf.iterrows():
 277            rt_min = row["scan_time_aligned_min"]
 278            rt_max = row["scan_time_aligned_max"]
 279
 280            # If expand_on_miss, also consider the allowed bounds
 281            if expand_on_miss:
 282                rt_min = min(rt_min, row["sta_min_allowed"])
 283                rt_max = max(rt_max, row["sta_max_allowed"])
 284
 285            # Apply buffer AFTER considering expand_on_miss bounds
 286            # This ensures buffer is added beyond even the expanded search window
 287            time_ranges.append((max(0, rt_min - buffer_rt), rt_max + buffer_rt))
 288
 289        # Merge overlapping time ranges to reduce number of separate loads
 290        time_ranges = sorted(time_ranges)
 291        merged_ranges = []
 292        if time_ranges:
 293            current_min, current_max = time_ranges[0]
 294            for rt_min, rt_max in time_ranges[1:]:
 295                if rt_min <= current_max:  # Overlapping or adjacent
 296                    current_max = max(current_max, rt_max)
 297                else:
 298                    merged_ranges.append((current_min, current_max))
 299                    current_min, current_max = rt_min, rt_max
 300            merged_ranges.append((current_min, current_max))
 301
 302        # time_range is forwarded to get_ms_raw; Thermo and mzML ignore it.
 303        collection.load_raw_data(sample_id, 1, time_range=merged_ranges)
 304
 305        # Get MS1 data
 306        ms1df = collection[sample_id]._ms_unprocessed[1].copy()
 307        scan_df = collection[sample_id].scan_df[["scan", "scan_time_aligned"]]
 308        ms1df = pd.merge(ms1df, scan_df, on="scan")
 309
 310        # Pre-extract all values from sampledf
 311        clusters = sampledf.cluster.values
 312        mz_mins = sampledf.mz_min.values
 313        mz_maxs = sampledf.mz_max.values
 314        st_mins = sampledf.scan_time_aligned_min.values
 315        st_maxs = sampledf.scan_time_aligned_max.values
 316
 317        if expand_on_miss:
 318            mz_mins_allowed = sampledf.mz_min_allowed.values
 319            mz_maxs_allowed = sampledf.mz_max_allowed.values
 320            st_mins_allowed = sampledf.sta_min_allowed.values
 321            st_maxs_allowed = sampledf.sta_max_allowed.values
 322
 323        # Pre-filter ms1df to reduce search space
 324        mz_global_min = mz_mins.min()
 325        mz_global_max = mz_maxs.max()
 326        st_global_min = st_mins.min()
 327        st_global_max = st_maxs.max()
 328
 329        if expand_on_miss:
 330            mz_global_min = min(mz_global_min, mz_mins_allowed.min())
 331            mz_global_max = max(mz_global_max, mz_maxs_allowed.max())
 332            st_global_min = min(st_global_min, st_mins_allowed.min())
 333            st_global_max = max(st_global_max, st_maxs_allowed.max())
 334
 335        ms1df_filtered = ms1df[
 336            (ms1df.mz >= mz_global_min)
 337            & (ms1df.mz <= mz_global_max)
 338            & (ms1df.scan_time_aligned >= st_global_min)
 339            & (ms1df.scan_time_aligned <= st_global_max)
 340        ].copy()
 341
 342        # Generate set_ids for all features
 343        set_ids = [f"c{clusters[i]}_{i}_i" for i in range(len(sampledf))]
 344
 345        # Use batch method to process all features at once
 346        if expand_on_miss:
 347            # First try with normal bounds
 348            peaks_dict = collection[sample_id].search_for_targeted_mass_features_batch(
 349                ms1df_filtered,
 350                mz_mins,
 351                mz_maxs,
 352                st_mins,
 353                st_maxs,
 354                set_ids,
 355                obj_idx=sample_id,
 356                st_aligned=True,
 357            )
 358
 359            # Retry failed features with expanded bounds
 360            failed_indices = [
 361                i for i, sid in enumerate(set_ids) if peaks_dict[sid].apex_scan == -99
 362            ]
 363            if failed_indices:
 364                failed_ids = [set_ids[i] for i in failed_indices]
 365                retry_peaks = collection[
 366                    sample_id
 367                ].search_for_targeted_mass_features_batch(
 368                    ms1df_filtered,
 369                    mz_mins_allowed[failed_indices],
 370                    mz_maxs_allowed[failed_indices],
 371                    st_mins_allowed[failed_indices],
 372                    st_maxs_allowed[failed_indices],
 373                    failed_ids,
 374                    obj_idx=sample_id,
 375                    st_aligned=True,
 376                )
 377                peaks_dict.update(retry_peaks)
 378        else:
 379            peaks_dict = collection[sample_id].search_for_targeted_mass_features_batch(
 380                ms1df_filtered,
 381                mz_mins,
 382                mz_maxs,
 383                st_mins,
 384                st_maxs,
 385                set_ids,
 386                obj_idx=sample_id,
 387                st_aligned=True,
 388            )
 389
 390        # Build induced_mass_features dict and update cluster_dict
 391        induced_mass_features = {}
 392        for i in range(len(sampledf)):
 393            peak = peaks_dict[set_ids[i]]
 394            induced_mass_features[peak.id] = peak
 395            cluster_dict[clusters[i]] += [set_ids[i]]
 396
 397        # Integrate mass features (don't fail on bad integration).
 398        # m/z, intensity, and retention time are already on the feature from the
 399        # targeted search. Representative selection reads mass_features_dataframe
 400        # only, so a gap-filled feature is never the cluster representative and
 401        # does not get a processed MS1 spectrum here.
 402        collection[sample_id].induced_mass_features = induced_mass_features
 403        # The duplicate scan only deletes keys from mass_features. Induced ids
 404        # are strings such as c12_0_i, so that scan does not remove them.
 405        collection[sample_id].integrate_mass_features(
 406            drop_if_fail=False, induced_features=True, drop_duplicates=False
 407        )
 408
 409        return collection[sample_id].induced_mass_features
 410
 411    def collect_results(self, sample_id, result, collection):
 412        """Collect induced mass features back into sample."""
 413        collection[sample_id].induced_mass_features = result
 414
 415
 416class ReloadFeaturesOperation(SampleOperation):
 417    """
 418    Reload mass features from HDF5 and optionally add MS1/MS2 spectra.
 419
 420    This is useful when the collection was loaded with load_light=True,
 421    which stores mass features only in the collection dataframe and not
 422    as LCMSMassFeature objects in individual samples.
 423
 424    Parameters
 425    ----------
 426    name : str
 427        Operation name
 428    add_ms1 : bool, optional
 429        If True, adds MS1 spectra to mass features. Automatically uses raw MS1 data
 430        if available (e.g., from gap-filling), otherwise uses parser. Spectrum mode
 431        is auto-detected. Default is False.
 432    add_ms2 : bool, optional
 433        If True, also loads and associates MS2 spectra. Spectrum mode is auto-detected.
 434        Default is False.
 435    auto_process_ms2 : bool, optional
 436        If True and add_ms2=True, auto-processes MS2 spectra. Default is True.
 437    ms2_scan_filter : str or None, optional
 438        Filter string for MS2 scans. Default is None.
 439
 440    Notes
 441    -----
 442    MS1 spectra association automatically uses raw MS1 data if loaded by a previous
 443    operation (e.g., GapFillOperation). This is efficient when multiple operations
 444    need MS1 data in the same pipeline. All spectrum modes are auto-detected from
 445    the data.
 446    """
 447
 448    @property
 449    def description(self):
 450        """Human-readable description for progress messages."""
 451        return "reloading features"
 452
 453    def needs_raw_ms_data(self):
 454        """This operation doesn't need raw data."""
 455        return False, None
 456
 457    def can_execute(self, sample, collection):
 458        """Check if collection parser is available."""
 459        return (
 460            hasattr(collection, "collection_parser")
 461            and collection.collection_parser is not None
 462        )
 463
 464    def execute(self, sample_id, collection, mf_ids_to_load=None, **runtime_params):
 465        """
 466        Execute feature reloading for a single sample.
 467
 468        Parameters
 469        ----------
 470        sample_id : int
 471            Sample ID to reload features for
 472        collection : LCMSBaseCollection
 473            The collection
 474        mf_ids_to_load : list of str, optional
 475            List of collection-level mf_ids to load
 476        **runtime_params
 477            Additional runtime parameters (ignored)
 478
 479        Returns
 480        -------
 481        dict
 482            Dictionary of reloaded mass features
 483        """
 484        # Get parameters
 485        add_ms1 = self.params.get("add_ms1", False)
 486        add_ms2 = self.params.get("add_ms2", False)
 487        auto_process_ms2 = self.params.get("auto_process_ms2", True)
 488        ms2_scan_filter = self.params.get("ms2_scan_filter", None)
 489
 490        sample = collection[sample_id]
 491        sample_name = collection.samples[sample_id]
 492
 493        # Auto-determine if we should use parser for MS1 (check if raw data is available)
 494        has_raw_ms1 = (
 495            1 in sample._ms_unprocessed and not sample._ms_unprocessed[1].empty
 496        )
 497        use_parser_for_ms1 = (
 498            not has_raw_ms1
 499        )  # Use parser only if raw data not available
 500
 501        # Spectrum modes will be auto-detected (None = auto-detect)
 502        spectrum_mode_ms1 = None
 503        ms2_spectrum_mode = None
 504
 505        # Check if we have a collection parser
 506        if (
 507            not hasattr(collection, "collection_parser")
 508            or collection.collection_parser is None
 509        ):
 510            print(
 511                f"Warning: Cannot reload mass features for {sample_name} - no collection_parser available"
 512            )
 513            return {}
 514
 515        # Get the HDF5 file for this sample
 516        hdf5_file = (
 517            collection.collection_parser.folder_location
 518            / f"{sample_name}.corems/{sample_name}.hdf5"
 519        )
 520
 521        if not hdf5_file.exists():
 522            print(f"Warning: HDF5 file not found for sample {sample_name}: {hdf5_file}")
 523            return {}
 524
 525        # Import here to avoid circular imports
 526        from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
 527
 528        # If specific mf_ids requested, extract the local mf_ids we need
 529        local_mf_ids_to_load = None
 530        if mf_ids_to_load is not None:
 531            # mf_ids_to_load is already a list of sample-level mf_ids (integers)
 532            # No parsing needed - they come from the mf_id column in the dataframe
 533            if len(mf_ids_to_load) == 0:
 534                # No features to load for this sample - return empty dict
 535                return {}
 536            local_mf_ids_to_load = set(mf_ids_to_load)
 537
 538        # Reload mass features from HDF5
 539        with ReadCoreMSHDFMassSpectra(hdf5_file) as parser:
 540            parser.import_mass_features(sample, mf_ids=local_mf_ids_to_load)
 541
 542        # If add_ms1, associate MS1 spectra with the loaded mass features
 543        if add_ms1 and len(sample.mass_features) > 0:
 544            # Check if raw MS1 data is already loaded (e.g., from gap-filling)
 545            has_raw_ms1 = (
 546                1 in sample._ms_unprocessed and not sample._ms_unprocessed[1].empty
 547            )
 548
 549            if has_raw_ms1 and not use_parser_for_ms1:
 550                # Use already-loaded raw data (more efficient)
 551                sample.add_associated_ms1(
 552                    auto_process=True, use_parser=False, spectrum_mode=spectrum_mode_ms1
 553                )
 554            else:
 555                # Use parser to get MS1 spectra
 556                sample.add_associated_ms1(
 557                    auto_process=True, use_parser=True, spectrum_mode=spectrum_mode_ms1
 558                )
 559
 560        # If add_ms2, associate MS2 spectra with the loaded mass features
 561        if add_ms2 and len(sample.mass_features) > 0:
 562            # Get the IDs of loaded mass features (use what was actually loaded)
 563            mf_ids_for_ms2 = list(sample.mass_features.keys())
 564
 565            collection._associate_ms2_with_mass_features(
 566                sample,
 567                mf_ids_for_ms2,
 568                auto_process=auto_process_ms2,
 569                spectrum_mode=ms2_spectrum_mode,
 570                scan_filter=ms2_scan_filter,
 571            )
 572
 573        # Return both mass_features and _ms so they can be collected in multiprocessing
 574        return {"mass_features": sample.mass_features, "_ms": sample._ms}
 575
 576    def collect_results(self, sample_id, result, collection):
 577        """
 578        Collect reloaded mass features back into sample.
 579
 580        This operation loads a subset of mass features (e.g., representatives)
 581        into sample.mass_features for processing, while preserving the full
 582        mass_features_dataframe at the collection level. Sets a lock flag to
 583        prevent automatic rebuilding of the collection dataframe from individual
 584        samples. Also collects loaded mass spectra.
 585
 586        Parameters
 587        ----------
 588        sample_id : int
 589            Sample ID that was processed
 590        result : dict
 591            Dictionary with 'mass_features' and '_ms' from execute()
 592        collection : LCMSBaseCollection
 593            The collection
 594        """
 595        # Update sample.mass_features with loaded features
 596        if isinstance(result, dict) and "mass_features" in result:
 597            collection[sample_id].mass_features = result["mass_features"]
 598            # Also collect the _ms dictionary (MS1 and MS2 spectra)
 599            if "_ms" in result:
 600                collection[sample_id]._ms.update(result["_ms"])
 601        else:
 602            # Backward compatibility - if result is just mass_features dict
 603            collection[sample_id].mass_features = result
 604
 605        # Lock the collection dataframe to prevent rebuilding from individual samples
 606        # (since we've only loaded a subset, rebuilding would lose data)
 607        collection._mass_features_locked = True
 608
 609
 610class MolecularFormulaSearchOperation(SampleOperation):
 611    """
 612    Perform molecular formula search on mass features using associated MS1 spectra.
 613
 614    This operation runs molecular formula search on all mass features in a sample
 615    that have associated MS1 spectra. Requires MS1 spectra to be loaded and
 616    processed before execution.
 617
 618    Parameters
 619    ----------
 620    name : str
 621        Operation name (for logging)
 622    **kwargs
 623        Additional parameters passed to parent class
 624
 625    Examples
 626    --------
 627    >>> op = MolecularFormulaSearchOperation('mf_search')
 628    >>> # Use in pipeline
 629    >>> results = collection.process_samples_pipeline([op])
 630
 631    Notes
 632    -----
 633    This operation requires that MS1 spectra have been associated with mass
 634    features (e.g., via ReloadFeaturesOperation with add_ms1=True). The
 635    molecular formula search uses parameters from the collection's
 636    parameters.mass_spectrum["ms1"].molecular_search settings.
 637    """
 638
 639    @property
 640    def description(self):
 641        """Human-readable description for progress messages."""
 642        return "molecular formula search"
 643
 644    def __init__(self, name="molecular_formula_search", **kwargs):
 645        super().__init__(name, **kwargs)
 646
 647    def needs_raw_ms_data(self):
 648        """
 649        This operation doesn't need raw data - it works on processed MS1 spectra
 650        that are already associated with mass features.
 651
 652        Returns
 653        -------
 654        tuple
 655            (False, None) - no raw data needed
 656        """
 657        return False, None
 658
 659    def can_execute(self, sample, collection, **runtime_params):
 660        """
 661        Check if molecular formula search can be executed.
 662
 663        Requires that the sample has mass features with associated MS1 spectra.
 664
 665        Parameters
 666        ----------
 667        sample : LCMSObject
 668            The sample object
 669        collection : LCMSCollection
 670            The collection containing the sample
 671        **runtime_params
 672            Runtime parameters (not used)
 673
 674        Returns
 675        -------
 676        bool
 677            True if sample has mass features with MS1 spectra
 678        """
 679        # Check if sample has mass features
 680        if not hasattr(sample, "mass_features") or not sample.mass_features:
 681            return False
 682
 683        # Check if at least some mass features have MS1 spectra
 684        has_ms1 = any(
 685            hasattr(mf, "mass_spectrum") and mf.mass_spectrum is not None
 686            for mf in sample.mass_features.values()
 687        )
 688
 689        return has_ms1
 690
 691    def execute(self, sample_id, collection, **runtime_params):
 692        """
 693        Execute molecular formula search on a sample.
 694
 695        Creates a SearchMolecularFormulasLC object and runs mass feature search,
 696        which annotates mass features with molecular formula assignments.
 697
 698        Parameters
 699        ----------
 700        sample_id : str
 701            Sample identifier
 702        collection : LCMSCollection
 703            The collection containing the sample
 704        **runtime_params
 705            Runtime parameters (not used)
 706
 707        Returns
 708        -------
 709        int
 710            Number of mass features that were searched
 711        """
 712        from corems.molecular_id.search.molecularFormulaSearch import (
 713            SearchMolecularFormulasLC,
 714        )
 715        import time
 716        import sqlalchemy.exc
 717        import sqlite3
 718
 719        sample = collection[sample_id]
 720
 721        # Verify that mass features exist
 722        if not hasattr(sample, "mass_features") or not sample.mass_features:
 723            return 0  # No mass features to search
 724
 725        # Verify that mass features have MS1 spectra associated
 726        if not hasattr(sample, "_ms") or not sample._ms:
 727            raise RuntimeError(
 728                f"Sample {sample_id} does not have MS1 spectra loaded in _ms dictionary. "
 729                "Molecular formula search requires MS1 spectra to be associated with mass features. "
 730                "Ensure add_ms1=True when reloading features."
 731            )
 732
 733        from corems.mass_spectra.calc.feature_grouping import (
 734            annotation_meta_for_sample_mf,
 735            constrain_annotation_active,
 736            should_skip_isotope_for_annotation,
 737        )
 738
 739        settings = collection.parameters.lcms_collection
 740        labels = getattr(collection, "feature_group_dataframe", None)
 741        constrain = constrain_annotation_active(
 742            labels, bool(getattr(settings, "feature_group_constrain_annotation", True))
 743        )
 744        mf_df = getattr(collection, "mass_features_dataframe", None)
 745
 746        # Prepare data for bulk molecular formula search
 747        # Group mass features by their apex scan
 748        scan_to_mf = {}
 749        for mf_id, mf in sample.mass_features.items():
 750            if constrain:
 751                meta = annotation_meta_for_sample_mf(mf_df, labels, sample_id, mf_id)
 752                if meta is not None and should_skip_isotope_for_annotation(
 753                    meta.get("ion_role"), True
 754                ):
 755                    continue
 756            apex_scan = mf.apex_scan
 757            if apex_scan not in scan_to_mf:
 758                scan_to_mf[apex_scan] = []
 759            scan_to_mf[apex_scan].append(mf)
 760
 761        # Build lists of mass spectra and corresponding peaks
 762        mass_spectrum_list = []
 763        ms_peaks_list = []
 764
 765        for scan_num, mf_list in scan_to_mf.items():
 766            # Get the mass spectrum for this scan
 767            if scan_num not in sample._ms:
 768                continue  # Skip if spectrum not loaded
 769
 770            mass_spectrum = sample._ms[scan_num]
 771
 772            # Verify spectrum is processed (has peaks)
 773            if not hasattr(mass_spectrum, "_mspeaks") or not mass_spectrum._mspeaks:
 774                continue  # Skip unprocessed spectra
 775
 776            # Get the MS1 peaks for each mass feature at this scan
 777            peaks_for_scan = []
 778            for mf in mf_list:
 779                try:
 780                    # Use the ms1_peak property which finds the closest peak
 781                    ms1_peak = mf.ms1_peak
 782                    peaks_for_scan.append(ms1_peak)
 783                except (AttributeError, IndexError):
 784                    # Skip if ms1_peak can't be determined
 785                    continue
 786
 787            if peaks_for_scan:
 788                mass_spectrum_list.append(mass_spectrum)
 789                ms_peaks_list.append(peaks_for_scan)
 790
 791        # Run molecular formula search if we have data, with retry logic for database locks
 792        if mass_spectrum_list and ms_peaks_list:
 793            max_retries = 10
 794            retry_delay = 2  # seconds
 795
 796            for attempt in range(max_retries):
 797                try:
 798                    mol_search = SearchMolecularFormulasLC(sample)
 799                    mol_search.bulk_run_molecular_formula_search(
 800                        mass_spectrum_list, ms_peaks_list
 801                    )
 802                    break  # Success, exit retry loop
 803                except (sqlalchemy.exc.OperationalError, sqlite3.OperationalError) as e:
 804                    if attempt < max_retries - 1:
 805                        # Database is locked, retry after delay
 806                        print(
 807                            f"Sample {sample_id}: Database locked during molecular formula search, retrying in {retry_delay}s (attempt {attempt + 1}/{max_retries})..."
 808                        )
 809                        time.sleep(retry_delay)
 810                    else:
 811                        # Max retries exceeded, re-raise the exception
 812                        raise RuntimeError(
 813                            f"Sample {sample_id}: Molecular formula search failed after {max_retries} attempts due to database lock. "
 814                            "Try reducing parallel cores or increasing database timeout."
 815                        ) from e
 816
 817        # Return count of features searched
 818        return len(sample.mass_features)
 819
 820    def collect_results(self, sample_id, result, collection):
 821        """
 822        Collect results (no-op as search modifies mass features in place).
 823
 824        The molecular formula search modifies mass features in place by adding
 825        molecular formula assignments, so no explicit result collection is needed.
 826
 827        Parameters
 828        ----------
 829        sample_id : str
 830            Sample identifier
 831        result : int
 832            Number of features searched
 833        collection : LCMSCollection
 834            The collection containing the sample
 835        """
 836        # Search modifies mass features in place, nothing to collect
 837        pass
 838
 839
 840class MS2SpectralSearchOperation(SampleOperation):
 841    """
 842    Perform MS2 spectral search using entropy-based matching.
 843
 844    This operation performs spectral library search on MS2 spectra associated
 845    with mass features using FlashEntropy for fast similarity scoring. Requires
 846    MS2 spectra to be loaded and processed before execution.
 847
 848    Parameters
 849    ----------
 850    name : str
 851        Operation name (for logging)
 852    ms2_scan_filter : str or None, optional
 853        Filter string for MS2 scans (e.g., 'hcd'). If None, uses all MS2 scans.
 854        Default is None.
 855    peak_sep_da : float, optional
 856        Peak separation in Daltons for spectral matching. Default is 0.01.
 857    **kwargs
 858        Additional parameters passed to parent class
 859
 860    Examples
 861    --------
 862    >>> op = MS2SpectralSearchOperation('ms2_search', ms2_scan_filter='hcd')
 863    >>> # Use in pipeline - requires fe_lib in runtime_params
 864    >>> results = collection.process_samples_pipeline([op])
 865
 866    Notes
 867    -----
 868    This operation requires:
 869    - MS2 spectra to be associated with mass features
 870    - FlashEntropy library (fe_lib) to be provided in runtime_params
 871    - MS2 spectra must be processed (centroided)
 872
 873    The spectral search modifies mass features in place by adding spectral
 874    match scores and metadata.
 875    """
 876
 877    @property
 878    def description(self):
 879        """Human-readable description for progress messages."""
 880        return "MS2 spectral search"
 881
 882    def __init__(self, name="ms2_spectral_search", ms2_scan_filter=None, **kwargs):
 883        super().__init__(name, **kwargs)
 884        self.params["ms2_scan_filter"] = ms2_scan_filter
 885
 886    def needs_raw_ms_data(self):
 887        """
 888        This operation doesn't need raw data - it works on processed MS2 spectra
 889        that are already associated with mass features.
 890
 891        Returns
 892        -------
 893        tuple
 894            (False, None) - no raw data needed
 895        """
 896        return False, None
 897
 898    def can_execute(self, sample, collection, **runtime_params):
 899        """
 900        Check if MS2 spectral search can be executed.
 901
 902        Requires that the sample has mass features with MS2 spectra associated.
 903
 904        Parameters
 905        ----------
 906        sample : LCMSObject
 907            The sample object
 908        collection : LCMSCollection
 909            The collection containing the sample
 910        **runtime_params
 911            Runtime parameters (not used)
 912
 913        Returns
 914        -------
 915        bool
 916            True if sample has mass features with MS2 spectra
 917        """
 918        # Check if sample has mass features
 919        if not hasattr(sample, "mass_features") or not sample.mass_features:
 920            return False
 921
 922        # Check if any mass features have MS2 spectra associated
 923        has_ms2 = any(
 924            hasattr(mf, "ms2_mass_spectra") and mf.ms2_mass_spectra
 925            for mf in sample.mass_features.values()
 926        )
 927
 928        return has_ms2
 929
 930    def execute(
 931        self,
 932        sample_id,
 933        collection,
 934        fe_lib=None,
 935        molecular_metadata=None,
 936        **runtime_params,
 937    ):
 938        """
 939        Execute MS2 spectral search on a sample.
 940
 941        Performs entropy-based spectral library search on all MS2 spectra
 942        in the sample that match the scan filter criteria.
 943
 944        Parameters
 945        ----------
 946        sample_id : str
 947            Sample identifier
 948        collection : LCMSCollection
 949            The collection containing the sample
 950        fe_lib : FlashEntropy library
 951            Pre-computed FlashEntropy library for spectral matching
 952        molecular_metadata : pd.DataFrame, optional
 953            Metadata for molecules in the spectral library
 954        **runtime_params
 955            Runtime parameters (not used)
 956
 957        Returns
 958        -------
 959        int
 960            Number of MS2 spectra searched
 961        """
 962        sample = collection[sample_id]
 963
 964        # Get parameters
 965        ms2_scan_filter = self.params.get("ms2_scan_filter")
 966
 967        # Verify that we have a spectral library
 968        if fe_lib is None:
 969            raise ValueError(
 970                f"Sample {sample_id}: MS2 spectral search requires fe_lib (FlashEntropy library) "
 971                "to be provided in runtime parameters. Create the library at the collection level "
 972                "and pass it to the pipeline."
 973            )
 974
 975        # Extract peak_sep_da from FlashEntropy library configuration
 976        # peak_sep_da should be 2 * max_ms2_tolerance_in_da to match the min_ms2_difference_in_da
 977        # used when creating the library
 978        tolerance_da = fe_lib.entropy_search.max_ms2_tolerance_in_da
 979        if tolerance_da is None:
 980            raise ValueError(
 981                f"Sample {sample_id}: Could not extract max_ms2_tolerance_in_da from FlashEntropy library. "
 982                "Ensure the library was created with this parameter specified."
 983            )
 984        peak_sep_da = 2 * tolerance_da
 985
 986        # Verify that sample has _ms dictionary
 987        if not hasattr(sample, "_ms") or not sample._ms:
 988            return 0  # No MS2 spectra to search
 989
 990        # Get MS2 scan numbers based on filter
 991        if ms2_scan_filter is not None:
 992            # Filter by scan text
 993            ms2_scan_df = sample.scan_df[
 994                sample.scan_df.scan_text.str.contains(ms2_scan_filter)
 995                & (sample.scan_df.ms_level == 2)
 996            ]
 997        else:
 998            # All MS2 scans
 999            ms2_scan_df = sample.scan_df[sample.scan_df.ms_level == 2]
1000
1001        # Get scans that are actually loaded in _ms
1002        ms2_scans_to_search = [
1003            scan for scan in ms2_scan_df.scan.tolist() if scan in sample._ms.keys()
1004        ]
1005
1006        if not ms2_scans_to_search:
1007            return 0  # No MS2 spectra to search
1008
1009        # Perform spectral search using the sample's fe_search method
1010        sample.fe_search(
1011            scan_list=ms2_scans_to_search, fe_lib=fe_lib, peak_sep_da=peak_sep_da
1012        )
1013
1014        from corems.mass_spectra.calc.feature_grouping import (
1015            annotation_meta_for_sample_mf,
1016            constrain_annotation_active,
1017            should_skip_isotope_for_annotation,
1018            allowed_ion_types_for_row,
1019            subset_hits_by_allowed_ion_types,
1020        )
1021
1022        settings = collection.parameters.lcms_collection
1023        labels = getattr(collection, "feature_group_dataframe", None)
1024        constrain = constrain_annotation_active(
1025            labels, bool(getattr(settings, "feature_group_constrain_annotation", True))
1026        )
1027        if constrain:
1028            mf_df = getattr(collection, "mass_features_dataframe", None)
1029            for mf_id, mf in list(sample.mass_features.items()):
1030                meta = annotation_meta_for_sample_mf(mf_df, labels, sample_id, mf_id)
1031                if meta is None:
1032                    continue
1033                if should_skip_isotope_for_annotation(meta.get("ion_role"), True):
1034                    mf.ms2_similarity_results = []
1035                    continue
1036                allowed = allowed_ion_types_for_row(
1037                    meta.get("ion_type"), meta.get("possible_ion_types")
1038                )
1039                if allowed is None:
1040                    continue
1041                results = getattr(mf, "ms2_similarity_results", None) or []
1042                filtered = []
1043                for res in results:
1044                    kept = subset_hits_by_allowed_ion_types(res, allowed)
1045                    if kept is not None:
1046                        filtered.append(kept)
1047                mf.ms2_similarity_results = filtered
1048
1049        # Return the spectral search results for collection
1050        # (needed for multiprocessing - results populated in worker need to be returned)
1051        return sample.spectral_search_results
1052
1053    def collect_results(self, sample_id, result, collection):
1054        """
1055        Collect spectral search results back into the sample.
1056
1057        In multiprocessing, the worker's modifications don't persist to the
1058        main process, so we need to explicitly collect and reassign the results.
1059        This also re-associates the results with mass features.
1060
1061        Parameters
1062        ----------
1063        sample_id : str
1064            Sample identifier
1065        result : dict
1066            Dictionary of spectral search results from execute()
1067        collection : LCMSCollection
1068            The collection containing the sample
1069        """
1070        # Assign the spectral search results back to the sample
1071        if result:
1072            collection[sample_id].spectral_search_results.update(result)
1073
1074            # Re-associate results with mass features (same logic as fe_search)
1075            sample = collection[sample_id]
1076            if len(sample.mass_features) > 0:
1077                for mass_feature_id, mass_feature in sample.mass_features.items():
1078                    scan_ids = mass_feature.ms2_scan_numbers
1079                    for ms2_scan_id in scan_ids:
1080                        precursor_mz = mass_feature.mz
1081                        try:
1082                            sample.spectral_search_results[ms2_scan_id][precursor_mz]
1083                        except KeyError:
1084                            pass
1085                        else:
1086                            sample.mass_features[
1087                                mass_feature_id
1088                            ].ms2_similarity_results.append(
1089                                sample.spectral_search_results[ms2_scan_id][
1090                                    precursor_mz
1091                                ]
1092                            )
1093
1094
1095class LoadEICsOperation(SampleOperation):
1096    """
1097    Load extracted ion chromatograms (EICs) from HDF5 for regular mass features.
1098
1099    Loads EICs for regular mass features that belong to consensus clusters from HDF5.
1100    Induced (gap-filled) features already have EICs from integrate_mass_features,
1101    so no additional loading is needed for them.
1102
1103    This operation enables downstream visualization and analysis of chromatographic
1104    peaks across all samples in a cluster.
1105
1106    Notes
1107    -----
1108    Requires that mass features have been loaded and cluster_index assigned.
1109    Regular mass feature EICs must have been previously saved to HDF5 with export_eics=True.
1110    Induced mass features already have EICs populated during gap-filling.
1111    """
1112
1113    @property
1114    def description(self):
1115        """Human-readable description for progress messages."""
1116        return "loading EICs"
1117
1118    def needs_raw_ms_data(self):
1119        """This operation doesn't need raw data - induced features already have EICs."""
1120        return False, None
1121
1122    def can_execute(self, sample, collection):
1123        """
1124        Check if EIC loading can be executed.
1125
1126        This operation can always execute if the sample exists - the actual work
1127        is determined by cluster_mz_dict in runtime_params. If cluster_mz_dict is
1128        empty or None, execute() will simply return 0 (no EICs loaded).
1129
1130        Returns
1131        -------
1132        bool
1133            True (always executable - runtime_params control actual work)
1134        """
1135        return True
1136
1137    def execute(self, sample_id, collection, cluster_mz_dict=None, **runtime_params):
1138        """
1139        Load EICs from HDF5 for a single sample.
1140
1141        Loads EICs for regular mass features that belong to consensus clusters.
1142        Induced (gap-filled) mass features already have EICs from integrate_mass_features,
1143        so no additional loading is needed for them.
1144
1145        The cluster_mz_dict parameter (passed from collection level) maps sample_id
1146        to a list of m/z values that belong to clusters for that sample.
1147
1148        Parameters
1149        ----------
1150        sample_id : int
1151            Sample index to process
1152        collection : LCMSBaseCollection
1153            The collection
1154        cluster_mz_dict : dict, optional
1155            Dictionary mapping sample_id to list of m/z values in clusters for that sample.
1156            If None, will not load any EICs. Default is None.
1157        **runtime_params
1158            Additional runtime parameters (ignored)
1159
1160        Returns
1161        -------
1162        dict
1163            Dictionary of loaded EIC_Data objects, keyed by m/z value
1164        """
1165        from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
1166
1167        sample = collection[sample_id]
1168
1169        # If no cluster info provided or no m/z values for this sample, return early
1170        if cluster_mz_dict is None or sample_id not in cluster_mz_dict:
1171            return {}
1172
1173        # Get m/z values for this sample that belong to clusters
1174        sample_cluster_mz = set(cluster_mz_dict[sample_id])
1175        if len(sample_cluster_mz) == 0:
1176            return {}
1177
1178        # Load EICs for each of the sample_cluster_mz
1179        hdf5_path = sample.file_location
1180        if hdf5_path and hdf5_path.exists():
1181            reader = ReadCoreMSHDFMassSpectra(str(hdf5_path))
1182            reader.import_eics(sample, mz_list=list(sample_cluster_mz))
1183            # Return the loaded EICs for multiprocessing collection
1184            # (modifications in worker process don't persist to main process)
1185            return sample.eics.copy()
1186
1187        return {}
1188
1189    def collect_results(self, sample_id, result, collection):
1190        """
1191        Collect loaded EICs back into sample.
1192
1193        In multiprocessing, the worker's modifications don't persist to the
1194        main process, so we need to explicitly collect and reassign the EICs.
1195        This also re-associates EICs with mass features.
1196
1197        Parameters
1198        ----------
1199        sample_id : int
1200            Sample ID that was processed
1201        result : dict
1202            Dictionary of EIC_Data objects keyed by m/z, returned from execute()
1203        collection : LCMSBaseCollection
1204            The collection
1205        """
1206        if result:
1207            # Update sample.eics with loaded EICs
1208            collection[sample_id].eics.update(result)
1209            # Note: EIC association with mass features happens after pipeline completes
1210            # to avoid multiprocessing issues (modifications in worker processes don't
1211            # persist to main process objects)