Improve SciMLStructures support

Fixes #1233

Author: ChrisRackauckas

src/gauss_adjoint.jl

@@ -421,16 +421,16 @@ function GaussIntegrand(sol, sensealg, checkpoints, dgdp = nothing)
     elseif sensealg.autojacvec isa MooncakeVJP
         pf = get_pf(sensealg.autojacvec, prob, f)
         paramjac_config = get_paramjac_config(
-            MooncakeLoaded(), sensealg.autojacvec, pf, p, f, y, tspan[2])
+            MooncakeLoaded(), sensealg.autojacvec, pf, tunables, f, y, tspan[2])
         pJ = nothing
     elseif isautojacvec # Zygote
         paramjac_config = nothing
         pf = nothing
         pJ = nothing
     else
-        pf = SciMLBase.ParamJacobianWrapper(unwrappedf, tspan[1], y)
+        pf = SciMLBase.ParamJacobianWrapper((du,u,p,t)->unwrappedf(du,u,repack(p),t), tspan[1], y)
         pJ = similar(u0, length(u0), numparams)
-        paramjac_config = build_param_jac_config(sensealg, pf, y, p)
+        paramjac_config = build_param_jac_config(sensealg, pf, y, tunables)
     end
 
     cpsol = sol
@@ -460,7 +460,7 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
         else
             pf.t = t
             pf.u = y
-            jacobian!(pJ, pf, p, f_cache, sensealg, paramjac_config)
+            jacobian!(pJ, pf, tunables, f_cache, sensealg, paramjac_config)
         end
         mul!(out', λ', pJ)
     elseif sensealg.autojacvec isa ReverseDiffVJP
@@ -610,7 +610,7 @@ function _adjoint_sensitivities(sol, sensealg::GaussAdjoint, alg; t = nothing,
             res .+= out
             iλ .= zero(eltype(iλ))
         end
-    end
+    end 
 
     return state_values(adj_sol)[end], __maybe_adjoint(res)
 end

test/scimlstructures_interface.jl

@@ -78,3 +78,87 @@ Zygote.gradient(0.1ones(length(SS.canonicalize(SS.Tunable(), p)[1]))) do tunable
     sol = solve(newprob, Tsit5())
     return sum(sol.u[end])
 end
+
+using OrdinaryDiffEq
+using StableRNGs, Lux
+using ComponentArrays
+using SciMLSensitivity
+import SciMLStructures as SS
+using Zygote
+using Parameters
+using ADTypes
+using Test
+
+mutable struct myparam{M,P,S}
+    model::M
+    ps ::P 
+    st ::S
+    α :: Float64
+    β :: Float64
+    γ :: Float64
+end
+
+SS.isscimlstructure(::myparam) = true
+SS.ismutablescimlstructure(::myparam) = true
+SS.hasportion(::SS.Tunable, ::myparam) = true
+function SS.canonicalize(::SS.Tunable, p::myparam)
+    buffer = copy(p.ps)
+    repack = let p = p
+        function repack(newbuffer)
+            SS.replace(SS.Tunable(), p, newbuffer)
+        end
+    end
+    return buffer, repack, false
+end
+function SS.replace(::SS.Tunable, p::myparam, newbuffer)
+    return myparam(p.model, newbuffer, p.st, p.α, p.β, p.γ)
+end
+function SS.replace!(::SS.Tunable, p::myparam, newbuffer)
+    p.ps = newbuffer
+    return p
+end
+function initialize()
+    # Defining the neural network
+    U = Lux.Chain(Lux.Dense(3,30,tanh),Lux.Dense(30,30,tanh),Lux.Dense(30,1))
+    rng = StableRNG(1111)
+    _para,st = Lux.setup(rng,U)
+    _para = ComponentArray(_para)
+    # Setting the parameters
+    α = 0.5
+    β = 0.1
+    γ = 0.01
+    return myparam(U,_para,st,α,β,γ)
+end
+function UDE_model!(du, u, p, t)
+    # Extracting parameters
+    Parameters.@unpack model, ps, st, α, β, γ = p
+    o = model(u,ps, st)[1][1]
+    du[1] = o * α * u[1] + β * u[2] + γ * u[3]
+    du[2] = -α * u[1] + β * u[2] - γ * u[3]
+    du[3] = α * u[1] - β * u[2] + γ * u[3]
+    nothing
+end
+   
+p = initialize()
+function run_diff(ps)
+    u01 = [1.0, 0.0, 0.0]
+    tspan = (0.0, 10.0)
+    prob = ODEProblem(UDE_model!, u01, tspan, ps)
+    sol = solve(prob, Rosenbrock23(), saveat = 0.1)
+    return sol.u |> last |> sum
+end
+
+run_diff(initialize())
+@test !iszero(Zygote.gradient(run_diff, initialize())[1].ps)
+
+function run_diff(ps,sensealg)
+    u01 = [1.0, 0.0, 0.0]
+    tspan = (0.0, 10.0)
+    prob = ODEProblem(UDE_model!, u01, tspan, ps)
+    sol = solve(prob, Rosenbrock23(), saveat = 0.1, sensealg=sensealg)
+    return sol.u |> last |> sum
+end
+
+run_diff(initialize())
+@test !iszero(Zygote.gradient(run_diff, initialize(), GaussAdjoint())[1].ps)
+@test !iszero(Zygote.gradient(run_diff, initialize(), GaussAdjoint(autojacvec=false))[1].ps)