如何正确地将 numpy vectorize 转换为 numba vectorize

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))

Answer 1

大多数内置 NumPy 函数已经矢量化，np.vectorize根本不需要装饰器。一般来说，numpy.vectorize装饰器会产生非常慢的结果（与 NumPy 相比）！正如文档在注释部分提到的：

提供矢量化函数主要是为了方便，而不是为了性能。该实现本质上是一个 for 循环。

f通过从、mz和中删除装饰器，您可以极大地提高代码效率mro。它会给出相同的结果，但运行速度更快（您的代码 10.4 秒，更改后的代码 0.014 秒）。

通过使用广播而不是 还可以改进该E1m功能（在性能方面）vectorize。

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

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

但如果您确实想使用该vectorize方法，请按照以下步骤操作f：

• 由于 ，您无法使用实例方法self，因此您需要静态方法或独立函数。因为您无权访问，self所以您需要传入delta或使其全局化。我决定将其作为一个论点：

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

• 然后，您需要找出您的参数的类型（或您想要支持的类型）以及签名返回的内容。你的 ro 是一个整数数组，rs 是一个浮点数组，delta 是一个整数，所以签名看起来像这样（语法是return_type(argument_1_type, argument_2_type, ....)）：

@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 倍。