Repository: lbl-anp/becquerel

Issue:
Spectrum `find_bin_index` should be able to accept multiple x-values
Spectrum `find_bin_index` should be able to accept multiple x-values and return their bin indices. This should be pretty simple to implement.

Discussion:

Patch:
diff --git a/becquerel/core/spectrum.py b/becquerel/core/spectrum.py
index b638b9b..88e9398 100644
--- a/becquerel/core/spectrum.py
+++ b/becquerel/core/spectrum.py
@@ -948,7 +948,7 @@ class Spectrum(object):
         return True
 
     def find_bin_index(self, x, use_kev=True):
-        """Find the Spectrum bin index that contains an x-axis value x.
+        """Find the Spectrum bin index or indices containing x-axis value(s) x.
 
         If the Spectrum has uniform binning, then we just solve the linear
         equation between bin indices and x-axis values. Otherwise, we use
@@ -956,15 +956,15 @@ class Spectrum(object):
         list of bin edges, then subtract 1 to get the index of the low edge.
 
         Args:
-          x: value whose bin to find
+          x: value(s) whose bin to find
           use_kev: check bin_edges_kev if True, otherwise bin_edges_raw
 
         Raises:
-          AssertionError: if x is outside the bin edges or equal to up edge
-          SpectrumError: if use_kev=True but Spectrum is not calibrated
+          SpectrumError: if use_kev=True but Spectrum is not calibrated; or if
+            x is outside the bin edges or equal to up edge
 
         Returns:
-          The integer bin index that contains x
+          The integer bin index or indices containing x
         """
 
         if use_kev and not self.is_calibrated:
@@ -973,12 +973,15 @@ class Spectrum(object):
 
         bin_edges = self.bin_edges_kev if use_kev else self.bin_edges_raw
         bin_widths = self.bin_widths_kev if use_kev else self.bin_widths_raw
+        x = np.asarray(x)
 
-        assert x >= bin_edges[0]
-        assert x < bin_edges[-1]
+        if np.any(x < bin_edges[0]):
+            raise SpectrumError('requested x is < lowest bin edge')
+        if np.any(x >= bin_edges[-1]):
+            raise SpectrumError('requested x is >= highest bin edge')
 
         if self.has_uniform_bins(use_kev=use_kev):
-            return int((x - bin_edges[0]) / bin_widths[0])
+            return ((x - bin_edges[0]) / bin_widths[0]).astype(np.int)
         else:
             return np.searchsorted(bin_edges, x, 'right') - 1
 
diff --git a/examples/spectrum.ipynb b/examples/spectrum.ipynb
index 268015a..b559f82 100644
--- a/examples/spectrum.ipynb
+++ b/examples/spectrum.ipynb
@@ -11,22 +11,13 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "becquerel/tools/xcom.py:13: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working\n",
-      "  from collections import Iterable\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Reload\n",
     "%load_ext autoreload\n",
     "%autoreload 2\n",
     "\n",
-    "from becquerel import Spectrum\n",
+    "from becquerel import Spectrum, SpectrumError\n",
     "import numpy as np\n",
     "np.random.seed(0)"
    ]
@@ -442,7 +433,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "becquerel/core/spectrum.py:724: SpectrumWarning: Subtraction of counts-based specta, spectra have been converted to CPS\n",
+      "becquerel/core/spectrum.py:726: SpectrumWarning: Subtraction of counts-based specta, spectra have been converted to CPS\n",
       "  'have been converted to CPS', SpectrumWarning)\n"
      ]
     },
@@ -656,8 +647,7 @@
     "plt.plot(spec7.bin_edges_raw[:-1], spec7.counts_vals)\n",
     "plt.show()\n",
     "# Note that the Spectrum.plot() method fails here because it works in bin mode.\n",
-    "# There might be a kwargs solution to that problem, but it should be fixed later.\n",
-    "# Note also that we still call bin_edges_kev, because bin_edges_adc has not yet been implemented."
+    "# There might be a kwargs solution to that problem, but it should be fixed later."
    ]
   },
   {
@@ -727,7 +717,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Find the bins corresponding to a particular channel"
+    "Find the bins corresponding to a particular channel(s)"
    ]
   },
   {
@@ -741,14 +731,18 @@
      "text": [
       "399\n",
       "0\n",
-      "999\n"
+      "999\n",
+      "[400 500 600]\n",
+      "[  1  58 105]\n"
      ]
     }
    ],
    "source": [
     "print(spec7.find_bin_index(3996, use_kev=False))\n",
     "print(spec7.find_bin_index(0, use_kev=False))\n",
-    "print(spec7.find_bin_index(9999, use_kev=False))"
+    "print(spec7.find_bin_index(9999, use_kev=False))\n",
+    "print(spec7.find_bin_index(np.array([4000, 5000, 6000]), use_kev=False))\n",
+    "print(spec8.find_bin_index(np.array([4000, 5000, 6000]), use_kev=False))"
    ]
   },
   {
@@ -768,7 +762,7 @@
     "# test that the bin up edge is out-of-range\n",
     "try:\n",
     "    print(spec7.find_bin_index(10000, use_kev=False))\n",
-    "except AssertionError:\n",
+    "except SpectrumError:\n",
     "    print(\"Failed successfully!\")"
    ]
   },


Test patch:
diff --git a/tests/spectrum_test.py b/tests/spectrum_test.py
index aee9e98..aafe678 100644
--- a/tests/spectrum_test.py
+++ b/tests/spectrum_test.py
@@ -253,9 +253,13 @@ def test_listmode():
     assert spec0.find_bin_index(XMIN, use_kev=False) == 0
     assert spec0.find_bin_index(XMIN + BW/4.0, use_kev=False) == 0
     assert spec0.find_bin_index(XMAX - BW/4.0, use_kev=False) == NBINS - 1
-    with pytest.raises(AssertionError):
+    assert np.all(spec0.find_bin_index(
+        spec0.bin_edges_raw[:-1], use_kev=False) == np.arange(NBINS))
+
+    # raise errors
+    with pytest.raises(bq.SpectrumError):
         spec0.find_bin_index(XMAX, use_kev=False)
-    with pytest.raises(AssertionError):
+    with pytest.raises(bq.SpectrumError):
         spec0.find_bin_index(XMIN - BW/4.0, use_kev=False)
 
     # test without args
@@ -271,6 +275,14 @@ def test_listmode():
     assert spec2.find_bin_index(1e1, use_kev=False) == 0
     assert spec2.find_bin_index(1e1 + 1e-9, use_kev=False) == 0
     assert spec2.find_bin_index(1e4 - 1e-9, use_kev=False) == NBINS - 1
+    assert np.all(spec2.find_bin_index(
+        spec2.bin_edges_raw[:-1], use_kev=False) == np.arange(NBINS))
+
+    # additional type checking
+    assert isinstance(
+        spec0.find_bin_index(XMIN, use_kev=False), (int, np.integer))
+    assert isinstance(
+        spec0.find_bin_index([XMIN], use_kev=False), np.ndarray)
 
 
 # ----------------------------------------------