Support for SparseDiffTools v2

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Rackauckas <[email protected]>", "Yingbo Ma <maying
 version = "6.49.4"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -69,7 +70,7 @@ SciMLNLSolve = "0.1"
 SimpleNonlinearSolve = "0.1.4"
 SimpleUnPack = "1"
 SnoopPrecompile = "1"
-SparseDiffTools = "1.26.2"
+SparseDiffTools = "2"
 StaticArrayInterface = "1.2"
 StaticArrays = "0.11, 0.12, 1.0"
 TruncatedStacktraces = "1.2"

src/OrdinaryDiffEq.jl

@@ -31,13 +31,16 @@ import DiffEqBase: solve!, step!, initialize!, isadaptive
 import DiffEqBase: ODE_DEFAULT_NORM, ODE_DEFAULT_ISOUTOFDOMAIN, ODE_DEFAULT_PROG_MESSAGE,
                    ODE_DEFAULT_UNSTABLE_CHECK
 
+# TODO: adjust all uses of the below two
 using DiffEqBase: DiffEqArrayOperator, DEFAULT_UPDATE_FUNC
 
 using DiffEqBase: TimeGradientWrapper, UJacobianWrapper, TimeDerivativeWrapper,
                   UDerivativeWrapper
 
 using DiffEqBase: DEIntegrator
 
+import SciMLBase: update_coefficients!
+
 import RecursiveArrayTools: chain, recursivecopy!
 
 using SimpleUnPack, ForwardDiff, RecursiveArrayTools,
@@ -90,6 +93,7 @@ import SparseDiffTools
 import SparseDiffTools: matrix_colors, forwarddiff_color_jacobian!,
                         forwarddiff_color_jacobian, ForwardColorJacCache,
                         default_chunk_size, getsize
+using ADTypes
 
 import Polyester
 using MacroTools, Adapt

src/alg_utils.jl

@@ -316,20 +316,34 @@ function DiffEqBase.prepare_alg(alg::CompositeAlgorithm, u0, p, prob)
     CompositeAlgorithm(algs, alg.choice_function)
 end
 
-function alg_autodiff(alg::OrdinaryDiffEqAlgorithm)
+# Extract AD type parameter from algorithm, returning as Val to ensure type stability for boolean options.
+function _alg_autodiff(alg::OrdinaryDiffEqAlgorithm)
     error("This algorithm does not have an autodifferentiation option defined.")
 end
-alg_autodiff(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS, AD}) where {CS, AD} = AD
-alg_autodiff(alg::DAEAlgorithm{CS, AD}) where {CS, AD} = AD
-alg_autodiff(alg::OrdinaryDiffEqImplicitAlgorithm{CS, AD}) where {CS, AD} = AD
-function alg_autodiff(alg::Union{OrdinaryDiffEqExponentialAlgorithm{CS, AD},
+_alg_autodiff(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(alg::DAEAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(alg::OrdinaryDiffEqImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+function _alg_autodiff(alg::Union{OrdinaryDiffEqExponentialAlgorithm{CS, AD},
                                  OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS, AD}}) where {
                                                                                              CS,
                                                                                              AD
                                                                                              }
-    AD
+    Val{AD}()
 end
 
+function alg_autodiff(alg)
+    autodiff = _alg_autodiff(alg)
+    if autodiff == Val(false)
+        return AutoFiniteDiff()
+    elseif autodiff == Val(true)
+        return AutoForwardDiff()
+    else
+        return _unwrap_val(autodiff)
+    end
+end
+
+# end
+
 # alg_autodiff(alg::CompositeAlgorithm) = alg_autodiff(alg.algs[alg.current_alg])
 get_current_alg_autodiff(alg, cache) = alg_autodiff(alg)
 function get_current_alg_autodiff(alg::CompositeAlgorithm, cache)

src/derivative_utils.jl

@@ -294,14 +294,14 @@ SciMLBase.isinplace(::WOperator{IIP}, i) where {IIP} = IIP
 Base.eltype(W::WOperator) = eltype(W.J)
 
 set_gamma!(W::WOperator, gamma) = (W.gamma = gamma; W)
-function DiffEqBase.update_coefficients!(W::WOperator, u, p, t)
+function update_coefficients!(W::WOperator, u, p, t)
     update_coefficients!(W.J, u, p, t)
     update_coefficients!(W.mass_matrix, u, p, t)
     W.jacvec !== nothing && update_coefficients!(W.jacvec, u, p, t)
     W
 end
 
-function DiffEqBase.update_coefficients!(J::SparseDiffTools.JacVec, u, p, t)
+function DiffEqBase.update_coefficients!(J::FunctionOperator{UJacobianWrapper}, u, p, t)
     copyto!(J.x, u)
     J.f.t = t
     J.f.p = p
@@ -451,6 +451,7 @@ function do_newJW(integrator, alg, nlsolver, repeat_step)::NTuple{2, Bool}
     isfreshJ = isJcurrent(nlsolver, integrator) && !integrator.u_modified
     iszero(nlsolver.fast_convergence_cutoff) && return isfs && !isfreshJ, isfs
     mm = integrator.f.mass_matrix
+    # TODO: adjust 
     is_varying_mm = mm isa DiffEqArrayOperator &&
                     mm.update_func !== SciMLBase.DEFAULT_UPDATE_FUNC
     if isfreshJ
@@ -678,8 +679,7 @@ function calc_W!(W, integrator, nlsolver::Union{Nothing, AbstractNLSolver}, cach
         isnewton(nlsolver) || DiffEqBase.update_coefficients!(W, uprev, p, t) # we will call `update_coefficients!` in NLNewton
         W.transform = W_transform
         set_gamma!(W, dtgamma)
-        if W.J !== nothing && !(W.J isa SparseDiffTools.JacVec) &&
-           !(W.J isa AbstractSciMLOperator)
+        if W.J !== nothing && !(W.J isa AbstractSciMLOperator)
             islin, isode = islinearfunction(integrator)
             islin ? (J = isode ? f.f : f.f1.f) :
             (new_jac && (calc_J!(W.J, integrator, lcache, next_step)))
@@ -854,7 +854,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
 
         _f = islin ? (isode ? f.f : f.f1.f) : f
         jacvec = SparseDiffTools.JacVec(UJacobianWrapper(_f, t, p), copy(u),
-                                        OrdinaryDiffEqTag(), autodiff = alg_autodiff(alg))
+                                        p, t; autodiff = alg_autodiff(alg))
         J = jacvec
         W = WOperator{IIP}(f.mass_matrix, dt, J, u, jacvec)
 
@@ -870,7 +870,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
             deepcopy(f.jac_prototype)
         end
         jacvec = SparseDiffTools.JacVec(UJacobianWrapper(_f, t, p), copy(u),
-                                        OrdinaryDiffEqTag(), autodiff = alg_autodiff(alg))
+                                        p, t; autodiff = alg_autodiff(alg))
         W = WOperator{IIP}(f.mass_matrix, dt, J, u, jacvec)
 
     elseif islin || (!IIP && DiffEqBase.has_jac(f))

src/derivative_wrappers.jl

@@ -79,8 +79,8 @@ function derivative!(df::AbstractArray{<:Number}, f,
                      x::Union{Number, AbstractArray{<:Number}}, fx::AbstractArray{<:Number},
                      integrator, grad_config)
     alg = unwrap_alg(integrator, true)
-    tmp = length(x) # We calculate derivative for all elements in gradient
-    if alg_autodiff(alg)
+    tmp = length(x) # We calculate derivtive for all elements in gradient
+    if alg_autodiff(alg) isa AutoForwardDiff
         T = if standardtag(alg)
             typeof(ForwardDiff.Tag(OrdinaryDiffEqTag(), eltype(df)))
         else
@@ -102,7 +102,7 @@ function derivative!(df::AbstractArray{<:Number}, f,
 
         df .= first.(ForwardDiff.partials.(grad_config))
         integrator.stats.nf += 1
-    else
+    elseif alg_autodiff(alg) isa AutoFiniteDiff
         FiniteDiff.finite_difference_gradient!(df, f, x, grad_config,
                                                dir = diffdir(integrator))
         fdtype = alg_difftype(alg)
@@ -113,6 +113,8 @@ function derivative!(df::AbstractArray{<:Number}, f,
             end
         end
         integrator.stats.nf += tmp
+    else
+        error("$alg_autodiff not yet supported in derivative! function")
     end
     nothing
 end
@@ -122,7 +124,7 @@ function derivative(f, x::Union{Number, AbstractArray{<:Number}},
     local d
     tmp = length(x) # We calculate derivative for all elements in gradient
     alg = unwrap_alg(integrator, true)
-    if alg_autodiff(alg)
+    if alg_autodiff(alg) isa AutoForwardDiff
         integrator.stats.nf += 1
         if integrator.iter == 1
             try
@@ -133,14 +135,16 @@ function derivative(f, x::Union{Number, AbstractArray{<:Number}},
         else
             d = ForwardDiff.derivative(f, x)
         end
-    else
+    elseif alg_autodiff(alg) isa AutoFiniteDiff
         d = FiniteDiff.finite_difference_derivative(f, x, alg_difftype(alg),
                                                     dir = diffdir(integrator))
         if alg_difftype(alg) === Val{:central} || alg_difftype(alg) === Val{:forward}
             tmp *= 2
         end
         integrator.stats.nf += tmp
         d
+    else
+        error("$alg_autodiff not yet supported in derivative function")
     end
 end
 
@@ -186,7 +190,7 @@ end
 function jacobian(f, x, integrator)
     alg = unwrap_alg(integrator, true)
     local tmp
-    if alg_autodiff(alg)
+    if alg_autodiff(alg) isa AutoForwardDiff
         if integrator.iter == 1
             try
                 J, tmp = jacobian_autodiff(f, x, integrator.f, alg)
@@ -196,12 +200,14 @@ function jacobian(f, x, integrator)
         else
             J, tmp = jacobian_autodiff(f, x, integrator.f, alg)
         end
-    else
+    elseif alg_autodiff(alg) isa AutoFiniteDiff
         jac_prototype = integrator.f.jac_prototype
         sparsity, colorvec = sparsity_colorvec(integrator.f, x)
         dir = diffdir(integrator)
         J, tmp = jacobian_finitediff(f, x, alg_difftype(alg), dir, colorvec, sparsity,
                                      jac_prototype)
+    else
+        bleh
     end
     integrator.stats.nf += tmp
     J
@@ -222,7 +228,7 @@ function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number},
                    fx::AbstractArray{<:Number}, integrator::DiffEqBase.DEIntegrator,
                    jac_config)
     alg = unwrap_alg(integrator, true)
-    if alg_autodiff(alg)
+    if alg_autodiff(alg) isa AutoForwardDiff 
         if integrator.iter == 1
             try
                 forwarddiff_color_jacobian!(J, f, x, jac_config)
@@ -233,7 +239,7 @@ function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number},
             forwarddiff_color_jacobian!(J, f, x, jac_config)
         end
         integrator.stats.nf += 1
-    else
+    elseif alg_autodiff(alg) isa AutoFiniteDiff 
         isforward = alg_difftype(alg) === Val{:forward}
         if isforward
             forwardcache = get_tmp_cache(integrator, alg, unwrap_cache(integrator, true))[2]
@@ -245,6 +251,8 @@ function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number},
             tmp = jacobian_finitediff!(J, f, x, jac_config, integrator)
         end
         integrator.stats.nf += tmp
+    else 
+        error("$alg_autodiff not yet supported in jacobian! function")
     end
     nothing
 end
@@ -272,7 +280,8 @@ function build_jac_config(alg, f::F1, uf::F2, du1, uprev, u, tmp, du2,
         end
 
         sparsity, colorvec = sparsity_colorvec(f, u)
-        if alg_autodiff(alg)
+        # TODO: more generc, do we need this?
+        if alg_autodiff(alg) isa AutoForwardDiff
             _chunksize = get_chunksize(alg) === Val(0) ? nothing : get_chunksize(alg) # SparseDiffEq uses different convection...
 
             T = if standardtag(alg)
@@ -282,7 +291,7 @@ function build_jac_config(alg, f::F1, uf::F2, du1, uprev, u, tmp, du2,
             end
             jac_config = ForwardColorJacCache(uf, uprev, _chunksize; colorvec = colorvec,
                                               sparsity = sparsity, tag = T)
-        else
+        elseif alg_autodiff(alg) isa AutoFiniteDiff
             if alg_difftype(alg) !== Val{:complex}
                 jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg_difftype(alg),
                                                       colorvec = colorvec,
@@ -294,6 +303,8 @@ function build_jac_config(alg, f::F1, uf::F2, du1, uprev, u, tmp, du2,
                                                       colorvec = colorvec,
                                                       sparsity = sparsity)
             end
+        else 
+            error("$alg_autodiff not yet supported in build_jac_config function")
         end
     else
         jac_config = nothing
@@ -343,7 +354,7 @@ end
 
 function build_grad_config(alg, f::F1, tf::F2, du1, t) where {F1, F2}
     if !DiffEqBase.has_tgrad(f)
-        if alg_autodiff(alg)
+        if alg_autodiff(alg) isa AutoForwardDiff
             T = if standardtag(alg)
                 typeof(ForwardDiff.Tag(OrdinaryDiffEqTag(), eltype(du1)))
             else
@@ -362,8 +373,10 @@ function build_grad_config(alg, f::F1, tf::F2, du1, t) where {F1, F2}
                                                                  (ForwardDiff.Partials((one(eltype(du1)),)),)) .*
                                                          false)
             end
-        else
+        elseif alg_autodiff(alg) isa AutoFiniteDiff
             grad_config = FiniteDiff.GradientCache(du1, t, alg_difftype(alg))
+        else 
+            error("$alg_autodiff not yet supported in build_grad_config function")
         end
     else
         grad_config = nothing

src/integrators/integrator_interface.jl

@@ -234,20 +234,15 @@ function resize_J_W!(cache, integrator, i)
         nf = nlsolve_f(f, integrator.alg)
         islin = f isa Union{ODEFunction, SplitFunction} && islinear(nf.f)
         if !islin
-            if isa(cache.J, AbstractSciMLOperator)
+            if cache.J isa AbstractSciMLOperator
                 resize!(cache.J, i)
             elseif f.jac_prototype !== nothing
                 J = similar(f.jac_prototype, i, i)
                 J = DiffEqArrayOperator(J; update_func = f.jac)
-            elseif cache.J isa SparseDiffTools.JacVec
-                resize!(cache.J.cache1, i)
-                resize!(cache.J.cache2, i)
-                resize!(cache.J.x, i)
             end
-            if cache.W.jacvec !== nothing
-                resize!(cache.W.jacvec.cache1, i)
-                resize!(cache.W.jacvec.cache2, i)
-                resize!(cache.W.jacvec.x, i)
+            if cache.W.jacvec isa AbstractSciMLOperator
+                # TODO: resize! will need to be implemented upstream to handle previously handled cases 
+                resize!(cache.W.jacvec, i)
             end
             cache.W = WOperator{DiffEqBase.isinplace(integrator.sol.prob)}(f.mass_matrix,
                                                                            integrator.dt,

src/misc_utils.jl

@@ -108,11 +108,13 @@ function dolinsolve(integrator, linsolve; A = nothing, linu = nothing, b = nothi
     # TODO: this ignores the add of the `f` count for add_steps!
     if integrator isa SciMLBase.DEIntegrator && _alg.linsolve !== nothing &&
        !LinearSolve.needs_concrete_A(_alg.linsolve) &&
-       linsolve.A isa WOperator && linsolve.A.J isa SparseDiffTools.JacVec
-        if alg_autodiff(_alg)
+       linsolve.A isa WOperator && linsolve.A.J isa AbstractSciMLOperator 
+        if alg_autodiff(_alg) isa AutoForwardDiff
             integrator.stats.nf += linres.iters
-        else
+        elseif alg_autodiff(_alg) isa AutoFiniteDiff
             integrator.stats.nf += 2 * linres.iters
+        else
+            error("$alg_autodiff not yet supported in dolinsolve function")
         end
     end
 

src/perform_step/linear_perform_step.jl

@@ -819,6 +819,7 @@ function perform_step!(integrator, cache::CayleyEulerConstantCache, repeat_step
         A = f.f
     end
 
+    # TODO: this is not in place, think about this
     L = update_coefficients(A, uprev, p, t)
     V = cay(L * dt)
     u = V * uprev * transpose(V)