Convenience API for derivatives and Jacobians.

SciML · Nov 8, 2017 · f6e7a6b · f6e7a6b
1 parent 31e05ab
commit f6e7a6b
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 156 deletions.
diff --git a/src/derivative_wrappers.jl b/src/derivative_wrappers.jl
@@ -46,3 +46,23 @@ mutable struct UDerivativeWrapper{F,tType} <: Function
   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)
+    end
+    nothing
+end
+
+function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, fx::AbstractArray{<:Number}, integrator::DEIntegrator, jac_config)
+    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)
+    end
+    nothing
+end
diff --git a/src/integrators/rosenbrock_integrators.jl b/src/integrators/rosenbrock_integrators.jl
@@ -29,12 +29,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -53,12 +48,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j] - γ*J[i,j]
@@ -148,12 +138,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -171,12 +156,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
         @inbounds W[i,j] = mass_matrix[i,j] - γ*J[i,j]
@@ -457,12 +437,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -481,12 +456,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]
@@ -667,12 +637,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -690,12 +655,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]
@@ -901,12 +861,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -925,12 +880,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]
@@ -1159,12 +1109,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -1185,12 +1130,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]
@@ -1546,12 +1486,7 @@ end
     else
       tf.vf.sizeu = sizeu
       tf.uprev = uprev
-      if alg_autodiff(integrator.alg)
-        ForwardDiff.derivative!(dT, tf, vec(du2), t)
-      else
-        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-        DiffEqDiffTools.finite_difference!(dT, tf, t, integrator.alg.diff_type, RealOrComplex, Val{:DiffEqDerivativeWrapper}, vec(du2))
-      end
+      derivative!(dT, tf, t, vec(du2), integrator)
     end
   end
 
@@ -1572,12 +1507,7 @@ end
       else
         uf.vfr.sizeu = sizeu
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J, uf, vec(du1), vec(uprev), jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
       for j in 1:length(u), i in 1:length(u)
           @inbounds W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]

diff --git a/src/integrators/sdirk_integrators.jl b/src/integrators/sdirk_integrators.jl
@@ -132,12 +132,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -362,12 +357,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -910,12 +900,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -1234,12 +1219,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -1558,12 +1538,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -1909,12 +1884,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -2423,12 +2393,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -3032,12 +2997,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -3608,12 +3568,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -4166,12 +4121,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -4822,12 +4772,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated
@@ -5576,12 +5521,7 @@ end
         f(Val{:jac},t,uprev,J)
       else
         uf.t = t
-        if alg_autodiff(integrator.alg)
-          ForwardDiff.jacobian!(J,uf,vec(du1),vec(uprev),jac_config)
-        else
-          RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
-          DiffEqDiffTools.finite_difference_jacobian!(J, uf, vec(uprev), integrator.alg.diff_type, RealOrComplex, Val{:JacobianWrapper}, vec(du1))
-        end
+        jacobian!(J, uf, vec(uprev), vec(du1), integrator, jac_config)
       end
     end
     # skip calculation of W if step is repeated