Merge branch 'master' into os/ForwardDiff1.0

Author: Oscar Smith

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLBase"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 authors = ["Chris Rackauckas <[email protected]> and contributors"]
-version = "2.79.1"
+version = "2.80.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

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,

src/scimlfunctions.jl

@@ -948,7 +948,8 @@ dt: the time step
 
 ```julia
 ImplicitDiscreteFunction{iip,specialize}(f;
-                                analytic = __has_analytic(f) ? f.analytic : nothing)
+                                analytic = __has_analytic(f) ? f.analytic : nothing, 
+                                resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing)
 ```
 
 Note that only the function `f` itself is required. This function should
@@ -973,12 +974,13 @@ For more details on this argument, see the ODEFunction documentation.
 
 The fields of the ImplicitDiscreteFunction type directly match the names of the inputs.
 """
-struct ImplicitDiscreteFunction{iip, specialize, F, Ta, O, SYS, ID} <:
+struct ImplicitDiscreteFunction{iip, specialize, F, Ta, O, SYS, RP, ID} <:
        AbstractDiscreteFunction{iip}
     f::F
     analytic::Ta
     observed::O
     sys::SYS
+    resid_prototype::RP
     initialization_data::ID
 end
 
@@ -3064,6 +3066,9 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
         observed = __has_observed(f) ?
                    f.observed :
                    DEFAULT_OBSERVED,
+        resid_prototype = __has_resid_prototype(f) ?
+                          f.resid_prototype : 
+                          nothing,
         sys = __has_sys(f) ? f.sys : nothing,
         initialization_data = __has_initialization_data(f) ? f.initialization_data :
                               nothing) where {
@@ -3074,39 +3079,40 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
     sys = sys_or_symbolcache(sys, syms, paramsyms, indepsym)
 
     if specialize === NoSpecialize
-        ImplicitDiscreteFunction{iip, specialize, Any, Any, Any, Any, Any}(_f,
+        ImplicitDiscreteFunction{iip, specialize, Any, Any, Any, Any, Any, Any}(_f,
             analytic,
             observed,
             sys,
+            resid_prototype, 
             initialization_data)
     else
         ImplicitDiscreteFunction{
-            iip, specialize, typeof(_f), typeof(analytic), typeof(observed), typeof(sys),
+                                 iip, specialize, typeof(_f), typeof(analytic), typeof(observed), typeof(sys), typeof(resid_prototype),
             typeof(initialization_data)}(
-            _f, analytic, observed, sys, initialization_data)
+            _f, analytic, observed, sys, resid_prototype, initialization_data)
     end
 end
 
 function ImplicitDiscreteFunction{iip}(f; kwargs...) where {iip}
     ImplicitDiscreteFunction{iip, FullSpecialize}(f; kwargs...)
 end
 ImplicitDiscreteFunction{iip}(f::ImplicitDiscreteFunction; kwargs...) where {iip} = f
-function ImplicitDiscreteFunction(f; kwargs...)
-    ImplicitDiscreteFunction{isinplace(f, 5), FullSpecialize}(f; kwargs...)
+function ImplicitDiscreteFunction(f; resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing, kwargs...)
+    ImplicitDiscreteFunction{isinplace(f, 5), FullSpecialize}(f; resid_prototype, kwargs...)
 end
 ImplicitDiscreteFunction(f::ImplicitDiscreteFunction; kwargs...) = f
 
 function unwrapped_f(f::ImplicitDiscreteFunction, newf = unwrapped_f(f.f))
     specialize = specialization(f)
 
     if specialize === NoSpecialize
-        ImplicitDiscreteFunction{isinplace(f, 6), specialize, Any, Any,
-            Any, Any, Any}(newf, f.analytic, f.observed, f.sys, f.initialization_data)
+        ImplicitDiscreteFunction{isinplace(f, 5), specialize, Any, Any, Any,
+            Any, Any, Any}(newf, f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     else
-        ImplicitDiscreteFunction{isinplace(f, 6), specialize, typeof(newf),
+        ImplicitDiscreteFunction{isinplace(f, 5), specialize, typeof(newf),
             typeof(f.analytic),
-            typeof(f.observed), typeof(f.sys), typeof(f.initialization_data)}(newf,
-            f.analytic, f.observed, f.sys, f.initialization_data)
+            typeof(f.observed), typeof(f.sys), typeof(resid_prototype), typeof(f.initialization_data)}(newf,
+            f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     end
 end
 

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