Define prepare_alg dispatch for OrdinaryDiffEq AD algorithms

Project.toml

@@ -31,7 +31,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Adapt = "1.1, 2.0, 3.0"
 ArrayInterface = "2.7, 3.0"
 DataStructures = "0.18"
-DiffEqBase = "6.38"
+DiffEqBase = "6.72"
 DocStringExtensions = "0.8"
 ExponentialUtilities = "1.2"
 FastClosures = "0.3"

src/OrdinaryDiffEq.jl

@@ -180,6 +180,7 @@ using DocStringExtensions
         du[3] = u[1]*u[2] - (8/3)*u[3]
       end
       lorenzprob = ODEProblem(lorenz,[1.0;0.0;0.0],(0.0,1.0))
+      solve(lorenzprob,BS3())
       solve(lorenzprob,Tsit5())
       solve(lorenzprob,Vern7())
       solve(lorenzprob,Vern9())
@@ -188,6 +189,16 @@ using DocStringExtensions
       solve(lorenzprob,Rodas4(autodiff=false))
       solve(lorenzprob,KenCarp4(autodiff=false))
       solve(lorenzprob,Rodas5())
+      solve(lorenzprob,QNDF())
+      solve(lorenzprob,QNDF(autodiff=false))
+      solve(lorenzprob,AutoTsit5(Rosenbrock23()))
+      solve(lorenzprob,AutoTsit5(Rosenbrock23(autodiff=false)))
+      solve(lorenzprob,AutoTsit5(TRBDF2(autodiff=false)))
+      solve(lorenzprob,AutoVern7(Rodas4(autodiff=false)))
+      solve(lorenzprob,AutoVern7(TRBDF2(autodiff=false)))
+      solve(lorenzprob,AutoVern9(Rodas5(autodiff=false)))
+      solve(lorenzprob,AutoVern9(KenCarp47(autodiff=false)))
+      solve(lorenzprob,AutoVern7(Rodas5()))
       break
     end
   end

src/alg_utils.jl

@@ -1,4 +1,4 @@
-isautodifferentiable(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = true
+isautodifferentiable(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = true
 
 DiffEqBase.isdiscrete(alg::FunctionMap) = true
 
@@ -63,6 +63,14 @@ end
 issplit(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
 issplit(alg::SplitAlgorithms) = true
 
+function _composite_beta1_default(algs::Tuple{T1,T2}, current, ::Type{QT}, beta2) where {T1, T2, QT}
+  if current == 1
+    return QT(beta1_default(algs[1], beta2))
+  elseif current == 2
+    return QT(beta1_default(algs[2], beta2))
+  end
+end
+
 @generated function _composite_beta1_default(algs::T, current, ::Type{QT}, beta2) where {T <: Tuple, QT}
   expr = Expr(:block)
   for i in 1:length(T.types)
@@ -74,6 +82,15 @@ issplit(alg::SplitAlgorithms) = true
   end
   return expr
 end
+
+function _composite_beta2_default(algs::Tuple{T1,T2}, current, ::Type{QT}) where {T1, T2, QT}
+  if current == 1
+    return QT(beta2_default(algs[1]))
+  elseif current == 2
+    return QT(beta2_default(algs[2]))
+  end
+end
+
 @generated function _composite_beta2_default(algs::T, current, ::Type{QT}) where {T <: Tuple, QT}
   expr = Expr(:block)
   for i in 1:length(T.types)
@@ -131,8 +148,36 @@ get_chunksize(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD}
 get_chunksize(alg::OrdinaryDiffEqImplicitAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
 get_chunksize(alg::DAEAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
 get_chunksize(alg::ExponentialAlgorithm) = Val(alg.chunksize)
+
+get_chunksize_int(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have a chunk size defined.")
+get_chunksize_int(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = CS
+get_chunksize_int(alg::OrdinaryDiffEqImplicitAlgorithm{CS,AD}) where {CS,AD} = CS
+get_chunksize_int(alg::DAEAlgorithm{CS,AD}) where {CS,AD} = CS
+get_chunksize_int(alg::ExponentialAlgorithm) = alg.chunksize
 # get_chunksize(alg::CompositeAlgorithm) = get_chunksize(alg.algs[alg.current_alg])
 
+function DiffEqBase.prepare_alg(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm{0,AD,FDT},
+                        OrdinaryDiffEqImplicitAlgorithm{0,AD,FDT},
+                        DAEAlgorithm{0,AD,FDT}},u0,p,prob) where {AD,FDT}
+    # If chunksize is zero, pick chunksize right at the start of solve and
+    # then do function barrier to infer the full solve
+    x = if prob.f.colorvec === nothing
+      length(u0)
+    else
+      maximum(prob.f.colorvec)
+    end
+
+    if typeof(alg) <: OrdinaryDiffEqImplicitExtrapolationAlgorithm
+      return alg # remake fails, should get fixed
+    else
+      remake(alg,chunk_size=ForwardDiff.pickchunksize(x))
+    end
+end
+
+function DiffEqBase.prepare_alg(alg::CompositeAlgorithm,u0,p,prob)
+    CompositeAlgorithm(Tuple(DiffEqBase.prepare_alg(alg,u0,p,prob) for alg in alg.algs),alg.choice_function)
+end
+
 alg_autodiff(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have an autodifferentiation option defined.")
 alg_autodiff(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = AD
 alg_autodiff(alg::DAEAlgorithm{CS,AD}) where {CS,AD} = AD

src/algorithms.jl

@@ -21,23 +21,24 @@ abstract type OrdinaryDiffEqAdamsVarOrderVarStepAlgorithm <: OrdinaryDiffEqAdapt
 abstract type OrdinaryDiffEqExtrapolationVarOrderVarStepAlgorithm <: OrdinaryDiffEqAdaptiveAlgorithm end
 abstract type OrdinaryDiffEqImplicitExtrapolationAlgorithm{CS,AD,FDT} <: OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT} end
 
+# DAE Specific Algorithms
+abstract type DAEAlgorithm{CS,AD,FDT} <: DiffEqBase.AbstractDAEAlgorithm end
+
 struct FunctionMap{scale_by_time} <: OrdinaryDiffEqAlgorithm end
 FunctionMap(;scale_by_time=false) = FunctionMap{scale_by_time}()
 
-function DiffEqBase.remake(thing::OrdinaryDiffEqAlgorithm, kwargs...)
+function DiffEqBase.remake(thing::OrdinaryDiffEqAlgorithm; kwargs...)
   T = DiffEqBase.remaker_of(thing)
   T(; DiffEqBase.struct_as_namedtuple(thing)...,kwargs...)
 end
 
-function DiffEqBase.remake(thing::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT}, kwargs...) where {CS, AD, FDT}
-  T = DiffEqBase.remaker_of(thing)
-  T(; chunk_size=CS,autodiff=AD,DiffEqBase.struct_as_namedtuple(thing)...,kwargs...)
+function DiffEqBase.remake(thing::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT},
+                        OrdinaryDiffEqImplicitAlgorithm{CS,AD,FDT},
+                        DAEAlgorithm{CS,AD,FDT}}; kwargs...) where {CS, AD, FDT}
+  T = SciMLBase.remaker_of(thing)
+  T(; chunk_size=CS,autodiff=AD,SciMLBase.struct_as_namedtuple(thing)...,kwargs...)
 end
 
-function DiffEqBase.remake(thing::OrdinaryDiffEqImplicitAlgorithm{CS,AD,FDT}, kwargs...) where {CS, AD, FDT}
-  T = DiffEqBase.remaker_of(thing)
-  T(; chunk_size=CS,autodiff=AD,DiffEqBase.struct_as_namedtuple(thing)...,kwargs...)
-end
 ###############################################################################
 
 # RK methods
@@ -1898,6 +1899,15 @@ SBDF(order;chunk_size=0,autodiff=true,diff_type=Val{:forward},
      typeof(κ),typeof(tol)}(
      linsolve,nlsolve,κ,tol,extrapolant,order)
 
+# All keyword form needed for remake
+SBDF(;chunk_size=0,autodiff=true,diff_type=Val{:forward},
+     linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),κ=nothing,tol=nothing,
+     extrapolant=:linear,
+     order) =
+     SBDF{chunk_size,autodiff,typeof(linsolve),typeof(nlsolve),diff_type,
+     typeof(κ),typeof(tol)}(
+     linsolve,nlsolve,κ,tol,extrapolant,order)
+
  """
  Uri M. Ascher, Steven J. Ruuth, Brian T. R. Wetton. Implicit-Explicit Methods for Time-
  Dependent Partial Differential Equations. 1995 Society for Industrial and Applied Mathematics
@@ -2833,10 +2843,6 @@ const MultistepAlgorithms = Union{IRKN3,IRKN4,
 const SplitAlgorithms = Union{CNAB2,CNLF2,IRKC,SBDF,
                               KenCarp3,KenCarp4,KenCarp47,KenCarp5,KenCarp58,CFNLIRK3}
 
-
-# DAE Specific Algorithms
-abstract type DAEAlgorithm{CS,AD,FDT} <: DiffEqBase.AbstractDAEAlgorithm end
-
 #=
 struct DBDF{CS,AD,F,F2,FDT} <: DAEAlgorithm{CS,AD,FDT}
   linsolve::F

src/dense/interpolants.jl

@@ -289,9 +289,14 @@ end
 @muladd function _ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache::Union{Tsit5ConstantCache,Tsit5Cache},idxs::Nothing,T::Type{Val{0}})
   @tsit5pre0
   @inbounds @.. out = y₀ + dt*(k[1]*b1Θ + k[2]*b2Θ + k[3]*b3Θ + k[4]*b4Θ + k[5]*b5Θ + k[6]*b6Θ + k[7]*b7Θ)
-  #@inbounds for i in eachindex(out)
-  #  out[i] = y₀[i] + dt*(k[1][i]*b1Θ + k[2][i]*b2Θ + k[3][i]*b3Θ + k[4][i]*b4Θ + k[5][i]*b5Θ + k[6][i]*b6Θ + k[7][i]*b7Θ)
-  #end
+  out
+end
+
+@muladd function _ode_interpolant!(out::Array,Θ,dt,y₀,y₁,k,cache::Union{Tsit5ConstantCache,Tsit5Cache},idxs::Nothing,T::Type{Val{0}})
+  @tsit5pre0
+  @inbounds @simd ivdep for i in eachindex(out)
+    out[i] = y₀[i] + dt*(k[1][i]*b1Θ + k[2][i]*b2Θ + k[3][i]*b3Θ + k[4][i]*b4Θ + k[5][i]*b5Θ + k[6][i]*b6Θ + k[7][i]*b7Θ)
+  end
   out
 end
 
@@ -304,6 +309,14 @@ end
   out
 end
 
+@muladd function _ode_interpolant!(out::Array,Θ,dt,y₀,y₁,k,cache::Union{Tsit5ConstantCache,Tsit5Cache},idxs,T::Type{Val{0}})
+  @tsit5pre0
+  @inbounds for (j,i) in enumerate(idxs)
+    out[j] = y₀[i] + dt*(k[1][i]*b1Θ + k[2][i]*b2Θ + k[3][i]*b3Θ + k[4][i]*b4Θ + k[5][i]*b5Θ + k[6][i]*b6Θ + k[7][i]*b7Θ)
+  end
+  out
+end
+
 @def tsit5pre1 begin
   @tsit5unpack
   b1Θdiff = @evalpoly(Θ, r11, 2*r12, 3*r13, 4*r14)

src/derivative_wrappers.jl

@@ -43,7 +43,7 @@ function jacobian_autodiff(f, x::AbstractArray, odefun, alg)
   jac_prototype = odefun.jac_prototype
   sparsity,colorvec = sparsity_colorvec(odefun,x)
   maxcolor = maximum(colorvec)
-  chunk_size = get_chunksize(alg)===Val(0) ? nothing : get_chunksize(alg) # SparseDiffEq uses different convection...
+  chunk_size = get_chunksize(alg)===Val(0) ? nothing : get_chunksize_int(alg) # SparseDiffEq uses different convection...
   num_of_chunks = chunk_size===nothing ? Int(ceil(maxcolor / getsize(ForwardDiff.pickchunksize(maxcolor)))) :
                                         Int(ceil(maxcolor / chunk_size))
   (forwarddiff_color_jacobian(f,x,colorvec = colorvec, sparsity = sparsity,