python - 使用 numba 中的数组索引数组 (njit) : variable-dimension ndarrays

Question

我有一个包含 D 元素的一维整数数组(即 idx = np.array([i0, i1, ...])，s.t. idx.size = D)，其中每个元素对应于具有 D 维度的 ND 数组的该维度的索引(即 data s.t. data.ndim = D)。如何使用索引数组 idx 索引 data 数组？

---

在 python 中，我会做 data[tuple(idx)]，但是 tuple 在 numba nopython 模式下不受支持。

我目前的解决方法是使用 data.ravel() 并将 ND 索引转换为展平数组的一维索引，但似乎必须有一个更简单(并且计算速度更快)的解决方案。某处是否有take_along_each_axis(data, idx) 方法？

Answer 1

让我们做一些时间测试:

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

这将近 1 Gb 的内存 - 不足以产生内存错误，但仍然应该是一个合理的测试。实际上，对于更小的数组，我得到了相同的基本索引时间。对于其他 idx 值

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

解释器将其转换为方法调用:

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

索引“平面”数组，可以通过以下方式完成:

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)

或将转换从循环中取出:

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)

我认为有些情况下平面索引速度更快，但这不是。

因此，对于一个独立的操作，我看不出编写一个 njit 版本有什么意义。我想如果它是某个更大操作的一部分，那是值得的。