Feat: Handle Adjoints through Initialization

.github/workflows/CI.yml

@@ -24,6 +24,7 @@ jobs:
           - Core5
           - Core6
           - Core7
+          - Core8
           - QA
           - SDE1
           - SDE2

src/adjoint_common.jl

@@ -379,11 +379,11 @@ function get_paramjac_config(autojacvec::ReverseDiffVJP, p, f, y, _p, _t;
         if !isRODE
             __p = p isa SciMLBase.NullParameters ? _p :
                   SciMLStructures.replace(Tunable(), p, _p)
-            tape = ReverseDiff.GradientTape((y, __p, [_t])) do u, p, t
+            tape = ReverseDiff.GradientTape((y, _p, [_t])) do u, p, t
                 du1 = (p !== nothing && p !== SciMLBase.NullParameters()) ?
                       similar(p, size(u)) : similar(u)
                 du1 .= false
-                f(du1, u, p, first(t))
+                f(du1, u, repack(p), first(t))
                 return vec(du1)
             end
         else

src/concrete_solve.jl

@@ -299,6 +299,7 @@ function DiffEqBase._concrete_solve_adjoint(
         tunables, repack = Functors.functor(p)
     end
 
+    u0 = state_values(prob) === nothing ? Float64[] : u0
     default_sensealg = automatic_sensealg_choice(prob, u0, tunables, verbose, repack)
     DiffEqBase._concrete_solve_adjoint(prob, alg, default_sensealg, u0, p,
         originator::SciMLBase.ADOriginator, args...; verbose,
@@ -425,6 +426,25 @@ function DiffEqBase._concrete_solve_adjoint(
             save_end = true, kwargs_fwd...)
     end
 
+    # Get gradients for the initialization problem if it exists
+    igs = if _prob.f.initialization_data != nothing
+        Zygote.gradient(tunables) do tunables
+            new_prob = remake(_prob, p = repack(tunables))
+            new_u0, new_p, _ = SciMLBase.get_initial_values(new_prob, new_prob, new_prob.f, SciMLBase.OverrideInit(), Val(true);
+                                                            abstol = 1e-6,
+                                                            reltol = 1e-6,
+                                                            sensealg = SteadyStateAdjoint(autojacvec = ZygoteVJP()))
+            new_tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), new_p)
+            if SciMLBase.initialization_status(_prob) == SciMLBase.OVERDETERMINED
+                sum(new_tunables)
+            else
+                sum(new_u0) + sum(new_tunables)
+            end
+        end[1] .- one(eltype(tunables))
+    else
+        nothing
+    end
+
     # Force `save_start` and `save_end` in the forward pass This forces the
     # solver to do the backsolve all the way back to `u0` Since the start aliases
     # `_prob.u0`, this doesn't actually use more memory But it cleans up the
@@ -642,6 +662,8 @@ function DiffEqBase._concrete_solve_adjoint(
         dp = p === nothing || p === DiffEqBase.NullParameters() ? nothing :
              dp isa AbstractArray ? reshape(dp', size(tunables)) : dp
 
+        dp = Zygote.accum(dp, igs)
+
         _, repack_adjoint = if p === nothing || p === DiffEqBase.NullParameters() ||
                                !isscimlstructure(p)
             nothing, x -> (x,)
@@ -1686,6 +1708,25 @@ function DiffEqBase._concrete_solve_adjoint(
         out = SciMLBase.sensitivity_solution(sol, sol[_save_idxs])
     end
 
+    # Get gradients for the initialization problem if it exists
+    igs = if _prob.f.initialization_data != nothing
+        Zygote.gradient(tunables) do tunables
+            new_prob = remake(_prob, p = repack(tunables))
+            new_u0, new_p, _ = SciMLBase.get_initial_values(new_prob, new_prob, new_prob.f, SciMLBase.OverrideInit(), Val(true);
+                                                            abstol = 1e-6,
+                                                            reltol = 1e-6,
+                                                            sensealg = SteadyStateAdjoint(autojacvec = ZygoteVJP()))
+            new_tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), new_p)
+            if SciMLBase.initialization_status(_prob) == SciMLBase.OVERDETERMINED
+                sum(new_tunables)
+            else
+                sum(new_u0) + sum(new_tunables)
+            end
+        end[1] .- one(eltype(tunables))
+    else
+        nothing
+    end
+
     function steadystatebackpass(Δ)
         # Δ = dg/dx or diffcache.dg_val
         # del g/del p = 0
@@ -1694,11 +1735,15 @@ function DiffEqBase._concrete_solve_adjoint(
                 _out[_save_idxs] = Δ[_save_idxs]
             elseif Δ isa Number
                 @. _out[_save_idxs] = Δ
-            else
+            elseif Δ isa AbstractArray{<:AbstractArray} || Δ isa AbstractVectorOfArray
                 @. _out[_save_idxs] = Δ[_save_idxs]
+            else
+                @. _out[_save_idxs] = Δ.u[_save_idxs]
             end
         end
+        # dp = adjoint_sensitivities(sol, alg; sensealg = sensealg, dgdu = df, dgdp = dp)
         dp = adjoint_sensitivities(sol, alg; sensealg = sensealg, dgdu = df)
+        dp = Zygote.accum(dp, SciMLStructures.replace(SciMLStructures.Tunable(), dp, igs))
 
         if originator isa SciMLBase.TrackerOriginator ||
            originator isa SciMLBase.ReverseDiffOriginator

src/derivative_wrappers.jl

@@ -456,9 +456,10 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::ReverseDiffVJP, dg
     elseif inplace_sensitivity(S)
         _y = eltype(y) === eltype(λ) ? y : convert.(promote_type(eltype(y), eltype(λ)), y)
         if W === nothing
-            tape = ReverseDiff.GradientTape((_y, _p, [t])) do u, p, t
+            _tunables, _repack, _ = canonicalize(Tunable(), _p)
+            tape = ReverseDiff.GradientTape((_y, _tunables, [t])) do u, p, t
                 du1 = similar(u, size(u))
-                f(du1, u, p, first(t))
+                f(du1, u, _repack(p), first(t))
                 return vec(du1)
             end
         else
@@ -1068,9 +1069,10 @@ end
 
 function build_param_jac_config(alg, pf, u, p)
     if alg_autodiff(alg)
-        jac_config = ForwardDiff.JacobianConfig(pf, u, p,
+        tunables, repack, aliases = canonicalize(Tunable(), p)
+        jac_config = ForwardDiff.JacobianConfig(pf, u, tunables,
             ForwardDiff.Chunk{
-                determine_chunksize(p,
+                determine_chunksize(tunables,
                 alg)}())
     else
         if diff_type(alg) != Val{:complex}

src/steadystate_adjoint.jl

@@ -50,13 +50,12 @@ end
     sense = SteadyStateAdjointSensitivityFunction(g, sensealg, alg, sol, dgdu, dgdp,
         f, f.colorvec, needs_jac)
     (; diffcache, y, sol, λ, vjp, linsolve) = sense
-
     if needs_jac
         if SciMLBase.has_jac(f)
             f.jac(diffcache.J, y, p, nothing)
         else
             if DiffEqBase.isinplace(sol.prob)
-                jacobian!(diffcache.J, diffcache.uf, y, diffcache.f_cache,
+                jacobian!(diffcache.J.du, diffcache.uf, y, diffcache.f_cache,
                     sensealg, diffcache.jac_config)
             else
                 diffcache.J .= jacobian(diffcache.uf, y, sensealg)
@@ -101,17 +100,14 @@ end
         linear_problem = LinearProblem(soperator, vec(dgdu_val); u0 = vec(λ))
         solve(linear_problem, linsolve; alias = LinearAliasSpecifier(alias_A = true), sensealg.linsolve_kwargs...)
     else
-        if linsolve === nothing && isempty(sensealg.linsolve_kwargs)
-            # For the default case use `\` to avoid any form of unnecessary cache allocation
-            vec(λ) .= diffcache.J' \ vec(dgdu_val)
-        else
-            linear_problem = LinearProblem(diffcache.J', vec(dgdu_val'); u0 = vec(λ))
-            solve(linear_problem, linsolve; alias = LinearAliasSpecifier(alias_A = true), sensealg.linsolve_kwargs...) # u is vec(λ)
-        end
+        linear_problem = LinearProblem(diffcache.J.du', vec(dgdu_val'); u0 = vec(λ))
+        solve(linear_problem, linsolve; alias = LinearAliasSpecifier(alias_A = true), sensealg.linsolve_kwargs...) # u is vec(λ)
     end
 
     try
-        vecjacobian!(vec(dgdu_val), y, λ, p, nothing, sense; dgrad = vjp, dy = nothing)
+        tunables, repack, aliases = canonicalize(Tunable(), p)
+        vjp_tunables, vjp_repack, vjp_aliases = canonicalize(Tunable(), vjp)
+        vecjacobian!(vec(dgdu_val), y, λ, tunables, nothing, sense; dgrad = vjp_tunables, dy = nothing)
     catch e
         if sense.sensealg.autojacvec === nothing
             @warn "Automatic AD choice of autojacvec failed in nonlinear solve adjoint, failing back to ODE adjoint + numerical vjp"

test/desauty_dae_mwe.jl

@@ -0,0 +1,64 @@
+using ModelingToolkit, OrdinaryDiffEq
+using ModelingToolkitStandardLibrary.Electrical
+using ModelingToolkitStandardLibrary.Blocks: Sine
+using NonlinearSolve
+import SciMLStructures as SS
+import SciMLSensitivity
+using Zygote
+
+function create_model(; C₁ = 3e-5, C₂ = 1e-6)
+    @variables t
+    @named resistor1 = Resistor(R = 5.0)
+    @named resistor2 = Resistor(R = 2.0)
+    @named capacitor1 = Capacitor(C = C₁)
+    @named capacitor2 = Capacitor(C = C₂)
+    @named source = Voltage()
+    @named input_signal = Sine(frequency = 100.0)
+    @named ground = Ground()
+    @named ampermeter = CurrentSensor()
+
+    eqs = [connect(input_signal.output, source.V)
+        connect(source.p, capacitor1.n, capacitor2.n)
+        connect(source.n, resistor1.p, resistor2.p, ground.g)
+        connect(resistor1.n, capacitor1.p, ampermeter.n)
+        connect(resistor2.n, capacitor2.p, ampermeter.p)]
+
+    @named circuit_model = ODESystem(eqs, t,
+        systems = [
+            resistor1, resistor2, capacitor1, capacitor2,
+            source, input_signal, ground, ampermeter,
+        ])
+end
+
+desauty_model = create_model()
+sys = structural_simplify(desauty_model)
+
+
+prob = ODEProblem(sys, [], (0.0, 0.1), guesses = [sys.resistor1.v => 1.])
+iprob = prob.f.initialization_data.initializeprob
+isys = iprob.f.sys
+
+tunables, repack, aliases = SS.canonicalize(SS.Tunable(), parameter_values(iprob))
+
+linsolve = LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization)
+sensealg = SciMLSensitivity.SteadyStateAdjoint(autojacvec = SciMLSensitivity.ZygoteVJP(), linsolve = linsolve)
+igs, = Zygote.gradient(tunables) do p
+    iprob2 = remake(iprob, p = repack(p))
+    sol = solve(iprob2,
+                sensealg = sensealg
+    )
+    sum(Array(sol))
+end
+
+@test !iszero(sum(igs))
+
+
+# tunable_parameters(isys) .=> gs
+
+# gradient_unk1_idx = only(findfirst(x -> isequal(x, Initial(sys.capacitor1.v)), tunable_parameters(isys)))
+
+# gs[gradient_unk1_idx]
+
+# prob.f.initialization_data.update_initializeprob!(iprob, prob)
+# prob.f.initialization_data.update_initializeprob!(iprob, ::Vector)
+# prob.f.initialization_data.update_initializeprob!(iprob, gs)

test/runtests.jl

@@ -102,6 +102,12 @@ end
         end
     end
 
+    if GROUP == "All" || GROUP == "Core8"
+        @testset "Core 8" begin
+            @time @safetestset "Initialization with MTK" include("desauty_dae_mwe.jl")
+        end
+    end
+
     if GROUP == "All" || GROUP == "QA"
         @time @safetestset "Quality Assurance" include("aqua.jl")
     end