提高高度重复的 numpy 数组索引操作的性能答案

【问题标题】：Increasing performance of highly repeated numpy array index operations提高高度重复的 numpy 数组索引操作的性能
【发布时间】：2018-02-16 08:51:43
【问题描述】：

在我的程序代码中，我有numpy 值数组和numpy 索引数组。两者都是在程序初始化期间预先分配和预定义的。
程序的每一部分都有一个数组values 用于执行计算，以及三个索引数组idx_from_exch、idx_values 和idx_to_exch。有一个全局值数组来交换几个部分的值：exch_arr.
大多数时候，索引数组有 2 到 5 个索引，很少（很可能永远不会）需要更多索引。 dtype=np.int32、shape 和值在整个程序运行期间保持不变。因此我在初始化后设置了ndarray.flags.writeable=False，但这是可选的。索引数组idx_values 和idx_to_exch 的索引值是按数字顺序排序的，idx_source 可以排序，但是没有办法定义它。一个值数组/部分对应的所有索引数组都具有相同的shape。
values 数组和 exch_arr 通常有 50 到 1000 个元素。 shape 和 dtype=np.float64 在整个程序运行期间是不变的，数组的值在每次迭代中都会改变。
以下是示例数组：

import numpy as np
import numba as nb

values = np.random.rand(100) * 100  # just some random numbers
exch_arr = np.random.rand(60) * 3  # just some random numbers
idx_values = np.array((0, 4, 55, -1), dtype=np.int32)  # sorted but varying steps
idx_to_exch = np.array((7, 8, 9, 10), dtype=np.int32)  # sorted and constant steps!
idx_from_exch = np.array((19, 4, 7, 43), dtype=np.int32)  # not sorted and varying steps

示例索引操作如下所示：

values[idx_values] = exch_arr[idx_from_exch]  # get values from exchange array
values *= 1.1  # some inplace array operations, this is just a dummy for more complex things
exch_arr[idx_to_exch] = values[idx_values]  # pass some values back to exchange array

由于这些操作每次迭代都会应用一次，持续数百万次迭代，因此速度至关重要。我一直在研究提高索引速度in my previous question 的许多不同方法，但考虑到我的应用程序（尤其是通过使用常量索引数组进行索引并将它们传递给另一个索引数组来获取值），我忘记了足够具体。
到目前为止，最好的方法似乎是花哨的索引。我目前也在尝试numbaguvectorize，但似乎不值得努力，因为我的数组很小。 memoryviews 会很好，但是由于索引数组不一定有一致的步骤，我知道没有办法使用memoryviews。

那么有没有更快的方法来进行重复索引？为每个索引操作预定义内存地址数组的某种方式，因为dtype 和shape 总是不变的？ ndarray.__array_interface__ 给了我一个内存地址，但我无法将它用于索引。我想到了类似的东西：

stride_exch = exch_arr.strides[0]
mem_address = exch_arr.__array_interface__['data'][0]
idx_to_exch = idx_to_exch * stride_exch + mem_address

这样可行吗？
我也一直在研究直接将strides 与as_strided 一起使用，但据我所知，只允许一致的步幅，我的问题需要不一致的strides。

任何帮助表示赞赏！提前致谢！

编辑：
我刚刚纠正了示例计算中的一个巨大错误！
操作values = values * 1.1改变了数组的内存地址。我在程序代码中的所有操作都布置为不更改数组的内存地址，因为许多其他操作都依赖于使用内存视图。因此，我用正确的就地操作替换了虚拟操作：values *= 1.1

【问题讨论】：

可以重新排序values 和exch_arr 吗？切片不会复制内存，因此如果您可以构建数据以便使用连续块，它会快得多。
不幸的是没有。有太多相互关联的部分都有自己的value 数组和索引数组，exch_arr 将它们连接起来。对于中间有很多连接的树状网络拓扑，排序是不可能的。 value 数组代表物理属性，因此也无法重新排序。至少存在某种重复/有序结构的所有其他数组已经转换为内存视图，但在提出的问题中，这很不幸是不可能的。
您有没有找到加快索引速度的方法？
@MrCheatak 不，我不得不接受使用花哨的索引。但我猜想使用 numba 的索引操作在过去三年中得到了改进。我目前看到的唯一解决方法是在列表/元组中存储单单元内存视图，例如values[4:5]，并在 jit 编译的函数中迭代列表/元组。我想这应该比重复花式索引要快得多。但你必须决定，是否值得付出努力......

标签： python arrays performance numpy memoryview

【解决方案1】：

使用 numpy 布尔数组绕过昂贵的花式索引的一种解决方案是使用 numba 并跳过 numpy 布尔数组中的 False 值。

示例实现：

@numba.guvectorize(['float64[:], float64[:,:], float64[:]'], '(n),(m,n)->(m)', nopython=True, target="cpu")
def test_func(arr1, arr2, inds, res):
    for i in range(arr1.shape[0]):
        if not inds[i]:
            continue
        for j in range(arr2.shape[0]):
            res[j, i] = arr1[i] + arr2[j, i]

当然，使用 numpy 数据类型（较小的字节大小会运行得更快）和目标是 "cpu" 或 "parallel"。

【讨论】：