Setup with DifferentialEquations.jl time steppers?

[ChrisRackauckas]

I looked through the package and it seems like things are already in CuArrays then

https://github.com/climate-machine/Oceananigans.jl/blob/master/src/time_steppers.jl#L42-L78

this should quite readily port over to using DifferentialEquations.jl. It looks like you're using an IMEXEuler scheme? I think there would be some pretty good performance gains, and it would be interesting to start being able to use this entire package as a benchmark.

[ali-ramadhan]

Hey @ChrisRackauckas thanks for having a look!

We do this operator-splitting method and the pressure term is treated implicitly (while all other terms are treated explicitly). To step from time n to n + 1 first an explicit 2nd-order Adams-Bashforth step is used to calculate the right-hand-side at time n + 1/2 and the pressure is calculated at n + 1/2 by solving a Poisson equation. Then we time step the solution from time n to n + 1 via a forward Euler step (using the right-hand-side evaluated at n + 1/2).

I'm not super familiar with IMEX schemes although perhaps this is equivalent to IMEXEuler?

PS: Yes our numerical methods are pretty pathetic :(

Could be a great idea to rely on DifferentialEquations.jl for time-stepping. We'd get more options and any performance boost would be awesome! In particular, we've been hoping to upgrade to a higher-order low-storage time stepping method.

Maybe a good first step would be to try and replicate the current time-stepping method we use and see if all the tests pass?

I can start looking into trying to do this. I don't think I saw AB2 as an option (only AB3 and up) but maybe another method can replicate what AB2 does for us (or I can code something up).

[glwagner]

@ChrisRackauckas just to build on what @ali-ramadhan said, I think the main thing we'd need is an abstraction for operator splitting, more or less. There's an elliptic solve buried in the time-stepping loop (which solves for pressure).

@ali-ramadhan I could be wrong, but I don't think our method is IMEX Euler. It's 2nd-order Adams-Bashforth with a 'fractional step' via the operator splitting you described. The fractional step method is implicit in nature, but to my eye it seems the algorithm differs from those employed standard split implicit-explicit schemes (which do not involve operator splitting / fractional steps).

That said, it would be interesting if we could write the AB2 method with operator splitting in such a way that would allow us to use standard IMEX algorithms. I don't quite see how that'd be done now.

Another possibility, which I think is promising, is to somehow abstract the fractional step algorithm into a single 'RHS evaluation' (explicit substep + pressure solve + pressure correction substep) which we could then (somehow) pass to DifferentialEquations.jl. Then from the perspective of DifferentialEquations.jl we would be using a purely explicit method.

[ChrisRackauckas]

I was incorrect about the method. Sorry, I am not sure why I thought that.

I could be wrong, but I don't think our method is IMEX Euler. It's 2nd-order Adams-Bashforth with a 'fractional step' via the operator splitting you described. The fractional step method is implicit in nature, but to my eye it seems the algorithm differs from those employed standard split implicit-explicit schemes (which do not involve operator splitting / fractional steps).

Yes, the interesting thing would be to see if these methods fit into the scope of horizontal or vertical splitting of the resulting ODE system:

http://docs.juliadiffeq.org/latest/solvers/dynamical_solve.html http://docs.juliadiffeq.org/latest/solvers/split_ode_solve.html

If they do, we can do the following:

• Setup a form where the time stepping is with an ODEProblem. This would give the standard methods, but not the extra things.
• Then set it up with either SplitODEProblem or PartitionedODEProblem, depending on how the system is being split, and benefit from the new methods that arise from those categories.
• Implement missing "PDE methods" as time stepping methods on these classes of ODEs. Anything you have or need should probably get implemented.

@ali-ramadhan I'll be back next week, and we should whiteboard what you're doing for the time stepping here to understand how to characterize the underlying ODE to see what splitting abstraction it needs.

[glwagner]

@ChrisRackauckas can the constant-dt Runge-Kutta methods be adapted to accept a changing time-step? This is important because we often have external criteria available (the CFL criterion, for example) that can be used to adapt time-step size.

As for the fractional step method, we can also write a function that performs the fractional step algorithm using two substeps and an implicit pressure solve. Multiple fractional steps can then be embedded in a multi-step algorithm like Runge-Kutta; this may provide a route to integration with DifferentialEquations.jl. Note that the implicit pressure solves that forms the second part of the fractional step method requires a specialized fast solver for the 3D elliptic problem. We use this method:

https://www.sciencedirect.com/science/article/pii/0021999188901027

Integration with DifferentialEquations.jl will require integration of this FFT-based Poisson/Helmholtz solver into the algorithm.

Our implicit solves are usually coupled, such that they require the use of fast methods for the solution of elliptic PDEs to time-step efficiently (either the FFT-based algorithm described above, or a fast batched tridiagonal solver for the GPU that we are currently working on). Can the user provide their own fast solver for implicit time stepping with the split ODE solver?

Our time-stepping method is roughly described here:

https://climate-machine.github.io/Oceananigans.jl/stable/manual/time_stepping/

[ChrisRackauckas]

@ChrisRackauckas can the constant-dt Runge-Kutta methods be adapted to accept a changing time-step? This is important because we often have external criteria available (the CFL criterion, for example) that can be used to adapt time-step size.

Yes, you can adjust dt in a DiscreteCallback.

http://docs.juliadiffeq.org/latest/features/callback_functions.html

Our implicit solves are usually coupled, such that they require the use of fast methods for the solution of elliptic PDEs to time-step efficiently (either the FFT-based algorithm described above, or a fast batched tridiagonal solver for the GPU that we are currently working on). Can the user provide their own fast solver for implicit time stepping with the split ODE solver?

Yes.

http://docs.juliadiffeq.org/latest/features/linear_nonlinear.html

[glwagner]

Thanks @ChrisRackauckas.

I think the .html ending on your links are spurious; I found info about callbacks here:

http://docs.juliadiffeq.org/latest/features/callback_functions/

and the split ODE solvers here:

http://docs.juliadiffeq.org/latest/features/linear_nonlinear/

[ali-ramadhan]

Just an update: @ChrisRackauckas and I chatted some months ago on Slack about this issue and tried to figure out how Oceananigans.jl could use DifferentialEquations.jl for time-stepping.

I think we stumbled upon papers like https://onlinelibrary.wiley.com/doi/abs/10.1002/fld.4083 that use the method of lines

The discretized Navier–Stokes equations form an index 2 system of differential algebraic equations, which are afterwards reduced to a system of ordinary differential equations (ODEs), using the discretized form of the continuity equation. The pressure field is computed solving a discrete pressure Poisson equation. Finally, the resulting ODEs are solved using the backward differentiation formulas.

which is a pretty different method from what Oceananigans currently uses.

Solving index-2 DAEs seemed quite niche (and maybe tricky) so probably not something to pursue for 3D simulations on GPUs especially when we already have a fast solver.

---

I'm going to close this issue since Oceananigans doesn't really have problems with time-stepping and I don't think higher-order time-steppers would help Oceananigans a huge amount.

If anyone thinks that it would be interesting to try time-stepping Oceananigans as a SplitODEProblem with the FFTBasedPoissonSolver as a non-linear solver following https://diffeq.sciml.ai/latest/features/linear_nonlinear/#linear_nonlinear then please feel free to reopen!