Refactor spectrum set

ms2query/benchmarking/AnnotatedSpectrumSet.py

@@ -0,0 +1,169 @@
+import os
+from collections import defaultdict
+from typing import Iterable, List, Mapping, Optional, Sequence
+from matchms import Spectrum
+from matchms.exporting import save_spectra
+from matchms.importing import load_spectra
+from ms2deepscore.models import SiameseSpectralModel
+from tqdm import tqdm
+from ms2query.benchmarking.Embeddings import Embeddings
+
+
+class AnnotatedSpectrumSet:
+    """Stores a spectrum dataset making it easy and fast to split on molecules"""
+
+    def __init__(
+        self,
+        spectra: Sequence[Spectrum],
+        spectrum_indices_per_inchikey: Mapping[str, Iterable[int]],
+        embeddings: Optional[Embeddings] = None,
+    ):
+        self._spectra = tuple([spectrum.clone() for spectrum in spectra])
+        self.spectrum_indices_per_inchikey: dict[str, tuple[int, ...]] = {
+            key: tuple(values) for key, values in spectrum_indices_per_inchikey.items()
+        }
+        self.embeddings = embeddings
+
+    @classmethod
+    def create_spectrum_set(cls, spectra: Sequence[Spectrum]) -> "AnnotatedSpectrumSet":
+        spectrum_indices_per_inchikey = defaultdict(list)
+        for spectrum_index, spectrum in enumerate(tqdm(spectra, desc="Create mapping from inchikey to spectrum")):
+            inchikey = spectrum.get("inchikey")
+            if inchikey is None:
+                raise ValueError("Annotated Spectrum set expects spectra that all have an inchikey")
+            spectrum_indices_per_inchikey[inchikey[:14]].append(spectrum_index)
+        return cls(spectra, spectrum_indices_per_inchikey)
+
+    def __add__(self, other) -> "AnnotatedSpectrumSet":
+        """Adds two spectrum sets together"""
+        if not isinstance(other, AnnotatedSpectrumSet):
+            return NotImplemented
+        spectra = self.spectra + other.spectra
+        # update spectrum_indices_per_inchikey
+        starting_index = len(self.spectra)
+        reindexed_indices_per_inchikey = {}
+        for inchikey, list_of_spectrum_indices in other.spectrum_indices_per_inchikey.items():
+            reindexed_indices_per_inchikey[inchikey] = [v + starting_index for v in list_of_spectrum_indices]
+        # combine indices
+        spectrum_indices_per_inchikey = defaultdict(list)
+        for indices_per_inchikey in (self.spectrum_indices_per_inchikey, reindexed_indices_per_inchikey):
+            for inchikey, indices in indices_per_inchikey.items():
+                spectrum_indices_per_inchikey[inchikey].extend(indices)
+
+        # combine embeddings
+        if self.has_embeddings != other.has_embeddings:
+            print("Only one of the two sets has an embeddings, so embeddings are not added")
+        embeddings = None
+        if self.has_embeddings and other.has_embeddings:
+            embeddings = self.embeddings + other.embeddings
+        return AnnotatedSpectrumSet(spectra, spectrum_indices_per_inchikey, embeddings=embeddings)
+
+    def subset_spectra(self, spectrum_indices) -> "AnnotatedSpectrumSet":
+        """Returns a new instance of a subset of the spectra"""
+        spectra = [self._spectra[index] for index in spectrum_indices]
+        new_instance = AnnotatedSpectrumSet.create_spectrum_set(spectra)
+        if self.has_embeddings:
+            new_instance._embeddings = self.embeddings.subset_embeddings(spectra)
+        return new_instance
+
+    def subset_spectra_on_metadata(self, metadata_key: str, values_to_keep: set) -> "AnnotatedSpectrumSet":
+        """Creates a subset from the spectra by checking for specific metadata keys
+
+        E.g. subset_spectra_on_metadata("ionmode", set(["positive"])) will return only the spectra in positive ion mode
+        """
+        spectrum_indexes_to_keep = []
+        for spectrum_index, spectrum in enumerate(tqdm(self.spectra, desc="Checking spectra for correct metadata")):
+            if spectrum.get(metadata_key) in values_to_keep:
+                spectrum_indexes_to_keep.append(spectrum_index)
+        return self.subset_spectra(spectrum_indexes_to_keep)
+
+    def spectra_per_inchikey(self, inchikey) -> List[Spectrum]:
+        matching_spectra = []
+        for index in self.spectrum_indices_per_inchikey[inchikey]:
+            matching_spectra.append(self._spectra[index])
+        return matching_spectra
+
+    def add_embeddings(self, model: SiameseSpectralModel):
+        self._embeddings = Embeddings.create_from_spectra(self._spectra, model)
+
+    @property
+    def has_embeddings(self) -> bool:
+        if self._embeddings is None:
+            return False
+        return True
+
+    @property
+    def spectra(self):
+        return self._spectra
+
+    @property
+    def embeddings(self) -> Embeddings:
+        if self._embeddings is None:
+            raise ValueError("First run the 'add_embeddings' method")
+        return self._embeddings
+
+    @embeddings.setter
+    def embeddings(self, embeddings: Optional[Embeddings]):
+        if embeddings is None:
+            self._embeddings = embeddings
+            return
+        if not embeddings.index_to_spectrum_hash == tuple(spectrum.__hash__() for spectrum in self.spectra):
+            raise ValueError(
+                "The embeddings spectrum hashes don't match the spectrum hashes, make sure you use matching embeddings"
+            )
+        self._embeddings = embeddings
+
+    @property
+    def inchikeys(self):
+        return tuple(self.spectrum_indices_per_inchikey.keys())
+
+    def __copy__(self):
+        return AnnotatedSpectrumSet(self.spectra, self.spectrum_indices_per_inchikey, self._embeddings)
+
+    def __eq__(self, other: object):
+        if not isinstance(other, AnnotatedSpectrumSet):
+            return NotImplemented("__Eq__ can only be done between two AnnotatedSpectrumSets")
+        if self.spectra != other.spectra:
+            return False
+        if self.spectrum_indices_per_inchikey != other.spectrum_indices_per_inchikey:
+            return False
+        if self._embeddings != other._embeddings:
+            return False
+        return True
+
+    def __len__(self):
+        return len(self._spectra)
+
+    def __repr__(self):
+        return (
+            f"AnnotatedSpectrumSet(nr_of_spectra = {len(self)},"
+            f"nr_of_unique_inchikeys = {len(self.inchikeys)}, "
+            f"has_embeddings={self.has_embeddings})"
+        )
+
+    def __str__(self):
+        with_embeddings = ""
+        if self.has_embeddings:
+            with_embeddings = "with embeddings"
+
+        return f"{len(self)} spectra for {len(self.inchikeys)} inchikeys {with_embeddings}"
+
+    def save(self, save_file: str) -> None:
+        """Save spectra to the specified path"""
+        save_spectra(list(self._spectra), save_file)
+
+        if self._embeddings is not None:
+            embedding_save_name = os.path.splitext(save_file)[0] + "_embeddings.npz"
+            print(f"Saving embeddings at {embedding_save_name}")
+            self._embeddings.save(embedding_save_name)
+
+    @classmethod
+    def load(cls, spectrum_file: str) -> "AnnotatedSpectrumSet":
+        """Load mass spectra into a AnnotatedSpectrmuSet, if embeddings are available they are loaded too"""
+        spectra = list(load_spectra(spectrum_file))
+
+        embedding_file_name = os.path.splitext(spectrum_file)[0] + "_embeddings.npz"
+        instance = cls.create_spectrum_set(spectra)
+        if os.path.exists(embedding_file_name):
+            instance.embeddings = Embeddings.load(embedding_file_name)
+        return instance

ms2query/benchmarking/Embeddings.py

@@ -0,0 +1,140 @@
+import json
+from pathlib import Path
+from typing import Sequence
+import numpy as np
+import pandas as pd
+from matchms import Spectrum
+from ms2deepscore.models import SiameseSpectralModel, compute_embedding_array
+from ms2deepscore.vector_operations import cosine_similarity_matrix
+from tqdm import tqdm
+
+
+class Embeddings:
+    """Stores Embeddings for a list of mass spectra"""
+
+    def __init__(self, embeddings: np.ndarray, spectrum_hashes: tuple, model_settings: dict):
+        if len(spectrum_hashes) != embeddings.shape[0]:
+            raise ValueError("Number of spectra hashes does not match number of embeddings")
+        self.index_to_spectrum_hash = spectrum_hashes
+        self._spectrum_hash_to_index = {
+            spectrum_hash: index for index, spectrum_hash in enumerate(self.index_to_spectrum_hash)
+        }
+        self.model_settings = model_settings
+        self._embeddings = embeddings
+
+    @classmethod
+    def create_from_spectra(cls, spectra: Sequence[Spectrum], model: SiameseSpectralModel) -> "Embeddings":
+        index_to_spectrum_hash = tuple(spectrum.__hash__() for spectrum in tqdm(spectra, desc="Hashing spectra"))
+        if len(set(index_to_spectrum_hash)) != len(spectra):
+            raise ValueError("There are duplicated spectra in the spectrum list")
+
+        model_settings = model.model_settings.get_dict()
+        embeddings: np.ndarray = compute_embedding_array(model, spectra)  # type: ignore
+        return cls(embeddings, index_to_spectrum_hash, model_settings)
+
+    def __add__(self, other: "Embeddings") -> "Embeddings":
+        if not isinstance(other, Embeddings):
+            return NotImplemented
+        if self.model_settings != other.model_settings:
+            raise ValueError("Model settings of merged embeddings do not match")
+        if not set(self.index_to_spectrum_hash).isdisjoint(other.index_to_spectrum_hash):
+            raise ValueError("There are repeated spectra in the embeddings that are added together")
+        combined_embeddings = np.vstack([self._embeddings, other._embeddings])
+        index_to_spectrum_hash = self.index_to_spectrum_hash + other.index_to_spectrum_hash
+        return Embeddings(combined_embeddings, index_to_spectrum_hash, self.model_settings)
+
+    def subset_embeddings(self, spectra):
+        spectrum_hashes = tuple(spectrum.__hash__() for spectrum in spectra)
+        try:
+            embedding_indexes = [self._spectrum_hash_to_index[spectrum_hash] for spectrum_hash in spectrum_hashes]
+        except KeyError:
+            raise ValueError("The given spectra are not stored in Embeddings")
+        embeddings = self._embeddings[embedding_indexes].copy()
+        return Embeddings(embeddings, spectrum_hashes, self.model_settings)
+
+    @property
+    def embeddings(self):
+        return self._embeddings.view()
+
+    @property
+    def model_settings(self):
+        return self._model_settings.copy()
+
+    @model_settings.setter
+    def model_settings(self, model_settings):
+        self._model_settings: dict = _to_json_serializable(model_settings)
+
+    def copy(self) -> "Embeddings":
+        return Embeddings(
+            embeddings=self._embeddings.copy(),
+            spectrum_hashes=tuple(self.index_to_spectrum_hash),
+            model_settings=dict(self.model_settings),
+        )
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, Embeddings):
+            return NotImplemented
+        if self.model_settings != other.model_settings:
+            print("Model setting not equal")
+            return False
+        if self.index_to_spectrum_hash != other.index_to_spectrum_hash:
+            print("index to spectrum hash not equal")
+            return False
+        return np.array_equal(self.embeddings, other.embeddings)
+
+    def save(self, path: str | Path) -> None:
+        """Save embeddings to a .npz file with metadata stored alongside.
+
+        Args:
+            path: File path. A '.npz' extension will be added if not present.
+        """
+        path = Path(path).with_suffix(".npz")
+        metadata = {
+            "model_settings": self.model_settings,
+            "index_to_spectrum_hash": list(self.index_to_spectrum_hash),
+        }
+        np.savez_compressed(
+            path,
+            embeddings=self._embeddings,
+            metadata=np.array(json.dumps(metadata)),
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> "Embeddings":
+        """Load embeddings from a saved .npz file.
+
+        Args:
+            path: Path to the saved .npz file.
+        """
+        path = Path(path).with_suffix(".npz")
+        with np.load(path, allow_pickle=False) as data:
+            embeddings = data["embeddings"]
+            metadata = json.loads(data["metadata"].item())
+        return cls(
+            embeddings=embeddings,
+            spectrum_hashes=tuple(metadata["index_to_spectrum_hash"]),
+            model_settings=metadata["model_settings"],
+        )
+
+
+def calculate_ms2deepscore_df(query_embeddings: Embeddings, library_embeddings: Embeddings):
+    """Returns a DF, where the indexes and column labels are the spectrum hashes"""
+    ms2deepscores = cosine_similarity_matrix(query_embeddings.embeddings, library_embeddings.embeddings)
+    return pd.DataFrame(
+        ms2deepscores, index=query_embeddings.index_to_spectrum_hash, columns=library_embeddings.index_to_spectrum_hash
+    )
+
+
+def _to_json_serializable(obj):
+    """Changes a dict to be json sericalizable, so it is the same when loaded"""
+    if isinstance(obj, dict):
+        return {key: _to_json_serializable(value) for key, value in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_to_json_serializable(item) for item in obj]
+    if isinstance(obj, np.integer):
+        return int(obj)
+    if isinstance(obj, np.floating):
+        return float(obj)
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    return obj