diffrax
diffrax copied to clipboard
SaveAt dense and fn
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
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.