python - 如何提高Python中for循环的性能

Question

我正在使用 python 进行大量模拟，模拟系统响应。

我目前一直在使用 Runge-Kutta 方案，但偶然发现了我一直在测试的另一个方案。

在 Matlab 中对此进行测试时，与我的 Runge-Kutta 相比，我获得了卓越的性能。然而，当我将其转移到 Python 时，速度明显慢了。

我不确定这是否就是事实，或者我是否可以改进我的编码方式，所以如果可能的话，我很想听听您的一些意见。

Matlab 中的代码，示例:

dt = 0.0001;
f = randn(1, (60 / dt));
ns = length(f);
yo = zeros(3,1);
P1 = [0; 0.001; 0];
F = [1 0.0001 0; 0.001 1 0.0001; 0.001 0 1];
y1 = zeros(3, ns);
tic
for i = 1:ns
    y1(:, i) = P1*f(:, i) + F*yo;
    yo = y1(:, i);
end
toc

循环执行时间为 0.55-0.61 秒。

Python 代码，示例:

dt = 0.0001
f = np.random.randn(1, int(60 / dt))
ns = np.size(f)
yo = np.zeros((3))
F = np.array([[1, 0.0001, 0], [0.001, 1, 0.0001], [0.001, 0, 1]])
P1 = np.transpose(np.array([[0, 0.0001, 0]]))
y1 = np.zeros((3, ns), order='F')
start_time = time.time()
for i in range(ns-1):
    y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
    yo = y1[: , i]
print("--- %s seconds ---" % (time.time() - start_time))

其中循环执行时间为 2.8 -3.1 秒。

我可以做些什么来改善这个吗？

感谢您考虑我的问题。

Answer 1

我建议在评论中使用numba。这是一个例子:

import numba
import numpy as np

def py_func(dt, F, P1):
    f = np.random.randn(1, int(60 / dt))
    ns = f.size
    yo = np.zeros((3))
    y1 = np.zeros((3, ns), order='F')
    for i in range(ns-1):
        y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
        yo = y1[: , i]
    return yo

@numba.jit(nopython=True)
def numba_func(dt, F, P1):
    f = np.random.randn(1, int(60 / dt))
    ns = f.size
    yo = np.zeros((3))
    y1 = np.zeros((3, ns))
    for i in range(ns-1):
        y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
        yo = y1[: , i]
    return yo

您不能将“F”顺序与 numba 一起使用，因为它使用 C 类型数组，而不是 FORTRAN 数组。

时间差异如下所示:

纯Python循环:

%%timeit
py_func(dt, F, P1)

结果:

2.88 s ± 100 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

努巴:

%%timeit
numba_func(dt, F, P1)

结果:

588 ms ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)