使用 numba、缓存或任何其他优化使 python 代码更快

Question

我定义了以下函数

def laplacian_2D_array(func_2D):
    func_2nd_derv_x_fin2D = np.zeros((N,N))
    for j in range (0,N):
        func_CD2_x_list = []
        for i in range (0,N):
            value = func_2D[i][j] #func_2D is a NxN matrix
            func_CD2_x_list.append(value)
        func_CD2_x_array = np.array (func_CD2_x_list)[np.newaxis]
        func_2nd_derv_x = matrix_R @ (np.transpose(func_CD2_x_array))
        func_2nd_derv_x_fin2D[j] = np.transpose (np.reshape(func_2nd_derv_x,[N,]))

    func_2nd_derv_y_fin2D = np.zeros((N,N))
    for i in range (0,N):
        func_CD2_y_list = []
        for j in range (0,N):
            value = func_2D[i][j]
            func_CD2_y_list.append(value)
        func_CD2_y_array = np.array (func_CD2_y_list)[np.newaxis]
        func_2nd_derv_y = matrix_S @ (np.transpose(func_CD2_y_array))
        func_2nd_derv_y_fin2D[i] = (np.reshape(func_2nd_derv_y,[N,]))    
    return (np.add(func_2nd_derv_x_fin2D, func_2nd_derv_y_fin2D))

在上面的代码中，一个二维矩阵 func_2D 每 j 行被提取为列向量并与 matrix_R 相乘并存储在 func_2nd_derv_x_fin2D 的第 j 列中，类似地存储在第 2 个块中代码和最终函数 returns 添加 func_2nd_derv_x_fin2D 和 func_2nd_derv_y_fin2D

在这个N = 401和matrix_S和matrix_R也是N X N矩阵。此函数在 while 循环中被多次调用，单次迭代的执行需要花费大量时间。我已尝试 @njit 使其更快，但我没有成功并出现错误。我也尝试过使用 cache.

我们如何在这方面优化列表和数组以及优化定义函数的其他方法是什么？

我正在展示使用该函数的代码。

while (time<timemax):

    #Analytical Solution----------------------------------------------------------------------------
    exact_time = time/a_sec
    omega_t = np.zeros((N,N))
    psi_t = np.zeros((N,N))
    for i in range (0,N):
        for j in range (0,N):
            psi_t[i][j] = np.sin(x_list[i]) * np.sin(y_list[j]) * np.exp((-2*exact_time)/Re)
            omega_t[i][j] = 2*np.sin (x_list[i]) * np.sin(y_list[j]) * np.exp((-2*exact_time)/Re)
.
.
.
.
.
.
.
    # BiCGSTAB algo

    x0 = psi_0 #initial guess--> psi of previous time step
    r0 = omega_0 - laplacian_2D_array(x0) # r0 = b-Ax0
    r0_hat = r0
    rho_0 = 1
    alpha = 1
    w0 = 1
    v0 = np.zeros((N,N))
    P_0 = np.zeros((N,N))
    tol = 10 ** (-7)
    iteration = 0
    while ((np.max(np.abs(laplacian_2D_array(x0) - omega_0))) < tol):
        rho_prev = rho_0
        rho = np.dot ((np.reshape(r0_hat,(N**2,1))),(np.reshape(r0,(N**2,1))))
        beta = (rho/rho_prev) * (alpha/w0)
        P_0 = r0 + beta * (P_0 - w0 * v0)
        v0 = laplacian_2D_array(P_0)
        alpha = rho / np.dot ((np.reshape(r0_hat,(N**2,1))),(np.reshape(v0,(N**2,1)))) 
        s = r0 - alpha * v0
        t = laplacian_2D_array(s)

非常感谢任何固定代码的建议。

Answer 1

原来laplacian_2D_array的复杂代码可以简化为如下实现：

def laplacian_2D_array(func_2D):
    return (matrix_R @ func_2D + matrix_S @ func_2D.T).T

根据您提供的输入，在我的机器上使用随机矩阵的速度提高了 63 倍。大部分时间应该花在矩阵乘法上（如果 Numpy/Python/BLAS 在目标平台上正确安装，则可以非常有效地并行执行）。

使用 numba、缓存或任何其他优化使 python 代码更快

making python code faster using numba, cache or any other optimization

python

list

performance

numba