SplitODEProblem implementation

[ba2tro]

Implemented SplitODEProblem for #1506

[ba2tro]

cd(@__DIR__)
using Pkg
Pkg.activate("ModelingToolkit")

using ModelingToolkit, OrdinaryDiffEq

@parameters μ
@variables t x1(t) x2(t)
D = Differential(t)

eqs = [D(x1) ~ x2,
       D(x2) ~ μ*(1-x1^2)*x2 - x1]

@named sys = ODESystem(eqs)

u0 = [4.0,
      6.0]

p = [μ => 0.2]

tspan = (0.0, 80.0)
prob = SplitODEProblem(sys, u0, tspan, p, jac=true)
sol = solve(prob, Rodas5(autodiff=false))

Test for Van Der Pol oscillator

[YingboMa]

Could you add tests in test/odesystem.jl? Also, please test the Jacobian generation as well.

[ChrisRackauckas]

Can you test this on BCR

https://benchmarks.sciml.ai/html/Bio/BCR.html

I'm curious if it has a significant compile time impact.

[ba2tro]

Sure

[ba2tro]

using ModelingToolkit
using Catalyst, OrdinaryDiffEq, ReactionNetworkImporters, Sundials, Plots, DiffEqDevTools, ODEInterface
using ODEInterfaceDiffEq, LSODA, TimerOutputs, LinearAlgebra, BenchmarkTools

datadir = joinpath(dirname(pathof(ReactionNetworkImporters)), "../data/bcr")
const to = TimerOutput()
tf       = 100000.0

#generate ModelingToolkit ODESystem
prnbng = loadrxnetwork(BNGNetwork(), joinpath(datadir, "bcr.net"))
rn = prnbng.rn
osys = convert(ODESystem, rn)

tspan = (0., tf)

@timeit to "ODEProb No Jac" oprob = ODEProblem(osys, Float64[], tspan, Float64[])
@timeit to "SplitODEProb No Jac" soprob = SplitODEProblem(osys, Float64[], tspan, Float64[])

@timeit to "ODEProb DenseJac" densejacprob = ODEProblem(osys, Float64[], tspan, Float64[], jac=true)
@timeit to "SplitODEProb Dense Jac" sdensejacprob = SplitODEProblem(osys, Float64[], tspan, Float64[], jac=true)

@timeit to "ODEProb SparseJac" sparsejacprob = ODEProblem(osys, Float64[], tspan, Float64[], jac=true, sparse=true)
@timeit to "SplitODEProb SparseJac" ssparsejacprob = SplitODEProblem(osys, Float64[], tspan, Float64[], jac=true, sparse=true)

show(to)

output:

────────────────────────────────────────────────────────────────────────────────
                                        Time                    Allocations
                               ───────────────────────   ────────────────────────
       Tot / % measured:             331s /  35.4%           16.1GiB /  59.5%

 Section               ncalls     time    %tot     avg     alloc    %tot      avg
 ────────────────────────────────────────────────────────────────────────────────
 ODEProb No Jac             1    65.9s   56.2%   65.9s   5.04GiB   52.6%  5.04GiB
 SplitODEProb No Jac        1    51.4s   43.8%   51.4s   4.53GiB   47.4%  4.53GiB
 ────────────────────────────────────────────────────────────────────────────────

[ba2tro]

For the one's with jacobian, it just doesn't fit in my system's ram(8 gigs). Due to using swap, it just doesn't finish executing.

[ChrisRackauckas]

If sparse=true it should first find the sparsity pattern and then remove the linear operator via that into a sparse matrix. That should make the BCR case a lot faster.

[ba2tro]

───────────────────────────────────────────────────────────────────────────────────
                                           Time                    Allocations
                                  ───────────────────────   ────────────────────────
         Tot / % measured:             2079s /  66.6%            210GiB /  96.9%

 Section                  ncalls     time    %tot     avg     alloc    %tot      avg
 ───────────────────────────────────────────────────────────────────────────────────
 ODEProb SparseJac             1    1324s   95.6%   1324s    198GiB   97.2%   198GiB
 SplitODEProb SparseJac        1    60.5s    4.4%   60.5s   5.61GiB    2.8%  5.61GiB
 ───────────────────────────────────────────────────────────────────────────────────

Got these benchmarks in my system for sparse jacobian case on BCR.

Edit: Running them in separate Julia sessions gives correct results. Running them back to back in same session somehow results in the latter problem using the computation of the former and then some of it is already compiled away. So, the following are more reliable benchmarks:

                                      Time                    Allocations      
                             ───────────────────────   ────────────────────────
      Tot / % measured:           2126s /  51.1%            204GiB /  96.8%

 Section             ncalls     time    %tot     avg     alloc    %tot      avg
 ──────────────────────────────────────────────────────────────────────────────
 ODEProb SparseJac        1    1087s  100.0%   1087s    198GiB  100.0%   198GiB
 ──────────────────────────────────────────────────────────────────────────────

───────────────────────────────────────────────────────────────────────────────────
                                           Time                    Allocations      
                                  ───────────────────────   ────────────────────────
         Tot / % measured:              414s /  69.4%           33.1GiB /  80.5%

 Section                  ncalls     time    %tot     avg     alloc    %tot      avg
 ───────────────────────────────────────────────────────────────────────────────────
 SplitODEProb SparseJac        1     287s  100.0%    287s   26.6GiB  100.0%  26.6GiB
 ───────────────────────────────────────────────────────────────────────────────────

[ba2tro]

ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 80.0)
u0: 2-element Vector{Float64}:
 4.0
 6.0

On building the SplitODEProblem it still returns this message. Where can we change it to say SplitODEProblem?

[ChrisRackauckas]

The original definition is captured in a prob.problem_type trait `https://github.com/SciML/SciMLBase.jl/blob/master/src/problems/ode_problems.jl#L86

IIRC, ODEProblem just uses the show fallback:

https://github.com/SciML/SciMLBase.jl/blob/9e5ba6f10a0d347d48ba4d2bae4a3c610a589bc1/src/problems/problem_utils.jl#L15-L29

so that could be specialized to show the problem type trait.

[xtalax]

Please ping me when this is merged :)