Aggregate Single Shooting Improvements from #150

.JuliaFormatter.toml

@@ -1,3 +1,4 @@
 style = "sciml"
 annotate_untyped_fields_with_any = false
 format_markdown = true
+format_docstrings = true

Project.toml

@@ -1,6 +1,6 @@
 name = "BoundaryValueDiffEq"
 uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
-version = "5.6.2"
+version = "5.6.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -10,6 +10,7 @@ BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
@@ -37,13 +38,14 @@ BoundaryValueDiffEqOrdinaryDiffEqExt = "OrdinaryDiffEq"
 
 [compat]
 ADTypes = "0.2"
-Adapt = "3"
-Aqua = "0.7"
+Adapt = "3, 4"
+Aqua = "0.8"
 ArrayInterface = "7"
 BandedMatrices = "1"
 ConcreteStructs = "0.2"
 DiffEqBase = "6.145"
 FastAlmostBandedMatrices = "0.1"
+FastClosures = "0.3"
 ForwardDiff = "0.10"
 LinearAlgebra = "1.9"
 LinearSolve = "2.20"
@@ -53,12 +55,16 @@ OrdinaryDiffEq = "6"
 PreallocationTools = "0.4"
 PrecompileTools = "1"
 Preferences = "1"
+Random = "1"
 RecursiveArrayTools = "2.38.10, 3"
 Reexport = "0.2, 1.0"
-SciMLBase = "2.5"
+SafeTestsets = "0.1"
+SciMLBase = "2.12"
 Setfield = "1"
 SparseArrays = "1.9"
 SparseDiffTools = "2.9"
+StaticArrays = "1.8.1"
+Test = "1"
 Tricks = "0.1"
 TruncatedStacktraces = "1"
 UnPack = "1"
@@ -71,9 +77,10 @@ LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua", "LinearSolve"]
+test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua", "LinearSolve", "RecursiveArrayTools"]

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

@@ -35,25 +35,25 @@ end
     bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))
 
     probs = [
-        BVProblem(f1!, bc1!, u0, tspan),
-        BVProblem(f1, bc1, u0, tspan),
-        TwoPointBVProblem(f1!, (bc1_a!, bc1_b!), u0, tspan; bcresid_prototype),
-        TwoPointBVProblem(f1, (bc1_a, bc1_b), u0, tspan; bcresid_prototype),
+        BVProblem(BVPFunction{true}(f1!, bc1!), u0, tspan; nlls = Val(false)),
+        BVProblem(BVPFunction{false}(f1, bc1), u0, tspan; nlls = Val(false)),
+        BVProblem(BVPFunction{true}(f1!, (bc1_a!, bc1_b!); bcresid_prototype,
+                twopoint = Val(true)), u0, tspan; nlls = Val(false)),
+        BVProblem(BVPFunction{false}(f1, (bc1_a, bc1_b); bcresid_prototype,
+                twopoint = Val(true)), u0, tspan; nlls = Val(false)),
     ]
 
     algs = []
 
-    if load_preference(BoundaryValueDiffEq, "PrecompileShooting", false)
+    if load_preference(BoundaryValueDiffEq, "PrecompileShooting", true)
         push!(algs,
             Shooting(Tsit5(); nlsolve = NewtonRaphson(),
                 jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))))
     end
 
-    if load_preference(BoundaryValueDiffEq, "PrecompileMultipleShooting", false)
+    if load_preference(BoundaryValueDiffEq, "PrecompileMultipleShooting", true)
         push!(algs,
-            MultipleShooting(10,
-                Tsit5();
-                nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2)),
+            MultipleShooting(10, Tsit5(); nlsolve = NewtonRaphson(),
                 jac_alg = BVPJacobianAlgorithm(;
                     bc_diffmode = AutoForwardDiff(; chunksize = 2),
                     nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))))
@@ -94,39 +94,38 @@ end
     bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))
 
     probs = [
-        BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
-            u0, tspan),
-        BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
-            u0, tspan),
-        TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
-            bcresid_prototype = bcresid_prototype2),
-        TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
-            bcresid_prototype = bcresid_prototype2),
+        BVProblem(BVPFunction{true}(f1_nlls!, bc1_nlls!;
+                bcresid_prototype = bcresid_prototype1), u0, tspan; nlls = Val(true)),
+        BVProblem(BVPFunction{false}(f1_nlls, bc1_nlls;
+                bcresid_prototype = bcresid_prototype1), u0, tspan; nlls = Val(true)),
+        BVProblem(BVPFunction{true}(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!);
+                bcresid_prototype = bcresid_prototype2, twopoint = Val(true)), u0, tspan;
+            nlls = Val(true)),
+        BVProblem(BVPFunction{false}(f1_nlls, (bc1_nlls_a, bc1_nlls_b);
+                bcresid_prototype = bcresid_prototype2, twopoint = Val(true)), u0, tspan;
+            nlls = Val(true)),
     ]
 
     algs = []
 
-    if load_preference(BoundaryValueDiffEq, "PrecompileShootingNLLS", false)
+    if load_preference(BoundaryValueDiffEq, "PrecompileShootingNLLS", true)
         append!(algs,
             [
-                Shooting(Tsit5(); nlsolve = LevenbergMarquardt(),
+                Shooting(Tsit5(); nlsolve = TrustRegion(),
                     jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))),
                 Shooting(Tsit5(); nlsolve = GaussNewton(),
                     jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))),
             ])
     end
 
-    if load_preference(BoundaryValueDiffEq, "PrecompileMultipleShootingNLLS", false)
+    if load_preference(BoundaryValueDiffEq, "PrecompileMultipleShootingNLLS", true)
         append!(algs,
             [
-                MultipleShooting(10, Tsit5();
-                    nlsolve = LevenbergMarquardt(;
-                        autodiff = AutoForwardDiff(chunksize = 2)),
+                MultipleShooting(10, Tsit5(); nlsolve = TrustRegion(),
                     jac_alg = BVPJacobianAlgorithm(;
                         bc_diffmode = AutoForwardDiff(; chunksize = 2),
                         nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
-                MultipleShooting(10, Tsit5();
-                    nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2)),
+                MultipleShooting(10, Tsit5(); nlsolve = GaussNewton(),
                     jac_alg = BVPJacobianAlgorithm(;
                         bc_diffmode = AutoForwardDiff(; chunksize = 2),
                         nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),

src/BoundaryValueDiffEq.jl

@@ -15,6 +15,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
         parameterless_type, undefmatrix, fast_scalar_indexing
     import ConcreteStructs: @concrete
     import DiffEqBase: solve
+    import FastClosures: @closure
     import ForwardDiff: pickchunksize
     import RecursiveArrayTools: ArrayPartition, DiffEqArray
     import SciMLBase: AbstractDiffEqInterpolation, StandardBVProblem, __solve, _unwrap_val

src/adaptivity.jl

@@ -259,14 +259,16 @@ end
     return z, z′
 end
 
+"""
+    interp_weights(τ, alg
+
+interp_weights: solver-specified interpolation weights and its first derivative
+"""
+function interp_weights end
+
 for order in (2, 3, 4, 5, 6)
     alg = Symbol("MIRK$(order)")
     @eval begin
-        """
-            interp_weights(τ, alg)
-
-        interp_weights: solver-specified interpolation weights and its first derivative
-        """
         function interp_weights(τ::T, ::$(alg)) where {T}
             if $(order == 2)
                 w = [0, τ * (1 - τ / 2), τ^2 / 2]

src/algorithms.jl

@@ -2,8 +2,15 @@
 abstract type BoundaryValueDiffEqAlgorithm <: SciMLBase.AbstractBVPAlgorithm end
 abstract type AbstractMIRK <: BoundaryValueDiffEqAlgorithm end
 
+## Disable the ugly verbose printing by default
+function Base.show(io::IO, alg::BoundaryValueDiffEqAlgorithm)
+    print(io, "$(nameof(typeof(alg)))()")
+end
+
 """
-    Shooting(ode_alg = nothing; nlsolve = nothing, jac_alg = BVPJacobianAlgorithm())
+    Shooting(ode_alg; kwargs...)
+    Shooting(ode_alg, nlsolve; kwargs...)
+    Shooting(; ode_alg = nothing, nlsolve = nothing, jac_alg = nothing)
 
 Single shooting method, reduces BVP to an initial value problem and solves the IVP.
 
@@ -12,53 +19,41 @@ Single shooting method, reduces BVP to an initial value problem and solves the I
   - `ode_alg`: ODE algorithm to use for solving the IVP. Any solver which conforms to the
     SciML `ODEProblem` interface can be used! (Defaults to `nothing` which will use
     poly-algorithm if `DifferentialEquations.jl` is loaded else this must be supplied)
-
-## Keyword Arguments
-
   - `nlsolve`: Internal Nonlinear solver. Any solver which conforms to the SciML
     `NonlinearProblem` interface can be used. Note that any autodiff argument for the solver
     will be ignored and a custom jacobian algorithm will be used.
-  - `jac_alg`: Jacobian Algorithm used for the nonlinear solver. Defaults to
-    `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to use based
-    on the input types and problem type. Only `diffmode` is used (defaults to
-    `AutoForwardDiff` if possible else `AutoFiniteDiff`).
-
-!!! note
-    For type-stability, the chunksizes for ForwardDiff ADTypes in `BVPJacobianAlgorithm`
-    must be provided.
+  - `jac_alg`: Jacobian Algorithm used for the Nonlinear Solver. If this is not set, we
+    check if `nlsolve.ad` exists and is not nothing. If it is, we use that to construct
+    the jacobian. If not, we try to use the best algorithm based on the input types
+    and problem type. If `BVPJacobianAlgorithm` is provided, only `diffmode` is used
+    (defaults to `AutoForwardDiff` if possible else `AutoFiniteDiff`).
 """
-struct Shooting{O, N, L <: BVPJacobianAlgorithm} <: BoundaryValueDiffEqAlgorithm
-    ode_alg::O
-    nlsolve::N
-    jac_alg::L
+@concrete struct Shooting{J <: BVPJacobianAlgorithm} <: BoundaryValueDiffEqAlgorithm
+    ode_alg
+    nlsolve
+    jac_alg::J
 end
 
-function concretize_jacobian_algorithm(alg::Shooting, prob)
-    jac_alg = alg.jac_alg
-    diffmode = jac_alg.diffmode === nothing ? __default_nonsparse_ad(prob.u0) :
-               jac_alg.diffmode
-    return Shooting(alg.ode_alg, alg.nlsolve, BVPJacobianAlgorithm(diffmode))
+function Shooting(; ode_alg = nothing, nlsolve = nothing, jac_alg = nothing)
+    return Shooting(ode_alg, nlsolve, __materialize_jacobian_algorithm(nlsolve, jac_alg))
 end
-
-function Shooting(ode_alg = nothing; nlsolve = nothing, jac_alg = nothing)
-    jac_alg === nothing && (jac_alg = __propagate_nlsolve_ad_to_jac_alg(nlsolve))
-    return Shooting(ode_alg, nlsolve, jac_alg)
+@inline Shooting(ode_alg; kwargs...) = Shooting(; ode_alg, kwargs...)
+@inline Shooting(ode_alg, nlsolve; kwargs...) = Shooting(; ode_alg, nlsolve, kwargs...)
+
+function Base.show(io::IO, alg::Shooting)
+    print(io, "Shooting(")
+    modifiers = String[]
+    alg.nlsolve !== nothing && push!(modifiers, "nlsolve = $(alg.nlsolve)")
+    alg.jac_alg !== nothing && push!(modifiers, "jac_alg = $(alg.jac_alg)")
+    alg.ode_alg !== nothing && push!(modifiers, "ode_alg = $(__nameof(alg.ode_alg))()")
+    print(io, join(modifiers, ", "))
+    print(io, ")")
 end
 
-Shooting(ode_alg, nlsolve; jac_alg = nothing) = Shooting(ode_alg; nlsolve, jac_alg)
-
-# This is a deprecation path. We forward the `ad` from nonlinear solver to `jac_alg`.
-# We will drop this function in
-function __propagate_nlsolve_ad_to_jac_alg(nlsolve::N) where {N}
-    # Defaults so no depwarn
-    nlsolve === nothing && return BVPJacobianAlgorithm()
-    ad = hasfield(N, :ad) ? nlsolve.ad : nothing
-    ad === nothing && return BVPJacobianAlgorithm()
-
-    Base.depwarn("Setting autodiff to the nonlinear solver in Shooting has been deprecated \
-                  and will have no effect from the next major release. Update to use \
-                  `BVPJacobianAlgorithm` directly", :Shooting)
-    return BVPJacobianAlgorithm(ad)
+@inline function concretize_jacobian_algorithm(alg::Shooting, prob)
+    alg.jac_alg.diffmode === nothing &&
+        (return @set alg.jac_alg.diffmode = __default_nonsparse_ad(prob.u0))
+    return alg
 end
 
 """
@@ -80,27 +75,31 @@ Significantly more stable than Single Shooting.
   - `nlsolve`: Internal Nonlinear solver. Any solver which conforms to the SciML
     `NonlinearProblem` interface can be used. Note that any autodiff argument for the solver
     will be ignored and a custom jacobian algorithm will be used.
+
   - `jac_alg`: Jacobian Algorithm used for the nonlinear solver. Defaults to
     `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to use based
     on the input types and problem type.
-    - For `TwoPointBVProblem`, only `diffmode` is used (defaults to
-      `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`).
-    - For `BVProblem`, `bc_diffmode` and `nonbc_diffmode` are used. For `nonbc_diffmode`
-      defaults to `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`. For
-      `bc_diffmode`, defaults to `AutoForwardDiff` if possible else `AutoFiniteDiff`.
+
+      + For `TwoPointBVProblem`, only `diffmode` is used (defaults to
+        `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`).
+      + For `BVProblem`, `bc_diffmode` and `nonbc_diffmode` are used. For `nonbc_diffmode`
+        defaults to `AutoSparseForwardDiff` if possible else `AutoSparseFiniteDiff`. For
+        `bc_diffmode`, defaults to `AutoForwardDiff` if possible else `AutoFiniteDiff`.
   - `grid_coarsening`: Coarsening the multiple-shooting grid to generate a stable IVP
     solution. Possible Choices:
-    - `true`: Halve the grid size, till we reach a grid size of 1.
-    - `false`: Do not coarsen the grid. Solve a Multiple Shooting Problem and finally
-      solve a Single Shooting Problem.
-    - `AbstractVector{<:Int}` or `Ntuple{N, <:Integer}`: Use the provided grid coarsening.
-      For example, if `nshoots = 10` and `grid_coarsening = [5, 2]`, then the grid will be
-      coarsened to `[5, 2]`. Note that `1` should not be present in the grid coarsening.
-    - `Function`: Takes the current number of shooting points and returns the next number
-      of shooting points. For example, if `nshoots = 10` and
-      `grid_coarsening = n -> n ÷ 2`, then the grid will be coarsened to `[5, 2]`.
+
+      + `true`: Halve the grid size, till we reach a grid size of 1.
+      + `false`: Do not coarsen the grid. Solve a Multiple Shooting Problem and finally
+        solve a Single Shooting Problem.
+      + `AbstractVector{<:Int}` or `Ntuple{N, <:Integer}`: Use the provided grid coarsening.
+        For example, if `nshoots = 10` and `grid_coarsening = [5, 2]`, then the grid will be
+        coarsened to `[5, 2]`. Note that `1` should not be present in the grid coarsening.
+      + `Function`: Takes the current number of shooting points and returns the next number
+        of shooting points. For example, if `nshoots = 10` and
+        `grid_coarsening = n -> n ÷ 2`, then the grid will be coarsened to `[5, 2]`.
 
 !!! note
+
     For type-stability, the chunksizes for ForwardDiff ADTypes in `BVPJacobianAlgorithm`
     must be provided.
 """
@@ -192,10 +191,12 @@ Fortran code for solving two-point boundary value problems. For detailed documen
 [ODEInterface.jl](https://github.com/luchr/ODEInterface.jl/blob/master/doc/SolverOptions.md#bvpm2).
 
 !!! warning
+
     Only supports inplace two-point boundary value problems, with very limited forms of
     input structures!
 
 !!! note
+
     Only available if the `ODEInterface` package is loaded.
 """
 Base.@kwdef struct BVPM2{S} <: BoundaryValueDiffEqAlgorithm
@@ -217,10 +218,12 @@ For detailed documentation, see
 [ODEInterface.jl](https://github.com/luchr/ODEInterface.jl/blob/master/doc/SolverOptions.md#bvpsol).
 
 !!! warning
+
     Only supports inplace two-point boundary value problems, with very limited forms of
     input structures!
 
 !!! note
+
     Only available if the `ODEInterface` package is loaded.
 """
 Base.@kwdef struct BVPSOL{O} <: BoundaryValueDiffEqAlgorithm