Accumulation for ODEProblem

ext/SciMLBaseZygoteExt.jl

@@ -1,13 +1,14 @@
 module SciMLBaseZygoteExt
 
 using Zygote
-using Zygote: @adjoint, pullback
-import Zygote: literal_getproperty
+using Zygote: @adjoint, pullback, @_adjoint_keepthunks, _project, pair
+import Zygote: literal_getproperty, literal_getfield
 import ChainRulesCore
 using SciMLBase
 using SciMLBase: ODESolution, remake, ODEFunction,
                  getobserved, build_solution, EnsembleSolution,
-                 NonlinearSolution, AbstractTimeseriesSolution
+                 NonlinearSolution, AbstractTimeseriesSolution,
+                 ODEProblem
 using SymbolicIndexingInterface: symbolic_type, NotSymbolic, variable_index, is_observed,
                                  observed, parameter_values, state_values, current_time
 using RecursiveArrayTools
@@ -299,4 +300,23 @@ end
     ∇responsible_map(__context__, f, args...)
 end
 
+@_adjoint_keepthunks function Zygote.literal_getfield(x::ODEProblem, ::Val{f}) where f
+  val = getfield(x, f)
+  function back(Δ)
+    # error()
+    Zygote.accum_param(__context__, val, Δ) === nothing && return
+    if isimmutable(x)
+      error()
+      dx = (; Zygote.nt_nothing(x)..., pair(Val(f), Δ, x)...)
+      (_project(x, dx), nothing)
+    else
+      dx = Zygote.grad_mut(__context__, x)
+      dx[] = (; dx[]..., pair(Val(f), Zygote.accum(getfield(dx[], f), Δ))...)
+      return (dx,nothing)
+    end
+  end
+  Zygote.unwrap(val), back
+end
+Zygote.accum(::Tuple{}, ::NamedTuple{}) = ()
+
 end

src/initialization.jl

@@ -258,50 +258,58 @@ function get_initial_values(prob, valp, f, alg::OverrideInit,
             initdata.update_initializeprob!(initprob, valp)
         end
     end
+    nlsol, success = solve_initialization(initdata, initprob, alg; reltol, abstol, nlsolve_alg )
 
-    if is_trivial_initialization(initdata)
-        nlsol = initdata
-        success = true
-    else
-        nlsolve_alg = something(nlsolve_alg, alg.nlsolve, Some(nothing))
-        if nlsolve_alg === nothing && state_values(initprob) !== nothing
-            throw(OverrideInitMissingAlgorithm())
-        end
-        if alg.abstol !== nothing
-            _abstol = alg.abstol
-        elseif abstol !== nothing
-            _abstol = abstol
-        else
-            throw(OverrideInitNoTolerance(:abstol))
-        end
-        if alg.reltol !== nothing
-            _reltol = alg.reltol
-        elseif reltol !== nothing
-            _reltol = reltol
-        else
-            throw(OverrideInitNoTolerance(:reltol))
-        end
-        nlsol = solve(initprob, nlsolve_alg; abstol = _abstol, reltol = _reltol, kwargs...)
-
-        success = if initprob isa NonlinearLeastSquaresProblem
-            # Do not accept StalledSuccess as a solution
-            # A good local minima is not a success 
-            resid = nlsol.resid
-            normresid = norm(resid)
-            SciMLBase.successful_retcode(nlsol) && normresid <= abstol
-        else
-            SciMLBase.successful_retcode(nlsol)
-        end
-    end
-
+    nlsol2 = prob.f.initialization_data.initializeprob
     if initdata.initializeprobmap !== nothing
-        u0 = initdata.initializeprobmap(choose_branch(nlsol))
+        u02 = initdata.initializeprobmap(nlsol2)
     end
     if initdata.initializeprobpmap !== nothing
-        p = initdata.initializeprobpmap(valp, choose_branch(nlsol))
+        p2 = initdata.initializeprobpmap(valp, nlsol)
     end
 
-    return u0, p, success
+    u03 = isnothing(initdata.initializeprobmap) ? u0 : u02
+    p3 = isnothing(initdata.initializeprobpmap) ? p : p2
+    return u03, p3, success
+end
+
+function solve_initialization(initdata::OverrideInitData{<:AbstractNonlinearProblem{Nothing}}, initprob, alg; kwargs...)
+    nlsol = initprob
+    success = true
+    return nlsol, success
+end
+
+function solve_initialization(initdata, initprob, alg; reltol, abstol, nlsolve_alg)
+    nlsolve_alg = something(nlsolve_alg, alg.nlsolve, Some(nothing))
+    if nlsolve_alg === nothing && state_values(initprob) !== nothing
+        throw(OverrideInitMissingAlgorithm())
+    end
+    if alg.abstol !== nothing
+        _abstol = alg.abstol
+    elseif abstol !== nothing
+        _abstol = abstol
+    else
+        throw(OverrideInitNoTolerance(:abstol))
+    end
+    if alg.reltol !== nothing
+        _reltol = alg.reltol
+    elseif reltol !== nothing
+        _reltol = reltol
+    else
+        throw(OverrideInitNoTolerance(:reltol))
+    end
+    nlsol = solve(initprob, nlsolve_alg; abstol = _abstol, reltol = _reltol, kwargs...)
+
+    success = if initprob isa NonlinearLeastSquaresProblem
+        # Do not accept StalledSuccess as a solution
+        # A good local minima is not a success 
+        resid = nlsol.resid
+        normresid = norm(resid)
+        SciMLBase.successful_retcode(nlsol) && normresid <= abstol
+    else
+        SciMLBase.successful_retcode(nlsol)
+    end
+    return nlsol, success
 end
 
 """
@@ -314,21 +322,6 @@ function get_initial_values(prob, integrator, f, ::NoInit, iip; kwargs...)
     return state_values(integrator), parameter_values(integrator), true
 end
 
-is_trivial_initialization(::Nothing) = true
-
-function is_trivial_initialization(initdata::OverrideInitData)
-    !(initdata.initializeprob isa NonlinearLeastSquaresProblem) &&
-        state_values(initdata.initializeprob) === nothing
-end
-
-function is_trivial_initialization(f::AbstractSciMLFunction)
-    has_initialization_data(f) && is_trivial_initialization(f.initialization_data)
-end
-
-function is_trivial_initialization(prob::AbstractSciMLProblem)
-    is_trivial_initialization(prob.f)
-end
-
 @enum DETERMINED_STATUS OVERDETERMINED FULLY_DETERMINED UNDERDETERMINED
 
 function initialization_status(prob::AbstractSciMLProblem)

src/remake.jl

@@ -1229,3 +1229,20 @@ function remake(thing::AbstractEnsembleProblem; kwargs...)
     en_kwargs = [k for k in kwargs if first(k) ∈ fieldnames(T)]
     T(remake(thing.prob; setdiff(kwargs, en_kwargs)...); en_kwargs...)
 end
+
+is_trivial_initialization(::Nothing) = true
+
+is_trivial_initialization(::OverrideInitData{<:NonlinearLeastSquaresProblem}) = false
+
+is_trivial_initialization(::OverrideInitData{<:NonlinearProblem{Nothing}}) = true
+
+is_trivial_initialization(::OverrideInitData) = false
+
+function is_trivial_initialization(f::AbstractSciMLFunction)
+    # has_initialization_data(f) && is_trivial_initialization(f.initialization_data)
+    is_trivial_initialization(f.initialization_data)
+end
+
+function is_trivial_initialization(prob::AbstractSciMLProblem)
+    is_trivial_initialization(prob.f)
+end

src/utils.jl

@@ -550,6 +550,3 @@ Strips a SciMLSolution object and its interpolation of their functions to better
 function strip_solution(sol::AbstractSciMLSolution)
     sol
 end
-
-choose_branch(x::OverrideInitData) = x.initializeprob
-choose_branch(sol::AbstractSciMLSolution) = sol