implement numba lookup

pyqtgraph/functions_numba.py

@@ -20,3 +20,23 @@ def rescaleData(data, scale, offset, dtype, clip):
         rescale_functions[key] = func
     func(data, scale, offset, clip[0], clip[1], out=data_out)
     return data_out
+
+@numba.jit(nopython=True)
+def rescale_and_lookup1d_function(xx, scale, offset, vmin, vmax, lut, yy):
+    for r in range(xx.shape[0]):
+        for c in range(xx.shape[1]):
+            val = (xx[r, c] - offset) * scale
+            val = min(max(val, vmin), vmax)
+            yy[r, c] = lut[int(val)]
+
+def rescale_and_lookup1d(data, scale, offset, clip, lut):
+    # data should be floating point and 2d
+    # lut is 1d
+    data_out = np.empty_like(data, dtype=lut.dtype)
+    rescale_and_lookup1d_function(data, float(scale), float(offset), float(clip[0]), float(clip[1]), lut, data_out)
+    return data_out
+
+@numba.jit(nopython=True)
+def numba_take(lut, data):
+    # numba supports only the 1st two arguments of np.take
+    return np.take(lut, data)

pyqtgraph/graphicsItems/ImageItem.py

@@ -512,6 +512,15 @@ class ImageItem(GraphicsObject):
             maxVal = xp.nextafter(maxVal, 2*maxVal)
         rng = maxVal - minVal
         rng = 1 if rng == 0 else rng
+
+        fn_numba = fn.getNumbaFunctions()
+        if xp == numpy and image.flags.c_contiguous and dtype == xp.uint16 and fn_numba is not None:
+            lut, augmented_alpha = self._convert_2dlut_to_1dlut(lut)
+            image = fn_numba.rescale_and_lookup1d(image, scale/rng, minVal, (0, num_colors-1), lut)
+            if image.dtype == xp.uint32:
+                image = image[..., xp.newaxis].view(xp.uint8)
+            return image, None, None, augmented_alpha
+        else:
             image = fn.rescaleData(image, scale/rng, offset=minVal, dtype=dtype, clip=(0, num_colors-1))
 
             levels = None
@@ -631,21 +640,39 @@ class ImageItem(GraphicsObject):
         # if we are contiguous, we can take a faster codepath where we
         # ensure that the lut is 1d
 
-        if lut.ndim == 2:
+        lut, augmented_alpha = self._convert_2dlut_to_1dlut(lut)
+
+        fn_numba = fn.getNumbaFunctions()
+        if xp == numpy and fn_numba is not None:
+            image = fn_numba.numba_take(lut, image)
+        else:
+            image = lut[image]
+
+        if image.dtype == xp.uint32:
+            image = image[..., xp.newaxis].view(xp.uint8)
+
+        return image, augmented_alpha
+
+    def _convert_2dlut_to_1dlut(self, lut):
+        # converts:
+        #   - uint8 (N, 1) to uint8 (N,)
+        #   - uint8 (N, 3) or (N, 4) to uint32 (N,)
+        # this allows faster lookup as 1d lookup is faster
+        xp = self._xp
+        augmented_alpha = False
+
+        if lut.ndim == 1:
+            return lut, augmented_alpha
+
         if lut.shape[1] == 3:   # rgb
             # convert rgb lut to rgba so that it is 32-bits
             lut = xp.column_stack([lut, xp.full(lut.shape[0], 255, dtype=xp.uint8)])
             augmented_alpha = True
         if lut.shape[1] == 4:   # rgba
             lut = lut.view(xp.uint32)
+        lut = lut.ravel()
 
-        image = lut.ravel()[image]
-        lut = None
-        # now both levels and lut are None
-        if image.dtype == xp.uint32:
-            image = image.view(xp.uint8).reshape(image.shape + (4,))
-
-        return image, augmented_alpha
+        return lut, augmented_alpha
 
     def _try_make_qimage(self, image, levels, lut, augmented_alpha):
         xp = self._xp