add usage of numba (for rescale) (#1695)

* implement rescale using numba

* workaround to pass test_makeARGB.py

* key on (input, output) manually

* remove minimum version check

* signature needs to be a list of signatures

numba considers it a mistake for single-item non-list but works around
it internally

* add numba test to test_makeARGB.py

* add numba as dependency in CI

* handle properly the case where useNumba was already True

* restore useCupy setting after test

* filter numba nan warning

* don't make changes to test_cupy_makeARGB_...

.github/workflows/main.yml

@@ -70,7 +70,7 @@ jobs:
     - name: Install Dependencies
       run: |
         pip install --upgrade pip
-        pip install ${{ matrix.qt-version }} numpy${{ matrix.numpy-version }} scipy pyopengl h5py matplotlib
+        pip install ${{ matrix.qt-version }} numpy${{ matrix.numpy-version }} scipy pyopengl h5py matplotlib numba
         pip install .
         pip install pytest pytest-cov pytest-xdist coverage
     - name: "Install Linux VirtualDisplay"

pyqtgraph/__init__.py

@@ -56,6 +56,7 @@ CONFIG_OPTIONS = {
                                  # The default is 'col-major' for backward compatibility, but this may
                                  # change in the future.
     'useCupy': False,  # When True, attempt to use cupy ( currently only with ImageItem and related functions )
+    'useNumba': False, # When True, use numba
 } 
 
 

pyqtgraph/functions.py

@@ -17,6 +17,7 @@ import math
 
 import numpy as np
 from .util.cupy_helper import getCupy
+from .util.numba_helper import getNumbaFunctions
 
 from . import debug, reload
 from .Qt import QtGui, QtCore, QT_LIB, QtVersion
@@ -1115,7 +1116,12 @@ def rescaleData(data, scale, offset, dtype=None, clip=None):
 
         # don't copy if no change in dtype
         return data_out.astype(out_dtype, copy=False)
-    else:
+
+    numba_fn = getNumbaFunctions()
+    if numba_fn and clip is not None:
+        # if we got here by makeARGB(), clip will not be None at this point
+        return numba_fn.rescaleData(data, scale, offset, out_dtype, clip)
+
     return _rescaleData_nditer(data, scale, offset, work_dtype, out_dtype, clip)
 
 

pyqtgraph/functions_numba.py

@@ -0,0 +1,22 @@
+import numpy as np
+import numba
+
+rescale_functions = {}
+
+def rescale_clip_source(xx, scale, offset, vmin, vmax, yy):
+    for i in range(xx.size):
+        val = (xx[i] - offset) * scale
+        yy[i] = min(max(val, vmin), vmax)
+
+def rescaleData(data, scale, offset, dtype, clip):
+    data_out = np.empty_like(data, dtype=dtype)
+    key = (data.dtype.name, data_out.dtype.name)
+    func = rescale_functions.get(key)
+    if func is None:
+        func = numba.guvectorize(
+            [f'{key[0]}[:],f8,f8,f8,f8,{key[1]}[:]'],
+            '(n),(),(),(),()->(n)',
+            nopython=True)(rescale_clip_source)
+        rescale_functions[key] = func
+    func(data, scale, offset, clip[0], clip[1], out=data_out)
+    return data_out

pyqtgraph/tests/test_makeARGB.py

@@ -4,7 +4,7 @@ import pytest
 import sys
 from typing import Dict, Any, Union, Type
 
-from pyqtgraph import getCupy, setConfigOption
+from pyqtgraph import getCupy, getConfigOption, setConfigOption
 from pyqtgraph.functions import makeARGB as real_makeARGB
 
 IN_2D_INT8 = np.array([[173, 48, 122, 41], [210, 192, 0, 5], [104, 56, 102, 115], [78, 19, 255, 6]], dtype=np.uint8)
@@ -4316,6 +4316,25 @@ def test_cupy_makeARGB_against_generated_references():
     _do_something_for_every_combo(assert_cupy_correct)
 
 
+@pytest.mark.filterwarnings("ignore:invalid value encountered")
+def test_numba_makeARGB_against_generated_references():
+    oldcfg_numba = getConfigOption("useNumba")
+    if not oldcfg_numba:
+        try:
+            import numba
+        except ImportError:
+            pytest.skip("Numba unavailable to test")
+
+    # useCupy needs to be set to False because it takes
+    # precedence over useNumba in rescaleData
+    oldcfg_cupy = getConfigOption("useCupy")
+    setConfigOption("useCupy", False)
+    setConfigOption("useNumba", not oldcfg_numba)
+    test_makeARGB_against_generated_references()
+    setConfigOption("useNumba", oldcfg_numba)
+    setConfigOption("useCupy", oldcfg_cupy)
+
+
 def test_makeARGB_with_human_readable_code():
     # Many parameters to test here:
     #  * data dtype (ubyte, uint16, float, others)

pyqtgraph/util/numba_helper.py

@@ -0,0 +1,16 @@
+from warnings import warn
+
+from .. import getConfigOption
+
+def getNumbaFunctions():
+    if getConfigOption("useNumba"):
+        try:
+            import numba
+        except ImportError:
+            warn("numba library could not be loaded, but 'useNumba' is set.")
+            return None
+
+        from .. import functions_numba
+        return functions_numba
+    else:
+        return None