Added an embedding setter with checks to AnnotatedSpectrumSet

Author: niekdejonge

ms2query/benchmarking/AnnotatedSpectrumSet.py

@@ -19,7 +19,7 @@ def __init__(
         self.spectrum_indices_per_inchikey: dict[str, tuple[int, ...]] = {
             key: tuple(values) for key, values in spectrum_indices_per_inchikey.items()
         }
-        self._embeddings = embeddings
+        self.embeddings = embeddings
 
     @classmethod
     def create_spectrum_set(cls, spectra: Sequence[Spectrum]) -> "AnnotatedSpectrumSet":
@@ -99,6 +99,17 @@ def embeddings(self) -> Embeddings:
             raise ValueError("First run the 'add_embeddings' method")
         return self._embeddings
 
+    @embeddings.setter
+    def embeddings(self, embeddings: Optional[Embeddings]):
+        if embeddings is None:
+            self._embeddings = embeddings
+            return
+        if not embeddings.index_to_spectrum_hash == tuple(spectrum.__hash__() for spectrum in self.spectra):
+            raise ValueError(
+                "The embeddings spectrum hashes don't match the spectrum hashes, make sure you use matching embeddings"
+            )
+        self._embeddings = embeddings
+
     @property
     def inchikeys(self):
         return tuple(self.spectrum_indices_per_inchikey.keys())

tests/test_benchmarking/test_SpectrumDataSet.py

@@ -1,3 +1,4 @@
+import pytest
 from ms2query.benchmarking.AnnotatedSpectrumSet import (
     AnnotatedSpectrumSet,
 )
@@ -60,3 +61,18 @@ def test_subset_on_metadata():
     spectrum_set.add_embeddings(model)
 
     assert correct_subsetted_set == spectrum_set.subset_spectra_on_metadata("ionmode", set(["positive"]))
+
+
+def test_add_embeddings():
+    test_spectra = create_test_spectra(nr_of_inchikeys=3, number_of_spectra_per_inchikey=3)
+
+    spectrum_set = AnnotatedSpectrumSet.create_spectrum_set(test_spectra)
+
+    subset_of_spectra = AnnotatedSpectrumSet.create_spectrum_set(test_spectra[:5])
+    spectrum_set = AnnotatedSpectrumSet.create_spectrum_set(test_spectra)
+    # with added embededings
+    model = ms2deepscore_model()
+    subset_of_spectra.add_embeddings(model)
+    with pytest.raises(ValueError):
+        # The spectra don't match so it should raise a valueerror
+        spectrum_set.embeddings = subset_of_spectra.embeddings