FEAT add missing columns to output

Author: GeorgWa

alphadia/plexscoring.py

@@ -738,6 +738,9 @@ def process(
             psm_proto_df.fragment_mass_error[
                 self.output_idx, : len(mass_error)
             ] = mass_error
+            psm_proto_df.fragment_position[
+                self.output_idx, : len(fragments.position)
+            ] = fragments.position
             psm_proto_df.fragment_number[
                 self.output_idx, : len(fragments.number)
             ] = fragments.number
@@ -1270,6 +1273,7 @@ class OuptutPsmDF:
     fragment_mass_error: nb.float32[:, ::1]
     fragment_correlation: nb.float32[:, ::1]
 
+    fragment_position: nb.uint8[:, ::1]
     fragment_number: nb.uint8[:, ::1]
     fragment_type: nb.uint8[:, ::1]
     fragment_charge: nb.uint8[:, ::1]
@@ -1296,6 +1300,7 @@ def __init__(self, n_psm, top_k_fragments):
         self.fragment_mass_error = np.zeros((n_psm, top_k_fragments), dtype=np.float32)
         self.fragment_correlation = np.zeros((n_psm, top_k_fragments), dtype=np.float32)
 
+        self.fragment_position = np.zeros((n_psm, top_k_fragments), dtype=np.uint8)
         self.fragment_number = np.zeros((n_psm, top_k_fragments), dtype=np.uint8)
         self.fragment_type = np.zeros((n_psm, top_k_fragments), dtype=np.uint8)
         self.fragment_charge = np.zeros((n_psm, top_k_fragments), dtype=np.uint8)
@@ -1313,6 +1318,8 @@ def to_fragment_df(self):
             self.fragment_intensity.flatten()[mask],
             self.fragment_mass_error.flatten()[mask],
             self.fragment_correlation.flatten()[mask],
+
+            self.fragment_position.flatten()[mask],
             self.fragment_number.flatten()[mask],
             self.fragment_type.flatten()[mask],
             self.fragment_charge.flatten()[mask],
@@ -1727,13 +1734,20 @@ def collect_candidates(
         ] + utils.get_isotope_column_names(self.precursors_flat_df.columns)
 
         precursor_df_columns += (
-            [self.rt_column] if self.rt_column not in precursor_df_columns else []
+            [self.rt_column] 
+            if self.rt_column not in precursor_df_columns 
+            else []
         )
         precursor_df_columns += (
             [self.mobility_column]
             if self.mobility_column not in precursor_df_columns
             else []
         )
+        precursor_df_columns += (
+            [self.precursor_mz_column]
+            if self.precursor_mz_column not in precursor_df_columns
+            else []
+        )
 
         df = utils.merge_missing_columns(
             df,
@@ -1789,6 +1803,7 @@ def collect_fragments(
             "intensity",
             "mass_error",
             "correlation",
+            "position",
             "number",
             "type",
             "charge",