clean up

wilhelm-lab · Apr 12, 2024 · f659213 · f659213
1 parent eda2dba
commit f659213
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/spectrum_io/search_result/xisearch.py b/spectrum_io/search_result/xisearch.py
@@ -17,6 +17,7 @@
 
 class Xisearch(SearchResults):
     """Handle search results from xisearch."""
+
     def read_result(self, tmt_labeled: str = "") -> pd.DataFrame:
         """
         Function to read a csv of CSMs and perform some basic formatting.
@@ -127,16 +128,24 @@ def update_columns_for_prosit(df: pd.DataFrame) -> pd.DataFrame:
         df["CROSSLINKER_POSITION_A"] = df["CROSSLINKER_POSITION_A"].astype("int")
         df["CROSSLINKER_POSITION_B"] = df["CROSSLINKER_POSITION_B"].astype("int")
 
-        df[["MODIFIED_SEQUENCE_A", "MODIFIED_SEQUENCE_B"]] = df.apply(
+        df["MODIFIED_SEQUENCE_A"] = df.apply(
             lambda row: xisearch_to_internal(
                 row["CROSSLINKER_TYPE"],
                 row["SEQUENCE_A"],
-                row["SEQUENCE_B"],
                 row["Modifications_A"],
-                row["Modifications_B"],
                 row["CROSSLINKER_POSITION_A"],
-                row["CROSSLINKER_POSITION_B"],
                 row["ModificationPositions1"],
+            ),
+            axis=1,
+            result_type="expand",
+        )
+
+        df["MODIFIED_SEQUENCE_B"] = df.apply(
+            lambda row: xisearch_to_internal(
+                row["CROSSLINKER_TYPE"],
+                row["SEQUENCE_B"],
+                row["Modifications_B"],
+                row["CROSSLINKER_POSITION_B"],
                 row["ModificationPositions2"],
             ),
             axis=1,

diff --git a/tests/unit_tests/test_xisearch.py b/tests/unit_tests/test_xisearch.py
@@ -14,9 +14,7 @@ def test_read_xisearch(self):
         """Test function for reading Xisearch results and transforming to Prosit format."""
         expected_xisearch_internal_path = Path(__file__).parent / "data" / "xisearch_output_internal.tsv"
 
-        internal_search_results_df = Xisearch(
-            path=Path(__file__).parent / "data" / "xisearch_output.tsv"
-        ).read_result()
+        internal_search_results_df = Xisearch(path=Path(__file__).parent / "data" / "xisearch_output.tsv").read_result()
         internal_search_results_df.reset_index(drop=True, inplace=True)
         expected_df = pd.read_csv(expected_xisearch_internal_path)
         expected_df["Modifications_A"] = expected_df["Modifications_A"].fillna("nan")