Using DifferentialEquations.jl for very large system

General Usage
1

Firstly, let me say that I think DifferentialEquations.jl is great, and I am very excited to try and use it.

I solve very large systems of ODEs (2^18 equations). I have usually used (in C++ and more recently in Julia) my own custom version of dopri5 for this purpose. The reason is I have a few requirements:

  • I usually pre-allocate the memory to save the solution at certain grid (t) points. While I don’t strictly need this if it doesn’t degrade performance, I do need:
  • The ability to save the solution at fixed grid (t) positions, but not others.
  • The ability to modify the solution at each step.

I’m having trouble with the last two when using DifferentialEquations.jl (and am willing to let go of the first).

I think the saving of the soultion should be straightforward, but doesn’t work:

zs = linspace(0, length, num_saves) # my 't' coordinate is z
prob = ODEProblem(g, AW0, (0.0, length))
solve(prob, Tsit5(), abstol=1e-10, dense=false, save_timeseries=true, tstops=zs)

with variations of this, I either get the solution saved at more points than I need, or it is not saved at every point I specify in tstops (I also tried saveat).

For the last point above I followed the discussion in https://discourse.julialang.org/t/get-values-of-symbols-inside-a-function/848/1 and tried:

my_callback = @ode_callback begin
    for c in cache
        for i = 1:length(f.L)
            c[i] *= <my function here>
        end
    end
    @ode_savevalues
end

This more or less does seem to work (I have other issues at the moment), but as this is not well documented, I wanted to check that this is the correct way to do this. Note that these are not events – I need to do this at every step – but simply the way the problem has to be cast (I am really solving a PDE, but can reduce it to an ODE with a transformation, but this needs to be updated every step to avoid loss of precision).

Thanks for any help/advice!

3 Likes
2

Let me get to the saving part last.

The documentation for the callbacks is here:

https://juliadiffeq.github.io/DiffEqDocs.jl/latest/features/callback_functions.html

It’s a little underdocumented because it’s going to be changed so that way it’s easier to make it to every solver on the common interface (currently it’s only implemented for OrdinaryDiffEq.jl. It’s not hard to add it to others, what’s hard is making it somewhat uniform while keeping the features).

If you’re just modifying the array values of the solution, you can just do:

my_callback = @ode_callback begin
    for i = 1:length(f.L)
      u[i] *= ...
    end
    @ode_savevalues
end

The changing of all of the cache is only needed if you’re changing non-index values or if you’re changing the size of the ODE system on the fly (maybe there are other odd cases like that).

However, note that if you are changing u like this at every step, the solution may no longer be continuous. This means that anything relying on interpolations (saveat and dense) will no longer work properly. If you do need this, you may want to have this just be an event which occurs every step, which is just by making event_f(args...) = true and then making an apply_event!. The downside here is that, to appropriately handle the discontinuity it has to save both before and after the event application.

May I ask why you are doing this post modification? If it’s to set some kind of normalization (multiplication, so that looks like what you may be doing?) or other condition on u, you probably should be using a DAE solver instead.

Pre-allocating in Julia doesn’t do too much because push! is amortized. The reason why I can make this statement and know it for a fact is that this part of the interface exists, it’s just undocumented. If you actually look at the definition for solve in OrdinaryDiffEq.jl, you will see that it is:

function solve{uType,tType,isinplace,T<:OrdinaryDiffEqAlgorithm,F}(
  prob::AbstractODEProblem{uType,tType,isinplace,F},
  alg::T,timeseries=[],ts=[],ks=[];<then it defines the kwargs>

so you can pass in timeseries, ts, and ks where timeseries is a Vector of the solutions at each timepoint, ts are the timepoints, and ks is a vector for saving interpolation information (no need here since dense=false). Note that it has to match what will be put in there, so timeseries is a Vector where each value matches the type of u0. So if u0 is a vector of size 2^18, timeseries should be a Vector{typeof(u0)}(num_timepoints). Note that this is the same setup as the solution, so you can just check sol.u on a similar problem if you aren’t sure.

I actually spent quite some time making this “safe”, as in it doesn’t even need to be the length of the final solution. It will fill the pre-allocated vector until it’s full, and then push! into it to finish (or if it doesn’t make it that far, it will resize! at the end to match the true length). So you don’t actually have to worry about making it the right size, just make sure it’s

But note that this isn’t implemented in other solvers yet. Actually, report back on whether it makes a difference. If it makes a difference on your system, maybe that means I just didn’t test it hard enough. Whether or not it makes a difference, please let me know. If it does make a substantial difference, open up an issue and we’ll get this implemented in the other solvers as well.

Now lastly:

The setup you want to use is the following:

solve(prob, Tsit5(), abstol=1e-10, dense=false,
                   save_timeseries=false, tstops=zs,
                   saveat = zs)

Maybe the save_timeseries argument should be renamed. Your setup (save_timeseries=true and tstops used) means is save every internal step of the method. tstops means that the solvers will stop at each value in tstops. So what this means is that it will try to step as large as possible (getting the tolerance), except if in that interval it has a value of tstops, in which case it will shorten the interval to hit tstops exactly, and save every step.

What the setup I am giving you means: “step as far as possible so that you hit the correct tolerance or a value in tstops, and ‘interpolate’ to save the values only at the timepoints in zs”. Note from earlier interpolation won’t actually work because your solution is discontinuous at each solution step, but since we make tstops the same as all of the values in saveat, there is no interpolation that is needed.

That said, I think you actually hit a bug: the interpolation won’t actually save anything if hits the value directly. I put a fix up on OrdinaryDiffEq.jl master and will put a patch out for this tonight.

There has been a discussion to change the argument to save_timeseries to save_everystep or something like that because this has come up in the chat before. If you want to pursue something like this, please open an issue. Also, since this saving behavior is modular (mix and matching tstops, saveat, save_timeseries, dense, and calck), I should probably make a tutorial notebook showing all of the different combinations and how to make it do things like this. (That would be a great contribution to DiffEqTutorials.jl if anyone finds the time!).

3

Also, since documentation of callbacks was mentioned here, if you want to push this to go quicker, please contribute your thoughts to this issue:

https://github.com/JuliaDiffEq/Roadmap/issues/10

OrdinaryDiffEq.jl has a full set of callbacks and events implemented so that way it’s all usable right now, but the question is how to set it up in a way that makes it easy for the same callback/events interface will work across the whole ecosystem. My proposal right now is to make event handling separate from callbacks, and then simplify what’s “common” in callbacks a bit.

The roadmap here is:

  • Agree on an interface which works for “the vast majority of use cases” but is easy to implement across the ecosystem (OrdinaryDiffEq.jl, Sundials.jl, ODE.jl, LSODA.jl, etc.)
  • Implement the new version into OrdinaryDiffEq.jl
  • Add these callbacks/events to Sundials.jl
  • Make sure it really is a “good universal” framework
  • Finish up by adding it to the others.

I can actually get this all done in a good weekend, but I don’t want to do it until I know we have it correct in some sense. Knowing whether it will work will be based on users like you trying to see if there’s holes in how they’d want to use callbacks/events, until all of the holes seemed to be patched up. So please feel free to contribute your input if possible!

4

Thanks very much for your help!
I am trying to implement your suggestions.

Running solve(prob, Tsit5(), abstol=1e-10, dense=false, save_timeseries=false, tstops=zs, saveat=zs) still only saves at two points (the start and end).

5

Yes, note that I mentioned that it was fixed on master. I’ll put a patch to METADATA right now for it.

Edit: here’s the PR if you want to track it:

https://github.com/JuliaLang/METADATA.jl/pull/7352

6

Well to make my system efficiently solvable I have to work on transformed variables G(t,x)=exp(-L(x) t) H(t,x), where G is what the ODE solver solves for, but H is what I am really interested in. The RHS of the ODE is a crazy function of G(t,x). Now this would be straightforward, apart from two things.

Firstly, L is complex and so for large t, the exponential can get very numerically challenging. My solution to this is to reset t every so often, by transforming the solution back to H, setting t = 0, and starting the ODE solver again. In my custom code I built this in, and I was trying to do this in your callback system with something like:

my_callback = @ode_callback begin
        zstep = t - f.czs
        f.czs = t
        for i = 1:length(f.L)
            u[i] *= exp(f.L[i]*zstep)
        end
        @ode_savevalues
    end

Essentially I save the “reset” value of t as czs in my function object.

The second problem I have is that sometimes L also depends on t. In that case I need to perform this “reset” procedure at every step. (Which is actually what the above code is doing). Usually, in this case we write a custom solver ( a so-called split-step routine) but it is much nicer to use a standard ODE package instead, if I can get this reset procedure to work (and I think it will).

Does that make it clearer?

7

Yes, it should work. And the new proposal for the event handling framework should handle this too. The only thing which is stopping this right now is the saving bug? If that’s the case and you’re anxious, do Pkg.checkout("OrdinaryDiffEq") with the options setup above to try out master and see if the patch fixes your problems (and then afterwards get back to release via Pkg.free("OrdinaryDiffEq")). Or just wait for the tag, either way I’ll be ready for the response.

8

Yes, the master branch works fine, and solves the saving issue.
I’ll continue porting and checking my code.
Thanks for a great package and great response!

1 Like