Numba apply mapping to ndarray

user_437296 0
Viewed 74

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

numba
1 Answers

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.

user_129395570
Related