python - 是否可以将类方法引用传递给 njit 函数？

Question

我尝试缩短一些代码的计算时间。所以我使用 numba 模块的 njit 装饰器来做到这一点。在此示例中:

import numpy as np
from numba import jitclass, jit, njit
from numba import int32, float64
import matplotlib.pyplot as plt
import time

spec = [('V_init' ,float64),
        ('a' ,float64),
        ('b' ,float64),
        ('g',float64),
        ('dt' ,float64),
        ('NbODEs',int32),
        ('dydx' ,float64[:]),
        ('time' ,float64[:]),
        ('V' ,float64[:]),
        ('W' ,float64[:]),
        ('y'    ,float64[:]) ]

@jitclass(spec, )
class FHNfunc:
    def __init__(self,):
        self.V_init = .04
        self.a= 0.25
        self.b=0.001
        self.g = 0.003
        self.dt = .01
        self.NbODEs = 2
        self.dydx    =np.zeros(self.NbODEs  )
        self.y    =np.zeros(self.NbODEs  )

    def Eul(self,):
        self.deriv()
        self.y += (self.dydx * self.dt)

    def deriv(self,):
        self.dydx[0]= self.V_init - self.y[0] *(self.a-(self.y[0]))*(1-(self.y[0]))-self.y[1]
        self.dydx[1]= self.b * self.y[0] - self.g * self.y[1]
        return



@njit(fastmath=True)
def solve1(FH1,FHEuler,tp):
    V = np.zeros(len(tp), )
    W = np.zeros(len(tp), )

    for idx, t in enumerate(tp):
        FHEuler
        V[idx] = FH1.y[0]
        W[idx] = FH1.y[1]
    return V,W


if __name__ == "__main__":

    FH1 = FHNfunc()
    FHEuler = FH1.Eul

    dt = .01
    tp = np.linspace(0, 1000, num = int((1000)/dt))

    t0 = time.time()
    [V1,W1] = solve1(FH1,FHEuler,tp)
    print(time.time()- t0)
    plt.figure()
    plt.plot(tp,V1)
    plt.plot(tp,W1)
    plt.show()

我想传递对名为 FHEuler = FH1.Eul 的类方法的引用，但它崩溃并给出此错误

This error may have been caused by the following argument(s):
- argument 1: cannot determine Numba type of <class 'method'>

那么是否可以传递对 njit 函数的引用？或者是否存在解决方法？

Answer 1

Numba 无法将函数作为参数处理。另一种方法是在使用内部函数处理其他参数之前编译该函数，并返回内部函数，并在其中运行已编译的输入函数。请尝试一下:

def solve1(FH1,FHEuler,tp):
    FHEuler_f = njit(FHEuler)
    @njit(fastmath=True)
    def inner(FH1_x, tp_x):
        V = np.zeros(len(tp_x), )
        W = np.zeros(len(tp_x), )
        for idx, t in enumerate(tp_x):
            FHEuler_f
            V[idx] = FH1_x.y[0]
            W[idx] = FH1_x.y[1]
        return V,W
    return inner(FH1, tp)

传递函数可能不是必需的。这个看起来不错

@njit(fastmath=True)
def solve1(FH1,tp):
    FHEuler = FH1.Eul
    V = np.zeros(len(tp), )
    W = np.zeros(len(tp), )

    for idx, t in enumerate(tp):
        FHEuler()
        V[idx] = FH1.y[0]
        W[idx] = FH1.y[1]
    return V,W