在numpy中加速循环

时间:2016-12-27 13:28:01

标签: python performance numpy vectorization

我正在寻找加速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是一个很大的数字。

问候。

1 个答案:

答案 0 :(得分:3)

以下是使用slicing rows/columns -

的方法
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