numpy中向量化函数的性能表现

Question

我想在Python中进行数学积分，具体步骤如下：

[1] 使用scipy.optimize.fsolve来解决一个隐式方程，以找到被积函数的最大值位置。

[2] 将被积函数的最大值移动到零，然后使用scipy.integrate.quad从-10到10进行积分。

因为被积函数有一个自由参数xi，所以我想用一系列的xi值来进行这个操作，借助numpy的功能，所以我使用了numpy.vectorize。

现在有两种方法可以将这个算法进行向量化：

[A] 分别对[1]和[2]进行向量化，然后将vec_[1]的结果作为输入传给vec_[2]。

[B] 向量化一个同时执行[1]和[2]的函数。

我注意到[A]的速度比[B]快很多。以下是代码：

from scipy import optimize, integrate
import numpy as np
from math import exp, sqrt, pi
import time

def integral_with_fsolve(xi):
    xc = optimize.fsolve(lambda x: x+1./(1+exp(-x))-xi,0.)
    def integrand(x,xi):
        return exp(-(x-xi+xc)**2)/(2.*sqrt(2.*pi))/(1.+exp(x+xc))
    integral = integrate.quad(integrand,-10.,10.,args=(xi,),epsabs=0.)
    return integral[0]

def integral(xi,xc):
    def integrand(x,xi):
        return exp(-(x-xi+xc)**2)/(2.*sqrt(2.*pi))/(1.+exp(x+xc))
    integral = integrate.quad(integrand,-10.,10.,args=(xi,),epsabs=0.)
    return integral[0]

def fsolve(xi):
    return optimize.fsolve(lambda x: x+1./(1+exp(-x))-xi,0.)

vec_integral_with_fsolve = np.vectorize(integral_with_fsolve)
vec_integral = np.vectorize(integral)
vec_fsolve = np.vectorize(fsolve)

xi = np.linspace(0.,2.,1000)

t0=time.time()
something = vec_integral_with_fsolve(xi)
dur=(time.time()-t0)
speed = xi.size/dur
print('Integrate and fsolve vectorized in one: speed = {} ints/sec'.format(speed))

t0=time.time()
xc = vec_fsolve(xi)
something = vec_integral(xi,xc)
dur=(time.time()-t0)
speed = xi.size/dur
print('Integrate and fsolve vectorized seperately: speed = {} ints/sec'.format(speed))

输出总是类似这样的内容：

同时向量化积分和fsolve：速度 = 298.151473998 ints/sec

分别向量化积分和fsolve：速度 = 2136.75134429 ints/sec

因为这只是我实际问题的简化版本，所以我想知道为什么会这样。有人能解释一下吗？谢谢！

Answer 1

总的来说，这个问题出在使用“同时处理”方法时的 xc 变量上。它是一个 ndarray 类型的数组，当你在 math.exp() 中用 x 或 xi（这两个都是浮点数）时，代码运行会变得非常慢。如果你在“同时处理”方法中把 xc=float(xc)，那么性能几乎和“分开处理”方法一样。

下面是关于如何发现这个问题的详细说明。

使用 cProfile 可以很容易地找到瓶颈所在：

AT ONCE:
ncalls  tottime  percall  cumtime  percall filename:lineno(function)
 1001    0.002    0.000    2.040    0.002 quadpack.py:135(quad)
 1001    0.001    0.000    2.038    0.002 quadpack.py:295(_quad)
 1001    0.002    0.000    2.143    0.002 tmp.py:15(integral_with_fsolve)
231231   1.776    0.000    1.925    0.000 tmp.py:17(integrand)
470780   0.118    0.000    0.118    0.000 {math.exp}
 1001    0.112    0.000    2.037    0.002 {scipy.integrate._quadpack._qagse}

SEPARATELY:
ncalls  tottime  percall  cumtime  percall filename:lineno(function)
 1001    0.001    0.000    0.340    0.000 quadpack.py:135(quad)
 1001    0.001    0.000    0.339    0.000 quadpack.py:295(_quad)
 1001    0.001    0.000    0.341    0.000 tmp.py:9(integral)
231231   0.200    0.000    0.278    0.000 tmp.py:10(integrand)
470780   0.054    0.000    0.054    0.000 {math.exp}
 1001    0.060    0.000    0.338    0.000 {scipy.integrate._quadpack._qagse}

整体的运行时间是：

AT ONCE:
1 loops, best of 3: 1.91 s per loop
SEPARATELY:
1 loops, best of 3: 312 ms per loop

最大的差别在于 integrand，在 integral_with_fsolve 中运行的时间几乎是7倍。数值积分 quad 也是如此。即使 math.exp 在“分开处理”方法中也快了两倍。

这表明在每种方法中评估的类型是不同的。实际上，这正是关键。当你使用“同时处理”时，可以打印 type(xc)，你会看到它是 numpy.ndarray, float64，而在“分开处理”方法中，它只是一个 float64。在 math.exp() 中对这两种类型进行求和似乎不是个好主意，如下所示：

xa = -0.389760856858
xc = np.array([[-0.389760856858]],dtype='float64')

timeit for i in range(1000000): exp(xc+xa)
#1 loops, best of 3: 1.96 s per loop

timeit for i in range(1000000): exp(xa+xa)
#10 loops, best of 3: 173 ms per loop

在这两种情况下，math.exp() 返回的都是 float。使用 numpy 中的 exp、sqrt 和 pi 可以减少差异，但会让你的代码变得更慢，可能是因为这些函数也可能返回 ndarray：

AT ONCE:
1 loops, best of 3: 4.46 s per loop
SEPARATELY:
1 loops, best of 3: 2.14 s per loop

所以在这种情况下，不转换为 ndarray 是个好主意。更好的做法是转换为 float，如下所示（在“同时处理”方法中是必要的）：

def integral_with_fsolve(xi):
    xc = optimize.fsolve(lambda x: x+1./(1+exp(-x))-xi,0.)
    xc = float(xc) # <-- SEE HERE
    def integrand(x,xi):
        return exp(-(x-xi+xc)**2)/(2.*sqrt(2.*pi))/(1.+exp(x+xc))
    integral = integrate.quad(integrand,-10.,10.,args=(xi,),epsabs=0.)
    return integral[0]

新的运行时间是：

AT ONCE:
1 loops, best of 3: 321 ms per loop
SEPARATELY:
1 loops, best of 3: 315 ms per loop