编写长代码作为嵌套列表理解

Question

我正在 Python 中编写代码以将离散拉普拉斯算子计算为二维中的稀疏矩阵。代码如下：

def discretizeLaplacian(N, step):
    NN = N**2
    dxx = (float) (1/((step)**2))
    i_loc, j_loc, vals = np.array([]), np.array([]), np.array([])
    for i in range(1, N-1):
        for j in range(1, N-1):      
            iv0 = j*N + i
            iv1 = (j - 1)*N + i
            iv2 = (j + 1)*N + i
            iv3 = j*N + i - 1
            iv4 = j*N + i + 1
            i_loc = np.concatenate((i_loc,[iv0, iv0, iv0, iv0, iv0]))
            j_loc = np.concatenate((j_loc,[iv0, iv1, iv2, iv3, iv4]))
            vals = np.concatenate((vals,dxx*np.array([-4, 1, 1, 1, 1])))
    sparseMatrix = csc_matrix((vals, (i_loc, j_loc)), shape = (NN, NN))
    return sparseMatrix

与 Matlab 相比，for 循环的运行时间要长得多。我想知道我是否可以将它写成嵌套列表理解以使其更快。

Answer 1

缓慢的性能来自算法的 糟糕的复杂性 。确实，原代码的复杂度是O(N**4)！这是由于 np.concatenate 通过复制旧数组并在新数组末尾添加一些项目来创建新数组 。这意味着执行了 3 个增长数组的 O(N**2) 个副本。通常，您应该避免在循环中使用 np.concatenate 来生成不断增长的数组。在这种情况下，您应该使用 Python 列表。

请注意，您可以使用 np.tile 重复数组的值并预先计算常量 dxx * np.array([-4, 1, 1, 1, 1])。

这是更正后的代码：

def discretizeLaplacian(N, step):
    NN = N**2
    dxx = (float) (1/((step)**2))
    tmp = dxx * np.array([-4, 1, 1, 1, 1])
    i_loc, j_loc = [], []
    for i in range(1, N-1):
        for j in range(1, N-1):   
            iv0 = j*N + I
            iv1 = (j - 1)*N + i
            iv2 = (j + 1)*N + i
            iv3 = j*N + i - 1
            iv4 = j*N + i + 1
            i_loc.extend([iv0, iv0, iv0, iv0, iv0])
            j_loc.extend([iv0, iv1, iv2, iv3, iv4])
    i_loc = np.array(i_loc, dtype=np.float64)
    j_loc = np.array(j_loc, dtype=np.float64)
    vals = np.tile(tmp, (N-2)**2)
    sparseMatrix = csc_matrix((vals, (i_loc, j_loc)), shape = (NN, NN))
    return sparseMatrix

以下是 N=200 在我的机器上的性能结果：

Original code:   22.510 s
Corrected code:   0.068 s

虽然您可以使用列表理解将 i_loc 设置为 [j*N+I for i in range(1, N-1) for j in range(1, N-1)]，但由于要追加多个项目，j_loc 并不容易。您可以构建一个 list/tuple 的列表，然后将其展平，但这会降低代码的可读性并且不会更快。

如果这还不够快，您可以使用 Numpy 向量化函数（例如 np.meshgrid、np.arange 和基本数学运算）来构建i_loc 和 j_loc。请考虑阅读 Numpy tutorials about that. Python loops are generally very slow since Python codes are generally interpreted. Matlab uses internally a just-in-time compiler (JIT) 以快速执行此类基于循环的代码。您可以使用 Numba 或 Cython 在 Python.

中做同样的事情

编写长代码作为嵌套列表理解

Writing long code as nested list comprehension

python

for-loop

nested-loops