almost

Author: ChrisRackauckas

ext/DiffEqBaseForwardDiffExt.jl

@@ -1,15 +1,15 @@
 module DiffEqBaseForwardDiffExt
 
 using DiffEqBase, ForwardDiff
+using DiffEqBase.ArrayInterface
 using DiffEqBase: Void, FunctionWrappersWrappers, OrdinaryDiffEqTag
 import DiffEqBase: hasdualpromote, wrapfun_oop, wrapfun_iip, promote_u0, prob2dtmin,
-promote_tspan, anyeltypedual, isdualtype, value, ODE_DEFAULT_NORM, InternalITP,
-nextfloat_tdir
+                   promote_tspan, anyeltypedual, isdualtype, value, ODE_DEFAULT_NORM,
+                   InternalITP,
+                   nextfloat_tdir, promote_dual
 
-const DUALCHECK_RECURSION_MAX = 10
-
-eltypedual(x)  = eltype(x) <: ForwardDiff.Dual
-isdualtype(::Type{<:ForwardDiff.Dual}) = true 
+eltypedual(x) = eltype(x) <: ForwardDiff.Dual
+isdualtype(::Type{<:ForwardDiff.Dual}) = true
 const dualT = ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, Float64}, Float64, 1}
 dualgen(::Type{T}) where {T} = ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, T}, T, 1}
 
@@ -24,12 +24,14 @@ function prob2dtmin(tspan, ::ForwardDiff.Dual, use_end_time)
     end
 end
 
-hasdualpromote(u0,t::Number) = hasmethod(ArrayInterface.promote_eltype,
-                    Tuple{Type{typeof(u0)}, Type{dualgen(eltype(u0))}}) &&
-                hasmethod(promote_rule,
-                    Tuple{Type{eltype(u0)}, Type{dualgen(eltype(u0))}}) &&
-                hasmethod(promote_rule,
-                    Tuple{Type{eltype(u0)}, Type{typeof(t)}})
+function hasdualpromote(u0, t::Number)
+    hasmethod(ArrayInterface.promote_eltype,
+            Tuple{Type{typeof(u0)}, Type{dualgen(eltype(u0))}}) &&
+        hasmethod(promote_rule,
+            Tuple{Type{eltype(u0)}, Type{dualgen(eltype(u0))}}) &&
+        hasmethod(promote_rule,
+            Tuple{Type{eltype(u0)}, Type{typeof(t)}})
+end
 
 const NORECOMPILE_IIP_SUPPORTED_ARGS = (
     Tuple{Vector{Float64}, Vector{Float64},
@@ -108,16 +110,6 @@ function wrapfun_iip(@nospecialize(ff))
     FunctionWrappersWrappers.FunctionWrappersWrapper{typeof(fwt), false}(fwt)
 end
 
-"""
-    promote_dual(::Type{T},::Type{T2})
-
-
-Is like the number promotion system, but always prefers a dual number type above
-anything else. For higher order differentiation, it returns the most dualiest of
-them all. This is then used to promote `u0` into the suspected highest differentiation
-space for solving the equation.
-"""
-promote_dual(::Type{T}, ::Type{T2}) where {T, T2} = T
 promote_dual(::Type{T}, ::Type{T2}) where {T <: ForwardDiff.Dual, T2} = T
 function promote_dual(::Type{T},
         ::Type{T2}) where {T <: ForwardDiff.Dual, T2 <: ForwardDiff.Dual}
@@ -227,7 +219,8 @@ end
 
 # Static Arrays don't support the `init` keyword argument for `sum`
 @inline __sum(f::F, args...; init, kwargs...) where {F} = sum(f, args...; init, kwargs...)
-@inline function __sum(f::F, a::DiffEqBase.StaticArraysCore.StaticArray...; init, kwargs...) where {F}
+@inline function __sum(
+        f::F, a::DiffEqBase.StaticArraysCore.StaticArray...; init, kwargs...) where {F}
     return mapreduce(f, +, a...; init, kwargs...)
 end
 
@@ -292,4 +285,4 @@ function SciMLBase.solve(
         right = ForwardDiff.Dual{T, V, P}(sol.right, partials))
 end
 
-end
+end

src/DiffEqBase.jl

@@ -7,7 +7,7 @@ end
 import PrecompileTools
 
 import FastPower
-@deprecate fastpow(x,y) FastPower.fastpower(x,y)
+@deprecate fastpow(x, y) FastPower.fastpower(x, y)
 
 using ArrayInterface
 
@@ -110,7 +110,7 @@ SciMLBase.isfunctionwrapper(x::FunctionWrapper) = true
 
 eltypedual(x) = false
 promote_u0(::Nothing, p, t0) = nothing
-isdualtype(::Type{T}) where T = true 
+isdualtype(::Type{T}) where {T} = true
 
 ## Types
 

src/internal_itp.jl

@@ -1,3 +1,20 @@
+"""
+  prevfloat_tdir(x, x0, x1)
+
+Move `x` one floating point towards x0.
+"""
+function prevfloat_tdir(x, x0, x1)
+    x1 > x0 ? prevfloat(x) : nextfloat(x)
+end
+
+function nextfloat_tdir(x, x0, x1)
+    x1 > x0 ? nextfloat(x) : prevfloat(x)
+end
+
+function max_tdir(a, b, x0, x1)
+    x1 > x0 ? max(a, b) : min(a, b)
+end
+
 """
 `InternalITP`: A non-allocating ITP method, internal to DiffEqBase for
 simpler dependencies.

src/norecompile.jl

@@ -24,4 +24,4 @@ end
 # Default dispatch assumes no ForwardDiff, gets added in the new dispatch
 function wrapfun_iip(ff, inputs)
     FunctionWrappersWrappers.FunctionWrappers.FunctionWrapper{Nothing, typeof(inputs)}(Void(ff))
-end
+end

src/solve.jl

@@ -574,7 +574,8 @@ function init_up(prob::AbstractDEProblem, sensealg, u0, p, args...; kwargs...)
         if tstops === nothing && has_kwargs(prob)
             tstops = get(prob.kwargs, :tstops, nothing)
         end
-        if !(tstops isa Union{Nothing, AbstractArray, Tuple, Real}) && !SciMLBase.allows_late_binding_tstops(alg)
+        if !(tstops isa Union{Nothing, AbstractArray, Tuple, Real}) &&
+           !SciMLBase.allows_late_binding_tstops(alg)
             throw(LateBindingTstopsNotSupportedError())
         end
         _prob = get_concrete_problem(prob, isadaptive(alg); u0 = u0, p = p, kwargs...)
@@ -1110,7 +1111,8 @@ function solve_up(prob::Union{AbstractDEProblem, NonlinearProblem}, sensealg, u0
         if tstops === nothing && has_kwargs(prob)
             tstops = get(prob.kwargs, :tstops, nothing)
         end
-        if !(tstops isa Union{Nothing, AbstractArray, Tuple, Real}) && !SciMLBase.allows_late_binding_tstops(alg)
+        if !(tstops isa Union{Nothing, AbstractArray, Tuple, Real}) &&
+           !SciMLBase.allows_late_binding_tstops(alg)
             throw(LateBindingTstopsNotSupportedError())
         end
         _prob = get_concrete_problem(prob, isadaptive(alg); u0 = u0, p = p, kwargs...)
@@ -1283,23 +1285,23 @@ function promote_f(f::F, ::Val{specialize}, u0, p, t) where {F, specialize}
     end
 
     f = if f isa ODEFunction && isinplace(f) && !(f.f isa AbstractSciMLOperator) &&
-            # Some reinitialization code still uses NLSolvers stuff which doesn't
-            # properly tag, so opt-out if potentially a mass matrix DAE
-            f.mass_matrix isa UniformScaling &&
-            # Jacobians don't wrap, so just ignore those cases
-            f.jac === nothing &&
-            ((specialize === SciMLBase.AutoSpecialize && eltype(u0) !== Any &&
-                RecursiveArrayTools.recursive_unitless_eltype(u0) === eltype(u0) &&
-                one(t) === oneunit(t) && hasdualpromote(u0,t)) ||
+           # Some reinitialization code still uses NLSolvers stuff which doesn't
+           # properly tag, so opt-out if potentially a mass matrix DAE
+           f.mass_matrix isa UniformScaling &&
+           # Jacobians don't wrap, so just ignore those cases
+           f.jac === nothing &&
+           ((specialize === SciMLBase.AutoSpecialize && eltype(u0) !== Any &&
+             RecursiveArrayTools.recursive_unitless_eltype(u0) === eltype(u0) &&
+             one(t) === oneunit(t) && hasdualpromote(u0, t)) ||
             (specialize === SciMLBase.FunctionWrapperSpecialize &&
-                !(f.f isa FunctionWrappersWrappers.FunctionWrappersWrapper)))
+             !(f.f isa FunctionWrappersWrappers.FunctionWrappersWrapper)))
         return unwrapped_f(f, wrapfun_iip(f.f, (u0, u0, p, t)))
     else
         return f
     end
 end
 
-hasdualpromote(u0,t) = true
+hasdualpromote(u0, t) = true
 
 function promote_f(f::SplitFunction, ::Val{specialize}, u0, p, t) where {specialize}
     typeof(f.cache) === typeof(u0) && isinplace(f) ? f : remake(f, cache = zero(u0))