On a minimalistic interface for DifferentialEquations.jl

I, like most people here, think that DifferentialEquations.jl is a truly fantastic package and one of the major reasons to use Julia over other langs. Still, I’ve found that it is often very bulky: instead of returning arrays, solutions to ODE problems have truly massive types like

EnsembleSolution Solution of length 1024 with uType:
RODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, NoiseProcess{Float64, 2, Float64, Vector{Float64}, Vector{Float64}, Vector{Vector{Float64}}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), Nothing, false, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Vector{Float64}}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Vector{Float64}}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, SciMLBase.NullParameters, Nothing, SDEFunction{false, SciMLBase.FullSpecialize, LorenzDrift, GeometricBrownianNoise, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, GeometricBrownianNoise, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, Nothing}, SRIW1, StochasticDiffEq.LinearInterpolationData{Vector{Vector{Float64}}, Vector{Float64}}, SciMLBase.DEStats, Nothing, Nothing}

Though likely useful for dispatch purposes, this often makes it hard to guarantee type stability and sanity when combining it with other packages. It also leads to comically large error messages whenever things break (as they often do in e.g., scientific machine learning); I am currently staring at a wall of text that is 120k characters long, the vast majority of which is mostly useless, inscrutable type information.

Is there an easy way to use the tooling of DifferentialEquations, but with more straightforward typing? I was hoping for something like scipy/torchdiffeq, where the solver takes in a function and an array of initial conditions and outputs an array of solutions.

6 Likes

If you want just an array, you can call Array(sol) to turn it into a multidimensional array. There are some good reasons to not do this but if it simplifies your life then feel free to use it.

This is a missing Julia feature to cut non-dispatching type parameters out. @jeff.bezanson knows my qualms there and there’s some things for us to try.

2 Likes

Thanks, Chris.

I know you can always extract an array after the integration has taken place; the question was more whether there were any methods we could use that would allow us to work at the lower level of arrays all throughout. In my case, there are a large number of array/tensor manipulations going on before/during/after integration is done, like operating over batch dimensions and reshaping, and it becomes a bit cumbersome to keep rebuilding the ODE problem at each step.

I assume internally DiffEq calls some simpler integrating procedure that does just this, but I haven’t been able to find it. If there isn’t, that’s ok, it should be easy enough to implement a minimal rk procedure and have it run reasonably fast (though maybe it wouldn’t be differentiable?)

Again, this is not at all criticism; it would just improve my workflow (and apparently that of some other people who liked the post too) if there were a way to opt out of some of the nice high-level features of the library in exchange for more direct access to the objects being manipulated (again, like torchdiffeq).

“keep rebuilding the ODE problem at each step” sounds very wrong. Can you describe a bit more what you’re doing? You should really only build the ODE problem once and integrate once, otherwise you’re discarding the caches and slowing it down.

Are you looking for the integrator interface?

This is a less minimal interface as it’s going deeper, but it seems to be what you’re asking for in the follow up?

1 Like