Computing ODE sensitivities with sparse Jacobians in solver

Specific Domains Numerics
, , ,
1

Hey there,

I’m trying to compute the derivatives of a function involving an ODE solution. This works if I use dense Jacobians in the ODE solver.

However, the Jacobian is actually sparse (PDE discretization, stiff) and I thought to exploit that:

using DifferentialEquations, ForwardDiff, SparseArrays, SciMLSensitivity

function rhs!(du, u, p, t)
    du[1] = p[1] * u[2]
    du[2] = -p[2] * u[1]
    du[3] = u[1] - u[2]
end

u0 = [1.0, 0.0, 0.0]
p0 = [0.5, 2.0]
tspan = (0.0, 5.0)

jac0 = spzeros(length(u0), length(u0))
jac0[1, 2] = 1.0
jac0[2, 1] = -1.0
jac0[3, 1] = 1.0
jac0[3, 2] = -1.0

function f_dense(p)
    prob = ODEProblem(ODEFunction(rhs!), u0, tspan)
    sol = solve(prob, QNDF(autodiff=true), p=p, saveat=1.0)
    return sol.u[end]
end

# This works:
ForwardDiff.jacobian(f_dense, p0)

function f_umfpack(p)
    prob = ODEProblem(ODEFunction(rhs!; jac_prototype=jac0), u0, tspan)
    sol = solve(prob, QNDF(autodiff=true, linsolve=UMFPACKFactorization()), p=p, saveat=1.0)
    return sol.u[end]
end

# This fails:
ForwardDiff.jacobian(f_umfpack, p0)

The error is

ERROR: MethodError: no method matching SparseArrays.UMFPACK.UmfpackLU(::SparseMatrixCSC{ForwardDiff.Dual{ForwardDiff.Tag{…}, Float64, 2}, Int64})

Closest candidates are:
  SparseArrays.UMFPACK.UmfpackLU(::SparseArrays.UMFPACK.Symbolic{Tv, Ti}, ::SparseArrays.UMFPACK.Numeric{Tv, Ti}, ::Int64, ::Int64, ::Vector{Ti}, ::Vector{Ti}, ::Vector{Tv}, ::Int64, ::SparseArrays.UMFPACK.UmfpackWS{Ti}, ::Vector{Float64}, ::Vector{Float64}, ::ReentrantLock) where {Tv<:Union{Float64, ComplexF64}, Ti<:Union{Int32, Int64}}
   @ SparseArrays ~/.julia/juliaup/julia-1.10.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.10/SparseArrays/src/solvers/umfpack.jl:225
  SparseArrays.UMFPACK.UmfpackLU(::SparseArrays.AbstractSparseMatrixCSC{Tv, Ti}; control) where {Tv<:Union{Float64, ComplexF64}, Ti<:Union{Int32, Int64}}
   @ SparseArrays ~/.julia/juliaup/julia-1.10.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.10/SparseArrays/src/solvers/umfpack.jl:239

Stacktrace:
  [1] init_cacheval(alg::UMFPACKFactorization, A::SparseMatrixCSC{…}, b::Vector{…}, u::Vector{…}, Pl::LinearSolve.InvPreconditioner{…}, Pr::LinearAlgebra.Diagonal{…}, maxiters::Int64, abstol::ForwardDiff.Dual{…}, reltol::ForwardDiff.Dual{…}, verbose::Bool, assumptions::OperatorAssumptions{…})
    @ LinearSolve ~/.julia/packages/LinearSolve/0Q5jw/src/factorization.jl:812
  [2] init(::LinearProblem{…}, ::UMFPACKFactorization; alias_A::Bool, alias_b::Bool, abstol::Float64, reltol::Float64, maxiters::Int64, verbose::Bool, Pl::LinearSolve.InvPreconditioner{…}, Pr::LinearAlgebra.Diagonal{…}, assumptions::OperatorAssumptions{…}, sensealg::LinearSolveAdjoint{…}, kwargs::@Kwargs{})
    @ LinearSolve ~/.julia/packages/LinearSolve/0Q5jw/src/common.jl:208
  [3] build_nlsolver(alg::QNDF{…}, nlalg::NLNewton{…}, u::Vector{…}, uprev::Vector{…}, p::Vector{…}, t::Float64, dt::Float64, f::ODEFunction{…}, rate_prototype::Vector{…}, ::Type{…}, ::Type{…}, ::Type{…}, γ::Rational{…}, c::Int64, α::Int64, ::Val{…})
    @ OrdinaryDiffEqNonlinearSolve ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/EXjDj/src/utils.jl:199
  [4] build_nlsolver
    @ ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/EXjDj/src/utils.jl:144 [inlined]
  [5] build_nlsolver
    @ ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/EXjDj/src/utils.jl:134 [inlined]
  [6] alg_cache(alg::QNDF{…}, u::Vector{…}, rate_prototype::Vector{…}, ::Type{…}, ::Type{…}, ::Type{…}, uprev::Vector{…}, uprev2::Vector{…}, f::ODEFunction{…}, t::Float64, dt::Float64, reltol::Float64, p::Vector{…}, calck::Bool, ::Val{…})
    @ OrdinaryDiffEqBDF ~/.julia/packages/OrdinaryDiffEqBDF/J0IGS/src/bdf_caches.jl:430
  [7] __init(prob::ODEProblem{…}, alg::QNDF{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Float64, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{…}, abstol::Nothing, reltol::Nothing, qmin::Rational{…}, qmax::Int64, qsteady_min::Int64, qsteady_max::Rational{…}, beta1::Nothing, beta2::Nothing, qoldinit::Rational{…}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(LinearAlgebra.opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEqCore.DefaultInit, kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/55UVY/src/solve.jl:349
  [8] __init (repeats 5 times)
    @ ~/.julia/packages/OrdinaryDiffEqCore/55UVY/src/solve.jl:11 [inlined]
  [9] __solve(::ODEProblem{…}, ::QNDF{…}; kwargs::@Kwargs{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/55UVY/src/solve.jl:6
 [10] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/55UVY/src/solve.jl:1 [inlined]
 [11] #solve_call#44
    @ ~/.julia/packages/DiffEqBase/vscDj/src/solve.jl:612 [inlined]
 [12] solve_up(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::Vector{…}, args::QNDF{…}; kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/vscDj/src/solve.jl:1084
 [13] solve_up
    @ ~/.julia/packages/DiffEqBase/vscDj/src/solve.jl:1070 [inlined]
 [14] solve(prob::ODEProblem{…}, args::QNDF{…}; sensealg::Nothing, u0::Nothing, p::Vector{…}, wrap::Val{…}, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/vscDj/src/solve.jl:1007
 [15] f_umfpack(p::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(f_umfpack), Float64}, Float64, 2}})
    @ Main ./REPL[13]:3
 [16] vector_mode_dual_eval!
    @ ~/.julia/packages/ForwardDiff/PcZ48/src/apiutils.jl:24 [inlined]
 [17] vector_mode_jacobian(f::typeof(f_umfpack), x::Vector{…}, cfg::ForwardDiff.JacobianConfig{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:125
 [18] jacobian(f::Function, x::Vector{…}, cfg::ForwardDiff.JacobianConfig{…}, ::Val{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:0
 [19] jacobian(f::Function, x::Vector{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{…}, Float64, 2, Vector{…}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/jacobian.jl:19
 [20] top-level scope
    @ REPL[14]:1

It does not depend on the linear solver since KLU also fails.

How can I use sparse Jacobians in the ODE solver while computing sensitivities?

Best,
Neodym

2

You’ll need to switch to SparsepakFactorization(), which it should do automatically?

3

This works! Thanks for the hint :slight_smile:

One more thing: I understand that the matrix element are Duals since I’m autodiff’ing through the solver, right?
I could also opt for sensealg=ForwardSensitivity(), which should allow me to use the Float64 Jacobians since we’re just extending the ODE system with the linearized ones for each parameter. However, this also fails with the same error as in the opening post.
Why is that?

4

Yes. Though that isn’t necessary, we just haven’t added an overload to LinearSolve.jl to handle that better.

That’s unexpected.