Change SpectrumSetBase to SpectrumSet

Author: niekdejonge

ms2query/benchmarking/EvaluateMethods.py

@@ -3,7 +3,7 @@
 import numpy as np
 from matchms.similarity.vector_similarity_functions import jaccard_similarity_matrix
 from tqdm import tqdm
-from ms2query.benchmarking.SpectrumDataSet import SpectraWithFingerprints, SpectrumSetBase
+from ms2query.benchmarking.SpectrumDataSet import SpectraWithFingerprints, SpectrumSet
 
 
 class EvaluateMethods:
@@ -107,7 +107,7 @@ def get_accuracy_recall_curve(self):
 
 
 def predict_between_two_sets(
-    library: SpectrumSetBase, query_set_1: SpectrumSetBase, query_set_2: SpectrumSetBase, prediction_function
+    library: SpectrumSet, query_set_1: SpectrumSet, query_set_2: SpectrumSet, prediction_function
 ):
     """Makes predictions between query sets and the library, with the other query set added.
 
@@ -123,7 +123,7 @@ def predict_between_two_sets(
     return predicted_inchikeys_1 + predicted_inchikeys_2
 
 
-def calculate_average_exact_match_accuracy(spectrum_set: SpectrumSetBase, predicted_inchikeys: List[str]):
+def calculate_average_exact_match_accuracy(spectrum_set: SpectrumSet, predicted_inchikeys: List[str]):
     if len(spectrum_set.spectra) != len(predicted_inchikeys):
         raise ValueError("The number of spectra should be equal to the number of predicted inchikeys ")
     exact_match_accuracy_per_inchikey = []
@@ -139,7 +139,7 @@ def calculate_average_exact_match_accuracy(spectrum_set: SpectrumSetBase, predic
     return sum(exact_match_accuracy_per_inchikey) / len(exact_match_accuracy_per_inchikey)
 
 
-def split_spectrum_set_per_inchikeys(spectrum_set: SpectrumSetBase) -> Tuple[SpectrumSetBase, SpectrumSetBase]:
+def split_spectrum_set_per_inchikeys(spectrum_set: SpectrumSet) -> Tuple[SpectrumSet, SpectrumSet]:
     """Splits a spectrum set into two.
     For each inchikey with more than one spectrum the spectra are divided over the two sets"""
     indexes_set_1 = []
@@ -157,8 +157,8 @@ def split_spectrum_set_per_inchikeys(spectrum_set: SpectrumSetBase) -> Tuple[Spe
 
 
 def split_spectrum_set_per_inchikey_across_ionmodes(
-    spectrum_set: SpectrumSetBase,
-) -> Tuple[SpectrumSetBase, SpectrumSetBase]:
+    spectrum_set: SpectrumSet,
+) -> Tuple[SpectrumSet, SpectrumSet]:
     """Splits a spectrum set in two sets on ionmode. Only uses spectra for inchikeys with at least 1 pos and 1 neg"""
     all_pos_indexes = []
     all_neg_indexes = []
@@ -190,7 +190,7 @@ def split_spectrum_set_per_inchikey_across_ionmodes(
     return pos_val_spectra, neg_val_spectra
 
 
-def subset_spectra_on_ionmode(spectrum_set: SpectrumSetBase, ionmode) -> SpectrumSetBase:
+def subset_spectra_on_ionmode(spectrum_set: SpectrumSet, ionmode) -> SpectrumSet:
     spectrum_indexes_to_keep = []
     for i, spectrum in enumerate(spectrum_set.spectra):
         if spectrum.get("ionmode") == ionmode:

ms2query/benchmarking/SpectrumDataSet.py

@@ -8,7 +8,8 @@
 from tqdm import tqdm
 
 
-class SpectrumSetBase:
+
+class SpectrumSet:
     """Stores a spectrum dataset making it easy and fast to split on molecules"""
 
     def __init__(self, spectra: List[Spectrum], progress_bars=False):
@@ -36,10 +37,10 @@ def _add_spectra_and_group_per_inchikey(self, spectra: List[Spectrum]):
                 ]
         return updated_inchikeys
 
-    def add_spectra(self, new_spectra: "SpectrumSetBase"):
+    def add_spectra(self, new_spectra: "SpectrumSet"):
         return self._add_spectra_and_group_per_inchikey(new_spectra.spectra)
 
-    def subset_spectra(self, spectrum_indexes) -> "SpectrumSetBase":
+    def subset_spectra(self, spectrum_indexes) -> "SpectrumSet":
         """Returns a new instance of a subset of the spectra"""
         new_instance = copy.copy(self)
         new_instance._spectra = []
@@ -65,7 +66,7 @@ def copy(self):
         return new_instance
 
 
-class SpectraWithFingerprints(SpectrumSetBase):
+class SpectraWithFingerprints(SpectrumSet):
     """Stores a spectrum dataset making it easy and fast to split on molecules"""
 
     def __init__(self, spectra: List[Spectrum], fingerprint_type="daylight", nbits=4096):

tests/test_SpectrumDataSet.py

@@ -3,21 +3,21 @@
 from ms2query.benchmarking.SpectrumDataSet import (
     SpectraWithFingerprints,
     SpectraWithMS2DeepScoreEmbeddings,
-    SpectrumSetBase,
+    SpectrumSet,
 )
 from tests.conftest import create_test_spectra, get_inchikey_inchi_pairs, ms2deepscore_model
 
 
 @pytest.mark.parametrize(
     "library",
     [
-        SpectrumSetBase(create_test_spectra()),
+        SpectrumSet(create_test_spectra()),
         SpectraWithFingerprints(create_test_spectra()),
         SpectraWithMS2DeepScoreEmbeddings(create_test_spectra(), ms2deepscore_model()),
     ],
 )
 def test_spectrum_set_base(library):
-    """Test all base functionality of SpectrumSetBase is implemented correctly
+    """Test all base functionality of SpectrumSet is implemented correctly
     also for all classes inheriting from it"""
     # test correct init
     assert len(library.spectra) == 9
@@ -74,7 +74,7 @@ def test_spectra_with_fingerprints(library):
         (get_inchikey_inchi_pairs(3), 3),  # Fully overlapping
         (get_inchikey_inchi_pairs(1), 3),  # Fully overlapping (but not all)
     ):
-        spectra_to_add = SpectrumSetBase(create_test_spectra(2, inchikey_inchi_pairs=inchikey_inchi_pairs))
+        spectra_to_add = SpectrumSet(create_test_spectra(2, inchikey_inchi_pairs=inchikey_inchi_pairs))
         new_copy = library.copy()
         new_copy.add_spectra(spectra_to_add)
         assert len(new_copy.inchikey_fingerprint_pairs) == expected_nr_of_inchikeys
@@ -125,6 +125,6 @@ def test_spectra_with_embeddings():
     for i, index in enumerate(subset_indexes):
         assert np.all(library.embeddings[index] == subset.embeddings[i])
 
-    # Check that subsetting on subset works. To make sure that a subset does not become of type SpectrumSetBase
+    # Check that subsetting on subset works. To make sure that a subset does not become of type SpectrumSet
     subsetted_subset = subset.subset_spectra([0, 1])
     assert subsetted_subset.embeddings.shape == (2, 100)