From 4d6a8e499860aa4f1099cf00599dc3cf3c09c4f0 Mon Sep 17 00:00:00 2001 From: pijyoi Date: Thu, 29 Apr 2021 13:29:09 +0800 Subject: [PATCH] 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 | 2 +- pyqtgraph/__init__.py | 1 + pyqtgraph/functions.py | 10 ++++++++-- pyqtgraph/functions_numba.py | 22 ++++++++++++++++++++++ pyqtgraph/tests/test_makeARGB.py | 21 ++++++++++++++++++++- pyqtgraph/util/numba_helper.py | 16 ++++++++++++++++ 6 files changed, 68 insertions(+), 4 deletions(-) create mode 100644 pyqtgraph/functions_numba.py create mode 100644 pyqtgraph/util/numba_helper.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ab4a2060..4782e8da 100644 --- a/.github/workflows/main.yml +++ b/.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" diff --git a/pyqtgraph/__init__.py b/pyqtgraph/__init__.py index bc96e720..916e937e 100644 --- a/pyqtgraph/__init__.py +++ b/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 } diff --git a/pyqtgraph/functions.py b/pyqtgraph/functions.py index 45043a4c..21f00bec 100644 --- a/pyqtgraph/functions.py +++ b/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,8 +1116,13 @@ 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: - return _rescaleData_nditer(data, scale, offset, work_dtype, out_dtype, clip) + + 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) def applyLookupTable(data, lut): diff --git a/pyqtgraph/functions_numba.py b/pyqtgraph/functions_numba.py new file mode 100644 index 00000000..d9e2f232 --- /dev/null +++ b/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 diff --git a/pyqtgraph/tests/test_makeARGB.py b/pyqtgraph/tests/test_makeARGB.py index d4b5d889..73b29f1b 100644 --- a/pyqtgraph/tests/test_makeARGB.py +++ b/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) diff --git a/pyqtgraph/util/numba_helper.py b/pyqtgraph/util/numba_helper.py new file mode 100644 index 00000000..6eb46680 --- /dev/null +++ b/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