Recursive use of DifferentialEquations.jl solutions

lieskjur · July 29, 2024, 3:11pm

Hi, I am currently in the process of implementing a two pass trajectory optimization algorithm where the solution (sol) of one IVP problem produced in the first pass is a parameter (part of p) in the IVP problem of the second pass, its solution then serving as a parameter for the first pass of the next iteration. The reason I am passing the whole solution is that I need to evaluate the dense output each time dynamics (f) are called, which isn’t fixed with adaptive time-stepping.

When using solve the main bottleneck seemed to be compilation time (99.9%) which occurred at every call, so instead I tried to first initialize an integrator with init and then just update p before every pass. Unfortunately, this produces an ERROR: LoadError: MethodError: Cannot convert an object of type.... The problem seems to be that even though I am using a method with “free” interpolation, a lot of information about the previous pass, such as the signature of f and the type of p are stored in sol, even though they shouldn’t be needed for interpolation (there might be some other reason for including them). This also explains why solve had to be compiled every time.

Is there some crafty way I could resolve this or should I bite the bullet and roll my own implementation of explicit RK integration? And perhaps, is something worth addressing in DifferentialEquations.jl itself?

ChrisRackauckas · July 29, 2024, 3:54pm

I do have an open issue for making this simpler:

github.com/SciML/SciMLBase.jl

Leaner serialization for solution types

opened 12:50PM - 22 Feb 24 UTC

ChrisRackauckas

The solution types contain functions, and things like JLD2 don't like functions.… So it would be nice to have a "strip_solution` function that generates a much leaner form of the solution for simply doing easier serialization. Here's how that should look: 1. [ ] we should make a new trait for `has_lazy_interpolation` in https://github.com/SciML/SciMLBase.jl/blob/master/src/alg_traits.jl which is then set to true on BS5 and VernX solvers. `strip_solution` should then give an informative error message that these problems are not supported by this function. Just some proper error handling, and with that out of the way the other solvers should be free to go. 2. [ ] https://github.com/SciML/SciMLBase.jl/blob/master/src/solutions/ode_solutions.jl#L102-L118 we can strip out any of the function information from this object. I.e. make a new problem where `prob` and `alg` are nothing. 3. [ ] Downstream, `sol.interp` holds some extra information as well. So in SciMLBase we need `strip_interp(::AbstractInterpolation)`. The SciMLBase ones are just the identity function, so that's easy https://github.com/SciML/SciMLBase.jl/blob/master/src/interpolation.jl . 4. [ ] At this point any non-OrdinaryDiffEq algorithm will strip well. The only downstream interp to handle is the OrdinaryDiffEq one. https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/src/interp_func.jl#L4-L15 . `f` is only in there for the lazy interpolations, so set it to `nothing`. You need to keep the `cache`s because that's used for dispatch, but at this point you're done with explicit methods. 5. [ ] For implicit methods, you need to strip https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/src/caches/rosenbrock_caches.jl#L30-L33 `jac_config` and `grad_config`, as those contain function information. As a simple thing, you can just create a new cache with everything nothing though since this is only used for dispatch (and `addsteps!`, but since you know these caches are non-lazy the interpolation `addsteps!` post solution is trivial so it won't use anything in the cache). We need to make good functionality for this anyways for better default solvers (CC @oscardssmith), so it might be a good time to do this now. With that done `strip_solution(sol)` should return a lean solution with no function information in it (or rather, any function information in it is in SciMLBase, which also defines the structure so that's a required import) and so it should work just fine with any BSON, JLD2, etc. package.

Though your case doesn’t need pretty much any of that. You could just put your p in a FunctionWrapper (FunctionWrappers.jl) in your case and that would eliminate all extra information and allow for caching the compilation in a non-f-specific way. It’s one line of code to put the wrapper on it, but on my phone so ask if you have an issue with that.

That should never be required, since at the end of the day SimpleDiffEq.jl has the straightforward loops if that’s what you need. Of course, that’s going to not have all of the error checking and adaptivity features, but it trades for being simpler and cheaper compile time.

lieskjur · July 29, 2024, 5:20pm

Thanks for the suggestions. How FunctionWrappers.jl works isn’t clear to me at first glance but I’ll have a closer look tomorrow. Using SimpleDiffEq.jl is certainly also a viable option, I’d just have add a function for interpolating f.

lieskjur · August 7, 2024, 11:50am

@ChrisRackauckas using FunctionWrappers.jl solved the performance issues, thank you for the recommendation.

In the future I will most likely want use DifferentialEquations.jl in more unusual ways, but before then I will hopefully publish a paper describing the algorithm.

ChrisRackauckas · August 8, 2024, 12:57am

No problem, glad it worked. Yes, we might want to create some nice tutorials with FunctionWrappers.jl to make it more accessible, it’s a good tool but underdocumented right now.