Feat: adjoints through observable functions

ext/SciMLBaseZygoteExt.jl

@@ -7,7 +7,7 @@
 using SciMLBase: ODESolution, remake,
                  getobserved, build_solution, EnsembleSolution,
                  NonlinearSolution, AbstractTimeseriesSolution
-using SymbolicIndexingInterface: symbolic_type, NotSymbolic, variable_index
+using SymbolicIndexingInterface: symbolic_type, NotSymbolic, variable_index, is_observed, observed, parameter_values
 using RecursiveArrayTools
 
 # This method resolves the ambiguity with the pullback defined in
@@ -109,7 +109,15 @@
 @adjoint function Base.getindex(VA::ODESolution, sym)
     function ODESolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
-        if i === nothing
+        if is_observed(VA, sym)
+            y, back = Zygote.pullback(VA) do sol
+                f = observed(sol, sym)
+                p = parameter_values(sol)
+                f.(sol.u,Ref(p), sol.t)
+            end
+            gs = back(Δ)
+            (gs[1], 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
             Δ′ = [[i == k ? Δ[j] : zero(x[1]) for k in 1:length(x)]
@@ -120,26 +128,49 @@
     VA[sym], ODESolution_getindex_pullback
 end
 
+function obs_grads(VA, sym, obss_idx, Δ)
+    y, back = Zygote.pullback(VA) do sol
+        getindex.(Ref(sol), sym[obss_idx])
+    end
+    Dprime = reduce(hcat, Δ)
+    Dobss = eachrow(Dprime[obss_idx, :])    
+    back(Dobss)
+end
+
+function obs_grads(VA, sym, ::Nothing, Δ)
+    Zygote.nt_nothing(VA)
+end
+
+function not_obs_grads(VA::ODESolution{T}, sym, not_obss_idx, i, Δ) where T
+    Δ′ = map(enumerate(VA.u)) do (t_idx, us)
+        map(enumerate(us)) do (u_idx, u)
+            if u_idx in i
+                idx = findfirst(isequal(u_idx), i)
+                Δ[t_idx][idx]
+            else
+                zero(T)
+            end
+        end
+    end
+
+    nt = Zygote.nt_nothing(VA)
+    Zygote.accum(nt, (u = Δ′,))
+end
+
 @adjoint function Base.getindex(
         VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}) where {T}
     function ODESolution_getindex_pullback(Δ)
         sym = sym isa Tuple ? collect(sym) : sym
         i = map(x -> symbolic_type(x) != NotSymbolic() ? variable_index(VA, x) : x, sym)
-        if 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
-            Δ′ = map(enumerate(VA.u)) do (t_idx, us)
-                map(enumerate(us)) do (u_idx, u)
-                    if u_idx in i
-                        idx = findfirst(isequal(u_idx), i)
-                        Δ[t_idx][idx]
-                    else
-                        zero(T)
-                    end
-                end
-            end
-            (Δ′, nothing)
-        end
+
+        obss_idx = findall(s -> is_observed(VA, s), sym)
+        not_obss_idx = setdiff(1:length(sym), obss_idx)
+
+        gs_obs = obs_grads(VA, sym, isempty(obss_idx) ? nothing : obss_idx, Δ)
+        gs_not_obs = not_obs_grads(VA, sym, not_obss_idx, i, Δ)
+
+        a = Zygote.accum(gs_obs[1], gs_not_obs)
+        (a, nothing)
     end
     VA[sym], ODESolution_getindex_pullback
 end
@@ -183,15 +214,15 @@
     NonlinearSolution{T, N, uType, R, P, A, O, uType2}(u, args...), NonlinearSolutionAdjoint
 end
 
-@adjoint function literal_getproperty(sol::AbstractTimeseriesSolution,
-        ::Val{:u})
-    function solu_adjoint(Δ)
-        zerou = zero(sol.prob.u0)
-        _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
-        (build_solution(sol.prob, sol.alg, sol.t, _Δ),)
-    end
-    sol.u, solu_adjoint
-end
+# @adjoint function literal_getproperty(sol::AbstractTimeseriesSolution,
+#         ::Val{:u})
+#     function solu_adjoint(Δ)
+#         zerou = zero(sol.prob.u0)
+#         _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
+#         (build_solution(sol.prob, sol.alg, sol.t, _Δ),)
+#     end
+#     sol.u, solu_adjoint
+# end
 
 @adjoint function literal_getproperty(sol::SciMLBase.AbstractNoTimeSolution,
         ::Val{:u})

test/downstream/observables_autodiff.jl

@@ -0,0 +1,34 @@
+using ModelingToolkit, OrdinaryDiffEq
+using Zygote
+
+@parameters σ ρ β
+@variables x(t) y(t) z(t) w(t)
+
+eqs = [D(D(x)) ~ σ * (y - x),
+    D(y) ~ x * (ρ - z) - y,
+    D(z) ~ x * y - β * z,
+    w ~ x + y + z + 2 * β]
+
+@mtkbuild sys = ODESystem(eqs, t)
+
+u0 = [D(x) => 2.0,
+    x => 1.0,
+    y => 0.0,
+    z => 0.0]
+
+p = [σ => 28.0,
+    ρ => 10.0,
+    β => 8 / 3]
+
+tspan = (0.0, 100.0)
+prob = ODEProblem(sys, u0, tspan, p, jac = true)
+sol = solve(prob, Tsit5())
+
+@testset "AutoDiff Observable Functions" begin
+    gs, = gradient(sol) do sol
+        sum(sol[sys.w])
+    end
+    du_ = [0., 1., 1., 1.]
+    du = [du_ for _ = sol.u]
+    @test du == gs.u
+end

test/runtests.jl

@@ -110,6 +110,9 @@ end
         @time @safetestset "Partial Functions" begin
             include("downstream/partial_functions.jl")
         end
+        @time @safetestset "Autodiff Observable Functions" begin
+            include("downstream/observables_autodiff.jl")
+        end
     end
 
     if !is_APPVEYOR && (GROUP == "Downstream" || GROUP == "SymbolicIndexingInterface")