Very simple question, which I have looked for but haven't found a clear answer. I would like to efficiently apply a mapping to an input ndarray, remapping each element in the array and returning the modified array.
A simple version using numpy would be something like:
def remap(a, mapping):
return np.vectorize(mapping.__getitem__)(a)
I am pretty sure I can't use Numba vectorize
but was hoping to be able to use guvectorize
. But it looks like I can't pass a Numba TypedDict into a guvectorize
function either. Any suggestions welcome
If you know the number of dimentions of the target array, then you just need a simple loop. AFAIK, guvectorize
does not support this use-case yet. However, you can solve this using a simple reshape.
import numba as nb
@nb.njit
def remap(a, mapping):
# Linearize the array (can cause a copy of `a`)
b = a.reshape(-1)
for i in range(b.size):
b[i] = mapping[b[i]]
return b.reshape(a.shape)
Note that the first reshape can cause a copy of a
if a
is not contiguously store in memory and Numpy cannot find a unified stride in 1D. Thus, the returned array may be a copy of a
. The second reshape is guaranteed not to copy the array. If you want to always return a copy of the input array (no mutation) like most Numpy function, then you can use flatten
instead of the first reshape
call.