遥か昔、とあるゴミ箱にJaxを使ったODEの解き方を書いたのですが、そのゴミ箱が更新されることもなさそうなので、情報を追加してここに書いておこうと思います。 コードをまとめたnotebookは以下。 github.com Jaxとは 環境と準備 扱う力学系 まずはScipyで計算してみる。 odeintバージョン solve_ivpバージョン numbaを利用する Jaxを使う 関数の準備 計算 実行速度まとめ Jaxとは docs.jax.dev JAX is a Python library for accelerator-oriented array computation and pr…