Using Jacobian sparsity structures of large ODEs/PDEs downstream

[Datseris]

Hi, I want to take advantage of the Jacobian sparsity pipeline that modelling toolkit offers, and use the sparse jacobian, and the generated function acting on this sparse Jacobian, in some packages downsteeam (specifically DynamicalSystems.lyapunovs). To this end I need both the Jacobian function, but also an initialized sparse matrix with the correct sparsity strucutre.

Here is my MWE, which uses a spatial discretization of the Kuramoto Shivashinky model:

using OrdinaryDiffEq, Random, ModelingToolkit

# Parameters
lx = 12    # length b of domain
T = 1000.0 # total integration time
dx = 0.1  # spatial descritization
nx = floor(Int, lx/dx) + 1  # no. of grid points

dx2 = dx + dx
dxx = dx*dx
dxxxx = dxx * dxx
p = [dx2, dxx, dxxxx]

xvec = range(0, lx; length = nx)
uinit = collect(xvec)
rng = MersenneTwister(42)
x = xvec
@. uinit = cos(x) + 0.1*sin(x/8) + 0.01*cos((2π/lx)*x)

function kse(du, u, p, t)
    D = length(u)
    dx2, dxx, dxxxx = p
    # indices that do not need modulo
    @inbounds for n in 3:length(u)-2
        np1 = n+1
        nm1 = n-1
        np2 = n+2
        nm2 = n-2
        du[n] = (
            -(u[np1] - 2u[n] + u[nm1])/dxx
            -(u[np2] - 4u[np1] + 6u[n] - 4u[nm1] + u[nm2])/dxxxx
            - u[n]*(u[np1] - u[nm1])/dx2
        )
    end
    # indices that needs modulo (off-diagonal entries)
    @inbounds for n in (1, 2, length(u)-1, length(u))
        np1 = mod1(n+1, D)
        nm1 = mod1(n-1, D)
        np2 = mod1(n+2, D)
        nm2 = mod1(n-2, D)
        du[n] = (
            -(u[np1] - 2u[n] + u[nm1])/dxx
            -(u[np2] - 4u[np1] + 6u[n] - 4u[nm1] + u[nm2])/dxxxx
            - u[n]*(u[np1] - u[nm1])/dx2
        )
    end
    nothing
end

prob = ODEProblem(kse, uinit, (0.0, T), p)

# Using modelling toolkit
using ModelingToolkit
prob = ODEProblem(kse, uinit, (0.0, T), p)
sys = modelingtoolkitize(prob)
jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
f1 = ODEFunction(kse, jac=jac) # from modeling toolkit
f2 = ODEFunction(kse, jac=true, sparse=true) # don't know from where
prob_jac = ODEProblem(f, uinit, (0.0, T), p)

the initial condition above generates a travelling wave

You will notice that f2 = will error, because sparse is not a valid keyword for ODEFunction. f1 = works for me. However, when I type f1.jac_prototype, the value of that field is nothing. So in conclusion I haven't found of a way to get an instance of a sparse Jacobian matrix so far.

[ChrisRackauckas]

using ModelingToolkit
prob = ODEProblem(kse, uinit, (0.0, T), p)
sys = modelingtoolkitize(prob)
prob_jac = ODEProblem(sys, uinit, (0.0, T), p, jac=true, sparse=true)
prob_jac.f.jac_prototype

[Datseris]

Cool, it worked. Here are the comparisons:

julia> @btime ds.jacobian($J0, $ds.u0, $ds.p, 0.0); # ForwardDiff
  81.700 μs (2 allocations: 2.12 KiB)

julia> @btime jac_acc($J0_acc, $ds.u0, $ds.p, 0.0); # Accelerated with ModelingToolkit
  186.509 ns (1 allocation: 48 bytes)

julia> @btime qr($J0); # non-sparse
  607.200 μs (6 allocations: 182.80 KiB)

julia> @btime qr($J0_acc); # sparse
  126.900 μs (119 allocations: 198.89 KiB)

These two are the main operations that take time in my algorithm (and also matrix-vector multiplication, but there it is quite obvious that the sparse will be faster).

The timings of calculating the 10 largest Lyapunov exponents went down from 1 hour to 6 minutes. ;) I'll also write a tutorial about this in the DynamicalSystems.jl, if you believe the 5 lines of code you pasted above are "stable".

[Datseris]

(compile times are much larger now, but of course this is totally worth it)

[ChrisRackauckas]

Yeah that code has been stable and I expect it to continue to be. And yeah, compile times are next 😅

[ChrisRackauckas]

Nothing more here than JuliaSimCompiler