优化numpy矩阵操作（目前使用for循环）

Question

我写了一些代码，用来根据列表中的n个项目计算n个矩阵，然后最后把所有矩阵乘在一起。

这段代码运行得比较慢，我想了解一下如何优化Python代码。我使用了性能分析工具，发现程序慢的原因主要是在这个矩阵乘法的循环里。

我想知道有没有人能给我一些建议，如何加快这个过程，或许可以利用Python/NumPy中基于C的内置函数？

def my_matrix(x):

    # Initialise overall matrix as an identity matrix
    # | M_11 M_12 |
    # | M_21 M_22 |
    M = np.matrix([[1, 0],[0, 1]])

    for z in z_all:
        param1 = func1(z)
        param2 = func2(x, z)
        param3 = func3(x, z)

        M_11 = param1 + param2
        M_12 = param1 - param2
        M_21 = param1 * param2
        M_22 = param1 / param2

        # Multiply matrix with overall master matrix
        M = M * np.matrix([[M_11, M_12],[M_21, M_22]])
    return M

根据我稍微查阅的一些资料，函数调用的计算开销比较大，所以，先计算好我的参数数组，然后在循环中访问这些数组，可能会比每次在循环中计算函数要高效一些……比如：

 param1s = funcs(z_all)
 param2s = funcs(x, z_all)
 etc

然后在for循环中：

for i, z in enumerate(z_all):
    param1 = params1[i]
    param2 = params2[i]
 etc.

这样做确实快了一点，但也就快了大约10%。因为减少函数调用带来的时间节省，被在循环中使用param1 = params1[i]访问数组所花的时间抵消了。

有没有人能给我一些建议呢？

Answer 1

你可以通过将计算 M_11, ... M_22 转换为向量化的方式来简化操作，比如可以写成 M_11s = params1 + params2 这样。

这样的话，你只需要在循环中进行矩阵的乘法运算就可以了：

import numpy as np

...

# compute your 'params' over vectors of z-values
param1s = func1(z_all)
param2s = func2(x, z_all)
param3s = func3(x, z_all)  # you don't seem to be using this for anything...

# compute 'M_11, ... M_22'
M_11 = param1s + param2s
M_12 = param1s - param2s
M_21 = param1s * param2s
M_22 = param1s / param2s

# we construct a (2, 2, nz) array from these
M_all = np.array([[M_11, M_12], [M_21, M_22]])

# roll the 'nz' axis to the front so that its shape is (nz, 2, 2)
M_all = np.rollaxis(M_all, -1, 0)

# initialize output with the identity
M_out = np.eye(2)

# loop over each (2, 2) subarray in 'M_all', update the output with the
# corresponding dot product
for mm in M_all:
    M_out = M_out.dot(mm)