From e8ae8755aa5e9f5085bbe83ecb9a3ff8b96752c0 Mon Sep 17 00:00:00 2001
From: Sebastian Achim Mueller <sebastian-achim.mueller@mpi-hd.mpg.de>
Date: Mon, 4 Mar 2024 00:29:48 +0100
Subject: [PATCH] new feature, query_ball

---
 binning_utils/__init__.py | 28 ++++++++++++++++++++++++++++
 binning_utils/version.py  |  2 +-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/binning_utils/__init__.py b/binning_utils/__init__.py
index a868ef9..cafa5b3 100644
--- a/binning_utils/__init__.py
+++ b/binning_utils/__init__.py
@@ -243,3 +243,31 @@ def edges_from_width_and_num(bin_width, num_bins, first_bin_center):
         num=num_bins + 1,
     )
     return bin_edges
+
+
+def query_ball(bin_edges, start, stop):
+    """
+    Returns the indices of the bins which are touching the range from
+    ``start`` to ``stop``.
+
+    Parameters
+    ----------
+    bin_edges : array_like (N + 1) floats
+        Edges of N bins.
+    start : float
+        Start of the range.
+    stop : float
+        Stop of the range.
+
+    Returns
+    -------
+    bin_indices : array_like ints
+        The indices touching the range from start to stop.
+    """
+    assert start <= stop
+    num_bins = len(bin_edges) - 1
+    bin_start = np.digitize(x=start, bins=bin_edges) - 1
+    bin_stop = np.digitize(x=stop, bins=bin_edges) - 1
+    ee = np.arange(bin_start, bin_stop + 1, 1)
+    mask = np.logical_and(ee >= 0, ee < num_bins)
+    return ee[mask]
diff --git a/binning_utils/version.py b/binning_utils/version.py
index 6e2648a..4ae81f3 100644
--- a/binning_utils/version.py
+++ b/binning_utils/version.py
@@ -1 +1 @@
-__version__ = "0.0.12"
+__version__ = "0.0.13"