Merge pull request #120 from avik-pal/ap/ms_improvements

Minor Updates to Multiple Shooting

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -10,13 +10,15 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
         group:
           - Shooting
           - MIRK
           - Others
         version:
           - '1'
+          - '~1.10.0-0'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

src/solve/multiple_shooting.jl

@@ -9,10 +9,10 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
     iip, bc, u0, u0_size = isinplace(prob), prob.f.bc, deepcopy(u0), size(u0)
 
     __alg = concretize_jacobian_algorithm(_alg, prob)
-    alg = if has_initial_guess && Nig != __alg.nshoots + 1
+    alg = if has_initial_guess && Nig != __alg.nshoots
         verbose &&
-            @warn "Initial guess length != `nshoots + 1`! Adapting to `nshoots = $(Nig - 1)`"
-        update_nshoots(__alg, Nig - 1)
+            @warn "Initial guess length != `nshoots + 1`! Adapting to `nshoots = $(Nig)`"
+        update_nshoots(__alg, Nig)
     else
         __alg
     end
@@ -57,13 +57,7 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
     compute_bc_residual! = if prob.problem_type isa TwoPointBVProblem
         @views function compute_bc_residual_tp!(resid_bc, us::ArrayPartition, p,
             cur_nshoots, nodes, resid_nodes::Union{Nothing, MaybeDiffCache} = nothing)
-            ua, ub0 = us.x
-            # Just Recompute the last ODE Solution
-            lastodeprob = ODEProblem{iip}(f, reshape(ub0, u0_size),
-                (nodes[end - 1], nodes[end]), p)
-            sol_ode_last = __solve(lastodeprob, alg.ode_alg; odesolve_kwargs..., verbose,
-                kwargs..., save_everystep = false, saveat = (), save_end = true)
-            ub = vec(sol_ode_last.u[end])
+            ua, ub = us.x
 
             resid_bc_a, resid_bc_b = if resid_bc isa ArrayPartition
                 resid_bc.x
@@ -147,7 +141,7 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
             resida, residb = resid_bc.x
             J_bc[1:length(resida), 1:N] .= J_bc′[1:length(resida), 1:N]
             idxᵢ = (length(resida) + 1):(length(resida) + length(residb))
-            J_bc[idxᵢ, (end - 2N + 1):(end - N)] .= J_bc′[idxᵢ, (end - N + 1):end]
+            J_bc[idxᵢ, (end - N + 1):end] .= J_bc′[idxᵢ, (end - N + 1):end]
 
             return nothing
         end
@@ -215,7 +209,8 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
 
             bc_jac_cache = (bc_jac_cache_partial, init_jacobian(bc_jac_cache_partial))
 
-            jac_prototype = if @isdefined(J_full)
+            jac_prototype = if alg.jac_alg.nonbc_diffmode isa AbstractSparseADType ||
+                               alg.jac_alg.bc_diffmode isa AbstractSparseADType
                 J_full
             else
                 __zeros_like(u_at_nodes, length(resid_prototype), length(u_at_nodes))
@@ -236,7 +231,7 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
                 nodes); resid_prototype, jac = jac_fn, jac_prototype)
         nlprob = NonlinearProblem(loss_function!, u_at_nodes, prob.p)
         sol_nlsolve = __solve(nlprob, alg.nlsolve; nlsolve_kwargs..., verbose, kwargs...)
-        # u_at_nodes = sol_nlsolve.u
+        u_at_nodes = sol_nlsolve.u::typeof(u0)
     end
 
     single_shooting_prob = remake(prob; u0 = reshape(u_at_nodes[1:N], u0_size))

src/sparse_jacobians.jl

@@ -198,7 +198,7 @@ function __generate_sparse_jacobian_prototype(alg::MultipleShooting, ::TwoPointB
 
     J_full[(L₁ + L₂ + 1):end, :] .= J_c.M
     J_full[1:L₁, 1:N] .= J_bc.M[1:L₁, 1:N]
-    J_full[(L₁ + 1):(L₁ + L₂), (end - 2N + 1):(end - N)] .= J_bc.M[(L₁ + 1):(L₁ + L₂),
+    J_full[(L₁ + 1):(L₁ + L₂), (end - N + 1):end] .= J_bc.M[(L₁ + 1):(L₁ + L₂),
         (N + 1):(2N)]
 
     return J_full, J_c, J_bc

test/interpolation_test.jl renamed to test/mirk/interpolation_test.jl

@@ -14,7 +14,8 @@ function prob_bvp_linear_bc!(res, u, p, t)
     res[1] = u[1][1] - 1
     res[2] = u[end][1]
 end
-prob_bvp_linear_function = ODEFunction(prob_bvp_linear_f!, analytic = prob_bvp_linear_analytic)
+prob_bvp_linear_function = ODEFunction(prob_bvp_linear_f!,
+    analytic = prob_bvp_linear_analytic)
 prob_bvp_linear_tspan = (0.0, 1.0)
 prob_bvp_linear = BVProblem(prob_bvp_linear_function, prob_bvp_linear_bc!,
     [1.0, 0.0], prob_bvp_linear_tspan, λ)
@@ -28,6 +29,6 @@ end
 @testset "Interpolation" begin
     @testset "MIRK$order" for order in (2, 3, 4, 5, 6)
         @time sol = solve(prob_bvp_linear, mirk_solver(Val(order)); dt = 0.001)
-        @test sol(0.001) ≈ [0.998687464, -1.312035941] atol=testTol
+        @test sol(0.001)≈[0.998687464, -1.312035941] atol=testTol
     end
 end

test/runtests.jl

@@ -11,8 +11,11 @@ const GROUP = uppercase(get(ENV, "GROUP", "ALL"))
             @time @safetestset "Ray Tracing BVP" begin
                 include("shooting/ray_tracing.jl")
             end
-            @time @safetestset "Orbital" begin
-                include("shooting/orbital.jl")
+            @static if VERSION ≥ v"1.10.0-beta2"
+                # Orbital Tests take extremely long to compile on Julia 1.9
+                @time @safetestset "Orbital" begin
+                    include("shooting/orbital.jl")
+                end
             end
         end
     end
@@ -28,6 +31,9 @@ const GROUP = uppercase(get(ENV, "GROUP", "ALL"))
             @time @safetestset "Vector of Vector" begin
                 include("mirk/vectorofvector_initials.jl")
             end
+            @time @safetestset "Interpolation Tests" begin
+                include("mirk/interpolation_test.jl")
+            end
         end
     end
 
@@ -46,10 +52,4 @@ const GROUP = uppercase(get(ENV, "GROUP", "ALL"))
             end
         end
     end
-    
-    @time @testset "Interpolation Tests" begin
-        @time @safetestset "MIRK Interpolation Test" begin
-            include("interpolation_test.jl")
-        end
-    end
 end