Repository: matchms/matchms
Pull request #749 (2025-03-11 15:43:44)

Embedding-based similarity / ANN


Patch:
diff --git a/matchms/similarity/BaseEmbeddingSimilarity.py b/matchms/similarity/BaseEmbeddingSimilarity.py
new file mode 100644
index 00000000..dff9751e
--- /dev/null
+++ b/matchms/similarity/BaseEmbeddingSimilarity.py
@@ -0,0 +1,407 @@
+import numpy as np
+from typing import List, Iterable, Union, Optional, Dict, Any, Tuple
+from pathlib import Path
+from abc import abstractmethod
+from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
+from matchms.similarity.BaseSimilarity import BaseSimilarity
+from matchms.typing import SpectrumType
+import pickle
+
+try:
+    import pynndescent
+except ImportError:
+    pynndescent = None
+
+
+class BaseEmbeddingSimilarity(BaseSimilarity):
+    """Base class for similarity measures that work with embeddings.
+
+    This class provides functionality for computing similarities between spectra based on their
+    embeddings (vector representations). It supports cosine and euclidean similarity metrics,
+    and includes approximate nearest neighbor (ANN) search capabilities.
+
+    Parameters
+    ----------
+    similarity : str
+        The similarity measure to use for comparing embeddings. Default is "cosine".
+        Options are "cosine" or "euclidean".
+
+    Attributes
+    ----------
+    index : object
+        The ANN index object; if built.
+    index_backend : str
+        The backend used for ANN indexing (currently only "pynndescent" supported); if index is built.
+    index_kwargs : dict
+        Additional arguments passed to the ANN index constructor; if index is built.
+    index_k : int
+        Number of nearest neighbors used in the ANN index; if index is built.
+    """
+
+    def __init__(self, similarity: str = "cosine"):  
+        self.similarity = similarity
+        self.index = None
+        self.index_backend = None
+        self.index_kwargs = None
+        self.index_k = None
+
+        if self.similarity == "cosine":
+            self.pairwise_similarity_fn = cosine_similarity
+        elif self.similarity == "euclidean":
+            self.pairwise_similarity_fn = lambda x, y: self._distances_to_similarities(euclidean_distances(x, y))
+        else:
+            raise ValueError(f"Only cosine and euclidean similarities are supported for now. Got {self.similarity}.")
+
+    @abstractmethod
+    def compute_embeddings(self, spectra: Iterable[SpectrumType]) -> np.ndarray:
+        """Compute embeddings for a list of spectra.
+        
+        Parameters
+        ----------
+        spectra:
+            List of spectra to compute embeddings for.
+            
+        Returns
+        -------
+        np.ndarray
+            Embeddings for the spectra. Shape: (n_spectra, n_embedding_features).
+        """
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def get_embeddings(
+            self,
+            spectra: Optional[Iterable[SpectrumType]] = None,
+            npy_path: Optional[Union[str, Path]] = None) -> np.ndarray:
+        """Get embeddings either by computing them or loading from disk.
+
+        Parameters
+        ----------
+        spectra:
+            List of spectra to compute embeddings for.
+        npy_path:
+            Path to load/save embeddings from/to. If provided, embeddings are loaded from disk if it exists,
+            otherwise they are computed and saved on disk to the provided path.
+
+        Returns
+        -------
+        np.ndarray
+            Embeddings array.
+
+        Raises
+        ------
+        ValueError
+            If neither spectra nor npy_path is provided.
+        """
+        if spectra is None and npy_path is None:
+            raise ValueError("Either spectra or npy_path must be provided.")
+
+        if npy_path is not None:
+            if Path(npy_path).exists():
+                # If file path is provided and exists, load embeddings
+                embs = self.load_embeddings(npy_path)
+            else:
+                # If file path is provided and does not exist, compute embeddings and store them
+                embs = self.compute_embeddings(spectra)
+                self.store_embeddings(npy_path, embs)
+        else:
+            # If no file path is provided, compute embeddings
+            embs = self.compute_embeddings(spectra)
+        return embs
+
+    def pair(self, reference: SpectrumType, query: SpectrumType) -> float:
+        """Compute similarity between a pair of spectra.
+
+        Parameters
+        ----------
+        reference : SpectrumType
+            Reference spectrum.
+        query : SpectrumType
+            Query spectrum.
+
+        Returns
+        -------
+        float
+            Similarity score between the spectra.
+        """
+        return self.matrix([reference], [query])[0, 0]
+            
+    def matrix(
+            self,
+            references: List[SpectrumType],
+            queries: List[SpectrumType],
+            array_type: str = "numpy",
+            is_symmetric: bool = True) -> np.ndarray:
+        """Compute similarity matrix between reference and query spectra.
+
+        Parameters
+        ----------
+        references:
+            List of reference spectra.
+        queries:
+            List of query spectra.
+        array_type:
+            Type of array to return. Must be "numpy".
+        is_symmetric:
+            Whether the matrix is symmetric. Must be True.
+
+        Returns
+        -------
+        np.ndarray
+            Similarity matrix.
+
+        Raises
+        ------
+        ValueError
+            If array_type is not "numpy" or is_symmetric is False.
+        """
+        if array_type != "numpy" or not is_symmetric:
+            raise ValueError("Any embedding base similarity matrix is supposed to be dense and symmetric.")
+
+        # Compute embeddings
+        embs_ref = self.compute_embeddings(references)
+        embs_query = self.compute_embeddings(queries)
+
+        # Compute pairwise similarities matrix                
+        return self.pairwise_similarity_fn(embs_ref, embs_query)
+
+    def build_ann_index(
+            self,
+            reference_spectra: Optional[Iterable[SpectrumType]] = None,
+            embeddings_path: Optional[Union[str, Path]] = None,
+            k: int = 100,
+            index_backend: str = "pynndescent",
+            **index_kwargs) -> Any:
+        """Build an ANN index for the reference spectra.
+
+        Parameters
+        ----------
+        reference_spectra : Optional[Iterable[SpectrumType]]
+            List of reference spectra to build the ANN index for.
+        embeddings_path : Optional[Union[str, Path]]
+            If embeddings are already computed, provide the path to the numpy file.
+        k : int, optional
+            Number of nearest neighbors to use for the ANN index.
+        index_backend : str, optional
+            Backend to use for ANN index. Currently only "pynndescent" is supported.
+        **index_kwargs
+            Additional keyword arguments passed to the index constructor.
+
+        Returns
+        -------
+        Any
+            The constructed ANN index.
+
+        Raises
+        ------
+        ImportError
+            If pynndescent is not installed.
+        ValueError
+            If an unsupported index_backend is specified.
+        """
+        # Compute reference embeddings
+        embs_ref = self.get_embeddings(reference_spectra, embeddings_path)
+
+        if index_backend == "pynndescent":
+            if not pynndescent:
+                raise ImportError("pynndescent is not installed. Please install it with `pip install pynndescent`.")
+            self.index_backend = index_backend
+            self.index_k = k
+            self.index_kwargs = index_kwargs
+
+            # Build ANN index
+            index = pynndescent.NNDescent(embs_ref, metric=self.similarity, n_neighbors=k, **index_kwargs)
+        else:
+            raise ValueError(f"Only pynndescent is supported for now. Got {index_backend}.")
+
+        # Keep index in memory
+        self.index = index
+        return self.index
+
+    def get_anns(
+            self,
+            query_spectra: Union[Iterable[SpectrumType], np.ndarray],
+            k: int = 100) -> Tuple[np.ndarray, np.ndarray]:
+        """Get approximate nearest neighbors for query spectra.
+
+        Parameters
+        ----------
+        query_spectra : Union[Iterable[SpectrumType], np.ndarray]
+            Query spectra or their embeddings.
+        k : int, optional
+            Number of nearest neighbors to return.
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray]
+            Neighbor indices and similarity scores.
+
+        Raises
+        ------
+        ValueError
+            If no index is built or k is larger than index k.
+        """
+        if self.index is None:
+            raise ValueError(
+                "No index built yet. Please call `build_ann_index` on your reference spectra or `load_ann_index` if it "
+                "was previously built and stored using `save_ann_index`."
+            )
+
+        if k > self.index_k:
+            raise ValueError(f"k ({k}) is larger than the k used to build the index ({self.index_k}).")
+    
+        if isinstance(query_spectra, np.ndarray):
+            embs_query = query_spectra
+            if embs_query.ndim != 2:
+                raise ValueError(f"Expected 2D embeddings array, got {embs_query.ndim}D array.")
+        else:
+            # Compute query embeddings
+            embs_query = self.compute_embeddings(query_spectra)
+
+        # Get ANN indices
+        if self.index_backend == "pynndescent":
+            neighbors, distances = self.index.query(embs_query, k=k)
+            similarities = self._distances_to_similarities(distances)
+        else:
+            raise ValueError(f"Only pynndescent is supported for now. Got {self.index_backend}.")
+        return neighbors, similarities
+
+    def get_index_anns(self) -> Tuple[np.ndarray, np.ndarray]:
+        """Get nearest neighbors for all points in the index.
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray]
+            Neighbor indices and similarity scores.
+
+        Raises
+        ------
+        ValueError
+            If unsupported index backend is used.
+        """
+        if self.index_backend == "pynndescent":
+            neighbors, distances = self.index.neighbor_graph
+            similarities = self._distances_to_similarities(distances)
+            return neighbors, similarities
+        raise ValueError(f"Only pynndescent is supported for now. Got {self.index_backend}.")
+
+    def _distances_to_similarities(self, distances: np.ndarray) -> np.ndarray:
+        """Convert distances to similarities based on similarity metric.
+
+        Parameters
+        ----------
+        distances : np.ndarray
+            Distance matrix.
+
+        Returns
+        -------
+        np.ndarray
+            Similarity matrix.
+
+        Raises
+        ------
+        ValueError
+            If unsupported similarity metric is used.
+        """
+        if self.similarity == "cosine":
+            return 1 - distances
+        elif self.similarity == "euclidean":
+            return -distances
+        raise ValueError(f"Only cosine and euclidean similarities are supported for now. Got {self.similarity}.")
+
+    @staticmethod
+    def load_embeddings(npy_path: Union[str, Path]) -> np.ndarray:
+        """Load embeddings from a numpy file.
+        
+        Parameters
+        ----------
+        npy_path : Union[str, Path]
+            Path to the numpy file.
+        
+        Returns
+        -------
+        np.ndarray
+            Embeddings array.
+
+        Raises
+        ------
+        ValueError
+            If loaded array is not 2D.
+        """
+        embs = np.load(npy_path)
+        if embs.ndim != 2:
+            raise ValueError(f"Expected 2D embeddings array, got {embs.ndim}D array.")
+        return embs
+
+    @staticmethod
+    def store_embeddings(npy_path: Union[str, Path], embs: np.ndarray) -> None:
+        """Store embeddings in a numpy file.
+        
+        Parameters
+        ----------
+        npy_path : Union[str, Path]
+            Path to save the embeddings to.
+        embs : np.ndarray
+            Embeddings array to store.
+        """
+        np.save(npy_path, embs)
+
+    def save_ann_index(self, path: Union[str, Path]) -> None:
+        """Save the ANN index to disk.
+        
+        Parameters
+        ----------
+        path : Union[str, Path]
+            Path to save the index to.
+
+        Raises
+        ------
+        ValueError
+            If no index exists to save.
+        """
+        if self.index is None:
+            raise ValueError("No index to save. Please build an index first using build_ann_index().")
+        
+        save_dict = {
+            'index': self.index,
+            'backend': self.index_backend,
+            'similarity': self.similarity,
+            'index_kwargs': self.index_kwargs,
+            'index_k': self.index_k
+        }
+        
+        with open(path, 'wb') as f:
+            pickle.dump(save_dict, f)
+
+    def load_ann_index(self, path: Union[str, Path]) -> Any:
+        """Load an ANN index from disk.
+        
+        Parameters
+        ----------
+        path : Union[str, Path]
+            Path to load the index from.
+            
+        Returns
+        -------
+        Any
+            The loaded ANN index.
+
+        Raises
+        ------
+        ValueError
+            If loaded index similarity metric doesn't match current metric.
+        """
+        with open(path, 'rb') as f:
+            load_dict = pickle.load(f)
+            
+        if load_dict['similarity'] != self.similarity:
+            raise ValueError(
+                f"Loaded index similarity metric ({load_dict['similarity']}) does not match "
+                f"current similarity metric ({self.similarity})"
+            )
+            
+        self.index = load_dict['index']
+        self.index_backend = load_dict['backend']
+        self.index_kwargs = load_dict['index_kwargs']
+        self.index_k = load_dict['index_k']
+        
+        return self.index
diff --git a/matchms/similarity/BinnedEmbeddingSimilarity.py b/matchms/similarity/BinnedEmbeddingSimilarity.py
new file mode 100644
index 00000000..b5a7df91
--- /dev/null
+++ b/matchms/similarity/BinnedEmbeddingSimilarity.py
@@ -0,0 +1,88 @@
+import numpy as np
+from matchms.similarity import BaseEmbeddingSimilarity
+from matchms.typing import SpectrumType
+from typing import Iterable
+
+
+class BinnedEmbeddingSimilarity(BaseEmbeddingSimilarity):
+    """A similarity measure that bins spectra into a fixed number of bins and uses the binned
+    intensities as embedding features. By default, the similarity between spectra is computed as the cosine
+    similarity between their binned representations.
+
+    Parameters
+    ----------
+    similarity : str, optional
+        The similarity measure to use for comparing embeddings. Default is "cosine".
+        Options are "cosine" or "euclidean".
+    max_mz : float, optional
+        The maximum m/z value to consider when binning. Default is 1005.
+    bin_width : float, optional
+        The width of each bin in m/z units. Default is 1.
+    """
+    def __init__(self, similarity: str = "cosine", max_mz: float = 1005, bin_width: float = 1):
+        super().__init__(similarity=similarity)
+        self.max_mz = max_mz
+        self.bin_width = bin_width
+
+    def _bin_spectrum(self, spectrum: SpectrumType) -> np.ndarray:
+        """Bin a spectrum's peaks into fixed-width m/z bins.
+
+        Parameters
+        ----------
+        spectrum : SpectrumType
+            The spectrum to bin.
+
+        Returns
+        -------
+        np.ndarray
+            Array of binned and normalized intensities.
+        """
+        # NOTE: copypaste from https://github.com/pluskal-lab/MassSpecGym/blob/f525a5e55a39ec4caa4f1a51e64acd046713179e/massspecgym/data/transforms.py#L97
+        mzs = spectrum.peaks.mz
+        intensities = spectrum.peaks.intensities
+
+        # Calculate the number of bins
+        num_bins = int(np.ceil(self.max_mz / self.bin_width))
+
+        # Calculate the bin indices for each mass
+        bin_indices = np.floor(mzs / self.bin_width).astype(int)
+
+        # Filter out mzs that exceed max_mz
+        valid_indices = bin_indices[mzs <= self.max_mz]
+        valid_intensities = intensities[mzs <= self.max_mz]
+
+        # Clip bin indices to ensure they are within the valid range
+        valid_indices = np.clip(valid_indices, 0, num_bins - 1)
+
+        # Initialize an array to store the binned intensities
+        binned_intensities = np.zeros(num_bins)
+
+        # Use np.add.at to sum intensities in the appropriate bins
+        np.add.at(binned_intensities, valid_indices, valid_intensities)
+
+        # Normalize the intensities to relative intensities
+        binned_intensities /= np.max(binned_intensities)
+
+        return binned_intensities
+
+    def compute_embeddings(self, spectra: Iterable[SpectrumType]) -> np.ndarray:
+        """Convert spectra into binned embeddings.
+
+        Parameters
+        ----------
+        spectra : Iterable[SpectrumType]
+            The spectra to convert into embeddings.
+
+        Returns
+        -------
+        np.ndarray
+            Array of shape (n_spectra, n_bins) containing the binned embeddings.
+        """
+        spectra_list = list(spectra)
+        embeddings = []
+
+        for spectrum in spectra_list:
+            binned_spectrum = self._bin_spectrum(spectrum)
+            embeddings.append(binned_spectrum)
+
+        return np.array(embeddings)
diff --git a/matchms/similarity/__init__.py b/matchms/similarity/__init__.py
index 8ce98be8..396bc0c6 100644
--- a/matchms/similarity/__init__.py
+++ b/matchms/similarity/__init__.py
@@ -27,8 +27,8 @@ from .ModifiedCosine import ModifiedCosine
 from .NeutralLossesCosine import NeutralLossesCosine
 from .ParentMassMatch import ParentMassMatch
 from .PrecursorMzMatch import PrecursorMzMatch
-
-
+from .BaseEmbeddingSimilarity import BaseEmbeddingSimilarity
+from .BinnedEmbeddingSimilarity import BinnedEmbeddingSimilarity
 __all__ = [
     "CosineGreedy",
     "CosineHungarian",
@@ -39,6 +39,8 @@ __all__ = [
     "NeutralLossesCosine",
     "ParentMassMatch",
     "PrecursorMzMatch",
+    "BaseEmbeddingSimilarity",
+    "BinnedEmbeddingSimilarity"
 ]
 
 
@@ -53,7 +55,8 @@ def get_similarity_function_by_name(similarity_function_name: str):
     """
     names = __all__
     functions = [CosineGreedy, CosineHungarian, FingerprintSimilarity, IntersectMz, MetadataMatch, ModifiedCosine,
-                 NeutralLossesCosine, ParentMassMatch, PrecursorMzMatch]
+                 NeutralLossesCosine, ParentMassMatch, PrecursorMzMatch, BaseEmbeddingSimilarity,
+                 BinnedEmbeddingSimilarity]
 
     assert similarity_function_name in names, f"Unknown similarity function: {similarity_function_name}"
     assert len(names) == len(functions), "Number of similarity functions and names do not match"
diff --git a/pyproject.toml b/pyproject.toml
index 4fe93e90..33138964 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,6 +44,7 @@ rdkit = "^2024.3.5"
 pyyaml = ">=6.0.1"
 deprecated = ">=1.2.14"
 pubchempy = "*"
+pynndescent = ">=0.5.13"
 
 [tool.poetry.group.dev.dependencies]
 decorator = "^5.1.1"


Test patch:
diff --git a/tests/similarity/test_binned_embedding_similarity.py b/tests/similarity/test_binned_embedding_similarity.py
new file mode 100644
index 00000000..c70aaa00
--- /dev/null
+++ b/tests/similarity/test_binned_embedding_similarity.py
@@ -0,0 +1,197 @@
+import numpy as np
+import pytest
+from matchms import Spectrum
+from matchms.similarity import BinnedEmbeddingSimilarity
+
+
+def test_binned_embedding_similarity_without_parameters():
+    """Test binned embedding similarity with default parameters."""
+    spectrum_1 = Spectrum(mz=np.array([100, 200, 300, 500, 510], dtype="float"),
+                         intensities=np.array([0.1, 0.2, 1.0, 0.3, 0.4], dtype="float"))
+
+    spectrum_2 = Spectrum(mz=np.array([100, 200, 290, 490, 510], dtype="float"),
+                         intensities=np.array([0.1, 0.2, 1.0, 0.3, 0.4], dtype="float"))
+
+    similarity = BinnedEmbeddingSimilarity()
+    score = similarity.pair(spectrum_1, spectrum_2)
+
+    # Calculate expected binned vectors
+    expected_vec1 = np.zeros(1005)
+    expected_vec2 = np.zeros(1005)
+    
+    # Fill bins for spectrum 1
+    expected_vec1[[100, 200, 300, 500, 510]] = [0.1, 0.2, 1.0, 0.3, 0.4]
+    expected_vec1 /= np.max(expected_vec1)
+    
+    # Fill bins for spectrum 2 
+    expected_vec2[[100, 200, 290, 490, 510]] = [0.1, 0.2, 1.0, 0.3, 0.4]
+    expected_vec2 /= np.max(expected_vec2)
+
+    # Calculate expected cosine similarity
+    expected_score = np.dot(expected_vec1, expected_vec2) / (np.linalg.norm(expected_vec1) * np.linalg.norm(expected_vec2))
+
+    assert score == pytest.approx(expected_score, rel=1e-6), "Expected different cosine score"
+
+
+def test_binned_embedding_similarity_matrix():
+    """Test binned embedding similarity matrix computation."""
+    spectrum_1 = Spectrum(mz=np.array([100, 200, 300], dtype="float"),
+                         intensities=np.array([0.1, 0.2, 1.0], dtype="float"))
+    spectrum_2 = Spectrum(mz=np.array([110, 190, 290], dtype="float"),
+                         intensities=np.array([0.5, 0.2, 1.0], dtype="float"))
+
+    spectra = [spectrum_1, spectrum_2]
+    similarity = BinnedEmbeddingSimilarity()
+    scores = similarity.matrix(spectra, spectra)
+
+    # Calculate expected binned vectors
+    expected_vec1 = np.zeros(1005)
+    expected_vec2 = np.zeros(1005)
+    
+    expected_vec1[[100, 200, 300]] = [0.1, 0.2, 1.0]
+    expected_vec1 /= np.max(expected_vec1)
+    
+    expected_vec2[[110, 190, 290]] = [0.5, 0.2, 1.0]
+    expected_vec2 /= np.max(expected_vec2)
+
+    # Calculate expected cosine similarity
+    expected_score = np.dot(expected_vec1, expected_vec2) / (np.linalg.norm(expected_vec1) * np.linalg.norm(expected_vec2))
+
+    assert scores[0, 1] == pytest.approx(expected_score, rel=1e-6), "Expected different cosine score"
+    assert np.allclose(scores[0, 0], scores[1, 1], atol=1e-6), "Expected perfect self-similarity"
+    assert np.allclose(scores[0, 1], scores[1, 0], atol=1e-6), "Expected symmetric matrix"
+
+
+def test_binned_embedding_similarity_parameters():
+    """Test binned embedding similarity with different parameters."""
+    spectrum = Spectrum(mz=np.array([100, 200, 300], dtype="float"),
+                       intensities=np.array([0.1, 0.2, 1.0], dtype="float"))
+
+    similarity_1 = BinnedEmbeddingSimilarity(max_mz=500, bin_width=1)
+    similarity_2 = BinnedEmbeddingSimilarity(max_mz=500, bin_width=10)
+
+    embedding_1 = similarity_1.compute_embeddings([spectrum])
+    embedding_2 = similarity_2.compute_embeddings([spectrum])
+
+    assert embedding_1.shape[1] == 500, "Expected 500 bins"
+    assert embedding_2.shape[1] == 50, "Expected 50 bins"
+
+    # Check if binning is correct for bin_width=10
+    expected_vec = np.zeros(50)
+    expected_vec[[10, 20, 30]] = [0.1, 0.2, 1.0]  # 100/10=10, 200/10=20, 300/10=30
+    expected_vec /= np.max(expected_vec)
+
+    assert np.allclose(embedding_2[0], expected_vec), "Expected different binned values"
+
+
+def test_binned_embedding_similarity_euclidean():
+    """Test binned embedding similarity with euclidean distance."""
+    spectrum_1 = Spectrum(mz=np.array([100, 200, 300], dtype="float"),
+                         intensities=np.array([0.1, 0.2, 1.0], dtype="float"))
+    spectrum_2 = Spectrum(mz=np.array([110, 210, 310], dtype="float"),
+                         intensities=np.array([0.1, 0.2, 1.0], dtype="float"))
+
+    similarity = BinnedEmbeddingSimilarity(similarity="euclidean")
+    score = similarity.pair(spectrum_1, spectrum_2)
+
+    # For shifted peaks, expect non-zero distance
+    assert score < 0, "Expected negative score for euclidean similarity"
+
+
+def test_binned_embedding_similarity_invalid_similarity():
+    """Test binned embedding similarity with invalid similarity measure."""
+    with pytest.raises(ValueError):
+        BinnedEmbeddingSimilarity(similarity="invalid")
+
+
+def test_binned_embedding_similarity_ann():
+    """Test approximate nearest neighbor search functionality."""
+    # Create 10 test spectra with slightly varying peaks
+    reference_spectra = []
+    for i in range(10):
+        mz = np.array([100 + i, 200 + i, 300 + i], dtype="float")
+        intensities = np.array([0.1 + i*0.01, 0.2 + i*0.01, 1.0 - i*0.01], dtype="float") 
+        spectrum = Spectrum(mz=mz, intensities=intensities)
+        reference_spectra.append(spectrum)
+
+    # Create query spectrum similar to first reference spectrum
+    query_spectrum = Spectrum(
+        mz=np.array([100, 200, 300], dtype="float"),
+        intensities=np.array([0.11, 0.21, 0.99], dtype="float")
+    )
+
+    similarity = BinnedEmbeddingSimilarity()
+    
+    # Build index
+    similarity.build_ann_index(reference_spectra, k=5)
+    
+    # Test get_anns
+    neighbors, similarities = similarity.get_anns([query_spectrum], k=3)
+    assert neighbors.shape == (1, 3), "Expected shape (1,3) for neighbors"
+    assert similarities.shape == (1, 3), "Expected shape (1,3) for similarities"
+    
+    # Query spectrum should be most similar to first reference spectrum
+    assert neighbors[0,0] == 0, "Expected first spectrum to be closest neighbor"
+    
+    # Test get_index_anns
+    neighbors, similarities = similarity.get_index_anns()
+    assert neighbors.shape == (10, 5), "Expected shape (10,5) for index neighbors"
+    assert similarities.shape == (10, 5), "Expected shape (10,5) for index similarities"
+
+    # Test against exact nearest neighbor
+    exact_neighbors, exact_similarities = similarity.get_anns([query_spectrum], k=1)
+    assert exact_neighbors.shape == (1, 1), "Expected shape (1,1) for exact neighbors"
+    assert exact_similarities.shape == (1, 1), "Expected shape (1,1) for exact similarities"
+    assert exact_neighbors[0,0] == 0, "Expected first spectrum to be exact nearest neighbor"
+
+
+def test_binned_embedding_similarity_ann_save_load(tmp_path):
+    """Test saving and loading ANN index."""
+    # Create 10 test spectra
+    reference_spectra = []
+    for i in range(10):
+        mz = np.array([100 + i, 200 + i, 300 + i], dtype="float")
+        intensities = np.array([0.1 + i*0.01, 0.2 + i*0.01, 1.0 - i*0.01], dtype="float")
+        spectrum = Spectrum(mz=mz, intensities=intensities)
+        reference_spectra.append(spectrum)
+    
+    similarity = BinnedEmbeddingSimilarity()
+    similarity.build_ann_index(reference_spectra, k=5)
+    
+    # Save index
+    index_path = tmp_path / "test_index.pkl"
+    similarity.save_ann_index(index_path)
+    
+    # Load index in new instance
+    similarity_2 = BinnedEmbeddingSimilarity()
+    similarity_2.load_ann_index(index_path)
+    
+    # Test loaded index with first spectrum
+    neighbors, similarities = similarity_2.get_anns([reference_spectra[0]], k=1)
+    assert neighbors.shape == (1, 1), "Expected shape (1,1) for neighbors"
+    assert similarities.shape == (1, 1), "Expected shape (1,1) for similarities"
+    assert neighbors[0,0] == 0, "Expected self as nearest neighbor"
+
+
+def test_binned_embedding_similarity_ann_errors():
+    """Test error handling in ANN functionality."""
+    # Create 10 test spectra
+    reference_spectra = []
+    for i in range(10):
+        mz = np.array([100 + i, 200 + i, 300 + i], dtype="float")
+        intensities = np.array([0.1 + i*0.01, 0.2 + i*0.01, 1.0 - i*0.01], dtype="float")
+        spectrum = Spectrum(mz=mz, intensities=intensities)
+        reference_spectra.append(spectrum)
+    
+    similarity = BinnedEmbeddingSimilarity()
+    
+    # Test error when getting ANNs without building index
+    with pytest.raises(ValueError, match="No index built yet"):
+        similarity.get_anns([reference_spectra[0]])
+    
+    # Build index
+    similarity.build_ann_index(reference_spectra, k=1)
+    
+    # Test error when k is larger than index k
+    with pytest.raises(ValueError, match="k .* is larger than"):
+        similarity.get_anns([reference_spectra[0]], k=2)