diffrax

<h1 align='center'>Diffrax</h1> <h2 align='center'>JAX中的数值微分方程求解器。自动微分和支持GPU。</h2>

Diffrax 是一个基于 JAX 的库，提供数值微分方程求解器。

功能包括：

ODE/SDE/CDE（常微分/随机微分/受控微分）求解器；
多种不同的求解器（包括 Tsit5、Dopri8、辛求解器、隐式求解器）；
vmappable 所有东西（包括积分区域）；
使用 PyTree 作为状态；
密集解；
多种反向传播的伴随方法；
支持神经微分方程。

从技术角度看，该库的内部结构非常酷——所有类型的方程（ODEs、SDEs、CDEs）都以统一的方式被求解（而不是分别处理），生成一个小巧紧凑的库。

安装

pip install diffrax

需要 Python 3.9+、JAX 0.4.13+ 和 Equinox 0.10.11+。

文档

可在 https://docs.kidger.site/diffrax 获取。

快速示例

from diffrax import diffeqsolve, ODETerm, Dopri5
import jax.numpy as jnp

def f(t, y, args):
    return -y

term = ODETerm(f)
solver = Dopri5()
y0 = jnp.array([2., 3.])
solution = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)

这里，Dopri5 是指 Dormand--Prince 5(4) 数值微分方程求解器，这是许多问题的标准选择。

引用

如果您在学术研究中发现这个库有用，请引用：(arXiv链接)

@phdthesis{kidger2021on,
    title={{O}n {N}eural {D}ifferential {E}quations},
    author={Patrick Kidger},
    year={2021},
    school={University of Oxford},
}

（也请考虑在GitHub上为该项目加星。）