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
536 stars 205 forks source link

Refactor OrdinaryDiffEq to use Subpackages for Solver Sets #2177

Closed ChrisRackauckas closed 1 month ago

ChrisRackauckas commented 5 months ago

From the timings on include:

  1.4346539974212646 => "perform_step/extrapolation_perform_step.jl"
  0.8758859634399414 => "dense/interpolants.jl"
   0.583730936050415 => "tableaus/rosenbrock_tableaus.jl"
  0.5598030090332031 => "perform_step/rosenbrock_perform_step.jl"
    0.50571608543396 => "algorithms.jl"
 0.41297292709350586 => "caches/rosenbrock_caches.jl"
  0.2902710437774658 => "algorithms/explicit_rk_pde.jl"
 0.22345685958862305 => "tableaus/verner_tableaus.jl"
  0.2234480381011963 => "perform_step/low_order_rk_perform_step.jl"
 0.20671916007995605 => "caches/low_storage_rk_caches.jl"

It's clear that this package has a long precompilation time simply because it's big. The suggestion then is to turn many of the solvers into subpackages. These can be subpackages like in a lib folder in this repo, kept in the same repository because it is expected that a lot of co-development will still be required between the solvers and the integrator type, but separated so that users can install a form without all of the solvers.

Now the question is how to get there? There are some pieces that can very clearly be separated. I think we can do this by classes of algorithms. For example, symplectic integrators, exponential integrators, low storage Runge-Kutta methods, and extrapolation methods are all not on the "main line" of the package, i.e. they are not part of the default algorithm and they are used in a minority of situations. Those clearly could be split to submodule packages and most users would benefit.

The edge case is what to do with methods on the default. The default set is being narrowed a bit with https://github.com/SciML/OrdinaryDiffEq.jl/pull/2103 to just some explicit RK methods, a few Rosenbrock methods, and FBDF. Do those and the default solver stay in OrdinaryDiffEq.jl, or do we move all of those out as well and make another subpackage just for the default solver?

The issue is that we need to make sure we get this "basically right" the first time, because this choice is breaking as some solvers would no longer be exported by default without pulling in new subpackages. The good thing though is that this won't be breaking to DifferentialEquations.jl, which keeps the simple users on a simpler update track.

Some questions to ask and answer:

  1. Why wait so long and why would we do something like this now? There was some hope that "tableau-ification", i.e. the refactoring to reduce perform_step! implementations into single tableaus, would be all that's needed. While we still have plans to do this refactoring, It's clear that some of the big times are just the tableau files themselves. This shows that the package is simply too big and you can only get the full improvements we want by splitting out some solvers, especially since many solvers are niche.
  2. What about moving some of the solvers to a different Github repository? The perform_step!, tableau, cache, etc. implementations are all relying on internals of this package. It would be unstable to keep them separate as any change to these internals would suddenly become breaking, and we expect from history that this is relatively common. To make this simpler to maintain, we can have the subpackages in lib folders so they can easily be bumped with any internals.
  3. What then becomes the difference between OrdinaryDiffEq core vs DiffEqBase? DiffEqBase is reused throughout all solvers, things like the solve.jl front end type promotions and callbacks, which are reused by things like Sundials. OrdinaryDiffEq core is only for
  4. Can this lead to dependency reduction? Indeed it would! For example, only the exponential integrators package actually needs ExponentialUtilities.jl. If FBDF is not in OrdinaryDiffEq.jl, then it would be possible to get the explicit methods without any of the AD or nonlinear solver stack. If the default methods are in OrdinaryDiffEq.jl then you would still need to have a lot of the stack in the core because of Rosenbrock and FBDF.
  5. Could it be non-breaking to do this to just the least used methods ASAP? I don't think we should do that, I think we have to do this all or nothing as silently breaking people's code is never a great time.
  6. What about one package per solver? Okay that's too far, that's too hard to maintain. In theory it's possible. In practice, we will probably want to do something like have Rosenbrock23 and Rodas5P together, and kick out all of the other Rosenbrocks so that the default can just grab the 2 most commonly used and the other 20 exist to still be pulled in? That's up for debate.
  7. What about just reducing the amount of code via refactoring the perform_step implementations? That's a different question, let's keep the discussion on that in the revived issue https://github.com/SciML/OrdinaryDiffEq.jl/issues/233.

Seeking comments from @devmotion @YingboMa @ranocha @KristofferC @oscardssmith

ranocha commented 5 months ago
  1. Why wait so long and why would we do something like this now? There was some hope that "tableau-ification", i.e. the refactoring to reduce perform_step! implementations into single tableaus, would be all that's needed. While we still have plans to do this refactoring, It's clear that some of the big times are just the tableau files themselves. This shows that the package is simply too big and you can only get the full improvements we want by splitting out some solvers, especially since many solvers are niche.

Sounds like an interesting idea 👍

  1. What about moving some of the solvers to a different Github repository? The perform_step!, tableau, cache, etc. implementations are all relying on internals of this package. It would be unstable to keep them separate as any change to these internals would suddenly become breaking, and we expect from history that this is relatively common. To make this simpler to maintain, we can have the subpackages in lib folders so they can easily be bumped with any internals.

I agree - it's easier if you keep everything in the same repo. Then, it's much easier for people to search for certain algorithms/caches etc. when they change something.

  1. What then becomes the difference between OrdinaryDiffEq core vs DiffEqBase? DiffEqBase is reused throughout all solvers, things like the solve.jl front end type promotions and callbacks, which are reused by things like Sundials. OrdinaryDiffEq core is only for

The last sentence is incomplete?

  1. Can this lead to dependency reduction? Indeed it would! For example, only the exponential integrators package actually needs ExponentialUtilities.jl. If FBDF is not in OrdinaryDiffEq.jl, then it would be possible to get the explicit methods without any of the AD or nonlinear solver stack. If the default methods are in OrdinaryDiffEq.jl then you would still need to have a lot of the stack in the core because of Rosenbrock and FBDF.

Having a package for explicit solvers only would be interesting for us in Trixi.jl - since we mainly use explicit methods (SSP and low-storage schemes).

  1. Could it be non-breaking to do this to just the least used methods ASAP? I don't think we should do that, I think we have to do this all or nothing as silently breaking people's code is never a great time.

Fully agree 👍 Even the least used methods are likely used by somebody.

  1. What about one package per solver? Okay that's too far, that's too hard to maintain. In theory it's possible. In practice, we will probably want to do something like have Rosenbrock23 and Rodas5P together, and kick out all of the other Rosenbrocks so that the default can just grab the 2 most commonly used and the other 20 exist to still be pulled in? That's up for debate.

This would make it much more tedious to experiment with different solvers. I think it's fine to keep similar solvers together.

  1. What about just reducing the amount of code via refactoring the perform_step implementations? That's a different question, let's keep the discussion on that in the revived issue Refactoring implementations of solver types to use tableau-based forms #233.

:+1: Just as a general comment: When we do that, we should also keep an eye on the readability of the code - using too many complicated macros etc. can make it much harder to understand the code for newcomers. I prefer having a bit of redundancy if it makes the code easier to read and understand.

oscardssmith commented 5 months ago

using too many complicated macros etc. can make it much harder to understand the code for newcomers. I prefer having a bit of redundancy if it makes the code easier to read and understand.

Fully agree on this, but it's also worth balancing this by the fact that redundant code significantly increases the bug/testing surface, and can obscure the differences between different methods.

ParamThakkar123 commented 3 months ago

the methods implemented in rkc_caches.jl show a Possible Method Call error image

@ChrisRackauckas

ChrisRackauckas commented 3 months ago

That should be caught in tests.

ChrisRackauckas commented 1 month ago
@time begin
    using OrdinaryDiffEqRosenbrock

    function lorenz(du, u, p, t)
        du[1] = 10.0(u[2] - u[1])
        du[2] = u[1] * (28.0 - u[3]) - u[2]
        du[3] = u[1] * u[2] - (8 / 3) * u[3]
    end

    lorenzprob = ODEProblem(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0))
    sol = solve(lorenzprob, Rodas5P());
end;
# 1.864242 seconds (3.40 M allocations: 217.743 MiB, 2.05% gc time, 6.16% compilation time: 2% of which was recompilation)

# Prior to the changes
@time begin
    using OrdinaryDiffEq

    function lorenz(du, u, p, t)
        du[1] = 10.0(u[2] - u[1])
        du[2] = u[1] * (28.0 - u[3]) - u[2]
        du[3] = u[1] * u[2] - (8 / 3) * u[3]
    end

    lorenzprob = ODEProblem(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0))
    sol = solve(lorenzprob, Rodas5P());
end;
# 3.419571 seconds (6.84 M allocations: 444.780 MiB, 4.16% gc time, 45.04% compilation time)
@time begin
    using OrdinaryDiffEqTsit5

    function lorenz(du, u, p, t)
        du[1] = 10.0(u[2] - u[1])
        du[2] = u[1] * (28.0 - u[3]) - u[2]
        du[3] = u[1] * u[2] - (8 / 3) * u[3]
    end

    lorenzprob = ODEProblem(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0))
    sol = solve(lorenzprob, Tsit5());
end;
# 0.564754 seconds (975.66 k allocations: 72.982 MiB, 1.42% gc time, 15.30% compilation time: 2% of which was recompilation)

# Before

@time begin
    using OrdinaryDiffEq

    function lorenz(du, u, p, t)
        du[1] = 10.0(u[2] - u[1])
        du[2] = u[1] * (28.0 - u[3]) - u[2]
        du[3] = u[1] * u[2] - (8 / 3) * u[3]
    end

    lorenzprob = ODEProblem(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0))
    sol = solve(lorenzprob, Tsit5());
end;

# 2.459902 seconds (5.46 M allocations: 349.671 MiB, 3.02% gc time, 22.39% compilation time)
ChrisRackauckas commented 1 month ago

Rosenbrock is kind of tricky. It has to have NonlinearSolve because it needs it if you have a DAE since it needs to run the initialization solve. That is about half of the remaining load time.