all tests work except test_hist_plot::test_plotgrid

Author: lgray

src/coffea/jitters/hist/hist_tools.py

@@ -487,9 +487,9 @@ class Bin(DenseAxis):
     def __init__(self, name, label, n_or_arr, lo=None, hi=None):
         super().__init__(name, label)
         self._lazy_intervals = None
-        if isinstance(n_or_arr, (list, numpy.ndarray)):
+        if isinstance(n_or_arr, (list, numpy.ndarray, cupy.ndarray)):
             self._uniform = False
-            self._bins = numpy.array(n_or_arr, dtype="d")
+            self._bins = cupy.array(n_or_arr, dtype="d")
             if not all(numpy.sort(self._bins) == self._bins):
                 raise ValueError("Binning not sorted!")
             self._lo = self._bins[0]
@@ -583,7 +583,10 @@ def index(self, identifier):
                     )
 
                 if isinstance(idx, (cupy.ndarray, numpy.ndarray)):
-                    _replace_nans(self.size - 1, idx)
+                    _replace_nans(
+                        self.size - 1,
+                        idx if idx.dtype.kind == "f" else idx.astype(cupy.float32),
+                    )
                     idx = idx.astype(int)
                 elif numpy.isnan(idx):
                     idx = self.size - 1
@@ -596,7 +599,13 @@ def index(self, identifier):
             if identifier.nan():
                 return self.size - 1
             for idx, interval in enumerate(self._intervals):
-                if interval._lo <= identifier._lo and interval._hi >= identifier._hi:
+                if (
+                    interval._lo <= identifier._lo
+                    or cupy.isclose(interval._lo, identifier._lo)
+                ) and (
+                    interval._hi >= identifier._hi
+                    or cupy.isclose(interval._hi, identifier._hi)
+                ):
                     return idx
             raise ValueError(
                 "Axis %r has no interval that fully contains identifier %r"
@@ -759,10 +768,10 @@ def edges(self, overflow="none"):
                 See `Hist.sum` description for the allowed values.
         """
         if self._uniform:
-            out = numpy.linspace(self._lo, self._hi, self._bins + 1)
+            out = cupy.linspace(self._lo, self._hi, self._bins + 1)
         else:
             out = self._bins[:-1].copy()
-        out = numpy.r_[
+        out = cupy.r_[
             2 * out[0] - out[1], out, 2 * out[-1] - out[-2], 3 * out[-1] - 2 * out[-2]
         ]
         return out[overflow_behavior(overflow)]

src/coffea/jitters/hist/plot.py

@@ -231,7 +231,7 @@ def plot1d(
     elif isinstance(axis, DenseAxis):
         ax.set_xlabel(axis.label)
         ax.set_ylabel(hist.label)
-        edges = axis.edges(overflow=overflow)
+        edges = axis.edges(overflow=overflow).get()
         if order is None:
             identifiers = (
                 hist.identifiers(overlay, overflow=overlay_overflow)
@@ -417,12 +417,14 @@ def plotratio(
     elif isinstance(axis, DenseAxis):
         ax.set_xlabel(axis.label)
         ax.set_ylabel(num.label)
-        edges = axis.edges(overflow=overflow)
-        centers = axis.centers(overflow=overflow)
+        edges = axis.edges(overflow=overflow).get()
+        centers = axis.centers(overflow=overflow).get()
         ranges = (edges[1:] - edges[:-1]) / 2 if xerr else None
 
         sumw_num, sumw2_num = num.values(sumw2=True, overflow=overflow)[()]
         sumw_denom, sumw2_denom = denom.values(sumw2=True, overflow=overflow)[()]
+        sumw_num, sumw2_num = sumw_num.get(), sumw2_num.get()
+        sumw_denom, sumw2_denom = sumw_denom.get(), sumw2_denom.get()
 
         rsumw = sumw_num / sumw_denom
         if unc == "clopper-pearson":
@@ -557,9 +559,10 @@ def plot2d(
     if isinstance(xaxis, SparseAxis) or isinstance(yaxis, SparseAxis):
         raise NotImplementedError("Plot a sparse axis (e.g. bar chart or labeled bins)")
     else:
-        xedges = xaxis.edges(overflow=xoverflow)
-        yedges = yaxis.edges(overflow=yoverflow)
+        xedges = xaxis.edges(overflow=xoverflow).get()
+        yedges = yaxis.edges(overflow=yoverflow).get()
         sumw, sumw2 = hist.values(sumw2=True, overflow="allnan")[()]
+        sumw, sumw2 = sumw.get(), sumw2.get()
         if transpose:
             sumw = sumw.T
             sumw2 = sumw2.T

tests/test_hist_plot.py

@@ -125,8 +125,8 @@ def test_plotratio():
 
     # Add some pseudodata to a pt histogram so we can make a nice data/mc plot
     pthist = lepton_kinematics.sum("eta")
-    bin_values = pthist.axis("pt").centers()
-    poisson_means = pthist.sum("flavor").values()[()]
+    bin_values = pthist.axis("pt").centers().get()
+    poisson_means = pthist.sum("flavor").values()[()].get()
     values = np.repeat(bin_values, np.random.poisson(poisson_means))
     pthist.fill(flavor="pseudodata", pt=values)
 

tests/test_hist_tools.py

@@ -283,7 +283,7 @@ def test_issue_247():
 def test_issue_333():
     axis = hist.Bin("channel", "Channel b1", 50, 0, 2000)
     temp = np.arange(0, 2000, 40, dtype=np.int16)
-    assert np.all(axis.index(temp) == np.arange(50) + 1)
+    assert np.all(axis.index(temp).get() == np.arange(50) + 1)
 
 
 def test_issue_394():