numba/jit,函数列表作为函数参数

hmae6n7t  于 2021-09-08  发布在  Java
关注(0)|答案(0)|浏览(324)

我试图通过对函数及其时间相关参数(速率函数)进行抖动来加速ode解决方案。我已经成功地JIT了速率函数,现在在处理速率函数列表时遇到了问题:

import numpy as np
from numba import njit
from scipy.integrate import solve_ivp

def makerates(): #generates a list of 3 jitted rate functions
    b_f = 1

    def rateBackward(n,b_f):

        @njit
        def  rate(t):
            return b_f*n*t
        return rate

    return [rateBackward(i,b_f) for i in range(3)]

rates=makerates() #gives the list  [r1(t), r2(t), r3(t)]

@njit
def dy2(t,y,*argv):

    ####this works:
    r1=argv[0](t)
    r2=argv[1](t)
    r3 = argv[2](t)
    print(r1,r2,r3)

    ####this doesn't:
    # for i in range(3):
    #     print(argv[i](t))

    return y

y0 = 1
t_span=[0,1]
sol2 = solve_ivp(dy2, t_span, [y0],args=(*rates,))

numba不喜欢异构列表,所以我通过解压列表来传递速率函数,这似乎很好。然而,我想以某种方式在dy2中迭代argv,但我不知道如何做到这一点。如果我在dy2中取消注解该部分,它将抛出错误

TypingError: Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (Tuple(type(CPUDispatcher(<function makerates.<locals>.rateBackward.<locals>.rate at 0x0000022C74F040D0>)), type(CPUDispatcher(<function makerates.<locals>.rateBackward.<locals>.rate at 0x0000022C74F04D08>)), type(CPUDispatcher(<function makerates.<locals>.rateBackward.<locals>.rate at 0x0000022C74F04840>))), int64)
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    All templates rejected with literals.
In definition 9:
    All templates rejected without literals.
In definition 10:
    All templates rejected with literals.
In definition 11:
    All templates rejected without literals.
In definition 12:
    All templates rejected with literals.
In definition 13:
    All templates rejected without literals.
In definition 14:
    All templates rejected with literals.
In definition 15:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.

暂无答案!

目前还没有任何答案,快来回答吧!

相关问题