add merge_fingerprint method

Author: florian-huber

ms2query/data_processing/__init__.py

@@ -1,5 +1,5 @@
 from .chemistry_utils import compute_morgan_fingerprints, inchikey14_from_full
-from .fingerprint_computation import compute_fingerprints_from_smiles
+from .fingerprint_computation import compute_fingerprints_from_smiles, merge_fingerprints
 from .merging_utils import cluster_block, get_merged_spectra
 from .spectra_processing import compute_spectra_embeddings, normalize_spectrum_sum
 
@@ -11,5 +11,6 @@
     "compute_spectra_embeddings",
     "get_merged_spectra",
     "inchikey14_from_full",
+    "merge_fingerprints",
     "normalize_spectrum_sum",
 ]

ms2query/data_processing/fingerprint_computation.py

@@ -1,6 +1,8 @@
+from typing import Optional, Sequence, Tuple
 import numba
 import numpy as np
 from numba import typed, types
+from numpy.typing import NDArray
 from rdkit import Chem
 from tqdm import tqdm
 
@@ -255,6 +257,105 @@ def count_fingerprint_keys(fingerprints):
     return unique_keys[order], count_arr[order], first_arr[order]
 
 
+def merge_fingerprints(
+    fingerprints: Sequence[Tuple[NDArray[np.integer], NDArray[np.floating]]],
+    weights: Optional[NDArray[np.floating]] = None,
+) -> Tuple[NDArray[np.integer], NDArray[np.floating]]:
+    """
+    Merge multiple sparse Morgan (count/TF-IDF) fingerprints into a single
+    weighted-average fingerprint.
+
+    Parameters
+    ----------
+    fingerprints :
+        Sequence of (bits, values) pairs.
+        - bits: 1D integer array of bit indices (non-zero entries)
+        - values: 1D float array of TF-IDF (or other) weights,
+          same length as `bits`.
+    weights :
+        Optional 1D array-like of length len(fingerprints) with one weight
+        per fingerprint. Each fingerprint's values are scaled by its weight,
+        then the merged fingerprint is normalized by the sum of all weights.
+
+        - If None, all fingerprints are weighted equally (weight = 1.0).
+
+    Returns
+    -------
+    merged_bits, merged_values :
+        - merged_bits: 1D integer array of unique bit indices
+        - merged_values: 1D float array of weighted-average values per bit
+          (sum over all weighted fingerprints, divided by sum(weights)).
+    """
+    n_fps = len(fingerprints)
+    if n_fps == 0:
+        # Return empty sparse fingerprint
+        return (
+            np.array([], dtype=np.int64),
+            np.array([], dtype=np.float64),
+        )
+
+    if weights is not None:
+        w = np.asarray(weights, dtype=np.float64).ravel()
+        if w.shape[0] != n_fps:
+            raise ValueError(
+                f"weights must have length {n_fps}, got {w.shape[0]}"
+            )
+        total_weight = float(w.sum())
+        if total_weight <= 0.0:
+            raise ValueError("Sum of weights must be positive.")
+    else:
+        # Equal weighting
+        w = None
+        total_weight = float(n_fps)
+
+    # Concatenate all indices and (weighted) values
+    bits_list = []
+    vals_list = []
+
+    for i, (bits, vals) in enumerate(fingerprints):
+        bits = np.asarray(bits)
+        vals = np.asarray(vals, dtype=np.float64)
+
+        if bits.shape[0] != vals.shape[0]:
+            raise ValueError(
+                f"Fingerprint {i}: bits and values must have same length, "
+                f"got {bits.shape[0]} and {vals.shape[0]}"
+            )
+
+        if w is not None:
+            vals = vals * w[i]
+
+        bits_list.append(bits)
+        vals_list.append(vals)
+
+    if not bits_list:
+        return (
+            np.array([], dtype=np.int64),
+            np.array([], dtype=np.float64),
+        )
+
+    all_bits = np.concatenate(bits_list)
+    all_vals = np.concatenate(vals_list)
+
+    if all_bits.size == 0:
+        return (
+            np.array([], dtype=np.int64),
+            np.array([], dtype=np.float64),
+        )
+
+    # Group by bit index and sum weighted values
+    unique_bits, inverse = np.unique(all_bits, return_inverse=True)
+    summed_vals = np.bincount(inverse, weights=all_vals)
+
+    # Weighted average: divide by sum of weights
+    avg_vals = summed_vals / total_weight
+
+    # Keep dtypes reasonably tight
+    merged_bits = unique_bits.astype(all_bits.dtype, copy=False)
+    merged_vals = avg_vals.astype(np.float32, copy=False)
+
+    return merged_bits, merged_vals
+
 ### ------------------------
 ### Bit Scaling and Weighing
 ### ------------------------

ms2query/ms2query_library.py

@@ -4,8 +4,9 @@
 import pandas as pd
 from matchms import Spectrum
 from ms2deepscore.models import load_model as _ms2ds_load_model
+from sklearn.metrics.pairwise import cosine_similarity
 from ms2query import MS2QueryDatabase
-from ms2query.data_processing import compute_spectra_embeddings
+from ms2query.data_processing import compute_spectra_embeddings, merge_fingerprints
 from ms2query.database import EmbeddingIndex, FingerprintSparseIndex
 
 
@@ -31,6 +32,7 @@ class MS2QueryLibrary:
     db: MS2QueryDatabase
     embedding_index: Optional[EmbeddingIndex] = None
     fingerprint_index: Optional[FingerprintSparseIndex] = None  # for now: reference spectra only
+    large_scale_fingerprint_index: Optional[FingerprintSparseIndex] = None  # for large body of reference compounds
     model_path: Optional[str] = None
 
     # internal: whether to apply spectrum normalization (sum=1) before embedding
@@ -319,13 +321,51 @@ def analogue_search(
             .set_index("spec_id")
         )
 
-        smiles = analogue_compounds["smiles"].tolist()
+        analogue_smiles = analogue_compounds["smiles"].tolist()
 
         # Step 3: fingerprint-based compound search
         top_compounds = self.query_compounds_by_compounds(
-            smiles, k_compounds=k_compounds
-        )
-        return top_compounds
+            smiles=analogue_smiles
+        ).set_index("query_ix")
+
+        # Step 4: for each query, pick the best matching spectrum among all spectra
+        fingerprints_merged = []
+        weighted_average_scores = []
+        embeddings_queries = self.compute_embeddings(spectra)  # TODO: this is now done twice! in step 1 and here
+        for i in range(len(analogue_smiles)):
+            comp_ids = top_compounds.loc[i].comp_id.to_list()
+
+            # Get chemically closest compounds
+            spec_ids_all = []
+            spec_ids_selected = []
+            embeddings_selected = []
+
+            all_spec_ids = self.db.spec_ids_by_comp_ids(comp_ids).set_index("comp_id")
+            for comp_id in comp_ids:
+                new_spec_ids = all_spec_ids.loc[comp_id].spec_id.to_list()
+
+                # Get most similar embedding from one of the top-10 compounds
+                embs = self.db.ref_sdb.get_embeddings(new_spec_ids)
+                similarities = cosine_similarity(embs[1], embeddings_queries[i].reshape(1, -1))
+                max_id = np.argmax(similarities)
+                spec_ids_selected.append(embs[0][max_id])
+                embeddings_selected.append(embs[1][max_id])
+                spec_ids_all.extend(new_spec_ids)
+            
+            top1_top10_similarities = cosine_similarity(embeddings_selected, embeddings_queries[i].reshape(1, -1))
+            fingerprints = self.db.ref_cdb.get_fingerprints(comp_ids)
+            fingerprints_merged.append(merge_fingerprints(fingerprints, weights=top1_top10_similarities))
+            weighted_average_scores.append(np.sum(top1_top10_similarities ** 2) / np.sum(top1_top10_similarities))
+            if self.large_scale_fingerprint_index:
+                analogue_predictions = self.large_scale_fingerprint_index.query(fingerprints_merged, k=k_compounds)
+            elif self.fingerprint_index:
+                analogue_predictions = self.fingerprint_index.query(fingerprints_merged, k=k_compounds)
+            else:
+                raise RuntimeError("No fingerprint index is set. Build or load it before querying.")       
+        return pd.DataFrame({
+            "analogue_predictions": analogue_predictions,
+            "weighted_average_scores": weighted_average_scores
+        })
 
     # ------------------------------------------------------------------
     # Helpers / glue