Add spectra to reference lib

README.md

@@ -5,8 +5,6 @@
 A first basic implementation is out now, more to follow soon...
 The new MS2Query appraoch has a higher accuracy and has a much simpler and faster underlying algorithm. We will hopefully soon share a first preprint as well, showing all the benchmarking. 
 
-The current runably version still requires to create the library files, which takes some time for the first run. 
-
 Soon this will be much easier and faster. We will add downloadable precomputed files, make MS2Query pip installable, add a database and allow faster MS2DeepScore searching. 
 
 The tutorial for the current prototype can be found in notebooks/tutorial. 

ms2query/ms2query_development/Embeddings.py

@@ -52,6 +52,11 @@ def subset_embeddings(self, spectra):
         embeddings = self._embeddings[embedding_indexes].copy()
         return Embeddings(embeddings, spectrum_hashes, self.model_settings)
 
+    def subset_embeddings_from_index(self, indexes):
+        spectrum_hashes = [self.index_to_spectrum_hash[index] for index in indexes]
+        embeddings = self._embeddings[indexes].copy()
+        return Embeddings(embeddings, tuple(spectrum_hashes), self.model_settings)
+
     @property
     def embeddings(self):
         return self._embeddings.view()

ms2query/ms2query_development/Fingerprints.py

@@ -34,6 +34,16 @@ def compute_fingerprints_from_inchi(cls, most_common_inchi_per_inchikey: dict[st
             fingerprints[inchikey_index, :] = fingerprint
         return cls(fingerprints, index_to_inchikey, fingerprint_type)
 
+    @classmethod
+    def from_dataframe(cls, dataframe: pd.DataFrame, fingerprint_type, nbits):
+        """From a dataframe with columns inchikey and inchi the Fingerprints are computed"""
+        most_common_inchi_per_inchikey = (
+            dataframe.groupby(dataframe["inchikey"].str[:14])["inchi"]
+            .agg(lambda x: x.value_counts().idxmax())
+            .to_dict()
+        )
+        return cls.compute_fingerprints_from_inchi(most_common_inchi_per_inchikey, fingerprint_type, nbits)
+
     @classmethod
     def from_spectrum_set(cls, spectrum_set: AnnotatedSpectrumSet, fingerprint_type, nbits):
         most_common_inchi_per_inchikey = {}

ms2query/ms2query_development/TopKTanimotoScores.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 import numpy as np
 import pandas as pd
+from tqdm import tqdm
 from ms2query.metrics import generalized_tanimoto_similarity_matrix
 from ms2query.ms2query_development.Fingerprints import Fingerprints
 
@@ -37,21 +38,49 @@ def _create_multi_index(
         return df
 
     @classmethod
-    def calculate_from_fingerprints(cls, query_fingerprints: Fingerprints, target_fingerprints: Fingerprints, k):
+    def calculate_from_fingerprints(
+        cls,
+        query_fingerprints: Fingerprints,
+        target_fingerprints: Fingerprints,
+        k: int,
+        batch_size: int = 1000,
+    ):
         """
-        Gets the top k highest inchikeys and scores for each inchikey in query_fingerprints from target_fingerprints
+        Gets the top k highest inchikeys and scores for each inchikey in query_fingerprints
+        from target_fingerprints.
+
+        Runs in batches over the query fingerprints so the full (n_queries x n_targets)
+        similarity matrix is never fully materialized in memory - only one batch's slice
+        is, which is then reduced down to just the top-k before moving to the next batch.
         """
-        if target_fingerprints.fingerprints.shape[0] < k:
+        n_targets = target_fingerprints.fingerprints.shape[0]
+        n_queries = query_fingerprints.fingerprints.shape[0]
+
+        if n_targets < k:
             raise ValueError("K cannot be larger than the number of fingerprints")
-        similarity_scores = generalized_tanimoto_similarity_matrix(
-            query_fingerprints.fingerprints, target_fingerprints.fingerprints
-        )
-        inchikey_indexes_of_top_k = np.argpartition(similarity_scores, -k, axis=1)[:, -k:]
-        top_k_inchikeys = np.array(target_fingerprints.inchikeys)[inchikey_indexes_of_top_k]
-        tanimoto_scores_for_top_k = similarity_scores[
-            np.arange(similarity_scores.shape[0])[:, None], inchikey_indexes_of_top_k
-        ]
-        return cls(tanimoto_scores_for_top_k, top_k_inchikeys, np.array(query_fingerprints.inchikeys))
+
+        target_inchikeys = np.array(target_fingerprints.inchikeys)
+
+        # Preallocate final outputs instead of concatenating per batch
+        top_k_scores = np.empty((n_queries, k), dtype=np.float32)
+        top_k_inchikeys = np.empty((n_queries, k), dtype=target_inchikeys.dtype)
+
+        n_batches = int(np.ceil(n_queries / batch_size))
+        for batch_idx in tqdm(range(n_batches), desc="Calculating top-k Tanimoto scores"):
+            start = batch_idx * batch_size
+            end = min(start + batch_size, n_queries)
+
+            query_batch = query_fingerprints.fingerprints[start:end]
+
+            # shape: (batch_size, n_targets) -- small enough to hold in memory
+            similarity_scores = generalized_tanimoto_similarity_matrix(query_batch, target_fingerprints.fingerprints)
+
+            top_k_idx = np.argpartition(similarity_scores, -k, axis=1)[:, -k:]
+
+            top_k_scores[start:end] = similarity_scores[np.arange(similarity_scores.shape[0])[:, None], top_k_idx]
+            top_k_inchikeys[start:end] = target_inchikeys[top_k_idx]
+
+        return cls(top_k_scores, top_k_inchikeys, np.array(query_fingerprints.inchikeys))
 
     def select_top_k_inchikeys_and_scores(self, inchikey) -> dict[str, float]:
         """Returns a dictionary with inchikeys and scores for the given inchikey"""
@@ -103,3 +132,10 @@ def load(cls, path: str | Path) -> "TopKTanimotoScores":
         instance.k = len(df.columns.get_level_values("result_rank").unique())
         instance.top_k_inchikeys_and_scores = df
         return instance
+
+    def __eq__(self, other):
+        if not self.k == other.k:
+            return False
+        if not self.top_k_inchikeys_and_scores.equals(other.top_k_inchikeys_and_scores):
+            return False
+        return True

ms2query/notebooks/tutorial.ipynb

@@ -385,21 +385,87 @@
     "reference_library.run_ms2query(test_spectra)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f5840421",
+   "metadata": {},
+   "source": [
+    "# Add spectra to an existing library\n",
+    "If you want to add some in house spectra to the existing library the code below does this without recreating all the files for the full library. \n",
+    "The last step, computing the top-8 closest tanimoto has to be fully recomputed, which on a laptop will take a few hours. However, if you have a GPU available it is much faster."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4aaf7fc0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ms2query.ms2query_development.ReferenceLibrary import ReferenceLibrary\n",
+    "from matchms.importing.load_spectra import load_spectra\n",
+    "from tqdm import tqdm\n",
+    "reference_library = ReferenceLibrary.load_from_files(os.path.join(folder_to_store_zenodo_files, \"ms2deepscore_model.pt\"),\n",
+    "                                                    os.path.join(folder_to_store_zenodo_files, \"embeddings.npz\"),\n",
+    "                                                    os.path.join(folder_to_store_zenodo_files, \"top_k_tanimoto_scores.parquet\"),\n",
+    "                                                    os.path.join(folder_to_store_zenodo_files, \"library_metadata.parquet\"))\n",
+    "\n",
+    "# Clean the spectra first, they will need annotations, smiles, inchi, inchikey. You can use matchms for this. \n",
+    "test_spectra = list(tqdm(load_spectra(\"../../tests/test_data/10_spectra.mgf\")))\n",
+    "\n",
+    "reference_library.add_spectra(test_spectra)\n",
+    "reference_library.save(store_file_directory=\"./model_with_extra_spectra\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "419cca3b",
+   "metadata": {},
+   "source": [
+    "# Run semi targeted search\n",
+    "If you already have a list of known compounds for which you want to find any analogues in your sample, you can use the semitargeted run method. \n",
+    "This takes in the refernce library, but only checks the molecules in the library that you specify. It will still use the rest of the library to compute the MS2Query reliability score. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "49cb74c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reference_library.run_semi_targeted_ms2query(test_spectra, inchikeys_to_check={\"UTXPDZFPPUZUBK\", \"OZRMEKAUZBKTTC\"})"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "cc3f2477",
    "metadata": {},
    "source": [
     "# Create a new MS2Query reference library\n",
-    "Here we show how to create the reference library that can be downloaded from Zenodo. But if you want to make your own of course replace with your own files (no need to download the library)"
+    "Here we show how to create the reference library that can be downloaded from Zenodo. But if you want to make your own of course replace with your own files (no need to download the library)\n",
+    "\n",
+    " For small libraries it is recommended to just add them to the existing libraries instead of creating a small (e.g. <10000 compounds) reference library, since we did not test the uncertainty estimation of MS2Query for small libraries. If you really only care about this small reference library, you can also consider just running MS2DeepScore, since the risk on false positives is much smaller for a small reference library. "
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "2501dc1e-9ab0-408d-b869-8cd9874a46d4",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'download_file' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mdownload_file\u001b[49m(\u001b[33m\"\u001b[39m\u001b[33mhttps://zenodo.org/records/16882111/files/merged_and_cleaned_libraries_1.mgf?download=1\u001b[39m\u001b[33m\"\u001b[39m, \n\u001b[32m      2\u001b[39m               os.path.join(folder_to_store_zenodo_files, \u001b[33m\"\u001b[39m\u001b[33mmerged_and_cleaned_libraries_1.mgf\u001b[39m\u001b[33m\"\u001b[39m))\n\u001b[32m      3\u001b[39m download_file(\u001b[33m\"\u001b[39m\u001b[33mhttps://zenodo.org/records/17826815/files/ms2deepscore_model.pt?download=1\u001b[39m\u001b[33m\"\u001b[39m, \n\u001b[32m      4\u001b[39m               os.path.join(folder_to_store_zenodo_files, \u001b[33m\"\u001b[39m\u001b[33mms2deepscore_model.pt\u001b[39m\u001b[33m\"\u001b[39m))\n",
+      "\u001b[31mNameError\u001b[39m: name 'download_file' is not defined"
+     ]
+    }
+   ],
    "source": [
     "download_file(\"https://zenodo.org/records/16882111/files/merged_and_cleaned_libraries_1.mgf?download=1\", \n",
     "              os.path.join(folder_to_store_zenodo_files, \"merged_and_cleaned_libraries_1.mgf\"))\n",
@@ -446,7 +512,8 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "reference_library = ReferenceLibrary.create_from_spectra(library_spectra, os.path.join(folder_to_store_zenodo_files, \"ms2deepscore_model.pt\"), store_file_directory = Path(os.path.join(folder_to_store_zenodo_files, \"full_library\")))"
+    "reference_library = ReferenceLibrary.create_from_spectra(library_spectra, os.path.join(folder_to_store_zenodo_files, \"ms2deepscore_model.pt\"))\n",
+    "reference_library.save(Path(os.path.join(folder_to_store_zenodo_files, \"full_library\")))"
    ]
   },
   {
@@ -1087,7 +1154,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "ms2query_2",
    "language": "python",
    "name": "python3"
   },

tests/test_ms2query_development/test_ReferenceLibrary.py

@@ -38,7 +38,8 @@ def test_create_library(tmp_path):
     lib_spectra = create_test_spectra(nr_of_inchikeys=10, number_of_spectra_per_inchikey=3)
     # save_as_mgf(lib_spectra, os.path.join(tmp_path, "library_spectra.mgf"))
     ms2deepscore_model_file = os.path.join(TEST_RESOURCES_PATH, "ms2deepscore_testmodel_v1.pt")
-    ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file, tmp_path)
+    library = ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file)
+    library.save(tmp_path)
     assert (tmp_path / ReferenceLibrary.embedding_file_name).exists()
     assert (tmp_path / ReferenceLibrary.top_k_tanimoto_scores_file_name).exists()
     assert (tmp_path / ReferenceLibrary.reference_metadata_file_name).exists()
@@ -47,7 +48,8 @@ def test_create_library(tmp_path):
 def test_create_and_use_library(tmp_path):
     lib_spectra = create_test_spectra(nr_of_inchikeys=10, number_of_spectra_per_inchikey=3)
     ms2deepscore_model_file = os.path.join(TEST_RESOURCES_PATH, "ms2deepscore_testmodel_v1.pt")
-    ms2query_library = ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file, tmp_path)
+    ms2query_library = ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file)
+    ms2query_library.save(tmp_path)
     test_spectra = create_test_spectra(1, nr_of_inchikeys=3)
     results = ms2query_library.run_ms2query(test_spectra)
 
@@ -60,3 +62,38 @@ def test_create_and_use_library(tmp_path):
 
     results_2 = ms2query_library_2.run_ms2query(test_spectra)
     pd.testing.assert_frame_equal(results, results_2)
+
+
+def test_add_spectra():
+    lib_spectra = create_test_spectra(nr_of_inchikeys=10, number_of_spectra_per_inchikey=3)
+    first_spectra = lib_spectra[:25]
+    later_spectra = lib_spectra[25:]
+    ms2deepscore_model_file = os.path.join(TEST_RESOURCES_PATH, "ms2deepscore_testmodel_v1.pt")
+    ms2query_library = ReferenceLibrary.create_from_spectra(first_spectra, ms2deepscore_model_file)
+    ms2query_library.add_spectra(later_spectra)
+
+    ms2query_library_2 = ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file)
+
+    assert ms2query_library.reference_embeddings == ms2query_library_2.reference_embeddings
+    assert ms2query_library.top_k_tanimoto_scores == ms2query_library_2.top_k_tanimoto_scores
+    pd.testing.assert_frame_equal(ms2query_library.reference_metadata, ms2query_library_2.reference_metadata)
+    test_spectra = create_test_spectra(1, nr_of_inchikeys=3)
+
+    results = ms2query_library.run_ms2query(test_spectra)
+    results_2 = ms2query_library_2.run_ms2query(test_spectra)
+
+    pd.testing.assert_frame_equal(results, results_2)
+
+
+def test_run_semi_targeted_search():
+    lib_spectra = create_test_spectra(nr_of_inchikeys=10, number_of_spectra_per_inchikey=3)
+    ms2deepscore_model_file = os.path.join(TEST_RESOURCES_PATH, "ms2deepscore_testmodel_v1.pt")
+    library = ReferenceLibrary.create_from_spectra(lib_spectra, ms2deepscore_model_file)
+    test_spectra = create_test_spectra(1, nr_of_inchikeys=3)
+    inchikeys = {spectrum.get("inchikey")[:14] for spectrum in lib_spectra}
+    results = library.run_semi_targeted_ms2query(test_spectra, inchikeys)
+
+    results_2 = library.run_ms2query(test_spectra)
+    print(results)
+    print(results_2)
+    pd.testing.assert_frame_equal(results, results_2)