merge im2deep into timsrescore branch

Dockerfile

@@ -1,8 +1,10 @@
-FROM ubuntu:focal
+FROM python:3.11
+
+# ARG DEBIAN_FRONTEND=noninteractive
 
 LABEL name="ms2rescore"
 
-ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/ms2rescore
+# ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/ms2rescore
 
 ADD pyproject.toml /ms2rescore/pyproject.toml
 ADD LICENSE /ms2rescore/LICENSE
@@ -11,8 +13,7 @@ ADD MANIFEST.in /ms2rescore/MANIFEST.in
 ADD ms2rescore /ms2rescore/ms2rescore
 
 RUN apt-get update \
-    && apt-get install -y python3-pip procps libglib2.0-0 libsm6 libxrender1 libxext6 \
-    && rm -rf /var/lib/apt/lists/* \
-    && pip3 install ms2rescore/
+    && apt install -y procps git-lfs \
+    && pip install /ms2rescore
 
 ENTRYPOINT [""]

docs/source/config_schema.md

@@ -10,6 +10,7 @@
     - **`deeplc`**: Refer to *[#/definitions/deeplc](#definitions/deeplc)*.
     - **`maxquant`**: Refer to *[#/definitions/maxquant](#definitions/maxquant)*.
     - **`ionmob`**: Refer to *[#/definitions/ionmob](#definitions/ionmob)*.
+    - **`im2deep`**: Refer to *[#/definitions/im2deep](#definitions/im2deep)*.
   - **`rescoring_engine`** *(object)*: Rescoring engine to use and its configuration. Leave empty to skip rescoring and write features to file. Default: `{"mokapot": {}}`.
     - **`.*`**: Refer to *[#/definitions/rescoring_engine](#definitions/rescoring_engine)*.
     - **`percolator`**: Refer to *[#/definitions/percolator](#definitions/percolator)*.
@@ -43,6 +44,14 @@
     - **One of**
       - *string*
       - *null*
+  - **`psm_id_rt_pattern`**: Regex pattern to extract index or scan number from spectrum file. Requires at least one capturing group. Default: `null`.
+    - **One of**
+      - *string*
+      - *null*
+  - **`psm_id_im_pattern`**: Regex pattern to extract index or scan number from spectrum file. Requires at least one capturing group. Default: `null`.
+    - **One of**
+      - *string*
+      - *null*
   - **`psm_id_pattern`**: Regex pattern to extract index or scan number from PSM file. Requires at least one capturing group. Default: `"(.*)"`.
     - **One of**
       - *string*
@@ -75,6 +84,8 @@
   - **`ionmob_model`** *(string)*: Path to Ionmob model directory. Default: `"GRUPredictor"`.
   - **`reference_dataset`** *(string)*: Path to Ionmob reference dataset file. Default: `"Meier_unimod.parquet"`.
   - **`tokenizer`** *(string)*: Path to tokenizer json file. Default: `"tokenizer.json"`.
+- <a id="definitions/im2deep"></a>**`im2deep`** *(object)*: Ion mobility feature generator configuration using IM2Deep. Can contain additional properties. Refer to *[#/definitions/feature_generator](#definitions/feature_generator)*.
+  - **`reference_dataset`** *(string)*: Path to IM2Deep reference dataset file. Default: `"Meier_unimod.parquet"`.
 - <a id="definitions/mokapot"></a>**`mokapot`** *(object)*: Mokapot rescoring engine configuration. Additional properties are passed to the Mokapot brew function. Can contain additional properties. Refer to *[#/definitions/rescoring_engine](#definitions/rescoring_engine)*.
   - **`write_weights`** *(boolean)*: Write Mokapot weights to a text file. Default: `false`.
   - **`write_txt`** *(boolean)*: Write Mokapot results to a text file. Default: `false`.

ms2rescore/core.py

@@ -55,8 +55,14 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
     )
 
     # TODO: avoid hard coding feature generators in some way
-    rt_required = "deeplc" in config["feature_generators"] and None in psm_list["retention_time"]
-    im_required = "ionmob" in config["feature_generators"] and None in psm_list["ion_mobility"]
+    rt_required = ("deeplc" in config["feature_generators"]) and (
+        None in psm_list["retention_time"]
+    )
+    im_required = ("ionmob" or "im2deep" in config["feature_generators"]) and (
+        None in psm_list["ion_mobility"]
+    )
+    logger.debug(f"RT required: {rt_required}, IM required: {im_required}")
+
     if rt_required or im_required:
         logger.info("Parsing missing retention time and/or ion mobility values from spectra...")
         get_missing_values(config, psm_list, missing_rt=rt_required, missing_im=im_required)

ms2rescore/feature_generators/__init__.py

@@ -7,11 +7,13 @@
 from ms2rescore.feature_generators.ionmob import IonMobFeatureGenerator
 from ms2rescore.feature_generators.maxquant import MaxQuantFeatureGenerator
 from ms2rescore.feature_generators.ms2pip import MS2PIPFeatureGenerator
+from ms2rescore.feature_generators.im2deep import IM2DeepFeatureGenerator
 
 FEATURE_GENERATORS = {
     "basic": BasicFeatureGenerator,
     "ms2pip": MS2PIPFeatureGenerator,
     "deeplc": DeepLCFeatureGenerator,
     "maxquant": MaxQuantFeatureGenerator,
     "ionmob": IonMobFeatureGenerator,
+    "im2deep": IM2DeepFeatureGenerator,
 }

ms2rescore/feature_generators/im2deep.py

@@ -0,0 +1,197 @@
+"""
+Placeholder
+"""
+
+import contextlib
+import logging
+import os
+from inspect import getfullargspec
+from itertools import chain
+from typing import List, Optional
+
+import numpy as np
+from im2deep.im2deep import predict_ccs
+from psm_utils import PSMList
+
+from ms2rescore.feature_generators.base import FeatureGeneratorBase
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+logger = logging.getLogger(__name__)
+
+
+class IM2DeepFeatureGenerator(FeatureGeneratorBase):
+    """IM2Deep collision cross section feature generator."""
+
+    def __init__(
+        self,
+        *args,
+        lower_score_is_better: bool = False,
+        spectrum_path: Optional[str] = None,
+        processes: int = 1,
+        calibrate_per_charge: bool = True,
+        **kwargs,
+    ):
+        """Placeholder"""  # TODO
+        super().__init__(*args, **kwargs)
+        self.lower_score_is_better = lower_score_is_better
+        self.spectrum_path = spectrum_path
+        self.processes = processes
+        self.deeplc_kwargs = kwargs or {}
+
+        self._verbose = logger.getEffectiveLevel() <= logging.DEBUG
+
+        # Lazy-load DeepLC
+        from deeplc import DeepLC
+
+        self.im2deep = DeepLC
+
+        # Remove any kwargs that are not DeepLC arguments
+        self.im2deep_kwargs = {
+            k: v for k, v in self.deeplc_kwargs.items() if k in getfullargspec(DeepLC).args
+        }
+        self.im2deep_kwargs.update({"config_file": None})
+
+        # TODO: Implement deeplc_retrain?
+
+        self.im2deep_predictor = None
+        self.calibrate_per_charge = calibrate_per_charge
+
+    @property
+    def feature_names(self) -> List[str]:
+        return [
+            "ccs_observed_im2deep",
+            "ccs_predicted_im2deep",
+            "ccs_error_im2deep",
+            "abs_ccs_error_im2deep",
+            "perc_ccs_error_im2deep",
+        ]
+
+    def add_features(self, psm_list: PSMList) -> None:
+        """Add IM2Deep-derived features to PSMs"""
+
+        logger.info("Adding IM2Deep-derived features to PSMs")
+
+        # Get easy-access nested version of PSMlist
+        psm_dict = psm_list.get_psm_dict()
+
+        # Run IM2Deep for each spectrum file
+        current_run = 1
+        total_runs = sum(len(runs) for runs in psm_dict.values())
+
+        for runs in psm_dict.values():
+            # Reset IM2Deep predictor for each collection of runs
+            self.im2deep_predictor = None
+            self.selected_model = None
+            for run, psms in runs.items():
+                logger.info(
+                    f"Running IM2Deep for PSMs from run ({current_run}/{total_runs}): `{run}`..."
+                )
+
+                # Disable wild logging to stdout by TensorFlow, unless in debug mode
+                with (
+                    contextlib.redirect_stdout(open(os.devnull, "w"))
+                    if not self._verbose
+                    else contextlib.nullcontext()
+                ):
+                    # Make new PSM list for this run (chain PSMs per spectrum to flat list)
+                    psm_list_run = PSMList(psm_list=list(chain.from_iterable(psms.values())))
+
+                    logger.debug("Calibrating IM2Deep...")
+
+                    # Convert ion mobility to CCS and calibrate CCS values
+                    psm_list_run_df = psm_list_run.to_dataframe()
+                    psm_list_run_df["charge"] = [
+                        peptidoform.precursor_charge
+                        for peptidoform in psm_list_run_df["peptidoform"]
+                    ]
+                    psm_list_run_df["ccs_observed"] = psm_list_run_df.apply(
+                        lambda x: self.im2ccs(
+                            x["ion_mobility"],
+                            x["precursor_mz"],  # TODO: Why does ionmob use calculated mz?
+                            x["charge"],
+                        ),
+                        axis=1,
+                    )
+
+                    # Create dataframe with high confidence hits for calibration
+                    cal_psm_df = self.make_cal_df(psm_list_run_df)
+
+                    # Make predictions with IM2Deep
+                    logger.debug("Predicting CCS values...")
+                    calibrated_predictions = predict_ccs(
+                        psm_list_run, cal_psm_df, write_output=False
+                    )
+
+                    # Add features to PSMs
+                    logger.debug("Adding features to PSMs...")
+                    predictions = calibrated_predictions
+                    observations = psm_list_run_df["ccs_observed"]
+                    ccs_diffs_run = np.abs(predictions - observations)
+                    for i, psm in enumerate(psm_list_run):
+                        psm["rescoring_features"].update(
+                            {
+                                "ccs_observed_im2deep": observations[i],
+                                "ccs_predicted_im2deep": predictions[i],
+                                "ccs_error_im2deep": ccs_diffs_run[i],
+                                "abs_ccs_error_im2deep": np.abs(ccs_diffs_run[i]),
+                                "perc_ccs_error_im2deep": np.abs(ccs_diffs_run[i])
+                                / observations[i]
+                                * 100,
+                            }
+                        )
+
+                current_run += 1
+
+    def im2ccs(self, reverse_im, mz, charge, mass_gas=28.013, temp=31.85, t_diff=273.15):
+        """
+        Convert ion mobility to CCS.  #TODO: Took this from ionmob. how to reference?
+
+        Parameters
+        ----------
+        reverse_im : float
+            Reduced ion mobility.
+        mz : float
+            Precursor m/z.
+        charge : int
+            Precursor charge.
+        mass_gas : float, optional
+            Mass of gas, by default 28.013
+        temp : float, optional
+            Temperature in Celsius, by default 31.85
+        t_diff : float, optional
+            Factor to convert Celsius to Kelvin, by default 273.15"""
+
+        SUMMARY_CONSTANT = 18509.8632163405
+        reduced_mass = (mz * charge * mass_gas) / (mz * charge + mass_gas)
+        return (SUMMARY_CONSTANT * charge) / (
+            np.sqrt(reduced_mass * (temp + t_diff)) * 1 / reverse_im
+        )
+
+    # TODO: Make threshold a parameter
+    def make_cal_df(self, psm_list_df, threshold=0.95):
+        """Make dataframe for calibration of IM2Deep predictions.
+
+        Parameters
+        ----------
+        psm_list_df : pd.DataFrame
+            DataFrame with PSMs.
+        threshold : float, optional
+            Threshold for high confidence hits, by default 0.95.
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with high confidence hits for calibration."""
+
+        psm_list_df = psm_list_df[
+            psm_list_df["charge"] < 5
+        ]  # predictions do not go higher for IM2Deep
+        high_conf_hits = list(
+            psm_list_df["spectrum_id"][psm_list_df["score"].rank(pct=True) > threshold]
+        )
+        logger.debug(
+            f"Number of high confidence hits for calculating shift: {len(high_conf_hits)}"
+        )
+        # Filter df for high_conf_hits
+        cal_psm_df = psm_list_df[psm_list_df["spectrum_id"].isin(high_conf_hits)]
+        return cal_psm_df

ms2rescore/gui/app.py

@@ -359,15 +359,20 @@ def __init__(self, *args, **kwargs):
         self.deeplc_config = DeepLCConfiguration(self)
         self.deeplc_config.grid(row=2, column=0, pady=(0, 20), sticky="nsew")
 
+        self.im2deep_config = Im2DeepConfiguration(self)
+        self.im2deep_config.grid(row=3, column=0, pady=(0, 20), sticky="nsew")
+
         self.ionmob_config = IonmobConfiguration(self)
-        self.ionmob_config.grid(row=3, column=0, pady=(0, 20), sticky="nsew")
+        self.ionmob_config.grid(row=4, column=0, pady=(0, 20), sticky="nsew")
 
     def get(self) -> Dict:
         """Return the configuration as a dictionary."""
         basic_enabled, basic_config = self.basic_config.get()
         ms2pip_enabled, ms2pip_config = self.ms2pip_config.get()
         deeplc_enabled, deeplc_config = self.deeplc_config.get()
+        im2deep_enabled, im2deep_config = self.im2deep_config.get()
         ionmob_enabled, ionmob_config = self.ionmob_config.get()
+
         config = {}
         if basic_enabled:
             config["basic"] = basic_config
@@ -522,6 +527,27 @@ def get(self) -> Dict:
         return enabled, config
 
 
+class Im2DeepConfiguration(ctk.CTkFrame):
+    def __init__(self, *args, **kwargs):
+        """IonMob configuration frame."""
+        super().__init__(*args, **kwargs)
+
+        self.configure(fg_color="transparent")
+        self.grid_columnconfigure(0, weight=1)
+
+        self.title = widgets.Heading(self, text="im2deep")
+        self.title.grid(row=0, column=0, columnspan=2, pady=(0, 5), sticky="ew")
+
+        self.enabled = widgets.LabeledSwitch(self, label="Enable im2deep", default=False)
+        self.enabled.grid(row=1, column=0, pady=(0, 10), sticky="nsew")
+
+    def get(self) -> Dict:
+        """Return the configuration as a dictionary."""
+        enabled = self.enabled.get()
+        config = {}
+        return enabled, config
+
+
 class RescoringEngineConfig(ctk.CTkFrame):
     def __init__(self, *args, **kwargs):
         """Rescoring engine configuration frame."""

ms2rescore/package_data/config_schema.json

@@ -29,6 +29,9 @@
                         },
                         "ionmob": {
                             "$ref": "#/definitions/ionmob"
+                        },
+                        "im2deep": {
+                            "$ref": "#/definitions/im2deep"
                         }
                     },
                     "default": {
@@ -236,6 +239,19 @@
                 }
             }
         },
+        "im2deep": {
+            "$ref": "#/definitions/feature_generator",
+            "description": "Ion mobility feature generator configuration using IM2Deep",
+            "type": "object",
+            "additionalProperties": true,
+            "properties": {
+                "reference_dataset": {
+                    "description": "Path to IM2Deep reference dataset file",
+                    "type": "string",
+                    "default": "Meier_unimod.parquet"
+                }
+            }
+        },
         "mokapot": {
             "$ref": "#/definitions/rescoring_engine",
             "description": "Mokapot rescoring engine configuration. Additional properties are passed to the Mokapot brew function.",