indexing arrays with arrays in numba (njit): variable-dimension ndarrays

Question

I have a 1D array of integers with D elements (i.e. idx = np.array([i0, i1, ...]), s.t. idx.size = D), where each element corresponds to the index along that dimension of an ND array with D dimensions (i.e. data s.t. data.ndim = D). How can I index the data array using the index array idx?

---

In python I would do data[tuple(idx)], but tuple aren't supported in numba nopython mode.

My current workaround is to use data.ravel() and convert from ND indices to 1D indices of the flattened array, but it seems like there must be an easier (and computationally faster) solution. Is there a take_along_each_axis(data, idx) method somewhere?

Answer 1

Lets do a bit of time testing:

In [135]: data = np.ones((100,100,100,100)); idx = (50,50,50,50)

That's nearly a Gb of memory - not huge enough to create a memory error, but still should be a reasonable test. Actually, I get the same time for basic indexing for much smaller arrays. And for other idx values

In [136]: timeit data[idx]
212 ns ± 9.25 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

the interpreter translates that into a method call:

In [137]: timeit data.__getitem__(idx)
283 ns ± 4.37 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

indexing the 'flat' array, can be done with:

In [138]: timeit data.flat[np.ravel_multi_index(idx,data.shape)]
6.65 µs ± 75.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

or taking the conversion out of the loop:

In [139]: %%timeit x=np.ravel_multi_index(idx,data.shape)
     ...: data.flat[x] 
574 ns ± 23.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [142]: %%timeit x=np.ravel_multi_index(idx,data.shape);df=data.flat
     ...: df[x]
345 ns ± 6.39 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

I think there are cases where flat indexing is faster, but this isn't one.

So a stand alone operation I don't see the point to writing a njit version. I suppose if it's part of some larger operation it could be worth it.