Merge pull request #1733 from SciML/exponential

Fix interface breaks with Exponential integrators

Author: ChrisRackauckas

src/alg_utils.jl

@@ -5,8 +5,8 @@ SciMLBase.allows_arbitrary_number_types(alg::Union{OrdinaryDiffEqAlgorithm,DAEAl
 SciMLBase.allowscomplex(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm,FunctionMap}) = true
 SciMLBase.isdiscrete(alg::FunctionMap) = true
 SciMLBase.forwarddiffs_model(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm,
-                                        DAEAlgorithm,OrdinaryDiffEqImplicitAlgorithm,
-                                        ExponentialAlgorithm}) = alg_autodiff(alg)
+    DAEAlgorithm,OrdinaryDiffEqImplicitAlgorithm,
+    ExponentialAlgorithm}) = alg_autodiff(alg)
 SciMLBase.forwarddiffs_model_time(alg::RosenbrockAlgorithm) = true
 
 # isadaptive is defined below.
@@ -162,21 +162,30 @@ get_chunksize(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not hav
 get_chunksize(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
 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)
+function get_chunksize(alg::Union{OrdinaryDiffEqExponentialAlgorithm{CS,AD},
+                             OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD}}) where {CS,AD}
+                             Val(CS)
+end
 
 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
+function get_chunksize_int(alg::Union{OrdinaryDiffEqExponentialAlgorithm{CS,AD},
+                             OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD}}) where {CS,AD}
+                             CS
+end
 # 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::AbstractArray{T}, p, prob) where {AD,FDT,T}
+        DAEAlgorithm{0,AD,FDT},OrdinaryDiffEqExponentialAlgorithm{0,AD,FDT}}, u0::AbstractArray{T},
+        p, prob) where {AD,FDT,T}
     alg isa OrdinaryDiffEqImplicitExtrapolationAlgorithm && return alg # remake fails, should get fixed
 
-    if alg.linsolve === nothing
+    if alg isa OrdinaryDiffEqExponentialAlgorithm
+        linsolve = nothing
+    elseif alg.linsolve === nothing
         if (prob.f isa ODEFunction && prob.f.f isa SciMLBase.AbstractDiffEqOperator)
             linsolve = LinearSolve.defaultalg(prob.f.f, u0)
         elseif (prob.f isa SplitFunction && prob.f.f1.f isa SciMLBase.AbstractDiffEqOperator)
@@ -187,8 +196,8 @@ function DiffEqBase.prepare_alg(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorit
                 linsolve = KrylovJL()
             end
         elseif (prob isa ODEProblem || prob isa DDEProblem) && (prob.f.mass_matrix === nothing ||
-                                       (prob.f.mass_matrix !== nothing &&
-                                        !(typeof(prob.f.jac_prototype) <: SciMLBase.AbstractDiffEqOperator)))
+                                                                (prob.f.mass_matrix !== nothing &&
+                                                                 !(typeof(prob.f.jac_prototype) <: SciMLBase.AbstractDiffEqOperator)))
             linsolve = LinearSolve.defaultalg(prob.f.jac_prototype, u0)
         else
             # If mm is a sparse matrix and A is a DiffEqArrayOperator, then let linear
@@ -202,8 +211,12 @@ function DiffEqBase.prepare_alg(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorit
     # If norecompile mode or very large bitsize, like a dual number u0 already, then
     # don't use a large chunksize as it will either error or not be beneficial
     if (isbitstype(T) && sizeof(T) > 24) || (prob.f isa ODEFunction && prob.f.f isa
-                                           FunctionWrappersWrappers.FunctionWrappersWrapper)
-        return remake(alg, chunk_size=Val{1}(), linsolve=linsolve)
+                                                                       FunctionWrappersWrappers.FunctionWrappersWrapper)
+        if alg isa OrdinaryDiffEqExponentialAlgorithm
+            return remake(alg, chunk_size=Val{1}())
+        else
+            return remake(alg, chunk_size=Val{1}(), linsolve=linsolve)
+        end
     end
 
     L = ArrayInterface.known_length(typeof(u0))
@@ -218,13 +231,29 @@ function DiffEqBase.prepare_alg(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorit
         end
 
         cs = ForwardDiff.pickchunksize(x)
-        remake(alg, chunk_size=Val{cs}(), linsolve=linsolve)
+
+        if alg isa OrdinaryDiffEqExponentialAlgorithm
+            return remake(alg, chunk_size=Val{cs}())
+        else
+            return remake(alg, chunk_size=Val{cs}(), linsolve=linsolve)
+        end
     else # statically sized
         cs = pick_static_chunksize(Val{L}())
-        remake(alg, chunk_size=cs, linsolve=linsolve)
+        if alg isa OrdinaryDiffEqExponentialAlgorithm
+            return remake(alg, chunk_size=cs)
+        else
+            return remake(alg, chunk_size=cs, linsolve=linsolve)
+        end
     end
 end
 
+# Linear Exponential doesn't have any of the AD stuff
+function DiffEqBase.prepare_alg(alg::Union{ETD2,SplitEuler,LinearExponential,
+                            OrdinaryDiffEqLinearExponentialAlgorithm}, u0::AbstractArray,
+                            p, prob)
+        alg
+end
+
 @generated function pick_static_chunksize(::Val{chunksize}) where {chunksize}
     x = ForwardDiff.pickchunksize(chunksize)
     :(Val{$x}())
@@ -239,24 +268,31 @@ alg_autodiff(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have
 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
-alg_autodiff(alg::ExponentialAlgorithm) = alg.autodiff
+function alg_autodiff(alg::Union{OrdinaryDiffEqExponentialAlgorithm{CS,AD},
+    OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD}}) where {CS,AD}
+    AD
+end
+
 # alg_autodiff(alg::CompositeAlgorithm) = alg_autodiff(alg.algs[alg.current_alg])
 get_current_alg_autodiff(alg, cache) = alg_autodiff(alg)
 get_current_alg_autodiff(alg::CompositeAlgorithm, cache) = alg_autodiff(alg.algs[cache.current])
 
 alg_difftype(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT,ST,CJ},
     OrdinaryDiffEqImplicitAlgorithm{CS,AD,FDT,ST,CJ},
-    OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ},
+    OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ},
+    OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD,FDT,ST,CJ},
     DAEAlgorithm{CS,AD,FDT,ST,CJ}}) where {CS,AD,FDT,ST,CJ} = FDT
 
 standardtag(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT,ST,CJ},
     OrdinaryDiffEqImplicitAlgorithm{CS,AD,FDT,ST,CJ},
-    OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ},
+    OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ},
+    OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD,FDT,ST,CJ},
     DAEAlgorithm{CS,AD,FDT,ST,CJ}}) where {CS,AD,FDT,ST,CJ} = ST
 
 concrete_jac(alg::Union{OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD,FDT,ST,CJ},
     OrdinaryDiffEqImplicitAlgorithm{CS,AD,FDT,ST,CJ},
-    OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ},
+    OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ},
+    OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD,FDT,ST,CJ},
     DAEAlgorithm{CS,AD,FDT,ST,CJ}}) where {CS,AD,FDT,ST,CJ} = CJ
 
 alg_extrapolates(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false

src/algorithms.jl

@@ -12,9 +12,9 @@ abstract type OrdinaryDiffEqRosenbrockAlgorithm{CS,AD,FDT,ST,CJ} <:  OrdinaryDif
 const NewtonAlgorithm = Union{OrdinaryDiffEqNewtonAlgorithm,OrdinaryDiffEqNewtonAdaptiveAlgorithm}
 const RosenbrockAlgorithm = Union{OrdinaryDiffEqRosenbrockAlgorithm,OrdinaryDiffEqRosenbrockAdaptiveAlgorithm}
 
-abstract type OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ} <: OrdinaryDiffEqAlgorithm end
-abstract type OrdinaryDiffEqAdaptiveExponentialAlgorithm{FDT,ST,CJ} <: OrdinaryDiffEqAdaptiveAlgorithm end
-abstract type OrdinaryDiffEqLinearExponentialAlgorithm <: OrdinaryDiffEqExponentialAlgorithm{Val{:forward},Val{true},nothing} end
+abstract type OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ} <: OrdinaryDiffEqAlgorithm end
+abstract type OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD,FDT,ST,CJ} <: OrdinaryDiffEqAdaptiveAlgorithm end
+abstract type OrdinaryDiffEqLinearExponentialAlgorithm <: OrdinaryDiffEqExponentialAlgorithm{0,false,Val{:forward},Val{true},nothing} end
 const ExponentialAlgorithm = Union{OrdinaryDiffEqExponentialAlgorithm,OrdinaryDiffEqAdaptiveExponentialAlgorithm}
 
 abstract type OrdinaryDiffEqAdamsVarOrderVarStepAlgorithm <: OrdinaryDiffEqAdaptiveAlgorithm end
@@ -3130,7 +3130,7 @@ end
 
 struct MagnusAdapt4 <: OrdinaryDiffEqAdaptiveAlgorithm end
 
-struct LinearExponential <: OrdinaryDiffEqExponentialAlgorithm{Val{:forward},Val{true},nothing}
+struct LinearExponential <: OrdinaryDiffEqExponentialAlgorithm{1,false,Val{:forward},Val{true},nothing}
   krylov::Symbol
   m::Int
   iop::Int
@@ -3856,50 +3856,47 @@ RosenbrockW6S4OS(;chunk_size=Val{0}(),autodiff=true, standardtag = Val{true}(),
 
 for Alg in [:LawsonEuler, :NorsettEuler, :ETDRK2, :ETDRK3, :ETDRK4, :HochOst4]
 
-  """
-  Hochbruck, Marlis, and Alexander Ostermann. “Exponential Integrators.” Acta
-    Numerica 19 (2010): 209–86. doi:10.1017/S0962492910000048.
-  """
-  @eval struct $Alg{FDT,ST,CJ} <: OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ}
-    krylov::Bool
-    m::Int
-    iop::Int
-    autodiff::Bool
-    chunksize::Int
-  end
-  @eval $Alg(;krylov=false, m=30, iop=0, autodiff=true, standardtag = Val{true}(), concrete_jac = nothing, chunksize=0,
-            diff_type = Val{:forward}) = $Alg{diff_type,_unwrap_val(standardtag),_unwrap_val(concrete_jac)}(krylov, m, iop, _unwrap_val(autodiff),
-            chunksize)
+    """
+    Hochbruck, Marlis, and Alexander Ostermann. “Exponential Integrators.” Acta
+      Numerica 19 (2010): 209–86. doi:10.1017/S0962492910000048.
+    """
+    @eval struct $Alg{CS,AD,FDT,ST,CJ} <: OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ}
+        krylov::Bool
+        m::Int
+        iop::Int
+    end
+    @eval $Alg(; krylov=false, m=30, iop=0, autodiff=true, standardtag=Val{true}(), concrete_jac=nothing, chunk_size=Val{0}(),
+        diff_type=Val{:forward}) = $Alg{_unwrap_val(chunk_size),_unwrap_val(autodiff),
+        diff_type,_unwrap_val(standardtag),_unwrap_val(concrete_jac)}(krylov, m, iop)
 end
 const ETD1 = NorsettEuler # alias
 for Alg in [:Exprb32, :Exprb43]
-  @eval struct $Alg{FDT,ST,CJ} <: OrdinaryDiffEqAdaptiveExponentialAlgorithm{FDT,ST,CJ}
+  @eval struct $Alg{CS,AD,FDT,ST,CJ} <: OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS,AD,FDT,ST,CJ}
     m::Int
     iop::Int
-    autodiff::Bool
-    chunksize::Int
   end
-  @eval $Alg(;m=30, iop=0, autodiff=true, standardtag = Val{true}(), concrete_jac = nothing, chunksize=0,
-            diff_type = Val{:forward}) = $Alg{diff_type,_unwrap_val(standardtag),_unwrap_val(concrete_jac)}(m, iop, _unwrap_val(autodiff), chunksize)
+  @eval $Alg(;m=30, iop=0, autodiff=true, standardtag = Val{true}(), concrete_jac = nothing, chunk_size=Val{0}(),
+            diff_type = Val{:forward}) = $Alg{_unwrap_val(chunk_size),_unwrap_val(autodiff),
+            diff_type,_unwrap_val(standardtag),
+            _unwrap_val(concrete_jac)}(m, iop)
 end
 for Alg in [:Exp4, :EPIRK4s3A, :EPIRK4s3B, :EPIRK5s3, :EXPRB53s3, :EPIRK5P1, :EPIRK5P2]
-  @eval struct $Alg{FDT,ST,CJ} <: OrdinaryDiffEqExponentialAlgorithm{FDT,ST,CJ}
+  @eval struct $Alg{CS,AD,FDT,ST,CJ} <: OrdinaryDiffEqExponentialAlgorithm{CS,AD,FDT,ST,CJ}
     adaptive_krylov::Bool
     m::Int
     iop::Int
-    autodiff::Bool
-    chunksize::Int
   end
   @eval $Alg(;adaptive_krylov=true, m=30, iop=0, autodiff=true, standardtag = Val{true}(), concrete_jac = nothing,
-              chunksize=0, diff_type = Val{:forward}) =
-              $Alg{diff_type,_unwrap_val(standardtag),_unwrap_val(concrete_jac)}(adaptive_krylov, m, iop, _unwrap_val(autodiff), chunksize)
+              chunk_size=Val{0}(), diff_type = Val{:forward}) =
+              $Alg{_unwrap_val(chunk_size),_unwrap_val(autodiff),diff_type,
+              _unwrap_val(standardtag),_unwrap_val(concrete_jac)}(adaptive_krylov, m, iop)
 end
-struct SplitEuler <: OrdinaryDiffEqExponentialAlgorithm{Val{:forward},Val{true},nothing} end
+struct SplitEuler <: OrdinaryDiffEqExponentialAlgorithm{0,false,Val{:forward},Val{true},nothing} end
 """
 ETD2: Exponential Runge-Kutta Method
   Second order Exponential Time Differencing method (in development).
 """
-struct ETD2 <: OrdinaryDiffEqExponentialAlgorithm{Val{:forward},Val{true},nothing} end
+struct ETD2 <: OrdinaryDiffEqExponentialAlgorithm{0,false,Val{:forward},Val{true},nothing} end
 
 #########################################
 

test/interface/norecompile.jl

@@ -28,6 +28,8 @@ if VERSION >= v"1.8"
     @test t2 < t4
 end
 
+solve(prob, EPIRK4s3A(), dt=1e-1)
+
 function f_oop(u, p, t)
     [0.2u[1], 0.4u[2]]
 end