Refactoring for MS2Rescore implementation

im2deep/__init__.py

@@ -1,3 +1,3 @@
 """IM2Deep: Deep learning framework for peptide collisional cross section prediction."""
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"

im2deep/__main__.py

@@ -8,17 +8,19 @@
 from typing import Optional
 
 import click
-# import pandas as pd
+import pandas as pd
+
 # from deeplc import DeepLC
-# from psm_utils.io import read_file
-# from psm_utils.io.exceptions import PSMUtilsIOException
-# from psm_utils.io.peptide_record import peprec_to_proforma
-# from psm_utils.psm import PSM
-# from psm_utils.psm_list import PSMList
+from psm_utils.io import read_file
+from psm_utils.io.exceptions import PSMUtilsIOException
+from psm_utils.io.peptide_record import peprec_to_proforma
+from psm_utils.psm import PSM
+from psm_utils.psm_list import PSMList
 from rich.logging import RichHandler
 
 from im2deep._exceptions import IM2DeepError
 from im2deep.im2deep import predict_ccs
+
 # from im2deep.calibrate import linear_calibration
 
 REFERENCE_DATASET_PATH = Path(__file__).parent / "reference_data" / "reference_ccs.zip"
@@ -48,7 +50,8 @@ def setup_logging(passed_level):
         handlers=[RichHandler()],
     )
 
-#Command line arguments TODO: Make config_parser script
+
+# Command line arguments TODO: Make config_parser script
 @click.command()
 @click.argument("psm_file", type=click.Path(exists=True, dir_okay=False))
 @click.option(
@@ -98,7 +101,6 @@ def setup_logging(passed_level):
     default=2,
     help="Charge state to use for calibration. Only used if calibrate_per_charge is set to False.",
 )
-
 def main(
     psm_file: str,
     calibration_file: Optional[str] = None,
@@ -111,11 +113,83 @@ def main(
 ):
     """Command line interface to IM2Deep."""
     setup_logging(log_level)
+
+    with open(psm_file) as f:
+        first_line_pred = f.readline().strip()
+    if calibration_file:
+        with open(calibration_file) as fc:
+            first_line_cal = fc.readline().strip()
+
+    if "modifications" in first_line_pred.split(",") and "seq" in first_line_pred.split(","):
+        # Read input file
+        df_pred = pd.read_csv(psm_file)
+        df_pred.fillna("", inplace=True)
+
+        list_of_psms = []
+        for seq, mod, charge, ident in zip(
+            df_pred["seq"], df_pred["modifications"], df_pred["charge"], df_pred.index
+        ):
+            list_of_psms.append(
+                PSM(peptidoform=peprec_to_proforma(seq, mod, charge), spectrum_id=ident)
+            )
+        psm_list_pred = PSMList(psm_list=list_of_psms)
+
+    else:
+        # psm_list_pred = read_file(file_pred)
+        try:
+            psm_list_pred = read_file(psm_file)
+        except PSMUtilsIOException:
+            LOGGER.error("Invalid input file. Please check the format of the input file.")
+            sys.exit(1)
+
+    psm_list_cal = []
+    if (
+        calibration_file
+        and "modifications" in first_line_cal.split(",")
+        and "seq" in first_line_cal.split(",")
+    ):
+        df_cal = pd.read_csv(calibration_file)
+        df_cal.fillna("", inplace=True)
+        del calibration_file
+
+        list_of_cal_psms = []
+        for seq, mod, charge, ident, CCS in zip(
+            df_cal["seq"],
+            df_cal["modifications"],
+            df_cal["charge"],
+            df_cal.index,
+            df_cal["CCS"],
+        ):
+            list_of_cal_psms.append(
+                PSM(peptidoform=peprec_to_proforma(seq, mod, charge), spectrum_id=ident)
+            )
+        psm_list_cal = PSMList(psm_list=list_of_cal_psms)
+        psm_list_cal_df = psm_list_cal.to_dataframe()
+        psm_list_cal_df["ccs_observed"] = df_cal["CCS"]
+        del df_cal
+
+    else:
+        LOGGER.error(
+            "Invalid calibration file. Please check the format of the calibration file."
+        )
+        sys.exit(1)
+
+    if not output_file:
+        output_file = Path(psm_file).parent / (Path(psm_file).stem + "_IM2Deep-predictions.csv")
     try:
-        predict_ccs(psm_file, calibration_file, REFERENCE_DATASET_PATH, output_file, model_name, calibrate_per_charge, use_charge_state, n_jobs)
+        predict_ccs(
+            psm_list_pred,
+            psm_list_cal_df,
+            output_file=output_file,
+            model_name=model_name,
+            calibrate_per_charge=calibrate_per_charge,
+            use_charge_state=use_charge_state,
+            n_jobs=n_jobs,
+        )
     except IM2DeepError as e:
         LOGGER.error(e)
         sys.exit(1)
 
+
 if __name__ == "__main__":
     main()

im2deep/calibrate.py

@@ -78,9 +78,10 @@ def get_ccs_shift(
         """Calculating CCS shift based on {} overlapping peptide-charge pairs
         between PSMs and reference dataset""".format(both.shape[0])
     )
+    LOGGER.debug(both.columns)
     # How much CCS in calibration data is larger than reference CCS, so predictions
     # need to be increased by this amount
-    return 0 if both.shape[0] == 0 else np.mean(both["observed_ccs"] - both["CCS"])
+    return 0 if both.shape[0] == 0 else np.mean(both["ccs_observed"] - both["CCS"])
 
 
 def get_ccs_shift_per_charge(cal_df: pd.DataFrame, reference_dataset: pd.DataFrame) -> ndarray:
@@ -123,7 +124,7 @@ def get_ccs_shift_per_charge(cal_df: pd.DataFrame, reference_dataset: pd.DataFra
         how="inner",
         suffixes=("_ref", "_data"),
     )
-    return both.groupby("charge").apply(lambda x: np.mean(x["observed_ccs"] - x["CCS"])).to_dict()
+    return both.groupby("charge").apply(lambda x: np.mean(x["ccs_observed"] - x["CCS"])).to_dict()
 
 
 def calculate_ccs_shift(
@@ -184,7 +185,7 @@ def linear_calibration(
         )
     else:
         shift_factor = calculate_ccs_shift(
-            preds_df, reference_dataset, per_charge=False, use_charge_state=use_charge_state
+            calibration_dataset, reference_dataset, per_charge=False, use_charge_state=use_charge_state
         )
         preds_df["predicted_ccs_calibrated"] = preds_df.apply(
             lambda x: x["predicted_ccs"] + shift_factor, axis=1

im2deep/im2deep.py

@@ -1,89 +1,32 @@
 import logging
-import sys
 from pathlib import Path
 
 import pandas as pd
 from deeplc import DeepLC
-from psm_utils.io import read_file
-from psm_utils.io.exceptions import PSMUtilsIOException
-from psm_utils.io.peptide_record import peprec_to_proforma
-from psm_utils.psm import PSM
 from psm_utils.psm_list import PSMList
-from rich.logging import RichHandler
 
 from im2deep.calibrate import linear_calibration
 
 LOGGER = logging.getLogger(__name__)
 REFERENCE_DATASET_PATH = Path(__file__).parent / "reference_data" / "reference_ccs.zip"
 
+
+# TODO: get file reading out of the function
 def predict_ccs(
-    file_pred,
-    file_cal=None,
+    psm_list_pred: PSMList,
+    psm_list_cal_df=None,
     file_reference=REFERENCE_DATASET_PATH,
-    file_pred_out=None,
+    output_file=None,
     model_name="tims",
     calibrate_per_charge=True,
     use_charge_state=2,
     n_jobs=None,
+    write_output=True,
 ):
     """Run IM2Deep."""
     LOGGER.info("IM2Deep started.")
     reference_dataset = pd.read_csv(file_reference)
 
-    with open(file_pred) as f:
-        first_line_pred = f.readline().strip()
-    if file_cal:
-        with open(file_cal) as fc:
-            first_line_cal = fc.readline().strip()
-
-    if "modifications" in first_line_pred.split(",") and "seq" in first_line_pred.split(","):
-        # Read input file
-        df_pred = pd.read_csv(file_pred)
-        df_pred.fillna("", inplace=True)
-
-        list_of_psms = []
-        for seq, mod, charge, ident in zip(
-            df_pred["seq"], df_pred["modifications"], df_pred["charge"], df_pred.index
-        ):
-            list_of_psms.append(
-                PSM(peptidoform=peprec_to_proforma(seq, mod, charge), spectrum_id=ident)
-            )
-        psm_list_pred = PSMList(psm_list=list_of_psms)
-
-    else:
-        # psm_list_pred = read_file(file_pred)
-        try:
-            psm_list_pred = read_file(file_pred)
-        except PSMUtilsIOException:
-            LOGGER.error("Invalid input file. Please check the format of the input file.")
-            sys.exit(1)
-
-    psm_list_cal = []
-    if (
-        file_cal
-        and "modifications" in first_line_cal.split(",")
-        and "seq" in first_line_cal.split(",")
-    ):
-        df_cal = pd.read_csv(file_cal)
-        df_cal.fillna("", inplace=True)
-        del file_cal
-
-        list_of_cal_psms = []
-        for seq, mod, charge, ident, CCS in zip(
-            df_cal["seq"], df_cal["modifications"], df_cal["charge"], df_cal.index, df_cal["CCS"]
-        ):
-            list_of_cal_psms.append(
-                PSM(peptidoform=peprec_to_proforma(seq, mod, charge), spectrum_id=ident)
-            )
-        psm_list_cal = PSMList(psm_list=list_of_cal_psms)
-        psm_list_cal_df = psm_list_cal.to_dataframe()
-        psm_list_cal_df["observed_ccs"] = df_cal["CCS"]
-        del df_cal
-
-    else:
-        LOGGER.error("Invalid calibration file. Please check the format of the calibration file.")
-        sys.exit(1)
-
     if model_name == "tims":
         path_model = Path(__file__).parent / "models" / "TIMS"
 
@@ -103,34 +46,17 @@ def predict_ccs(
         per_charge=calibrate_per_charge,
         use_charge_state=use_charge_state,
     )
-
-    LOGGER.info("Writing output file...")
-    if file_pred_out:
-        file_pred_out = open(file_pred_out, "w")
-        file_pred_out.write("seq,modifications,charge,predicted CCS\n")
-        for seq, mod, charge, ident, CCS in zip(
-            df_pred["seq"],
-            df_pred["modifications"],
-            df_pred["charge"],
-            df_pred.index,
-            calibrated_psm_list_pred_df["predicted_ccs_calibrated"],
-        ):
-            file_pred_out.write(f"{seq},{mod},{charge},{CCS}\n")
-        file_pred_out.close()
-    else:
-        #Get path of psm file
-        output_file = Path(file_pred).parent / (Path(file_pred).stem + "_IM2Deep-predictions.csv")
-        LOGGER.info("Writing output file to %s", output_file)
+    if write_output:
+        LOGGER.info("Writing output file...")
         output_file = open(output_file, "w")
         output_file.write("seq,modifications,charge,predicted CCS\n")
-        for seq, mod, charge, ident, CCS in zip(
-            df_pred["seq"],
-            df_pred["modifications"],
-            df_pred["charge"],
-            df_pred.index,
+        for peptidoform, charge, CCS in zip(
+            calibrated_psm_list_pred_df["peptidoform"],
+            calibrated_psm_list_pred_df["charge"],
             calibrated_psm_list_pred_df["predicted_ccs_calibrated"],
         ):
-            output_file.write(f"{seq},{mod},{charge},{CCS}\n")
+            output_file.write(f"{peptidoform},{charge},{CCS}\n")
         output_file.close()
 
     LOGGER.info("IM2Deep finished!")
+    return calibrated_psm_list_pred_df["predicted_ccs_calibrated"]