Merge pull request #979 from DhairyaLGandhi/dg/paramgrad

feat: pass parameter gradients from `getindex`

Author: ChrisRackauckas

ext/SciMLBaseZygoteExt.jl

@@ -93,15 +93,13 @@ end
         if is_observed(VA, sym)
             f = observed(VA, sym)
             p = parameter_values(VA)
-            tunables, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
             u = state_values(VA)
             t = current_time(VA)
-            y, back = Zygote.pullback(u, tunables) do u, tunables
-                _p = repack(tunables)
-                f.(u, Ref(_p), t)
+            y, back = Zygote.pullback(u, p) do u, p
+                f.(u, Ref(p), t)
             end
             gs = back(Δ)
-            (gs[1], nothing)
+            (u = gs[1], prob = (p = gs[2],)), nothing
         elseif i === nothing
             throw(error("Zygote AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))
         else
@@ -163,14 +161,14 @@ end
         gs_obs = obs_grads(VA, sym, isempty(obs_idx) ? nothing : obs_idx, Δ)
         gs_not_obs = not_obs_grads(VA, sym, not_obs_idx, i, Δ)
 
-        a = Zygote.accum(gs_obs[1], gs_not_obs)
+        a = Zygote.accum(gs_obs[1], (u = gs_not_obs,))
 
         (a, nothing)
     end
     VA[sym], ODESolution_getindex_pullback
 end
 
-@adjoint function Base.getindex(VA::SciMLBase.NonlinearSolution, sym)
+@adjoint function Base.getindex(VA::SciMLBase.AbstractNonlinearSolution, sym)
     function NonlinearSolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if is_observed(VA, sym)
@@ -181,7 +179,7 @@ end
                 f.f_oop(u, p)
             end
             gs = back(Δ)
-            ((u = gs[1], prob = (p = gs[2],),), nothing)
+            ((u = gs[1], prob = (p = gs[2],)), nothing)
         elseif i === nothing
             throw(error("Zygote AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))
         elseif i isa Int && VA.u isa Number

test/downstream/adjoints.jl

@@ -52,7 +52,7 @@ idx_vecsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lor
 true_grad_vecsym = zeros(length(ModelingToolkit.unknowns(sys)))
 true_grad_vecsym[idx_vecsym] .= 1.0
 
-@test all(map(x -> x == true_grad_vecsym, gs_vec))
+@test all(map(x -> x == true_grad_vecsym, gs_vec.u))
 
 gs_tup, = Zygote.gradient(sol) do sol
     sum(sum.(collect.(sol[(lorenz1.x, lorenz2.x)])))
@@ -61,7 +61,7 @@ idx_tupsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lor
 true_grad_tupsym = zeros(length(ModelingToolkit.unknowns(sys)))
 true_grad_tupsym[idx_tupsym] .= 1.0
 
-@test all(map(x -> x == true_grad_tupsym, gs_tup))
+@test all(map(x -> x == true_grad_tupsym, gs_tup.u))
 
 gs_ts, = Zygote.gradient(sol) do sol
     sum(sum.(sol[[lorenz1.x, lorenz2.x], :]))

test/downstream/observables_autodiff.jl

@@ -26,7 +26,7 @@ p = [σ => 28.0,
     β => 8 / 3]
 
 tspan = (0.0, 100.0)
-prob = ODEProblem(sys, u0, tspan, p, jac = true)
+prob = ODEProblem(sys, u0, tspan, p)
 sol = solve(prob, Tsit5())
 
 @testset "AutoDiff Observable Functions" begin
@@ -35,25 +35,36 @@ sol = solve(prob, Tsit5())
     end
     du_ = [0.0, 1.0, 1.0, 1.0]
     du = [du_ for _ in sol.u]
-    @test du == gs
+    @test du == gs.u
 
     # Observable in a vector
     gs, = gradient(sol) do sol
         sum(sum.(sol[[sys.w, sys.x]]))
     end
     du_ = [0.0, 1.0, 1.0, 2.0]
     du = [du_ for _ in sol.u]
-    @test du == gs
+    @test du == gs.u
 end
 
-# @testset "AD Observable Functions for Initialization" begin
-#     iprob = prob.f.initialization_data.initalizeprob
-#     isol = solve(iprob)
-#     gs, = gradient(isol) do isol
-#         isol[w]
-#     end
+@testset "AD Observable Functions for Initialization" begin
+    iprob = prob.f.initialization_data.initializeprob
+    isol = solve(iprob)
+    gs, = Zygote.gradient(isol) do isol
+        isol[w]
+    end
 
-# end
+    @test gs isa NamedTuple
+    @test isempty(setdiff(fieldnames(typeof(gs)), fieldnames(typeof(isol))))
+
+    # Compare gradient for parameters match from observed function
+    # to ensure parameter gradients are passed through the observed function
+    f = SII.observed(iprob.f.sys, w)
+    gu0, gp = gradient(SII.state_values(iprob), SII.parameter_values(iprob)) do u0, p
+        f(u0, p)
+    end
+
+    @test gs.prob.p == gp
+end
 
 # DAE
 
@@ -92,7 +103,7 @@ end
     end
     du_ = [0.2, 1.0]
     du = [du_ for _ in sol.u]
-    @test gs == du
+    @test gs.u == du
 
     @testset "DAE Initialization Observable function AD" begin
         iprob = prob.f.initialization_data.initializeprob