Optimization based solvers

docs/src/basics/solve.md

@@ -5,6 +5,7 @@
   - `controller`: Error controller for collocation methods, default as `DefectControl()`, more controller options in [Error Control Adaptivity](@ref error_control).
   - `defect_threshold`: Monitor of the size of defect norm. Defaults to `0.1`.
   - `odesolve_kwargs`: OrdinaryDiffEq.jl solvers kwargs for passing to ODE solving in shooting methods. For more information, see the documentation for OrdinaryDiffEq: [Common Solver Options](https://docs.sciml.ai/DiffEqDocs/latest/basics/common_solver_opts/).
-  - `nlsolve_kwargs`: NonlinearSolve.jl solvers kwargs for passing to nonlinear solving in collocation methods and shooting methods. For more information, see the documentation for NonlinearSolve: [Commom Solver Options](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/). The default absolute tolerance of nonlinear solving in collocaio
+  - `nlsolve_kwargs`: NonlinearSolve.jl solvers kwargs for passing to nonlinear solving in collocation methods and shooting methods. For more information, see the documentation for NonlinearSolve: [Commom Solver Options](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/). The default internal nonlinear solver is [customized polyalgorithm](https://github.com/SciML/BoundaryValueDiffEq.jl/blob/master/lib/BoundaryValueDiffEqCore/src/default_nlsolve.jl) and the default absolute tolerance of nonlinear solving in collocation methods is `1e-6`.
+  - `optimize_kwargs`: Optimization.jl solvers kwargs for passing to optimizing in collocation methods and shooting methods. For more information, see the documentation for Optimization: [Commom Solver Options](https://docs.sciml.ai/Optimization/stable/API/solve/). The internal optimization solver should be specified and the default absolute tolerance of optimization problem solving in collocation methods is `1e-6`.
   - `verbose`:  Toggles whether warnings are thrown when the solver exits early. Defaults to `true`.
   - `ensemblealg`: Whether `MultipleShooting` uses multithreading, default as `EnsembleThreads()`. For more information, see the documentation for OrdinaryDiffEq: [EnsembleAlgorithms](https://docs.sciml.ai/DiffEqDocs/latest/features/ensemble/#EnsembleAlgorithms).

lib/BoundaryValueDiffEqAscher/src/BoundaryValueDiffEqAscher.jl

@@ -6,11 +6,11 @@ using AlmostBlockDiagonals: AlmostBlockDiagonals, IntermediateAlmostBlockDiagona
 using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
                                AbstractBoundaryValueDiffEqCache, BVPJacobianAlgorithm,
                                __extract_problem_details, concrete_jacobian_algorithm,
-                               __Fix3, __concrete_nonlinearsolve_algorithm,
+                               __Fix3, __concrete_solve_algorithm,
                                __internal_nlsolve_problem, __vec, __vec_f, __vec_f!,
                                __vec_bc, __vec_bc!, __extract_mesh, get_dense_ad,
-                               __get_bcresid_prototype, __split_kwargs,
-                               __default_nonsparse_ad
+                               __get_bcresid_prototype, __split_kwargs, __concrete_kwargs,
+                               __default_nonsparse_ad, __construct_internal_problem
 
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase

lib/BoundaryValueDiffEqAscher/src/algorithms.jl

@@ -14,6 +14,9 @@ for stage in (1, 2, 3, 4, 5, 6, 7)
           - `nlsolve`: Internal Nonlinear solver. Any solver which conforms to the SciML
             `NonlinearProblem` interface can be used. Note that any autodiff argument for
             the solver will be ignored and a custom jacobian algorithm will be used.
+          - `optimize`: Internal Optimization solver. Any solver which conforms to the SciML
+            `OptimizationProblem` interface can be used. Note that any autodiff argument for
+            the solver will be ignored and a custom jacobian algorithm will be used.
           - `max_num_subintervals`: Number of maximal subintervals, default as 3000.
           - `zeta`: side condition points, should always be provided.
 
@@ -46,8 +49,9 @@ for stage in (1, 2, 3, 4, 5, 6, 7)
         }
         ```
         """
-        @kwdef struct $(alg){N, J <: BVPJacobianAlgorithm} <: AbstractAscher
+        @kwdef struct $(alg){N, O, J <: BVPJacobianAlgorithm} <: AbstractAscher
             nlsolve::N = nothing
+            optimize::O = nothing
             zeta::Vector{Float64} = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             max_num_subintervals::Int = 3000

lib/BoundaryValueDiffEqAscher/src/ascher.jl

@@ -39,6 +39,7 @@
     TU
     valstr
     nlsolve_kwargs
+    optimize_kwargs
     kwargs
 end
 
@@ -60,7 +61,8 @@ end
 
 function SciMLBase.__init(
         prob::BVProblem, alg::AbstractAscher; dt = 0.0, controller = GlobalErrorControl(),
-        adaptive = true, abstol = 1e-4, nlsolve_kwargs = (; abstol = abstol), kwargs...)
+        adaptive = true, abstol = 1e-4, nlsolve_kwargs = (; abstol = abstol),
+        optimize_kwargs = (; abstol = abstol), kwargs...)
     (; tspan, p) = prob
     _, T, ncy, n, u0 = __extract_problem_details(prob; dt, check_positive_dt = true)
     t₀, t₁ = tspan
@@ -147,9 +149,9 @@ function SciMLBase.__init(
     g = build_almost_block_diagonals(zeta, ncomp, mesh, T)
     cache = AscherCache{iip, T}(
         prob, f, jac, bc, bcjac, k, copy(mesh), mesh, mesh_dt, ncomp, ny, p, zeta,
-        fixpnt, alg, prob.problem_type, bcresid_prototype, residual, zval, yval,
-        gval, err, g, w, v, lz, ly, dmz, delz, deldmz, dqdmz, dmv, pvtg, pvtw, TU,
-        valst, nlsolve_kwargs, (; abstol, dt, adaptive, controller, kwargs...))
+        fixpnt, alg, prob.problem_type, bcresid_prototype, residual, zval, yval, gval,
+        err, g, w, v, lz, ly, dmz, delz, deldmz, dqdmz, dmv, pvtg, pvtw, TU, valst,
+        nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...))
     return cache
 end
 
@@ -176,8 +178,10 @@ function __perform_ascher_iteration(
         cache::AscherCache{iip, T}, abstol, adaptive::Bool) where {iip, T}
     info::ReturnCode.T = ReturnCode.Success
     nlprob = __construct_nlproblem(cache)
-    nlsolve_alg = __concrete_nonlinearsolve_algorithm(nlprob, cache.alg.nlsolve)
-    nlsol = __solve(nlprob, nlsolve_alg; cache.nlsolve_kwargs...)
+    solve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve, cache.alg.optimize)
+    kwargs = __concrete_kwargs(
+        cache.alg.nlsolve, cache.alg.optimize, cache.nlsolve_kwargs, cache.optimize_kwargs)
+    nlsol = solve(nlprob, solve_alg; kwargs...)
     error_norm = 2 * abstol
     info = nlsol.retcode
 
@@ -203,7 +207,7 @@ function __perform_ascher_iteration(
         __expand_cache_for_error!(cache)
 
         _nlprob = __construct_nlproblem(cache)
-        nlsol = __solve(_nlprob, nlsolve_alg; cache.nlsolve_kwargs...)
+        nlsol = solve(_nlprob, solve_alg; kwargs...)
 
         error_norm = error_estimate!(cache)
         if norm(error_norm) > abstol
@@ -346,9 +350,9 @@ function __construct_nlproblem(cache::AscherCache{iip, T}) where {iip, T}
             jac_prototype, u, diffmode, jac_cache, loss, cache.p)
     end
 
-    nlf = NonlinearFunction{iip}(
-        loss; jac = jac, resid_prototype = resid_prototype, jac_prototype = jac_prototype)
-    return __internal_nlsolve_problem(cache.prob, similar(lz), lz, nlf, lz, cache.p)
+    return __construct_internal_problem(
+        cache.prob, alg, loss, jac, jac_prototype, resid_prototype,
+        lz, cache.p, cache.ncomp, length(cache.mesh))
 end
 
 function __ascher_mpoint_jacobian!(J, x, diffmode, diffcache, loss, resid, p)

lib/BoundaryValueDiffEqAscher/src/utils.jl

@@ -112,7 +112,7 @@ end
     return nothing
 end
 
-@views function recursive_flatten!(y::Vector, x::Vector{Vector{T}}) where {T}
+@views function recursive_flatten!(y::AbstractArray, x::Vector{Vector{T}}) where {T}
     i = 0
     for xᵢ in x
         copyto!(y[(i + 1):(i + length(xᵢ))], xᵢ)

lib/BoundaryValueDiffEqCore/src/BoundaryValueDiffEqCore.jl

@@ -29,7 +29,7 @@ include("utils.jl")
 include("algorithms.jl")
 include("abstract_types.jl")
 include("alg_utils.jl")
-include("default_nlsolve.jl")
+include("default_internal_solve.jl")
 include("calc_errors.jl")
 
 function SciMLBase.__solve(prob::AbstractBVProblem,

lib/BoundaryValueDiffEqCore/src/algorithms.jl

@@ -26,3 +26,12 @@ function Base.show(io::IO, alg::AbstractBoundaryValueDiffEqAlgorithm)
     print(io, join(modifiers, ", "))
     print(io, ")")
 end
+
+# Check what's the internal solver, nonlinear or optimization?
+function __internal_solver(alg::AbstractBoundaryValueDiffEqAlgorithm)
+    # We don't allow both `nlsolve` and `optimize` to be specified at the same time
+    (isnothing(alg.nlsolve) && isnothing(alg.optimize)) &&
+        error("Either `nlsolve` or `optimize` must be specified in the algorithm, but not both.")
+    isnothing(alg.nlsolve) && return alg.optimize
+    isnothing(alg.optimize) && return alg.nlsolve
+end

lib/BoundaryValueDiffEqCore/src/default_nlsolve.jl → lib/BoundaryValueDiffEqCore/src/default_internal_solve.jl

@@ -48,8 +48,28 @@ function __FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = noth
     return NonlinearSolvePolyAlgorithm(algs)
 end
 
-@inline __concrete_nonlinearsolve_algorithm(prob, alg) = alg
-@inline function __concrete_nonlinearsolve_algorithm(prob, ::Nothing)
+"""
+    __concrete_solve_algorithm(prob, nlsolve_alg, optimize_alg)
+
+Automatic solver choosing according to the input solver.
+If none of the solvers are specified, we use nonlinear solvers from NonlinearSolve.jl.
+If both of the nonlinear solver and optimization solver are specified, we throw an error.
+If only one of the nonlinear solver and optimization solver is specified, we use that solver.
+"""
+@inline __concrete_solve_algorithm(prob, alg) = alg
+@inline __concrete_solve_algorithm(prob, alg, ::Nothing) = alg
+@inline __concrete_solve_algorithm(prob, ::Nothing, alg) = alg
+@inline __concrete_solve_algorithm(prob,
+    alg1,
+    alg2) = error("Both `nlsolve` and `optimize` are specified in the algorithm, but only one of them is allowed. Please specify only one of them.")
+@inline function __concrete_solve_algorithm(prob, ::Nothing)
+    if prob isa NonlinearLeastSquaresProblem
+        return __FastShortcutBVPCompatibleNLLSPolyalg(eltype(prob.u0))
+    else
+        return __FastShortcutBVPCompatibleNonlinearPolyalg(eltype(prob.u0))
+    end
+end
+@inline function __concrete_solve_algorithm(prob, ::Nothing, ::Nothing)
     if prob isa NonlinearLeastSquaresProblem
         return __FastShortcutBVPCompatibleNLLSPolyalg(eltype(prob.u0))
     else

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -468,6 +468,19 @@ end
     return NonlinearProblem(args...; kwargs...)
 end
 
+# Construct the internal OptimizationProblem
+@inline function __internal_optimization_problem(
+        ::BVProblem{uType, tType, iip}, args...; kwargs...) where {uType, tType, iip}
+    prob = OptimizationProblem(args...; kwargs...)
+    return prob
+end
+
+@inline function __internal_optimization_problem(::SecondOrderBVProblem{uType, tType, iip},
+        args...; kwargs...) where {uType, tType, iip}
+    prob = OptimizationProblem(args...; kwargs...)
+    return prob
+end
+
 # Handling Initial Guesses
 """
     __extract_u0(u₀, t₀)
@@ -563,10 +576,17 @@ end
 end
 
 # Construct BVP Solution
-function __build_solution(prob::AbstractBVProblem, odesol, nlsol)
+function __build_solution(
+        prob::AbstractBVProblem, odesol, nlsol::SciMLBase.NonlinearSolution)
     retcode = ifelse(SciMLBase.successful_retcode(nlsol), odesol.retcode, nlsol.retcode)
     return SciMLBase.solution_new_original_retcode(odesol, nlsol, retcode, nlsol.resid)
 end
+function __build_solution(
+        prob::AbstractBVProblem, odesol, optsol::SciMLBase.OptimizationSolution)
+    retcode = ifelse(SciMLBase.successful_retcode(optsol), odesol.retcode, optsol.retcode)
+    return SciMLBase.solution_new_original_retcode(
+        odesol, optsol, retcode, zeros(length(first(odesol)))) # Need a patch in SciMLBase
+end
 
 # Fix3
 @concrete struct __Fix3
@@ -600,3 +620,106 @@ end
 function __split_kwargs(; abstol, adaptive, controller, kwargs...)
     return ((abstol, adaptive, controller), (; abstol, adaptive, kwargs...))
 end
+
+@inline __concrete_kwargs(nlsolve,
+    ::Nothing,
+    nlsolve_kwargs,
+    optimize_kwargs) = (
+    nlsolve_kwargs..., alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = true))
+@inline __concrete_kwargs(::Nothing, optimize, nlsolve_kwargs, optimize_kwargs) = (;) # Doesn't support for now
+@inline __concrete_kwargs(::Nothing,
+    ::Nothing,
+    nlsolve_kwargs,
+    optimize_kwargs) = (
+    nlsolve_kwargs..., alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = true))
+
+## Optimization solver related utils ##
+
+@inline __default_cost(::Nothing) = (x, p) -> 0.0
+@inline __default_cost(f) = f
+@inline __default_cost(fun::BVPFunction) = __default_cost(fun.cost)
+
+"""
+    __construct_internal_problem
+
+Constructs the internal problem based on the type of the boundary value problem and the
+algorithm used. It returns either a `NonlinearProblem` or an `OptimizationProblem`.
+"""
+function __construct_internal_problem(prob::AbstractBVProblem, alg, loss, jac,
+        jac_prototype, resid_prototype, y, p, M::Int, N::Int)
+    T = eltype(y)
+    # multiple shooting always use iip
+    iip = SciMLBase.isinplace(prob)
+    if !isnothing(alg.nlsolve) || (isnothing(alg.nlsolve) && isnothing(alg.optimize))
+        nlf = NonlinearFunction{iip}(loss; jac = jac, resid_prototype = resid_prototype,
+            jac_prototype = jac_prototype)
+        return __internal_nlsolve_problem(prob, resid_prototype, y, nlf, y, p)
+    else
+        optf = OptimizationFunction{true}(__default_cost(prob.f), AutoFiniteDiff(), # Need to investigate the ForwardDiff dual problem
+            cons = loss,
+            cons_j = jac, cons_jac_prototype = jac_prototype)
+        lcons = zeros(T, N*M)
+        ucons = zeros(T, N*M)
+        return __internal_optimization_problem(
+            prob, optf, y, p; lcons = lcons, ucons = ucons)
+    end
+end
+
+function __construct_internal_problem(prob::TwoPointBVProblem, alg, loss, jac,
+        jac_prototype, resid_prototype, y, p, M::Int, N::Int)
+    T = eltype(y)
+    iip = SciMLBase.isinplace(prob)
+    if !isnothing(alg.nlsolve) || (isnothing(alg.nlsolve) && isnothing(alg.optimize))
+        nlf = NonlinearFunction{iip}(loss; jac = jac, resid_prototype = resid_prototype,
+            jac_prototype = jac_prototype)
+        return __internal_nlsolve_problem(prob, resid_prototype, y, nlf, y, p)
+    else
+        optf = OptimizationFunction{true}(
+            __default_cost(prob.f), get_dense_ad(alg.jac_alg.diffmode),
+            cons = loss, cons_j = jac, cons_jac_prototype = jac_prototype)
+        lcons = zeros(T, N*M)
+        ucons = zeros(T, N*M)
+
+        return __internal_optimization_problem(
+            prob, optf, y, p; lcons = lcons, ucons = ucons)
+    end
+end
+# Multiple shooting only use inplace version internal problem constructor
+function __construct_internal_problem(prob, alg, loss, jac, jac_prototype,
+        resid_prototype, y, p, M::Int, N::Int, ::Nothing)
+    T = eltype(y)
+    if !isnothing(alg.nlsolve) || (isnothing(alg.nlsolve) && isnothing(alg.optimize))
+        nlf = NonlinearFunction{true}(loss; jac = jac, resid_prototype = resid_prototype,
+            jac_prototype = jac_prototype)
+        return __internal_nlsolve_problem(prob, resid_prototype, y, nlf, y, p)
+    else
+        optf = OptimizationFunction{true}(
+            __default_cost(prob.f), get_dense_ad(alg.jac_alg.diffmode),
+            cons = loss, cons_j = jac, cons_jac_prototype = jac_prototype)
+        lcons = zeros(T, N*M)
+        ucons = zeros(T, N*M)
+
+        return __internal_optimization_problem(
+            prob, optf, y, p; lcons = lcons, ucons = ucons)
+    end
+end
+function __construct_internal_problem(
+        prob::TwoPointBVProblem, alg, loss, jac, jac_prototype,
+        resid_prototype, y, p, M::Int, N::Int, ::Nothing)
+    T = eltype(y)
+    iip = SciMLBase.isinplace(prob)
+    if !isnothing(alg.nlsolve) || (isnothing(alg.nlsolve) && isnothing(alg.optimize))
+        nlf = NonlinearFunction{iip}(loss; jac = jac, resid_prototype = resid_prototype,
+            jac_prototype = jac_prototype)
+        return __internal_nlsolve_problem(prob, resid_prototype, y, nlf, y, p)
+    else
+        optf = OptimizationFunction{true}(
+            __default_cost(prob.f), get_dense_ad(alg.jac_alg.nonbc_diffmode),
+            cons = loss, cons_j = jac, cons_jac_prototype = jac_prototype)
+        lcons = zeros(T, N*M)
+        ucons = zeros(T, N*M)
+
+        return __internal_optimization_problem(
+            prob, optf, y, p; lcons = lcons, ucons = ucons)
+    end
+end

lib/BoundaryValueDiffEqFIRK/src/BoundaryValueDiffEqFIRK.jl

@@ -6,23 +6,24 @@ using BandedMatrices: BandedMatrix, Ones
 using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
                                AbstractBoundaryValueDiffEqCache, BVPJacobianAlgorithm,
                                recursive_flatten, recursive_flatten!, recursive_unflatten!,
-                               __concrete_nonlinearsolve_algorithm, diff!, EvalSol,
+                               __concrete_solve_algorithm, diff!, EvalSol,
                                concrete_jacobian_algorithm, eval_bc_residual, interval,
                                eval_bc_residual!, get_tmp, __maybe_matmul!, __resize!,
                                __extract_problem_details, __initial_guess, nodual_value,
                                __maybe_allocate_diffcache, __restructure_sol,
                                __get_bcresid_prototype, __vec, __vec_f, __vec_f!, __vec_bc,
-                               __vec_bc!, recursive_flatten_twopoint!,
+                               __vec_bc!, recursive_flatten_twopoint!, __concrete_kwargs,
                                __internal_nlsolve_problem, __extract_mesh, __extract_u0,
                                __default_coloring_algorithm, __maybe_allocate_diffcache,
                                __restructure_sol, __get_bcresid_prototype, safe_similar,
                                __vec, __vec_f, __vec_f!, __vec_bc, __vec_bc!, __cache_trait,
                                recursive_flatten_twopoint!, __internal_nlsolve_problem,
                                __extract_mesh, __extract_u0, DiffCacheNeeded,
                                NoDiffCacheNeeded, __has_initial_guess,
-                               __initial_guess_length, __initial_guess_on_mesh,
-                               __flatten_initial_guess, __build_solution, __Fix3,
-                               __split_kwargs, _sparse_like, get_dense_ad
+                               __construct_internal_problem, __initial_guess_length,
+                               __initial_guess_on_mesh, __flatten_initial_guess,
+                               __build_solution, __Fix3, __split_kwargs, _sparse_like,
+                               get_dense_ad, __internal_optimization_problem
 
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase

lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl

@@ -106,7 +106,7 @@ end
     end
     τ = (t - mesh[j])
 
-    nest_nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, alg.nlsolve)
+    nest_nlsolve_alg = __concrete_solve_algorithm(nest_prob, alg.nlsolve)
     nestprob_p = zeros(T, cache.M + 2)
 
     yᵢ = copy(cache.y[j].du)
@@ -152,7 +152,7 @@ end
     end
     τ = (t - mesh[j])
 
-    nest_nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, alg.nlsolve)
+    nest_nlsolve_alg = __concrete_solve_algorithm(nest_prob, alg.nlsolve)
     nestprob_p = zeros(T, cache.M + 2)
 
     yᵢ = copy(cache.y[j])
@@ -457,7 +457,7 @@ end
     (; f, mesh, mesh_dt, defect, ITU, nest_prob, alg) = cache
     (; q_coeff, τ_star) = ITU
 
-    nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, cache.alg.nlsolve)
+    nlsolve_alg = __concrete_solve_algorithm(nest_prob, cache.alg.nlsolve)
     nestprob_p = zeros(T, cache.M + 2)
 
     for i in 1:(length(mesh) - 1)
@@ -509,7 +509,7 @@ end
     (; f, mesh, mesh_dt, defect, ITU, nest_prob, alg) = cache
     (; q_coeff, τ_star) = ITU
 
-    nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, cache.alg.nlsolve)
+    nlsolve_alg = __concrete_solve_algorithm(nest_prob, cache.alg.nlsolve)
     nestprob_p = zeros(T, cache.M + 2)
 
     for i in 1:(length(mesh) - 1)