Split select_inchikeys_with_highest_ms2deepscores to make more modular and to select top k highest ms2deepscores

Author: niekdejonge

ms2query/benchmarking/reference_methods/predict_using_closest_tanimoto.py

@@ -14,34 +14,47 @@ def predict_using_closest_tanimoto(
     (simplified version of old MS2Query)
     """
     inchikeys_of_best_match = []
-    single_highest_score = []
+    highest_scores = []
     for spectrum_idx in range(len(query_spectra.spectra)):
         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)
         inchikeys_of_best_match.append(inchikey_of_best_match)
-        single_highest_score.append(score)
-    return inchikeys_of_best_match, single_highest_score
+        highest_scores.append(score)
+    return inchikeys_of_best_match, highest_scores
 
 
 def predict_using_closest_tanimoto_single_spectrum(spectra_with_embeddings, single_spectrum_with_embeddings,
-                                                   nr_of_closest_inchikeys_to_select) -> Tuple[str, float]:
+                                                   nr_of_closest_inchikeys_to_select,
+                                                   nr_of_inchikeys_with_highest_ms2deepscore_to_select) -> 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,
                                              spectra_with_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, spectrum_indexes in spectra_with_embeddings.spectrum_indexes_per_inchikey.items():
-        all_ms2deepscores_for_inchikey = ms2deepscores[spectrum_indexes]
-        if max(all_ms2deepscores_for_inchikey) > 0.7:
-            top_k_inchikeys, _ = get_inchikey_and_tanimoto_scores_for_top_k(
-                spectra_with_embeddings, inchikey, nr_of_closest_inchikeys_to_select)
-            average_predicted_score = get_average_predictions_for_closely_related_metabolites(
-                spectra_with_embeddings, top_k_inchikeys, ms2deepscores)
-            average_predicted_scores[inchikey] = average_predicted_score
+    for inchikey in top_inchikeys:
+        top_k_inchikeys, _ = get_inchikey_and_tanimoto_scores_for_top_k(
+            spectra_with_embeddings, inchikey, nr_of_closest_inchikeys_to_select)
+        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, ms2deepscores, nr_of_inchikeys_to_select=10):
+    highest_score_for_inchikey = []
+    for inchikey, spectrum_indexes in spectra_with_embeddings.spectrum_indexes_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:]
+
+    all_inchikeys = list(spectra_with_embeddings.inchikey_fingerprint_pairs.keys())
+    top_inchikeys = [all_inchikeys[inchikey_index] for inchikey_index in inchikey_indexes_with_highest_ms2deepscore]
+    return top_inchikeys
+
 def get_average_predictions_for_closely_related_metabolites(spectra_with_embeddings, top_k_inchikeys,
                                                             all_ms2deepscores):
     """Calculates the average ms2deepscore predictions for top k closest inchikeys"""