DiffEqArrayOperator for sensitivity problems

[ArnoStrouwen]

Is it worthwhile to set up sensitivity problems with DiffEqArrayOperator ? I made a forward prototype that does not exploit the structure of A, i.e. it is a dense matrix. There is also AffineDiffEqOperator, but there is no example in the docs, so I could not get this to work. Also interesting for Adjoint mode.

using OrdinaryDiffEq, DiffEqSensitivity, ForwardDiff
using Plots, BenchmarkTools
function fiip(du,u,p,t)
    du[1] = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = -p[3]*u[2] + p[4]*u[1]*u[2]
end
u0 = [1.0;1.0]
p = [1.5,1.0,3.0,1.0]
n_u = length(u0)
n_p = length(p)
prob = ODEForwardSensitivityProblem(fiip,u0,(0.0,10.0),p)
sol = solve(prob,Tsit5(),dt=0.1)
@btime solve(prob,Tsit5()) # 99.002 μs (488 allocations: 78.42 KiB)
plot(sol)

function update_func(A,u,p,t,cfgJu,cfgJp,uf,pf)
    n_p = length(p)
    n_u = (length(u)-1)÷(n_p+1)  
    uf.t = t
    pf.t = t
    copyto!(pf.u,@view(u[1:n_u]))
    ForwardDiff.jacobian!(@view(A[1+n_u:2*n_u,1+n_u:2*n_u]),uf, @view(A[1:n_u,end]),@view(u[1:n_u]),cfgJu)
    for i = 2:n_p
        A[1+i*n_u:(i+1)*n_u,1+i*n_u:(i+1)*n_u] .= @view(A[1+n_u:2*n_u,1+n_u:2*n_u])
    end
    ForwardDiff.jacobian!(reshape(@view(A[1 + n_u:end-1,end]),n_u*n_p), pf,@view(A[1:n_u,end]),p,cfgJp)
    return nothing
end
A = zeros(n_u + n_u*n_p + 1,n_u + n_u*n_p + 1)
u0_ = vcat(u0,zeros(n_p*n_u),1.0)

const uf = prob.f.uf
const pf = prob.f.pf
const cfgJu = ForwardDiff.JacobianConfig(uf,A[1:n_u,end], u0)
const cfgJp = ForwardDiff.JacobianConfig(pf,A[1:n_u,end], p)

A = DiffEqArrayOperator(A,update_func=(A,u,p,t)->update_func(A,u,p,t,cfgJu,cfgJp,uf,pf))
@btime A.update_func(A.A,u0_,p,1.0) # 869.729 ns (10 allocations: 608 bytes)

prob = ODEProblem(A,u0_,(0.0,10.0),p)
sol = solve(prob,LieRK4(),dt=0.1)
plot(sol)
@btime solve(prob,LieRK4(),dt=0.1) # 14.277 ms (23312 allocations: 16.04 MiB)

[ArnoStrouwen]

By defining n_p additional constant states and ordering them in smart way I think A has the structure of a block diagonal matrix. Do block diagonal matrices have good properties for these types of solvers? E.g. can the matrix exponential be calculated from matrix exponential of each block?

using BlockDiagonals
function update_func(A,u,p,t,cfgJu,cfgJp,uf,pf,cachePj)
    n_p = length(p)
    n_u = (length(u)-1)÷(n_p+1)  
    uf.t = t
    pf.t = t
    copyto!(pf.u,@view(u[2:n_u+1]))
    ForwardDiff.jacobian!(@view(A.blocks[2][2:end,2:end]),uf, @view(A[2:n_u+1,1]),@view(u[2:n_u+1]),cfgJu)
    for i = 3:n_p+1
        @view(A.blocks[i][2:end,2:end]) .= @view(A.blocks[2][2:end,2:end])
    end
    ForwardDiff.jacobian!(cachePj, pf, @view(A[2:n_u+1,1]),p,cfgJp)
    for i = 2:n_p+1
        @view(A.blocks[i][2:end,1]) .= @view(cachePj[:,i-1])
    end
    return nothing
end
A = [zeros(n_u+1,n_u+1) for i in 1:n_p+1]
A = BlockDiagonal(A)
u0_ = vcat(1.0,u0,repeat(vcat(1.0,zeros(n_u)),n_p))
const cachePj = zeros(n_u,n_p)

A = DiffEqArrayOperator(A,update_func=(A,u,p,t)->update_func(A,u,p,t,cfgJu,cfgJp,uf,pf,cachePj))
@btime A.update_func(A.A,u0_,p,1.0) # 2.326 μs (48 allocations: 4.52 KiB)

prob = ODEProblem(A,u0_,(0.0,10.0),p)
sol = solve(prob,LieRK4(),dt=0.1)
plot(sol)
@btime solve(prob,LieRK4(),dt=0.1) # 32.233 ms (122677 allocations: 37.77 MiB)

[ChrisRackauckas]

E.g. can the matrix exponential be calculated from matrix exponential of each block?

I think ExponentialUtilities.jl will need a specialization there.

[ArnoStrouwen]

That won't be enough because of stuff like: https://github.com/SciML/OrdinaryDiffEq.jl/blob/cdc5ef291d37f31ce7746e9f27aeeeae5d806393/src/perform_step/linear_perform_step.jl#L50 The solvers do not seem to use the AbstractDiffEqOperator interface? From docs: "Note that many other AbstractDiffEqOperators can be used and DiffEqArrayOperator is just one version that represents A via a matrix (other choices are matrix-free)." https://diffeq.sciml.ai/latest/types/nonautonomous_linear_ode/

[ChrisRackauckas]

Those methods need more optimizations. Right now they concretize because we would need a better way of building a Krylov-subspace for the commutators, or commutators to be defined. I don't think that's a great direction for general equations without more information about the Lie groups.