查找每个唯一子数组出现的位置

Question

情况：

我正在填充一个形状为 (2N, 2N) 的 narray，其中 N 接近 8000，将其称为 A，其值是我通过使用嵌套的 for 循环从函数中获取的值来调用一个函数，该函数将形状 (2,) 来自形状 (N,N,2) 数组的最后一维，称为 B。

这显然代价高昂，虽然我没有成功对这部分代码进行向量化（也非常欢迎这方面的任何帮助），但我知道 B 在最后一维中有许多重复的子数组。所以，我想要的是找出唯一的子数组以及每个子数组出现的位置。然后通过迭代每个唯一的子数组并用函数返回的值填充它出现的所有位置来加速 A 的填充，该值将只计算一次。

我所做的如下，但它似乎既不是最直接的方式，也不是最 Numpy 的方式。

我一直在使用填充矩阵的代码如下：

translat_avg_disloc_matrix = np.zeros([2*n, 2*n])

for i in range(n):
    for alpha in range(2):

        for j in range(n):
            for beta in range(2):

                    translat_avg_disloc_matrix[2*i+alpha,2*j+beta] = average_lat_pos(alpha, beta, b_matrix[i][j])

虽然我可以通过执行类似于此处所做的操作来找到唯一子数组：Efficiently count the number of occurrences of unique subarrays in NumPy?)，但我在查找每个子数组出现的索引时遇到了问题。

我尝试过的是这样的：

1) 通过norm = (B*B).sum(axis=2)计算B最后一维子数组的范数，计算B-1最后一维子数组的范数

norm_ = ((B-1)*(B-1)).sum(axis=2)

2) 使用norm.reshape((norm.size,1)) 为这两个规范重塑这些 narrays

3) 创建瓦片矩阵为tile_norm = np.tile(norm.T, (len(norm),1))

4) 然后执行np.unique(np.non_zero(np.abs(tile_norm - norm)+np.abs(tile_norm_-norm_) == 0), axis=0)，这给了我们类似：array([[0, 0, 0, 4], [4, 4, 4, 0]]) 其中每行中的零表示这些索引对应于 B 矩阵中的相同 (2,) 向量。

换句话说，我发现 (2,) 数组的范数按原样一致，并且当从它们中减去 1 时也一致 - 两个方程，两个变量。

我正在寻找的是一种查找每个唯一子数组在 B 中出现位置的方法，这样使用一些花哨的索引将允许我填充矩阵而无需重复调用该函数average_lat_pos（这里重复意味着调用相同的（alpha，beta，（2，）数组）有序对）。

Answer 1

一个简单的技巧是将average_lat_pos 参数存储在字典中。因此，如果键在字典中，则返回它的值，而无需为重复参数计算 average_lat_pos。 由于average_lat_pos 函数不可用，我使用一个虚拟函数作为：

def average_lat_pos(alpha, betha, td):
    result = 0
    for i in range(1000):
        result += 1
    return result

以及最终代码：

import numpy as np

def average_lat_pos_1(alpha, betha, td):
    get_result = dictionary.get((alpha, betha, td[0], td[1]), 'N')
    if get_result == 'N':
        result = 0
        for i in range(1000):
            result += 1
        dictionary[(alpha, betha, td[0], td[1])] = result
        return result
    else:
        return get_result

def average_lat_pos_2(alpha, betha, td):
    result = 0
    for i in range(1000):
        result += 1
    return result


N = 500
B = np.random.rand(N*N*2) * 100
B = np.floor(B)
B = B.reshape(N,N,2)

dictionary = dict()
A = np.empty([2*N, 2*N])
for i in range(2*N):
    for j in range(2*N):
        A[i, j] = average_lat_pos_1(i%2, j%2, B[i//2, j//2])

对于 N = 500，average_lat_pos_1 的执行速度大约比 average_lat_pos_2 快 X10

Answer 2

矢量化通常更好。稍微调整一下你的功能有助于这项工作：

import numpy as np
def average_lat_pos(a,b,x,y):  # all arguments are scalars
    return a*x+2*b*y           # as example

n=1000    
B=np.random.rand(n,n,2)


def loops():
    A=np.empty((2*n,2*n))
    for i in range(n):
        for alpha in range(2):

          for j in range(n):
            for beta in range(2):

              A[2*i+alpha,2*j+beta] = average_lat_pos(alpha, beta, 
              B[i,j,0],B[i,j,1])
    return A

要矢量化，只需在相应维度上变换循环级别：

Z=average_lat_pos(np.arange(2).reshape(1,2,1,1),np.arange(2).reshape(1,1,1,2),
B[:,:,0].reshape(n,1,n,1),B[:,:,1].reshape(n,1,n,1)).reshape(2*n,2*n)

测试：

np.allclose(loops(),Z)
Out[105]: True

%timeit loops()
3.73 s ± 9.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit average_lat_pos(np.arange(2).reshape(1,2,1,1),np.arange(2).reshape(1,1,1,2),
   B[:,:,0].reshape(n,1,n,1),B[:,:,1].reshape(n,1,n,1)).reshape(2*n,2*n)
38.7 ms ± 211 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.unique(B,axis=2)
1.31 s ± 6.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

np.unique 意味着 O(n ln n) 操作，这将扼杀任何潜在的收益。

但这里更好，使用 numba，装饰两个函数：

@numba.njit
def average_lat_pos(a,b,x,y):
   ....

@numba.njit
def loops(B):
   ....

然后

%timeit loops(B)
3.04 ms ± 34.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


it's now 1000x faster.