Added save and load models for embeddings

Author: niekdejonge

ms2query/benchmarking/Embeddings.py

@@ -1,3 +1,5 @@
+import json
+from pathlib import Path
 from typing import Sequence
 import numpy as np
 import pandas as pd
@@ -17,7 +19,7 @@ def __init__(self, embeddings: np.ndarray, spectrum_hashes: tuple, model_setting
         self._spectrum_hash_to_index = {
             spectrum_hash: index for index, spectrum_hash in enumerate(self.index_to_spectrum_hash)
         }
-        self._model_settings = model_settings
+        self.model_settings = model_settings
         self._embeddings = embeddings
 
     @classmethod
@@ -39,7 +41,7 @@ def __add__(self, other: "Embeddings") -> "Embeddings":
             raise ValueError("There are repeated spectra in the embeddings that are added together")
         combined_embeddings = np.vstack([self._embeddings, other._embeddings])
         index_to_spectrum_hash = self.index_to_spectrum_hash + other.index_to_spectrum_hash
-        return Embeddings(combined_embeddings, index_to_spectrum_hash, self._model_settings)
+        return Embeddings(combined_embeddings, index_to_spectrum_hash, self.model_settings)
 
     def subset_embeddings(self, spectra):
         spectrum_hashes = tuple(spectrum.__hash__() for spectrum in spectra)
@@ -58,11 +60,15 @@ def embeddings(self):
     def model_settings(self):
         return self._model_settings.copy()
 
+    @model_settings.setter
+    def model_settings(self, model_settings):
+        self._model_settings: dict = _to_json_serializable(model_settings)
+
     def copy(self) -> "Embeddings":
         return Embeddings(
             embeddings=self._embeddings.copy(),
             spectrum_hashes=tuple(self.index_to_spectrum_hash),
-            model_settings=dict(self._model_settings),
+            model_settings=dict(self.model_settings),
         )
 
     def __eq__(self, other) -> bool:
@@ -76,10 +82,59 @@ def __eq__(self, other) -> bool:
             return False
         return np.array_equal(self.embeddings, other.embeddings)
 
+    def save(self, path: str | Path) -> None:
+        """Save embeddings to a .npz file with metadata stored alongside.
+
+        Args:
+            path: File path. A '.npz' extension will be added if not present.
+        """
+        path = Path(path).with_suffix(".npz")
+        metadata = {
+            "model_settings": self.model_settings,
+            "index_to_spectrum_hash": list(self.index_to_spectrum_hash),
+        }
+        np.savez_compressed(
+            path,
+            embeddings=self._embeddings,
+            metadata=np.array(json.dumps(metadata)),
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> "Embeddings":
+        """Load embeddings from a saved .npz file.
+
+        Args:
+            path: Path to the saved .npz file.
+        """
+        path = Path(path).with_suffix(".npz")
+        with np.load(path, allow_pickle=False) as data:
+            embeddings = data["embeddings"]
+            metadata = json.loads(data["metadata"].item())
+        return cls(
+            embeddings=embeddings,
+            spectrum_hashes=tuple(metadata["index_to_spectrum_hash"]),
+            model_settings=metadata["model_settings"],
+        )
+
 
 def calculate_ms2deepscore_df(query_embeddings: Embeddings, library_embeddings: Embeddings):
     """Returns a DF, where the indexes and column labels are the spectrum hashes"""
     ms2deepscores = cosine_similarity_matrix(query_embeddings.embeddings, library_embeddings.embeddings)
     return pd.DataFrame(
         ms2deepscores, index=query_embeddings.index_to_spectrum_hash, columns=library_embeddings.index_to_spectrum_hash
     )
+
+
+def _to_json_serializable(obj):
+    """Changes a dict to be json sericalizable, so it is the same when loaded"""
+    if isinstance(obj, dict):
+        return {key: _to_json_serializable(value) for key, value in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_to_json_serializable(item) for item in obj]
+    if isinstance(obj, np.integer):
+        return int(obj)
+    if isinstance(obj, np.floating):
+        return float(obj)
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    return obj

tests/test_benchmarking/test_embeddings.py

@@ -1,3 +1,4 @@
+import os
 import pytest
 from ms2query.benchmarking.Embeddings import Embeddings, calculate_ms2deepscore_df
 from tests.helper_functions import create_test_spectra, get_library_and_test_spectra_not_identical, ms2deepscore_model
@@ -34,3 +35,13 @@ def test_add_embeddings():
 def test_calculate_ms2deepscore_df():
     library_spectra, query_spectra = get_library_and_test_spectra_not_identical()
     calculate_ms2deepscore_df(library_spectra.embeddings, query_spectra.embeddings)
+
+
+def test_save_and_load(tmp_path):
+    test_spectra = create_test_spectra()
+    model = ms2deepscore_model()
+    embeddings = Embeddings.create_from_spectra(test_spectra[:4], model)
+    file_name = os.path.join(tmp_path, "embeddings.npz")
+    embeddings.save(file_name)
+    loaded_embeddings = Embeddings.load(file_name)
+    assert embeddings == loaded_embeddings