DifferentialEquations.jl : supplied Jacobian solution takes more time and memory

Specific Domains Modelling & Simulations
1

Hello,
I want to supply the Jacobian to the ODE solver, it seems to me that on supplying the Jacobian the time and memory used by the solver is more than if I hadn’t supplied.
I am using sparse matrices for memory efficiency (for the example given below tri-diagonal would be better).

using DifferentialEquations,Plots, SparseArrays, LinearAlgebra

function my_fun!(du,u,p,t)
 for i = 2:length(u)-1
    du[i] = u[i+1] - 2*u[i] + u[i-1]
 end
 du[1] = 0
 du[end] = - 2*u[end] + u[end-1]
    nothing
end

function jacobian_!(J,u,p,t)
    for i = 2:length(u)-1
        J[i,i-1] = 1
        J[i,i] = -2
        J[i,i+1] = 1
    end
    J[1,1] = 0
    J[1,2] = 0
    J[end,end] = -2
    J[end,end-1] = 1
    nothing  
end

N = 1000;
u0 = zeros(N)
u0[1] =1.0
tspan = (0.0,1000.0)
p = []

abc = spzeros(N,N)
f! = ODEFunction(my_fun!,jac = jacobian_!,jac_prototype = jacobian_!(abc,u0,p,0))
prob_jac = ODEProblem(f!,u0,tspan,p)
println("Solving with Analytical Jacobian")
@time sol_jac = solve(prob_jac,TRBDF2());  
nothing
println("Solving without Analytical Jacobian")
prob = ODEProblem(my_fun!,u0,tspan,p)
@time sol_no_jac = solve(prob,TRBDF2());
nothing

Solving with Analytical Jacobian
  1.067418 seconds (1.22 M allocations: 81.478 MiB, 6.94% gc time, 56.08% compilation time: 100% of which was recompilation)
Solving without Analytical Jacobian
  0.708462 seconds (32.45 k allocations: 26.289 MiB, 40.18% gc time, 4.56% compilation time: 100% of which was recompilation)

What is it that I might be doing wrong?
Note: The condition number is very high for this example

julia> cond(Array(abc),2)
8.207431449069317e15
2

is this just compile time? if you run it a second time is it fast?

3

No I tried it a few times and similar result as above.
The real difference I think is in the amount of allocations.

4

Could the non concrete element type of the trivial vector p be problematic?

Also, it looks like you’re passing nothing as the jac_prototype as that is what your jacobian function returns

1 Like
5

make abc const ?

Actually jacobian_! returns nothing

1 Like
6

Thank you for the reply

I do not think so because I get the same behaviour when using p=[1,0].

If my understanding is correct the jacobian_! should only be used to update the already preallocated matrix in this case abc. Hence it should return nothing. I tried return J which in turn takes more memory and time.

7

abc is the preallocated matrix which will be updated. Hence making it const will probably not help. I tried it and the code fails to run

8

Jac_prototype is likely used as a type dispatch in the methods. Pass a sparse matrix to it with proper sparsity structure and it will use sparse linear solver. You are right to return nothing for your Jacobian but the prototype does not have the right type imm

1 Like
9

Having an in-place Jacobian computation is indeed good for performance, better than allocating a new (sparse) matrix each time the Jacobian is evaluated.

The problem is that because jacobian_! returns nothing, and that is what you pass as jac_prototype to ODEFunction, this is interpreted as the default case where the Jacobian is dense. There are 2 ways to fix this:

  1. First compute the Jacobian prototype separately by jacobian_!(abc,u0,p,0) and then pass jac_prototype = abc;
  2. Leave your ODEProblem as is but have jacoban_! return the Jacobian. This doesn’t introduce any new allocations.
10

One possibility that you could rule out is that your analytical Jacobian is not completely correct, which would explain why the algorithm has a hard time converging. Try comparing with an autodiff Jocobian

11

Also, since your Jacobian is state-independent (since your problem is linear), maybe the second problem detects linearity which is optimized for which you overrule in the first problem?

12

This iteration structure is slow. Iterate along columns not rows.

Note that sparse AD will absolutely murder your handcode in performance though, since it would use a coloring vector with chunked ForwardDiff to SIMD multiple elements along the diagonal at the same time. I wouldn’t even want to show you the equivalent code because it would be nasty to write out by hand, but you’d effectively clump chunks of 8 columns at the same time and iterate down those columns, and then have another loop on top that blocks it and SIMD from that outer loop, into a denseified matrix which matches the Tridiagonal structure.

1 Like