Out of place QuadratureAdjoint for Working with StaticArrays

Project.toml

@@ -34,6 +34,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

src/SciMLSensitivity.jl

@@ -14,6 +14,7 @@ import ZygoteRules, Zygote, ReverseDiff
 import ArrayInterfaceCore, ArrayInterfaceTracker
 import Enzyme
 import GPUArraysCore
+using StaticArrays
 
 import PreallocationTools: dualcache, get_tmp, DiffCache
 
@@ -24,7 +25,8 @@ using EllipsisNotation
 using Markdown
 
 using Reexport
-import ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented
+import ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented, Tangent, ProjectTo,
+                       project_type, _eltype_projectto, rrule
 abstract type SensitivityFunction end
 abstract type TransformedFunction end
 
@@ -45,6 +47,7 @@ include("concrete_solve.jl")
 include("second_order.jl")
 include("steadystate_adjoint.jl")
 include("sde_tools.jl")
+include("staticarrays.jl")
 
 # AD Extensions
 include("reversediff.jl")

src/adjoint_common.jl

@@ -401,7 +401,7 @@ inplace_sensitivity(S::SensitivityFunction) = isinplace(getprob(S))
 
 struct ReverseLossCallback{λType, timeType, yType, RefType, FMType, AlgType, dg1Type,
                            dg2Type,
-                           cacheType}
+                           cacheType, solType}
     isq::Bool
     λ::λType
     t::timeType
@@ -413,6 +413,7 @@ struct ReverseLossCallback{λType, timeType, yType, RefType, FMType, AlgType, dg
     dgdu::dg1Type
     dgdp::dg2Type
     diffcache::cacheType
+    sol::solType
 end
 
 function ReverseLossCallback(sensefun, λ, t, dgdu, dgdp, cur_time)
@@ -422,13 +423,17 @@ function ReverseLossCallback(sensefun, λ, t, dgdu, dgdp, cur_time)
     @unpack factorized_mass_matrix = sensefun.diffcache
     prob = getprob(sensefun)
     idx = length(prob.u0)
-
-    return ReverseLossCallback(isq, λ, t, y, cur_time, idx, factorized_mass_matrix,
-                               sensealg, dgdu, dgdp, sensefun.diffcache)
+    if ArrayInterfaceCore.ismutable(y)
+        return ReverseLossCallback(isq, λ, t, y, cur_time, idx, factorized_mass_matrix,
+                                   sensealg, dgdu, dgdp, sensefun.diffcache, nothing)
+    else
+        return ReverseLossCallback(isq, λ, t, y, cur_time, idx, factorized_mass_matrix,
+                                   sensealg, dgdu, dgdp, sensefun.diffcache, sensefun.sol)
+    end
 end
 
 function (f::ReverseLossCallback)(integrator)
-    @unpack isq, λ, t, y, cur_time, idx, F, sensealg, dgdu, dgdp = f
+    @unpack isq, λ, t, y, cur_time, idx, F, sensealg, dgdu, dgdp, sol = f
     @unpack diffvar_idxs, algevar_idxs, issemiexplicitdae, J, uf, f_cache, jac_config = f.diffcache
 
     p, u = integrator.p, integrator.u
@@ -437,16 +442,23 @@ function (f::ReverseLossCallback)(integrator)
         copyto!(y, integrator.u[(end - idx + 1):end])
     end
 
-    # Warning: alias here! Be careful with λ
-    gᵤ = isq ? λ : @view(λ[1:idx])
-    if dgdu !== nothing
-        dgdu(gᵤ, y, p, t[cur_time[]], cur_time[])
-        # add discrete dgdp contribution
-        if dgdp !== nothing && !isq
-            gp = @view(λ[(idx + 1):end])
-            dgdp(gp, y, p, t[cur_time[]], cur_time[])
-            u[(idx + 1):length(λ)] .+= gp
+    if ArrayInterfaceCore.ismutable(u)
+        # Warning: alias here! Be careful with λ
+        gᵤ = isq ? λ : @view(λ[1:idx])
+        if dgdu !== nothing
+            dgdu(gᵤ, y, p, t[cur_time[]], cur_time[])
+            # add discrete dgdp contribution
+            if dgdp !== nothing && !isq
+                gp = @view(λ[(idx + 1):end])
+                dgdp(gp, y, p, t[cur_time[]], cur_time[])
+                u[(idx + 1):length(λ)] .+= gp
+            end
         end
+    else
+        @assert sensealg isa QuadratureAdjoint
+        outtype = DiffEqBase.parameterless_type(λ)
+        y = sol(t[cur_time[]])
+        gᵤ = dgdu(y, p, t[cur_time[]], cur_time[]; outtype = outtype)
     end
 
     if issemiexplicitdae
@@ -468,7 +480,12 @@ function (f::ReverseLossCallback)(integrator)
         F !== I && F !== (I, I) && ldiv!(F, Δλd)
     end
 
-    u[diffvar_idxs] .+= Δλd
+    if ArrayInterfaceCore.ismutable(u)
+        u[diffvar_idxs] .+= Δλd
+    else
+        @assert sensealg isa QuadratureAdjoint
+        integrator.u += Δλd
+    end
     u_modified!(integrator, true)
     cur_time[] -= 1
     return nothing

src/concrete_solve.jl

@@ -343,7 +343,7 @@ function DiffEqBase._concrete_solve_adjoint(prob::Union{SciMLBase.AbstractODEPro
     _save_idxs = save_idxs === nothing ? Colon() : save_idxs
 
     function adjoint_sensitivity_backpass(Δ)
-        function df(_out, u, p, t, i)
+        function df_iip(_out, u, p, t, i)
             outtype = typeof(_out) <: SubArray ?
                       DiffEqBase.parameterless_type(_out.parent) :
                       DiffEqBase.parameterless_type(_out)
@@ -404,16 +404,82 @@ function DiffEqBase._concrete_solve_adjoint(prob::Union{SciMLBase.AbstractODEPro
             end
         end
 
+        function df_oop(u, p, t, i; outtype = nothing)
+            if only_end
+                eltype(Δ) <: NoTangent && return
+                if typeof(Δ) <: AbstractArray{<:AbstractArray} && length(Δ) == 1 && i == 1
+                    # user did sol[end] on only_end
+                    if typeof(_save_idxs) <: Number
+                        x = vec(Δ[1])
+                        _out = adapt(outtype, @view(x[_save_idxs]))
+                    elseif _save_idxs isa Colon
+                        _out = adapt(outtype, vec(Δ[1]))
+                    else
+                        _out = adapt(outtype,
+                                     vec(Δ[1])[_save_idxs])
+                    end
+                else
+                    Δ isa NoTangent && return
+                    if typeof(_save_idxs) <: Number
+                        x = vec(Δ)
+                        _out = adapt(outtype, @view(x[_save_idxs]))
+                    elseif _save_idxs isa Colon
+                        _out = adapt(outtype, vec(Δ))
+                    else
+                        x = vec(Δ)
+                        _out = adapt(outtype, @view(x[_save_idxs]))
+                    end
+                end
+            else
+                !Base.isconcretetype(eltype(Δ)) &&
+                    (Δ[i] isa NoTangent || eltype(Δ) <: NoTangent) && return
+                if typeof(Δ) <: AbstractArray{<:AbstractArray} || typeof(Δ) <: DESolution
+                    x = Δ[i]
+                    if typeof(_save_idxs) <: Number
+                        _out = @view(x[_save_idxs])
+                    elseif _save_idxs isa Colon
+                        _out = vec(x)
+                    else
+                        _out = vec(@view(x[_save_idxs]))
+                    end
+                else
+                    if typeof(_save_idxs) <: Number
+                        _out = adapt(outtype,
+                                     reshape(Δ, prod(size(Δ)[1:(end - 1)]),
+                                             size(Δ)[end])[_save_idxs, i])
+                    elseif _save_idxs isa Colon
+                        _out = vec(adapt(outtype,
+                                         reshape(Δ, prod(size(Δ)[1:(end - 1)]),
+                                                 size(Δ)[end])[:, i]))
+                    else
+                        _out = vec(adapt(outtype,
+                                         reshape(Δ,
+                                                 prod(size(Δ)[1:(end - 1)]),
+                                                 size(Δ)[end])[:, i]))
+                    end
+                end
+            end
+            return _out
+        end
+
         if haskey(kwargs_adj, :callback_adj)
             cb2 = CallbackSet(cb, kwargs[:callback_adj])
         else
             cb2 = cb
         end
-
-        du0, dp = adjoint_sensitivities(sol, alg, args...; t = ts, dgdu_discrete = df,
-                                        sensealg = sensealg,
-                                        callback = cb2,
-                                        kwargs_adj...)
+        if ArrayInterfaceCore.ismutable(eltype(sol.u))
+            du0, dp = adjoint_sensitivities(sol, alg, args...; t = ts,
+                                            dgdu_discrete = df_iip,
+                                            sensealg = sensealg,
+                                            callback = cb2,
+                                            kwargs_adj...)
+        else
+            du0, dp = adjoint_sensitivities(sol, alg, args...; t = ts,
+                                            dgdu_discrete = df_oop,
+                                            sensealg = sensealg,
+                                            callback = cb2,
+                                            kwargs_adj...)
+        end
 
         du0 = reshape(du0, size(u0))
         dp = p === nothing || p === DiffEqBase.NullParameters() ? nothing :

src/derivative_wrappers.jl

@@ -222,6 +222,12 @@ function vecjacobian!(dλ, y, λ, p, t, S::TS;
     return
 end
 
+function vecjacobian(y, λ, p, t, S::TS;
+                     dgrad = nothing, dy = nothing,
+                     W = nothing) where {TS <: SensitivityFunction}
+    return _vecjacobian(y, λ, p, t, S, S.sensealg.autojacvec, dgrad, dy, W)
+end
+
 function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::Bool, dgrad, dy,
                        W) where {TS <: SensitivityFunction}
     @unpack sensealg, f = S
@@ -588,6 +594,43 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::ZygoteVJP, dgrad,
     return
 end
 
+function _vecjacobian(y, λ, p, t, S::TS, isautojacvec::ZygoteVJP, dgrad, dy,
+                      W) where {TS <: SensitivityFunction}
+    @unpack sensealg, f = S
+    prob = getprob(S)
+
+    isautojacvec = get_jacvec(sensealg)
+
+    if W === nothing
+        _dy, back = Zygote.pullback(y, p) do u, p
+            vec(f(u, p, t))
+        end
+    else
+        _dy, back = Zygote.pullback(y, p) do u, p
+            vec(f(u, p, t, W))
+        end
+    end
+
+    # Grab values from `_dy` before `back` in case mutated
+    dy !== nothing && (dy[:] .= vec(_dy))
+
+    tmp1, tmp2 = back(λ)
+    if tmp1 === nothing && !sensealg.autojacvec.allow_nothing
+        throw(ZygoteVJPNothingError())
+    elseif tmp1 !== nothing
+        (dλ = vec(tmp1))
+    end
+
+    if dgrad !== nothing
+        if tmp2 === nothing && !sensealg.autojacvec.allow_nothing
+            throw(ZygoteVJPNothingError())
+        elseif tmp2 !== nothing
+            (dgrad[:] .= vec(tmp2))
+        end
+    end
+    return dy, dλ, dgrad
+end
+
 function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::EnzymeVJP, dgrad, dy,
                        W) where {TS <: SensitivityFunction}
     @unpack sensealg = S
@@ -923,6 +966,19 @@ function accumulate_cost!(dλ, y, p, t, S::TS,
     return nothing
 end
 
+function accumulate_cost(dλ, y, p, t, S::TS,
+                         dgrad = nothing) where {TS <: SensitivityFunction}
+    @unpack dgdu, dgdp = S.diffcache
+
+    dλ -= dgdu(y, p, t)
+    if dgdp !== nothing
+        if dgrad !== nothing
+            dgrad -= dgdp(y, p, t)
+        end
+    end
+    return dλ, dgrad
+end
+
 function build_jac_config(alg, uf, u)
     if alg_autodiff(alg)
         jac_config = ForwardDiff.JacobianConfig(uf, u, u,