Construct explicit second-order AD for optimization BVPs

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/BoundaryValueDiffEqCore/Project.toml

@@ -1,6 +1,6 @@
 name = "BoundaryValueDiffEqCore"
 uuid = "56b672f2-a5fe-4263-ab2d-da677488eb3a"
-version = "2.7.4"
+version = "2.7.5"
 authors = ["Qingyu Qu <erikqqy123@gmail.com>"]
 
 [deps]
@@ -9,6 +9,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
 LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
@@ -35,6 +36,7 @@ Adapt = "4.1.1"
 ArrayInterface = "7.18"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.213, 7"
+DifferentiationInterface = "0.7.13"
 ForwardDiff = "0.10.38, 1"
 Integrals = "4.7.1, 5"
 LineSearch = "0.1.4"

lib/BoundaryValueDiffEqCore/src/BoundaryValueDiffEqCore.jl

@@ -2,10 +2,11 @@ module BoundaryValueDiffEqCore
 
 using Adapt: adapt
 using ADTypes: ADTypes, AbstractADType, AutoSparse, AutoForwardDiff, AutoFiniteDiff,
-    AutoPolyesterForwardDiff
+    AutoPolyesterForwardDiff, AutoSymbolics, AutoZygote
 using ArrayInterface: parameterless_type
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase, solve
+using DifferentiationInterface: SecondOrder
 using ForwardDiff: ForwardDiff, pickchunksize
 using Integrals: Integrals, IntegralProblem
 using LinearAlgebra: LinearAlgebra, mul!

lib/BoundaryValueDiffEqCore/src/internal_problems.jl

@@ -1,6 +1,19 @@
 @inline __default_cost(::Nothing) = (x, p) -> 0.0
 @inline __default_cost(f) = f
 
+@inline __optimization_second_order_ad(ad) = SecondOrder(ad, ad)
+@inline __optimization_second_order_ad(ad::SecondOrder) = ad
+@inline __optimization_second_order_ad(ad::AutoZygote) = SecondOrder(AutoForwardDiff(), ad)
+@inline __optimization_second_order_ad(ad::AutoSymbolics) = ad
+@inline __optimization_second_order_ad(ad::SciMLBase.NoAD) = ad
+
+@inline function __optimization_ad(diffmode, detector_diffmode = diffmode)
+    return AutoSparse(
+        __optimization_second_order_ad(get_dense_ad(diffmode)),
+        sparsity_detector = __default_sparsity_detector(detector_diffmode)
+    )
+end
+
 """
     __build_cost(fun, cache, mesh, M; tune_parameters = false, p = nothing)
 
@@ -119,10 +132,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{true}(
             cost_fun,
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.nonbc_diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.nonbc_diffmode, alg.jac_alg.diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -151,10 +161,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{true}(
             cost_fun,
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -184,10 +191,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{iip}(
             __default_cost(prob.f.cost),
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -227,10 +231,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{true}(
             __default_cost(prob.f.cost),
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.nonbc_diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.nonbc_diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.nonbc_diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -258,10 +259,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{true}(
             __default_cost(prob.f.cost),
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.nonbc_diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.diffmode, alg.jac_alg.nonbc_diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -291,10 +289,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{iip}(
             __default_cost(prob.f.cost),
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.nonbc_diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.nonbc_diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.nonbc_diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)
@@ -322,10 +317,7 @@ function __construct_internal_problem(
     else
         optf = OptimizationFunction{true}(
             __default_cost(prob.f),
-            AutoSparse(
-                get_dense_ad(alg.jac_alg.diffmode),
-                sparsity_detector = __default_sparsity_detector(alg.jac_alg.nonbc_diffmode)
-            ),
+            __optimization_ad(alg.jac_alg.diffmode, alg.jac_alg.nonbc_diffmode),
             cons = loss,
             cons_j = jac,
             cons_jac_prototype = sparse(jac_prototype)

lib/BoundaryValueDiffEqMIRK/test/Core/mirk_basic_tests.jl

@@ -643,7 +643,8 @@ end
         simplependulum!, bc!, [pi / 2, pi / 2], tspan,
         lcons = [-10.0, -10.0], ucons = [10.0, 10.0]
     )
-    @test_nowarn sol = solve(prob, MIRK4(; optimize = IpoptOptimizer()), dt = 0.05)
+    sol = @test_nowarn solve(prob, MIRK4(; optimize = IpoptOptimizer()), dt = 0.05)
+    @test SciMLBase.successful_retcode(sol)
 end
 
 # https://github.com/SciML/BoundaryValueDiffEq.jl/pull/473
@@ -668,7 +669,8 @@ end
     prob = BVProblem(
         simplependulum!, bc!, [pi / 2, pi / 2], tspan
     )
-    @test_nowarn solve(prob, MIRK4(; optimize = IpoptOptimizer()), dt = 0.05)
+    sol = @test_nowarn solve(prob, MIRK4(; optimize = IpoptOptimizer()), dt = 0.05)
+    @test SciMLBase.successful_retcode(sol)
 end
 
 @testset "Test initial guess" begin