Merge pull request #109 from yueqixuan/dev

Fix the quantms LFQ/TMT modules

Author: ypriverol

quantmsio/commands/convert/quantms.py

@@ -88,7 +88,10 @@ def convert_quantms_feature_cmd(
         indexer = None
         if database_path and Path(database_path).exists():
             logger.info(f"Opening existing MzTabIndexer at {database_path}")
-            indexer = MzTabIndexer.open(str(database_path))
+            indexer = MzTabIndexer.open(
+                database_path=str(database_path),
+                sdrf_path=sdrf_file,
+            )
         elif database_path and mztab_path:
             logger.info(
                 f"Creating new MzTabIndexer at {database_path} from {mztab_path}"
@@ -177,7 +180,9 @@ def convert_quantms_psm_cmd(
         # Determine how to open or create the indexer
         if database_path and Path(database_path).exists():
             logger.info(f"Opening existing MzTabIndexer at {database_path}")
-            indexer = MzTabIndexer.open(str(database_path))
+            indexer = MzTabIndexer.open(
+                database_path=str(database_path),
+            )
         elif database_path and mztab_path:
             logger.info(
                 f"Creating new MzTabIndexer at {database_path} from {mztab_path}"
@@ -363,7 +368,10 @@ def convert_quantms_pg_cmd(
         indexer = None
         if database_path and Path(database_path).exists():
             logger.info(f"Opening existing MzTabIndexer at {database_path}")
-            indexer = MzTabIndexer.open(str(database_path))
+            indexer = MzTabIndexer.open(
+                database_path=str(database_path),
+                sdrf_path=sdrf_file,
+            )
         elif database_path and mztab_path:
             logger.info(
                 f"Creating new MzTabIndexer at {database_path} from {mztab_path}"
@@ -403,7 +411,7 @@ def convert_quantms_pg_cmd(
             # Convert to parquet and write
             table = mztab_pg._convert_to_parquet_format(result_df)
             pq.write_table(table, str(output_file))
-            logger.info("Successfully wrote protein groups to parquet file")
+            logger.info(f"[Writer] Successfully wrote protein groups to: {output_file}")
 
     except Exception as e:
         logger.exception(f"Error in mzTab protein group conversion: {str(e)}")

quantmsio/core/duckdb.py

@@ -47,7 +47,7 @@ def __init__(
             },
         )
         self.logger.info(
-            f"Time to initialize duckdb {time.time() - start_time} seconds"
+            f"Time to initialize duckdb {(time.time() - start_time):.2f} seconds"
         )
 
     def destroy_database(self):
@@ -56,6 +56,27 @@ def destroy_database(self):
             self._duckdb.close()
             self._duckdb = None
 
+    def cleanup_duckdb(self):
+        """Check if DuckDB connection is closed, then delete the database file."""
+        # Close connection if it is still open
+        if self._duckdb:
+            try:
+                self._duckdb.close()
+                self.logger.info("[Check] DuckDB connection closed.")
+            except Exception as e:
+                self.logger.info(f"Failed to close DuckDB connection: {e}")
+            finally:
+                self._duckdb = None
+
+        db_file = Path(self._database_path)
+        # Delete the database file using pathlib
+        if db_file.exists():
+            try:
+                db_file.unlink()
+                self.logger.info(f"[CleanUp] Database file deleted: {db_file}")
+            except Exception as e:
+                self.logger.info(f"Failed to delete database file: {e}")
+
     def query_to_df(self, query: str) -> pd.DataFrame:
         """Execute query and return result as DataFrame."""
         return self._duckdb.execute(query).df()

quantmsio/core/quantms/feature.py

@@ -1,11 +1,7 @@
 import logging
-import tempfile
 from pathlib import Path
-from typing import Union
-
 import pandas as pd
 import pyarrow as pa
-import pyarrow.parquet as pq
 
 from quantmsio.core.common import FEATURE_SCHEMA
 
@@ -202,7 +198,7 @@ def generate_modifications_details(self, peptidoform, modifications_dict):
                                 [
                                     {
                                         "score_name": "localization_probability",
-                                        "score_value": 1.0,
+                                        "score_value": None,
                                     }
                                 ]
                                 if mod_name in select_mods
@@ -350,10 +346,6 @@ def _create_file_metadata(self):
         }
 
     def transform_msstats_in(self, file_num=10, protein_str=None):
-        # Check if msstats data is already loaded in the indexer
-        # if not self._indexer._msstats_path:
-        #     # Add msstats data to the existing indexer
-        #     self._indexer.add_msstats_table(self._msstats_in)
 
         # Determine experiment type (LFQ vs TMT)
         experiment_type = self._indexer.get_msstats_experiment_type()
@@ -382,17 +374,15 @@ def _aggregate_msstats_to_features(self, msstats_batch, experiment_type):
         """
 
         # Group by feature identifier (peptidoform + charge + reference file + protein)
-        grouping_cols = ["PeptideSequence", "ProteinName", "reference_file_name"]
+        grouping_cols = ["peptidoform", "pg_accessions", "reference_file_name"]
 
         # Add charge column if available, otherwise use default
-        if "Charge" in msstats_batch.columns:
-            grouping_cols.append("Charge")
-        elif "PrecursorCharge" in msstats_batch.columns:
-            grouping_cols.append("PrecursorCharge")
+        if "charge" in msstats_batch.columns:
+            grouping_cols.append("charge")
         else:
             # Add a default charge if not available
-            msstats_batch["Charge"] = 3
-            grouping_cols.append("Charge")
+            msstats_batch["charge"] = 3
+            grouping_cols.append("charge")
 
         features_list = []
 
@@ -430,7 +420,7 @@ def _aggregate_msstats_to_features(self, msstats_batch, experiment_type):
                 "intensities": intensities,
                 "pg_accessions": [protein_name] if protein_name else [],
                 "anchor_protein": protein_name or "",
-                "rt": first_row.get("RetentionTime", None),
+                "rt": first_row.get("rt", None),
                 # Will add more fields in subsequent processing steps
             }
 
@@ -588,40 +578,42 @@ def write_feature_to_file(
             batch_writer.close()
 
             if Path(output_path).exists():
-                self.logger.info(f"Feature file written to {output_path}")
+                self.logger.info(
+                    f"[Writer] Successfully wrote Feature to: {output_path}"
+                )
 
             # Clean up the temporary MzTabIndexer
-            self._indexer.destroy_database()
-
-    def write_features_to_file(
-        self,
-        output_folder,
-        filename,
-        partitions,
-        file_num=10,
-        protein_file=None,
-        duckdb_max_memory="16GB",
-        duckdb_threads=4,
-    ):
-        logger = logging.getLogger("quantmsio.core.feature")
-
-        # Log input and output paths
-        logger.info(f"Input mzTab file: {self._indexer._mztab_path}")
-        logger.info(f"Output folder: {output_folder}")
-        logger.info(f"Base filename: {filename}")
-        if protein_file:
-            logger.info(f"Protein filter file: {protein_file}")
-
-        pqwriters = {}
-        protein_list = extract_protein_list(protein_file) if protein_file else None
-        protein_str = "|".join(protein_list) if protein_list else None
-        for key, feature in self.generate_slice_feature(
-            partitions, file_num, protein_str, duckdb_max_memory, duckdb_threads
-        ):
-            pqwriters = save_slice_file(
-                feature, pqwriters, output_folder, key, filename
-            )
-        close_file(pqwriters)
+            self._indexer.cleanup_duckdb()
+
+    # def write_features_to_file(
+    #     self,
+    #     output_folder,
+    #     filename,
+    #     partitions,
+    #     file_num=10,
+    #     protein_file=None,
+    #     duckdb_max_memory="16GB",
+    #     duckdb_threads=4,
+    # ):
+    #     logger = logging.getLogger("quantmsio.core.feature")
+
+    #     # Log input and output paths
+    #     logger.info(f"Input mzTab file: {self._indexer._mztab_path}")
+    #     logger.info(f"Output folder: {output_folder}")
+    #     logger.info(f"Base filename: {filename}")
+    #     if protein_file:
+    #         logger.info(f"Protein filter file: {protein_file}")
+
+    #     pqwriters = {}
+    #     protein_list = extract_protein_list(protein_file) if protein_file else None
+    #     protein_str = "|".join(protein_list) if protein_list else None
+    #     for key, feature in self.generate_slice_feature(
+    #         partitions, file_num, protein_str, duckdb_max_memory, duckdb_threads
+    #     ):
+    #         pqwriters = save_slice_file(
+    #             feature, pqwriters, output_folder, key, filename
+    #         )
+    #     close_file(pqwriters)
 
     @staticmethod
     def generate_best_scan(rows, pep_dict):