加快numpy中的循环

Question

我正在寻找加速numpy中的循环，但我发现没有明显的模式可以这样做：

for index1 in range(1, len_route):
        time_diff_matrix[index1, (index1+1):len_route] = \
            M[(index1-1):(len_route-2)] - \
            M[index1-1] + \
            N[index1-1, index1:(len_route-1)] + \
            N[index1, (index1+1):len_route] - \
            P[index1:(len_route-1)]

time_diff_matrix的其余部分填充零。 首先是双循环。 我摆脱了一个循环，但我不知道如何摆脱另一个循环。 len_route是一个大数字。

问候。

Answer 1

这是使用slicing行/列的方法 -

n = len_route
vals = M[:n-2] - M[:n-2,None] + N[:n-2,1:n-1] + N[1:n-1,2:n] - P[1:n-1]
r,c = np.triu_indices(len_route-1,1)
time_diff_matrix[r+1,c+1] = vals[r,c-1]

另一种方法，以避免使用np.triu_indices和使用np.triu代替-

time_diff_matrix[1:n-1,2:n] = np.triu(vals)

验证结果-

In [265]: # Setup inputs
     ...: S = 10
     ...: M = np.random.randint(11,99,(S))
     ...: N = np.random.randint(11,99,(S,S))
     ...: P = np.random.randint(11,99,(S))
     ...: time_diff_matrix = np.zeros((S,S), dtype=int)
     ...: len_route = N.shape[0]
     ...: 

In [266]: # Original approach
     ...: for index1 in range(1, len_route):
     ...:     time_diff_matrix[index1, (index1+1):len_route] = \
     ...:         M[(index1-1):(len_route-2)] - \
     ...:         M[index1-1] + \
     ...:         N[index1-1, index1:(len_route-1)] + \
     ...:         N[index1, (index1+1):len_route] - \
     ...:         P[index1:(len_route-1)]
     ...:     

In [267]: # Proposed approach
     ...: time_diff_matrix_out = np.zeros_like(time_diff_matrix)
     ...: n = len_route
     ...: vals = M[:n-2] - M[:n-2,None] + N[:n-2,1:n-1] + N[1:n-1,2:n] - P[1:n-1]
     ...: time_diff_matrix_out[1:n-1,2:n] = np.triu(vals)
     ...: 

In [268]: np.allclose(time_diff_matrix, time_diff_matrix_out)
Out[268]: True