在 numpy 中创建索引数组 - 消除双循环

Question

我有一些物理模拟代码，用 python 编写并使用 numpy/scipy。分析代码显示 38% 的 CPU 时间花费在单个双重嵌套 for 循环中 - 这似乎过多，所以我一直在努力减少它。

循环的目标是创建一个索引数组，显示一维数组的哪些元素与二维数组的元素相等。

indices[i,j] = where(1D_array == 2D_array[i,j])

例如，如果 1D_array = [7.2, 2.5, 3.9] 和

2D_array = [[7.2, 2.5] 
            [3.9, 7.2]]

我们应该

indices = [[0, 1]
           [2, 0]]

我目前已将其实施为

for i in range(ni):
    for j in range(nj):
        out[i, j] = (1D_array - 2D_array[i, j]).argmin()

我处理浮点数时需要 argmin，因此等式不一定准确。我知道一维数组中的每个数字都是唯一的，而二维数组中的每个元素都有一个匹配项，所以这种方法给出了正确的结果。

有什么办法可以消除双for循环吗？

注:

我需要索引数组来执行以下操作：

f = complex_function(1D_array)
output = f[indices]

这比备选方案更快，因为二维数组的大小为 NxN，而一维数组的大小为 1xN，而且二维数组有很多重复值。如果有人可以提出一种不通过索引数组即可获得相同输出的不同方法，那也可能是一个解决方案

Answer 1

在纯 Python 中，您可以在 O(N) 时间内使用字典完成此操作，唯一的时间损失将是涉及的 Python 循环：

>>> arr1 = np.array([7.2, 2.5, 3.9])
>>> arr2 = np.array([[7.2, 2.5], [3.9, 7.2]])
>>> indices = dict(np.hstack((arr1[:, None], np.arange(3)[:, None])))
>>> np.fromiter((indices[item] for item in arr2.ravel()), dtype=arr2.dtype).reshape(arr2.shape)
array([[ 0.,  1.],
       [ 2.,  0.]])

Answer 2

要摆脱两个 Python for 循环，您可以通过向数组添加新轴（使它们可以相互广播）来进行所有相等比较 "in one go" ).

请记住，这会生成一个包含 len(arr1)*len(arr2) 值的新数组。如果这是一个非常大的数字，根据您的内存限制，这种方法可能不可行。否则，它应该相当快：

>>> (arr1[:,np.newaxis] == arr2[:,np.newaxis]).argmax(axis=1)
array([[0, 1],
       [2, 0]], dtype=int32)

如果您需要获取 arr1 中 最接近的 匹配值的索引，请使用：

np.abs(arr1[:,np.newaxis] - arr2[:,np.newaxis]).argmin(axis=1)

Answer 3

其他一些人建议的字典方法可能有效，但它需要您提前知道目标数组（二维数组）中的每个元素在您的搜索数组（您的一维数组）中都有完全匹配.即使这在原则上应该是正确的，你仍然必须处理浮点精度问题，例如试试这个 .1 * 3 == .3.

另一种方法是使用 numpy 的 searchsorted 函数。 searchsorted 采用排序的一维搜索数组，然后任何目标数组都会为目标数组中的每个项目找到搜索数组中最接近的元素。我已根据您的情况对 answer 进行了调整，请查看它以了解 find_closest 函数的工作原理。

import numpy as np

def find_closest(A, target):
    order = A.argsort()
    A = A[order]

    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return order[idx]

array1d = np.array([7.2, 2.5, 3.9])
array2d = np.array([[7.2, 2.5],
                    [3.9, 7.2]])

indices = find_closest(array1d, array2d)
print(indices)
# [[0 1]
#  [2 0]]

在 numpy 中创建索引数组 - 消除双循环

Creating index array in numpy - eliminating double for loop

python

arrays

indexing

performance

numpy