ODE solvers and Autodiff

Help
1

Hi,

I’ve seen multiple discussions in the forum but they are all 3-5 years old (at least the ones I spotted). Do you guys have any recommendations for ODE solvers (anything close to DifferentialEquations.jl in performance and broadness). Similarly, any autodiff packages that might be compatible?

I don’t really code Julia; I come from Python, but scipy’s solvers can be notably slow (mostly due to the Python side generating the equations). On the other hand, using something like jax requires that I rewrite most of the code for my use case, which I’m currently doing. I wanted to try Fortran and make a comparison, in terms of speed, and also how long it takes me to get this simple code working. But since Julia seems to have the upper hand in differential equations, I was thinking perhaps I should go there directly, even though I’ve been avoiding Julia for a few years now

I saw some recommendations here like FLINT and FOODIE, but it seems like those are unmaintained.

2

I’ve developed an automatic differentiation package in Fortran:

It can be used up to the 3d order derivatives and there is no limit in terms of independent variables. You can get gradient, jacobian…

3

I checked Beliavsky’s list of Fortran code on GitHub and I could only see one ODE solver library that mentioned anything which sounded like AD: FATODE. I have never used this library and it appears to be unmaintained.

Many of these libraries were developed well before automatic differentiation was popularized, so it’s not likely many are compatible with Fortran derived type based AD.

My suggestion would be to pick an AD library you like and write your own minimal ODE solver that works with it. The ODE solver part will be a lot easier than the AD part. This is roughly what I did in my Nerf blaster simulator, which uses a forward-mode AD library I wrote for another project, a new RK4 integrator (which is pretty deeply embedded and can not really be extracted into a standalone library), and a form of event detection.

Developing a Fortran ODE solver library that works with a particular AD library would be a good idea in my view. The various AD libraries I’ve seen seem to have broadly compatible operators, but incompatible ways of initializing variables. There would be benefit to standardizing around a particular way to initialize AD variables in Fortran to encourage interoperability of AD libraries.

4

Hi,

FOODIE is essentially frozen; it was an experiment, and its results were only partially successful. In a nutshell, the abstract pattern is a real success for prototyping, whereas the performance penalty due to the abstraction is really blocking me from adopting it in my other projects.

If FOODIE fits your goals, I can try to help you use it (and do the probable necessary “maintenance”, it being frozen in recent years).

My best regards,
Stefano

5

Maybe you can find some useful solvers in this repo:

The solvers are compatible with solve_ivp, in case you need to use event handling.

You can also go “halfway” using the Python wrappers of DifferentialEquations.jl: GitHub - SciML/diffeqpy: Solving differential equations in Python using DifferentialEquations.jl and the SciML Scientific Machine Learning organization · GitHub

It does come with the limitation,

Autodiff does not work on Python functions. …

so you’d need to move (at least) your callback functions to pure Julia if autodiff is a must.

There is rklib (GitHub - jacobwilliams/rklib: Fixed and variable-step Runge-Kutta solvers in Modern Fortran · GitHub) which has a big collection of variable-step explicit Runge-Kutta methods.

Otherwise I’m not aware of anything comparable in terms of composability of different ODE solvers, linear solvers, non-linear solvers, event handling facilities and other bells and whistles. The investment in the SciML libraries is enormous. On the other hand we have LLMs nowadays so I believe it’s feasible to create your own set of SciPy-compatible or custom set ODE solvers, as long as you’ve already got your autodiff library figured out (PyTorch, JAX, …?).

I have investigated a number of ODE solving methods in the past (see Time-stepping method for cylindrical diffusion equation). Maybe if you provide more information about your problem we can give a better recommendation.

Here are few points to consider:

  • size of the ODE system: 10, 100, 1000, …
  • stiff or non-stiff?
  • differential-algebraic equations
  • mass matrix
  • event handling
  • symbolic Jacobian, numerical differencing, autodiff…
  • dense, banded, sparse, or other Jacobian pattern
  • discrete PDE problem

I’ve written Python bindings to the RKC ODE solver in Fortran before (https://github.com/ivan-pi/rkc/tree/python/python). It took me ~1 day of work and this was without LLMs. I did have a decent understanding of C/Python/Fortran interfacing, which made things easier. What I found the most challenging was learning to use Meson and figuring out how to create a pip-compatible Python Package.

When it comes to speed, the scipy.integrate.OdeSolver expects the solver to perform a single-step. This is so the abstract parent class can drive the integration and do the event-handling (if used). But it is not optimal from a performance point of view, because it means at each ODE solver step you go the full route Python → C → Fortran → C → Python, unpacking and re-packing the argument tuples, switching from interpreted to compiled code and vice-versa, which adds overhead.

To avoid this overhead you’d want to write your own solve_ivp like driver, and do all the time-stepping in the compiled language, ideally using a callback function that is also compiled (either calling C/Fortran or using Numba). I noticed this issue when designing the RKC bindings (https://github.com/ivan-pi/rkc/blob/python/python/rkc/rkc_proc.py), but I didn’t get around to “pushing” the time loop down into the C layer yet.

6

FWIW, I just came across this article - https://www.researchgate.net/publication/404975619_DJ4Earth_Differentiable_and_Performance-Portable_Earth_System_Modeling_via_Program_Transformations/references

7

You can give a try to DASPK3.1 ( Software | Computational Science and Engineering Research Group ), mind the copyright restrictions.
It was designed to work with ADIFOR to perform sensitivity analysis.

8

Another option, staying in the JAX ecosystem:

9

Thank you very much for such a detailed response!

Sorry for the late reply, I spent the last few days, tweeking things to port the code I had into Jax and forgot to check replies here.

This is sort of what I was trying before jax, I was using cython, however I should say that this was more about generating the ODEs than actually solving it. The reason I started considering other languages is that I thought that getting performance this way seems to be as much work as implementing the same in another language. I initially did some experiments with Julia, but since I’m still planning to use pure Python to prototype any new project, and only plan to use the “optimized” code I thought I might just use fortran instead. But I’d still like to test fortran, though I fear I might run into troubles implementing routines that work for both dense and sparse arrays, but LLMs should help.

That is what I did :sweat_smile: However, I’m still struggling with optimizing the code, but now it is fast enough that I might not want to go to julia, or a compiled language.

Other info:

  • The size of the ODE systems I typically study are around: 64 to ~8000 complex equations
  • Most likely non-stiff

The equations are always obtained from the equation:

\frac{d}{dt} \rho(t) = i \sum_{\omega,\omega'} \tilde{\mathcal{S}}(\omega,\omega',t) [ \rho(t), A^{\dagger} (\omega) A (\omega')]+ \sum_{\omega,\omega'} \tilde \gamma(\omega, \omega', t) \left( A (\omega') \rho(t) A^{\dagger}(\omega) -\frac{1}{2} \{A^{\dagger} (\omega) A (\omega'), \rho(t) \} \right)

where the A(\\omega) are matrices of size n \\times n.I guess the biggest drawback of my previous implementation and perhaps of the new one that uses jax is the generation of the RHS, so I thought going to another programming language would help.

I asked for autodiff because It would be nice to have, but I don’t actually need it for now.

EDIT: Problems rendering the equation

10

Did you try ZVODE?

It happens that I was testing fpm dynamical library with that (though it was not ready back then)

but there is a “modern fortran” test for zvode there fpm-dependencies/src/modern_testzvode.f90 at main · ettaka/fpm-dependencies · GitHub

For the reasons @ivanpribec mentioned, it will outperform scipy version (for small systems like yours) by a great margin.

11

I agree with @eelis, I also wanted to suggest ZVODE as a good fit due to the native complex number support.

The way the equations are written it looks like you could benefit from using ZGEMM for matrix-matrix multiplication, perhaps also with re-grouping and diagonalization of some of the terms. Moreover, the system appears to be linear in \rho(t), and the Jacobian matrix, \partial F / \partial \rho, has a dependence on t only.

(A colleague has told me these are the Gorini-Kossakowski-Sudarshan-Lindblad (GKSL) equations.)

12

Yes.

Adams method probably works well with non-stiff.

One technical aspect is to vectorize the density matrix so that the superoperators can be described with matmul operation.

It might be good also to use type for the liouviillian and type bound procedure for the call back function (so you don’t need to pass parameters through the function interface). Though nvfortran does not like it (in case you plan to use it but maybe GPU is not relevant for small system). Gfortran will work with it.

13

Indeed, not technically the GKLS, but closely related to them. in my case the coefficients are time dependent and do not form a positive semi-definite matrix, this kind of spoils diagonalization since I would need to do it for every t (not sure if there’s a smart way to go around it)

Thans for the suggestions ! I haven’t tried ZVODE, for the moment I haven’t tried any fortran implementation just considering it for the moment, I had pure python with scipy and now I have a jax implementation.

Indeed, I’m using vectorization, and actually don’t use the same equation I wrote down, but use the Redfield tensor plus a few extra terms that are neglected in the wikipedia article. I’m currently using jax/numpy’s eisum, not sure what the fortran equivalent would be

14

Something like this: if you store A(i,j,k,l) and rho(k,l), then

integer :: d=10 ! dimension
complex :: A(d,d,d,d)
complex :: rho(d,d)
! Populate A and rho somehow
! then vectorize rho and reshape A correspondingly
rho_v=reshape(rho,[d*d])
A_m=reshape(A,[d*d,d*d])
! Calculate the operation
rho_vo=matmul(A_m,rho_v)
! reshape to density matrix
rho_o=reshape(rho_vo,[d,d])
15

I just realized that VODE and ZVODE are available in the older scipy.integrate.ode integrator interface.

If I’m not mistaken, the C translation found here is called: