如何在 python 中对以下四个 for 循环进行矢量化？

Question

我必须构造一个大小为(Nx*Ny, Nx*Ny) 的矩阵，其中Nx 和Ny 可以大于100。目前我正在使用四个for 循环来初始化我的最终矩阵“matrix_result”（具有维度(Nx*Ny, Nx*Ny)) 很慢。

前两个循环遍历数组 xs 和 ys 中的所有元素。后两个循环再次遍历 xs 和 ys 中的相同元素。然后我用matrix_idx = idx_y1 + Ny * idx_x和matrix_idy = idx_y2 + Ny * idx_x2构造matrix_result的x-index和y-index。

这是完整的代码。如何向量化矩阵“matrix_result”的这些初始化？

import numpy as np

Nx = 100
Ny = 100

xs = np.linspace(0.0, 2.0, Nx)
ys = np.linspace(0.0, 2.0, Ny)

matrix_result = np.zeros((Nx * Ny, Nx * Ny))

for idx_x1 in range(Nx):
    for idx_y1 in range(Ny):

    # Get values of the arrays xs and ys
    x1 = xs[idx_x1]
    y1 = ys[idx_y1]

    # Compute arctan2 of y1 and x1
    argument1 = np.arctan2(y1, x1)

    for idx_x2 in range(Nx):
        for idx_y2 in range(Ny):

            if idx_x1 != idx_x2 or idx_y1 != idx_y2:

                # Get values of the arrays xs and ys
                x2 = xs[idx_x2]
                y2 = ys[idx_y2]

                # Compute arctan2 of y1 and x1
                argument2 = np.arctan2(y2, x2)

                # Compute the elements of the matrix matrix_results.
                distance_12 = np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
                matrix_element = np.cos(argument2 - argument1) * np.exp(distance_12)

                # Construct indices of the matrix matrix_result with dimension (Nx * Ny, Nx * Ny)
                matrix_idx = idx_y1 + Ny * idx_x1
                matrix_idy = idx_y2 + Ny * idx_x2

                # Insert elements into matrix
                matrix_result[matrix_idx, matrix_idy] = matrix_element

Answer 1

def vectorised_calculation(xs, xy):
    Nx = len(xs)
    Ny = len(ys)
    matrix_result = np.zeros((Nx * Ny, Nx * Ny))

    indices = np.indices([Nx,Ny,Nx,Ny])
    #Remove those sets of indices that do not satisfy the following condition
    mask = np.logical_or(indices[0] != indices[2], indices[1] != indices[3])
    indices = indices[:, mask]

    x1 = xs[indices[0]]
    y1 = ys[indices[1]]
    x2 = xs[indices[2]]
    y2 = ys[indices[3]]
    argument1 = np.arctan2(y1,x1)
    argument2 = np.arctan2(y2,x2)
    
    # Compute the elements of the matrix matrix_results.
    distance_12 = np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
    matrix_element = np.cos(argument2 - argument1) * np.exp(distance_12)
        
    # Construct indices of the matrix matrix_result with dimension (Nx * Ny,     Nx *     Ny)
    matrix_idx = indices[1] + Ny * indices[0]
    matrix_idy = indices[3] + Ny * indices[2]
        
    # Insert elements into matrix
    matrix_result[matrix_idx, matrix_idy] = matrix_element

    return matrix_result

我发现很难解释这个问题，但据我所知，这就是你要问的。

一般来说，您可以使用 np.indices 向量化嵌套 for 循环集。

xs = np.random.rand(40)
ys = np.random.rand(40)
s = time.time()
func1(xs, ys)
print("time taken for original:", time.time() - s)
s = time.time()
func2(xs, ys)
print("time taken for vectorised:", time.time() - s)

原创时间：19.610620737075806

矢量化所用时间：0.3943142890930176

与：

xs = np.linspace(0.0, 2.0, 100)
ys = np.linspace(0.0, 2.0, 100)

使用 plt.imshow 时会得到以下输出：