Search code examples
numba

Numba apply mapping to ndarray


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


Solution

  • 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.