pyqtgraph/pyqtgraph/functions_numba.py
pijyoi 4d6a8e4998
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_...
2021-04-28 22:29:09 -07:00

23 lines
727 B
Python

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