Merge pull request #47 from JuliaDiffEq/difftools

update for difftools v0.2.0

Author: ChrisRackauckas

REQUIRE

@@ -13,4 +13,4 @@ StaticArrays
 Reexport
 RandomNumbers
 MuladdMacro
-DiffEqDiffTools 0.1.0
+DiffEqDiffTools 0.2.0

src/caches/sdirk_caches.jl

@@ -32,13 +32,8 @@ function alg_cache(alg::ImplicitEM,prob,u,ΔW,ΔZ,rate_prototype,noise_rate_prot
   fsalfirst = zeros(rate_prototype)
   k = zeros(rate_prototype)
 
-  uf = UJacobianWrapper(f,t,tmp,dz)
-  if alg_autodiff(alg)
-    jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,
-                    ForwardDiff.Chunk{determine_chunksize(u,alg)}())
-  else
-    jac_config = nothing
-  end
+  uf = DiffEqDiffTools.UJacobianWrapper(f,t)
+  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,dz)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing
@@ -73,7 +68,7 @@ end
 
 function alg_cache(alg::ImplicitEM,prob,u,ΔW,ΔZ,rate_prototype,noise_rate_prototype,
                    uEltypeNoUnits,tTypeNoUnits,uprev,f,t,::Type{Val{false}})
-  uf = UDerivativeWrapper(f,t)
+  uf = DiffEqDiffTools.UDerivativeWrapper(f,t)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing
@@ -125,13 +120,8 @@ function alg_cache(alg::ImplicitEulerHeun,prob,u,ΔW,ΔZ,rate_prototype,noise_ra
   fsalfirst = zeros(rate_prototype)
   k = zeros(rate_prototype)
 
-  uf = UJacobianWrapper(f,t,tmp,dz)
-  if alg_autodiff(alg)
-    jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,
-                    ForwardDiff.Chunk{determine_chunksize(u,alg)}())
-  else
-    jac_config = nothing
-  end
+  uf = DiffEqDiffTools.UJacobianWrapper(f,t)
+  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,dz)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing
@@ -162,7 +152,7 @@ end
 
 function alg_cache(alg::ImplicitEulerHeun,prob,u,ΔW,ΔZ,rate_prototype,noise_rate_prototype,
                    uEltypeNoUnits,tTypeNoUnits,uprev,f,t,::Type{Val{false}})
-  uf = UDerivativeWrapper(f,t)
+  uf = DiffEqDiffTools.UDerivativeWrapper(f,t)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing
@@ -215,13 +205,8 @@ function alg_cache(alg::ImplicitRKMil,prob,u,ΔW,ΔZ,rate_prototype,noise_rate_p
   fsalfirst = zeros(rate_prototype)
   k = zeros(rate_prototype)
 
-  uf = UJacobianWrapper(f,t,tmp,dz)
-  if alg_autodiff(alg)
-    jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,
-                    ForwardDiff.Chunk{determine_chunksize(u,alg)}())
-  else
-    jac_config = nothing
-  end
+  uf = DiffEqDiffTools.UJacobianWrapper(f,t)
+  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,dz)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing
@@ -253,7 +238,7 @@ end
 
 function alg_cache(alg::ImplicitRKMil,prob,u,ΔW,ΔZ,rate_prototype,noise_rate_prototype,
                    uEltypeNoUnits,tTypeNoUnits,uprev,f,t,::Type{Val{false}})
-  uf = UDerivativeWrapper(f,t)
+  uf = DiffEqDiffTools.UDerivativeWrapper(f,t)
   ηold = one(uEltypeNoUnits)
 
   if alg.κ != nothing

src/derivative_wrappers.jl

@@ -1,39 +1,9 @@
-mutable struct TimeGradientWrapper{fType,uType} <: Function
-  f::fType
-  uprev::uType
-  fx1::uType
-end
-(p::TimeGradientWrapper)(t) = (du2 = similar(p.uprev); p.f(t,p.uprev,du2); du2)
-(p::TimeGradientWrapper)(du2,t) = p.f(t,p.uprev,du2)
-
-mutable struct UJacobianWrapper{fType,tType,CacheType} <: Function
-  f::fType
-  t::tType
-  x1::CacheType
-  fx1::CacheType
-end
-
-(p::UJacobianWrapper)(du1,uprev) = p.f(p.t,uprev,du1)
-(p::UJacobianWrapper)(uprev) = (du1 = similar(uprev); p.f(p.t,uprev,du1); du1)
-
-mutable struct TimeDerivativeWrapper{F,uType} <: Function
-  f::F
-  u::uType
-end
-(p::TimeDerivativeWrapper)(t) = p.f(t,p.u)
-
-mutable struct UDerivativeWrapper{F,tType} <: Function
-  f::F
-  t::tType
-end
-(p::UDerivativeWrapper)(u) = p.f(p.t,u)
-
 function derivative!(df::AbstractArray{<:Number}, f, x::Union{Number,AbstractArray{<:Number}}, fx::AbstractArray{<:Number}, integrator::DEIntegrator)
     if alg_autodiff(integrator.alg)
         ForwardDiff.derivative!(df, f, fx, x)
     else
         RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(df, f, x, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, fx)
+        DiffEqDiffTools.finite_difference!(df, f, x, integrator.alg.diff_type, RealOrComplex, fx)
     end
     nothing
 end
@@ -42,8 +12,24 @@ function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, f
     if alg_autodiff(integrator.alg)
       ForwardDiff.jacobian!(J, f, fx, x, jac_config)
     else
-      RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-      DiffEqDiffTools.finite_difference_jacobian!(J, f, x, integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, fx)
+      DiffEqDiffTools.finite_difference_jacobian!(J, f, x, jac_config)
     end
     nothing
 end
+
+function build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
+  if !has_jac(f)
+    if alg_autodiff(alg)
+      jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,ForwardDiff.Chunk{determine_chunksize(u,alg)}())
+    else
+      if alg.diff_type != Val{:complex}
+        jac_config = DiffEqDiffTools.JacobianCache(tmp,du1,du2,alg.diff_type)
+      else
+        jac_config = DiffEqDiffTools.JacobianCache(Complex{eltype(tmp)}.(tmp),Complex{eltype(du1)}.(du1),nothing,alg.diff_type)
+      end
+    end
+  else
+    jac_config = nothing
+  end
+  jac_config
+end