SciML / OrdinaryDiffEq.jl

High performance ordinary differential equation (ODE) and differential-algebraic equation (DAE) solvers, including neural ordinary differential equations (neural ODEs) and scientific machine learning (SciML)
https://diffeq.sciml.ai/latest/
Other
535 stars 203 forks source link

Post-stage callbacks #1292

Open ranocha opened 3 years ago

ranocha commented 3 years ago

It would be very nice to have (pre- and) post-stage callbacks, similar to PreStage and PostStage in PETSc

Motivation

I'm rewriting Trixi.jl, our tree-based numerical simulation framework for hyperbolic PDEs, to be more modular and use the SciML ecosystem for ODE solvers. There, we would like to use positivity-preserving limiters (and multi-physics coupling), which can be implemented conveniently using such post-stage callbacks, see https://github.com/trixi-framework/Trixi.jl/issues/235.

Existing infrastructure in OrdinaryDiffEq

Right now, we have stage_limiter! for SSPRK methods and in-place problems: https://github.com/SciML/OrdinaryDiffEq.jl/blob/0c2048a502101647ac35faabd80da8a5645beac7/src/perform_step/ssprk_perform_step.jl#L43 These are basically what I have in mind for post-stage callbacks: Some function that is called after a stage value of an (explicit) Runge-Kutta method has been computed but before the RHS function is evaluated at that stage; the callback is allowed to modify the recently computed stage value. For example, search for TSPostStage in https://www.mcs.anl.gov/petsc/petsc-current/src/ts/impls/explicit/rk/rk.c.html.

For OrdinaryDiffEq, I would propose to use the same calling signature as for stage_limiter!(stage_u, f, p, stage_t).

Questions

Would such a feature be okay in general? How shall we implement it? Similar to stage_limiters! in SSPRK methods, i.e. by adding appropriate functions as optional parameters to the algorithms, defaulting to trivial_limiter! which does nothing? Then, we could add these additional features algorithm by algorithm, starting by low-storage methods that are most important for us right now.

ChrisRackauckas commented 3 years ago

Would such a feature be okay in general?

Yes, it would be a great feature to have.

Similar to stage_limiters! in SSPRK methods, i.e. by adding appropriate functions as optional parameters to the algorithms, defaulting to trivial_limiter! which does nothing?

Definitely agree this is the right way to do it.

Then, we could add these additional features algorithm by algorithm, starting by low-storage methods that are most important for us right now.

I think it needs to be done algorithm by algorithm since only certain algorithms can support it. For example, there's no way any of the wrapped methods like Sundials could ever do this. But we can certainly get all of OrdinaryDiffEq.jl working with it. It's a big and tedious refactor but not particularly difficult. Maybe it's better to add it to the common interface though and just error if it's used on an algorithm that doesn't support it, given the sheer number of algorithms that will support it.

ranocha commented 3 years ago

Thanks for the feedback! I will start looking into this soon.