如何在这段简短代码中提高numpy性能？

Question

我正在尝试找出我的一个Python脚本为什么比gfortran慢大约4倍，经过一些分析，我发现了以下内容：

import numpy as np

nvar_x=40
nvar_y=10

def fn_tst(x):
    for i in range(int(1e7)):
        y=np.repeat(x,1+nvar_y)
    return y

x = np.arange(40)
y = fn_tst(x)

print y.min(),y.max()

这个Python脚本的速度大约是下面这段Fortran代码的13倍慢。

module test
integer,parameter::nvar_x=40,nvar_y=10
contains
subroutine fn_tst(x,y)
real,dimension(nvar_x)::x
real,dimension(nvar_x*(1+nvar_y))::y

do i = 1,10000000
   do k = 1,nvar_x
      y(k)=x(k)
      ibeg=nvar_x+(k-1)*nvar_y+1
      iend=ibeg+nvar_y-1
      y(ibeg:iend)=x(k)
   enddo
enddo

end subroutine fn_tst
end module test

program tst_cp
use test
real,dimension(nvar_x)::x
real,dimension(nvar_x*(1+nvar_y))::y
do k = 1,nvar_x
   x(k)=k-1
enddo

call fn_tst(x,y)

print *,minval(y),maxval(y)

stop
end

你能给我一些建议，让我的Python脚本运行得更快吗？另外，如果有关于numpy性能优化的建议也非常欢迎。我更倾向于继续使用Python，而不是为Fortran的例程编写Python的封装。

谢谢！

@isedev，所以，这就是结果吗？gfortran用时1.2秒，而Python用时6.3秒？这是我第一次关注性能问题，但正如我所说，我在尝试加速的代码中，Python的速度只能达到gfortran的四分之一。

对了，抱歉之前的代码没有做同样的事情。实际上，你在循环中提到的内容更像是我原始代码中的内容。

如果我没有理解错，我不同意最后的说法：我必须在fn_tst中创建y，而np.repeat只是右侧表达式中的一个部分（直接把输出放到现有数组中）。如果我注释掉np.repeat这一项，速度就会快很多……

rhs_slow = rhs[:J]
rhs_fast = rhs[J:]

rhs_fast[:] = c* ( b*in2[3:-1] * ( in2[1:-3] - in2[4:]  ) - fast) + hc_ovr_b * np.repeat(slow,K) #slow

Answer 1

首先，Python代码和Fortran代码的输出结果是不一样的。在Fortran程序中，y的序列是从0到39，然后是十个0，十个1，一直到十个39。而Python代码的输出是十一组0，十一组1，一直到十一组39。

这段代码的输出和你原来的代码是一样的，而且内存分配的次数也差不多：

import numpy as np

nvar_x = 40
nvar_y = 10

def fn_tst(x):
    for i in range(10000000):
        y = np.empty(nvar_x*(1+nvar_y))
        y[0:nvar_x] = x[0:nvar_x]
        y[nvar_x:] = np.repeat(x,nvar_y)
    return y

x = np.arange(40)
fn_tst(x)

print y.min(), y.max()

在我的系统上（只运行了1,000,000次循环），Fortran代码运行了1.2秒，而上面的Python代码运行了8.6秒。

不过，这样比较并不公平：Fortran代码中的y只分配了一次（在fn_tst函数外面），而Python代码中的y是在fn_tst函数里面分配的。

所以，把Python代码改成下面这样可以得到更好的比较：

import numpy as np

nvar_x = 40
nvar_y = 10

def fn_tst(x,y):
    for i in range(10000000):
        y[0:nvar_x] = x[0:nvar_x]
        y[nvar_x:] = np.repeat(x,nvar_y)
    return y

x = np.arange(40)
y = np.empty(nvar_x*(1+nvar_y))
fn_tst(x,y)

print y.min(), y.max()

在我的系统上，上面的代码运行了6.3秒（同样是1,000,000次循环）。所以已经快了大约25%。

不过，这里主要的性能问题是numpy.repeat()生成了一个数组，然后还需要把这个数组复制回y。如果numpy.repeat()能直接把结果放到一个已有的数组里（也就是这里的y），那就会快很多……但看起来这是不可能的。