Feature request: Minimal step size for adaptive solvers.

rtqichen / torchdiffeq

Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation.

MIT License

5.54k stars 922 forks source link

Feature request: Minimal step size for adaptive solvers. #159

Closed jambo6 closed 3 years ago

jambo6 commented 3 years ago

Would it be straightforward to implement a dt_min argument in odeint analogous to that implemented in torchsde?

For a project I am working on using an adaptive solver is taking an order of magnitude longer than a fixed solver, however I only require the min step of the adaptive solver to be the same as that of the fixed grid.

rtqichen commented 3 years ago

Yeah sounds good. It should simply be adding a min_step kwarg to https://github.com/rtqichen/torchdiffeq/blob/master/torchdiffeq/_impl/rk_common.py#L117 and thresholding the adaptive step sizes. Scipy solvers also have a min_step argument that our wrapper can pass through (https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html).

I can take a PR for this, or I can add it some time in the future.

jambo6 commented 3 years ago

Nice one, I'll sort out a PR over the weekend.