[WIP] A Faster (Experimental) Version of Multiple Shooting

Project.toml

@@ -13,6 +13,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
@@ -22,6 +23,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
@@ -40,7 +42,7 @@ Aqua = "0.7"
 ArrayInterface = "7"
 BandedMatrices = "1"
 ConcreteStructs = "0.2"
-DiffEqBase = "6.135"
+DiffEqBase = "6.138"
 ForwardDiff = "0.10"
 LinearAlgebra = "1.9"
 LinearSolve = "2"
@@ -56,6 +58,7 @@ SciMLBase = "2.5"
 Setfield = "1"
 SparseArrays = "1.9"
 SparseDiffTools = "2.9"
+Tricks = "0.1"
 TruncatedStacktraces = "1"
 UnPack = "1"
 julia = "1.9"

README.md

@@ -48,6 +48,7 @@ Precompilation can be controlled via `Preferences.jl`
   - `PrecompileMIRK` -- Precompile the MIRK2 - MIRK6 algorithms (default: `true`).
   - `PrecompileShooting` -- Precompile the single shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
   - `PrecompileMultipleShooting` -- Precompile the multiple shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
+    <<<<<<< HEAD
   - `PrecompileMIRKNLLS` -- Precompile the MIRK2 - MIRK6 algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above).
   - `PrecompileShootingNLLS` -- Precompile the single shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.
   - `PrecompileMultipleShootingNLLS` -- Precompile the multiple shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.

src/BoundaryValueDiffEq.jl

@@ -5,7 +5,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
 @recompile_invalidations begin
     using ADTypes, Adapt, BandedMatrices, DiffEqBase, ForwardDiff, LinearAlgebra,
         NonlinearSolve, PreallocationTools, Preferences, RecursiveArrayTools, Reexport,
-        SciMLBase, Setfield, SparseArrays, SparseDiffTools
+        SciMLBase, Setfield, SparseArrays, SparseDiffTools, Tricks
 
     import ADTypes: AbstractADType
     import ArrayInterface: matrix_colors,

src/algorithms.jl

@@ -62,7 +62,8 @@
 
 """
     MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(),
-        grid_coarsening = true, jac_alg = BVPJacobianAlgorithm())
+        grid_coarsening = true, jac_alg = BVPJacobianAlgorithm(),
+        auto_static_nodes::Val = Val(false))
 
 Multiple Shooting method, reduces BVP to an initial value problem and solves the IVP.
 Significantly more stable than Single Shooting.
@@ -98,21 +99,28 @@
       of shooting points. For example, if `nshoots = 10` and
       `grid_coarsening = n -> n ÷ 2`, then the grid will be coarsened to `[5, 2]`.
 
+## Experimental Features
+
+  - `auto_static_nodes`: Automatically detect the timepoints used in the boundary condition
+    and use a faster version of the algorithm! This particular keyword argument should be
+    considered experimental and should be used with care! (Note that we ignore
+    `grid_coarsening` if this is set to `Val(true)`. We plan to support this in the future.)
+
 !!! note
     For type-stability, the chunksizes for ForwardDiff ADTypes in `BVPJacobianAlgorithm`
     must be provided.
 """
-@concrete struct MultipleShooting{J <: BVPJacobianAlgorithm}
+@concrete struct MultipleShooting{S, J <: BVPJacobianAlgorithm}
     ode_alg
     nlsolve
     jac_alg::J
     nshoots::Int
     grid_coarsening
 end
 
-function concretize_jacobian_algorithm(alg::MultipleShooting, prob)
+function concretize_jacobian_algorithm(alg::MultipleShooting{S}, prob) where {S}
     jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
-    return MultipleShooting(alg.ode_alg, alg.nlsolve, jac_alg, alg.nshoots,
+    return MultipleShooting{S}(alg.ode_alg, alg.nlsolve, jac_alg, alg.nshoots,
         alg.grid_coarsening)
 end
 
@@ -121,17 +129,29 @@
         alg.grid_coarsening)
 end
 
+function __without_static_nodes(ms::MultipleShooting{S}) where {S}
+    return MultipleShooting{false}(ms.ode_alg, ms.nlsolve, ms.jac_alg, ms.nshoots,
+        ms.grid_coarsening)
+end
+
 function MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(),
-        grid_coarsening = true, jac_alg = BVPJacobianAlgorithm())
-    @assert grid_coarsening isa Bool || grid_coarsening isa Function ||
-            grid_coarsening isa AbstractVector{<:Integer} ||
-            grid_coarsening isa NTuple{N, <:Integer} where {N}
+        grid_coarsening = missing, jac_alg = BVPJacobianAlgorithm(),
+        auto_static_nodes::Val{S} = Val(false)) where {S}
+    @assert S isa Bool "`auto_static_nodes` must be either `Val(true)` or `Val(false)`."
+    if S
+        @assert grid_coarsening === missing||(grid_coarsening isa Bool && !grid_coarsening) "`auto_static_nodes` doesn't support grid_coarsening."
+    else
+        grid_coarsening === missing && (grid_coarsening = false)
+        @assert grid_coarsening isa Bool || grid_coarsening isa Function ||
+                grid_coarsening isa AbstractVector{<:Integer} ||
+                grid_coarsening isa NTuple{N, <:Integer} where {N}
+    end
     grid_coarsening isa Tuple && (grid_coarsening = Vector(grid_coarsening...))
     if grid_coarsening isa AbstractVector
         sort!(grid_coarsening; rev = true)
         @assert all(grid_coarsening .> 0) && 1 ∉ grid_coarsening
     end
-    return MultipleShooting(ode_alg, nlsolve, jac_alg, nshoots, grid_coarsening)
+    return MultipleShooting{S}(ode_alg, nlsolve, jac_alg, nshoots, grid_coarsening)
 end
 
 for order in (2, 3, 4, 5, 6)

src/solve/mirk.jl

@@ -35,18 +35,20 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0,
         abstol = 1e-3, adaptive = true, kwargs...)
     @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
     iip = isinplace(prob)
+
     _, T, M, n, X = __extract_problem_details(prob; dt, check_positive_dt = true)
+    # NOTE: Assumes the user provided initial guess is on a uniform mesh
+    mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1))
+
+    mesh_dt = diff(mesh)
+
     chunksize = pickchunksize(M * (n + 1))
 
     __alloc = x -> __maybe_allocate_diffcache(vec(x), chunksize, alg.jac_alg)
 
     fᵢ_cache = __alloc(similar(X))
     fᵢ₂_cache = vec(similar(X))
 
-    # NOTE: Assumes the user provided initial guess is on a uniform mesh
-    mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1))
-    mesh_dt = diff(mesh)
-
     defect_threshold = T(0.1)  # TODO: Allow user to specify these
     MxNsub = 3000              # TODO: Allow user to specify these
 
@@ -100,7 +102,9 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0,
         vecf, vecbc
     end
 
-    return MIRKCache{iip, T}(alg_order(alg), stage, M, size(X), f, bc, prob,
+    prob_ = !(prob.u0 isa AbstractArray) ? remake(prob; u0 = X) : prob
+
+    return MIRKCache{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, k_interp, y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, new_stages,
         resid₁_size, (; defect_threshold, MxNsub, abstol, dt, adaptive, kwargs...))