MIRKN work fine

ErikQQY · ErikQQY · commit b7f5762c6a19 · 2025-07-21T16:17:28.000+08:00
diff --git a/lib/BoundaryValueDiffEqCore/src/utils.jl b/lib/BoundaryValueDiffEqCore/src/utils.jl
@@ -471,6 +471,12 @@ end
     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₀)
@@ -630,15 +636,15 @@ Constructs the internal problem based on the type of the boundary value problem
 algorithm used. It returns either a `NonlinearProblem` or an `OptimizationProblem`.
 """
 function __construct_internal_problem(
-        prob::BVProblem, alg, loss, jac, jac_prototype, resid_prototype, y, p, M, N)
+        prob::AbstractBVProblem, alg, loss, jac, jac_prototype, resid_prototype, y, p, M, N)
     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), AutoFiniteDiff(), # Need to investigate the ForwardDiff dual problem
+        optf = OptimizationFunction{true}((x, p) -> 0.0, AutoFiniteDiff(), # Need to investigate the ForwardDiff dual problem
             cons = loss,
             cons_j = jac, cons_jac_prototype = jac_prototype)
         lcons = zeros(T, N*M)
@@ -658,7 +664,7 @@ function __construct_internal_problem(
         return __internal_nlsolve_problem(prob, resid_prototype, y, nlf, y, p)
     else
         optf = OptimizationFunction{true}(
-            __default_cost(prob.f), get_dense_ad(diffmode), cons = loss,
+            (x, p) -> 0.0, get_dense_ad(diffmode), cons = loss,
             cons_j = jac, cons_jac_prototype = jac_prototype)
         lcons = zeros(T, N*M)
         ucons = zeros(T, N*M)
diff --git a/lib/BoundaryValueDiffEqFIRK/test/expanded/firk_basic_tests.jl b/lib/BoundaryValueDiffEqFIRK/test/expanded/firk_basic_tests.jl
@@ -307,30 +307,6 @@ end
     testTol = 1e-6
     nested = false
 
-    for stage in (2, 3, 5, 7)
-        s = Symbol("RadauIIa$(stage)")
-        @eval radau_solver(::Val{$stage}) = $(s)(
-            NewtonRaphson(), BVPJacobianAlgorithm(); nested)
-    end
-
-    for stage in (3, 4, 5)
-        s = Symbol("LobattoIIIa$(stage)")
-        @eval lobattoIIIa_solver(::Val{$stage}) = $(s)(
-            NewtonRaphson(), BVPJacobianAlgorithm(); nested)
-    end
-
-    for stage in (3, 4, 5)
-        s = Symbol("LobattoIIIb$(stage)")
-        @eval lobattoIIIb_solver(::Val{$stage}) = $(s)(
-            NewtonRaphson(), BVPJacobianAlgorithm(); nested)
-    end
-
-    for stage in (3, 4, 5)
-        s = Symbol("LobattoIIIc$(stage)")
-        @eval lobattoIIIc_solver(::Val{$stage}) = $(s)(
-            NewtonRaphson(), BVPJacobianAlgorithm(); nested)
-    end
-
     @testset "Radau interpolations" begin
         @testset "Interpolation tests for RadauIIa$stage" for stage in (2, 3, 5, 7)
             @time sol = solve(prob_bvp_linear, radau_solver(Val(stage)); dt = 0.001)
@@ -356,10 +332,11 @@ end
             for (id, lobatto_solver) in zip(("a", "b", "c"),
                 (lobattoIIIa_solver, lobattoIIIb_solver, lobattoIIIc_solver))
                 begin
-                    @testset "Interpolation tests for LobattoIII$(id)$stage" for stage in
-                                                                                 (3, 4, 5)
-                        @time sol = solve(
-                            prob_bvp_linear, lobatto_solver(Val(stage)); dt = 0.001)
+                    @testset "Interpolation tests for LobattoIII$(id)$stage" for stage in (
+                        2, 3, 4, 5)
+                        adaptive = ifelse(stage == 2, false, true) # LobattoIIIa2 is not adaptive
+                        @time sol = solve(prob_bvp_linear, lobatto_solver(Val(stage));
+                            dt = 0.001, adaptive = adaptive)
                         @test sol(0.001)≈[0.998687464, -1.312035941] atol=testTol
                         @test sol(0.001; idxs = [1, 2])≈[0.998687464, -1.312035941] atol=testTol
                         @test sol(0.001; idxs = 1)≈0.998687464 atol=testTol
@@ -368,9 +345,10 @@ end
 
                     @testset "Derivative Interpolation tests for lobatto$(id)$stage" for stage in
                                                                                          (
-                        3, 4, 5)
-                        @time sol = solve(
-                            prob_bvp_linear, lobatto_solver(Val(stage)); dt = 0.001)
+                        2, 3, 4, 5)
+                        adaptive = ifelse(stage == 2, false, true) # LobattoIIIa2 is not adaptive
+                        @time sol = solve(prob_bvp_linear, lobatto_solver(Val(stage));
+                            dt = 0.001, adaptive = adaptive)
                         sol_analytic = prob_bvp_linear_analytic(nothing, λ, 0.04)
                         dsol_analytic = prob_bvp_linear_analytic_derivative(
                             nothing, λ, 0.04)
diff --git a/lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl b/lib/BoundaryValueDiffEqMIRKN/src/BoundaryValueDiffEqMIRKN.jl
@@ -19,7 +19,8 @@ using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
                                __default_sparse_ad, __default_nonsparse_ad, get_dense_ad,
                                concrete_jacobian_algorithm, __default_coloring_algorithm,
                                __default_sparsity_detector, interval, __split_kwargs,
-                               NoErrorControl
+                               NoErrorControl, __construct_internal_problem,
+                               __concrete_kwargs
 
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase
diff --git a/lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl b/lib/BoundaryValueDiffEqMIRKN/src/algorithms.jl
@@ -6,7 +6,7 @@ for order in (4, 6)
 
     @eval begin
         """
-              $($alg)(; nlsolve = NewtonRaphson(), jac_alg = BVPJacobianAlgorithm(),
+              $($alg)(; nlsolve = NewtonRaphson(), optimize = nothing, jac_alg = BVPJacobianAlgorithm(),
                       defect_threshold = 0.1, max_num_subintervals = 3000)
 
         $($order)th order Monotonic Implicit Runge Kutta Nyström method.
@@ -16,6 +16,10 @@ for order in (4, 6)
             - `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. Optimization
+              solvers should first be loaded to allow this functionality.
             - `jac_alg`: Jacobian Algorithm used for the nonlinear solver. Defaults to
               `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to
               use based on the input types and problem type.
@@ -46,8 +50,9 @@ for order in (4, 6)
         }
         ```
         """
-        @kwdef struct $(alg){N, J <: BVPJacobianAlgorithm, T} <: AbstractMIRKN
+        @kwdef struct $(alg){N, O, J <: BVPJacobianAlgorithm, T} <: AbstractMIRKN
             nlsolve::N = nothing
+            optimize::O = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             defect_threshold::T = 0.1
             max_num_subintervals::Int = 3000
diff --git a/lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl b/lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl
@@ -21,16 +21,19 @@
     fᵢ₂_cache
     resid_size
     nlsolve_kwargs
+    optimize_kwargs
     kwargs
 end
 
 Base.eltype(::MIRKNCache{iip, T}) where {iip, T} = T
 
-function SciMLBase.__init(prob::SecondOrderBVProblem, alg::AbstractMIRKN; dt = 0.0,
-        adaptive = false, abstol = 1e-6, controller = NoErrorControl(),
-        nlsolve_kwargs = (; abstol = abstol), kwargs...)
+function SciMLBase.__init(prob::SecondOrderBVProblem, alg::AbstractMIRKN;
+        dt = 0.0, adaptive = false, abstol = 1e-6,
+        controller = NoErrorControl(), nlsolve_kwargs = (; abstol = abstol),
+        optimize_kwargs = (; abstol = abstol), kwargs...)
     @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
     iip = isinplace(prob)
+    @assert (iip || isnothing(alg.optimize)) "Out-of-place constraints don't allow optimization solvers "
     t₀, t₁ = prob.tspan
     ig, T, M, Nig, X = __extract_problem_details(prob; dt, check_positive_dt = true)
     mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
@@ -93,9 +96,10 @@ function SciMLBase.__init(prob::SecondOrderBVProblem, alg::AbstractMIRKN; dt = 0
     prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
 
     return MIRKNCache{iip, T}(
-        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p, alg, TU,
-        bcresid_prototype, mesh, mesh_dt, k_discrete, y, y₀, residual, fᵢ_cache, fᵢ₂_cache,
-        resid_size, nlsolve_kwargs, (; abstol, dt, adaptive, controller, kwargs...))
+        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type,
+        prob.p, alg, TU, bcresid_prototype, mesh, mesh_dt, k_discrete,
+        y, y₀, residual, fᵢ_cache, fᵢ₂_cache, resid_size, nlsolve_kwargs,
+        optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...))
 end
 
 function SciMLBase.solve!(cache::MIRKNCache{iip, T}) where {iip, T}
@@ -112,9 +116,11 @@ function SciMLBase.solve!(cache::MIRKNCache{iip, T}) where {iip, T}
 end
 
 function __perform_mirkn_iteration(cache::MIRKNCache)
-    nlprob::NonlinearProblem = __construct_nlproblem(cache, vec(cache.y₀), copy(cache.y₀))
-    nlsolve_alg = __concrete_solve_algorithm(nlprob, cache.alg.nlsolve)
-    sol_nlprob = __solve(nlprob, nlsolve_alg; cache.nlsolve_kwargs..., alias_u0 = true)
+    nlprob = __construct_nlproblem(cache, vec(cache.y₀), copy(cache.y₀))
+    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)
+    sol_nlprob = solve(nlprob, solve_alg, kwargs...)
     recursive_unflatten!(cache.y₀, sol_nlprob.u)
 
     return sol_nlprob, sol_nlprob.retcode
@@ -227,9 +233,8 @@ function __construct_nlproblem(cache::MIRKNCache{iip}, y, loss_bc::BC, loss_coll
             cache_collocation, loss_bc, loss_collocation, L, cache.p)
     end
     resid_prototype = vcat(resid_bc, resid_collocation)
-    nlf = NonlinearFunction{iip}(
-        loss; jac = jac, resid_prototype = resid_prototype, jac_prototype = jac_prototype)
-    __internal_nlsolve_problem(cache.prob, resid_prototype, y, nlf, y, cache.p)
+    return __construct_internal_problem(cache.prob, cache.alg, loss, jac, jac_prototype,
+        resid_prototype, y, cache.p, cache.M, 2 * N)
 end
 
 function __construct_nlproblem(cache::MIRKNCache{iip}, y, loss_bc::BC, loss_collocation::C,
@@ -270,9 +275,8 @@ function __construct_nlproblem(cache::MIRKNCache{iip}, y, loss_bc::BC, loss_coll
     end
 
     resid_prototype = copy(resid)
-    nlf = NonlinearFunction{iip}(
-        loss; jac = jac, resid_prototype = resid_prototype, jac_prototype = jac_prototype)
-    return __internal_nlsolve_problem(cache.prob, resid_prototype, y, nlf, y, cache.p)
+    return __construct_internal_problem(cache.prob, cache.alg, loss, jac, jac_prototype,
+        resid_prototype, y, cache.p, cache.M, 2 * N)
 end
 
 function __mirkn_2point_jacobian!(J, x, diffmode, diffcache, loss_fn::L, resid, p) where {L}
diff --git a/lib/BoundaryValueDiffEqShooting/src/multiple_shooting.jl b/lib/BoundaryValueDiffEqShooting/src/multiple_shooting.jl
@@ -7,6 +7,7 @@ function SciMLBase.__solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwa
         throw(ArgumentError("Currently MultipleShooting only supports `EnsembleSerial` and \
                              `EnsembleThreads`!"))
     end
+    @assert (iip || isnothing(_alg.optimize)) "Out-of-place constraints don't allow optimization solvers "
 
     ig, T, N, Nig, u0 = __extract_problem_details(prob; dt = 0.1)
     has_initial_guess = _unwrap_val(ig)
@@ -221,7 +222,7 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_
     nlprob = __construct_internal_problem(prob, alg, loss_fn, jac_fn, jac_prototype,
         resid_prototype, u_at_nodes, prob.p, M, length(nodes))
     nlsolve_alg = __concrete_solve_algorithm(nlprob, alg.nlsolve, alg.optimize)
-    kwargs = __concrete_kwargs(alg.nlsolve, alg.optimize, kwargs...)
+    kwargs = __concrete_kwargs(alg.nlsolve, alg.optimize, nlsolve_kwargs, optimize_kwargs)
     solve(nlprob, nlsolve_alg; kwargs...)
 
     return nothing
diff --git a/lib/BoundaryValueDiffEqShooting/src/single_shooting.jl b/lib/BoundaryValueDiffEqShooting/src/single_shooting.jl
@@ -7,6 +7,7 @@ function SciMLBase.__solve(prob::BVProblem, alg_::Shooting; odesolve_kwargs = (;
         verbose && @warn "Initial guess provided, but will be ignored for Shooting."
         u0 = __extract_u0(prob.u0, prob.p, first(prob.tspan))
     end
+    @assert (iip || isnothing(alg_.optimize)) "Out-of-place constraints don't allow optimization solvers "
 
     alg = concretize_jacobian_algorithm(alg_, prob)
     (; diffmode) = alg.jac_alg
