increase version 0.0.9 -> 0.0.10

Author: ypriverol

pyproject.toml

@@ -3,7 +3,7 @@ name = "quantms-rescoring"
 description = "quantms-rescoring: Python scripts and helpers for the quantMS workflow"
 readme = "README.md"
 license = "MIT"
-version = "0.0.9"
+version = "0.0.10"
 authors = [
     "Yasset Perez-Riverol <ypriverol@gmail.com>",
     "Dai Chengxin <chengxin2024@126.com>",

quantmsrescore/annotator.py

@@ -23,24 +23,12 @@ class FeatureAnnotator:
     from MS2PIP and DeepLC models to improve rescoring.
     """
 
-    def __init__(
-        self,
-        feature_generators: str,
-        only_features: Optional[str] = None,
-        ms2pip_model: str = "HCD2021",
-        force_model: bool = False,
-        ms2pip_model_path: str = "models",
-        ms2_tolerance: float = 0.05,
-        calibration_set_size: float = 0.2,
-        valid_correlations_size: float = 0.7,
-        skip_deeplc_retrain: bool = False,
-        processes: int = 2,
-        log_level: str = "INFO",
-        spectrum_id_pattern: str = "(.*)",  # default for openms idXML
-        psm_id_pattern: str = "(.*)",  # default for openms idXML
-        remove_missing_spectra: bool = True,
-        ms2_only: bool = True
-    ):
+    def __init__(self, feature_generators: str, only_features: Optional[str] = None, ms2pip_model: str = "HCD2021",
+                 force_model: bool = False, ms2pip_model_path: str = "models", ms2_tolerance: float = 0.05,
+                 calibration_set_size: float = 0.2, valid_correlations_size: float = 0.7,
+                 skip_deeplc_retrain: bool = False, processes: int = 2, log_level: str = "INFO",
+                 spectrum_id_pattern: str = "(.*)", psm_id_pattern: str = "(.*)", remove_missing_spectra: bool = True,
+                 ms2_only: bool = True, find_best_model: bool = False) -> None:
         """
         Initialize the Annotator with configuration parameters.
 
@@ -74,6 +62,10 @@ def __init__(
             Process only MS2-level PSMs (default: True).
         find_best_ms2pip_model : bool, optional
             Find best MS2PIP model for the dataset (default: False).
+        force_model : bool, optional
+            Force the use of the provided MS2PIP model (default: False).
+        find_best_model : bool, optional
+            Force the use of the best MS2PIP model (default: False).
 
         Raises
         ------
@@ -117,6 +109,7 @@ def __init__(
         self._remove_missing_spectra = remove_missing_spectra
         self._ms2_only = ms2_only
         self._force_model = force_model
+        self._find_best_model = find_best_model
 
     def build_idxml_data(
         self, idxml_file: Union[str, Path], spectrum_path: Union[str, Path]
@@ -233,20 +226,53 @@ def _run_ms2pip_annotation(self) -> None:
             logger.error(f"Failed to initialize MS2PIP: {e}")
             raise
 
-        # Apply MS2PIP annotation
+        # Get PSM list
         psm_list = self._idxml_reader.psms
+
         try:
-            ms2pip_generator.add_features(psm_list)
-            self._idxml_reader.psms = psm_list
-            logger.info("MS2PIP annotations added to PSMs")
-        except Ms2pipIncorrectModelException:
-            if not self._force_model:
-                self._find_and_apply_best_ms2pip_model(psm_list)
+            # Save original model for reference
+            original_model = ms2pip_generator.model
+
+            # Determine which model to use based on configuration and validation
+            model_to_use = original_model
+
+            # Case 1: Force specific model regardless of validation
+            if self._force_model:
+                model_to_use = original_model
+                logger.info(f"Using forced model: {model_to_use}")
+
+            # Case 2: Find best model if requested and not forcing original
+            elif self._find_best_model:
+                best_model = ms2pip_generator._find_best_ms2pip_model(psm_list)
+                if ms2pip_generator.validate_features(psm_list=psm_list, model=best_model[0]):
+                    model_to_use = best_model[0]
+                    logger.info(f"Using best model: {model_to_use}")
+                else:
+                    # Fallback to original model if best model doesn't validate
+                    if ms2pip_generator.validate_features(psm_list, model=original_model):
+                        logger.warning("Best model validation failed, falling back to original model")
+                    else:
+                        logger.error("Both best model and original model validation failed")
+                        return  # Exit early since no valid model is available
+
+            # Case 3: Use original model but validate it first
             else:
-                logger.error("MS2PIP model not suitable for this data. Skip MS2PIP annotations.")
+                if not ms2pip_generator.validate_features(psm_list):
+                    logger.error("Original model validation failed. No features added.")
+                    return  # Exit early since validation failed
+                logger.info(f"Using original model: {model_to_use}")
+
+            # Apply the selected model
+            ms2pip_generator.model = model_to_use
+            ms2pip_generator.add_features(psm_list)
+            logger.info(f"Successfully applied MS2PIP annotation using model: {model_to_use}")
 
         except Exception as e:
-            logger.error(f"Failed to add MS2PIP features: {e}")
+            logger.error(f"Failed to apply MS2PIP annotation: {e}")
+            return  # Indicate failure through early return
+
+        return  # Successful completion
+
 
     def _create_ms2pip_annotator(
         self, model: Optional[str] = None, tolerance: Optional[float] = None
@@ -580,4 +606,4 @@ def _get_mae_from_psm_list(self, psm_list: PSMList) -> float:
             logger.warning("No valid retention time differences for MAE calculation")
             return float("inf")
 
-        return total_error / count
+        return total_error / count