From 687514d13372268cf77f0d4dd730e8451424c6b0 Mon Sep 17 00:00:00 2001
From: rodvrees <robbe.devreese@hotmail.com>
Date: Tue, 19 Nov 2024 16:09:58 +0100
Subject: [PATCH] Move IM conversion outside output writing

---
 im2deep/im2deep.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/im2deep/im2deep.py b/im2deep/im2deep.py
index a4ac496..66f8f16 100644
--- a/im2deep/im2deep.py
+++ b/im2deep/im2deep.py
@@ -57,6 +57,14 @@ def predict_ccs(
             use_charge_state=use_charge_state,
         )
 
+    if ion_mobility:
+        LOGGER.info("Converting CCS to IM values...")
+        psm_list_pred_df["predicted_im"] = ccs2im(
+            psm_list_pred_df["predicted_ccs"],
+            psm_list_pred_df["peptidoform"].apply(lambda x: x.theoretical_mz),
+            psm_list_pred_df["charge"],
+        )
+
     if write_output:
         if not ion_mobility:
             LOGGER.info("Writing output file...")
@@ -70,12 +78,6 @@ def predict_ccs(
                 output_file.write(f"{peptidoform},{charge},{CCS}\n")
             output_file.close()
         else:
-            LOGGER.info("Converting CCS to IM values...")
-            psm_list_pred_df["predicted_im"] = ccs2im(
-                psm_list_pred_df["predicted_ccs"],
-                psm_list_pred_df["peptidoform"].apply(lambda x: x.theoretical_mz),
-                psm_list_pred_df["charge"],
-            )
             LOGGER.info("Writing output file...")
             output_file = open(output_file, "w")
             output_file.write("modified_seq,charge,predicted IM\n")
@@ -89,4 +91,7 @@ def predict_ccs(
 
     LOGGER.info("IM2Deep finished!")
 
-    return psm_list_pred_df["predicted_ccs"]
+    if not ion_mobility:
+        return psm_list_pred_df["predicted_ccs"]
+    else:
+        return psm_list_pred_df["predicted_im"]