diffrax icon indicating copy to clipboard operation
diffrax copied to clipboard

SaveAt dense and fn

Open dv-ai opened this issue 2 years ago • 1 comments

When using dense=True in SaveAt and a custom fn, I was expected that the interpolation of the sol will respect the custom fn but it's not the case.

For now, it seems impossible to have an interpolation only on "statisctics" which have be saved with a custom fn without saved all the info.

    def f(t, x, args):
        y = x
        return y

    solver = diffrax.Kvaerno5()
    stepsize_controller = diffrax.PIDController(rtol=1e-8, atol=1e-8)  # rtol=1e-8, atol=1e-8


    t0 = 0.
    t1 = 10
    y0 = jnp.array([3.0,2.0,10.0])
    dt0 = 0.0002

    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(f),
        solver,
        t0,
        t1,
        dt0,
        y0,
        saveat=diffrax.SaveAt(dense=True,steps=True,fn=lambda t,y,args: y[:2]),
        stepsize_controller=stepsize_controller,
        max_steps=100000,
    )

    print(sol.interpolation.infos) # dim 3 ko
    print(sol.ys) # -> dim 2 ok

dv-ai avatar Aug 29 '23 09:08 dv-ai

Yup, right now this is working as expected. Dense interpolation just saves absolutely everything.

I can see it might be nice to change that, but it'd take some effort to rewrite the internals appropriately, so it's not a priority for me. If it's important to someone then I'd be happy to advise how to write a PR on it.

patrick-kidger avatar Aug 29 '23 18:08 patrick-kidger