Use DifferentiationInterface for AD in Implicit Solvers

Project.toml

@@ -68,14 +68,13 @@ SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 
 [compat]
-ADTypes = "0.2, 1"
+ADTypes = "1.13"
 Adapt = "3.0, 4"
 ArrayInterface = "7"
 DataStructures = "0.18"
@@ -138,7 +137,6 @@ SciMLOperators = "0.3"
 SciMLStructures = "1"
 SimpleNonlinearSolve = "1, 2"
 SimpleUnPack = "1"
-SparseDiffTools = "2"
 Static = "0.8, 1"
 StaticArrayInterface = "1.2"
 StaticArrays = "1.0"
@@ -153,6 +151,7 @@ ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
@@ -167,11 +166,13 @@ RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Calculus", "ComponentArrays", "Symbolics", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DiffEqDevTools", "ODEProblemLibrary", "ElasticArrays", "InteractiveUtils", "ParameterizedFunctions", "PoissonRandom", "Printf", "Random", "ReverseDiff", "SafeTestsets", "SparseArrays", "Statistics", "StructArrays", "Test", "Unitful", "ModelingToolkit", "Pkg", "NLsolve", "RecursiveFactorization"]
+test = ["Calculus", "ComponentArrays", "Symbolics", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DiffEqDevTools", "ODEProblemLibrary", "ElasticArrays", "InteractiveUtils", "ParameterizedFunctions", "PoissonRandom", "Printf", "Random", "ReverseDiff", "SafeTestsets", "SparseArrays", "Statistics", "StructArrays", "Test", "Unitful", "ModelingToolkit", "Pkg", "NLsolve", "RecursiveFactorization", "Enzyme", "SparseConnectivityTracer", "SparseMatrixColorings"]

lib/OrdinaryDiffEqBDF/Project.toml

@@ -50,12 +50,14 @@ julia = "1.10"
 
 [extras]
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DiffEqDevTools", "ForwardDiff", "Random", "SafeTestsets", "Test", "ODEProblemLibrary", "StaticArrays"]
+test = ["DiffEqDevTools", "ForwardDiff", "Random", "SafeTestsets", "Test", "ODEProblemLibrary", "StaticArrays", "Enzyme", "LinearSolve"]

lib/OrdinaryDiffEqBDF/test/bdf_convergence_tests.jl

@@ -1,5 +1,5 @@
 # This definitely needs cleaning
-using OrdinaryDiffEqBDF, ODEProblemLibrary, DiffEqDevTools
+using OrdinaryDiffEqBDF, ODEProblemLibrary, DiffEqDevTools, ADTypes, Enzyme, LinearSolve
 using OrdinaryDiffEqNonlinearSolve: NLFunctional, NLAnderson, NonlinearSolveAlg
 using Test, Random
 Random.seed!(100)
@@ -39,6 +39,27 @@ dts3 = 1 .// 2 .^ (12:-1:7)
     @test sim.𝒪est[:l2]≈1 atol=testTol
     @test sim.𝒪est[:l∞]≈1 atol=testTol
 
+    sim = test_convergence(dts, prob, QNDF1(autodiff = AutoFiniteDiff()))
+    @test sim.𝒪est[:final]≈1 atol=testTol
+    @test sim.𝒪est[:l2]≈1 atol=testTol
+    @test sim.𝒪est[:l∞]≈1 atol=testTol
+
+    sim = test_convergence(dts,
+        prob,
+        QNDF1(autodiff = AutoEnzyme(mode = set_runtime_activity(Enzyme.Forward),
+            function_annotation = Enzyme.Const)))
+    @test sim.𝒪est[:final]≈1 atol=testTol
+    @test sim.𝒪est[:l2]≈1 atol=testTol
+    @test sim.𝒪est[:l∞]≈1 atol=testTol
+
+    sim = test_convergence(dts,
+        prob,
+        QNDF1(autodiff = AutoEnzyme(mode = set_runtime_activity(Enzyme.Forward),
+            function_annotation = Enzyme.Const), linsolve = LinearSolve.KrylovJL()))
+    @test sim.𝒪est[:final]≈1 atol=testTol
+    @test sim.𝒪est[:l2]≈1 atol=testTol
+    @test sim.𝒪est[:l∞]≈1 atol=testTol
+
     sim = test_convergence(dts3, prob, QNDF2())
     @test sim.𝒪est[:final]≈2 atol=testTol
     @test sim.𝒪est[:l2]≈2 atol=testTol

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -72,7 +72,7 @@ import DiffEqBase: calculate_residuals,
 
 import Polyester
 using MacroTools, Adapt
-import ADTypes: AutoFiniteDiff, AutoForwardDiff, AbstractADType
+import ADTypes: AutoFiniteDiff, AutoForwardDiff, AbstractADType, AutoSparse
 import Accessors: @reset
 
 using SciMLStructures: canonicalize, Tunable, isscimlstructure

lib/OrdinaryDiffEqCore/src/alg_utils.jl

@@ -176,6 +176,7 @@ _get_fwd_chunksize(::Type{<:AutoForwardDiff{CS}}) where {CS} = Val(CS)
 _get_fwd_chunksize_int(::Type{<:AutoForwardDiff{CS}}) where {CS} = CS
 _get_fwd_chunksize(AD) = Val(0)
 _get_fwd_chunksize_int(AD) = 0
+_get_fwd_chunksize_int(::AutoForwardDiff{CS}) where {CS} = CS
 _get_fwd_tag(::AutoForwardDiff{CS, T}) where {CS, T} = T
 
 _get_fdtype(::AutoFiniteDiff{T1}) where {T1} = T1

lib/OrdinaryDiffEqCore/src/algorithms.jl

@@ -59,9 +59,16 @@ function DiffEqBase.remake(
             },
             DAEAlgorithm{CS, AD, FDT, ST, CJ}};
         kwargs...) where {CS, AD, FDT, ST, CJ}
+
+    if haskey(kwargs, :autodiff) && kwargs[:autodiff] isa AutoForwardDiff
+        chunk_size = _get_fwd_chunksize(kwargs[:autodiff])
+    else
+        chunk_size = Val{CS}()
+    end
+
     T = SciMLBase.remaker_of(thing)
     T(; SciMLBase.struct_as_namedtuple(thing)...,
-        chunk_size = Val{CS}(), autodiff = thing.autodiff, standardtag = Val{ST}(),
+        chunk_size = chunk_size, autodiff = thing.autodiff, standardtag = Val{ST}(),
         concrete_jac = CJ === nothing ? CJ : Val{CJ}(),
         kwargs...)
 end

lib/OrdinaryDiffEqCore/src/misc_utils.jl

@@ -158,10 +158,11 @@ end
 function _process_AD_choice(
         ad_alg::AutoForwardDiff{CS}, ::Val{CS2}, ::Val{FD}) where {CS, CS2, FD}
     # Non-default `chunk_size`
-    if CS2 != 0
+    if (CS2 != 0) && (isnothing(CS) || (CS2 !== CS))
         @warn "The `chunk_size` keyword is deprecated. Please use an `ADType` specifier. For now defaulting to using `AutoForwardDiff` with `chunksize=$(CS2)`."
         return _bool_to_ADType(Val{true}(), Val{CS2}(), Val{FD}()), Val{CS2}(), Val{FD}()
     end
+
     _CS = CS === nothing ? 0 : CS
     return ad_alg, Val{_CS}(), Val{FD}()
 end
@@ -186,3 +187,12 @@ function _process_AD_choice(
     end
     return ad_alg, Val{CS}(), ad_alg.fdtype
 end
+
+function _process_AD_choice(ad_alg::AutoSparse, cs2::Val{CS2}, fd::Val{FD}) where {CS2, FD}
+    _, cs, fd = _process_AD_choice(ad_alg.dense_ad, cs2, fd)
+    ad_alg, cs, fd
+end
+
+function _process_AD_choice(ad_alg, cs2, fd)
+    ad_alg, cs2, fd
+end

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -6,7 +6,10 @@ version = "1.4.0"
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -15,16 +18,18 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-ADTypes = "1.11"
+ADTypes = "1.14"
 ArrayInterface = "7"
 DiffEqBase = "6"
 DiffEqDevTools = "2.44.4"
+DifferentiationInterface = "0.6.48"
 FastBroadcast = "0.3"
 FiniteDiff = "2"
 ForwardDiff = "0.10"
@@ -36,7 +41,6 @@ Random = "<0.0.1, 1"
 SafeTestsets = "0.1.0"
 SciMLBase = "2"
 SparseArrays = "1"
-SparseDiffTools = "2"
 StaticArrayInterface = "1"
 StaticArrays = "1"
 Test = "<0.0.1, 1"

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -1,10 +1,7 @@
 module OrdinaryDiffEqDifferentiation
 
-import ADTypes: AutoFiniteDiff, AutoForwardDiff, AbstractADType
-
-import SparseDiffTools: SparseDiffTools, matrix_colors, forwarddiff_color_jacobian!,
-                        forwarddiff_color_jacobian, ForwardColorJacCache,
-                        default_chunk_size, getsize, JacVec
+import ADTypes
+import ADTypes: AutoFiniteDiff, AutoForwardDiff, AbstractADType, AutoSparse
 
 import ForwardDiff, FiniteDiff
 import ForwardDiff.Dual
@@ -16,7 +13,7 @@ using DiffEqBase
 import LinearAlgebra
 import LinearAlgebra: Diagonal, I, UniformScaling, diagind, mul!, lmul!, axpby!, opnorm, lu
 import LinearAlgebra: LowerTriangular, UpperTriangular
-import SparseArrays: SparseMatrixCSC, AbstractSparseMatrix, nonzeros
+import SparseArrays: SparseMatrixCSC, AbstractSparseMatrix, nonzeros, sparse
 import ArrayInterface
 
 import StaticArrayInterface
@@ -27,7 +24,9 @@ import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray, MMatrix, S
 using DiffEqBase: TimeGradientWrapper,
                   UJacobianWrapper, TimeDerivativeWrapper,
                   UDerivativeWrapper
-using SciMLBase: AbstractSciMLOperator, constructorof
+using SciMLBase: AbstractSciMLOperator, constructorof, @set
+using SciMLOperators
+import SparseMatrixColorings
 import OrdinaryDiffEqCore
 using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplicitAlgorithm,
                           DAEAlgorithm,
@@ -44,11 +43,16 @@ using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplici
                           FastConvergence, Convergence, SlowConvergence,
                           VerySlowConvergence, Divergence, NLStatus, MethodType, constvalue
 
-import OrdinaryDiffEqCore: get_chunksize, resize_J_W!, resize_nlsolver!, alg_autodiff,
-                           _get_fwd_tag
+import OrdinaryDiffEqCore: get_chunksize, resize_J_W!, resize_nlsolver!, alg_autodiff, _get_fwd_tag
+
+using ConstructionBase
+
+import DifferentiationInterface as DI
 
 using FastBroadcast: @..
 
+using ConcreteStructs: @concrete
+
 @static if isdefined(DiffEqBase, :OrdinaryDiffEqTag)
     import DiffEqBase: OrdinaryDiffEqTag
 else
@@ -59,5 +63,6 @@ include("alg_utils.jl")
 include("linsolve_utils.jl")
 include("derivative_utils.jl")
 include("derivative_wrappers.jl")
+include("operators.jl")
 
 end