Reuse ODE Solver Cache

Author: avik-pal

Project.toml

@@ -1,6 +1,6 @@
 name = "BoundaryValueDiffEq"
 uuid = "764a87c0-6b3e-53db-9096-fe964310641d"
-version = "5.4.0"
+version = "6.0.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

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)

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,33 @@ 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
+
+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
 
-Shooting(ode_alg; nlsolve = NewtonRaphson()) = Shooting(ode_alg, nlsolve)
+function Shooting(ode_alg; nlsolve = NewtonRaphson(), jac_alg = BVPJacobianAlgorithm())
+    return Shooting(ode_alg, nlsolve, jac_alg)
+end
 
 """
     MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(),

src/solve/single_shooting.jl

@@ -1,66 +1,141 @@
-function __solve(prob::BVProblem, alg::Shooting; odesolve_kwargs = (;),
+function __solve(prob::BVProblem, alg_::Shooting; odesolve_kwargs = (;),
         nlsolve_kwargs = (;), verbose = true, kwargs...)
     ig, T, N, _, u0 = __extract_problem_details(prob; dt = 0.1)
     _unwrap_val(ig) && verbose &&
         @warn "Initial guess provided, but will be ignored for Shooting!"
 
+    alg = concretize_jacobian_algorithm(alg_, prob)
+
     bcresid_prototype, resid_size = __get_bcresid_prototype(prob, u0)
     iip, bc, u0, u0_size = isinplace(prob), prob.f.bc, deepcopy(u0), size(u0)
     resid_prototype = __vec(bcresid_prototype)
 
+    # Construct the residual function
     ode_kwargs = (; kwargs..., verbose, odesolve_kwargs...)
+    internal_prob = ODEProblem{iip}(prob.f, u0, prob.tspan, prob.p)
+    ode_cache_loss_fn = SciMLBase.__init(internal_prob, alg.ode_alg; ode_kwargs...)
+
     loss_fn = if iip
-        (du, u, p) -> __single_shooting_loss!(du, u, p, prob.f, bc, u0_size, prob.tspan,
-            prob.problem_type, resid_size, alg, ode_kwargs)
+        (du, u, p) -> __single_shooting_loss!(du, u, p, ode_cache_loss_fn, bc, u0_size,
+            prob.problem_type, resid_size)
+    else
+        (u, p) -> __single_shooting_loss(u, p, ode_cache_loss_fn, bc, u0_size,
+            prob.problem_type)
+    end
+
+    # Construct the jacobian function
+    # NOTE: We pass in a separate Jacobian Function because that allows us to cache the
+    #       the internal ode solve cache. This cache needs to be distinct from the regular
+    #       residual function cache
+    loss_fnₚ = ifelse(iip, (du, u) -> loss_fn(du, u, prob.p), (u) -> loss_fn(u, prob.p))
+
+    # TODO: We probably won't be able to support Symbolics through ODE Solver but we should
+    #       be able to allow prespecified coloring.
+    if alg.jac_alg.diffmode isa AbstractSparseADType
+        error("Single Shooting doesn't support sparse AD yet!")
+    end
+    sd = NoSparsityDetection()
+    y_ = similar(resid_prototype)
+
+    jac_cache = if iip
+        sparse_jacobian_cache(alg.jac_alg.diffmode, sd, loss_fnₚ, y_, vec(u0))
+    else
+        sparse_jacobian_cache(alg.jac_alg.diffmode, sd, loss_fnₚ, vec(u0); fx = y_)
+    end
+
+    ode_cache_jac_fn = __single_shooting_jacobian_ode_cache(internal_prob, jac_cache,
+        alg.jac_alg.diffmode, vec(u0), alg.ode_alg; ode_kwargs...)
+
+    jac_prototype = init_jacobian(jac_cache)
+
+    loss_fn2ₚ = if iip
+        (du, u) -> __single_shooting_loss!(du, u, prob.p, ode_cache_jac_fn, bc, u0_size,
+            prob.problem_type, resid_size)
+    else
+        (u) -> __single_shooting_loss(u, prob.p, ode_cache_jac_fn, bc, u0_size,
+            prob.problem_type)
+    end
+
+    jac_fn = if iip
+        (J, u, p) -> __single_shooting_jacobian!(J, u, jac_cache, alg.jac_alg.diffmode,
+            loss_fn2ₚ, y_)
     else
-        (u, p) -> __single_shooting_loss(u, p, prob.f, bc, u0_size, prob.tspan,
-            prob.problem_type, alg, ode_kwargs)
+        (u, p) -> __single_shooting_jacobian(jac_prototype, u, jac_cache,
+            alg.jac_alg.diffmode, loss_fn2ₚ)
     end
 
-    nlf = NonlinearFunction{iip}(loss_fn; prob.f.jac_prototype, resid_prototype)
+    nlf = NonlinearFunction{iip}(loss_fn; prob.f.jac_prototype, resid_prototype,
+        jac = jac_fn)
     nlprob = if length(resid_prototype) == length(u0)
         NonlinearProblem(nlf, vec(u0), prob.p)
     else
         NonlinearLeastSquaresProblem(nlf, vec(u0), prob.p)
     end
     opt = __solve(nlprob, alg.nlsolve; nlsolve_kwargs..., verbose, kwargs...)
 
-    newprob = ODEProblem{iip}(prob.f, reshape(opt.u, u0_size), prob.tspan, prob.p)
-    sol = __solve(newprob, alg.ode_alg; odesolve_kwargs..., verbose, kwargs...)
+    SciMLBase.reinit!(ode_cache_loss_fn, reshape(opt.u, u0_size))
+    sol = solve!(ode_cache_loss_fn)
 
     !SciMLBase.successful_retcode(opt) &&
         return SciMLBase.solution_new_retcode(sol, ReturnCode.Failure)
     return sol
 end
 
-function __single_shooting_loss!(resid_, u0_, p, f::F, bc::BC, u0_size, tspan,
-        pt::TwoPointBVProblem, (resida_size, residb_size), alg::Shooting,
-        kwargs) where {F, BC}
+function __single_shooting_loss!(resid_, u0_, p, cache, bc::BC, u0_size,
+        pt::TwoPointBVProblem, (resida_size, residb_size)) where {BC}
     resida = @view resid_[1:prod(resida_size)]
     residb = @view resid_[(prod(resida_size) + 1):end]
     resid = (reshape(resida, resida_size), reshape(residb, residb_size))
 
-    odeprob = ODEProblem{true}(f, reshape(u0_, u0_size), tspan, p)
-    odesol = __solve(odeprob, alg.ode_alg; kwargs...)
+    SciMLBase.reinit!(cache, reshape(u0_, u0_size))
+    odesol = solve!(cache)
+
     eval_bc_residual!(resid, pt, bc, odesol, p)
 
     return nothing
 end
 
-function __single_shooting_loss!(resid_, u0_, p, f::F, bc::BC, u0_size, tspan,
-        pt::StandardBVProblem, resid_size, alg::Shooting, kwargs) where {F, BC}
+function __single_shooting_loss!(resid_, u0_, p, cache, bc::BC, u0_size,
+        pt::StandardBVProblem, resid_size) where {BC}
     resid = reshape(resid_, resid_size)
 
-    odeprob = ODEProblem{true}(f, reshape(u0_, u0_size), tspan, p)
-    odesol = __solve(odeprob, alg.ode_alg; kwargs...)
+    SciMLBase.reinit!(cache, reshape(u0_, u0_size))
+    odesol = solve!(cache)
+
     eval_bc_residual!(resid, pt, bc, odesol, p)
 
     return nothing
 end
 
-function __single_shooting_loss(u0_, p, f::F, bc::BC, u0_size, tspan, pt, alg::Shooting,
-        kwargs) where {F, BC}
-    odeprob = ODEProblem{false}(f, reshape(u0_, u0_size), tspan, p)
-    odesol = __solve(odeprob, alg.ode_alg; kwargs...)
+function __single_shooting_loss(u, p, cache, bc::BC, u0_size, pt) where {BC}
+    SciMLBase.reinit!(cache, reshape(u, u0_size))
+    odesol = solve!(cache)
     return __safe_vec(eval_bc_residual(pt, bc, odesol, p))
 end
+
+function __single_shooting_jacobian!(J, u, jac_cache, diffmode, loss_fn::L, fu) where {L}
+    sparse_jacobian!(J, diffmode, jac_cache, loss_fn, fu, vec(u))
+    return J
+end
+
+function __single_shooting_jacobian(J, u, jac_cache, diffmode, loss_fn::L) where {L}
+    sparse_jacobian!(J, diffmode, jac_cache, loss_fn, vec(u))
+    return J
+end
+
+function __single_shooting_jacobian_ode_cache(prob, jac_cache, alg, u0, ode_alg; kwargs...)
+    prob_ = remake(prob; u0)
+    return SciMLBase.__init(prob_, ode_alg; kwargs...)
+end
+
+function __single_shooting_jacobian_ode_cache(prob, jac_cache,
+        ::Union{AutoForwardDiff, AutoSparseForwardDiff}, u0, ode_alg; kwargs...)
+    cache = jac_cache.cache
+    if cache isa ForwardDiff.JacobianConfig
+        xduals = cache.duals isa Tuple ? cache.duals[2] : cache.duals
+        prob_ = remake(prob; u0 = xduals)
+        return SciMLBase.__init(prob_, ode_alg; kwargs...)
+    else
+        error("Single Shooting doesn't support sparse AD yet!")
+    end
+end