Refactor embedding computation method(s)

ms2deepscore/MS2DeepScore.py

@@ -2,8 +2,7 @@
 import numpy as np
 from matchms import Spectrum
 from matchms.similarity.BaseSimilarity import BaseSimilarity
-from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel,
-                                                      compute_embedding_array)
+from ms2deepscore.models.SiameseSpectralModel import SiameseSpectralModel
 from .vector_operations import cosine_similarity, cosine_similarity_matrix
 
 
@@ -55,8 +54,14 @@ def __init__(self, model: SiameseSpectralModel, progress_bar: bool = True):
         self.output_vector_dim = self.model.model_settings.embedding_dim
         self.progress_bar = progress_bar
 
-    def get_embedding_array(self, spectrums):
-        return compute_embedding_array(self.model, spectrums, progress_bar=self.progress_bar)
+    def get_embedding_array(self, spectrums, datatype: str = "numpy", batch_size: int = 1024) -> np.ndarray:
+        """Calculate the spectrum embeddings for a list of spectrums."""
+        return self.model.compute_embedding_array(
+            spectrums,
+            datatype=datatype,
+            progress_bar=self.progress_bar,
+            batch_size=batch_size
+            )
 
     def pair(self, reference: Spectrum, query: Spectrum) -> float:
         """Calculate the MS2DeepScore similaritiy between a reference and a query spectrum.

ms2deepscore/MS2DeepScoreEvaluated.py

@@ -4,8 +4,7 @@
 from matchms.similarity.BaseSimilarity import BaseSimilarity
 from ms2deepscore.models.LinearEmbeddingEvaluation import \
     compute_error_predictions
-from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel,
-                                                      compute_embedding_array)
+from ms2deepscore.models.SiameseSpectralModel import SiameseSpectralModel
 from ms2deepscore.vector_operations import (cosine_similarity,
                                             cosine_similarity_matrix)
 
@@ -69,8 +68,13 @@ def __init__(self, model: SiameseSpectralModel,
         self.output_vector_dim = self.model.model_settings.embedding_dim
         self.progress_bar = progress_bar
 
-    def get_embedding_array(self, spectrums, datatype="numpy"):
-        return compute_embedding_array(self.model, spectrums, datatype)
+    def get_embedding_array(self, spectrums, datatype="numpy", batch_size=1024):
+        return self.model.compute_embedding_array(
+            spectrums,
+            datatype=datatype,
+            batch_size=batch_size,
+            progress_bar=self.progress_bar,
+        )
 
     def get_embedding_evaluations(self, embeddings):
         """Compute the RMSE.

ms2deepscore/models/SiameseSpectralModel.py

@@ -1,10 +1,10 @@
 import os
-from typing import Optional, Union, Dict, Any
+from typing import Optional, Union, Dict, Any, Literal
 from pathlib import Path
 import numpy as np
 
-from torch import save, cat, zeros, cuda, no_grad
-from torch import device as torch_device
+import torch
+from torch import save, cat, no_grad
 from torch.nn.functional import relu
 from torch.optim import Adam
 from torch import nn
@@ -75,6 +75,91 @@ def save(self, filepath: Union[str, Path]) -> None:
         # Important: no custom objects outside tensors/strings/primitives.
         save(checkpoint, str(filepath))
 
+    def compute_embedding_array(
+        self,
+        spectra,
+        datatype: Literal["numpy", "pytorch"] = "numpy",
+        device: Optional[torch.device | str] = None,
+        batch_size: int = 1024,
+        progress_bar: bool = True,
+    ):
+        """
+        Compute embeddings for a list of matchms Spectrum objects.
+
+        Parameters
+        ----------
+        spectra:
+            List of spectra to embed.
+        datatype:
+            "numpy" returns a NumPy array.
+            "pytorch" returns a CPU torch tensor.
+        device:
+            Device used for inference. If None, CUDA is used when available.
+        batch_size:
+            Number of spectra processed per encoder forward pass.
+        progress_bar:
+            Show progress bar.
+
+        Returns
+        -------
+        np.ndarray or torch.Tensor
+            Embedding array with shape (n_spectra, embedding_dim).
+        """
+        datatype = datatype.lower()
+        if datatype not in {"numpy", "pytorch"}:
+            raise ValueError("datatype can only be 'numpy' or 'pytorch'.")
+
+        if batch_size <= 0:
+            raise ValueError("batch_size must be a positive integer.")
+
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            device = torch.device(device)
+
+        n_spectra = len(spectra)
+        embedding_dim = self.model_settings.embedding_dim
+
+        self.to(device)
+        self.eval()
+
+        if datatype == "numpy":
+            embeddings = np.empty((n_spectra, embedding_dim), dtype=np.float32)
+        else:
+            embeddings = torch.empty((n_spectra, embedding_dim), dtype=torch.float32)
+
+        batch_starts = range(0, n_spectra, batch_size)
+
+        with no_grad():
+            with tqdm(
+                total=n_spectra,
+                desc="Computing spectral embeddings ...",
+                unit="spectrum",
+                disable=not progress_bar,
+            ) as progress:
+                for start in batch_starts:
+                    stop = min(start + batch_size, n_spectra)
+                    batch_spectra = spectra[start:stop]
+
+                    spectra_tensors, metadata_tensors = tensorize_spectra(
+                        batch_spectra,
+                        self.model_settings,
+                    )
+
+                    batch_embeddings = self.encoder(
+                        spectra_tensors.to(device),
+                        metadata_tensors.to(device),
+                    ).detach().cpu()
+
+                    if datatype == "numpy":
+                        embeddings[start:stop, :] = batch_embeddings.numpy()
+                    else:
+                        embeddings[start:stop, :] = batch_embeddings
+
+                    progress.update(stop - start)
+
+        return embeddings
+
 
 class PeakBinner(nn.Module):
     """
@@ -351,45 +436,17 @@ def compute_embedding_array(
         datatype="numpy",
         device=None,
         progress_bar: bool = True,
-        ):
+        batch_size: int = 1024):
     """
-    Compute the embeddings of all given spectra (list of matchms Spectrum objects).
+    Compatibility wrapper.
 
-    Parameters
-    ----------
-    model:
-        A trained SiameseSpectralModel used to compute spectral embeddings.
-    spectra:
-        A list (or other iterable) of spectra to be embedded.
-    datatype:
-        Determines the output type of the embedding array:
-        - "numpy": returns a NumPy array of shape (n_spectra, embedding_dim).
-        - "pytorch": returns a PyTorch tensor of shape (n_spectra, embedding_dim).
-    device:
-        The device on which to perform the computation.
-        If None, it automatically uses CUDA if available, otherwise CPU.
-    progress_bar:
-        Whether to display a progress bar during embedding computation.
+    Prefer:
+        model.compute_embedding_array(...)
     """
-    if datatype.lower() not in ["numpy", "pytorch"]:
-        raise ValueError("datatype can only be 'numpy' or 'pytorch'.")
-    if datatype.lower() == "numpy":
-        embeddings = np.zeros((len(spectra), model.model_settings.embedding_dim))
-    else:
-        embeddings = zeros((len(spectra), model.model_settings.embedding_dim))
-
-    if device is None:
-        device = torch_device("cuda" if cuda.is_available() else "cpu")
-    model.to(device)
-    for i, spec in tqdm(
-            enumerate(spectra),
-            total=len(spectra),
-            desc="Computing spectral embeddings ...",
-            disable=not progress_bar):
-        X = tensorize_spectra([spec], model.model_settings)
-        with no_grad():
-            if datatype.lower() == "numpy":
-                embeddings[i, :] = model.encoder(X[0].to(device), X[1].to(device)).cpu().detach().numpy()
-            else:
-                embeddings[i, :] = model.encoder(X[0].to(device), X[1].to(device)).cpu().detach()
-    return embeddings
+    return model.compute_embedding_array(
+        spectra,
+        datatype=datatype,
+        device=device,
+        progress_bar=progress_bar,
+        batch_size=batch_size,
+    )

tests/test_ms2deepscore.py

@@ -22,11 +22,15 @@ def get_test_ms2deepscore_instance():
     return spectrums, model, similarity_measure
 
 
-def test_MS2DeepScore_vector_creation():
+@pytest.mark.parametrize("batch_size", [16, None])
+def test_MS2DeepScore_vector_creation(batch_size):
     """Test embeddings creation.
     """
     spectrums, _, similarity_measure = get_test_ms2deepscore_instance()
-    embeddings = similarity_measure.get_embedding_array(spectrums)
+    if batch_size is None:
+        embeddings = similarity_measure.get_embedding_array(spectrums)
+    else:
+        embeddings = similarity_measure.get_embedding_array(spectrums, batch_size=batch_size)
     assert embeddings.shape == (76, 100), "Expected different embeddings shape"
     assert isinstance(embeddings, np.ndarray), "Expected embeddings to be numpy array"
 

tests/test_siamese_spectra_model.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+import torch
 from matchms import Spectrum
 from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel,
                                                       train)
@@ -151,3 +152,66 @@ def test_model_training(simple_training_spectra):
     # Check if bias in data is handled correctly
     assert (np.array(history["collection_targets"]) == 1).sum() == 200
     assert (np.array(history["collection_targets"]) < .2).sum() == 200
+
+
+def test_siamese_model_compute_embedding_array_batched(dummy_spectra):
+    """Test that batched embedding computation matches single-spectrum batching."""
+    settings = SettingsMS2Deepscore(
+        mz_bin_width=1.0,
+        base_dims=(100,),
+        embedding_dim=20,
+        train_binning_layer=False,
+    )
+    model = SiameseSpectralModel(settings)
+
+    embeddings_batch_size_1 = model.compute_embedding_array(
+        dummy_spectra,
+        datatype="numpy",
+        batch_size=1,
+        progress_bar=False,
+    )
+    embeddings_batch_size_2 = model.compute_embedding_array(
+        dummy_spectra,
+        datatype="numpy",
+        batch_size=2,
+        progress_bar=False,
+    )
+
+    assert isinstance(embeddings_batch_size_1, np.ndarray)
+    assert isinstance(embeddings_batch_size_2, np.ndarray)
+    assert embeddings_batch_size_1.shape == (len(dummy_spectra), settings.embedding_dim)
+    assert embeddings_batch_size_2.shape == (len(dummy_spectra), settings.embedding_dim)
+    assert np.allclose(embeddings_batch_size_1, embeddings_batch_size_2, atol=1e-6)
+
+    embeddings_torch = model.compute_embedding_array(
+        dummy_spectra,
+        datatype="pytorch",
+        batch_size=2,
+        progress_bar=False,
+    )
+
+    assert isinstance(embeddings_torch, torch.Tensor)
+    assert embeddings_torch.shape == (len(dummy_spectra), settings.embedding_dim)
+    assert embeddings_torch.device.type == "cpu"
+    assert not embeddings_torch.requires_grad
+    assert np.allclose(
+        embeddings_batch_size_2,
+        embeddings_torch.numpy(),
+        atol=1e-6,
+    )
+
+
+def test_siamese_model_compute_embedding_array_invalid_datatype(dummy_spectra):
+    settings = SettingsMS2Deepscore(
+        mz_bin_width=1.0,
+        base_dims=(100,),
+        embedding_dim=20,
+    )
+    model = SiameseSpectralModel(settings)
+
+    with pytest.raises(ValueError, match="datatype"):
+        model.compute_embedding_array(
+            dummy_spectra,
+            datatype="invalid",
+            progress_bar=False,
+        )