Make MultipleShooting more robust with different initial guess

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -301,18 +301,22 @@ end
 
 function __initial_guess(f::F, p::P, t::T; fit_parameters = false) where {F, P, T}
     if hasmethod(f, Tuple{P, T})
-        p isa SciMLBase.NullParameters &&
-            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
-        fit_parameters && return vcat(f(p, t), p)
+        if fit_parameters
+            p isa SciMLBase.NullParameters &&
+                throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+            return vcat(f(p, t), p)
+        end
         return f(p, t)
     elseif hasmethod(f, Tuple{T})
         Base.depwarn("initial guess function must take 2 inputs `(p, t)` instead of just \
                      `t`. The single argument version has been deprecated and will be \
                      removed in the next major release of SciMLBase.",
             :__initial_guess)
-        p isa SciMLBase.NullParameters &&
-            throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
-        fit_parameters && return vcat(f(t), p)
+        if fit_parameters
+            p isa SciMLBase.NullParameters &&
+                throw(ArgumentError("`fit_parameters` is true but `prob.p` is not set."))
+            return vcat(f(t), p)
+        end
         return f(t)
     else
         throw(ArgumentError("`initial_guess` must be a function of the form `f(p, t)`"))

lib/BoundaryValueDiffEqFIRK/test/expanded/firk_basic_tests.jl

@@ -471,7 +471,7 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, RadauIIa5(), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 
     tspan = (0.0, pi)
     function f!(du, u, p, t)
@@ -490,5 +490,5 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, RadauIIa5(), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 end

lib/BoundaryValueDiffEqFIRK/test/nested/firk_basic_tests.jl

@@ -540,7 +540,7 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, RadauIIa5(; nested_nlsolve = true), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 
     tspan = (0.0, pi)
     function f!(du, u, p, t)
@@ -559,5 +559,5 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, RadauIIa5(; nested_nlsolve = true), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 end

lib/BoundaryValueDiffEqMIRK/test/mirk_basic_tests.jl

@@ -441,7 +441,7 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, MIRK4(), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 
     tspan = (0.0, pi)
     function f!(du, u, p, t)
@@ -460,5 +460,5 @@ end
         bcresid_prototype = (zeros(2), zeros(1)), fit_parameters = true)
     sol = solve(bvp, MIRK4(), dt = 0.05)
 
-    @test sol.prob.p ≈ [17.09658] atol=1e-5
+    @test sol.prob.p≈[17.09658] atol=1e-5
 end

lib/BoundaryValueDiffEqShooting/src/multiple_shooting.jl

@@ -406,18 +406,15 @@ end
 # Problem has initial guess
 @views function __multiple_shooting_initialize!(
         nodes, prob, alg, ::Val{true}, nshoots::Int, odecache; kwargs...)
-    (; u0, tspan) = prob
+    (; u0, tspan, p) = prob
 
     resize!(nodes, nshoots + 1)
     nodes .= range(tspan[1], tspan[2]; length = nshoots + 1)
-
-    # NOTE: We don't check `u0 isa Function` since `u0` in-principle can be a callable
-    #       struct
-    u0_ = u0 isa VectorOfArray ? u0 : [__initial_guess(u0, prob.p, t) for t in nodes]
+    u0_ = __initial_guess_on_mesh(u0, nodes, p)
 
     N = length(first(u0_))
     u_at_nodes = similar(first(u0_), (nshoots + 1) * N)
-    recursive_flatten!(u_at_nodes, u0_)
+    recursive_flatten!(u_at_nodes, u0_.u)
 
     return u_at_nodes
 end

lib/BoundaryValueDiffEqShooting/test/basic_problems_tests.jl

@@ -366,3 +366,50 @@ end
         @test norm(sol.resid, Inf) < 1e-6
     end
 end
+
+@testitem "Shooting with heterogeneous initial guess" begin
+    using BoundaryValueDiffEqShooting, OrdinaryDiffEqVerner, LinearAlgebra
+    g = 9.81
+    L = 1.0
+    tspan = (0.0, pi / 2.0)
+    function simplependulum!(du, u, p, t)
+        θ = u[1]
+        dθ = u[2]
+        du[1] = dθ
+        du[2] = -(g / L) * sin(θ)
+    end
+
+    function bc2a!(resid_a, u_a, p)
+        resid_a[1] = u_a[1] + pi / 2
+    end
+    function bc2b!(resid_b, u_b, p)
+        resid_b[1] = u_b[1] - pi / 2
+    end
+
+    function initialGuess(p, t)
+        u1 = t * (-3.1444 * t + 6.9895) - 1.5708
+        u2 = t * (t * (6.3342 * t - 19.208) + 9.6822) + 3.9882
+        return [u1, u2]
+    end
+
+    u0 = [pi / 2, pi / 2]
+    bvp2 = TwoPointBVProblem(simplependulum!, (bc2a!, bc2b!), u0, tspan;
+        bcresid_prototype = (zeros(1), zeros(1)))
+    sol2 = solve(bvp2, MultipleShooting(5, Vern7()))
+    @test SciMLBase.successful_retcode(sol2)
+
+    bvp3 = TwoPointBVProblem(simplependulum!, (bc2a!, bc2b!), sol2, tspan;
+        bcresid_prototype = (zeros(1), zeros(1)))
+    sol3 = solve(bvp3, MultipleShooting(5, Vern7()))
+    @test SciMLBase.successful_retcode(sol3)
+
+    bvp4 = TwoPointBVProblem(simplependulum!, (bc2a!, bc2b!), sol2.u, tspan;
+        bcresid_prototype = (zeros(1), zeros(1)))
+    sol4 = solve(bvp4, MultipleShooting(5, Vern7()))
+    @test SciMLBase.successful_retcode(sol4)
+
+    bvp5 = TwoPointBVProblem(simplependulum!, (bc2a!, bc2b!), initialGuess,
+        tspan; bcresid_prototype = (zeros(1), zeros(1)))
+    sol5 = solve(bvp5, MultipleShooting(5, Vern7()))
+    @test SciMLBase.successful_retcode(sol5)
+end