more efficient metadata fetching and removed unused dynamic size fields from vector store

Author: MichelDierckx

TreeMS2/similarity_matrix/filters/precursor_mz_filter.py

@@ -25,10 +25,12 @@ def __init__(self, groups: Groups, vector_store: VectorStore,
     def construct_mask(self, similarity_matrix: SimilarityMatrix) -> SpectraMatrix:
         rows, cols = similarity_matrix.matrix.nonzero()
 
-        precursor_mz_rows = self.vector_store.get_data(rows=rows, columns=["precursor_mz"])["precursor_mz"].to_numpy(
-            dtype=np.float32)
-        precursor_mz_cols = self.vector_store.get_data(rows=cols, columns=["precursor_mz"])["precursor_mz"].to_numpy(
-            dtype=np.float32)
+        # precursor mz values for all entries
+        precursor_mz = self.vector_store.get_col("precursor_mz").to_numpy(dtype=np.float32).ravel()
+
+        # Vectorized lookup using rows and cols
+        precursor_mz_rows = precursor_mz[rows]
+        precursor_mz_cols = precursor_mz[cols]
 
         mask_data = np.abs(precursor_mz_rows - precursor_mz_cols) > self.precursor_mz_window
         # Filter the rows and columns based on the mask

TreeMS2/similarity_sets.py

@@ -25,8 +25,10 @@ def update_similarity_sets(self, similarity_matrix: SimilarityMatrix):
         rows, cols = similarity_matrix.matrix.nonzero()
         total_spectra = self.groups.total_spectra
 
-        row_ids = self.vector_store.get_data(rows, ["global_id"])["global_id"].to_numpy(dtype=np.int32)
-        col_ids = self.vector_store.get_data(cols, ["global_id"])["global_id"].to_numpy(dtype=np.int32)
+        global_ids = self.vector_store.get_col("global_id").to_numpy(dtype=np.uint32).ravel()
+        row_ids = global_ids[rows]
+        col_ids = global_ids[cols]
+
         m = csr_matrix((similarity_matrix.matrix.data, (row_ids, col_ids)), shape=(total_spectra, total_spectra),
                        dtype=np.bool_)
 

TreeMS2/spectrum/group_spectrum.py

@@ -34,11 +34,11 @@ def to_dict(self) -> Dict:
             "spectrum_id": self._spectrum_id,
             "file_id": self._file_id,
             "group_id": self._group_id,
-            "identifier": self.spectrum.identifier,
+            #self "identifier": self.spectrum.identifier,
             "precursor_mz": self.spectrum.precursor_mz,
             "precursor_charge": self.spectrum.precursor_charge,
-            "mz": self.spectrum.mz,
-            "intensity": self.spectrum.intensity,
+            # "mz": self.spectrum.mz,
+            # "intensity": self.spectrum.intensity,
             "retention_time": self.spectrum.retention_time,
             "vector": self.vector,
         }

TreeMS2/vector_store/vector_store.py

@@ -37,13 +37,13 @@ def __init__(self, name: str, directory: str, vector_dim: int):
             pa.field("spectrum_id", pa.uint16()),
             pa.field("file_id", pa.uint16()),
             pa.field("group_id", pa.uint16()),
-            pa.field("identifier", pa.string()),
+            # pa.field("identifier", pa.string()),
             pa.field("precursor_mz", pa.float32()),
             pa.field("precursor_charge", pa.int8()),
-            pa.field("mz", pa.list_(pa.float32())),
-            pa.field("intensity", pa.list_(pa.float32())),
+            # pa.field("mz", pa.list_(pa.float32())),
+            # pa.field("intensity", pa.list_(pa.float32())),
             pa.field("retention_time", pa.float32()),
-            pa.field("vector", pa.list_(pa.float32())),
+            pa.field("vector", pa.list_(pa.float32(), vector_dim))            ,
         ])
 
     def _get_dataset(self) -> Optional[LanceDataset]:
@@ -115,6 +115,12 @@ def get_data(self, rows: List[int], columns: List[str]) -> pd.DataFrame:
             return pd.DataFrame(columns=columns)
         return ds.take(indices=rows, columns=columns).to_pandas()
 
+    def get_col(self, column) -> pd.DataFrame:
+        ds = self._get_dataset()
+        if ds is None:
+            return pd.DataFrame(columns=column)
+        return ds.to_table(columns=[column]).to_pandas()
+
     def add_global_ids(self, groups: Groups) -> None:
         def compute_global_id(row: pd.Series) -> int:
             offset = groups.get_group(row['group_id']).get_peak_file(row['file_id']).begin