Use values computed in ODETerms inside of SaveAt

[Ricky5389]

Hi, I would like to save some values that are computed inside of my ODETerms. Is it possible to have the term output other things than the differential equations parameters ? For example, if my term looks like this.

def terms(t,y,args):
    y_dot = y + args
    args_changed = args / y_dot
    return y_dot

Would it be possible to save args_changed and use it inside of SaveAt? (This is a dummy example but it captures the idea)

[patrick-kidger]

You can use SaveAt(fn=...). This is a separate function to the one used in the vector field.

In general we don't allow saving extra values from the vector field. The reason is that the times at which this is evaluated are frequently not the same as the times at which we output values.

[Ricky5389]

Is there a way to have access to y_dot inside of SaveAt ?

[patrick-kidger]

No. As above this is intended -- the values at which you save may not be the values at which you evaluate the vector field. So there's no way to map the values from one to the other.

[etienney]

Is there a way to have access to y_dot inside of SaveAt ?

Maybe what you want to do is putting the option "steps=True" to sync the values at which you save to the values that you evaluate.

[BenjaminDAnjou]

In diffrax, it is possible to work around this directly for First-Same-As-Last (FSAL) solvers, which are the most popular ODE solvers (Tsit5, Dopri5, etc). For FSAL solvers, the vector field is stored as the second element of the solver_state tuple, with the first element being a bool that indicates whether it's the first step). So you can do this:

MORE DETAILS ON HOW TO DO THIS (CLICK TO EXPAND)

• Break up your diffeqsolve into many sub-diffeqsolve.

• Each sub-diffeqsolve should take as input the solver_state at t0 (with diffeqsolve(..., solver_state=solver_state)) and return the solver_state at t1 (with SaveAt(..., solver_state=True))

• The solver_state is then initialized as (False, init_vector_field), and the sub-diffeqsolve are put in a jax loop (like while_loop or scan), repassing the solver_state each time.

For best results, you may also want to pass controller_state and made_jump at each sub-diffeqsolve.

It's a bit of a hack. It would be better if we could always access the vector field computed at t0 and then save that (Edit: XLA deduplication could achieve this in practice).