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/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() diff --git a/ms2query/ms2query_development/Fingerprints.py b/ms2query/ms2query_development/Fingerprints.py index 4e20cce..931847b 100644 --- a/ms2query/ms2query_development/Fingerprints.py +++ b/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 = {} diff --git a/ms2query/ms2query_development/ReferenceLibrary.py b/ms2query/ms2query_development/ReferenceLibrary.py index e12c91c..372b8b4 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 @@ -26,6 +26,7 @@ class ReferenceLibrary: "collision_energy", "compound_name", "smiles", + "inchi", "inchikey", ] fingerprint_type = "daylight" @@ -43,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 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") - # 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 @@ -93,93 +107,172 @@ 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) + 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) + + 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( - library_spectrum_set, - cls.metadata_to_store, + new_spectrum_set, + self.metadata_to_store, ) + self.reference_metadata = pd.concat([self.reference_metadata, reference_metadata], ignore_index=True) - 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) - return cls(ms2deepscore_model, library_spectrum_set.embeddings, top_k_tanimoto_scores, reference_metadata) + # 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) + 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], 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, highest_ms2deepscore = get_highest_ms2deepscore( + query_embeddings.embeddings, reference_embeddings, batch_size + ) - library_index_highest_ms2deepscore = np.zeros((num_of_query_embeddings), dtype=int) - highest_ms2deepscore_values = np.zeros((num_of_query_embeddings), dtype=float) - ms2query_scores = [] + # construct results df + 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 + # 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 - 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) - highest_score_values = np.max(score_matrix, axis=1) - - library_index_highest_ms2deepscore[start_idx:end_idx] = highest_score_idx - highest_ms2deepscore_values[start_idx:end_idx] = highest_score_values - - # 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, + def run_semi_targeted_ms2query( + self, + query_spectra: Sequence[Spectrum], + inchikeys_to_check: set[str], + batch_size: int = 1000, + ) -> pd.DataFrame: + """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.""" + inchikey_14_to_check = {inchikey[:14] for inchikey in inchikeys_to_check} + + spectrum_indices_inchikeys_to_check = [] + for inchikey in inchikey_14_to_check: + 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 + ) + 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) + + def get_ms2query_reliability_prediction( + self, + predicted_inchikeys: Iterable[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 - # construct results df - results = self.reference_metadata.iloc[library_index_highest_ms2deepscore].copy() - results["predicted_tanimoto"] = highest_ms2deepscore_values - results["ms2query_reliability_prediction"] = ms2query_scores - return results + +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( @@ -202,25 +295,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"] = [] diff --git a/ms2query/ms2query_development/TopKTanimotoScores.py b/ms2query/ms2query_development/TopKTanimotoScores.py index c930ecc..4a63dc2 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""" @@ -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 diff --git a/ms2query/notebooks/tutorial.ipynb b/ms2query/notebooks/tutorial.ipynb index 7afc892..64bfae2 100644 --- a/ms2query/notebooks/tutorial.ipynb +++ b/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" }, diff --git a/tests/test_ms2query_development/test_ReferenceLibrary.py b/tests/test_ms2query_development/test_ReferenceLibrary.py index 77b5bf9..928bad8 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) @@ -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)