Support for ODE simulation with diffrax

[johannahaffner]

Hi, just found your library! Very exciting.

Were you planning to add support for ODE simulation with diffrax?
I use diffrax to simulate my models and optimistix to fit them, and more generally work in a hierarchical modeling context where I compose a larger model out of many smaller equinox modules.

I see a lot of similarities here - and I think it would make sense to serialize to an ODE system that can be used as a diffrax.ODETerm, which requires the signature to be t, y, args. The step function used in your example would then not be needed.

I'd also like to take advantage of diffrax.LinearInterpolation, I use it to define my experimental inputs.
I believe that input/condition handling is currently out-of-scope for vanilla SBML and only supported in PETab, right?

Let me know if this is on your To-Do list already - or if you'd be happy to take a PR on this.

[mayalenE]

Hi Johanna, thanks for your interest!

Yes I think that adding support for ODE simulation with diffrax would be great and happy to integrate a PR on that! However I'm not sure whether we can and should discard the step function: for some models they are parameter updates made by assignment rules after each integration step so we need to check whether this can be integrated within diffrax (i'm not familiar with it) and I think it's nice for users to be able to intervene on the system dynamics through time. Ideally we should support both solving one step at a time and solving for an array of timesteps.

Regarding your second point I'm not sure I understand, could you specify what you have in mind?

[johannahaffner]

Hi! Nice, I'm glad to hear it!

for some models they are parameter updates made by assignment rules after each integration step so we need to check whether this can be integrated within diffrax

These can usually be integrated into the ODE, no? I am working with a model that does contain algebraic relations for some parameters, too, it looks something like this

class ODESystem(eqx.Module):
    
     def __call__(self, time, states, args):
           x1, x2 = states
           p1, p2 = args
 
           pa = p1 * (x1 - p2) / x2    # Some algebraic relation that depends on the system state
           
           d_x1 = pa * x1    # Regular ODE system
           d_x2 = - p2 * x2 + p1 * x1

           d_state = d_x1, d_x2
           return jnp.array(d_state)

That looks like something that can be integrated into a Serializer to me.

Regarding your second point I'm not sure I understand, could you specify what you have in mind?

I have time-varying inputs, which I represent with a dfx.LinearInterpolation and then evaluate inside of the ODE system, like this

ts = jnp.array([0, 1, 1+0.001, 2])  # A step function
ys = jnp.array([1, 1, 2, 2])

u = dfx.LinearInterpolation(ts=ts, ys=ys)

def i1ffl(time, states, args):
       x, y = states
       params, u = args
       
       d_x = u.evaluate(time) - x
       d_y = u.evaluate(time)/x - y
       
       d_state = d_x, d_y
       return jnp.array(d_state)

I've been working with diffrax for a while and it truly is a joy to use. It also supports interactively stepping through a solve, similar to your step function, but you could use any one of its 19 ODE solvers to do it. A new library for nonlinear optimization in that same ecosystem (and by the same author), optimistix, already interfaces with optax (and also supports interactively stepping through a solve).

Both diffrax and optimistix are built on top of equinox, which you already use to register models as PyTrees.
Diffrax also has support for steady-state and custom events and is written to be easily extendable.

It seems to me that there are some synergies here!

[mayalenE]

Great yes looks like we should then be able to include it easily, akin your example! I agree that having the 19 ODE solvers of diffrax would be a great addition, and would indeed enable users to couple it with optimistix or the steady state event.

Do you wanna start the PR? I think that simply rewriting the ModelStep function should make the job (at the end of https://github.com/flowersteam/sbmltoodejax/blob/main/sbmltoodejax/modulegeneration.py) using interactive stepping through diffrax solve.

Then to check whether this is functioning as expected you can try to reproduce examples from https://developmentalsystems.org/sbmltoodejax/tutorials/biomodels_curation.html.

Regarding the potential use of LinearInterpolation, optimistix, SteadyStateEvent I think we should keep them outside of the main library (which is intended to stay minimal) but we could definitely add more tutorials on how to combine them with sbmltoodejax.

What do you think?