Sparsity in DifferentialEquations: how to exploit?

Sorry for bothering you with a small request for help.
I have the following code, but I would expect the solution to the one-timestep
ODEProblem (using implicit Euler) to be computed in a time similar to AA\b or (I+AA)\b, which is not the case on my computer. How can this be fixed?
Any help (or constructive pointers to my stupidness) is highly appreciated!
Thanks.

using BlockBandedMatrices, BandedMatrices
using SparseArrays
using LinearAlgebra

## create matrix
Nx = Ny = 100
KK = Nx*Ny

FM = BandedMatrix(-1=>ones(KK-1), 0=>-2*ones(KK), 1=>ones(KK-1))
A = BandedBlockBandedMatrix(FM, [Nx for i in 1:Ny], [Ny for i in 1:Nx], (1,1), (1,1))
for k = 2:Nx-1
    view(A, Block(k-1,k)) .+= I(Nx)
end
for k = 1:Nx-1
    view(A, Block(k,k+1)) .+= -I(Nx)
end
view(A, Block(1,1))
view(A .+ A, Block(1,1))

##
@time AA = sparse(A)
println(nnz(AA))

##
b = randn(KK)
@time A*b
@time AA\b

## solve DiffEq
using DifferentialEquations

function dummy_prob(du, u, p, t) 
    du .= AA*u
end

tspan = (0., 1.)
dt = 1.
prob = ODEProblem(dummy_prob, b, tspan, [], dt=dt, adaptive=false)
@time sol = solve(prob, ImplicitEuler(), save_everystep=true)

A lot of it is in DiffEq doc

Thank you for the link.
However, when I do what’s below (i.e. specifying jac as an argument to ODEFunction), the solve(…) time went down from 30 to around 10 seconds, while

(I+AA)\b

still only takes less than one hectosecond! Would you have any clue why solve(…) is so slow?

## solve DiffEq
using DifferentialEquations

function dummy_prob(du, u, p, t) 
    du = AA*u
end

function f_jac(J,u,p,t)
    J .= AA
    nothing
end

tspan = (0., 1.)
dt = 1.

ff = ODEFunction(dummy_prob; jac=f_jac)

prob = ODEProblem(ff, b, tspan, [], dt=dt, adaptive=false)
@time sol = solve(prob, ImplicitEuler(), save_everystep=false)

you use global variables

Thanks… Yeah, so I looked into this and I added the matrix AA to the parameters p.
As you will see, this did not lead to a significant speedup:
Would you have yet another idea? :slight_smile:

using BlockBandedMatrices, BandedMatrices
using SparseArrays
using LinearAlgebra
using SparsityDetection
using BenchmarkTools

## construct matrix
Nx = Ny = 100
KK = Nx*Ny

FM = BandedMatrix(-1=>ones(KK-1), 0=>-2*ones(KK), 1=>ones(KK-1))
A = BandedBlockBandedMatrix(FM, [Nx for i in 1:Ny], [Ny for i in 1:Nx], (1,1), (1,1))
for k = 2:Nx-1
    view(A, Block(k-1,k)) .+= I(Nx)
end
for k = 1:Nx-1
    view(A, Block(k,k+1)) .+= -I(Nx)
end
view(A, Block(1,1))
view(A .+ A, Block(1,1))

## sparsify matrix
@time AA = sparse(A)
println("number nonzero elements = ", nnz(AA))

## some timings
b = randn(KK)
@btime A*b
@btime AA\b
@btime (I+AA)\b

## solve DiffEq
using DifferentialEquations

p = (AA)

function dummy_prob(du, u, p, t) 
    AA = p
    du .= AA*u
end

function f_jac(J,u,p,t)
    AA = p
    J .= AA
    nothing
end

ff = ODEFunction(dummy_prob; jac=f_jac)

prob = ODEProblem(ff, b, (0., 1.), p, dt=1., adaptive=false)
@btime sol = solve(prob, ImplicitEuler(), save_everystep=false)

Output:
0.002654 seconds (2.74 k allocations: 12.995 MiB)
number nonzero elements = 88804
845.376 μs (605 allocations: 125.20 KiB)
56.594 ms (73 allocations: 17.23 MiB)
7.093 ms (85 allocations: 12.20 MiB)
8.650 s (153 allocations: 1.49 GiB)

Iffyou look at the docs I linked, you only provide the jacobian function but it cannot know the sparsity. You should provide jac_prototype. Also, you can use mul! and friends in dummy_prob

Oh thank you @rveltz !
Adding the jac_prototype indeed did wonders!

For those readers who would be interested in the solution:

## solve DiffEq
using DifferentialEquations

p = (AA)

function dummy_prob(du, u, p, t) 
    AA = p
    du .= mul!(du, AA, u)
end

function f_jac(J,u,p,t)
    AA = p
    J .= AA
    nothing
end

ff = ODEFunction(dummy_prob; jac=f_jac, jac_prototype=AA)

prob = ODEProblem(ff, b, (0., 1.), p, dt=1., adaptive=false)
@btime sol = solve(prob, ImplicitEuler(), save_everystep=false)

To avoid aliases, I usually do ODEFunction(dummy_prob; jac=f_jac, jac_prototype=copy(AA)) but I am confident @ChrisRackauckas save guarded against it.

Yeah there’s a defensive copy.