Fix initial guess in collocation solvers

lib/BoundaryValueDiffEqCore/Project.toml

@@ -33,7 +33,7 @@ Adapt = "4.1.1"
 Aqua = "0.8"
 ArrayInterface = "7.18"
 ConcreteStructs = "0.2.3"
-DiffEqBase = "6.183"
+DiffEqBase = "6.213"
 ForwardDiff = "0.10.38, 1"
 Integrals = "4.7.1, 5"
 InteractiveUtils = "<0.0.1, 1"

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -323,8 +323,6 @@ function __extract_problem_details(
     _u0 = first(u0.u)
     _t = u0.t
     if tune_parameters
-        prob.p isa SciMLBase.NullParameters &&
-            throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
         new_u = vcat(_u0, __tunable_part(prob.p))
         return Val(false), eltype(new_u), length(new_u), Int(cld(t₁ - t₀, dt)), new_u
     end
@@ -333,11 +331,7 @@ end
 
 function __initial_guess(f::F, p::P, t::T; tune_parameters = false) where {F, P, T}
     if hasmethod(f, Tuple{P, T})
-        if tune_parameters
-            p isa SciMLBase.NullParameters &&
-                throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
-            return vcat(f(p, t), __tunable_part(p))
-        end
+        tune_parameters && return vcat(f(p, t), __tunable_part(p))
         return f(p, t)
     elseif hasmethod(f, Tuple{T})
         Base.depwarn(
@@ -346,11 +340,7 @@ function __initial_guess(f::F, p::P, t::T; tune_parameters = false) where {F, P,
                      removed in the next major release of SciMLBase.",
             :__initial_guess
         )
-        if tune_parameters
-            p isa SciMLBase.NullParameters &&
-                throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
-            return vcat(f(t), __tunable_part(p))
-        end
+        tune_parameters && return vcat(f(t), __tunable_part(p))
         return f(t)
     else
         throw(ArgumentError("`initial_guess` must be a function of the form `f(p, t)`"))
@@ -621,27 +611,46 @@ initial guess, it returns `vec(u₀)`.
 Returns the initial guess on the mesh. For `DiffEqArray` assumes that the mesh is the same
 as the mesh of the `DiffEqArray`.
 """
-@inline function __initial_guess_on_mesh(u₀::AbstractVector{<:AbstractArray}, _, p)
+@inline function __initial_guess_on_mesh(u₀::AbstractVector{<:AbstractArray}, mesh, p; tune_parameters = false)
+    tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀])
     return VectorOfArray([copy(vec(u)) for u in u₀])
 end
-@inline function __initial_guess_on_mesh(u₀::VectorOfArray, _, p)
+@inline function __initial_guess_on_mesh(u₀::VectorOfArray, mesh, p; tune_parameters = false)
+    tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
     return copy(u₀)
 end
-@inline function __initial_guess_on_mesh(u₀::DiffEqArray, mesh, p)
+@inline function __initial_guess_on_mesh(u₀::DiffEqArray, mesh, p; tune_parameters = false)
+    tune_parameters && return DiffEqArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
     return copy(u₀)
 end
-@inline function __initial_guess_on_mesh(u₀::SciMLBase.ODESolution, mesh, p)
+@inline function __initial_guess_on_mesh(u₀::SciMLBase.ODESolution, mesh, p; tune_parameters = false)
+    tune_parameters && return VectorOfArray([vcat(vec(u), __tunable_part(p)) for u in u₀.u])
     return copy(VectorOfArray(u₀.u))
 end
-@inline function __initial_guess_on_mesh(u₀::AbstractArray, mesh, p)
+@inline function __initial_guess_on_mesh(u₀::AbstractArray, mesh, p; tune_parameters = false)
+    tune_parameters && return VectorOfArray([vcat(vec(u₀), __tunable_part(p)) for _ in mesh])
     return VectorOfArray([copy(vec(u₀)) for _ in mesh])
 end
-@inline function __initial_guess_on_mesh(u₀::F, mesh, p) where {F}
+@inline function __initial_guess_on_mesh(u₀::F, mesh, p; tune_parameters = false) where {F}
+    tune_parameters && return VectorOfArray([vcat(vec(__initial_guess(u₀, p, t)), __tunable_part(p)) for t in mesh])
     return VectorOfArray([vec(__initial_guess(u₀, p, t)) for t in mesh])
 end
+@inline function __initial_guess_on_mesh(u₀::Number, mesh, p; tune_parameters = false)
+    tune_parameters && return VectorOfArray([vcat([u₀], __tunable_part(p)) for _ in mesh])
+    return VectorOfArray([copy([u₀]) for _ in mesh])
+end
 @inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::AbstractArray, Nig, p)
     return VectorOfArray([copy(vec(u₀)) for _ in 1:(2 * (Nig + 1))])
 end
+@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::AbstractVector{<:AbstractVector}, _, p)
+    return VectorOfArray(vcat([copy(vec(u)) for u in u₀], [copy(vec(u)) for u in u₀]))
+end
+@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::VectorOfArray, _, p)
+    return VectorOfArray(vcat(copy(u₀.u), copy(u₀.u)))
+end
+@inline function __initial_guess_on_mesh(prob::SecondOrderBVProblem, u₀::SciMLBase.ODESolution, Nig, p)
+    return VectorOfArray(vcat(copy(VectorOfArray(u₀.u)), copy(VectorOfArray(u₀.u))))
+end
 
 # Construct BVP Solution
 function __build_solution(prob::AbstractBVProblem, odesol, nlsol::SciMLBase.NonlinearSolution)

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -126,6 +126,11 @@ function init_nested(
     end
     diffcache = __cache_trait(alg.jac_alg)
     tune_parameters = haskey(prob.kwargs, :tune_parameters)
+    if tune_parameters
+        prob.p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
+    end
+
     constraint = (!isnothing(prob.f.inequality)) ||
         (!isnothing(prob.f.equality)) ||
         (!isnothing(prob.lb)) ||
@@ -135,18 +140,18 @@ function init_nested(
     ig, T,
         M,
         Nig,
-        X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
+        u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
     mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
     mesh_dt = diff(mesh)
 
     chunksize = pickchunksize(M * (Nig - 1))
     __alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
 
-    fᵢ_cache = __alloc(zero(X))
-    fᵢ₂_cache = vec(zero(X))
+    fᵢ_cache = __alloc(zero(u0))
+    fᵢ₂_cache = vec(zero(u0))
 
     # Don't flatten this here, since we need to expand it later if needed
-    y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
+    y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p; tune_parameters = tune_parameters)
 
     y = __alloc.(copy.(y₀.u))
     TU, ITU = constructRK(alg, T)
@@ -156,17 +161,17 @@ function init_nested(
 
     k_discrete = if !constraint
         [
-            __maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
+            __maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
                 for _ in 1:Nig
         ]
     else
         [
-            __maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
+            __maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
                 for _ in 1:Nig
         ]
     end
 
-    bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, X)
+    bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, u0)
 
     residual = if iip
         if !constraint
@@ -192,12 +197,12 @@ function init_nested(
         nothing
     end
 
-    defect = VectorOfArray([safe_similar(X, ifelse(adaptive, M, 0)) for _ in 1:Nig])
+    defect = VectorOfArray([safe_similar(u0, ifelse(adaptive, M, 0)) for _ in 1:Nig])
 
     # Transform the functions to handle non-vector inputs
     bcresid_prototype = __vec(bcresid_prototype)
     f,
-        bc = if X isa AbstractVector
+        bc = if u0 isa AbstractVector
         if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
             tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
             l_parameters = length(tunable_part)
@@ -220,39 +225,39 @@ function init_nested(
             prob.f, prob.f.bc
         end
     elseif iip
-        vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
+        vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
         vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
-            @closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(X))
+            @closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
         else
             (
                 @closure(
                     (
                         r, u,
                         p,
-                    ) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(X))
+                    ) -> __vec_bc!(r, u, p, first(prob.f.bc), resid₁_size[1], size(u0))
                 ),
                 @closure(
                     (
                         r, u, p,
-                    ) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(X))
+                    ) -> __vec_bc!(r, u, p, last(prob.f.bc), resid₁_size[2], size(u0))
                 ),
             )
         end
         vecf!, vecbc!
     else
-        vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
+        vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
         vecbc = if !(prob.problem_type isa TwoPointBVProblem)
-            @closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
+            @closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
         else
             (
-                @closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(X))),
-                @closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(X))),
+                @closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
+                @closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
             )
         end
         vecf, vecbc
     end
 
-    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
+    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = u0) : prob
 
     # Somewhat arbitrary initialization of K
     K0 = __K0_on_u0(prob, stage; tune_parameters = tune_parameters)
@@ -266,7 +271,7 @@ function init_nested(
     end
 
     return FIRKCacheNested{iip, T, typeof(diffcache), tune_parameters}(
-        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p,
+        alg_order(alg), stage, M, size(u0), f, bc, prob_, prob.problem_type, prob.p,
         alg, TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete,
         y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, nestprob, resid₁_size, prob.singular_term,
         nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...), verbose_spec
@@ -287,6 +292,10 @@ function init_expanded(
     end
     diffcache = __cache_trait(alg.jac_alg)
     tune_parameters = haskey(prob.kwargs, :tune_parameters)
+    if tune_parameters
+        prob.p isa SciMLBase.NullParameters &&
+            throw(ArgumentError("`tune_parameters` is true but `prob.p` is not set."))
+    end
     constraint = (!isnothing(prob.f.inequality)) ||
         (!isnothing(prob.f.equality)) ||
         (!isnothing(prob.lb)) ||
@@ -296,7 +305,7 @@ function init_expanded(
     ig, T,
         M,
         Nig,
-        X = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
+        u0 = __extract_problem_details(prob; dt, check_positive_dt = true, tune_parameters = tune_parameters)
     mesh = __extract_mesh(prob.u0, t₀, t₁, Nig)
     mesh_dt = diff(mesh)
 
@@ -308,27 +317,27 @@ function init_expanded(
     chunksize = pickchunksize(M + M * Nig * (stage + 1))
     __alloc = @closure x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
 
-    fᵢ_cache = __alloc(zero(X)) # Runtime dispatch
-    fᵢ₂_cache = vec(zero(X))
+    fᵢ_cache = __alloc(zero(u0)) # Runtime dispatch
+    fᵢ₂_cache = vec(zero(u0))
 
     # Don't flatten this here, since we need to expand it later if needed
-    _y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
+    _y₀ = __initial_guess_on_mesh(prob.u0, mesh, prob.p; tune_parameters = tune_parameters)
     y₀ = extend_y(_y₀, Nig + 1, stage)
     y = __alloc.(copy.(y₀.u)) # Runtime dispatch
 
     k_discrete = if !constraint
         [
-            __maybe_allocate_diffcache(safe_similar(X, M, stage), chunksize, alg.jac_alg)
+            __maybe_allocate_diffcache(safe_similar(u0, M, stage), chunksize, alg.jac_alg)
                 for _ in 1:Nig
         ] # Runtime dispatch
     else
         [
-            __maybe_allocate_diffcache(safe_similar(X, L_f_prototype, stage), chunksize, alg.jac_alg)
+            __maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
                 for _ in 1:Nig
         ] # Runtime dispatch
     end
 
-    bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, X)
+    bcresid_prototype, resid₁_size = __get_bcresid_prototype(prob.problem_type, prob, u0)
 
     residual = if iip
         if !constraint
@@ -354,12 +363,12 @@ function init_expanded(
         nothing
     end
 
-    defect = VectorOfArray([similar(X, ifelse(adaptive, M, 0)) for _ in 1:Nig])
+    defect = VectorOfArray([similar(u0, ifelse(adaptive, M, 0)) for _ in 1:Nig])
 
     # Transform the functions to handle non-vector inputs
     bcresid_prototype = __vec(bcresid_prototype)
     f,
-        bc = if X isa AbstractVector
+        bc = if u0 isa AbstractVector
         if tune_parameters && SciMLStructures.isscimlstructure(prob.p)
             tunable_part, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), prob.p)
             l_parameters = length(tunable_part)
@@ -382,43 +391,43 @@ function init_expanded(
             prob.f, prob.f.bc
         end
     elseif iip
-        vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(X))
+        vecf! = @closure (du, u, p, t) -> __vec_f!(du, u, p, t, prob.f, size(u0))
         vecbc! = if !(prob.problem_type isa TwoPointBVProblem)
-            @closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(X))
+            @closure (r, u, p, t) -> __vec_bc!(r, u, p, t, prob.f.bc, resid₁_size, size(u0))
         else
             (
                 @closure(
                     (
                         r, u,
                         p,
-                    ) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(X))
+                    ) -> __vec_bc!(r, u, p, first(prob.f.bc)[1], resid₁_size[1], size(u0))
                 ),
                 @closure (
                     (
                         r, u,
                         p,
-                    ) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(X))
+                    ) -> __vec_bc!(r, u, p, last(prob.f.bc)[2], resid₁_size[2], size(u0))
                 )
             )
         end
         vecf!, vecbc!
     else
-        vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(X))
+        vecf = @closure (u, p, t) -> __vec_f(u, p, t, prob.f, size(u0))
         vecbc = if !(prob.problem_type isa TwoPointBVProblem)
-            @closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(X))
+            @closure (u, p, t) -> __vec_bc(u, p, t, prob.f.bc, size(u0))
         else
             (
-                @closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(X))),
-                @closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(X))),
+                @closure((u, p) -> __vec_bc(u, p, first(prob.f.bc), size(u0))),
+                @closure((u, p) -> __vec_bc(u, p, last(prob.f.bc), size(u0))),
             )
         end
         vecf, vecbc
     end
 
-    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
+    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = u0) : prob
 
     return FIRKCacheExpand{iip, T, typeof(diffcache), tune_parameters}(
-        alg_order(alg), stage, M, size(X), f, bc, prob_, prob.problem_type, prob.p,
+        alg_order(alg), stage, M, size(u0), f, bc, prob_, prob.problem_type, prob.p,
         alg, TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete,
         y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, resid₁_size, prob.singular_term, nlsolve_kwargs,
         optimize_kwargs, (; abstol, dt, adaptive, controller, kwargs...), verbose_spec

lib/BoundaryValueDiffEqFIRK/test/expanded/firk_basic_tests.jl

@@ -525,3 +525,46 @@ end
     @test sol_struct.prob.p.params ≈ [17.09658] atol = 1.0e-5
     @test sol_struct.prob.p.params ≈ sol_vec.prob.p atol = 1.0e-10
 end
+
+#=
+# The initial guess for expanded FIRK just stall the CI, need to find out why.
+@testitem "Test initial guess" begin
+    tspan = (0.0, 1.0)
+    function f!(du, u, p, t)
+        cond = 0.002
+        vol_heat = 0.2
+        du[1] = -u[2] / cond
+        du[2] = vol_heat
+        du[3] = 0.0
+    end
+    function bca!(res_a, u_a, p)
+        res_a[1] = u_a[2]
+        res_a[2] = u_a[1] - 100.0
+    end
+    function bcb!(res_b, u_b, p)
+        tref = 20.0
+        res_b[1] = u_b[3] * (u_b[1] - tref) - u_b[2]
+    end
+    u_guess = [
+        [100.0, 0.0, 0.006666666666666668],
+        [99.5, 0.020000000000000004, 0.006666666666666668],
+        [98.0, 0.04000000000000001, 0.006666666666666668],
+        [95.5, 0.060000000000000005, 0.006666666666666668],
+        [92.0, 0.08000000000000002, 0.006666666666666668],
+        [87.5, 0.1, 0.006666666666666668],
+        [82.0, 0.12000000000000001, 0.006666666666666668],
+        [75.5, 0.14, 0.006666666666666668],
+        [68.0, 0.16000000000000003, 0.006666666666666668],
+        [59.49999999999999, 0.18000000000000002, 0.006666666666666668],
+        [50.0, 0.2, 0.006666666666666668],
+    ]
+
+    bvp1 = TwoPointBVProblem(f!, (bca!, bcb!), u_guess, tspan; bcresid_prototype = (zeros(2), zeros(1)))
+    sol1 = solve(bvp1, LobattoIIIa3(), dt = 0.1, adaptive = false, nlsolve_kwargs = (; maxiters = 0))
+    @test sol1.u == u_guess
+
+    bvp2 = TwoPointBVProblem(f!, (bca!, bcb!), sol1, tspan; bcresid_prototype = (zeros(2), zeros(1)))
+    sol2 = solve(bvp2, LobattoIIIa3(), dt = 0.1, adaptive = false, nlsolve_kwargs = (; maxiters = 0))
+    @test sol2.u == u_guess
+end
+=#