Do correct subsetting of inchikey sets

Author: niekdejonge

ms2query/benchmarking/SpectrumDataSet.py

@@ -108,6 +108,15 @@ def copy(self):
         new_instance.inchikey_fingerprint_pairs = copy.copy(self.inchikey_fingerprint_pairs)
         return new_instance
 
+    def subset_spectra(self, spectrum_indexes) -> "SpectraWithFingerprints":
+        """Returns a new instance of a subset of the spectra"""
+        new_instance = super().subset_spectra(spectrum_indexes)
+        # Only keep the fingerprints for which we have inchikeys.
+        # Important note: This is not a deep copy!
+        # And the fingerprint is not reset (so it is not always actually matching the most common inchi)
+        new_instance.inchikey_fingerprint_pairs = {inchikey: self.inchikey_fingerprint_pairs[inchikey] for inchikey in new_instance.spectrum_indexes_per_inchikey.keys()}
+        return new_instance
+
 
 class SpectraWithMS2DeepScoreEmbeddings(SpectraWithFingerprints):
     def __init__(self, spectra: List[Spectrum], ms2deepscore_model: SiameseSpectralModel, **kwargs):

tests/test_SpectrumDataSet.py

@@ -91,9 +91,9 @@ def test_spectra_with_fingerprints(library):
     )
 
     # test correct subsetting
-    subset_indexes = [1, 4, 6, 7]
+    subset_indexes = [1, 6, 7]
     subset = library.subset_spectra(subset_indexes)
-    assert len(subset.inchikey_fingerprint_pairs) == 3
+    assert len(subset.inchikey_fingerprint_pairs) == 2
     assert all(
         np.array_equal(library.inchikey_fingerprint_pairs[key], value)
         for key, value in subset.inchikey_fingerprint_pairs.items()