Add variable-step integrators (adaptive time step)

[diegoferigo]

This PR:

• Introduces new variable-step integrator based on Embedded Runge-Kutta schemes (Section II.4 of Hairer93 ¹).
• These integrators are a particular implementation of the Explicit Runge-Kutta schemes introduced in #88. Given the support of any explicit Butcher tableau, embedded schemes can be built from an ExplicitRungeKutta class just by passing a multi-row $\mathbf{b}^\top$ matrix, that was already supported.
• The new EmbeddedRungeKutta class just calls the underlying explicit scheme in a while loop implementing a logic that either accepts or rejects an integration step based on an estimation of the integration error computed from two different outputs (corresponding to the $\mathbf{b}^\top$ rows) and corresponding to different integration orders (in other words, corresponding to a different number of terms considered in the Taylor expansion).
• Contrarily to #72, the new variable-step integrators now support dynamic states modelled as generic pytrees as opposed to just flat arrays.
• Note that contrarily to the fixed-step integrators, these variable step integrators use jax.lax.while_loop. As a consequence, these integrators are not reverse-mode differentiable.

---

📚 Documentation preview 📚: https://jaxsim--115.org.readthedocs.build//115/

Footnotes

1. Hairer, Norsett, Wanner, Solving Ordinary Differential Equations I - Nonstiff Problems, 1993. ↩

[diegoferigo]

cc @S-Dafarra

[diegoferigo]

I still need to port the logic to export the data of the actual variable steps, quite useful to debug these integrators. Since I want to move forward faster and this PR is not merging in main, I will add it in another PR.

@flferretti you can review the changes in #108.