Merge pull request #305 from SciML/qqy/nested_nlsolve_kwargs

Author: ErikQQY

docs/src/solvers/firk.md

@@ -12,9 +12,11 @@ solve(prob::BVProblem, alg, dt; kwargs...)
 solve(prob::TwoPointBVProblem, alg, dt; kwargs...)
 ```
 
-!!! note "Nested nonlinear solving in FIRK methods"
-    
-    When encountered with large BVP system, setting `nested_nlsolve` to `true` enables FIRK methods to use nested nonlinear solving for the implicit FIRK step instead of solving as a part of the global residual(when default as `nested_nlsolve=false`),
+## Nested nonlinear solving in FIRK methods
+
+When working with large boundary value problems, especially those involving stiff systems, computational efficiency and solver robustness become critical concerns. To improve the efficiency of FIRK methods on large BVPs, we can use nested nonlinear solving to obtain the implicit FIRK step instead of solving them as part of the global residual. In BoundaryValueDiffEq.jl, we can set `nested_nlsolve` as `true` to enable FIRK methods to compute the implicit FIRK steps using nested nonlinear solving(default option in FIRK methods is `nested_nlsolve=false`).
+
+Moreover, the nested nonlinear problem solver can be finely tuned to meet specific accuracy requirements by providing detailed keyword arguments through the `nested_nlsolve_kwargs` option in any FIRK solver, for example, `RadauIIa5(; nested_nlsolve = true, nested_nlsolve_kwargs = (; abstol = 1e-6, reltol = 1e-6))`, where `nested_nlsolve_kwargs` can be any common keyword arguments in NonlinearSolve.jl, see [Common Solver Options in NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/).
 
 ## Full List of Methods
 

lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl

@@ -49,7 +49,7 @@ end
 
 @views function interp_eval!(
         y::AbstractArray, cache::FIRKCacheNested{iip, T}, t, mesh, mesh_dt) where {iip, T}
-    (; f, ITU, nest_prob, nest_tol, alg) = cache
+    (; f, ITU, nest_prob, alg) = cache
     (; q_coeff) = ITU
 
     j = interval(mesh, t)
@@ -82,7 +82,7 @@ end
     nestprob_p[3:end] .= yᵢ
 
     _nestprob = remake(nest_prob, p = nestprob_p)
-    nestsol = __solve(_nestprob, nest_nlsolve_alg; abstol = nest_tol)
+    nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
     K = nestsol.u
 
     z₁, z₁′ = eval_q(yᵢ, 0.5, h, q_coeff, K) # Evaluate q(x) at midpoints
@@ -325,7 +325,7 @@ an interpolant
 end
 
 @views function defect_estimate!(cache::FIRKCacheNested{iip, T}) where {iip, T}
-    (; f, mesh, mesh_dt, defect, ITU, nest_prob, nest_tol) = cache
+    (; f, mesh, mesh_dt, defect, ITU, nest_prob, alg) = cache
     (; q_coeff, τ_star) = ITU
 
     nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, cache.alg.nlsolve)
@@ -347,7 +347,7 @@ end
         nestprob_p[3:end] .= yᵢ₁
 
         _nestprob = remake(nest_prob, p = nestprob_p)
-        nest_sol = __solve(_nestprob, nlsolve_alg; abstol = nest_tol)
+        nest_sol = __solve(_nestprob, nlsolve_alg; alg.nested_nlsolve_kwargs...)
 
         # Defect estimate from q(x) at y_i + τ* * h
         z₁, z₁′ = eval_q(yᵢ₁, τ_star, h, q_coeff, nest_sol.u)

lib/BoundaryValueDiffEqFIRK/src/algorithms.jl

@@ -84,13 +84,13 @@ for stage in (1, 2, 3, 5, 7)
             nlsolve::N = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             nested_nlsolve::Bool = false
-            nest_tol::Union{Number, Nothing} = nothing
+            nested_nlsolve_kwargs::NamedTuple = (;)
             defect_threshold::T = 0.1
             max_num_subintervals::Int = 3000
         end
-        $(alg)(nlsolve::N, jac_alg::J; nested = false, nest_tol::Union{Number, Nothing} = nothing, defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
-            N, J, T}(
-            nlsolve, jac_alg, nested, nest_tol, defect_threshold, max_num_subintervals)
+        $(alg)(nlsolve::N, jac_alg::J; nested = false, nested_nlsolve_kwargs::NamedTuple = (;), defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
+            N, J, T}(nlsolve, jac_alg, nested, nested_nlsolve_kwargs,
+            defect_threshold, max_num_subintervals)
     end
 end
 
@@ -178,13 +178,13 @@ for stage in (2, 3, 4, 5)
             nlsolve::N = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             nested_nlsolve::Bool = false
-            nest_tol::Union{Number, Nothing} = nothing
+            nested_nlsolve_kwargs::NamedTuple = (;)
             defect_threshold::T = 0.1
             max_num_subintervals::Int = 3000
         end
-        $(alg)(nlsolve::N, jac_alg::J; nested = false, nest_tol::Union{Number, Nothing} = nothing, defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
-            N, J, T}(
-            nlsolve, jac_alg, nested, nest_tol, defect_threshold, max_num_subintervals)
+        $(alg)(nlsolve::N, jac_alg::J; nested = false, nested_nlsolve_kwargs::NamedTuple = (;), defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
+            N, J, T}(nlsolve, jac_alg, nested, nested_nlsolve_kwargs,
+            defect_threshold, max_num_subintervals)
     end
 end
 
@@ -272,13 +272,13 @@ for stage in (2, 3, 4, 5)
             nlsolve::N = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             nested_nlsolve::Bool = false
-            nest_tol::Union{Number, Nothing} = nothing
+            nested_nlsolve_kwargs::NamedTuple = (;)
             defect_threshold::T = 0.1
             max_num_subintervals::Int = 3000
         end
-        $(alg)(nlsolve::N, jac_alg::J; nested = false, nest_tol::Union{Number, Nothing} = nothing, defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
-            N, J, T}(
-            nlsolve, jac_alg, nested, nest_tol, defect_threshold, max_num_subintervals)
+        $(alg)(nlsolve::N, jac_alg::J; nested = false, nested_nlsolve_kwargs::NamedTuple = (;), defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
+            N, J, T}(nlsolve, jac_alg, nested, nested_nlsolve_kwargs,
+            defect_threshold, max_num_subintervals)
     end
 end
 
@@ -366,13 +366,13 @@ for stage in (2, 3, 4, 5)
             nlsolve::N = nothing
             jac_alg::J = BVPJacobianAlgorithm()
             nested_nlsolve::Bool = false
-            nest_tol::Union{Number, Nothing} = nothing
+            nested_nlsolve_kwargs::NamedTuple = (;)
             defect_threshold::T = 0.1
             max_num_subintervals::Int = 3000
         end
-        $(alg)(nlsolve::N, jac_alg::J; nested = false, nest_tol::Union{Number, Nothing} = nothing, defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
-            N, J, T}(
-            nlsolve, jac_alg, nested, nest_tol, defect_threshold, max_num_subintervals)
+        $(alg)(nlsolve::N, jac_alg::J; nested = false, nested_nlsolve_kwargs::NamedTuple = (;), defect_threshold::T = 0.1, max_num_subintervals::Int = 3000) where {N, J, T} = $(alg){
+            N, J, T}(nlsolve, jac_alg, nested, nested_nlsolve_kwargs,
+            defect_threshold, max_num_subintervals)
     end
 end
 

lib/BoundaryValueDiffEqFIRK/src/collocation.jl

@@ -83,11 +83,11 @@ end
 @views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::FIRKTableau{true},
         y, u, p, mesh, mesh_dt, stage::Int, cache)
     (; b) = TU
-    (; nest_prob, nest_tol) = cache
+    (; nest_prob, alg) = cache
 
     T = eltype(u)
     nestprob_p = vcat(T(mesh[1]), T(mesh_dt[1]), get_tmp(y[1], u))
-    nest_nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, cache.alg.nlsolve)
+    nest_nlsolve_alg = __concrete_nonlinearsolve_algorithm(nest_prob, alg.nlsolve)
 
     for i in eachindex(k_discrete)
         residᵢ = residual[i]
@@ -103,7 +103,7 @@ end
         K = get_tmp(k_discrete[i], u)
 
         _nestprob = remake(nest_prob, p = nestprob_p)
-        nestsol = solve(_nestprob, nest_nlsolve_alg; abstol = nest_tol)
+        nestsol = solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
         @. K = nestsol.u
         @. residᵢ = yᵢ₊₁ - yᵢ
         __maybe_matmul!(residᵢ, nestsol.u, b, -h, T(1))
@@ -159,7 +159,7 @@ end
 @views function Φ(fᵢ_cache, k_discrete, f!, TU::FIRKTableau{true},
         y, u, p, mesh, mesh_dt, stage::Int, cache)
     (; b) = TU
-    (; nest_prob, alg, nest_tol) = cache
+    (; nest_prob, alg) = cache
 
     residuals = [safe_similar(yᵢ) for yᵢ in y[1:(end - 1)]]
 
@@ -179,7 +179,7 @@ end
         nestprob_p[3:end] = yᵢ
 
         _nestprob = remake(nest_prob, p = nestprob_p)
-        nestsol = solve(_nestprob, nest_nlsolve_alg, abstol = nest_tol)
+        nestsol = solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
 
         @. residᵢ = yᵢ₊₁ - yᵢ
         __maybe_matmul!(residᵢ, nestsol.u, b, -h, T(1))

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -24,7 +24,6 @@
     fᵢ₂_cache
     defect
     nest_prob
-    nest_tol
     resid_size
     kwargs
 end
@@ -165,7 +164,6 @@ function init_nested(prob::BVProblem, alg::AbstractFIRK; dt = 0.0, abstol = 1e-3
     K0 = __K0_on_u0(prob.u0, stage) # Somewhat arbitrary initialization of K
 
     nestprob_p = zeros(T, M + 2)
-    nest_tol = alg.nest_tol
 
     if iip
         nestprob = NonlinearProblem(
@@ -176,10 +174,9 @@ function init_nested(prob::BVProblem, alg::AbstractFIRK; dt = 0.0, abstol = 1e-3
     end
 
     return FIRKCacheNested{iip, T}(
-        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type,
-        prob.p, alg, TU, ITU, bcresid_prototype, mesh, mesh_dt,
-        k_discrete, y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, nestprob,
-        nest_tol, resid₁_size, (; abstol, dt, adaptive, kwargs...))
+        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p, alg,
+        TU, ITU, bcresid_prototype, mesh, mesh_dt, k_discrete, y, y₀, residual, fᵢ_cache,
+        fᵢ₂_cache, defect, nestprob, resid₁_size, (; abstol, dt, adaptive, kwargs...))
 end
 
 function init_expanded(prob::BVProblem, alg::AbstractFIRK; dt = 0.0, abstol = 1e-3,

lib/BoundaryValueDiffEqFIRK/src/interpolation.jl

@@ -180,7 +180,7 @@ nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
 # Nested FIRK
 function (s::EvalSol{C})(tval::Number) where {C <: FIRKCacheNested}
     (; t, u, cache) = s
-    (; f, nest_prob, nest_tol, alg, mesh_dt, p, ITU) = cache
+    (; f, nest_prob, alg, mesh_dt, p, ITU) = cache
     (; q_coeff) = ITU
     stage = alg_stage(alg)
     # Quick handle for the case where tval is at the boundary
@@ -212,7 +212,7 @@ function (s::EvalSol{C})(tval::Number) where {C <: FIRKCacheNested}
     nestprob_p[3:end] .= nodual_value(yᵢ)
 
     _nestprob = remake(nest_prob, p = nestprob_p)
-    nestsol = __solve(_nestprob, nest_nlsolve_alg; abstol = nest_tol)
+    nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
     K = nestsol.u
 
     z₁, z₁′ = eval_q(yᵢ, 0.5, h, q_coeff, K) # Evaluate q(x) at midpoints

lib/BoundaryValueDiffEqFIRK/test/nested/firk_basic_tests.jl

@@ -443,3 +443,50 @@ end =#
         (0, pi / 2), pi / 2; bcresid_prototype = (zeros(1), zeros(1)))
     SciMLBase.successful_retcode(solve(bvp5, RadauIIa5(; nested_nlsolve = true), dt = 0.05))
 end
+
+@testitem "Nested nlsolve kwargs in FIRK" setup=[FIRKNestedConvergenceTests] begin
+    tspan = (0.0, π / 2)
+    function simplependulum!(du, u, p, t)
+        g, L, θ, dθ = 9.81, 1.0, u[1], u[2]
+        du[1] = dθ
+        du[2] = -(g / L) * sin(θ)
+    end
+
+    function bc_pendulum!(residual, u, p, t)
+        residual[1] = u(pi / 4)[1] + π / 2
+        residual[2] = u(pi / 2)[1] - π / 2
+    end
+
+    u0 = [pi / 2, pi / 2]
+    prob = BVProblem(simplependulum!, bc_pendulum!, u0, tspan)
+    nested = true
+    nested_nlsolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)
+
+    @testset "RadauIIa$stage" for stage in (2, 3, 5, 7)
+        @test_nowarn solve(prob,
+            radau_solver(Val(stage); nested_nlsolve = nested,
+                nested_nlsolve_kwargs = nested_nlsolve_kwargs);
+            dt = 0.005)
+    end
+
+    @testset "LobattoIIIa$stage" for stage in (3, 4, 5)
+        @test_nowarn solve(prob,
+            lobattoIIIa_solver(Val(stage); nested_nlsolve = nested,
+                nested_nlsolve_kwargs = nested_nlsolve_kwargs);
+            dt = 0.005)
+    end
+
+    @testset "LobattoIIIb$stage" for stage in (3, 4, 5)
+        @test_nowarn solve(prob,
+            lobattoIIIb_solver(Val(stage); nested_nlsolve = nested,
+                nested_nlsolve_kwargs = nested_nlsolve_kwargs);
+            dt = 0.005)
+    end
+
+    @testset "LobattoIIIc$stage" for stage in (3, 4, 5)
+        @test_nowarn solve(prob,
+            lobattoIIIb_solver(Val(stage); nested_nlsolve = nested,
+                nested_nlsolve_kwargs = nested_nlsolve_kwargs);
+            dt = 0.005)
+    end
+end