Merge pull request #952 from DhairyaLGandhi/dg/nonlinear

Adjoint through `getindex(::NonlinearSolution, sym)`

Author: ChrisRackauckas

ext/SciMLBaseZygoteExt.jl

@@ -116,6 +116,14 @@ function obs_grads(VA, sym, obs_idx, Δ)
     back(Δobs)
 end
 
+function obs_grads2(VA::SciMLBase.NonlinearSolution, sym, obs_idx, Δ)
+    y, back = Zygote.pullback(VA) do sol
+        getindex.(Ref(sol), sym[obs_idx])
+    end
+    Δobs = Δ[obs_idx, :]
+    back(Δobs)
+end
+
 function obs_grads(VA, sym, ::Nothing, Δ)
     Zygote.nt_nothing(VA)
 end
@@ -154,6 +162,31 @@ end
     VA[sym], ODESolution_getindex_pullback
 end
 
+@adjoint function Base.getindex(VA::SciMLBase.NonlinearSolution, sym)
+    function NonlinearSolution_getindex_pullback(Δ)
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
+        if is_observed(VA, sym)
+            f = observed(VA, sym)
+            p = parameter_values(VA)
+            u = state_values(VA)
+            _, back = Zygote.pullback(u, p) do u, p
+                f.f_oop(u, p)
+            end
+            gs = back(Δ)
+            ((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
+            VA = recursivecopy(VA)
+            recursivefill!(VA, zero(eltype(VA)))
+            v = view(VA, i, ntuple(_ -> :, ndims(VA) - 1)...)
+            copyto!(v, Δ)
+            (VA, nothing)
+        end
+    end
+    VA[sym], NonlinearSolution_getindex_pullback
+end
+
 @adjoint function ODESolution{
         T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15}(u,
         args...) where {T1, T2, T3, T4, T5, T6, T7, T8,

test/downstream/observables_autodiff.jl

@@ -46,6 +46,15 @@ sol = solve(prob, Tsit5())
     @test du == gs
 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
+
+# end
+
 # DAE
 
 function create_model(; C₁ = 3e-5, C₂ = 1e-6)
@@ -84,18 +93,29 @@ end
     du_ = [0.2, 1.0]
     du = [du_ for _ in sol.u]
     @test gs == du
+
+    @testset "DAE Initialization Observable function AD" begin
+        iprob = prob.f.initialization_data.initializeprob
+        isol = solve(iprob)
+        tunables, repack, _ = SS.canonicalize(SS.Tunable(), SII.parameter_values(iprob))
+        gs, = gradient(isol) do isol
+            isol[sys.ampermeter.i]
+        end
+        gt = gs.prob.p.tunable
+        @test length(findall(!iszero, gt)) == 1
+    end
 end
 
 # @testset "Adjoints with DAE" begin
-#     gs_mtkp, gs_p_new = gradient(mtkparams, p_new) do p, new_tunables
-#         new_p = SciMLStructures.replace(SciMLStructures.Tunable(), p, new_tunables)
+#     gs_mtkp, gs_p_new = gradient(prob.p, prob.p.tunable) do p, new_tunables
+#         new_p = SS.replace(SS.Tunable(), p, new_tunables)
 #         new_prob = remake(prob, p = new_p)
 #         sol = solve(new_prob, Rodas4())
 #         @show size(sol)
 #         # mean(abs.(sol[sys.ampermeter.i] .- gt))
 #         sum(sol[sys.ampermeter.i])
 #     end
-#
+# 
 #     @test isnothing(gs_mtkp)
 #     @test length(gs_p_new) == length(p_new)
 # end