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    Uses time range filtering for efficient data loading - only loads the 
 184    retention time windows where gaps need to be filled, plus a buffer 
 185    (controlled by eic_buffer_time parameter) for complete EIC extraction. 
 186    Multiple time ranges are automatically merged if they overlap.
 187    
 188    Parameters
 189    ----------
 190    name : str
 191        Operation name
 192    expand_on_miss : bool, optional
 193        If True, expands search window when no peak is found. Default is False.
 194        
 195    Notes
 196    -----
 197    Requires that add_consensus_mass_features() has been run on the collection.
 198    This operation loads raw MS1 data which will be available for subsequent operations.
 199    Time range filtering significantly reduces memory usage and loading time for
 200    large datasets with sparse gaps.
 201    """
 202    
 203    @property
 204    def description(self):
 205        """Human-readable description for progress messages."""
 206        return "gap-filling"
 207    
 208    def needs_raw_ms_data(self):
 209        """This operation needs raw MS1 data."""
 210        return True, 1
 211    
 212    def can_execute(self, sample, collection):
 213        """Check if cluster summary exists."""
 214        return hasattr(collection, 'cluster_summary_dataframe') and \
 215               collection.cluster_summary_dataframe is not None
 216    
 217    def execute(self, sample_id, collection, **runtime_params):
 218        """
 219        Execute gap-filling for a single sample.
 220        
 221        Parameters
 222        ----------
 223        sample_id : int
 224            Sample index to process
 225        collection : LCMSBaseCollection
 226            The collection
 227        **runtime_params
 228            Runtime parameters including:
 229            - missingdf : pd.DataFrame - Cluster information and missing samples (optional)
 230            - cluster_dict : dict - Cluster feature dictionary (optional)
 231            - expand_on_miss : bool - Whether to expand search window on miss (optional)
 232            If these are not provided, returns empty dict (no gaps to fill).
 233            
 234        Returns
 235        -------
 236        dict
 237            Dictionary of induced mass features (empty if no gaps to fill)
 238        """
 239        # Extract gap-fill parameters from runtime_params
 240        # If not present, there are no gaps to fill, so return early
 241        if 'missingdf' not in runtime_params:
 242            return {}
 243        
 244        missingdf = runtime_params['missingdf']
 245        cluster_dict = runtime_params['cluster_dict']
 246        expand_on_miss = runtime_params['expand_on_miss']
 247        
 248        # This is essentially the same logic as _search_for_targeted_mass_features_in_sample
 249        # but extracted into an operation
 250        
 251        # Get clusters missing data for this sample
 252        sampledf = missingdf[
 253            missingdf.missing_samples.apply(lambda x: sample_id in x)
 254        ].reset_index(drop=True).copy()
 255
 256        # Skip if no missing features for this sample
 257        if len(sampledf) == 0:
 258            return {}
 259
 260        # Get buffer time from LCMS parameters for EIC extraction
 261        # This ensures we capture the full chromatographic peak beyond cluster bounds
 262        buffer_rt = collection[sample_id].parameters.lc_ms.eic_buffer_time
 263        
 264        # Calculate time ranges for efficient loading with buffer for EIC extraction
 265        time_ranges = []
 266        
 267        for _, row in sampledf.iterrows():
 268            rt_min = row['scan_time_aligned_min']
 269            rt_max = row['scan_time_aligned_max']
 270            
 271            # If expand_on_miss, also consider the allowed bounds
 272            if expand_on_miss:
 273                rt_min = min(rt_min, row['sta_min_allowed'])
 274                rt_max = max(rt_max, row['sta_max_allowed'])
 275            
 276            # Apply buffer AFTER considering expand_on_miss bounds
 277            # This ensures buffer is added beyond even the expanded search window
 278            time_ranges.append((max(0, rt_min - buffer_rt), rt_max + buffer_rt))
 279        
 280        # Merge overlapping time ranges to reduce number of separate loads
 281        time_ranges = sorted(time_ranges)
 282        merged_ranges = []
 283        if time_ranges:
 284            current_min, current_max = time_ranges[0]
 285            for rt_min, rt_max in time_ranges[1:]:
 286                if rt_min <= current_max:  # Overlapping or adjacent
 287                    current_max = max(current_max, rt_max)
 288                else:
 289                    merged_ranges.append((current_min, current_max))
 290                    current_min, current_max = rt_min, rt_max
 291            merged_ranges.append((current_min, current_max))
 292
 293        # Load raw data for this sample with time range filtering
 294        collection.load_raw_data(sample_id, 1, time_range=merged_ranges)
 295        
 296        # Get MS1 data
 297        ms1df = collection[sample_id]._ms_unprocessed[1].copy()
 298        scan_df = collection[sample_id].scan_df[['scan', 'scan_time_aligned']]
 299        ms1df = pd.merge(ms1df, scan_df, on='scan')
 300
 301        # Pre-extract all values from sampledf
 302        clusters = sampledf.cluster.values
 303        mz_mins = sampledf.mz_min.values
 304        mz_maxs = sampledf.mz_max.values
 305        st_mins = sampledf.scan_time_aligned_min.values
 306        st_maxs = sampledf.scan_time_aligned_max.values
 307        
 308        if expand_on_miss:
 309            mz_mins_allowed = sampledf.mz_min_allowed.values
 310            mz_maxs_allowed = sampledf.mz_max_allowed.values
 311            st_mins_allowed = sampledf.sta_min_allowed.values
 312            st_maxs_allowed = sampledf.sta_max_allowed.values
 313
 314        # Pre-filter ms1df to reduce search space
 315        mz_global_min = mz_mins.min()
 316        mz_global_max = mz_maxs.max()
 317        st_global_min = st_mins.min()
 318        st_global_max = st_maxs.max()
 319        
 320        if expand_on_miss:
 321            mz_global_min = min(mz_global_min, mz_mins_allowed.min())
 322            mz_global_max = max(mz_global_max, mz_maxs_allowed.max())
 323            st_global_min = min(st_global_min, st_mins_allowed.min())
 324            st_global_max = max(st_global_max, st_maxs_allowed.max())
 325        
 326        ms1df_filtered = ms1df[
 327            (ms1df.mz >= mz_global_min) & 
 328            (ms1df.mz <= mz_global_max) &
 329            (ms1df.scan_time_aligned >= st_global_min) &
 330            (ms1df.scan_time_aligned <= st_global_max)
 331        ].copy()
 332
 333        # Generate set_ids for all features
 334        set_ids = [f'c{clusters[i]}_{i}_i' for i in range(len(sampledf))]
 335        
 336        # Use batch method to process all features at once
 337        if expand_on_miss:
 338            # First try with normal bounds
 339            peaks_dict = collection[sample_id].search_for_targeted_mass_features_batch(
 340                ms1df_filtered,
 341                mz_mins,
 342                mz_maxs,
 343                st_mins,
 344                st_maxs,
 345                set_ids,
 346                obj_idx=sample_id,
 347                st_aligned=True
 348            )
 349            
 350            # Retry failed features with expanded bounds
 351            failed_indices = [i for i, sid in enumerate(set_ids) if peaks_dict[sid].apex_scan == -99]
 352            if failed_indices:
 353                failed_ids = [set_ids[i] for i in failed_indices]
 354                retry_peaks = collection[sample_id].search_for_targeted_mass_features_batch(
 355                    ms1df_filtered,
 356                    mz_mins_allowed[failed_indices],
 357                    mz_maxs_allowed[failed_indices],
 358                    st_mins_allowed[failed_indices],
 359                    st_maxs_allowed[failed_indices],
 360                    failed_ids,
 361                    obj_idx=sample_id,
 362                    st_aligned=True
 363                )
 364                peaks_dict.update(retry_peaks)
 365        else:
 366            peaks_dict = collection[sample_id].search_for_targeted_mass_features_batch(
 367                ms1df_filtered,
 368                mz_mins,
 369                mz_maxs,
 370                st_mins,
 371                st_maxs,
 372                set_ids,
 373                obj_idx=sample_id,
 374                st_aligned=True
 375            )
 376        
 377        # Build induced_mass_features dict and update cluster_dict
 378        induced_mass_features = {}
 379        for i in range(len(sampledf)):
 380            peak = peaks_dict[set_ids[i]]
 381            induced_mass_features[peak.id] = peak
 382            cluster_dict[clusters[i]] += [set_ids[i]]
 383        
 384        # Integrate mass features (don't fail on bad integration)
 385        collection[sample_id].induced_mass_features = induced_mass_features
 386        collection[sample_id].integrate_mass_features(drop_if_fail=False, induced_features=True)
 387        
 388        # Add MS1 spectra and peak metrics to successfully detected induced features
 389        # Only process features that were successfully detected (apex_scan != -99)
 390        # This is critical for having m/z values in the pivot table for gap-filled features
 391        successful_induced = {k: v for k, v in induced_mass_features.items() if v.apex_scan != -99}
 392        
 393        if len(successful_induced) > 0:
 394            # Use the already-loaded raw data (use_parser=False) for efficiency
 395            collection[sample_id].add_associated_ms1(
 396                auto_process=True,
 397                use_parser=False,
 398                spectrum_mode=None,
 399                induced_features=True
 400            )
 401        
 402        # Return the induced features (some may have been filtered out)
 403        return collection[sample_id].induced_mass_features
 404        
 405    def collect_results(self, sample_id, result, collection):
 406        """Collect induced mass features back into sample."""
 407        collection[sample_id].induced_mass_features = result
 408
 409
 410class ReloadFeaturesOperation(SampleOperation):
 411    """
 412    Reload mass features from HDF5 and optionally add MS1/MS2 spectra.
 413    
 414    This is useful when the collection was loaded with load_light=True,
 415    which stores mass features only in the collection dataframe and not
 416    as LCMSMassFeature objects in individual samples.
 417    
 418    Parameters
 419    ----------
 420    name : str
 421        Operation name
 422    add_ms1 : bool, optional
 423        If True, adds MS1 spectra to mass features. Automatically uses raw MS1 data
 424        if available (e.g., from gap-filling), otherwise uses parser. Spectrum mode
 425        is auto-detected. Default is False.
 426    add_ms2 : bool, optional
 427        If True, also loads and associates MS2 spectra. Spectrum mode is auto-detected.
 428        Default is False.
 429    auto_process_ms2 : bool, optional
 430        If True and add_ms2=True, auto-processes MS2 spectra. Default is True.
 431    ms2_scan_filter : str or None, optional
 432        Filter string for MS2 scans. Default is None.
 433        
 434    Notes
 435    -----
 436    MS1 spectra association automatically uses raw MS1 data if loaded by a previous
 437    operation (e.g., GapFillOperation). This is efficient when multiple operations
 438    need MS1 data in the same pipeline. All spectrum modes are auto-detected from
 439    the data.
 440    """
 441    
 442    @property
 443    def description(self):
 444        """Human-readable description for progress messages."""
 445        return "reloading features"
 446    
 447    def needs_raw_ms_data(self):
 448        """This operation doesn't need raw data."""
 449        return False, None
 450    
 451    def can_execute(self, sample, collection):
 452        """Check if collection parser is available."""
 453        return hasattr(collection, 'collection_parser') and \
 454               collection.collection_parser is not None
 455    
 456    def execute(self, sample_id, collection, mf_ids_to_load=None, **runtime_params):
 457        """
 458        Execute feature reloading for a single sample.
 459        
 460        Parameters
 461        ----------
 462        sample_id : int
 463            Sample ID to reload features for
 464        collection : LCMSBaseCollection
 465            The collection
 466        mf_ids_to_load : list of str, optional
 467            List of collection-level mf_ids to load
 468        **runtime_params
 469            Additional runtime parameters (ignored)
 470            
 471        Returns
 472        -------
 473        dict
 474            Dictionary of reloaded mass features
 475        """
 476        # Get parameters
 477        add_ms1 = self.params.get('add_ms1', False)
 478        add_ms2 = self.params.get('add_ms2', False)
 479        auto_process_ms2 = self.params.get('auto_process_ms2', True)
 480        ms2_scan_filter = self.params.get('ms2_scan_filter', None)
 481        
 482        sample = collection[sample_id]
 483        sample_name = collection.samples[sample_id]
 484        
 485        # Auto-determine if we should use parser for MS1 (check if raw data is available)
 486        has_raw_ms1 = 1 in sample._ms_unprocessed and not sample._ms_unprocessed[1].empty
 487        use_parser_for_ms1 = not has_raw_ms1  # Use parser only if raw data not available
 488        
 489        # Spectrum modes will be auto-detected (None = auto-detect)
 490        spectrum_mode_ms1 = None
 491        ms2_spectrum_mode = None
 492        
 493        # Check if we have a collection parser
 494        if not hasattr(collection, 'collection_parser') or collection.collection_parser is None:
 495            print(f"Warning: Cannot reload mass features for {sample_name} - no collection_parser available")
 496            return {}
 497        
 498        # Get the HDF5 file for this sample
 499        hdf5_file = collection.collection_parser.folder_location / f"{sample_name}.corems/{sample_name}.hdf5"
 500        
 501        if not hdf5_file.exists():
 502            print(f"Warning: HDF5 file not found for sample {sample_name}: {hdf5_file}")
 503            return {}
 504        
 505        # Import here to avoid circular imports
 506        from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
 507        
 508        # If specific mf_ids requested, extract the local mf_ids we need
 509        local_mf_ids_to_load = None
 510        if mf_ids_to_load is not None:
 511            # mf_ids_to_load is already a list of sample-level mf_ids (integers)
 512            # No parsing needed - they come from the mf_id column in the dataframe
 513            if len(mf_ids_to_load) == 0:
 514                # No features to load for this sample - return empty dict
 515                return {}
 516            local_mf_ids_to_load = set(mf_ids_to_load)
 517        
 518        # Reload mass features from HDF5
 519        with ReadCoreMSHDFMassSpectra(hdf5_file) as parser:
 520            parser.import_mass_features(sample, mf_ids=local_mf_ids_to_load)
 521        
 522        # If add_ms1, associate MS1 spectra with the loaded mass features
 523        if add_ms1 and len(sample.mass_features) > 0:
 524            # Check if raw MS1 data is already loaded (e.g., from gap-filling)
 525            has_raw_ms1 = 1 in sample._ms_unprocessed and not sample._ms_unprocessed[1].empty
 526            
 527            if has_raw_ms1 and not use_parser_for_ms1:
 528                # Use already-loaded raw data (more efficient)
 529                sample.add_associated_ms1(
 530                    auto_process=True,
 531                    use_parser=False,
 532                    spectrum_mode=spectrum_mode_ms1
 533                )
 534            else:
 535                # Use parser to get MS1 spectra
 536                sample.add_associated_ms1(
 537                    auto_process=True,
 538                    use_parser=True,
 539                    spectrum_mode=spectrum_mode_ms1
 540                )
 541        
 542        # If add_ms2, associate MS2 spectra with the loaded mass features
 543        if add_ms2 and len(sample.mass_features) > 0:
 544            # Get the IDs of loaded mass features (use what was actually loaded)
 545            mf_ids_for_ms2 = list(sample.mass_features.keys())
 546            
 547            collection._associate_ms2_with_mass_features(
 548                sample, 
 549                mf_ids_for_ms2,
 550                auto_process=auto_process_ms2,
 551                spectrum_mode=ms2_spectrum_mode,
 552                scan_filter=ms2_scan_filter
 553            )
 554        
 555        # Return both mass_features and _ms so they can be collected in multiprocessing
 556        return {'mass_features': sample.mass_features, '_ms': sample._ms}
 557        
 558    def collect_results(self, sample_id, result, collection):
 559        """
 560        Collect reloaded mass features back into sample.
 561        
 562        This operation loads a subset of mass features (e.g., representatives)
 563        into sample.mass_features for processing, while preserving the full
 564        mass_features_dataframe at the collection level. Sets a lock flag to
 565        prevent automatic rebuilding of the collection dataframe from individual
 566        samples. Also collects loaded mass spectra.
 567        
 568        Parameters
 569        ----------
 570        sample_id : int
 571            Sample ID that was processed
 572        result : dict
 573            Dictionary with 'mass_features' and '_ms' from execute()
 574        collection : LCMSBaseCollection
 575            The collection
 576        """
 577        # Update sample.mass_features with loaded features
 578        if isinstance(result, dict) and 'mass_features' in result:
 579            collection[sample_id].mass_features = result['mass_features']
 580            # Also collect the _ms dictionary (MS1 and MS2 spectra)
 581            if '_ms' in result:
 582                collection[sample_id]._ms.update(result['_ms'])
 583        else:
 584            # Backward compatibility - if result is just mass_features dict
 585            collection[sample_id].mass_features = result
 586        
 587        # Lock the collection dataframe to prevent rebuilding from individual samples
 588        # (since we've only loaded a subset, rebuilding would lose data)
 589        collection._mass_features_locked = True
 590
 591
 592class MolecularFormulaSearchOperation(SampleOperation):
 593    """
 594    Perform molecular formula search on mass features using associated MS1 spectra.
 595    
 596    This operation runs molecular formula search on all mass features in a sample
 597    that have associated MS1 spectra. Requires MS1 spectra to be loaded and
 598    processed before execution.
 599    
 600    Parameters
 601    ----------
 602    name : str
 603        Operation name (for logging)
 604    **kwargs
 605        Additional parameters passed to parent class
 606        
 607    Examples
 608    --------
 609    >>> op = MolecularFormulaSearchOperation('mf_search')
 610    >>> # Use in pipeline
 611    >>> results = collection.process_samples_pipeline([op])
 612    
 613    Notes
 614    -----
 615    This operation requires that MS1 spectra have been associated with mass
 616    features (e.g., via ReloadFeaturesOperation with add_ms1=True). The
 617    molecular formula search uses parameters from the collection's 
 618    parameters.mass_spectrum["ms1"].molecular_search settings.
 619    """
 620    
 621    @property
 622    def description(self):
 623        """Human-readable description for progress messages."""
 624        return "molecular formula search"
 625    
 626    def __init__(self, name='molecular_formula_search', **kwargs):
 627        super().__init__(name, **kwargs)
 628    
 629    def needs_raw_ms_data(self):
 630        """
 631        This operation doesn't need raw data - it works on processed MS1 spectra
 632        that are already associated with mass features.
 633        
 634        Returns
 635        -------
 636        tuple
 637            (False, None) - no raw data needed
 638        """
 639        return False, None
 640    
 641    def can_execute(self, sample, collection, **runtime_params):
 642        """
 643        Check if molecular formula search can be executed.
 644        
 645        Requires that the sample has mass features with associated MS1 spectra.
 646        
 647        Parameters
 648        ----------
 649        sample : LCMSObject
 650            The sample object
 651        collection : LCMSCollection
 652            The collection containing the sample
 653        **runtime_params
 654            Runtime parameters (not used)
 655            
 656        Returns
 657        -------
 658        bool
 659            True if sample has mass features with MS1 spectra
 660        """        
 661        # Check if sample has mass features
 662        if not hasattr(sample, 'mass_features') or not sample.mass_features:
 663            return False
 664        
 665        # Check if at least some mass features have MS1 spectra
 666        has_ms1 = any(
 667            hasattr(mf, 'mass_spectrum') and mf.mass_spectrum is not None
 668            for mf in sample.mass_features.values()
 669        )
 670        
 671        return has_ms1
 672    
 673    def execute(self, sample_id, collection, **runtime_params):
 674        """
 675        Execute molecular formula search on a sample.
 676        
 677        Creates a SearchMolecularFormulasLC object and runs mass feature search,
 678        which annotates mass features with molecular formula assignments.
 679        
 680        Parameters
 681        ----------
 682        sample_id : str
 683            Sample identifier
 684        collection : LCMSCollection
 685            The collection containing the sample
 686        **runtime_params
 687            Runtime parameters (not used)
 688            
 689        Returns
 690        -------
 691        int
 692            Number of mass features that were searched
 693        """
 694        from corems.molecular_id.search.molecularFormulaSearch import SearchMolecularFormulasLC
 695        import time
 696        import sqlalchemy.exc
 697        import sqlite3
 698        
 699        sample = collection[sample_id]
 700        
 701        # Verify that mass features exist
 702        if not hasattr(sample, 'mass_features') or not sample.mass_features:
 703            return 0  # No mass features to search
 704        
 705        # Verify that mass features have MS1 spectra associated
 706        if not hasattr(sample, '_ms') or not sample._ms:
 707            raise RuntimeError(
 708                f"Sample {sample_id} does not have MS1 spectra loaded in _ms dictionary. "
 709                "Molecular formula search requires MS1 spectra to be associated with mass features. "
 710                "Ensure add_ms1=True when reloading features."
 711            )
 712        
 713        # Prepare data for bulk molecular formula search
 714        # Group mass features by their apex scan
 715        scan_to_mf = {}
 716        for mf_id, mf in sample.mass_features.items():
 717            apex_scan = mf.apex_scan
 718            if apex_scan not in scan_to_mf:
 719                scan_to_mf[apex_scan] = []
 720            scan_to_mf[apex_scan].append(mf)
 721        
 722        # Build lists of mass spectra and corresponding peaks
 723        mass_spectrum_list = []
 724        ms_peaks_list = []
 725        
 726        for scan_num, mf_list in scan_to_mf.items():
 727            # Get the mass spectrum for this scan
 728            if scan_num not in sample._ms:
 729                continue  # Skip if spectrum not loaded
 730                
 731            mass_spectrum = sample._ms[scan_num]
 732            
 733            # Verify spectrum is processed (has peaks)
 734            if not hasattr(mass_spectrum, '_mspeaks') or not mass_spectrum._mspeaks:
 735                continue  # Skip unprocessed spectra
 736            
 737            # Get the MS1 peaks for each mass feature at this scan
 738            peaks_for_scan = []
 739            for mf in mf_list:
 740                try:
 741                    # Use the ms1_peak property which finds the closest peak
 742                    ms1_peak = mf.ms1_peak
 743                    peaks_for_scan.append(ms1_peak)
 744                except (AttributeError, IndexError):
 745                    # Skip if ms1_peak can't be determined
 746                    continue
 747            
 748            if peaks_for_scan:
 749                mass_spectrum_list.append(mass_spectrum)
 750                ms_peaks_list.append(peaks_for_scan)
 751        
 752        # Run molecular formula search if we have data, with retry logic for database locks
 753        if mass_spectrum_list and ms_peaks_list:
 754            max_retries = 10
 755            retry_delay = 2  # seconds
 756            
 757            for attempt in range(max_retries):
 758                try:
 759                    mol_search = SearchMolecularFormulasLC(sample)
 760                    mol_search.bulk_run_molecular_formula_search(mass_spectrum_list, ms_peaks_list)
 761                    break  # Success, exit retry loop
 762                except (sqlalchemy.exc.OperationalError, sqlite3.OperationalError) as e:
 763                    if attempt < max_retries - 1:
 764                        # Database is locked, retry after delay
 765                        print(f"Sample {sample_id}: Database locked during molecular formula search, retrying in {retry_delay}s (attempt {attempt + 1}/{max_retries})...")
 766                        time.sleep(retry_delay)
 767                    else:
 768                        # Max retries exceeded, re-raise the exception
 769                        raise RuntimeError(
 770                            f"Sample {sample_id}: Molecular formula search failed after {max_retries} attempts due to database lock. "
 771                            "Try reducing parallel cores or increasing database timeout."
 772                        ) from e
 773        
 774        # Return count of features searched
 775        return len(sample.mass_features)
 776    
 777    def collect_results(self, sample_id, result, collection):
 778        """
 779        Collect results (no-op as search modifies mass features in place).
 780        
 781        The molecular formula search modifies mass features in place by adding
 782        molecular formula assignments, so no explicit result collection is needed.
 783        
 784        Parameters
 785        ----------
 786        sample_id : str
 787            Sample identifier
 788        result : int
 789            Number of features searched
 790        collection : LCMSCollection
 791            The collection containing the sample
 792        """
 793        # Search modifies mass features in place, nothing to collect
 794        pass
 795
 796
 797class MS2SpectralSearchOperation(SampleOperation):
 798    """
 799    Perform MS2 spectral search using entropy-based matching.
 800    
 801    This operation performs spectral library search on MS2 spectra associated
 802    with mass features using FlashEntropy for fast similarity scoring. Requires
 803    MS2 spectra to be loaded and processed before execution.
 804    
 805    Parameters
 806    ----------
 807    name : str
 808        Operation name (for logging)
 809    ms2_scan_filter : str or None, optional
 810        Filter string for MS2 scans (e.g., 'hcd'). If None, uses all MS2 scans.
 811        Default is None.
 812    peak_sep_da : float, optional
 813        Peak separation in Daltons for spectral matching. Default is 0.01.
 814    **kwargs
 815        Additional parameters passed to parent class
 816        
 817    Examples
 818    --------
 819    >>> op = MS2SpectralSearchOperation('ms2_search', ms2_scan_filter='hcd')
 820    >>> # Use in pipeline - requires fe_lib in runtime_params
 821    >>> results = collection.process_samples_pipeline([op])
 822    
 823    Notes
 824    -----
 825    This operation requires:
 826    - MS2 spectra to be associated with mass features
 827    - FlashEntropy library (fe_lib) to be provided in runtime_params
 828    - MS2 spectra must be processed (centroided)
 829    
 830    The spectral search modifies mass features in place by adding spectral
 831    match scores and metadata.
 832    """
 833    
 834    @property
 835    def description(self):
 836        """Human-readable description for progress messages."""
 837        return "MS2 spectral search"
 838    
 839    def __init__(self, name='ms2_spectral_search', ms2_scan_filter=None, **kwargs):
 840        super().__init__(name, **kwargs)
 841        self.params['ms2_scan_filter'] = ms2_scan_filter
 842    
 843    def needs_raw_ms_data(self):
 844        """
 845        This operation doesn't need raw data - it works on processed MS2 spectra
 846        that are already associated with mass features.
 847        
 848        Returns
 849        -------
 850        tuple
 851            (False, None) - no raw data needed
 852        """
 853        return False, None
 854    
 855    def can_execute(self, sample, collection, **runtime_params):
 856        """
 857        Check if MS2 spectral search can be executed.
 858        
 859        Requires that the sample has mass features with MS2 spectra associated.
 860        
 861        Parameters
 862        ----------
 863        sample : LCMSObject
 864            The sample object
 865        collection : LCMSCollection
 866            The collection containing the sample
 867        **runtime_params
 868            Runtime parameters (not used)
 869            
 870        Returns
 871        -------
 872        bool
 873            True if sample has mass features with MS2 spectra
 874        """        
 875        # Check if sample has mass features
 876        if not hasattr(sample, 'mass_features') or not sample.mass_features:
 877            return False
 878        
 879        # Check if any mass features have MS2 spectra associated
 880        has_ms2 = any(
 881            hasattr(mf, 'ms2_mass_spectra') and mf.ms2_mass_spectra
 882            for mf in sample.mass_features.values()
 883        )
 884        
 885        return has_ms2
 886    
 887    def execute(self, sample_id, collection, fe_lib=None, molecular_metadata=None, **runtime_params):
 888        """
 889        Execute MS2 spectral search on a sample.
 890        
 891        Performs entropy-based spectral library search on all MS2 spectra
 892        in the sample that match the scan filter criteria.
 893        
 894        Parameters
 895        ----------
 896        sample_id : str
 897            Sample identifier
 898        collection : LCMSCollection
 899            The collection containing the sample
 900        fe_lib : FlashEntropy library
 901            Pre-computed FlashEntropy library for spectral matching
 902        molecular_metadata : pd.DataFrame, optional
 903            Metadata for molecules in the spectral library
 904        **runtime_params
 905            Runtime parameters (not used)
 906            
 907        Returns
 908        -------
 909        int
 910            Number of MS2 spectra searched
 911        """
 912        sample = collection[sample_id]
 913        
 914        # Get parameters
 915        ms2_scan_filter = self.params.get('ms2_scan_filter')
 916        
 917        # Verify that we have a spectral library
 918        if fe_lib is None:
 919            raise ValueError(
 920                f"Sample {sample_id}: MS2 spectral search requires fe_lib (FlashEntropy library) "
 921                "to be provided in runtime parameters. Create the library at the collection level "
 922                "and pass it to the pipeline."
 923            )
 924        
 925        # Extract peak_sep_da from FlashEntropy library configuration
 926        # peak_sep_da should be 2 * max_ms2_tolerance_in_da to match the min_ms2_difference_in_da
 927        # used when creating the library
 928        tolerance_da = fe_lib.entropy_search.max_ms2_tolerance_in_da
 929        if tolerance_da is None:
 930            raise ValueError(
 931                f"Sample {sample_id}: Could not extract max_ms2_tolerance_in_da from FlashEntropy library. "
 932                "Ensure the library was created with this parameter specified."
 933            )
 934        peak_sep_da = 2 * tolerance_da
 935        
 936        # Verify that sample has _ms dictionary
 937        if not hasattr(sample, '_ms') or not sample._ms:
 938            return 0  # No MS2 spectra to search
 939        
 940        # Get MS2 scan numbers based on filter
 941        if ms2_scan_filter is not None:
 942            # Filter by scan text
 943            ms2_scan_df = sample.scan_df[
 944                sample.scan_df.scan_text.str.contains(ms2_scan_filter) &
 945                (sample.scan_df.ms_level == 2)
 946            ]
 947        else:
 948            # All MS2 scans
 949            ms2_scan_df = sample.scan_df[sample.scan_df.ms_level == 2]
 950        
 951        # Get scans that are actually loaded in _ms
 952        ms2_scans_to_search = [
 953            scan for scan in ms2_scan_df.scan.tolist()
 954            if scan in sample._ms.keys()
 955        ]
 956        
 957        if not ms2_scans_to_search:
 958            return 0  # No MS2 spectra to search
 959        
 960        # Perform spectral search using the sample's fe_search method
 961        sample.fe_search(
 962            scan_list=ms2_scans_to_search,
 963            fe_lib=fe_lib,
 964            peak_sep_da=peak_sep_da
 965        )
 966        
 967        # Return the spectral search results for collection
 968        # (needed for multiprocessing - results populated in worker need to be returned)
 969        return sample.spectral_search_results
 970    
 971    def collect_results(self, sample_id, result, collection):
 972        """
 973        Collect spectral search results back into the sample.
 974        
 975        In multiprocessing, the worker's modifications don't persist to the
 976        main process, so we need to explicitly collect and reassign the results.
 977        This also re-associates the results with mass features.
 978        
 979        Parameters
 980        ----------
 981        sample_id : str
 982            Sample identifier
 983        result : dict
 984            Dictionary of spectral search results from execute()
 985        collection : LCMSCollection
 986            The collection containing the sample
 987        """
 988        # Assign the spectral search results back to the sample
 989        if result:
 990            collection[sample_id].spectral_search_results.update(result)
 991            
 992            # Re-associate results with mass features (same logic as fe_search)
 993            sample = collection[sample_id]
 994            if len(sample.mass_features) > 0:
 995                for mass_feature_id, mass_feature in sample.mass_features.items():
 996                    scan_ids = mass_feature.ms2_scan_numbers
 997                    for ms2_scan_id in scan_ids:
 998                        precursor_mz = mass_feature.mz
 999                        try:
1000                            sample.spectral_search_results[ms2_scan_id][precursor_mz]
1001                        except KeyError:
1002                            pass
1003                        else:
1004                            sample.mass_features[
1005                                mass_feature_id
1006                            ].ms2_similarity_results.append(
1007                                sample.spectral_search_results[ms2_scan_id][precursor_mz]
1008                            )
1009
1010
1011class LoadEICsOperation(SampleOperation):
1012    """
1013    Load extracted ion chromatograms (EICs) from HDF5 for regular mass features.
1014    
1015    Loads EICs for regular mass features that belong to consensus clusters from HDF5.
1016    Induced (gap-filled) features already have EICs from integrate_mass_features,
1017    so no additional loading is needed for them.
1018    
1019    This operation enables downstream visualization and analysis of chromatographic
1020    peaks across all samples in a cluster.
1021    
1022    Notes
1023    -----
1024    Requires that mass features have been loaded and cluster_index assigned.
1025    Regular mass feature EICs must have been previously saved to HDF5 with export_eics=True.
1026    Induced mass features already have EICs populated during gap-filling.
1027    """
1028    
1029    @property
1030    def description(self):
1031        """Human-readable description for progress messages."""
1032        return "loading EICs"
1033    
1034    def needs_raw_ms_data(self):
1035        """This operation doesn't need raw data - induced features already have EICs."""
1036        return False, None
1037    
1038    def can_execute(self, sample, collection):
1039        """
1040        Check if EIC loading can be executed.
1041        
1042        This operation can always execute if the sample exists - the actual work
1043        is determined by cluster_mz_dict in runtime_params. If cluster_mz_dict is
1044        empty or None, execute() will simply return 0 (no EICs loaded).
1045        
1046        Returns
1047        -------
1048        bool
1049            True (always executable - runtime_params control actual work)
1050        """
1051        return True
1052    
1053    def execute(self, sample_id, collection, cluster_mz_dict=None, **runtime_params):
1054        """
1055        Load EICs from HDF5 for a single sample.
1056        
1057        Loads EICs for regular mass features that belong to consensus clusters.
1058        Induced (gap-filled) mass features already have EICs from integrate_mass_features,
1059        so no additional loading is needed for them.
1060        
1061        The cluster_mz_dict parameter (passed from collection level) maps sample_id
1062        to a list of m/z values that belong to clusters for that sample.
1063        
1064        Parameters
1065        ----------
1066        sample_id : int
1067            Sample index to process
1068        collection : LCMSBaseCollection
1069            The collection
1070        cluster_mz_dict : dict, optional
1071            Dictionary mapping sample_id to list of m/z values in clusters for that sample.
1072            If None, will not load any EICs. Default is None.
1073        **runtime_params
1074            Additional runtime parameters (ignored)
1075            
1076        Returns
1077        -------
1078        dict
1079            Dictionary of loaded EIC_Data objects, keyed by m/z value
1080        """
1081        from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
1082        
1083        sample = collection[sample_id]
1084        
1085        # If no cluster info provided or no m/z values for this sample, return early
1086        if cluster_mz_dict is None or sample_id not in cluster_mz_dict:
1087            return {}
1088        
1089        # Get m/z values for this sample that belong to clusters
1090        sample_cluster_mz = set(cluster_mz_dict[sample_id])
1091        if len(sample_cluster_mz) == 0:
1092            return {}
1093
1094        # Load EICs for each of the sample_cluster_mz
1095        hdf5_path = sample.file_location
1096        if hdf5_path and hdf5_path.exists():
1097            reader = ReadCoreMSHDFMassSpectra(str(hdf5_path))
1098            reader.import_eics(sample, mz_list=list(sample_cluster_mz))
1099            # Return the loaded EICs for multiprocessing collection
1100            # (modifications in worker process don't persist to main process)
1101            return sample.eics.copy()
1102        
1103        return {}
1104    
1105    def collect_results(self, sample_id, result, collection):
1106        """
1107        Collect loaded EICs back into sample.
1108        
1109        In multiprocessing, the worker's modifications don't persist to the
1110        main process, so we need to explicitly collect and reassign the EICs.
1111        This also re-associates EICs with mass features.
1112        
1113        Parameters
1114        ----------
1115        sample_id : int
1116            Sample ID that was processed
1117        result : dict
1118            Dictionary of EIC_Data objects keyed by m/z, returned from execute()
1119        collection : LCMSBaseCollection
1120            The collection
1121        """
1122        if result:
1123            # Update sample.eics with loaded EICs
1124            collection[sample_id].eics.update(result)
1125            # Note: EIC association with mass features happens after pipeline completes
1126            # to avoid multiprocessing issues (modifications in worker processes don't
1127            # persist to main process objects)