首先请注意,JiTCODE 不会像您这样的常规函数learned_fn
作为输入。它采用符号表达式的迭代或返回符号表达式的生成器函数。这就是您的示例代码可能会产生错误的原因。
你要什么
f
您可以通过更改属性并告诉它无法编译实际的衍生产品,将任何具有正确签名的衍生产品“注入”到 JiTCODE 中。这是一个最小的例子:
from jitcode import jitcode, y
ODE = jitcode([0])
ODE.f = lambda t,y: y[0]
ODE.compile_attempt = False
ODE.set_integrator("dopri5")
ODE.set_initial_value([1],0.0)
for time in range(30):
print(time,*ODE.integrate(time))
为什么你可能不想这样做
暂时忽略 Lyapunov 指数,JiTCODE 的全部意义在于为您硬编码您的导数并将其传递给 SciPyode
或solve_ivp
执行实际集成的人。因此,上面的示例代码只是将函数传递给 SciPy 的标准集成器(此处ode
)的一种过于复杂的方式,没有任何优势。如果您NN_model
一开始就非常有效地实施,那么您甚至可能无法从 JiTCODE 的自动编译中获得速度提升。
使用 JiTCODE 的李雅普诺夫指数功能的主要原因是它自动从导数的符号表示中获得切线向量演化的雅可比和常微分方程(Benettin 方法所需)。没有符号输入,它不可能做到这一点。从理论上讲,您也可以注入一个切向量 ODE,但是再一次,JiTCODE 几乎没有什么可做的,您可能最好使用 SciPyode
或solve_ivp
直接使用 SciPy。
你可能需要什么
如果要使用 JiTCODE,则需要编写一小段代码,将神经网络训练的输出转换为 JiTCODE 所需的 ODE 的符号表示。这可能远没有听起来那么可怕。您只需要获取训练好的系数并将其插入到神经网络一般形式的方程中即可。
如果你很幸运并且你NN_model
完全支持鸭子类型(和),你可以这样做:
from jitcode import t,y
n = 10 # dimension of your ODE
NN_input = [y(i) for i in range(n)]
learned_fn = NN_model(t,NN_input)[1]
这个想法是你NN_model
用抽象的符号输入(t
和NN_input
)喂一次。NN_model
然后一旦作用于这个抽象输入,为你提供一个抽象结果(这里你需要鸭子类型的支持)。如果我NN_model
正确解释了您的输出,则此结果的第二个组成部分应该是 JiTCODE 要求的抽象导数作为输入。
请注意,您NN_model
似乎期望维度是索引,但 JiTCODE y
期望维度是函数参数。因此,您不能只选择NN_input = y
,而必须如上所述对其进行转换。