From 85b869e4d19c93fbdc21d55474736519fb4049e7 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 13:29:58 +0200
Subject: [PATCH 01/20] Add fingerprints from dataframe method

---
 ms2query/ms2query_development/Fingerprints.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/ms2query/ms2query_development/Fingerprints.py b/ms2query/ms2query_development/Fingerprints.py
index 4e20cce..d5770a9 100644
--- a/ms2query/ms2query_development/Fingerprints.py
+++ b/ms2query/ms2query_development/Fingerprints.py
@@ -34,6 +34,14 @@ 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("inchikey")["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 = {}

From 6e6dcd9fe2d1f53fab503322c94145a937e16454 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 13:31:16 +0200
Subject: [PATCH 02/20] Add a save option to ReferenceLibrary

---
 .../ms2query_development/ReferenceLibrary.py  | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index d9fa9b4..dd93391 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -131,6 +131,33 @@ def create_from_spectra(
             library_spectrum_set.embeddings.save(store_file_directory / cls.embedding_file_name)
         return cls(ms2deepscore_model, library_spectrum_set.embeddings, top_k_tanimoto_scores, reference_metadata)
 
+    def save(self, store_file_directory: str | Path):
+        store_file_directory = Path(store_file_directory)
+        if not store_file_directory.is_file():
+            raise ValueError(
+                f"The store_file_directory is a file, please use a directory instead, given dir: {store_file_directory}"
+            )
+        store_file_directory.mkdir(parents=True, exist_ok=True)
+
+        def file_does_not_exist_yet(file_name: Path):
+            if file_name.exists:
+                print(f"The file: {file_name} already exists, not saved.")
+                return False
+            return True
+
+        # Save files after checking it does not exist yet
+        if file_does_not_exist_yet(store_file_directory / self.reference_metadata_file_name):
+            self.reference_metadata.to_parquet(store_file_directory / self.reference_metadata_file_name)
+
+        if file_does_not_exist_yet(store_file_directory / self.top_k_tanimoto_scores_file_name):
+            self.top_k_tanimoto_scores.save(store_file_directory / self.top_k_tanimoto_scores_file_name)
+
+        if file_does_not_exist_yet(store_file_directory / self.embedding_file_name):
+            self.reference_embeddings.save(store_file_directory / self.embedding_file_name)
+
+        if file_does_not_exist_yet(store_file_directory / self.ms2deepscore_model_file_name):
+            self.ms2deepscore_model.save(store_file_directory / self.ms2deepscore_model_file_name)
+
     def run_ms2query(
         self,
         query_spectra: Sequence[Spectrum],

From 959ccb62ac38f4a59943f8f30ab081ce7e1c4c97 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 13:32:42 +0200
Subject: [PATCH 03/20] Remove saving from create_from_spectra

---
 .../ms2query_development/ReferenceLibrary.py  | 24 +------------------
 ms2query/notebooks/tutorial.ipynb             |  3 ++-
 .../test_ReferenceLibrary.py                  |  6 +++--
 3 files changed, 7 insertions(+), 26 deletions(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index dd93391..5afab61 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -93,29 +93,11 @@ def create_from_spectra(
         cls,
         library_spectra: Sequence[Spectrum],
         ms2deepscore_model_file_name: str,
-        store_file_directory=None,
-        store_files=True,
     ) -> "ReferenceLibrary":
         """Creates all the files needed for MS2Query and stores them"""
-        if store_file_directory is None:
-            store_file_directory = Path(ms2deepscore_model_file_name).parent
-        else:
-            store_file_directory = Path(store_file_directory)
-        if store_files:
-            # Check the files don't exist yet
-            for file in (
-                store_file_directory / cls.embedding_file_name,
-                store_file_directory / cls.top_k_tanimoto_scores_file_name,
-                store_file_directory / cls.reference_metadata_file_name,
-            ):
-                if file.exists():
-                    raise FileExistsError(f"There is already a file stored with the name {file}")
-
         # library_spectra = list(tqdm(load_spectra(library_spectra_file), "Loading library spectra"))
         library_spectrum_set = AnnotatedSpectrumSet.create_spectrum_set(library_spectra)
         ms2deepscore_model = load_model(ms2deepscore_model_file_name)
-        library_spectrum_set.add_embeddings(ms2deepscore_model)
-
         fingerprints = Fingerprints.from_spectrum_set(library_spectrum_set, cls.fingerprint_type, cls.fingerprint_nbits)
         top_k_tanimoto_scores = TopKTanimotoScores.calculate_from_fingerprints(
             fingerprints, fingerprints, cls.top_k_inchikeys
@@ -124,11 +106,7 @@ def create_from_spectra(
             library_spectrum_set,
             cls.metadata_to_store,
         )
-
-        if store_files:
-            reference_metadata.to_parquet(store_file_directory / cls.reference_metadata_file_name)
-            top_k_tanimoto_scores.save(store_file_directory / cls.top_k_tanimoto_scores_file_name)
-            library_spectrum_set.embeddings.save(store_file_directory / cls.embedding_file_name)
+        library_spectrum_set.add_embeddings(ms2deepscore_model)
         return cls(ms2deepscore_model, library_spectrum_set.embeddings, top_k_tanimoto_scores, reference_metadata)
 
     def save(self, store_file_directory: str | Path):
diff --git a/ms2query/notebooks/tutorial.ipynb b/ms2query/notebooks/tutorial.ipynb
index 7afc892..12af41f 100644
--- a/ms2query/notebooks/tutorial.ipynb
+++ b/ms2query/notebooks/tutorial.ipynb
@@ -446,7 +446,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\")))"
    ]
   },
   {
diff --git a/tests/test_ms2query_development/test_ReferenceLibrary.py b/tests/test_ms2query_development/test_ReferenceLibrary.py
index 77b5bf9..3768029 100644
--- a/tests/test_ms2query_development/test_ReferenceLibrary.py
+++ b/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)
 

From 9ac5676ab6c39a678bc2f5c1ed0fcc13429816a2 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 13:35:09 +0200
Subject: [PATCH 04/20] Add add_spectra functionality

---
 .../ms2query_development/ReferenceLibrary.py  | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 5afab61..870a8e2 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -26,6 +26,7 @@ class ReferenceLibrary:
         "collision_energy",
         "compound_name",
         "smiles",
+        "inchi",
         "inchikey",
     ]
     fingerprint_type = "daylight"
@@ -109,6 +110,39 @@ def create_from_spectra(
         library_spectrum_set.add_embeddings(ms2deepscore_model)
         return cls(ms2deepscore_model, library_spectrum_set.embeddings, top_k_tanimoto_scores, reference_metadata)
 
+    def add_spectra(self, new_library_spectra):
+        """Add spectra to the already existing library (the ms2deepscore model won't be retrained)"""
+        # Check that no duplicates are added
+        hashes = [spectrum.__hash__() for spectrum in tqdm(new_library_spectra, desc="Hashing spectra")]
+        if len(hashes) != len(set(hashes)):
+            raise ValueError("There are duplicated spectra, please make sure there are no duplicates")
+        # Only add spectra not already in the library
+        existing_hashes = set(self.reference_metadata["spectrum_hashes"])
+        new_spectra = [spectrum for spectrum, h in zip(new_library_spectra, hashes) if h not in existing_hashes]
+        if len(new_spectra) != len(new_library_spectra):
+            print(f"{len(new_library_spectra) - len(new_spectra)} were not added, since already in the library")
+
+        new_spectrum_set = AnnotatedSpectrumSet.create_spectrum_set(new_spectra)
+
+        # Add spectrum metadata
+        reference_metadata = extract_metadata_from_library(
+            new_spectrum_set,
+            self.metadata_to_store,
+        )
+        self.reference_metadata = pd.concat([self.reference_metadata, reference_metadata], ignore_index=True)
+
+        # Add embeddings
+        new_spectrum_set.add_embeddings(self.ms2deepscore_model)
+        self.reference_embeddings = self.reference_embeddings + new_spectrum_set.embeddings
+
+        # Recompute top k tanimoto scores
+        fingerprints = Fingerprints.from_dataframe(
+            self.reference_metadata, self.fingerprint_type, self.fingerprint_nbits
+        )
+        self.top_k_tanimoto_scores = TopKTanimotoScores.calculate_from_fingerprints(
+            fingerprints, fingerprints, self.top_k_inchikeys
+        )
+
     def save(self, store_file_directory: str | Path):
         store_file_directory = Path(store_file_directory)
         if not store_file_directory.is_file():

From 15a0bc8f7918f2570202c7388498a0b831cc48f1 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 13:55:53 +0200
Subject: [PATCH 05/20] Fix bug testing file is a file

---
 ms2query/ms2query_development/ReferenceLibrary.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 870a8e2..919de58 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -145,10 +145,6 @@ def add_spectra(self, new_library_spectra):
 
     def save(self, store_file_directory: str | Path):
         store_file_directory = Path(store_file_directory)
-        if not store_file_directory.is_file():
-            raise ValueError(
-                f"The store_file_directory is a file, please use a directory instead, given dir: {store_file_directory}"
-            )
         store_file_directory.mkdir(parents=True, exist_ok=True)
 
         def file_does_not_exist_yet(file_name: Path):

From f2dabaf2922ce9f04fc2de9d886befd740d40a7b Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 14:23:37 +0200
Subject: [PATCH 06/20] Fix bug

---
 ms2query/ms2query_development/ReferenceLibrary.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 919de58..30d2373 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -148,7 +148,7 @@ def save(self, store_file_directory: str | Path):
         store_file_directory.mkdir(parents=True, exist_ok=True)
 
         def file_does_not_exist_yet(file_name: Path):
-            if file_name.exists:
+            if file_name.exists():
                 print(f"The file: {file_name} already exists, not saved.")
                 return False
             return True

From cfb1912d872019b99d2318577f0b35662e82eb80 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 14:24:38 +0200
Subject: [PATCH 07/20] Compute tanimoto scores in batches to reduce memory
 footprint

---
 .../TopKTanimotoScores.py                     | 53 ++++++++++++++-----
 1 file changed, 41 insertions(+), 12 deletions(-)

diff --git a/ms2query/ms2query_development/TopKTanimotoScores.py b/ms2query/ms2query_development/TopKTanimotoScores.py
index c930ecc..6bc1512 100644
--- a/ms2query/ms2query_development/TopKTanimotoScores.py
+++ b/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"""

From 2664d86a69cf50d9843e463ec358ad42a8b8bb0c Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Fri, 3 Jul 2026 15:44:44 +0200
Subject: [PATCH 08/20] Add add spectra to tutorial

---
 ms2query/notebooks/tutorial.ipynb | 54 ++++++++++++++++++++++++++++---
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/ms2query/notebooks/tutorial.ipynb b/ms2query/notebooks/tutorial.ipynb
index 12af41f..3888048 100644
--- a/ms2query/notebooks/tutorial.ipynb
+++ b/ms2query/notebooks/tutorial.ipynb
@@ -385,21 +385,67 @@
     "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": "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",
@@ -1088,7 +1134,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "ms2query_2",
    "language": "python",
    "name": "python3"
   },

From 34afddea30cd4b35a8d8f73a16c7f70d77f55a61 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 11:12:11 +0200
Subject: [PATCH 09/20] Use first 14 inchikey characters for grouping from
 dataframe

---
 ms2query/ms2query_development/Fingerprints.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ms2query/ms2query_development/Fingerprints.py b/ms2query/ms2query_development/Fingerprints.py
index d5770a9..931847b 100644
--- a/ms2query/ms2query_development/Fingerprints.py
+++ b/ms2query/ms2query_development/Fingerprints.py
@@ -38,7 +38,9 @@ def compute_fingerprints_from_inchi(cls, most_common_inchi_per_inchikey: dict[st
     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("inchikey")["inchi"].agg(lambda x: x.value_counts().idxmax()).to_dict()
+            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)
 

From e963380ae3efa03d38ec5142b8c8f7f48290588e Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 11:13:19 +0200
Subject: [PATCH 10/20] Clean up init and make setter for refence metadata

---
 .../ms2query_development/ReferenceLibrary.py  | 39 ++++++++++++-------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 30d2373..ab58724 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -44,28 +44,41 @@ def __init__(
         self.reference_embeddings = reference_embeddings
         self.top_k_tanimoto_scores = top_k_tanimoto_scores
         self.reference_metadata = reference_metadata
+        self._validate()
 
+    @property
+    def reference_metadata(self):
+        return self._reference_metadata
+
+    @reference_metadata.setter
+    def reference_metadata(self, reference_metadata: pd.DataFrame):
+        self._reference_metadata = reference_metadata
+        # Get the spectrum_indices_per_inchikey
+        self.spectrum_indices_per_inchikey = defaultdict(list)
+        for lib_spec_index, inchikey in enumerate(self.reference_metadata["inchikey"]):
+            self.spectrum_indices_per_inchikey[inchikey[:14]].append(lib_spec_index)
+
+    def _validate(self):
         # Check that the loaded files match
-        if _to_json_serializable(ms2deepscore_model.model_settings.get_dict()) != reference_embeddings.model_settings:
+        if (
+            _to_json_serializable(self.ms2deepscore_model.model_settings.get_dict())
+            != self.reference_embeddings.model_settings
+        ):
             raise ValueError(
                 "The settings of the ms2deepscore model do not match the model used for creating the library embeddings"
             )
         if list(self.reference_metadata["spectrum_hashes"]) != [
-            str(spectrum_hash) for spectrum_hash in reference_embeddings.index_to_spectrum_hash
+            str(spectrum_hash) for spectrum_hash in self.reference_embeddings.index_to_spectrum_hash
         ]:
             raise ValueError("The loaded metadata does not match the used embeddings")
-        if {inchikey[:14] for inchikey in reference_metadata["inchikey"]} != set(
-            top_k_tanimoto_scores.top_k_inchikeys_and_scores.index
+        if {inchikey[:14] for inchikey in self.reference_metadata["inchikey"]} != set(
+            self.top_k_tanimoto_scores.top_k_inchikeys_and_scores.index
         ):
             raise ValueError("The inchikeys in the metadata and in the top_k_tanimoto_scores do not match")
 
-        # Get the spectrum_indices_per_inchikey
-        self.spectrum_indices_per_inchikey = defaultdict(list)
-        for lib_spec_index, inchikey in enumerate(reference_metadata["inchikey"]):
-            self.spectrum_indices_per_inchikey[inchikey[:14]].append(lib_spec_index)
-
     @classmethod
     def load_from_directory(cls, library_file_directory) -> "ReferenceLibrary":
+        library_file_directory = Path(library_file_directory)
         reference_embeddings_file = library_file_directory / cls.embedding_file_name
         top_k_tanimoto_scores_file = library_file_directory / cls.top_k_tanimoto_scores_file_name
         reference_metadata_file = library_file_directory / cls.reference_metadata_file_name
@@ -99,14 +112,14 @@ def create_from_spectra(
         # library_spectra = list(tqdm(load_spectra(library_spectra_file), "Loading library spectra"))
         library_spectrum_set = AnnotatedSpectrumSet.create_spectrum_set(library_spectra)
         ms2deepscore_model = load_model(ms2deepscore_model_file_name)
-        fingerprints = Fingerprints.from_spectrum_set(library_spectrum_set, cls.fingerprint_type, cls.fingerprint_nbits)
-        top_k_tanimoto_scores = TopKTanimotoScores.calculate_from_fingerprints(
-            fingerprints, fingerprints, cls.top_k_inchikeys
-        )
         reference_metadata = extract_metadata_from_library(
             library_spectrum_set,
             cls.metadata_to_store,
         )
+        fingerprints = Fingerprints.from_dataframe(reference_metadata, cls.fingerprint_type, cls.fingerprint_nbits)
+        top_k_tanimoto_scores = TopKTanimotoScores.calculate_from_fingerprints(
+            fingerprints, fingerprints, cls.top_k_inchikeys
+        )
         library_spectrum_set.add_embeddings(ms2deepscore_model)
         return cls(ms2deepscore_model, library_spectrum_set.embeddings, top_k_tanimoto_scores, reference_metadata)
 

From 095a061b2639bd8a480dc83ae2b3cc27afeb19f9 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 11:14:42 +0200
Subject: [PATCH 11/20] Add __eq__ to top_k_tanimoto

---
 ms2query/ms2query_development/TopKTanimotoScores.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/ms2query/ms2query_development/TopKTanimotoScores.py b/ms2query/ms2query_development/TopKTanimotoScores.py
index 6bc1512..4a63dc2 100644
--- a/ms2query/ms2query_development/TopKTanimotoScores.py
+++ b/ms2query/ms2query_development/TopKTanimotoScores.py
@@ -132,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

From d1ecd16b6d8639ff04500417067d5bbb76de8895 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 11:14:49 +0200
Subject: [PATCH 12/20] Add test_add_spectra

---
 .../test_ReferenceLibrary.py                  | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/test_ms2query_development/test_ReferenceLibrary.py b/tests/test_ms2query_development/test_ReferenceLibrary.py
index 3768029..a44dbc4 100644
--- a/tests/test_ms2query_development/test_ReferenceLibrary.py
+++ b/tests/test_ms2query_development/test_ReferenceLibrary.py
@@ -62,3 +62,24 @@ 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)

From d3369c91b4b7e7207e75e4290520ce7510ede5be Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 11:21:36 +0200
Subject: [PATCH 13/20] Slight speedup of init

---
 ms2query/ms2query_development/ReferenceLibrary.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index ab58724..41720de 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -71,7 +71,7 @@ def _validate(self):
             str(spectrum_hash) for spectrum_hash in self.reference_embeddings.index_to_spectrum_hash
         ]:
             raise ValueError("The loaded metadata does not match the used embeddings")
-        if {inchikey[:14] for inchikey in self.reference_metadata["inchikey"]} != set(
+        if set(self.spectrum_indices_per_inchikey.keys()) != set(
             self.top_k_tanimoto_scores.top_k_inchikeys_and_scores.index
         ):
             raise ValueError("The inchikeys in the metadata and in the top_k_tanimoto_scores do not match")

From e9272adbac2b7cbae955cb1c6f37b355f6e21b6d Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Mon, 6 Jul 2026 14:24:21 +0200
Subject: [PATCH 14/20] Add ms2deepscore to results and query spectrum metadata

---
 ms2query/ms2query_development/ReferenceLibrary.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 41720de..0c0509c 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -190,6 +190,7 @@ def run_ms2query(
         num_of_query_embeddings = query_embeddings.embeddings.shape[0]
 
         library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int)
+        highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=float)
         ms2query_scores = []
         for start_idx in tqdm(
             range(0, num_of_query_embeddings, batch_size),
@@ -203,6 +204,7 @@ def run_ms2query(
             score_matrix = cosine_similarity_matrix(selected_query_embeddings, self.reference_embeddings.embeddings)
             highest_score_idx = np.argmax(score_matrix, axis=1)
             library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx
+            highest_ms2deepscore[start_idx:end_idx] = np.max(score_matrix, axis=1)
 
             # get predicted inchikeys
             predicted_inchikeys = self.reference_metadata.iloc[highest_score_idx]["inchikey"]
@@ -216,6 +218,12 @@ def run_ms2query(
         # construct results df
         results = self.reference_metadata.iloc[library_index_highest_ms2deepscore]
         results["ms2query_reliability_prediction"] = ms2query_scores
+        results["highest_ms2deepscore"] = highest_ms2deepscore
+
+        # spectrum metadata
+        results["query_precursor_mz"] = [spectrum.get("precursor_mz") for spectrum in query_spectra]
+        results["query_retention_time"] = [spectrum.get("retention_time") for spectrum in query_spectra]
+
         return results
 
 

From 59f62cbf6c6f802c0db21adcc121f92752b84da8 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Tue, 7 Jul 2026 11:38:43 +0200
Subject: [PATCH 15/20] Add subset embeddings from index

---
 ms2query/ms2query_development/Embeddings.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/ms2query/ms2query_development/Embeddings.py b/ms2query/ms2query_development/Embeddings.py
index 2b3612e..141ea01 100644
--- a/ms2query/ms2query_development/Embeddings.py
+++ b/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()

From 1a0e8b21013aaf2c4a8bf72bf4df60421bd339c3 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Tue, 7 Jul 2026 11:39:31 +0200
Subject: [PATCH 16/20] Add run_semi_targeted_ms2query

---
 .../ms2query_development/ReferenceLibrary.py  | 85 ++++++++++++++++++-
 1 file changed, 84 insertions(+), 1 deletion(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 0c0509c..9d5c9bd 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -216,7 +216,7 @@ def run_ms2query(
             )
 
         # construct results df
-        results = self.reference_metadata.iloc[library_index_highest_ms2deepscore]
+        results = self.reference_metadata.iloc[library_index_highest_ms2deepscore].copy()
         results["ms2query_reliability_prediction"] = ms2query_scores
         results["highest_ms2deepscore"] = highest_ms2deepscore
 
@@ -226,6 +226,89 @@ def run_ms2query(
 
         return results
 
+    def run_semi_targeted_ms2query(
+        self,
+        query_spectra: Sequence[Spectrum],
+        inchikeys_to_check: set[str],
+        batch_size: int = 1000,
+    ) -> pd.DataFrame:
+        """This does a search in just a subset of the inchikeys in the library.
+
+        This is a semi targeted analogue search, where you specifically search for a subset of inchikeys.
+        The rest of the library is still used to compute the MS2Query reliability score."""
+
+        query_embeddings = Embeddings.create_from_spectra(query_spectra, self.ms2deepscore_model)
+
+        num_of_query_embeddings = query_embeddings.embeddings.shape[0]
+
+        library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int)
+        highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=float)
+
+        inchikey_14_to_check = {inchikey[:14] for inchikey in inchikeys_to_check}
+
+        spectrum_indices_inchikeys_to_check = []
+        for inchikey in inchikey_14_to_check:
+            spectrum_indices_inchikeys_to_check.extend(self.spectrum_indices_per_inchikey[inchikey])
+        embeddings_inchikeys_to_check = self.reference_embeddings.subset_embeddings_from_index(
+            spectrum_indices_inchikeys_to_check
+        )
+        metadata_subset = self.reference_metadata.iloc[spectrum_indices_inchikeys_to_check]
+
+        ms2query_scores = []
+        for start_idx in tqdm(
+            range(0, num_of_query_embeddings, batch_size),
+            desc="Predicting highest ms2deepscore per batch of "
+            + str(min(batch_size, num_of_query_embeddings))
+            + " embeddings",
+        ):
+            # Do MS2DeepScore predictions for batch
+            end_idx = min(start_idx + batch_size, num_of_query_embeddings)
+            selected_query_embeddings = query_embeddings.embeddings[start_idx:end_idx]
+            score_matrix = cosine_similarity_matrix(selected_query_embeddings, embeddings_inchikeys_to_check.embeddings)
+            highest_score_idx = np.argmax(score_matrix, axis=1)
+            library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx
+            highest_ms2deepscore[start_idx:end_idx] = np.max(score_matrix, axis=1)
+
+            # get predicted inchikeys
+            predicted_inchikeys = metadata_subset.iloc[highest_score_idx]["inchikey"]
+            # Compute MS2Query reliability score
+            ms2query_scores.extend(
+                self.get_ms2query_reliability_prediction_recompute_scores(
+                    predicted_inchikeys,
+                    selected_query_embeddings,
+                )
+            )
+
+        # construct results df
+        results = metadata_subset.iloc[library_index_highest_ms2deepscore].copy()
+        results["ms2query_reliability_prediction"] = ms2query_scores
+        results["highest_ms2deepscore"] = highest_ms2deepscore
+
+        # spectrum metadata
+        results["query_precursor_mz"] = [spectrum.get("precursor_mz") for spectrum in query_spectra]
+        results["query_retention_time"] = [spectrum.get("retention_time") for spectrum in query_spectra]
+
+        return results
+
+    def get_ms2query_reliability_prediction_recompute_scores(
+        self,
+        predicted_inchikeys: list[str],
+        query_embeddings,
+    ) -> list[float]:
+        ms2query_scores = []
+        for query_spectrum_index, library_inchikey in enumerate(predicted_inchikeys):
+            top_k_inchikeys = self.top_k_tanimoto_scores.select_top_k_inchikeys(library_inchikey[:14])
+            maximum_ms2deepscores = np.zeros(self.top_k_tanimoto_scores.k, dtype=float)
+            for i, inchikey in enumerate(top_k_inchikeys):
+                spectrum_indexes = self.spectrum_indices_per_inchikey[inchikey]
+                score_matrix = cosine_similarity_matrix(
+                    query_embeddings[[query_spectrum_index], :], self.reference_embeddings.embeddings[spectrum_indexes]
+                )
+                highest_ms2deepscore = np.max(score_matrix)
+                maximum_ms2deepscores[i] = highest_ms2deepscore
+            ms2query_scores.append(np.mean(maximum_ms2deepscores))
+        return ms2query_scores
+
 
 def run_ms2query_from_files(
     query_spectrum_file,

From 89177bc6ac66b4919d18dff4f79c42a9079f2e66 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Tue, 7 Jul 2026 11:39:45 +0200
Subject: [PATCH 17/20] Add test_run-semi_targeted_search

---
 .../test_ReferenceLibrary.py                       | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tests/test_ms2query_development/test_ReferenceLibrary.py b/tests/test_ms2query_development/test_ReferenceLibrary.py
index a44dbc4..928bad8 100644
--- a/tests/test_ms2query_development/test_ReferenceLibrary.py
+++ b/tests/test_ms2query_development/test_ReferenceLibrary.py
@@ -83,3 +83,17 @@ def test_add_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)

From 36e82386ed69b429cee97a42c063b3d151fe207c Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Wed, 8 Jul 2026 11:17:14 +0200
Subject: [PATCH 18/20] Restructure code to reduce code duplication

---
 .../ms2query_development/ReferenceLibrary.py  | 141 ++++++------------
 1 file changed, 42 insertions(+), 99 deletions(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index 9d5c9bd..bba0e74 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -1,6 +1,6 @@
 from collections import defaultdict
 from pathlib import Path
-from typing import Sequence
+from typing import Iterable, Sequence
 import numpy as np
 import pandas as pd
 from matchms.importing import load_spectra
@@ -184,46 +184,30 @@ def run_ms2query(
         query_spectra: Sequence[Spectrum],
         batch_size: int = 1000,
     ) -> pd.DataFrame:
+        return self._run_ms2query(
+            query_spectra, self.reference_embeddings.embeddings, self.reference_metadata, batch_size
+        )
 
+    def _run_ms2query(
+        self, query_spectra, reference_embeddings: np.ndarray, reference_metadata: pd.DataFrame, batch_size
+    ):
         query_embeddings = Embeddings.create_from_spectra(query_spectra, self.ms2deepscore_model)
 
-        num_of_query_embeddings = query_embeddings.embeddings.shape[0]
-
-        library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int)
-        highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=float)
-        ms2query_scores = []
-        for start_idx in tqdm(
-            range(0, num_of_query_embeddings, batch_size),
-            desc="Predicting highest ms2deepscore per batch of "
-            + str(min(batch_size, num_of_query_embeddings))
-            + " embeddings",
-        ):
-            # Do MS2DeepScore predictions for batch
-            end_idx = min(start_idx + batch_size, num_of_query_embeddings)
-            selected_query_embeddings = query_embeddings.embeddings[start_idx:end_idx]
-            score_matrix = cosine_similarity_matrix(selected_query_embeddings, self.reference_embeddings.embeddings)
-            highest_score_idx = np.argmax(score_matrix, axis=1)
-            library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx
-            highest_ms2deepscore[start_idx:end_idx] = np.max(score_matrix, axis=1)
-
-            # get predicted inchikeys
-            predicted_inchikeys = self.reference_metadata.iloc[highest_score_idx]["inchikey"]
-            # Compute MS2Query reliability score
-            ms2query_scores.extend(
-                get_ms2query_reliability_prediction(
-                    predicted_inchikeys, self.spectrum_indices_per_inchikey, self.top_k_tanimoto_scores, score_matrix
-                )
-            )
+        library_index_highest_ms2deepscore, highest_ms2deepscore = get_highest_ms2deepscore(
+            query_embeddings.embeddings, reference_embeddings, batch_size
+        )
 
         # construct results df
-        results = self.reference_metadata.iloc[library_index_highest_ms2deepscore].copy()
-        results["ms2query_reliability_prediction"] = ms2query_scores
+        results = reference_metadata.iloc[library_index_highest_ms2deepscore].copy()
+        # Compute MS2Query reliability score
+        results["ms2query_reliability_prediction"] = self.get_ms2query_reliability_prediction(
+            results["inchikey"],
+            query_embeddings.embeddings,
+        )
         results["highest_ms2deepscore"] = highest_ms2deepscore
-
-        # spectrum metadata
+        # Add spectrum metadata
         results["query_precursor_mz"] = [spectrum.get("precursor_mz") for spectrum in query_spectra]
         results["query_retention_time"] = [spectrum.get("retention_time") for spectrum in query_spectra]
-
         return results
 
     def run_semi_targeted_ms2query(
@@ -232,18 +216,10 @@ def run_semi_targeted_ms2query(
         inchikeys_to_check: set[str],
         batch_size: int = 1000,
     ) -> pd.DataFrame:
-        """This does a search in just a subset of the inchikeys in the library.
+        """Does a search in just a subset of the inchikeys in the library.
 
         This is a semi targeted analogue search, where you specifically search for a subset of inchikeys.
         The rest of the library is still used to compute the MS2Query reliability score."""
-
-        query_embeddings = Embeddings.create_from_spectra(query_spectra, self.ms2deepscore_model)
-
-        num_of_query_embeddings = query_embeddings.embeddings.shape[0]
-
-        library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int)
-        highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=float)
-
         inchikey_14_to_check = {inchikey[:14] for inchikey in inchikeys_to_check}
 
         spectrum_indices_inchikeys_to_check = []
@@ -253,46 +229,11 @@ def run_semi_targeted_ms2query(
             spectrum_indices_inchikeys_to_check
         )
         metadata_subset = self.reference_metadata.iloc[spectrum_indices_inchikeys_to_check]
+        return self._run_ms2query(query_spectra, embeddings_inchikeys_to_check.embeddings, metadata_subset, batch_size)
 
-        ms2query_scores = []
-        for start_idx in tqdm(
-            range(0, num_of_query_embeddings, batch_size),
-            desc="Predicting highest ms2deepscore per batch of "
-            + str(min(batch_size, num_of_query_embeddings))
-            + " embeddings",
-        ):
-            # Do MS2DeepScore predictions for batch
-            end_idx = min(start_idx + batch_size, num_of_query_embeddings)
-            selected_query_embeddings = query_embeddings.embeddings[start_idx:end_idx]
-            score_matrix = cosine_similarity_matrix(selected_query_embeddings, embeddings_inchikeys_to_check.embeddings)
-            highest_score_idx = np.argmax(score_matrix, axis=1)
-            library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx
-            highest_ms2deepscore[start_idx:end_idx] = np.max(score_matrix, axis=1)
-
-            # get predicted inchikeys
-            predicted_inchikeys = metadata_subset.iloc[highest_score_idx]["inchikey"]
-            # Compute MS2Query reliability score
-            ms2query_scores.extend(
-                self.get_ms2query_reliability_prediction_recompute_scores(
-                    predicted_inchikeys,
-                    selected_query_embeddings,
-                )
-            )
-
-        # construct results df
-        results = metadata_subset.iloc[library_index_highest_ms2deepscore].copy()
-        results["ms2query_reliability_prediction"] = ms2query_scores
-        results["highest_ms2deepscore"] = highest_ms2deepscore
-
-        # spectrum metadata
-        results["query_precursor_mz"] = [spectrum.get("precursor_mz") for spectrum in query_spectra]
-        results["query_retention_time"] = [spectrum.get("retention_time") for spectrum in query_spectra]
-
-        return results
-
-    def get_ms2query_reliability_prediction_recompute_scores(
+    def get_ms2query_reliability_prediction(
         self,
-        predicted_inchikeys: list[str],
+        predicted_inchikeys: Iterable[str],
         query_embeddings,
     ) -> list[float]:
         ms2query_scores = []
@@ -310,6 +251,27 @@ def get_ms2query_reliability_prediction_recompute_scores(
         return ms2query_scores
 
 
+def get_highest_ms2deepscore(query_embeddings: np.ndarray, reference_embeddings: np.ndarray, batch_size):
+    num_of_query_embeddings = query_embeddings.shape[0]
+
+    library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int)
+    highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=float)
+    for start_idx in tqdm(
+        range(0, num_of_query_embeddings, batch_size),
+        desc="Predicting highest ms2deepscore per batch of "
+        + str(min(batch_size, num_of_query_embeddings))
+        + " embeddings",
+    ):
+        # Do MS2DeepScore predictions for batch
+        end_idx = min(start_idx + batch_size, num_of_query_embeddings)
+        selected_query_embeddings = query_embeddings[start_idx:end_idx]
+        score_matrix = cosine_similarity_matrix(selected_query_embeddings, reference_embeddings)
+        highest_score_idx = np.argmax(score_matrix, axis=1)
+        library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx
+        highest_ms2deepscore[start_idx:end_idx] = np.max(score_matrix, axis=1)
+    return library_index_highest_ms2deepscore, highest_ms2deepscore
+
+
 def run_ms2query_from_files(
     query_spectrum_file,
     ms2deepscore_model_file_name,
@@ -330,25 +292,6 @@ def run_ms2query_from_files(
     results_df.to_csv(save_file_location)
 
 
-def get_ms2query_reliability_prediction(
-    predicted_inchikeys: list[str],
-    spectrum_indices_per_inchikey,
-    top_k_tanimoto_scores: TopKTanimotoScores,
-    ms2deepscore_score_matrix,
-) -> list[float]:
-    ms2query_scores = []
-    for query_spectrum_index, library_inchikey in enumerate(predicted_inchikeys):
-        top_k_inchikeys = top_k_tanimoto_scores.select_top_k_inchikeys(library_inchikey[:14])
-        maximum_ms2deepscores = np.zeros(top_k_tanimoto_scores.k, dtype=float)
-        for i, inchikey in enumerate(top_k_inchikeys):
-            spectrum_indexes = spectrum_indices_per_inchikey[inchikey]
-            highest_ms2deepscore = np.max(ms2deepscore_score_matrix[query_spectrum_index, spectrum_indexes])
-            maximum_ms2deepscores[i] = highest_ms2deepscore
-        ms2query_scores.append(np.mean(maximum_ms2deepscores))
-    # todo get the spectrum hashes instead of the indexes for lookup later.
-    return ms2query_scores
-
-
 def extract_metadata_from_library(spectra: AnnotatedSpectrumSet, metadata_to_collect: list):
     collected_metadata = {key: [] for key in metadata_to_collect}
     collected_metadata["spectrum_hashes"] = []

From e3b8c7cf9f34ac88d2517209449a4c7e0e16e141 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Wed, 8 Jul 2026 12:52:01 +0200
Subject: [PATCH 19/20] Add check that inchikey in library

---
 ms2query/ms2query_development/ReferenceLibrary.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py
index bba0e74..372b8b4 100644
--- a/ms2query/ms2query_development/ReferenceLibrary.py
+++ b/ms2query/ms2query_development/ReferenceLibrary.py
@@ -224,7 +224,10 @@ def run_semi_targeted_ms2query(
 
         spectrum_indices_inchikeys_to_check = []
         for inchikey in inchikey_14_to_check:
-            spectrum_indices_inchikeys_to_check.extend(self.spectrum_indices_per_inchikey[inchikey])
+            if inchikey in self.spectrum_indices_per_inchikey:
+                spectrum_indices_inchikeys_to_check.extend(self.spectrum_indices_per_inchikey[inchikey])
+            else:
+                print(f"The inchikey: {inchikey} is not in the reference library, so it won't be searched for")
         embeddings_inchikeys_to_check = self.reference_embeddings.subset_embeddings_from_index(
             spectrum_indices_inchikeys_to_check
         )

From c626a4301e8e9aa85803e3093dfec7efb2e610f0 Mon Sep 17 00:00:00 2001
From: niekdejonge <niekdejonge@gmail.com>
Date: Wed, 8 Jul 2026 12:58:47 +0200
Subject: [PATCH 20/20] Update readme and tutorial

---
 README.md                         |  2 --
 ms2query/notebooks/tutorial.ipynb | 20 ++++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 5c1ec45..2715f04 100644
--- a/README.md
+++ b/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. 
diff --git a/ms2query/notebooks/tutorial.ipynb b/ms2query/notebooks/tutorial.ipynb
index 3888048..64bfae2 100644
--- a/ms2query/notebooks/tutorial.ipynb
+++ b/ms2query/notebooks/tutorial.ipynb
@@ -417,6 +417,26 @@
     "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",