Added a class to support fast tolerance-window lookups

src/tools/__init__.py

@@ -2,5 +2,6 @@
 from tools.elements import Compound, Element, Isotope, IsotopeDB
 from tools.peak import Peak, Peaks
 from tools.spectra import Spectra
+from tools.utils import SortedValueIndex
 
-__all__ = ["Compound", "Database", "Element", "Isotope", "IsotopeDB", "Peak", "Peaks", "Spectra"]
+__all__ = ["Compound", "Database", "Element", "Isotope", "IsotopeDB", "Peak", "Peaks", "SortedValueIndex", "Spectra"]

src/tools/peak.py

@@ -7,7 +7,7 @@
 import pandas as pd
 
 from tools import Compound, IsotopeDB
-from tools.utils import get_ppm_range
+from tools.utils import SortedValueIndex, get_ppm_range
 
 
 class Peaks:
@@ -29,9 +29,7 @@ def __init__(self, filepath: str | Path, isotope_filepath: str | Path | None = N
         self._row_by_id = {peak_id: pos for pos, peak_id in enumerate(self._df["peak_id"])}
         # m/z sorted once so match queries can binary-search their window
         # (O(log N) per query) instead of scanning every peak.
-        mz = self._df["mz"].to_numpy()
-        self._mz_sort_order = np.argsort(mz, kind="stable")
-        self._sorted_mz = mz[self._mz_sort_order]
+        self._mz_index = SortedValueIndex(self._df["mz"].to_numpy(), get_ppm_range)
 
     def __len__(self) -> int:
         return len(self._df)
@@ -288,10 +286,7 @@ def match_mzs(self, mzs: Iterable[float], ppm_error: float) -> pd.DataFrame:
         """
         mask = np.zeros(len(self._df), dtype=bool)
         for mz in mzs:
-            lower_bound, upper_bound = get_ppm_range(mz, ppm_error)
-            lo = np.searchsorted(self._sorted_mz, lower_bound, side="left")
-            hi = np.searchsorted(self._sorted_mz, upper_bound, side="right")
-            mask[self._mz_sort_order[lo:hi]] = True
+            mask[self._mz_index.search(mz, ppm_error)] = True
 
         return self._df[mask].copy()
 

src/tools/utils/__init__.py

@@ -0,0 +1,34 @@
+from .core import (
+    aggregate_dict_values,
+    get_adducts,
+    get_charge,
+    get_decoy_info,
+    get_element_count,
+    get_file_delimiter,
+    get_file_info,
+    get_formula,
+    get_ppm_range,
+    modify_charge,
+    modify_formula_dict,
+    remove_noise,
+    str_to_dict,
+)
+from .sort_value_index import SortedValueIndex, symmetric_window
+
+__all__ = [
+    "SortedValueIndex",
+    "symmetric_window",
+    "aggregate_dict_values",
+    "get_adducts",
+    "get_charge",
+    "get_decoy_info",
+    "get_element_count",
+    "get_file_delimiter",
+    "get_file_info",
+    "get_formula",
+    "get_ppm_range",
+    "modify_charge",
+    "modify_formula_dict",
+    "remove_noise",
+    "str_to_dict",
+]

src/tools/utils/sort_value_index.py

@@ -0,0 +1,69 @@
+from typing import Callable, Tuple
+
+import numpy as np
+
+
+def symmetric_window(value: float, tolerance: float) -> Tuple[float, float]:
+    """Return the symmetric absolute window ``(value - tolerance, value + tolerance)``."""
+    return value - tolerance, value + tolerance
+
+
+class SortedValueIndex:
+    """A sorted numeric index supporting fast tolerance-window lookups.
+
+    Sorting the reference values once lets each query binary-search its
+    tolerance window in O(log N + k) instead of scanning all N values, which
+    pays off when many queries are matched against the same reference set.
+
+    The window shape is delegated to ``bounds_func``, so the same index serves
+    m/z matching (a parts-per-million window via :func:`~tools.utils.get_ppm_range`)
+    and retention-time matching (a symmetric absolute window) alike.
+    """
+
+    def __init__(
+        self,
+        values: np.ndarray,
+        bounds_func: Callable[[float, float], Tuple[float, float]] = symmetric_window,
+    ):
+        """
+        Build the index from a set of reference values.
+
+        Parameters
+        ----------
+        values : np.ndarray
+            Reference values (e.g. m/z or retention times). NaNs sort to the
+            end and so never fall inside a finite window.
+        bounds_func : Callable
+            A function that maps a query to its tolerance window. Called as
+            ``bounds_func(value, tolerance)`` and must return the
+            ``(lower_bound, upper_bound)`` of the (inclusive) window. Defaults
+            to a symmetric absolute window.
+        """
+        self._bounds_func = bounds_func
+        values = np.asarray(values, dtype=float)
+        self._order = np.argsort(values)
+        self._sorted_values = values[self._order]
+
+    def search(self, value: float, tolerance: float) -> np.ndarray:
+        """
+        Return the positions of reference values within ``tolerance`` of ``value``.
+
+        Parameters
+        ----------
+        value : float
+            Query value.
+        tolerance : float
+            Tolerance defining the (inclusive) window; its meaning is set by
+            ``bounds_func`` (e.g. parts-per-million for a ppm window, or the
+            absolute half-width for a symmetric window).
+
+        Returns
+        -------
+        np.ndarray
+            Integer positions into the original (unsorted) reference array,
+            in ascending order.
+        """
+        lower_bound, upper_bound = self._bounds_func(value, tolerance)
+        lo = np.searchsorted(self._sorted_values, lower_bound, side="left")
+        hi = np.searchsorted(self._sorted_values, upper_bound, side="right")
+        return np.sort(self._order[lo:hi])