python - 如何正确地将 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 相比)!作为documentation mentions in the Notes section :

The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.

您可以通过从 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 是一个 float 组，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

可能还有一些可以通过循环提升或更智能的计算方法进行优化，但在我的计算机上它已经相当快了:大约 100 微秒与上面的 14 毫秒相比，所以又快了 100 倍。