use solve_up for SciMLSensitivity integration

Author: jClugstor

lib/OptimizationBase/src/OptimizationBase.jl

@@ -13,7 +13,10 @@ import SciMLBase: solve, init, solve!, __init, __solve,
                   allowsconstraints, requiresconstraints,
                   allowscallback, requiresgradient,
                   requireshessian, requiresconsjac,
-                  requiresconshess
+                  requiresconshess, wrap_sol, has_kwargs,
+                  get_root_indp, get_updated_symbolic_problem,
+                  get_concrete_p, get_concrete_u0, promote_u0,
+                  promote_p, KeywordArgError, checkkwargs
 
 export ObjSense, MaxSense, MinSense
 export allowsbounds, requiresbounds, allowsconstraints, requiresconstraints,

lib/OptimizationBase/src/solve.jl

@@ -90,18 +90,33 @@ from NLopt for an example. The common local optimizer arguments are:
   - `local_reltol`: relative tolerance  in changes of the objective value
   - `local_options`: `NamedTuple` of keyword arguments for local optimizer
 """
-function solve(prob::SciMLBase.OptimizationProblem, alg, args...;
-        kwargs...)::SciMLBase.AbstractOptimizationSolution
-    if SciMLBase.has_init(alg)
-        solve!(init(prob, alg, args...; kwargs...))
+function solve(prob::SciMLBase.OptimizationProblem, args...; sensealg = nothing,
+        u0 = nothing, p = nothing, wrap = Val(true), kwargs...)::SciMLBase.AbstractOptimizationSolution
+    if sensealg === nothing && haskey(prob.kwargs, :sensealg)
+        sensealg = prob.kwargs[:sensealg]
+    end
+
+    u0 = u0 !== nothing ? u0 : prob.u0
+    p = p !== nothing ? p : prob.p
+
+    if wrap isa Val{true}
+        wrap_sol(solve_up(prob,
+            sensealg,
+            u0,
+            p,
+            args...;
+            originator = SciMLBase.ChainRulesOriginator(),
+            kwargs...))
     else
-        if prob.u0 !== nothing && !isconcretetype(eltype(prob.u0))
-            throw(SciMLBase.NonConcreteEltypeError(eltype(prob.u0)))
-        end
-        _check_opt_alg(prob, alg; kwargs...)
-        __solve(prob, alg, args...; kwargs...)
+        solve_up(prob,
+            sensealg,
+            u0,
+            p,
+            args...;
+            originator = SciMLBase.ChainRulesOrginator(),
+            kwargs...)
     end
-end
+end 
 
 function solve(
         prob::SciMLBase.EnsembleProblem{T}, args...; kwargs...) where {T <:
@@ -216,3 +231,101 @@ end
 function __solve(prob::SciMLBase.OptimizationProblem, alg, args...; kwargs...)
     throw(OptimizerMissingError(alg))
 end
+
+function solve_up(prob::SciMLBase.OptimizationProblem, sensealg, u0, p, args...; originator = SciMLBase.ChainRulesOriginator(),
+    kwargs...)
+    alg = extract_opt_alg(args, kwargs, has_kwargs(prob) ? prob.kwargs : kwargs)
+    _prob = get_concrete_problem(prob; u0 = u0, p = p, kwargs...)
+    if length(args) < 1
+        solve_call(_prob, alg, Base.tails(args)..., kwargs...)
+    else
+        solve_call(_prob, alg; kwargs...)
+    end
+end
+
+function solve_call(_prob, alg, args...; merge_callbacks = true, kwargshandle = nothing,
+    kwargs...)
+    kwargshandle = kwargshandle === nothing ? KeywordArgError : kwargshandle
+    kwargshandle = has_kwargs(_prob) && haskey(_prob.kwargs, :kwargshandle) ?
+                   _prob.kwargs[:kwargshandle] : kwargshandle
+
+    if has_kwargs(_prob)
+        kwargs = isempty(_prob.kwargs) ? kwargs : merge(values(_prob.kwargs), kwargs)
+    end
+
+    checkkwargs(kwargshandle; kwargs...)
+
+    if SciMLBase.has_init(alg)
+        solve!(init(_prob, alg, args...; kwargs...))
+    else
+        if _prob.u0 !== nothing && !isconcretetype(eltype(_prob.u0))
+            throw(SciMLBase.NonConcreteEltypeError(eltype(_prob.u0)))
+        end
+        _check_opt_alg(prob, alg; kwargs...)
+        __solve(_prob, alg, args...; kwargs...)
+    end
+end
+
+function get_concrete_problem(prob::OptimizationProblem; kwargs...)
+    oldprob = prob
+    prob = get_updated_symbolic_problem(get_root_indp(prob), prob; kwargs...)
+    if prob !== oldprob
+        kwargs = (;kwargs..., u0 = SII.state_values(prob), p = SII.parameter_values(prob))
+    end
+    p = get_concrete_p(prob, kwargs)
+    u0 = get_concrete_u0(prob, false, nothing, kwargs)
+    u0 = promote_u0(u0, p, nothing)
+    remake(prob; u0 = u0, p = p)
+
+end
+
+
+@inline function extract_opt_alg(solve_args, solve_kwargs, prob_kwargs)
+    if isempty(solve_args) || isnothing(first(solve_args))
+        if haskey(solve_kwargs, :alg)
+            solve_kwargs[:alg]
+        elseif haskey(prob_kwargs, :alg)
+            prob_kwargs[:alg]
+        else
+            nothing
+        end
+    else
+        first(solve_args)
+    end
+end
+
+
+function _solve_forward(prob, sensealg, u0, p, originator, args...; merge_callbacks = true,
+    kwargs...)
+    alg = extract_opt_alg(args, kwargs, prob.kwargs)
+    _prob = get_concrete_problem(prob; u0 = u0, p = p, kwargs...)
+
+    if has_kwargs(_prob)
+        kwargs = isempty(_porb.kwargs) ? kwargs : merge(values(_prob.kwargs), kwargs)
+    end
+
+    if length(args) > 1
+        _concrete_solve_forward(_prob, alg, sensealg, u0, p, originator,
+            Base.tail(args)...; kwargs...)
+    else
+        _concrete_solve_forward(_prob, alg, sensealg, u0, p, originator; kwargs...)
+    end
+end
+
+function _solve_adjoint(_prob, sensealg, u0, p, originator, args...; merge_callbacks = true, 
+    kwargs...)
+    alg = extract_alg(args, kwargs, prob.kwargs)
+    
+    _prob = get_concrete_problem(prob; u0 = u0, p = p, kwargs...)
+
+    if has_kwargs(_prob)
+        kwargs = isempty(_prob.kwargs) ? kwargs : merge(values(_prob.kwargs), kwargs)
+    end
+    
+    if length(args) > 1
+        _concrete_solve_adjoint(_prob, alg, sensealg, u0, p, originator,
+            Base.tail(args)...; kwargs...)
+    else
+        _concrete_solve_adjoint(_prob, alg, sensealg, u0, p, originator; kwargs...)
+    end
+end