如何正确将numpy向量化转换为numba向量化

Question

我正在尝试将我的程序转换为Numba，但在一个函数嵌套另一个函数时遇到了问题。我的方法基于NumPy的vectorize，但使用numba做不到同样的效果。您知道任何类似的示例可以供我参考吗？

这是我的程序：

import numpy as np
import scipy
import functools
from multiprocessing import Pool
import lib as mlib
from tqdm import tqdm

class vectorize(np.vectorize):
    def __get__(self, obj, objtype):
        return functools.partial(self.__call__, obj)

class stability:
    def __init__(self):

        self.r1 = 20
        self.r2 = 50
        self.rz= 20

        self.zeta = self.r2/self.r1
        self.beta = self.rz/self.r1
        self.Ms = 0.956e6
        self.delta = 1

    @vectorize
    def f(self,ro,rs):
        # print("delta=",self.delta)
        return rs/ro*np.exp( (-1/self.delta)*(ro-rs))

    @vectorize
    def mz(self,ro,rs):
        return ( 1-self.f(ro,rs)**2 ) / ( 1+self.f(ro,rs)**2 )

    @vectorize
    def mro(self,ro,rs):
        return ( 2*self.f(ro,rs) ) / ( 1+self.f(ro,rs)**2 )

    @vectorize
    def E1m(self,a, b, N,rs,d):     
        r = np.linspace(a+(b-a)/(2*N), b-(b-a)/(2*N), N)
        fx = r* ((1/self.delta+1/r)**2 * self.mro(r,rs)**2 + (1/r**2 + 1)*self.mro(r,rs)**2+d*(-(1/self.delta + 1/r) * self.mro(r,rs) + 1/r * self.mro(r,rs)*self.mz(r,rs) ))
        area = np.sum(fx)*(b-a)/N
        return area

if __name__ == "__main__":
    rs = np.arange(0,100,1)
    model = stability()
    print(model.E1m(0,20,300,rs,2))

- kking

---

这些方法中没有一个需要用 np.vectorize() 进行包装，因为它们本身已经被向量化了。 - Nils Werner

@MSeifert 您是正确的，抱歉，在进行了一些小的更正之后，代码可以正常工作。我并不想直接询问解决我的问题，而是希望引导我找到解决方案。NilsWerner 是的，您是正确的，但总的来说，我想优化我的代码，因为从数百个参数进行计算需要很长时间。 - kking

@kking 使用 vectorize 装饰器并不是一种优化方法，它只是一个方便的方法，并不是为了提高性能而设计的。事实上，它可能会降低性能。 - juanpa.arrivillaga

Answer 1

大多数内置的NumPy函数已经向量化，根本不需要使用np.vectorize修饰符。一般来说，numpy.vectorize修饰符将产生非常慢的结果（与NumPy相比）！正如文档在备注中提到： 向量化函数主要是为了方便而提供的，而不是为了性能。该实现本质上是一个for循环。

从f、mz和mro中删除修饰符可以极大地提高代码效率。它将产生相同的结果，但运行速度更快（您的代码10.4秒，更改后的代码0.014秒）。

也可以通过使用广播而不是vectorize来提高E1m函数（在性能方面）。

然而，由于您的问题是关于如何在这些函数上使用numba.vectorize，我有一些坏消息：不可能在实例方法上使用numba.vectorize——因为numba需要类型信息，而自定义Python类没有这些信息。

一般来说，对于numba，最好从在NumPy数组上无需向量化的纯循环代码开始，然后使用numba njit装饰器（或jit（nopython=True）。这对方法也无效，但更容易传递标量参数并仅迭代所需的数组。

但是，如果您真的想使用向量化方法，以下是您将如何为f执行它的方法：

由于存在self，因此无法使用实例方法，因此您需要静态方法或独立函数。由于您无法访问self，因此必须传递delta或使其全局。我决定将其作为参数：

def f(ro, rs, delta):
    return rs / ro * np.exp((-1 / delta) * (ro - rs))

• 你需要确定你的参数类型是什么（或者你想要支持哪些类型），以及签名返回的类型。你的 ro 是整数数组，rs 是浮点数数组，delta 是整数，所以签名看起来像这样（语法是 返回类型(参数1类型，参数2类型，....)）：

@nb.vectorize('f8(i8, f8, f8)')
def f(ro, rs, delta):
    return rs / ro * np.exp((-1 / delta) * (ro - rs))

基本上就是这样。

对于mz和mro，您可以做相同的操作（请记住您还需要那里的delta）：

@nb.vectorize('f8(i8, f8, f8)')
def mz(ro, rs, delta):
    return (1 - f(ro, rs, delta)**2) / (1 + f(ro, rs, delta)**2)

@nb.vectorize('f8(i8, f8, f8)')
def mro(ro, rs, delta):
    return (2 * f(ro, rs, delta) ) / (1 + f(ro, rs, delta)**2)

转换E1m函数似乎有点棘手（我没有尝试过），我把它留给读者作为练习。

---

如果您想知道我如何在没有使用vectorize的情况下解决它：

import numpy as np
import numba as nb

@nb.njit
def f(ro, rs, delta):
    return rs / ro * np.exp((-1 / delta) * (ro - rs))

@nb.njit
def mz(ro, rs, delta):
    f_2 = f(ro, rs, delta) ** 2
    return (1 - f_2) / (1 + f_2)

@nb.njit
def mro(ro, rs, delta):
    f_ = f(ro, rs, delta)
    return (2 * f_ ) / (1 + f_**2)

@nb.njit(parallel=True)
def E1m(a, b, N, rs, d):
    delta = 1
    r = np.linspace(a + (b - a) / (2 * N), b - (b - a) / (2 * N), N)
    result = np.empty(rs.size)
    for idx in nb.prange(rs.size):
        rs_item = rs[idx]
        sum_ = 0.
        for r_item in r:
            mro_ = mro(r_item, rs_item, delta)
            sum_ += r_item * ((1 / delta + 1 / r_item)**2 * mro_**2  
                              + (1 / r_item**2 + 1) * mro_**2  
                              + d * (-(1 / delta + 1 / r_item) * mro_ 
                                     + 1 / r_item * mro_ * mz(r_item, rs_item, delta)))
        result[idx] = sum_ * (b - a) / N
    return result

可能还有一些可以通过循环提升或更智能的计算方法进行优化，但在我的电脑上，它已经非常快了：与上面的14毫秒相比仅需约100微秒，因此又快了100倍。

- MSeifert

---

从答案来看，我觉得numpy的本地向量化（即去除装饰器）已经至少与numba.vectorize表现相当。在大多数情况下，这是正确的吗？ - undefined

Answer 2

MSeifert 谢谢您！现在我有了一个快40倍的解决方案。

@numba.vectorize('f8(f8, f8, f8)')
def f(ro, rs, delta):
    return rs / ro * np.exp((-1 / delta) * (ro - rs))

@numba.vectorize('f8(f8, f8, f8)')
def mz(ro, rs, delta):
    return (1 - f(ro, rs, delta)**2) / (1 + f(ro, rs, delta)**2)

@numba.vectorize('f8(f8, f8, f8)')
def mro(ro, rs, delta):
    return (2 * f(ro, rs, delta) ) / (1 + f(ro, rs, delta)**2)

@numba.vectorize(nopython=True)
def E1m(a, b, N,rs,d):   
    r = np.linspace(a+(b-a)/(2*N), b-(b-a)/(2*N), N)
    fx = r* ((1/delta+1/r)**2 * mro(r,rs,delta)**2 + (1/r**2 + 1)*mro(r,rs,delta)**2+d*(-(1/delta + 1/r) * mro(r,rs,delta) + 1/r * mro(r,rs,delta)*mz(r,rs,delta) ))
    area = np.sum(fx)*(b-a)/N
    return area

x=np.arange(0,100,1)
%timeit E1m(0,20,300,x,2)

571 µs ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

- kking

---

如果我的回答对您有帮助：请不要忘记点赞。如果您认为这是最有帮助的答案，请考虑接受它。另请参阅当有人回答我的问题时我该怎么做？。 - MSeifert