patrick-kidger / diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Apache License 2.0
1.39k stars 123 forks source link

[time-scale separation] Adaptive step size for different variables? #505

Open davidschlegel opened 3 days ago

davidschlegel commented 3 days ago

In some scenarios, the different variables in a differential equation can have vastly different time-scales at which dynamics occurs. This is why often, it can be useful to select different integration steps for the variables, such that the variable with fast dynamics has a smaller time-step than the variable with the slow dynamics. I am wondering if such a method is in principle possible using the AbstractAdaptiveStepSizeController class.

patrick-kidger commented 2 days ago

So I think these kinds of multiple-step-size approaches often involves:

  1. picking a step size that is suited to the 'slow' timescale
  2. picking your favourite solver for the slow dynamics using that step size
  3. defining a solver for the fast dynamics that, for every one of those large steps: 3a. freezes the value of the slow dynamics 3b. runs a differential equation solver for the fast dynamics over the interval defined by the large step, and which will internally use smaller step sizes.

You should be able to do the same thing here: subclass AbstractSolver to implement a solver that exhibits the above behaviour. You will find that your .step method will make a call to diffeqsolve in order to solve the fast dynamics over that interval.

I hope that helps!