Make predict_using_closest_tanimoto use TopKTanimotoScores

Author: niekdejonge

ms2query/benchmarking/reference_methods/predict_using_closest_tanimoto.py

@@ -4,6 +4,8 @@
 from tqdm import tqdm
 from ms2query.benchmarking.AnnotatedSpectrumSet import AnnotatedSpectrumSet
 from ms2query.benchmarking.Fingerprints import Fingerprints
+from ms2query.benchmarking.reference_methods.predict_top_ms2deepscores import select_inchikeys_with_highest_ms2deepscore
+from ms2query.benchmarking.TopKTanimotoScores import TopKTanimotoScores
 from ms2query.metrics import generalized_tanimoto_similarity_matrix
 
 
@@ -17,66 +19,34 @@ def predict_using_closest_tanimoto(
     """Predict best inchikey, by taking the average score over all spectra for the 10 closest related library inchikeys.
     (simplified version of old MS2Query)
     """
+    top_k_tanimoto_scores = TopKTanimotoScores.calculate_from_fingerprints(
+        library_fingerprints,
+        library_fingerprints,
+        k=nr_of_closest_inchikeys_to_select,
+    )
+    ms2deepscores = cosine_similarity_matrix(query_spectra.embeddings.embeddings, library_spectra.embeddings.embeddings)
+    inchikeys_with_highest_ms2deepscores = select_inchikeys_with_highest_ms2deepscore(
+        query_spectra, library_spectra, nr_of_inchikeys_with_highest_ms2deepscore_to_select, ms2deepscores=ms2deepscores
+    )
+
     inchikeys_of_best_match = []
     highest_scores = []
     for spectrum_idx in tqdm(range(len(query_spectra.spectra)), "Predicting using closest tanimoto"):
-        inchikey_of_best_match, score = predict_using_closest_tanimoto_single_spectrum(
-            library_spectra,
-            query_spectra.subset_spectra([spectrum_idx]),
-            nr_of_closest_inchikeys_to_select,
-            nr_of_inchikeys_with_highest_ms2deepscore_to_select,
-            library_fingerprints,
-        )
-        inchikeys_of_best_match.append(inchikey_of_best_match)
-        highest_scores.append(score)
-    return inchikeys_of_best_match, highest_scores
+        average_predicted_scores = {}
+        for inchikey in inchikeys_with_highest_ms2deepscores[spectrum_idx]:
+            top_k_inchikeys = top_k_tanimoto_scores.select_top_k_inchikeys(inchikey)
 
+            average_predicted_score = get_average_predictions_for_closely_related_metabolites(
+                library_spectra, top_k_inchikeys, ms2deepscores[spectrum_idx]
+            )
+            average_predicted_scores[inchikey] = average_predicted_score
 
-def predict_using_closest_tanimoto_single_spectrum(
-    spectra_with_embeddings: AnnotatedSpectrumSet,
-    single_spectrum_with_embeddings: AnnotatedSpectrumSet,
-    nr_of_closest_inchikeys_to_select,
-    nr_of_inchikeys_with_highest_ms2deepscore_to_select,
-    fingerprints,
-) -> Tuple[str, float]:
-    if len(single_spectrum_with_embeddings.spectra) != 1:
-        raise ValueError("expected a single spectrum")
-    ms2deepscores = cosine_similarity_matrix(
-        single_spectrum_with_embeddings.embeddings.embeddings, spectra_with_embeddings.embeddings.embeddings
-    )[0]
-    top_inchikeys = select_inchikeys_with_highest_ms2deepscore(
-        spectra_with_embeddings, ms2deepscores, nr_of_inchikeys_with_highest_ms2deepscore_to_select
-    )
-    average_predicted_scores = {}
-    for inchikey in top_inchikeys:
-        top_k_inchikeys, _ = get_inchikey_and_tanimoto_scores_for_top_k(
-            fingerprints, inchikey, nr_of_closest_inchikeys_to_select
+        inchikey_with_highest_average_prediction, score = max(
+            average_predicted_scores.items(), key=lambda item: item[1]
         )
-        average_predicted_score = get_average_predictions_for_closely_related_metabolites(
-            spectra_with_embeddings, top_k_inchikeys, ms2deepscores
-        )
-        average_predicted_scores[inchikey] = average_predicted_score
-
-    inchikey_with_highest_average_prediction, score = max(average_predicted_scores.items(), key=lambda item: item[1])
-    return inchikey_with_highest_average_prediction, score
-
-
-def select_inchikeys_with_highest_ms2deepscore(
-    spectra_with_embeddings: AnnotatedSpectrumSet, ms2deepscores, nr_of_inchikeys_to_select=10
-):
-    highest_score_for_inchikey = []
-    for inchikey, spectrum_indexes in spectra_with_embeddings.spectrum_indices_per_inchikey.items():
-        all_ms2deepscores_for_inchikey = ms2deepscores[spectrum_indexes,]
-        highest_score_for_inchikey.append(max(all_ms2deepscores_for_inchikey))
-    inchikey_indexes_with_highest_ms2deepscore = np.argpartition(
-        np.array(highest_score_for_inchikey), -nr_of_inchikeys_to_select
-    )[-nr_of_inchikeys_to_select:]
-
-    top_inchikeys = [
-        spectra_with_embeddings.inchikeys[inchikey_index]
-        for inchikey_index in inchikey_indexes_with_highest_ms2deepscore
-    ]
-    return top_inchikeys
+        inchikeys_of_best_match.append(inchikey_with_highest_average_prediction)
+        highest_scores.append(score)
+    return inchikeys_of_best_match, highest_scores
 
 
 def get_average_predictions_for_closely_related_metabolites(

tests/test_benchmarking/test_predict_using_closest_tanimoto.py

@@ -6,8 +6,6 @@
     get_average_predictions_for_closely_related_metabolites,
     get_inchikey_and_tanimoto_scores_for_top_k,
     predict_using_closest_tanimoto,
-    predict_using_closest_tanimoto_single_spectrum,
-    select_inchikeys_with_highest_ms2deepscore,
 )
 from tests.helper_functions import create_test_spectra, ms2deepscore_model
 
@@ -28,38 +26,6 @@ def test_predict_using_closest_tanimoto():
     assert len(scores) == 3
 
 
-def test_predict_using_closest_tanimoto_single_spectrum():
-    """Only very basic test that the function runs and that the output is the right format"""
-    model = ms2deepscore_model()
-    library_spectra = AnnotatedSpectrumSet.create_spectrum_set(create_test_spectra(nr_of_inchikeys=7))
-    test_spectra = AnnotatedSpectrumSet.create_spectrum_set(create_test_spectra(1, nr_of_inchikeys=1))
-    library_spectra.add_embeddings(model)
-    test_spectra.add_embeddings(model)
-    fingerprints = Fingerprints.from_spectrum_set(library_spectra, "daylight", 2048)
-
-    predicted_inchikey, score = predict_using_closest_tanimoto_single_spectrum(
-        library_spectra, test_spectra, 3, 3, fingerprints
-    )
-
-    assert isinstance(predicted_inchikey, str)
-    assert len(predicted_inchikey) == 14
-    assert isinstance(score, float)
-
-
-def test_select_inchikeys_with_highest_ms2deepscore():
-    test_spectra = create_test_spectra(nr_of_inchikeys=7)
-    spectra = AnnotatedSpectrumSet.create_spectrum_set(test_spectra)
-
-    ms2deepscores = np.zeros(len(test_spectra))
-    ms2deepscores[2] = 0.4
-    ms2deepscores[5] = 0.9
-    ms2deepscores[7] = 0.6
-    inchikeys_with_highest_ms2deepscore = select_inchikeys_with_highest_ms2deepscore(spectra, ms2deepscores, 3)
-    expected_inchikeys = list(spectra.spectrum_indices_per_inchikey.keys())[:3]
-    assert set(expected_inchikeys) == set(inchikeys_with_highest_ms2deepscore)
-    print(inchikeys_with_highest_ms2deepscore)
-
-
 def test_get_average_predictions_for_closely_related_metabolites():
     test_spectra = create_test_spectra(nr_of_inchikeys=7)
     # Select different number per inchikey (only one for the first) to check that it is correctly weighted.