Merge pull request #131 from avik-pal/ap/shooting_cache

Author: avik-pal

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

@@ -44,8 +44,8 @@ end
 
     if @load_preference("PrecompileShooting", true)
         push!(algs,
-            Shooting(Tsit5();
-                nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2))))
+            Shooting(Tsit5(); nlsolve = NewtonRaphson(),
+                jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))))
     end
 
     if @load_preference("PrecompileMultipleShooting", true)
@@ -108,11 +108,10 @@ end
     if @load_preference("PrecompileShootingNLLS", VERSION≥v"1.10-")
         append!(algs,
             [
-                Shooting(Tsit5();
-                    nlsolve = LevenbergMarquardt(;
-                        autodiff = AutoForwardDiff(chunksize = 2))),
-                Shooting(Tsit5();
-                    nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2))),
+                Shooting(Tsit5(); nlsolve = LevenbergMarquardt(),
+                    jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))),
+                Shooting(Tsit5(); nlsolve = GaussNewton(),
+                    jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))),
             ])
     end
 

src/algorithms.jl

@@ -3,7 +3,7 @@ abstract type BoundaryValueDiffEqAlgorithm <: SciMLBase.AbstractBVPAlgorithm end
 abstract type AbstractMIRK <: BoundaryValueDiffEqAlgorithm end
 
 """
-    Shooting(ode_alg; nlsolve = NewtonRaphson())
+    Shooting(ode_alg; nlsolve = NewtonRaphson(), jac_alg = BVPJacobianAlgorithm())
 
 Single shooting method, reduces BVP to an initial value problem and solves the IVP.
 
@@ -15,19 +15,50 @@ Single shooting method, reduces BVP to an initial value problem and solves the I
 ## Keyword Arguments
 
   - `nlsolve`: Internal Nonlinear solver. Any solver which conforms to the SciML
-    `NonlinearProblem` interface can be used.
+    `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, you need to specify the chunksize for autodiff. This can be done
-    via `NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = <chunksize>))`.
-    Alternatively, you can use other ADTypes!
+    For type-stability, the chunksizes for ForwardDiff ADTypes in `BVPJacobianAlgorithm`
+    must be provided.
 """
-struct Shooting{O, N} <: BoundaryValueDiffEqAlgorithm
+struct Shooting{O, N, L <: BVPJacobianAlgorithm} <: BoundaryValueDiffEqAlgorithm
     ode_alg::O
     nlsolve::N
+    jac_alg::L
 end
 
-Shooting(ode_alg; nlsolve = NewtonRaphson()) = Shooting(ode_alg, nlsolve)
+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))
+end
+
+function Shooting(ode_alg; nlsolve = NewtonRaphson(), jac_alg = nothing)
+    jac_alg === nothing && (jac_alg = __propagate_nlsolve_ad_to_jac_alg(nlsolve))
+    return Shooting(ode_alg, nlsolve, jac_alg)
+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)
+end
 
 """
     MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(),